Imputation with trained STIMP in the Pearl River Estuary

import torch
import numpy as np
import sys
import os
from dataset.dataset_imputation import PRE8dDataset

import argparse
parser = argparse.ArgumentParser(description='Imputation')

# args for area and methods
parser.add_argument('--area', type=str, default='PRE', help='which bay area we focus')

# basic args
parser.add_argument('--epochs', type=int, default=500, help='epochs')
parser.add_argument('--batch_size', type=int, default=16, help='batch size')
parser.add_argument('--lr', type=float, default=1e-3, help='learning rate')
parser.add_argument('--wd', type=float, default=1e-4, help='weight decay')
parser.add_argument('--test_freq', type=int, default=500, help='test per n epochs')
parser.add_argument('--embedding_size', type=int, default=32)
parser.add_argument('--hidden_channels', type=int, default=32)
parser.add_argument('--diffusion_embedding_size', type=int, default=64)
parser.add_argument('--side_channels', type=int, default=1)

# args for tasks
parser.add_argument('--in_len', type=int, default=46)
parser.add_argument('--out_len', type=int, default=46)
parser.add_argument('--missing_ratio', type=float, default=0.1)

# args for diffusion
parser.add_argument('--beta_start', type=float, default=0.0001, help='beta start from this')
parser.add_argument('--beta_end', type=float, default=0.2, help='beta end to this')
parser.add_argument('--num_steps', type=float, default=50, help='denoising steps')
parser.add_argument('--num_samples', type=int, default=10, help='n datasets')
parser.add_argument('--schedule', type=str, default='quad', help='noise schedule type')
parser.add_argument('--target_strategy', type=str, default='random', help='mask')

# args for mae
parser.add_argument('--num_heads', type=int, default=8, help='n heads for self attention')
config = parser.parse_args([])

Fig 3b: Missing rate is equal to 0.1

from torch.utils.data import DataLoader
test_dloader = DataLoader(PRE8dDataset(config, mode='test'), 1, shuffle=False)
datas, data_ob_masks, data_gt_masks, labels, label_masks = next(iter(test_dloader))
device = "cuda"
datas, data_ob_masks, data_gt_masks, labels, label_masks = datas.float().to(device), data_ob_masks.to(device), data_gt_masks.to(device), labels.to(device), label_masks.to(device)

# load model
model = torch.load("./log_bak/imputation/PRE/STIMP/best_0.1.pt")
model = model.to(device)
cond_mask = data_gt_masks
adj = np.load("./data/{}/adj.npy".format(config.area))
adj = torch.from_numpy(adj).float().to(device)

imputed_data_our = model.impute(datas, cond_mask, adj, 10)
imputed_data_our = imputed_data_our.median(1).values
mask = data_ob_masks - cond_mask
imputed_our = imputed_data_our[0][mask.bool().cpu()[0]]
truth = datas[0][mask.bool()[0]].cpu()

from einops import rearrange
from model.dineof_per_step import DINEOF
model = DINEOF(10, [60, 96], keep_non_negative_only=False)
is_sea = np.load("./data/PRE/is_sea.npy")
datas_image = torch.zeros(1,46,1,60,96)
datas_image = datas_image.to(device)
datas_image[:,:,:,is_sea.astype(bool)]=datas
cond_mask_image = torch.zeros(1,46,1,60,96)
cond_mask_image = cond_mask_image.to(device)
cond_mask_image[:,:,:,is_sea.astype(bool)]=cond_mask
ob_mask_image = torch.zeros(1,46,1,60,96)
ob_mask_image = ob_mask_image.to(device)
ob_mask_image[:,:,:,is_sea.astype(bool)]=data_ob_masks

impute_data_list = []
for t in range(datas_image.shape[1]):
    data = datas_image[:, t, :, :, :].squeeze()

    tmp_data = torch.where(cond_mask_image[:,t].cpu().squeeze()==0, float("nan"), data.cpu())
    model.fit(tmp_data.numpy())

    imputed_data = model.predict()
    imputed_data = rearrange(imputed_data, "(b t c h) w->b t c h w", b=1, t=1, c=1, h=data.shape[-2], w=data.shape[-1])
    impute_data_list.append(torch.from_numpy(imputed_data))

imputed_data = torch.cat(impute_data_list, dim=1).numpy()

2025-05-11 14:43:54.149 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 45: 0.0014910557074472308, 8.841627277433872e-06
2025-05-11 14:43:54.394 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 161: 0.0020145191811025143, 9.89530235528946e-06
2025-05-11 14:43:54.670 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 180: 0.0018620046321302652, 9.817071259021759e-06
2025-05-11 14:43:54.934 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 177: 0.002274477621540427, 9.705079719424248e-06
2025-05-11 14:43:55.162 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 152: 0.002120467135682702, 9.96631570160389e-06
2025-05-11 14:43:55.324 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 106: 0.004458919167518616, 9.90321859717369e-06
2025-05-11 14:43:55.610 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 194: 0.002096793847158551, 9.964453056454659e-06
2025-05-11 14:43:55.789 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 121: 0.001970689045265317, 9.945826604962349e-06
2025-05-11 14:43:56.002 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 144: 0.0014974757796153426, 9.927665814757347e-06
2025-05-11 14:43:56.169 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 113: 0.0014182495651766658, 9.937561117112637e-06
2025-05-11 14:43:56.343 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 117: 0.0020637093111872673, 9.991228580474854e-06
2025-05-11 14:43:56.526 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 124: 0.0016064144438132644, 9.940005838871002e-06
2025-05-11 14:43:56.677 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 101: 0.0014993631048128009, 9.976793080568314e-06
2025-05-11 14:43:56.851 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 120: 0.0014136313693597913, 9.938492439687252e-06
2025-05-11 14:43:57.048 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 126: 0.0026658896822482347, 9.852228686213493e-06
2025-05-11 14:43:57.180 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 90: 0.0006393097573891282, 9.605777449905872e-06
2025-05-11 14:43:57.331 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 104: 0.0014122845605015755, 9.945360943675041e-06
2025-05-11 14:43:57.467 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 92: 0.0018238031771034002, 9.910902008414268e-06
2025-05-11 14:43:57.708 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 164: 0.0015448576305061579, 9.923707693815231e-06
2025-05-11 14:43:57.952 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 167: 0.0015645741950720549, 9.648152627050877e-06
2025-05-11 14:43:58.099 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 100: 0.002514046151190996, 9.73953865468502e-06
2025-05-11 14:43:58.293 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 133: 0.0019750131759792566, 9.955372661352158e-06
2025-05-11 14:43:58.452 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 106: 0.0019571008160710335, 9.746290743350983e-06
2025-05-11 14:43:58.610 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 107: 0.0011961341369897127, 9.978190064430237e-06
2025-05-11 14:43:58.706 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 67: 0.0007462473004125059, 9.821378625929356e-06
2025-05-11 14:43:58.913 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 141: 0.002045384841039777, 9.592156857252121e-06
2025-05-11 14:43:59.052 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 98: 0.0007599854725413024, 9.92434797808528e-06
2025-05-11 14:43:59.461 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 299: nan, nan
2025-05-11 14:43:59.601 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 95: 0.0016603971598669887, 9.63511411100626e-06
2025-05-11 14:43:59.734 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 91: 0.0007316667470149696, 9.928364306688309e-06
2025-05-11 14:43:59.952 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 150: 0.002029009163379669, 9.980518370866776e-06
2025-05-11 14:44:00.182 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 157: 0.0019494262523949146, 9.899260476231575e-06
2025-05-11 14:44:00.302 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 81: 0.0007577738142572343, 9.739713277667761e-06
2025-05-11 14:44:00.529 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 155: 0.0022052577696740627, 9.93930734694004e-06
2025-05-11 14:44:00.707 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 122: 0.001389348995871842, 9.925337508320808e-06
2025-05-11 14:44:00.993 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 196: 0.0022645005956292152, 9.95071604847908e-06
2025-05-11 14:44:01.115 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 84: 0.0009352410561405122, 9.918527211993933e-06
2025-05-11 14:44:01.330 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 148: 0.0012891022488474846, 9.6586300060153e-06
2025-05-11 14:44:01.638 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 198: 0.0025061771739274263, 9.871786460280418e-06
2025-05-11 14:44:01.912 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 181: 0.0025455141440033913, 9.87248495221138e-06
2025-05-11 14:44:02.108 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 131: 0.0013296843972057104, 9.854091331362724e-06
2025-05-11 14:44:02.323 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 146: 0.001996497390791774, 9.90252010524273e-06
2025-05-11 14:44:02.481 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 108: 0.001363650313578546, 9.863520972430706e-06
2025-05-11 14:44:02.651 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 117: 0.0012905821204185486, 9.911716915667057e-06
2025-05-11 14:44:02.770 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 82: 0.0009908095235005021, 9.638257324695587e-06
2025-05-11 14:44:03.011 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 168: 0.0031585937831550837, 9.780516847968102e-06

