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 andParsedFeaturesparser._nwp— NWP upsampling, processing, and power/NWP join helpers._lags— power-lag, weather-lag (dual-strategy), and leaky-lag nullification.tabular_feature_engineer—_engineer_featuresorchestrator 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 | |
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_timegiven): power-centric, joins exactly one NWP run (nwp_init_time, or derived frompower_fcst_init_timeminusnwp_publication_delay_hoursif omitted) and stamps a constantpower_fcst_init_timeon 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 |
required |
time_series_metadata
|
DataFrame[TimeSeriesMetadata]
|
Per-time-series metadata (carries |
required |
nwp
|
LazyFrame[Nwp]
|
Gridded NWP in physical units, keyed by |
required |
power_fcst_init_time
|
datetime | None
|
|
None
|
nwp_init_time
|
datetime | None
|
The NWP run to join in single-run mode. Must be |
None
|
nwp_publication_delay_hours
|
int
|
Delay used to derive whichever of
|
NWP_PUBLICATION_DELAY_HOURS
|
Returns:
| Type | Description |
|---|---|
LazyFrame[AllFeatures]
|
A lazy |
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 | |
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 | |
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 | |
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 | |
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 | |
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 | |
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 | |
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 | |
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 | |
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
|
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 | |
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 ( |
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 | |
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 |
'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 | |
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 | |
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:
- Joins predictions to observed
poweron(time_series_id, valid_time). - Assigns each row a
horizon_slicefrom its lead time (valid_time − power_fcst_init_time) — see_horizon_slice_exprfor the bands. - 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 empiricalDELIVERY_QUANTILES. Each run covering avalid_timeis scored independently, exactly as a production consumer would experience it — runs at different lead times are never pooled. - 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. - Joins
time_series_typefrommetadataonto each row. - Returns one row per
(time_series_id, fold_id, power_fcst_model_name, horizon_slice, metric_name, metric_param)in the tallMetricsformat.
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; |
required |
actuals
|
LazyFrame[PowerTimeSeries]
|
Observed half-hourly power (lazy — only the joined subset is collected).
Deduplicated on |
required |
metadata
|
DataFrame[TimeSeriesMetadata]
|
Substation metadata used to join |
required |
capacity
|
DataFrame[EffectiveCapacity]
|
Pre-computed per-series effective capacity; |
required |
Returns:
| Type | Description |
|---|---|
DataFrame[Metrics]
|
A validated tall |
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 ( |
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 | |
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 theforecast_metricsDelta table, as does the full 13-quantile / 6-band parametric detail.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
metrics_df
|
DataFrame
|
Per-series |
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 | |
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 |
required |
experiment_name
|
str
|
Experiment that produced these forecasts. |
required |
evaluation_scope
|
EvalScopeType
|
|
required |
window_start
|
datetime
|
Inclusive start of the evaluated |
required |
window_end
|
datetime
|
Inclusive end of the evaluated |
required |
window_label
|
str
|
Human-readable label ( |
required |
computed_at
|
datetime
|
UTC timestamp when this metric batch was computed. |
required |
mlflow_run_id
|
str | None
|
MLflow fold run ID; |
required |
Returns:
| Type | Description |
|---|---|
DataFrame[Metrics]
|
A validated |
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 | |