YetAnotherCoupler 3.2.0
Loading...
Searching...
No Matches
yac_replay.py
Go to the documentation of this file.
1#!/usr/bin/env python3
2# Copyright (c) 2024 The YAC Authors
3#
4# SPDX-License-Identifier: BSD-3-Clause
5
6from math import prod
7import numpy as np
8import xarray as xr
9import grid_utils
10import yac
11import argparse
12import logging
13
14
16 import isodate
17 print(s)
18 dt = isodate.parse_duration(s)
19 return np.timedelta64(int(dt.total_seconds()*10e6), "us")
20
21
22parser = argparse.ArgumentParser("yac_replay",
23 description="""Replay simulation from a dataset.
24The given files are loaded with `xarray.open_mfdataset`.
25Currently only cell based fields are supported that have 'cell' as dimension.
26It is currently not supported to run this component in parallel.""")
27parser.add_argument("files", type=str, nargs='+',
28 default="xarray URL for the data file")
29parser.add_argument("--engine", type=str, required=False, default=None,
30 help="xarray backend engine (see xarray.open_mfdataset)")
31parser.add_argument("--fallback_grid", type=str, required=False,
32 help="Description of grid (default: try to deduce from file")
33parser.add_argument("--compname", type=str, default="replay", required=False,
34 help="Name for the yac component (default: replay)")
35parser.add_argument("--gridprefix", type=str, default="replay_", required=False,
36 help="Prefix for the grid name (default: replay_)")
37parser.add_argument("--start_date", type=np.datetime64, required=False,
38 help="start time for date")
39parser.add_argument("--end_date", type=np.datetime64, required=False,
40 help="end time for date")
41parser.add_argument("--dt", type=parse_duration, required=False,
42 help="timestep")
43parser.add_argument('--debug', help="Print lots of debugging statements",
44 action="store_const", dest="loglevel", const=logging.DEBUG,
45 default=logging.WARNING)
46parser.add_argument('-v', '--verbose', help="Be verbose",
47 action="store_const", dest="loglevel", const=logging.INFO)
48
49args = parser.parse_args()
50log_handler = logging.StreamHandler()
51log_handler.setFormatter(logging.Formatter("%(asctime)-10s %(name)-10s %(levelname)-8s %(message)s"))
52logging.basicConfig(level=args.loglevel, handlers=[log_handler])
53
54logging.debug("set calendar to PROLEPTIC_GREGORIAN")
55yac.def_calendar(yac.Calendar.PROLEPTIC_GREGORIAN)
56logging.debug("Instantiate yac")
58
59logging.debug(f"open dataset: {args.files}")
60ds = xr.open_mfdataset(args.files, engine=args.engine)
61
62start_date = args.start_date or np.datetime64(ds.time[0].data)
63end_date = args.end_date or np.datetime64(ds.time[-1].data)
64
65if args.dt:
66 dt = args.dt
67else:
68 dt = np.diff(ds.coords["time"])
69 assert np.all(dt[0] == dt), "Time coordinates are not equidistant"
70 dt = dt[0]
71
72logging.info(f"define datetime: {start_date} - {end_date}")
73y.def_datetime(np.datetime_as_string(start_date)[:23],
74 np.datetime_as_string(end_date)[:23])
75
76logging.debug(f"define component: {args.compname}")
77comp = y.def_comp(args.compname)
78
79logging.info(f"Time interval: {y.start_datetime} - {y.end_datetime}")
80logging.info(f"dt: {dt} timesteps: {(end_date - start_date)/dt}")
81
82
83grids = {}
84def grid_cache(grid_spec):
85 global grids
86 if grid_spec not in grids:
87 logging.info(f"Defining grid {args.gridprefix+grid_spec}")
88 grids[grid_spec] = grid_utils.get_grid(grid_spec).def_yac(args.gridprefix+grid_spec, [yac.Location.CELL, yac.Location.CORNER])
89 return grids[grid_spec]
90
91
92def get_grid(variable):
93 global ds
94 if {"lat", "lon"} <= set(variable.dims):
95 lat = ds.lat
96 lon = ds.lon
97 return grid_cache(f"g{len(lon)},{len(lat)},{float(lon.min())},{float(lon.max())},{float(lat.min())},{float(lat.max())}")
98 if "grid_mapping" in variable.attrs:
99 crs = variable.attrs["grid_mapping"]
100 if ds[crs].attrs["grid_mapping_name"] == "healpix":
101 nside = ds[crs].attrs['healpix_nside']
102 nest = "n" if ds[crs].attrs['healpix_order'] == "nest" else "r"
103 return grid_cache(f"h{nside}{nest}")
104 if args.fallback_grid:
105 return grid_cache(args.fallback_grid)
106 raise Exception(f"Cannot not determine grid for variable {variable.name}. Please specify a fallback grid.")
107
108
109fields = []
110
111for varname in ds.variables:
112 var = ds[varname]
113 logging.debug(f"checking variable {varname}")
114 if "time" not in var.dims:
115 continue
116 if "cell" in var.dims or {"lat", "lon"} <= set(var.dims):
117 cells, points = get_grid(var)
118 point_id = cells
119 else:
120 continue
121 collection_size = prod(var.sizes[d] for d in var.dims if d != "time" and d != "cell")
122 logging.debug(f"define field with collection size: {collection_size}")
123 fields.append((yac.Field.create(varname, comp, point_id, collection_size,
124 str(int(dt / np.timedelta64(1, 's'))), yac.TimeUnit.SECOND),
125 var))
126
127logging.info("calling enddef")
128y.enddef()
129
130fields = [(field, var) for field, var in fields if field.role is yac.ExchangeType.SOURCE]
131logging.info(f"have {len(fields)} with role source")
132
133for t in np.arange(start_date, end_date+dt, dt):
134 logging.info(f"Time: {t}")
135 for field, var in fields:
136 logging.info(f"calling put for field {var.name}")
137 field.put(var.sel(time=t).data)
138
139logging.info("done")
create(cls, str field_name, Component comp, points, collection_size, str timestep, TimeUnit timeunit, masks=None)
Definition yac.pyx:1153
Initializies a YAC instance and provides further functionality.
Definition yac.pyx:471
get_grid(grid)
grid_cache(grid_spec)
Definition yac_replay.py:84
def_calendar(Calendar calendar)
Definition yac.pyx:395