imputed_data_dineof = imputed_data[:,:,:,is_sea.astype(bool)]
imputed_dineof = imputed_data_dineof[0][mask.bool().cpu()[0]]
imputed_our = imputed_our[:]
imputed_dineof = imputed_dineof[:]
truth = truth[:]

import matplotlib.pyplot as plt
import seaborn as sns
import scipy.stats as stat
import pandas as pd
method = []
method.extend(['STIMP' for i in range(imputed_our.shape[0])])
method.extend(['DINEOF' for i in range(imputed_our.shape[0])])
data = {'truth': np.concatenate([truth.numpy() for i in range(2)]),
        'imputed':np.concatenate([imputed_our.numpy(), imputed_dineof]),
        'method':method}
data = pd.DataFrame.from_dict(data)
color =["#9F0000","#576fa0"]
g=sns.jointplot(data=data, x="truth", y="imputed", hue="method", kind='kde', palette=color, marginal_kws={'common_norm':False,'shade':True})
xpoints = ypoints = plt.xlim()
plt.ylim(ypoints)
plt.plot(xpoints, ypoints, 'k-', alpha=0.75, zorder=0)
plt.ylabel("imputed", size=24)
plt.xlabel("truth", size=24)
pcc1=stat.pearsonr(truth[:], imputed_our).statistic.item()
pcc2=stat.pearsonr(truth[:], imputed_dineof).statistic.item()
plt.text(0.40, 0.92, 'PCC={:.4f}'.format(pcc1), transform=plt.gca().transAxes, fontsize=12, verticalalignment='top', horizontalalignment='left', color="#9F0000")
plt.text(0.40, 0.88, 'PCC={:.4f}'.format(pcc2), transform=plt.gca().transAxes, fontsize=12, verticalalignment='top', horizontalalignment='left', color="#576fa0")
plt.xticks(fontsize=15)
plt.yticks(fontsize=15)

(array([-2. , -1.5, -1. , -0.5,  0. ,  0.5,  1. ,  1.5,  2. ,  2.5]),
 [Text(0, -2.0, '−2.0'),
  Text(0, -1.5, '−1.5'),
  Text(0, -1.0, '−1.0'),
  Text(0, -0.5, '−0.5'),
  Text(0, 0.0, '0.0'),
  Text(0, 0.5, '0.5'),
  Text(0, 1.0, '1.0'),
  Text(0, 1.5, '1.5'),
  Text(0, 2.0, '2.0'),
  Text(0, 2.5, '2.5')])

Fig 3b: Missing rate is equal to 0.3

config.missing_ratio=0.3
from torch.utils.data import DataLoader
test_dloader = DataLoader(PRE8dDataset(config, mode='test'), 1, shuffle=False)
datas, data_ob_masks, data_gt_masks, labels, label_masks = next(iter(test_dloader))
device = "cuda"
datas, data_ob_masks, data_gt_masks, labels, label_masks = datas.float().to(device), data_ob_masks.to(device), data_gt_masks.to(device), labels.to(device), label_masks.to(device)

model = torch.load("./log_bak/imputation/PRE/STIMP/best_0.3.pt")
model = model.to(device)
cond_mask = data_gt_masks
imputed_data_our = model.impute(datas, cond_mask, adj, 10)
imputed_data_our = imputed_data_our.median(1).values
mask = data_ob_masks - cond_mask
imputed_our = imputed_data_our[0][mask.bool().cpu()[0]]
truth = datas[0][mask.bool()[0]].cpu()

from einops import rearrange
from model.dineof_per_step import DINEOF
model = DINEOF(10, [60, 96], keep_non_negative_only=False)
is_sea = np.load("./data/PRE/is_sea.npy")
datas_image = torch.zeros(1,46,1,60,96)
datas_image = datas_image.to(device)
datas_image[:,:,:,is_sea.astype(bool)]=datas
cond_mask_image = torch.zeros(1,46,1,60,96)
cond_mask_image = cond_mask_image.to(device)
cond_mask_image[:,:,:,is_sea.astype(bool)]=cond_mask
ob_mask_image = torch.zeros(1,46,1,60,96)
ob_mask_image = ob_mask_image.to(device)
ob_mask_image[:,:,:,is_sea.astype(bool)]=data_ob_masks

impute_data_list = []
for t in range(datas_image.shape[1]):
    data = datas_image[:, t, :, :, :].squeeze()

    tmp_data = torch.where(cond_mask_image[:,t].cpu().squeeze()==0, float("nan"), data.cpu())
    model.fit(tmp_data.numpy())

    imputed_data = model.predict()
    imputed_data = rearrange(imputed_data, "(b t c h) w->b t c h w", b=1, t=1, c=1, h=data.shape[-2], w=data.shape[-1])
    impute_data_list.append(torch.from_numpy(imputed_data))

imputed_data = torch.cat(impute_data_list, dim=1).numpy()

2025-05-11 15:05:50.678 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 44: 0.005508442874997854, 8.286908268928528e-06
2025-05-11 15:05:50.924 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 162: 0.001973017817363143, 9.736744686961174e-06
2025-05-11 15:05:51.152 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 148: 0.002047701505944133, 9.890878573060036e-06
2025-05-11 15:05:51.398 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 163: 0.0021193658467382193, 9.947223588824272e-06
2025-05-11 15:05:51.625 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 151: 0.001720884582027793, 9.93930734694004e-06
2025-05-11 15:05:51.884 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 173: 0.002218881156295538, 9.912531822919846e-06
2025-05-11 15:05:52.104 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 144: 0.0018800323596224189, 9.852694347500801e-06
2025-05-11 15:05:52.306 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 134: 0.001716161728836596, 9.738607332110405e-06
2025-05-11 15:05:52.473 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 110: 0.0016324546886608005, 9.94082074612379e-06
2025-05-11 15:05:52.639 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 111: 0.001017803093418479, 9.90019179880619e-06
2025-05-11 15:05:52.837 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 130: 0.0014110855991020799, 9.903335012495518e-06
2025-05-11 15:05:52.992 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 104: 0.0014970683259889483, 9.936047717928886e-06
2025-05-11 15:05:53.176 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 123: 0.001167680136859417, 9.986921213567257e-06
2025-05-11 15:05:53.358 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 123: 0.00117063382640481, 9.872368536889553e-06
2025-05-11 15:05:53.585 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 151: 0.0017393412999808788, 9.719515219330788e-06
2025-05-11 15:05:53.707 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 84: 0.0005303223733790219, 9.860959835350513e-06
2025-05-11 15:05:53.871 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 111: 0.0011724329087883234, 9.718118235468864e-06
2025-05-11 15:05:54.061 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 128: 0.0013448175741359591, 9.94920264929533e-06
2025-05-11 15:05:54.263 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 135: 0.0014139995910227299, 9.837327525019646e-06
2025-05-11 15:05:54.444 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 120: 0.001958453794941306, 9.831506758928299e-06
2025-05-11 15:05:54.627 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 123: 0.0015418785624206066, 9.981682524085045e-06
2025-05-11 15:05:54.869 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 162: 0.001466072746552527, 9.8812161013484e-06
2025-05-11 15:05:55.063 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 129: 0.0012096288846805692, 9.991112165153027e-06
2025-05-11 15:05:55.199 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 91: 0.0009915523696690798, 9.787734597921371e-06
2025-05-11 15:05:55.289 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 62: 0.00033707829425111413, 9.886251064017415e-06
2025-05-11 15:05:55.533 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 166: 0.0016458609607070684, 9.957468137145042e-06
2025-05-11 15:05:55.624 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 63: 0.000758700945880264, 9.71485860645771e-06
2025-05-11 15:05:56.029 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 299: nan, nan
2025-05-11 15:05:56.230 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 136: 0.0017096907831728458, 9.979703463613987e-06
2025-05-11 15:05:56.335 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 72: 0.0008733951253816485, 9.948504157364368e-06
2025-05-11 15:05:56.531 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 130: 0.0018548358930274844, 9.798095561563969e-06
2025-05-11 15:05:56.708 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 119: 0.00184019620064646, 9.799376130104065e-06
2025-05-11 15:05:56.851 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 98: 0.000935693911742419, 9.845243766903877e-06
2025-05-11 15:05:57.067 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 145: 0.0018997704610228539, 9.892857633531094e-06
2025-05-11 15:05:57.278 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 144: 0.002222516806796193, 9.878305718302727e-06
2025-05-11 15:05:57.534 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 172: 0.0018725380068644881, 9.985757060348988e-06
2025-05-11 15:05:57.658 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 85: 0.0010096868500113487, 9.61078330874443e-06
2025-05-11 15:05:57.822 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 109: 0.001317586051300168, 9.927898645401001e-06
2025-05-11 15:05:58.076 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 167: 0.0021525188349187374, 9.770505130290985e-06
2025-05-11 15:05:58.302 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 150: 0.0015805780421942472, 9.531737305223942e-06
2025-05-11 15:05:58.484 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 122: 0.0010255500674247742, 9.659910574555397e-06
2025-05-11 15:05:58.723 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 159: 0.0014296522131189704, 9.834999218583107e-06
2025-05-11 15:05:58.896 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 116: 0.0011437289649620652, 9.891926310956478e-06
2025-05-11 15:05:59.052 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 105: 0.0014254780253395438, 9.890296496450901e-06
2025-05-11 15:05:59.191 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 94: 0.0008255378343164921, 9.934243280440569e-06
2025-05-11 15:05:59.467 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 185: 0.0018339669331908226, 9.786570444703102e-06

