Skip to content

ML Core API

ml_core

Unified ML model interface and shared utilities for substation forecasting.

Why this package exists

To enable rapid experimentation with different model architectures (XGBoost, GNNs, etc.), we need a common interface that shields the model developer from the underlying data engineering.

ml_core.features

Feature engineering package.

Public interface:

  • FeatureEngineer — abstract base; implement to swap the full feature pipeline.
  • TabularFeatureEngineer — default implementation: nearest-cell NWP spatial join followed by the declarative tabular pipeline.

Private sub-modules (importable by tests, not part of the public API):

  • _parsed_features — typed feature descriptors and ParsedFeatures parser.
  • _nwp — NWP upsampling, processing, and power/NWP join helpers.
  • _lags — power-lag, weather-lag (dual-strategy), and leaky-lag nullification.
  • tabular_feature_engineer_engineer_features orchestrator and tabular pipeline helpers.

Attributes

__all__ = ['FeatureEngineer', 'TabularFeatureEngineer'] module-attribute

Classes

FeatureEngineer

Bases: ABC

Turns raw power + NWP + metadata into a model-ready AllFeatures frame.

Source code in packages/ml_core/src/ml_core/features/feature_engineer.py
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
class FeatureEngineer(ABC):
    """Turns raw power + NWP + metadata into a model-ready ``AllFeatures`` frame."""

    @abstractmethod
    def engineer(
        self,
        *,
        selected_features: set[str],
        power_time_series: pt.LazyFrame[PowerTimeSeries],
        time_series_metadata: pt.DataFrame[TimeSeriesMetadata],
        nwp: pt.LazyFrame[Nwp],
        power_fcst_init_time: datetime | None = None,
        nwp_init_time: datetime | None = None,
        nwp_publication_delay_hours: int = NWP_PUBLICATION_DELAY_HOURS,
    ) -> pt.LazyFrame[AllFeatures]:
        """Engineer features for training/inference.

        Two operating modes, matching ``tabular_feature_engineer._engineer_features`` (see
        that function for the full detail this docstring summarises):

        - **Bulk mode** (``power_fcst_init_time=None``, the default): NWP-centric, vectorised
          over every NWP run in the input, one forecast per ``(nwp_init_time, valid_time)``
          pair. Used for training and multi-run backtesting.
        - **Single-run mode** (``power_fcst_init_time`` given): power-centric, joins exactly
          one NWP run (``nwp_init_time``, or derived from ``power_fcst_init_time`` minus
          ``nwp_publication_delay_hours`` if omitted) and stamps a constant
          ``power_fcst_init_time`` on every row. Used for production inference and replay
          backfilling.

        Args:
            selected_features: The feature names to produce.
            power_time_series: Observed power, one row per ``(time_series_id, time)``.
            time_series_metadata: Per-time-series metadata (carries ``h3_res_5`` for the
                spatial mapping).
            nwp: Gridded NWP in physical units, keyed by ``h3_index``.
            power_fcst_init_time: ``None`` for bulk mode; a single datetime for single-run
                mode. See ``_engineer_features`` for the full contract.
            nwp_init_time: The NWP run to join in single-run mode. Must be ``None`` in bulk
                mode. See ``_engineer_features``.
            nwp_publication_delay_hours: Delay used to derive whichever of
                ``power_fcst_init_time``/``nwp_init_time`` is not supplied.

        Returns:
            A lazy ``AllFeatures`` frame, ready to hand to ``BaseForecaster.train``/``predict``.
        """
Methods:
engineer(*, selected_features, power_time_series, time_series_metadata, nwp, power_fcst_init_time=None, nwp_init_time=None, nwp_publication_delay_hours=NWP_PUBLICATION_DELAY_HOURS) abstractmethod

Engineer features for training/inference.

Two operating modes, matching tabular_feature_engineer._engineer_features (see that function for the full detail this docstring summarises):

  • Bulk mode (power_fcst_init_time=None, the default): NWP-centric, vectorised over every NWP run in the input, one forecast per (nwp_init_time, valid_time) pair. Used for training and multi-run backtesting.
  • Single-run mode (power_fcst_init_time given): power-centric, joins exactly one NWP run (nwp_init_time, or derived from power_fcst_init_time minus nwp_publication_delay_hours if omitted) and stamps a constant power_fcst_init_time on every row. Used for production inference and replay backfilling.

Parameters:

Name Type Description Default
selected_features set[str]

The feature names to produce.

required
power_time_series LazyFrame[PowerTimeSeries]

Observed power, one row per (time_series_id, time).

required
time_series_metadata DataFrame[TimeSeriesMetadata]

Per-time-series metadata (carries h3_res_5 for the spatial mapping).

required
nwp LazyFrame[Nwp]

Gridded NWP in physical units, keyed by h3_index.

required
power_fcst_init_time datetime | None

None for bulk mode; a single datetime for single-run mode. See _engineer_features for the full contract.

None
nwp_init_time datetime | None

The NWP run to join in single-run mode. Must be None in bulk mode. See _engineer_features.

None
nwp_publication_delay_hours int

Delay used to derive whichever of power_fcst_init_time/nwp_init_time is not supplied.

NWP_PUBLICATION_DELAY_HOURS

Returns:

Type Description
LazyFrame[AllFeatures]

A lazy AllFeatures frame, ready to hand to BaseForecaster.train/predict.

Source code in packages/ml_core/src/ml_core/features/feature_engineer.py
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
@abstractmethod
def engineer(
    self,
    *,
    selected_features: set[str],
    power_time_series: pt.LazyFrame[PowerTimeSeries],
    time_series_metadata: pt.DataFrame[TimeSeriesMetadata],
    nwp: pt.LazyFrame[Nwp],
    power_fcst_init_time: datetime | None = None,
    nwp_init_time: datetime | None = None,
    nwp_publication_delay_hours: int = NWP_PUBLICATION_DELAY_HOURS,
) -> pt.LazyFrame[AllFeatures]:
    """Engineer features for training/inference.

    Two operating modes, matching ``tabular_feature_engineer._engineer_features`` (see
    that function for the full detail this docstring summarises):

    - **Bulk mode** (``power_fcst_init_time=None``, the default): NWP-centric, vectorised
      over every NWP run in the input, one forecast per ``(nwp_init_time, valid_time)``
      pair. Used for training and multi-run backtesting.
    - **Single-run mode** (``power_fcst_init_time`` given): power-centric, joins exactly
      one NWP run (``nwp_init_time``, or derived from ``power_fcst_init_time`` minus
      ``nwp_publication_delay_hours`` if omitted) and stamps a constant
      ``power_fcst_init_time`` on every row. Used for production inference and replay
      backfilling.

    Args:
        selected_features: The feature names to produce.
        power_time_series: Observed power, one row per ``(time_series_id, time)``.
        time_series_metadata: Per-time-series metadata (carries ``h3_res_5`` for the
            spatial mapping).
        nwp: Gridded NWP in physical units, keyed by ``h3_index``.
        power_fcst_init_time: ``None`` for bulk mode; a single datetime for single-run
            mode. See ``_engineer_features`` for the full contract.
        nwp_init_time: The NWP run to join in single-run mode. Must be ``None`` in bulk
            mode. See ``_engineer_features``.
        nwp_publication_delay_hours: Delay used to derive whichever of
            ``power_fcst_init_time``/``nwp_init_time`` is not supplied.

    Returns:
        A lazy ``AllFeatures`` frame, ready to hand to ``BaseForecaster.train``/``predict``.
    """

