credit.datasets.gen_2.grid_utils ================================ .. py:module:: credit.datasets.gen_2.grid_utils .. autoapi-nested-parse:: grid_utils.py ------------- Everything about horizontal grid geometry for gen2: coordinate-pair detection, rectilinear-vs-curvilinear classification, and GridSchema (the real-coordinate contract for output, mirroring ChannelSchema in channel_utils.py). find_coord_pair / infer_grid_type ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Small, dependency-free helpers for locating a lon/lat coordinate pair in an ``xr.Dataset`` (by name) and classifying it as rectilinear (1D) or curvilinear (2D). Used by the per-dataset grid detection in ``local.py``/``era5.py``, and by ``credit.grid.scrip_from_netcdf`` (SCRIP-format grid generation for ESMF regridding) — this is the shared home for both rather than duplicating the logic in each. GridSchema ~~~~~~~~~~ ``ForecastWriter`` previously fabricated output lat/lon from ``model.image_height``/``image_width`` via a global ``[-90, 90] x [0, 360)`` linspace — wrong for regional domains and for curvilinear sources (e.g. HRRR), which need a real 2D lat/lon field, not a fabricated 1D one. There is no fabricated fallback anywhere in this pipeline anymore: if a real grid can't be found or resolved, callers raise rather than guess. Two distinct grid concepts: * **Native input grid** — each dataset class (``LocalDataset``, ``era5.py``, ``goes.py``, ``hrrr.py``) exposes the real lat/lon it read from its own files via ``self.static_metadata["grid"]`` (a debugging aid, inspectable directly on the live dataset — not necessarily what ends up in the output file). The same dataset class also best-effort persists it to ``{save_loc}/{source}_grid_schema.nc`` the moment it's known, via ``write_source_grid_schema_if_missing`` — including from inside a DataLoader worker subprocess, which has filesystem access even though it can't propagate Python object state back to the main training process. * **Resolved output grid** — what ``ForecastWriter`` actually writes. The model produces one flat tensor at one fixed ``(H, W)``, so there is exactly one output grid per run: the (single, in practice) source's native grid, overridden by an active ``Regridder`` preblock's real destination grid when regridding is in play (the common case where regridding is *not* used is the default: the native grid passes straight through unchanged). Lifecycle ^^^^^^^^^ Training/rollout setup resolves the schema once — via ``GridSchema.resolve`` using the live dataset + preblocks (with a disk fallback to each source's ``{source}_grid_schema.nc`` when the live process never saw the data itself) — and, only when an active ``Regridder`` actually changes the grid, saves the result to ``{save_loc}/output_grid_schema.nc`` (skipped when no regridder is active: the single source's own file already *is* the effective output grid, so a second identical copy would be pure duplication). Later runs (or a re-run without training) load whichever file is present instead of re-resolving, via ``GridSchema.load_or_resolve``. Scope: rectilinear and curvilinear only (no unstructured). Projection/CRS metadata (e.g. HRRR's Lambert Conformal Conic) is deferred to a follow-on — plain lat/lon coordinates are CF-valid on their own. Attributes ---------- .. autoapisummary:: credit.datasets.gen_2.grid_utils.logger credit.datasets.gen_2.grid_utils.SOURCE_GRID_SCHEMA_FILENAME credit.datasets.gen_2.grid_utils.OUTPUT_GRID_SCHEMA_FILENAME credit.datasets.gen_2.grid_utils.GridType Classes ------- .. autoapisummary:: credit.datasets.gen_2.grid_utils.GridSchema Functions --------- .. autoapisummary:: credit.datasets.gen_2.grid_utils.find_coord_pair credit.datasets.gen_2.grid_utils.infer_grid_type credit.datasets.gen_2.grid_utils.write_source_grid_schema_if_missing Module Contents --------------- .. py:data:: logger .. py:function:: find_coord_pair(ds) Find a lon/lat coordinate pair in an xr.Dataset. Searches _COORD_CANDIDATES in order across both ds.coords and ds.data_vars. :rtype: (lon_array, lat_array, lon_name, lat_name) :raises ValueError if no recognised pair is found.: .. py:function:: infer_grid_type(lat, lon) Classify a lat/lon coordinate pair as "rectilinear" or "curvilinear". 1D lat and lon -> rectilinear. 2D lat and lon -> curvilinear. :returns: **str** :rtype: "rectilinear" or "curvilinear" :raises ValueError if lat/lon are not both 1D or both 2D.: .. py:data:: SOURCE_GRID_SCHEMA_FILENAME :value: '{source}_grid_schema.nc' .. py:data:: OUTPUT_GRID_SCHEMA_FILENAME :value: 'output_grid_schema.nc' .. py:data:: GridType .. py:function:: write_source_grid_schema_if_missing(source_name: str, grid: dict[str, Any] | None, save_loc: str | None) -> None Best-effort persist one source's native grid to ``{save_loc}/{source}_grid_schema.nc``. Called from each dataset class right after ``static_metadata["grid"]`` is first populated — including from inside a DataLoader worker subprocess under ``num_workers > 0``, since workers have their own filesystem access even though they can't propagate Python object state back to the main process. Guarded by file existence, so redundant calls (e.g. multiple workers independently resolving the same grid) are cheap after the first successful write. Failures (no ``save_loc``, read-only filesystem, ...) are logged and swallowed — this must never break the data-loading path it's piggybacked onto. A source can report a native ``grid_type`` this module doesn't support as a resolved output grid (e.g. ``"unstructured"`` — see the module docstring's "Scope" note): that's expected, not a failure, so it's logged at info level and skipped here rather than attempting (and swallowing the inevitable failure of) a ``GridSchema`` construction. Such a source still needs an active ``Regridder`` preblock for ``GridSchema.resolve`` to produce a resolvable output grid (see its docstring). .. py:class:: GridSchema(grid_type: GridType, lat: numpy.ndarray, lon: numpy.ndarray, origin: Literal['native', 'regridded'] = 'native') The resolved horizontal output grid: shared across every variable/source in one output file (the model produces one flat tensor at one fixed shape). :param grid_type: ``"rectilinear"`` or ``"curvilinear"``. :param lat: 1D (rectilinear) or 2D ``(y, x)`` (curvilinear) latitude array. :param lon: 1D (rectilinear) or 2D ``(y, x)`` (curvilinear) longitude array. :param origin: ``"native"`` (a source's own grid, unchanged) or ``"regridded"`` (an active ``Regridder`` preblock's destination grid). Set by ``.resolve()``; defaults to ``"native"`` for direct construction. Used by callers (e.g. the trainer) to decide whether this schema needs its own ``output_grid_schema.nc`` — a ``"native"`` schema is already fully covered by the source's own ``{source}_grid_schema.nc``. .. py:attribute:: grid_type :type: GridType .. py:attribute:: lat .. py:attribute:: lon .. py:attribute:: origin :value: 'native' .. py:method:: resolve(dataset: Any, ic_preblocks=None, step_preblocks=None, save_loc: str | None = None) -> GridSchema :classmethod: Resolve the effective output grid from a live dataset + preblocks. Starts from the dataset's native grid (``static_metadata["grid"]``, falling back to each source's persisted ``{source}_grid_schema.nc`` in *save_loc* when the live process doesn't have it); if an active ``Regridder`` preblock is found, its real destination grid is used instead. When no regridder is active (the common case), the native grid passes through unchanged and the returned schema's ``origin`` is ``"native"``. :raises ValueError: if no native grid is available, if grids disagree (see ``_native_grid`` / ``_find_regridder``), or if the native grid_type (e.g. ``"unstructured"``) isn't directly resolvable and no active Regridder preblock provides a rectilinear/curvilinear destination. .. py:method:: save(path: str) -> None Write the schema as NetCDF (atomically: temp file + rename). The staging file is per-process. Every rank — and every DataLoader worker under ``num_workers > 0`` — persists the grid it resolved (see ``write_source_grid_schema_if_missing``), so several processes routinely write the same *path* at once. A single shared ``.tmp`` made them collide inside HDF5, surfacing as ``[Errno 13] Permission denied`` for every writer but the first. With one temp file each, the writes are independent and the renames are atomic; the content is identical, so last-writer-wins is harmless. .. py:method:: load(path: str) -> GridSchema :classmethod: .. py:method:: load_or_resolve(dataset: Any, ic_preblocks=None, step_preblocks=None, save_loc: str | None = None) -> GridSchema | None :classmethod: Load ``output_grid_schema.nc`` from ``save_loc`` (only ever written when regridding was used), else resolve live from *dataset* (native grid, with a disk fallback to each source's own ``{source}_grid_schema.nc``, overridden by an active Regridder). Returns ``None`` (with a warning) when neither is possible — callers should treat that as fatal; there is no fabricated fallback.