Dst Forecasting with Keras 3 — PyTorch Backend

This complete workflow implements the experiment in the Dst Forecasting chapter. The notebook forms gap-safe windows inside fixed year partitions and evaluates a one-hour-ahead forecast against true forecast-origin persistence.

In Colab, select Runtime → Run all; the data cell downloads and verifies only the archived OMNI file when no local checkout is available. To run the example more quickly, set EPOCHS to 1 or 2 in the data-preparation cell.

Runtime dependency check

import importlib.util
import subprocess
import sys
from pathlib import Path

REQUIRED_RUNTIME = {'keras': 'keras', 'torch': 'torch'}
COLAB_EXTRAS = {}
missing_required = [
    package for module, package in REQUIRED_RUNTIME.items()
    if importlib.util.find_spec(module) is None
]
missing_extras = [
    package for module, package in COLAB_EXTRAS.items()
    if importlib.util.find_spec(module) is None
]
if missing_extras and "google.colab" in sys.modules:
    subprocess.check_call(
        [sys.executable, "-m", "pip", "install", "-q", *missing_extras]
    )
    missing_extras = []
if missing_required or missing_extras:
    missing = ", ".join(missing_required + missing_extras)
    raise RuntimeError(
        f"Missing notebook dependencies: {missing}. Locally run "
        "`uv sync --group notebooks`; in Colab restart the runtime if an "
        "installation cell just changed the environment."
    )
print("runtime dependency check passed")

runtime dependency check passed

%matplotlib inline

Keras-to-PyTorch crosswalk

This optional implementation uses the same Torch runtime as the canonical native PyTorch path. Keras compile() selects the optimizer and loss, fit() owns the explicit batch and epoch loop, and callbacks provide high-level training control. Data, splits, budgets, evidence, and scientific conclusions remain aligned with the canonical PyTorch workflow.

Imports and reproducibility

These imports provide the numerical, plotting, and evaluation tools used below. The fixed seed makes repeated runs easier to compare.

import hashlib
import json
import os
import random
from pathlib import Path
from urllib.parse import quote
from urllib.request import urlopen

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.preprocessing import StandardScaler

Resolve the archived dataset

The example uses the hourly OMNI2 archive for 2010–2015. The file is verified before it is read so the two implementations use exactly the same measurements.

DATASET_ID = "dst-omni-2010-2015"
DATA_FILENAME = "omni2_2010-2015.dat"
DATA_RELATIVE_PATH = "data/dst-omni-2010-2015/omni2_2010-2015.dat"
DATA_SHA256 = "18a4ce192bdcc481bdef699a6e11f7f0441b4e933def8dd0c9cd25fc766bcecf"


def file_sha256(path):
    digest = hashlib.sha256()
    with Path(path).open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def resolve_data_file():
    candidates = []
    override = os.getenv("HELIO_DATA_DIR")
    if override:
        root = Path(override).expanduser()
        candidates.extend([root / DATASET_ID / DATA_FILENAME, root / DATA_FILENAME])
    for root in [Path.cwd(), *Path.cwd().parents]:
        candidates.append(root / DATA_RELATIVE_PATH)
    for candidate in candidates:
        if candidate.is_file() and file_sha256(candidate) == DATA_SHA256:
            return candidate

    cache = (
        Path(os.getenv("HELIO_DATA_CACHE", Path.home() / ".cache" / "helio-data-methods"))
        / "datasets"
        / DATASET_ID
        / DATA_FILENAME
    )
    if not cache.is_file() or file_sha256(cache) != DATA_SHA256:
        cache.parent.mkdir(parents=True, exist_ok=True)
        ref = os.getenv("HELIO_DATA_REF", "main")
        url = (
            "https://raw.githubusercontent.com/SavvasRaptis/helio-data-methods/"
            f"{quote(ref, safe='')}/{quote(DATA_RELATIVE_PATH, safe='/')}"
        )
        try:
            with urlopen(url, timeout=60) as response, cache.open("wb") as output:
                output.write(response.read())
        except Exception as exc:
            cache.unlink(missing_ok=True)
            raise RuntimeError(
                "Dst data could not be downloaded. Check network access or set "
                "HELIO_DATA_DIR to the archived data directory."
            ) from exc
    if file_sha256(cache) != DATA_SHA256:
        cache.unlink(missing_ok=True)
        raise ValueError("Dst dataset checksum mismatch; the invalid file was removed.")
    return cache