imputed_data_dineof = imputed_data[:,:,:,is_sea.astype(bool)]
imputed_dineof = imputed_data_dineof[0][mask.bool().cpu()[0]]
imputed_our = imputed_our[:]
imputed_dineof = imputed_dineof[:]
truth = truth[:]

import matplotlib.pyplot as plt
import seaborn as sns
import scipy.stats as stat
import pandas as pd
method = []
method.extend(['STIMP' for i in range(imputed_our.shape[0])])
method.extend(['DINEOF' for i in range(imputed_our.shape[0])])
data = {'truth': np.concatenate([truth.numpy() for i in range(2)]),
        'imputed':np.concatenate([imputed_our.numpy(), imputed_dineof]),
        'method':method}
data = pd.DataFrame.from_dict(data)
color =["#9F0000","#576fa0"]
g=sns.jointplot(data=data, x="truth", y="imputed", hue="method", kind='kde', palette=color, marginal_kws={'common_norm':False,'shade':True})
xpoints = ypoints = plt.xlim()
plt.ylim(ypoints)
plt.plot(xpoints, ypoints, 'k-', alpha=0.75, zorder=0)
plt.ylabel("imputed", size=24)
plt.xlabel("truth", size=24)
pcc1=stat.pearsonr(truth[:], imputed_our).statistic.item()
pcc2=stat.pearsonr(truth[:], imputed_dineof).statistic.item()
plt.text(0.40, 0.92, 'PCC={:.4f}'.format(pcc1), transform=plt.gca().transAxes, fontsize=12, verticalalignment='top', horizontalalignment='left', color="#9F0000")
plt.text(0.40, 0.88, 'PCC={:.4f}'.format(pcc2), transform=plt.gca().transAxes, fontsize=12, verticalalignment='top', horizontalalignment='left', color="#576fa0")
plt.xticks(fontsize=15)
plt.yticks(fontsize=15)

(array([-2. , -1.5, -1. , -0.5,  0. ,  0.5,  1. ,  1.5,  2. ,  2.5]),
 [Text(0, -2.0, '−2.0'),
  Text(0, -1.5, '−1.5'),
  Text(0, -1.0, '−1.0'),
  Text(0, -0.5, '−0.5'),
  Text(0, 0.0, '0.0'),
  Text(0, 0.5, '0.5'),
  Text(0, 1.0, '1.0'),
  Text(0, 1.5, '1.5'),
  Text(0, 2.0, '2.0'),
  Text(0, 2.5, '2.5')])

Fig 3b: Missing rate is equal to 0.5

config.missing_ratio=0.5
from torch.utils.data import DataLoader
test_dloader = DataLoader(PRE8dDataset(config, mode='test'), 1, shuffle=False)
datas, data_ob_masks, data_gt_masks, labels, label_masks = next(iter(test_dloader))
device = "cuda"
datas, data_ob_masks, data_gt_masks, labels, label_masks = datas.float().to(device), data_ob_masks.to(device), data_gt_masks.to(device), labels.to(device), label_masks.to(device)

model = torch.load("./log_bak/imputation/PRE/STIMP/best_0.5.pt")
model = model.to(device)
cond_mask = data_gt_masks
imputed_data_our = model.impute(datas, cond_mask, adj, 10)
imputed_data_our = imputed_data_our.median(1).values
mask = data_ob_masks - cond_mask
imputed_our = imputed_data_our[0][mask.bool().cpu()[0]]
truth = datas[0][mask.bool()[0]].cpu()

from einops import rearrange
from model.dineof_per_step import DINEOF
model = DINEOF(10, [60, 96, config.in_len], keep_non_negative_only=False)
is_sea = np.load("./data/PRE/is_sea.npy")
datas_image = torch.zeros(1,46,1,60,96)
datas_image = datas_image.to(device)
datas_image[:,:,:,is_sea.astype(bool)]=datas
cond_mask_image = torch.zeros(1,46,1,60,96)
cond_mask_image = cond_mask_image.to(device)
cond_mask_image[:,:,:,is_sea.astype(bool)]=cond_mask
ob_mask_image = torch.zeros(1,46,1,60,96)
ob_mask_image = ob_mask_image.to(device)
ob_mask_image[:,:,:,is_sea.astype(bool)]=data_ob_masks

impute_data_list = []
for t in range(datas_image.shape[1]):
    data = datas_image[:, t, :, :, :].squeeze()

    tmp_data = torch.where(cond_mask_image[:,t].cpu().squeeze()==0, float("nan"), data.cpu())
    model.fit(tmp_data.numpy())

    imputed_data = model.predict()
    imputed_data = rearrange(imputed_data, "(b t c h) w->b t c h w", b=1, t=1, c=1, h=data.shape[-2], w=data.shape[-1])
    impute_data_list.append(torch.from_numpy(imputed_data))

imputed_data = torch.cat(impute_data_list, dim=1).numpy()

2025-05-11 14:47:29.862 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 67: 0.0019250796176493168, 7.151509635150433e-06
2025-05-11 14:47:30.102 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 147: 0.0013824701309204102, 9.963405318558216e-06
2025-05-11 14:47:30.326 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 145: 0.0017284167697653174, 9.819050319492817e-06
2025-05-11 14:47:30.512 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 120: 0.001543168444186449, 9.693088941276073e-06
2025-05-11 14:47:30.719 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 134: 0.0016107747796922922, 9.767943993210793e-06
2025-05-11 14:47:30.946 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 148: 0.0017730435356497765, 9.90892294794321e-06
2025-05-11 14:47:31.175 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 145: 0.0017320731421932578, 9.867013432085514e-06
2025-05-11 14:47:31.412 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 155: 0.0014214974362403154, 9.721145033836365e-06
2025-05-11 14:47:31.611 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 131: 0.0014656129060313106, 9.95188020169735e-06
2025-05-11 14:47:31.761 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 99: 0.0007779413717798889, 9.909563232213259e-06
2025-05-11 14:47:31.939 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 115: 0.0014224384212866426, 9.657465852797031e-06
2025-05-11 14:47:32.083 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 95: 0.001020384719595313, 9.709387086331844e-06
2025-05-11 14:47:32.235 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 101: 0.0006059801089577377, 9.987445082515478e-06
2025-05-11 14:47:32.390 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 104: 0.0010464841034263372, 9.85804945230484e-06
2025-05-11 14:47:32.614 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 146: 0.001764750457368791, 9.876443073153496e-06
2025-05-11 14:47:32.748 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 91: 0.0005355632165446877, 9.935407433658838e-06
2025-05-11 14:47:32.898 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 100: 0.001080273650586605, 9.836163371801376e-06
2025-05-11 14:47:33.038 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 92: 0.0012320331297814846, 9.879004210233688e-06
2025-05-11 14:47:33.235 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 129: 0.0015168003737926483, 9.96515154838562e-06
2025-05-11 14:47:33.407 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 112: 0.0016239298274740577, 9.749084711074829e-06
2025-05-11 14:47:33.601 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 128: 0.0014908276498317719, 9.697279892861843e-06
2025-05-11 14:47:33.789 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 124: 0.0013005960499867797, 9.93441790342331e-06
2025-05-11 14:47:33.936 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 97: 0.0009410178172402084, 9.914510883390903e-06
2025-05-11 14:47:34.072 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 90: 0.0005773325683549047, 9.771727491170168e-06
2025-05-11 14:47:34.125 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 35: 6.185279198689386e-05, 9.9055832833983e-06
2025-05-11 14:47:34.407 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 188: 0.0017637049313634634, 9.859679266810417e-06
2025-05-11 14:47:34.482 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 50: 0.00015991614782251418, 9.699928341433406e-06
2025-05-11 14:47:34.888 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 299: nan, nan
2025-05-11 14:47:35.120 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 157: 0.0014785535167902708, 9.812298230826855e-06
2025-05-11 14:47:35.215 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 65: 0.0003315122448839247, 9.771960321813822e-06
2025-05-11 14:47:35.437 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 144: 0.001408361946232617, 9.834766387939453e-06
2025-05-11 14:47:35.619 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 119: 0.0014464706182479858, 9.915791451931e-06
2025-05-11 14:47:35.747 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 85: 0.0008855302585288882, 9.816722013056278e-06
2025-05-11 14:47:36.023 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 182: 0.001519430661574006, 9.882496669888496e-06
2025-05-11 14:47:36.302 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 188: 0.001653132727369666, 9.796698577702045e-06
2025-05-11 14:47:36.531 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 150: 0.0016559042269364, 9.79460310190916e-06
2025-05-11 14:47:36.679 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 102: 0.0006868503405712545, 9.83185600489378e-06
2025-05-11 14:47:36.819 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 92: 0.001118299551308155, 9.978655725717545e-06
2025-05-11 14:47:37.065 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 159: 0.001680708839558065, 9.894603863358498e-06
2025-05-11 14:47:37.274 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 132: 0.001894683577120304, 9.999959729611874e-06
2025-05-11 14:47:37.426 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 99: 0.000961434212513268, 9.998329915106297e-06
2025-05-11 14:47:37.657 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 150: 0.0013353353133425117, 9.960494935512543e-06
2025-05-11 14:47:37.819 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 107: 0.0007167813018895686, 9.999494068324566e-06
2025-05-11 14:47:37.988 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 112: 0.0008635912090539932, 9.9594471976161e-06
2025-05-11 14:47:38.139 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 103: 0.0006369164912030101, 9.875744581222534e-06
2025-05-11 14:47:38.342 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 132: 0.0015423235017806292, 9.952578693628311e-06