TabularFeatureEngineer

Bases: FeatureEngineer

Nearest res-5 NWP cell per time series, then the declarative tabular pipeline.

Source code in packages/ml_core/src/ml_core/features/tabular_feature_engineer.py
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
class TabularFeatureEngineer(FeatureEngineer):
    """Nearest res-5 NWP cell per time series, then the declarative tabular pipeline."""

    def engineer(
        self,
        *,
        selected_features: set[str],
        power_time_series: pt.LazyFrame[PowerTimeSeries],
        time_series_metadata: pt.DataFrame[TimeSeriesMetadata],
        nwp: pt.LazyFrame[Nwp],
        power_fcst_init_time: datetime | None = None,
        nwp_init_time: datetime | None = None,
        nwp_publication_delay_hours: int = NWP_PUBLICATION_DELAY_HOURS,
    ) -> pt.LazyFrame[AllFeatures]:
        nwp_per_time_series = _attach_nearest_nwp_cell(nwp, time_series_metadata)
        return _engineer_features(
            selected_features,
            power_time_series,
            time_series_metadata,
            nwp=nwp_per_time_series,
            power_fcst_init_time=power_fcst_init_time,
            nwp_init_time=nwp_init_time,
            nwp_publication_delay_hours=nwp_publication_delay_hours,
        )
Methods:
engineer(*, selected_features, power_time_series, time_series_metadata, nwp, power_fcst_init_time=None, nwp_init_time=None, nwp_publication_delay_hours=NWP_PUBLICATION_DELAY_HOURS)
Source code in packages/ml_core/src/ml_core/features/tabular_feature_engineer.py
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
def engineer(
    self,
    *,
    selected_features: set[str],
    power_time_series: pt.LazyFrame[PowerTimeSeries],
    time_series_metadata: pt.DataFrame[TimeSeriesMetadata],
    nwp: pt.LazyFrame[Nwp],
    power_fcst_init_time: datetime | None = None,
    nwp_init_time: datetime | None = None,
    nwp_publication_delay_hours: int = NWP_PUBLICATION_DELAY_HOURS,
) -> pt.LazyFrame[AllFeatures]:
    nwp_per_time_series = _attach_nearest_nwp_cell(nwp, time_series_metadata)
    return _engineer_features(
        selected_features,
        power_time_series,
        time_series_metadata,
        nwp=nwp_per_time_series,
        power_fcst_init_time=power_fcst_init_time,
        nwp_init_time=nwp_init_time,
        nwp_publication_delay_hours=nwp_publication_delay_hours,
    )

ml_core.base_forecaster

Classes

BaseForecasterConfig

Bases: BaseModel

Universal configuration for all forecasting models.

Subclasses add model-specific hyperparameters. Having a shared base ensures that every forecaster carries its own feature list and optional MLflow experiment id in one serialisable object, simplifying save/load and Hydra config wiring.

