| """Create small physical-coordinate tiles without allocating a global dense sample.""" |
|
|
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import yaml |
|
|
| import sys |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
| from model.fuxi_ocean import enumerate_global_tiles |
|
|
|
|
| def main(): |
| config = yaml.safe_load((ROOT / "conf/config.yaml").read_text()) |
| data = config["data"] |
| rng = np.random.default_rng(config["seed"]) |
| count = data["train_samples"] + data["test_samples"] |
| th, tw = data["tile_height"], data["tile_width"] |
| history, channels, atmosphere_channels = data["history_steps"], data["ocean_channels"], data["atmosphere_channels"] |
| global_records = enumerate_global_tiles((data["global_height"], data["global_width"]), (th, tw), data["tile_overlap"]) |
| selected_indices = np.linspace(0, len(global_records) - 1, count, dtype=np.int64) |
| selected_records = global_records[selected_indices] |
| origins = selected_records[:, (0, 2)] |
| ocean = np.empty((count, history, channels, th, tw), dtype=np.float32) |
| targets = np.empty((count, channels, th, tw), dtype=np.float32) |
| atmosphere = np.empty((count, atmosphere_channels, th, tw), dtype=np.float32) |
| latitude = np.empty((count, th), dtype=np.float32) |
| longitude = np.empty((count, tw), dtype=np.float32) |
| bathymetry = np.empty((count, 1, th, tw), dtype=np.float32) |
| depth_mask = np.empty((count, 26, th, tw), dtype=np.float32) |
| time_features = np.empty((count, 3), dtype=np.float32) |
| yy, xx = np.meshgrid(np.arange(th), np.arange(tw), indexing="ij") |
| channel_scale = np.linspace(0.2, 1.0, channels, dtype=np.float32)[:, None, None] |
| for sample, (y0, x0) in enumerate(origins): |
| lat = 90.0 - (y0 + np.arange(th) + 0.5) * 180.0 / data["global_height"] |
| lon = (x0 + np.arange(tw) + 0.5) * 360.0 / data["global_width"] |
| latitude[sample], longitude[sample] = lat, lon |
| wave = np.sin(np.deg2rad(lat))[:, None] + 0.5 * np.cos(np.deg2rad(lon))[None, :] |
| for history_index in range(history): |
| ocean[sample, history_index] = channel_scale * (wave + 0.025 * history_index) + rng.normal(0, 0.005, (channels, th, tw)) |
| atmosphere[sample] = np.stack([wave + 0.03 * k for k in range(atmosphere_channels)]) |
| bathymetry[sample, 0] = 200 + 1400 * (0.5 + 0.5 * np.sin((xx + x0) / 12)) |
| for depth, depth_m in enumerate(data["depth_levels_m"]): |
| depth_mask[sample, depth] = bathymetry[sample, 0] >= depth_m |
| targets[sample] = ocean[sample, -1] + channel_scale * (0.015 * atmosphere[sample, 0] - 0.006 * atmosphere[sample, 1]) |
| time_features[sample] = ((sample * 6) % 24, 100 + sample // 4, 0) |
| channel_depth = np.asarray(list(range(26)) * 4 + [0]) |
| targets *= depth_mask[:, channel_depth] |
| path = ROOT / data["path"] |
| path.parent.mkdir(parents=True, exist_ok=True) |
| input_shape = np.asarray([history, channels, data["global_height"], data["global_width"]]) |
| atmosphere_shape = np.asarray([atmosphere_channels, data["global_height"], data["global_width"]]) |
| output_shape = np.asarray([channels, data["global_height"], data["global_width"]]) |
| owned_pixels = sum((r[5] - r[4]) * (r[7] - r[6]) for r in selected_records) |
| coverage_fraction = owned_pixels / (data["global_height"] * data["global_width"]) |
| np.savez_compressed(path, format_version=data["format_version"], input_shape=input_shape, |
| atmosphere_shape=atmosphere_shape, output_shape=output_shape, |
| global_tile_records=global_records, selected_tile_indices=selected_indices, |
| selected_tile_records=selected_records, tile_origins=origins, |
| tile_record_fields=np.asarray(["y0", "y1", "x0", "x1", "crop_top", "crop_bottom", "crop_left", "crop_right"]), |
| tile_order="row-major", overlap_crop_semantics="midpoint ownership; crops partition global grid exactly", |
| coverage_fraction=coverage_fraction, is_complete_global=False, synthetic=True, |
| ocean=ocean, atmosphere=atmosphere, targets=targets, latitude_deg=latitude, |
| longitude_deg=longitude, bathymetry_m=bathymetry, depth_mask=depth_mask, |
| time_features=time_features, train_count=data["train_samples"]) |
| print(json.dumps({"path": str(path.relative_to(ROOT)), "input_shape": input_shape.tolist(), |
| "tile_shape": list(ocean.shape), "global_tile_count": len(global_records), |
| "coverage_fraction": coverage_fraction, "bytes": path.stat().st_size})) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