imputed_data_dineof = imputed_data[:,:,:,is_sea.astype(bool)]
imputed_dineof = imputed_data_dineof[0][mask.bool().cpu()[0]]
imputed_our = imputed_our[:]
imputed_dineof = imputed_dineof[:]
truth = truth[:]

import matplotlib.pyplot as plt
import seaborn as sns
import scipy.stats as stat
import pandas as pd
method = []
method.extend(['STIMP' for i in range(imputed_our.shape[0])])
method.extend(['DINEOF' for i in range(imputed_our.shape[0])])
data = {'truth': np.concatenate([truth.numpy() for i in range(2)]),
        'imputed':np.concatenate([imputed_our.numpy(), imputed_dineof]),
        'method':method}
data = pd.DataFrame.from_dict(data)
color =["#9F0000","#576fa0"]
g=sns.jointplot(data=data, x="truth", y="imputed", hue="method", kind='kde', palette=color, marginal_kws={'common_norm':False,'shade':True})
# plt.scatter(truth, imputed_our,s=10, c="#9F0000",label="Our")
# plt.scatter(truth, imputed_mae.cpu().detach().numpy(),s=10, c="#e3b87f", label="MAE")
# plt.scatter(truth, imputed_itp.cpu(),s=10, c="#576fa0", label="Lin-ITP")
xpoints = ypoints = plt.xlim()
plt.ylim(ypoints)
plt.plot(xpoints, ypoints, 'k-', alpha=0.75, zorder=0)
plt.ylabel("imputed", size=24)
plt.xlabel("truth", size=24)
print(stat.pearsonr(truth[:], imputed_our))
print(stat.pearsonr(truth[:], imputed_dineof))
pcc1=stat.pearsonr(truth[:], imputed_our).statistic.item()
pcc2=stat.pearsonr(truth[:], imputed_dineof).statistic.item()
plt.text(0.40, 0.92, 'PCC={:.4f}'.format(pcc1), transform=plt.gca().transAxes, fontsize=12, verticalalignment='top', horizontalalignment='left', color="#9F0000")
plt.text(0.40, 0.88, 'PCC={:.4f}'.format(pcc2), transform=plt.gca().transAxes, fontsize=12, verticalalignment='top', horizontalalignment='left', color="#576fa0")
plt.xticks(fontsize=15)
plt.yticks(fontsize=15)

PearsonRResult(statistic=0.9868376743082833, pvalue=0.0)
PearsonRResult(statistic=0.9018553339227753, pvalue=0.0)

(array([-2. , -1.5, -1. , -0.5,  0. ,  0.5,  1. ,  1.5,  2. ,  2.5]),
 [Text(0, -2.0, '−2.0'),
  Text(0, -1.5, '−1.5'),
  Text(0, -1.0, '−1.0'),
  Text(0, -0.5, '−0.5'),
  Text(0, 0.0, '0.0'),
  Text(0, 0.5, '0.5'),
  Text(0, 1.0, '1.0'),
  Text(0, 1.5, '1.5'),
  Text(0, 2.0, '2.0'),
  Text(0, 2.5, '2.5')])

Fig 3b: Missing rate is equal to 0.7

config.missing_ratio=0.7
from torch.utils.data import DataLoader
test_dloader = DataLoader(PRE8dDataset(config, mode='test'), 1, shuffle=False)
datas, data_ob_masks, data_gt_masks, labels, label_masks = next(iter(test_dloader))
device = "cuda"
datas, data_ob_masks, data_gt_masks, labels, label_masks = datas.float().to(device), data_ob_masks.to(device), data_gt_masks.to(device), labels.to(device), label_masks.to(device)

adj = np.load("./data/{}/adj.npy".format(config.area))
adj = torch.from_numpy(adj).float().to(device)
model = torch.load("./log/imputation/PRE/STIMP/best_0.9.pt")
model = model.to(device)
cond_mask = data_gt_masks
imputed_data_our = model.impute(datas, cond_mask, adj, 10)
imputed_data_our = imputed_data_our.median(1).values
mask = data_ob_masks - cond_mask
imputed_our = imputed_data_our[0][mask.bool().cpu()[0]]
truth = datas[0][mask.bool()[0]].cpu()

from einops import rearrange
from model.dineof_per_step import DINEOF
model = DINEOF(10, [60, 96], keep_non_negative_only=False)
is_sea = np.load("./data/PRE/is_sea.npy")
datas_image = torch.zeros(1,46,1,60,96)
datas_image = datas_image.to(device)
datas_image[:,:,:,is_sea.astype(bool)]=datas
cond_mask_image = torch.zeros(1,46,1,60,96)
cond_mask_image = cond_mask_image.to(device)
cond_mask_image[:,:,:,is_sea.astype(bool)]=cond_mask
ob_mask_image = torch.zeros(1,46,1,60,96)
ob_mask_image = ob_mask_image.to(device)
ob_mask_image[:,:,:,is_sea.astype(bool)]=data_ob_masks

impute_data_list = []
for t in range(datas_image.shape[1]):
    data = datas_image[:, t, :, :, :].squeeze()

    tmp_data = torch.where(cond_mask_image[:,t].cpu().squeeze()==0, float("nan"), data.cpu())
    model.fit(tmp_data.numpy())

    imputed_data = model.predict()
    imputed_data = rearrange(imputed_data, "(b t c h) w->b t c h w", b=1, t=1, c=1, h=data.shape[-2], w=data.shape[-1])
    impute_data_list.append(torch.from_numpy(imputed_data))

imputed_data = torch.cat(impute_data_list, dim=1).numpy()

