#!/usr/bin/env python3 """Reproduce Control-F's univariate sensor forecasting benchmark on NASA C-MAPSS FD001. With Python 3.12, install the inference packages in a virtual environment: pip install chronos-forecasting==2.2.2 torch==2.7.1 transformers==4.57.6 requests Then run: python reproduce.py --output results The script downloads the official NASA archive and the pinned model and verifies their hashes. Use --dataset and --model-dir to reuse verified local copies. """ import argparse import csv from collections import defaultdict from hashlib import sha256 import io import json from pathlib import Path from statistics import fmean from tempfile import TemporaryDirectory from zipfile import ZipFile from chronos import BaseChronosPipeline import requests import torch DATASET_URL = "https://data.nasa.gov/dataset/cmapss-jet-engine-simulated-data" ARCHIVE_URL = "https://data.nasa.gov/docs/legacy/CMAPSSData.zip" ARCHIVE_SHA256 = "74bef434a34db25c7bf72e668ea4cd52afe5f2cf8e44367c55a82bfd91a5a34f" TEST_FILE = "test_FD001.txt" MODEL_ID = "amazon/chronos-bolt-tiny" MODEL_REVISION = "a0e552de83495b5c28c14c71c374f3e33280b340" MODEL_WEIGHTS_SHA256 = "75068728d376d2bec670379eeef4bfb4d24c0cfe24d957451f8d19b447030a32" MODEL_CONFIG_SHA256 = "278f0086733031635fb1c861cb01c1bad6477420c7fcb19381a2993e335785e0" CONTEXT = 48 HORIZONS = (5, 10, 20) FORECAST_LENGTH = max(HORIZONS) SENSORS = { 2: ("T24", "Temperatur am Austritt des Niederdruckverdichters", "°R"), 4: ("T50", "Temperatur am Austritt der Niederdruckturbine", "°R"), 11: ("Ps30", "Statischer Druck am Austritt des Hochdruckverdichters", "psia"), } QUANTILES = [0.1, 0.5, 0.9] def read_engines(archive_path): archive_bytes = archive_path.read_bytes() if sha256(archive_bytes).hexdigest() != ARCHIVE_SHA256: raise ValueError("C-MAPSS archive does not match the published benchmark hash.") engines = defaultdict(list) with ZipFile(io.BytesIO(archive_bytes)) as archive: for line in archive.read(TEST_FILE).splitlines(): columns = line.split() if not columns: continue if len(columns) != 26: raise ValueError(f"Expected 26 columns, found {len(columns)}.") engine, cycle = map(int, columns[:2]) sensors = {sensor: float(columns[4 + sensor]) for sensor in SENSORS} engines[engine].append((cycle, sensors)) if len(engines) != 100: raise ValueError("Expected the 100 official FD001 test trajectories.") return dict(sorted(engines.items())) def make_windows(engines): windows = [] for engine, rows in engines.items(): if len(rows) < CONTEXT + FORECAST_LENGTH: continue tail = rows[-(CONTEXT + FORECAST_LENGTH):] cycles = [cycle for cycle, _ in tail] if cycles != list(range(cycles[0], cycles[-1] + 1)): raise ValueError(f"Engine {engine} has a gap in the selected window.") for sensor in SENSORS: values = [sensors[sensor] for _, sensors in tail] windows.append({ "engine": engine, "sensor": sensor, "cycles": cycles, "history": values[:CONTEXT], "actual": values[CONTEXT:], }) return windows def download(url, destination): try: with requests.get(url, timeout=(10, 30), stream=True) as response: response.raise_for_status() with destination.open("wb") as file: for chunk in response.iter_content(chunk_size=1024 * 1024): file.write(chunk) except requests.RequestException: raise RuntimeError(f"Could not download {url}. Try again or use a verified local copy.") from None def load_model(model_dir, temp): files = { "config.json": MODEL_CONFIG_SHA256, "model.safetensors": MODEL_WEIGHTS_SHA256, } if model_dir is None: model_dir = temp / "model" model_dir.mkdir() for filename in files: download( f"https://huggingface.co/{MODEL_ID}/resolve/{MODEL_REVISION}/{filename}", model_dir / filename, ) for filename, expected in files.items(): if sha256((model_dir / filename).read_bytes()).hexdigest() != expected: raise ValueError(f"{filename} does not match the pinned model revision.") return BaseChronosPipeline.from_pretrained( str(model_dir), device_map="cpu", dtype=torch.float32, local_files_only=True ) def add_forecasts(windows, pipeline): torch.set_num_threads(4) batch_size = 32 for start in range(0, len(windows), batch_size): batch = windows[start:start + batch_size] history = torch.tensor([window["history"] for window in batch], dtype=torch.float32) with torch.inference_mode(): quantiles, _ = pipeline.predict_quantiles( history, prediction_length=FORECAST_LENGTH, quantile_levels=QUANTILES, ) for window, forecast in zip(batch, quantiles.tolist(), strict=True): window["p10"], window["chronos"], window["p90"] = ( [round(value, 6) for value in series] for series in zip(*forecast, strict=True) ) window["baseline"] = [window["history"][-1]] * FORECAST_LENGTH def summarise(windows, horizon): def mae(key): return round(fmean( abs(actual - predicted) for window in windows for actual, predicted in zip(window["actual"][:horizon], window[key][:horizon], strict=True) ), 6) return { "chronosMae": mae("chronos"), "baselineMae": mae("baseline"), "observations": len(windows) * horizon, } def write_output(output, engines, windows): channels = [] for sensor, (symbol, description, unit) in SENSORS.items(): channel_windows = [window for window in windows if window["sensor"] == sensor] channels.append({ "id": f"s{sensor}", "label": f"Sensor {sensor}", "symbol": symbol, "description": description, "unit": unit, "metrics": {str(h): summarise(channel_windows, h) for h in HORIZONS}, "examples": [ {key: value for key, value in window.items() if key != "sensor"} for window in channel_windows ], }) report = { "dataset": { "name": "NASA C-MAPSS FD001", "url": DATASET_URL, "archiveUrl": ARCHIVE_URL, "split": TEST_FILE, "sha256": ARCHIVE_SHA256, "testEngineCount": len(engines), }, "model": { "name": "Chronos-Bolt Tiny", "id": MODEL_ID, "url": f"https://huggingface.co/{MODEL_ID}", "revision": MODEL_REVISION, "weightsSha256": MODEL_WEIGHTS_SHA256, }, "contextLength": CONTEXT, "horizons": HORIZONS, "engineCount": len(windows) // len(SENSORS), "protocol": { "sampling": "One trailing 68-cycle window per FD001 test engine with at least 68 observed cycles.", "forecast": "Univariate zero-shot forecasts from the first 48 cycles, withholding the last 20 cycles.", "baseline": "Repeat the last observed value across all forecast steps.", "metric": "Mean absolute error per sensor across all eligible engines and the first H forecast steps.", "quantiles": QUANTILES, }, "channels": channels, } output.mkdir(parents=True, exist_ok=True) (output / "benchmark.json").write_text( json.dumps(report, separators=(",", ":"), allow_nan=False) + "\n", encoding="utf-8" ) with (output / "benchmark.csv").open("w", newline="", encoding="utf-8") as file: writer = csv.writer(file) writer.writerow(("engine", "sensor", "cycle", "step", "actual", "chronos_p50", "last_value", "chronos_p10", "chronos_p90")) for window in windows: for step, actual in enumerate(window["actual"], start=1): index = step - 1 writer.writerow(( window["engine"], f"s{window['sensor']}", window["cycles"][CONTEXT + index], step, actual, window["chronos"][index], window["baseline"][index], window["p10"][index], window["p90"][index], )) def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--dataset", type=Path, help="Optional local copy of NASA's CMAPSSData.zip.") parser.add_argument("--model-dir", type=Path, help="Optional local copy of the pinned Hugging Face model.") parser.add_argument("--output", type=Path, default=Path(__file__).resolve().parent) args = parser.parse_args() with TemporaryDirectory(prefix="control-f-benchmark-") as temp: temp_path = Path(temp) archive_path = args.dataset if archive_path is None: archive_path = temp_path / "CMAPSSData.zip" download(ARCHIVE_URL, archive_path) engines = read_engines(archive_path) windows = make_windows(engines) add_forecasts(windows, load_model(args.model_dir, temp_path)) write_output(args.output, engines, windows) if __name__ == "__main__": main()