From 57d7a6d690bec569f9b94bc45fdbc91584b735d7 Mon Sep 17 00:00:00 2001 From: Xee authors Date: Sun, 3 Dec 2023 21:40:42 -0800 Subject: [PATCH] Efficient way to convert a structured numpy array into a 3d array. This CL uses a more efficient way to convert the numpy structured array removing the need to copy the array into memory as a list. Using `.view()` doesn't seem to change the data buffer and should be more memory efficient. closes #9 PiperOrigin-RevId: 587594128 --- xee/ext.py | 13 +++++++++---- xee/ext_integration_test.py | 1 + 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/xee/ext.py b/xee/ext.py index 8f6399a..a385e9d 100644 --- a/xee/ext.py +++ b/xee/ext.py @@ -464,10 +464,15 @@ def image_to_array( pixels_getter, params, catch=ee.ee_exception.EEException ) - # TODO(#9): Find a way to make this more efficient. This is needed because - # `raw` is a structured array of all the same dtype (i.e. number of images). - arr = np.array(raw.tolist(), dtype=dtype) - data = arr.T + x_size = kwargs['grid']['dimensions']['width'] + y_size = kwargs['grid']['dimensions']['height'] + + # This is needed because `raw` is a structured array of all the same dtype + # (i.e. number of images). + arr = raw.view(dtype).reshape(-1, x_size, y_size,) + arr = np.moveaxis(arr, 0, -1) + + data = arr current_mask_value = np.array(self.mask_value, dtype=data.dtype) # Sets EE nodata masked value to NaNs. data = np.where(data == current_mask_value, np.nan, data) diff --git a/xee/ext_integration_test.py b/xee/ext_integration_test.py index d1eafc0..4c52dd9 100644 --- a/xee/ext_integration_test.py +++ b/xee/ext_integration_test.py @@ -85,6 +85,7 @@ def setUp(self): ) def test_creates_lat_long_array(self): + init_ee_for_tests() arr = xee.EarthEngineBackendArray('longitude', self.lnglat_store) self.assertEqual((1, 360, 180), arr.shape)