Grouped basis and state metadata#

Basis functions can be constructed separately for different parts of the full domain—for example, a detailed inner region and a coarser outer region—and then combined into a single basis array. The combined array remembers which group and partition each basis region came from, so states can still be selected and analysed by group after they have been assembled.

BasisLayout provides this assembly step. It combines disjoint, partition-local label maps into the one flat basis consumed by a bucket operator. Alongside that flat map it produces the stable state metadata basis_group, basis_partition, and region_in_partition.

This is currently a lower-level Python construction API. Import it from openghg_inversions.basis.layout when an application has already generated or loaded the partition label arrays. The implementation is eager and in-memory, and it does not infer a remainder partition.

The cells on this page are executed when the documentation is built. You can download them as a Jupyter notebook to modify and run locally.

Build an inner/outer layout#

Each partition uses positive integers for its own local region labels. Zero, negative, and NaN values mean that a cell is outside that partition. The two arrays below are disjoint, and the explicitly named remainder covers every cell outside the inner partition.

import numpy as np
import xarray as xr

from openghg_inversions.basis.basis_functions import BasisFunctions
from openghg_inversions.basis.layout import BasisLayout, BasisPartition


coords = {"lat": [51.0, 52.0], "lon": [-2.0, -1.0]}
inner_labels = xr.DataArray(
    [[1, 2], [0, 0]],
    dims=("lat", "lon"),
    coords=coords,
    name="inner_labels",
)
remainder_labels = xr.DataArray(
    [[0, 0], [1, 1]],
    dims=("lat", "lon"),
    coords=coords,
    name="remainder_labels",
)

inner = BasisPartition(name="generated_inner", labels=inner_labels, group="inner")
remainder = BasisPartition(
    name="explicit_remainder",
    labels=remainder_labels,
    group="outer",
)
layout = BasisLayout(partitions=(inner, remainder), state_dim="region")
result = layout.to_flat_basis()
result.basis_flat.values

array([[1, 2],
       [3, 3]])

The combined map has raw basis_label values 1, 2, and 3. The metadata is keyed by those raw labels before an operator applies its final state-label policy. Selecting only the small state-aligned variables keeps the spatial coordinates out of the display.

result.state_metadata[["basis_group", "basis_partition", "region_in_partition"]]

<xarray.Dataset> Size: 96B
Dimensions:              (basis_label: 3)
Coordinates:
  * basis_label          (basis_label) int64 24B 1 2 3
Data variables:
    basis_group          (basis_label) object 24B 'inner' 'inner' 'outer'
    basis_partition      (basis_label) object 24B 'generated_inner' ... 'expl...
    region_in_partition  (basis_label) int64 24B 1 2 1
Attributes:
    state_dim:  region

The exact byte counts in an xarray display can vary by platform; the important contract is the coordinate values and their order.

Remainders are explicit#

Leaving out the remainder is an error rather than a request for BasisLayout to invent one.

try:
    BasisLayout(partitions=(inner,), state_dim="region").to_flat_basis()
except ValueError as error:
    uncovered_error = str(error)
    print(uncovered_error)

BasisLayout partitions leave 2 grid cells unmapped.

Attach the metadata to an operator#

Pass both outputs to BasisFunctions.from_flat_basis. The state_metadata option is forwarded to BucketBasisOperator.

flux = xr.DataArray(
    np.ones((2, 2, 1)),
    dims=("lat", "lon", "time"),
    coords={**coords, "time": ["2020-01-01"]},
    name="flux",
)
basis_functions = BasisFunctions.from_flat_basis(
    result.basis_flat,
    flux,
    region_labels="range0",
    operator_kwargs={
        "state_dim": "region",
        "state_metadata": result.state_metadata,
    },
)
matrix = basis_functions.operator.basis_matrix
print("Raw basis labels:", result.state_metadata.basis_label.values)
print("Operator state labels:", matrix.region.values)
matrix.coords.to_dataset()[
    ["basis_group", "basis_partition", "region_in_partition"]
].compute()

Raw basis labels: [1 2 3]
Operator state labels: [0 1 2]

<xarray.Dataset> Size: 96B
Dimensions:              (region: 3)
Coordinates:
  * region               (region) int64 24B 0 1 2
    basis_group          (region) object 24B 'inner' 'inner' 'outer'
    basis_partition      (region) object 24B 'generated_inner' ... 'explicit_...
    region_in_partition  (region) int64 24B 1 2 1
Data variables:
    *empty*

Here the raw basis labels are [1, 2, 3], while region_labels="range0" makes the final operator state coordinate [0, 1, 2]. Metadata follows the states through that relabelling.

Select states by group#

Because the grouping fields are coordinates on the state dimension, normal xarray operations can select an inner or outer subset.

state = xr.DataArray(
    [0.9, 1.1, 1.0],
    dims="region",
    coords={
        "region": matrix.region.values,
        "basis_group": ("region", matrix.basis_group.values),
        "basis_partition": ("region", matrix.basis_partition.values),
        "region_in_partition": ("region", matrix.region_in_partition.values),
    },
    name="scaling",
)
inner_state = state.where(state.basis_group == "inner", drop=True)
outer_state = state.where(state.basis_group == "outer", drop=True)
print("Inner states:", inner_state.region.values)
print("Outer states:", outer_state.region.values)
print("Outer partition:", outer_state.basis_partition.item())

Inner states: [0 1]
Outer states: [2]
Outer partition: explicit_remainder

Round-trip preservation#

BasisFunctions preserves the state metadata in its versioned DataTree representation. save and load use the same representation for NetCDF or Zarr artifacts.

restored = BasisFunctions.from_datatree(basis_functions.to_datatree())
restored_matrix = restored.operator.basis_matrix
{
    "region": restored_matrix.region.equals(matrix.region),
    "basis_group": restored_matrix.basis_group.equals(matrix.basis_group),
    "region_in_partition": restored_matrix.region_in_partition.equals(
        matrix.region_in_partition
    ),
}

{'region': True, 'basis_group': True, 'region_in_partition': True}

Stability boundary#

This page documents the implemented layout artifact, state metadata, operator attachment, and serialization contract. BasisLayout remains an eager, in-memory, lower-level Python API and does not infer a remainder.

Separate group-specific weight or sensitivity construction, grouped priors, public RHIME configuration, and a full extension API are outside this contract. Those broader changes are outside the implemented layout API described here.