def __init__(self, topoconfig: TopoConfig, grid: Gridlines, dtype: str):
dtype = DummyDtype.validator(dtype)
dem, _ = read_cf(topoconfig.file, [topoconfig.key])
# copy to a nplike.ndarray
vert = nplike.array(dem[topoconfig.key][:])
# see if we need to do interpolation
try:
interp = not (nplike.allclose(grid.x.vert, nplike.array(dem["x"]))
and nplike.allclose(grid.y.vert,
nplike.array(dem["y"])))
except ValueError: # assume thie excpetion means a shape mismatch
interp = True
# unfortunately, we need to do interpolation in such a situation
if interp:
interpolator = RectBivariateSpline(dem["x"], dem["y"], vert.T)
vert = nplike.array(interpolator(
grid.x.vert, grid.y.vert).T) # it uses vanilla numpy
# cast to desired float type
vert = vert.astype(dtype)
# topography elevation at cell centers through linear interpolation
cntr = vert[:-1, :-1] + vert[:-1, 1:] + vert[1:, :-1] + vert[1:, 1:]
cntr /= 4
# topography elevation at cell faces' midpoints through linear interpolation
xface = (vert[:-1, :] + vert[1:, :]) / 2.
yface = (vert[:, :-1] + vert[:, 1:]) / 2.
# gradient at cell centers through central difference; here allows nonuniform grids
# this function does not assume constant cell sizes, so we re-calculate dx, dy
# the `delta`s in grid.x and y are constants (current solver only supports uniform grid)
xgrad = (xface[:, 1:] - xface[:, :-1]) / (grid.x.vert[1:] -
grid.x.vert[:-1])[None, :]
ygrad = (yface[1:, :] - yface[:-1, :]) / (grid.y.vert[1:] -
grid.y.vert[:-1])[:, None]
# initialize DataModel and let pydantic validates data
super().__init__(nx=grid.x.n,
ny=grid.y.n,
dtype=dtype,
vert=vert,
cntr=cntr,
xface=xface,
yface=yface,
xgrad=xgrad,
ygrad=ygrad)
def _val_linspace(cls, v, values):
"""Make sure the linspace is working correctly."""
diff = v[1:] - v[:-1]
assert nplike.all(diff > 0), "Not in monotonically increasing order."
assert nplike.allclose(diff, values["delta"],
atol=1e-10), "Delta does not match."
return v
def create_ic(comm, ic_config, grid, topo, dtype):
"""Create initial conditions.
When the x_cntr and y_cntr have different resolutions from the x and y in the NetCDF file, an
bi-cubic spline interpolation will take place.
Arguments
---------
comm : mpi4py.MPI.Comm
The communicator.
ic_config : torchswe.utils.config.ICConfig
grid : torchswe.utils.data.Gridlines
topo : torchswe.utils.data.Topography
dtype : str; either "float32" or "float64"
Returns
-------
torchswe.utils.data.WHUHVModel
"""
# special case: constant I.C.
if ic_config.values is not None:
return _WHUHVModel(nx=grid.x.n,
ny=grid.y.n,
dtype=dtype,
w=_nplike.maximum(
topo.cntr, _nplike.array(ic_config.values[0])),
hu=_nplike.full(topo.cntr.shape,
ic_config.values[1],
dtype=topo.dtype),
hv=_nplike.full(topo.cntr.shape,
ic_config.values[2],
dtype=topo.dtype))
# otherwise, read data from a NetCDF file
icdata, _ = _ncread(
ic_config.file,
ic_config.keys,
[grid.x.cntr[0], grid.x.cntr[-1], grid.y.cntr[0], grid.y.cntr[-1]],
parallel=True,
comm=comm)
# see if we need to do interpolation
try:
interp = not (_nplike.allclose(grid.x.cntr, icdata["x"])
and _nplike.allclose(grid.y.cntr, icdata["y"]))
except ValueError: # assume thie excpetion means a shape mismatch
interp = True
# unfortunately, we need to do interpolation in such a situation
if interp:
_logger.warning("Grids do not match. Doing spline interpolation.")
w = _nplike.array(
_interpolate(icdata["x"], icdata["y"], icdata[ic_config.keys[0]].T,
grid.x.cntr, grid.y.cntr).T)
hu = _nplike.array(
_interpolate(icdata["x"], icdata["y"], icdata[ic_config.keys[1]].T,
grid.x.cntr, grid.y.cntr).T)
hv = _nplike.array(
_interpolate(icdata["x"], icdata["y"], icdata[ic_config.keys[2]].T,
grid.x.cntr, grid.y.cntr).T)
else:
w = icdata[ic_config.keys[0]]
hu = icdata[ic_config.keys[1]]
hv = icdata[ic_config.keys[2]]
# make sure the w can not be smaller than topopgraphy elevation
w = _nplike.maximum(w, topo.cntr)
return _WHUHVModel(nx=grid.x.n,
ny=grid.y.n,
dtype=dtype,
w=w,
hu=hu,
hv=hv)
def get_topography(topofile: Union[str, os.PathLike], key: str,
grid_xv: nplike.ndarray, grid_yv: nplike.ndarray,
dtype: str):
"""Get a Topography object from a config object.
Arguments
---------
topofile : str or PathLike
key : str
grid_xv, grid_yv : nplike.ndarray
dtype : str, nplike.float32, nplike.float64
Returns
-------
topo : Topography
"""
dtype = DummyDtype.validator(dtype)
assert dtype == grid_xv.dtype
assert dtype == grid_yv.dtype
dem, _ = ncread(topofile, [key])
vert = dem[key]
# see if we need to do interpolation
try:
interp = not (nplike.allclose(grid_xv, dem["x"])
and nplike.allclose(grid_yv, dem["y"]))
except ValueError: # assume thie excpetion means a shape mismatch
interp = True
# unfortunately, we need to do interpolation in such a situation
if interp:
logger.warning("Grids do not match. Doing spline interpolation.")
vert = nplike.array(
_interpolate(dem["x"], dem["y"], vert.T, grid_xv, grid_yv).T)
# cast to desired float type
vert = vert.astype(dtype)
# topography elevation at cell centers through linear interpolation
cntr = vert[:-1, :-1] + vert[:-1, 1:] + vert[1:, :-1] + vert[1:, 1:]
cntr /= 4
# topography elevation at cell faces' midpoints through linear interpolation
xface = (vert[:-1, :] + vert[1:, :]) / 2.
yface = (vert[:, :-1] + vert[:, 1:]) / 2.
# gradient at cell centers through central difference; here allows nonuniform grids
# this function does not assume constant cell sizes
xgrad = (xface[:, 1:] - xface[:, :-1]) / (grid_xv[1:] -
grid_xv[:-1])[None, :]
ygrad = (yface[1:, :] - yface[:-1, :]) / (grid_yv[1:] - grid_yv[:-1])[:,
None]
# initialize DataModel and let pydantic validates data
return Topography(nx=len(grid_xv) - 1,
ny=len(grid_yv) - 1,
dtype=dtype,
vert=vert,
cntr=cntr,
xface=xface,
yface=yface,
xgrad=xgrad,
ygrad=ygrad)
def create_ic(ic_config, gridlines, topo, dtype):
"""Create initial conditions.
When the x_cntr and y_cntr have different resolutions from the x and y in the NetCDF file, an
bi-cubic spline interpolation will take place.
Arguments
---------
ic_config : torchswe.utils.config.ICConfig
gridlines : torchswe.utils.data.Gridlines
topo : torchswe.utils.data.Topography
dtype : str; either "float32" or "float64"
Returns
-------
torchswe.utils.data.WHUHVModel
"""
# special case: constant I.C.
if ic_config.values is not None:
return WHUHVModel(gridlines.x.n,
gridlines.y.n,
dtype,
w=nplike.maximum(topo.cntr,
nplike.array(ic_config.values[0])),
hu=nplike.full(topo.cntr.shape,
ic_config.values[1],
dtype=topo.dtype),
hv=nplike.full(topo.cntr.shape,
ic_config.values[2],
dtype=topo.dtype))
# otherwise, read data from a NetCDF file
icdata, _ = read_cf(ic_config.file, ic_config.keys)
# see if we need to do interpolation
try:
interp = not (
nplike.allclose(gridlines.x.cntr, nplike.array(icdata["x"]))
and nplike.allclose(gridlines.y.cntr, nplike.array(icdata["y"])))
except ValueError: # assume thie excpetion means a shape mismatch
interp = True
# unfortunately, we need to do interpolation in such a situation
if interp:
interpolator = RectBivariateSpline(icdata["x"], icdata["y"],
icdata[ic_config.keys[0]][:].T)
w = interpolator(gridlines.x.cntr, gridlines.y.cntr).T
# get an interpolator for conserv_q_ic[1], use the default 3rd order spline
interpolator = RectBivariateSpline(icdata["x"], icdata["y"],
icdata[ic_config.keys[1]][:].T)
hu = interpolator(gridlines.x.cntr, gridlines.y.cntr).T
# get an interpolator for conserv_q_ic[2], use the default 3rd order spline
interpolator = RectBivariateSpline(icdata["x"], icdata["y"],
icdata[ic_config.keys[2]][:].T)
hv = interpolator(gridlines.x.cntr, gridlines.y.cntr).T
else:
w = nplike.array(icdata[ic_config.keys[0]][:].copy())
hu = nplike.array(icdata[ic_config.keys[1]][:].copy())
hv = nplike.array(icdata[ic_config.keys[2]][:].copy())
# make sure the w can not be smaller than topopgraphy elevation
w = nplike.maximum(w, topo.cntr)
return WHUHVModel(gridlines.x.n, gridlines.y.n, dtype, w=w, hu=hu, hv=hv)