openghg_inversions.basis.operators#

Labelled basis geometry operators for retained scaling states.

A BasisOperator separates bucket geometry from flux weighting and native covariance. For one source, BasisOperator.basis_matrix is the bucket prolongation U_bucket with native-grid rows and retained-state columns. A gathered multisource matrix is the spatial membership template from which BasisOperator.native_prolongation() expands a canonical source-native U_bucket. Its transpose is not automatically the retained restriction Pi.

FluxWeightedBasis handles flux weighting for sensitivity projection and flux reconstruction. Native covariance actions, retained restrictions, and covariance product transforms instead belong to openghg_inversions.native_covariance, openghg_inversions.source_covariance, and openghg_inversions.basis.covariance_products.

Operators expose one canonical state dimension, named "state" by default. Multisource operators represent ragged per-source region counts with a MultiIndex over source and region-within-source, avoiding padded state arrays. BasisMeta records only the grid dimensions and state-dimension name; concrete operators perform any source-label alignment.

Serialization uses self-describing xarray.DataTree objects with schema, kind, and version metadata. Multisource operators store one source-labelled basis_flat array while readers retain compatibility with the earlier per-source child layout.

Examples

Construct operators, compute a sensitivity, and round-trip one operator:

op = BucketBasisOperator(basis_flat)
multisource_op = MultiSourceBucketBasisOperator({"a": bf_a, "b": bf_b})
sensitivity = op.sensitivity(fp_x_flux)
restored = BasisOperator.decode_datatree(op.to_datatree())

fp_x_flux contains the configured grid dimensions and typically time. In multisource workflows it may also have a source dimension, whose labels the multisource operator aligns to the state MultiIndex.

Note

Basis geometry cannot currently vary through time. A singleton time dimension is dropped; a longer time dimension raises ValueError.

class openghg_inversions.basis.operators.BasisMeta(grid_dims: tuple[str, ...] = ('lat', 'lon'), state_dim: str = 'state')#

Bases: object

Metadata describing how to apply a basis operator.

The intent is to keep this minimal: we only store what is needed for the default implementations of BasisOperator.sensitivity and BasisOperator.interpolate.

Variables:

grid_dims (tuple[str, ...]) – Dimensions to dot over when reducing a gridded quantity to the reduced state (typically (“lat”, “lon”)).
state_dim (str) – Canonical dimension name for the reduced state axis.

grid_dims: tuple[str, ...] = ('lat', 'lon')#

state_dim: str = 'state'#

class openghg_inversions.basis.operators.BasisOperator#

Bases: ABC

Abstract basis operator.

Concrete subclasses provide operator metadata, a bucket prolongation U_bucket whose ordered dimensions are the grid dimensions followed by the state dimension, and DataTree serialization methods.

The default sensitivity implementation requires fp_x_flux to contain the grid dimensions and preserves extra dimensions such as time or source.

abstract property basis_matrix: DataArray#

Return the bucket prolongation U_bucket from state to grid.

Its ordered dimensions are the configured grid dimensions followed by the state dimension, which may carry a MultiIndex coordinate.

Multiplying this matrix by a state vector reconstructs a native scaling field. Although its transpose has the shape of a restriction, it is not generally the covariance-compatible retained restriction Pi.

classmethod decode_datatree(dt: DataTree) → BasisOperator#

Dispatches a DataTree to the correct registered operator subclass.

Parameters:

dt – DataTree representation of a basis operator.

Returns:

A concrete BasisOperator instance.

Raises:

ValueError – If the schema or schema version is unsupported.
KeyError – If the kind is not registered.

abstractmethod classmethod from_datatree(dt: DataTree) → Self#

Constructs an operator instance from an xarray.DataTree.

Concrete subclasses implement this to load whatever canonical representation they write in to_datatree().

Parameters:: dt – DataTree created by to_datatree().
Returns:: An instance of the operator.

interpolate(state: DataArray, weights: DataArray | None = None) → DataArray#

Interpolates/reconstructs a gridded field from a state vector.

This maps from the reduced basis space back to the grid by multiplying the basis dummy matrix by a state vector.

If weights is provided (e.g. a flux field on the grid), it is multiplied elementwise with the basis matrix before interpolation. This corresponds to using a flux-weighted interpolation operator.

Parameters:

state – State vector with dimension meta.state_dim.
weights – Optional gridded weights with dimensions matching meta.grid_dims (and broadcastable to basis_matrix).

Returns:

Reconstructed gridded field with dimensions including meta.grid_dims.

Raises:

ValueError – If meta.state_dim is not a dimension of state.

kind: ClassVar[str]#