2025-05-11 14:49:25.657 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 24: 0.0032221016008406878, 4.917383193969727e-07
2025-05-11 14:49:25.832 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 114: 0.0009344807476736605, 9.96061135083437e-06
2025-05-11 14:49:26.020 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 121: 0.0013941103825345635, 9.853276424109936e-06
2025-05-11 14:49:26.229 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 136: 0.0011940064141526818, 9.989133104681969e-06
2025-05-11 14:49:26.428 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 130: 0.0011690087849274278, 9.83069185167551e-06
2025-05-11 14:49:26.635 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 133: 0.0011790381977334619, 9.811134077608585e-06
2025-05-11 14:49:26.837 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 131: 0.0013803131878376007, 9.956187568604946e-06
2025-05-11 14:49:27.005 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 111: 0.000692429777700454, 9.835464879870415e-06
2025-05-11 14:49:27.162 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 103: 0.0007930059218779206, 9.861832950264215e-06
2025-05-11 14:49:27.283 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 79: 0.00036197120789438486, 9.77562740445137e-06
2025-05-11 14:49:27.442 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 104: 0.0008035101927816868, 9.778770618140697e-06
2025-05-11 14:49:27.557 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 75: 0.000344570551533252, 9.949901141226292e-06
2025-05-11 14:49:27.655 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 64: 0.0003343119751662016, 9.818613762035966e-06
2025-05-11 14:49:27.785 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 86: 0.0005215647397562861, 9.848386980593204e-06
2025-05-11 14:49:27.977 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 125: 0.0010331820230931044, 9.882263839244843e-06
2025-05-11 14:49:28.099 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 83: 0.0001935218897415325, 9.684896213002503e-06
2025-05-11 14:49:28.229 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 86: 0.0005104696028865874, 9.768933523446321e-06
2025-05-11 14:49:28.360 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 87: 0.0004002027853857726, 9.719340596348047e-06
2025-05-11 14:49:28.510 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 99: 0.0007080034702084959, 9.662064258009195e-06
2025-05-11 14:49:28.673 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 106: 0.0007541952072642744, 9.828596375882626e-06
2025-05-11 14:49:28.839 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 110: 0.0007024802616797388, 9.858864359557629e-06
2025-05-11 14:49:28.994 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 102: 0.0007986794807948172, 9.92469722405076e-06
2025-05-11 14:49:29.111 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 77: 0.0005171731463633478, 9.891751687973738e-06
2025-05-11 14:49:29.172 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 39: 0.0006039125146344304, 2.3249885998666286e-06
2025-05-11 14:49:29.215 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 26: 3.8699177821399644e-05, 8.95281118573621e-06
2025-05-11 14:49:29.434 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 145: 0.0008976078242994845, 9.966141078621149e-06
2025-05-11 14:49:29.478 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 29: 3.5128756280755624e-05, 9.401148417964578e-06
2025-05-11 14:49:29.885 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 299: nan, nan
2025-05-11 14:49:30.012 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 85: 0.00019694474758580327, 9.756549843586981e-06
2025-05-11 14:49:30.064 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 34: 0.00010439470497658476, 9.46773070609197e-06
2025-05-11 14:49:30.260 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 128: 0.0012598203029483557, 9.89111140370369e-06
2025-05-11 14:49:30.417 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 103: 0.0006935808924026787, 9.913346730172634e-06
2025-05-11 14:49:30.523 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 70: 0.00021651925635524094, 9.704424883238971e-06
2025-05-11 14:49:30.721 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 132: 0.0009047722560353577, 9.799550753086805e-06
2025-05-11 14:49:30.976 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 174: 0.0010753634851425886, 9.926967322826385e-06
2025-05-11 14:49:31.210 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 154: 0.0011085536098107696, 9.785755537450314e-06
2025-05-11 14:49:31.280 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 47: 9.556740405969322e-05, 9.799754479900002e-06
2025-05-11 14:49:31.405 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 82: 0.000460479554021731, 9.82446363195777e-06
2025-05-11 14:49:31.613 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 136: 0.0011028937296941876, 9.840819984674454e-06
2025-05-11 14:49:31.805 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 123: 0.001306606107391417, 9.897863492369652e-06
2025-05-11 14:49:31.933 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 82: 0.0003792727366089821, 9.94509900920093e-06
2025-05-11 14:49:32.102 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 111: 0.0007979529909789562, 9.922892786562443e-06
2025-05-11 14:49:32.219 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 77: 0.0002964245213661343, 9.707669960334897e-06
2025-05-11 14:49:32.346 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 83: 0.0004223940777592361, 9.755021892488003e-06
2025-05-11 14:49:32.426 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 53: 0.00012010508362436667, 9.443036105949432e-06
2025-05-11 14:49:32.641 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 141: 0.0012422060826793313, 9.99518670141697e-06

imputed_data_dineof = imputed_data[:,:,:,is_sea.astype(bool)]
imputed_dineof = imputed_data_dineof[0][mask.bool().cpu()[0]]
imputed_our = imputed_our[:]
imputed_dineof = imputed_dineof[:]
truth = truth[:]

import matplotlib.pyplot as plt
import seaborn as sns
import scipy.stats as stat
import pandas as pd
method = []
method.extend(['STIMP' for i in range(imputed_our.shape[0])])
method.extend(['DINEOF' for i in range(imputed_our.shape[0])])
data = {'truth': np.concatenate([truth.numpy() for i in range(2)]),
        'imputed':np.concatenate([imputed_our.numpy(), imputed_dineof]),
        'method':method}
data = pd.DataFrame.from_dict(data)
color =["#9F0000","#576fa0"]
g=sns.jointplot(data=data, x="truth", y="imputed", hue="method", kind='kde', palette=color, marginal_kws={'common_norm':False,'shade':True})
# plt.scatter(truth, imputed_our,s=10, c="#9F0000",label="Our")
# plt.scatter(truth, imputed_mae.cpu().detach().numpy(),s=10, c="#e3b87f", label="MAE")
# plt.scatter(truth, imputed_itp.cpu(),s=10, c="#576fa0", label="Lin-ITP")
xpoints = ypoints = plt.xlim()
plt.ylim(ypoints)
plt.plot(xpoints, ypoints, 'k-', alpha=0.75, zorder=0)
plt.ylabel("imputed", size=24)
plt.xlabel("truth", size=24)
print(stat.pearsonr(truth[:], imputed_our))
print(stat.pearsonr(truth[:], imputed_dineof))
pcc1=stat.pearsonr(truth[:], imputed_our).statistic.item()
pcc2=stat.pearsonr(truth[:], imputed_dineof).statistic.item()
plt.text(0.40, 0.92, 'PCC={:.4f}'.format(pcc1), transform=plt.gca().transAxes, fontsize=12, verticalalignment='top', horizontalalignment='left', color="#9F0000")
plt.text(0.40, 0.88, 'PCC={:.4f}'.format(pcc2), transform=plt.gca().transAxes, fontsize=12, verticalalignment='top', horizontalalignment='left', color="#576fa0")
plt.xticks(fontsize=15)
plt.yticks(fontsize=15)

PearsonRResult(statistic=0.9551173753689999, pvalue=0.0)
PearsonRResult(statistic=0.7227783409335321, pvalue=0.0)

(array([-2. , -1.5, -1. , -0.5,  0. ,  0.5,  1. ,  1.5,  2. ,  2.5]),
 [Text(0, -2.0, '−2.0'),
  Text(0, -1.5, '−1.5'),
  Text(0, -1.0, '−1.0'),
  Text(0, -0.5, '−0.5'),
  Text(0, 0.0, '0.0'),
  Text(0, 0.5, '0.5'),
  Text(0, 1.0, '1.0'),
  Text(0, 1.5, '1.5'),
  Text(0, 2.0, '2.0'),
  Text(0, 2.5, '2.5')])

Fig 3b: Missing rate is equal to 0.9

config.missing_ratio=0.9
from torch.utils.data import DataLoader
test_dloader = DataLoader(PRE8dDataset(config, mode='test'), 1, shuffle=False)
datas, data_ob_masks, data_gt_masks, labels, label_masks = next(iter(test_dloader))
device = "cuda"
datas, data_ob_masks, data_gt_masks, labels, label_masks = datas.float().to(device), data_ob_masks.to(device), data_gt_masks.to(device), labels.to(device), label_masks.to(device)

adj = np.load("./data/{}/adj.npy".format(config.area))
adj = torch.from_numpy(adj).float().to(device)
model = torch.load("./log/imputation/PRE/STIMP/best_0.9.pt")
model = model.to(device)
cond_mask = data_gt_masks
imputed_data_our = model.impute(datas, cond_mask, adj, 10)
imputed_data_our = imputed_data_our.median(1).values
mask = data_ob_masks - cond_mask
imputed_our = imputed_data_our[0][mask.bool().cpu()[0]]
truth = datas[0][mask.bool()[0]].cpu()

from einops import rearrange
from model.dineof_per_step import DINEOF
model = DINEOF(10, [60, 96], keep_non_negative_only=False)
is_sea = np.load("./data/PRE/is_sea.npy")
datas_image = torch.zeros(1,46,1,60,96)
datas_image = datas_image.to(device)
datas_image[:,:,:,is_sea.astype(bool)]=datas
cond_mask_image = torch.zeros(1,46,1,60,96)
cond_mask_image = cond_mask_image.to(device)
cond_mask_image[:,:,:,is_sea.astype(bool)]=cond_mask
ob_mask_image = torch.zeros(1,46,1,60,96)
ob_mask_image = ob_mask_image.to(device)
ob_mask_image[:,:,:,is_sea.astype(bool)]=data_ob_masks

impute_data_list = []
for t in range(datas_image.shape[1]):
    data = datas_image[:, t, :, :, :].squeeze()

    tmp_data = torch.where(cond_mask_image[:,t].cpu().squeeze()==0, float("nan"), data.cpu())
    model.fit(tmp_data.numpy())

    imputed_data = model.predict()
    imputed_data = rearrange(imputed_data, "(b t c h) w->b t c h w", b=1, t=1, c=1, h=data.shape[-2], w=data.shape[-1])
    impute_data_list.append(torch.from_numpy(imputed_data))

imputed_data = torch.cat(impute_data_list, dim=1).numpy()

