openghg_inversions.models#

Reusable model components, coordinates, priors, and state APIs.

Linear sensitivity preparation removes structural zero columns once before graph construction; state activity then controls scientific fixing among the full labelled state.

Importing this package configures PyTensor before re-exporting reusable model primitives. RHIME-specific recipes and contracts live in openghg_inversions.rhime.

class openghg_inversions.models.CoordRegistry(pymc_coords: dict[str, ~numpy.ndarray]=<factory>, original_coords: dict[str, ~typing.Any]=<factory>, auxiliary_coords: dict[str, ~xarray.core.dataarray.DataArray]=<factory>)#

Bases: object

Track scientific and PyMC-safe coordinates for a model.

Variables:

  • pymc_coords (dict[str, numpy.ndarray]) – Sanitized coordinates actually registered with PyMC.

  • original_coords (dict[str, Any]) – Original scientific coordinates keyed by model dimension name.

  • auxiliary_coords (dict[str, xarray.core.dataarray.DataArray]) – Additional non-dimension coordinates attached to model dimensions, such as exploded time or site coordinates derived from a stacked nmeasure MultiIndex.

add(coords: dict[str, Any] | Coordinates, *, model_dims: tuple[str, ...] | list[str] | set[str] | None = None) → None#

Register model and auxiliary coordinates with consistency checks.

Parameters:

  • coords – Coordinate mapping or xarray coordinate container to register.

  • model_dims – Optional subset of model dimensions represented by the current data variable. Auxiliary coordinates attached to these dimensions are also preserved when possible.

Raises:

ValueError – If the same coordinate name is registered more than once with conflicting lengths, shapes, or values.

auxiliary_coords: dict[str, DataArray]#
original_coords: dict[str, Any]#
pymc_coords: dict[str, ndarray]#
class openghg_inversions.models.CorrelatedLognormalPrior(mean: DataArray, arithmetic_covariance: DataArray | ndarray, *, covariance_dim: str | None = None)#

Bases: object

Validated labelled moments for one correlated positive state vector.

The arithmetic covariance is an already-reduced dense matrix. It uses a distinct second dimension whose entries follow the primary state coordinate in the same order. Only the primary axis owns the rich scientific coordinate, avoiding duplicate MultiIndex level coordinates while retaining an explicit labelled-row contract.

Variables:

  • mean – Arithmetic mean with dimension (state_dim,). Access returns an independent deep copy.

  • arithmetic_covariance – Dense arithmetic covariance with dimensions (state_dim, covariance_dim). Access returns an independent deep copy.

  • latent_mean – Derived Gaussian mean with dimension (state_dim,). Access returns an independent deep copy.

  • latent_covariance – Derived Gaussian covariance with dimensions (state_dim, covariance_dim). Access returns an independent deep copy.

  • latent_cholesky – Lower Cholesky factor with dimensions (state_dim, covariance_dim). Access returns an independent deep copy.

  • state_dim (str) – Name of the scientific state dimension.

  • covariance_dim (str) – Name of the same-order covariance column dimension.

Notes

Construction owns all supplied values. Array properties return deep copies so callers cannot mutate the cached validated moments through the public API.

property arithmetic_covariance: DataArray#

Return the validated arithmetic covariance.

Returns:

An independent deep copy of the dense labelled covariance.

covariance_dim: str#
classmethod from_dataset(dataset: Dataset) → CorrelatedLognormalPrior#

Reload and revalidate a serialized correlated-state contract.

Parameters:

dataset – Dataset produced by to_dataset().

Returns:

A newly validated prior derived from the stored arithmetic moments.

Raises:

ValueError – If schema metadata, variables, dimensions, labels, or cached latent moments are missing or inconsistent.

Warns:

UserWarning – If the stored reduced state contains more than 1000 components. This is an operational threshold, not a mathematical limit.

Notes

Stored latent values are checked against moments recomputed from the eagerly materialized arithmetic inputs; they are never trusted as an alternative construction path.

property latent_cholesky: DataArray#

Return the Cholesky factor of the latent covariance.

Returns:

An independent deep copy of the dense lower-triangular factor.

property latent_covariance: DataArray#

Return the derived latent Gaussian covariance.

Returns:

An independent deep copy of the dense latent covariance.

property latent_mean: DataArray#

Return the derived latent Gaussian mean.

