openghg_inversions.basis.basis_functions#

Flux-weighted basis objects and their durable representation.

FluxWeightedBasis (also exported as BasisFunctions) pairs a retained flux field with a basis operator. The operator owns the grid-to-state mapping, including region and optional source ordering; the flux remains separate so the same object can project sensitivities and reconstruct gridded fluxes.

NetCDF and Zarr persistence uses a versioned DataTree with separate basis and flux children. Only namespaced OpenGHG Inversions metadata is retained. Loading reconstructs the operator and preserves its grid, state, region, and source coordinates. A shared flux may omit the source dimension; when source labels are present, they must exactly match the operator’s labels. Source-specific flux arrays are stacked only when their native time indexes match exactly; callers must explicitly resample heterogeneous frequencies first.

Example

Project a cached footprint-times-flux field into state space:

def apply_basis_functions(ds: xr.Dataset, bf: BasisFunctions) -> xr.Dataset:
    if "fp_x_flux" not in ds:
        return ds
    return bf.sensitivity(ds.fp_x_flux).rename("H").to_dataset()

openghg_inversions.basis.basis_functions.BasisFunctions#: alias of FluxWeightedBasis

class openghg_inversions.basis.basis_functions.FluxWeightedBasis(operator: BasisOperator, flux: DataArray, metadata: dict[str, ~typing.Any]=<factory>)#

Bases: object

A thin wrapper pairing a BasisOperator with a flux field.

This class is intentionally lightweight: it stores the operator and the flux, and provides constructors that build the appropriate operator from a basis representation.

Variables:

operator (openghg_inversions.basis.operators.BasisOperator) – Basis operator that owns the state/grid mapping.
flux (xarray.core.dataarray.DataArray) – Representative flux field associated with the basis.
metadata (dict[str, Any]) – Namespaced metadata carried through DataTree serialization.

Notes

A flux with a source dimension can still be paired with a BucketBasisOperator; in that case the same basis is applied to each source when broadcasting occurs in downstream operations.
Alignment/broadcasting details are currently delegated to the operator implementations.

property basis_artifact_path: str | None#: Return the basis artifact path used to create/load this basis, when recorded.

property basis_artifact_source: str | None#: Return the source used to create/load this basis, when recorded.

flat_basis() → DataArray | dict[str, DataArray]#: Return the legacy flattened basis view used by fp_sensitivity.

flux: DataArray#

for_source(source: str, *, state_dim: str | None = None) → Self#

Return a source-specific basis object for postprocessing.

Shared-basis operators are reused with the selected source flux. Operators that own source-specific basis maps can expose an operator_for_source method to return a single-source operator without going through the legacy flat-basis view.

Parameters:

source – Flux source label to select from source-specific flux and basis operators.
state_dim – Optional state dimension for the selected source operator. When omitted, the operator chooses its default.

Returns:

A basis object with the selected source flux and a compatible operator.

Raises:

ValueError – If the underlying source-specific operator does not contain source.

classmethod from_datatree(dt: DataTree) → FluxWeightedBasis#

Deserialise from a DataTree produced by to_datatree.

Parameters:

dt – DataTree with basis and flux groups.

Returns:

A FluxWeightedBasis instance.

Raises:

KeyError – If required groups are missing.
ValueError – If the schema/version is unsupported, the flux variable is missing, or stored multi-source flux labels do not match the basis operator.

classmethod from_flat_basis(basis_flat: DataArray, flux: DataArray, *, region_labels: Literal['range0', 'range1', 'basis_values'] = 'range0', operator_kwargs: Mapping[str, Any] | None = None, metadata: Mapping[str, Any] | None = None) → FluxWeightedBasis#

Construct from a single-source (standard) flattened basis array.

Parameters:

basis_flat – Flattened basis labels on the inversion grid, e.g. dims like (lat, lon) (or whatever grid dims the operator expects). Values should label regions.
flux – Flux on the same grid dims as basis_flat. May optionally contain extra dims such as time or source; these will be carried along by downstream operations.
region_labels – Policy for the output state coordinate labels: - “range0”: 0..N-1 (legacy-friendly) - “range1”: 1..N - “basis_values”: use the unique positive labels found in basis_flat.
operator_kwargs – Optional kwargs forwarded to BucketBasisOperator.
metadata – Optional namespaced metadata to carry on the basis object.

Returns:

A FluxWeightedBasis pairing a BucketBasisOperator with the provided flux.

classmethod from_multi_source_flat_basis(basis_flat: Mapping[str, DataArray], flux: DataArray | Mapping[str, DataArray], *, operator_kwargs: Mapping[str, Any] | None = None, metadata: Mapping[str, Any] | None = None) → FluxWeightedBasis#

Construct from a multi-source flattened basis mapping.

Parameters:

basis_flat – Mapping from source name to that source’s flattened basis labels on the grid. Each value should be a DataArray on the inversion grid (e.g. dims like (lat, lon)).
flux – Flux on the inversion grid. A shared DataArray may omit the source dimension. When source is present, its unique labels must match basis_flat. A mapping is stacked along source after requiring identical time indexes.
operator_kwargs – Optional kwargs forwarded to MultiSourceBucketBasisOperator.
metadata – Optional namespaced metadata to carry on the basis object.

Returns:

A FluxWeightedBasis pairing a MultiSourceBucketBasisOperator with the provided flux.

Raises:

ValueError – If present flux source labels are duplicated or do not match the basis operator, or mapped fluxes have incompatible time indexes.

interpolate(state: DataArray, *, flux: bool = False) → DataArray#

Interpolate from state vector to the grid.

Parameters:

state – State vector values with dim matching the operator’s state dim (usually state).
flux – If True, apply flux-weighting using self.flux as weights. If False, returns the unweighted basis interpolation.

Returns:

Interpolated gridded array on the operator grid dims, with any non-dot dims preserved.

classmethod load(file_path: str | Path) → Self#

Load a basis object saved with BasisFunctions.save.

The DataTree is loaded while the xarray file context is open, so the returned BasisFunctions object does not retain references to closed file handles.

Parameters:

file_path – Path to a BasisFunctions DataTree artifact.

Returns:

The reconstructed basis object.

Raises:

OSError – If the artifact cannot be opened.
KeyError – If a required DataTree group or basis variable is missing.
ValueError – If the serialized schema, metadata, labels, or operator state is invalid.

metadata: dict[str, Any]#

operator: BasisOperator#

plot(*, shuffle: bool = False, **plot_kwargs: Any) → Any#

Plot basis labels.

For single-source basis, returns the result of xarray’s plot call.
For multi-source basis, creates one subplot per source and returns (fig, axes).

Parameters:

shuffle – If True, randomly permute region labels for visual separation.
**plot_kwargs – Forwarded to xarray’s .plot().

Returns:

Plotting object(s). For multi-source, returns (fig, axes).

save(output_file: str | Path) → None#

Save the basis object as a DataTree artifact.

Parameters:: output_file – Path to a .nc or .zarr artifact. NetCDF is used for .nc paths and Zarr is used for .zarr paths.
Raises:: ValueError – If output_file does not end with .nc or .zarr.

select_sources(sources: Sequence[str]) → Self#

Return a copy restricted to source-specific basis functions.

This is used when a retained artifact contains more sources than the current run requested. Single-source/shared-basis operators are kept as is, with only flux restricted when it has a source dimension.

Parameters:: sources – Source names and order to keep.
Returns:: A basis object restricted to sources.
Raises:: ValueError – If a source-specific basis is missing a requested source.

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

Compute sensitivity (grid -> state) via the underlying operator.

This is typically used to form the reduced Jacobian:: H = operator.sensitivity(fp_x_flux)

Parameters:

fp_x_flux – Footprints multiplied by flux, on the inversion grid dims. May contain extra dims (e.g. time, site, source).
fillna – if True, fill NaNs in fp_x_flux with 0.0.

Returns:

Sensitivity with a state-like dimension as defined by the operator.

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

Return the authoritative source labels when this basis is source-aware.

Source-specific operators own their source order. A shared operator may instead obtain source labels from its retained flux. A completely source-independent basis returns None.

Returns:: Ordered source labels, or None when neither the operator nor flux is source-indexed.

to_datatree() → DataTree#

Serialise to a DataTree with basis and flux groups.

Returns:

basis: BasisOperator DataTree (via operator.to_datatree()).
flux: a Dataset containing the flux DataArray as variable flux.

Return type:

A DataTree with

validated() → Self#

Return a copy that revalidates mutable xarray-backed state.

Returns:: A basis object reconstructed from the current operator, flux, and metadata values.

with_flux(flux: DataArray) → Self#: Return a copy with the same operator and metadata but a different flux.

with_metadata(metadata: Mapping[str, Any]) → Self#: Return a copy with additional metadata.

openghg_inversions.basis.basis_functions.basis_functions_from_fp_all_flat_basis(*, fp_all: dict, basis_flat: DataArray | Mapping[str, DataArray], metadata: Mapping[str, Any] | None = None) → FluxWeightedBasis#

Compatibility adapter that constructs a basis object from flat basis data.

This is also the public handoff for project-owned flat-basis algorithms in copied RHIME runners. A single two-dimensional array defines a shared spatial basis. A source-keyed mapping defines source-specific bases and is restricted and ordered to match runtime flux sources when fp_all is sector resolved. In both cases the retained flux is reconstructed from the current fp_all input rather than from the flat-basis producer.

Parameters:

fp_all – Legacy merged-data dictionary containing a ".flux" side channel and optional ".split_by_sectors" flag.
basis_flat – Two-dimensional shared flat basis array, or source-keyed two-dimensional flat basis arrays, produced by a compatible basis algorithm or legacy loading path.
metadata – Optional namespaced metadata to carry on the basis object.

Returns:

A retained BasisFunctions object with flux reconstructed from fp_all.

Raises:

TypeError – If a runtime flux entry or basis input has an unsupported type.
ValueError – If runtime flux is missing or inconsistent, a source-specific basis omits a runtime source, or basis labels and dimensions violate the retained-basis contract.
xarray.AlignmentError – If basis and current-run flux grids are not physically compatible.

openghg_inversions.basis.basis_functions.flux_from_fp_all(fp_all: dict) → DataArray#

Legacy adapter that builds representative flux from fp_all.

This is a temporary compatibility helper for the current fp_sensitivity path. Remove it when issue #429 makes BasisFunctions.sensitivity the source of prepared sensitivities.

Parameters:

fp_all – Legacy merged-data dictionary containing fp_all[".flux"] and optional fp_all[".split_by_sectors"] metadata.

Returns:

A combined flux array for non-sectoral workflows, or a source-stacked flux array for sectoral workflows. Source-stacked output follows the insertion order of fp_all[".flux"].

Raises:

ValueError – If fp_all[".flux"] is missing or empty, or a sectoral source mapping mixes timed and timeless arrays or contains unequal, missing, or duplicate native time coordinates.
TypeError – If a flux entry cannot be converted to a DataArray.