2025-05-11 14:54:32.351 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 65: 0.000706246355548501, 7.287424523383379e-06
2025-05-11 14:54:32.433 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 53: 0.0002873373159673065, 9.953597327694297e-06
2025-05-11 14:54:32.530 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 62: 0.0001635895314393565, 9.94822767097503e-06
2025-05-11 14:54:32.636 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 68: 0.00018984131747856736, 9.668496204540133e-06
2025-05-11 14:54:32.725 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 57: 0.00013677547394763678, 9.794297511689365e-06
2025-05-11 14:54:32.839 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 74: 0.00017466300050728023, 9.6262083388865e-06
2025-05-11 14:54:32.933 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 61: 0.0001767123758327216, 9.803450666368008e-06
2025-05-11 14:54:33.061 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 85: 0.00044902495574206114, 9.67431697063148e-06
2025-05-11 14:54:33.140 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 50: 0.0001354789565084502, 9.745985153131187e-06
2025-05-11 14:54:33.244 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 68: 0.00017697038128972054, 9.69927350524813e-06
2025-05-11 14:54:33.329 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 54: 0.00015260894724633545, 9.498631698079407e-06
2025-05-11 14:54:33.408 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 51: 0.00015384849393740296, 9.882336598820984e-06
2025-05-11 14:54:33.463 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 33: 0.0010427890811115503, 9.971437975764275e-06
2025-05-11 14:54:33.490 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 13: 0.008921586908400059, 7.611699402332306e-06
2025-05-11 14:54:33.588 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 64: 0.00016728427726775408, 9.733557817526162e-06
2025-05-11 14:54:33.689 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 62: 0.0007253197254613042, 7.278867997229099e-07
2025-05-11 14:54:33.770 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 51: 0.0002318904153071344, 9.505951311439276e-06
2025-05-11 14:54:33.844 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 47: 0.00023805000819265842, 8.193019311875105e-06
2025-05-11 14:54:33.932 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 57: 0.000141343756695278, 9.603551006875932e-06
2025-05-11 14:54:34.024 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 59: 0.000186920675332658, 9.765935828909278e-06
2025-05-11 14:54:34.123 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 63: 0.0002919269318226725, 9.981915354728699e-06
2025-05-11 14:54:34.188 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 39: 0.0008461946272291243, 2.434360794723034e-06
2025-05-11 14:54:34.277 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 48: 0.0008443274418823421, 2.4886103346943855e-06
2025-05-11 14:54:34.347 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 46: 0.00012933972175233066, 9.6485746325925e-06
2025-05-11 14:54:34.371 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 15: 2.013596531469375e-05, 5.8406130847288296e-06
2025-05-11 14:54:34.452 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 53: 9.276653145207092e-05, 9.398841939400882e-06
2025-05-11 14:54:34.478 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 16: 0.00011357730545569211, 2.7240748750045896e-06
2025-05-11 14:54:34.884 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 299: nan, nan
2025-05-11 14:54:34.963 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 52: 8.175906987162307e-05, 9.267205314245075e-06
2025-05-11 14:54:35.036 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 49: 2.1827143427799456e-05, 6.52942071610596e-06
2025-05-11 14:54:35.129 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 60: 0.00017567579925525934, 9.94287256617099e-06
2025-05-11 14:54:35.227 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 65: 0.0001758216822054237, 9.985320502892137e-06
2025-05-11 14:54:35.292 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 42: 0.00014831697626505047, 9.908224456012249e-06
2025-05-11 14:54:35.387 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 62: 0.00014934840146452188, 9.473427780903876e-06
2025-05-11 14:54:35.481 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 63: 0.00014056880900170654, 9.6159492386505e-06
2025-05-11 14:54:35.582 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 66: 0.00016137976490426809, 9.578638128004968e-06
2025-05-11 14:54:35.618 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 22: 3.2870542781893164e-05, 8.16872125142254e-06
2025-05-11 14:54:35.687 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 44: 0.00022198198712430894, 8.95107223186642e-06
2025-05-11 14:54:35.787 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 65: 0.00019465263176243752, 9.896379197016358e-06
2025-05-11 14:54:35.890 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 67: 0.0002210488310083747, 9.602386853657663e-06
2025-05-11 14:54:35.976 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 55: 0.00017595324607100338, 9.957919246517122e-06
2025-05-11 14:54:36.060 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 54: 0.0001307199418079108, 9.745810530148447e-06
2025-05-11 14:54:36.115 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 33: 0.002378571778535843, 5.875946953892708e-06
2025-05-11 14:54:36.196 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 52: 0.0004807891964446753, 8.983566658571362e-06
2025-05-11 14:54:36.252 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 36: 0.00016774308460298926, 9.97130700852722e-06
2025-05-11 14:54:36.360 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 71: 0.0002451837353873998, 9.888550266623497e-06

imputed_data_dineof = imputed_data[:,:,:,is_sea.astype(bool)]
imputed_dineof = imputed_data_dineof[0][mask.bool().cpu()[0]]
imputed_our = imputed_our[:]
imputed_dineof = imputed_dineof[:]
truth = truth[:]

import matplotlib.pyplot as plt
import seaborn as sns
import scipy.stats as stat
import pandas as pd
method = []
method.extend(['STIMP' for i in range(imputed_our.shape[0])])
method.extend(['DINEOF' for i in range(imputed_our.shape[0])])
data = {'truth': np.concatenate([truth.numpy() for i in range(2)]),
        'imputed':np.concatenate([imputed_our.numpy(), imputed_dineof]),
        'method':method}
data = pd.DataFrame.from_dict(data)
color =["#9F0000","#576fa0"]
g=sns.jointplot(data=data, x="truth", y="imputed", hue="method", kind='kde', palette=color, marginal_kws={'common_norm':False,'shade':True})
# plt.scatter(truth, imputed_our,s=10, c="#9F0000",label="Our")
# plt.scatter(truth, imputed_mae.cpu().detach().numpy(),s=10, c="#e3b87f", label="MAE")
# plt.scatter(truth, imputed_itp.cpu(),s=10, c="#576fa0", label="Lin-ITP")
xpoints = ypoints = plt.xlim()
plt.ylim(ypoints)
plt.plot(xpoints, ypoints, 'k-', alpha=0.75, zorder=0)
plt.ylabel("imputed", size=24)
plt.xlabel("truth", size=24)
print(stat.pearsonr(truth[:], imputed_our))
print(stat.pearsonr(truth[:], imputed_dineof))
pcc1=stat.pearsonr(truth[:], imputed_our).statistic.item()
pcc2=stat.pearsonr(truth[:], imputed_dineof).statistic.item()
plt.text(0.40, 0.92, 'PCC={:.4f}'.format(pcc1), transform=plt.gca().transAxes, fontsize=12, verticalalignment='top', horizontalalignment='left', color="#9F0000")
plt.text(0.40, 0.88, 'PCC={:.4f}'.format(pcc2), transform=plt.gca().transAxes, fontsize=12, verticalalignment='top', horizontalalignment='left', color="#576fa0")
plt.xticks(fontsize=15)
plt.yticks(fontsize=15)

PearsonRResult(statistic=0.9462578040774265, pvalue=0.0)
PearsonRResult(statistic=0.3536718379145061, pvalue=0.0)

(array([-2. , -1.5, -1. , -0.5,  0. ,  0.5,  1. ,  1.5,  2. ,  2.5]),
 [Text(0, -2.0, '−2.0'),
  Text(0, -1.5, '−1.5'),
  Text(0, -1.0, '−1.0'),
  Text(0, -0.5, '−0.5'),
  Text(0, 0.0, '0.0'),
  Text(0, 0.5, '0.5'),
  Text(0, 1.0, '1.0'),
  Text(0, 1.5, '1.5'),
  Text(0, 2.0, '2.0'),
  Text(0, 2.5, '2.5')])

Fig 3c: case study

from torch.utils.data import DataLoader
config.missing_ratio=0.1
test_dloader = DataLoader(PRE8dDataset(config, mode='test'), 1, shuffle=False)

datas, data_ob_masks, data_gt_masks, labels, label_masks = next(iter(test_dloader))
device = "cuda"
datas, data_ob_masks, data_gt_masks, labels, label_masks = datas.float().to(device), data_ob_masks.to(device), data_gt_masks.to(device), labels.to(device), label_masks.to(device)

adj = np.load("./data/{}/adj.npy".format(config.area))
adj = torch.from_numpy(adj).float().to(device)
model = torch.load("./log/imputation/PRE/STIMP/best_0.1.pt")
model = model.to(device)

cond_mask = data_ob_masks
imputed_data_our = model.impute(datas, cond_mask, adj, 10)
imputed_data_our = imputed_data_our.median(1).values
truth = datas[0]
truth[~data_ob_masks[0].bool()]=np.nan