Returns:

An independent deep copy of the labelled latent mean.

property mean: DataArray#

Return the validated arithmetic mean.

Returns:

An independent deep copy of the labelled arithmetic mean.

state_dim: str#
to_dataset() → Dataset#

Serialize the validated arithmetic and latent moments.

Returns:

A new labelled dataset containing deep copies of both moment parameterizations and the correlated-state schema metadata.

class openghg_inversions.models.CorrelatedStateResult(latent: TensorVariable, state: TensorVariable, prior: CorrelatedLognormalPrior)#

Bases: object

Objects created by add_correlated_lognormal_state.

Variables:

  • latent (pytensor.tensor.variable.TensorVariable) – Standard-normal whitened state used by the sampler.

  • state (pytensor.tensor.variable.TensorVariable) – Positive user-facing state with the requested arithmetic LogNormal moments.

  • prior (openghg_inversions.correlated_state.CorrelatedLognormalPrior) – Validated backend-neutral moment contract used to build the graph.

latent: TensorVariable#
prior: CorrelatedLognormalPrior#
state: TensorVariable#
class openghg_inversions.models.LinearComponentResult(data: TensorVariable, latent: TensorVariable | None, state: TensorVariable, output: TensorVariable, activity: ResolvedStateActivity)#

Bases: object

Objects created by add_linear_component().

activity: ResolvedStateActivity#
data: TensorVariable#
latent: TensorVariable | None#
output: TensorVariable#
state: TensorVariable#
class openghg_inversions.models.PreparedLinearSensitivity(sensitivity: DataArray, removed: DataArray, output_dim: str)#

Bases: object

A retained sensitivity matrix and its lossless full-state mapping.

output_dim: str#
removed: DataArray#
property retained_indices: ndarray#

Return full-state positions retained by sensitivity.

sensitivity: DataArray#
property state_dim: str#

Return the sole state dimension.

class openghg_inversions.models.ResolvedStateActivity(state_dim: str, active: DataArray, fixed_value: DataArray, zero_sensitivity: DataArray)#

Bases: object

A state-activity policy aligned to one detected sensitivity matrix.

Variables:
active: DataArray#
property active_indices: ndarray#

Return positional indices of active states in canonical order.

property fixed_indices: ndarray#

Return positional indices of inactive states in canonical order.

fixed_value: DataArray#
property n_active: int#

Return the number of sampled states.

property n_state: int#

Return the full state-vector length.

state_dim: str#
zero_sensitivity: DataArray#
class openghg_inversions.models.StateActivity(active: bool | ndarray | DataArray = True, fixed_value: float | int | ndarray | DataArray = 1.0, fixed_groups: tuple[str, ...] = (), group_coord: str = 'basis_group')#

Bases: object

Describe which states are sampled and how inactive states are fixed.

Parameters:

  • active – Boolean scalar, positional one-dimensional mask in canonical state order, or labelled one-dimensional mask. A labelled mask is aligned to the canonical state coordinate carried by the zero-sensitivity mask, so its input order need not match the canonical order. Set this to False to freeze a whole component or sector.

  • fixed_value – Scalar or one-dimensional state-aligned values used for inactive states. Multiplicative scaling states default to one.

  • fixed_groups – basis_group labels to freeze. Group selection is by coordinate value, never by state-number ranges.

  • group_coord – Name of the state coordinate containing group labels.

Explicit active masks and fixed groups are combined with structural zero-column removal using logical AND. Structural removal is owned by sensitivity preparation and cannot be disabled as an activity policy.

active: bool | ndarray | DataArray = True#
fixed_groups: tuple[str, ...] = ()#
fixed_value: float | int | ndarray | DataArray = 1.0#
group_coord: str = 'basis_group'#
openghg_inversions.models.active_prior_args(prior_args: dict[str, Any], activity: ResolvedStateActivity) → dict[str, Any]#

Return prior arguments sliced to active states in canonical order.

Scalar parameters are retained. One-dimensional NumPy state parameters are interpreted in canonical order and must describe the full state vector. Labelled xarray parameters are first aligned by the canonical state coordinate and then subset, allowing input label order to differ safely. Array-backed parameters are materialized at this model-building boundary. The support arrays of an Interpolated prior are preserved rather than interpreted as state-valued parameters.

