openghg_inversions.models.components#

Reusable PyMC model graph helpers.

These helpers operate on the active PyMC model context and are designed to be xarray-first. They return explicit component results or PyTensor/PyMC tensors and should not implement their own coordinate sanitization policy; coordinate handling lives in openghg_inversions.models.coords.

All component helpers operate inside an active PyMC model context. add_state_vector consumes an already resolved activity contract; add_linear_component consumes a sensitivity matrix inspected by prepare_linear_sensitivity before constructing that graph.

Naming conventions:

  • data_name: name for registered pm.Data

  • var_name: name for the latent random variable

  • output_name: name for the aligned deterministic output

  • plain name is reserved for helpers that truly create only one semantic object or where a base name is the clearest API

Frequency indicators may be supplied explicitly or derived from observation coordinates using shared helper logic based on openghg_inversions.inversion_inputs.make_freq_indicator.

class openghg_inversions.models.components.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.components.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.components.StateVectorResult(latent: TensorVariable | None, state: TensorVariable, activity: ResolvedStateActivity)#

Bases: object

Objects created by add_state_vector.

Variables:

  • latent (pytensor.tensor.variable.TensorVariable | None) – Effective sampled latent variable, or None when every state is fixed.

  • state (pytensor.tensor.variable.TensorVariable) – Full ordered state vector, including fixed values.

  • activity (openghg_inversions.models.state_activity.ResolvedStateActivity) – Resolved state-activity contract in canonical state order.

activity: ResolvedStateActivity#
latent: TensorVariable | None#
state: TensorVariable#
openghg_inversions.models.components.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.components.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.components.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.components.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.components.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.components.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.components.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.components.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.components.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.components.add_state_vector(activity: ResolvedStateActivity, /, prior_args: dict[str, Any], var_name: str) → StateVectorResult#

Construct an active/fixed state graph from a resolved activity contract.

When every state is active, this creates the same base prior graph as add_linear_component. Partial activity creates an active-only prior and restores it into a full deterministic state vector. An all-fixed policy creates no random variable and exposes the fixed values as the full deterministic state.

Parameters:

  • activity – Resolved activity and state-coordinate contract. Linear design inspection must be completed before calling this helper.

  • prior_args – Prior specification. Distribution parameters may be scalar, full-state arrays, or labelled state DataArray objects.

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

Returns:

The effective latent, full state vector, and supplied activity.

Raises:

  • KeyError – If the prior specification omits a required parameter.

  • TypeError – If the prior specification contains an unsupported value.

  • ValueError – If state-valued prior parameters are invalid.

Notes

This helper registers state variables and state coordinates, but it does not inspect or register a sensitivity matrix and does not construct a forward-model output. The registered activity mask is immutable build-time metadata in semantic terms; changing it with pm.set_data would not rebuild the latent state layout. Call this helper inside an active pm.Model context.

openghg_inversions.models.components.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.components.get_model_latent(variable: TensorVariable, base_name: str) → TensorVariable#

Return the effective latent variable for a named model component.

Parameters:

  • variable – User-facing variable returned by parse_prior.

  • base_name – Base model variable name used to look up a reparameterized latent variable.

Returns:

The reparameterized latent variable {base_name}_latent when it is present on the active model, otherwise variable.

openghg_inversions.models.components.resolve_model_variable(model: Model, base_name: str) → TensorVariable | None#

Return a named model variable, preferring the reparameterised latent form.

Parameters:

  • model – PyMC model to inspect.

  • base_name – Base variable name to resolve.

Returns:

The reparameterised latent variable {base_name}_latent when it is present on model, otherwise the user-facing variable named base_name. Returns None if neither variable exists.