feat(stilt): JacobianBuilder.build_from_grid (0.1.0b4)

CHANGELOG.md

@@ -7,6 +7,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.1.0b4] - 2026-06-17
+
+### Added
+- **`JacobianBuilder.build_from_grid`** (`fips.problems.flux.transport.stilt`): Build the Jacobian from an xarray target grid (`lon`/`lat` or `x`/`y`; `NaN`-masked cells) or a coords list, conservatively regridding each footprint onto it. `build_from_coords` is now a thin wrapper that forwards to it. Advances the "build Jacobian from geometries or nested grid" roadmap item.
+
+### Changed
+- **PYSTILT dependency bumped to `>=0.1.0a3`** (`pyproject.toml`): requires the conservative `Footprint.aggregate` (the old nearest-point sampling undercounted coarse-grid sensitivities), which `build_from_grid` and `build_from_coords` rely on.
+
 ### Fixed
 - **Parallel footprint loading in JacobianBuilder** (`build_from_coords`): Footprints are now loaded inside each parallel worker instead of serially before dispatch, eliminating a major NFS I/O bottleneck. Hours filtering is applied up-front by parsing the sim_id to avoid dispatching workers for footprints that would be discarded.
 - **Unnamed time column causing MatrixBlock validation error** (`_build_jacobian_row_from_path`): `Footprint.aggregate` returns a DataFrame whose column index has `name=None`; this caused "All levels in the columns must be named." Fixed by setting `agg.columns.name = "time"` before stacking.

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "fips"
-version = "0.1.0b3"
+version = "0.1.0b4"
 description = "Flexible Inverse Problem Solver (FIPS)"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -35,7 +35,7 @@ flux = [
     "cartopy>=0.24.0",  # numpy 2.0 support
     "h5py>=3.11.0",  # numpy 2.0 support
     "matplotlib>=3.6.0",
-    "pystilt>=0.1.0a1",
+    "pystilt>=0.1.0a3",  # conservative Footprint.aggregate + Grid.to_xarray
 ]
 
 [dependency-groups]

src/fips/problems/flux/transport/stilt/builder.py

@@ -6,9 +6,12 @@
 footprints over specified time bins and spatial resolutions.
 """
 
+from __future__ import annotations
+
 import logging
 from collections import defaultdict
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import pandas as pd
 from joblib import Parallel, delayed
@@ -17,6 +20,12 @@
 
 from fips.matrix import MatrixBlock
 
+if TYPE_CHECKING:
+    import xarray as xr
+
+    # An aggregation target: an (x, y) coords list or an xarray grid.
+    Target = list[tuple[float, float]] | xr.DataArray | xr.Dataset
+
 logger = logging.getLogger(__name__)
 
 
@@ -66,6 +75,43 @@ def build_from_coords(
         coords: list[tuple[float, float]] | dict[str, list[tuple[float, float]]],
         flux_times: pd.IntervalIndex,
         footprint: str,
+        **kwargs,
+    ) -> MatrixBlock | dict[str, MatrixBlock]:
+        """
+        Build the Jacobian H from output-grid coordinates and flux time bins.
+
+        Convenience wrapper over :meth:`build_from_grid` for callers that have a
+        plain list of ``(x, y)`` cell centers (a regular grid is assumed; PYSTILT
+        infers the cell size from the coordinate spacing). Prefer
+        :meth:`build_from_grid` when you already have an xarray grid, since it
+        carries resolution, bounds, CRS, and any mask explicitly.
+
+        Parameters
+        ----------
+        coords : list[tuple[float, float]] | dict[str, list[tuple[float, float]]]
+            Output grid cell centers as (x, y) tuples. Pass a dict to build
+            multiple Jacobians over different coordinate sets.
+        flux_times : pd.IntervalIndex
+            Time bins for the fluxes.
+        footprint : str
+            Name of the footprint to load from each simulation.
+        **kwargs
+            Forwarded to :meth:`build_from_grid` (``mets``, ``time_range``,
+            ``location_ids``, ``subset_hours``, ``num_processes``,
+            ``location_mapper``, ``timeout``, ``threshold``, ``sparse``).
+
+        Returns
+        -------
+        MatrixBlock | dict[str, MatrixBlock]
+            Single MatrixBlock when coords is a list; dict when coords is a dict.
+        """
+        return self.build_from_grid(coords, flux_times, footprint, **kwargs)
+
+    def build_from_grid(
+        self,
+        grid: Target | dict[str, Target],
+        flux_times: pd.IntervalIndex,
+        footprint: str,
         *,
         mets: str | list[str] | None = None,
         time_range: tuple | None = None,
@@ -76,15 +122,21 @@ def build_from_coords(
         timeout: float | int | None = None,
         threshold: float | None = 1e-15,
         sparse: bool = False,
-    ) -> "MatrixBlock | dict[str, MatrixBlock]":
+    ) -> MatrixBlock | dict[str, MatrixBlock]:
         """