data_path = resolve_data_file()
print(f"data file: {DATA_FILENAME} (checksum verified)")
print(f"dataset SHA-256: {file_sha256(data_path)}")

data file: omni2_2010-2015.dat (checksum verified)
dataset SHA-256: 18a4ce192bdcc481bdef699a6e11f7f0441b4e933def8dd0c9cd25fc766bcecf

Parse data, form windows, and fit training-only preprocessing

Each sample contains the preceding three hourly values of the solar-wind and Dst inputs. The target is Dst one hour after the latest input, and windows are formed only across contiguous hourly measurements.

HEADERS = [
    "year", "day", "hour", "Bartels", "IMF_spacecraft", "plasma_spacecraft",
    "IMF_av_npoints", "plasma_av_npoints", "av_|B|", "|av_B|",
    "lat_av_B_GSE", "lon_av_B_GSE", "Bx", "By_GSE", "Bz_GSE", "By_GSM",
    "Bz_GSM", "sigma_|B|", "sigma_B", "sigma_Bx", "sigma_By", "sigma_Bz",
    "Tp", "Np", "V_plasma", "phi_V_angle", "theta_V_angle", "Na/Np",
    "P_dyn", "sigma_Tp", "sigma_Np", "sigma_V", "sigma_phi_V",
    "sigma_theta_V", "sigma_Na/Np", "E", "beta", "Ma", "Kp", "R", "Dst",
    "AE", "p_flux_>1MeV", "p_flux_>2MeV", "p_flux_>4MeV",
    "p_flux_>10MeV", "p_flux_>30MeV", "p_flux_>60MeV", "flag", "Ap",
    "f10.7", "PC", "AL", "AU", "M_ms",
]
FILL_VALUES = {
    "av_|B|": 999.9,
    "Bz_GSM": 999.9,
    "V_plasma": 9999.0,
    "Dst": 99999.0,
}
INPUT_COLUMNS = ["V_plasma", "Bz_GSM", "av_|B|", "Dst"]


def read_omni(path):
    frame = pd.read_csv(path, sep=r"\s+", header=None, names=HEADERS)
    frame["timestamp"] = (
        pd.to_datetime(frame["year"].astype(str), format="%Y")
        + pd.to_timedelta(frame["day"] - 1, unit="D")
        + pd.to_timedelta(frame["hour"], unit="h")
    )
    for column, fill_value in FILL_VALUES.items():
        frame.loc[frame[column] == fill_value, column] = np.nan
    return frame


def make_windows(frame, years, history_hours, horizon_hours):
    selected = frame.loc[frame["year"].isin(years)].reset_index(drop=True)
    times = selected["timestamp"].to_numpy(dtype="datetime64[h]")
    values = selected[INPUT_COLUMNS].to_numpy(dtype=np.float32)
    dst = selected["Dst"].to_numpy(dtype=np.float32)
    features, targets, persistence, target_times = [], [], [], []
    for origin in range(history_hours - 1, len(selected) - horizon_hours):
        first = origin - history_hours + 1
        target_index = origin + horizon_hours
        history_times = times[first : origin + 1]
        if not np.all(np.diff(history_times) == np.timedelta64(1, "h")):
            continue
        if times[target_index] - times[origin] != np.timedelta64(horizon_hours, "h"):
            continue
        history = values[first : origin + 1]
        target = dst[target_index]
        if not np.isfinite(history).all() or not np.isfinite(target):
            continue
        features.append(history.reshape(-1))
        targets.append(target)
        persistence.append(history[-1, INPUT_COLUMNS.index("Dst")])
        target_times.append(times[target_index])
    return (
        np.asarray(features, dtype=np.float32),
        np.asarray(targets, dtype=np.float32),
        np.asarray(persistence, dtype=np.float32),
        np.asarray(target_times),
    )


SEED = 42
HISTORY_HOURS = 3
HORIZON_HOURS = 1
EPOCHS = 10  # Reduce to 1 or 2 for a quicker run.
BATCH_SIZE = 128  # Number of hourly windows used for each parameter update.

random.seed(SEED)
np.random.seed(SEED)
# Build each year partition independently so information never crosses a split.
frame = read_omni(data_path)
x_train_raw, y_train, persistence_train, time_train = make_windows(
    frame, range(2010, 2014), HISTORY_HOURS, HORIZON_HOURS
)
x_validation_raw, y_validation, persistence_validation, time_validation = make_windows(
    frame, [2014], HISTORY_HOURS, HORIZON_HOURS
)
x_test_raw, y_test, persistence_test, time_test = make_windows(
    frame, [2015], HISTORY_HOURS, HORIZON_HOURS
)

# Estimate the scaling only from 2010–2013.
scaler = StandardScaler().fit(x_train_raw)
x_train = scaler.transform(x_train_raw).astype(np.float32)
x_validation = scaler.transform(x_validation_raw).astype(np.float32)
x_test = scaler.transform(x_test_raw).astype(np.float32)

split_signature = hashlib.sha256(
    time_test.astype("<M8[h]").astype("<i8").tobytes()
    + np.asarray([HISTORY_HOURS, HORIZON_HOURS], dtype="<i8").tobytes()
).hexdigest()[:16]

assert time_train.max() < time_validation.min() < time_test.min()
print(f"split signature: {split_signature}")
print(
    f"train={len(y_train):,}, validation={len(y_validation):,}, "
    f"test={len(y_test):,}, history={HISTORY_HOURS} h, horizon={HORIZON_HOURS} h"
)

split signature: 58385675268180d1
train=35,043, validation=8,757, test=8,757, history=3 h, horizon=1 h

Inspect coverage and the training target

The time series shows the fixed year partitions, while the histogram summarizes the Dst values available for fitting the model.

fig, axes = plt.subplots(1, 2, figsize=(11, 3.5))
axes[0].plot(frame["timestamp"], frame["Dst"], linewidth=0.5)
axes[0].axvspan(pd.Timestamp("2014-01-01"), pd.Timestamp("2015-01-01"), alpha=0.15)
axes[0].axvspan(pd.Timestamp("2015-01-01"), frame["timestamp"].max(), alpha=0.15)
axes[0].set(title="Archived hourly Dst and fixed year partitions", ylabel="Dst [nT]")
axes[1].hist(y_train, bins=50, alpha=0.8)
axes[1].set(title="Training-target distribution", xlabel="Dst at forecast target [nT]")
plt.tight_layout()
plt.show()

Establish the persistence baseline

For a one-hour forecast, persistence assumes that Dst remains at its value at the forecast origin. This is a demanding and physically meaningful reference for a short-horizon forecast.

def regression_metrics(y_true, y_prediction):
    return {
        "mae": float(mean_absolute_error(y_true, y_prediction)),
        "rmse": float(mean_squared_error(y_true, y_prediction) ** 0.5),
        "r2": float(r2_score(y_true, y_prediction)),
    }


persistence_metrics = regression_metrics(y_test, persistence_test)
print("persistence baseline:", persistence_metrics)

persistence baseline: {'mae': 3.0668036937713623, 'rmse': 4.765770881796388, 'r2': 0.9529839158058167}

Define the Keras model

This small network maps the three-hour input state to one Dst prediction. The two hidden layers retain the structure of the original example.

os.environ["KERAS_BACKEND"] = "torch"
import keras
import torch
from keras import layers

keras.utils.set_random_seed(SEED)
torch.use_deterministic_algorithms(True)  # Prefer repeatable operations when available.
assert keras.backend.backend() == "torch"
# Define the neural network used for the Dst forecast.
model = keras.Sequential(
    [
        keras.Input(shape=(x_train.shape[1],)),
        layers.Dense(50, activation="relu"),
        layers.Dense(30, activation="relu"),
        layers.Dense(1),
    ],
    name="dst_forecast",
)
model.compile(
    optimizer=keras.optimizers.Adam(),  # Adam updates the model weights.
    loss="mse",  # Mean-squared error for the continuous Dst target.
)
model.summary()

Model: "dst_forecast"

┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Layer (type)                    ┃ Output Shape           ┃       Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ dense (Dense)                   │ (None, 50)             │           650 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ dense_1 (Dense)                 │ (None, 30)             │         1,530 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ dense_2 (Dense)                 │ (None, 1)              │            31 │
└─────────────────────────────────┴────────────────────────┴───────────────┘

 Total params: 2,211 (8.64 KB)

 Trainable params: 2,211 (8.64 KB)

 Non-trainable params: 0 (0.00 B)

