| """Validate NetCDF variables, dimensions, and statistics for FV3GFS inputs.""" |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import xarray as xr |
|
|
|
|
| REQUIRED_VARIABLES = { |
| "PRESsfc", |
| "surface_temperature", |
| "DSWRFtoa", |
| "HGTsfc", |
| "ocean_fraction", |
| *{f"air_temperature_{i}" for i in range(8)}, |
| *{f"specific_total_water_{i}" for i in range(8)}, |
| *{f"eastward_wind_{i}" for i in range(8)}, |
| *{f"northward_wind_{i}" for i in range(8)}, |
| } |
|
|
|
|
| def validate_data(data_path: str, latitude: int, longitude: int) -> int: |
| files = sorted(Path(data_path).rglob("*.nc")) |
| if not files: |
| print(f"ERROR: no NetCDF files found under {data_path}") |
| return 1 |
| failed = False |
| for path in files: |
| with xr.open_dataset(path) as dataset: |
| variables = set(dataset.data_vars) |
| missing = sorted(REQUIRED_VARIABLES - variables) |
| dimensions = {name: int(size) for name, size in dataset.sizes.items()} |
| print(f"{path}: dimensions={dimensions}, variables={len(variables)}") |
| if missing: |
| print(f"ERROR: missing required variables: {', '.join(missing)}") |
| failed = True |
| if latitude not in dimensions.values() or longitude not in dimensions.values(): |
| print(f"ERROR: expected spatial dimensions containing {latitude} and {longitude}") |
| failed = True |
| time_sizes = [dimensions[name] for name in dimensions if "time" in name.lower()] |
| if not time_sizes or max(time_sizes) < 7: |
| print("ERROR: expected a time dimension with at least 7 frames") |
| failed = True |
| return int(failed) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--data-dir", required=True) |
| parser.add_argument("--latitude", type=int, default=180) |
| parser.add_argument("--longitude", type=int, default=360) |
| args = parser.parse_args() |
| raise SystemExit(validate_data(args.data_dir, args.latitude, args.longitude)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