from einops import rearrange
from model.dineof import DINEOF
model = DINEOF(10, [60, 96, config.in_len], keep_non_negative_only=False)
datas_image = torch.zeros(1,46,1,60,96)
datas_image = datas_image.to(device)
is_sea = np.load("./data/PRE/is_sea.npy")
datas_image[:,:,:,is_sea.astype(bool)]=datas
cond_mask_image = torch.zeros(1,46,1,60,96)
cond_mask_image = cond_mask_image.to(device)
cond_mask_image[:,:,:,is_sea.astype(bool)]=cond_mask
ob_mask_image = torch.zeros(1,46,1,60,96)
ob_mask_image = ob_mask_image.to(device)
ob_mask_image[:,:,:,is_sea.astype(bool)]=data_ob_masks

tmp_data = torch.where(ob_mask_image.cpu()==0, float("nan"), datas_image.cpu())
tmp_data = rearrange(tmp_data, "b t c h w -> (b h) (w c) t")
tmp_data = tmp_data.cpu().numpy()
model.fit(tmp_data)

imputed_data = model.predict()
imputed_data = rearrange(imputed_data, "(b h w c) t->b t c h w", b=1, t=config.in_len, c=1, h=datas_image.shape[-2], w=datas_image.shape[-1])

2025-05-11 14:59:31.080 | INFO     | model.dineof:_fit:103 - Error/Relative Error at iteraion 167: 0.0024101396556943655, 9.974697604775429e-06

from einops import rearrange
from model.dineof_per_step import DINEOF
model = DINEOF(5, [60, 96], keep_non_negative_only=False)
is_sea = np.load("./data/PRE/is_sea.npy")
datas_image = torch.zeros(1,46,1,60,96)
datas_image = datas_image.to(device)
datas_image[:,:,:,is_sea.astype(bool)]=datas
cond_mask_image = torch.zeros(1,46,1,60,96)
cond_mask_image = cond_mask_image.to(device)
cond_mask_image[:,:,:,is_sea.astype(bool)]=cond_mask
ob_mask_image = torch.zeros(1,46,1,60,96)
ob_mask_image = ob_mask_image.to(device)
ob_mask_image[:,:,:,is_sea.astype(bool)]=data_ob_masks

impute_data_list = []
for t in range(datas_image.shape[1]):
    data = datas_image[:, t, :, :, :].squeeze()

    tmp_data = torch.where(ob_mask_image[:,t].cpu().squeeze()==0, float("nan"), data.cpu())
    model.fit(tmp_data.numpy())

    imputed_data_per = model.predict()
    imputed_data_per = rearrange(imputed_data_per, "(b t c h) w->b t c h w", b=1, t=1, c=1, h=data.shape[-2], w=data.shape[-1])
    impute_data_list.append(torch.from_numpy(imputed_data_per))

imputed_data_per = torch.cat(impute_data_list, dim=1).numpy()

2025-05-11 14:59:31.244 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 56: 0.0041714743711054325, 5.664769560098648e-06
2025-05-11 14:59:31.450 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 152: 0.002235795371234417, 9.55723226070404e-06
2025-05-11 14:59:31.649 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 153: 0.002534731524065137, 9.534414857625961e-06
2025-05-11 14:59:31.927 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 190: 0.00366205838508904, 9.81055200099945e-06
2025-05-11 14:59:32.212 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 208: 0.0030324955005198717, 9.872019290924072e-06
2025-05-11 14:59:32.555 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 259: 0.004456556402146816, 9.830575436353683e-06
2025-05-11 14:59:32.774 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 167: 0.0037735402584075928, 9.89413820207119e-06
2025-05-11 14:59:32.998 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 165: 0.0030917623080313206, 9.970273822546005e-06
2025-05-11 14:59:33.181 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 138: 0.0017133936053141952, 9.71183180809021e-06
2025-05-11 14:59:33.428 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 181: 0.0019672042690217495, 9.860843420028687e-06
2025-05-11 14:59:33.481 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 37: 0.006010887678712606, 6.908085197210312e-06
2025-05-11 14:59:33.682 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 148: 0.0019257558742538095, 9.873299859464169e-06
2025-05-11 14:59:33.898 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 160: 0.0017487411387264729, 9.995419532060623e-06
2025-05-11 14:59:34.103 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 150: 0.002298995852470398, 9.957468137145042e-06
2025-05-11 14:59:34.348 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 187: 0.0025810913648456335, 9.90135595202446e-06
2025-05-11 14:59:34.461 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 84: 0.0011646002531051636, 9.960494935512543e-06
2025-05-11 14:59:34.684 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 171: 0.0015168560203164816, 9.78854950517416e-06
2025-05-11 14:59:34.850 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 121: 0.002321934560313821, 9.641516953706741e-06
2025-05-11 14:59:34.930 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 57: 0.006378342863172293, 9.470153599977493e-06
2025-05-11 14:59:35.188 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 194: 0.0020481988321989775, 9.415671229362488e-06
2025-05-11 14:59:35.358 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 124: 0.003460838459432125, 9.816372767090797e-06
2025-05-11 14:59:35.537 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 133: 0.00247183651663363, 9.868526831269264e-06
2025-05-11 14:59:35.744 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 153: 0.0017135960515588522, 9.670271538197994e-06
2025-05-11 14:59:35.864 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 88: 0.0014925585128366947, 9.725801646709442e-06
2025-05-11 14:59:36.002 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 102: 0.0013137105852365494, 9.886454790830612e-06
2025-05-11 14:59:36.267 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 199: 0.002770154271274805, 9.961659088730812e-06
2025-05-11 14:59:36.369 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 75: 0.0020057098008692265, 9.726034477353096e-06
2025-05-11 14:59:36.757 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 299: nan, nan
2025-05-11 14:59:36.958 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 151: 0.002161961980164051, 9.780749678611755e-06
2025-05-11 14:59:37.020 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 45: 0.0017848352435976267, 9.826966561377048e-06
2025-05-11 14:59:37.267 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 189: 0.002233447739854455, 9.871786460280418e-06
2025-05-11 14:59:37.523 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 194: 0.0021927852649241686, 9.64268110692501e-06
2025-05-11 14:59:37.601 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 58: 0.0016406765207648277, 9.615905582904816e-06
2025-05-11 14:59:37.865 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 198: 0.0034097570460289717, 9.979121387004852e-06
2025-05-11 14:59:38.055 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 144: 0.0015246866969391704, 9.599956683814526e-06
2025-05-11 14:59:38.335 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 216: 0.002692047506570816, 9.89111140370369e-06
2025-05-11 14:59:38.431 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 71: 0.0011859011137858033, 9.493203833699226e-06
2025-05-11 14:59:38.689 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 195: 0.002132347784936428, 9.97842289507389e-06
2025-05-11 14:59:38.971 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 214: 0.0024574596900492907, 9.776325896382332e-06
2025-05-11 14:59:39.082 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 83: 0.005691788624972105, 9.865965694189072e-06
2025-05-11 14:59:39.291 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 155: 0.0019555736798793077, 9.680632501840591e-06
2025-05-11 14:59:39.527 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 179: 0.0023386424873024225, 9.715091437101364e-06
2025-05-11 14:59:39.743 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 162: 0.002585723763331771, 9.723007678985596e-06
2025-05-11 14:59:39.947 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 153: 0.0015087324427440763, 9.932788088917732e-06
2025-05-11 14:59:40.058 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 83: 0.001302882912568748, 9.691226296126842e-06
2025-05-11 14:59:40.337 | INFO     | model.dineof_per_step:_fit:102 - Error/Relative Error at iteraion 218: 0.002734414068982005, 9.823590517044067e-06

imputed_dineof = imputed_data[0,:,0]
imputed_dineof_per = imputed_data_per[0,:,0]
imputed_stimp = torch.zeros(46,60,96)
imputed_stimp[:,is_sea.astype(bool)]=imputed_data_our[0,:,0]

observed = torch.zeros(46,60,96)
observed[:,is_sea.astype(bool)]=truth.cpu().squeeze()

ob_mask_image = ob_mask_image.to(device)
ob_mask_image[:,:,:,is_sea.astype(bool)]=data_ob_masks
imputed_stimp = torch.where(ob_mask_image[0,:,0].cpu().bool(), observed, imputed_stimp)
imputed_dineof = torch.where(ob_mask_image[0,:,0].cpu().bool(), observed, torch.from_numpy(imputed_dineof))
imputed_dineof_per = torch.where(ob_mask_image[0,:,0].cpu().bool(), observed, torch.from_numpy(imputed_dineof_per))

#oral observation
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import cartopy.io.shapereader as shpreader
import matplotlib.ticker as mticker
from copy import deepcopy

from matplotlib.cm import ScalarMappable
from matplotlib.colors import Normalize

t=0
is_sea =np.load("./data/PRE/is_sea.npy")
lon = np.load("./data/PRE/lon.npy")
lati = np.load("./data/PRE/lati.npy")

fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})
lon1, lon2, lati1, lati2 = lon.min(), lon.max(), lati.min(), lati.max()
axes.set_extent([lon1, lon2, lati1, lati2], crs=ccrs.PlateCarree())

