XGBoost Forecaster API

XGBoost Substation Forecaster

This package implements an XGBoost-based model to forecast power flows at NGED primary substations using numerical weather prediction (NWP) forecasts. It implements the BaseForecaster interface defined in ml_core.

How it works

One xgb.Booster is trained per time_series_id, so each substation's model can learn its own relationship between weather and power. Features are passed via the AllFeatures schema (see contracts), which joins NWP variables, power lag and rolling features, and static metadata. A feature is one input column offered to the model: a lag feature is the power measured a fixed interval before the target time, and a rolling feature is a statistic over a recent window of that power. Categorical and string columns are encoded as integer codes before being handed to XGBoost; all features are cast to Float32, and missing values are left as NaN so XGBoost handles them natively. ECMWF publishes ~51 members of the same weather forecast — one unperturbed control member plus ~50 perturbed members — and an ensemble forecast is one power prediction per member. The model is deterministic, and an ensemble forecast still comes out of that deterministic model — the XGBoostForecaster class below says how.

Both train() and predict() collect their input once, so keeping that collect bounded is the caller's job. The dominant cost is the multi-tens-of-GB NWP scan, which has to be pruned at the inputs and streamed. Feature engineering joins and upsamples the raw inputs into AllFeatures, and plans that work lazily, so filtering the engineered output cannot prune the NWP scan. The train and predict docstrings below give the mechanics. For the dataset sizes and the table of which predicates actually prune the NWP scan, see Bounding feature-engineering memory.

Save format

XGBoostForecaster.save(path) writes:

• {time_series_id}.ubj — one XGBoost native binary model per trained substation
• meta.json — the full XGBoostConfig serialised via Pydantic, the trained time_series_id population, and the model class, so load() is completely self-contained

Configuration

