Customising RHIME#

RHIME is built from reusable functions for data preparation, basis construction, model construction, sampling, and output. This is a deliberate design choice: when the standard workflow does not express the science you need, you should be able to start from a working example and replace the smallest relevant part.

The examples on this page use a procedural style. Data is passed from one function to the next in the order that the inversion runs, so readers do not need to design a new class hierarchy or framework before changing a model. Some library functions return structured objects that group related values, but the examples show how to use those objects directly.

Reusing the components does not make every custom workflow automatic. A replacement function is responsible for the scientific choices it introduces, and a deeper change may also need different data preparation. The remaining RHIME stages can only validate the inputs and assumptions covered by their documented interfaces.

Choose the smallest starting point that fits the change:

  • To change only the likelihood, pass a Python function to run_rhime.

  • To resume from externally supplied merged observations, footprints, and fluxes, pass a borrowed RhimeMergedData object as merged_data.

  • To change a preparation stage such as basis construction, copy the visible runner and replace that stage.

  • To start from prepared data or replace the complete model, use run_rhime_from_prepared_inputs.

Resume from cached or external scientific data#

run_rhime and run_rhime_multisector accept merged_data as a Python-only handoff. It bypasses OpenGHG acquisition and merged-cache I/O, checks the single- or multi-sector layout, and then re-enters the visible recipe at filtering:

result = run_rhime(
    config_file="config.ini",
    merged_data=my_merged_data,
)

The supplied RhimeMergedData and its xarray or Dask arrays remain borrowed. Filtering returns a replacement handoff when it changes observations; basis construction and labelled assembly consume that result without mutating the external object. A normal reload_merged_data request instead belongs to the same retrieval stage and may read a configured artifact from disk.

Change the likelihood with a Python function#

For an ordinary likelihood variation, keep RHIME’s complete acquisition-to-output workflow and pass one Python function to run_rhime. The function is deliberately unavailable to INI files: it is not imported from configuration or stored in a run or model specification.

The complete integration is one named argument:

from my_project.likelihoods import likelihood_builder
from openghg_inversions.rhime import run_rhime

result = run_rhime(
    config_file="config.ini",
    mismatch_model=None,
    likelihood_builder=likelihood_builder,
)

RHIME calls the function with explicit keyword arguments while constructing the PyMC model: the prepared observations, completed forward-model mean, reported observation error, selected aggregation error, and output dimension. Pollution-event-only terms remain inside the built-in pollution-event component. The function adds epsilon and the canonical observed variable y to the active model and returns y. There is no framework context or likelihood-result record to construct.

Options owned only by a custom likelihood can be supplied separately with likelihood_kwargs. RHIME expands that mapping into the callable without hiding the common scientific arrays in an opaque object:

result = run_rhime(
    config_file="config.ini",
    mismatch_model=None,
    likelihood_builder=likelihood_builder,
    likelihood_kwargs={"degrees_of_freedom": 4.0},
)

mismatch_model=None explicitly opts out of the built-in selected by the configuration template. Passing both a custom callable and either built-in selection is rejected as ambiguous.

RHIME passes these options directly and records them beside the likelihood identity in result metadata and any saved inversion output. A non-empty mapping is rejected when no likelihood_builder is active.

Editable likelihood#

The following example replaces RHIME’s Gaussian observation distribution with a Student-t distribution while reusing RHIME’s current mean and error construction:

 1"""Editable likelihoods using RHIME's explicit scientific inputs.
 2
 3This project-owned function changes RHIME's observation distribution while
 4reusing its current mean and fixed-error construction.
 5
 6``likelihood_builder`` is the exported runner seam. Importing this module does
 7not build a model, retrieve data, sample, or write outputs.
 8"""
 9
10import pymc as pm
11import pytensor.tensor as pt
12import xarray as xr
13from pytensor.tensor.variable import TensorVariable
14
15from openghg_inversions.observation_error import (
16    AggregationError,
17    validate_observation_error_arrays,
18)
19
20
21def likelihood_builder(
22    *,
23    observations: xr.DataArray,
24    observation_error: xr.DataArray,
25    aggregation_error: AggregationError,
26    mean: TensorVariable,
27    output_dim: str,
28    degrees_of_freedom: float = 4.0,
29) -> TensorVariable:
30    """Build a fixed-error Student-t model from RHIME's common inputs.
31
32    Degrees of freedom default to 4.0. RHIME's ``epsilon`` is passed as the
33    Student-t scale, not its marginal standard deviation; at the default the
34    marginal standard deviation is ``epsilon * sqrt(2)``.
35
36    Args:
37        observations: Observed mole fractions.
38        observation_error: Reported observation-error standard deviations.
39        aggregation_error: Validated fixed aggregation-error representation.
40        mean: Completed forward-model concentration.
41        output_dim: Observation dimension used by named PyMC variables.
42        degrees_of_freedom: Positive Student-t degrees of freedom supplied as
43            a custom likelihood option.
44
45    Returns:
46        The observed Student-t variable, named ``y``.
47
48    Raises:
49        ValueError: If dense or low-rank aggregation error would require a
50            multivariate Student-t likelihood.
51
52    Notes:
53        On success, this adds RHIME's canonical observation-state nodes and
54        ``y`` to the active PyMC model. It performs no sampling or output
55        writes. Unsupported aggregation modes are rejected before any nodes
56        are added.
57    """
58    if aggregation_error.mode not in {"none", "diagonal"}:
59        raise ValueError("This Student-t model assumes independent observations.")
60    if degrees_of_freedom <= 0:
61        raise ValueError("Student-t degrees of freedom must be positive.")
62    validate_observation_error_arrays(
63        observations,
64        observation_error,
65        None,
66        owner="Custom Student-t likelihood",
67        output_dim=output_dim,
68    )
69    reported_error = pm.Data(
70        "error",
71        pm.floatX(observation_error.transpose(output_dim).compute().values),
72        dims=output_dim,
73    )
74    aggregation_variance = pm.Data(
75        "aggregation_error_marginal_variance",
76        pm.floatX(aggregation_error.marginal_variance),
77        dims=output_dim,
78    )
79    epsilon = pm.Deterministic(
80        "epsilon",
81        pt.sqrt(reported_error**2 + aggregation_variance),
82        dims=output_dim,
83    )
84    observed = pm.StudentT(
85        "y",
86        nu=degrees_of_freedom,
87        mu=mean,
88        sigma=epsilon,
89        observed=pm.floatX(observations.transpose(output_dim).compute().values),
90        dims=output_dim,
91    )
92    return observed
93
94
95__all__ = ["likelihood_builder"]