Train with validation-based early stopping

The model is fitted on 2010–2013 and monitored on 2014. Early stopping restores the state with the lowest validation loss before the 2015 evaluation.

callbacks = [
    keras.callbacks.EarlyStopping(
        monitor="val_loss", patience=3, restore_best_weights=True
    )
]
history = model.fit(
    x_train,
    y_train,
    validation_data=(x_validation, y_validation),
    epochs=EPOCHS,
    batch_size=BATCH_SIZE,
    callbacks=callbacks,
    verbose=2,
)
fig, ax = plt.subplots(figsize=(6, 3.5))
ax.plot(history.history["loss"], marker="o", label="training")
ax.plot(history.history["val_loss"], marker="o", label="validation")
ax.set(title="Mean-squared error", xlabel="Epoch", ylabel="MSE")
ax.legend()
ax.grid(alpha=0.25)
plt.show()

Epoch 1/10

  torch._foreach_mul_(m_list, self.beta_1)

274/274 - 13s - 47ms/step - loss: 110.9146 - val_loss: 19.5035

Epoch 2/10

274/274 - 13s - 46ms/step - loss: 14.4926 - val_loss: 12.7329

Epoch 3/10

274/274 - 12s - 45ms/step - loss: 10.6775 - val_loss: 10.6261

Epoch 4/10

274/274 - 12s - 44ms/step - loss: 9.6549 - val_loss: 10.2272

Epoch 5/10

274/274 - 15s - 55ms/step - loss: 9.1338 - val_loss: 9.7334

Epoch 6/10

274/274 - 13s - 47ms/step - loss: 8.9038 - val_loss: 9.9476

Epoch 7/10

274/274 - 12s - 45ms/step - loss: 8.7841 - val_loss: 9.3828

Epoch 8/10

274/274 - 13s - 46ms/step - loss: 8.6363 - val_loss: 9.3690

Epoch 9/10

274/274 - 12s - 46ms/step - loss: 8.6212 - val_loss: 9.1956

Epoch 10/10

274/274 - 12s - 43ms/step - loss: 8.5584 - val_loss: 9.0904

Evaluate the untouched test year

The final metrics use 2015 only after training and model selection are complete. Persistence skill is reported without assuming that the neural model must win.

predictions = model.predict(x_test, batch_size=BATCH_SIZE, verbose=0).reshape(-1)

model_metrics = regression_metrics(y_test, predictions)
persistence_skill = 1.0 - (
    model_metrics["rmse"] ** 2 / persistence_metrics["rmse"] ** 2
)
print("model:", model_metrics)
print(f"persistence skill: {persistence_skill:.4f}")
print(
    "HELIO_RESULT "
    + json.dumps(
        {
            "split_signature": split_signature,
            "model_rmse": model_metrics["rmse"],
            "persistence_rmse": persistence_metrics["rmse"],
            "persistence_skill": persistence_skill,
            "prediction_shape": list(predictions.shape),
        },
        sort_keys=True,
    )
)
assert predictions.shape == y_test.shape
assert np.isfinite(predictions).all()

model: {'mae': 2.479313611984253, 'rmse': 3.6077599813637664, 'r2': 0.973056435585022}
persistence skill: 0.4269
HELIO_RESULT {"model_rmse": 3.6077599813637664, "persistence_rmse": 4.765770881796388, "persistence_skill": 0.42692830970020057, "prediction_shape": [8757], "split_signature": "58385675268180d1"}

Diagnose timing and residual errors

The time traces show whether the model follows the evolution of Dst, while the residual plot exposes systematic errors. The final panel focuses on the January 2015 minimum, where timing and amplitude are easier to inspect directly.

display_count = min(1000, len(y_test))
worst = np.argsort(np.abs(y_test - predictions))[-8:][::-1]

