ENH: nested grid preserves zonation (#438)

Author: lahus-0

src/fmu/tools/nestedhybridgrid/nestedhybrid.py

@@ -310,6 +310,28 @@ def _generate_layer_mappings(
     return (lmap1, lmap2)
 
 
+def _set_zonation(subgrid: dict, lmap: np.ndarray, nlay: int) -> dict:
+    """Create an updated subgrid dictionary for the merged grid.
+    Args:
+        subgrid: subgrid dictionary from input grid
+        lmap: layer mapping for grid1 - unrefined input grid
+        nlay: total number of layers in merged grid
+    """
+
+    updated_subgrid = {}
+
+    # sorted list of zones to add
+    zl = sorted(subgrid, key=lambda x: subgrid[x][0])
+
+    for zi in range(len(zl)):
+        zn = zl[zi]
+        zmin = lmap[subgrid[zn][0] - 1] + 1
+        zmax = nlay + 1 if zi == len(zl) - 1 else lmap[subgrid[zl[zi + 1]][0] - 1] + 1
+        updated_subgrid[zn] = range(zmin, zmax)
+
+    return updated_subgrid
+
+
 # ---------------------------------------------------------------------------
 # Public API
 # ---------------------------------------------------------------------------
@@ -377,6 +399,9 @@ def create_nested_hybrid_grid(
     grid = grid.copy()
     region = region.copy()
 
+    # Save input zonation
+    subgrid = grid.subgrids
+
     # Attach the region property to the grid.
     grid.append_prop(region)
 
@@ -422,6 +447,8 @@ def create_nested_hybrid_grid(
     # 7. Merge the two grids.
     merged = xtgeo.grid_merge(grid, refined, lmap1, lmap2)
     _logger.info("Merged grid dimensions: %s", merged.dimensions)
+    if subgrid is not None:
+        merged.subgrids = _set_zonation(subgrid=subgrid, lmap=lmap1, nlay=merged.nlay)
 
     return merged, nnc_table
 

tests/nestedhybridgrid/test_nestedhybrid.py

@@ -198,6 +198,37 @@ def test_lmap_generation_ref10(self):
         assert np.array_equal(lmap1, np.array([0, 1, 11]))
         assert np.array_equal(lmap2, np.arange(20) + 1)
 
+    def test_zonation(self):
+        """Input grid with zonation returns a valid zonation"""
+        grid, region, rid = _make_box_grid_with_region(dimension=(6, 6, 2))
+        grid.subgrids = {"ZONE1": [1], "ZONE2": [2]}
+        merged, nnc_table = create_nested_hybrid_grid(
+            grid, region, rid, refinement=(2, 2, 2)
+        )
+        subgrids_nlay = merged.get_subgrids()
+        assert subgrids_nlay["ZONE1"] == 2
+        assert subgrids_nlay["ZONE2"] == 2
+
+        assert merged.subgrids["ZONE1"] == range(1, 3)
+        assert merged.subgrids["ZONE2"] == range(3, 5)
+
+    def test_zonation_with_layer_offset(self):
+        """Input grid with zonation returns a valid zonation with offset"""
+        grid, region, rid = _make_box_grid_with_region(dimension=(6, 6, 3))
+        grid.subgrids = {"ZONE1": [1], "ZONE2": [2, 3]}
+
+        region.values[:, :, 0] = 1  # set value in first layer to not be rid
+
+        merged, nnc_table = create_nested_hybrid_grid(
+            grid, region, rid, refinement=(2, 2, 2)
+        )
+        subgrids_nlay = merged.get_subgrids()
+        assert subgrids_nlay["ZONE1"] == 1
+        assert subgrids_nlay["ZONE2"] == 4
+
+        assert merged.subgrids["ZONE1"] == range(1, 2)
+        assert merged.subgrids["ZONE2"] == range(2, 6)
+
 
 # ---------------------------------------------------------------------------
 # Tests for get_transmissibilities with nested hybrid NNCs