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 minmod_slope_x_one_comp(q: nplike.ndarray, dx: float, ngh: int,
theta: float, tol: float):
"""Minmod slope in x direction for only one conservative quantity.
Arguments
---------
q : nplike.ndarray
(ny+2*ngh, nx+2*ngh) array of a conservative quantity.
dx : float
Cell size in x-direction. Assume an uniform grid.
ngh : int
Number of ghost cell layers outside each boundary.
theta: float
Parameter to controll oscillation and dispassion. 1 <= theta <= 2.
tol : float
To control how small can be treat as zero.
Returns
-------
slpx : nplike.ndarray
(ny, nx+2) array of the slopes in x-direction, including one ghost layer at west and east.
"""
# pylint: disable=invalid-name
cells = slice(ngh, -ngh) # non-ghost cells, length ny or nx
i = slice(ngh - 1, q.shape[1] - ngh +
1) # i = one ghost outside each bound; length nx+2
ip1 = slice(ngh, q.shape[1] - ngh + 2) # i + 1; length nx+2
im1 = slice(ngh - 2, q.shape[1] - ngh) # i - 1; length nx+2
denominator = q[cells, ip1] - q[cells,
i] # q_{j, i+1} - q_{j, i} for all j
zeros = nplike.nonzero(
nplike.logical_and(denominator > -tol, denominator < tol))
with nplike.errstate(divide="ignore", invalid="ignore"):
slpx = (q[cells, i] - q[cells, im1]) / denominator
slpx[zeros] = 0. # where q_{j, i+1} - q_{j, i} = 0
slpx = nplike.maximum(
nplike.minimum(nplike.minimum(theta * slpx, (1. + slpx) / 2.),
nplike.array(theta)), nplike.array(0.))
slpx *= denominator
slpx /= dx
return slpx
def get_local_speed(states: States, gravity: float) -> States:
"""Calculate local speeds on the two sides of cell faces.
Arguments
---------
states : torchswe.utils.data.States
gravity : float
Gravity in m / s^2.
Returns
-------
states : torchswe.utils.data.States
The same object as the input. Changed inplace. Returning it just for coding style.
"""
# faces normal to x-direction
sqrt_gh_plus = nplike.sqrt(gravity * states.face.x.plus.h)
sqrt_gh_minus = nplike.sqrt(gravity * states.face.x.minus.h)
# for convenience
zero = nplike.array(0.)
states.face.x.plus.a = nplike.maximum(
nplike.maximum(states.face.x.plus.u + sqrt_gh_plus,
states.face.x.minus.u + sqrt_gh_minus), zero)
states.face.x.minus.a = nplike.minimum(
nplike.minimum(states.face.x.plus.u - sqrt_gh_plus,
states.face.x.minus.u - sqrt_gh_minus), zero)
# faces normal to y-direction
sqrt_gh_plus = nplike.sqrt(gravity * states.face.y.plus.h)
sqrt_gh_minus = nplike.sqrt(gravity * states.face.y.minus.h)
states.face.y.plus.a = nplike.maximum(
nplike.maximum(states.face.y.plus.v + sqrt_gh_plus,
states.face.y.minus.v + sqrt_gh_minus), zero)
states.face.y.minus.a = nplike.minimum(
nplike.minimum(states.face.y.plus.v - sqrt_gh_plus,
states.face.y.minus.v - sqrt_gh_minus), zero)
return states
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 get_uv(h, hu, hv):
# pylint: disable=invalid-name
h4 = nplike.power(h, 4)
coeff = h * sqrt2 / nplike.sqrt(
h4 + nplike.maximum(h4, nplike.array(epsilon)))
return coeff * hu, coeff * hv
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)