implement setter for stress period data property (#280)

followup on #277 which began with a read-only property, give it a setter. dims are found on 1) the parent if exists, 2) other arrays if present, otherwise an error is raised. just supports structured grids for now. all still rough, needs cleanup at some point

Author: wpbonelli

flopy4/mf6/converter/structure.py

@@ -16,7 +16,9 @@ def structure_keyword(value, field) -> str | None:
     return field.name if value else None
 
 
-def _resolve_dimensions(self_, field) -> tuple[list[str], list[int], dict]:
+def _resolve_dimensions(
+    self_, field, *, dims: dict | None = None
+) -> tuple[list[str], list[int], dict]:
     """
     Get expected dimensions, shape, and resolved dimension values.
 
@@ -26,6 +28,9 @@ def _resolve_dimensions(self_, field) -> tuple[list[str], list[int], dict]:
         Parent object containing dimension context
     field : object
         Field specification with dims, dtype, default
+    dims : dict, optional
+        Explicit dimension sizes to use. If provided, takes precedence over
+        dims from parent or self_.__dict__.
 
     Returns
     -------
@@ -42,10 +47,15 @@ def _resolve_dimensions(self_, field) -> tuple[list[str], list[int], dict]:
         raise ValueError(f"Field {field} missing dims")
 
     # Resolve dims from model context
-    explicit_dims = self_.__dict__.get("dims", {})
+    # Priority: 1) explicit dims parameter, 2) self_.__dict__, 3) parent
     inherited_dims = dict(self_.parent.data.dims) if self_.parent else {}
+    explicit_dims = self_.__dict__.get("dims", {})
     dim_dict = inherited_dims | explicit_dims
 
+    # Override with explicitly provided dims (highest priority)
+    if dims is not None:
+        dim_dict.update(dims)
+
     # Check object attributes directly for dimension values
     # These override inherited dims (important during initialization when dims are passed as kwargs)
     for dim_name in field.dims:
@@ -482,7 +492,13 @@ def _parse_dict_format(
 
 
 def structure_array(
-    value, self_, field, *, return_xarray: bool = False, sparse_threshold: int | None = None
+    value,
+    self_,
+    field,
+    *,
+    return_xarray: bool = False,
+    sparse_threshold: int | None = None,
+    dims: dict | None = None,
 ) -> xr.DataArray | NDArray | sparse.COO:
     """
     Convert various array representations to structured arrays.
@@ -507,14 +523,17 @@ def structure_array(
         If True, return xr.DataArray; otherwise return raw array (for backward compatibility)
     sparse_threshold : int | None
         Override default sparse threshold for COO vs dense
+    dims : dict | None
+        Explicit dimension sizes (e.g., {'nper': 10, 'nodes': 100}).
+        If provided, takes precedence over dims from parent or self_.
 
     Returns
     -------
     xr.DataArray | np.ndarray | sparse.COO
         Structured array with proper shape and metadata
     """
     # Resolve dimensions
-    dims, shape, dim_dict = _resolve_dimensions(self_, field)
+    dims_names, shape, dim_dict = _resolve_dimensions(self_, field, dims=dims)
     threshold = sparse_threshold if sparse_threshold is not None else SPARSE_THRESHOLD
 
     # Handle different input types
@@ -526,7 +545,7 @@ def structure_array(
 
     if isinstance(value, dict):
         # Parse dict format with fill-forward logic
-        parsed_dict = _parse_dict_format(value, dims, tuple(shape), dim_dict, field, self_)
+        parsed_dict = _parse_dict_format(value, dims_names, tuple(shape), dim_dict, field, self_)
 
         # Build array using sparse or dense approach
         if np.prod(shape) > threshold:
@@ -553,15 +572,15 @@ def structure_array(
                         )
                         value_data = row[-1]
                         nn = get_nn(cellid, **dim_dict)
-                        if "nper" in dims:
+                        if "nper" in dims_names:
                             coords_dict[(key, nn)] = value_data
                         else:
                             coords_dict[(nn,)] = value_data
                 elif isinstance(val, dict):
                     # Nested dict: {cellid: value}
                     for cellid, v in val.items():
                         nn = get_nn(cellid, **dim_dict)
-                        if "nper" in dims:
+                        if "nper" in dims_names:
                             coords_dict[(key, nn)] = v
                         else:
                             coords_dict[(nn,)] = v
@@ -602,7 +621,7 @@ def structure_array(
                 val = parsed_dict[key]
 
                 # Determine fill range (current key to next key or end)
-                if "nper" in dims:
+                if "nper" in dims_names:
                     next_key = (
                         sorted_keys[idx + 1]
                         if idx + 1 < len(sorted_keys)
@@ -629,27 +648,27 @@ def structure_array(
                             )
                             value_data = row[-1]
                             nn = get_nn(cellid, **dim_dict)
-                            if "nper" in dims:
+                            if "nper" in dims_names:
                                 result[kper, nn] = value_data
                             else:
                                 result[nn] = value_data
                     elif isinstance(val, dict):
                         # Nested dict: {cellid: value}
                         for cellid, v in val.items():
                             nn = get_nn(cellid, **dim_dict)
-                            if "nper" in dims:
+                            if "nper" in dims_names:
                                 result[kper, nn] = v
                             else:
                                 result[nn] = v
                     elif isinstance(val, np.ndarray):
                         # Array value
-                        if "nper" in dims:
+                        if "nper" in dims_names:
                             result[kper] = val
                         else:
                             result = val
                     elif isinstance(val, xr.DataArray):
                         # xarray value
-                        if "nper" in dims:
+                        if "nper" in dims_names:
                             result[kper] = val.values
                         else:
                             result = val.values
@@ -667,15 +686,15 @@ def structure_array(
 
     elif isinstance(value, list):
         # List format
-        result = _parse_list_format(value, dims, tuple(shape), field)
+        result = _parse_list_format(value, dims_names, tuple(shape), field)
 
     elif isinstance(value, (xr.DataArray, np.ndarray)):
         # Duck array - validate and reshape if needed
-        result = _validate_duck_array(value, dims, tuple(shape), dim_dict)
+        result = _validate_duck_array(value, dims_names, tuple(shape), dim_dict)
 
         # Handle time fill-forward
-        if "nper" in dims and "nper" in dim_dict:
-            result = _fill_forward_time(result, dims, dim_dict["nper"])
+        if "nper" in dims_names and "nper" in dim_dict:
+            result = _fill_forward_time(result, dims_names, dim_dict["nper"])
 
     elif isinstance(value, (int, float)):
         # Scalar - broadcast to full shape
@@ -689,10 +708,10 @@ def structure_array(
     if return_xarray and not isinstance(result, xr.DataArray):
         # Build coordinates
         xr_coords: dict[str, Any] = {}
-        for dim in dims:
+        for dim in dims:  # type: ignore
             if dim in dim_dict:
                 xr_coords[dim] = np.arange(dim_dict[dim])
 
-        result = _to_xarray(result, dims, xr_coords)
+        result = _to_xarray(result, dims, xr_coords)  # type: ignore
 
     return result

flopy4/mf6/package.py

@@ -194,3 +194,104 @@ def stress_period_data(self) -> pd.DataFrame:
             df = df.groupby(coord_columns, as_index=False).first()
 
         return df
+
+    @stress_period_data.setter
+    def stress_period_data(self, value: pd.DataFrame) -> None:
+        """
+        Set stress period data from a DataFrame.
+
+        Parameters
+        ----------
+        value : pd.DataFrame
+            DataFrame with columns: 'kper' (stress period), spatial coordinates
+            (either 'layer'/'row'/'col' or 'node'), and field value columns.
+
+        Examples
+        --------
+        >>> # Modify existing package data
+        >>> chd = Chd(parent=gwf, head={0: {(0, 0, 0): 1.0}})
+        >>> df = chd.stress_period_data
+        >>> df['head'] = df['head'] * 2  # Double all values
+        >>> chd.stress_period_data = df  # Apply changes
+
+        >>> # Create new data from scratch
+        >>> df = pd.DataFrame({
+        ...     'kper': [0, 0, 1],
+        ...     'layer': [0, 0, 0],
+        ...     'row': [0, 5, 0],
+        ...     'col': [0, 5, 5],
+        ...     'head': [10.0, 8.0, 9.0]
+        ... })
+        >>> chd.stress_period_data = df
+        """
+        import xarray as xr
+        from xattree import get_xatspec
+
+        from flopy4.mf6.converter.structure import structure_array
+
+        if not isinstance(value, pd.DataFrame):
+            raise TypeError(f"Expected DataFrame, got {type(value)}")
+
+        # Get xattree field specifications
+        spec = get_xatspec(type(self)).flat
+
+        # Find all period block fields
+        period_fields = []
+        field_objects = {}
+        for field_name, field_spec in spec.items():
+            if field_spec.metadata.get("block") == "period" and hasattr(field_spec, "dims"):  # type: ignore
+                period_fields.append(field_name)
+                field_objects[field_name] = field_spec
+
+        if not period_fields:
+            raise TypeError("No period block fields found in package")
+
+        # Check which fields are present in the DataFrame
+        available_fields = [f for f in period_fields if f in value.columns]
+        if not available_fields:
+            raise ValueError(
+                f"DataFrame must contain at least one period field column. "
+                f"Expected one of {period_fields}, got {value.columns.tolist()}"
+            )
+
+        # Build dimension context for the converter
+        # Priority: 1) parent model dims, 2) existing array data
+        dim_dict = {}
+
+        # 1. Get dims from parent if available (most common case)
+        if hasattr(self, "parent") and self.parent is not None and hasattr(self.parent, "data"):
+            dim_dict.update(dict(self.parent.data.dims))
+
+        # 2. Extract dimensions from existing field data
+        for field_name in period_fields:
+            field_data = getattr(self, field_name, None)
+            if field_data is not None and isinstance(field_data, xr.DataArray):
+                # xarray stores dimension sizes
+                dim_dict.update(dict(field_data.sizes))
+                break  # One field is enough to get dimensions
+
+        # 3. Check if DataFrame requires structured grid dims (nrow, ncol, nlay)
+        #    but they're not available - provide helpful error
+        has_structured_coords = all(col in value.columns for col in ["layer", "row", "col"])
+        if has_structured_coords:
+            missing_dims = [d for d in ["nrow", "ncol", "nlay"] if d not in dim_dict]
+            if missing_dims:
+                raise ValueError(
+                    f"DataFrame has structured coordinates (layer/row/col) but package "
+                    f"is missing required dimensions: {missing_dims}. "
+                    f"Attach the package to a parent model with these dimensions, or use "
+                    f"node-based coordinates in the DataFrame instead."
+                )
+
+        # Update each field present in the DataFrame
+        # Pass dims explicitly to converter - no __dict__ manipulation needed
+        for field_name in available_fields:
+            field_obj = field_objects[field_name]
+
+            # Call converter with explicit dims parameter
+            converted_value = structure_array(
+                value, self, field_obj, dims=dim_dict if dim_dict else None
+            )
+
+            # Set the attribute, which will trigger on_setattr hooks (e.g., update_maxbound)
+            setattr(self, field_name, converted_value)