Add engineering reproduction package

.gitattributes

@@ -1,35 +1,4 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
 *.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,159 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- OneScience
+- Earth Science
+- Multi-Source Earth Observation
+- Remote Sensing Foundation Model
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">Clay Foundation Model</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+Clay Foundation Model is an independent engineering reproduction of the public Clay v1.5 specification. It jointly encodes satellite imagery from different sensors, band center wavelengths, ground sample distance, time, and latitude/longitude into Earth observation embeddings, and reconstructs input bands with a Masked Autoencoder. It can serve as a representation backbone for land-cover classification, regression, change detection, and other downstream remote-sensing tasks.
+
+Official documentation: Clay Foundation Model v1.5  
+https://clay-foundation.github.io/model/release-notes/specification.html
+
+# Model Description
+
+Clay Foundation Model was proposed by the Clay Foundation. The model was trained on multi-sensor Earth observation data including Sentinel-2, Landsat, Sentinel-1, NAIP, LINZ, and MODIS. It is suitable for remote-sensing image representation, multispectral reconstruction, land-cover classification, regression, change detection, and related tasks.
+
+# Use Cases
+
+| Use Case | Description |
+| :---: | :--- |
+| Multi-sensor representation | Use one model to process imagery with different numbers of bands from sensors such as Sentinel-2, Landsat, and Sentinel-1. |
+| Dynamic spectral encoding | Dynamically generate Patch Embeddings from the center wavelengths of input bands to validate integration methods for new sensors. |
+| Remote-sensing image reconstruction | Use a Masked Autoencoder to reconstruct masked patches in multispectral imagery. |
+| Local engineering validation | Use a small amount of synthetic data to validate training, inference, evaluation, visualization, and checkpoint workflows. |
+| Multi-GPU training | Launch distributed data-parallel training with `torchrun`. |
+
+# Usage Instructions
+
+## 1.OneCode
+
+Experience intelligent, one-click AI4S programming through the OneCode online environment:
+
+[Try intelligent, one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2.Download and Installation
+
+```bash
+hf download OneScience-Group/ClayFoundation --local-dir ./ClayFoundation
+cd ClayFoundation
+```
+
+### Environment Dependencies
+
+**Hardware Requirements**
+
+- A GPU or DCU is recommended.
+- A CPU can be used to verify connectivity with the default small-sample configuration; training the official-size model requires large-scale acceleration resources.
+- DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.
+
+**DCU Environment**
+
+```bash
+# Activate DTK and CONDA first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# uv installation is supported
+pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+
+```bash
+# Activate CONDA first
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+# uv installation is supported
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+### Training Data
+
+This repository uses a small number of synthetic samples to validate the engineering workflow. The training and test data are stored in `data/train.npz` and `data/test.npz`, respectively. The synthetic data contains 10-band Sentinel-2, 6-band Landsat, and 2-band Sentinel-1 imagery, with center wavelengths and ground sample distances corresponding to the official metadata for each sensor. Each sample also includes week, hour, latitude, and longitude metadata; a synthetic teacher representation target; and classification and regression targets for lightweight evaluation.
+
+The default input size is 64×64, with four training samples and two test samples. This scale is intended only to validate the dynamic-band interface, spatiotemporal metadata encoding, MAE training, and multi-sensor inference workflow. It does not represent the data distribution or training scale of the approximately 70 million global remote-sensing chips used by the official model.
+
+```bash
+python scripts/fake_data.py
+```
+
+### Training
+
+```bash
+python scripts/train.py
+```
+
+For multi-GPU training, use:
+
+```bash
+torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py
+```
+
+The default configuration is intended for rapid workflow validation. Formal experiments should use real remote-sensing data, the official-size configuration, a DINOv2 teacher, and a complete training schedule.
+
+```text
+result/checkpoints/clayfoundation.pt
+result/training/metrics.json
+```
+
+### Trained Weights
+
+This repository does not include synthetic or official weights in `weight/`. The Clay Foundation has released the official Clay v1.5 checkpoint:
+
+https://huggingface.co/made-with-clay/Clay/resolve/main/v1.5/clay-v1.5.ckpt
+
+This repository is a scaled-down independent engineering implementation. Its model parameter names and dimensions are incompatible with the official Encoder weights of approximately 1.25 GB. To use the official weights, use the model implementation and data preprocessing workflow provided by the official Clay repository.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference loads the training checkpoint and separately processes Sentinel-2, Landsat, and Sentinel-1 test imagery. It generates Encoder embeddings, L2-normalized projected embeddings, complete multispectral MAE reconstructions, reconstruction loss, and teacher representation alignment loss, and saves them to:
+
+```text
+result/output/predictions.npz
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+The evaluation measures image reconstruction and representation quality across sensors and reports reconstruction error, embedding norms, and representation similarity. It also generates comparison plots of the input imagery, reconstructed results, and absolute errors. Results on synthetic data are intended only to validate the engineering workflow and do not represent the official Clay v1.5 training loss, real downstream-task performance, or cross-region generalization capability.
+
+```text
+result/evaluation/metrics.json
+result/evaluation/comparison.png
+```
+
+# Official OneScience Information
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation and License
+
+This repository is an independent engineering reproduction of the public Clay Foundation Model v1.5 specification. The official Clay source code and model weights are licensed under Apache-2.0.
+
+Official Clay repository:  
+https://github.com/Clay-foundation/model
+
+Use of this repository's code, the official model weights, and the data remains subject to the licenses and terms of use of their respective projects.

README_zh.md

@@ -0,0 +1,162 @@
+---
+license: Apache-2.0
+language:
+- en
+- zh
+tags:
+- OneScience
+- 地球科学
+- 多源地球观测
+- 遥感基础模型
+frameworks: PyTorch
+datasets:
+  - Synthetic Multi-Sensor Earth Observation Chips
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">Clay Foundation Model</span>
+  </strong>
+</p>
+
+# 模型介绍
+
+Clay Foundation Model 是 Clay v1.5 公开规格的独立工程复现模型，将不同传感器的卫星影像、波段中心波长、地面采样距离、时间和经纬度统一编码为地球观测 embedding，并通过 Masked Autoencoder 重建输入波段，可作为地物分类、回归、变化检测和其他下游遥感任务的表征骨干网络。
+
+官方文档：Clay Foundation Model v1.5  
+https://clay-foundation.github.io/model/release-notes/specification.html
+
+# 模型描述
+
+Clay Foundation Model 由 Clay Foundation 提出。模型使用 Sentinel-2、Landsat、Sentinel-1、NAIP、LINZ 和 MODIS 等多传感器地球观测数据训练。模型适用于遥感影像表征、多波段重建、地物分类、回归和变化检测等任务。
+
+# 适用场景
+
+| 场景 | 说明 |
+| :---: | :--- |
+| 多传感器表征 | 使用同一模型处理 Sentinel-2、Landsat 和 Sentinel-1 等不同波段数量的影像。 |
+| 动态光谱编码 | 根据输入波段的中心波长动态生成 Patch Embedding，验证新传感器接入方式。 |
+| 遥感影像重建 | 使用 Masked Autoencoder 对被遮挡的多波段影像 Patch 进行重建。 |
+| 本地工程验证 | 使用少量虚拟数据检查训练、推理、评估、可视化和 checkpoint 流程。 |
+| 多卡训练 | 通过 `torchrun` 启动分布式数据并行训练。 |
+
+# 使用说明
+
+## 1.OneCode
+
+可通过 OneCode 在线环境体验智能化一键式 AI4S 编程：
+
+[点击体验智能化一键式 AI4S 编程](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2.下载安装
+
+```bash
+modelscope download --model OneScience/ClayFoundation --local_dir ./ClayFoundation
+cd ClayFoundation
+```
+
+### 环境依赖
+
+**硬件要求**
+
+- 推荐使用 GPU 或 DCU 运行。
+- CPU 可用于默认小样本配置的连通性验证，官方尺寸模型训练需要大规模加速资源。
+- DCU 用户需预先安装 DTK，建议使用 DTK 25.04.2 以上版本或与当前集群匹配的 OneScience 推荐版本。
+
+**DCU环境**
+
+```bash
+# 请首先激活DTK及CONDA
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# 支持uv安装
+pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+**GPU环境**
+
+```bash
+# 请首先激活CONDA
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+# 支持uv安装
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+### 训练数据介绍
+
+本仓库使用少量虚拟样本验证工程流程，训练数据和测试数据分别保存为 `data/train.npz` 和 `data/test.npz`。虚拟数据包含 Sentinel-2 10 波段、Landsat 6 波段和 Sentinel-1 2 波段影像，并为每个传感器提供官方元数据对应的中心波长和地面采样距离。样本同时包含 week、hour、latitude、longitude 元数据、合成 teacher 表征目标，以及用于轻量评估的分类和回归目标。
+
+默认输入尺寸为 64×64，训练集包含 4 个样本，测试集包含 2 个样本。该规模只用于验证动态波段接口、时空元数据编码、MAE 训练和多传感器推理流程，不代表官方约 7000 万全球遥感 chips 的数据分布与训练规模。
+
+```bash
+python scripts/fake_data.py
+```
+
+### 训练
+
+```bash
+python scripts/train.py
+```
+
+多卡训练可使用：
+
+```bash
+torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py
+```
+
+默认配置面向快速流程验证；开展正式实验时，应使用真实遥感数据、官方尺寸配置、DINOv2 teacher 和完整训练周期。
+
+```text
+result/checkpoints/clayfoundation.pt
+result/training/metrics.json
+```
+
+### 训练权重
+
+本仓库不在 `weight/` 中内置虚拟权重或官方权重。Clay Foundation 已公开 Clay v1.5 官方 checkpoint：
+
+https://huggingface.co/made-with-clay/Clay/resolve/main/v1.5/clay-v1.5.ckpt
+
+本仓库是缩小的独立工程实现，模型参数名称和尺寸不与官方约 1.25 GB Encoder 权重兼容。需要使用官方权重时，应采用 Clay 官方仓库提供的模型实现和数据预处理流程。
+
+### 推理
+
+```bash
+python scripts/inference.py
+```
+
+推理加载训练 checkpoint，分别处理 Sentinel-2、Landsat 和 Sentinel-1 测试影像，生成 Encoder embedding、L2 归一化投影 embedding、完整多波段 MAE 重建结果，以及重建损失和 teacher 表征对齐损失，并保存到：
+
+```text
+result/output/predictions.npz
+```
+
+### 评估和可视化
+
+```bash
+python scripts/result.py
+```
+
+评估模型在不同传感器上的影像重建与表征效果，并统计重建误差、embedding 范数和表征相似度。评估过程同时生成输入影像、重建结果和绝对误差对比图。虚拟数据结果仅用于验证工程流程，不代表官方 Clay v1.5 的训练损失、真实下游任务性能或跨区域泛化能力。
+
+```text
+result/evaluation/metrics.json
+result/evaluation/comparison.png
+```
+
+# OneScience官方信息
+
+| 平台 | OneScience 主仓库 | Skills 仓库 |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# 引用与许可证
+
+本仓库为 Clay Foundation Model v1.5 公开规格的独立工程复现版本。Clay 官方源码和模型权重采用 Apache-2.0 许可证。
+
+Clay 官方仓库：  
+https://github.com/Clay-foundation/model
+
+本仓库代码、官方模型权重和数据的使用仍应以各自项目中的许可证及使用条款为准。

conf/config.yaml

@@ -0,0 +1,63 @@
+seed: 42
+data:
+  root: data
+  format_version: clay_synthetic_v1_5_engineering_v1
+  train_samples: 4
+  test_samples: 2
+  image_size: 64
+  sensors:
+    sentinel-2-l2a:
+      channels: 10
+      gsd: 10.0
+      wavelengths_nm: [493, 560, 665, 704, 740, 783, 842, 865, 1610, 2190]
+      rgb_indices: [2, 1, 0]
+    landsat-c2l2-sr:
+      channels: 6
+      gsd: 30.0
+      wavelengths_nm: [650, 560, 480, 860, 1600, 2200]
+      rgb_indices: [0, 1, 2]
+    sentinel-1-rtc:
+      channels: 2
+      gsd: 10.0
+      wavelengths_nm: [3500, 4000]
+      rgb_indices: [0, 1, 0]
+model:
+  patch_size: 8
+  mask_ratio: 0.75
+  encoder_dim: 96
+  encoder_depth: 2
+  encoder_heads: 4
+  decoder_dim: 64
+  decoder_depth: 1
+  decoder_heads: 4
+  wave_dim: 32
+  wave_latents: 16
+  teacher_dim: 128
+  norm_pix_loss: false
+paper_model:
+  patch_size: 8
+  mask_ratio: 0.75
+  encoder_dim: 1024
+  encoder_depth: 24
+  encoder_heads: 16
+  decoder_dim: 512
+  decoder_depth: 4
+  decoder_heads: 4
+  parameter_count: 632M_including_DINOv2_teacher
+  encoder_parameter_count: 311M
+  teacher: DINOv2_304M
+train:
+  epochs: 1
+  batch_size: 2
+  learning_rate: 0.0005
+  weight_decay: 0.000001
+  reconstruction_weight: 0.95
+  representation_weight: 0.05
+  num_workers: 0
+runtime:
+  device: auto
+paths:
+  checkpoint: result/checkpoints/clayfoundation.pt
+  training_metrics: result/training/metrics.json
+  inference_dir: result/output
+  evaluation_dir: result/evaluation

config.json

@@ -0,0 +1,110 @@
+{
+  "model_name": "ClayFoundation",
+  "model_type": "clayfoundation",
+  "architectures": [
+    "ClayFoundation",
+    "DynamicEmbedding",
+    "Transformer",
+    "DynamicDecoder"
+  ],
+  "framework": "PyTorch",
+  "domain": "earth-observation",
+  "task": "multi-sensor-remote-sensing-representation-learning",
+  "implementation": {
+    "entry_point": "model/clayfoundation.py",
+    "scope": "scaled-down independent engineering reproduction of the Clay v1.5 multi-sensor masked-autoencoder specification",
+    "train_script": "scripts/train.py",
+    "inference_script": "scripts/inference.py",
+    "evaluation_script": "scripts/result.py",
+    "synthetic_data_script": "scripts/fake_data.py"
+  },
+  "architecture": {
+    "family": "multi-sensor masked autoencoder with wavelength-conditioned dynamic embeddings",
+    "patch_size": 8,
+    "mask_ratio": 0.75,
+    "encoder_dim": 96,
+    "encoder_depth": 2,
+    "encoder_heads": 4,
+    "decoder_dim": 64,
+    "decoder_depth": 1,
+    "decoder_heads": 4,
+    "wave_dim": 32,
+    "wave_latents": 16,
+    "teacher_dim": 128,
+    "norm_pix_loss": false,
+    "metadata": [
+      "week",
+      "hour",
+      "latitude",
+      "longitude",
+      "ground_sample_distance",
+      "band_center_wavelengths"
+    ],
+    "paper_model": {
+      "patch_size": 8,
+      "mask_ratio": 0.75,
+      "encoder_dim": 1024,
+      "encoder_depth": 24,
+      "encoder_heads": 16,
+      "decoder_dim": 512,
+      "decoder_depth": 4,
+      "decoder_heads": 4,
+      "parameter_count": "632M including DINOv2 teacher",
+      "encoder_parameter_count": "311M",
+      "teacher": "DINOv2 304M"
+    }
+  },
+  "data": {
+    "datasets": [
+      "Synthetic Multi-Sensor Earth Observation Chips"
+    ],
+    "protocol": "clay_synthetic_v1_5_engineering_v1",
+    "format": "NPZ",
+    "train_file": "data/train.npz",
+    "test_file": "data/test.npz",
+    "image_size": 64,
+    "train_samples": 4,
+    "test_samples": 2,
+    "sensors": {
+      "sentinel-2-l2a": {
+        "channels": 10,
+        "gsd_m": 10.0,
+        "wavelengths_nm": [493, 560, 665, 704, 740, 783, 842, 865, 1610, 2190],
+        "rgb_indices": [2, 1, 0]
+      },
+      "landsat-c2l2-sr": {
+        "channels": 6,
+        "gsd_m": 30.0,
+        "wavelengths_nm": [650, 560, 480, 860, 1600, 2200],
+        "rgb_indices": [0, 1, 2]
+      },
+      "sentinel-1-rtc": {
+        "channels": 2,
+        "gsd_m": 10.0,
+        "wavelengths_nm": [3500, 4000],
+        "rgb_indices": [0, 1, 0]
+      }
+    },
+    "required_metadata": [
+      "format_version",
+      "data_source",
+      "time",
+      "latlon",
+      "class_target",
+      "regression_target",
+      "teacher_target",
+      "pixels",
+      "valid",
+      "wavelengths",
+      "gsd"
+    ]
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/clayfoundation.py",
+    "scripts/fake_data.py",
+    "scripts/train.py",
+    "scripts/inference.py",
+    "scripts/result.py"
+  ]
+}

configuration.json

@@ -0,0 +1,12 @@
+{
+  "framework": "PyTorch",
+  "task": "earth_observation_representation_learning",
+  "model": "Clay Foundation Model",
+  "input_format": "BCHW",
+  "protocol": "engineering_mae_training_and_inference",
+  "default_config": "conf/config.yaml",
+  "training": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/clayfoundation.py

@@ -0,0 +1,222 @@
+"""Pure-PyTorch engineering reproduction of the Clay v1.5 model specification."""
+
+import math
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+def fourier_encode(values, dim, max_frequency=10000.0):
+    """Encode scalar metadata while preserving exactly ``dim`` output features."""
+    if dim < 1:
+        return values.new_zeros(*values.shape, 0)
+    pairs = (dim + 1) // 2
+    frequencies = torch.exp(
+        torch.linspace(0, math.log(max_frequency), pairs, device=values.device, dtype=values.dtype)
+    )
+    angles = values.unsqueeze(-1) * frequencies
+    return torch.cat((angles.sin(), angles.cos()), dim=-1)[..., :dim]
+
+
+def position_encoding_2d(height, width, dim, gsd, device, dtype):
+    if dim % 4:
+        raise ValueError("spatial position dimension must be divisible by four")
+    y, x = torch.meshgrid(
+        torch.arange(height, device=device, dtype=dtype),
+        torch.arange(width, device=device, dtype=dtype),
+        indexing="ij",
+    )
+    scale = torch.as_tensor(gsd, device=device, dtype=dtype) / 10.0
+    quarter = dim // 4
+    frequencies = torch.exp(
+        torch.arange(quarter, device=device, dtype=dtype) * (-math.log(10000.0) / max(quarter, 1))
+    )
+    x_angles = x.reshape(-1, 1) * scale * frequencies
+    y_angles = y.reshape(-1, 1) * scale * frequencies
+    return torch.cat((x_angles.sin(), x_angles.cos(), y_angles.sin(), y_angles.cos()), dim=-1)
+
+
+def patchify(pixels, patch_size):
+    batch, channels, height, width = pixels.shape
+    if height % patch_size or width % patch_size:
+        raise ValueError("image dimensions must be divisible by patch_size")
+    return pixels.reshape(
+        batch, channels, height // patch_size, patch_size, width // patch_size, patch_size
+    ).permute(0, 2, 4, 1, 3, 5).reshape(batch, -1, channels * patch_size * patch_size)
+
+
+def unpatchify(patches, channels, height, width, patch_size):
+    batch = patches.shape[0]
+    return patches.reshape(
+        batch, height // patch_size, width // patch_size, channels, patch_size, patch_size
+    ).permute(0, 3, 1, 4, 2, 5).reshape(batch, channels, height, width)
+
+
+class Transformer(nn.Module):
+    def __init__(self, dim, depth, heads, mlp_ratio=4):
+        super().__init__()
+        layer = nn.TransformerEncoderLayer(
+            dim, heads, int(dim * mlp_ratio), activation="gelu", batch_first=True, norm_first=True
+        )
+        self.layers = nn.TransformerEncoder(layer, depth)
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, values):
+        return self.norm(self.layers(values))
+
+
+class DynamicEmbedding(nn.Module):
+    """Create sensor-agnostic patches by conditioning per-band kernels on wavelength."""
+
+    def __init__(self, patch_size, embed_dim, wave_dim, wave_latents):
+        super().__init__()
+        self.patch_size = patch_size
+        self.wave_dim = wave_dim
+        self.wave_mlp = nn.Sequential(nn.Linear(wave_dim, wave_dim), nn.GELU(), nn.Linear(wave_dim, wave_dim))
+        self.latents = nn.Parameter(torch.randn(wave_latents, wave_dim) * 0.02)
+        self.cross_attention = nn.MultiheadAttention(wave_dim, 4, batch_first=True)
+        self.kernel = nn.Linear(wave_dim, patch_size * patch_size * embed_dim)
+        self.bias = nn.Parameter(torch.zeros(embed_dim))
+        self.embed_dim = embed_dim
+
+    def forward(self, pixels, wavelengths):
+        batch, channels, height, width = pixels.shape
+        if wavelengths.ndim == 1:
+            wavelengths = wavelengths[None].expand(batch, -1)
+        if wavelengths.shape != (batch, channels):
+            raise ValueError(f"wavelengths must have shape {(batch, channels)}, got {tuple(wavelengths.shape)}")
+        wave_features = fourier_encode(wavelengths / 1000.0, self.wave_dim)
+        wave_features = self.wave_mlp(wave_features)
+        queries = self.latents[None].expand(batch, -1, -1)
+        context = self.cross_attention(queries, wave_features, wave_features, need_weights=False)[0].mean(dim=1)
+        conditioned = wave_features + context[:, None]
+        kernels = self.kernel(conditioned).reshape(
+            batch, channels, self.embed_dim, self.patch_size, self.patch_size
+        )
+        patches = []
+        for index in range(batch):
+            patches.append(F.conv2d(pixels[index:index + 1], kernels[index].permute(1, 0, 2, 3),
+                                    stride=self.patch_size) + self.bias[None, :, None, None])
+        return torch.cat(patches).flatten(2).transpose(1, 2), conditioned
+
+
+class DynamicDecoder(nn.Module):
+    def __init__(self, patch_size, decoder_dim, wave_dim):
+        super().__init__()
+        self.patch_size = patch_size
+        self.wave_mlp = nn.Sequential(nn.Linear(wave_dim, decoder_dim), nn.GELU(), nn.Linear(decoder_dim, decoder_dim))
+        self.output = nn.Linear(decoder_dim, patch_size * patch_size)
+
+    def forward(self, tokens, wave_features):
+        wave_context = self.wave_mlp(wave_features)
+        joint = tokens[:, :, None, :] + wave_context[:, None, :, :]
+        return self.output(joint).permute(0, 1, 2, 3).flatten(2)
+
+
+class ClayFoundation(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = dict(config)
+        self.patch_size = int(config["patch_size"])
+        self.mask_ratio = float(config["mask_ratio"])
+        enc_dim, dec_dim = int(config["encoder_dim"]), int(config["decoder_dim"])
+        if enc_dim < 12 or (enc_dim - 8) % 4:
+            raise ValueError("encoder_dim - 8 must be positive and divisible by four")
+        if dec_dim < 12 or (dec_dim - 8) % 4:
+            raise ValueError("decoder_dim - 8 must be positive and divisible by four")
+        self.dynamic_embedding = DynamicEmbedding(
+            self.patch_size, enc_dim, int(config["wave_dim"]), int(config["wave_latents"])
+        )
+        self.cls_token = nn.Parameter(torch.randn(1, 1, enc_dim) * 0.02)
+        self.encoder = Transformer(enc_dim, int(config["encoder_depth"]), int(config["encoder_heads"]))
+        self.encoder_to_decoder = nn.Linear(enc_dim, dec_dim)
+        self.mask_token = nn.Parameter(torch.randn(1, 1, dec_dim) * 0.02)
+        self.decoder = Transformer(dec_dim, int(config["decoder_depth"]), int(config["decoder_heads"]))
+        self.wave_to_decoder = nn.Linear(int(config["wave_dim"]), int(config["wave_dim"]))
+        self.dynamic_decoder = DynamicDecoder(self.patch_size, dec_dim, int(config["wave_dim"]))
+        self.representation_head = nn.Linear(enc_dim, int(config["teacher_dim"]))
+        self.norm_pix_loss = bool(config.get("norm_pix_loss", False))
+
+    @staticmethod
+    def _metadata_encoding(time, latlon, dim):
+        values = torch.cat((time, latlon), dim=1)
+        widths = [dim // 4] * 4
+        for index in range(dim % 4):
+            widths[index] += 1
+        return torch.cat([fourier_encode(values[:, index], widths[index]) for index in range(4)], dim=1)
+
+    def _add_encoding(self, tokens, time, latlon, gsd):
+        batch, length, dim = tokens.shape
+        grid = int(math.sqrt(length))
+        if grid * grid != length:
+            raise ValueError("Clay reproduction requires a square patch grid")
+        spatial = position_encoding_2d(grid, grid, dim - 8, gsd, tokens.device, tokens.dtype)
+        metadata = self._metadata_encoding(time, latlon, 8)
+        encoding = torch.cat((spatial[None].expand(batch, -1, -1), metadata[:, None].expand(-1, length, -1)), dim=-1)
+        return tokens + encoding
+
+    def encode(self, pixels, time, latlon, gsd, wavelengths):
+        patches, _ = self.dynamic_embedding(pixels, wavelengths)
+        patches = self._add_encoding(patches, time, latlon, gsd)
+        cls = self.cls_token.expand(len(pixels), -1, -1)
+        encoded = self.encoder(torch.cat((cls, patches), dim=1))
+        return encoded[:, 0], encoded[:, 1:]
+
+    def forward(self, pixels, time, latlon, gsd, wavelengths, teacher_target=None, mask_ratio=None):
+        ratio = self.mask_ratio if mask_ratio is None else float(mask_ratio)
+        patches, wave_features = self.dynamic_embedding(pixels, wavelengths)
+        patches = self._add_encoding(patches, time, latlon, gsd)
+        batch, length, _ = patches.shape
+        keep = max(1, length - int(length * ratio))
+        noise = torch.rand(batch, length, device=pixels.device)
+        ordering = noise.argsort(dim=1)
+        unmasked_indices, masked_indices = ordering[:, :keep], ordering[:, keep:]
+        gather = unmasked_indices[:, :, None].expand(-1, -1, patches.shape[-1])
+        visible = patches.gather(1, gather)
+        encoded = self.encoder(torch.cat((self.cls_token.expand(batch, -1, -1), visible), dim=1))
+        embedding = encoded[:, 0]
+
+        decoded_visible = self.encoder_to_decoder(encoded[:, 1:])
+        decoder_tokens = self.mask_token.expand(batch, length, -1).clone()
+        decoder_tokens.scatter_(1, unmasked_indices[:, :, None].expand(-1, -1, decoded_visible.shape[-1]), decoded_visible)
+        grid = int(math.sqrt(length))
+        spatial = position_encoding_2d(grid, grid, decoder_tokens.shape[-1] - 8, gsd,
+                                       decoder_tokens.device, decoder_tokens.dtype)
+        metadata = self._metadata_encoding(time, latlon, 8)
+        decoder_tokens = decoder_tokens + torch.cat((spatial[None].expand(batch, -1, -1),
+                                                       metadata[:, None].expand(-1, length, -1)), dim=-1)
+        decoded = self.decoder(decoder_tokens)
+        predicted_patches = self.dynamic_decoder(decoded, self.wave_to_decoder(wave_features))
+        target_patches = patchify(pixels, self.patch_size)
+        if self.norm_pix_loss:
+            mean = target_patches.mean(dim=-1, keepdim=True)
+            variance = target_patches.var(dim=-1, keepdim=True)
+            target_patches = (target_patches - mean) / (variance + 1e-6).sqrt()
+        mask = torch.zeros(batch, length, device=pixels.device)
+        mask.scatter_(1, masked_indices, 1.0)
+        patch_loss = (predicted_patches - target_patches).abs().mean(dim=-1)
+        reconstruction_loss = (patch_loss * mask).sum() / mask.sum().clamp_min(1)
+        projected = F.normalize(self.representation_head(embedding), dim=1)
+        if teacher_target is None:
+            representation_loss = embedding.new_zeros(())
+        else:
+            representation_loss = 1.0 - (projected * F.normalize(teacher_target, dim=1)).sum(dim=1).mean()
+        reconstruction = unpatchify(predicted_patches, pixels.shape[1], pixels.shape[2], pixels.shape[3], self.patch_size)
+        return {
+            "embedding": embedding,
+            "projected_embedding": projected,
+            "reconstruction": reconstruction,
+            "mask": mask,
+            "reconstruction_loss": reconstruction_loss,
+            "representation_loss": representation_loss,
+        }
+
+
+def compute_loss(outputs, reconstruction_weight=0.95, representation_weight=0.05):
+    total = reconstruction_weight * outputs["reconstruction_loss"] + representation_weight * outputs["representation_loss"]
+    return total, {
+        "reconstruction": outputs["reconstruction_loss"],
+        "representation": outputs["representation_loss"],
+        "total": total,
+    }

scripts/fake_data.py

@@ -0,0 +1,91 @@
+"""Generate deterministic multi-sensor chips for the Clay engineering workflow."""
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def load_config():
+    return yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+
+
+def make_split(path, count, config, seed):
+    rng = np.random.default_rng(seed)
+    size = int(config["data"]["image_size"])
+    y, x = np.mgrid[-1:1:complex(size), -1:1:complex(size)].astype(np.float32)
+    payload = {
+        "format_version": np.asarray(config["data"]["format_version"]),
+        "data_source": np.asarray("synthetic"),
+        "time": np.empty((count, 2), np.float32),
+        "latlon": np.empty((count, 2), np.float32),
+        "class_target": np.empty(count, np.int64),
+        "regression_target": np.empty(count, np.float32),
+        "teacher_target": np.empty((count, config["model"]["teacher_dim"]), np.float32),
+    }
+    projection = rng.normal(size=(6, config["model"]["teacher_dim"])).astype(np.float32)
+    for sample in range(count):
+        phase = rng.uniform(0, 2 * np.pi)
+        week = rng.uniform(0, 2 * np.pi)
+        hour = rng.uniform(0, 2 * np.pi)
+        latitude = rng.uniform(-math_pi_over_two(), math_pi_over_two())
+        longitude = rng.uniform(-np.pi, np.pi)
+        payload["time"][sample] = (week, hour)
+        payload["latlon"][sample] = (latitude, longitude)
+        payload["class_target"][sample] = int(np.sin(phase) > 0)
+        payload["regression_target"][sample] = np.cos(phase) + 0.2 * np.sin(latitude)
+        descriptor = np.asarray([
+            np.sin(phase), np.cos(phase), np.sin(week), np.cos(week),
+            np.sin(latitude), np.cos(longitude),
+        ], np.float32)
+        target = descriptor @ projection
+        payload["teacher_target"][sample] = target / np.linalg.norm(target).clip(1e-6)
+
+    for sensor_index, (name, spec) in enumerate(config["data"]["sensors"].items()):
+        channels = int(spec["channels"])
+        pixels = np.empty((count, channels, size, size), np.float32)
+        valid = np.ones((count, 1, size, size), np.float32)
+        for sample in range(count):
+            phase = np.arctan2(payload["teacher_target"][sample, 0], payload["teacher_target"][sample, 1])
+            base = np.sin((2.0 + sensor_index) * np.pi * x + phase) * np.cos(2 * np.pi * y - phase)
+            base += 0.35 * x + 0.15 * y
+            for channel in range(channels):
+                spectral = 0.12 * channel + 0.2 * np.sin(phase + channel / max(channels, 1))
+                pixels[sample, channel] = base + spectral + rng.normal(0, 0.03, (size, size))
+            valid[sample, :, :2] = 0
+        pixels -= pixels.mean(axis=(0, 2, 3), keepdims=True)
+        pixels /= pixels.std(axis=(0, 2, 3), keepdims=True).clip(1e-6)
+        payload[f"pixels_{name}"] = np.clip(pixels, -6, 6).astype(np.float32)
+        payload[f"valid_{name}"] = valid
+        payload[f"wavelengths_{name}"] = np.tile(np.asarray(spec["wavelengths_nm"], np.float32), (count, 1))
+        payload[f"gsd_{name}"] = np.full(count, float(spec["gsd"]), np.float32)
+    np.savez_compressed(path, **payload)
+
+
+def math_pi_over_two():
+    return np.pi / 2
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--force", action="store_true")
+    args = parser.parse_args()
+    config = load_config()
+    data_dir = ROOT / config["data"]["root"]
+    data_dir.mkdir(parents=True, exist_ok=True)
+    for offset, (filename, count) in enumerate((
+        ("train.npz", config["data"]["train_samples"]),
+        ("test.npz", config["data"]["test_samples"]),
+    )):
+        path = data_dir / filename
+        if args.force or not path.exists():
+            make_split(path, int(count), config, int(config["seed"]) + offset)
+        print(f"generated={path.relative_to(ROOT)} samples={count} format={config['data']['format_version']}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,54 @@
+"""Generate Clay embeddings and MAE reconstructions for every configured sensor."""
+
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.clayfoundation import ClayFoundation
+from train import ClayDataset, device_from_config
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    device = device_from_config(config)
+    checkpoint = torch.load(ROOT / config["paths"]["checkpoint"], map_location=device, weights_only=True)
+    if checkpoint["format_version"] != config["data"]["format_version"]:
+        raise ValueError("checkpoint and data formats are incompatible")
+    model = ClayFoundation(checkpoint["model_config"]).to(device)
+    model.load_state_dict(checkpoint["model"])
+    model.eval()
+    dataset = ClayDataset(ROOT / config["data"]["root"] / "test.npz", config)
+    raw = dataset.data
+    payload = {
+        "format_version": np.asarray(config["data"]["format_version"]),
+        "class_target": raw["class_target"],
+        "regression_target": raw["regression_target"],
+    }
+    time = torch.from_numpy(raw["time"]).to(device)
+    latlon = torch.from_numpy(raw["latlon"]).to(device)
+    teacher = torch.from_numpy(raw["teacher_target"]).to(device)
+    with torch.no_grad():
+        for name, spec in config["data"]["sensors"].items():
+            pixels = torch.from_numpy(raw[f"pixels_{name}"]).to(device)
+            waves = torch.from_numpy(raw[f"wavelengths_{name}"]).to(device)
+            outputs = model(pixels, time, latlon, float(spec["gsd"]), waves, teacher, mask_ratio=0.0)
+            payload[f"embedding_{name}"] = outputs["embedding"].cpu().numpy()
+            payload[f"projected_embedding_{name}"] = outputs["projected_embedding"].cpu().numpy()
+            payload[f"reconstruction_{name}"] = outputs["reconstruction"].cpu().numpy()
+            payload[f"pixels_{name}"] = raw[f"pixels_{name}"]
+            payload[f"reconstruction_loss_{name}"] = np.asarray(float(outputs["reconstruction_loss"]))
+            payload[f"representation_loss_{name}"] = np.asarray(float(outputs["representation_loss"]))
+    output = ROOT / config["paths"]["inference_dir"] / "predictions.npz"
+    output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(output, **payload)
+    print(f"predictions={output.relative_to(ROOT)}")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,64 @@
+"""Evaluate reconstruction and cross-sensor representation consistency."""
+
+import json
+from pathlib import Path
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import yaml
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    predictions = np.load(ROOT / config["paths"]["inference_dir"] / "predictions.npz")
+    sensors = list(config["data"]["sensors"])
+    metrics = {"samples": int(len(predictions["class_target"])), "sensors": {}}
+    projected = []
+    for name in sensors:
+        embedding = predictions[f"embedding_{name}"]
+        projection = predictions[f"projected_embedding_{name}"]
+        projected.append(projection)
+        metrics["sensors"][name] = {
+            "reconstruction_mae": float(np.abs(predictions[f"reconstruction_{name}"] - predictions[f"pixels_{name}"]).mean()),
+            "embedding_norm": float(np.linalg.norm(embedding, axis=1).mean()),
+            "projected_embedding_norm": float(np.linalg.norm(projection, axis=1).mean()),
+            "representation_loss": float(predictions[f"representation_loss_{name}"]),
+        }
+    similarities = []
+    for left in range(len(projected)):
+        for right in range(left + 1, len(projected)):
+            similarities.append((projected[left] * projected[right]).sum(axis=1))
+    metrics["mean_cross_sensor_cosine_similarity"] = float(np.concatenate(similarities).mean())
+    output_dir = ROOT / config["paths"]["evaluation_dir"]
+    output_dir.mkdir(parents=True, exist_ok=True)
+    (output_dir / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
+
+    figure, axes = plt.subplots(len(sensors), 3, figsize=(9, 3 * len(sensors)), squeeze=False)
+    for row, name in enumerate(sensors):
+        rgb = config["data"]["sensors"][name]["rgb_indices"]
+        source = predictions[f"pixels_{name}"][0, rgb].transpose(1, 2, 0)
+        reconstruction = predictions[f"reconstruction_{name}"][0, rgb].transpose(1, 2, 0)
+        for image in (source, reconstruction):
+            image -= image.min()
+            image /= image.max() + 1e-6
+        axes[row, 0].imshow(source)
+        axes[row, 0].set_title(f"{name} input")
+        axes[row, 1].imshow(reconstruction)
+        axes[row, 1].set_title("MAE reconstruction")
+        axes[row, 2].imshow(np.abs(source - reconstruction).mean(axis=2), cmap="magma")
+        axes[row, 2].set_title("absolute error")
+        for axis in axes[row]:
+            axis.axis("off")
+    figure.tight_layout()
+    figure.savefig(output_dir / "comparison.png", dpi=150)
+    plt.close(figure)
+    print(f"metrics={output_dir.relative_to(ROOT) / 'metrics.json'}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,121 @@
+"""Train the reduced Clay MAE on deterministic multi-sensor synthetic chips."""
+
+import json
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, Dataset, DistributedSampler
+
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.clayfoundation import ClayFoundation, compute_loss
+
+
+class ClayDataset(Dataset):
+    def __init__(self, path, config):
+        self.data = np.load(path)
+        self.sensors = config["data"]["sensors"]
+        if str(self.data["format_version"]) != config["data"]["format_version"]:
+            raise ValueError("incompatible synthetic data format")
+        size = int(config["data"]["image_size"])
+        for name, spec in self.sensors.items():
+            expected = (int(spec["channels"]), size, size)
+            if self.data[f"pixels_{name}"].shape[1:] != expected:
+                raise ValueError(f"{name} shape does not match {expected}")
+
+    def __len__(self):
+        return len(self.data["time"])
+
+    def __getitem__(self, index):
+        item = {
+            "time": torch.from_numpy(self.data["time"][index]),
+            "latlon": torch.from_numpy(self.data["latlon"][index]),
+            "teacher_target": torch.from_numpy(self.data["teacher_target"][index]),
+        }
+        for name in self.sensors:
+            item[f"pixels_{name}"] = torch.from_numpy(self.data[f"pixels_{name}"][index])
+            item[f"valid_{name}"] = torch.from_numpy(self.data[f"valid_{name}"][index])
+            item[f"waves_{name}"] = torch.from_numpy(self.data[f"wavelengths_{name}"][index])
+        return item
+
+
+def device_from_config(config, local_rank=0):
+    requested = config["runtime"]["device"]
+    if requested == "auto":
+        return torch.device("cuda", local_rank) if torch.cuda.is_available() else torch.device("cpu")
+    return torch.device(requested)
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    torch.manual_seed(int(config["seed"]))
+    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    if distributed:
+        torch.distributed.init_process_group("nccl" if torch.cuda.is_available() else "gloo")
+    rank = torch.distributed.get_rank() if distributed else 0
+    device = device_from_config(config, local_rank)
+    if device.type == "cuda":
+        torch.cuda.set_device(device)
+    dataset = ClayDataset(ROOT / config["data"]["root"] / "train.npz", config)
+    sampler = DistributedSampler(dataset, shuffle=True) if distributed else None
+    loader = DataLoader(dataset, batch_size=int(config["train"]["batch_size"]), sampler=sampler,
+                        shuffle=sampler is None, num_workers=int(config["train"]["num_workers"]))
+    model = ClayFoundation(config["model"]).to(device)
+    if distributed:
+        model = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=float(config["train"]["learning_rate"]),
+                                  weight_decay=float(config["train"]["weight_decay"]), betas=(0.9, 0.95))
+    history = []
+    for epoch in range(int(config["train"]["epochs"])):
+        if sampler:
+            sampler.set_epoch(epoch)
+        model.train()
+        totals = {"reconstruction": 0.0, "representation": 0.0, "total": 0.0}
+        steps = 0
+        for batch in loader:
+            optimizer.zero_grad(set_to_none=True)
+            sensor_losses = []
+            components = []
+            for name, spec in config["data"]["sensors"].items():
+                pixels = batch[f"pixels_{name}"].to(device) * batch[f"valid_{name}"].to(device)
+                outputs = model(pixels, batch["time"].to(device), batch["latlon"].to(device),
+                                float(spec["gsd"]), batch[f"waves_{name}"].to(device),
+                                batch["teacher_target"].to(device))
+                loss, values = compute_loss(outputs, float(config["train"]["reconstruction_weight"]),
+                                            float(config["train"]["representation_weight"]))
+                sensor_losses.append(loss)
+                components.append(values)
+            loss = torch.stack(sensor_losses).mean()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+            for key in totals:
+                totals[key] += sum(float(item[key].detach()) for item in components) / len(components)
+            steps += 1
+        metrics = {key: value / max(steps, 1) for key, value in totals.items()}
+        history.append({"epoch": epoch + 1, **metrics})
+        if rank == 0:
+            print(f"epoch={epoch + 1} total_loss={metrics['total']:.6f} reconstruction={metrics['reconstruction']:.6f}")
+    if rank == 0:
+        checkpoint = ROOT / config["paths"]["checkpoint"]
+        metrics_path = ROOT / config["paths"]["training_metrics"]
+        checkpoint.parent.mkdir(parents=True, exist_ok=True)
+        metrics_path.parent.mkdir(parents=True, exist_ok=True)
+        state = model.module.state_dict() if distributed else model.state_dict()
+        torch.save({"model": state, "model_config": config["model"], "sensors": config["data"]["sensors"],
+                    "format_version": config["data"]["format_version"]}, checkpoint)
+        metrics_path.write_text(json.dumps({"history": history}, indent=2) + "\n")
+        print(f"checkpoint={checkpoint.relative_to(ROOT)}")
+    if distributed:
+        torch.distributed.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()