XGBoostConfig extends BaseForecasterConfig with XGBoost hyperparameters (n_estimators, learning_rate, max_depth, etc.). BaseForecasterConfig contributes selected_features, random_seed (threaded into XGBoost's own seed parameter for deterministic training), the experiment-identity fields experiment_name and ml_flow_experiment_id, and the leaderboard tag fields weather_source and training_strategy, which label a run so the leaderboard can compare like with like. Model-family identity — MODEL_NAME ("xgboost") and MODEL_VERSION — lives on the XGBoostForecaster class itself, not in the config; both a config's experiment identity and the class's model-family identity are stamped onto every row of the PowerForecast output, so the Delta Lake table is self-describing. XGBoostConfig inherits extra="forbid" from BaseForecasterConfig, so a key that names neither an XGBoostConfig field nor an inherited BaseForecasterConfig field raises ValidationError rather than being silently ignored.

xgboost_forecaster.forecaster

XGBoost-based power forecasting model.

Classes

XGBoostConfig

Bases: BaseForecasterConfig

Configuration for XGBoostForecaster.

Inherits the universal fields (selected_features, MLflow experiment id) from BaseForecasterConfig and adds XGBoost-specific hyperparameters.

Source code in packages/xgboost_forecaster/src/xgboost_forecaster/forecaster.py

class XGBoostConfig(BaseForecasterConfig):
    """Configuration for XGBoostForecaster.

    Inherits the universal fields (selected_features, MLflow experiment id) from
    BaseForecasterConfig and adds XGBoost-specific hyperparameters.
    """

    n_estimators: int = 1000
    learning_rate: float = 0.05
    max_depth: int = 6
    min_child_weight: int = 1
    subsample: float = 0.8
    colsample_bytree: float = 0.8
    device: str = "cpu"
    objective: str = "reg:squarederror"

    def to_xgb_params(self) -> dict[str, Any]:
        """Return the params dict accepted by xgb.train()."""
        return {
            "eta": self.learning_rate,
            "max_depth": self.max_depth,
            "min_child_weight": self.min_child_weight,
            "subsample": self.subsample,
            "colsample_bytree": self.colsample_bytree,
            "device": self.device,
            "objective": self.objective,
            "seed": self.random_seed,
        }

Attributes

n_estimators = 1000 class-attribute instance-attribute

learning_rate = 0.05 class-attribute instance-attribute

max_depth = 6 class-attribute instance-attribute

min_child_weight = 1 class-attribute instance-attribute

subsample = 0.8 class-attribute instance-attribute

colsample_bytree = 0.8 class-attribute instance-attribute

device = 'cpu' class-attribute instance-attribute

objective = 'reg:squarederror' class-attribute instance-attribute

model_config = ConfigDict(extra='forbid') class-attribute instance-attribute

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

Methods:

to_xgb_params()

Return the params dict accepted by xgb.train().

Source code in packages/xgboost_forecaster/src/xgboost_forecaster/forecaster.py

def to_xgb_params(self) -> dict[str, Any]:
    """Return the params dict accepted by xgb.train()."""
    return {
        "eta": self.learning_rate,
        "max_depth": self.max_depth,
        "min_child_weight": self.min_child_weight,
        "subsample": self.subsample,
        "colsample_bytree": self.colsample_bytree,
        "device": self.device,
        "objective": self.objective,
        "seed": self.random_seed,
    }

XGBoostForecaster

Bases: BaseForecaster

Trains and serves one XGBoost Booster per time_series_id.

All lead times for a given time_series_id are handled by a single Booster. The model is deterministic; ensemble forecasts arise because each NWP ensemble member's weather is a separate row through the relevant Booster. predict scores every ensemble member present in its input in one call. predict groups by time_series_id and dispatches each group to its Booster.

Save layout: a directory holding one {time_series_id}.ubj file per trained Booster, plus a meta.json. meta.json stores the full XGBoostConfig, the trained time_series_id population, and the model class, so that load() is self-contained.

Source code in packages/xgboost_forecaster/src/xgboost_forecaster/forecaster.py

class XGBoostForecaster(BaseForecaster):
    """Trains and serves one XGBoost Booster per time_series_id.

    All lead times for a given time_series_id are handled by a single Booster. The model is
    deterministic; ensemble forecasts arise because each NWP ensemble member's weather is a
    separate row through the relevant Booster. ``predict`` scores every ensemble member present
    in its input in one call. ``predict`` groups by ``time_series_id`` and dispatches each group
    to its Booster.

    Save layout: a directory holding one ``{time_series_id}.ubj`` file per trained Booster, plus
    a ``meta.json``. ``meta.json`` stores the full XGBoostConfig, the trained ``time_series_id``
    population, and the model class, so that load() is self-contained.
    """

    MODEL_NAME = "xgboost"
    MODEL_VERSION = 1
    CONFIG_CLASS: ClassVar[type[XGBoostConfig]] = XGBoostConfig

    model_params: XGBoostConfig  # narrows the base class annotation for type checkers

    def __init__(self, model_params: XGBoostConfig) -> None:
        """Start with no trained boosters; ``train`` populates one per ``time_series_id``."""
        super().__init__(model_params)
        self._models: dict[int, xgb.Booster] = {}

    @property
    def _feature_cols(self) -> list[str]:
        return sorted(self.model_params.selected_features)

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

        For ``XGBoostForecaster``, the population is exactly the set of series it holds a trained
        Booster for, one Booster per ``time_series_id``. ``predict`` ignores rows for any other
        series. ``save()`` records this set in ``meta.json``, which is what ``load()`` reads it
        back from. A model scores exactly the series it was trained on, no more and no fewer:
        that rule is the train==predict population invariant. See
        ``BaseForecaster.trained_time_series_ids`` for the model-agnostic contract this
        implements.
        """
        return sorted(self._models.keys())

    def train(self, data: pt.LazyFrame[AllFeatures], time_series_ids: list[int]) -> None:
        """Fit one Booster per ``time_series_id`` in ``time_series_ids``.

        ``data`` is collected once and grouped in memory by ``time_series_id``. Each group's rows
        feed an ``xgb.QuantileDMatrix``, which compresses them to 8-bit quantile bins rather than
        holding an uncompressed Float32 copy. Keeping the collect bounded is the *caller's* job.
        The NWP scan must be pruned at the inputs: the control member, the relevant H3 cells, and
        the window's ``init_time`` partitions. An ECMWF ensemble run holds one unperturbed
        control member plus ~50 perturbed members, and training uses the control member alone.
        Filtering the engineered output cannot prune the upstream join/upsample. See
        ``load_engineering_inputs`` and "Bounding feature-engineering memory: prune the inputs,
        not the output" in
        <https://openclimatefix.github.io/nged-substation-forecast/architecture/performance/#bounding-feature-engineering-memory-prune-the-inputs-not-the-output>.

        Only the requested ``time_series_ids`` are trained. A requested series with no non-null
        ``power`` rows does not appear, and gets no Booster. A series can lack those rows because
        none fall in the training window.
        """
        feature_cols = self._feature_cols
        requested = set(time_series_ids)
        # Stream the collect. init_time prunes whole NWP partitions. Rows are written sorted by a
        # fixed column order, and a filter on a column early in that order can be answered from each
        # row group's recorded minimum and maximum. The member-early sort
        # (delta_store.nwp.NWP_SORT_COLS) therefore lets row-group stats skip most of each partition
        # for the control-member read. h3_index is not a sort-early column, so cell filtering is
        # still decode-then-filter within the surviving row groups. The streaming engine applies
        # those predicates per morsel, holding peak memory to a few GB where the in-memory engine
        # would materialise every surviving row first. See
        # <https://openclimatefix.github.io/nged-substation-forecast/architecture/performance/#bounding-feature-engineering-memory-prune-the-inputs-not-the-output>.
        df = data.drop_nulls(subset=["power"]).collect(engine="streaming")
        for group_key, group in df.group_by(["time_series_id"]):
            ts_id = int(group_key[0])
            if ts_id not in requested:
                continue
            group = cast(pt.DataFrame[AllFeatures], group)
            features = _prepare_features(group, feature_cols)
            label = group["power"].cast(pl.Float32)
            dtrain = xgb.QuantileDMatrix(features, label=label)
            booster = xgb.train(
                self.model_params.to_xgb_params(),
                dtrain,
                num_boost_round=self.model_params.n_estimators,
            )
            self._models[ts_id] = booster

    def predict(
        self, data: pt.LazyFrame[AllFeatures], *, fold_id: str = "live"
    ) -> pt.DataFrame[PowerForecast]:
        """Generate one power_fcst per row, dispatching by time_series_id to the right Booster.

        ``data`` is collected once and grouped in memory by ``time_series_id``. This model
        ignores rows for any ``time_series_id`` it was not trained on, and only scores its own
        trained population (see ``trained_time_series_ids``). Keeping the collect bounded is the
        caller's job. Training reads the control member alone, but at validation the full
        ~51-member NWP ensemble is present, so the caller predicts one ``init_time`` chunk at a
        time. The caller appends each chunk to Delta as it goes. ``init_time`` is one of the NWP
        table's two partition columns, and it is the axis that fans the output out across runs.
        Chunking on ``init_time`` therefore bounds each iteration's forecast frame, and still
        reads every partition exactly once. Looping per H3 cell instead runs out of memory on the
        busiest cell. See ``cv_power_forecasts`` and
        <https://openclimatefix.github.io/nged-substation-forecast/architecture/performance/#bounding-feature-engineering-memory-prune-the-inputs-not-the-output>.

        ``fold_id`` is stamped onto every output row (the model has no inherent fold; the caller
        supplies it). Defaults to the ``"live"`` production sentinel.
        """
        feature_cols = self._feature_cols
        cfg = self.model_params

        def _build_part(
            group_df: pt.DataFrame[AllFeatures], predictions: np.ndarray
        ) -> pl.DataFrame:
            return group_df.select(
                [
                    "valid_time",
                    "time_series_id",
                    "ensemble_member",
                    "nwp_init_time",
                    "power_fcst_init_time",
                ]
            ).with_columns(
                pl.Series("power_fcst", predictions, dtype=pl.Float32),
                power_fcst_model_name=pl.lit(self.MODEL_NAME),
                power_fcst_model_version=pl.lit(self.MODEL_VERSION, dtype=pl.Int16),
                ml_flow_experiment_id=pl.lit(cfg.ml_flow_experiment_id, dtype=pl.Int32),
                experiment_name=pl.lit(cfg.experiment_name),
                fold_id=pl.lit(fold_id),
            )

        # Stream the NWP scan, for the reasons on ``train`` above. See
        # <https://openclimatefix.github.io/nged-substation-forecast/architecture/performance/#bounding-feature-engineering-memory-prune-the-inputs-not-the-output>.
        df = data.collect(engine="streaming")
        parts: list[pl.DataFrame] = []
        for group_key, group in df.group_by(["time_series_id"]):
            booster = self._models.get(int(group_key[0]))
            if booster is None:
                continue  # ignore series this model was not trained on
            group = cast(pt.DataFrame[AllFeatures], group)
            X = _prepare_features(group, feature_cols)
            predictions: np.ndarray = booster.predict(xgb.DMatrix(X))
            parts.append(_build_part(group, predictions))

        if not parts:
            empty = cast(pt.DataFrame[AllFeatures], df.head(0))
            parts.append(_build_part(empty, np.empty(0, dtype=np.float32)))

        # The identity columns (power_fcst_model_name / experiment_name / fold_id) are built as
        # ``pl.lit(str)`` above, so they are already ``String``. ``String`` is the dtype
        # PowerForecast declares, so these columns need no cast before validation.
        return PowerForecast.validate(pl.concat(parts))

    def save(self, path: Path) -> None:
        """Save all Boosters as .ubj files plus a meta.json with the full config.

        Clears ``path`` first. The dropped series' boosters can then never be left behind on
        disk, even when re-saving a model trained on fewer series over a directory that already
        holds the ``.ubj`` files of a model trained on more (issue #197). The same
        replace-don't-merge property holds through MLflow, because
        ``BaseForecaster.save_to_mlflow`` uploads this directory as a single archive artifact.
        """
        shutil.rmtree(path, ignore_errors=True)
        path.mkdir(parents=True, exist_ok=True)
        for ts_id, booster in self._models.items():
            booster.save_model(str(path / f"{ts_id}.ubj"))
        (path / "meta.json").write_text(
            json.dumps(
                {
                    "model_params": self.model_params.model_dump(mode="json"),
                    "trained_time_series_ids": self.trained_time_series_ids,
                    "model_class": class_target(self),
                }
            )
        )

    @classmethod
    def load(cls, path: Path) -> Self:
        """Reconstruct an XGBoostForecaster from a saved directory.

        ``meta.json``'s ``trained_time_series_ids`` decides the population, not whatever ``.ubj``
        files happen to be in the directory. A model's population is the model's own frozen
        record, not a directory listing (issue #197). A directory can hold files this model did
        not write. ``ml_core.base_forecaster.save_to_mlflow`` adds a
        ``time_series_metadata.parquet``. ``ml_core.production_helpers.fetch_model_artifacts``
        adds a ``promotion.json``. A hand-assembled directory can hold any file at all. Globbing
        would let one of those files enlarge the population. The model would then silently score
        a series it was never trained for, breaking the train==predict invariant (see
        ``BaseForecaster.trained_time_series_ids``).
        """
        meta = json.loads((path / "meta.json").read_text())
        config = cls.CONFIG_CLASS.model_validate(meta["model_params"])
        instance = cls(config)
        for ts_id in meta["trained_time_series_ids"]:
            booster = xgb.Booster()
            booster.load_model(str(path / f"{ts_id}.ubj"))
            instance._models[int(ts_id)] = booster
        return instance

Attributes

MODEL_NAME = 'xgboost' class-attribute instance-attribute

MODEL_VERSION = 1 class-attribute instance-attribute

CONFIG_CLASS = XGBoostConfig class-attribute

model_params instance-attribute

trained_time_series_ids property

The sorted time_series_ids this forecaster will serve a predict for.

For XGBoostForecaster, the population is exactly the set of series it holds a trained Booster for, one Booster per time_series_id. predict ignores rows for any other series. save() records this set in meta.json, which is what load() reads it back from. A model scores exactly the series it was trained on, no more and no fewer: that rule is the train==predict population invariant. See BaseForecaster.trained_time_series_ids for the model-agnostic contract this implements.

feature_engineer = TabularFeatureEngineer() class-attribute

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

the forecaster references a feature engineer rather than

implementing feature engineering. A forecaster can therefore swap the whole pipeline by overriding this attribute with a different FeatureEngineer. The default produces the tabular AllFeatures

frame that train/predict consume.

Methods:

__init__(model_params)

Start with no trained boosters; train populates one per time_series_id.

Source code in packages/xgboost_forecaster/src/xgboost_forecaster/forecaster.py

def __init__(self, model_params: XGBoostConfig) -> None:
    """Start with no trained boosters; ``train`` populates one per ``time_series_id``."""
    super().__init__(model_params)
    self._models: dict[int, xgb.Booster] = {}

train(data, time_series_ids)

Fit one Booster per time_series_id in time_series_ids.

data is collected once and grouped in memory by time_series_id. Each group's rows feed an xgb.QuantileDMatrix, which compresses them to 8-bit quantile bins rather than holding an uncompressed Float32 copy. Keeping the collect bounded is the caller's job. The NWP scan must be pruned at the inputs: the control member, the relevant H3 cells, and the window's init_time partitions. An ECMWF ensemble run holds one unperturbed control member plus ~50 perturbed members, and training uses the control member alone. Filtering the engineered output cannot prune the upstream join/upsample. See load_engineering_inputs and "Bounding feature-engineering memory: prune the inputs, not the output" in https://openclimatefix.github.io/nged-substation-forecast/architecture/performance/#bounding-feature-engineering-memory-prune-the-inputs-not-the-output.

Only the requested time_series_ids are trained. A requested series with no non-null power rows does not appear, and gets no Booster. A series can lack those rows because none fall in the training window.

Source code in packages/xgboost_forecaster/src/xgboost_forecaster/forecaster.py

def train(self, data: pt.LazyFrame[AllFeatures], time_series_ids: list[int]) -> None:
    """Fit one Booster per ``time_series_id`` in ``time_series_ids``.

    ``data`` is collected once and grouped in memory by ``time_series_id``. Each group's rows
    feed an ``xgb.QuantileDMatrix``, which compresses them to 8-bit quantile bins rather than
    holding an uncompressed Float32 copy. Keeping the collect bounded is the *caller's* job.
    The NWP scan must be pruned at the inputs: the control member, the relevant H3 cells, and
    the window's ``init_time`` partitions. An ECMWF ensemble run holds one unperturbed
    control member plus ~50 perturbed members, and training uses the control member alone.
    Filtering the engineered output cannot prune the upstream join/upsample. See
    ``load_engineering_inputs`` and "Bounding feature-engineering memory: prune the inputs,
    not the output" in
    <https://openclimatefix.github.io/nged-substation-forecast/architecture/performance/#bounding-feature-engineering-memory-prune-the-inputs-not-the-output>.

    Only the requested ``time_series_ids`` are trained. A requested series with no non-null
    ``power`` rows does not appear, and gets no Booster. A series can lack those rows because
    none fall in the training window.
    """
    feature_cols = self._feature_cols
    requested = set(time_series_ids)
    # Stream the collect. init_time prunes whole NWP partitions. Rows are written sorted by a
    # fixed column order, and a filter on a column early in that order can be answered from each
    # row group's recorded minimum and maximum. The member-early sort
    # (delta_store.nwp.NWP_SORT_COLS) therefore lets row-group stats skip most of each partition
    # for the control-member read. h3_index is not a sort-early column, so cell filtering is
    # still decode-then-filter within the surviving row groups. The streaming engine applies
    # those predicates per morsel, holding peak memory to a few GB where the in-memory engine
    # would materialise every surviving row first. See
    # <https://openclimatefix.github.io/nged-substation-forecast/architecture/performance/#bounding-feature-engineering-memory-prune-the-inputs-not-the-output>.
    df = data.drop_nulls(subset=["power"]).collect(engine="streaming")
    for group_key, group in df.group_by(["time_series_id"]):
        ts_id = int(group_key[0])
        if ts_id not in requested:
            continue
        group = cast(pt.DataFrame[AllFeatures], group)
        features = _prepare_features(group, feature_cols)
        label = group["power"].cast(pl.Float32)
        dtrain = xgb.QuantileDMatrix(features, label=label)
        booster = xgb.train(
            self.model_params.to_xgb_params(),
            dtrain,
            num_boost_round=self.model_params.n_estimators,
        )
        self._models[ts_id] = booster

predict(data, *, fold_id='live')

Generate one power_fcst per row, dispatching by time_series_id to the right Booster.

data is collected once and grouped in memory by time_series_id. This model ignores rows for any time_series_id it was not trained on, and only scores its own trained population (see trained_time_series_ids). Keeping the collect bounded is the caller's job. Training reads the control member alone, but at validation the full ~51-member NWP ensemble is present, so the caller predicts one init_time chunk at a time. The caller appends each chunk to Delta as it goes. init_time is one of the NWP table's two partition columns, and it is the axis that fans the output out across runs. Chunking on init_time therefore bounds each iteration's forecast frame, and still reads every partition exactly once. Looping per H3 cell instead runs out of memory on the busiest cell. See cv_power_forecasts and https://openclimatefix.github.io/nged-substation-forecast/architecture/performance/#bounding-feature-engineering-memory-prune-the-inputs-not-the-output.

fold_id is stamped onto every output row (the model has no inherent fold; the caller supplies it). Defaults to the "live" production sentinel.

Source code in packages/xgboost_forecaster/src/xgboost_forecaster/forecaster.py

def predict(
    self, data: pt.LazyFrame[AllFeatures], *, fold_id: str = "live"
) -> pt.DataFrame[PowerForecast]:
    """Generate one power_fcst per row, dispatching by time_series_id to the right Booster.

    ``data`` is collected once and grouped in memory by ``time_series_id``. This model
    ignores rows for any ``time_series_id`` it was not trained on, and only scores its own
    trained population (see ``trained_time_series_ids``). Keeping the collect bounded is the
    caller's job. Training reads the control member alone, but at validation the full
    ~51-member NWP ensemble is present, so the caller predicts one ``init_time`` chunk at a
    time. The caller appends each chunk to Delta as it goes. ``init_time`` is one of the NWP
    table's two partition columns, and it is the axis that fans the output out across runs.
    Chunking on ``init_time`` therefore bounds each iteration's forecast frame, and still
    reads every partition exactly once. Looping per H3 cell instead runs out of memory on the
    busiest cell. See ``cv_power_forecasts`` and
    <https://openclimatefix.github.io/nged-substation-forecast/architecture/performance/#bounding-feature-engineering-memory-prune-the-inputs-not-the-output>.

    ``fold_id`` is stamped onto every output row (the model has no inherent fold; the caller
    supplies it). Defaults to the ``"live"`` production sentinel.
    """
    feature_cols = self._feature_cols
    cfg = self.model_params

    def _build_part(
        group_df: pt.DataFrame[AllFeatures], predictions: np.ndarray
    ) -> pl.DataFrame:
        return group_df.select(
            [
                "valid_time",
                "time_series_id",
                "ensemble_member",
                "nwp_init_time",
                "power_fcst_init_time",
            ]
        ).with_columns(
            pl.Series("power_fcst", predictions, dtype=pl.Float32),
            power_fcst_model_name=pl.lit(self.MODEL_NAME),
            power_fcst_model_version=pl.lit(self.MODEL_VERSION, dtype=pl.Int16),
            ml_flow_experiment_id=pl.lit(cfg.ml_flow_experiment_id, dtype=pl.Int32),
            experiment_name=pl.lit(cfg.experiment_name),
            fold_id=pl.lit(fold_id),
        )

    # Stream the NWP scan, for the reasons on ``train`` above. See
    # <https://openclimatefix.github.io/nged-substation-forecast/architecture/performance/#bounding-feature-engineering-memory-prune-the-inputs-not-the-output>.
    df = data.collect(engine="streaming")
    parts: list[pl.DataFrame] = []
    for group_key, group in df.group_by(["time_series_id"]):
        booster = self._models.get(int(group_key[0]))
        if booster is None:
            continue  # ignore series this model was not trained on
        group = cast(pt.DataFrame[AllFeatures], group)
        X = _prepare_features(group, feature_cols)
        predictions: np.ndarray = booster.predict(xgb.DMatrix(X))
        parts.append(_build_part(group, predictions))

    if not parts:
        empty = cast(pt.DataFrame[AllFeatures], df.head(0))
        parts.append(_build_part(empty, np.empty(0, dtype=np.float32)))

    # The identity columns (power_fcst_model_name / experiment_name / fold_id) are built as
    # ``pl.lit(str)`` above, so they are already ``String``. ``String`` is the dtype
    # PowerForecast declares, so these columns need no cast before validation.
    return PowerForecast.validate(pl.concat(parts))

save(path)

Save all Boosters as .ubj files plus a meta.json with the full config.

Clears path first. The dropped series' boosters can then never be left behind on disk, even when re-saving a model trained on fewer series over a directory that already holds the .ubj files of a model trained on more (issue #197). The same replace-don't-merge property holds through MLflow, because BaseForecaster.save_to_mlflow uploads this directory as a single archive artifact.

Source code in packages/xgboost_forecaster/src/xgboost_forecaster/forecaster.py

def save(self, path: Path) -> None:
    """Save all Boosters as .ubj files plus a meta.json with the full config.

    Clears ``path`` first. The dropped series' boosters can then never be left behind on
    disk, even when re-saving a model trained on fewer series over a directory that already
    holds the ``.ubj`` files of a model trained on more (issue #197). The same
    replace-don't-merge property holds through MLflow, because
    ``BaseForecaster.save_to_mlflow`` uploads this directory as a single archive artifact.
    """
    shutil.rmtree(path, ignore_errors=True)
    path.mkdir(parents=True, exist_ok=True)
    for ts_id, booster in self._models.items():
        booster.save_model(str(path / f"{ts_id}.ubj"))
    (path / "meta.json").write_text(
        json.dumps(
            {
                "model_params": self.model_params.model_dump(mode="json"),
                "trained_time_series_ids": self.trained_time_series_ids,
                "model_class": class_target(self),
            }
        )
    )

load(path) classmethod

Reconstruct an XGBoostForecaster from a saved directory.

meta.json's trained_time_series_ids decides the population, not whatever .ubj files happen to be in the directory. A model's population is the model's own frozen record, not a directory listing (issue #197). A directory can hold files this model did not write. ml_core.base_forecaster.save_to_mlflow adds a time_series_metadata.parquet. ml_core.production_helpers.fetch_model_artifacts adds a promotion.json. A hand-assembled directory can hold any file at all. Globbing would let one of those files enlarge the population. The model would then silently score a series it was never trained for, breaking the train==predict invariant (see BaseForecaster.trained_time_series_ids).

Source code in packages/xgboost_forecaster/src/xgboost_forecaster/forecaster.py

@classmethod
def load(cls, path: Path) -> Self:
    """Reconstruct an XGBoostForecaster from a saved directory.

    ``meta.json``'s ``trained_time_series_ids`` decides the population, not whatever ``.ubj``
    files happen to be in the directory. A model's population is the model's own frozen
    record, not a directory listing (issue #197). A directory can hold files this model did
    not write. ``ml_core.base_forecaster.save_to_mlflow`` adds a
    ``time_series_metadata.parquet``. ``ml_core.production_helpers.fetch_model_artifacts``
    adds a ``promotion.json``. A hand-assembled directory can hold any file at all. Globbing
    would let one of those files enlarge the population. The model would then silently score
    a series it was never trained for, breaking the train==predict invariant (see
    ``BaseForecaster.trained_time_series_ids``).
    """
    meta = json.loads((path / "meta.json").read_text())
    config = cls.CONFIG_CLASS.model_validate(meta["model_params"])
    instance = cls(config)
    for ts_id in meta["trained_time_series_ids"]:
        booster = xgb.Booster()
        booster.load_model(str(path / f"{ts_id}.ubj"))
        instance._models[int(ts_id)] = booster
    return instance

save_to_mlflow(run_id, *, time_series_metadata)

Upload this trained model to the given MLflow run, as one replaceable archive.

Writes the model to a temporary directory via save (the subclass's own format), then adds the frozen metadata copy (write_trained_metadata). Packs that directory into a single model.tar.gz, and logs that one file to the run's artifact root. Logging one archive rather than a directory of files is what makes a re-upload replace the previous model instead of merging with it — see _MLFLOW_MODEL_ARTIFACT. 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 artifact to.	required
time_series_metadata	DataFrame[TimeSeriesMetadata]	The roster rows this model was engineered against. Required, because a model uploaded without them cannot be promoted — see TRAINED_METADATA_FILENAME. Narrowed to trained_time_series_ids before the rows are written. Callers engineer over a wider population than they end up training, because an eligible series with no usable power gets no model. Every consumer of these rows reads only the NWP cells the rows' own H3 cells cover, so carrying rows for series the model cannot forecast would widen that NWP read for nothing.	required

Source code in packages/ml_core/src/ml_core/base_forecaster.py

def save_to_mlflow(
    self, run_id: str, *, time_series_metadata: pt.DataFrame[TimeSeriesMetadata]
) -> None:
    """Upload this trained model to the given MLflow run, as one replaceable archive.

    Writes the model to a temporary directory via ``save`` (the subclass's own format), then
    adds the frozen metadata copy (``write_trained_metadata``). Packs that directory into a
    single ``model.tar.gz``, and logs *that one file* to the run's artifact root. Logging one
    archive rather than a directory of files is what makes a re-upload **replace** the previous
    model instead of merging with it — see ``_MLFLOW_MODEL_ARTIFACT``. The caller is responsible
    for setting the tracking URI (``mlflow.set_tracking_uri``) beforehand.

    Args:
        run_id: The MLflow run to attach the artifact to.
        time_series_metadata: The roster rows this model was engineered against. Required,
            because a model uploaded without them cannot be promoted — see
            ``TRAINED_METADATA_FILENAME``. Narrowed to ``trained_time_series_ids`` before the
            rows are written. Callers engineer over a wider population than they end up
            training, because an eligible series with no usable power gets no model. Every
            consumer of these rows reads only the NWP cells the rows' own H3 cells cover, so
            carrying rows for series the model cannot forecast would widen that NWP read for
            nothing.
    """
    with tempfile.TemporaryDirectory() as tmp_dir:
        model_dir = Path(tmp_dir) / "model"
        model_dir.mkdir()
        self.save(model_dir)
        write_trained_metadata(
            model_dir=model_dir,
            time_series_metadata=time_series_metadata.filter(
                pl.col("time_series_id").is_in(self.trained_time_series_ids)
            ),
        )
        archive_path = Path(tmp_dir) / _MLFLOW_MODEL_ARTIFACT
        _archive_model_dir(model_dir, archive_path)
        with mlflow.start_run(run_id=run_id):
            mlflow.log_artifact(str(archive_path))

load_from_mlflow(run_id) classmethod

Download a trained model's archive from an MLflow run, unpack it and load it.

Downloads into a temporary directory and loads from there. There is deliberately no local-disk cache. A CV fold run is reused across re-materialisations: get_or_create_fold_run resolves the same run for every re-run of a fold's partition, and each re-run overwrites that run's single model archive. One run_id therefore names different model bytes at different times, so a cache keyed by run_id would not be unique for its contents. Production inference makes no MLflow call at all either. Full rationale: https://openclimatefix.github.io/nged-substation-forecast/architecture/ml-orchestration/#why-there-is-no-local-cache.

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 the model was saved under.	required

Returns:

Type	Description
Self	The reconstructed, trained forecaster.

Source code in packages/ml_core/src/ml_core/base_forecaster.py

@classmethod
def load_from_mlflow(cls, run_id: str) -> Self:
    """Download a trained model's archive from an MLflow run, unpack it and load it.

    Downloads into a temporary directory and loads from there. There is deliberately no
    local-disk cache. A CV fold run is **reused** across re-materialisations:
    ``get_or_create_fold_run`` resolves the same run for every re-run of a fold's partition, and
    each re-run overwrites that run's single model archive. One ``run_id`` therefore names
    different model bytes at different times, so a cache keyed by ``run_id`` would not be unique
    for its contents. Production inference makes no MLflow call at all either. Full rationale:
    <https://openclimatefix.github.io/nged-substation-forecast/architecture/ml-orchestration/#why-there-is-no-local-cache>.

    The caller is responsible for setting the tracking URI (``mlflow.set_tracking_uri``)
    beforehand.

    Args:
        run_id: The MLflow run the model was saved under.

    Returns:
        The reconstructed, trained forecaster.
    """
    with tempfile.TemporaryDirectory() as tmp_dir:
        return cls.load(
            _download_and_unpack_model(
                run_id=run_id,
                work_dir=Path(tmp_dir),
                remedy="re-materialise `trained_cv_model` for this fold.",
            )
        )

Functions: