Publish DINCAE reproduction
Browse files- .gitattributes +2 -34
- README.md +117 -0
- README_zh.md +120 -0
- conf/config.yaml +46 -0
- config.json +45 -0
- configuration.json +12 -0
- model/dincae.py +122 -0
- scripts/fake_data.py +65 -0
- scripts/inference.py +51 -0
- scripts/result.py +89 -0
- scripts/train.py +108 -0
- weight/.gitkeep +0 -0
.gitattributes
CHANGED
|
@@ -1,35 +1,3 @@
|
|
| 1 |
-
*.7z filter=lfs diff=lfs merge=lfs -text
|
| 2 |
-
*.arrow filter=lfs diff=lfs merge=lfs -text
|
| 3 |
-
*.bin filter=lfs diff=lfs merge=lfs -text
|
| 4 |
-
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
| 5 |
-
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
| 6 |
-
*.ftz filter=lfs diff=lfs merge=lfs -text
|
| 7 |
-
*.gz filter=lfs diff=lfs merge=lfs -text
|
| 8 |
-
*.h5 filter=lfs diff=lfs merge=lfs -text
|
| 9 |
-
*.joblib filter=lfs diff=lfs merge=lfs -text
|
| 10 |
-
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
| 11 |
-
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
| 12 |
-
*.model filter=lfs diff=lfs merge=lfs -text
|
| 13 |
-
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
| 14 |
-
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 15 |
-
*.npz filter=lfs diff=lfs merge=lfs -text
|
| 16 |
-
*.onnx filter=lfs diff=lfs merge=lfs -text
|
| 17 |
-
*.ot filter=lfs diff=lfs merge=lfs -text
|
| 18 |
-
*.parquet filter=lfs diff=lfs merge=lfs -text
|
| 19 |
-
*.pb filter=lfs diff=lfs merge=lfs -text
|
| 20 |
-
*.pickle filter=lfs diff=lfs merge=lfs -text
|
| 21 |
-
*.pkl filter=lfs diff=lfs merge=lfs -text
|
| 22 |
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 23 |
-
*.
|
| 24 |
-
*.
|
| 25 |
-
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 26 |
-
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
| 27 |
-
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
| 28 |
-
*.tar filter=lfs diff=lfs merge=lfs -text
|
| 29 |
-
*.tflite filter=lfs diff=lfs merge=lfs -text
|
| 30 |
-
*.tgz filter=lfs diff=lfs merge=lfs -text
|
| 31 |
-
*.wasm filter=lfs diff=lfs merge=lfs -text
|
| 32 |
-
*.xz filter=lfs diff=lfs merge=lfs -text
|
| 33 |
-
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
-
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
-
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 2 |
+
*.npz filter=lfs diff=lfs merge=lfs -text
|
| 3 |
+
*.png filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
README.md
ADDED
|
@@ -0,0 +1,117 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: apache-2.0
|
| 3 |
+
language:
|
| 4 |
+
- en
|
| 5 |
+
tags:
|
| 6 |
+
- OneScience
|
| 7 |
+
- Earth-science
|
| 8 |
+
- sea-surface-temperature
|
| 9 |
+
- gap-filling
|
| 10 |
+
frameworks: PyTorch
|
| 11 |
+
---
|
| 12 |
+
|
| 13 |
+
<p align="center"><strong><span style="font-size: 30px;">DINCAE</span></strong></p>
|
| 14 |
+
|
| 15 |
+
# Model Introduction
|
| 16 |
+
|
| 17 |
+
DINCAE reconstructs daily sea-surface temperature obscured by clouds with a probabilistic convolutional autoencoder and provides pixel-wise reconstruction uncertainty.
|
| 18 |
+
|
| 19 |
+
Paper: DINCAE 1.0: a convolutional neural network with error estimates to reconstruct sea surface temperature satellite observations
|
| 20 |
+
https://doi.org/10.5194/gmd-13-1609-2020
|
| 21 |
+
|
| 22 |
+
# Model Description
|
| 23 |
+
|
| 24 |
+
The method was proposed by research teams from the University of Liege and the National Institute of Biology in Slovenia. The paper trains and validates the method on AVHRR Pathfinder daily sea-surface temperature from 1985-2009. The model reconstructs cloud-obscured SST while estimating reconstruction-error variance.
|
| 25 |
+
|
| 26 |
+
# Use Cases
|
| 27 |
+
|
| 28 |
+
| Use Case | Description |
|
| 29 |
+
| :---: | :--- |
|
| 30 |
+
| SST gap reconstruction | Reconstruct cloud-obscured SST from current, previous-day, and next-day observations. |
|
| 31 |
+
| Probabilistic error estimation | Jointly output reconstruction means and pixel-wise error variances. |
|
| 32 |
+
| Local engineering validation | Validate training, inference, and evaluation with structured data at reduced spatial scale. |
|
| 33 |
+
| ModelScope/OneCode execution | Validate structured data, training, inference, SST reconstruction metrics, and visualization in ModelScope or OneCode environments. |
|
| 34 |
+
| Multi-GPU training | Validate distributed training and the checkpoint workflow through `torchrun`. |
|
| 35 |
+
|
| 36 |
+
# Usage
|
| 37 |
+
|
| 38 |
+
## 1.OneCode
|
| 39 |
+
|
| 40 |
+
[Try one-click AI4S programming in OneCode](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
|
| 41 |
+
|
| 42 |
+
## 2. Download and Installation
|
| 43 |
+
|
| 44 |
+
```bash
|
| 45 |
+
hf download OneScience-Group/DINCAE --local-dir ./DINCAE
|
| 46 |
+
cd DINCAE
|
| 47 |
+
```
|
| 48 |
+
|
| 49 |
+
### Environment Dependencies
|
| 50 |
+
|
| 51 |
+
**Hardware Requirements**
|
| 52 |
+
|
| 53 |
+
- A GPU or DCU is recommended.
|
| 54 |
+
- A CPU can be used for connectivity validation with the default small-sample configuration.
|
| 55 |
+
- DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.
|
| 56 |
+
|
| 57 |
+
### Training Data
|
| 58 |
+
|
| 59 |
+
```bash
|
| 60 |
+
python scripts/fake_data.py
|
| 61 |
+
```
|
| 62 |
+
|
| 63 |
+
The paper data contain 5,266 daily `112×112` AVHRR SST time steps, with current, previous-day, and next-day observations, precision, coordinates, and seasonal inputs. The default uses eight structured synthetic `32×32` time steps for efficient CPU validation, while paper dimensions and parameters remain recorded under `paper_model`. The synthetic data validate engineering only and do not represent the real AVHRR distribution, training scale, or paper performance.
|
| 64 |
+
|
| 65 |
+
### Training
|
| 66 |
+
|
| 67 |
+
For single-GPU training, use:
|
| 68 |
+
|
| 69 |
+
```bash
|
| 70 |
+
python scripts/train.py
|
| 71 |
+
```
|
| 72 |
+
|
| 73 |
+
For multi-GPU training, use:
|
| 74 |
+
|
| 75 |
+
```bash
|
| 76 |
+
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
|
| 77 |
+
```
|
| 78 |
+
|
| 79 |
+
Training uses random cloud masks and masked Gaussian NLL; the default reduces grid size, samples, convolution widths, bottleneck size, and epochs while retaining the paper's `112×112` setup under `paper_model`. Formal experiments require complete AVHRR data and paper-scale training, with artifacts saved to:
|
| 80 |
+
|
| 81 |
+
```text
|
| 82 |
+
result/checkpoints/dincae.pt
|
| 83 |
+
result/training/metrics.json
|
| 84 |
+
```
|
| 85 |
+
|
| 86 |
+
### Trained Weights
|
| 87 |
+
|
| 88 |
+
The paper does not provide directly loadable official model weights, and this repository bundles no weights under `weight/`. The locally trained checkpoint is saved to `result/checkpoints/dincae.pt` and must not be represented as an official pretrained weight.
|
| 89 |
+
|
| 90 |
+
### Inference
|
| 91 |
+
|
| 92 |
+
```bash
|
| 93 |
+
python scripts/inference.py
|
| 94 |
+
```
|
| 95 |
+
|
| 96 |
+
This writes reconstructed means, variances, targets, missing masks, and dates to `result/output/predictions.npz`.
|
| 97 |
+
|
| 98 |
+
### Evaluation and Visualization
|
| 99 |
+
|
| 100 |
+
```bash
|
| 101 |
+
python scripts/result.py
|
| 102 |
+
```
|
| 103 |
+
|
| 104 |
+
Evaluation reports RMSE, centered RMSE (CRMSE), bias, standardized-residual calibration, and a DINEOF-like iterative rank-13 baseline. Outputs are `result/evaluation/metrics.json` and `result/evaluation/comparison.png`. Synthetic-data metrics demonstrate connectivity only and are not paper results.
|
| 105 |
+
|
| 106 |
+
# Official OneScience Information
|
| 107 |
+
|
| 108 |
+
| Platform | OneScience | OneSkills |
|
| 109 |
+
| --- | --- | --- |
|
| 110 |
+
| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
|
| 111 |
+
| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
|
| 112 |
+
|
| 113 |
+
# Citation and License
|
| 114 |
+
|
| 115 |
+
This repository is an independent engineering reproduction of the public DINCAE specifications.
|
| 116 |
+
|
| 117 |
+
Use of this repository's code, official model weights, and data remains subject to the licenses and terms of their respective projects.
|
README_zh.md
ADDED
|
@@ -0,0 +1,120 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: apache-2.0
|
| 3 |
+
language:
|
| 4 |
+
- en
|
| 5 |
+
- zh
|
| 6 |
+
tags:
|
| 7 |
+
- OneScience
|
| 8 |
+
- 地球科学
|
| 9 |
+
- 海表温度
|
| 10 |
+
- 缺测重建
|
| 11 |
+
frameworks: PyTorch
|
| 12 |
+
---
|
| 13 |
+
|
| 14 |
+
<p align="center"><strong><span style="font-size: 30px;">DINCAE</span></strong></p>
|
| 15 |
+
|
| 16 |
+
# 模型介绍
|
| 17 |
+
|
| 18 |
+
DINCAE 通过概率卷积自编码器重建云遮挡造成的逐日海表温度缺测,并提供像素级重建不确定性。
|
| 19 |
+
|
| 20 |
+
论文:DINCAE 1.0: a convolutional neural network with error estimates to reconstruct sea surface temperature satellite observations
|
| 21 |
+
https://doi.org/10.5194/gmd-13-1609-2020
|
| 22 |
+
|
| 23 |
+
# 模型描述
|
| 24 |
+
|
| 25 |
+
该方法由列日大学和斯洛文尼亚国家生物研究所的研究团队提出。论文使用 1985-2009 年 AVHRR Pathfinder 逐日海表温度资料开展训练与验证。模型适用于重建云遮挡区域的 SST,并同时估计重建误差方差。
|
| 26 |
+
|
| 27 |
+
# 适用场景
|
| 28 |
+
|
| 29 |
+
| 场景 | 说明 |
|
| 30 |
+
| :---: | :--- |
|
| 31 |
+
| 海表温度缺测重建 | 根据当前、前一日和后一日观测重建云遮挡 SST。 |
|
| 32 |
+
| 概率误差估计 | 联合输出重建均值和像素级误差方差。 |
|
| 33 |
+
| 本地工程验证 | 使用缩小空间尺度的结构化虚拟数据验证训练、推理和评估流程。 |
|
| 34 |
+
| ModelScope/OneCode 运行 | 在 ModelScope 或 OneCode 环境中验证结构化数据、训练、推理、SST 重建指标和可视化流程。 |
|
| 35 |
+
| 多卡训练 | 通过 `torchrun` 验证分布式训练和 checkpoint 流程。 |
|
| 36 |
+
|
| 37 |
+
# 使用说明
|
| 38 |
+
|
| 39 |
+
## 1.OneCode
|
| 40 |
+
|
| 41 |
+
可通过 OneCode 在线环境体验智能化一键式 AI4S 编程:
|
| 42 |
+
|
| 43 |
+
[点击体验智能化一键式 AI4S 编程](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
|
| 44 |
+
|
| 45 |
+
## 2.下载安装
|
| 46 |
+
|
| 47 |
+
```bash
|
| 48 |
+
modelscope download --model OneScience/DINCAE --local_dir ./DINCAE
|
| 49 |
+
cd DINCAE
|
| 50 |
+
```
|
| 51 |
+
|
| 52 |
+
### 环境依赖
|
| 53 |
+
|
| 54 |
+
**硬件要求**
|
| 55 |
+
|
| 56 |
+
- 推荐使用 GPU 或 DCU 运行。
|
| 57 |
+
- CPU 可用于默认小样本配置的连通性验证。
|
| 58 |
+
- DCU 用户需预先安装 DTK,建议使用 DTK 25.04.2 以上版本或与当前集群匹配的 OneScience 推荐版本。
|
| 59 |
+
|
| 60 |
+
### 训练数据介绍
|
| 61 |
+
|
| 62 |
+
```bash
|
| 63 |
+
python scripts/fake_data.py
|
| 64 |
+
```
|
| 65 |
+
|
| 66 |
+
论文数据包含 5266 个 `112×112` AVHRR 逐日 SST 时刻,模型输入当前、前一日和后一日观测及精度、坐标和季节信息。本仓库默认使用 8 个 `32×32` 结构化虚拟时刻以控制 CPU 验证成本,论文真实网格和模型参数单独记录在 `paper_model`。虚拟数据仅用于验证工程流程,不代表 AVHRR 的真实数据分布、训练规模或论文性能。
|
| 67 |
+
|
| 68 |
+
### 训练
|
| 69 |
+
|
| 70 |
+
单卡训练可使用:
|
| 71 |
+
|
| 72 |
+
```bash
|
| 73 |
+
python scripts/train.py
|
| 74 |
+
```
|
| 75 |
+
|
| 76 |
+
多卡训练可使用:
|
| 77 |
+
|
| 78 |
+
```bash
|
| 79 |
+
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
|
| 80 |
+
```
|
| 81 |
+
|
| 82 |
+
训练使用随机云掩码和 masked Gaussian NLL;默认缩小网格、样本数、卷积宽度、瓶颈和训练轮数,论文配置保留 `112×112` 网格与 1000 epochs。正式实验需要完整 AVHRR 数据和论文规模模型,训练产物保存到:
|
| 83 |
+
|
| 84 |
+
```text
|
| 85 |
+
result/checkpoints/dincae.pt
|
| 86 |
+
result/training/metrics.json
|
| 87 |
+
```
|
| 88 |
+
|
| 89 |
+
### 训练权重
|
| 90 |
+
|
| 91 |
+
论文未提供可直接加载的官方模型权重,本仓库不在 `weight/` 中内置权重。本地训练生成的 checkpoint 保存到 `result/checkpoints/dincae.pt`,不得描述为官方预训练权重。
|
| 92 |
+
|
| 93 |
+
### 推理
|
| 94 |
+
|
| 95 |
+
```bash
|
| 96 |
+
python scripts/inference.py
|
| 97 |
+
```
|
| 98 |
+
|
| 99 |
+
输出 `result/output/predictions.npz`,包含重建均值、方差、目标、缺测掩膜和日期。
|
| 100 |
+
|
| 101 |
+
### 评估和可视化
|
| 102 |
+
|
| 103 |
+
```bash
|
| 104 |
+
python scripts/result.py
|
| 105 |
+
```
|
| 106 |
+
|
| 107 |
+
评估输出 RMSE、中心化 RMSE(CRMSE)、bias、标准化残差 calibration,以及 DINEOF-like rank-13 迭代低秩基线。结果写入 `result/evaluation/metrics.json` 和 `result/evaluation/comparison.png`。少样本指标只证明流程连通,不代表论文性能。
|
| 108 |
+
|
| 109 |
+
# OneScience官方信息
|
| 110 |
+
|
| 111 |
+
| 平台 | OneScience 主仓库 | Skills 仓库 |
|
| 112 |
+
| --- | --- | --- |
|
| 113 |
+
| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
|
| 114 |
+
| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
|
| 115 |
+
|
| 116 |
+
# 引用与许可证
|
| 117 |
+
|
| 118 |
+
本仓库为 DINCAE 公开规格的独立工程复现版本。
|
| 119 |
+
|
| 120 |
+
本仓库代码、官方模型权重和数据的使用仍应以各自项目中的许可证及使用条款为准。
|
conf/config.yaml
ADDED
|
@@ -0,0 +1,46 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
seed: 2026
|
| 2 |
+
data:
|
| 3 |
+
root: data
|
| 4 |
+
file: avhrr_daily_sst.npz
|
| 5 |
+
format_version: dincae_avhrr_v1
|
| 6 |
+
samples: 8
|
| 7 |
+
train_samples: 6
|
| 8 |
+
start_date: "2008-01-01"
|
| 9 |
+
grid: [32, 32]
|
| 10 |
+
paper_period: "1985-2009"
|
| 11 |
+
paper_time_steps: 5266
|
| 12 |
+
random_cloud_fraction: 0.20
|
| 13 |
+
model:
|
| 14 |
+
in_channels: 10
|
| 15 |
+
out_channels: 2
|
| 16 |
+
filters: [8, 12, 18, 27]
|
| 17 |
+
bottleneck_features: 22
|
| 18 |
+
dropout: 0.3
|
| 19 |
+
kernel_size: 3
|
| 20 |
+
grid_size: 32
|
| 21 |
+
training:
|
| 22 |
+
epochs: 1
|
| 23 |
+
batch_size: 2
|
| 24 |
+
learning_rate: 0.001
|
| 25 |
+
gamma: 10.0
|
| 26 |
+
delta: 0.001
|
| 27 |
+
paper_epochs: 1000
|
| 28 |
+
paper_batch_size: 50
|
| 29 |
+
runtime:
|
| 30 |
+
device: cpu
|
| 31 |
+
num_workers: 0
|
| 32 |
+
paths:
|
| 33 |
+
checkpoint: result/checkpoints/dincae.pt
|
| 34 |
+
training_metrics: result/training/metrics.json
|
| 35 |
+
inference: result/output/predictions.npz
|
| 36 |
+
evaluation_dir: result/evaluation
|
| 37 |
+
evaluation:
|
| 38 |
+
dineof_rank: 13
|
| 39 |
+
calibration_bins: 8
|
| 40 |
+
paper_model:
|
| 41 |
+
grid: [112, 112]
|
| 42 |
+
time_steps: 5266
|
| 43 |
+
filters: [16, 24, 36, 54]
|
| 44 |
+
bottleneck_features: 529
|
| 45 |
+
epochs: 1000
|
| 46 |
+
batch_size: 50
|
config.json
ADDED
|
@@ -0,0 +1,45 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_name": "DINCAE",
|
| 3 |
+
"model_type": "dincae",
|
| 4 |
+
"architectures": ["DINCAE"],
|
| 5 |
+
"framework": "PyTorch",
|
| 6 |
+
"domain": "ocean",
|
| 7 |
+
"task": "probabilistic-sst-gap-reconstruction",
|
| 8 |
+
"implementation": {
|
| 9 |
+
"entry_point": "model/dincae.py",
|
| 10 |
+
"scope": "core-method reduced-spatial-scale engineering reproduction with paper dimensions recorded separately",
|
| 11 |
+
"train_script": "scripts/train.py",
|
| 12 |
+
"inference_script": "scripts/inference.py",
|
| 13 |
+
"evaluation_script": "scripts/result.py",
|
| 14 |
+
"synthetic_data_script": "scripts/fake_data.py"
|
| 15 |
+
},
|
| 16 |
+
"architecture": {
|
| 17 |
+
"engineering_input_shape": ["B", 10, 32, 32],
|
| 18 |
+
"engineering_output_shape": ["B", 2, 32, 32],
|
| 19 |
+
"paper_input_shape": ["B", 10, 112, 112],
|
| 20 |
+
"paper_output_shape": ["B", 2, 112, 112],
|
| 21 |
+
"engineering_filters": [8, 12, 18, 27],
|
| 22 |
+
"paper_filters": [16, 24, 36, 54],
|
| 23 |
+
"paper_bottleneck": [2646, 529, 2646],
|
| 24 |
+
"decoder": "nearest-neighbor upsampling with encoder skips"
|
| 25 |
+
},
|
| 26 |
+
"data": {
|
| 27 |
+
"dataset": "AVHRR Pathfinder daily SST",
|
| 28 |
+
"format_version": "dincae_avhrr_v1",
|
| 29 |
+
"paper_period": "1985-2009",
|
| 30 |
+
"paper_time_steps": 5266,
|
| 31 |
+
"paper_grid": [112, 112],
|
| 32 |
+
"engineering_grid": [32, 32],
|
| 33 |
+
"input_channels": 10,
|
| 34 |
+
"output_channels": 2,
|
| 35 |
+
"synthetic": true
|
| 36 |
+
},
|
| 37 |
+
"configuration_sources": [
|
| 38 |
+
"conf/config.yaml",
|
| 39 |
+
"model/dincae.py",
|
| 40 |
+
"scripts/fake_data.py",
|
| 41 |
+
"scripts/train.py",
|
| 42 |
+
"scripts/inference.py",
|
| 43 |
+
"scripts/result.py"
|
| 44 |
+
]
|
| 45 |
+
}
|
configuration.json
ADDED
|
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"framework": "PyTorch",
|
| 3 |
+
"task": "sea surface temperature gap reconstruction with uncertainty",
|
| 4 |
+
"model": "DINCAE",
|
| 5 |
+
"input_format": "BCHW: engineering [B,10,32,32], paper [B,10,112,112]",
|
| 6 |
+
"protocol": "probabilistic daily SST reconstruction from cloudy observations",
|
| 7 |
+
"default_config": "conf/config.yaml",
|
| 8 |
+
"training": "scripts/train.py",
|
| 9 |
+
"inference": "scripts/inference.py",
|
| 10 |
+
"evaluation": "scripts/result.py",
|
| 11 |
+
"visualization": "scripts/result.py"
|
| 12 |
+
}
|
model/dincae.py
ADDED
|
@@ -0,0 +1,122 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Independent PyTorch implementation of DINCAE for daily SST gap filling."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
from typing import Dict, Tuple
|
| 6 |
+
|
| 7 |
+
import numpy as np
|
| 8 |
+
import torch
|
| 9 |
+
from torch import Tensor, nn
|
| 10 |
+
from torch.nn import functional as F
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
class ConvBlock(nn.Sequential):
|
| 14 |
+
def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3):
|
| 15 |
+
padding = kernel_size // 2
|
| 16 |
+
super().__init__(
|
| 17 |
+
nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding),
|
| 18 |
+
nn.LeakyReLU(0.1, inplace=True),
|
| 19 |
+
nn.Conv2d(out_channels, out_channels, kernel_size, padding=padding),
|
| 20 |
+
nn.LeakyReLU(0.1, inplace=True),
|
| 21 |
+
)
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
class DINCAE(nn.Module):
|
| 25 |
+
"""Four-level convolutional autoencoder with the paper-sized dense bottleneck."""
|
| 26 |
+
|
| 27 |
+
def __init__(self, in_channels: int = 10, out_channels: int = 2,
|
| 28 |
+
filters=(16, 24, 36, 54),
|
| 29 |
+
bottleneck_features: int = 529, dropout: float = 0.3,
|
| 30 |
+
kernel_size: int = 3, grid_size: int = 112):
|
| 31 |
+
super().__init__()
|
| 32 |
+
if len(filters) != 4 or out_channels != 2:
|
| 33 |
+
raise ValueError("DINCAE requires four encoder levels and two output channels")
|
| 34 |
+
if grid_size < 16 or grid_size % 16:
|
| 35 |
+
raise ValueError("grid_size must be at least 16 and divisible by 16")
|
| 36 |
+
self.input_shape = (in_channels, grid_size, grid_size)
|
| 37 |
+
self.encoders = nn.ModuleList()
|
| 38 |
+
previous = in_channels
|
| 39 |
+
for width in filters:
|
| 40 |
+
self.encoders.append(ConvBlock(previous, width, kernel_size))
|
| 41 |
+
previous = width
|
| 42 |
+
self.pool = nn.AvgPool2d(2)
|
| 43 |
+
latent_size = grid_size // 16
|
| 44 |
+
flattened = filters[-1] * latent_size * latent_size
|
| 45 |
+
self.latent_size = latent_size
|
| 46 |
+
self.dense_encode = nn.Linear(flattened, bottleneck_features)
|
| 47 |
+
self.dropout = nn.Dropout(dropout)
|
| 48 |
+
self.dense_decode = nn.Linear(bottleneck_features, flattened)
|
| 49 |
+
decoder_inputs = (filters[3] + filters[3], filters[3] + filters[2],
|
| 50 |
+
filters[2] + filters[1], filters[1] + filters[0])
|
| 51 |
+
decoder_outputs = (filters[3], filters[2], filters[1], filters[0])
|
| 52 |
+
self.decoders = nn.ModuleList([
|
| 53 |
+
ConvBlock(cin, cout, kernel_size)
|
| 54 |
+
for cin, cout in zip(decoder_inputs, decoder_outputs)
|
| 55 |
+
])
|
| 56 |
+
self.output = nn.Conv2d(filters[0], out_channels, 1)
|
| 57 |
+
|
| 58 |
+
def forward(self, x: Tensor) -> Tensor:
|
| 59 |
+
if tuple(x.shape[-2:]) != self.input_shape[1:] or x.shape[1] != self.input_shape[0]:
|
| 60 |
+
raise ValueError(f"expected [B,{self.input_shape[0]},{self.input_shape[1]},{self.input_shape[2]}], got {tuple(x.shape)}")
|
| 61 |
+
skips = []
|
| 62 |
+
for encoder in self.encoders:
|
| 63 |
+
x = encoder(x)
|
| 64 |
+
skips.append(x)
|
| 65 |
+
x = self.pool(x)
|
| 66 |
+
x = x.flatten(1)
|
| 67 |
+
x = self.dense_decode(self.dropout(F.leaky_relu(self.dense_encode(x), 0.1)))
|
| 68 |
+
x = x.view(x.shape[0], -1, self.latent_size, self.latent_size)
|
| 69 |
+
for decoder, skip in zip(self.decoders, reversed(skips)):
|
| 70 |
+
x = F.interpolate(x, size=skip.shape[-2:], mode="nearest")
|
| 71 |
+
x = decoder(torch.cat((x, skip), dim=1))
|
| 72 |
+
return self.output(x)
|
| 73 |
+
|
| 74 |
+
|
| 75 |
+
def output_distribution(output: Tensor, gamma: float = 10.0,
|
| 76 |
+
delta: float = 1e-3) -> Tuple[Tensor, Tensor, Tensor]:
|
| 77 |
+
"""Recover precision, variance, and mean from [log precision, temperature*precision]."""
|
| 78 |
+
log_precision = output[:, 0]
|
| 79 |
+
precision = torch.maximum(torch.exp(torch.minimum(log_precision, log_precision.new_tensor(gamma))),
|
| 80 |
+
log_precision.new_tensor(delta))
|
| 81 |
+
variance = precision.reciprocal()
|
| 82 |
+
mean = output[:, 1] / precision
|
| 83 |
+
return mean, variance, precision
|
| 84 |
+
|
| 85 |
+
|
| 86 |
+
def masked_gaussian_nll(output: Tensor, target: Tensor, mask: Tensor,
|
| 87 |
+
gamma: float = 10.0, delta: float = 1e-3) -> Tensor:
|
| 88 |
+
mean, _, precision = output_distribution(output, gamma, delta)
|
| 89 |
+
target = target[:, 0]
|
| 90 |
+
valid = mask[:, 0].bool() & torch.isfinite(target)
|
| 91 |
+
if not torch.any(valid):
|
| 92 |
+
return output.sum() * 0.0
|
| 93 |
+
error = target[valid] - mean[valid]
|
| 94 |
+
return (0.5 * (precision[valid] * error.square() - torch.log(precision[valid]))).mean()
|
| 95 |
+
|
| 96 |
+
|
| 97 |
+
def reconstruction_metrics(target: np.ndarray, prediction: np.ndarray,
|
| 98 |
+
mask: np.ndarray) -> Dict[str, float]:
|
| 99 |
+
valid = mask.astype(bool) & np.isfinite(target) & np.isfinite(prediction)
|
| 100 |
+
if not np.any(valid):
|
| 101 |
+
return {"rmse": float("nan"), "crmse": float("nan"), "bias": float("nan")}
|
| 102 |
+
error = prediction[valid].astype(np.float64) - target[valid].astype(np.float64)
|
| 103 |
+
bias = float(error.mean())
|
| 104 |
+
return {
|
| 105 |
+
"rmse": float(np.sqrt(np.mean(error ** 2))),
|
| 106 |
+
"crmse": float(np.sqrt(np.mean((error - bias) ** 2))),
|
| 107 |
+
"bias": bias,
|
| 108 |
+
}
|
| 109 |
+
|
| 110 |
+
|
| 111 |
+
def build_input(observed: np.ndarray, precision: np.ndarray, index: int,
|
| 112 |
+
longitude: np.ndarray, latitude: np.ndarray, day_of_year: int) -> np.ndarray:
|
| 113 |
+
"""Build 10 channels: three anomaly*precision, three precision, and four auxiliaries."""
|
| 114 |
+
frames = []
|
| 115 |
+
for offset in (0, -1, 1):
|
| 116 |
+
j = min(max(index + offset, 0), observed.shape[0] - 1)
|
| 117 |
+
anomaly = np.nan_to_num(observed[j], nan=0.0)
|
| 118 |
+
frames.extend((anomaly * precision[j], precision[j]))
|
| 119 |
+
angle = 2.0 * np.pi * day_of_year / 365.25
|
| 120 |
+
frames.extend((longitude, latitude, np.full_like(longitude, np.cos(angle)),
|
| 121 |
+
np.full_like(longitude, np.sin(angle))))
|
| 122 |
+
return np.stack(frames).astype(np.float32)
|
scripts/fake_data.py
ADDED
|
@@ -0,0 +1,65 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Generate compact AVHRR-like daily SST data for engineering validation."""
|
| 2 |
+
|
| 3 |
+
from datetime import date, timedelta
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
|
| 6 |
+
import numpy as np
|
| 7 |
+
import yaml
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
def main():
|
| 14 |
+
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
|
| 15 |
+
data_config = config["data"]
|
| 16 |
+
n = int(data_config["samples"])
|
| 17 |
+
height, width = map(int, data_config["grid"])
|
| 18 |
+
if height != width or height < 16 or height % 16 or n < 4:
|
| 19 |
+
raise ValueError("DINCAE engineering grids must be square, divisible by 16, and contain at least four days")
|
| 20 |
+
rng = np.random.default_rng(int(config["seed"]))
|
| 21 |
+
yy, xx = np.mgrid[-1:1:complex(height), -1:1:complex(width)]
|
| 22 |
+
longitude = (12.0 + 6.0 * xx).astype(np.float32)
|
| 23 |
+
latitude = (38.0 + 6.0 * yy).astype(np.float32)
|
| 24 |
+
ocean = ((xx + 0.28) ** 2 + (yy - 0.10) ** 2 > 0.055) & (xx - 0.75 * yy < 1.12)
|
| 25 |
+
sst = np.empty((n, height, width), np.float32)
|
| 26 |
+
observed = np.empty_like(sst)
|
| 27 |
+
precision = np.empty_like(sst)
|
| 28 |
+
start = date.fromisoformat(data_config["start_date"])
|
| 29 |
+
timestamps = []
|
| 30 |
+
for t in range(n):
|
| 31 |
+
current = start + timedelta(days=t)
|
| 32 |
+
timestamps.append(current.isoformat())
|
| 33 |
+
phase = 2 * np.pi * t / max(n, 16)
|
| 34 |
+
eddy = 1.9 * np.exp(-((xx - 0.35 * np.sin(phase)) ** 2 +
|
| 35 |
+
(yy - 0.25 * np.cos(phase)) ** 2) / 0.055)
|
| 36 |
+
field = 18.0 - 3.1 * yy + 0.7 * np.sin(3 * xx + phase) + eddy
|
| 37 |
+
field += rng.normal(0, 0.06, (height, width))
|
| 38 |
+
field[~ocean] = np.nan
|
| 39 |
+
cloud_noise = rng.normal(size=(height, width))
|
| 40 |
+
cloud = (np.sin(5 * xx + phase) + np.cos(4 * yy - phase) + cloud_noise > 1.05)
|
| 41 |
+
valid = ocean & ~cloud
|
| 42 |
+
sigma = (0.22 + 0.08 * (1 + np.sin(2 * xx - yy + phase))).astype(np.float32)
|
| 43 |
+
obs = field + rng.normal(0, sigma)
|
| 44 |
+
obs[~valid] = np.nan
|
| 45 |
+
sst[t], observed[t] = field, obs
|
| 46 |
+
precision[t] = np.where(valid, 1.0 / sigma ** 2, 0.0)
|
| 47 |
+
climatology = np.zeros((height, width), np.float32)
|
| 48 |
+
climatology[ocean] = np.mean(sst[:, ocean], axis=0)
|
| 49 |
+
anomaly = sst - climatology
|
| 50 |
+
observed_anomaly = observed - climatology
|
| 51 |
+
output = ROOT / data_config["root"]
|
| 52 |
+
output.mkdir(parents=True, exist_ok=True)
|
| 53 |
+
np.savez_compressed(output / data_config["file"],
|
| 54 |
+
format_version=np.array(data_config["format_version"]),
|
| 55 |
+
sst_anomaly=anomaly, observed_anomaly=observed_anomaly,
|
| 56 |
+
precision=precision, climatology=climatology,
|
| 57 |
+
ocean_mask=ocean, longitude=longitude, latitude=latitude,
|
| 58 |
+
timestamps=np.asarray(timestamps), units=np.array("degree_Celsius"),
|
| 59 |
+
grid=np.array([height, width], np.int32),
|
| 60 |
+
source=np.array("synthetic AVHRR-like engineering data"))
|
| 61 |
+
print(f"generated={n} grid={height}x{width} observed_fraction={np.isfinite(observed).mean():.4f}")
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
if __name__ == "__main__":
|
| 65 |
+
main()
|
scripts/inference.py
ADDED
|
@@ -0,0 +1,51 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Run probabilistic DINCAE reconstruction on held-out daily SST fields."""
|
| 2 |
+
|
| 3 |
+
from datetime import date
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
import sys
|
| 6 |
+
|
| 7 |
+
import numpy as np
|
| 8 |
+
import torch
|
| 9 |
+
import yaml
|
| 10 |
+
|
| 11 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 12 |
+
sys.path.insert(0, str(ROOT))
|
| 13 |
+
|
| 14 |
+
from model.dincae import DINCAE, build_input, output_distribution
|
| 15 |
+
|
| 16 |
+
|
| 17 |
+
def main():
|
| 18 |
+
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
|
| 19 |
+
device = torch.device("cuda" if torch.cuda.is_available() and config["runtime"]["device"] != "cpu" else "cpu")
|
| 20 |
+
data_path = ROOT / config["data"]["root"] / config["data"]["file"]
|
| 21 |
+
with np.load(data_path) as loaded:
|
| 22 |
+
data = {key: loaded[key] for key in loaded.files}
|
| 23 |
+
checkpoint = torch.load(ROOT / config["paths"]["checkpoint"], map_location=device, weights_only=False)
|
| 24 |
+
model = DINCAE(**checkpoint["model_config"]).to(device)
|
| 25 |
+
model.load_state_dict(checkpoint["model"])
|
| 26 |
+
model.eval()
|
| 27 |
+
start = min(int(config["data"]["train_samples"]), len(data["timestamps"]) - 1)
|
| 28 |
+
means, variances, targets, missing_masks, inputs_observed = [], [], [], [], []
|
| 29 |
+
with torch.no_grad():
|
| 30 |
+
for index in range(start, len(data["timestamps"])):
|
| 31 |
+
timestamp = date.fromisoformat(str(data["timestamps"][index]))
|
| 32 |
+
model_input = build_input(data["observed_anomaly"], data["precision"], index,
|
| 33 |
+
data["longitude"], data["latitude"], timestamp.timetuple().tm_yday)
|
| 34 |
+
output = model(torch.from_numpy(model_input[None]).to(device))
|
| 35 |
+
mean, variance, _ = output_distribution(output, checkpoint["gamma"], checkpoint["delta"])
|
| 36 |
+
means.append(mean[0].cpu().numpy() + data["climatology"])
|
| 37 |
+
variances.append(variance[0].cpu().numpy())
|
| 38 |
+
targets.append(data["sst_anomaly"][index] + data["climatology"])
|
| 39 |
+
missing_masks.append((data["precision"][index] == 0) & data["ocean_mask"])
|
| 40 |
+
inputs_observed.append(data["observed_anomaly"][index] + data["climatology"])
|
| 41 |
+
output_path = ROOT / config["paths"]["inference"]
|
| 42 |
+
output_path.parent.mkdir(parents=True, exist_ok=True)
|
| 43 |
+
np.savez_compressed(output_path, prediction=np.asarray(means), variance=np.asarray(variances),
|
| 44 |
+
target=np.asarray(targets), missing_mask=np.asarray(missing_masks),
|
| 45 |
+
observed=np.asarray(inputs_observed), ocean_mask=data["ocean_mask"],
|
| 46 |
+
timestamps=data["timestamps"][start:], units=data["units"])
|
| 47 |
+
print(f"predictions={len(means)} shape={np.asarray(means).shape} output={output_path}")
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
if __name__ == "__main__":
|
| 51 |
+
main()
|
scripts/result.py
ADDED
|
@@ -0,0 +1,89 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Evaluate DINCAE, residual calibration, and a rank-13 DINEOF-like baseline."""
|
| 2 |
+
|
| 3 |
+
import json
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
import sys
|
| 6 |
+
|
| 7 |
+
import matplotlib
|
| 8 |
+
matplotlib.use("Agg")
|
| 9 |
+
import matplotlib.pyplot as plt
|
| 10 |
+
import numpy as np
|
| 11 |
+
import yaml
|
| 12 |
+
|
| 13 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 14 |
+
sys.path.insert(0, str(ROOT))
|
| 15 |
+
|
| 16 |
+
from model.dincae import reconstruction_metrics
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
def dineof_like(observed, ocean_mask, requested_rank):
|
| 20 |
+
matrix = observed.reshape(observed.shape[0], -1).astype(np.float64)
|
| 21 |
+
sea = ocean_mask.ravel()
|
| 22 |
+
values = matrix[:, sea]
|
| 23 |
+
missing = ~np.isfinite(values)
|
| 24 |
+
column_mean = np.nanmean(values, axis=0)
|
| 25 |
+
column_mean = np.nan_to_num(column_mean, nan=float(np.nanmean(values)))
|
| 26 |
+
filled = np.where(missing, column_mean[None], values)
|
| 27 |
+
rank = min(int(requested_rank), min(filled.shape) - 1)
|
| 28 |
+
for _ in range(4):
|
| 29 |
+
mean = filled.mean(axis=0, keepdims=True)
|
| 30 |
+
u, singular, vt = np.linalg.svd(filled - mean, full_matrices=False)
|
| 31 |
+
reconstructed = (u[:, :rank] * singular[:rank]) @ vt[:rank] + mean
|
| 32 |
+
filled[missing] = reconstructed[missing]
|
| 33 |
+
result = np.full_like(matrix, np.nan)
|
| 34 |
+
result[:, sea] = filled
|
| 35 |
+
return result.reshape(observed.shape).astype(np.float32), rank
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
def main():
|
| 39 |
+
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
|
| 40 |
+
with np.load(ROOT / config["paths"]["inference"]) as loaded:
|
| 41 |
+
pred = {key: loaded[key] for key in loaded.files}
|
| 42 |
+
with np.load(ROOT / config["data"]["root"] / config["data"]["file"]) as loaded:
|
| 43 |
+
full = {key: loaded[key] for key in loaded.files}
|
| 44 |
+
mask = pred["missing_mask"]
|
| 45 |
+
metrics = {"dincae": reconstruction_metrics(pred["target"], pred["prediction"], mask)}
|
| 46 |
+
std_residual = (pred["prediction"] - pred["target"]) / np.sqrt(np.maximum(pred["variance"], 1e-12))
|
| 47 |
+
residual = std_residual[mask & np.isfinite(std_residual)]
|
| 48 |
+
bins = int(config["evaluation"]["calibration_bins"])
|
| 49 |
+
edges = np.linspace(-4, 4, bins + 1)
|
| 50 |
+
counts, _ = np.histogram(residual, bins=edges)
|
| 51 |
+
metrics["standardized_residual_calibration"] = {
|
| 52 |
+
"mean": float(residual.mean()), "std": float(residual.std()),
|
| 53 |
+
"within_1sigma": float(np.mean(np.abs(residual) <= 1)),
|
| 54 |
+
"within_2sigma": float(np.mean(np.abs(residual) <= 2)),
|
| 55 |
+
"histogram_edges": edges.tolist(), "histogram_counts": counts.tolist(),
|
| 56 |
+
}
|
| 57 |
+
baseline_all, effective_rank = dineof_like(full["observed_anomaly"] + full["climatology"],
|
| 58 |
+
full["ocean_mask"], config["evaluation"]["dineof_rank"])
|
| 59 |
+
start = min(int(config["data"]["train_samples"]), len(full["timestamps"]) - 1)
|
| 60 |
+
baseline = baseline_all[start:]
|
| 61 |
+
metrics["dineof_like"] = reconstruction_metrics(pred["target"], baseline, mask)
|
| 62 |
+
metrics["dineof_like"]["requested_rank"] = int(config["evaluation"]["dineof_rank"])
|
| 63 |
+
metrics["dineof_like"]["effective_rank"] = effective_rank
|
| 64 |
+
out_dir = ROOT / config["paths"]["evaluation_dir"]
|
| 65 |
+
out_dir.mkdir(parents=True, exist_ok=True)
|
| 66 |
+
(out_dir / "metrics.json").write_text(json.dumps(metrics, indent=2))
|
| 67 |
+
index = 0
|
| 68 |
+
fields = (pred["observed"][index], pred["target"][index], pred["prediction"][index],
|
| 69 |
+
baseline[index], pred["prediction"][index] - pred["target"][index])
|
| 70 |
+
titles = ("Cloudy AVHRR-like input", "Complete target", "DINCAE", "DINEOF-like rank 13", "DINCAE error")
|
| 71 |
+
fig, axes = plt.subplots(1, 5, figsize=(17, 3.6), constrained_layout=True)
|
| 72 |
+
lower, upper = np.nanpercentile(pred["target"][index], (2, 98))
|
| 73 |
+
for axis, field, title in zip(axes, fields, titles):
|
| 74 |
+
if title.endswith("error"):
|
| 75 |
+
limit = max(float(np.nanpercentile(np.abs(field), 98)), 1e-3)
|
| 76 |
+
image = axis.imshow(field, cmap="RdBu_r", vmin=-limit, vmax=limit)
|
| 77 |
+
else:
|
| 78 |
+
image = axis.imshow(field, cmap="turbo", vmin=lower, vmax=upper)
|
| 79 |
+
axis.set_title(title, fontsize=9)
|
| 80 |
+
axis.set_axis_off()
|
| 81 |
+
fig.colorbar(image, ax=axis, fraction=0.046)
|
| 82 |
+
fig.savefig(out_dir / "comparison.png", dpi=150)
|
| 83 |
+
plt.close(fig)
|
| 84 |
+
print(json.dumps(metrics, indent=2))
|
| 85 |
+
print(f"comparison={out_dir / 'comparison.png'}")
|
| 86 |
+
|
| 87 |
+
|
| 88 |
+
if __name__ == "__main__":
|
| 89 |
+
main()
|
scripts/train.py
ADDED
|
@@ -0,0 +1,108 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Train DINCAE with random cloud masking and masked Gaussian NLL."""
|
| 2 |
+
|
| 3 |
+
import json
|
| 4 |
+
import os
|
| 5 |
+
import sys
|
| 6 |
+
from datetime import date
|
| 7 |
+
from pathlib import Path
|
| 8 |
+
|
| 9 |
+
import numpy as np
|
| 10 |
+
import torch
|
| 11 |
+
import yaml
|
| 12 |
+
from torch.nn.parallel import DistributedDataParallel
|
| 13 |
+
from torch.utils.data import DataLoader, Dataset, DistributedSampler
|
| 14 |
+
|
| 15 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 16 |
+
sys.path.insert(0, str(ROOT))
|
| 17 |
+
|
| 18 |
+
from model.dincae import DINCAE, build_input, masked_gaussian_nll
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
class SSTDataset(Dataset):
|
| 22 |
+
def __init__(self, arrays, indices, cloud_fraction, seed):
|
| 23 |
+
self.data, self.indices = arrays, list(indices)
|
| 24 |
+
self.cloud_fraction, self.seed = float(cloud_fraction), int(seed)
|
| 25 |
+
|
| 26 |
+
def __len__(self):
|
| 27 |
+
return len(self.indices)
|
| 28 |
+
|
| 29 |
+
def __getitem__(self, item):
|
| 30 |
+
index = self.indices[item]
|
| 31 |
+
timestamp = date.fromisoformat(str(self.data["timestamps"][index]))
|
| 32 |
+
observed = self.data["observed_anomaly"].copy()
|
| 33 |
+
precision = self.data["precision"].copy()
|
| 34 |
+
rng = np.random.default_rng(self.seed + index)
|
| 35 |
+
original_valid = np.isfinite(observed[index]) & self.data["ocean_mask"]
|
| 36 |
+
withheld = original_valid & (rng.random(original_valid.shape) < self.cloud_fraction)
|
| 37 |
+
observed[index, withheld] = np.nan
|
| 38 |
+
precision[index, withheld] = 0.0
|
| 39 |
+
inputs = build_input(observed, precision, index, self.data["longitude"],
|
| 40 |
+
self.data["latitude"], timestamp.timetuple().tm_yday)
|
| 41 |
+
target = self.data["sst_anomaly"][index][None].astype(np.float32)
|
| 42 |
+
return torch.from_numpy(inputs), torch.from_numpy(target), torch.from_numpy(withheld[None])
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def main():
|
| 46 |
+
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
|
| 47 |
+
rank = int(os.environ.get("RANK", "0"))
|
| 48 |
+
world_size = int(os.environ.get("WORLD_SIZE", "1"))
|
| 49 |
+
local_rank = int(os.environ.get("LOCAL_RANK", "0"))
|
| 50 |
+
distributed = world_size > 1
|
| 51 |
+
if distributed:
|
| 52 |
+
backend = "nccl" if torch.cuda.is_available() else "gloo"
|
| 53 |
+
torch.distributed.init_process_group(backend=backend)
|
| 54 |
+
requested = config["runtime"]["device"]
|
| 55 |
+
use_cuda = torch.cuda.is_available() and requested != "cpu"
|
| 56 |
+
device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu")
|
| 57 |
+
if use_cuda:
|
| 58 |
+
torch.cuda.set_device(device)
|
| 59 |
+
torch.manual_seed(int(config["seed"]) + rank)
|
| 60 |
+
path = ROOT / config["data"]["root"] / config["data"]["file"]
|
| 61 |
+
with np.load(path) as loaded:
|
| 62 |
+
arrays = {key: loaded[key] for key in loaded.files}
|
| 63 |
+
train_count = min(int(config["data"]["train_samples"]), len(arrays["timestamps"]) - 1)
|
| 64 |
+
dataset = SSTDataset(arrays, range(train_count), config["data"]["random_cloud_fraction"], config["seed"])
|
| 65 |
+
sampler = DistributedSampler(dataset, shuffle=True, seed=int(config["seed"])) if distributed else None
|
| 66 |
+
loader = DataLoader(dataset, batch_size=int(config["training"]["batch_size"]),
|
| 67 |
+
shuffle=sampler is None, sampler=sampler,
|
| 68 |
+
num_workers=int(config["runtime"]["num_workers"]))
|
| 69 |
+
model = DINCAE(**config["model"]).to(device)
|
| 70 |
+
if distributed:
|
| 71 |
+
model = DistributedDataParallel(model, device_ids=[local_rank] if use_cuda else None)
|
| 72 |
+
optimizer = torch.optim.Adam(model.parameters(), lr=float(config["training"]["learning_rate"]))
|
| 73 |
+
history = []
|
| 74 |
+
for epoch in range(int(config["training"]["epochs"])):
|
| 75 |
+
if sampler is not None:
|
| 76 |
+
sampler.set_epoch(epoch)
|
| 77 |
+
model.train()
|
| 78 |
+
total, batches = 0.0, 0
|
| 79 |
+
for inputs, target, mask in loader:
|
| 80 |
+
inputs, target, mask = inputs.to(device), target.to(device), mask.to(device)
|
| 81 |
+
optimizer.zero_grad(set_to_none=True)
|
| 82 |
+
output = model(inputs)
|
| 83 |
+
loss = masked_gaussian_nll(output, target, mask,
|
| 84 |
+
config["training"]["gamma"], config["training"]["delta"])
|
| 85 |
+
loss.backward()
|
| 86 |
+
optimizer.step()
|
| 87 |
+
total, batches = total + float(loss.detach()), batches + 1
|
| 88 |
+
history.append({"epoch": epoch + 1, "masked_gaussian_nll": total / max(batches, 1)})
|
| 89 |
+
if rank == 0:
|
| 90 |
+
print(f"epoch={epoch + 1} loss={history[-1]['masked_gaussian_nll']:.6f}")
|
| 91 |
+
if rank == 0:
|
| 92 |
+
raw_model = model.module if distributed else model
|
| 93 |
+
checkpoint = ROOT / config["paths"]["checkpoint"]
|
| 94 |
+
checkpoint.parent.mkdir(parents=True, exist_ok=True)
|
| 95 |
+
torch.save({"model": raw_model.state_dict(), "model_config": config["model"],
|
| 96 |
+
"gamma": config["training"]["gamma"], "delta": config["training"]["delta"],
|
| 97 |
+
"format_version": "dincae_checkpoint_v1"}, checkpoint)
|
| 98 |
+
metrics = ROOT / config["paths"]["training_metrics"]
|
| 99 |
+
metrics.parent.mkdir(parents=True, exist_ok=True)
|
| 100 |
+
metrics.write_text(json.dumps({"history": history, "world_size": world_size,
|
| 101 |
+
"parameter_count": sum(p.numel() for p in raw_model.parameters())}, indent=2))
|
| 102 |
+
print(f"checkpoint={checkpoint}")
|
| 103 |
+
if distributed:
|
| 104 |
+
torch.distributed.destroy_process_group()
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
if __name__ == "__main__":
|
| 108 |
+
main()
|
weight/.gitkeep
ADDED
|
File without changes
|