Coverage for hml/rasterizeShapefile.py: 2%

1#!/usr/bin/env python3 

2 

3# Define function ... 

4def rasterizeShapefile(sfObj, /, *, nx = 1024, ny = 1024, px = 1024.0): 

5 """ 

6 Rasterize a ShapeFile. 

7 

8 Arguments: 

9 sfObj -- a shapefile.Reader of a ShapeFile 

10 

11 Keyword arguments: 

12 px -- pixel size (default 1024.0) 

13 nx -- number of x pixels (default 1024) 

14 ny -- number of y pixels (default 1024) 

15 

16 Note: 

17 This function only works for ShapeFiles that solely exist in the (positive, 

18 positive) quadrant. 

19 """ 

20 

21 # Import standard modules ... 

22 import multiprocessing 

23 

24 # Import special modules ... 

25 try: 

26 import numpy 

27 except: 

28 raise Exception("\"numpy\" is not installed; run \"pip install --user numpy\"") from None 

29 try: 

30 import shapefile 

31 except: 

32 raise Exception("\"shapefile\" is not installed; run \"pip install --user pyshp\"") from None 

33 try: 

34 import shapely 

35 import shapely.geometry 

36 import shapely.validation 

37 except: 

38 raise Exception("\"shapely\" is not installed; run \"pip install --user Shapely\"") from None 

39 

40 # Import sub-functions ... 

41 from .rasterizePolygon import rasterizePolygon 

42 

43 # Check argument ... 

44 if not isinstance(sfObj, shapefile.Reader): 

45 raise TypeError("\"sfObj\" is not a shapefile.Reader") 

46 

47 # Initialize counter and global grid ... 

48 n = 0 # [#] 

49 globalGrid = numpy.zeros((ny, nx), dtype = numpy.float32) # [m2] 

50 

51 # Create a pool of workers ... 

52 with multiprocessing.Pool(maxtasksperchild = 1) as pObj: 

53 # Initialize list ... 

54 results = [] 

55 

56 # Loop over shape+record pairs ... 

57 for shapeRecord in sfObj.iterShapeRecords(): 

58 # Crash if this shape+record is not a shapefile polygon ... 

59 if shapeRecord.shape.shapeType != shapefile.POLYGON: 

60 raise Exception("\"shape\" is not a POLYGON") from None 

61 

62 # Convert shapefile.Shape to shapely.geometry.polygon.Polygon or 

63 # shapely.geometry.multipolygon.MultiPolygon ... 

64 poly = shapely.geometry.shape(shapeRecord.shape) 

65 if not poly.is_valid: 

66 print(f"WARNING: Skipping a shape as it is not valid ({shapely.validation.explain_validity(poly)}).") 

67 n += 1 # [#] 

68 continue 

69 if poly.is_empty: 

70 n += 1 # [#] 

71 continue 

72 

73 # Add rasterization job to worker pool ... 

74 results.append( 

75 pObj.apply_async( 

76 rasterizePolygon, 

77 (poly,), 

78 { 

79 "px" : px, 

80 }, 

81 ) 

82 ) 

83 

84 print(f"INFO: {n:,d} records were skipped because they were invalid") 

85 

86 # Loop over results ... 

87 for result in results: 

88 # Get result ... 

89 ix1, iy1, localGrid = result.get() 

90 

91 # Check result ... 

92 if not result.successful(): 

93 raise Exception("\"multiprocessing.Pool().apply_async()\" was not successful") from None 

94 

95 # Loop over x-axis ... 

96 for ix in range(localGrid.shape[1]): 

97 # Loop over y-axis ... 

98 for iy in range(localGrid.shape[0]): 

99 # Add local grid to global grid ... 

100 globalGrid[iy1 + iy, ix1 + ix] += localGrid[iy, ix] # [m2] 

101 

102 # Close the pool of worker processes and wait for all of the tasks to 

103 # finish ... 

104 # NOTE: The "__exit__()" call of the context manager for 

105 # "multiprocessing.Pool()" calls "terminate()" instead of 

106 # "join()", so I must manage the end of the pool of worker 

107 # processes myself. 

108 pObj.close() 

109 pObj.join() 

110 

111 # Return answer ... 

112 return globalGrid