tmp = deepcopy(observed[t].numpy())
tmp[~is_sea.astype(bool)]= np.nan

axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())

#STIMP
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})

tmp = deepcopy(imputed_stimp[t])
mask = np.zeros_like(tmp)
mask[is_sea.astype(bool)]=data_ob_masks[0,t,0].cpu()
tmp[~is_sea.astype(bool)]= np.nan

# 添加基础地理特征
axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())

#DINEOF_per
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})

tmp = deepcopy(imputed_dineof_per[t])
mask = np.zeros_like(tmp)
mask[is_sea.astype(bool)]=data_ob_masks[0,t,0].cpu()
tmp[~is_sea.astype(bool)]= np.nan

# 添加基础地理特征
axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())

t=13

fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})
lon1, lon2, lati1, lati2 = lon.min(), lon.max(), lati.min(), lati.max()
axes.set_extent([lon1, lon2, lati1, lati2], crs=ccrs.PlateCarree())

tmp = deepcopy(observed[t].numpy())
tmp[~is_sea.astype(bool)]= np.nan

axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())

fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})

tmp = deepcopy(imputed_stimp[t])
mask = np.zeros_like(tmp)
mask[is_sea.astype(bool)]=data_ob_masks[0,t,0].cpu()
tmp[~is_sea.astype(bool)]= np.nan

# 添加基础地理特征
axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())

# DINEOF_per
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})

axes.set_extent([lon1, lon2, lati1, lati2], crs=ccrs.PlateCarree())
tmp = deepcopy(imputed_dineof_per[t])
mask = np.zeros_like(tmp)
mask[is_sea.astype(bool)]=data_ob_masks[0,t,0].cpu()
tmp = np.where(mask, observed[t].numpy(), tmp)
tmp[~is_sea.astype(bool)]= np.nan

# 添加基础地理特征
axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())

t=33

fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})
lon1, lon2, lati1, lati2 = lon.min(), lon.max(), lati.min(), lati.max()
axes.set_extent([lon1, lon2, lati1, lati2], crs=ccrs.PlateCarree())

tmp = deepcopy(observed[t].numpy())
tmp[~is_sea.astype(bool)]= np.nan

axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())

fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})

tmp = deepcopy(imputed_stimp[t])
mask = np.zeros_like(tmp)
mask[is_sea.astype(bool)]=data_ob_masks[0,t,0].cpu()
tmp[~is_sea.astype(bool)]= np.nan

# 添加基础地理特征
axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())

# DINEOF_per
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})

axes.set_extent([lon1, lon2, lati1, lati2], crs=ccrs.PlateCarree())
tmp = deepcopy(imputed_dineof_per[t])
mask = np.zeros_like(tmp)
mask[is_sea.astype(bool)]=data_ob_masks[0,t,0].cpu()
tmp = np.where(mask, observed[t].numpy(), tmp)
tmp[~is_sea.astype(bool)]= np.nan

# 添加基础地理特征
axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())

Response Letter DINEOF(t=1) vs DINEOF(t=46)

t=42

fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})
lon1, lon2, lati1, lati2 = lon.min(), lon.max(), lati.min(), lati.max()
axes.set_extent([lon1, lon2, lati1, lati2], crs=ccrs.PlateCarree())

tmp = deepcopy(observed[t].numpy())
tmp[~is_sea.astype(bool)]= np.nan

axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())

# DINEOF
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})

axes.set_extent([lon1, lon2, lati1, lati2], crs=ccrs.PlateCarree())
tmp = deepcopy(imputed_dineof[t])
mask = np.zeros_like(tmp)
mask[is_sea.astype(bool)]=data_ob_masks[0,t,0].cpu()
tmp = np.where(mask, observed[t].numpy(), tmp)
tmp[~is_sea.astype(bool)]= np.nan

# 添加基础地理特征
axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())

#STIMP
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})

tmp = deepcopy(imputed_stimp[t])
mask = np.zeros_like(tmp)
mask[is_sea.astype(bool)]=data_ob_masks[0,t,0].cpu()
tmp[~is_sea.astype(bool)]= np.nan

# 添加基础地理特征
axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())

# DINEOF_per
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})

axes.set_extent([lon1, lon2, lati1, lati2], crs=ccrs.PlateCarree())
tmp = deepcopy(imputed_dineof_per[t])
mask = np.zeros_like(tmp)
mask[is_sea.astype(bool)]=data_ob_masks[0,t,0].cpu()
tmp = np.where(mask, observed[t].numpy(), tmp)
tmp[~is_sea.astype(bool)]= np.nan

# 添加基础地理特征
axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())

t=43

fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})
lon1, lon2, lati1, lati2 = lon.min(), lon.max(), lati.min(), lati.max()
axes.set_extent([lon1, lon2, lati1, lati2], crs=ccrs.PlateCarree())

tmp = deepcopy(observed[t].numpy())
tmp[~is_sea.astype(bool)]= np.nan

axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())

# DINEOF
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})

axes.set_extent([lon1, lon2, lati1, lati2], crs=ccrs.PlateCarree())
tmp = deepcopy(imputed_dineof[t])
mask = np.zeros_like(tmp)
mask[is_sea.astype(bool)]=data_ob_masks[0,t,0].cpu()
tmp = np.where(mask, observed[t].numpy(), tmp)
tmp[~is_sea.astype(bool)]= np.nan

# 添加基础地理特征
axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())

#STIMP
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})

tmp = deepcopy(imputed_stimp[t])
mask = np.zeros_like(tmp)
mask[is_sea.astype(bool)]=data_ob_masks[0,t,0].cpu()
tmp[~is_sea.astype(bool)]= np.nan

# 添加基础地理特征
axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())

# DINEOF_per
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})

axes.set_extent([lon1, lon2, lati1, lati2], crs=ccrs.PlateCarree())
tmp = deepcopy(imputed_dineof_per[t])
mask = np.zeros_like(tmp)
mask[is_sea.astype(bool)]=data_ob_masks[0,t,0].cpu()
tmp = np.where(mask, observed[t].numpy(), tmp)
tmp[~is_sea.astype(bool)]= np.nan

# 添加基础地理特征
axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())

t=44

fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})
lon1, lon2, lati1, lati2 = lon.min(), lon.max(), lati.min(), lati.max()
axes.set_extent([lon1, lon2, lati1, lati2], crs=ccrs.PlateCarree())

tmp = deepcopy(observed[t].numpy())
tmp[~is_sea.astype(bool)]= np.nan

axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())

# DINEOF
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})

axes.set_extent([lon1, lon2, lati1, lati2], crs=ccrs.PlateCarree())
tmp = deepcopy(imputed_dineof[t])
mask = np.zeros_like(tmp)
mask[is_sea.astype(bool)]=data_ob_masks[0,t,0].cpu()
tmp = np.where(mask, observed[t].numpy(), tmp)
tmp[~is_sea.astype(bool)]= np.nan

# 添加基础地理特征
axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())

#STIMP
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})

tmp = deepcopy(imputed_stimp[t])
mask = np.zeros_like(tmp)
mask[is_sea.astype(bool)]=data_ob_masks[0,t,0].cpu()
tmp[~is_sea.astype(bool)]= np.nan

# 添加基础地理特征
axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())

#DINEOF_per
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})

tmp = deepcopy(imputed_dineof_per[t])
mask = np.zeros_like(tmp)
mask[is_sea.astype(bool)]=data_ob_masks[0,t,0].cpu()
tmp[~is_sea.astype(bool)]= np.nan

# 添加基础地理特征
axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())

#DINEOF_per
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 5),
                         subplot_kw={'projection': ccrs.PlateCarree()})

tmp = deepcopy(imputed_dineof_per[t])
mask = np.zeros_like(tmp)
mask[is_sea.astype(bool)]=data_ob_masks[0,t,0].cpu()
tmp[~is_sea.astype(bool)]= np.nan

# 添加基础地理特征
axes.add_feature(cfeature.LAND, facecolor='#f0f0f0')
axes.add_feature(cfeature.COASTLINE.with_scale('10m'), linewidth=0.8)
axes.add_feature(cfeature.OCEAN, facecolor='#d0e0f0')
axes.add_feature(cfeature.LAKES.with_scale('10m'), facecolor='#d0e0f0')
axes.add_feature(cfeature.RIVERS.with_scale('10m'), edgecolor='#004d99')

contourf = axes.contourf(lon, lati, tmp,
                        levels=np.linspace(-1.5, 1.5, 40),cmap="rainbow",extend='both',transform=ccrs.PlateCarree())
cbar = fig.colorbar(contourf, ax=axes, ticks=np.linspace(-1.5, 1.5, 5), orientation='vertical', shrink=1)
cbar.set_label('Chl_a (Log(ug/L))')