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,156 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- OneScience
+- Earth Science
+- Multi-Temporal Earth Observation
+- Remote Sensing Foundation Model
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">Prithvi-EO</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+Prithvi-EO-2.0 is a foundation model for multi-temporal Earth observation data. It encodes HLS multispectral time series, image acquisition dates, and geographic locations into unified representations and reconstructs masked spatiotemporal patches through a Masked Autoencoder. It can be used for remote-sensing classification, semantic segmentation, regression, and environmental change monitoring.
+
+Paper: Prithvi-EO-2.0: A Versatile Multi-Temporal Foundation Model for Earth Observation Applications  
+https://arxiv.org/abs/2412.02732
+
+# Model Description
+
+Prithvi-EO-2.0 was jointly proposed by IBM, NASA, and the Jülich Supercomputing Centre. The model was trained on 4.2 million global four-timestamp samples from NASA's Harmonized Landsat Sentinel-2 dataset. It is suitable for disaster response, land-cover and crop mapping, ecosystem dynamics monitoring, and remote-sensing image classification, segmentation, and regression tasks.
+
+# Use Cases
+
+| Use Case | Description |
+| :---: | :--- |
+| Multi-temporal representation | Uses a Transformer to jointly encode spatial and temporal information from four time steps. |
+| Remote-sensing image reconstruction | Uses a 3D Masked Autoencoder to reconstruct masked multispectral spatiotemporal patches. |
+| Spatiotemporal metadata modeling | Integrates year, day of year, latitude, and longitude, with support for randomly dropping metadata during training. |
+| Local engineering validation | Uses a small number of synthetic HLS samples to verify the training, inference, evaluation, visualization, and checkpoint workflows. |
+| Multi-GPU training | Launches distributed data-parallel training through `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/PrithviEO --local-dir ./PrithviEO
+cd PrithviEO
+```
+
+### 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 an official-size model requires large-scale accelerated computing 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
+# Installation with uv 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
+# Installation with uv 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 preserve the four time steps and six HLS common bands used in the paper. The band order is Blue, Green, Red, Narrow NIR, SWIR1, and SWIR2, and normalization uses the officially published means and standard deviations.
+
+The synthetic data preserve the `4×224×224` spatiotemporal dimensions used for training in the paper, and each sample's complete image tensor has dimensions `6×4×224×224`. Only the number of samples, model width, model depth, and number of training epochs are reduced. This setup validates the 3D Patch Embedding, spatiotemporal positional encoding, temporal and location metadata encoding, and MAE training workflow; it does not represent the data distribution or training scale of the official 4.2 million HLS samples.
+
+```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 HLS time series, the official 300M or 600M configuration, and the complete training schedule.
+
+```text
+result/checkpoints/prithvi_eo.pt
+result/training/metrics.json
+```
+
+### Trained Weights
+
+This repository does not include synthetic or official weights in `weight/`. IBM and NASA have published the Prithvi-EO-2.0 tiny, 100M, 300M, and 600M weights on Hugging Face, including the TL versions with temporal and location encodings:
+
+https://huggingface.co/ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL
+
+This repository is a reduced, independent engineering implementation. Its model parameter names and dimensions are not compatible with the official weights. To use the official weights, use TerraTorch or the implementation and data preprocessing workflow provided by the official repository.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference loads the training checkpoint, generates the CLS embedding, spatiotemporal patch embeddings, multi-temporal reconstructions of masked patches, and the mask, and saves them to:
+
+```text
+result/output/predictions.npz
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+The evaluation measures masked-patch MSE, full spatiotemporal reconstruction MAE, reconstruction error for each time step, and embedding norms. It also generates comparison plots of the input images, reconstructions, and absolute errors for all four time steps. Results on synthetic data only validate the engineering workflow and do not represent GEO-Bench performance or performance on real downstream tasks reported in the paper.
+
+```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 based on the publicly available specifications from the Prithvi-EO-2.0 paper.
+
+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,159 @@
+---
+license: Apache-2.0
+language:
+- en
+- zh
+tags:
+- OneScience
+- 地球科学
+- 多时相地球观测
+- 遥感基础模型
+frameworks: PyTorch
+datasets:
+  - Harmonized Landsat Sentinel-2
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">Prithvi-EO</span>
+  </strong>
+</p>
+
+# 模型介绍
+
+Prithvi-EO-2.0 是面向多时相地球观测数据的基础模型，将 HLS 多光谱时间序列、影像获取日期和地理位置编码为统一表征，并通过 Masked Autoencoder 重建被遮挡的时空 Patch，可用于遥感分类、语义分割、回归和环境变化监测。
+
+论文：Prithvi-EO-2.0: A Versatile Multi-Temporal Foundation Model for Earth Observation Applications  
+https://arxiv.org/abs/2412.02732
+
+# 模型描述
+
+Prithvi-EO-2.0 由 IBM、NASA 和 Jülich Supercomputing Centre 联合提出。模型使用 NASA Harmonized Landsat Sentinel-2 数据集中 420 万个全球四时相样本训练。模型适用于灾害响应、土地覆盖与作物制图、生态系统动态监测，以及遥感影像分类、分割和回归任务。
+
+# 适用场景
+
+| 场景 | 说明 |
+| :---: | :--- |
+| 多时相表征 | 使用 Transformer 同时编码四个时间步的空间和时间信息。 |
+| 遥感影像重建 | 使用 3D Masked Autoencoder 重建被遮挡的多光谱时空 Patch。 |
+| 时空元数据建模 | 融合年份、年积日、纬度和经度，并支持训练时随机丢弃元数据。 |
+| 本地工程验证 | 使用少量虚拟 HLS 样本检查训练、推理、评估、可视化和 checkpoint 流程。 |
+| 多卡训练 | 通过 `torchrun` 启动分布式数据并行训练。 |
+
+# 使用说明
+
+## 1.OneCode
+
+可通过 OneCode 在线环境体验智能化一键式 AI4S 编程：
+
+[点击体验智能化一键式 AI4S 编程](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2.下载安装
+
+```bash
+modelscope download --model OneScience/Prithvi-EO --local_dir ./PrithviEO
+cd PrithviEO
+```
+
+### 环境依赖
+
+**硬件要求**
+
+- 推荐使用 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`。虚拟数据保持论文的四时间步与六个 HLS 公共波段，波段顺序为 Blue、Green、Red、Narrow NIR、SWIR1 和 SWIR2，并使用官方公开的均值与标准差完成归一化。
+
+虚拟数据保持论文训练使用的 4×224×224 时空尺寸，每个样本的完整影像张量为 6×4×224×224。当前仅缩小样本数量、模型宽度、模型深度和训练周期，用于验证 3D Patch Embedding、时空位置编码、时间与位置元数据编码和 MAE 训练流程，不代表官方 420 万个 HLS 样本的数据分布与训练规模。
+
+```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
+```
+
+默认配置面向快速流程验证；开展正式实验时，应使用真实 HLS 时间序列、官方 300M 或 600M 配置和完整训练周期。
+
+```text
+result/checkpoints/prithvi_eo.pt
+result/training/metrics.json
+```
+
+### 训练权重
+
+本仓库不在 `weight/` 中内置虚拟权重或官方权重。IBM 和 NASA 已在 Hugging Face 公开 Prithvi-EO-2.0 的 tiny、100M、300M 和 600M 权重，包括带时间与位置编码的 TL 版本：
+
+https://huggingface.co/ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL
+
+本仓库是缩小的独立工程实现，模型参数名称和尺寸不与官方权重兼容。需要使用官方权重时，应采用 TerraTorch 或官方仓库提供的实现和数据预处理流程。
+
+### 推理
+
+```bash
+python scripts/inference.py
+```
+
+推理加载训练 checkpoint，生成 CLS embedding、时空 Patch embedding、被遮挡 Patch 的多时相重建结果和 Mask，并保存到：
+
+```text
+result/output/predictions.npz
+```
+
+### 评估和可视化
+
+```bash
+python scripts/result.py
+```
+
+评估模型的 masked-patch MSE、完整时空重建 MAE、各时间步重建误差和 embedding 范数。评估过程同时生成四个时间步的输入影像、重建结果和绝对误差对比图。虚拟数据结果仅用于验证工程流程，不代表论文中的 GEO-Bench 或真实下游任务性能。
+
+```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 |
+
+# 引用与许可证
+
+本仓库为 Prithvi-EO-2.0 论文公开规格的独立工程复现版本。
+
+本仓库代码、官方模型权重和数据的使用仍应以各自项目中的许可证及使用条款为准。

conf/config.yaml

@@ -0,0 +1,68 @@
+seed: 42
+data:
+  root: data
+  format_version: prithvi_eo_2_synthetic_engineering_v1
+  train_samples: 4
+  test_samples: 2
+  channels: 6
+  frames: 4
+  image_size: 224
+  band_names: [Blue, Green, Red, Narrow_NIR, SWIR1, SWIR2]
+  band_codes: [B02, B03, B04, B8A, B11, B12]
+  mean: [1087.0, 1342.0, 1433.0, 2734.0, 1958.0, 1363.0]
+  std: [2248.0, 2179.0, 2178.0, 1850.0, 1242.0, 1049.0]
+model:
+  input_size: [4, 224, 224]
+  patch_size: [1, 16, 16]
+  channels: 6
+  mask_ratio: 0.75
+  encoder_dim: 96
+  encoder_depth: 2
+  encoder_heads: 4
+  decoder_dim: 64
+  decoder_depth: 1
+  decoder_heads: 4
+  mlp_ratio: 4
+  metadata_dropout: 0.1
+  norm_pix_loss: false
+paper_model_300M_TL:
+  input_size: [4, 224, 224]
+  patch_size: [1, 16, 16]
+  encoder_dim: 1024
+  encoder_depth: 24
+  encoder_heads: 16
+  decoder_dim: 512
+  decoder_depth: 8
+  decoder_heads: 16
+  parameter_count: 300M
+paper_model_600M_TL:
+  input_size: [4, 224, 224]
+  patch_size: [1, 14, 14]
+  encoder_dim: 1280
+  encoder_depth: 32
+  encoder_heads: 16
+  decoder_dim: 512
+  decoder_depth: 8
+  decoder_heads: 16
+  parameter_count: 600M
+paper_training:
+  samples: 4.2M
+  validation_samples: 46k
+  epochs: 400
+  global_batch_size: 3840
+  maximum_learning_rate: 0.0005
+  warmup_epochs: 40
+  weight_decay: 0.05
+train:
+  epochs: 1
+  batch_size: 2
+  learning_rate: 0.0005
+  weight_decay: 0.000001
+  num_workers: 0
+runtime:
+  device: auto
+paths:
+  checkpoint: result/checkpoints/prithvi_eo.pt
+  training_metrics: result/training/metrics.json
+  inference_dir: result/output
+  evaluation_dir: result/evaluation

config.json

@@ -0,0 +1,44 @@
+{
+  "model_name": "PrithviEO",
+  "model_type": "prithvi_eo",
+  "architectures": ["PrithviEO2", "CoordinateEncoder", "Transformer"],
+  "framework": "PyTorch",
+  "domain": "earth-observation",
+  "task": "multi-temporal-remote-sensing-representation-learning",
+  "implementation": {
+    "entry_point": "model/prithvi_eo.py",
+    "scope": "engineering reproduction of the Prithvi-EO-2.0 temporal-location masked autoencoder",
+    "train_script": "scripts/train.py",
+    "inference_script": "scripts/inference.py",
+    "evaluation_script": "scripts/result.py",
+    "synthetic_data_script": "scripts/fake_data.py"
+  },
+  "architecture": {
+    "family": "three-dimensional vision-transformer masked autoencoder",
+    "input_size": [4, 224, 224],
+    "patch_size": [1, 16, 16],
+    "in_channels": 6,
+    "mask_ratio": 0.75,
+    "encoder_dim": 96,
+    "encoder_depth": 2,
+    "encoder_heads": 4,
+    "decoder_dim": 64,
+    "decoder_depth": 1,
+    "decoder_heads": 4,
+    "metadata": ["year", "day_of_year", "latitude", "longitude"],
+    "outputs": ["loss", "embedding", "patch_embeddings", "reconstruction", "mask"]
+  },
+  "data": {
+    "datasets": ["Harmonized Landsat Sentinel-2"],
+    "protocol": "prithvi_eo_2_synthetic_engineering_v1",
+    "format": "NPZ",
+    "train_file": "data/train.npz",
+    "test_file": "data/test.npz",
+    "input_shape": ["N", 6, 4, 224, 224],
+    "temporal_shape": ["N", 4, 2],
+    "location_shape": ["N", 2],
+    "bands": ["B02", "B03", "B04", "B8A", "B11", "B12"],
+    "required_metadata": ["format_version", "data_source", "temporal_coords", "location_coords"]
+  },
+  "configuration_sources": ["conf/config.yaml", "model/prithvi_eo.py", "scripts/fake_data.py", "scripts/train.py", "scripts/inference.py", "scripts/result.py"]
+}

configuration.json

@@ -0,0 +1,12 @@
+{
+  "framework": "PyTorch",
+  "task": "multi_temporal_earth_observation_representation_learning",
+  "model": "Prithvi-EO-2.0",
+  "input_format": "BCTHW",
+  "protocol": "engineering_3d_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/prithvi_eo.py

@@ -0,0 +1,175 @@
+"""Pure-PyTorch engineering reproduction of the Prithvi-EO-2.0 TL MAE."""
+
+import math
+
+import torch
+from torch import nn
+
+
+def sincos_1d(positions, dim):
+    if dim % 2:
+        raise ValueError("sine/cosine dimensions must be even")
+    frequencies = torch.exp(
+        torch.arange(dim // 2, device=positions.device, dtype=positions.dtype)
+        * (-math.log(10000.0) / max(dim // 2, 1))
+    )
+    angles = positions.unsqueeze(-1) * frequencies
+    return torch.cat((angles.sin(), angles.cos()), dim=-1)
+
+
+def sincos_3d(frames, height, width, dim, device, dtype):
+    if dim % 16:
+        raise ValueError("3D position dimension must be divisible by 16")
+    width_dim, height_dim, time_dim = 6 * dim // 16, 6 * dim // 16, 4 * dim // 16
+    time, row, column = torch.meshgrid(
+        torch.arange(frames, device=device, dtype=dtype),
+        torch.arange(height, device=device, dtype=dtype),
+        torch.arange(width, device=device, dtype=dtype),
+        indexing="ij",
+    )
+    return torch.cat((
+        sincos_1d(column.reshape(-1), width_dim),
+        sincos_1d(row.reshape(-1), height_dim),
+        sincos_1d(time.reshape(-1), time_dim),
+    ), dim=-1)
+
+
+def patchify(values, patch_size):
+    batch, channels, frames, height, width = values.shape
+    pt, ph, pw = patch_size
+    if frames % pt or height % ph or width % pw:
+        raise ValueError("input dimensions must be divisible by patch_size")
+    return values.reshape(
+        batch, channels, frames // pt, pt, height // ph, ph, width // pw, pw
+    ).permute(0, 2, 4, 6, 3, 5, 7, 1).reshape(batch, -1, pt * ph * pw * channels)
+
+
+def unpatchify(patches, channels, output_size, patch_size):
+    batch = patches.shape[0]
+    frames, height, width = output_size
+    pt, ph, pw = patch_size
+    return patches.reshape(
+        batch, frames // pt, height // ph, width // pw, pt, ph, pw, channels
+    ).permute(0, 7, 1, 4, 2, 5, 3, 6).reshape(batch, channels, frames, height, width)
+
+
+class Transformer(nn.Module):
+    def __init__(self, dim, depth, heads, mlp_ratio):
+        super().__init__()
+        layer = nn.TransformerEncoderLayer(
+            dim, heads, int(dim * mlp_ratio), activation="gelu", batch_first=True, norm_first=True
+        )
+        self.blocks = nn.TransformerEncoder(layer, depth)
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, values):
+        return self.norm(self.blocks(values))
+
+
+class CoordinateEncoder(nn.Module):
+    def __init__(self, dim, scale=0.1):
+        super().__init__()
+        if dim % 4:
+            raise ValueError("coordinate embedding dimension must be divisible by four")
+        self.dim = dim
+        self.scale = nn.Parameter(torch.tensor(float(scale)))
+
+    def forward(self, coordinates):
+        return self.scale * torch.cat((
+            sincos_1d(coordinates[..., 0], self.dim // 2),
+            sincos_1d(coordinates[..., 1], self.dim // 2),
+        ), dim=-1)
+
+
+class PrithviEO2(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = dict(config)
+        self.input_size = tuple(int(value) for value in config["input_size"])
+        self.patch_size = tuple(int(value) for value in config["patch_size"])
+        self.channels = int(config["channels"])
+        self.mask_ratio = float(config["mask_ratio"])
+        self.metadata_dropout = float(config["metadata_dropout"])
+        self.norm_pix_loss = bool(config.get("norm_pix_loss", False))
+        enc_dim, dec_dim = int(config["encoder_dim"]), int(config["decoder_dim"])
+        self.patch_embed = nn.Conv3d(
+            self.channels, enc_dim, kernel_size=self.patch_size, stride=self.patch_size
+        )
+        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"]),
+                                   float(config["mlp_ratio"]))
+        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"]),
+                                   float(config["mlp_ratio"]))
+        patch_volume = math.prod(self.patch_size) * self.channels
+        self.decoder_prediction = nn.Linear(dec_dim, patch_volume)
+        self.time_encoder = CoordinateEncoder(enc_dim)
+        self.location_encoder = CoordinateEncoder(enc_dim)
+        self.decoder_time_encoder = CoordinateEncoder(dec_dim)
+        self.decoder_location_encoder = CoordinateEncoder(dec_dim)
+
+    def _grid(self, pixels):
+        return tuple(size // patch for size, patch in zip(pixels.shape[-3:], self.patch_size))
+
+    def _metadata(self, temporal, location, grid, encoder=True):
+        frames, height, width = grid
+        time_encoder = self.time_encoder if encoder else self.decoder_time_encoder
+        location_encoder = self.location_encoder if encoder else self.decoder_location_encoder
+        temporal_embedding = time_encoder(temporal)
+        temporal_embedding = temporal_embedding[:, :, None, :].expand(-1, -1, height * width, -1).reshape(
+            len(temporal), frames * height * width, -1
+        )
+        location_embedding = location_encoder(location)[:, None, :].expand(-1, frames * height * width, -1)
+        if self.training and self.metadata_dropout:
+            time_keep = (torch.rand(len(temporal), 1, 1, device=temporal.device) >= self.metadata_dropout).to(temporal.dtype)
+            location_keep = (torch.rand(len(location), 1, 1, device=location.device) >= self.metadata_dropout).to(location.dtype)
+            temporal_embedding = temporal_embedding * time_keep
+            location_embedding = location_embedding * location_keep
+        return temporal_embedding + location_embedding
+
+    def _encoded_tokens(self, pixels, temporal, location):
+        grid = self._grid(pixels)
+        tokens = self.patch_embed(pixels).flatten(2).transpose(1, 2)
+        position = sincos_3d(*grid, tokens.shape[-1], tokens.device, tokens.dtype)
+        tokens = tokens + position[None] + self._metadata(temporal, location, grid, encoder=True)
+        return tokens, grid
+
+    def encode(self, pixels, temporal, location):
+        tokens, _ = self._encoded_tokens(pixels, temporal, location)
+        cls = self.cls_token.expand(len(pixels), -1, -1)
+        encoded = self.encoder(torch.cat((cls, tokens), dim=1))
+        return encoded[:, 0], encoded[:, 1:]
+
+    def forward(self, pixels, temporal, location, mask_ratio=None):
+        ratio = self.mask_ratio if mask_ratio is None else float(mask_ratio)
+        tokens, grid = self._encoded_tokens(pixels, temporal, location)
+        batch, length, dim = tokens.shape
+        keep = max(1, int(length * (1.0 - ratio)))
+        ordering = torch.rand(batch, length, device=pixels.device).argsort(dim=1)
+        visible_indices, masked_indices = ordering[:, :keep], ordering[:, keep:]
+        visible = tokens.gather(1, visible_indices[:, :, None].expand(-1, -1, dim))
+        encoded = self.encoder(torch.cat((self.cls_token.expand(batch, -1, -1), visible), dim=1))
+        embedding = encoded[:, 0]
+        visible_decoder = self.encoder_to_decoder(encoded[:, 1:])
+        decoder_tokens = self.mask_token.expand(batch, length, -1).clone()
+        decoder_tokens.scatter_(1, visible_indices[:, :, None].expand(-1, -1, visible_decoder.shape[-1]), visible_decoder)
+        position = sincos_3d(*grid, decoder_tokens.shape[-1], decoder_tokens.device, decoder_tokens.dtype)
+        decoder_tokens = decoder_tokens + position[None] + self._metadata(temporal, location, grid, encoder=False)
+        predictions = self.decoder_prediction(self.decoder(decoder_tokens))
+        targets = patchify(pixels, self.patch_size)
+        if self.norm_pix_loss:
+            mean, variance = targets.mean(dim=-1, keepdim=True), targets.var(dim=-1, keepdim=True)
+            targets = (targets - mean) / (variance + 1e-6).sqrt()
+        mask = torch.zeros(batch, length, device=pixels.device)
+        mask.scatter_(1, masked_indices, 1.0)
+        patch_mse = (predictions - targets).pow(2).mean(dim=-1)
+        loss = (patch_mse * mask).sum() / mask.sum().clamp_min(1)
+        reconstruction = unpatchify(predictions, self.channels, pixels.shape[-3:], self.patch_size)
+        return {
+            "loss": loss,
+            "embedding": embedding,
+            "patch_embeddings": encoded[:, 1:],
+            "reconstruction": reconstruction,
+            "mask": mask,
+        }

scripts/fake_data.py

@@ -0,0 +1,68 @@
+"""Generate deterministic HLS-like four-timestamp samples for engineering validation."""
+
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def make_split(path, count, config, seed):
+    rng = np.random.default_rng(seed)
+    data = config["data"]
+    channels, frames, size = int(data["channels"]), int(data["frames"]), int(data["image_size"])
+    means = np.asarray(data["mean"], np.float32)
+    stds = np.asarray(data["std"], np.float32)
+    y, x = np.mgrid[-1:1:complex(size), -1:1:complex(size)].astype(np.float32)
+    pixels = np.empty((count, channels, frames, size, size), np.float32)
+    temporal = np.empty((count, frames, 2), np.float32)
+    location = np.empty((count, 2), np.float32)
+    class_target = np.empty(count, np.int64)
+    regression_target = np.empty(count, np.float32)
+    for sample in range(count):
+        latitude, longitude = rng.uniform(-70, 70), rng.uniform(-180, 180)
+        start_day = int(rng.integers(1, 80))
+        days = np.clip(start_day + np.arange(frames) * int(rng.integers(45, 100)), 1, 365)
+        temporal[sample, :, 0] = 2018 + sample % 5
+        temporal[sample, :, 1] = days
+        location[sample] = (latitude, longitude)
+        phase = rng.uniform(0, 2 * np.pi)
+        class_target[sample] = int(np.sin(phase) > 0)
+        regression_target[sample] = np.cos(phase) + latitude / 180
+        for step, day in enumerate(days):
+            seasonal = np.sin(2 * np.pi * day / 365 + phase)
+            landscape = np.sin(2.5 * np.pi * x + phase) * np.cos(2 * np.pi * y - phase)
+            landscape += 0.35 * x + 0.2 * y + 0.25 * seasonal
+            for channel in range(channels):
+                normalized = landscape + 0.12 * channel + rng.normal(0, 0.04, (size, size))
+                pixels[sample, channel, step] = normalized * stds[channel] + means[channel]
+    payload = {
+        "format_version": np.asarray(data["format_version"]),
+        "data_source": np.asarray("synthetic_hls_like"),
+        "pixels": pixels,
+        "temporal_coords": temporal,
+        "location_coords": location,
+        "class_target": class_target,
+        "regression_target": regression_target,
+    }
+    np.savez_compressed(path, **payload)
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    output = ROOT / config["data"]["root"]
+    output.mkdir(parents=True, exist_ok=True)
+    for offset, (filename, count) in enumerate((
+        ("train.npz", int(config["data"]["train_samples"])),
+        ("test.npz", int(config["data"]["test_samples"])),
+    )):
+        path = output / filename
+        if not path.exists():
+            make_split(path, 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 @@
+"""Run multi-temporal reconstruction and embedding inference."""
+
+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.prithvi_eo import PrithviEO2
+from train import PrithviDataset, 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 = PrithviEO2(checkpoint["model_config"]).to(device)
+    model.load_state_dict(checkpoint["model"])
+    model.eval()
+    dataset = PrithviDataset(ROOT / config["data"]["root"] / "test.npz", config)
+    pixels = torch.stack([dataset[index]["pixels"] for index in range(len(dataset))]).to(device)
+    temporal = torch.stack([dataset[index]["temporal"] for index in range(len(dataset))]).to(device)
+    location = torch.stack([dataset[index]["location"] for index in range(len(dataset))]).to(device)
+    torch.manual_seed(int(config["seed"]))
+    with torch.no_grad():
+        output = model(pixels, temporal, location)
+        cls_embedding, patch_embeddings = model.encode(pixels, temporal, location)
+    target = ROOT / config["paths"]["inference_dir"] / "predictions.npz"
+    target.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(
+        target,
+        format_version=np.asarray(config["data"]["format_version"]),
+        pixels=pixels.cpu().numpy(),
+        reconstruction=output["reconstruction"].cpu().numpy(),
+        mask=output["mask"].cpu().numpy(),
+        embedding=cls_embedding.cpu().numpy(),
+        patch_embeddings=patch_embeddings.cpu().numpy(),
+        temporal_coords=temporal.cpu().numpy(),
+        location_coords=location.cpu().numpy(),
+        class_target=dataset.data["class_target"],
+        regression_target=dataset.data["regression_target"],
+        masked_patch_mse=np.asarray(float(output["loss"])),
+    )
+    print(f"predictions={target.relative_to(ROOT)}")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,56 @@
+"""Evaluate Prithvi reconstruction and visualize temporal HLS samples."""
+
+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())
+    prediction = np.load(ROOT / config["paths"]["inference_dir"] / "predictions.npz")
+    pixels, reconstruction = prediction["pixels"], prediction["reconstruction"]
+    error = np.abs(reconstruction - pixels)
+    per_frame = error.mean(axis=(0, 1, 3, 4))
+    embeddings = prediction["embedding"]
+    metrics = {
+        "samples": int(len(pixels)),
+        "masked_patch_mse": float(prediction["masked_patch_mse"]),
+        "reconstruction_mae": float(error.mean()),
+        "per_frame_reconstruction_mae": [float(value) for value in per_frame],
+        "mean_embedding_norm": float(np.linalg.norm(embeddings, axis=1).mean()),
+    }
+    output = ROOT / config["paths"]["evaluation_dir"]
+    output.mkdir(parents=True, exist_ok=True)
+    (output / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
+    figure, axes = plt.subplots(3, int(config["data"]["frames"]), figsize=(12, 8))
+    for frame in range(int(config["data"]["frames"])):
+        source = pixels[0, [2, 1, 0], frame].transpose(1, 2, 0)
+        rebuilt = reconstruction[0, [2, 1, 0], frame].transpose(1, 2, 0)
+        low, high = np.percentile(source, (2, 98))
+        source = np.clip((source - low) / max(high - low, 1e-6), 0, 1)
+        rebuilt = np.clip((rebuilt - low) / max(high - low, 1e-6), 0, 1)
+        axes[0, frame].imshow(source)
+        axes[1, frame].imshow(rebuilt)
+        axes[2, frame].imshow(error[0, :, frame].mean(axis=0), cmap="magma")
+        axes[0, frame].set_title(f"time {frame + 1}")
+        for axis in axes[:, frame]:
+            axis.axis("off")
+    axes[0, 0].set_ylabel("input")
+    axes[1, 0].set_ylabel("reconstruction")
+    axes[2, 0].set_ylabel("absolute error")
+    figure.tight_layout()
+    figure.savefig(output / "comparison.png", dpi=150)
+    plt.close(figure)
+    print(f"metrics={output.relative_to(ROOT) / 'metrics.json'}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,107 @@
+"""Train the reduced Prithvi-EO-2.0 temporal-location MAE."""
+
+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.prithvi_eo import PrithviEO2
+
+
+class PrithviDataset(Dataset):
+    def __init__(self, path, config):
+        self.data = np.load(path)
+        self.config = config
+        if str(self.data["format_version"]) != config["data"]["format_version"]:
+            raise ValueError("incompatible data format")
+        expected = (
+            int(config["data"]["channels"]), int(config["data"]["frames"]),
+            int(config["data"]["image_size"]), int(config["data"]["image_size"]),
+        )
+        if self.data["pixels"].shape[1:] != expected:
+            raise ValueError(f"pixels have shape {self.data['pixels'].shape[1:]}, expected {expected}")
+        self.mean = torch.tensor(config["data"]["mean"], dtype=torch.float32)[:, None, None, None]
+        self.std = torch.tensor(config["data"]["std"], dtype=torch.float32)[:, None, None, None]
+
+    def __len__(self):
+        return len(self.data["pixels"])
+
+    def __getitem__(self, index):
+        pixels = torch.from_numpy(self.data["pixels"][index]).float()
+        return {
+            "pixels": (pixels - self.mean) / self.std,
+            "temporal": torch.from_numpy(self.data["temporal_coords"][index]).float(),
+            "location": torch.from_numpy(self.data["location_coords"][index]).float(),
+        }
+
+
+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 = PrithviDataset(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 = PrithviEO2(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()
+        total, steps = 0.0, 0
+        for batch in loader:
+            output = model(batch["pixels"].to(device), batch["temporal"].to(device), batch["location"].to(device))
+            optimizer.zero_grad(set_to_none=True)
+            output["loss"].backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+            total += float(output["loss"].detach())
+            steps += 1
+        metrics = {"epoch": epoch + 1, "masked_patch_mse": total / max(steps, 1)}
+        history.append(metrics)
+        if rank == 0:
+            print(f"epoch={epoch + 1} masked_patch_mse={metrics['masked_patch_mse']:.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"],
+                    "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()