WIP: Working implementation in serial #494

src/ocgis/ocli.py

@@ -94,9 +94,12 @@ def ocli():
 @click.option('--loglvl', default="INFO", help='Verbosity level for standard out logging. Default is '
               '"INFO". See Python logging level docs for additional values: https://docs.python.org/3/howto/logging.html')
 @click.option('--weightfilemode', default="BASIC", help=M5)
+@click.option('--esmf_global_index/--not_esmf_global_index', default=True)
 def chunked_rwg(source, destination, weight, nchunks_dst, merge, esmf_src_type, esmf_dst_type, genweights,
                 esmf_regrid_method, spatial_subset, src_resolution, dst_resolution, buffer_distance, wd, persist,
-                eager, ignore_degenerate, data_variables, spatial_subset_path, verbose, loglvl, weightfilemode):
+                eager, ignore_degenerate, data_variables, spatial_subset_path, verbose, loglvl, weightfilemode,
+                esmf_global_index):
+    #tdk:doc: esmf_global_index
 
     # Used for creating the history string.
     the_locals = locals()
@@ -164,7 +167,8 @@ def chunked_rwg(source, destination, weight, nchunks_dst, merge, esmf_src_type,
             spatial_subset_path = os.path.join(wd, 'spatial_subset.nc')
         msg = "Executing spatial subset. Output path is: {}".format(spatial_subset_path)
         ocgis_lh(msg=msg, level=logging.INFO, logger=CRWG_LOG)
-        _write_spatial_subset_(rd_src, rd_dst, spatial_subset_path, src_resmax=src_resolution)
+        _write_spatial_subset_(rd_src, rd_dst, spatial_subset_path, src_resmax=src_resolution,
+                               esmf_global_index=esmf_global_index)
     # Only split grids if a spatial subset is not requested.
     else:
         # Update the paths to use for the grid.
@@ -321,10 +325,10 @@ def _is_subdir_(path, potential_subpath):
     return not relative.startswith(os.pardir + os.sep)
 
 
-def _write_spatial_subset_(rd_src, rd_dst, spatial_subset_path, src_resmax=None):
+def _write_spatial_subset_(rd_src, rd_dst, spatial_subset_path, src_resmax=None, esmf_global_index=False):
     src_field = rd_src.create_field()
     dst_field = rd_dst.create_field()
-    sso = SpatialSubsetOperation(src_field)
+    sso = SpatialSubsetOperation(src_field, add_esmf_index=esmf_global_index)
 
     with grid_abstraction_scope(dst_field.grid, Topology.POLYGON):
         dst_field_extent = dst_field.grid.extent_global

src/ocgis/regrid/base.py

@@ -10,7 +10,7 @@
 from ocgis import env
 from ocgis.base import AbstractOcgisObject, get_dimension_names, iter_dict_slices
 from ocgis.collection.field import Field
-from ocgis.constants import DMK, GridChunkerConstants, DecompositionType
+from ocgis.constants import DMK, GridChunkerConstants, DecompositionType, AttributeName
 from ocgis.exc import RegriddingError, CornersInconsistentError
 from ocgis.spatial.grid import Grid, expand_grid
 from ocgis.spatial.spatial_subset import SpatialSubsetOperation
@@ -926,6 +926,13 @@ def create_esmf_grid_fromfile(filename, grid, esmf_kwargs):
     else:
         meshname = str(grid.dimension_map.get_variable(DMK.ATTRIBUTE_HOST))
         ret = klass(filename=filename, filetype=filetype, meshname=meshname)
+
+    vc = RequestDataset(uri=filename, driver='netcdf', decomp_type=DecompositionType.ESMF).create_field()
+    poss = vc.find_by_attribute(key=AttributeName.ESMF_GLOBAL_INDICES, value=1)
+    if len(poss) > 0:
+        assert(len(poss) == 1)
+        ret._set_arb_indices_(poss[0].v().flatten())
+
     return ret
 
 

src/ocgis/spatial/spatial_subset.py

@@ -255,8 +255,8 @@ def _prepare_geometry_(self, geom):
                 prepared.wrap()
 
         # Add the ESMF index variable if requested
-        #tdk:todo: implement find by attribute instead of checking for this name in keys
-        if self.add_esmf_index and AttributeName.ESMF_GLOBAL_INDICES not in self.field.keys():
+        possvars = self.field.find_by_attribute(key=AttributeName.ESMF_GLOBAL_INDICES, value=1)
+        if self.add_esmf_index and len(possvars) == 0:
             grid = self.field.grid
             vals = grid._gc_create_global_indices_(grid.shape_global)
             var = Variable(name=AttributeName.ESMF_GLOBAL_INDICES, value=vals, dimensions=grid.dimensions)

src/ocgis/test/test_ocgis/test_ocli.py

@@ -266,7 +266,7 @@ def test_system_mock_combinations(self, mRequestDataset, mGridChunker, m_mkdtemp
                 m.reset_mock()
 
     @attr('esmf')
-    def test_chunked_rwg_spatial_subset(self):
+    def test_system_chunked_rwg_spatial_subset(self):
         env.CLOBBER_UNITS_ON_BOUNDS = False
 
         src_grid = create_gridxy_global(crs=Spherical())
@@ -297,7 +297,7 @@ def test_chunked_rwg_spatial_subset(self):
         runner = CliRunner()
         cli_args = ['chunked-rwg', '--source', source, '--destination', destination, '--wd', wd, '--spatial_subset',
                     '--spatial_subset_path', spatial_subset, '--weight', weight, '--esmf_regrid_method', 'BILINEAR',
-                    '--persist']
+                    '--persist', '--esmf_global_index']
         result = runner.invoke(ocli, args=cli_args, catch_exceptions=False)
         self.assertEqual(result.exit_code, 0)
 
@@ -311,3 +311,6 @@ def test_chunked_rwg_spatial_subset(self):
         self.assertTrue(os.path.exists(weight))
         actual = RequestDataset(weight, driver='netcdf').create_field()
         self.assertIn('history', actual.attrs)
+
+        # Test that the arbitrary sequence indices are used.
+        self.assertGreater(actual['col'].v().min(), 4000)

src/ocgis/test/test_ocgis/test_spatial/test_spatial_subset.py

@@ -108,6 +108,19 @@ def test_init_output_crs(self):
                 if isinstance(k.target, Field):
                     self.assertEqual(ss.sdim.crs, k.target.spatial.crs)
 
+    def test_system_spatial_subset_with_esmf_index(self):
+        """Test adding an ESMF index to a spatial subset."""
+        #tdk:test: with unstructured grid
+        #tdk:test: parallel
+
+        grid = create_gridxy_global()
+        subset_geom = box(10, 30, 20, 40)
+        sso = SpatialSubsetOperation(grid.parent, add_esmf_index=True)
+        sub = sso.get_spatial_subset('intersects', subset_geom)
+        self.assertIn(AttributeName.ESMF_GLOBAL_INDICES, sub.keys())
+        self.assertNotEqual(sub[AttributeName.ESMF_GLOBAL_INDICES].v()[0, 0], 1)
+        self.assertEqual(sub[AttributeName.ESMF_GLOBAL_INDICES].attrs[AttributeName.ESMF_GLOBAL_INDICES], 1)
+
     def test_get_buffered_geometry(self):
         proj4 = '+proj=aea +lat_1=20 +lat_2=60 +lat_0=40 +lon_0=-96 +x_0=0 +y_0=0 +ellps=GRS80 +datum=NAD83 +units=m +no_defs'
         buffer_crs_list = [None, CoordinateReferenceSystem(proj4=proj4)]
@@ -275,22 +288,6 @@ def test_get_spatial_subset_wrap(self):
         self.assertEqual(ret.wrapped_state, WrappedState.UNWRAPPED)
         self.assertAlmostEqual(ret.grid.get_value_stacked()[1].mean(), 260.15625)
 
-    def test_system_spatial_subset_with_esmf_index(self):
-        """Test adding an ESMF index to a spatial subset."""
-        #tdk:test: with unstructured grid
-        #tdk:test: parallel
-        #tdk:order
-
-        grid = create_gridxy_global()
-        subset_geom = box(10, 30, 20, 40)
-        sso = SpatialSubsetOperation(grid.parent, add_esmf_index=True)
-        sub = sso.get_spatial_subset('intersects', subset_geom)
-        print(sub.grid.extent) #tdk:p
-        self.assertIn(AttributeName.ESMF_GLOBAL_INDICES, sub.keys())
-        self.assertNotEqual(sub[AttributeName.ESMF_GLOBAL_INDICES].v()[0, 0], 1)
-        print(sub[AttributeName.ESMF_GLOBAL_INDICES].v()) #tdk:p
-        self.assertIn(AttributeName.ESMF_GLOBAL_INDICES, sub[AttributeName.ESMF_GLOBAL_INDICES].attrs)
-
     @attr('data')
     def test_prepare_target(self):
         for ss, k in self: