zhangrenchao commited on
Commit
43dce84
·
verified ·
1 Parent(s): 36e4eb5

Publish DINCAE reproduction

Browse files
.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
- *.pth filter=lfs diff=lfs merge=lfs -text
24
- *.rar filter=lfs diff=lfs merge=lfs -text
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