Parameters:

  • prior_args – PyMC prior specification.

  • activity – Resolved activity contract in canonical state order.

Returns:

A copy of prior_args containing scalar or active-state parameters.

Raises:

ValueError – If an array-valued distribution parameter is not scalar or full-state one-dimensional data.

openghg_inversions.models.add_coherent_affine_component(fixed_contribution: DataArray, linear_signal: TensorVariable, /, *, output_name: str) → TensorVariable#

Add a labelled fixed contribution to an already composed linear signal.

Parameters:

  • fixed_contribution – Labelled affine intercept mu_prior - H_alpha @ m_alpha, not the full prior-forward mean. Its name owns the registered pm.Data name and its coordinates own the output axis labels.

  • linear_signal – Existing linear signal composed by the calling model recipe.

  • output_name – Name for the affine deterministic output.

Returns:

The deterministic sum of the registered fixed contribution and linear signal.

Notes

This component uses the equivalent coherent-reduction identity mu = mu_prior + H_alpha @ (x - m_alpha) = (mu_prior - H_alpha @ m_alpha) + H_alpha @ x.

This helper does not construct states, sensitivities, ratios, or channel signals. The calling recipe owns those scientific choices.

openghg_inversions.models.add_coords(coords: dict[str, ndarray] | Coordinates, *, model_dims: tuple[str, ...] | list[str] | set[str] | None = None) → None#

Register coordinates on the active model and capture scientific metadata.

Parameters:

  • coords – Coordinate mapping or xarray coordinate container to register.

  • model_dims – Optional subset of model dimensions represented by the current data variable. When provided, auxiliary coordinates attached to those dimensions are also stored in the registry.

This helper must be called inside an active pm.Model context.

openghg_inversions.models.add_correlated_lognormal_state(prior: CorrelatedLognormalPrior, /, *, var_name: str) → CorrelatedStateResult#

Add a whitened correlated LogNormal state to the active PyMC model.

Parameters:

  • prior – Validated labelled arithmetic and latent moment contract.

  • var_name – Name of the positive user-facing state. The whitened standard normal is named {var_name}_latent.

Returns:

The whitened latent, positive state, and supplied prior contract.

Raises:

ValueError – If the arithmetic mean, latent moments, Cholesky diagonal, or exponentiated central state is not finite and positive where required in PyMC’s configured floating-point dtype.

Notes

This function must run in an active pm.Model context. After backend dtype validation completes, it mutates that model by registering the length-p state coordinate, {var_name}_latent random variable, and length-p {var_name} deterministic state.

prior should contain reduced arithmetic moments produced together with the matching forward operator and Gaussian unresolved-error term. The coherent covariance, transformed-forward-model, and aggregation-error identities are exact only for a jointly Gaussian state. Reusing those first two moments with a LogNormal retained state and Gaussian unresolved contribution is a moment-matched closure, not exact LogNormal marginalization. Known-exact state fixing is handled separately by StateActivity in the state-linear component builders.

openghg_inversions.models.add_correlated_lognormal_state_with_activity(activity: ResolvedStateActivity, prior: CorrelatedLognormalPrior, /, *, var_name: str) → StateVectorResult#

Construct a correlated LogNormal state with exact active/fixed values.

The arithmetic-moment prior is subset to sampled states before its LogNormal transformation. Inactive states keep their exact fixed values in the full public vector. This is the correlated counterpart of add_state_vector().

Parameters:

  • activity – Resolved activity in canonical full-state order.

  • prior – Validated labelled arithmetic-moment LogNormal prior for the full state.

  • var_name – Name of the full user-facing state vector.

Returns:

Effective whitened latent, full state vector, and supplied activity.

openghg_inversions.models.add_inferpymc_likelihood_component(data: Dataset, /, mu: TensorVariable, mu_bc: TensorVariable | None, sigprior: dict, sigma_alignment: SigmaAlignment, offset: TensorVariable | None = None, power: dict | float = 1.99, pollution_events_from_obs: bool = False, no_model_error: bool = False, output_dim: str = 'nmeasure') → TensorVariable#

Add the inferpymc observation model.

mu is the non-baseline forward-model contribution. mu_bc is the baseline contribution, usually H_bc @ bc, plus offset if applicable.