# Locate the strongest Dst decrease around 7 January 2015, then show a
# focused 20-hour interval around the observed minimum.
event_search = (
    (time_test >= np.datetime64("2015-01-05"))
    & (time_test < np.datetime64("2015-01-10"))
)
event_candidates = np.flatnonzero(event_search)
event_center_index = event_candidates[np.argmin(y_test[event_candidates])]
event_center = time_test[event_center_index]
event_window = (
    (time_test >= event_center - np.timedelta64(10, "h"))
    & (time_test <= event_center + np.timedelta64(10, "h"))
)
event_model_metrics = regression_metrics(y_test[event_window], predictions[event_window])
event_persistence_metrics = regression_metrics(
    y_test[event_window], persistence_test[event_window]
)

fig, axes = plt.subplots(3, 1, figsize=(12, 10.5))
axes[0].plot(time_test[:display_count], y_test[:display_count], label="observed", linewidth=1)
axes[0].plot(time_test[:display_count], predictions[:display_count], label="neural model")
axes[0].plot(
    time_test[:display_count],
    persistence_test[:display_count],
    label="persistence",
    linestyle=":",
)
axes[0].set(title="First test interval", ylabel="Dst [nT]")
axes[0].legend()
axes[1].scatter(predictions, y_test - predictions, s=8, alpha=0.35)
axes[1].axhline(0, color="black", linewidth=1)
axes[1].set(xlabel="Predicted Dst [nT]", ylabel="Residual [nT]", title="Test residuals")
axes[2].plot(time_test[event_window], y_test[event_window], label="observed", marker="o")
axes[2].plot(
    time_test[event_window], predictions[event_window], label="neural model", marker="o"
)
axes[2].plot(
    time_test[event_window],
    persistence_test[event_window],
    label="persistence",
    marker="o",
    linestyle=":",
)
axes[2].set(
    title=f"Dst minimum near {str(event_center)[:13]} (±10 hours)",
    xlabel="Target time",
    ylabel="Dst [nT]",
)
axes[2].legend()
axes[2].grid(alpha=0.25)
axes[2].text(
    0.01,
    0.03,
    (
        f"Neural: MAE={event_model_metrics['mae']:.1f}, "
        f"RMSE={event_model_metrics['rmse']:.1f} nT\n"
        f"Persistence: MAE={event_persistence_metrics['mae']:.1f}, "
        f"RMSE={event_persistence_metrics['rmse']:.1f} nT"
    ),
    transform=axes[2].transAxes,
    fontsize=9,
    bbox={"facecolor": "white", "alpha": 0.85, "edgecolor": "0.8"},
)
plt.tight_layout()
plt.show()

print(f"Focused interval center: {event_center}")
print("Focused neural-model metrics:", event_model_metrics)
print("Focused persistence metrics:", event_persistence_metrics)

print("Largest absolute test errors:")
for index in worst:
    print(
        str(time_test[index]),
        f"observed={y_test[index]:.1f}",
        f"predicted={predictions[index]:.1f}",
        f"persistence={persistence_test[index]:.1f}",
    )

Focused interval center: 2015-01-07T11
Focused neural-model metrics: {'mae': 7.409087181091309, 'rmse': 9.66505243262501, 'r2': 0.9461441040039062}
Focused persistence metrics: {'mae': 10.714285850524902, 'rmse': 16.211400595795038, 'r2': 0.8484814167022705}
Largest absolute test errors:
2015-11-03T08 observed=-21.0 predicted=29.2 persistence=34.0
2015-06-23T01 observed=-131.0 predicted=-93.1 persistence=-101.0
2015-03-17T22 observed=-223.0 predicted=-186.5 persistence=-187.0
2015-01-07T08 observed=-21.0 predicted=10.4 persistence=18.0
2015-03-17T19 observed=-165.0 predicted=-134.8 persistence=-143.0
2015-06-22T18 observed=-8.0 predicted=-36.5 persistence=-41.0
2015-03-17T05 observed=56.0 predicted=27.6 persistence=25.0
2015-06-22T19 observed=-53.0 predicted=-25.5 persistence=-8.0

Try it yourself in Keras

Change one choice at a time and keep the data split and evaluation unchanged:

• change HORIZON_HOURS from 1 to 3 and then 6 while keeping the three-hour input history and 2015 as the final test year.

• compare three- and six-hour input histories while keeping the one-hour forecast horizon fixed.

• add one solar-wind variable and fit its scaling on training years only.

• go one step further and rebuild the data-loading stage with NASA CDAWeb’s official cdasws Python API, then reproduce the hourly OMNI variables and time range used here.