add lkt/lke to codegen

Author: mjreno

flopy4/mf6/codec/writer/filters.py

@@ -256,13 +256,21 @@ def dataset2list(value: xr.Dataset):
         return
 
     combined_mask: Any = None
+    spatial_da = None  # a non-aux DataArray for deriving spatial dims/mask
     for name, first in value.data_vars.items():
-        mask = nonempty(first)
-        combined_mask = mask if combined_mask is None else combined_mask | mask
+        if "naux" in first.dims:
+            # nonempty gives (spatial..., naux) bool array; collapse naux with any()
+            mask = nonempty(first).any(axis=-1)
+        else:
+            mask = nonempty(first)
+            spatial_da = first
+        combined_mask = mask if combined_mask is None else (combined_mask | mask)
     if combined_mask is None or not np.any(combined_mask):
         return
 
-    spatial_dims = [d for d in first.dims if d in ("nlay", "nrow", "ncol", "nodes")]
+    if spatial_da is None:
+        spatial_da = first
+    spatial_dims = [d for d in spatial_da.dims if d in ("nlay", "nrow", "ncol", "nodes")]
     has_spatial_dims = len(spatial_dims) > 0
     indices = np.where(combined_mask)
     for i in range(len(indices[0])):
@@ -287,7 +295,10 @@ def dataset2list(value: xr.Dataset):
                         else:
                             row2.append(val + 1)
                     else:
-                        row2.append(val)
+                        if hasattr(val, "ndim") and val.ndim > 0:
+                            row2.extend(float(v) for v in np.asarray(val).flat)
+                        else:
+                            row2.append(val)
             if has_spatial_dims:
                 cellid = tuple(idx[i] + 1 for idx in indices)
                 yield tuple(cellid) + tuple(row2)

flopy4/mf6/component.py

@@ -15,7 +15,7 @@
 from flopy4.mf6.write_context import WriteContext
 from flopy4.uio import IO, Loader, Writer
 
-COMPONENTS = {}
+COMPONENTS: dict[str, type] = {}
 """MF6 component registry."""
 
 

flopy4/mf6/converter/egress/unstructure.py

@@ -101,12 +101,21 @@ def _hack_structured_grid_dims(
             "nlay": range(structured_grid_dims["nlay"]),
             "ncpl": range(structured_grid_dims["ncpl"]),
         }
+    else:
+        # No structured grid decomposition available — keep nodes dim as-is.
+        return value
 
     if "nper" in value.dims:
         shape.insert(0, value.sizes["nper"])
         dims.insert(0, "nper")
         coords = {"nper": value.coords["nper"], **coords}
 
+    if "naux" in value.dims:
+        naux = value.sizes["naux"]
+        shape.append(naux)
+        dims.append("naux")
+        coords["naux"] = range(naux)
+
     return xr.DataArray(
         value.data.reshape(shape),
         dims=dims,
@@ -352,12 +361,31 @@ def _unstructure_array_component(value: Component) -> dict[str, Any]:
         for kper, block in period_blocks.items():
             key = f"period {kper + 1}"
             for arr_name, val in block.items():
-                if not np.all(val == FILL_DNODATA):
+                # G/A variant aux: split naux-dimensioned array into one readarray
+                # block per auxiliary variable, named by value.auxiliary.
+                if arr_name == "aux" and isinstance(val, xr.DataArray) and "naux" in val.dims:
+                    _aux = getattr(value, "auxiliary", None)
+                    if _aux is not None and hasattr(_aux, "values"):
+                        aux_names = [str(n) for n in _aux.values.tolist()]
+                    else:
+                        aux_names = list(_aux or [])
+                    for k in range(val.sizes["naux"]):
+                        aux_slice = val.isel(naux=k)
+                        if not np.all(aux_slice.values == FILL_DNODATA):
+                            if key not in blocks:
+                                blocks[key] = {}
+                            name = aux_names[k] if k < len(aux_names) else f"aux{k}"
+                            blocks[key][name] = aux_slice
+                elif not np.all(val == FILL_DNODATA):
                     if key not in blocks:
                         blocks[key] = {}
                     blocks[key][arr_name] = val
 
-    return {name: block for name, block in blocks.items() if name != "period"}
+    return {
+        name: block
+        for name, block in blocks.items()
+        if name != "period" and not name.startswith("__")
+    }
 
 
 # Block names that MF6 rejects if present but empty.
@@ -467,6 +495,17 @@ def _unstructure_component(value: Component) -> dict[str, Any]:
         if block_name in _LIST_BLOCK_NAMES:
             current_block = blocks.get(block_name, {})
             if current_block:
+                # Expand any 2D DataArrays (e.g. aux with shape (nlakes, naux)) into
+                # separate per-column 1D DataArrays so the Dataset stays uniformly 1D.
+                expanded: dict[str, Any] = {}
+                for name, v in current_block.items():
+                    if isinstance(v, xr.DataArray) and v.ndim == 2:
+                        for j in range(v.shape[1]):
+                            expanded[f"{name}_{j}"] = v.isel({v.dims[1]: j})
+                    else:
+                        expanded[name] = v
+                current_block = expanded
+
                 das = [v for v in current_block.values() if isinstance(v, xr.DataArray)]
                 if das and len(das) == len(current_block):
                     first_dim = das[0].dims[0] if das[0].dims else None