RHIME records the function’s module and name, together with its explicit likelihood_kwargs, in the result and in saved output. It does not copy the Python source code into the output, so a project should keep the source and environment used for an inversion. Ordinary likelihood builders retain the canonical y and epsilon names used by sampling and output code.

The example rejects dense and low-rank aggregation covariance because it uses an independent Student-t distribution. Supporting those aggregation-error modes would require a multivariate likelihood.

Built-in mismatch equations are direct model components, not examples of this custom-callback contract. In particular, models.additive_sigma.add_additive_sigma_likelihood adds an absolute concentration-scale variance, while models.pollution_event owns the pollution-enhancement-scaled equation. Select the absolute additive equation directly with mismatch_model="additive_sigma" in run_rhime or run_rhime_multisector; sigma_prior, sigma_freq, and sigma_per_site then configure its scale. Additive sigma does not load the historical min_error floor unless use_minimum_error_floor=True is selected explicitly. A custom likelihood loads only mf and mf_error from the universal observation inputs and does not require min_error.

Legacy run_hbmcmc additive configuration is translated at that script’s entry point into the same explicit likelihood settings. Parameter resolution stores built-in mismatch science in RhimeModelSpec before the standard or multisector recipe is called; the recipes themselves select no default.

Built-in aggregation covariance relies on the guarantees of its construction pipeline. A custom pipeline that assembles its own covariance may optionally call openghg_inversions.observation_error.validate_complete_observation_covariance() with its fixed independent variance. Model builders do not run this eager diagnostic automatically.

Optional project CLI#

The short wrapper below packages the same one-call integration as a reusable Python function and command-line entry point:

  1"""Run ordinary RHIME with a project-owned Student-t likelihood.
  2
  3This is the preferred low-ceremony customization route. It keeps RHIME's
  4complete acquisition-to-output pipeline and changes only the direct-Python
  5likelihood callable passed to :func:`openghg_inversions.rhime.run_rhime`.
  6Use :func:`run_with_likelihood` from Python or :func:`main` from the command
  7line with a normal RHIME configuration. The workflow may retrieve or reload
  8data, materializes related model arrays together at the named PyMC boundary,
  9samples, and writes requested outputs while canonical xarray/Dask inputs remain
 10borrowed.
 11"""
 12
 13from __future__ import annotations
 14
 15import argparse
 16from collections.abc import Sequence
 17import json
 18from pathlib import Path
 19from typing import Any
 20
 21from openghg_inversions.rhime import RhimeResult, run_rhime
 22
 23from .likelihoods import likelihood_builder
 24
 25
 26def run_with_likelihood(
 27    *,
 28    config_file: str | Path | None = None,
 29    **kwargs: Any,
 30) -> RhimeResult:
 31    """Run standard RHIME with a project-owned Student-t likelihood.
 32
 33    Args:
 34        config_file: Optional RHIME INI configuration file.
 35        **kwargs: Standard RHIME option names that override values from
 36            ``config_file``.
 37
 38    Returns:
 39        The sampled result and any outputs requested by the RHIME options.
 40
 41    Raises:
 42        TypeError: If the likelihood builder does not return a PyTensor
 43            variable.
 44        ValueError: If required options are missing, aggregation error is
 45            unsupported, or the likelihood omits canonical ``y`` or
 46            ``epsilon`` variables.
 47
 48    Notes:
 49        This workflow may retrieve or reload data, materializes related model
 50        arrays together at the named PyMC boundary without mutating canonical
 51        prepared inputs, runs sampling, and writes configured outputs.
 52    """
 53    kwargs["mismatch_model"] = None
 54    return run_rhime(
 55        config_file=config_file,
 56        likelihood_builder=likelihood_builder,
 57        **kwargs,
 58    )
 59
 60
 61def _json_object(value: str) -> dict[str, Any]:
 62    """Parse one command-line JSON object containing RHIME option overrides.
 63
 64    Args:
 65        value: JSON text to parse.
 66
 67    Returns:
 68        The decoded JSON object.
 69
 70    Raises:
 71        argparse.ArgumentTypeError: If the text is invalid JSON or decodes to
 72            a non-object value.
 73    """
 74    try:
 75        parsed = json.loads(value)
 76    except json.JSONDecodeError as exc:
 77        raise argparse.ArgumentTypeError(f"invalid JSON: {exc.msg}") from exc
 78    if not isinstance(parsed, dict):
 79        raise argparse.ArgumentTypeError("--kwargs must decode to a JSON object")
 80    return parsed
 81
 82
 83def main(argv: Sequence[str] | None = None) -> RhimeResult:
 84    """Parse command-line options and run the customized RHIME workflow.
 85
 86    Args:
 87        argv: Command-line tokens excluding the program name. ``None`` reads
 88            tokens from :data:`sys.argv`.
 89
 90    Returns:
 91        The RHIME result returned by :func:`run_with_likelihood`.
 92
 93    Raises:
 94        SystemExit: If argument parsing fails or help is requested.
 95    """
 96    parser = argparse.ArgumentParser(description=__doc__)
 97    parser.add_argument("config_file", nargs="?", type=Path, help="RHIME INI configuration file")
 98    parser.add_argument(
 99        "--kwargs",
100        type=_json_object,
101        default={},
102        metavar="JSON",
103        help="additional RHIME options as a JSON object",
104    )
105    args = parser.parse_args(argv)
106    return run_with_likelihood(config_file=args.config_file, **args.kwargs)
107
108
109if __name__ == "__main__":
110    main()

