From 481f15574a98f375b5566fa621aec9e3175e5c2f Mon Sep 17 00:00:00 2001
From: Ana Gainaru <ana.gainaru@gmail.com>
Date: Mon, 2 Mar 2026 10:10:47 -0500
Subject: [PATCH 1/9] Include the AERIS dataset

---
 examples/utils.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/examples/utils.py b/examples/utils.py
index 0cde7b7..392a845 100644
--- a/examples/utils.py
+++ b/examples/utils.py
@@ -15,6 +15,10 @@ def get_example(cfg: Config) -> BaseModelHarness:
         from examples.imagenet.model import IMAGENET_VISION
 
         return IMAGENET_VISION(cfg=cfg)
+    elif cfg.data.name == "aeris_dataset.csv":
+        from examples.aeris.model import AERIS
+
+        return AERIS(cfg=cfg)
     else:
         raise NotImplementedError(
             f"Example for dataset {cfg.data.name} is not implemented."

From 5213164b4f64096553ae0921970c1982f188b916 Mon Sep 17 00:00:00 2001
From: Ana Gainaru <ana.gainaru@gmail.com>
Date: Mon, 2 Mar 2026 10:16:15 -0500
Subject: [PATCH 2/9] AERIS model harness

---
 examples/aeris/__init__.py |   0
 examples/aeris/aeris.toml  |  52 ++++++
 examples/aeris/model.py    | 228 +++++++++++++++++++++++
 examples/aeris/utils.py    | 357 +++++++++++++++++++++++++++++++++++++
 4 files changed, 637 insertions(+)
 create mode 100644 examples/aeris/__init__.py
 create mode 100644 examples/aeris/aeris.toml
 create mode 100644 examples/aeris/model.py
 create mode 100644 examples/aeris/utils.py

diff --git a/examples/aeris/__init__.py b/examples/aeris/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/examples/aeris/aeris.toml b/examples/aeris/aeris.toml
new file mode 100644
index 0000000..7d16c7c
--- /dev/null
+++ b/examples/aeris/aeris.toml
@@ -0,0 +1,52 @@
+# aeris.toml  —  AERIS continuous-learning
+seed = 42
+device = "auto"
+multi_gpu = false
+verbosity = "INFO"
+
+[model]
+name = "aeris_model.pt"
+pretrained_path = "examples/aeris/model"
+
+[data]
+name = "aeris_dataset.csv"
+path = "examples/aeris/data"
+
+[train]
+batch_size = 500
+num_workers = 4
+init_lr = 1e-6
+max_iter = 4000
+grad_accumulation_steps = 1
+
+[continual_learning]
+update_mode = "base"
+
+# JVP regularization (used when update_mode = "jvp_reg")
+jvp_lambda = 10
+jvp_deltax_norm = 1
+
+# EWC (used when update_mode = "ewc_online")
+ewc_lambda = 1000.0
+ewc_ema_decay = 0.95
+
+# KFAC (used when update_mode = "kfac_online")
+kfac_lambda = 1e-2
+kfac_ema_decay = 0.95
+
+[drift_detection]
+detector_name = "ADWINDetector"
+detection_interval = 10
+aggregation = "mean"
+metric_index = 0
+reset_after_learning = false
+max_stream_updates = 20
+
+# ADWIN hyperparameters
+adwin_delta = 0.002
+adwin_minor_threshold = 0.3
+adwin_moderate_threshold = 0.6
+
+[logging]
+backend = "wandb"
+experiment_name = "aeris-cl"
diff --git a/examples/aeris/model.py b/examples/aeris/model.py
new file mode 100644
index 0000000..2dfdb80
--- /dev/null
+++ b/examples/aeris/model.py
@@ -0,0 +1,228 @@
+# examples/aeris/model.py
+"""AERIS model harness for the BaseSim continuous-learning framework.
+
+This harness wraps a 8-layer neural network trained to predict enthalpy per atom from a given fuel material."""
+
+import gc
+import torch
+from typing import Tuple, Optional, List, Any, Mapping, cast
+from torch import nn, Tensor
+from torch.optim import Optimizer
+from torch.utils.data import DataLoader, ConcatDataset, TensorDataset
+
+from model.torch_model_harness import BaseModelHarness
+from config.configuration import Config
+
+from examples.aeris.utils import (
+    load_datasets,
+    make_loader,
+    load_pretrained_model,
+    split_into_windows,
+)
+
+
+# Aeris model architecture used for prediction
+class AerisFullStructure(nn.Module):
+    def __init__(self, input_dim, dropout=0.3):
+        super().__init__()
+        first_layer = min(1024, max(512, input_dim * 2))
+        self.layers = nn.Sequential(
+            nn.Linear(input_dim, first_layer),
+            nn.ReLU(),
+            nn.BatchNorm1d(first_layer),
+            nn.Linear(first_layer, first_layer),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(first_layer, 512),
+            nn.ReLU(),
+            nn.BatchNorm1d(512),
+            nn.Linear(512, 512),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(512, 256),
+            nn.ReLU(),
+            nn.BatchNorm1d(256),
+            nn.Linear(256, 256),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(256, 128),
+            nn.ReLU(),
+            nn.BatchNorm1d(128),
+            nn.Linear(128, 64),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(64, 32),
+            nn.ReLU(),
+            nn.Linear(32, 1),
+        )
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+# Fraction of each time window reserved for validation
+_VAL_FRACTION: float = 0.2
+
+
+class AERIS(BaseModelHarness):
+    """
+    Continuous-learning harness for the AERIS prediction model.
+    """
+
+    def __init__(self, cfg: Config):
+        # ----- build model ---------------------------------------------------
+        ckpt = load_pretrained_model(
+            cfg.model.pretrained_path, cfg.model.name, device=cfg.device
+        )
+
+        # Checkpoint is a dict saved via torch.save(model_info, ...)
+        input_dim_raw = ckpt.get("input_dim")
+        if input_dim_raw is None:
+            raise KeyError("Checkpoint missing required key: 'input_dim'")
+        input_dim = int(cast(int, input_dim_raw))
+
+        feature_names_raw = ckpt.get("feature_names")
+        if feature_names_raw is None:
+            raise KeyError("Checkpoint missing required key: 'feature_names'")
+        feature_names = cast(List[str], feature_names_raw)
+
+        scaler_raw = ckpt.get("scaler")
+        if scaler_raw is None:
+            raise KeyError("Checkpoint missing required key: 'scaler'")
+        scaler = cast(Any, scaler_raw)
+
+        state_raw = ckpt.get("model_state_dict")
+        if state_raw is None:
+            raise KeyError("Checkpoint missing required key: 'model_state_dict'")
+        state = cast(Mapping[str, Any], state_raw)
+
+        model = AerisFullStructure(input_dim=input_dim)
+        model.load_state_dict(state)
+        model.to(cfg.device)
+
+        super().__init__(cfg=cfg, model=model)
+
+        # ----- eval metrics (prediction) -------------------------------------
+        self.eval_metrics = {"mae": self.mae_metric(), "loss": self.get_criterion()}
+        self.higher_is_better = {"accuracy": False, "loss": False}
+
+        # ----- data loaders  -------------------------------------
+        X, y = load_datasets(cfg.data.path, cfg.data.name, feature_names)
+        X_raw = torch.tensor(X, dtype=torch.float32)
+        X_scaled: Tensor = scaler.transform(X_raw)
+        y_raw = torch.tensor(y, dtype=torch.float32)
+        # y is a 1D array of shape (N,), but model outputs (N, 1):
+        if y_raw.ndim == 1:
+            y_raw = y_raw.unsqueeze(1)
+
+        self.windows = split_into_windows(X_scaled, y_raw)
+
+        # ----- streaming state -----------------------------------------------
+        self.window_idx: int = 0
+        self.history_windows: List[Tuple[Tensor, Tensor]] = []
+
+        self._cur_train_loader: Optional[DataLoader] = None
+        self._cur_val_loader: Optional[DataLoader] = None
+
+    def get_optmizer(self) -> Optimizer:  # noqa: D102  (spelling kept for ABC)
+        return torch.optim.Adam(self.model.parameters(), lr=self.cfg.train.init_lr)
+
+    def get_criterion(self):  # noqa: D102
+        return nn.MSELoss()
+
+    def mae_metric(self):
+        return nn.L1Loss()
+
+    def get_cur_data_loaders(self) -> Tuple[DataLoader, DataLoader]:  # noqa: D102
+        assert self._cur_train_loader is not None and self._cur_val_loader is not None
+        return self._cur_train_loader, self._cur_val_loader
+
+    def get_hist_data_loaders(
+        self,
+    ) -> Tuple[Optional[DataLoader], Optional[DataLoader]]:
+        """Return loaders over all previously-seen time windows.
+
+        Returns ``(None, None)`` until at least two windows have been served.
+        """
+        if self.window_idx <= 1:
+            return None, None
+
+        # Concatenate all history windows
+        hist_train_views: List[TensorDataset] = []
+        hist_val_views: List[TensorDataset] = []
+
+        for X_w, y_w in self.history_windows:
+            n = X_w.shape[0]
+            n_val = max(1, int(n * _VAL_FRACTION))
+            n_train = n - n_val
+            hist_train_views.append(TensorDataset(X_w[:n_train], y_w[:n_train]))
+            hist_val_views.append(TensorDataset(X_w[n_train:], y_w[n_train:]))
+
+        ds_hist_train: ConcatDataset[Any] = ConcatDataset(hist_train_views)
+        ds_hist_val: ConcatDataset[Any] = ConcatDataset(hist_val_views)
+
+        bs = self.cfg.train.batch_size
+        nw = self.cfg.train.num_workers
+        pin = torch.cuda.is_available()
+        return (
+            make_loader(
+                ds_hist_train, bs, shuffle=True, num_workers=nw, pin_memory=pin
+            ),
+            make_loader(ds_hist_val, bs, shuffle=False, num_workers=nw, pin_memory=pin),
+        )
+
+    def update_data_stream(self) -> None:
+        """Advance to the next chronological time window.
+
+        The current window is added to the history, and new train/val loaders
+        are built from the upcoming window.
+        """
+        self._dispose_current_loaders()
+
+        if self.window_idx >= len(self.windows):
+            print(
+                f"Warning: All {len(self.windows)} time windows exhausted; "
+                "wrapping around to the first window."
+            )
+            self.window_idx = 0
+
+        X_w, y_w = self.windows[self.window_idx]
+
+        # Archive previous window in history (skip the very first call)
+        if self.window_idx > 0:
+            prev_X, prev_y = self.windows[self.window_idx - 1]
+            # Only add if not already stored (idempotency guard)
+            if len(self.history_windows) < self.window_idx:
+                self.history_windows.append((prev_X, prev_y))
+        # Train / val split (last _VAL_FRACTION chronologically)
+        n = X_w.shape[0]
+        n_val = max(1, int(n * _VAL_FRACTION))
+        n_train = n - n_val
+
+        ds_train = TensorDataset(X_w[:n_train], y_w[:n_train])
+        ds_val = TensorDataset(X_w[n_train:], y_w[n_train:])
+
+        bs = self.cfg.train.batch_size
+        nw = self.cfg.train.num_workers
+        pin = torch.cuda.is_available()
+
+        self._cur_train_loader = make_loader(
+            ds_train, bs, shuffle=True, num_workers=nw, pin_memory=pin
+        )
+        self._cur_val_loader = make_loader(
+            ds_val, bs, shuffle=False, num_workers=nw, pin_memory=pin
+        )
+
+        self.window_idx += 1
+
+    # --------------------------------------------------------------------- #
+    # Helpers
+    # --------------------------------------------------------------------- #
+    def _dispose_current_loaders(self) -> None:
+        if self._cur_train_loader is not None:
+            del self._cur_train_loader
+            self._cur_train_loader = None
+        if self._cur_val_loader is not None:
+            del self._cur_val_loader
+            self._cur_val_loader = None
+        gc.collect()
diff --git a/examples/aeris/utils.py b/examples/aeris/utils.py
new file mode 100644
index 0000000..ac2b253
--- /dev/null
+++ b/examples/aeris/utils.py
@@ -0,0 +1,357 @@
+# examples/aeris/utils.py
+"""Utility functions for the AERIS continuous-learning example.
+
+Expected directory layout (pointed to by ``cfg.data.path``)::
+
+    <data_dir>/
+        dataset.csv            # data that will be parsed by the SIM framework
+        aeris_model.pt        # AERIS pre-trained model
+"""
+
+import os
+import glob
+import re
+from typing import Dict, List, Tuple, Any
+
+import numpy as np
+import pandas as pd
+import torch
+from torch import Tensor
+from torch.utils.data import DataLoader, Dataset
+
+from pymatgen.core.composition import Composition
+from matminer.featurizers.base import MultipleFeaturizer
+from matminer.featurizers import composition as cf
+
+
+def load_pretrained_model(
+    data_path: str, model_name: str, device: str = "cpu"
+) -> dict[str, Any]:
+    """Load the pretrained AERIS model.
+
+    Parameters
+    ----------
+    data_path:
+        Directory containing the model.
+    model_name:
+        The name of the pretrained model.
+    device:
+        Device to map the scalers to.
+
+    Returns
+    -------
+    model_info = {
+        'model_state_dict': model.state_dict(),
+        'input_dim': input_dim,
+        'feature_names': feature_names,
+        'scaler': scaler,
+        'metrics': {'mae': mae, 'rmse': rmse, 'r2': r2},
+        'history': history
+    }
+    """
+    ckpt = None
+    if os.path.exists(data_path):
+        ckpt = torch.load(
+            os.path.join(data_path, model_name), map_location=device, weights_only=False
+        )
+    if ckpt is None:
+        raise FileNotFoundError("No model found at path: " + data_path)
+    return ckpt
+
+
+def _parse_formula(s: str) -> Dict[str, float]:
+    parts = re.findall(r"([A-Z][a-z]?)([0-9]*\.?[0-9]*)", str(s).strip())
+    if not parts:
+        raise ValueError(f"Could not parse formula: {s}")
+    comp: Dict[str, float] = {}
+    for el, num in parts:
+        comp[el] = float(num) if num else 1.0
+    return comp
+
+
+def _apply_df_parse_formula_num(val):
+    try:
+        if pd.isna(val):
+            return None
+        parsed = _parse_formula(str(val))
+        return int(sum(parsed.values()))
+    except Exception:
+        return None
+
+
+def _apply_df_parse_formula_str(val):
+    try:
+        if pd.isna(val):
+            return None
+        parsed = _parse_formula(str(val))
+        return "".join(f"{k}{v}" for k, v in sorted(parsed.items()))
+    except Exception:
+        return None
+
+
+def _parse_structure_string(struct_str: str) -> Dict[str, float]:
+    # minimal lattice extractor (compatible with training utils)
+    result = {
+        "lattice_a": np.nan,
+        "lattice_b": np.nan,
+        "lattice_c": np.nan,
+        "lattice_alpha": np.nan,
+        "lattice_beta": np.nan,
+        "lattice_gamma": np.nan,
+        "volume": np.nan,
+        "density": np.nan,
+        "nsites": np.nan,
+        "spacegroup_number": np.nan,
+    }
+    if struct_str is None:
+        return result
+    s = str(struct_str)
+    abc_pattern = r"abc\s*:\s*([\d.]+)\s+([\d.]+)\s+([\d.]+)"
+    angles_pattern = r"angles\s*:\s*([\d.]+)\s+([\d.]+)\s+([\d.]+)"
+    abc = re.search(abc_pattern, s)
+    ang = re.search(angles_pattern, s)
+    if abc:
+        result["lattice_a"] = float(abc.group(1))
+        result["lattice_b"] = float(abc.group(2))
+        result["lattice_c"] = float(abc.group(3))
+    if ang:
+        result["lattice_alpha"] = float(ang.group(1))
+        result["lattice_beta"] = float(ang.group(2))
+        result["lattice_gamma"] = float(ang.group(3))
+    # try volume
+    vol_match = re.search(r"volume\s*[:=]\s*([\d.]+)", s)
+    if vol_match:
+        result["volume"] = float(vol_match.group(1))
+    dens_match = re.search(r"density\s*[:=]\s*([\d.]+)", s)
+    if dens_match:
+        result["density"] = float(dens_match.group(1))
+    sg_match = re.search(r"spacegroup(?:_number)?\s*[:=]\s*(\d+)", s)
+    if sg_match:
+        result["spacegroup_number"] = int(sg_match.group(1))
+    nsites_match = re.search(r"nsites\s*[:=]\s*(\d+)", s)
+    if nsites_match:
+        result["nsites"] = int(nsites_match.group(1))
+    return result
+
+
+def _build_feature_vector(
+    composition: str, features: Dict, feature_names: List[str]
+) -> np.ndarray:
+    comp = _parse_formula(composition)
+    total_atoms = float(sum(comp.values()))
+    # prepare composition fractions
+    elem_frac = {k: v / total_atoms for k, v in comp.items()}
+
+    # parse structure if string/dict
+    struct_vals = {}
+    if features is not None:
+        for feature in features:
+            if feature not in feature_names:
+                continue
+            struct_vals[feature] = features[feature]
+        if "structure" in features:
+            parsed_struct = _parse_structure_string(features["structure"])
+            struct_vals.update(parsed_struct)
+
+    # magpie
+    feature_calculators = MultipleFeaturizer(
+        [
+            cf.Stoichiometry(),
+            cf.ElementProperty.from_preset("magpie"),
+            cf.ValenceOrbital(props=["avg"]),
+            cf.IonProperty(fast=True),
+        ]
+    )
+
+    comp_obj = Composition(composition)
+    data = pd.DataFrame([{"comp_obj": comp_obj, "composition_reduced": composition}])
+
+    # Calculate Magpie features.
+    # IMPORTANT: when running under MPI, do NOT let matminer spawn multiprocessing pools
+    # inside each rank (oversubscription/hangs). Force single-process.
+    magpie_features_dict = {}
+    try:
+        # Some matminer versions support n_jobs; if yours does, keep it at 1.
+        magpie_features = feature_calculators.featurize_dataframe(
+            data, col_id="comp_obj", ignore_errors=True, pbar=False, n_jobs=1
+        )
+        magpie_features_dict = magpie_features.iloc[0].to_dict()
+    except Exception:
+        try:
+            feats = feature_calculators.featurize_many(
+                [Composition(composition)], n_jobs=1
+            )
+            magpie_features = pd.DataFrame(feats)
+            magpie_features.index = [0]
+            magpie_features_dict = magpie_features.iloc[0].to_dict()
+        except Exception as e:
+            print("Magpie featurizer failed, falling back to empty features:", repr(e))
+            magpie_features_dict = {}
+
+    vec = np.zeros(len(feature_names), dtype=np.float32)
+    for i, name in enumerate(feature_names):
+        # elemental features (assume single element name)
+        if re.match(r"^[A-Z][a-z]?$", name) and name in elem_frac:
+            vec[i] = float(elem_frac.get(name, 0.0))
+            continue
+
+        # structural features
+        if name in struct_vals:
+            vec[i] = float(struct_vals[name])
+            continue
+
+        # magpie features
+        if name in magpie_features_dict:
+            vec[i] = float(magpie_features_dict[name])
+            continue
+
+        # try numeric keys in struct_vals
+        val = struct_vals.get(name)
+        if val is None:
+            v = struct_vals.get(name, 0.0)
+            try:
+                vec[i] = float(v)
+            except Exception:
+                vec[i] = 0.0
+        else:
+            try:
+                vec[i] = float(val)
+            except Exception:
+                vec[i] = 0.0
+
+    X = vec.reshape(1, -1)
+    # if there are nan values in the feature vector
+    return np.nan_to_num(X, nan=0.0, posinf=1e6, neginf=-1e6)
+
+
+def load_datasets(data_path: str, dataset_name: str, feature_names: List[str]):
+    """Load the dataset that will be parsed, return features and ground truth.
+
+    Parameters
+    ----------
+    data_path:
+        Directory containing the datasets.
+    dataset_name:
+        The name or regular expression for the datasets
+    feature_names:
+        The features used by the model for prediction
+
+    Returns
+    -------
+    input festures, output target values
+    """
+    dfs = []
+    dataset_pattern = os.path.join(data_path, dataset_name)
+    dataset_files: List[str] = glob.glob(dataset_pattern)
+    if not dataset_files:
+        raise FileNotFoundError(f"No dataset files matched pattern: {dataset_pattern}")
+    for file_path in dataset_files:
+        dfs.append(pd.read_csv(file_path, low_memory=False))
+    dataset: pd.DataFrame = pd.concat(dfs, ignore_index=True)
+
+    # Filter all entries that do not have a target value
+    dataset = dataset.dropna(subset=["formation_energy_per_atom"]).copy()
+
+    # Replace NaN/+inf/-inf in numeric columns (keep DataFrame type)
+    num_cols = dataset.select_dtypes(include=[np.number]).columns
+    dataset[num_cols] = dataset[num_cols].replace([np.inf, -np.inf], np.nan).fillna(0.0)
+
+    y = dataset["formation_energy_per_atom"].values.astype(np.float32)
+    X = []
+    for _, row in dataset.iterrows():
+        composition = row["composition_reduced"]
+        features = {
+            "composition": row["composition"],
+            "structure": row["structure"],
+            "spacegroup_number": row["spacegroup_number"],
+            "density_atomic": row["density_atomic"],
+            "CN_max": row["CN_max"],
+            "CN_min": row["CN_min"],
+            "CN_avg": row["CN_avg"],
+        }
+        X.append(_build_feature_vector(composition, features, feature_names))
+
+    assert len(X) == len(y), (
+        "The feature and target vectors do not have the same lenght"
+    )
+    return X, y
+
+
+# Default number of samples per time window.  Can be overridden by the caller.
+DEFAULT_WINDOW_SIZE: int = 5000
+
+
+def split_into_windows(
+    X: Tensor,
+    y: Tensor,
+    window_size: int = DEFAULT_WINDOW_SIZE,
+) -> List[Tuple[Tensor, Tensor]]:
+    """Split chronologically-ordered tensors into non-overlapping windows.
+
+    Any leftover samples that don't fill a complete window are appended as
+    a final (smaller) window so no data is discarded.
+
+    Parameters
+    ----------
+    X:
+        Input features ``[N, D]``.
+    y:
+        Targets ``[N, T]``.
+    window_size:
+        Number of samples per window.
+
+    Returns
+    -------
+    List of ``(X_chunk, y_chunk)`` tuples.
+    """
+    n = X.shape[0]
+    windows: List[Tuple[Tensor, Tensor]] = []
+    for start in range(0, n, window_size):
+        end = min(start + window_size, n)
+        windows.append((X[start:end], y[start:end]))
+    return windows
+
+
+def make_loader(
+    ds: Dataset,
+    batch_size: int,
+    shuffle: bool,
+    num_workers: int = 4,
+    pin_memory: bool = True,
+    persistent_workers: bool = True,
+    prefetch_factor: int = 2,
+) -> DataLoader:
+    """Build a ``DataLoader`` from a ``Dataset``.
+
+    Parameters
+    ----------
+    ds:
+        The base dataset.
+    batch_size:
+        Batch size.
+    shuffle:
+        Whether to shuffle.
+    num_workers:
+        Number of data-loading workers.
+    pin_memory:
+        Pin CUDA memory for faster transfers.
+    persistent_workers:
+        Keep worker processes alive between iterations.
+    prefetch_factor:
+        Samples to prefetch per worker.
+
+    Returns
+    -------
+    DataLoader
+    """
+    kwargs: dict = dict(batch_size=batch_size, shuffle=shuffle, drop_last=False)
+    if num_workers > 0:
+        kwargs.update(
+            dict(
+                num_workers=num_workers,
+                pin_memory=pin_memory,
+                persistent_workers=persistent_workers,
+                prefetch_factor=prefetch_factor,
+            )
+        )
+    return DataLoader(ds, **kwargs)  # type: ignore[arg-type]

From 481392cdaeca8b5d4a4e9543281af9ee3c481536 Mon Sep 17 00:00:00 2001
From: Ana Gainaru <ana.gainaru@gmail.com>
Date: Mon, 2 Mar 2026 16:52:59 -0500
Subject: [PATCH 3/9] wip

---
 examples/aeris/model.py |  18 ++++---
 examples/aeris/utils.py | 115 ++++++++++++++++++++++++++++++++++++----
 2 files changed, 118 insertions(+), 15 deletions(-)

diff --git a/examples/aeris/model.py b/examples/aeris/model.py
index 2dfdb80..4ef6619 100644
--- a/examples/aeris/model.py
+++ b/examples/aeris/model.py
@@ -5,6 +5,7 @@
 
 import gc
 import torch
+import numpy as np
 from typing import Tuple, Optional, List, Any, Mapping, cast
 from torch import nn, Tensor
 from torch.optim import Optimizer
@@ -108,14 +109,19 @@ def __init__(self, cfg: Config):
 
         # ----- data loaders  -------------------------------------
         X, y = load_datasets(cfg.data.path, cfg.data.name, feature_names)
-        X_raw = torch.tensor(X, dtype=torch.float32)
-        X_scaled: Tensor = scaler.transform(X_raw)
+        # X shape: (n_samples, 1, 245) y shape: (n_samples,1)
+
+        # apply scaler if present
+        if scaler is not None:
+            try:
+                X = scaler.transform(X)
+            except Exception:
+                pass
+        with torch.no_grad():
+            X_scaled = torch.FloatTensor(X).to(cfg.device)
         y_raw = torch.tensor(y, dtype=torch.float32)
-        # y is a 1D array of shape (N,), but model outputs (N, 1):
-        if y_raw.ndim == 1:
-            y_raw = y_raw.unsqueeze(1)
 
-        self.windows = split_into_windows(X_scaled, y_raw)
+        self.windows = split_into_windows(X_scaled, y_raw, cfg.train.batch_size)
 
         # ----- streaming state -----------------------------------------------
         self.window_idx: int = 0
diff --git a/examples/aeris/utils.py b/examples/aeris/utils.py
index ac2b253..78cb86a 100644
--- a/examples/aeris/utils.py
+++ b/examples/aeris/utils.py
@@ -178,9 +178,8 @@ def _build_feature_vector(
         magpie_features_dict = magpie_features.iloc[0].to_dict()
     except Exception:
         try:
-            feats = feature_calculators.featurize_many(
-                [Composition(composition)], n_jobs=1
-            )
+            feature_calculators.set_n_jobs(1)
+            feats = feature_calculators.featurize_many([Composition(composition)])
             magpie_features = pd.DataFrame(feats)
             magpie_features.index = [0]
             magpie_features_dict = magpie_features.iloc[0].to_dict()
@@ -219,12 +218,106 @@ def _build_feature_vector(
             except Exception:
                 vec[i] = 0.0
 
-    X = vec.reshape(1, -1)
-    # if there are nan values in the feature vector
-    return np.nan_to_num(X, nan=0.0, posinf=1e6, neginf=-1e6)
+    # Return a 1D feature vector (D,) instead of (1, D)
+    vec = np.nan_to_num(vec, nan=0.0, posinf=1e6, neginf=-1e6)
+    return vec
 
 
 def load_datasets(data_path: str, dataset_name: str, feature_names: List[str]):
+    """Load the dataset used by the model.
+
+    This function attempts to *prefer* loading the exact columns listed in
+    `feature_names` (in the same order). If those columns are present in the
+    CSV(s), they are used directly (fast, deterministic). If not all feature
+    columns are present, the function falls back to building feature vectors
+    row-by-row using _build_feature_vector to preserve compatibility with older
+    or alternate CSV formats.
+
+    The function returns:
+        X: numpy.ndarray of shape (n_samples, n_features) dtype float32
+        y: numpy.ndarray of shape (n_samples,) dtype float32
+
+    Note: scaling is intentionally NOT applied here. The caller (model harness)
+    will apply the saved scaler from the checkpoint (if any) via scaler.transform().
+    """
+
+    # collect files
+    dataset_pattern = os.path.join(data_path, dataset_name)
+    dataset_files: List[str] = glob.glob(dataset_pattern)
+    if not dataset_files:
+        raise FileNotFoundError(f"No dataset files matched pattern: {dataset_pattern}")
+
+    # read & concatenate CSV files
+    dfs = []
+    for file_path in dataset_files:
+        dfs.append(pd.read_csv(file_path, low_memory=False))
+    dataset: pd.DataFrame = pd.concat(dfs, ignore_index=True)
+
+    # Ensure target present
+    if "formation_energy_per_atom" not in dataset.columns:
+        raise KeyError("Required target column 'formation_energy_per_atom' not found in dataset")
+
+    # Drop rows missing target
+    dataset = dataset.dropna(subset=["formation_energy_per_atom"]).copy()
+
+    # If all feature_names are present as columns, take that branch (preferred)
+    all_present = all((fn in dataset.columns) for fn in feature_names)
+
+    if all_present:
+        # Select columns in the exact saved order
+        X = dataset[feature_names].to_numpy(dtype=np.float32)
+
+        # Replace infs / NaNs in numeric columns with column means (same as training)
+        # (do not change dtype or drop rows here; keep alignment with model)
+        numeric_mask = np.isfinite(X)
+        # For each column replace non-finite with column mean (computed over finite rows)
+        col_means = np.nanmean(np.where(np.isfinite(X), X, np.nan), axis=0)
+        # Where a column is completely NaN/inf, set mean to 0.0
+        col_means = np.where(np.isnan(col_means), 0.0, col_means)
+        inds = np.where(~np.isfinite(X))
+        if inds[0].size > 0:
+            X[inds] = np.take(col_means, inds[1])
+
+    else:
+        # Fall back to building feature vectors row-by-row using the helper
+        # This creates the same ordering as feature_names when possible (elemental
+        # names are interpreted by _build_feature_vector).
+        X_rows = []
+        # Build a minimal features dict per row (this mirrors training's inputs)
+        for _, row in dataset.iterrows():
+            comp = row.get("composition_reduced", row.get("composition", None))
+            features = {
+                "composition": row.get("composition", None),
+                "structure": row.get("structure", None),
+                "spacegroup_number": row.get("spacegroup_number", None),
+                "density_atomic": row.get("density_atomic", None),
+                "CN_max": row.get("CN_max", None),
+                "CN_min": row.get("CN_min", None),
+                "CN_avg": row.get("CN_avg", None),
+            }
+            try:
+                vec = _build_feature_vector(comp, features, feature_names)
+            except Exception:
+                # on failure, append zeros to avoid mismatched shapes
+                vec = np.zeros(len(feature_names), dtype=np.float32)
+            X_rows.append(vec)
+        X = np.vstack(X_rows).astype(np.float32)
+
+        # Replace any remaining nan/inf with column means
+        col_means = np.nanmean(np.where(np.isfinite(X), X, np.nan), axis=0)
+        col_means = np.where(np.isnan(col_means), 0.0, col_means)
+        inds = np.where(~np.isfinite(X))
+        if inds[0].size > 0:
+            X[inds] = np.take(col_means, inds[1])
+
+    # Prepare target vector shape (N,)
+    y = dataset["formation_energy_per_atom"].to_numpy(dtype=np.float32).reshape(-1, 1)
+    assert X.shape[0] == y.shape[0], "Feature matrix and target vector must have same number of rows"
+
+    return X, y
+
+
+def load_datasets2(data_path: str, dataset_name: str, feature_names: List[str]):
     """Load the dataset that will be parsed, return features and ground truth.
 
     Parameters
@@ -256,7 +349,11 @@ def load_datasets(data_path: str, dataset_name: str, feature_names: List[str]):
     num_cols = dataset.select_dtypes(include=[np.number]).columns
     dataset[num_cols] = dataset[num_cols].replace([np.inf, -np.inf], np.nan).fillna(0.0)
 
-    y = dataset["formation_energy_per_atom"].values.astype(np.float32)
+    y = (
+        dataset["formation_energy_per_atom"]
+        .values.astype(np.float32)
+        .reshape(-1, 1)
+    )
     X = []
     for _, row in dataset.iterrows():
         composition = row["composition_reduced"]
@@ -271,6 +368,7 @@ def load_datasets(data_path: str, dataset_name: str, feature_names: List[str]):
         }
         X.append(_build_feature_vector(composition, features, feature_names))
 
+    print("X shape:", np.array(X).shape, "y shape:", y.shape)
     assert len(X) == len(y), (
         "The feature and target vectors do not have the same lenght"
     )
@@ -278,8 +376,7 @@ def load_datasets(data_path: str, dataset_name: str, feature_names: List[str]):
 
 
 # Default number of samples per time window.  Can be overridden by the caller.
-DEFAULT_WINDOW_SIZE: int = 5000
-
+DEFAULT_WINDOW_SIZE: int = 100
 
 def split_into_windows(
     X: Tensor,

From 3c707bf798abcf79f1029eaa5b9391e36048df48 Mon Sep 17 00:00:00 2001
From: Ana Gainaru <ana.gainaru@gmail.com>
Date: Tue, 3 Mar 2026 15:46:24 -0500
Subject: [PATCH 4/9] Updates to build the features in the same order as the
 training

---
 examples/aeris/model.py |  88 +++-------
 examples/aeris/utils.py | 355 +++++++++++++++-------------------------
 2 files changed, 159 insertions(+), 284 deletions(-)

diff --git a/examples/aeris/model.py b/examples/aeris/model.py
index 4ef6619..ed1c767 100644
--- a/examples/aeris/model.py
+++ b/examples/aeris/model.py
@@ -21,40 +21,21 @@
     split_into_windows,
 )
 
-
 # Aeris model architecture used for prediction
 class AerisFullStructure(nn.Module):
     def __init__(self, input_dim, dropout=0.3):
         super().__init__()
         first_layer = min(1024, max(512, input_dim * 2))
         self.layers = nn.Sequential(
-            nn.Linear(input_dim, first_layer),
-            nn.ReLU(),
-            nn.BatchNorm1d(first_layer),
-            nn.Linear(first_layer, first_layer),
-            nn.ReLU(),
-            nn.Dropout(dropout),
-            nn.Linear(first_layer, 512),
-            nn.ReLU(),
-            nn.BatchNorm1d(512),
-            nn.Linear(512, 512),
-            nn.ReLU(),
-            nn.Dropout(dropout),
-            nn.Linear(512, 256),
-            nn.ReLU(),
-            nn.BatchNorm1d(256),
-            nn.Linear(256, 256),
-            nn.ReLU(),
-            nn.Dropout(dropout),
-            nn.Linear(256, 128),
-            nn.ReLU(),
-            nn.BatchNorm1d(128),
-            nn.Linear(128, 64),
-            nn.ReLU(),
-            nn.Dropout(dropout),
-            nn.Linear(64, 32),
-            nn.ReLU(),
-            nn.Linear(32, 1),
+            nn.Linear(input_dim, first_layer), nn.ReLU(), nn.BatchNorm1d(first_layer),
+            nn.Linear(first_layer, first_layer), nn.ReLU(), nn.Dropout(dropout),
+            nn.Linear(first_layer, 512), nn.ReLU(), nn.BatchNorm1d(512),
+            nn.Linear(512, 512), nn.ReLU(), nn.Dropout(dropout),
+            nn.Linear(512, 256), nn.ReLU(), nn.BatchNorm1d(256),
+            nn.Linear(256, 256), nn.ReLU(), nn.Dropout(dropout),
+            nn.Linear(256, 128), nn.ReLU(), nn.BatchNorm1d(128),
+            nn.Linear(128, 64), nn.ReLU(), nn.Dropout(dropout),
+            nn.Linear(64, 32), nn.ReLU(), nn.Linear(32, 1)
         )
 
     def forward(self, x):
@@ -75,31 +56,14 @@ def __init__(self, cfg: Config):
         ckpt = load_pretrained_model(
             cfg.model.pretrained_path, cfg.model.name, device=cfg.device
         )
-
-        # Checkpoint is a dict saved via torch.save(model_info, ...)
-        input_dim_raw = ckpt.get("input_dim")
-        if input_dim_raw is None:
-            raise KeyError("Checkpoint missing required key: 'input_dim'")
-        input_dim = int(cast(int, input_dim_raw))
-
-        feature_names_raw = ckpt.get("feature_names")
-        if feature_names_raw is None:
-            raise KeyError("Checkpoint missing required key: 'feature_names'")
-        feature_names = cast(List[str], feature_names_raw)
-
-        scaler_raw = ckpt.get("scaler")
-        if scaler_raw is None:
-            raise KeyError("Checkpoint missing required key: 'scaler'")
-        scaler = cast(Any, scaler_raw)
-
-        state_raw = ckpt.get("model_state_dict")
-        if state_raw is None:
-            raise KeyError("Checkpoint missing required key: 'model_state_dict'")
-        state = cast(Mapping[str, Any], state_raw)
+        feature_names: List[str] = ckpt["feature_names"]
+        scaler = ckpt["scaler"]
+        input_dim = int(ckpt["input_dim"])
 
         model = AerisFullStructure(input_dim=input_dim)
-        model.load_state_dict(state)
+        model.load_state_dict(ckpt["model_state_dict"])
         model.to(cfg.device)
+        model.eval()
 
         super().__init__(cfg=cfg, model=model)
 
@@ -108,20 +72,16 @@ def __init__(self, cfg: Config):
         self.higher_is_better = {"accuracy": False, "loss": False}
 
         # ----- data loaders  -------------------------------------
-        X, y = load_datasets(cfg.data.path, cfg.data.name, feature_names)
-        # X shape: (n_samples, 1, 245) y shape: (n_samples,1)
-
-        # apply scaler if present
-        if scaler is not None:
-            try:
-                X = scaler.transform(X)
-            except Exception:
-                pass
-        with torch.no_grad():
-            X_scaled = torch.FloatTensor(X).to(cfg.device)
-        y_raw = torch.tensor(y, dtype=torch.float32)
-
-        self.windows = split_into_windows(X_scaled, y_raw, cfg.train.batch_size)
+        X, y = load_datasets(cfg.data.path, cfg.data.name, feature_names, input_dim)
+        # X shape: (n_samples, 245) y shape: (n_samples,1)
+
+        # scale (must match training)
+        X_scaled = scaler.transform(X).astype(np.float32)
+        X_tensor = torch.tensor(X_scaled, dtype=torch.float32)
+        y_tensor = torch.tensor(y, dtype=torch.float32)
+
+        self.windows = split_into_windows(X_tensor, y_tensor, cfg.train.batch_size)
+        #print(f"Prepared {len(self.windows)} time windows for streaming. Each window has ~{self.windows[0][0].shape[0]} samples.")
 
         # ----- streaming state -----------------------------------------------
         self.window_idx: int = 0
diff --git a/examples/aeris/utils.py b/examples/aeris/utils.py
index 78cb86a..d02eeb3 100644
--- a/examples/aeris/utils.py
+++ b/examples/aeris/utils.py
@@ -11,7 +11,7 @@
 import os
 import glob
 import re
-from typing import Dict, List, Tuple, Any
+from typing import Dict, List, Tuple, Any, Optional
 
 import numpy as np
 import pandas as pd
@@ -69,26 +69,6 @@ def _parse_formula(s: str) -> Dict[str, float]:
     return comp
 
 
-def _apply_df_parse_formula_num(val):
-    try:
-        if pd.isna(val):
-            return None
-        parsed = _parse_formula(str(val))
-        return int(sum(parsed.values()))
-    except Exception:
-        return None
-
-
-def _apply_df_parse_formula_str(val):
-    try:
-        if pd.isna(val):
-            return None
-        parsed = _parse_formula(str(val))
-        return "".join(f"{k}{v}" for k, v in sorted(parsed.items()))
-    except Exception:
-        return None
-
-
 def _parse_structure_string(struct_str: str) -> Dict[str, float]:
     # minimal lattice extractor (compatible with training utils)
     result = {
@@ -134,96 +114,144 @@ def _parse_structure_string(struct_str: str) -> Dict[str, float]:
     return result
 
 
-def _build_feature_vector(
-    composition: str, features: Dict, feature_names: List[str]
-) -> np.ndarray:
-    comp = _parse_formula(composition)
-    total_atoms = float(sum(comp.values()))
-    # prepare composition fractions
-    elem_frac = {k: v / total_atoms for k, v in comp.items()}
-
-    # parse structure if string/dict
-    struct_vals = {}
-    if features is not None:
-        for feature in features:
-            if feature not in feature_names:
-                continue
-            struct_vals[feature] = features[feature]
-        if "structure" in features:
-            parsed_struct = _parse_structure_string(features["structure"])
-            struct_vals.update(parsed_struct)
-
-    # magpie
-    feature_calculators = MultipleFeaturizer(
-        [
-            cf.Stoichiometry(),
-            cf.ElementProperty.from_preset("magpie"),
-            cf.ValenceOrbital(props=["avg"]),
-            cf.IonProperty(fast=True),
-        ]
-    )
-
-    comp_obj = Composition(composition)
-    data = pd.DataFrame([{"comp_obj": comp_obj, "composition_reduced": composition}])
-
-    # Calculate Magpie features.
-    # IMPORTANT: when running under MPI, do NOT let matminer spawn multiprocessing pools
-    # inside each rank (oversubscription/hangs). Force single-process.
-    magpie_features_dict = {}
+# -----------------------------
+# Build X,y in *checkpoint feature order*
+# -----------------------------
+# optional numeric columns (if present in CSV) that we will include as features
+OPTIONAL_NUMERIC_COLS = [
+    'density_atomic', 'CN_max', 'CN_min', 'CN_avg',
+    # add more if you know they exist & are useful
+]
+
+def _make_magpie_featurizer() -> MultipleFeaturizer:
+    return MultipleFeaturizer([
+        cf.Stoichiometry(),
+        cf.ElementProperty.from_preset("magpie"),
+        cf.ValenceOrbital(props=['avg']),
+        cf.IonProperty(fast=True),
+    ])
+
+def _compute_magpie_df(compositions: pd.Series) -> pd.DataFrame:
+    featurizer = _make_magpie_featurizer()
+
+    comp_objs = []
+    for s in compositions.astype(str).tolist():
+        try:
+            comp_objs.append(Composition(s))
+        except Exception:
+            comp_objs.append(None)
+
+    base = pd.DataFrame({"comp_obj": comp_objs}, index=compositions.index)
+
     try:
-        # Some matminer versions support n_jobs; if yours does, keep it at 1.
-        magpie_features = feature_calculators.featurize_dataframe(
-            data, col_id="comp_obj", ignore_errors=True, pbar=False, n_jobs=1
+        feat_df = featurizer.featurize_dataframe(
+            base, col_id="comp_obj", ignore_errors=True, pbar=False, n_jobs=1
         )
-        magpie_features_dict = magpie_features.iloc[0].to_dict()
-    except Exception:
+    except TypeError:
         try:
-            feature_calculators.set_n_jobs(1)
-            feats = feature_calculators.featurize_many([Composition(composition)])
-            magpie_features = pd.DataFrame(feats)
-            magpie_features.index = [0]
-            magpie_features_dict = magpie_features.iloc[0].to_dict()
-        except Exception as e:
-            print("Magpie featurizer failed, falling back to empty features:", repr(e))
-            magpie_features_dict = {}
-
-    vec = np.zeros(len(feature_names), dtype=np.float32)
-    for i, name in enumerate(feature_names):
-        # elemental features (assume single element name)
-        if re.match(r"^[A-Z][a-z]?$", name) and name in elem_frac:
-            vec[i] = float(elem_frac.get(name, 0.0))
-            continue
-
-        # structural features
-        if name in struct_vals:
-            vec[i] = float(struct_vals[name])
-            continue
-
-        # magpie features
-        if name in magpie_features_dict:
-            vec[i] = float(magpie_features_dict[name])
-            continue
-
-        # try numeric keys in struct_vals
-        val = struct_vals.get(name)
-        if val is None:
-            v = struct_vals.get(name, 0.0)
+            featurizer.set_n_jobs(1)
+        except Exception:
+            pass
+        feat_df = featurizer.featurize_dataframe(
+            base, col_id="comp_obj", ignore_errors=True, pbar=False
+        )
+
+    feat_df = feat_df.drop(columns=[c for c in feat_df.columns if c == "comp_obj"], errors="ignore")
+    return feat_df
+
+def _build_X_y_in_ckpt_order(
+    df: pd.DataFrame,
+    feature_names: List[str],
+    target_col: Optional[str],
+) -> Tuple[np.ndarray, Optional[np.ndarray]]:
+    required = ["composition", "structure"]
+    for c in required:
+        if c not in df.columns:
+            raise KeyError(f"Missing required column '{c}'")
+
+    if target_col is not None and target_col not in df.columns:
+        raise KeyError(f"Missing target column '{target_col}'")
+
+    magpie_df = _compute_magpie_df(df["composition"])
+    n = len(df)
+    X = np.zeros((n, len(feature_names)), dtype=np.float32)
+    y: Optional[np.ndarray] = None
+    if target_col is not None:
+        y = np.zeros((n, 1), dtype=np.float32)
+
+    df2 = df.reset_index(drop=True)
+
+    for i, row in df2.iterrows():
+        comp_str = str(row["composition"])
+
+        # 1) element fractions
+        try:
+            parsed = _parse_formula(comp_str)
+            total = float(sum(parsed.values())) if parsed else 0.0
+        except Exception:
+            parsed, total = {}, 0.0
+
+        elem_frac: Dict[str, float] = {}
+        if total > 0:
+            for el, cnt in parsed.items():
+                elem_frac[el] = float(cnt) / total
+
+        # 2) structure features
+        struct_vals = _parse_structure_string(row.get("structure"))
+
+        # 3) optional numeric cols
+        opt_vals: Dict[str, float] = {}
+        for c in OPTIONAL_NUMERIC_COLS:
+            if c in df2.columns:
+                v = row.get(c)
+                try:
+                    opt_vals[c] = float(v)
+                except Exception:
+                    opt_vals[c] = np.nan
+
+        # 4) magpie row
+        magpie_row = magpie_df.iloc[i].to_dict()
+
+        # single lookup dict, then assemble in EXACT feature_names order
+        value_by_name: Dict[str, float] = {}
+        for el, frac in elem_frac.items():
+            value_by_name[el] = float(frac)
+        for k, v in struct_vals.items():
+            try:
+                value_by_name[k] = float(v)
+            except Exception:
+                pass
+        for k, v in opt_vals.items():
+            try:
+                value_by_name[k] = float(v)
+            except Exception:
+                pass
+        for k, v in magpie_row.items():
             try:
-                vec[i] = float(v)
+                value_by_name[k] = float(v)
             except Exception:
-                vec[i] = 0.0
-        else:
+                pass
+
+        X[i, :] = np.array([value_by_name.get(name, 0.0) for name in feature_names], dtype=np.float32)
+
+        if y is not None:
             try:
-                vec[i] = float(val)
+                y[i, 0] = float(row[target_col])  # type: ignore[arg-type]
             except Exception:
-                vec[i] = 0.0
+                y[i, 0] = np.nan
 
-    # Return a 1D feature vector (D,) instead of (1, D)
-    vec = np.nan_to_num(vec, nan=0.0, posinf=1e6, neginf=-1e6)
-    return vec
+    X = np.nan_to_num(X, nan=0.0, posinf=1e6, neginf=-1e6).astype(np.float32)
 
+    if y is not None:
+        y = y.astype(np.float32)
+        # drop rows where y is nan
+        mask = ~np.isnan(y[:, 0])
+        X = X[mask]
+        y = y[mask]
+    return X, y
 
-def load_datasets(data_path: str, dataset_name: str, feature_names: List[str]):
+
+def load_datasets(data_path: str, dataset_name: str, feature_names: List[str], input_dim: int) -> Tuple[np.ndarray, Optional[np.ndarray]]:
     """Load the dataset used by the model.
 
     This function attempts to *prefer* loading the exact columns listed in
@@ -253,127 +281,14 @@ def load_datasets(data_path: str, dataset_name: str, feature_names: List[str]):
         dfs.append(pd.read_csv(file_path, low_memory=False))
     dataset: pd.DataFrame = pd.concat(dfs, ignore_index=True)
 
-    # Ensure target present
-    if "formation_energy_per_atom" not in dataset.columns:
-        raise KeyError("Required target column 'formation_energy_per_atom' not found in dataset")
-
-    # Drop rows missing target
-    dataset = dataset.dropna(subset=["formation_energy_per_atom"]).copy()
-
-    # If all feature_names are present as columns, take that branch (preferred)
-    all_present = all((fn in dataset.columns) for fn in feature_names)
-
-    if all_present:
-        # Select columns in the exact saved order
-        X = dataset[feature_names].to_numpy(dtype=np.float32)
-
-        # Replace infs / NaNs in numeric columns with column means (same as training)
-        # (do not change dtype or drop rows here; keep alignment with model)
-        numeric_mask = np.isfinite(X)
-        # For each column replace non-finite with column mean (computed over finite rows)
-        col_means = np.nanmean(np.where(np.isfinite(X), X, np.nan), axis=0)
-        # Where a column is completely NaN/inf, set mean to 0.0
-        col_means = np.where(np.isnan(col_means), 0.0, col_means)
-        inds = np.where(~np.isfinite(X))
-        if inds[0].size > 0:
-            X[inds] = np.take(col_means, inds[1])
-
-    else:
-        # Fall back to building feature vectors row-by-row using the helper
-        # This creates the same ordering as feature_names when possible (elemental
-        # names are interpreted by _build_feature_vector).
-        X_rows = []
-        # Build a minimal features dict per row (this mirrors training's inputs)
-        for _, row in dataset.iterrows():
-            comp = row.get("composition_reduced", row.get("composition", None))
-            features = {
-                "composition": row.get("composition", None),
-                "structure": row.get("structure", None),
-                "spacegroup_number": row.get("spacegroup_number", None),
-                "density_atomic": row.get("density_atomic", None),
-                "CN_max": row.get("CN_max", None),
-                "CN_min": row.get("CN_min", None),
-                "CN_avg": row.get("CN_avg", None),
-            }
-            try:
-                vec = _build_feature_vector(comp, features, feature_names)
-            except Exception:
-                # on failure, append zeros to avoid mismatched shapes
-                vec = np.zeros(len(feature_names), dtype=np.float32)
-            X_rows.append(vec)
-        X = np.vstack(X_rows).astype(np.float32)
-
-        # Replace any remaining nan/inf with column means
-        col_means = np.nanmean(np.where(np.isfinite(X), X, np.nan), axis=0)
-        col_means = np.where(np.isnan(col_means), 0.0, col_means)
-        inds = np.where(~np.isfinite(X))
-        if inds[0].size > 0:
-            X[inds] = np.take(col_means, inds[1])
-
-    # Prepare target vector shape (N,)
-    y = dataset["formation_energy_per_atom"].to_numpy(dtype=np.float32).reshape(-1, 1)
-    assert X.shape[0] == y.shape[0], "Feature matrix and target vector must have same number of rows"
-
-    return X, y
-
-
-def load_datasets2(data_path: str, dataset_name: str, feature_names: List[str]):
-    """Load the dataset that will be parsed, return features and ground truth.
-
-    Parameters
-    ----------
-    data_path:
-        Directory containing the datasets.
-    dataset_name:
-        The name or regular expression for the datasets
-    feature_names:
-        The features used by the model for prediction
-
-    Returns
-    -------
-    input festures, output target values
-    """
-    dfs = []
-    dataset_pattern = os.path.join(data_path, dataset_name)
-    dataset_files: List[str] = glob.glob(dataset_pattern)
-    if not dataset_files:
-        raise FileNotFoundError(f"No dataset files matched pattern: {dataset_pattern}")
-    for file_path in dataset_files:
-        dfs.append(pd.read_csv(file_path, low_memory=False))
-    dataset: pd.DataFrame = pd.concat(dfs, ignore_index=True)
+    target_col = 'formation_energy_per_atom'
+    X_raw, y = _build_X_y_in_ckpt_order(dataset, feature_names=feature_names, target_col=target_col)
+    #print("Prepared X:", X_raw.shape, "y:", None if y is None else y.shape, "num_features:", len(feature_names))
 
-    # Filter all entries that do not have a target value
-    dataset = dataset.dropna(subset=["formation_energy_per_atom"]).copy()
-
-    # Replace NaN/+inf/-inf in numeric columns (keep DataFrame type)
-    num_cols = dataset.select_dtypes(include=[np.number]).columns
-    dataset[num_cols] = dataset[num_cols].replace([np.inf, -np.inf], np.nan).fillna(0.0)
-
-    y = (
-        dataset["formation_energy_per_atom"]
-        .values.astype(np.float32)
-        .reshape(-1, 1)
-    )
-    X = []
-    for _, row in dataset.iterrows():
-        composition = row["composition_reduced"]
-        features = {
-            "composition": row["composition"],
-            "structure": row["structure"],
-            "spacegroup_number": row["spacegroup_number"],
-            "density_atomic": row["density_atomic"],
-            "CN_max": row["CN_max"],
-            "CN_min": row["CN_min"],
-            "CN_avg": row["CN_avg"],
-        }
-        X.append(_build_feature_vector(composition, features, feature_names))
-
-    print("X shape:", np.array(X).shape, "y shape:", y.shape)
-    assert len(X) == len(y), (
-        "The feature and target vectors do not have the same lenght"
-    )
-    return X, y
+    if X_raw.shape[1] != input_dim:
+        raise ValueError(f"Checkpoint input_dim={input_dim} but built X has {X_raw.shape[1]} features.")
 
+    return X_raw, y
 
 # Default number of samples per time window.  Can be overridden by the caller.
 DEFAULT_WINDOW_SIZE: int = 100

From cc260a13306558cedcef183a67500bae5faf66c7 Mon Sep 17 00:00:00 2001
From: Ana Gainaru <ana.gainaru@gmail.com>
Date: Wed, 4 Mar 2026 18:10:22 -0500
Subject: [PATCH 5/9] wip

---
 examples/aeris/aeris.toml | 17 +++++++++--------
 examples/aeris/model.py   |  6 ++++--
 examples/aeris/utils.py   |  2 +-
 3 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/examples/aeris/aeris.toml b/examples/aeris/aeris.toml
index 7d16c7c..99bc171 100644
--- a/examples/aeris/aeris.toml
+++ b/examples/aeris/aeris.toml
@@ -2,7 +2,7 @@
 seed = 42
 device = "auto"
 multi_gpu = false
-verbosity = "INFO"
+verbosity = "INFO:2"
 
 [model]
 name = "aeris_model.pt"
@@ -11,12 +11,13 @@ pretrained_path = "examples/aeris/model"
 [data]
 name = "aeris_dataset.csv"
 path = "examples/aeris/data"
+batch_size = 1
 
 [train]
-batch_size = 500
+batch_size = 256
 num_workers = 4
-init_lr = 1e-6
-max_iter = 4000
+init_lr = 1e-3
+max_iter = 100
 grad_accumulation_steps = 1
 
 [continual_learning]
@@ -36,15 +37,15 @@ kfac_ema_decay = 0.95
 
 [drift_detection]
 detector_name = "ADWINDetector"
-detection_interval = 10
+detection_interval = 1
 aggregation = "mean"
 metric_index = 0
 reset_after_learning = false
-max_stream_updates = 20
+max_stream_updates = 250
 
 # ADWIN hyperparameters
-adwin_delta = 0.002
-adwin_minor_threshold = 0.3
+adwin_delta = 0.2
+adwin_minor_threshold = 0.1
 adwin_moderate_threshold = 0.6
 
 [logging]
diff --git a/examples/aeris/model.py b/examples/aeris/model.py
index ed1c767..4de7fc4 100644
--- a/examples/aeris/model.py
+++ b/examples/aeris/model.py
@@ -80,8 +80,8 @@ def __init__(self, cfg: Config):
         X_tensor = torch.tensor(X_scaled, dtype=torch.float32)
         y_tensor = torch.tensor(y, dtype=torch.float32)
 
-        self.windows = split_into_windows(X_tensor, y_tensor, cfg.train.batch_size)
-        #print(f"Prepared {len(self.windows)} time windows for streaming. Each window has ~{self.windows[0][0].shape[0]} samples.")
+        self.windows = split_into_windows(X_tensor, y_tensor)
+        print(f"Prepared {len(self.windows)} time windows for streaming. Each window has ~{self.windows[0][0].shape[0]} samples.")
 
         # ----- streaming state -----------------------------------------------
         self.window_idx: int = 0
@@ -167,6 +167,8 @@ def update_data_stream(self) -> None:
 
         ds_train = TensorDataset(X_w[:n_train], y_w[:n_train])
         ds_val = TensorDataset(X_w[n_train:], y_w[n_train:])
+        #print(f"Window {self.window_idx}: {n_train} train samples, {n_val} val samples.")
+        #print(len(ds_train), len(ds_val))
 
         bs = self.cfg.train.batch_size
         nw = self.cfg.train.num_workers
diff --git a/examples/aeris/utils.py b/examples/aeris/utils.py
index d02eeb3..dc86135 100644
--- a/examples/aeris/utils.py
+++ b/examples/aeris/utils.py
@@ -291,7 +291,7 @@ def load_datasets(data_path: str, dataset_name: str, feature_names: List[str], i
     return X_raw, y
 
 # Default number of samples per time window.  Can be overridden by the caller.
-DEFAULT_WINDOW_SIZE: int = 100
+DEFAULT_WINDOW_SIZE: int = 10
 
 def split_into_windows(
     X: Tensor,

From 84c28252afaebd27df23ce9e1e07ad25b5d12935 Mon Sep 17 00:00:00 2001
From: Ana Gainaru <againaru@pn2612783.lanl.gov>
Date: Wed, 2 Sep 2026 11:41:15 -0600
Subject: [PATCH 6/9] updates to the aeris model harness

---
 examples/aeris/aeris.toml |  49 ++++---
 examples/aeris/model.py   | 245 ++++++++++++++++++++++++++++++---
 examples/aeris/utils.py   | 283 ++++++++++----------------------------
 examples/mnist/mnist.toml |   2 +-
 examples/utils.py         |   2 +-
 5 files changed, 335 insertions(+), 246 deletions(-)

diff --git a/examples/aeris/aeris.toml b/examples/aeris/aeris.toml
index 99bc171..31a4c2b 100644
--- a/examples/aeris/aeris.toml
+++ b/examples/aeris/aeris.toml
@@ -1,31 +1,36 @@
+[root]
 # aeris.toml  —  AERIS continuous-learning
 seed = 42
 device = "auto"
 multi_gpu = false
-verbosity = "INFO:2"
+verbosity = "INFO:1"
 
 [model]
-name = "aeris_model.pt"
+name = "aeris_init.pt"
 pretrained_path = "examples/aeris/model"
+max_ckpts = 100
+ckpts_path = "output/aeris/"
 
 [data]
-name = "aeris_dataset.csv"
-path = "examples/aeris/data"
-batch_size = 1
+name = "aeris"
+path = "examples/aeris/data/aeris_dataset.csv"
+batch_size = 16
 
 [train]
-batch_size = 256
+batch_size = 300
 num_workers = 4
 init_lr = 1e-3
 max_iter = 100
 grad_accumulation_steps = 1
 
 [continual_learning]
-update_mode = "base"
+update_mode = "ewc_online"
+mix_historic_data = true
 
 # JVP regularization (used when update_mode = "jvp_reg")
-jvp_lambda = 10
-jvp_deltax_norm = 1
+jvp_rho_theta = 0.05
+jvp_rho_x = 1.0
+jvp_data_sign = 1.0
 
 # EWC (used when update_mode = "ewc_online")
 ewc_lambda = 1000.0
@@ -36,18 +41,30 @@ kfac_lambda = 1e-2
 kfac_ema_decay = 0.95
 
 [drift_detection]
-detector_name = "ADWINDetector"
-detection_interval = 1
+detector_name = "KSWINDetector"
+detection_interval = 2
 aggregation = "mean"
-metric_index = 0
-reset_after_learning = false
-max_stream_updates = 250
+metric_index = 1
+reset_after_learning = true
+max_stream_updates = 84
 
-# ADWIN hyperparameters
-adwin_delta = 0.2
+# --- ADWINDetector ---
+adwin_delta = 0.9
 adwin_minor_threshold = 0.1
 adwin_moderate_threshold = 0.6
 
+# --- KSWINDetector ---
+kswin_alpha = 0.05
+kswin_window_size = 100
+kswin_stat_size = 30
+
+# --- PageHinkleyDetector ---
+ph_min_instances = 30
+ph_delta = 0.005
+ph_threshold = 0.6
+ph_alpha = 0.9999
+
 [logging]
 backend = "wandb"
 experiment_name = "aeris-cl"
+# mlflow_tracking_uri = "http://127.0.0.1:5000"
diff --git a/examples/aeris/model.py b/examples/aeris/model.py
index 4de7fc4..a2f284c 100644
--- a/examples/aeris/model.py
+++ b/examples/aeris/model.py
@@ -4,15 +4,21 @@
 This harness wraps a 8-layer neural network trained to predict enthalpy per atom from a given fuel material."""
 
 import gc
+import os
+import math
+from examples.cifar import model
 import torch
 import numpy as np
 from typing import Tuple, Optional, List, Any, Mapping, cast
 from torch import nn, Tensor
 from torch.optim import Optimizer
+from pathlib import Path
 from torch.utils.data import DataLoader, ConcatDataset, TensorDataset
 
-from model.torch_model_harness import BaseModelHarness
-from config.configuration import Config
+from apeiron.model.torch_model_harness import BaseModelHarness
+from apeiron.config.configuration import Config
+
+from apeiron.evaluation.metrics import accuracy
 
 from examples.aeris.utils import (
     load_datasets,
@@ -23,7 +29,7 @@
 
 # Aeris model architecture used for prediction
 class AerisFullStructure(nn.Module):
-    def __init__(self, input_dim, dropout=0.3):
+    def __init__(self, input_dim, dropout=0.1):
         super().__init__()
         first_layer = min(1024, max(512, input_dim * 2))
         self.layers = nn.Sequential(
@@ -41,10 +47,97 @@ def __init__(self, input_dim, dropout=0.3):
     def forward(self, x):
         return self.layers(x)
 
+    def train(self, mode: bool = True):
+        """Set training mode, but always keep BatchNorm layers in eval mode.
+
+        The ``jvp_reg`` continual-learning updater runs this model through
+        ``torch.func.jvp``/``grad`` via ``functional_call``. In training mode
+        ``BatchNorm1d`` performs an in-place ``num_batches_tracked.add_(1)`` on a
+        captured buffer, which functorch transforms forbid. Keeping BatchNorm in
+        eval mode (frozen running stats) both avoids that crash and prevents the
+        small, drifted CL batches from corrupting the normalization statistics.
+        Dropout still follows ``mode`` normally.
+        """
+        super().train(mode)
+        for m in self.modules():
+            if isinstance(m, nn.BatchNorm1d):
+                m.eval()
+        return self
+
 
 # Fraction of each time window reserved for validation
 _VAL_FRACTION: float = 0.2
 
+# Reference harness checkpoints (looked up in cfg.model.pretrained_path) used to
+# recover feature_names/scaler when the selected checkpoint is a bare state_dict
+# -- e.g. after_cl.pt, which the CL run saves as weights only. aeris_drift_init.pt
+# is the model CL starts from, so its scaler matches the preprocessing used
+# during the run; the rest are fallbacks.
+_REFERENCE_CKPTS = (
+    "aeris_drift_init.pt", "aeris_init.pt", "aeris_full.pt", "aeris_infer.pt",
+)
+
+
+def _is_harness_ckpt(obj: Any) -> bool:
+    """True for a full harness checkpoint (has metadata, not just weights)."""
+    return isinstance(obj, dict) and "model_state_dict" in obj and "feature_names" in obj
+
+
+def _infer_input_dim(state_dict: Mapping[str, Tensor]) -> Optional[int]:
+    """Read input_dim from the first Linear layer's weight (out, in)."""
+    w = state_dict.get("layers.0.weight")
+    return int(w.shape[1]) if w is not None else None
+
+
+def _resolve_checkpoint(
+    ckpt: Any, pretrained_path: str, device: str, ckpt_name: str
+) -> Tuple[Mapping[str, Tensor], List[str], Any, int]:
+    """Normalize a loaded checkpoint into (state_dict, feature_names, scaler, input_dim).
+
+    Accepts either:
+      * a full harness checkpoint (dict with model_state_dict / feature_names /
+        scaler / input_dim), or
+      * a bare state_dict of weights (what the CL run writes, e.g. after_cl.pt),
+        optionally wrapped under a "state_dict" key. In that case feature_names
+        and the scaler are borrowed from the first available reference harness
+        checkpoint, and input_dim is inferred from the first layer's weights.
+    """
+    if _is_harness_ckpt(ckpt):
+        return (ckpt["model_state_dict"], ckpt["feature_names"],
+                ckpt["scaler"], int(ckpt["input_dim"]))
+
+    state_dict = (
+        ckpt["state_dict"] if isinstance(ckpt, dict) and "state_dict" in ckpt else ckpt
+    )
+    if not isinstance(state_dict, Mapping) or not all(
+        isinstance(v, torch.Tensor) for v in state_dict.values()
+    ):
+        raise KeyError(
+            f"Checkpoint '{ckpt_name}' is neither a harness checkpoint (missing "
+            "'feature_names') nor a recognizable state_dict of weights."
+        )
+
+    for ref_name in _REFERENCE_CKPTS:
+        if ref_name == ckpt_name:
+            continue
+        ref_path = os.path.join(pretrained_path, ref_name)
+        if os.path.exists(ref_path):
+            ref = torch.load(ref_path, map_location=device, weights_only=False)
+            if _is_harness_ckpt(ref):
+                input_dim = _infer_input_dim(state_dict) or int(ref["input_dim"])
+                print(
+                    f"[AERIS] '{ckpt_name}' is a bare state_dict; borrowing "
+                    f"feature_names/scaler from reference '{ref_name}' "
+                    f"(input_dim={input_dim})."
+                )
+                return state_dict, ref["feature_names"], ref["scaler"], input_dim
+
+    raise KeyError(
+        f"'{ckpt_name}' is a bare state_dict but no reference harness checkpoint "
+        f"({', '.join(_REFERENCE_CKPTS)}) was found in {pretrained_path} to supply "
+        "feature_names/scaler. Add one, or point [model].name at a full checkpoint."
+    )
+
 
 class AERIS(BaseModelHarness):
     """
@@ -56,54 +149,136 @@ def __init__(self, cfg: Config):
         ckpt = load_pretrained_model(
             cfg.model.pretrained_path, cfg.model.name, device=cfg.device
         )
-        feature_names: List[str] = ckpt["feature_names"]
-        scaler = ckpt["scaler"]
-        input_dim = int(ckpt["input_dim"])
+        # Accept either a full harness checkpoint or a bare state_dict (e.g.
+        # after_cl.pt from the CL run) -- see _resolve_checkpoint.
+        state_dict, feature_names, scaler, input_dim = _resolve_checkpoint(
+            ckpt, cfg.model.pretrained_path, cfg.device, cfg.model.name
+        )
 
         model = AerisFullStructure(input_dim=input_dim)
-        model.load_state_dict(ckpt["model_state_dict"])
+        model.load_state_dict(state_dict)
         model.to(cfg.device)
         model.eval()
 
         super().__init__(cfg=cfg, model=model)
 
-        # ----- eval metrics (prediction) -------------------------------------
-        self.eval_metrics = {"mae": self.mae_metric(), "loss": self.get_criterion()}
-        self.higher_is_better = {"accuracy": False, "loss": False}
+        self._feature_names = feature_names
+        self._scaler = scaler
+        self._input_dim = input_dim
 
         # ----- data loaders  -------------------------------------
         X, y = load_datasets(cfg.data.path, cfg.data.name, feature_names, input_dim)
-        # X shape: (n_samples, 245) y shape: (n_samples,1)
+        # X shape: (n_samples, 245) y shape: (n_samples, 1)
 
         # scale (must match training)
         X_scaled = scaler.transform(X).astype(np.float32)
         X_tensor = torch.tensor(X_scaled, dtype=torch.float32)
-        y_tensor = torch.tensor(y, dtype=torch.float32)
+        y_tensor = torch.tensor(y, dtype=torch.float32).view(-1, 1)
 
         self.windows = split_into_windows(X_tensor, y_tensor)
         print(f"Prepared {len(self.windows)} time windows for streaming. Each window has ~{self.windows[0][0].shape[0]} samples.")
 
+        # ----- optional base/initial training data ---------------------------
+        # The pre-drift training split the model was originally fit on. When
+        # cfg.data.base_train_path is set it is blended into full retrains via
+        # get_base_train_dataloaders() so the from-scratch model does not forget
+        # the base distribution. Featurized + scaled exactly like the stream.
+        self._base_train_ds: Optional[TensorDataset] = None
+        self._base_val_ds: Optional[TensorDataset] = None
+        base_path = getattr(cfg.data, "base_train_path", "")
+        if base_path:
+            Xb, yb = load_datasets(base_path, cfg.data.name, feature_names, input_dim)
+            Xb_scaled = scaler.transform(Xb).astype(np.float32)
+            Xb_t = torch.tensor(Xb_scaled, dtype=torch.float32)
+            yb_t = torch.tensor(yb, dtype=torch.float32).view(-1, 1)
+            # Shuffle once (seeded) so the val slice isn't a biased tail.
+            gen = torch.Generator().manual_seed(cfg.seed)
+            perm = torch.randperm(Xb_t.shape[0], generator=gen)
+            Xb_t, yb_t = Xb_t[perm], yb_t[perm]
+            n = Xb_t.shape[0]
+            n_val = max(1, int(n * _VAL_FRACTION))
+            n_train = n - n_val
+            self._base_train_ds = TensorDataset(Xb_t[:n_train], yb_t[:n_train])
+            self._base_val_ds = TensorDataset(Xb_t[n_train:], yb_t[n_train:])
+            print(
+                f"Loaded {n} base training samples from {base_path} "
+                "(blended into full retrains)."
+            )
+
+        # ----- eval metrics (prediction) -------------------------------------
+        self._y_var_ref = self._reference_variance(y_tensor)
+        self.eval_metrics = {
+            "mse": self.mse_metric(),
+            "mae": self.get_criterion(),
+            "r2": self.r2_metric(),
+            "nrmse": self.nrmse_metric(),
+        }
+        self.higher_is_better = {
+            "mse": False, "mae": False, "r2": True, "nrmse": False,
+        }
+
         # ----- streaming state -----------------------------------------------
         self.window_idx: int = 0
         self.history_windows: List[Tuple[Tensor, Tensor]] = []
 
         self._cur_train_loader: Optional[DataLoader] = None
         self._cur_val_loader: Optional[DataLoader] = None
+        self._cur_stream_loader: Optional[DataLoader] = None
+
+    def _reference_variance(self, y_stream: Tensor) -> float:
+        """Fixed denominator for R^2 / NRMSE metrics."""
+        y_ref = (
+            self._base_train_ds.tensors[1]
+            if self._base_train_ds is not None
+            else y_stream
+        )
+        var = float(y_ref.float().var(unbiased=False).item())
+        if var <= 0.0:
+            raise ValueError("Reference target variance is 0; R^2/NRMSE undefined.")
+        return var
+
+    def r2_metric(self):
+        """1 - MSE/var_ref against a constant denominator.
+
+        Being affine in MSE means the sample-weighted mean that
+        BaseModelHarness.eval() computes is exactly the pooled R^2 -- which
+        would NOT hold if each batch normalized by its own variance.
+        """
+        var_ref = self._y_var_ref
+
+        def _r2(y_hat: Tensor, y: Tensor) -> Tensor:
+            return 1.0 - torch.mean((y_hat - y) ** 2) / var_ref
+
+        return _r2
+
+    def nrmse_metric(self):
+        """RMSE / std_ref, same reference as r2_metric."""
+        std_ref = math.sqrt(self._y_var_ref)
+
+        def _nrmse(y_hat: Tensor, y: Tensor) -> Tensor:
+            return torch.sqrt(torch.mean((y_hat - y) ** 2)) / std_ref
+
+        return _nrmse
 
     def get_optmizer(self) -> Optimizer:  # noqa: D102  (spelling kept for ABC)
-        return torch.optim.Adam(self.model.parameters(), lr=self.cfg.train.init_lr)
+        weight_decay = 1e-7
+        return torch.optim.AdamW(self.model.parameters(), lr=self.cfg.train.init_lr, weight_decay=weight_decay)
 
-    def get_criterion(self):  # noqa: D102
+    def mse_metric(self):  # noqa: D102
         return nn.MSELoss()
 
-    def mae_metric(self):
+    def get_criterion(self):
         return nn.L1Loss()
 
-    def get_cur_data_loaders(self) -> Tuple[DataLoader, DataLoader]:  # noqa: D102
+    def get_stream_dataloader(self):
+        assert self._cur_stream_loader is not None
+        return self._cur_stream_loader
+
+    def get_train_dataloaders(self) -> Tuple[DataLoader, DataLoader]:  # noqa: D102
         assert self._cur_train_loader is not None and self._cur_val_loader is not None
         return self._cur_train_loader, self._cur_val_loader
 
-    def get_hist_data_loaders(
+    def get_hist_dataloaders(
         self,
     ) -> Tuple[Optional[DataLoader], Optional[DataLoader]]:
         """Return loaders over all previously-seen time windows.
@@ -137,6 +312,33 @@ def get_hist_data_loaders(
             make_loader(ds_hist_val, bs, shuffle=False, num_workers=nw, pin_memory=pin),
         )
 
+    def get_base_train_dataloaders(
+        self,
+    ) -> Tuple[Optional[DataLoader], Optional[DataLoader]]:
+        """Return (train, val) loaders over the initial training split.
+
+        Returns ``(None, None)`` unless ``cfg.data.base_train_path`` was set.
+        Loaders are built on demand (only during a full retrain) so no worker
+        processes are held open for the rest of the run.
+        """
+        if self._base_train_ds is None:
+            return None, None
+
+        bs = self.cfg.train.batch_size
+        nw = self.cfg.train.num_workers
+        pin = torch.cuda.is_available()
+        train_loader = make_loader(
+            self._base_train_ds, bs, shuffle=True, num_workers=nw, pin_memory=pin
+        )
+        val_loader = (
+            make_loader(
+                self._base_val_ds, bs, shuffle=False, num_workers=nw, pin_memory=pin
+            )
+            if self._base_val_ds is not None
+            else None
+        )
+        return train_loader, val_loader
+
     def update_data_stream(self) -> None:
         """Advance to the next chronological time window.
 
@@ -167,8 +369,6 @@ def update_data_stream(self) -> None:
 
         ds_train = TensorDataset(X_w[:n_train], y_w[:n_train])
         ds_val = TensorDataset(X_w[n_train:], y_w[n_train:])
-        #print(f"Window {self.window_idx}: {n_train} train samples, {n_val} val samples.")
-        #print(len(ds_train), len(ds_val))
 
         bs = self.cfg.train.batch_size
         nw = self.cfg.train.num_workers
@@ -181,6 +381,10 @@ def update_data_stream(self) -> None:
             ds_val, bs, shuffle=False, num_workers=nw, pin_memory=pin
         )
 
+        bs = self.cfg.data.batch_size
+        self._cur_stream_loader = make_loader(
+            ds_train, bs, shuffle=True, num_workers=nw, pin_memory=pin
+        )
         self.window_idx += 1
 
     # --------------------------------------------------------------------- #
@@ -193,4 +397,7 @@ def _dispose_current_loaders(self) -> None:
         if self._cur_val_loader is not None:
             del self._cur_val_loader
             self._cur_val_loader = None
+        if self._cur_stream_loader is not None:
+            del self._cur_stream_loader
+            self._cur_stream_loader = None
         gc.collect()
diff --git a/examples/aeris/utils.py b/examples/aeris/utils.py
index dc86135..46c3221 100644
--- a/examples/aeris/utils.py
+++ b/examples/aeris/utils.py
@@ -6,6 +6,13 @@
     <data_dir>/
         dataset.csv            # data that will be parsed by the SIM framework
         aeris_model.pt        # AERIS pre-trained model
+
+Featurization uses the *fast path* and must stay byte-for-byte consistent with
+``examples/aeris/scripts/make_drift_split.py:build_features`` -- the same code
+that produced the training features. 227 of the 233 features are pre-computed
+columns pulled directly from the CSV; the remaining 6 lattice params are parsed
+from the ``structure`` string. No matminer recompute (that path could disagree
+with the pre-computed columns the model was trained on).
 """
 
 import os
@@ -19,10 +26,6 @@
 from torch import Tensor
 from torch.utils.data import DataLoader, Dataset
 
-from pymatgen.core.composition import Composition
-from matminer.featurizers.base import MultipleFeaturizer
-from matminer.featurizers import composition as cf
-
 
 def load_pretrained_model(
     data_path: str, model_name: str, device: str = "cpu"
@@ -59,239 +62,101 @@ def load_pretrained_model(
     return ckpt
 
 
-def _parse_formula(s: str) -> Dict[str, float]:
-    parts = re.findall(r"([A-Z][a-z]?)([0-9]*\.?[0-9]*)", str(s).strip())
-    if not parts:
-        raise ValueError(f"Could not parse formula: {s}")
-    comp: Dict[str, float] = {}
-    for el, num in parts:
-        comp[el] = float(num) if num else 1.0
-    return comp
-
-
-def _parse_structure_string(struct_str: str) -> Dict[str, float]:
-    # minimal lattice extractor (compatible with training utils)
-    result = {
-        "lattice_a": np.nan,
-        "lattice_b": np.nan,
-        "lattice_c": np.nan,
-        "lattice_alpha": np.nan,
-        "lattice_beta": np.nan,
-        "lattice_gamma": np.nan,
-        "volume": np.nan,
-        "density": np.nan,
-        "nsites": np.nan,
-        "spacegroup_number": np.nan,
-    }
-    if struct_str is None:
-        return result
-    s = str(struct_str)
-    abc_pattern = r"abc\s*:\s*([\d.]+)\s+([\d.]+)\s+([\d.]+)"
-    angles_pattern = r"angles\s*:\s*([\d.]+)\s+([\d.]+)\s+([\d.]+)"
-    abc = re.search(abc_pattern, s)
-    ang = re.search(angles_pattern, s)
-    if abc:
-        result["lattice_a"] = float(abc.group(1))
-        result["lattice_b"] = float(abc.group(2))
-        result["lattice_c"] = float(abc.group(3))
-    if ang:
-        result["lattice_alpha"] = float(ang.group(1))
-        result["lattice_beta"] = float(ang.group(2))
-        result["lattice_gamma"] = float(ang.group(3))
-    # try volume
-    vol_match = re.search(r"volume\s*[:=]\s*([\d.]+)", s)
-    if vol_match:
-        result["volume"] = float(vol_match.group(1))
-    dens_match = re.search(r"density\s*[:=]\s*([\d.]+)", s)
-    if dens_match:
-        result["density"] = float(dens_match.group(1))
-    sg_match = re.search(r"spacegroup(?:_number)?\s*[:=]\s*(\d+)", s)
-    if sg_match:
-        result["spacegroup_number"] = int(sg_match.group(1))
-    nsites_match = re.search(r"nsites\s*[:=]\s*(\d+)", s)
-    if nsites_match:
-        result["nsites"] = int(nsites_match.group(1))
-    return result
-
-
 # -----------------------------
-# Build X,y in *checkpoint feature order*
+# Fast-path featurization
+# (mirrors scripts/make_drift_split.py so the harness featurizes inputs exactly
+#  the way the model was trained)
 # -----------------------------
-# optional numeric columns (if present in CSV) that we will include as features
-OPTIONAL_NUMERIC_COLS = [
-    'density_atomic', 'CN_max', 'CN_min', 'CN_avg',
-    # add more if you know they exist & are useful
+LATTICE_KEYS = [
+    "lattice_a", "lattice_b", "lattice_c",
+    "lattice_alpha", "lattice_beta", "lattice_gamma",
 ]
 
-def _make_magpie_featurizer() -> MultipleFeaturizer:
-    return MultipleFeaturizer([
-        cf.Stoichiometry(),
-        cf.ElementProperty.from_preset("magpie"),
-        cf.ValenceOrbital(props=['avg']),
-        cf.IonProperty(fast=True),
-    ])
-
-def _compute_magpie_df(compositions: pd.Series) -> pd.DataFrame:
-    featurizer = _make_magpie_featurizer()
-
-    comp_objs = []
-    for s in compositions.astype(str).tolist():
-        try:
-            comp_objs.append(Composition(s))
-        except Exception:
-            comp_objs.append(None)
-
-    base = pd.DataFrame({"comp_obj": comp_objs}, index=compositions.index)
-
-    try:
-        feat_df = featurizer.featurize_dataframe(
-            base, col_id="comp_obj", ignore_errors=True, pbar=False, n_jobs=1
-        )
-    except TypeError:
-        try:
-            featurizer.set_n_jobs(1)
-        except Exception:
-            pass
-        feat_df = featurizer.featurize_dataframe(
-            base, col_id="comp_obj", ignore_errors=True, pbar=False
-        )
 
-    feat_df = feat_df.drop(columns=[c for c in feat_df.columns if c == "comp_obj"], errors="ignore")
-    return feat_df
+def _parse_lattice(struct_str: Any) -> Dict[str, float]:
+    """Extract the 6 lattice params from a pymatgen structure string."""
+    r = {k: 0.0 for k in LATTICE_KEYS}
+    s = str(struct_str)
+    abc = re.search(r"abc\s*:\s*([\d.]+)\s+([\d.]+)\s+([\d.]+)", s)
+    ang = re.search(r"angles\s*:\s*([\d.]+)\s+([\d.]+)\s+([\d.]+)", s)
+    if abc:
+        r["lattice_a"], r["lattice_b"], r["lattice_c"] = map(float, abc.groups())
+    if ang:
+        r["lattice_alpha"], r["lattice_beta"], r["lattice_gamma"] = map(float, ang.groups())
+    return r
 
-def _build_X_y_in_ckpt_order(
-    df: pd.DataFrame,
-    feature_names: List[str],
-    target_col: Optional[str],
-) -> Tuple[np.ndarray, Optional[np.ndarray]]:
-    required = ["composition", "structure"]
-    for c in required:
-        if c not in df.columns:
-            raise KeyError(f"Missing required column '{c}'")
 
-    if target_col is not None and target_col not in df.columns:
-        raise KeyError(f"Missing target column '{target_col}'")
+def _build_X_fast(df: pd.DataFrame, feature_names: List[str]) -> np.ndarray:
+    """Assemble the ``(N, len(feature_names))`` matrix via the fast path.
 
-    magpie_df = _compute_magpie_df(df["composition"])
+    227 features are pulled directly from pre-computed CSV columns; the 6 lattice
+    params are parsed from the ``structure`` string. Any feature not found is
+    left at 0. Identical to ``make_drift_split.build_features``.
+    """
     n = len(df)
     X = np.zeros((n, len(feature_names)), dtype=np.float32)
-    y: Optional[np.ndarray] = None
-    if target_col is not None:
-        y = np.zeros((n, 1), dtype=np.float32)
-
-    df2 = df.reset_index(drop=True)
-
-    for i, row in df2.iterrows():
-        comp_str = str(row["composition"])
-
-        # 1) element fractions
-        try:
-            parsed = _parse_formula(comp_str)
-            total = float(sum(parsed.values())) if parsed else 0.0
-        except Exception:
-            parsed, total = {}, 0.0
-
-        elem_frac: Dict[str, float] = {}
-        if total > 0:
-            for el, cnt in parsed.items():
-                elem_frac[el] = float(cnt) / total
-
-        # 2) structure features
-        struct_vals = _parse_structure_string(row.get("structure"))
-
-        # 3) optional numeric cols
-        opt_vals: Dict[str, float] = {}
-        for c in OPTIONAL_NUMERIC_COLS:
-            if c in df2.columns:
-                v = row.get(c)
-                try:
-                    opt_vals[c] = float(v)
-                except Exception:
-                    opt_vals[c] = np.nan
-
-        # 4) magpie row
-        magpie_row = magpie_df.iloc[i].to_dict()
-
-        # single lookup dict, then assemble in EXACT feature_names order
-        value_by_name: Dict[str, float] = {}
-        for el, frac in elem_frac.items():
-            value_by_name[el] = float(frac)
-        for k, v in struct_vals.items():
-            try:
-                value_by_name[k] = float(v)
-            except Exception:
-                pass
-        for k, v in opt_vals.items():
-            try:
-                value_by_name[k] = float(v)
-            except Exception:
-                pass
-        for k, v in magpie_row.items():
-            try:
-                value_by_name[k] = float(v)
-            except Exception:
-                pass
-
-        X[i, :] = np.array([value_by_name.get(name, 0.0) for name in feature_names], dtype=np.float32)
-
-        if y is not None:
-            try:
-                y[i, 0] = float(row[target_col])  # type: ignore[arg-type]
-            except Exception:
-                y[i, 0] = np.nan
-
-    X = np.nan_to_num(X, nan=0.0, posinf=1e6, neginf=-1e6).astype(np.float32)
-
-    if y is not None:
-        y = y.astype(np.float32)
-        # drop rows where y is nan
-        mask = ~np.isnan(y[:, 0])
-        X = X[mask]
-        y = y[mask]
-    return X, y
+    col_idx = {f: j for j, f in enumerate(feature_names)}
+
+    present = [f for f in feature_names if f in df.columns]
+    sub = df[present].apply(pd.to_numeric, errors="coerce").to_numpy(np.float32)
+    for k, f in enumerate(present):
+        X[:, col_idx[f]] = sub[:, k]
+
+    if "structure" in df.columns:
+        lat = np.array(
+            [list(_parse_lattice(s).values()) for s in df["structure"].tolist()],
+            dtype=np.float32,
+        )
+        for k, f in enumerate(LATTICE_KEYS):
+            if f in col_idx:
+                X[:, col_idx[f]] = lat[:, k]
+
+    return np.nan_to_num(X, nan=0.0, posinf=1e6, neginf=-1e6)
 
 
-def load_datasets(data_path: str, dataset_name: str, feature_names: List[str], input_dim: int) -> Tuple[np.ndarray, Optional[np.ndarray]]:
+def load_datasets(
+    data_path: str, dataset_name: str, feature_names: List[str], input_dim: int
+) -> Tuple[np.ndarray, Optional[np.ndarray]]:
     """Load the dataset used by the model.
 
-    This function attempts to *prefer* loading the exact columns listed in
-    `feature_names` (in the same order). If those columns are present in the
-    CSV(s), they are used directly (fast, deterministic). If not all feature
-    columns are present, the function falls back to building feature vectors
-    row-by-row using _build_feature_vector to preserve compatibility with older
-    or alternate CSV formats.
+    Features are assembled in the exact ``feature_names`` order via the fast
+    path (pre-computed columns + parsed lattice params), matching how the model
+    was trained. Rows with a missing target are dropped.
 
-    The function returns:
-        X: numpy.ndarray of shape (n_samples, n_features) dtype float32
-        y: numpy.ndarray of shape (n_samples,) dtype float32
+    Returns
+    -------
+    X: numpy.ndarray of shape (n_samples, n_features) dtype float32 (unscaled)
+    y: numpy.ndarray of shape (n_samples, 1) dtype float32
 
     Note: scaling is intentionally NOT applied here. The caller (model harness)
-    will apply the saved scaler from the checkpoint (if any) via scaler.transform().
+    applies the saved scaler from the checkpoint via scaler.transform().
     """
-
-    # collect files
-    dataset_pattern = os.path.join(data_path, dataset_name)
+    dataset_pattern = os.path.join(data_path)
     dataset_files: List[str] = glob.glob(dataset_pattern)
     if not dataset_files:
         raise FileNotFoundError(f"No dataset files matched pattern: {dataset_pattern}")
 
-    # read & concatenate CSV files
-    dfs = []
-    for file_path in dataset_files:
-        dfs.append(pd.read_csv(file_path, low_memory=False))
+    dfs = [pd.read_csv(fp, low_memory=False) for fp in dataset_files]
     dataset: pd.DataFrame = pd.concat(dfs, ignore_index=True)
 
-    target_col = 'formation_energy_per_atom'
-    X_raw, y = _build_X_y_in_ckpt_order(dataset, feature_names=feature_names, target_col=target_col)
-    #print("Prepared X:", X_raw.shape, "y:", None if y is None else y.shape, "num_features:", len(feature_names))
+    target_col = "formation_energy_per_atom"
+    if target_col not in dataset.columns:
+        raise KeyError(f"Missing target column '{target_col}'")
+    dataset = dataset[dataset[target_col].notna()].reset_index(drop=True)
+
+    X = _build_X_fast(dataset, feature_names)
+    y = dataset[target_col].to_numpy(np.float32).reshape(-1, 1)
 
-    if X_raw.shape[1] != input_dim:
-        raise ValueError(f"Checkpoint input_dim={input_dim} but built X has {X_raw.shape[1]} features.")
+    if X.shape[1] != input_dim:
+        raise ValueError(
+            f"Checkpoint input_dim={input_dim} but built X has {X.shape[1]} features."
+        )
+
+    return X, y
 
-    return X_raw, y
 
 # Default number of samples per time window.  Can be overridden by the caller.
-DEFAULT_WINDOW_SIZE: int = 10
+DEFAULT_WINDOW_SIZE: int = 500
 
 def split_into_windows(
     X: Tensor,
diff --git a/examples/mnist/mnist.toml b/examples/mnist/mnist.toml
index 12449b1..d7b288d 100644
--- a/examples/mnist/mnist.toml
+++ b/examples/mnist/mnist.toml
@@ -26,7 +26,7 @@ grad_accumulation_steps = 2
 
 [continual_learning]
 update_mode = "kfac_online"
-mix_historic_data = false
+mix_historic_data = true
 
 # jvp_reg
 jvp_rho_theta = 0.05
diff --git a/examples/utils.py b/examples/utils.py
index 392a845..915d970 100644
--- a/examples/utils.py
+++ b/examples/utils.py
@@ -15,7 +15,7 @@ def get_example(cfg: Config) -> BaseModelHarness:
         from examples.imagenet.model import IMAGENET_VISION
 
         return IMAGENET_VISION(cfg=cfg)
-    elif cfg.data.name == "aeris_dataset.csv":
+    elif cfg.data.name == "aeris":
         from examples.aeris.model import AERIS
 
         return AERIS(cfg=cfg)

From a89e049759fe43b4d21442d2844910f26c2dbdb6 Mon Sep 17 00:00:00 2001
From: Ana Gainaru <ana.gainaru@gmail.com>
Date: Fri, 4 Sep 2026 13:50:05 -0600
Subject: [PATCH 7/9] wip

---
 examples/aeris/aeris.toml | 1 -
 1 file changed, 1 deletion(-)

diff --git a/examples/aeris/aeris.toml b/examples/aeris/aeris.toml
index 31a4c2b..e3489d8 100644
--- a/examples/aeris/aeris.toml
+++ b/examples/aeris/aeris.toml
@@ -1,4 +1,3 @@
-[root]
 # aeris.toml  —  AERIS continuous-learning
 seed = 42
 device = "auto"

From 011e8ea2f60b5d5479c9762c17a1d7c4a21f3142 Mon Sep 17 00:00:00 2001
From: Ana Gainaru <ana.gainaru@gmail.com>
Date: Sat, 5 Sep 2026 07:47:21 -0600
Subject: [PATCH 8/9] Expect the same model structure between the init model
 and the checkpoints of Apeiron

---
 examples/aeris/model.py | 108 ++++++++++------------------------------
 1 file changed, 27 insertions(+), 81 deletions(-)

diff --git a/examples/aeris/model.py b/examples/aeris/model.py
index a2f284c..361953f 100644
--- a/examples/aeris/model.py
+++ b/examples/aeris/model.py
@@ -4,22 +4,17 @@
 This harness wraps a 8-layer neural network trained to predict enthalpy per atom from a given fuel material."""
 
 import gc
-import os
 import math
-from examples.cifar import model
 import torch
 import numpy as np
-from typing import Tuple, Optional, List, Any, Mapping, cast
+from typing import Tuple, Optional, List, Any
 from torch import nn, Tensor
 from torch.optim import Optimizer
-from pathlib import Path
 from torch.utils.data import DataLoader, ConcatDataset, TensorDataset
 
 from apeiron.model.torch_model_harness import BaseModelHarness
 from apeiron.config.configuration import Config
 
-from apeiron.evaluation.metrics import accuracy
-
 from examples.aeris.utils import (
     load_datasets,
     make_loader,
@@ -68,76 +63,6 @@ def train(self, mode: bool = True):
 # Fraction of each time window reserved for validation
 _VAL_FRACTION: float = 0.2
 
-# Reference harness checkpoints (looked up in cfg.model.pretrained_path) used to
-# recover feature_names/scaler when the selected checkpoint is a bare state_dict
-# -- e.g. after_cl.pt, which the CL run saves as weights only. aeris_drift_init.pt
-# is the model CL starts from, so its scaler matches the preprocessing used
-# during the run; the rest are fallbacks.
-_REFERENCE_CKPTS = (
-    "aeris_drift_init.pt", "aeris_init.pt", "aeris_full.pt", "aeris_infer.pt",
-)
-
-
-def _is_harness_ckpt(obj: Any) -> bool:
-    """True for a full harness checkpoint (has metadata, not just weights)."""
-    return isinstance(obj, dict) and "model_state_dict" in obj and "feature_names" in obj
-
-
-def _infer_input_dim(state_dict: Mapping[str, Tensor]) -> Optional[int]:
-    """Read input_dim from the first Linear layer's weight (out, in)."""
-    w = state_dict.get("layers.0.weight")
-    return int(w.shape[1]) if w is not None else None
-
-
-def _resolve_checkpoint(
-    ckpt: Any, pretrained_path: str, device: str, ckpt_name: str
-) -> Tuple[Mapping[str, Tensor], List[str], Any, int]:
-    """Normalize a loaded checkpoint into (state_dict, feature_names, scaler, input_dim).
-
-    Accepts either:
-      * a full harness checkpoint (dict with model_state_dict / feature_names /
-        scaler / input_dim), or
-      * a bare state_dict of weights (what the CL run writes, e.g. after_cl.pt),
-        optionally wrapped under a "state_dict" key. In that case feature_names
-        and the scaler are borrowed from the first available reference harness
-        checkpoint, and input_dim is inferred from the first layer's weights.
-    """
-    if _is_harness_ckpt(ckpt):
-        return (ckpt["model_state_dict"], ckpt["feature_names"],
-                ckpt["scaler"], int(ckpt["input_dim"]))
-
-    state_dict = (
-        ckpt["state_dict"] if isinstance(ckpt, dict) and "state_dict" in ckpt else ckpt
-    )
-    if not isinstance(state_dict, Mapping) or not all(
-        isinstance(v, torch.Tensor) for v in state_dict.values()
-    ):
-        raise KeyError(
-            f"Checkpoint '{ckpt_name}' is neither a harness checkpoint (missing "
-            "'feature_names') nor a recognizable state_dict of weights."
-        )
-
-    for ref_name in _REFERENCE_CKPTS:
-        if ref_name == ckpt_name:
-            continue
-        ref_path = os.path.join(pretrained_path, ref_name)
-        if os.path.exists(ref_path):
-            ref = torch.load(ref_path, map_location=device, weights_only=False)
-            if _is_harness_ckpt(ref):
-                input_dim = _infer_input_dim(state_dict) or int(ref["input_dim"])
-                print(
-                    f"[AERIS] '{ckpt_name}' is a bare state_dict; borrowing "
-                    f"feature_names/scaler from reference '{ref_name}' "
-                    f"(input_dim={input_dim})."
-                )
-                return state_dict, ref["feature_names"], ref["scaler"], input_dim
-
-    raise KeyError(
-        f"'{ckpt_name}' is a bare state_dict but no reference harness checkpoint "
-        f"({', '.join(_REFERENCE_CKPTS)}) was found in {pretrained_path} to supply "
-        "feature_names/scaler. Add one, or point [model].name at a full checkpoint."
-    )
-
 
 class AERIS(BaseModelHarness):
     """
@@ -149,11 +74,19 @@ def __init__(self, cfg: Config):
         ckpt = load_pretrained_model(
             cfg.model.pretrained_path, cfg.model.name, device=cfg.device
         )
-        # Accept either a full harness checkpoint or a bare state_dict (e.g.
-        # after_cl.pt from the CL run) -- see _resolve_checkpoint.
-        state_dict, feature_names, scaler, input_dim = _resolve_checkpoint(
-            ckpt, cfg.model.pretrained_path, cfg.device, cfg.model.name
-        )
+
+        # Expect full harness checkpoint format
+        if not isinstance(ckpt, dict) or "model_state_dict" not in ckpt:
+            raise ValueError(
+                f"Checkpoint '{cfg.model.name}' must be a full harness checkpoint "
+                "with 'model_state_dict', 'feature_names', 'scaler', and 'input_dim'. "
+                "Legacy weight-only checkpoints are no longer supported."
+            )
+
+        state_dict = ckpt["model_state_dict"]
+        feature_names = ckpt["feature_names"]
+        scaler = ckpt["scaler"]
+        input_dim = int(ckpt["input_dim"])
 
         model = AerisFullStructure(input_dim=input_dim)
         model.load_state_dict(state_dict)
@@ -387,6 +320,19 @@ def update_data_stream(self) -> None:
         )
         self.window_idx += 1
 
+    def build_checkpoint_payload(self) -> dict[str, Any]:
+        """Save full harness checkpoint matching load_pretrained_model format.
+
+        Overrides BaseModelHarness to include feature names, scaler, and input_dim
+        so that saved checkpoints can be loaded directly without reference checkpoints.
+        """
+        return {
+            "model_state_dict": self.model.state_dict(),
+            "feature_names": self._feature_names,
+            "scaler": self._scaler,
+            "input_dim": self._input_dim,
+        }
+
     # --------------------------------------------------------------------- #
     # Helpers
     # --------------------------------------------------------------------- #

From b14587aa4b4d7c2cca198e19b4a4cf67ad4f1116 Mon Sep 17 00:00:00 2001
From: Ana Gainaru <ana.gainaru@gmail.com>
Date: Sat, 5 Sep 2026 07:48:29 -0600
Subject: [PATCH 9/9] Allow model harnesses to define how Apeiron saves
 checkpoints

---
 src/apeiron/model/torch_model_harness.py | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/src/apeiron/model/torch_model_harness.py b/src/apeiron/model/torch_model_harness.py
index 07c2c29..1768ee0 100644
--- a/src/apeiron/model/torch_model_harness.py
+++ b/src/apeiron/model/torch_model_harness.py
@@ -222,13 +222,27 @@ def task_diagonals(self) -> List[List[float]]:
     def ckpts_enabled(self) -> bool:
         return self.cfg.model.max_ckpts > 0 and bool(self.cfg.model.ckpts_path)
 
+    def build_checkpoint_payload(self) -> Any:
+        """Build the checkpoint object to save.
+
+        Subclasses can override this to include additional metadata beyond weights
+        (e.g., preprocessing scalers, feature names, architecture parameters)
+        so that saved checkpoints match the format expected by the loader.
+
+        Returns
+        -------
+        By default, returns ``model.state_dict()`` (weights only).
+        """
+        return self.model.state_dict()
+
     def save_ckpt(self, event: int) -> str:
         """Persist model state, evict oldest when over budget."""
         d = Path(self.cfg.model.ckpts_path)
         d.mkdir(parents=True, exist_ok=True)
 
         fname = f"drift_adaptation_{event}.pt"
-        torch.save(self.model.state_dict(), d / fname)
+        payload = self.build_checkpoint_payload()
+        torch.save(payload, d / fname)
         (d / "latest").write_text(fname)
 
         # Guillotine the oldest survivors