openghg_inversions.array_ops#

General methods for xarray Datasets and DataArrays.

The functions here are not specific to OpenGHG inversions: they add functionality missing from xarray. These functions should accept xarray Datasets and DataArrays, and return either a Dataset or a DataArray.

Functions#

get_xr_dummies

Applies pandas get_dummies to xarray DataArrays.

sparse_xr_dot

Multiplies a Dataset or DataArray by a DataArray with sparse underlying array. The built-in xarray functionality doesn’t work correctly.

Coordinate Alignment Policy#

Spatial coordinates (e.g. lat and lon) are treated as structural invariants of the inversion workflow. Although grids are expected to be identical, small floating-point differences can arise from I/O, reprojection, or sparse/dask operations.

Because xarray uses strict, label-based alignment, even negligible coordinate differences can trigger unintended reindexing, interpolation, or alignment errors.

To avoid this, we:

  1. Validate numerical equivalence within tolerance.

  2. Force coordinate identity via assign_coords when validation passes.

This ensures deterministic arithmetic, prevents silent interpolation, and keeps grid equivalence an explicit, testable invariant.

The main function that does this is force_align; there is a legacy function align_sparse_lat_lon that was introduced as a work-around to an xarray issue, which now is written in terms of force_align.

openghg_inversions.array_ops.align_sparse_lat_lon(sparse_da: DataArray, other_array: DataArray | Dataset, *, rtol: float = 1e-05, atol: float = 1e-08) → DataArray#

Align lat/lon coordinates of sparse_da with those from other_array.

NOTE: Workaround for xarray issue #3445: pydata/xarray#3445

Validates numerical closeness of coordinates before forcing coordinate identity. No interpolation or reindexing is performed.

Parameters:

  • sparse_da – DataArray with sparse backend.

  • other_array – Dataset or DataArray providing canonical lat/lon coords.

  • rtol – Relative tolerance for coordinate comparison.

  • atol – Absolute tolerance for coordinate comparison.

Returns:

Copy of sparse_da with lat/lon coords replaced by those from other_array.

Raises:

ValueError – If sizes differ or coordinates differ beyond tolerance.

openghg_inversions.array_ops.align_to_multi_index_level_values(da: DataArray, *, multi_index: DataArray, multi_dim: str, level: str, other_dim: str) → DataArray#

Broadcast da(other_dim, …) onto multi_dim using a MultiIndex level.

Raises a ValueError if any other MultiIndex level names are already present as coordinate variables on da, because this creates ambiguous/conflicting index ownership when the MultiIndex is assigned.

Parameters:

  • da – DataArray to align, e.g. fp_x_flux with dimension other_dim=”source”.

  • multi_index – Coordinate for the MultiIndex dimension (e.g. basis_matrix[“state”]). Must be a MultiIndex coordinate that includes the level named by level.

  • multi_dim – Name of the MultiIndex dimension to align onto (e.g. “state”).

  • level – Name of the MultiIndex level within multi_index to use for alignment (e.g. “source”).

  • other_dim – Dimension on da that corresponds to level (e.g. “source”).

Returns:

A DataArray aligned to multi_dim, with a canonical MultiIndex coordinate installed on multi_dim.

Raises:

ValueError – If level is not a level of multi_index, or if other level names are already present as coordinates on da.

openghg_inversions.array_ops.concat_data_arrays(da_dict: Mapping[str, DataArray], key_dim: str, **concat_kwargs) → DataArray#
openghg_inversions.array_ops.concat_gather_data_arrays(da_dict: Mapping[str, DataArray], key_dim: str, ragged_dim: str, stack_dim: str | None = None, **concat_kwargs) → DataArray#

Concatenate DataArrays by gathering along a ragged coordinate.

For example, if the keys are site codes and the ragged dimension is time, then the stacked dimension is the usual nmeasure coordinate with (site, time) levels. The same operation can gather source-specific state blocks into (source, region_in_source) without rectangular padding. A native MultiIndex on ragged_dim is flattened after the key level rather than stored as tuple-valued labels. Any alignment policy for dimensions other than ragged_dim should be passed explicitly through concat_kwargs.

Parameters:

  • da_dict – DataArrays keyed by the values to record in key_dim.

  • key_dim – Dimension name for the mapping keys.

  • ragged_dim – Name of the ragged dimension.

  • stack_dim – Name for the gathered MultiIndex dimension.

  • **concat_kwargs – Arguments passed to xarray.concat().

Returns:

Combined DataArray with a new gathered dimension.

openghg_inversions.array_ops.concat_gather_datasets(ds_dict: Mapping[str, Dataset], key_dim: str, ragged_dim: str, stack_dim: str | None = None, missing_data_vars: Literal['error', 'drop'] = 'error', **concat_kwargs) → Dataset#

Concatenate dictionary of xr.Datasets by gathering ragged coordinates.

Parameters:

  • ds_dict – Mapping from key values to datasets to concatenate.

  • key_dim – Name for the key dimension level in the output MultiIndex.

  • ragged_dim – Ragged coordinate dimension to gather across datasets.

  • stack_dim – Optional name for the gathered dimension in the output.

  • missing_data_vars – Policy for handling data variables that are not shared by every dataset. Use "error" to raise a ValueError listing the per-dataset differences, or "drop" to keep only the intersection of shared data variables and emit a warning when any are dropped.

  • **concat_kwargs – Additional keyword arguments passed to xr.concat.

Returns:

Dataset containing gathered versions of the selected data variables.

Raises:

ValueError – If missing_data_vars="error" and the datasets do not all have identical data-variable sets, or if missing_data_vars is not recognised.

openghg_inversions.array_ops.concat_gather_datatree(dt: DataTree, key_dim: str, ragged_dim: str, stack_dim: str | None = None, missing_data_vars: Literal['error', 'drop'] = 'error', **concat_kwargs) → Dataset#

Concatenate xr.DataTree children by gathering ragged coordinates.

Parameters:

  • dt – DataTree whose children will be converted to datasets and gathered.

  • key_dim – Name for the key dimension level in the output MultiIndex.

  • ragged_dim – Ragged coordinate dimension to gather across children.

  • stack_dim – Optional name for the gathered dimension in the output.

  • missing_data_vars – Policy for handling child data variables that are not shared by every dataset. See concat_gather_datasets.

  • **concat_kwargs – Additional keyword arguments passed to xr.concat.

Returns:

Dataset containing gathered versions of the selected data variables.

openghg_inversions.array_ops.force_align(obj: T, reference: DataArray | Dataset, *, dims: Iterable[Hashable], rtol: float = 1e-05, atol: float = 1e-08) → T#

Force coordinate identity with a reference object along given dimensions.

This function verifies that obj and reference share numerically equivalent coordinates (within tolerance) along the specified dimensions. If validation succeeds, the coordinates of obj are reassigned to be identical (object identity) to those of reference.

No interpolation or reindexing is performed. If coordinates differ beyond tolerance, a ValueError is raised.

Parameters:

  • obj – The xarray DataArray or Dataset whose coordinates will be validated and reassigned.

  • reference – The object providing canonical coordinates. Must contain the specified dimensions.

  • dims – Dimensions along which to validate and enforce coordinate identity (e.g., [“lat”, “lon”]).

  • rtol – Relative tolerance for coordinate comparison. Defaults to NumPy’s default (1e-5).

  • atol – Absolute tolerance for coordinate comparison. Defaults to NumPy’s default (1e-8).

Returns:

The same type as obj, with validated coordinates reassigned.

Raises:

ValueError – If a dimension is missing, sizes differ, or coordinates differ beyond tolerance.

openghg_inversions.array_ops.get_xr_dummies(da: DataArray, categories: Sequence[Any] | Index | DataArray | ndarray | None = None, cat_dim: str = 'categories', return_sparse: bool = True) → DataArray#

Create 0-1 dummy matrix from DataArray with values that correspond to categories.

If the values of da are integers 0-N, then the result has N + 1 columns, and the (i, j) coordiante of the result is 1 if da[i] == j, and is 0 otherwise.

This function works like the pandas function get_dummies, but preserves the coordinates of the input data, and allowing the user to specify coordinates for the categories used to make the “dummies” (or “one-hot encoding”).

Parameters:

  • da – DataArray encoding categories.

  • categories – optional coordinates for categories.

  • cat_dim – dimension for categories coordinate

  • sparse – if True, store values in sparse.COO matrix

Returns:

Dummy matrix corresponding to the input vector. Its dimensions are the same as the

input DataArray, plus an additional “categories” dimension, which has one value for each distinct value in the input DataArray.

openghg_inversions.array_ops.iter_multi_index_level_slices(da: DataArray, *, multi_index: DataArray, multi_dim: str, level: str, array_dim: str) → Iterator[tuple[Hashable, ndarray, DataArray]]#

Pair unique-level values with matching gathered-axis slices.