Run it with a normal RHIME configuration and optional JSON overrides:

python -m package_name.run_with_likelihood config.ini \
    --kwargs '{"output_path": "outputs", "output_format": "inv_out"}'

run_rhime_multisector accepts the same Python-only builder contract. The standard multi-sector model retains sector flux components and roles, then passes their combined observation mean to the likelihood, so no special case or semantic compromise is required.

Use the seam from a generated project#

Create a normal downstream package with the current OpenGHG project cookiecutter:

uvx cookiecutter gh:openghg/openghg-project-cookiecutter

The template already declares openghg and openghg_inversions as dependencies and creates a src layout. Add the project-owned files beside the generated package’s existing modules, without copying any OpenGHG Inversions implementation:

src/my_inversion/
  __init__.py
  likelihoods.py
  runner.py
tests/
  test_runner.py

Put only the scientific change in likelihoods.py:

 1"""Project-owned scientific variation for a cookiecutter-generated package."""
 2
 3import pymc as pm
 4import pytensor.tensor as pt
 5import xarray as xr
 6from pytensor.tensor.variable import TensorVariable
 7
 8from openghg_inversions.observation_error import (
 9    AggregationError,
10    validate_observation_error_arrays,
11)
12
13
14def likelihood_builder(
15    *,
16    observations: xr.DataArray,
17    observation_error: xr.DataArray,
18    aggregation_error: AggregationError,
19    mean: TensorVariable,
20    output_dim: str,
21    degrees_of_freedom: float = 4.0,
22) -> TensorVariable:
23    """Build an independent Student-t observation likelihood for RHIME.
24
25    Args:
26        observations: Observed mole fractions.
27        observation_error: Reported observation-error standard deviations.
28        aggregation_error: Validated fixed aggregation-error representation.
29        mean: Completed forward-model concentration.
30        output_dim: Observation dimension used by named PyMC variables.
31        degrees_of_freedom: Positive Student-t degrees of freedom supplied as
32            a custom likelihood option.
33
34    Returns:
35        The observed Student-t variable, named ``y``.
36
37    Raises:
38        ValueError: If aggregation error requires a multivariate likelihood.
39    """
40    if aggregation_error.mode not in {"none", "diagonal"}:
41        raise ValueError("This Student-t model assumes independent observations.")
42    if degrees_of_freedom <= 0:
43        raise ValueError("Student-t degrees of freedom must be positive.")
44
45    validate_observation_error_arrays(
46        observations,
47        observation_error,
48        None,
49        owner="Custom Student-t likelihood",
50        output_dim=output_dim,
51    )
52    reported_error = pm.Data(
53        "error",
54        pm.floatX(observation_error.transpose(output_dim).compute().values),
55        dims=output_dim,
56    )
57    aggregation_variance = pm.Data(
58        "aggregation_error_marginal_variance",
59        pm.floatX(aggregation_error.marginal_variance),
60        dims=output_dim,
61    )
62    epsilon = pm.Deterministic(
63        "epsilon",
64        pt.sqrt(reported_error**2 + aggregation_variance),
65        dims=output_dim,
66    )
67    observed = pm.StudentT(
68        "y",
69        nu=degrees_of_freedom,
70        mu=mean,
71        sigma=epsilon,
72        observed=pm.floatX(observations.transpose(output_dim).compute().values),
73        dims=output_dim,
74    )
75    return observed
76
77
78__all__ = ["likelihood_builder"]

Keep project-level invocation in runner.py. Its one library call owns no retrieval, filtering, basis, input assembly, sampling, predictive selection, or output implementation:

 1"""Run RHIME from a cookiecutter-generated project package."""
 2
 3from __future__ import annotations
 4
 5import argparse
 6from collections.abc import Sequence
 7import json
 8from pathlib import Path
 9from typing import Any
10
11from openghg_inversions.rhime import RhimeResult, run_rhime
12
13from .likelihoods import likelihood_builder
14
15
16def run(
17    *,
18    config_file: str | Path | None = None,
19    **kwargs: Any,
20) -> RhimeResult:
21    """Run ordinary RHIME with this project's likelihood.
22
23    Args:
24        config_file: Optional RHIME INI configuration file.
25        **kwargs: Standard RHIME option names overriding the configuration.
26
27    Returns:
28        The library-owned RHIME result and any requested supported output.
29    """
30    kwargs["mismatch_model"] = None
31    return run_rhime(
32        config_file=config_file,
33        likelihood_builder=likelihood_builder,
34        **kwargs,
35    )
36
37
38def _json_object(value: str) -> dict[str, Any]:
39    """Parse JSON command-line overrides as an object."""
40    try:
41        parsed = json.loads(value)
42    except json.JSONDecodeError as exc:
43        raise argparse.ArgumentTypeError(f"invalid JSON: {exc.msg}") from exc
44    if not isinstance(parsed, dict):
45        raise argparse.ArgumentTypeError("--kwargs must decode to a JSON object")
46    return parsed
47
48
49def main(argv: Sequence[str] | None = None) -> RhimeResult:
50    """Parse project CLI arguments and run the customized inversion."""
51    parser = argparse.ArgumentParser(description=__doc__)
52    parser.add_argument("config_file", type=Path, help="RHIME INI configuration file")
53    parser.add_argument(
54        "--kwargs",
55        type=_json_object,
56        default={},
57        metavar="JSON",
58        help="additional RHIME options as a JSON object",
59    )
60    args = parser.parse_args(argv)
61    return run(config_file=args.config_file, **args.kwargs)
62
63
64if __name__ == "__main__":
65    main()