The tag fields (weather_source, training_strategy) are stamped onto MLflow runs for leaderboard grouping. They live here so that hydra.utils.instantiate(model_cfg.model_params) validates them at load time — they are present in every conf/model/*.yaml file under model_params.

experiment_name is the per-experiment key (set to the MLflow experiment name at registration and stored in the saved config); it is stamped onto every PowerForecast row, distinct from the model-family MODEL_NAME. random_seed is threaded into each model's training so that re-training a fold reproduces the same model — keeping retries and the leaderboard stable.

Model identity (name and version) lives on the BaseForecaster class itself as MODEL_NAME and MODEL_VERSION — those are properties of the implementation, not the experiment config.

Source code in packages/ml_core/src/ml_core/base_forecaster.py
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
class BaseForecasterConfig(BaseModel):
    """Universal configuration for all forecasting models.

    Subclasses add model-specific hyperparameters. Having a shared base ensures that every
    forecaster carries its own feature list and optional MLflow experiment id in one
    serialisable object, simplifying save/load and Hydra config wiring.

    The tag fields (weather_source, training_strategy) are stamped onto MLflow runs for
    leaderboard grouping. They live here so that
    ``hydra.utils.instantiate(model_cfg.model_params)`` validates them at load time — they
    are present in every ``conf/model/*.yaml`` file under ``model_params``.

    ``experiment_name`` is the per-experiment key (set to the MLflow experiment name at
    registration and stored in the saved config); it is stamped onto every ``PowerForecast``
    row, distinct from the model-family ``MODEL_NAME``. ``random_seed`` is threaded into each
    model's training so that re-training a fold reproduces the same model — keeping retries
    and the leaderboard stable.

    Model identity (name and version) lives on the ``BaseForecaster`` class itself as
    ``MODEL_NAME`` and ``MODEL_VERSION`` — those are properties of the implementation, not
    the experiment config.
    """

    selected_features: set[str]
    ml_flow_experiment_id: int | None = None
    experiment_name: str = ""
    weather_source: str = ""
    training_strategy: str = ""
    random_seed: int = 0
Attributes
selected_features instance-attribute
ml_flow_experiment_id = None class-attribute instance-attribute
experiment_name = '' class-attribute instance-attribute
weather_source = '' class-attribute instance-attribute
training_strategy = '' class-attribute instance-attribute
random_seed = 0 class-attribute instance-attribute

BaseForecaster

Bases: ABC

Defines the universal interface for all energy forecasting ML models.

Every forecasting model subclasses this abstract base to allow shared Dagster assets and evaluation code to remain completely agnostic to the underlying model implementation.

Subclasses must define MODEL_NAME and MODEL_VERSION as class-level constants. These are stamped onto every PowerForecast row at predict time and used as the MLflow experiment name. Bumping MODEL_VERSION requires a code change (intentional), not a config edit.

Lazy evaluation contract: train and predict both accept a pt.LazyFrame[AllFeatures]. Callers must not collect before passing data in; doing so wastes memory and prevents Polars from optimising the full query plan. A subclass materialises the data at the model boundary (typically a single .collect(), streamed). Keeping that bounded is the caller's responsibility: it prunes the inputs (NWP control member, the relevant H3 cells, the window's init_time partitions) and, where the full ensemble is needed, processes one init_time chunk at a time — filtering the engineered output cannot prune the upstream join/upsample. See the NWP scan-pruning notes in docs/architecture/overview.md.

Persistence has two layers. Subclasses implement save/load for their own on-disk format and need know nothing about MLflow. The concrete save_to_mlflow/load_from_mlflow methods, shared by all subclasses, wrap that disk format with MLflow's artifact store and a local-disk cache, so the same trained model can be shared across machines and served offline.

Source code in packages/ml_core/src/ml_core/base_forecaster.py
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
class BaseForecaster(ABC):
    """Defines the universal interface for all energy forecasting ML models.

    Every forecasting model subclasses this abstract base to allow shared Dagster assets and
    evaluation code to remain completely agnostic to the underlying model implementation.

    Subclasses must define ``MODEL_NAME`` and ``MODEL_VERSION`` as class-level constants.
    These are stamped onto every ``PowerForecast`` row at predict time and used as the MLflow
    experiment name. Bumping ``MODEL_VERSION`` requires a code change (intentional), not a
    config edit.

    Lazy evaluation contract: `train` and `predict` both accept a `pt.LazyFrame[AllFeatures]`.
    Callers must not collect before passing data in; doing so wastes memory and prevents Polars
    from optimising the full query plan. A subclass materialises the data at the model boundary
    (typically a single `.collect()`, streamed). Keeping that bounded is the *caller's*
    responsibility: it prunes the inputs (NWP control member, the relevant H3 cells, the window's
    `init_time` partitions) and, where the full ensemble is needed, processes one `init_time` chunk
    at a time — filtering the engineered output cannot prune the upstream join/upsample. See the NWP
    scan-pruning notes in `docs/architecture/overview.md`.

    Persistence has two layers. Subclasses implement ``save``/``load`` for their own on-disk
    format and need know nothing about MLflow. The concrete ``save_to_mlflow``/``load_from_mlflow``
    methods, shared by all subclasses, wrap that disk format with MLflow's artifact store and a
    local-disk cache, so the same trained model can be shared across machines and served offline.
    """

    MODEL_NAME: ClassVar[str]
    MODEL_VERSION: ClassVar[int]

    feature_engineer: ClassVar[FeatureEngineer] = TabularFeatureEngineer()
    """The feature pipeline this forecaster's data is engineered through.

    Associated by composition (the forecaster *references* a feature engineer rather than
    *implementing* feature engineering), so a forecaster can swap the whole pipeline by overriding
    this with a different ``FeatureEngineer``. The default produces the tabular ``AllFeatures``
    frame that ``train``/``predict`` consume.
    """

    def __init__(self, model_params: BaseForecasterConfig) -> None:
        self.model_params = model_params

    @property
    @abstractmethod
    def trained_time_series_ids(self) -> list[int]:
        """The sorted ``time_series_id`` population this model will serve a ``predict`` for.

        This is the model's own frozen record of who it trained on — *not* a statement about how
        the model is internally structured. A model may hold one sub-model per series, a single
        model spanning many series, or anything in between; the contract is only about which
        ``time_series_id``s it will score.

        **Why this is load-bearing.** The **train==predict population invariant**: the
        population a model is scored on must equal the population it was trained on, *even if* the
        live eligibility set has drifted since training (power coverage changes, so a series may
        newly qualify or drop out). Consumers (``cv_power_forecasts``, ``live_forecasts``) filter
        their inputs to this set, so a model scores exactly the population it learned — never
        whatever eligibility says *today*. That is what keeps the leaderboard apples-to-apples and
        stops production forecasting a series the model never saw.

        Subclasses persist and reconstruct this set through their own ``save``/``load`` (e.g. in
        ``meta.json``), so it survives a round-trip through MLflow.
        """
        ...

    @abstractmethod
    def save(self, path: Path) -> None:
        """Save the trained model state to a directory.

        The saved directory must contain a ``meta.json`` with a ``model_class`` field — the
        fully-qualified ``{module}.{qualname}`` of the concrete subclass (e.g.
        ``"xgboost_forecaster.forecaster.XGBoostForecaster"``) — so that production inference
        (``ml_core._production_helpers.load_forecaster_from_dir``) can reconstruct the correct
        class from a plain model directory with no other context (issue #221).
        """
        pass

    @classmethod
    @abstractmethod
    def load(cls, path: Path) -> Self:
        """Reconstruct a trained instance from a previously saved directory."""
        pass

    def save_to_mlflow(self, run_id: str) -> None:
        """Upload this trained model's artifacts to the given MLflow run.

        Writes the model to a temporary directory via ``save`` (the subclass's own format), then
        uploads that directory to the run's artifact store under ``model/``. The caller is
        responsible for setting the tracking URI (``mlflow.set_tracking_uri``) beforehand.

        Args:
            run_id: The MLflow run to attach the artifacts to.
        """
        with tempfile.TemporaryDirectory() as tmp_dir:
            self.save(Path(tmp_dir))
            with mlflow.start_run(run_id=run_id):
                mlflow.log_artifacts(tmp_dir, artifact_path=_MLFLOW_ARTIFACT_PATH)

    @classmethod
    def load_from_mlflow(cls, run_id: str, cache_base_path: Path) -> Self:
        """Load a trained model for ``run_id``, serving from the local cache when possible.

        On a cache hit (``{cache_base_path}/{run_id}/model`` already exists) the model is loaded
        straight from disk and MLflow is never contacted — this is what lets the live service keep
        serving during an MLflow outage. On a cache miss the artifacts are downloaded from the run
        into the cache, then loaded. The cache key is the immutable run ID, so a cached model never
        goes stale. The caller sets the tracking URI (``mlflow.set_tracking_uri``) beforehand.

        Args:
            run_id: The MLflow run the model was saved under.
            cache_base_path: Root of the local cache; the model lives at
                ``{cache_base_path}/{run_id}``.

        Returns:
            The reconstructed, trained forecaster.
        """
        run_cache_dir = cache_base_path / run_id
        model_dir = run_cache_dir / _MLFLOW_ARTIFACT_PATH
        if not model_dir.exists():
            run_cache_dir.mkdir(parents=True, exist_ok=True)
            mlflow.artifacts.download_artifacts(
                run_id=run_id,
                artifact_path=_MLFLOW_ARTIFACT_PATH,
                dst_path=str(run_cache_dir),
            )
        return cls.load(model_dir)

    @abstractmethod
    def train(self, data: pt.LazyFrame[AllFeatures], time_series_ids: list[int]) -> None:
        """Fit the model on ``data`` for the given ``time_series_id`` population.

        Args:
            data: The engineered features (lazy; the caller must not pre-collect).
            time_series_ids: The population the model must train on — the caller's eligible set. The
                model decides how to map it onto boosters: one booster per id (``XGBoostForecaster``),
                or one booster per group of ids (e.g. all solar sites) for a future model. It is the
                model's frozen record of who it trained on (the train==predict invariant), and bounds
                ``predict`` to that same population.
        """
        pass

    @abstractmethod
    def predict(
        self, data: pt.LazyFrame[AllFeatures], *, fold_id: str = "live"
    ) -> pt.DataFrame[PowerForecast]:
        """Return power forecasts for all rows in the given AllFeatures data.

        Args:
            data: The engineered features to forecast from.
            fold_id: The value stamped onto every row's ``fold_id`` column. The model has no
                inherent notion of which CV fold it is serving (``fold_id`` is orchestration
                context), so the caller supplies it: ``cv_power_forecasts`` passes the fold's
                label, while production inference keeps the ``"live"`` default.
        """
        pass
Attributes
MODEL_NAME class-attribute
MODEL_VERSION class-attribute
feature_engineer = TabularFeatureEngineer() class-attribute

The feature pipeline this forecaster's data is engineered through.

Associated by composition (the forecaster references a feature engineer rather than implementing feature engineering), so a forecaster can swap the whole pipeline by overriding this with a different FeatureEngineer. The default produces the tabular AllFeatures frame that train/predict consume.

model_params = model_params instance-attribute
trained_time_series_ids abstractmethod property

The sorted time_series_id population this model will serve a predict for.

This is the model's own frozen record of who it trained on — not a statement about how the model is internally structured. A model may hold one sub-model per series, a single model spanning many series, or anything in between; the contract is only about which time_series_ids it will score.

Why this is load-bearing. The train==predict population invariant: the population a model is scored on must equal the population it was trained on, even if the live eligibility set has drifted since training (power coverage changes, so a series may newly qualify or drop out). Consumers (cv_power_forecasts, live_forecasts) filter their inputs to this set, so a model scores exactly the population it learned — never whatever eligibility says today. That is what keeps the leaderboard apples-to-apples and stops production forecasting a series the model never saw.

Subclasses persist and reconstruct this set through their own save/load (e.g. in meta.json), so it survives a round-trip through MLflow.

Methods:
__init__(model_params)
Source code in packages/ml_core/src/ml_core/base_forecaster.py
87
88
def __init__(self, model_params: BaseForecasterConfig) -> None:
    self.model_params = model_params
save(path) abstractmethod

Save the trained model state to a directory.

The saved directory must contain a meta.json with a model_class field — the fully-qualified {module}.{qualname} of the concrete subclass (e.g. "xgboost_forecaster.forecaster.XGBoostForecaster") — so that production inference (ml_core._production_helpers.load_forecaster_from_dir) can reconstruct the correct class from a plain model directory with no other context (issue #221).

Source code in packages/ml_core/src/ml_core/base_forecaster.py
113
114
115
116
117
118
119
120
121
122
123
@abstractmethod
def save(self, path: Path) -> None:
    """Save the trained model state to a directory.

    The saved directory must contain a ``meta.json`` with a ``model_class`` field — the
    fully-qualified ``{module}.{qualname}`` of the concrete subclass (e.g.
    ``"xgboost_forecaster.forecaster.XGBoostForecaster"``) — so that production inference
    (``ml_core._production_helpers.load_forecaster_from_dir``) can reconstruct the correct
    class from a plain model directory with no other context (issue #221).
    """
    pass
load(path) abstractmethod classmethod

Reconstruct a trained instance from a previously saved directory.

Source code in packages/ml_core/src/ml_core/base_forecaster.py
125
126
127
128
129
@classmethod
@abstractmethod
def load(cls, path: Path) -> Self:
    """Reconstruct a trained instance from a previously saved directory."""
    pass
save_to_mlflow(run_id)

Upload this trained model's artifacts to the given MLflow run.

Writes the model to a temporary directory via save (the subclass's own format), then uploads that directory to the run's artifact store under model/. The caller is responsible for setting the tracking URI (mlflow.set_tracking_uri) beforehand.

Parameters:

Name Type Description Default
run_id str

The MLflow run to attach the artifacts to.

required
Source code in packages/ml_core/src/ml_core/base_forecaster.py
131
132
133
134
135
136
137
138
139
140
141
142
143
144
def save_to_mlflow(self, run_id: str) -> None:
    """Upload this trained model's artifacts to the given MLflow run.

    Writes the model to a temporary directory via ``save`` (the subclass's own format), then
    uploads that directory to the run's artifact store under ``model/``. The caller is
    responsible for setting the tracking URI (``mlflow.set_tracking_uri``) beforehand.

    Args:
        run_id: The MLflow run to attach the artifacts to.
    """
    with tempfile.TemporaryDirectory() as tmp_dir:
        self.save(Path(tmp_dir))
        with mlflow.start_run(run_id=run_id):
            mlflow.log_artifacts(tmp_dir, artifact_path=_MLFLOW_ARTIFACT_PATH)
load_from_mlflow(run_id, cache_base_path) classmethod

Load a trained model for run_id, serving from the local cache when possible.

On a cache hit ({cache_base_path}/{run_id}/model already exists) the model is loaded straight from disk and MLflow is never contacted — this is what lets the live service keep serving during an MLflow outage. On a cache miss the artifacts are downloaded from the run into the cache, then loaded. The cache key is the immutable run ID, so a cached model never goes stale. The caller sets the tracking URI (mlflow.set_tracking_uri) beforehand.

Parameters:

Name Type Description Default
run_id str

The MLflow run the model was saved under.

required
cache_base_path Path

Root of the local cache; the model lives at {cache_base_path}/{run_id}.

required

Returns:

Type Description
Self

The reconstructed, trained forecaster.

Source code in packages/ml_core/src/ml_core/base_forecaster.py
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
@classmethod
def load_from_mlflow(cls, run_id: str, cache_base_path: Path) -> Self:
    """Load a trained model for ``run_id``, serving from the local cache when possible.

    On a cache hit (``{cache_base_path}/{run_id}/model`` already exists) the model is loaded
    straight from disk and MLflow is never contacted — this is what lets the live service keep
    serving during an MLflow outage. On a cache miss the artifacts are downloaded from the run
    into the cache, then loaded. The cache key is the immutable run ID, so a cached model never
    goes stale. The caller sets the tracking URI (``mlflow.set_tracking_uri``) beforehand.

    Args:
        run_id: The MLflow run the model was saved under.
        cache_base_path: Root of the local cache; the model lives at
            ``{cache_base_path}/{run_id}``.

    Returns:
        The reconstructed, trained forecaster.
    """
    run_cache_dir = cache_base_path / run_id
    model_dir = run_cache_dir / _MLFLOW_ARTIFACT_PATH
    if not model_dir.exists():
        run_cache_dir.mkdir(parents=True, exist_ok=True)
        mlflow.artifacts.download_artifacts(
            run_id=run_id,
            artifact_path=_MLFLOW_ARTIFACT_PATH,
            dst_path=str(run_cache_dir),
        )
    return cls.load(model_dir)
train(data, time_series_ids) abstractmethod

Fit the model on data for the given time_series_id population.

Parameters:

Name Type Description Default
data LazyFrame[AllFeatures]

The engineered features (lazy; the caller must not pre-collect).

required
time_series_ids list[int]

The population the model must train on — the caller's eligible set. The model decides how to map it onto boosters: one booster per id (XGBoostForecaster), or one booster per group of ids (e.g. all solar sites) for a future model. It is the model's frozen record of who it trained on (the train==predict invariant), and bounds predict to that same population.

required
Source code in packages/ml_core/src/ml_core/base_forecaster.py
175
176
177
178
179
180
181
182
183
184
185
186
187
@abstractmethod
def train(self, data: pt.LazyFrame[AllFeatures], time_series_ids: list[int]) -> None:
    """Fit the model on ``data`` for the given ``time_series_id`` population.

    Args:
        data: The engineered features (lazy; the caller must not pre-collect).
        time_series_ids: The population the model must train on — the caller's eligible set. The
            model decides how to map it onto boosters: one booster per id (``XGBoostForecaster``),
            or one booster per group of ids (e.g. all solar sites) for a future model. It is the
            model's frozen record of who it trained on (the train==predict invariant), and bounds
            ``predict`` to that same population.
    """
    pass
predict(data, *, fold_id='live') abstractmethod

Return power forecasts for all rows in the given AllFeatures data.

Parameters:

Name Type Description Default
data LazyFrame[AllFeatures]

The engineered features to forecast from.

required
fold_id str

The value stamped onto every row's fold_id column. The model has no inherent notion of which CV fold it is serving (fold_id is orchestration context), so the caller supplies it: cv_power_forecasts passes the fold's label, while production inference keeps the "live" default.

'live'
Source code in packages/ml_core/src/ml_core/base_forecaster.py
189
190
191
192
193
194
195
196
197
198
199
200
201
202
@abstractmethod
def predict(
    self, data: pt.LazyFrame[AllFeatures], *, fold_id: str = "live"
) -> pt.DataFrame[PowerForecast]:
    """Return power forecasts for all rows in the given AllFeatures data.

    Args:
        data: The engineered features to forecast from.
        fold_id: The value stamped onto every row's ``fold_id`` column. The model has no
            inherent notion of which CV fold it is serving (``fold_id`` is orchestration
            context), so the caller supplies it: ``cv_power_forecasts`` passes the fold's
            label, while production inference keeps the ``"live"`` default.
    """
    pass

ml_core.metrics

Metric computation for cross-validation results.

Attributes

Classes

NoOverlappingActualsError

Bases: ValueError

Raised by compute_metrics when no forecast row joins to any observed actual.

A distinct subclass so callers that score a fold in per-series batches can treat "this batch's series have no overlapping actuals" as skippable — mirroring how such series silently vanish from the inner join when the whole fold is scored in one call — while every other ValueError (negative lead times, missing capacity) still propagates.

Source code in packages/ml_core/src/ml_core/metrics.py
27
28
29
30
31
32
33
34
class NoOverlappingActualsError(ValueError):
    """Raised by ``compute_metrics`` when no forecast row joins to any observed actual.

    A distinct subclass so callers that score a fold in per-series batches can treat "this
    batch's series have no overlapping actuals" as skippable — mirroring how such series
    silently vanish from the inner join when the whole fold is scored in one call — while
    every other ``ValueError`` (negative lead times, missing capacity) still propagates.
    """

Functions:

compute_metrics(cv_forecasts, actuals, metadata, capacity)

Compute evaluation metrics from CV predictions and observed power.

Full metric definitions, equations, and design rationale: https://openclimatefix.github.io/nged-substation-forecast/techniques/evaluation-metrics/

For each (time_series_id, fold_id, power_fcst_model_name) group:

  1. Joins predictions to observed power on (time_series_id, valid_time).
  2. Assigns each row a horizon_slice from its lead time (valid_time − power_fcst_init_time) — see _horizon_slice_expr for the bands.
  3. Collapses the ensemble members within each forecast run (per power_fcst_init_time) into per-timestamp quantities: the deterministic ensemble mean, the fair CRPS, the Fortin-corrected ensemble variance, and the empirical DELIVERY_QUANTILES. Each run covering a valid_time is scored independently, exactly as a production consumer would experience it — runs at different lead times are never pooled.
  4. Aggregates per horizon_slice, plus the "all" aggregate over every lead time: MAE, NMAE, RMSE, and MBE on the ensemble mean; CRPS and the spread-skill ratio from the member-aware quantities; pinball loss at each delivery quantile (plus their mean); and PICP and interval width for each symmetric quantile band.
  5. Joins time_series_type from metadata onto each row.
  6. Returns one row per (time_series_id, fold_id, power_fcst_model_name, horizon_slice, metric_name, metric_param) in the tall Metrics format.

NMAE is normalised by the pre-computed full-history effective_capacity_mw (joined per time_series_id from capacity) — a capacity-like denominator, computed over the full history so it stays stable across folds. This time_series_id-only join is correct while capacity is one scalar row per series (the v0.1 shape). When the v0.7 upgrade makes capacity time-varying (one row per (time_series_id, time)), this join must become a temporal as-of join on (time_series_id, valid_time).

Single-member "ensembles" (e.g. a deterministic baseline forecaster) are scored unconditionally: their fair CRPS equals their MAE, their spread-skill ratio is 0, and their quantile bands are degenerate (all quantiles coincide, so PICP ≈ 0 and interval width = 0). Those are honest descriptions of a deterministic forecast, not errors.

Parameters:

Name Type Description Default
cv_forecasts DataFrame[PowerForecast]

CV predictions to evaluate. Currently only CV fold rows are handled; fold_id="live" support will be added with the production_monitoring scope in Phase 8.

required
actuals LazyFrame[PowerTimeSeries]

Observed half-hourly power (lazy — only the joined subset is collected). Deduplicated on (time_series_id, time) before the join — a duplicated actual would otherwise double the ensemble members and corrupt the member-aware metrics.

required
metadata DataFrame[TimeSeriesMetadata]

Substation metadata used to join time_series_type onto each metric row. Series absent from metadata receive a null time_series_type.

required
capacity DataFrame[EffectiveCapacity]

Pre-computed per-series effective capacity; effective_capacity_mw is the NMAE denominator. Must cover every scored time_series_id.

required

Returns:

Type Description
DataFrame[Metrics]

A validated tall Metrics DataFrame with time_series_type populated.

Raises:

Type Description
NoOverlappingActualsError

If no rows survive the inner join (forecasts cover a period with no observed data).

ValueError

If any forecast row has a negative lead time (valid_time before power_fcst_init_time — an undeliverable hindcast row; see issue #346), if any scored series has no row in capacity, or if any computed metric value is non-finite (NaN/inf — which Metrics.validate would otherwise accept, since NaN is not null, and which would poison the MLflow aggregate means).

Source code in packages/ml_core/src/ml_core/metrics.py
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
def compute_metrics(
    cv_forecasts: pt.DataFrame[PowerForecast],
    actuals: pt.LazyFrame[PowerTimeSeries],
    metadata: pt.DataFrame[TimeSeriesMetadata],
    capacity: pt.DataFrame[EffectiveCapacity],
) -> pt.DataFrame[Metrics]:
    """Compute evaluation metrics from CV predictions and observed power.

    Full metric definitions, equations, and design rationale:
    <https://openclimatefix.github.io/nged-substation-forecast/techniques/evaluation-metrics/>

    For each ``(time_series_id, fold_id, power_fcst_model_name)`` group:

    1. Joins predictions to observed ``power`` on ``(time_series_id, valid_time)``.
    2. Assigns each row a ``horizon_slice`` from its lead time
       (``valid_time − power_fcst_init_time``) — see ``_horizon_slice_expr`` for the bands.
    3. Collapses the ensemble members *within each forecast run* (per
       ``power_fcst_init_time``) into per-timestamp quantities: the deterministic ensemble
       mean, the fair CRPS, the Fortin-corrected ensemble variance, and the empirical
       ``DELIVERY_QUANTILES``. Each run covering a ``valid_time`` is scored independently,
       exactly as a production consumer would experience it — runs at different lead times
       are never pooled.
    4. Aggregates per ``horizon_slice``, plus the ``"all"`` aggregate over every lead time:
       MAE, NMAE, RMSE, and MBE on the ensemble mean; CRPS and the spread-skill ratio from
       the member-aware quantities; pinball loss at each delivery quantile (plus their
       mean); and PICP and interval width for each symmetric quantile band.
    5. Joins ``time_series_type`` from ``metadata`` onto each row.
    6. Returns one row per ``(time_series_id, fold_id, power_fcst_model_name,
       horizon_slice, metric_name, metric_param)`` in the tall ``Metrics`` format.

    NMAE is normalised by the pre-computed full-history ``effective_capacity_mw`` (joined per
    ``time_series_id`` from ``capacity``) — a capacity-like denominator, computed over the full
    history so it stays stable across folds. This ``time_series_id``-only join is correct while
    ``capacity`` is one scalar row per series (the v0.1 shape). When the v0.7
    upgrade makes capacity time-varying (one row per ``(time_series_id,
    time)``), this join must become a temporal as-of join on ``(time_series_id, valid_time)``.

    Single-member "ensembles" (e.g. a deterministic baseline forecaster) are scored
    unconditionally: their fair CRPS equals their MAE, their spread-skill ratio is 0, and
    their quantile bands are degenerate (all quantiles coincide, so PICP ≈ 0 and interval
    width = 0). Those are honest descriptions of a deterministic forecast, not errors.

    Args:
        cv_forecasts: CV predictions to evaluate.
            Currently only CV fold rows are handled; ``fold_id="live"`` support will
            be added with the ``production_monitoring`` scope in Phase 8.
        actuals: Observed half-hourly power (lazy — only the joined subset is collected).
            Deduplicated on ``(time_series_id, time)`` before the join — a duplicated actual
            would otherwise double the ensemble members and corrupt the member-aware metrics.
        metadata: Substation metadata used to join ``time_series_type`` onto each metric row.
            Series absent from ``metadata`` receive a null ``time_series_type``.
        capacity: Pre-computed per-series effective capacity; ``effective_capacity_mw`` is the
            NMAE denominator. Must cover every scored ``time_series_id``.

    Returns:
        A validated tall ``Metrics`` DataFrame with ``time_series_type`` populated.

    Raises:
        NoOverlappingActualsError: If no rows survive the inner join (forecasts cover a
            period with no observed data).
        ValueError: If any forecast row has a negative lead time (``valid_time`` before
            ``power_fcst_init_time`` — an undeliverable hindcast row; see issue #346), if
            any scored series has no row in ``capacity``, or if any computed metric value is
            non-finite (NaN/inf — which ``Metrics.validate`` would otherwise accept, since
            NaN is not null, and which would poison the MLflow aggregate means).
    """
    # A negative lead time means hindcast rows — valid times already in the past at
    # power_fcst_init_time, which a live forecast could never deliver. Scoring them would
    # silently flatter the model (they'd land in "intraday" via the left-closed bands), so
    # fail loudly. The CV inference pass currently emits such rows for valid times inside
    # the NWP publication-delay window; issue #346 tracks removing them at the source.
    n_negative = cv_forecasts.filter(pl.col("valid_time") < pl.col("power_fcst_init_time")).height
    if n_negative > 0:
        raise ValueError(
            f"{n_negative} forecast row(s) have valid_time before power_fcst_init_time "
            "(negative lead time). These are undeliverable hindcast rows and must not be "
            "scored — regenerate the forecasts without them (see issue #346)."
        )

    # Join forecasts to actuals; rename power → power_actual to avoid shadowing.
    # Strip the Patito model subclass from actuals so that Polars' cross-subclass
    # type check (assert_same_type) doesn't reject a join between two differently-typed
    # pt.LazyFrame objects.
    # Dedupe actuals on the join key: a duplicated (time_series_id, time) row would double
    # every ensemble member through the join, silently corrupting the member-aware metrics
    # (CRPS, spread, quantiles) while leaving the deterministic ones — which only see the
    # mean — untouched. Nothing enforces this uniqueness upstream, so guard it here.
    actuals_plain = pl.LazyFrame._from_pyldf(actuals._ldf)
    joined = cv_forecasts.lazy().join(
        actuals_plain.select(["time_series_id", "time", "power"])
        .unique(subset=["time_series_id", "time"], keep="any")
        .rename({"power": "power_actual"}),
        left_on=["time_series_id", "valid_time"],
        right_on=["time_series_id", "time"],
        how="inner",
    )

    # Collapse the ensemble members of each forecast run into per-timestamp quantities: the
    # deterministic ensemble mean plus the member-aware values (fair CRPS, Fortin-corrected
    # variance, empirical delivery quantiles). power_fcst_init_time is a group key so that
    # runs covering the same valid_time at different lead times are scored independently,
    # never pooled into a lagged-ensemble blend. horizon_slice is constant within a
    # (power_fcst_init_time, valid_time) group.
    quantile_aggs = {
        _quantile_column(q): pl.col("power_fcst").cast(pl.Float64).quantile(q, "linear")
        for q in DELIVERY_QUANTILES
    }
    per_run = (
        joined.with_columns(horizon_slice=_horizon_slice_expr())
        .group_by(
            [
                "time_series_id",
                "fold_id",
                "power_fcst_model_name",
                "power_fcst_init_time",
                "valid_time",
            ]
        )
        .agg(
            power_fcst=pl.col("power_fcst").mean(),
            power_actual=pl.col("power_actual").first(),
            horizon_slice=pl.col("horizon_slice").first(),
            crps=_fair_crps_expr(),
            corrected_var=_corrected_variance_expr(),
            **quantile_aggs,
        )
    )

    # Compute the ensemble-mean error column once, reuse in all aggregations.
    with_error = per_run.with_columns(error=pl.col("power_fcst") - pl.col("power_actual"))

    # Wide metrics: one row per (time_series_id, fold_id, power_fcst_model_name, horizon_slice),
    # where horizon_slice covers the four lead-time bands plus the "all" aggregate over every
    # lead time.
    base_keys = ["time_series_id", "fold_id", "power_fcst_model_name"]
    wide_columns = [*base_keys, "horizon_slice", *_wide_metric_columns()]
    per_slice = _wide_metrics(with_error, [*base_keys, "horizon_slice"])
    all_slice = _wide_metrics(with_error, base_keys).with_columns(
        horizon_slice=pl.lit("all").cast(pl.Enum(HORIZON_SLICES))
    )
    metrics_wide = pl.concat([per_slice.select(wide_columns), all_slice.select(wide_columns)])

    # Join the pre-computed full-history effective capacity — the NMAE denominator. Strip the
    # Patito model so Polars' cross-subclass join type check doesn't reject the join.
    capacity_denom = pl.LazyFrame._from_pyldf(capacity.lazy()._ldf).select(
        ["time_series_id", "effective_capacity_mw"]
    )
    wide = metrics_wide.join(capacity_denom, on="time_series_id", how="left").collect()

    # Every scored series must have a capacity row; fail loudly rather than silently emitting a
    # null NMAE (which Metrics.metric_value, a non-nullable Float32, would reject anyway).
    missing = wide.filter(pl.col("effective_capacity_mw").is_null())["time_series_id"]
    if missing.len() > 0:
        raise ValueError(
            f"No effective_capacity row for time_series_id(s) {sorted(set(missing.to_list()))}; "
            "materialise the effective_capacity asset for these series before scoring."
        )
    wide = wide.with_columns(nmae=pl.col("mae") / pl.col("effective_capacity_mw"))

    # Pivot to tall format, then split the "name:param" encoding of the parametric wide
    # columns into (metric_name, metric_param); scalar metrics have no separator, so their
    # split field_1 is null and fills to "all". The leftover effective_capacity_mw column is
    # dropped by the unpivot.
    name_parts = pl.col("metric_name").str.split_exact(":", 1)
    metrics_tall = (
        wide.unpivot(
            on=[*_wide_metric_columns(), "nmae"],
            index=[*base_keys, "horizon_slice"],
            variable_name="metric_name",
            value_name="metric_value",
        )
        .with_columns(
            metric_name=name_parts.struct.field("field_0"),
            metric_param=name_parts.struct.field("field_1").fill_null("all"),
        )
        .cast(
            {
                "metric_name": pl.Enum(METRIC_NAMES),
                "metric_param": pl.Enum(METRIC_PARAMS),
                "metric_value": pl.Float32,
            }
        )
    )

    if metrics_tall.is_empty():
        raise NoOverlappingActualsError(
            "No rows in the joined forecast/actuals data. "
            "Check that cv_forecasts and actuals overlap in time."
        )

    # NaN/inf are not null, so Metrics.validate would accept them — and a single non-finite
    # value poisons every downstream MLflow mean. Fail loudly instead, naming the offenders.
    non_finite = metrics_tall.filter(~pl.col("metric_value").is_finite())
    if not non_finite.is_empty():
        offenders = non_finite.select(
            ["time_series_id", "horizon_slice", "metric_name", "metric_param", "metric_value"]
        )
        raise ValueError(
            f"{non_finite.height} metric row(s) have a non-finite metric_value, which would "
            f"silently poison the MLflow aggregate means. First offenders:\n{offenders.head(10)}"
        )

    # Join time_series_type from metadata (left — keeps all metric rows; unmatched → null).
    type_map = metadata.select(["time_series_id", "time_series_type"])
    metrics_tall = metrics_tall.join(type_map, on="time_series_id", how="left").with_columns(
        time_series_type=pl.col("time_series_type").cast(pl.Enum(TIME_SERIES_TYPE_SLICES))
    )

    return Metrics.validate(metrics_tall, allow_superfluous_columns=True)

build_mlflow_aggregate_metrics(metrics_df)

Return a flat {metric_key: value} dict for mlflow.log_metrics.

Computes mean metric_value across all series ("all" aggregate), per time_series_type, and per horizon_slice. The key token is {metric_name} for metric_param="all" metrics and {metric_name}_{metric_param} for parametric ones (e.g. pinball_loss_p10, picp_p10_p90); parametric metrics are restricted to the _MLFLOW_LOGGED_PARAMETRIC headline subset. Key formats:

  • "{token}__all" — overall aggregate (horizon_slice="all").
  • "{token}__{type_slug}" — per-type aggregates (horizon_slice="all").
  • "{token}__all__{horizon_slice}" — overall aggregate per lead-time band (e.g. "nmae__all__day_ahead"). Per-type sliced aggregates are deliberately not logged — that detail stays queryable in the forecast_metrics Delta table, as does the full 13-quantile / 6-band parametric detail.

Parameters:

Name Type Description Default
metrics_df DataFrame

Per-series Metrics rows with time_series_type populated.

required

Returns:

Type Description
dict[str, float]

Flat dict of MLflow metric key → mean float value.

Source code in packages/ml_core/src/ml_core/metrics.py
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
def build_mlflow_aggregate_metrics(
    metrics_df: pl.DataFrame,
) -> dict[str, float]:
    """Return a flat ``{metric_key: value}`` dict for ``mlflow.log_metrics``.

    Computes mean ``metric_value`` across all series (``"all"`` aggregate), per
    ``time_series_type``, and per ``horizon_slice``. The key token is ``{metric_name}`` for
    ``metric_param="all"`` metrics and ``{metric_name}_{metric_param}`` for parametric ones
    (e.g. ``pinball_loss_p10``, ``picp_p10_p90``); parametric metrics are restricted to the
    ``_MLFLOW_LOGGED_PARAMETRIC`` headline subset. Key formats:

    - ``"{token}__all"`` — overall aggregate (``horizon_slice="all"``).
    - ``"{token}__{type_slug}"`` — per-type aggregates (``horizon_slice="all"``).
    - ``"{token}__all__{horizon_slice}"`` — overall aggregate per lead-time band
      (e.g. ``"nmae__all__day_ahead"``). Per-type sliced aggregates are deliberately not
      logged — that detail stays queryable in the ``forecast_metrics`` Delta table, as does
      the full 13-quantile / 6-band parametric detail.

    Args:
        metrics_df: Per-series ``Metrics`` rows with ``time_series_type`` populated.

    Returns:
        Flat dict of MLflow metric key → mean float value.
    """
    logged = metrics_df.filter(_mlflow_logged_expr()).with_columns(
        metric_key=_metric_key_token_expr()
    )
    base = logged.filter(pl.col("horizon_slice") == "all")

    result: dict[str, float] = {}

    # Per-type aggregates (only for non-null time_series_type values).
    if "time_series_type" in base.columns:
        per_type = (
            base.filter(pl.col("time_series_type").is_not_null())
            .group_by(["metric_key", "time_series_type"])
            .agg(mean_value=pl.col("metric_value").mean())
        )
        for row in per_type.iter_rows(named=True):
            key = f"{row['metric_key']}__{_type_slug(str(row['time_series_type']))}"
            result[key] = float(row["mean_value"])

    # Overall "all" aggregate (includes all series regardless of type).
    overall = base.group_by("metric_key").agg(mean_value=pl.col("metric_value").mean())
    for row in overall.iter_rows(named=True):
        result[f"{row['metric_key']}__all"] = float(row["mean_value"])

    # Overall aggregate per lead-time band (includes all series regardless of type).
    per_slice = (
        logged.filter(pl.col("horizon_slice") != "all")
        .group_by(["metric_key", "horizon_slice"])
        .agg(mean_value=pl.col("metric_value").mean())
    )
    for row in per_slice.iter_rows(named=True):
        result[f"{row['metric_key']}__all__{row['horizon_slice']}"] = float(row["mean_value"])

    return result

enrich_metrics_rows(per_series_metrics, experiment_name, evaluation_scope, window_start, window_end, window_label, computed_at, mlflow_run_id)

Add scope and evaluation-window provenance columns to a per-series Metrics frame.

Called by the metrics Dagster asset after compute_metrics() returns, once the window bounds and MLflow run ID are known. Kept here so the enrichment logic is unit-testable without Dagster.

Parameters:

Name Type Description Default
per_series_metrics DataFrame[Metrics]

Frame produced by compute_metrics().

required
experiment_name str

Experiment that produced these forecasts.

required
evaluation_scope EvalScopeType

"leaderboard" or "ad_hoc".

required
window_start datetime

Inclusive start of the evaluated valid_time window.

required
window_end datetime

Inclusive end of the evaluated valid_time window.

required
window_label str

Human-readable label (fold_id for leaderboard; "ad_hoc").

required
computed_at datetime

UTC timestamp when this metric batch was computed.

required
mlflow_run_id str | None

MLflow fold run ID; None for ad_hoc.

required

Returns:

Type Description
DataFrame[Metrics]

A validated Metrics DataFrame with all columns fully populated.

Source code in packages/ml_core/src/ml_core/metrics.py
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
def enrich_metrics_rows(
    per_series_metrics: pt.DataFrame[Metrics],
    experiment_name: str,
    evaluation_scope: EvalScopeType,
    window_start: datetime,
    window_end: datetime,
    window_label: str,
    computed_at: datetime,
    mlflow_run_id: str | None,
) -> pt.DataFrame[Metrics]:
    """Add scope and evaluation-window provenance columns to a per-series Metrics frame.

    Called by the ``metrics`` Dagster asset after ``compute_metrics()`` returns, once the
    window bounds and MLflow run ID are known. Kept here so the enrichment logic is
    unit-testable without Dagster.

    Args:
        per_series_metrics: Frame produced by ``compute_metrics()``.
        experiment_name: Experiment that produced these forecasts.
        evaluation_scope: ``"leaderboard"`` or ``"ad_hoc"``.
        window_start: Inclusive start of the evaluated ``valid_time`` window.
        window_end: Inclusive end of the evaluated ``valid_time`` window.
        window_label: Human-readable label (``fold_id`` for leaderboard; ``"ad_hoc"``).
        computed_at: UTC timestamp when this metric batch was computed.
        mlflow_run_id: MLflow fold run ID; ``None`` for ``ad_hoc``.

    Returns:
        A validated ``Metrics`` DataFrame with all columns fully populated.
    """
    return Metrics.validate(
        per_series_metrics.with_columns(
            experiment_name=pl.lit(experiment_name, dtype=pl.String),
            evaluation_scope=pl.lit(evaluation_scope).cast(pl.Enum(EVALUATION_SCOPES)),
            window_start=pl.lit(window_start).cast(UTC_DATETIME_DTYPE),
            window_end=pl.lit(window_end).cast(UTC_DATETIME_DTYPE),
            window_label=pl.lit(window_label, dtype=pl.String),
            computed_at=pl.lit(computed_at).cast(UTC_DATETIME_DTYPE),
            mlflow_run_id=pl.lit(mlflow_run_id, dtype=pl.String),
        )
    )