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 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 (Legate version).
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
# legate currently has no `logical_and`, so we use element-wise multiplication
zero_locs = (denominator > -tol).astype(int) * (denominator <
tol).astype(int)
with nplike.errstate(divide="ignore", invalid="ignore"):
slpx = nplike.where(zero_locs.astype(bool), 0.,
(q[cells, i] - q[cells, im1]) / denominator)
slpx = nplike.maximum(
nplike.minimum(nplike.minimum(theta * slpx, (1. + slpx) / 2.), theta),
0.)
slpx *= denominator
slpx /= dx
return slpx
def minmod_slope_y_one_comp(q: nplike.ndarray, dy: 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.
dy : float
Cell size in y-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
-------
slpy : nplike.ndarray
(ny+2, nx) array of the slopes in y-direction, including one ghost layer at west and east.
"""
# pylint: disable=invalid-name
cells = slice(ngh, -ngh) # non-ghost cells, length ny or nx
j = slice(ngh - 1, q.shape[0] - ngh +
1) # j = one ghost outside each bound; length ny+2
jp1 = slice(ngh, q.shape[0] - ngh + 2) # j + 1; length ny+2
jm1 = slice(ngh - 2, q.shape[0] - ngh) # j - 1; length ny+2
denominator = q[jp1, cells] - q[j,
cells] # q_{j+1, i} - q_{j, i} for all i
zeros = nplike.nonzero(
nplike.logical_and(denominator > -tol, denominator < tol))
with nplike.errstate(divide="ignore", invalid="ignore"):
slpy = (q[j, cells] - q[jm1, cells]) / denominator
slpy[zeros] = 0. # where q_{j+1, i} - q_{j, i} = 0
slpy = nplike.maximum(
nplike.minimum(nplike.minimum(theta * slpy, (1. + slpy) / 2.),
nplike.array(theta)), nplike.array(0.))
slpy *= denominator
slpy /= dy
return slpy
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_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 fvm(states: States, grid: Gridlines, topo: Topography, config: Config, runtime: DummyDict):
"""Get the right-hand-side of a time-marching step with finite volume method.
Arguments
---------
states : torchswe.utils.data.States
grid : torchswe.utils.data.Gridlines
topo : torchswe.utils.data.Topography
config : torchswe.utils.config.Config
runtime : torchswe.utils.dummy.DummyDict
Returns:
--------
states : torchswe.utils.data.States
The same object as the input. Updated in-place. Returning it just for coding style.
max_dt : float
A scalar indicating the maximum safe time-step size.
"""
# pylint: disable=invalid-name
# calculate source term contributed from topography gradients
states = topography_gradient(states, topo, config.params.gravity)
# calculate slopes of piecewise linear approximation
states = minmod_slope(states, grid, config.params.theta, runtime.tol)
# interpolate to get discontinuous conservative quantities at cell faces
states = get_discontinuous_cnsrv_q(states, grid)
# fix non-physical negative depth
states = correct_negative_depth(states, topo)
# get non-conservative variables at cell faces
states = decompose_variables(states, topo, runtime.epsilon)
# get local speed at cell faces
states = get_local_speed(states, config.params.gravity)
# get discontinuous PDE flux at cell faces
states = get_discontinuous_flux(states, topo, config.params.gravity)
# get common/continuous numerical flux at cell faces
states = central_scheme(states, runtime.tol)
# get final right hand side
states.rhs.w = \
(states.face.x.num_flux.w[:, :-1] - states.face.x.num_flux.w[:, 1:]) / grid.x.delta + \
(states.face.y.num_flux.w[:-1, :] - states.face.y.num_flux.w[1:, :]) / grid.y.delta + \
states.src.w
states.rhs.hu = \
(states.face.x.num_flux.hu[:, :-1] - states.face.x.num_flux.hu[:, 1:]) / grid.x.delta + \
(states.face.y.num_flux.hu[:-1, :] - states.face.y.num_flux.hu[1:, :]) / grid.y.delta + \
states.src.hu
states.rhs.hv = \
(states.face.x.num_flux.hv[:, :-1] - states.face.x.num_flux.hv[:, 1:]) / grid.x.delta + \
(states.face.y.num_flux.hv[:-1, :] - states.face.y.num_flux.hv[1:, :]) / grid.y.delta + \
states.src.hv
# remove rounding errors
states.rhs = remove_rounding_errors(states.rhs, runtime.tol)
# obtain the maximum safe dt
amax = nplike.max(nplike.maximum(states.face.x.plus.a, -states.face.x.minus.a))
bmax = nplike.max(nplike.maximum(states.face.y.plus.a, -states.face.y.minus.a))
max_dt = min(0.25*grid.x.delta/amax, 0.25*grid.y.delta/bmax)
return states, max_dt
def inflow_bc_velocity(conserv_q):
depth = nplike.maximum(conserv_q[w_idx] - topo_cache[bcslc], 0.)
delta = (const * depth - conserv_q[anchor]) * 2
conserv_q[slc] = conserv_q[anchor] + seq * delta
return conserv_q
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)