After uv sync --extra dev, run the module with a normal RHIME INI file:

uv run python -m my_inversion.runner inversion.ini \
    --kwargs '{"output_path": "outputs", "output_format": "inv_out"}'

An optional console command can point directly at the same main function. Add this table to the generated pyproject.toml:

[project.scripts]
my-inversion = "my_inversion.runner:main"

Then the equivalent command is:

uv run my-inversion inversion.ini \
    --kwargs '{"output_path": "outputs", "output_format": "inv_out"}'

The generated runner uses documented names from openghg_inversions.rhime. Its likelihood module imports reusable components from their documented owner modules: models.pollution_event, observation_error, and sigma. The dependency direction is therefore the generated project to OpenGHG Inversions; OpenGHG Inversions does not import the consumer package. Pin the release or commit used for a scientific run in the downstream project’s lockfile. An optional RHIME recipe or generated-project CI in the generic cookiecutter would be a separate cross-repository change and is not required for this workflow.

Copy the complete runner#

The copied runner below is an advanced, version-coupled escape hatch for scientific changes beyond the likelihood seam. It makes the major stages visible while importing their implementations from the supported public RHIME API. Prefer run_rhime_from_prepared_inputs when replaying prepared inputs, replacing the complete model, or deliberately starting from a different preparation graph.

The deliberate change is the likelihood passed to build_standard_rhime_model_result: the example selects the same project-owned Student-t builder as the preferred form. Acquisition, filtering, basis construction, labelled input assembly, conversion of delayed arrays for PyMC, sampling, predictive selection, filenames, and output handling remain library-owned.

In a project created with the OpenGHG project cookiecutter, copy the preferred modules to src/<package_name>/likelihoods.py and src/<package_name>/run_with_likelihood.py. Copy the complete module below to src/<package_name>/rhime_runner.py only when the project needs to own a deeper orchestration change. Import scientific stage implementations from openghg_inversions.rhime rather than copying those implementations.

Run it with a normal RHIME configuration and optional overrides:

python -m package_name.rhime_runner config.ini \
    --start-date 2020-01-01 --end-date 2020-02-01 \
    --output-path outputs --draws 1000 --tune 1000 --chains 4

Less common Python/config options can be supplied as a JSON object:

python -m package_name.rhime_runner config.ini \
    --kwargs '{"reload_merged_data": true, "output_format": "inv_out"}'

The test suite imports all three sources, exercises both runners, and validates the likelihood contract directly, so the documentation and runnable examples cannot drift apart.

  1"""Advanced copy-and-modify runner for a standard RHIME inversion.
  2
  3This example deliberately replaces RHIME's default likelihood with a
  4project-owned Student-t builder. The orchestration is intentionally copied so
  5that a project can make deeper scientific changes while continuing to reuse
  6the supported acquisition, preparation, model, sampling, and output stages.
  7Use :func:`run_custom_rhime` from Python or :func:`main` from the command line.
  8The standard single-sector stages are preserved, model inputs materialize only
  9at the explicit PyMC boundary, and a run may acquire or reload data, sample a
 10model, and write its configured outputs.
 11"""
 12
 13from __future__ import annotations
 14
 15import argparse
 16from collections.abc import Sequence
 17import json
 18from pathlib import Path
 19from time import perf_counter
 20from typing import Any
 21
 22from openghg_inversions.rhime import (
 23    RhimeResult,
 24    assemble_rhime_inputs,
 25    build_rhime_basis,
 26    build_rhime_sensitivities,
 27    build_standard_rhime_model_result,
 28    filter_rhime_observations,
 29    make_standard_rhime_result,
 30    make_standard_rhime_outputs,
 31    materialize_pymc_inputs,
 32    params_from_config,
 33    resolve_rhime_options,
 34    retrieve_or_reload_rhime_data,
 35    sample_rhime_model,
 36    standard_model_input_names,
 37    with_prepared_rhime_sites,
 38)
 39
 40from .likelihoods import likelihood_builder
 41
 42
 43def run_custom_rhime(
 44    *,
 45    config_file: str | Path | None = None,
 46    **kwargs: Any,
 47) -> RhimeResult:
 48    """Run standard RHIME with a project-owned Student-t likelihood.
 49
 50    Args:
 51        config_file: Optional RHIME INI configuration file.
 52        **kwargs: Standard RHIME option names that override values from
 53            ``config_file``.
 54
 55    Returns:
 56        The sampled result and any outputs requested by the RHIME options.
 57
 58    Raises:
 59        TypeError: If the likelihood builder returns an invalid result type.
 60        ValueError: If required options are missing, aggregation or output is
 61            unsupported, or likelihood roles or metadata are invalid.
 62
 63    Notes:
 64        This workflow may retrieve or reload data, materializes related model
 65        arrays together at the named PyMC boundary without mutating canonical
 66        prepared inputs, runs sampling, and writes outputs requested by the
 67        resolved RHIME options.
 68    """
 69    params = (
 70        params_from_config(config_file, extra_kwargs=kwargs, normalise=False)
 71        if config_file is not None
 72        else dict(kwargs)
 73    )
 74    params["mismatch_model"] = None
 75    setup = resolve_rhime_options(params=params, multisector=False)
 76
 77    merged = retrieve_or_reload_rhime_data(setup.data_args, multisector=False)
 78    filtered = filter_rhime_observations(merged, setup.data_args)
 79    basis_functions = build_rhime_basis(filtered, setup.data_args)
 80    site_data = build_rhime_sensitivities(
 81        filtered,
 82        basis_functions,
 83        setup.data_args,
 84        multisector=False,
 85    )
 86    prepared = assemble_rhime_inputs(filtered, basis_functions, site_data, setup.data_args)
 87    run_spec = with_prepared_rhime_sites(setup.run_spec, prepared)
 88
 89    model_inputs = materialize_pymc_inputs(
 90        prepared,
 91        variable_names=standard_model_input_names(
 92            prepared,
 93            run_spec.model,
 94        ),
 95    )
 96    build_and_sample_start = perf_counter()
 97    model_build_result = build_standard_rhime_model_result(
 98        prepared=prepared,
 99        model_inputs=model_inputs,
100        run_spec=run_spec,
101        # This is the deliberate scientific replacement in the copied runner.
102        likelihood_builder=likelihood_builder,
103    )
104    idata = sample_rhime_model(model_build_result, setup.sampler)
105
106    result = make_standard_rhime_result(
107        prepared=prepared,
108        run_spec=run_spec,
109        sampler=setup.sampler,
110        model_build_result=model_build_result,
111        idata=idata,
112        build_and_sample_seconds=perf_counter() - build_and_sample_start,
113        likelihood_builder=likelihood_builder,
114    )
115    make_standard_rhime_outputs(result=result, prepared=prepared)
116    return result
117
118
119def _json_object(value: str) -> dict[str, Any]:
120    """Parse one command-line JSON object containing additional RHIME options.
121
122    Args:
123        value: JSON text to parse.
124
125    Returns:
126        The decoded JSON object.
127
128    Raises:
129        argparse.ArgumentTypeError: If the text is invalid JSON or decodes to
130            a non-object value.
131    """
132    try:
133        parsed = json.loads(value)
134    except json.JSONDecodeError as exc:
135        raise argparse.ArgumentTypeError(f"invalid JSON: {exc.msg}") from exc
136    if not isinstance(parsed, dict):
137        raise argparse.ArgumentTypeError("--kwargs must decode to a JSON object")
138    return parsed
139
140
141def main(argv: Sequence[str] | None = None) -> RhimeResult:
142    """Parse command-line options and run the customized RHIME workflow.
143
144    This performs the same acquisition, materialization, sampling, and output
145    side effects as :func:`run_custom_rhime`.
146
147    Args:
148        argv: Command-line tokens excluding the program name. ``None`` reads
149            tokens from :data:`sys.argv`.
150
151    Returns:
152        The RHIME result returned by :func:`run_custom_rhime`.
153
154    Raises:
155        SystemExit: If argument parsing fails or help is requested.
156    """
157    parser = argparse.ArgumentParser(description=__doc__)
158    parser.add_argument("config_file", nargs="?", type=Path, help="RHIME INI configuration file")
159    parser.add_argument("--start-date")
160    parser.add_argument("--end-date")
161    parser.add_argument("--output-path", type=Path)
162    parser.add_argument("--output-name")
163    parser.add_argument("--draws", type=int)
164    parser.add_argument("--tune", type=int)
165    parser.add_argument("--chains", type=int)
166    parser.add_argument(
167        "--kwargs",
168        type=_json_object,
169        default={},
170        metavar="JSON",
171        help="additional RHIME options as a JSON object",
172    )
173    args = parser.parse_args(argv)
174
175    overrides = dict(args.kwargs)
176    for name in ("start_date", "end_date", "output_path", "output_name", "draws", "tune", "chains"):
177        value = getattr(args, name)
178        if value is not None:
179            overrides[name] = value
180    return run_custom_rhime(config_file=args.config_file, **overrides)
181
182
183if __name__ == "__main__":
184    main()

