Add engineering reproduction package

.gitattributes

@@ -1,35 +1,13 @@
-*.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
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl 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

README.md

@@ -0,0 +1,151 @@
+---
+license: cc-by-4.0
+language:
+- en
+tags:
+- OneScience
+- Earth Science
+- Multi-Source Earth Observation
+- Geospatial Embeddings
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">AlphaEarthFoundations</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+AlphaEarthFoundations is a multi-source spatiotemporal embedding field model for global Earth observation mapping. It jointly encodes sparse and asynchronous data, including optical imagery, radar, LiDAR, climate, elevation, land cover, and geographic text, into 64-dimensional unit-sphere embeddings that support classification, regression, and change detection with limited labels.
+
+Paper: AlphaEarth Foundations: An embedding field model for accurate and efficient global mapping from sparse label data  
+https://arxiv.org/abs/2507.22291
+
+# Model Description
+
+AlphaEarth Foundations was proposed by research teams at Google DeepMind and Google. The model was trained on more than three billion observations. It takes Sentinel-2, Sentinel-1, and Landsat-8/9 time-series imagery as input and uses data including PALSAR-2, ERA5-Land, GEDI, GRACE, Copernicus DEM, NLCD, and geographic text as learning targets. The model is suitable for multi-source Earth observation representation learning, thematic mapping with sparse labels, biophysical variable estimation, and time-series change detection.
+
+# Use Cases
+
+| Use Case | Description |
+| :---: | :--- |
+| Multi-source spatiotemporal representation | Fuse asynchronous time-series observations from Sentinel-2, Sentinel-1, and Landsat-8/9. |
+| Mapping with sparse labels | Train kNN or linear predictors using 64-dimensional embeddings and a small number of point labels. |
+| Biophysical variable estimation | Regress continuous variables such as land surface emissivity and evapotranspiration from embeddings. |
+| Land surface change detection | Compare unit-sphere embeddings from different validity periods to perform supervised or unsupervised change detection. |
+| Local engineering validation | Use a small amount of synthetic data that preserves the paper's actual dimensions to validate training, inference, quantization, evaluation, and visualization 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/AlphaEarthFoundations --local-dir ./AlphaEarthFoundations
+cd AlphaEarthFoundations
+```
+
+### Environment Dependencies
+
+**Hardware Requirements**
+
+- A GPU or DCU is recommended.
+- A CPU can be used to verify connectivity with the default small-sample configuration; full training and global inference require 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 preserves the paper's 1.28 km × 1.28 km extent, 10-meter grid, 65 Sentinel-2 frames, 17 Sentinel-1 frames, 21 Landsat frames, and the explicitly disclosed input-source and training-target channel specifications; only the number of samples and the default internal model width are reduced.
+
+```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 the multi-temporal data, complete task labels, model configuration, and training schedule corresponding to the paper.
+
+```text
+result/checkpoints/alphaearthfoundations.pt
+result/training/metrics.json
+```
+
+### Trained Weights
+
+This repository does not include synthetic or trained weights. Training generates `result/checkpoints/alphaearthfoundations.pt`; Google and Google DeepMind have not released the paper's v2.0/v2.1 model weights or checkpoints.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference loads the training checkpoint and generates float32 unit-sphere embeddings, the paper's `s8²` signed int8 quantized embeddings, and reconstruction results for nine target categories. The outputs are saved to:
+
+```text
+result/output/predictions.npz
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+The evaluation reports reconstruction MAE or classification error rates for each training source, the mean embedding norm, and `s8²` quantization error. It also follows the paper's lightweight transfer approach by running kNN with `k=1`, kNN with `k=3`, unregularized linear classification, and linear regression; reporting Balanced Accuracy and R²; and generating comparison plots of the A01, A16, and A09 embedding axes against the land-cover target. Results on synthetic data are intended only to validate the engineering workflow and do not represent the paper's complete performance on 15 real downstream datasets.
+
+```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 engineering reproduction of the original AlphaEarth Foundations paper.
+
+Use of the code and data in this repository remains subject to the licenses and terms of use of their respective projects.

README_zh.md

