cleanup

Author: wpbonelli

flopy4/mf6/codec/writer/filters.py

@@ -201,11 +201,15 @@ def dataset2list(value: xr.Dataset):
     if value is None or not any(value.data_vars):
         return
 
-    first = next(iter(value.data_vars.values()))
-    is_union = first.dtype.type is np.str_
+    # special case OC for now.
+    is_oc = all(
+        str(v.name).startswith("save_") or str(v.name).startswith("print_")
+        for v in value.data_vars.values()
+    )
 
-    if first.ndim == 0:  # handle scalar
-        if is_union:
+    # handle scalar
+    if (first := next(iter(value.data_vars.values()))).ndim == 0:
+        if is_oc:
             for name in value.data_vars.keys():
                 val = value[name]
                 val = val.item() if val.shape == () else val
@@ -230,7 +234,7 @@ def dataset2list(value: xr.Dataset):
     has_spatial_dims = len(spatial_dims) > 0
     indices = np.where(combined_mask)
     for i in range(len(indices[0])):
-        if is_union:
+        if is_oc:
             for name in value.data_vars.keys():
                 val = value[name][tuple(idx[i] for idx in indices)]
                 val = val.item() if val.shape == () else val

flopy4/mf6/converter.py

@@ -110,6 +110,7 @@ def unstructure_component(value: Component) -> dict[str, Any]:
             blocks[block_name] = {}
         period_data = {}
         period_blocks = {}  # type: ignore
+        period_block_name = None
 
         for field_name in block.keys():
             # Skip child components that have been processed as bindings
@@ -150,51 +151,43 @@ def unstructure_component(value: Component) -> dict[str, Any]:
                         field_value,
                         structured_grid_dims=value.parent.data.dims,  # type: ignore
                     )
-
+                if "period" in block_name:
+                    period_block_name = block_name
                     period_data[field_name] = {
                         kper: field_value.isel(nper=kper)
                         for kper in range(field_value.sizes["nper"])
                     }
                 else:
-                    # TODO why not putting in block here but doing below? how does this even work
-                    if np.issubdtype(field_value.dtype, np.str_):
-                        period_data[field_name] = {
-                            kper: field_value[kper] for kper in range(field_value.sizes["nper"])
-                        }
-                    else:
-                        if block_name not in period_data:
-                            period_data[block_name] = {}
-                        period_data[block_name][field_name] = field_value  # type: ignore
+                    blocks[block_name][field_name] = field_value
             else:
                 if field_value is not None:
+                    # only include boolean fields (keywords) if true
                     if isinstance(field_value, bool):
                         if field_value:
                             blocks[block_name][field_name] = field_value
                     else:
                         blocks[block_name][field_name] = field_value
 
-        if block_name in period_data and isinstance(period_data[block_name], dict):
-            dataset = xr.Dataset(period_data[block_name])
-            blocks[block_name] = {block_name: dataset}
-            del period_data[block_name]
-
         for arr_name, periods in period_data.items():
             for kper, arr in periods.items():
                 if kper not in period_blocks:
                     period_blocks[kper] = {}
                 period_blocks[kper][arr_name] = arr
 
         for kper, block in period_blocks.items():
-            dataset = xr.Dataset(block)
-            blocks[f"{block_name} {kper + 1}"] = {block_name: dataset}
+            assert isinstance(period_block_name, str)
+            blocks[f"{period_block_name} {kper + 1}"] = {
+                period_block_name: xr.Dataset(block, coords=block[arr_name].coords)
+            }
 
-    # total temporary hack! manually set solutiongroup 1. still need to support multiple..
+    # total temporary hack! manually set solutiongroup 1.
+    # TODO still need to support multiple..
     if "solutiongroup" in blocks:
         sg = blocks["solutiongroup"]
         blocks["solutiongroup 1"] = sg
         del blocks["solutiongroup"]
 
-    return {name: block for name, block in blocks.items() if name != "period"}
+    return {name: block for name, block in blocks.items() if name != period_block_name}
 
 
 def _make_converter() -> Converter:

test/test_codec.py

@@ -1,7 +1,5 @@
 from pprint import pprint
 
-import pytest
-
 from flopy4.mf6.codec import dumps, loads
 from flopy4.mf6.converter import COMPONENT_CONVERTER
 
@@ -55,7 +53,6 @@ def test_dumps_ic():
     pprint(loaded)
 
 
-@pytest.mark.xfail(reason="TODO")
 def test_dumps_oc():
     from flopy4.mf6.gwf import Oc