NHIRS-8 - Improve memory usage for large raster input (#41)

c-mckenna · web-flow · commit a230e0876d89 · 2021-02-26T16:28:48.000+11:00
* NHIRS-8: Improve memory usage for large raster input

* No longer reads the entire raster into memory, reads only the cells defined in the exposure data
* Removed the ability to pass an in-memory array via 'from_array'
* Added a 'ThreadPoolExecutor' for some performance improvement when reading hazard data for large exposure datasets

* NHIRS-8: Update docstring

Co-authored-by: Callum McKenna &lt;callum.mckenna@ga.gov.au&gt;
diff --git a/hazimp/jobs/jobs.py b/hazimp/jobs/jobs.py
@@ -612,13 +612,9 @@ def __init__(self):
         super(LoadRaster, self).__init__()
         self.call_funct = LOADRASTER
 
-    # R0913:326: Too many arguments (9/6)
-    # pylint: disable=R0913
-    def __call__(self, context, attribute_label,
+    def __call__(self, context, attribute_label, file_list,
                  clip_exposure2all_hazards=False,
-                 file_list=None, file_format=None, variable=None,
-                 raster=None, upper_left_x=None, upper_left_y=None,
-                 cell_size=None, no_data_value=None):
+                 file_format=None, variable=None, no_data_value=None):
         """
         Load one or more files and get the value for all the
         exposure points. All files have to be of the same attribute.
@@ -628,79 +624,45 @@ def __call__(self, context, attribute_label,
         :param attribute_label: The string to be associated with this data.
         :param clip_exposure2all_hazards: True if the exposure data is
             clippped to the hazard data, so no hazard values are ignored.
-
         :param file_list: A list of files or a single file to be loaded.
-        OR
-        :param raster: A 2D numeric array of the raster values, North is up.
-        :param upper_left_x: The longitude at the upper left corner.
-        :param upper_left_y: The latitude at the upper left corner.
-        :param cell_size: The cell size.
         :param no_data_value: Values in the raster that represent no data.
 
-
         Context return:
            exposure_att: Add the file values into this dictionary.
                key: column titles
                value: column values, except the title
         """
 
-        # We need a file or a full set of raster info.
-        if file_list is None:
-            # The raster info is being passed as an array
-            assert raster is not None
-            assert upper_left_x is not None
-            assert upper_left_y is not None
-            assert cell_size is not None
-            assert no_data_value is not None
-            a_raster = raster_module.Raster.from_array(
-                raster, upper_left_x,
-                upper_left_y,
-                cell_size,
-                no_data_value)
-
-            if clip_exposure2all_hazards:
-                # Reduce the context to the hazard area
-                # before the raster info has been added to the context
-                extent = a_raster.extent()
-                context.clip_exposure(*extent)
-
-            file_data = a_raster.raster_data_at_points(
-                context.exposure_long,
-                context.exposure_lat)
-            file_data = np.where(file_data == no_data_value, np.NAN,
-                                 file_data)
-            context.exposure_att[attribute_label] = file_data
-        else:
-            if isinstance(file_list, str):
-                file_list = [file_list]
-
-            for f in file_list:
-                f = misc.download_file_from_s3_if_needed(f)
-                dt = misc.get_file_mtime(f)
-                atts = {"dcterms:title": "Source hazard data",
-                        "prov:type": "prov:Dataset",
-                        "prov:atLocation": os.path.basename(f),
-                        "prov:format": os.path.splitext(f)[1].replace('.', ''),
-                        "prov:generatedAtTime": dt, }
-                if file_format == 'nc' and variable:
-                    atts['prov:variable'] = variable
-                hazent = context.prov.entity(":Hazard data", atts)
-                context.prov.used(context.provlabel, hazent)
+        if isinstance(file_list, str):
+            file_list = [file_list]
 
+        for f in file_list:
+            f = misc.download_file_from_s3_if_needed(f)
+            dt = misc.get_file_mtime(f)
+            atts = {"dcterms:title": "Source hazard data",
+                    "prov:type": "prov:Dataset",
+                    "prov:atLocation": os.path.basename(f),
+                    "prov:format": os.path.splitext(f)[1].replace('.', ''),
+                    "prov:generatedAtTime": dt, }
             if file_format == 'nc' and variable:
-                file_list = misc.mod_file_list(file_list, variable)
+                atts['prov:variable'] = variable
+            hazent = context.prov.entity(":Hazard data", atts)
+            context.prov.used(context.provlabel, hazent)
+
+        if file_format == 'nc' and variable:
+            file_list = misc.mod_file_list(file_list, variable)
 
-            file_data, extent = raster_module.files_raster_data_at_points(
-                context.exposure_long,
-                context.exposure_lat, file_list)
-            file_data[file_data == no_data_value] = np.NAN
+        file_data, extent = raster_module.files_raster_data_at_points(
+            context.exposure_long,
+            context.exposure_lat, file_list)
+        file_data[file_data == no_data_value] = np.NAN
 
-            context.exposure_att[attribute_label] = file_data
+        context.exposure_att[attribute_label] = file_data
 
-            if clip_exposure2all_hazards:
-                # Clipping the exposure points after the data has been added.
-                # Not optimised for speed, but easy to implement.
-                context.clip_exposure(*extent)
+        if clip_exposure2all_hazards:
+            # Clipping the exposure points after the data has been added.
+            # Not optimised for speed, but easy to implement.
+            context.clip_exposure(*extent)
 
 
 class AggregateLoss(Job):
diff --git a/hazimp/raster.py b/hazimp/raster.py
@@ -17,16 +17,13 @@
 
 """
 Manipulate raster data.
-
-Currently loads the entire raster layer into memory, which can blow out memory
-usage in some cases.
-
-# TODO: optimise raster loading and reading
 """
+import threading
+from concurrent.futures.thread import ThreadPoolExecutor
 
-import numpy
 import gdal
-from gdalconst import GA_ReadOnly, GDT_Float32
+import numpy
+from gdalconst import GA_ReadOnly
 
 
 class Raster(object):
@@ -43,13 +40,11 @@ class Raster(object):
     # R0913: 34:Raster.__init__: Too many arguments (9/6)
     # pylint: disable=R0902, R0913
 
-    def __init__(self, raster, upper_left_x, upper_left_y,
+    def __init__(self, filename, upper_left_x, upper_left_y,
                  x_pixel, y_pixel, no_data_value, x_size, y_size):
         """
 
-        :param raster: A 2D numeric array of the raster values, North is up.
-                       The values are listed in 'English reading order' i.e.
-                       left-right and top-down.
+        :param filename: The raster file path string.
         :param upper_left_x: The longitude at the upper left corner of the
                              top left pixel.
         :param upper_left_y: The latitude at the upper left corner of the
@@ -62,23 +57,22 @@ def __init__(self, raster, upper_left_x, upper_left_y,
         :param y_size: Number of rows.
         :param no_data_value: Values in the raster that represent no data.
         """
-        self.raster = raster
+        self.filename = filename
         self.ul_x = upper_left_x
         self.ul_y = upper_left_y
         self.x_pixel = x_pixel
         self.y_pixel = y_pixel
         self.no_data_value = no_data_value
         self.x_size = x_size
         self.y_size = y_size
-        self.raster[self.raster == self.no_data_value] = numpy.NAN
 
     @classmethod
     def from_file(cls, filename):
         """
         Load a file in a raster file format known to GDAL.
         Note, image must be 'North up'.
 
-        :param filename: The csv file path string.
+        :param filename: The raster file path string.
         :returns: A Raster instance.
         """
 
@@ -99,40 +93,9 @@ def from_file(cls, filename):
         y_size = dataset.RasterYSize  # This will be a negative value.
         band = dataset.GetRasterBand(1)
         no_data_value = band.GetNoDataValue()
-        raster = band.ReadAsArray(0, 0, x_size, y_size, buf_type=GDT_Float32)
-        instance = cls(raster, upper_left_x, upper_left_y,
+        instance = cls(filename, upper_left_x, upper_left_y,
                        x_pixel, y_pixel, no_data_value, x_size, y_size)
-        return instance
-
-    @classmethod
-    def from_array(cls, raster, upper_left_x, upper_left_y,
-                   cell_size, no_data_value, dtype='float'):
-        """
-        Convert numeric array of raster data and info to a raster instance.
-        The values are listed in 'English reading order' i.e.
-        left-right and top-down.
-
-        :param raster: A 2D numeric array of the raster values, North is up.
-        :param upper_left_x: The longitude at the upper left corner.
-        :param upper_left_y: The latitude at the upper left corner.
-        :param cell_size: The cell size.
-        :param no_data_value: Values in the raster that represent no data.
-        :param dtype: Data type for the raster values (default float).
-        :returns: A Raster instance
-        """
-        raster = numpy.array(raster, dtype=dtype, copy=False)
-        if not len(raster.shape) == 2:
-            msg = ('Bad Raster shape %s' % (str(raster.shape)))
-            raise TypeError(msg)
-
-        x_size = raster.shape[1]
-        y_size = raster.shape[0]
-
-        x_pixel = cell_size
-        y_pixel = -cell_size
 
-        instance = cls(raster, upper_left_x, upper_left_y,
-                       x_pixel, y_pixel, no_data_value, x_size, y_size)
         return instance
 
     def raster_data_at_points(self, lon, lat):
@@ -168,11 +131,25 @@ def raster_data_at_points(self, lon, lat):
             raw_row_offset = (lat - self.ul_y) / self.y_pixel
             row_offset = numpy.trunc(raw_row_offset).astype(int)
 
-            values[good_indexes] = self.raster[row_offset[good_indexes],
-                                               col_offset[good_indexes]]
-            # Change NODATA_value to NAN
-            values = numpy.where(values == self.no_data_value, numpy.NAN,
-                                 values)
+            data = threading.local()
+
+            def read_cell(i, x, y):
+                if 'band' not in data.__dict__:
+                    data.dataset = gdal.Open(self.filename, GA_ReadOnly)
+                    data.band = data.dataset.GetRasterBand(1)
+
+                values[i] = data.band.ReadAsArray(x, y, 1, 1)[0]
+
+            with ThreadPoolExecutor() as executor:
+                for index in good_indexes:
+                    executor.submit(read_cell, index,
+                                    col_offset[index].item(),
+                                    row_offset[index].item())
+
+        # Change NODATA_value to NAN
+        values = numpy.where(values == self.no_data_value, numpy.NAN,
+                             values)
+
         return values
 
     def extent(self):
diff --git a/tests/data/basic_raster.aai b/tests/data/basic_raster.aai
@@ -0,0 +1,8 @@
+ncols 3
+nrows 2
+xllcorner +0.
+yllcorner +8.
+cellsize 1
+NODATA_value -9999
+1 2 -9999
+4 5 6
diff --git a/tests/jobs/test_jobs.py b/tests/jobs/test_jobs.py
@@ -46,6 +46,7 @@
 from hazimp import context
 from hazimp import misc
 from hazimp import parallel
+from tests import CWD
 from tests.jobs.test_vulnerability_model import build_example1
 
 prov = mock.MagicMock(name='prov.model')
@@ -569,20 +570,11 @@ def test_load_raster_clippingIII(self, mock_used):
         inst(con_in, **test_kwargs)
         os.remove(f.name)
 
-        raster = array([[1, 2, -9999], [4, 5, 6]])
-        upper_left_x = 0
-        upper_left_y = 10
-        cell_size = 1
-        no_data_value = -9999
         haz_v = 'haz_v'
         inst = JOBS[LOADRASTER]
         test_kwargs = {'attribute_label': haz_v,
                        'clip_exposure2all_hazards': True,
-                       'raster': raster,
-                       'upper_left_x': upper_left_x,
-                       'upper_left_y': upper_left_y,
-                       'cell_size': cell_size,
-                       'no_data_value': no_data_value}
+                       'file_list': [str(CWD / 'data/basic_raster.aai')]}
         inst(con_in, **test_kwargs)
 
         # There should be only no exposure points
diff --git a/tests/test_raster.py b/tests/test_raster.py
@@ -31,9 +31,9 @@
 Test the raster module.
 """
 
-import unittest
-import tempfile
 import os
+import tempfile
+import unittest
 
 import numpy
 from scipy import asarray, allclose, nan
@@ -179,42 +179,6 @@ def test2_files_raster_data_at_points(self):
         for a_file in files:
             os.remove(a_file)
 
-    def test3_raster_data_from_array(self):
-        # A test based on this info;
-        # http://en.wikipedia.org/wiki/Esri_grid
-        # Let's hope no one edits the data....
-        raster = [[-9999, -9999, 5, 2], [-9999, 20, 100, 36],
-                  [3, 8, 35, 10], [32, 42, 50, 6],
-                  [88, 75, 27, 9], [13, 5, 1, -9999]]
-        upper_left_x = 0.
-        upper_left_y = 300.
-        cell_size = 50.0
-        no_data_value = -9999
-
-        # Just outside the midpoint of all sides
-        lon = asarray([125, 125, 125, 125, 125, 125])
-        lat = asarray([275, 225, 175, 125, 75, 25])
-
-        raster = Raster.from_array(raster, upper_left_x, upper_left_y,
-                                   cell_size, no_data_value)
-        self.assertEqual(raster.ul_x, 0)
-        self.assertEqual(raster.ul_y, 300)
-        self.assertEqual(raster.x_pixel, 50)
-        self.assertEqual(raster.y_pixel, -50)
-        self.assertEqual(raster.x_size, 4)
-        self.assertEqual(raster.y_size, 6)
-
-        data = raster.raster_data_at_points(lon, lat)
-        self.assertTrue(allclose(data, asarray([5.0, 100.0, 35.0,
-                                                50.0, 27.0, 1.0])))
-
-        # testing extent
-        min_long, min_lat, max_long, max_lat = raster.extent()
-        self.assertEqual(min_long, 0)
-        self.assertEqual(min_lat, 0)
-        self.assertEqual(max_long, 200)
-        self.assertEqual(max_lat, 300)
-
     def test3_recalc_max(self):
         max_extent = (0, 0, 0, 0)
         extent = [-10, -20, 20, 40]
@@ -232,7 +196,6 @@ def test3_recalc_max(self):
         self.assertEqual(old_max_extent, max_extent)
 
 
-# -------------------------------------------------------------
 if __name__ == "__main__":
     Suite = unittest.makeSuite(TestRaster, 'test')
     Runner = unittest.TextTestRunner()
diff --git a/tests/test_templates.py b/tests/test_templates.py
