| import os |
| import h5py |
| import numpy as np |
| import xarray as xr |
| from onescience.utils.YParams import YParams |
|
|
|
|
| |
| def get_dims(cfg_model, cfg_data): |
| H, W = int(cfg_model.n_lats_px), int(cfg_model.n_lons_px) |
| if tuple(map(int, cfg_data.dataset.img_size)) != (H, W): |
| raise ValueError("model grid and datapipe.dataset.img_size must match") |
| input_steps, output_steps = int(cfg_model.input_size_time), 1 |
| samples = max(int(cfg_data.dataloader.batch_size), 2) |
| T = input_steps + output_steps + samples - 1 |
| return { |
| "T": T, "H": H, "W": W, "time_step": 6, |
| "input_steps": input_steps, "output_steps": output_steps, |
| "static_channels": int(cfg_model.in_channels_static), |
| } |
|
|
|
|
| def generate_fake_h5(data_dir, var_names, years, dims): |
| """ |
| 为每个年份生成一个空 h5 文件。 |
| 利用 HDF5 chunked 数据集未写入 chunk 即返回 fill_value=0 的特性， |
| 文件实际只含元数据，极小，但 shape 与真实数据完全一致。 |
| 均值/标准差也作为数据集内嵌进每年的 h5，与 era5.py 新版读取方式对应。 |
| """ |
| os.makedirs(os.path.join(data_dir, "data"), exist_ok=True) |
| T, C = dims["T"], len(var_names) |
| H, W = dims["H"], dims["W"] |
|
|
| means = np.zeros((1, C, 1, 1), dtype=np.float32) |
| stds = np.ones((1, C, 1, 1), dtype=np.float32) |
|
|
| for year in years: |
| path = os.path.join(data_dir, "data", f"{year}.h5") |
| with h5py.File(path, "w") as f: |
| ds = f.create_dataset( |
| "fields", |
| shape=(T, C, H, W), |
| dtype="float32", |
| chunks=(1, C, H, W), |
| fillvalue=0.0, |
| ) |
| ds.attrs["variables"] = var_names |
| ds.attrs["time_step"] = dims["time_step"] |
| f.create_dataset("global_means", data=means) |
| f.create_dataset("global_stds", data=stds) |
|
|
| size_kb = os.path.getsize(path) / 1024 |
| print(f" {year}.h5 shape=({T},{C},{H},{W}) " |
| f"logical={T*C*H*W*4/1024**3:.1f}GB actual={size_kb:.1f}KB") |
|
|
|
|
| def get_static(data_dir, H, W, channels): |
| os.makedirs(data_dir, exist_ok=True) |
| lat = np.linspace(90, -90, H, dtype=np.float32) |
| lon = np.linspace(0, 360 - 360 / W, W, dtype=np.float32) |
| lat_grid = np.broadcast_to(lat[:, None], (H, W)) / 90.0 |
| lon_grid = np.broadcast_to(lon[None, :], (H, W)) / 180.0 - 1.0 |
| land_mask = (np.sin(np.deg2rad(lat_grid * 90)) > 0).astype(np.float32) |
| topography = np.cos(np.deg2rad(lat_grid * 90)).astype(np.float32) |
| base = [lat_grid, lon_grid, land_mask, topography] |
| static = np.stack((base * ((channels + 3) // 4))[:channels]).astype(np.float32) |
|
|
| ds = xr.Dataset( |
| data_vars={ |
| "z": (("valid_time", "latitude", "longitude"), static[-1:]), |
| "lsm": (("valid_time", "latitude", "longitude"), static[min(2, channels - 1):min(2, channels - 1) + 1]), |
| }, |
| coords={ |
| "valid_time": ["2015-12-31"], |
| "latitude": lat.astype(np.float64), |
| "longitude": lon.astype(np.float64), |
| "number": 0, |
| "expver": "", |
| }, |
| attrs={ |
| "GRIB_centre": "ecmf", |
| "GRIB_centreDescription": "European Centre for Medium-Range Weather Forecasts", |
| "GRIB_subCentre": "0", |
| "Conventions": "CF-1.7", |
| "institution": "European Centre for Medium-Range Weather Forecasts", |
| "history": "Generated manually", |
| } |
| ) |
|
|
| ds[["z"]].to_netcdf(f"{data_dir}/geopotential.nc") |
| ds[["lsm"]].to_netcdf(f"{data_dir}/land_sea_mask.nc") |
| np.save(f'{data_dir}/static.npy', static) |
| np.save(f'{data_dir}/land_mask.npy', land_mask) |
| np.save(f'{data_dir}/soil_type.npy', np.zeros((H, W), dtype=np.float32)) |
| np.save(f'{data_dir}/topography.npy', topography) |
| print(f"✅ Static data: {static.shape}, dtype: {static.dtype}, save to {data_dir}") |
|
|
|
|
| if __name__ == "__main__": |
| cfg_model = YParams("conf/config.yaml", "model") |
| cfg_datapipe = YParams("conf/config.yaml", "datapipe") |
|
|
| if cfg_datapipe.dataset.data_dir.startswith("/public/") or cfg_datapipe.dataset.data_dir.startswith("/work2/"): |
| print("请检查 config，确保各 *_dir 指向本地测试路径而非生产路径。") |
| exit() |
|
|
| years = cfg_datapipe.dataset.train_time + cfg_datapipe.dataset.val_time + cfg_datapipe.dataset.test_time |
| atm_vars = cfg_datapipe.dataset.channels |
| if len(atm_vars) != int(cfg_model.in_channels): |
| raise ValueError("channel count must match model.in_channels") |
|
|
| dims = get_dims(cfg_model, cfg_datapipe) |
| generate_fake_h5(cfg_datapipe.dataset.data_dir, atm_vars, years, dims) |
|
|
| static_dir = os.path.join(cfg_datapipe.dataset.data_dir, "static") |
| get_static(static_dir, dims["H"], dims["W"], dims["static_channels"]) |
|
|
| print("\n✅ Fake datasets generated.") |
|
|
