File size: 4,684 Bytes
d3e46b7 | 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 | """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()
|