Parameters:

  • data – Canonical inferpymc input dataset.

  • mu – Non-baseline forward-model contribution.

  • mu_bc – Baseline contribution, if present.

  • sigprior – Prior specification for sigma.

  • sigma_alignment – Backend-neutral site and period alignment for sigma.

  • offset – Optional aligned offset term.

  • power – Scalar or prior specification controlling pollution-event scaling.

  • pollution_events_from_obs – Whether to derive pollution events from the observations instead of mu.

  • no_model_error – Whether to bypass the model-error term.

  • output_dim – Observation/output dimension name.

Returns:

The epsilon deterministic variable used by the observation model.

openghg_inversions.models.add_linear_component(prepared: PreparedLinearSensitivity, /, data_name: str, prior_args: dict, var_name: str, output_name: str, state_activity: StateActivity | None = None, output_dim: str = 'nmeasure', compute_deterministic: bool = True) → LinearComponentResult#

Add one independent labelled linear component.

Parameters:

  • prepared – Retained sensitivity and full-state mapping produced by prepare_linear_sensitivity().

  • data_name – Name used when registering the data as pm.Data.

  • prior_args – Prior specification for the latent random variable.

  • var_name – Name for the latent random variable.

  • output_name – Name for the aligned deterministic output.

  • output_dim – Observation/output dimension name.

  • compute_deterministic – Whether to wrap the aligned output in pm.Deterministic.

  • state_activity – Optional active/fixed policy over the full scientific state. None samples every retained column.

Returns:

The registered sensitivity, effective latent, full state, aligned forward contribution, and resolved activity.

openghg_inversions.models.add_linked_linear_component(prepared: PreparedLinearSensitivity, linked_state: TensorVariable, /, *, data_name: str, output_name: str) → TensorVariable#

Apply a prepared sensitivity to an already constructed linked state.

Parameters:

  • prepared – Sensitivity prepared by prepare_linear_sensitivity(). It owns retained-state selection and the output dimension.

  • linked_state – Existing full state expression in the state order owned by prepared. The caller owns any ratio or other transformation used to construct this expression.

  • data_name – Name used to register the prepared sensitivity as pm.Data.

  • output_name – Name for the output-dimension-aligned deterministic.

Returns:

The linked linear signal registered as output_name.

Notes

This helper creates neither a state nor a multiplier. It only registers the prepared sensitivity and owns the resulting deterministic output.

openghg_inversions.models.add_model_data(data: DataArray, name: str | None = None) → TensorVariable#

Add labelled xarray data to the active PyMC model.

Parameters:

  • data – Xarray data to register as pm.Data.

  • name – Optional PyMC variable name. If omitted, data.name is used.

Returns:

The registered pm.Data tensor for data.

Raises:

ValueError – If no name can be determined for the data variable.

openghg_inversions.models.add_offset_component(observations: DataArray, /, prior_args: dict, offset_freq_indicator: DataArray | ndarray | None = None, offset_freq: str | None = None, var_name: str = 'offset_latent', output_name: str = 'offset', output_dim: str = 'nmeasure', drop_first: bool = False, per_site: bool = True) → TensorVariable#

Add a global, site-only, or site-by-period offset component.

Parameters:

  • observations – Observation data carrying an observation-aligned site coordinate.

  • prior_args – Prior specification for the offset latent variable.

  • offset_freq_indicator – Optional explicit observation-aligned offset frequency indicator.

  • offset_freq – Optional frequency string used to derive an indicator when offset_freq_indicator is not provided.

  • var_name – Name for the latent offset variable.

  • output_name – Name for the aligned deterministic offset output.

  • output_dim – Observation/output dimension name.

  • drop_first – Whether to omit the first site indicator column.

  • per_site – Whether to create site-specific terms. If false, create one global scalar latent offset and broadcast it over observations.

Returns:

The aligned offset deterministic variable.

Raises:

ValueError – If observations does not have an observation-aligned site coordinate, or if global-offset options conflict.

openghg_inversions.models.add_sigma_component(alignment: SigmaAlignment, /, prior_args: dict, compute_deterministic: bool = False) → TensorVariable#

Register a latent sigma component and align it to observations.

Parameters:

  • alignment – Backend-neutral site and period alignment for the component.

  • prior_args – Prior specification for the sigma random variable.

  • compute_deterministic – Whether to register the aligned sigma term as a deterministic variable.

Returns:

The observation-aligned sigma tensor or deterministic variable.

openghg_inversions.models.apply_linear_sensitivity(prepared: PreparedLinearSensitivity, state: TensorVariable, /, *, data_name: str, output_name: str, compute_deterministic: bool = True) → TensorVariable#

Apply a prepared sensitivity to an already-built full state vector.

openghg_inversions.models.attach_coord_registry(model: Model, registry: CoordRegistry) → None#

Attach a coordinate registry to a PyMC model.

openghg_inversions.models.detect_zero_sensitivity(sensitivity: DataArray, *, output_dim: str = 'nmeasure') → DataArray#

Return the labelled mask of exactly-zero sensitivity columns.

Parameters:

  • sensitivity – Finite two-dimensional sensitivity matrix containing output_dim and one uniquely labelled state dimension.

  • output_dim – Name of the observation/output dimension.

Returns:

A materialized boolean mask over the state dimension. State labels and auxiliary state coordinates are retained from sensitivity.

Raises:

ValueError – If sensitivity is not a finite two-dimensional design with the required output dimension and unique state labels.

Notes

Finite validation and exact-zero reduction materialize lazy design data during model construction.

openghg_inversions.models.get_coord_registry(model: Model) → CoordRegistry | None#

Return the coordinate registry attached to a PyMC model, if any.

openghg_inversions.models.parse_prior(name: str, prior_params: dict[str, Any], **kwargs) → TensorVariable#

Create a continuous PyMC prior from a prior-parameter dictionary.

Parameters:

  • name – Name of the user-facing PyMC variable to create.

  • prior_params – Prior specification including pdf and any distribution parameters accepted by the chosen PyMC distribution.

  • **kwargs – Additional keyword arguments forwarded to the created PyMC variable, such as dims.

Returns:

The created PyMC random variable or deterministic transform.

Raises:

ValueError – If prior_params["pdf"] does not name a supported PyMC continuous distribution.

This helper must be called inside an active pm.Model context because it registers the created variable with the current model.

openghg_inversions.models.prepare_linear_sensitivity(sensitivity: DataArray, *, output_dim: str = 'nmeasure') → PreparedLinearSensitivity#

Remove exact-zero columns once and retain their full-state mapping.

This is the eager inspection boundary for a labelled sensitivity matrix. The returned sensitivity keeps borrowed array data and contains only nonzero columns; removed retains the complete scientific state coordinate and auxiliary state metadata for reconstruction and provenance.

openghg_inversions.models.registered_model(*args: Any, **kwargs: Any) → Model#

Construct a PyMC model with RHIME coordinate tracking attached.

Parameters:

  • *args – Positional arguments forwarded to pymc.Model.

  • **kwargs – Keyword arguments forwarded to pymc.Model. A scientific coords mapping is sanitized for PyMC and retained in the attached registry.

Returns:

A new PyMC model with a seeded CoordRegistry attached.

openghg_inversions.models.resolve_state_activity(zero_sensitivity: DataArray, policy: StateActivity | None = None) → ResolvedStateActivity#

Resolve an active/fixed policy against a labelled zero-state mask.

Parameters:

  • zero_sensitivity – One-dimensional, strictly boolean mask identifying exactly-zero design columns. It must have a unique labelled state coordinate and may carry auxiliary state coordinates used by the policy.

  • policy – Optional activity policy. When omitted, every retained column is active and structurally removed states use inactive value one.

Returns:

A policy aligned to the mask’s canonical state coordinate.

Raises:

ValueError – If the zero mask is not one-dimensional, boolean, and uniquely labelled; supplied arrays cannot be aligned; or requested group metadata is absent or not state-aligned.

Notes

Policy vectors are materialized during model construction. Use detect_zero_sensitivity to validate and reduce a sensitivity matrix before calling this function.

openghg_inversions.models.restore_inferencedata_coords(idata: InferenceData, coords_or_registry: CoordRegistry | dict[str, Any]) → InferenceData#

Restore saved scientific coordinates onto matching InferenceData groups.

Parameters:

  • idata – Inference data object returned by sampling.

  • coords_or_registry – Either a CoordRegistry or a legacy mapping of original coordinates keyed by dimension name.

Returns:

The same InferenceData object with compatible original coordinates and auxiliary coordinates restored onto its xarray groups.

On this page
Show Source