abstract property meta: BasisMeta#: Operator metadata (must be provided by subclasses).

native_prolongation(native_layout: DataArray, *, native_dims: tuple[str, ...]) → DataArray#

Return the bucket prolongation on a canonical labelled native grid.

This basis-side operation aligns geometry only. It preserves the basis matrix’s lazy or sparse storage and does not choose an execution or covariance materialization policy.

Parameters:

native_layout – Array carrying the canonical native coordinates and compatible auxiliary coordinates. It may contain additional non-native dimensions.
native_dims – Ordered native dimensions required by the covariance action.

Returns:

Potentially lazy U_bucket with dimensions (*native_dims, meta.state_dim).

Raises:

ValueError – If this single-source basis does not span exactly the requested native dimensions or their indexes do not match.

schema: ClassVar[str] = 'openghg_inversions.basis_operator'#

schema_version: ClassVar[int] = 1#

sensitivity(fp_x_flux: DataArray, fillna: bool = True) → DataArray#

Computes the sensitivity matrix (“H”) by dotting over the grid.

This implements the common bucket-basis forward operator H U_bucket:

fp_x_flux is a gridded quantity with dimensions that include meta.grid_dims (typically lat and lon) and usually a time dimension.
basis_matrix is a one-hot/dummy matrix that maps each grid cell to exactly one basis region/state.

The returned array keeps all non-grid dimensions from fp_x_flux and includes the reduced state dimension meta.state_dim.

Parameters:

fp_x_flux – Footprint x flux array to reduce. Must contain all meta.grid_dims.
fillna – if True, fill NaNs in fp_x_flux with 0.0.

Returns:

Sensitivity matrix with dimension meta.state_dim and any remaining non-grid dimensions (e.g. time).

abstractmethod to_datatree() → DataTree#

Serialises this operator to an xarray.DataTree.

The returned DataTree is intended to be self-describing. It must include enough information to allow round-tripping via BasisOperator.decode_datatree().

Returns:: DataTree representation of the operator.

class openghg_inversions.basis.operators.BucketBasisOperator(basis_flat: DataArray, *, meta: BasisMeta | None = None, state_dim: str | None = None, region_labels: Literal['range0', 'range1', 'basis_values'] = 'range0', state_metadata: Dataset | BasisStateMetadata | None = None, chunks: dict[str, int] | None = None)#

Bases: BasisOperator

Single flat bucket basis: basis_flat(lat, lon) with integer region labels.

Stores basis_flat and constructs basis_matrix via get_xr_dummies.

property basis_matrix: DataArray#

Return U_bucket, mapping retained scalings to native scalings.

The dimensions are native grid by retained state. Its transpose is not generally the compatible retained restriction Pi.

classmethod from_datatree(dt: DataTree) → Self#

Deserialises a BucketBasisOperator from a DataTree.

Parameters:

dt – DataTree produced by BucketBasisOperator.to_datatree().

Returns:

A reconstructed BucketBasisOperator.

Raises:

KeyError – If the serialized basis_flat variable is missing.
ValueError – If serialized labels, metadata, or state coordinates are invalid.

kind: ClassVar[str] = 'bucket'#

property meta: BasisMeta#: Basis metadata.

property state_metadata: Dataset | None#

Semantic metadata coordinates carried on the state dimension.

Returns:: Dataset containing basis_group, basis_partition, and region_in_partition indexed by meta.state_dim, or None when no grouped metadata was supplied.

to_datatree() → DataTree#

Serialises the operator to a DataTree.

Returns:: A DataTree with a dataset containing basis_flat and attributes sufficient to reconstruct the operator.

class openghg_inversions.basis.operators.MultiSourceBucketBasisOperator(basis_flat: dict[str, DataArray], *, meta: BasisMeta | None = None, source_dim: str = 'source', region_in_source_dim: str = 'region_in_source', state_dim: str | None = None, chunks: dict[str, int] | None = None)#

Bases: BasisOperator

Multiple flat bases keyed by source, with potentially ragged region counts.

The canonical state dimension is a ragged MultiIndex over (source, region_in_source).

property basis_matrix: DataArray#

Return the gathered spatial template for multisource U_bucket.

The dimensions are spatial grid by retained state; source identity is carried by the ragged state coordinate. native_prolongation() expands an explicit source-native dimension and zeros cross-source columns. This template’s transpose is not the compatible retained restriction Pi.

classmethod from_datatree(dt: DataTree) → Self#

Deserialises a MultiSourceBucketBasisOperator from a DataTree.