This supports split -> operate -> combine algorithms without first broadcasting da(array_dim, ...) onto every entry of multi_dim. That distinction is essential when the non-index dimensions are large: vectorized selection with repeated level values can create a temporary of shape (multi_dim, ...) before the expensive dimensions are reduced.

Target values are visited in first-seen MultiIndex order. Positions are grouped in one pass over the gathered index, and callers retain the positions needed to restore canonical gathered order after combining operation-specific results.

Parameters:

  • da – Array indexed once per unique level value.

  • multi_index – Coordinate backed by the gathered pandas MultiIndex.

  • multi_dim – Gathered dimension name.

  • level – MultiIndex level paired with array_dim.

  • array_dim – Unique-level dimension on da.

Yields:

The level value, its integer positions in multi_dim, and the matching scalar selection from da.

Raises:

ValueError – If the gathered coordinate is not the requested MultiIndex, or the input labels are absent, duplicated, missing, or extra relative to the target level.

openghg_inversions.array_ops.select_gathered_data_array(da: DataArray, *, key: Any, key_dim: str, ragged_dim: str, stack_dim: str) → DataArray#

Select one key from a gathered DataArray and restore its ragged index.

This is the labelled inverse of one branch of concat_gather_data_arrays(). The gathered dimension remains named stack_dim, but its key level is removed. An ordinary ragged index is restored directly; flattened native MultiIndex levels are restored as a MultiIndex.

Parameters:

  • da – DataArray containing a gathered MultiIndex dimension.

  • key – Key value to select from the gathered index.

  • key_dim – Name of the key level in the gathered MultiIndex.

  • ragged_dim – Name of the original ragged dimension.

  • stack_dim – Name of the gathered dimension.

Returns:

Selected DataArray indexed by the original ragged labels.

Raises:

  • TypeError – If the gathered dimension does not use a pandas MultiIndex.

  • ValueError – If the gathered levels or requested key are invalid.

openghg_inversions.array_ops.sparse_xr_dot(da1: DataArray, da2: DataArray, dim: list[str] | None = None) → DataArray#
openghg_inversions.array_ops.sparse_xr_dot(da1: DataArray, da2: Dataset, dim: list[str] | None = None) → Dataset

Compute the matrix “dot” of a tuple of DataArrays with sparse.COO values.

This multiplies and sums over all common dimensions of the input DataArrays, and preserves the coordinates and dimensions that are not summed over.

Common dimensions are automatically selected by name. The input arrays must have at least one dimension in common. All matching dimensions will be used for multiplication.

Compared to just using da1 @ da2, this function has two advantages: 1. if da1 is sparse but not a dask array, then da1 @ da2 will fail if da2 is a dask array 2. da2 can be a Dataset, and current DataArray @ Dataset is not allowed by xarray

Parameters:

  • da1 – xr.DataArrays to multiply and sum along common dimensions.

  • da2 – xr.DataArrays to multiply and sum along common dimensions.

  • dim – optional list of dimensions to sum over; if None, then all common dimensions are summed over.

Returns:

containing the result of matrix/tensor multiplication.

The type that is returned will be the same as the type of da2.

Return type:

xr.Dataset or xr.DataArray

openghg_inversions.array_ops.to_dense(da: DataArray) → DataArray#

Convert sparse array payloads to dense NumPy arrays.

For an eager sparse input, this returns an eager NumPy-backed DataArray. For a Dask array with sparse chunks, the outer Dask collection and its chunking are preserved while a lazy block operation converts each chunk to NumPy when executed. An array whose Dask chunks are already dense is returned unchanged.

Parameters:

da – DataArray whose eager data or Dask chunk payloads may be sparse.

Returns:

DataArray with dense NumPy data or lazily densified Dask chunks.

openghg_inversions.array_ops.validate_covariance_coordinates(covariance: DataArray, *, dim: str, covariance_dim: str | None = None) → None#

Require a square covariance matrix to repeat one coordinate exactly.

The conventional column dimension is f"{dim}_cov". PyMC treats the row and column dimensions as distinct and does not verify that their labels denote the same ordered values, so that scientific invariant is checked before labels reach the model backend.

Parameters:

  • covariance – Labelled covariance matrix to validate.

  • dim – Row dimension and canonical coordinate name.

  • covariance_dim – Column dimension. Defaults to f"{dim}_cov".

Raises:

ValueError – If dimensions, shape, coordinates, or coordinate ordering do not follow the covariance convention.