-        Build the Jacobian matrix H from specified coordinates and flux time bins.
+        Build the Jacobian matrix H over a target grid and flux time bins.
+
+        Each footprint is conservatively regridded onto ``grid`` (see
+        :meth:`stilt.Footprint.aggregate`) and its time-binned sensitivities
+        become one Jacobian row.
 
         Parameters
         ----------
-        coords : list[tuple[float, float]] | dict[str, list[tuple[float, float]]]
-            Coordinates of the output grid points as (x, y) tuples. Pass a
-            dict to build multiple Jacobians over different coordinate sets.
+        grid : xr.DataArray | xr.Dataset | list[tuple[float, float]] | dict
+            The target grid: a CF xarray grid (``lon``/``lat`` or ``x``/``y``
+            coordinates; ``NaN`` cells in a 2-D DataArray are masked out) or a
+            plain list of ``(x, y)`` cell centers. Pass a dict to build multiple
+            Jacobians over different grids.
         flux_times : pd.IntervalIndex
             Time bins for the fluxes.
         footprint : str
@@ -113,12 +165,11 @@ def build_from_coords(
         Returns
         -------
         MatrixBlock | dict[str, MatrixBlock]
-            Single MatrixBlock when coords is a list; dict when coords is a dict.
+            Single MatrixBlock when ``grid`` is a single grid; dict when a dict.
         """
         logger.info("Building Jacobian matrix...")
 
-        if not isinstance(coords, dict):
-            coords = {"DEFAULT": coords}
+        targets = grid if isinstance(grid, dict) else {"DEFAULT": grid}
 
         if time_range is None:
             time_range = (flux_times[0].left, flux_times[-1].right)
@@ -148,7 +199,7 @@ def build_from_coords(
         results = Parallel(n_jobs=num_processes, timeout=timeout)(
             delayed(_build_jacobian_row_from_path)(
                 path=path,
-                coords=coords,
+                targets=targets,
                 location_dim=self.location_dim,
                 time_dim=self.time_dim,
                 flux_times=flux_times,
@@ -201,32 +252,27 @@ def build_from_coords(
         return H_dict
 
 
-def _build_jacobian_row_from_path(  # must be top-level for multiprocessing
-    path: Path,
-    coords: dict[str, list[tuple[float, float]]],
+def _build_jacobian_row(
+    fp: Footprint,
+    targets: dict[str, Target],
     location_dim: str,
     time_dim: str,
     flux_times: pd.IntervalIndex,
-) -> "dict[str, pd.DataFrame] | None":
+) -> dict[str, pd.DataFrame] | None:
     """
-    Load one footprint from disk and build its Jacobian row.
+    Build one footprint's Jacobian row by aggregating it onto each target.
 
-    Loading inside the worker avoids serial NFS reads and large object
-    pickling that would occur if footprints were pre-loaded before dispatch.
+    Returns ``None`` when the footprint does not overlap any target (all-zero
+    aggregate), so it contributes no row.
     """
-    try:
-        fp = Footprint.from_netcdf(path)
-    except Exception:
-        return None
-
     obs_index = pd.MultiIndex.from_arrays(
         [[fp.receptor.location_id], [fp.receptor.time]],
         names=[location_dim, time_dim],
     )
 
     rows: dict[str, pd.DataFrame] = {}
-    for key, coord_list in coords.items():
-        agg = fp.aggregate(coord_list, flux_times)
+    for key, target in targets.items():
+        agg = fp.aggregate(target, flux_times)
         if not agg.values.any():
             continue
 
@@ -237,3 +283,24 @@ def _build_jacobian_row_from_path(  # must be top-level for multiprocessing
         rows[key] = row
 
     return rows or None
+
+
+def _build_jacobian_row_from_path(  # must be top-level for multiprocessing
+    path: Path,
+    targets: dict[str, Target],
+    location_dim: str,
+    time_dim: str,
+    flux_times: pd.IntervalIndex,
+) -> dict[str, pd.DataFrame] | None:
+    """
+    Load one footprint from disk and build its Jacobian row.
+
+    Loading inside the worker avoids serial NFS reads and large object
+    pickling that would occur if footprints were pre-loaded before dispatch.
+    """
+    try:
+        fp = Footprint.from_netcdf(path)
+    except Exception:
+        return None
+
+    return _build_jacobian_row(fp, targets, location_dim, time_dim, flux_times)