Compose a custom basis stage#

The second complete runner replaces one call in the preparation spine: build_project_basis replaces build_rhime_basis. The project function composes public basis primitives instead of selecting a built-in basis algorithm. It:

  • derives and normalises a two-dimensional weight field with basis_weights_from_fp_all;

  • loads the public country grid and reduces positive country codes to land and the remaining cells to ocean;

  • creates physical north-south/east-west coordinates with LatLonGridGeometry;

  • uses balanced inertial splits, decomposes every proposed child into four-neighbour connected components, and rejects splits whose children exceed the configured PCA eccentricity guard;

  • generates class-safe labels with region_constrained_basis; and

  • wraps the flat labels and current run flux in retained BasisFunctions with basis_functions_from_fp_all_flat_basis.

The nested split strategy follows the latest selected-country guarded-basis variant in the verification-games project. That variant gives the UK, Ireland, France, Germany, Italy, Belgium, and the Netherlands separate classes, while remaining land and ocean form two more. This smaller example deliberately uses the land/ocean class variant so the composition stays readable. To adopt the selected-country policy, replace _land_ocean_classes with a project function that maps the loaded integer country codes or country names to those classes. Keep this classification outside OpenGHG Inversions unless it becomes a broadly supported policy.

The weighting is deliberately not identical to that later verification-games preparation step. Verification-games sums absolute cached fp_x_flux after those sensitivities exist. At this earlier visible-runner basis boundary they have not been constructed yet, so the example uses the public basis_weights_from_fp_all field while preserving the guarded split strategy and class composition.

The flat labels and retained object record namespaced provenance for the class policy, normalised weight source, connected balanced-inertial strategy, connectivity, and eccentricity threshold. Those fields travel with a saved BasisFunctions artifact and make the project choice inspectable later.

Acquisition, filtering, sensitivity construction, labelled input assembly, the standard likelihood and model, conversion of delayed arrays for PyMC, sampling, predictive selection, filenames, and output handling remain library-owned.

The project owns the scientific validity of its classification, coverage, region count, eccentricity threshold, and split policy. BasisFunctions and the unchanged downstream stages validate the grid, coordinates, sources, site alignment, state layout, and model inputs they consume; they cannot certify that the project’s scientific partition is appropriate. Malformed artifacts fail while loading; structurally valid artifacts fail later if they violate a downstream alignment contract.

The example exposes project_basis_path separately from standard RHIME options. Point it at a self-contained .nc or .zarr artifact written by BasisFunctions.save to bypass fitting:

python -m package_name.custom_basis_runner config.ini \
    --project-basis-path cache/project-basis.zarr

