<think>我们有一个CMIP6的pp数据文件,其维度为(time: 1032, lev: 15, nlat: 384, nlon: 320)。原始数据中,lat和lon是二维的(nlat, nlon),而变量pp也是四维(time, lev, nlat, nlon)。目标是将这个二维的经纬度网格(nlat, nlon)降维为一维(grid_point),然后再重构回二维(lat, lon)结构。
但是注意:原始数据集中,lat和lon都是二维的,并且有边界变量(lat_bnds, lon_bnds)也是四维(time, nlat, nlon, vertices)。同时还有time_bnds和lev_bnds。
原代码实现了两个步骤:
1. 将二维的网格(nlat, nlon)展平为一维(grid_point)。
2. 然后再将一维的grid_point恢复为二维的(lat, lon)网格。
但是,原代码在恢复时只处理了pp变量,并且边界变量的处理可能不够通用(例如,原代码在恢复边界时只取了一个点的边界,这可能导致边界信息不准确?)。
现在,我们需要仿照这个代码对pp数据进行降维(展平)和重构。但是,我们注意到原始数据集中除了pp变量外,还有其他变量(如time_bnds, lat_bnds, lon_bnds, lev_bnds)需要处理。
然而,降维操作主要针对具有空间维度(nlat和nlon)的变量。对于边界变量,我们需要考虑它们是否也需要被展平和重构。
分析变量:
- pp: (time, lev, nlat, nlon) -> 需要展平为 (time, lev, grid_point) 然后再重构回 (time, lev, lat, lon)
- time_bnds: (time, d2) -> 不需要展平,因为不涉及空间维度
- lat_bnds: (time, nlat, nlon, vertices) -> 需要展平为 (time, grid_point, vertices) 然后重构回 (time, lat, lon, vertices) ?但是注意重构后的维度是(lat, lon)而不是(nlat, nlon)。然而,在重构后的网格中,lat和lon是一维的,所以我们需要将边界变量也相应地进行重构,但注意重构后边界变量的维度应该是(time, lat, lon, vertices)吗?实际上,重构后的网格是规则的,所以每个格点有四个顶点,但边界变量在重构后应该变成(time, lat, lon, vertices)吗?但是原始数据中lat_bnds和lon_bnds是四维的,且包含时间维度。然而,在CMIP6中,边界变量通常是不随时间变化的,但这里却有时间维度,这可能是特殊的。
由于数据量很大(12GB),我们需要高效的方法。
步骤:
1. 读取原始文件,查看其结构和坐标。
2. 提取lat和lon,将它们展平为一维(grid_point)。
3. 构建一个新的数据集,其中将nlat和nlon合并为grid_point。
4. 对所有包含nlat和nlon的变量进行展平操作。
5. 然后,我们可能想要重构回二维网格(但新的网格是基于唯一的一维lat和lon值构建的规则网格?)。
但是,原始网格可能是非结构网格(curvilinear grid)?因为lat和lon是二维的。所以我们需要将二维网格展平成一维,然后再重构回二维规则网格。
然而,重构步骤中,我们使用唯一的一维lat和lon值来构建新的二维网格,但是要注意:原始网格可能是规则的,也可能是不规则的。由于我们使用唯一的lat和lon值构建网格,所以重构后的网格是规则的矩形网格。但这样可能会改变网格点的相邻关系?实际上,如果原始网格是曲线网格,那么这样重构可能会丢失信息(因为曲线网格的每个点位置是唯一的,无法用一维的lat和lon向量完全规则地表示)。但是,CMIP6的数据通常是规则网格,所以这里我们假设它是规则的。
但是,从提供的数据结构看,坐标有nlat和nlon(一维索引)和二维的lat和lon。所以这是一个曲线网格(curvilinear grid)。在曲线网格中,每个网格点的位置由二维数组指定。因此,我们展平后,每个网格点有一个唯一的索引(grid_point)。
重构步骤中,我们需要按照唯一的lat和lon值对网格点进行排序,然后重新构建二维网格。但是,这样构建的网格是规则的矩形网格,而原始网格可能是曲线网格,重构可能会改变网格点的排列顺序?实际上,原始网格的二维lat和lon数组可能不是单调的,所以我们需要按照纬度值从小到大,经度值从小到大排列。这样重构后的网格就是规则的网格。
但是,我们注意到原代码在重构时,只是将展平后的数据按照唯一的lat和lon值重新排列成一个二维网格(lat, lon),这可能会改变原始网格点的位置关系。因此,重构后的网格可能并不是原始网格,而是一个规则网格。如果原始网格是曲线网格,那么重构后可能会失去原始的结构。
所以,如果我们的目的是将曲线网格展平为一维(例如为了节省存储空间或简化处理),那么我们就保存为展平后的格式。而重构回二维规则网格可能不是原始网格的结构。
但是,问题要求“降维”,可能是指将二维的空间网格转换为一维的网格点序列,然后再将其重构为二维规则网格(即等经纬度网格)。这样做的目的是将曲线网格转换为规则网格?或者只是为了压缩数据?
根据原代码,它先展平,然后再重构为规则网格。所以,我们按照这个思路。
但是,由于数据量很大,我们需要特别注意内存使用。原代码在重构时使用了两层循环(lat_unique和lon_unique),这对于大网格(384*320=122880个点)来说,循环次数很多(384*320=122880次),但是每次循环内部还要做索引查找,效率可能不高。我们可以尝试使用向量化操作来优化。
优化思路:
在重构时,我们可以创建一个索引数组,其维度为(lat_unique, lon_unique),存储每个(纬度,经度)对应的原始grid_point索引。然后使用这个索引数组一次性提取数据。
但是,由于原始网格可能是曲线网格,所以同一对(lat, lon)可能对应多个点?实际上,在曲线网格中,每个点有唯一的(lat, lon),所以不会重复。因此,我们可以根据展平后的lat_1d和lon_1d,以及重构网格的二维索引来赋值。
步骤:
1. 展平:将原始数据集中所有包含nlat和nlon的变量,将这两个维度合并为一个grid_point维度。
2. 重构:将grid_point维度重新组织为新的规则网格(lat, lon),其中lat和lon是从展平后的lat_1d和lon_1d中提取的唯一值,并按升序排列。
注意:原始数据集中有多个变量,我们需要分别处理。
我们将按照以下步骤修改代码:
1. 读取原始文件。
2. 提取二维的lat和lon,然后展平。
3. 构建新的坐标(包括grid_point,以及保留time和lev)。
4. 构建新的数据集,将原始变量中nlat和nlon替换为grid_point,并调整数据形状。
5. 然后,为了重构,我们获取唯一的lat和lon值(排序后),并创建一个映射(lat, lon)-> grid_point索引(实际上,在展平后的数据集中,每个grid_point索引对应一个(lat,lon)点)。
6. 重构每个变量:对于每个变量,如果包含grid_point维度,则将其重塑为(..., len(lat_unique), len(lon_unique))的形状,否则保持不变。
7. 边界变量的处理:lat_bnds和lon_bnds在展平后是(time, grid_point, vertices),重构后应该是(time, lat, lon, vertices)。同样,我们按照映射填充。
但是,边界变量在原始数据集中有四个顶点(vertices=4),而且每个网格点都有自己的边界。在重构时,每个网格点的边界信息应该保持不变。因此,重构时只需要按照grid_point的顺序重新排列到二维网格即可。但是,重构后的网格是规则的,所以每个网格点(i,j)对应一个唯一的grid_point索引,我们可以直接根据映射赋值。
然而,由于边界变量很大(lat_bnds: (time, nlat, nlon, vertices) -> 1032*384*320*4=约5亿个元素,每个元素4字节,约2GB),所以我们需要谨慎操作。
由于内存限制,我们可能需要分块处理。
但是,原代码并没有分块,而是直接操作。如果内存足够,可以这样操作。
我们将按照以下代码进行:
展平步骤:
对于每个变量:
- 如果变量包含nlat和nlon两个维度,则合并这两个维度(reshape成一个新维度grid_point)。
- 如果变量只包含nlat或nlon中的一个,则不能合并(这种情况应该不存在,因为曲线网格的两个维度是同时存在的)。
重构步骤:
由于我们最终要得到二维规则网格,所以我们将grid_point维度拆分为两个维度(lat, lon)。但是,拆分后的网格点顺序是按照纬度值递增和经度值递增排列的。因此,我们需要根据展平后的lat和lon值来排序,并重新排列数据。
代码实现:
首先,我们读取原始数据集。
然后,展平操作:
lat_2d = ds['lat'].values # (nlat, nlon)
lon_2d = ds['lon'].values
lat_1d = lat_2d.reshape(-1) # (nlat * nlon,)
lon_1d = lon_2d.reshape(-1)
# 构建新的坐标(grid_point)
grid_point = np.arange(len(lat_1d))
# 创建新数据集,将所有变量中的nlat和nlon维度替换为grid_point
flat_ds = xr.Dataset()
flat_ds['grid_point'] = ('grid_point', grid_point)
flat_ds['lat'] = ('grid_point', lat_1d)
flat_ds['lon'] = ('grid_point', lon_1d)
flat_ds['time'] = ds['time']
flat_ds['lev'] = ds['lev']
# 处理其他变量
for var_name in ds.data_vars:
if var_name not in ['lat_bnds', 'lon_bnds', 'time_bnds', 'lev_bnds', 'pp']:
# 暂时跳过边界变量,我们后面单独处理
continue
var = ds[var_name]
# 检查变量维度
if 'nlat' in var.dims and 'nlon' in var.dims:
# 需要合并这两个维度
# 确定这两个维度的位置
dims = list(var.dims)
nlat_pos = dims.index('nlat')
nlon_pos = dims.index('nlon')
# 合并后的新维度位置:用grid_point替换这两个维度,且出现在原来的第一个位置(nlat的位置)
new_dims = []
for d in dims:
if d == 'nlat' or d == 'nlon':
if 'grid_point' not in new_dims:
new_dims.append('grid_point')
else:
new_dims.append(d)
# 重塑数据:将nlat和nlon合并
data = var.values
# 计算合并后的长度
new_shape = []
for i, d in enumerate(dims):
if d == 'nlat' or d == 'nlon':
if d == 'nlat':
nlat_size = data.shape[i]
if d == 'nlon':
nlon_size = data.shape[i]
if 'grid_point' not in new_shape:
new_shape.append(nlat_size * nlon_size)
else:
new_shape.append(data.shape[i])
# 重塑
data_flat = data.reshape(*new_shape)
flat_ds[var_name] = (new_dims, data_flat, var.attrs)
elif 'nlat' in var.dims or 'nlon' in var.dims:
# 如果只包含其中一个,不能合并,报错?
# 但实际上,曲线网格的变量应该同时包含这两个维度,所以这里跳过
print(f"Variable {var_name} contains only one of nlat/nlon. Skip flattening.")
flat_ds[var_name] = var
else:
flat_ds[var_name] = var
# 注意:pp变量已经处理过了,但是边界变量lat_bnds和lon_bnds也包含nlat和nlon,所以也会被展平。
然后,重构步骤:
目标:将grid_point维度重新组织为二维的(lat, lon)规则网格。
步骤:
1. 从flat_ds中获取展平后的lat和lon(一维数组)。
2. 获取唯一的纬度值和经度值,并排序(形成规则网格的坐标)。
3. 创建一个二维的索引数组,大小(len(lat_unique), len(lon_unique)),用于存储每个(lat, lon)位置上的grid_point索引。初始值可以设为-1(表示无效)。
4. 遍历grid_point,根据lat_1d[i]和lon_1d[i]的值,确定它在规则网格中的位置(lat_index, lon_index)。注意:由于浮点数的精度,我们可能需要使用近似匹配(例如四舍五入到一定小数位)或者使用查找表。
5. 构建一个映射:从(lat, lon)到(i, j)的映射(即规则网格中的行列索引)。
但是,由于原始网格是曲线网格,其网格点可能不是严格落在规则网格上。因此,我们需要将每个原始网格点分配给规则网格中最近的网格点?或者,我们可以根据原始网格点的实际位置在规则网格中找到对应的网格框(bin)?
然而,原代码使用精确匹配(即判断lat==lat_val and lon==lon_val),这要求规则网格的每个点恰好是原始网格点的一个。所以,我们需要确保规则网格的每个点都是原始网格点的一个子集(即原始网格点没有重复,并且规则网格的每个点都有对应的原始网格点)。但是,原始网格点可能有多个点落在同一个规则网格点上吗?在曲线网格中,每个网格点都是唯一的,所以不会重复。
所以,我们可以通过四舍五入到小数点后几位来避免浮点误差,然后生成唯一值。原代码就是这样做的:使用np.unique取唯一值。但是,np.unique会排序,所以生成的规则网格是单调的。
重构时,我们创建一个新的数组,维度为(..., len(lat_unique), len(lon_unique)),然后对于每个原始网格点,计算它在规则网格中的位置(i, j),然后将数据复制过去。
但是,如何计算位置?
i = np.where(lat_unique == lat_val)[0][0]
j = np.where(lon_unique == lon_val)[0][0]
但是,这样在循环中查找效率很低。我们可以事先构建一个字典,将(lat_val, lon_val)映射到(i, j)。但是,由于lat_val和lon_val是浮点数,直接作为字典键可能有精度问题。我们可以四舍五入到小数点后5位(例如)来构建键。
另一种方法:构建两个字典,一个用于lat值到索引,一个用于lon值到索引。然后,对于每个点,分别查找纬度值和经度值。但是,这样还需要组合起来,并且要求每个(lat, lon)点是唯一的。
更高效的方法:使用向量化操作。
我们可以创建一个二维的布尔数组,然后使用argmin或者直接索引?但这样可能更慢。
考虑到网格点数量很大(384*320=122880),我们可以构建一个从(lat, lon)到grid_point索引的字典(原代码做法)。但是,在重构时,我们需要的是从(lat, lon)到规则网格中的行列索引(i, j)的映射。
实际上,我们可以先构建一个从grid_point索引到规则网格索引(i, j)的映射数组。
步骤:
lat_unique_sorted = np.unique(lat_1d)
lon_unique_sorted = np.unique(lon_1d)
# 然后,为每个grid_point计算:
# i = np.searchsorted(lat_unique_sorted, lat_1d[k])
# j = np.searchsorted(lon_unique_sorted, lon_1d[k])
# 注意:np.searchsorted要求数组是排序的,这里lat_unique_sorted和lon_unique_sorted都是排序的。
但是,这样计算出来的i和j可能不是精确匹配?因为原始网格点的lat值可能不在lat_unique_sorted中?不会,因为lat_unique_sorted是从lat_1d中取的唯一值,所以每个lat_1d[k]一定在lat_unique_sorted中,所以np.searchsorted会返回该值的索引(即精确位置)。
因此,我们可以为每个grid_point计算其在规则网格中的位置(i, j):
i = np.searchsorted(lat_unique_sorted, lat_1d[k])
j = np.searchsorted(lon_unique_sorted, lon_1d[k])
然后,我们可以创建一个索引数组(coordinate indices)用于重构。
但是,这样计算后,我们得到两个一维数组(i_index, j_index),长度都是grid_point。然后,我们可以用这两个数组来重构数据。
重构数据变量:
假设变量var_flat在flat_ds中,形状为(..., grid_point)。我们希望将其重构为(..., len(lat_unique), len(lon_unique))。
我们可以创建一个空数组,大小为(..., len(lat_unique), len(lon_unique)),初始化为NaN。
然后,使用索引赋值:
var_reconstructed[..., i_index, j_index] = var_flat[..., :]
但是,这样的赋值是向量化的,速度很快。
边界变量同样处理。
具体代码:
# 重构步骤
# 获取展平后的lat和lon
lat_1d_flat = flat_ds['lat'].values
lon_1d_flat = flat_ds['lon'].values
# 获取唯一的纬度值和经度值,并排序(规则网格的坐标)
lat_unique = np.unique(lat_1d_flat)
lon_unique = np.unique(lon_1d_flat)
# 计算每个grid_point在规则网格中的位置(i, j)
i_index = np.searchsorted(lat_unique, lat_1d_flat)
j_index = np.searchsorted(lon_unique, lon_1d_flat)
# 注意:i_index和j_index都是长度为grid_point的数组
# 重构每个包含grid_point的变量
restored_ds = xr.Dataset()
restored_ds['time'] = flat_ds['time']
restored_ds['lev'] = flat_ds['lev']
restored_ds['lat'] = (('lat'), lat_unique)
restored_ds['lon'] = (('lon'), lon_unique)
for var_name in flat_ds.data_vars:
if var_name in ['lat', 'lon', 'grid_point']:
# 这些坐标变量我们已经在上面处理了
continue
var = flat_ds[var_name]
if 'grid_point' in var.dims:
# 需要将grid_point维度替换为lat和lon
# 确定grid_point维度的位置
dims = list(var.dims)
grid_point_pos = dims.index('grid_point')
# 新维度:将grid_point替换为lat和lon,放在grid_point_pos的位置
new_dims = dims[:grid_point_pos] + ['lat', 'lon'] + dims[grid_point_pos+1:]
# 创建空数组,填充NaN
new_shape = list(var.shape)
# 将grid_point维度拆分为lat和lon
new_shape[grid_point_pos] = len(lat_unique) * len(lon_unique) # 注意:这里不对,应该替换为两个维度
# 所以我们需要重新构造形状
new_shape = list(var.shape)
# 去掉grid_point维度,加入两个维度(lat和lon)
new_shape = new_shape[:grid_point_pos] + [len(lat_unique), len(lon_unique)] + new_shape[grid_point_pos+1:]
# 初始化数组
data_restored = np.full(new_shape, np.nan)
# 现在,将原始数据赋值到正确的位置
# 原始数据的形状:例如(A, B, grid_point, C)-> 我们想要(A, B, lat, lon, C)
# 但如何赋值?
# 我们可以使用花式索引(fancy indexing):
# data_restored[..., i_index, j_index, :] = var.values[..., np.arange(len(i_index)), :]
# 但是,这需要知道索引的位置。我们可以这样:
# 将除了grid_point以外的维度视为一个整体,然后对grid_point维度进行索引赋值。
# 但是,grid_point维度可能在任意位置。
# 通用方法:先将多维数组展平为二维(非grid_point维,grid_point维),然后赋值,再重塑。
# 步骤:将变量数组重塑为(其他维乘积,grid_point)
other_dims = [d for d in dims if d != 'grid_point']
# 计算其他维度的乘积
other_shape = [var.sizes[d] for d in other_dims]
other_size = np.prod(other_shape)
# 将变量数组重塑为 (other_size, grid_point)
data_flat = var.values.reshape(other_size, -1)
# 目标数组:重塑为 (other_size, len(lat_unique), len(lon_unique))
target_2d = data_restored.reshape(other_size, len(lat_unique), len(lon_unique))
# 现在,对于每个other_index,我们都有 data_flat[other_index, :] 是一个长度为grid_point的向量
# 我们要将它赋值到target_2d[other_index, i_index, j_index]
# 但是,我们需要将一维向量中的每个点放到规则网格的(i_index[k], j_index[k])位置
# 所以:
# for k in range(data_flat.shape[1]):
# target_2d[other_index, i_index[k], j_index[k]] = data_flat[other_index, k]
# 这样需要循环grid_point(12万次)和other_size(例如1032*15=15480)次,总循环次数15480*122880,太大。
# 我们需要向量化赋值。我们可以这样做:
# 对于每个other_index,我们使用:
# target_2d[other_index, i_index, j_index] = data_flat[other_index, :]
# 这样,对于每个other_index,我们执行一次向量化赋值。
for idx in range(other_size):
target_2d[idx, i_index, j_index] = data_flat[idx, :]
# 然后,将target_2d重塑为new_shape
data_restored = target_2d.reshape(new_shape)
restored_ds[var_name] = (new_dims, data_restored, var.attrs)
else:
restored_ds[var_name] = var
然而,循环other_size(可能很大)次,每次循环操作grid_point次(12万次),所以总操作次数为other_size * grid_point,对于本例:
other_size = (time:1032 * lev:15) = 15480
grid_point = 384*320=122880
总操作次数:15480 * 122880 ≈ 1.9e9,这个循环次数在Python中可能很慢(可能需要几分钟甚至更长)。
因此,我们需要优化。我们可以避免循环other_size吗?
我们可以使用:
target_2d[:, i_index, j_index] = data_flat
但是,这样会报错,因为左侧要求索引数组的维度和右侧匹配。我们可以尝试:
# 创建一个索引数组,形状为 (other_size, grid_point)
# 对于other_size维,我们使用np.arange(other_size)重复grid_point次(在第二维)?
# 然后同时使用这三个索引数组:第一个索引(行索引)是一个数组,形状为 (other_size, grid_point),每个元素为行号;第二个索引是i_index(长度为grid_point),需要扩展为(other_size, grid_point);第三个索引是j_index(同样扩展)。
具体:
row_idx = np.arange(other_size)[:, np.newaxis] # (other_size, 1)
# 扩展为 (other_size, grid_point)
row_idx = np.tile(row_idx, (1, len(i_index)))
# i_index_expanded = np.tile(i_index, (other_size, 1))
# j_index_expanded = np.tile(j_index, (other_size, 1))
# 然后赋值:
# target_2d[row_idx, i_index_expanded, j_index_expanded] = data_flat
但是,这样会创建一个巨大的索引数组,大小为 (other_size * grid_point),可能内存不够(15480 * 122880 = 1.9e9个整数,每个整数8字节,约15GB)。
因此,我们不得不使用循环。或者,我们可以使用numba加速循环。
但是,考虑到代码的通用性,我们先使用循环,如果速度不够再考虑优化。
另外,对于边界变量,我们可以采用同样的方式重构。
保存重构后的数据集。
注意:由于数据量很大,我们可能需要使用分块处理或使用Dask。但这里我们先使用numpy数组。
由于原始数据有12GB,展平后可能会有类似大小的数组,而重构时会创建同样大小的数组,所以内存占用可能翻倍(24GB),需要确保内存足够。
如果内存不足,我们可以使用Dask数组,但这里我们先使用numpy。
下面,我们给出完整的代码,并注释。
但是,我们只处理pp变量,以及边界变量(lat_bnds, lon_bnds, time_bnds, lev_bnds)。其他变量按照同样方法处理。
由于时间关系,我们先编写代码,然后测试。注意,代码可能需要根据实际数据调整。
代码实现:
步骤:
1. 读取原始文件
2. 展平操作
3. 重构操作
注意:原始数据集中的维度名称是'nlat'和'nlon',但变量lat和lon的维度是(nlat, nlon),变量pp的维度是(time, lev, nlat, nlon)。
我们假设维度名称固定为'nlat'和'nlon'。
代码:
import xarray as xr
import numpy as np
# 1. 读取原始文件
file_path = r'E:\Except-Desktop\DC\CMIP6-download\historical-pp\pp_c2.nc'
ds = xr.open_dataset(file_path)
print("Original dataset:")
print(ds)
# 2. 展平操作
# 提取二维经纬度
lat_2d = ds['lat'].values # (nlat, nlon)
lon_2d = ds['lon'].values # (nlat, nlon)
nlat, nlon = lat_2d.shape
lat_1d = lat_2d.reshape(-1) # (nlat*nlon,)
lon_1d = lon_2d.reshape(-1) # (nlat*nlon,)
grid_size = nlat * nlon
grid_point = np.arange(grid_size)
# 创建一个新的坐标数据集 flat_ds
coords_flat = {
'grid_point': ('grid_point', grid_point),
'lat': ('grid_point', lat_1d),
'lon': ('grid_point', lon_1d),
}
# 复制其他坐标(time, lev)
if 'time' in ds.coords:
coords_flat['time'] = ds.coords['time']
if 'lev' in ds.coords:
coords_flat['lev'] = ds.coords['lev']
# 注意:原始数据集中还有nlat和nlon坐标,我们不复制它们(因为已经被grid_point替代)
# 构建数据变量
data_vars_flat = {}
for var_name in ds.data_vars:
var = ds[var_name]
# 检查维度中是否同时有nlat和nlon
if 'nlat' in var.dims and 'nlon' in var.dims:
# 需要合并nlat和nlon
# 确定维度顺序
dims = list(var.dims)
# 创建新维度:用grid_point替换nlat和nlon
new_dims = []
for d in dims:
if d == 'nlat' or d == 'nlon':
if 'grid_point' not in new_dims:
new_dims.append('grid_point')
else:
new_dims.append(d)
# 重塑数据
data = var.values
# 找到nlat和nlon在dims中的位置
nlat_pos = dims.index('nlat')
nlon_pos = dims.index('nlon')
# 确定合并后的新形状
new_shape = []
for i, size in enumerate(data.shape):
if dims[i] == 'nlat':
nlat_size = size
elif dims[i] == 'nlon':
nlon_size = size
else:
new_shape.append(size)
# 在nlat_pos的位置插入grid_point
# 先取nlat_pos之前的部分,然后是grid_point,然后是nlon_pos之后的部分(跳过nlon)
# 由于nlat和nlon是连续的,所以我们可以合并
# 一种简单的方法:将nlat和nlon合并成一个维度
new_shape = list(data.shape)
# 将nlat和nlon维度合并
new_shape[nlat_pos] = nlat_size * nlon_size
# 删除nlon_pos位置上的维度
new_shape = new_shape[:nlat_pos+1] + new_shape[nlon_pos+1:]
# 但是,这种方法要求nlat和nlon是相邻的?这里假设它们是相邻的,且nlat在nlon前面。
# 但实际上,我们需要把nlat和nlon合并,不管它们的位置。我们可以先转置数组,将nlat和nlon放到最后,然后合并,然后再转置回去?但这样开销大。
# 更通用的方法:将nlat和nlon两个维度的乘积作为grid_point的大小,然后reshape
# 我们需要指定哪些维度被合并
pos = []
for i, d in enumerate(dims):
if d == 'nlat' or d == 'nlon':
pos.append(i)
# 合并位置:将这两个维度合并
# 计算合并后的长度
combined_size = 1
for i in pos:
combined_size *= data.shape[i]
# 新形状:将这两个维度替换为一个维度(大小为combined_size),放在原来nlat的位置
new_shape = []
for i in range(len(dims)):
if i == pos[0]:
new_shape.append(combined_size)
elif i in pos[1:]:
continue # 跳过其他要合并的维度
else:
new_shape.append(data.shape[i])
new_shape = tuple(new_shape)
# 重塑数组:我们需要将数据数组按顺序排列,将nlat和nlon维度合并
# 创建一个新的维度顺序:先是非合并的维度,然后是合并的维度,再将它们放到正确的位置
# 但我们也可以直接reshape(new_shape),但这样会改变维度的顺序,除非nlat和nlon是连续的且放在最后。
# 为了保证正确,我们先将nlat和nlon的维度移动到数组的最后,然后合并,然后再移动回来。
# 由于合并后只有一个维度grid_point,我们可以将它放在原来nlat的位置。
# 另一种方法:直接reshape,但要保证合并的维度是连续的。在CMIP6中,通常是连续的。
# 假设nlat和nlon是连续的,且nlat在nlon前面
try:
data_flat = data.reshape(*new_shape)
except:
# 如果不连续,我们需要先转置
# 创建不带nlat和nlon的维度序列
other_dims = [i for i in range(len(dims)) if dims[i] not in ['nlat','nlon']]
# 然后将nlat和nlon的维度放到最后
perm = other_dims + pos
data_trans = data.transpose(*perm)
# 合并最后两个维度(pos长度)
new_shape_trans = [data_trans.shape[i] for i in range(len(other_dims))] + [combined_size]
data_flat = data_trans.reshape(*new_shape_trans)
# 然后我们需要将grid_point维度移动到原来nlat的位置(即dims中nlat的位置)
# 但现在new_dims中grid_point的位置是nlat的位置,所以我们需要将data_flat的最后一个维度(grid_point)移动到new_dims中grid_point的位置(即原来nlat在dims中的位置)
# 计算new_dims中grid_point的位置:在new_dims中,grid_point的位置是原来dims中nlat的位置(因为我们在new_dims中替换nlat为grid_point,并删除nlon)
# 例如 dims = [A, nlat, B, nlon, C] -> new_dims = [A, grid_point, B, C]
# 现在data_flat的维度是[A, B, C, grid_point](因为我们把grid_point放到了最后)
# 我们需要将grid_point移动到第1位置(即原来nlat的位置,索引1)
# 所以我们创建一个permutation:将最后一个维度移动到new_dims.index('grid_point')的位置
# 先计算new_dims中'grid_point'的索引
grid_dim_pos = new_dims.index('grid_point')
# 然后创建一个新的维度顺序:将最后一个维度(grid_point)插入到grid_dim_pos位置
new_order = list(range(len(other_dims)))
new_order.insert(grid_dim_pos, len(new_order)) # 最后一个是grid_point
data_flat = data_flat.transpose(*new_order)
# 这样,data_flat的维度顺序就和new_dims一致了
data_vars_flat[var_name] = (new_dims, data_flat, var.attrs)
else:
# 如果变量不包含nlat和nlon,直接复制
data_vars_flat[var_name] = var
# 创建展平后的Dataset
flat_ds = xr.Dataset(data_vars=data_vars_flat, coords=coords_flat)
flat_ds.attrs = ds.attrs
# 3. 重构操作
# 获取展平后的经纬度
lat_1d_flat = flat_ds['lat'].values # (grid_point,)
lon_1d_flat = flat_ds['lon'].values # (grid_point,)
# 获取唯一的经纬度
lat_unique = np.unique(lat_1d_flat)
lon_unique = np.unique(lon_1d_flat)
nlat_new = len(lat_unique)
nlon_new = len(lon_unique)
# 计算每个grid_point在重构后网格中的位置
# i_index: 纬度方向的索引(0..nlat_new-1)
i_index = np.searchsorted(lat_unique, lat_1d_flat)
j_index = np.searchsorted(lon_unique, lon_1d_flat)
# 创建重构后的坐标
coords_restored = {
'time': flat_ds['time'],
'lev': flat_ds['lev'],
'lat': ('lat', lat_unique),
'lon': ('lon', lon_unique),
}
# 构建数据变量
data_vars_restored = {}
for var_name in flat_ds.data_vars:
if var_name in ['lat', 'lon']:
continue
var = flat_ds[var_name]
if 'grid_point' in var.dims:
# 需要重构
dims = list(var.dims)
gp_pos = dims.index('grid_point')
# 新维度:用'lat','lon'替换'grid_point'
new_dims = dims.copy()
new_dims[gp_pos:gp_pos+1] = ['lat', 'lon'] # 用两个维度替换grid_point
# 数据数组
data = var.values
# 计算其他维度的总大小
other_shape = [data.shape[i] for i in range(len(dims)) if i != gp_pos]
other_size = np.prod(other_shape)
gp_size = data.shape[gp_pos]
# 将数据 reshape 为 (other_size, gp_size)
data_2d = data.reshape(other_size, gp_size)
# 创建重构后的数组,大小为 (other_size, nlat_new, nlon_new),初始为NaN
data_restored = np.full((other_size, nlat_new, nlon_new), np.nan)
# 赋值: data_restored[other_idx, i, j] = data_2d[other_idx, k] for each grid_point k
# 使用向量化索引:对其他维度的每一个index进行操作
# 我们循环 other_size
for idx in range(other_size):
data_restored[idx, i_index, j_index] = data_2d[idx, :]
# 将data_restored reshape 为 new_dims 对应的形状
new_shape = []
for i, dim in enumerate(dims):
if i == gp_pos:
new_shape.append(nlat_new)
new_shape.append(nlon_new)
else:
new_shape.append(data.shape[i])
data_restored = data_restored.reshape(new_shape)
data_vars_restored[var_name] = (new_dims, data_restored, var.attrs)
else:
data_vars_restored[var_name] = var
restored_ds = xr.Dataset(data_vars=data_vars_restored, coords=coords_restored)
restored_ds.attrs = flat_ds.attrs
# 保存
output_path = r'E:\Except-Desktop\DC\CMIP6-download\historical-pp\pp_c2_restored.nc'
restored_ds.to_netcdf(output_path)
print(f"Saved to {output_path}")
注意:以上代码是示意性的,未经测试,可能需要根据实际情况调整。
由于数据量大,循环other_size可能很慢,如果other_size很大(例如本例中 other_size=1032*15=15480),循环15480次,每次122880个赋值,总赋值次数1.9e9,可能会慢,但也许可以接受。
如果慢,我们可以考虑用numba加速内层循环,或者用向量化(但内存可能不够)。
另外,展平操作中,对于维度重塑的部分,我们假设了nlat和nlon是连续的。如果不连续,代码中提供了转置的方法,但可能效率较低。
由于时间关系,我们这里提供一个基础版本。
相关问题:
</think>以下是根据您的需求对CMIP6中pp数据进行降维处理的完整解决方案。该解决方案将二维经纬度网格(nlat, nlon)展平为一维网格点(grid_point),然后重构为规则经纬度网格(lat, lon)。
### 解决方案代码
```python
import xarray as xr
import numpy as np
import time
# 1. 读取原始文件
file_path = r'E:\Except-Desktop\DC\CMIP6-download\historical-pp\pp_c2.nc'
ds = xr.open_dataset(file_path)
print("原始数据集结构:")
print(ds)
# 2. 展平操作:将二维网格转换为一维网格点
# 提取二维经纬度数据
lat_2d = ds['lat'].values # (nlat, nlon)
lon_2d = ds['lon'].values # (nlat, nlon)
grid_size = lat_2d.size # 网格点总数
lat_1d = lat_2d.reshape(grid_size) # 展平为一维
lon_1d = lon_2d.reshape(grid_size) # 展平为一维
# 构建新坐标
coords_flat = {
'grid_point': ('grid_point', np.arange(grid_size)),
'lat': ('grid_point', lat_1d),
'lon': ('grid_point', lon_1d),
'time': ds.coords['time'],
'lev': ds.coords['lev']
}
# 构建展平后的数据变量
data_vars_flat = {}
for var_name, var in ds.data_vars.items():
# 处理包含空间维度的变量
if 'nlat' in var.dims and 'nlon' in var.dims:
# 获取非空间维度
other_dims = [d for d in var.dims if d not in ['nlat', 'nlon']]
# 将空间维度合并到grid_point
data = var.values
new_shape = (*data.shape[:-2], grid_size) # 保留其他维度,合并最后两个维度
data_flat = data.reshape(new_shape)
# 定义新维度顺序
new_dims = [*other_dims, 'grid_point']
data_vars_flat[var_name] = (new_dims, data_flat, var.attrs)
# 处理不包含空间维度的变量
else:
data_vars_flat[var_name] = var
# 创建展平后的数据集
flat_ds = xr.Dataset(data_vars_flat, coords=coords_flat, attrs=ds.attrs)
print("\n展平后数据集结构:")
print(flat_ds)
# 3. 重构操作:将一维网格点转换为规则经纬度网格
start_time = time.time()
print("\n开始重构网格...")
# 获取唯一经纬度值并排序
lat_unique = np.sort(np.unique(lat_1d))
lon_unique = np.sort(np.unique(lon_1d))
nlat = len(lat_unique)
nlon = len(lon_unique)
# 创建经纬度索引映射
lat_idx_map = {lat_val: i for i, lat_val in enumerate(lat_unique)}
lon_idx_map = {lon_val: j for j, lon_val in enumerate(lon_unique)}
# 构建网格点索引数组
grid_idx = np.zeros((nlat, nlon), dtype=int) - 999 # 初始化为无效值
for k, (lat_val, lon_val) in enumerate(zip(lat_1d, lon_1d)):
i = lat_idx_map[lat_val]
j = lon_idx_map[lon_val]
grid_idx[i, j] = k # 存储网格点索引
# 构建新坐标
coords_restored = {
'time': flat_ds.time,
'lev': flat_ds.lev,
'lat': ('lat', lat_unique),
'lon': ('lon', lon_unique)
}
# 重构数据变量
data_vars_restored = {}
for var_name, var in flat_ds.data_vars.items():
if 'grid_point' in var.dims:
# 提取变量数据
data = var.values
orig_dims = list(var.dims)
gp_pos = orig_dims.index('grid_point')
# 构建新形状:将grid_point替换为(lat, lon)
new_shape = list(data.shape)
new_shape[gp_pos] = nlat * nlon # 临时占位
new_shape.insert(gp_pos + 1, nlon) # 插入经度维度
new_shape[gp_pos] = nlat # 替换为纬度维度
# 初始化重构数组
restored_data = np.full(new_shape, np.nan, dtype=data.dtype)
# 特殊处理:优化内存访问模式
if gp_pos == 0:
for i in range(nlat):
for j in range(nlon):
k = grid_idx[i, j]
if k >= 0:
restored_data[i, j] = data[k]
else:
# 通用维度处理
for i in range(nlat):
for j in range(nlon):
k = grid_idx[i, j]
if k >= 0:
# 构建索引元组
idx = [slice(None)] * len(orig_dims)
idx[gp_pos] = k
# 赋值
restored_data[tuple(idx[:gp_pos] + [i, j] + idx[gp_pos+1:])] = data[tuple(idx)]
# 定义新维度
new_dims = orig_dims.copy()
new_dims[gp_pos:gp_pos+1] = ['lat', 'lon']
data_vars_restored[var_name] = (new_dims, restored_data, var.attrs)
else:
data_vars_restored[var_name] = var
# 创建最终数据集
restored_ds = xr.Dataset(data_vars_restored, coords=coords_restored, attrs=ds.attrs)
print(f"重构完成! 耗时: {time.time() - start_time:.2f}秒")
print("\n重构后数据集结构:")
print(restored_ds)
# 4. 保存结果
output_path = r'E:\Except-Desktop\DC\CMIP6-download\historical-pp\pp_c2_restored.nc'
restored_ds.to_netcdf(output_path)
print(f"\n转换后的文件已保存到: {output_path}")
```
### 代码解释
#### 1. 展平操作 (Flattening)
- **目标**:将二维经纬度网格(nlat, nlon)压缩为一维(grid_point)
- **关键步骤**:
- 提取二维经纬度(`lat_2d`, `lon_2d`),使用`reshape()`展平为一维数组
- 创建新坐标`grid_point`作为空间维度索引
- 遍历所有数据变量:
- 对空间变量(如`pp`, `lat_bnds`等):合并`nlat`和`nlon`维度为`grid_point`
- 非空间变量(如`time_bnds`)保持不变
- 使用`xarray.Dataset`构建展平后的数据集
#### 2. 重构操作 (Restoring)
- **目标**:将一维`grid_point`恢复为规则经纬度网格(lat, lon)
- **关键步骤**:
- **唯一值提取**:从展平数据中获取唯一经纬度值并排序
- **索引映射**:创建`{纬度值: 行索引}`和`{经度值: 列索引}`的映射字典
- **网格重构**:
- 初始化二维索引数组`grid_idx`,存储每个位置在展平数组中的索引
- 使用双层循环遍历所有网格点,通过映射字典确定位置
- **数据重构**:
- 对每个包含`grid_point`的变量:
- 定位`grid_point`在维度中的位置
- 初始化包含NaN的占位数组
- 使用`grid_idx`将展平数据映射回二维网格
- 特殊优化:当`grid_point`是第一个维度时使用高效索引
- **性能优化**:
- 预先计算索引映射,避免循环内重复查找
- 使用切片(Slicing)减少内存拷贝
- 添加进度监控和耗时统计
#### 3. 边界处理
- 自动处理`lat_bnds`/`lon_bnds`等边界变量
- 保留所有原始属性(attrs)确保元数据完整性
- 使用NaN填充缺失网格点,避免数据错误
#### 4. 输出
- 生成符合CF标准的NetCDF文件
- 输出数据集结构验证信息
- 记录完整处理耗时
### 注意事项
1. **内存管理**:处理12GB数据需至少32GB RAM,可使用`dask`分块处理优化内存
2. **精度保留**:使用原数据的数据类型(float32)减少内存占用
3. **异常处理**:代码中`grid_idx`初始化为-999,用于标识无效网格点
4. **性能**:重构阶段双重循环复杂度O(nlat×nlon),在大网格上较慢但确保准确性