OneScience-Group/Pangu_Weather

1import os2import h5py3import numpy as np4import xarray as xr5from onescience.utils.YParams import YParams6 7 8# 各数据集固定的空间和时间维度9DATASET_DIMS = {"T": 10, "H": 721, "W": 1440, "time_step": 6}10 11 12def generate_fake_h5(data_dir, var_names, years, dims):13    """14    为每个年份生成一个空 h5 文件。15    利用 HDF5 chunked 数据集未写入 chunk 即返回 fill_value=0 的特性，16    文件实际只含元数据，极小，但 shape 与真实数据完全一致。17    均值/标准差也作为数据集内嵌进每年的 h5，与 era5.py 新版读取方式对应。18    """19    os.makedirs(os.path.join(data_dir, "data"), exist_ok=True)20    T, C = dims["T"], len(var_names)21    H, W = dims["H"], dims["W"]22 23    means = np.zeros((1, C, 1, 1), dtype=np.float32)24    stds  = np.ones((1, C, 1, 1), dtype=np.float32)25 26    for year in years:27        path = os.path.join(data_dir, "data", f"{year}.h5")28        with h5py.File(path, "w") as f:29            ds = f.create_dataset(30                "fields",31                shape=(T, C, H, W),32                dtype="float32",33                chunks=(1, C, H, W),34                fillvalue=0.0,35            )36            ds.attrs["variables"] = var_names37            ds.attrs["time_step"] = dims["time_step"]38            f.create_dataset("global_means", data=means)39            f.create_dataset("global_stds",  data=stds)40 41        size_kb = os.path.getsize(path) / 102442        print(f"  {year}.h5  shape=({T},{C},{H},{W})  "43              f"logical={T*C*H*W*4/1024**3:.1f}GB  actual={size_kb:.1f}KB")44 45 46def get_static(data_dir, var, name):47    os.makedirs(data_dir, exist_ok=True)48    ds = xr.Dataset(49        data_vars={50            f"{var}": (("valid_time", "latitude", "longitude"),51                np.random.rand(1, 721, 1440).astype(np.float32))52        },53        coords={54            "valid_time": ["2015-12-31"],55            "latitude": np.linspace(90, -90, 721, dtype=np.float64),56            "longitude": np.linspace(0, 359.75, 1440, dtype=np.float64),57            "number": 0,58            "expver": "",59        },60        attrs={61            "GRIB_centre": "ecmf",62            "GRIB_centreDescription": "European Centre for Medium-Range Weather Forecasts",63            "GRIB_subCentre": "0",64            "Conventions": "CF-1.7",65            "institution": "European Centre for Medium-Range Weather Forecasts",66            "history": "Generated manually",67        }68    )69 70    ds.to_netcdf(f"{data_dir}/{name}.nc")71    arr = np.random.randn(721, 1440).astype(np.float32)72    np.save(f'{data_dir}/land_mask.npy', arr)73    np.save(f'{data_dir}/soil_type.npy', arr)74    np.save(f'{data_dir}/topography.npy', arr)75    print(f"✅ Static data: {arr.shape}, dtype: {arr.dtype}, save to {data_dir}")76 77 78if __name__ == "__main__":79    cfg_datapipe = YParams("conf/config.yaml", "datapipe")80 81    if cfg_datapipe.dataset.data_dir.startswith("/public/") or cfg_datapipe.dataset.data_dir.startswith("/work2/"):82        print("请检查 config，确保各 *_dir 指向本地测试路径而非生产路径。")83        exit()84 85    years    = cfg_datapipe.dataset.train_time + cfg_datapipe.dataset.val_time + cfg_datapipe.dataset.test_time86    atm_vars = cfg_datapipe.dataset.channels87 88    generate_fake_h5(cfg_datapipe.dataset.data_dir, atm_vars, years, DATASET_DIMS)89 90    static_dir = os.path.join(cfg_datapipe.dataset.data_dir, "static")91    get_static(static_dir, 'z', 'geopotential')92    get_static(static_dir, 'lsm', 'land_sea_mask')93 94 95    print("\n✅ Fake datasets generated.")96