The eccentricity guard is also a visible project-owned option, rather than an opaque library default. It defaults to 10 and is removed before standard RHIME option resolution:

python -m package_name.custom_basis_runner config.ini \
    --max-child-pca-eccentricity 10

BasisFunctions.load loads the saved operator, metadata, and flux into memory, then closes the artifact. The saved flux is deliberately retained. Use the standard public load_basis_functions helper instead when loading a named RHIME basis cache that should take its retained flux from the current fp_all acquisition.

Without an artifact, the basis-building function computes the arrays needed by the splitting algorithm. Before that point, the filtered xarray objects may still defer their calculations with Dask. The function derives and normalises the two-dimensional weights, loads the country classes, constructs the geometry, and creates the region labels. It then combines the labels with the current run’s flux in a BasisFunctions object. That object flows unchanged through sensitivity construction and labelled assembly. The concrete recipe’s adjacent standard_model_input_names declaration selects the arrays it needs. materialize_pymc_inputs converts only those arrays, together, immediately before model construction; project-owned lazy extensions in RhimePreparedInputs remain untouched.

The customisation is concentrated in _guarded_basis. The runner’s one deliberate substitution is marked by an inline comment where build_project_basis replaces the standard build_rhime_basis call. The complete source remains below because the test suite executes the same file that the documentation displays.

Copy examples/rhime_customisation/custom_basis_runner.py to src/<package_name>/custom_basis_runner.py and keep the project basis rule beside the copied orchestration spine. This source is imported and executed by integration tests and rendered here, so the documentation and runnable example cannot drift apart.

