-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_dataset.py
More file actions
151 lines (119 loc) · 5.91 KB
/
Copy pathtest_dataset.py
File metadata and controls
151 lines (119 loc) · 5.91 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
#%%
import torch
from torch.utils.data import Dataset
import numpy as np
import xarray as xr
import pickle
#%%
class FSTRDataset(Dataset):
def __init__(self, cfg, fvarLst, pvarLst, ovarLst, prevLst, fileidx, tRange, is_train=True):
super(FSTRDataset, self).__init__()
# 添加以下三行
self.fvarLst = fvarLst # 动态变量名列表
self.pvarLst = pvarLst # 静态参数名列表
self.ovarLst = ovarLst # 目标变量名列表
self.prevLst = prevLst # 新增：前序变量列表
self.dataPath = cfg.input_path
self.seq_length = cfg.frames_input # 读取模型输入序列长度（ConvLSTM的“时间帧数量”）
self.df = xr.open_dataset(cfg.input_path + 'forcings_day.nc') #强迫数据
self.dp = xr.open_dataset(cfg.input_path + 'par2000.nc') #静态参数
self.do = xr.open_dataset(cfg.input_path + 'sm2000.nc') #目标变量
self.nfile = fileidx[1]-fileidx[0] #文件总数
self.fileidx = fileidx #文件范围
self.tRange = tRange #时间范围
self.nt = len(self.do.time.loc[tRange[0]:tRange[1]]) #有效时间帧
self.is_train = is_train #训练标记
if self.is_train:
self.stat = {}
else:
scaler_file = cfg.out_dir + "train_data_scaler.bin"
with open(scaler_file, mode='rb') as fp:
self.stat = pickle.load(fp)
self.forcings = self.getDataTs(self.df, fvarLst)
self.pars = self.getDataConst(self.dp, pvarLst)
self.target = self.getDataTs(self.do, ovarLst)
self.prev = self.getDataTs(self.do, prevLst)
if self.is_train: # 训练集：保存统计量到文件（供测试集使用）
file_path = cfg.out_dir + "train_data_scaler.bin"
with open(file_path, mode='wb') as fp:
pickle.dump(self.stat, fp)
# 构建样本索引表（ConvLSTM核心适配点！）
lookup = [(i, k) for i in range(self.nfile) for k in range(self.seq_length, self.nt)]
self.lookup_table = {i: elem for i, elem in enumerate(lookup)}
def __len__(self):
return len(self.lookup_table) #样本总数
def __getitem__(self, idx):
file, indices = self.lookup_table[idx]
# ✅ 正确：按 file 取 forcing
input = np.concatenate([
self.forcings[file, indices - self.seq_length + 1 : indices + 1, :, :, :],
self.pars[file, indices - self.seq_length + 1 : indices + 1, :, :, :]
], axis=1)
output = self.target[file, indices - self.seq_length + 1 : indices + 1, :, :, :]
prevs = self.prev[file, indices - self.seq_length : indices, :, :, :]
input = torch.as_tensor(input, dtype=torch.float32)
output = torch.as_tensor(output, dtype=torch.float32)
prevs = torch.as_tensor(prevs, dtype=torch.float32)
return input, output, prevs
def getDataTs(self, ds, varLst):
# 通过检查第一个变量的维度来判断数据是否包含 file 维度
first_var = ds[varLst[0]]
has_file_dim = 'nfiles' in first_var.dims
nvar = len(varLst) #变量个数
nt = self.nt #有效步长
# 始终分配 self.nfile 大小以保持 __getitem__ 索引兼容
nfile = self.nfile
data = np.ndarray([nfile, nt, nvar, len(ds.lat), len(ds.lon)])
for k in range(nvar):
var = ds[varLst[k]]
if has_file_dim:
dataTemp = var.loc[self.fileidx[0]:self.fileidx[1],self.tRange[0]:self.tRange[1],0,:,:].values ## only use the first layer of output data
else:
dataTemp = var.loc[self.tRange[0]:self.tRange[1],:,:].values
if self.is_train: #归一化
mean = np.nanmean(dataTemp)
std = np.nanstd(dataTemp)
self.stat[varLst[k]+'_mean'] = mean
self.stat[varLst[k]+'_std'] = std
else:
mean = self.stat[varLst[k]+'_mean']
std = self.stat[varLst[k]+'_std']
dataTemp = (dataTemp-mean)/std
if has_file_dim:
data[:, :, k, :, :] = dataTemp
else:
# 气象强迫数据没有 file 维度，广播到所有文件
for f in range(nfile):
data[f, :, k, :, :] = dataTemp
return data # 返回处理后的五维时空序列数据（[nfile, nt, nvar, nlat, nlon]）
def getDataConst(self, ds, varLst):
nvar = len(varLst)
nfile = self.nfile
nt = self.nt
data = np.ndarray([nfile, nvar, len(ds.lat), len(ds.lon)])
for k in range(nvar): #特殊变量的层筛选
if varLst[k] == 'expt':
dataTemp = ds[varLst[k]].values
dataTemp = dataTemp[self.fileidx[0]:self.fileidx[1],1,:,:] # second layer
elif varLst[k] == 'd2':
dataTemp = ds['depth'].values
dataTemp = dataTemp[self.fileidx[0]:self.fileidx[1],1,:,:] # second layer
elif varLst[k] == 'd3':
dataTemp = ds['depth'].values
dataTemp = dataTemp[self.fileidx[0]:self.fileidx[1],2,:,:] # third layer
else:
dataTemp = ds[varLst[k]].values
dataTemp = dataTemp[self.fileidx[0]:self.fileidx[1],:,:]
if self.is_train:
mean = np.nanmean(dataTemp)
std = np.nanstd(dataTemp)
self.stat[varLst[k]+'_mean'] = mean
self.stat[varLst[k]+'_std'] = std
else:
mean = self.stat[varLst[k]+'_mean']
std = self.stat[varLst[k]+'_std']
dataTemp = (dataTemp-mean)/std
data[:, k, :, :] = dataTemp
out = np.repeat(np.reshape(data, [nfile, 1, nvar, len(ds.lat), len(ds.lon)]), nt, axis=1) #扩维
return out #（[nfile, nt, nvar, nlat, nlon]）
# %%