@@ -0,0 +1,154 @@
+---
+license: CC-BY-4.0
+language:
+- en
+- zh
+tags:
+- OneScience
+- 地球科学
+- 多源地球观测
+- 地理空间嵌入
+frameworks: PyTorch
+datasets:
+  - Google Satellite Embedding V1 Annual
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">AlphaEarthFoundations</span>
+  </strong>
+</p>
+
+# 模型介绍
+
+AlphaEarthFoundations 是面向全球地球观测制图的多源时空嵌入场模型，将光学、雷达、LiDAR、气候、高程、土地覆盖和地理文本等稀疏且异步的数据统一编码为 64 维单位球嵌入，以少量标签支持分类、回归和变化检测。
+
+论文：AlphaEarth Foundations: An embedding field model for accurate and efficient global mapping from sparse label data  
+https://arxiv.org/abs/2507.22291
+
+# 模型描述
+
+AlphaEarth Foundations 由 Google DeepMind 与 Google 的研究团队提出。模型使用超过 30 亿次观测训练，输入 Sentinel-2、Sentinel-1 和 Landsat-8/9 时序影像，并以 PALSAR-2、ERA5-Land、GEDI、GRACE、Copernicus DEM、NLCD 和地理文本等数据作为学习目标。模型适用于多源地球观测表征学习、稀疏标签专题制图、生物物理变量估计和时序变化检测。
+
+# 适用场景
+
+| 场景 | 说明 |
+| :---: | :--- |
+| 多源时空表征 | 融合 Sentinel-2、Sentinel-1 和 Landsat-8/9 的异步时序观测。 |
+| 稀疏标签制图 | 使用 64 维 embedding 和少量点标签训练 kNN 或线性预测器。 |
+| 生物物理变量估计 | 基于 embedding 回归地表发射率、蒸散量等连续变量。 |
+| 地表变化检测 | 比较不同有效时段的单位球 embedding，执行监督或无监督变化检测。 |
+| 本地工程验证 | 使用保持论文真实维度的少量虚拟数据检查训练、推理、量化、评估和可视化流程。 |
+| 多卡训练 | 通过 `torchrun` 启动分布式数据并行训练。 |
+
+# 使用说明
+
+## 1.OneCode
+
+可通过 OneCode 在线环境体验智能化一键式 AI4S 编程：
+
+[点击体验智能化一键式 AI4S 编程](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2.下载安装
+
+```bash
+modelscope download --model OneScience/AlphaEarthFoundations --local_dir ./AlphaEarthFoundations
+cd AlphaEarthFoundations
+```
+
+### 环境依赖
+
+**硬件要求**
+
+- 推荐使用 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`。虚拟数据保持论文的 1.28 km × 1.28 km、10 米网格、Sentinel-2 65 帧、Sentinel-1 17 帧、Landsat 21 帧，以及论文明确公开的输入源和训练目标通道规格；仅样本数量和默认模型内部宽度被缩小。
+
+```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
+```
+
+默认配置面向快速流程验证；开展正式实验时，应使用论文对应的多时相数据、完整任务标签、模型配置和训练周期。
+
+```text
+result/checkpoints/alphaearthfoundations.pt
+result/training/metrics.json
+```
+
+### 训练权重
+
+本仓库不内置虚拟权重或训练权重。执行训练后生成 `result/checkpoints/alphaearthfoundations.pt`；Google 和 Google DeepMind 尚未公开论文 v2.0/v2.1 模型权重或 checkpoint。
+
+### 推理
+
+```bash
+python scripts/inference.py
+```
+
+推理加载训练 checkpoint，生成 float32 单位球 embedding、论文 `s8²` 有符号 int8 量化 embedding 和九类目标重建结果，并保存到：
+
+```text
+result/output/predictions.npz
+```
+
+### 评估和可视化
+
+```bash
+python scripts/result.py
+```
+
+评估报告各训练源的重建 MAE 或分类错误率、embedding 平均范数和 `s8²` 量化误差；同时按论文的轻量迁移方式执行 kNN `k=1`、kNN `k=3`、无正则线性分类和线性回归，报告 Balanced Accuracy 与 R²，并生成 A01/A16/A09 三个 embedding 轴与土地覆盖目标的对比图。虚拟数据结���仅用于验证工程流程，不代表论文在 15 个真实下游数据集上的完整性能。
+
+```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 |
+
+# 引用与许可证
+
+本仓库为 AlphaEarth Foundations 原始论文的工程复现版本。
+
+本仓库代码和数据的使用仍应以各自项目中的许可证及使用条款为准。

conf/config.yaml

@@ -0,0 +1,57 @@
+seed: 42
+data:
+  root: data
+  format_version: aef_synthetic_paper_spec_v3
+  train_samples: 2
+  test_samples: 1
+  image_size: 128
+  input_sources:
+    sentinel2: {timesteps: 65, channels: 5, geometry_dim: 0}
+    sentinel1: {timesteps: 17, channels: 5, geometry_dim: 2}
+    landsat: {timesteps: 21, channels: 7, geometry_dim: 0}
+  target_sources:
+    sentinel2: {channels: 5, type: continuous, weight: 1.0, loss_grid_m: 10, shift_pixels: 2, geometry_dim: 0}
+    sentinel1: {channels: 5, type: continuous, weight: 1.0, loss_grid_m: 10, shift_pixels: 2, geometry_dim: 2}
+    landsat: {channels: 7, type: continuous, weight: 1.0, loss_grid_m: 30, geometry_dim: 0}
+    palsar2: {channels: 3, type: continuous, weight: 1.0, loss_grid_m: 30, geometry_dim: 2}
+    era5_land: {channels: 12, type: continuous, weight: 1.0, loss_grid_m: 10, geometry_dim: 0}
+    gedi: {channels: 101, type: continuous, weight: 1.0, loss_grid_m: 20, geometry_dim: 0}
+    grace: {channels: 1, type: continuous, weight: 0.5, loss_grid_m: 1280, geometry_dim: 0}
+    glo30: {channels: 1, type: continuous, weight: 1.0, loss_grid_m: 30, geometry_dim: 0, derived_channels: undisclosed}
+    nlcd: {channels: 16, type: categorical, weight: 0.5, loss_grid_m: 30, geometry_dim: 0}
+model:
+  precision_dim: 8
+  time_dim: 16
+  space_dim: 32
+  embedding_dim: 64
+  num_blocks: 1
+  num_heads: 4
+  decoder_hidden_dim: 32
+  vmf_kappa: 8000.0
+  max_geometry_dim: 2
+paper_model:
+  precision_dim: 128
+  time_dim: 512
+  space_dim: 1024
+  embedding_dim: 64
+  num_blocks: 15
+  parameter_count: approximately_480M
+train:
+  epochs: 1
+  batch_size: 2
+  learning_rate: 0.0005
+  weight_decay: 0.000001
+  reconstruction_weight: 1.0
+  uniformity_weight: 0.05
+  consistency_weight: 0.02
+  text_weight: 0.001
+  source_dropout: {sentinel2: 0.0, sentinel1: 0.3, landsat: 0.3}
+  frame_dropout: {sentinel2: 0.5, sentinel1: 0.3, landsat: 0.3}
+  num_workers: 0
+runtime:
+  device: auto
+paths:
+  checkpoint: result/checkpoints/alphaearthfoundations.pt
+  training_metrics: result/training/metrics.json
+  inference_dir: result/output
+  evaluation_dir: result/evaluation

config.json

@@ -0,0 +1,109 @@
+{
+  "model_name": "AlphaEarthFoundations",
+  "model_type": "alphaearthfoundations",
+  "architectures": [
+    "AlphaEarthFoundations",
+    "STPBlock",
+    "ConditionalDecoder"
+  ],
+  "framework": "PyTorch",
+  "domain": "earth-observation",
+  "task": "multi-source-geospatial-embedding",
+  "implementation": {
+    "entry_point": "model/alphaearthfoundations.py",
+    "scope": "engineering reproduction of a multi-source spatiotemporal embedding field for sparse-label global Earth observation mapping",
+    "train_script": "scripts/train.py",
+    "inference_script": "scripts/inference.py",
+    "evaluation_script": "scripts/result.py",
+    "synthetic_data_script": "scripts/fake_data.py"
+  },
+  "architecture": {
+    "family": "spatiotemporal precision-time-space embedding field",
+    "input_sources": [
+      "sentinel2",
+      "sentinel1",
+      "landsat"
+    ],
+    "precision_dim": 8,
+    "time_dim": 16,
+    "space_dim": 32,
+    "embedding_dim": 64,
+    "num_blocks": 1,
+    "num_heads": 4,
+    "decoder_hidden_dim": 32,
+    "vmf_kappa": 8000.0,
+    "max_geometry_dim": 2,
+    "embedding_normalization": "L2 unit sphere",
+    "paper_model": {
+      "precision_dim": 128,
+      "time_dim": 512,
+      "space_dim": 1024,
+      "embedding_dim": 64,
+      "num_blocks": 15,
+      "parameter_count": "approximately 480M"
+    }
+  },
+  "data": {
+    "datasets": [
+      "Google Satellite Embedding V1 Annual"
+    ],
+    "protocol": "aef_synthetic_paper_spec_v3",
+    "format": "NPZ",
+    "train_file": "data/train.npz",
+    "test_file": "data/test.npz",
+    "image_size": 128,
+    "spatial_extent_km": [
+      1.28,
+      1.28
+    ],
+    "grid_resolution_m": 10,
+    "input_sources": {
+      "sentinel2": {
+        "timesteps": 65,
+        "channels": 5,
+        "geometry_dim": 0
+      },
+      "sentinel1": {
+        "timesteps": 17,
+        "channels": 5,
+        "geometry_dim": 2
+      },
+      "landsat": {
+        "timesteps": 21,
+        "channels": 7,
+        "geometry_dim": 0
+      }
+    },
+    "target_sources": {
+      "sentinel2": {"channels": 5, "type": "continuous", "loss_grid_m": 10},
+      "sentinel1": {"channels": 5, "type": "continuous", "loss_grid_m": 10},
+      "landsat": {"channels": 7, "type": "continuous", "loss_grid_m": 30},
+      "palsar2": {"channels": 3, "type": "continuous", "loss_grid_m": 30},
+      "era5_land": {"channels": 12, "type": "continuous", "loss_grid_m": 10},
+      "gedi": {"channels": 101, "type": "continuous", "loss_grid_m": 20},
+      "grace": {"channels": 1, "type": "continuous", "loss_grid_m": 1280},
+      "glo30": {"channels": 1, "type": "continuous", "loss_grid_m": 30},
+      "nlcd": {"channels": 16, "type": "categorical", "loss_grid_m": 30}
+    },
+    "required_metadata": [
+      "format_version",
+      "data_source",
+      "timestamps",
+      "frame_available",
+      "channel_available",
+      "pixel_valid",
+      "geometry",
+      "support_period",
+      "valid_period",
+      "text_target"
+    ]
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/alphaearthfoundations.py",
+    "scripts/fake_data.py",
+    "scripts/train.py",
+    "scripts/inference.py",
+    "scripts/result.py"
+  ]
+}

configuration.json

@@ -0,0 +1,11 @@
+{
+  "framework": "PyTorch",
+  "task": "geospatial-embedding-field",
+  "model": "AlphaEarthFoundations",
+  "paper": "https://arxiv.org/abs/2507.22291",
+  "dataset": "GOOGLE/SATELLITE_EMBEDDING/V1/ANNUAL",
+  "license": "CC-BY-4.0",
+  "allow_remote": true,
+  "model": {"type": "AlphaEarthFoundations", "repo_id": "OneScience/AlphaEarthFoundations"},
+  "pipeline": {"type": "earth-observation-embedding"}
+}

model/alphaearthfoundations.py

@@ -0,0 +1,234 @@
+"""Engineering reproduction of AlphaEarth Foundations from the paper specification."""
+
+import math
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+def sinusoidal_timecode(timestamps, dim, origin=None, scale=365.25 * 24 * 3600 * 1000):
+    if origin is None:
+        origin = timestamps.amin(dim=1, keepdim=True)
+    values = (timestamps.double() - origin.double()) / scale
+    values = values.float()
+    frequencies = torch.exp(
+        torch.arange(0, dim, 2, device=timestamps.device) * (-math.log(10000.0) / dim)
+    )
+    angles = values.unsqueeze(-1) * frequencies
+    code = torch.zeros(*timestamps.shape, dim, device=timestamps.device)
+    code[..., 0::2] = angles.sin()
+    code[..., 1::2] = angles.cos()
+    return code
+
+
+class STPBlock(nn.Module):
+    """Parallel precision, time and space operators with learned pyramid exchange."""
+
+    def __init__(self, precision_dim, time_dim, space_dim, num_heads):
+        super().__init__()
+        self.precision = nn.Sequential(
+            nn.GroupNorm(1, precision_dim),
+            nn.Conv2d(precision_dim, precision_dim, 3, padding=1),
+            nn.GELU(),
+            nn.Conv2d(precision_dim, precision_dim, 3, padding=1),
+        )
+        self.time_norm = nn.LayerNorm(time_dim)
+        self.time_attention = nn.MultiheadAttention(time_dim, num_heads, batch_first=True)
+        self.space_norm = nn.LayerNorm(space_dim)
+        self.space_attention = nn.MultiheadAttention(space_dim, num_heads, batch_first=True)
+        self.to_precision = nn.ModuleList([nn.Conv2d(time_dim, precision_dim, 1), nn.Conv2d(space_dim, precision_dim, 1)])
+        self.to_time = nn.Conv2d(precision_dim, time_dim, 1)
+        self.to_space = nn.Conv2d(precision_dim, space_dim, 1)
+
+    def forward(self, precision, time, space, frame_available):
+        batch, frames = precision.shape[:2]
+        p_size, t_size, s_size = precision.shape[-2:], time.shape[-2:], space.shape[-2:]
+        p = precision.flatten(0, 1)
+        p = p + self.precision(p)
+
+        sequence = time.permute(0, 3, 4, 1, 2).reshape(-1, frames, time.shape[2])
+        normalized = self.time_norm(sequence)
+        time_mask = (~frame_available.bool())[:, None, None, :].expand(batch, *t_size, frames).reshape(-1, frames)
+        sequence = sequence + self.time_attention(
+            normalized, normalized, normalized, key_padding_mask=time_mask, need_weights=False
+        )[0]
+        time = sequence.reshape(batch, *t_size, frames, -1).permute(0, 3, 4, 1, 2)
+
+        available = frame_available[:, :, None, None, None].to(space.dtype)
+        spatial = (space * available).sum(dim=1) / available.sum(dim=1).clamp_min(1)
+        spatial = spatial.flatten(2).transpose(1, 2)
+        normalized = self.space_norm(spatial)
+        spatial = spatial + self.space_attention(normalized, normalized, normalized, need_weights=False)[0]
+        spatial = spatial.transpose(1, 2).reshape(batch, -1, *s_size)
+        space = space + spatial[:, None]
+
+        t_flat, s_flat = time.flatten(0, 1), space.flatten(0, 1)
+        precision = p + self.to_precision[0](F.interpolate(t_flat, p_size, mode="bilinear", align_corners=False))
+        precision = precision + self.to_precision[1](F.interpolate(s_flat, p_size, mode="bilinear", align_corners=False))
+        time = time + self.to_time(F.interpolate(p, t_size, mode="bilinear", align_corners=False)).unflatten(0, (batch, frames))
+        space = space + self.to_space(F.interpolate(p, s_size, mode="bilinear", align_corners=False)).unflatten(0, (batch, frames))
+        return precision.unflatten(0, (batch, frames)), time, space
+
+
+class ConditionalDecoder(nn.Module):
+    def __init__(self, embedding_dim, condition_dim, hidden_dim, output_dim):
+        super().__init__()
+        self.condition = nn.Linear(condition_dim, hidden_dim)
+        self.network = nn.Sequential(
+            nn.Conv2d(embedding_dim + hidden_dim, hidden_dim, 1),
+            nn.GELU(),
+            nn.Conv2d(hidden_dim, hidden_dim, 1),
+            nn.GELU(),
+            nn.Conv2d(hidden_dim, output_dim, 1),
+        )
+
+    def forward(self, embedding, condition):
+        context = self.condition(condition)[:, :, None, None].expand(-1, -1, *embedding.shape[-2:])
+        return self.network(torch.cat([embedding, context], dim=1))
+
+
+class AlphaEarthFoundations(nn.Module):
+    def __init__(self, input_sources, target_sources, config):
+        super().__init__()
+        p_dim, t_dim, s_dim = config["precision_dim"], config["time_dim"], config["space_dim"]
+        self.input_names = list(input_sources)
+        self.target_sources = target_sources
+        self.embedding_dim = config["embedding_dim"]
+        self.vmf_kappa = float(config["vmf_kappa"])
+        self.projectors = nn.ModuleDict({
+            name: nn.Sequential(nn.Conv2d(spec["channels"], p_dim, 3, stride=2, padding=1), nn.GELU())
+            for name, spec in input_sources.items()
+        })
+        self.time_projector = nn.Conv2d(p_dim, t_dim, 3, stride=4, padding=1)
+        self.space_projector = nn.Conv2d(p_dim, s_dim, 3, stride=8, padding=1)
+        self.time_context = nn.Linear(t_dim, t_dim)
+        self.blocks = nn.ModuleList([
+            STPBlock(p_dim, t_dim, s_dim, config["num_heads"]) for _ in range(config["num_blocks"])
+        ])
+        self.summary_query = nn.Linear(t_dim * 2, p_dim)
+        self.embedding_head = nn.Conv2d(p_dim, self.embedding_dim, 1)
+        self.embedding_upsample = nn.ConvTranspose2d(p_dim, p_dim, 4, stride=2, padding=1)
+        condition_dim = t_dim + config["max_geometry_dim"]
+        self.decoders = nn.ModuleDict({
+            name: ConditionalDecoder(self.embedding_dim, condition_dim, config["decoder_hidden_dim"], spec["channels"])
+            for name, spec in target_sources.items()
+        })
+
+    def _summarize(self, precision, availability, period, origin):
+        period_codes = sinusoidal_timecode(period, self.time_context.in_features, origin)
+        query = self.summary_query(period_codes.flatten(1))
+        scores = (precision * query[:, None, :, None, None]).sum(dim=2).mean(dim=(-1, -2))
+        scores = scores.masked_fill(~availability.bool(), torch.finfo(scores.dtype).min)
+        summary = (precision * scores.softmax(dim=1)[:, :, None, None, None]).sum(dim=1)
+        return F.normalize(self.embedding_head(self.embedding_upsample(summary)), dim=1)
+
+    def forward(self, sources, timestamps, valid_period, frame_available, target_times=None,
+                target_geometry=None, target_periods=None):
+        precision_parts, code_parts = [], []
+        origin = torch.cat(list(timestamps.values()), dim=1).amin(dim=1, keepdim=True)
+        for name in self.input_names:
+            values = sources[name]
+            batch, frames = values.shape[:2]
+            projected = self.projectors[name](values.flatten(0, 1)).unflatten(0, (batch, frames))
+            precision_parts.append(projected)
+            code_parts.append(sinusoidal_timecode(timestamps[name], self.time_context.in_features, origin))
+        availability = torch.cat([frame_available[name] for name in self.input_names], dim=1)
+        precision = torch.cat(precision_parts, dim=1)
+        codes = torch.cat(code_parts, dim=1)
+        time = self.time_projector(precision.flatten(0, 1)).unflatten(0, precision.shape[:2])
+        time = time + self.time_context(codes)[:, :, :, None, None]
+        space = self.space_projector(precision.flatten(0, 1)).unflatten(0, precision.shape[:2])
+        for block in self.blocks:
+            precision, time, space = block(precision, time, space, availability)
+
+        embedding = self._summarize(precision, availability, valid_period, origin)
+        outputs = {"embedding": embedding}
+        if target_times is not None:
+            outputs["reconstructions"] = {}
+            for name in self.target_sources:
+                source_embedding = self._summarize(precision, availability, target_periods[name], origin)
+                if self.training:
+                    source_embedding = F.normalize(
+                        source_embedding + torch.randn_like(source_embedding) / math.sqrt(self.vmf_kappa), dim=1
+                    )
+                relative_time = (
+                    (target_times[name] - target_periods[name][:, 0]).float()
+                    / (target_periods[name][:, 1] - target_periods[name][:, 0]).float().clamp_min(1)
+                )
+                time_code = sinusoidal_timecode(
+                    relative_time[:, None], self.time_context.in_features,
+                    torch.zeros_like(relative_time[:, None]), scale=1.0
+                )[:, 0]
+                geometry = target_geometry[name]
+                outputs["reconstructions"][name] = self.decoders[name](source_embedding, torch.cat([time_code, geometry], dim=1))
+        return outputs
+
+
+def _pool_continuous(values, grid_m):
+    if grid_m == 10:
+        return values
+    size = max(1, round(values.shape[-1] * 10 / grid_m))
+    return F.adaptive_avg_pool2d(values, (size, size))
+
+
+def _shift_invariant_l1(prediction, target, mask, radius):
+    losses = []
+    for dy in range(-radius, radius + 1):
+        for dx in range(-radius, radius + 1):
+            shifted = torch.roll(prediction, (dy, dx), dims=(-2, -1))
+            valid = mask.clone()
+            if dy > 0: valid[..., :dy, :] = 0
+            if dy < 0: valid[..., dy:, :] = 0
+            if dx > 0: valid[..., :, :dx] = 0
+            if dx < 0: valid[..., :, dx:] = 0
+            losses.append((torch.abs(shifted - target) * valid).sum() / valid.sum().clamp_min(1))
+    return torch.stack(losses).amin()
+
+
+def compute_losses(teacher, student, targets, masks, text_target, target_sources, weights):
+    reconstruction = teacher["embedding"].new_zeros(())
+    components = {}
+    for name, spec in target_sources.items():
+        prediction, target, mask = teacher["reconstructions"][name], targets[name], masks[name]
+        grid_m = int(spec["loss_grid_m"])
+        if spec["type"] == "categorical":
+            size = max(1, round(prediction.shape[-1] * 10 / grid_m))
+            prediction = F.adaptive_avg_pool2d(prediction, (size, size))
+            one_hot = F.one_hot(target.long(), num_classes=prediction.shape[1]).permute(0, 3, 1, 2).float()
+            target = F.adaptive_avg_pool2d(one_hot, (size, size)).argmax(dim=1)
+            mask = F.adaptive_avg_pool2d(mask, (size, size))
+            value = F.cross_entropy(prediction, target, reduction="none")
+            value = (value * mask[:, 0]).sum() / mask[:, 0].sum().clamp_min(1)
+        else:
+            if spec.get("shift_pixels", 0):
+                value = _shift_invariant_l1(prediction, target, mask, int(spec["shift_pixels"]))
+            else:
+                prediction, target, mask = (_pool_continuous(item, grid_m) for item in (prediction, target, mask))
+                value = (torch.abs(prediction - target) * mask).sum() / mask.sum().clamp_min(1)
+        components[f"reconstruction_{name}"] = value
+        reconstruction = reconstruction + float(spec["weight"]) * value
+    flat = teacher["embedding"].permute(0, 2, 3, 1).reshape(-1, teacher["embedding"].shape[1])
+    rotated = torch.roll(flat, max(1, flat.shape[0] // 2), dims=0)
+    uniformity = (flat * rotated).sum(dim=1).abs().mean()
+    consistency = 1.0 - (teacher["embedding"] * student["embedding"]).sum(dim=1).mean()
+    pooled = F.normalize(teacher["embedding"].mean(dim=(2, 3)), dim=1)
+    normalized_text = F.normalize(text_target, dim=1)
+    logits = pooled @ normalized_text.transpose(0, 1)
+    labels = torch.arange(len(logits), device=logits.device)
+    text_alignment = 0.5 * (F.cross_entropy(logits, labels) + F.cross_entropy(logits.transpose(0, 1), labels))
+    total = (weights["reconstruction"] * reconstruction + weights["uniformity"] * uniformity
+             + weights["consistency"] * consistency + weights["text"] * text_alignment)
+    components.update(reconstruction=reconstruction, uniformity=uniformity,
+                      consistency=consistency, text_alignment=text_alignment, total=total)
+    return total, components
+
+
+def quantize_embeddings(embedding, power=2, scale=127.5):
+    transformed = embedding.abs().pow(1.0 / power) * embedding.sign()
+    return torch.round(transformed * scale).clamp(-127, 127).to(torch.int8)
+
+
+def dequantize_embeddings(quantized, power=2, scale=127.5):
+    values = quantized.float() / scale
+    return values.abs().pow(power) * values.sign()

scripts/fake_data.py

@@ -0,0 +1,118 @@
+"""Generate sparse synthetic observations following the paper-confirmed data contract."""
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+
+ROOT = Path(__file__).resolve().parents[1]
+DAY_MS = 86_400_000
+TARGET_ONLY_SLOTS = 3
+
+
+def load_config():
+    return yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+
+
+def standardized_signal(rng, count, frames, channels, size, source_phase):
+    y, x = np.mgrid[-1:1:complex(size), -1:1:complex(size)].astype(np.float32)
+    base = np.sin(3 * np.pi * x) * np.cos(2 * np.pi * y) + 0.4 * x + 0.2 * y
+    values = np.empty((count, frames, channels, size, size), np.float32)
+    for sample in range(count):
+        phase = rng.uniform(0, 2 * np.pi) + source_phase
+        for step in range(frames):
+            seasonal = np.sin(2 * np.pi * step / max(frames, 1) + phase)
+            for channel in range(channels):
+                values[sample, step, channel] = base + 0.08 * channel + 0.25 * seasonal
+        values[sample] += rng.normal(0, 0.04, values[sample].shape)
+    values -= values.mean(axis=(0, 1, 3, 4), keepdims=True)
+    values /= values.std(axis=(0, 1, 3, 4), keepdims=True).clip(1e-6)
+    return np.clip(values, -6, 6).astype(np.float32)
+
+
+def quality_mask(rng, count, frames, channels, size, sparse=False):
+    probability = 0.08 if not sparse else 0.92
+    mask = rng.random((count, frames, 1, size, size)) > probability
+    if not sparse:
+        mask[:, :, :, :2] = False
+    return np.repeat(mask, channels, axis=2).astype(np.float32)
+
+
+def make_split(path, count, config, seed):
+    rng = np.random.default_rng(seed)
+    data, size = config["data"], config["data"]["image_size"]
+    start = np.datetime64("2020-01-01", "ms").astype(np.int64)
+    payload = {"format_version": np.asarray(data["format_version"]), "data_source": np.asarray("synthetic")}
+
+    for source_index, (name, spec) in enumerate(data["input_sources"].items()):
+        frames, channels = spec["timesteps"], spec["channels"]
+        payload[name] = standardized_signal(rng, count, frames, channels, size, source_index)
+        days = np.linspace(0, 364, frames, dtype=np.int64)
+        payload[f"timestamps_{name}"] = np.tile(start + days * DAY_MS + source_index, (count, 1))
+        available = np.ones((count, frames), np.bool_)
+        available[:, -max(1, frames // 10):] = False
+        payload[f"frame_available_{name}"] = available
+        channel_available = np.ones((count, frames, channels), np.bool_)
+        if name == "sentinel1":
+            channel_available[..., :4] = False
+            for sample in range(count):
+                for step in range(frames):
+                    pair = (0, 1) if (sample + step) % 2 == 0 else (2, 3)
+                    channel_available[sample, step, list(pair)] = True
+        payload[f"channel_available_{name}"] = channel_available
+        payload[f"pixel_valid_{name}"] = quality_mask(rng, count, frames, channels, size)
+        geometry = np.zeros((count, frames, config["model"]["max_geometry_dim"]), np.float32)
+        if name == "sentinel1":
+            geometry[..., 0] = rng.uniform(-np.pi, np.pi, (count, frames))
+            geometry[..., 1] = rng.uniform(1.6, 1.8, (count, frames))
+        payload[f"geometry_{name}"] = geometry
+
+    for target_index, (name, spec) in enumerate(data["target_sources"].items()):
+        if name in data["input_sources"]:
+            continue
+        channels = 1 if spec["type"] == "categorical" else spec["channels"]
+        values = standardized_signal(rng, count, TARGET_ONLY_SLOTS, channels, size, target_index + 3)
+        if spec["type"] == "categorical":
+            bins = np.linspace(values.min(), values.max(), spec["channels"] + 1)[1:-1]
+            values = np.digitize(values[:, :, 0], bins).astype(np.int64)
+        payload[f"target_sequence_{name}"] = values
+        payload[f"target_timestamps_{name}"] = np.tile(
+            start + np.asarray([30, 182, 330], np.int64) * DAY_MS + target_index, (count, 1)
+        )
+        mask_channels = 1 if spec["type"] == "categorical" else spec["channels"]
+        payload[f"target_pixel_valid_{name}"] = quality_mask(
+            rng, count, TARGET_ONLY_SLOTS, mask_channels, size, sparse=name == "gedi"
+        )
+        geometry = np.zeros((count, TARGET_ONLY_SLOTS, config["model"]["max_geometry_dim"]), np.float32)
+        if name == "palsar2":
+            geometry[..., 0] = rng.integers(0, 2, (count, TARGET_ONLY_SLOTS))
+            geometry[..., 1] = rng.integers(0, 2, (count, TARGET_ONLY_SLOTS))
+        payload[f"target_geometry_{name}"] = geometry
+        payload[f"target_frame_available_{name}"] = np.ones((count, TARGET_ONLY_SLOTS), np.bool_)
+
+    payload["support_period"] = np.tile(np.array([start, start + 365 * DAY_MS], np.int64), (count, 1))
+    payload["valid_period"] = np.tile(np.array([start + 60 * DAY_MS, start + 300 * DAY_MS], np.int64), (count, 1))
+    text = rng.normal(size=(count, config["model"]["embedding_dim"])).astype(np.float32)
+    payload["text_target"] = text / np.linalg.norm(text, axis=1, keepdims=True)
+    np.savez_compressed(path, **payload)
+
+
+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, (name, count) in enumerate((("train.npz", config["data"]["train_samples"]),
+                                            ("test.npz", config["data"]["test_samples"]))):
+        target = data_dir / name
+        if args.force or not target.exists():
+            make_split(target, count, config, config["seed"] + offset)
+        print(f"generated={target.relative_to(ROOT)} samples={count} format={config['data']['format_version']}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,60 @@
+"""Generate float and paper-style signed-int8 annual embedding fields."""
+
+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))
+sys.path.insert(0, str(ROOT / "scripts"))
+from model.alphaearthfoundations import AlphaEarthFoundations, dequantize_embeddings, quantize_embeddings
+from train import AEFDataset, device_from_config, unpack
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    torch.manual_seed(config["seed"])
+    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 do not match")
+    model = AlphaEarthFoundations(checkpoint["input_sources"], checkpoint["target_sources"], checkpoint["model_config"]).to(device)
+    model.load_state_dict(checkpoint["model"])
+    model.eval()
+    dataset = AEFDataset(ROOT / config["data"]["root"] / "test.npz", config)
+    embeddings, quantized, restored = [], [], []
+    reconstruction = {name: [] for name in config["data"]["target_sources"]}
+    selected_targets = {name: [] for name in config["data"]["target_sources"]}
+    selected_masks = {name: [] for name in config["data"]["target_sources"]}
+    with torch.no_grad():
+        for index in range(len(dataset)):
+            batch = {key: value.unsqueeze(0) for key, value in dataset[index].items()}
+            (sources, timestamps, frame_available, targets, masks, target_times,
+             target_periods, geometry) = unpack(batch, config, device)
+            output = model(sources, timestamps, batch["valid_period"].to(device), frame_available,
+                           target_times, geometry, target_periods)
+            q = quantize_embeddings(output["embedding"])
+            embeddings.append(output["embedding"].cpu().numpy())
+            quantized.append(q.cpu().numpy())
+            restored.append(dequantize_embeddings(q).cpu().numpy())
+            for name, values in output["reconstructions"].items():
+                reconstruction[name].append(values.cpu().numpy())
+                selected_targets[name].append(targets[name].cpu().numpy())
+                selected_masks[name].append(masks[name].cpu().numpy())
+    output_dir = ROOT / config["paths"]["inference_dir"]
+    output_dir.mkdir(parents=True, exist_ok=True)
+    payload = {"embedding": np.concatenate(embeddings), "embedding_s8_power2": np.concatenate(quantized),
+               "embedding_dequantized": np.concatenate(restored)}
+    payload.update({f"reconstruction_{name}": np.concatenate(values) for name, values in reconstruction.items()})
+    payload.update({f"target_{name}": np.concatenate(values) for name, values in selected_targets.items()})
+    payload.update({f"mask_{name}": np.concatenate(values) for name, values in selected_masks.items()})
+    np.savez_compressed(output_dir / "predictions.npz", **payload)
+    print(f"predictions={(output_dir / 'predictions.npz').relative_to(ROOT)}")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,99 @@
+"""Evaluate reconstruction, spherical embeddings, quantization and low-shot transfer."""
+
+import json
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import yaml
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def balanced_accuracy(target, prediction):
+    scores = [(prediction[target == label] == label).mean() for label in np.unique(target)]
+    return float(np.mean(scores))
+
+
+def transfer_metrics(embedding, class_target, regression_target, seed):
+    rng = np.random.default_rng(seed)
+    features = embedding.transpose(0, 2, 3, 1).reshape(-1, embedding.shape[1])
+    classes = class_target.reshape(-1)
+    regression = regression_target.reshape(-1)
+    train_indices, test_indices = [], []
+    for label in np.unique(classes):
+        indices = np.flatnonzero(classes == label)
+        rng.shuffle(indices)
+        split = min(10, max(1, len(indices) // 3))
+        train_indices.extend(indices[:split])
+        test_indices.extend(indices[split:])
+    train_indices, test_indices = np.asarray(train_indices), np.asarray(test_indices)
+    x_train, x_test = features[train_indices], features[test_indices]
+    y_train, y_test = classes[train_indices], classes[test_indices]
+    distances = ((x_test[:, None] - x_train[None]) ** 2).sum(axis=-1)
+    transfer = {}
+    for k in (1, 3):
+        neighbors = np.argpartition(distances, min(k, len(x_train)) - 1, axis=1)[:, :k]
+        votes = y_train[neighbors]
+        prediction = np.asarray([np.bincount(row).argmax() for row in votes])
+        transfer[f"knn_k{k}_balanced_accuracy"] = balanced_accuracy(y_test, prediction)
+    labels = np.unique(classes)
+    one_hot = np.stack([np.where(y_train == label, 1.0, -1.0) for label in labels], axis=1)
+    design = np.column_stack([x_train, np.ones(len(x_train))])
+    coefficients = np.linalg.lstsq(design, one_hot, rcond=None)[0]
+    class_prediction = labels[np.argmax(np.column_stack([x_test, np.ones(len(x_test))]) @ coefficients, axis=1)]
+    transfer["linear_balanced_accuracy"] = balanced_accuracy(y_test, class_prediction)
+    regression_coefficients = np.linalg.lstsq(design, regression[train_indices], rcond=None)[0]
+    regression_prediction = np.column_stack([x_test, np.ones(len(x_test))]) @ regression_coefficients
+    residual = ((regression[test_indices] - regression_prediction) ** 2).sum()
+    total = ((regression[test_indices] - regression[test_indices].mean()) ** 2).sum()
+    transfer["linear_regression_r2"] = float(1.0 - residual / max(total, 1e-12))
+    transfer["train_pixels"] = int(len(train_indices))
+    transfer["test_pixels"] = int(len(test_indices))
+    return transfer
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    predictions = np.load(ROOT / config["paths"]["inference_dir"] / "predictions.npz")
+    embedding, restored = predictions["embedding"], predictions["embedding_dequantized"]
+    metrics = {
+        "samples": int(len(embedding)),
+        "mean_embedding_norm": float(np.linalg.norm(embedding, axis=1).mean()),
+        "s8_power2_quantization_mae": float(np.abs(embedding - restored).mean()),
+        "reconstruction_mae": {},
+        "low_shot_transfer": transfer_metrics(
+            embedding, predictions["target_nlcd"], predictions["target_sentinel2"][:, 0], config["seed"]
+        ),
+    }
+    for name, spec in config["data"]["target_sources"].items():
+        prediction = predictions[f"reconstruction_{name}"]
+        target = predictions[f"target_{name}"]
+        mask = predictions[f"mask_{name}"]
+        if spec["type"] == "categorical":
+            metrics["reconstruction_mae"][name] = float(
+                (((prediction.argmax(axis=1) != target) * mask[:, 0]).sum()) / max(mask[:, 0].sum(), 1)
+            )
+        else:
+            metrics["reconstruction_mae"][name] = float((np.abs(prediction - target) * mask).sum() / max(mask.sum(), 1))
+    evaluation_dir = ROOT / config["paths"]["evaluation_dir"]
+    evaluation_dir.mkdir(parents=True, exist_ok=True)
+    (evaluation_dir / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
+    rgb = embedding[0, [1, 16, 9]].transpose(1, 2, 0)
+    rgb = np.clip((rgb + 0.3) / 0.6, 0, 1)
+    figure, axes = plt.subplots(1, 2, figsize=(8, 4))
+    axes[0].imshow(rgb)
+    axes[0].set_title("AEF axes A01/A16/A09")
+    axes[1].imshow(predictions["target_nlcd"][0], cmap="tab20", vmin=0, vmax=15)
+    axes[1].set_title("Synthetic NLCD target")
+    for axis in axes:
+        axis.axis("off")
+    figure.tight_layout()
+    figure.savefig(evaluation_dir / "comparison.png", dpi=160)
+    plt.close(figure)
+    print(json.dumps(metrics, indent=2))
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,199 @@
+"""Train AEF reconstruction, uniformity, consistency and text objectives."""
+
+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.alphaearthfoundations import AlphaEarthFoundations, compute_losses
+
+
+class AEFDataset(Dataset):
+    def __init__(self, path, config):
+        self.data = np.load(path)
+        self.input_sources = config["data"]["input_sources"]
+        self.target_sources = config["data"]["target_sources"]
+        format_version = str(self.data["format_version"])
+        if format_version != config["data"]["format_version"]:
+            raise ValueError(f"Expected format {config['data']['format_version']}, got {format_version}")
+        for name, spec in self.input_sources.items():
+            expected = (spec["timesteps"], spec["channels"], config["data"]["image_size"], config["data"]["image_size"])
+            if self.data[name].shape[1:] != expected:
+                raise ValueError(f"{name} has shape {self.data[name].shape[1:]}, expected {expected}")
+
+    def __len__(self):
+        return len(self.data["valid_period"])
+
+    def __getitem__(self, index):
+        item = {name: torch.from_numpy(self.data[name][index]).float() for name in self.input_sources}
+        item.update({f"timestamps_{name}": torch.from_numpy(self.data[f"timestamps_{name}"][index]).long()
+                     for name in self.input_sources})
+        for name in self.input_sources:
+            item[f"frame_available_{name}"] = torch.from_numpy(self.data[f"frame_available_{name}"][index])
+            item[f"channel_available_{name}"] = torch.from_numpy(self.data[f"channel_available_{name}"][index])
+            item[f"pixel_valid_{name}"] = torch.from_numpy(self.data[f"pixel_valid_{name}"][index]).float()
+            item[f"geometry_{name}"] = torch.from_numpy(self.data[f"geometry_{name}"][index]).float()
+        for name, spec in self.target_sources.items():
+            if name in self.input_sources:
+                continue
+            values = torch.from_numpy(self.data[f"target_sequence_{name}"][index])
+            item[f"target_sequence_{name}"] = values.long() if spec["type"] == "categorical" else values.float()
+            item[f"target_timestamps_{name}"] = torch.from_numpy(self.data[f"target_timestamps_{name}"][index]).long()
+            item[f"target_pixel_valid_{name}"] = torch.from_numpy(self.data[f"target_pixel_valid_{name}"][index]).float()
+            item[f"target_geometry_{name}"] = torch.from_numpy(self.data[f"target_geometry_{name}"][index]).float()
+            item[f"target_frame_available_{name}"] = torch.from_numpy(self.data[f"target_frame_available_{name}"][index])
+        item["support_period"] = torch.from_numpy(self.data["support_period"][index]).long()
+        item["valid_period"] = torch.from_numpy(self.data["valid_period"][index]).long()
+        item["text_target"] = torch.from_numpy(self.data["text_target"][index]).float()
+        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 _select_indices(available):
+    selected = []
+    for row in available:
+        candidates = torch.nonzero(row, as_tuple=False).flatten()
+        choice = torch.randint(len(candidates), (), device=candidates.device)
+        selected.append(candidates[choice])
+    return torch.stack(selected).to(available.device)
+
+
+def _gather(values, indices):
+    return values[torch.arange(len(values), device=values.device), indices]
+
+
+def unpack(batch, config, device, remove_input_targets=True):
+    input_names, target_specs = config["data"]["input_sources"], config["data"]["target_sources"]
+    sources, frame_available = {}, {}
+    timestamps = {name: batch[f"timestamps_{name}"].to(device) for name in input_names}
+    targets, masks, target_times, target_periods, geometry = {}, {}, {}, {}, {}
+    for name in input_names:
+        values = batch[name].to(device)
+        available = batch[f"frame_available_{name}"].to(device).clone()
+        channel_available = batch[f"channel_available_{name}"].to(device)
+        values = values * channel_available[:, :, :, None, None]
+        indices = _select_indices(available)
+        targets[name] = _gather(values, indices)
+        masks[name] = _gather(batch[f"pixel_valid_{name}"].to(device), indices) * _gather(channel_available, indices)[:, :, None, None]
+        target_times[name] = _gather(timestamps[name], indices)
+        target_periods[name] = torch.stack([
+            target_times[name] - 5 * 86_400_000, target_times[name] + 5 * 86_400_000
+        ], dim=1)
+        geometry[name] = _gather(batch[f"geometry_{name}"].to(device), indices)
+        if remove_input_targets:
+            available[torch.arange(len(available), device=device), indices] = False
+        sources[name] = values * available[:, :, None, None, None]
+        frame_available[name] = available
+    for name in target_specs:
+        if name in input_names:
+            continue
+        available = batch[f"target_frame_available_{name}"].to(device)
+        indices = _select_indices(available)
+        targets[name] = _gather(batch[f"target_sequence_{name}"].to(device), indices)
+        masks[name] = _gather(batch[f"target_pixel_valid_{name}"].to(device), indices)
+        target_times[name] = _gather(batch[f"target_timestamps_{name}"].to(device), indices)
+        target_periods[name] = torch.stack([
+            target_times[name] - 5 * 86_400_000, target_times[name] + 5 * 86_400_000
+        ], dim=1)
+        geometry[name] = _gather(batch[f"target_geometry_{name}"].to(device), indices)
+    return sources, timestamps, frame_available, targets, masks, target_times, target_periods, geometry
+
+
+def perturb_sources(sources, frame_available, timestamps, support_period, config):
+    perturbed, perturbed_available = {}, {}
+    source_dropout, frame_dropout = config["train"]["source_dropout"], config["train"]["frame_dropout"]
+    for name, values in sources.items():
+        output, available = values + 0.01 * torch.randn_like(values), frame_available[name].clone()
+        source_mask = torch.rand(values.shape[0], 1, 1, 1, 1, device=values.device) < source_dropout[name]
+        strategy = torch.randint(3, (1,), device=values.device).item()
+        if strategy == 0:
+            dropped = torch.rand_like(available.float()) < frame_dropout[name]
+        else:
+            midpoint = support_period[:, :1] + (support_period[:, 1:] - support_period[:, :1]) // 2
+            dropped = timestamps[name] >= midpoint if strategy == 1 else timestamps[name] < midpoint
+        available = available & ~dropped & ~source_mask[:, 0, 0, 0]
+        perturbed[name] = output * available[:, :, None, None, None]
+        perturbed_available[name] = available
+    return perturbed, perturbed_available
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    torch.manual_seed(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 = AEFDataset(ROOT / config["data"]["root"] / "train.npz", config)
+    sampler = DistributedSampler(dataset, shuffle=True) if distributed else None
+    loader = DataLoader(dataset, batch_size=config["train"]["batch_size"], sampler=sampler,
+                        shuffle=sampler is None, num_workers=config["train"]["num_workers"])
+    model = AlphaEarthFoundations(config["data"]["input_sources"], config["data"]["target_sources"], config["model"]).to(device)
+    if distributed:
+        model = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None)
+    optimizer = torch.optim.Adam(model.parameters(), lr=config["train"]["learning_rate"],
+                                 weight_decay=config["train"]["weight_decay"])
+    weights = {name: config["train"][f"{name}_weight"] for name in ("reconstruction", "uniformity", "consistency", "text")}
+    history = []
+    for epoch in range(config["train"]["epochs"]):
+        if sampler:
+            sampler.set_epoch(epoch)
+        model.train()
+        totals = {}
+        for batch in loader:
+            sources, timestamps, frame_available, targets, masks, target_times, target_periods, geometry = unpack(batch, config, device)
+            valid_period = batch["valid_period"].to(device)
+            teacher = model(sources, timestamps, valid_period, frame_available, target_times, geometry, target_periods)
+            perturbed, perturbed_available = perturb_sources(
+                sources, frame_available, timestamps, batch["support_period"].to(device), config
+            )
+            student = model(perturbed, timestamps, valid_period, perturbed_available)
+            loss, components = compute_losses(teacher, student, targets, masks, batch["text_target"].to(device),
+                                               config["data"]["target_sources"], weights)
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+            for name, value in components.items():
+                totals[name] = totals.get(name, 0.0) + float(value.detach())
+        metrics = {name: value / len(loader) for name, 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"], "input_sources": config["data"]["input_sources"],
+                    "target_sources": config["data"]["target_sources"], "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()