209 lines
11 KiB
Python
209 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import math
|
|
from pathlib import Path
|
|
import re
|
|
|
|
import numpy as np
|
|
|
|
GENERATOR_VERSION = "2"
|
|
_RESOLUTIONS = {"T21": 32, "T42": 64, "T63": 96, "T85": 128}
|
|
_SURFACE = {
|
|
"seed": {"type": "integer", "min": 0, "max": 4294967295, "default": 1},
|
|
"land_fraction": {"type": "number", "min": 0.0, "max": 1.0, "default": 0.3},
|
|
"relief": {"type": "number", "min": 0.0, "max": 12000.0, "default": 4000.0, "unit": "m"},
|
|
"scale": {"type": "number", "min": 1.0, "max": 12.0, "default": 4.0},
|
|
}
|
|
|
|
|
|
def surface_schema():
|
|
return json.loads(json.dumps(_SURFACE))
|
|
|
|
|
|
def _validate_surface(surface):
|
|
if not isinstance(surface, dict) or set(surface) - set(_SURFACE):
|
|
raise ValueError("surface must contain only supported generator fields")
|
|
result = {key: item["default"] for key, item in _SURFACE.items()}
|
|
result.update(surface)
|
|
for key, item in _SURFACE.items():
|
|
value = result[key]
|
|
if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value):
|
|
raise ValueError(f"surface.{key} must be finite numeric")
|
|
if value < item["min"] or value > item["max"] or (key == "seed" and int(value) != value):
|
|
raise ValueError(f"surface.{key} outside supported range")
|
|
if key == "seed":
|
|
result[key] = int(value)
|
|
return result
|
|
|
|
|
|
def write_map_sra(path, resolution, field, destination, gravity, landmask=None):
|
|
try:
|
|
from .mapcatalog import FIELDS
|
|
except ImportError:
|
|
from mapcatalog import FIELDS
|
|
contract = FIELDS[field]
|
|
code = contract['native_code']
|
|
raw = np.fromfile(path, dtype='<f4')
|
|
nlat = _RESOLUTIONS[resolution]
|
|
if raw.size != nlat * nlat * 2 or not np.isfinite(raw).all():
|
|
raise ValueError('snapshot map has invalid dimensions or values')
|
|
grid = raw.reshape((nlat, nlat * 2))
|
|
if not np.isfinite(grid).all() or np.any(grid < contract['min']) or np.any(grid > contract['max']):
|
|
raise ValueError(f'{field} values outside supported range')
|
|
if contract['categorical'] and not np.isin(grid, (0.0, 1.0)).all():
|
|
raise ValueError('land mask values must be 0 or 1')
|
|
if field == 'height':
|
|
if landmask is None:
|
|
land = np.ones(grid.shape, dtype=bool)
|
|
else:
|
|
mask_path = Path(landmask)
|
|
if mask_path.suffix == '.f32':
|
|
values = np.fromfile(mask_path, dtype='<f4')
|
|
if values.size != grid.size or not np.isfinite(values).all():
|
|
raise ValueError('land mask grid has invalid dimensions or values')
|
|
land = values.reshape(grid.shape) > 0.5
|
|
else:
|
|
values = np.fromstring(mask_path.read_text(), sep=' ')[8:]
|
|
if values.size != grid.size:
|
|
raise ValueError('land mask grid has invalid dimensions')
|
|
land = values.reshape(grid.shape) > 0.5
|
|
grid = np.where(land & (grid > 0.0), grid * gravity, 0.0)
|
|
destination = Path(destination)
|
|
months = contract.get('months', 1)
|
|
with destination.open('w', encoding='ascii') as stream:
|
|
for _ in range(months):
|
|
stream.write(''.join(f'{value:10d}' for value in (code, 0, 0, 0, nlat * 2, nlat, 0, 0)) + '\n')
|
|
stream.write(''.join(f'{value:16.6E}' for value in grid.ravel()) + '\n')
|
|
return str(destination)
|
|
|
|
|
|
def _gaussian_grid(resolution):
|
|
nlat = _RESOLUTIONS[resolution]
|
|
x, weights = np.polynomial.legendre.leggauss(nlat)
|
|
lat = np.rad2deg(np.arcsin(x[::-1]))
|
|
lon = np.arange(nlat * 2, dtype=np.float64) * (360.0 / (nlat * 2))
|
|
return lat, lon, weights[::-1]
|
|
|
|
|
|
def _terrain(surface, lat, lon, lat_weights):
|
|
rng = np.random.default_rng(surface['seed'])
|
|
lonr, latr = np.meshgrid(np.deg2rad(lon), np.deg2rad(lat))
|
|
xyz = np.stack((np.cos(latr) * np.cos(lonr), np.cos(latr) * np.sin(lonr), np.sin(latr)), axis=-1)
|
|
field = np.zeros(latr.shape)
|
|
for degree in range(1, 9):
|
|
count = max(4, 2 * degree + 1)
|
|
centers = rng.normal(size=(count, 3))
|
|
centers /= np.linalg.norm(centers, axis=1)[:, None]
|
|
amplitudes = rng.normal(size=count)
|
|
field += (np.einsum('...c,kc->...k', xyz, centers) ** degree @ amplitudes) / degree ** 1.25
|
|
field -= np.average(field, weights=np.broadcast_to(lat_weights[:, None], field.shape))
|
|
spread = np.std(field)
|
|
if spread:
|
|
field /= spread
|
|
fine = np.sin(surface['scale'] * 2.0 * xyz[..., 0] + surface['scale'] * 3.0 * xyz[..., 1] + surface['scale'] * xyz[..., 2])
|
|
return .75 * field + .25 * fine
|
|
|
|
|
|
def _reference_gravity(config):
|
|
path = Path(config['_paths']['catalog'])
|
|
fields = json.loads(path.read_text())['fields']
|
|
return next(field['implicit_default'] for field in fields if field['name'] == 'gravity')
|
|
|
|
|
|
def prepare_surface(surface, model, config, destination):
|
|
surface = _validate_surface(surface)
|
|
resolution = model.get("resolution")
|
|
if resolution not in _RESOLUTIONS:
|
|
raise ValueError("unsupported model resolution")
|
|
if not isinstance(config, dict):
|
|
raise ValueError("config must be a mapping")
|
|
lat, lon, lat_weights = _gaussian_grid(resolution)
|
|
field = _terrain(surface, lat, lon, lat_weights)
|
|
cell_weights = np.broadcast_to(lat_weights[:, None], field.shape)
|
|
fraction = surface["land_fraction"]
|
|
if fraction == 0:
|
|
land = np.zeros(field.shape, dtype=bool)
|
|
threshold = float(np.max(field))
|
|
elif fraction == 1:
|
|
land = np.ones(field.shape, dtype=bool)
|
|
threshold = float(np.min(field))
|
|
else:
|
|
order = np.argsort(field.ravel())[::-1]
|
|
cumulative = np.cumsum(cell_weights.ravel()[order]) / np.sum(cell_weights)
|
|
position = int(np.searchsorted(cumulative, fraction, side="left"))
|
|
threshold = float(field.ravel()[order[position]])
|
|
land = field >= threshold
|
|
# Boundary ties are resolved by deterministic cell order to attain closest discrete area.
|
|
equal = np.flatnonzero(field.ravel() == threshold)
|
|
desired = fraction * np.sum(cell_weights)
|
|
above = np.sum(cell_weights[field > threshold])
|
|
choose = max(0, min(len(equal), int(round((desired - above) / cell_weights.flat[0]))))
|
|
land.ravel()[equal] = False
|
|
land.ravel()[equal[:choose]] = True
|
|
bathymetry = np.where(land, np.maximum(field - threshold, 0.0), np.minimum(field - threshold, 0.0))
|
|
maximum = float(np.max(np.abs(bathymetry)))
|
|
elevation = bathymetry * (surface["relief"] / maximum) if maximum else bathymetry
|
|
bathymetry = np.minimum(elevation, 0.0)
|
|
gravity = config.get('parameters', {}).get('gravity', config.get('gravity'))
|
|
gravity = float(_reference_gravity(config) if gravity is None else gravity)
|
|
geopotential = np.where(land, np.maximum(elevation, 0.0) * gravity, 0.0)
|
|
root = Path(destination)
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
token = hashlib.sha256(json.dumps({"version": GENERATOR_VERSION, "surface": surface, "resolution": resolution}, sort_keys=True).encode()).hexdigest()[:32]
|
|
npz = root / f"surface-{token}.npz"
|
|
metadata = root / f"surface-{token}_metadata.npz"
|
|
np.savez_compressed(npz, lat=lat, lon=lon, time=np.array([0.0]), elevation=elevation[None], bathymetry=bathymetry[None], landmask=land.astype(np.float32)[None])
|
|
np.savez(metadata, elevation=np.array(["Elevation", "Procedural elevation", "m"]), bathymetry=np.array(["Bathymetry", "Signed procedural bathymetry", "m"]), landmask=np.array(["Land mask", "Procedural land fraction", "1"]))
|
|
nlat = len(lat)
|
|
prefix = f"N{nlat:03d}_surf_"
|
|
topopath, landpath = root / f"{prefix}0129.sra", root / f"{prefix}0172.sra"
|
|
def write_sra(path, code, values):
|
|
with path.open("w", encoding="ascii") as stream:
|
|
stream.write("".join(f"{value:10d}" for value in (code, 0, 0, 0, len(lon), nlat, 0, 0)) + "\n")
|
|
stream.write("".join(f"{value:16.6E}" for value in np.asarray(values).ravel()) + "\n")
|
|
write_sra(topopath, 129, geopotential)
|
|
write_sra(landpath, 172, land.astype(float))
|
|
return {"landmap": str(landpath), "topomap": str(topopath), "preview": str(npz), "metadata": str(metadata), "land_fraction_actual": float(np.sum(cell_weights * land) / np.sum(cell_weights)), "latitude": lat.tolist(), "longitude": lon.tolist()}
|
|
|
|
|
|
def generate_planet(recipe, config, destination):
|
|
if not isinstance(recipe, dict) or set(recipe) != {"model", "parameters", "run", "surface"}:
|
|
raise ValueError("recipe must contain model, parameters, run, and surface")
|
|
spec = json.loads(json.dumps(recipe))
|
|
surface = _validate_surface(spec["surface"])
|
|
spec["surface"] = surface
|
|
model = spec["model"]
|
|
if model.get("resolution") not in _RESOLUTIONS:
|
|
raise ValueError("unsupported model resolution")
|
|
rng = np.random.default_rng(surface["seed"])
|
|
params = spec["parameters"]
|
|
# Mass is retained as generator evidence; ExoPlaSim accepts radius and gravity instead.
|
|
mass = float(rng.uniform(0.6, 1.4) ** 2.7)
|
|
radius = mass ** 0.27
|
|
gravity = _reference_gravity(config) * mass / radius ** 2
|
|
params.setdefault("radius", radius)
|
|
params.setdefault("gravity", gravity)
|
|
params.setdefault("startemp", float(rng.uniform(4000, 6200)))
|
|
params.setdefault("flux", float(rng.uniform(700, 1400)))
|
|
params.setdefault("rotationperiod", float(rng.uniform(0.5, 2.0)))
|
|
params.setdefault("obliquity", float(rng.uniform(0, 35)))
|
|
params.setdefault("eccentricity", float(rng.uniform(0, 0.15)))
|
|
params.setdefault("pN2", 1.0)
|
|
params.setdefault("pCO2", float(rng.uniform(1e-4, 2e-2)))
|
|
surface_config = dict(config)
|
|
surface_config['parameters'] = params
|
|
dataset = prepare_surface(surface, model, surface_config, destination)
|
|
dataset_id = hashlib.sha256(json.dumps({"version": GENERATOR_VERSION, "recipe": spec}, sort_keys=True, separators=(",", ":")).encode()).hexdigest()[:32]
|
|
if __package__:
|
|
from .results import prepare_results
|
|
else:
|
|
from results import prepare_results
|
|
index = prepare_results(dataset["preview"], Path(destination) / "preview", dataset_id, config=config)
|
|
manifest = {"generator_version": GENERATOR_VERSION, "seed": surface["seed"], "sampled_mass_earth_units": mass, "recipe": spec, "dataset_id": dataset_id, "land_fraction_actual": dataset["land_fraction_actual"], "surface_files": {"landmap": Path(dataset["landmap"]).name, "topomap": Path(dataset["topomap"]).name, "preview": Path(dataset["preview"]).name}}
|
|
(Path(destination) / "planet.json").write_text(json.dumps(manifest, sort_keys=True, indent=2), encoding="utf-8")
|
|
dataset["index"] = index
|
|
dataset["id"] = dataset_id
|
|
return {"spec": spec, "dataset": dataset}
|