Parameters:: dt – DataTree produced by MultiSourceBucketBasisOperator.to_datatree().
Returns:: A reconstructed MultiSourceBucketBasisOperator. The canonical representation obtains order from the source coordinate. Earlier per-source child artifacts remain readable; for those, legacy source_order metadata controls insertion/state order when present.
Raises:: ValueError – If required basis data or source coordinates are missing, or legacy source-order metadata is malformed or inconsistent with stored source children.

interpolate(state: DataArray, weights: DataArray | None = None) → DataArray#

Interpolate/reconstruct a gridded field from a state vector.

For MultiSourceBucketBasisOperator, weights may include a source_dim that is also a level name in the gathered MultiIndex on meta.state_dim. In that case we broadcast weights along the gathered state axis (repeating per-source weights across all regions within that source) by replacing source_dim with meta.state_dim.

The state vector itself is expected to be defined on meta.state_dim and should not include a separate coordinate named like a MultiIndex level (e.g. source_dim).

Parameters:

state – State vector defined on meta.state_dim.
weights – Optional gridded weights (e.g. prior fluxes) on meta.grid_dims. May optionally include source_dim for per-source weights.

Returns:

Gridded reconstructed field on meta.grid_dims.

kind: ClassVar[str] = 'multisource_bucket'#

property meta: BasisMeta#: Basis metadata.

native_prolongation(native_layout: DataArray, *, native_dims: tuple[str, ...]) → DataArray#

Expand the gathered basis to a labelled source-native prolongation.

Parameters:

native_layout – Canonical native layout carrying an explicit source dimension and native-compatible auxiliary coordinates.
native_dims – Ordered native dimensions. The first must be this operator’s source dimension and the remainder its grid dims.

Returns:

Potentially lazy U_bucket with zero cross-source columns and dimensions (*native_dims, meta.state_dim).

Raises:

ValueError – If dimensions, source labels, or native indexes are incompatible with this multisource basis.

operator_for_source(source: str, *, state_dim: str | None = None) → BucketBasisOperator#

Return a single-source bucket operator for one source.

This keeps source-specific basis selection at the operator boundary, avoiding direct use of the legacy flat-basis compatibility view in modern postprocessing code.

Parameters:

source – Source label to select from the source-specific basis mapping.
state_dim – Optional state dimension for the returned single-source operator. If omitted, the per-source region dimension is used.

Returns:

A single-source bucket operator for source.

Raises:

ValueError – If source is not present in this operator.

sensitivity(fp_x_flux: DataArray, fillna: bool = True) → DataArray#

Compute sensitivity for multisource fp_x_flux.

Fuse source selection with the sparse spatial prolongation before contracting. This avoids broadcasting the full spatial cache across every retained state while keeping the source pairing sparse.

property source_labels: tuple[str, ...]#

Return canonical source labels in operator/state insertion order.

Returns:: Source labels in the same order used by basis_flat and the ragged state MultiIndex.

to_datatree() → DataTree#

Serialises the multisource operator to a DataTree.

The returned DataTree stores one source-labelled basis_flat array. Its source coordinate is the sole source-order representation, avoiding source names in storage paths and redundant JSON metadata.

Returns:: DataTree representation of the operator.
Raises:: ValueError – If source bases do not have compatible labeled grids.

openghg_inversions.basis.operators.drop_singleton_time(da: DataArray, *, name: str = 'basis_flat') → DataArray#

Drop a singleton time dimension if present; otherwise raise.

This is a strict helper intended for basis operators that assume a 2D basis over the grid dims. It avoids silently discarding time-varying basis information.

Parameters:

da – Input DataArray which may or may not have a time dimension.
name – Label used in error messages to identify what is being checked.

Returns:

da with time removed if it exists and has length 1, otherwise da unchanged.

Raises:

ValueError – If time exists and has length not equal to 1.

openghg_inversions.basis.operators.get_basis_operator_class(kind: str) → type[BasisOperator]#

Looks up a registered BasisOperator subclass.

Parameters:: kind – Registry key for the operator type (e.g. “bucket”).
Returns:: The registered BasisOperator subclass.
Raises:: KeyError – If kind is not registered.

openghg_inversions.basis.operators.register_basis_operator(kind: str) → Callable[[type[BasisOperatorT]], type[BasisOperatorT]]#

Registers a BasisOperator subclass for DataTree deserialisation.

This decorator builds a small module-level registry mapping a stable kind string (stored in dt.attrs[“kind”]) to a concrete BasisOperator subclass.

Parameters:: kind – Stable key identifying the operator type on disk. This is written to dt.attrs[“kind”] by BasisOperator.to_datatree() and used by BasisOperator.decode_datatree().
Returns:: A class decorator that registers the decorated class under kind.
Raises:: ValueError – If kind is already registered to a different class.