This guarded composition is intentionally project-specific and is not yet a common, stable strategy that warrants another lower-ceremony run_rhime option.

  1"""Run standard RHIME with project-owned emissions basis construction.
  2
  3This copy-and-modify example keeps RHIME's supported acquisition, filtering,
  4sensitivity, labelled-input, model, sampling, and output stages.  Only the
  5emissions-basis stage is replaced by :func:`build_project_basis`.  The example
  6either loads a self-contained retained ``BasisFunctions`` artifact or creates
  7a guarded region-constrained basis whose regions stay within land/ocean
  8classes, remain connected, and satisfy a project eccentricity threshold.
  9The verification-games workflow weights that strategy with summed absolute
 10cached ``fp_x_flux``. At this earlier runner stage projected sensitivities do
 11not exist, so the example deliberately substitutes the public
 12``basis_weights_from_fp_all`` field while preserving the guarded strategy and
 13class composition.
 14
 15Use :func:`run_custom_rhime` from Python or :func:`main` from the command line.
 16Acquired xarray objects are treated as borrowed.  Generated-basis fitting is
 17an explicit eager algorithm boundary, artifact loading materializes the saved
 18artifact, and model arrays materialize later at the named PyMC boundary.
 19"""
 20
 21from __future__ import annotations
 22
 23import argparse
 24from collections.abc import Mapping, Sequence
 25import json
 26from pathlib import Path
 27from time import perf_counter
 28from typing import Any
 29
 30import numpy as np
 31import xarray as xr
 32
 33from openghg_inversions.basis import (
 34    BasisFunctions,
 35    basis_functions_from_fp_all_flat_basis,
 36    basis_weights_from_fp_all,
 37    load_country_region_classes,
 38)
 39from openghg_inversions.basis.algorithms import (
 40    ConnectedComponentPartitionStep,
 41    ConnectedComponentSplitStrategy,
 42    GreedySplitStrategy,
 43    InertialSplitStep,
 44    LatLonGridGeometry,
 45    MaxChildPCAEccentricity,
 46    region_constrained_basis,
 47)
 48from openghg_inversions.inversion_data import RhimeMergedData
 49from openghg_inversions.rhime import (
 50    RhimeResult,
 51    assemble_rhime_inputs,
 52    build_rhime_sensitivities,
 53    build_standard_rhime_model_result,
 54    filter_rhime_observations,
 55    make_standard_rhime_result,
 56    make_standard_rhime_outputs,
 57    materialize_pymc_inputs,
 58    params_from_config,
 59    resolve_rhime_options,
 60    retrieve_or_reload_rhime_data,
 61    sample_rhime_model,
 62    standard_model_input_names,
 63    with_prepared_rhime_sites,
 64)
 65
 66
 67DEFAULT_MAX_CHILD_PCA_ECCENTRICITY = 10.0
 68
 69
 70def _land_ocean_classes(
 71    *,
 72    domain: str,
 73    country_directory: str | Path | None = None,
 74) -> xr.DataArray:
 75    """Reduce the public country grid to the project's land/ocean classes.
 76
 77    Args:
 78        domain: RHIME domain used to select the country grid.
 79        country_directory: Optional directory containing the domain country
 80            file.
 81
 82    Returns:
 83        An eager object-valued ``(lat, lon)`` class grid preserving the loaded
 84        coordinates. Positive country codes are ``land``; null, zero, and
 85        negative codes are ``ocean``.
 86
 87    Raises:
 88        FileNotFoundError: If the requested country grid does not exist.
 89        KeyError: If the country-grid artifact does not contain ``country``.
 90    """
 91    country_classes = load_country_region_classes(
 92        domain,
 93        country_directory,
 94    )
 95    return xr.where(country_classes > 0, "land", "ocean").astype(object).rename("basis_class")
 96
 97
 98def _guarded_basis(
 99    merged: RhimeMergedData,
100    data_args: Mapping[str, Any],
101    *,
102    max_child_pca_eccentricity: float,
103) -> BasisFunctions:
104    """Build the project's connected, eccentricity-guarded retained basis.
105
106    The public weight adapter intentionally materializes a derived 2-D weight
107    field from borrowed footprint and flux inputs here. The project then owns
108    normalization, class composition, guarded splitting, and conversion of the
109    flat labels to retained ``BasisFunctions`` without mutating ``merged``.
110
111    Args:
112        merged: Filtered, borrowed RHIME observations and flux data.
113        data_args: Resolved standard RHIME data and basis options. This stage
114            consumes required ``domain`` and optional ``flux_sources``,
115            ``nbasis``, and ``country_directory``.
116        max_child_pca_eccentricity: Project limit on each proposed child's
117            physical-coordinate PCA eccentricity.
118
119    Returns:
120        Retained basis functions fitted within connected land/ocean regions.
121
122    Raises:
123        FileNotFoundError: If the requested country grid does not exist.
124        KeyError: If ``domain`` or country-grid content is missing.
125        ValueError: If zero-filled weights do not have a positive finite
126            maximum, or geometry, allocation, or guarded splitting is invalid.
127    """
128    # 1. Turn the filtered footprints and flux into one spatial importance map.
129    weights = basis_weights_from_fp_all(
130        merged.fp_all,
131        data_args.get("flux_sources"),
132        abs_flux=True,
133    )
134    finite_weights = weights.fillna(0.0).astype(np.float64)
135    maximum = float(finite_weights.max())
136    if not np.isfinite(maximum) or maximum <= 0.0:
137        raise ValueError("Project basis weights must have a positive finite maximum.")
138    normalized_weights = finite_weights / maximum
139
140    # 2. Define scientific boundaries that no basis region may cross.
141    region_classes = _land_ocean_classes(
142        domain=data_args["domain"],
143        country_directory=data_args.get("country_directory"),
144    )
145    geometry = LatLonGridGeometry.from_dataarray(normalized_weights)
146    # 3. Compose the split algorithm and its shape/connectivity safeguards.
147    strategy = ConnectedComponentSplitStrategy(
148        split_strategy=GreedySplitStrategy(
149            split_step=ConnectedComponentPartitionStep(
150                split_step=InertialSplitStep(
151                    balanced=True,
152                    geometry=geometry,
153                ),
154                connectivity=1,
155            ),
156            split_acceptance=MaxChildPCAEccentricity(
157                max_child_pca_eccentricity=max_child_pca_eccentricity,
158                geometry=geometry,
159            ),
160        ),
161        connectivity=1,
162    )
163    # 4. Allocate and split regions within the land/ocean classes.
164    basis_flat = (
165        region_constrained_basis(
166            normalized_weights,
167            region_classes,
168            int(data_args.get("nbasis", 100)),
169            allocation="weight",
170            min_regions_per_class=1,
171            split_strategy=strategy,
172        )
173        .astype(np.int16)
174        .rename("basis")
175    )
176    provenance: dict[str, str | int | float] = {
177        "openghg_inversions:basis_artifact_source": "project-guarded",
178        "openghg_inversions:project_basis_strategy": "connected_component_balanced_inertial",
179        "openghg_inversions:project_basis_connectivity": 1,
180        "openghg_inversions:project_basis_max_child_pca_eccentricity": float(max_child_pca_eccentricity),
181        "openghg_inversions:project_basis_class_policy": "land_ocean",
182        "openghg_inversions:project_basis_weights": "basis_weights_from_fp_all_abs_flux_normalized",
183    }
184    basis_flat.attrs.update(provenance)
185    # 5. Attach the current flux so standard RHIME sensitivity code can use it.
186    return basis_functions_from_fp_all_flat_basis(
187        fp_all=merged.fp_all,
188        basis_flat=basis_flat,
189        metadata=provenance,
190    )
191
192
193def build_project_basis(
194    merged: RhimeMergedData,
195    data_args: Mapping[str, Any],
196    *,
197    project_basis_path: str | Path | None = None,
198    max_child_pca_eccentricity: float = DEFAULT_MAX_CHILD_PCA_ECCENTRICITY,
199) -> BasisFunctions:
200    """Return the project-selected retained emissions basis.
201
202    Args:
203        merged: Filtered, borrowed RHIME observations and flux data.
204        data_args: Resolved standard RHIME data and basis options.
205        project_basis_path: Optional ``.nc`` or ``.zarr`` artifact previously
206            written by :meth:`BasisFunctions.save`.  The artifact is
207            self-contained, so its serialized operator, metadata, and flux are
208            retained rather than replaced with flux from ``merged``.
209        max_child_pca_eccentricity: Project limit passed to the guarded split
210            policy when generating a basis. Defaults to ``10`` and is ignored
211            when ``project_basis_path`` is supplied.
212
213    Returns:
214        Loaded or newly fitted retained basis functions.
215
216    Raises:
217        OSError: If the requested artifact cannot be opened.
218        KeyError: If required artifact content or generated-basis options are
219            missing.
220        ValueError: If the artifact or generated-basis configuration is
221            invalid.
222
223    Notes:
224        Loading eagerly materializes the artifact so no open file handle is
225        retained.  Without a path, project basis fitting is the named eager
226        basis-generation boundary.
227    """
228    if project_basis_path is not None:
229        return BasisFunctions.load(project_basis_path)
230    return _guarded_basis(
231        merged,
232        data_args,
233        max_child_pca_eccentricity=max_child_pca_eccentricity,
234    )
235
236
237def run_custom_rhime(
238    *,
239    config_file: str | Path | None = None,
240    project_basis_path: str | Path | None = None,
241    max_child_pca_eccentricity: float | None = None,
242    **kwargs: Any,
243) -> RhimeResult:
244    """Run standard single-sector RHIME with a project-owned basis stage.
245
246    Args:
247        config_file: Optional RHIME INI configuration file.
248        project_basis_path: Optional self-contained ``BasisFunctions`` artifact
249            to use instead of fitting the project guarded basis.
250        max_child_pca_eccentricity: Optional project split-policy threshold.
251            When omitted, the merged config/keyword value is used, falling back
252            to ``10``. This option is removed before standard RHIME resolution.
253        **kwargs: Standard RHIME option names that override values from
254            ``config_file``.
255
256    Returns:
257        The sampled result and any outputs requested by the RHIME options.
258
259    Raises:
260        OSError: If configured artifact, input, or output I/O fails.
261        ValueError: If RHIME options, the custom basis, or prepared inputs are
262            invalid.
263
264    Notes:
265        This workflow may retrieve or reload data, may eagerly fit or load a
266        basis, eagerly materializes PyMC model inputs, runs sampling, and writes
267        outputs requested by the resolved RHIME options.
268    """
269    params = (
270        params_from_config(config_file, extra_kwargs=kwargs, normalise=False)
271        if config_file is not None
272        else dict(kwargs)
273    )
274    configured_project_basis_path = params.pop("project_basis_path", None)
275    if project_basis_path is None:
276        project_basis_path = configured_project_basis_path
277    configured_eccentricity = params.pop(
278        "max_child_pca_eccentricity",
279        DEFAULT_MAX_CHILD_PCA_ECCENTRICITY,
280    )
281    if max_child_pca_eccentricity is None:
282        max_child_pca_eccentricity = float(configured_eccentricity)
283    setup = resolve_rhime_options(params=params, multisector=False)
284
285    merged = retrieve_or_reload_rhime_data(setup.data_args, multisector=False)
286    filtered = filter_rhime_observations(merged, setup.data_args)
287
288    # CUSTOMISATION POINT: the standard runner calls build_rhime_basis here.
289    # This project function either loads a saved basis or composes the guarded
290    # basis-building tools above. Every stage after this call is standard RHIME.
291    basis_functions = build_project_basis(
292        filtered,
293        dict(setup.data_args),
294        project_basis_path=project_basis_path,
295        max_child_pca_eccentricity=max_child_pca_eccentricity,
296    )
297    site_data = build_rhime_sensitivities(
298        filtered,
299        basis_functions,
300        setup.data_args,
301        multisector=False,
302    )
303    prepared = assemble_rhime_inputs(filtered, basis_functions, site_data, setup.data_args)
304    run_spec = with_prepared_rhime_sites(setup.run_spec, prepared)
305
306    # Cross the explicit eager PyMC boundary without changing canonical inputs.
307    model_inputs = materialize_pymc_inputs(
308        prepared,
309        variable_names=standard_model_input_names(prepared, run_spec.model),
310    )
311    build_and_sample_start = perf_counter()
312    model_build_result = build_standard_rhime_model_result(
313        prepared=prepared,
314        model_inputs=model_inputs,
315        run_spec=run_spec,
316    )
317    idata = sample_rhime_model(model_build_result, setup.sampler)
318
319    result = make_standard_rhime_result(
320        prepared=prepared,
321        run_spec=run_spec,
322        sampler=setup.sampler,
323        model_build_result=model_build_result,
324        idata=idata,
325        build_and_sample_seconds=perf_counter() - build_and_sample_start,
326    )
327    make_standard_rhime_outputs(result=result, prepared=prepared)
328    return result
329
330
331def _json_object(value: str) -> dict[str, Any]:
332    """Parse one command-line JSON object containing additional RHIME options.
333
334    Args:
335        value: JSON text to parse.
336
337    Returns:
338        The decoded JSON object.
339
340    Raises:
341        argparse.ArgumentTypeError: If the text is invalid JSON or decodes to
342            a non-object value.
343    """
344    try:
345        parsed = json.loads(value)
346    except json.JSONDecodeError as exc:
347        raise argparse.ArgumentTypeError(f"invalid JSON: {exc.msg}") from exc
348    if not isinstance(parsed, dict):
349        raise argparse.ArgumentTypeError("--kwargs must decode to a JSON object")
350    return parsed
351
352
353def main(argv: Sequence[str] | None = None) -> RhimeResult:
354    """Parse command-line options and run the custom-basis RHIME workflow.
355
356    Args:
357        argv: Command-line tokens excluding the program name. ``None`` reads
358            tokens from :data:`sys.argv`.
359
360    Returns:
361        The RHIME result returned by :func:`run_custom_rhime`.
362
363    Raises:
364        SystemExit: If argument parsing fails or help is requested.
365
366    Notes:
367        Data access, eager basis and model materialization, sampling, and
368        requested output writes are delegated to :func:`run_custom_rhime`; its
369        exceptions propagate unchanged.
370    """
371    parser = argparse.ArgumentParser(description=__doc__)
372    parser.add_argument("config_file", nargs="?", type=Path, help="RHIME INI configuration file")
373    parser.add_argument(
374        "--project-basis-path",
375        type=Path,
376        help="self-contained BasisFunctions .nc or .zarr artifact",
377    )
378    parser.add_argument(
379        "--max-child-pca-eccentricity",
380        type=float,
381        help="project guarded-basis threshold (default: 10)",
382    )
383    parser.add_argument("--start-date")
384    parser.add_argument("--end-date")
385    parser.add_argument("--output-path", type=Path)
386    parser.add_argument("--output-name")
387    parser.add_argument("--draws", type=int)
388    parser.add_argument("--tune", type=int)
389    parser.add_argument("--chains", type=int)
390    parser.add_argument(
391        "--kwargs",
392        type=_json_object,
393        default={},
394        metavar="JSON",
395        help="additional standard RHIME options as a JSON object",
396    )
397    args = parser.parse_args(argv)
398
399    overrides = dict(args.kwargs)
400    for name in ("start_date", "end_date", "output_path", "output_name", "draws", "tune", "chains"):
401        value = getattr(args, name)
402        if value is not None:
403            overrides[name] = value
404    return run_custom_rhime(
405        config_file=args.config_file,
406        project_basis_path=args.project_basis_path,
407        max_child_pca_eccentricity=args.max_child_pca_eccentricity,
408        **overrides,
409    )
410
411
412if __name__ == "__main__":
413    main()