def get_empty_huvmodel(nx: int, ny: int, dtype: str):
"""Get an empty (i.e., zero arrays) HUVModel.
Arguments
---------
nx, ny : int
dtype : str, nplike.float32, nplike.float64
Returns
-------
A HUVModel with zero arrays.
"""
dtype = DummyDtype.validator(dtype)
h = nplike.zeros((ny, nx), dtype=dtype)
u = nplike.zeros((ny, nx), dtype=dtype)
v = nplike.zeros((ny, nx), dtype=dtype)
return WHUHVModel(nx=nx, ny=ny, dtype=dtype, h=h, u=u, v=v)
def get_empty_whuhvmodel(nx: int, ny: int, dtype: str):
"""Get an empty (i.e., zero arrays) WHUHVModel.
Arguments
---------
nx, ny : int
dtype : str, nplike.float32, nplike.float64
Returns
-------
A WHUHVModel with zero arrays.
"""
dtype = DummyDtype.validator(dtype)
w = nplike.zeros((ny, nx), dtype=dtype)
hu = nplike.zeros((ny, nx), dtype=dtype)
hv = nplike.zeros((ny, nx), dtype=dtype)
return WHUHVModel(nx=nx, ny=ny, dtype=dtype, w=w, hu=hu, hv=hv)
def get_gridline(direction: str, n: int, start: float, end: float, dtype: str):
"""Get a Gridline object.
Arguments
---------
direction : str
Either "x" or "y".
n : int
Number of cells.
start, end : float
Lower and upper bound of this axis.
dtype : str, nplike.float32, or nplike.float64
Returns
-------
gridline : Gridline
"""
dtype = DummyDtype.validator(dtype)
delta = (end - start) / n
vert = nplike.linspace(start, end, n + 1, dtype=dtype)
cntr = nplike.linspace(start + delta / 2.,
end - delta / 2.,
n,
dtype=dtype)
if direction == "x":
xface = copy.deepcopy(vert)
yface = copy.deepcopy(cntr)
else: # if this is not "y", pydantic will let me know
xface = copy.deepcopy(cntr)
yface = copy.deepcopy(vert)
# pydantic will validate the data here
return Gridline(direction=direction,
n=n,
start=start,
end=end,
delta=delta,
dtype=dtype,
vert=vert,
cntr=cntr,
xface=xface,
yface=yface)
def get_empty_facetwosidemodel(nx: int, ny: int, dtype: str):
"""Get an empty (i.e., zero arrays) FaceTwoSideModel.
Arguments
---------
nx, ny : int
dtype : str, nplike.float32, nplike.float64
Returns
-------
A FaceTwoSideModel with zero arrays.
"""
dtype = DummyDtype.validator(dtype)
return FaceTwoSideModel(nx=nx,
ny=ny,
dtype=dtype,
plus=get_empty_faceonesidemodel(nx, ny, dtype),
minus=get_empty_faceonesidemodel(nx, ny, dtype),
num_flux=get_empty_whuhvmodel(nx, ny, dtype))
def get_empty_slopes(nx: int, ny: int, dtype: str):
"""Get an empty (i.e., zero arrays) Slopes.
Arguments
---------
nx, ny : int
dtype : str, nplike.float32, nplike.float64
Returns
-------
A Slopes with zero arrays.
"""
dtype = DummyDtype.validator(dtype)
return Slopes(
nx=nx,
ny=ny,
dtype=dtype,
x=get_empty_whuhvmodel(nx + 2, ny, dtype),
y=get_empty_whuhvmodel(nx, ny + 2, dtype),
)
def get_empty_facequantitymodel(nx: int, ny: int, dtype: str):
"""Get an empty (i.e., zero arrays) FaceQuantityModel.
Arguments
---------
nx, ny : int
dtype : str, nplike.float32, nplike.float64
Returns
-------
A FaceQuantityModel with zero arrays.
"""
dtype = DummyDtype.validator(dtype)
return FaceQuantityModel(
nx=nx,
ny=ny,
dtype=dtype,
x=get_empty_facetwosidemodel(nx + 1, ny, dtype),
y=get_empty_facetwosidemodel(nx, ny + 1, dtype),
)
def get_empty_states(nx: int, ny: int, ngh: int, dtype: str):
"""Get an empty (i.e., zero arrays) States.
Arguments
---------
nx, ny : int
ngh : int
dtype : str, nplike.float32, nplike.float64
Returns
-------
A States with zero arrays.
"""
dtype = DummyDtype.validator(dtype)
return States(nx=nx,
ny=ny,
ngh=ngh,
dtype=dtype,
q=get_empty_whuhvmodel(nx + 2 * ngh, ny + 2 * ngh, dtype),
src=get_empty_whuhvmodel(nx, ny, dtype),
slp=get_empty_slopes(nx, ny, dtype),
rhs=get_empty_whuhvmodel(nx, ny, dtype),
face=get_empty_facequantitymodel(nx, ny, dtype))
def get_empty_faceonesidemodel(nx: int, ny: int, dtype: str):
"""Get an empty (i.e., zero arrays) FaceOneSideModel.
Arguments
---------
nx, ny : int
dtype : str, nplike.float32, nplike.float64
Returns
-------
A FaceOneSideModel with zero arrays.
"""
dtype = DummyDtype.validator(dtype)
return FaceOneSideModel(nx=nx,
ny=ny,
dtype=dtype,
w=nplike.zeros((ny, nx), dtype=dtype),
hu=nplike.zeros((ny, nx), dtype=dtype),
hv=nplike.zeros((ny, nx), dtype=dtype),
h=nplike.zeros((ny, nx), dtype=dtype),
u=nplike.zeros((ny, nx), dtype=dtype),
v=nplike.zeros((ny, nx), dtype=dtype),
a=nplike.zeros((ny, nx), dtype=dtype),
flux=get_empty_whuhvmodel(nx, ny, dtype))
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_ghost_cell_updaters(bcs, nx, ny, ngh, dtype, topo=None): # pylint: disable=invalid-name
"""Get a function that updates all ghost cells.
This is a function factory. The return of this funciton is a function with signature:
torchswe.utils.data.States = func(torchswe.utils.data.States)
The update happens in-place, so the return of this function is not important. We return it
just to comform the coding style.
Arguments
---------
bcs : torchswe.utils.config.BCConfig
The configuration instance of boundary conditions.
nx, ny : int
Numbers of non-ghost cells along x and y directions.
ngh : int
Number of ghost cell layers outside each boundary.
dtype : str, nplike.float32, or nplike.float64
Floating number precision.
topo : torchswe.tuils.data.Topography
Topography instance. Some boundary conditions require topography elevations.
Returns
-------
A callable with signature `torchswe.utils.data.States = func(torchswe.utils.data.States)`.
"""
bcs.check()
dtype = _DummyDtype.validator(dtype)
nngh = {"west": ny, "east": ny, "south": nx, "north": nx}
funcs = {"w": {}, "hu": {}, "hv": {}}
for i, key in enumerate(["w", "hu", "hv"]):
for ornt in ["west", "east", "south", "north"]:
# periodic BC
if bcs[ornt].types[i] == "periodic":
funcs[key][ornt] = periodic_factory(ngh, ornt)
# constant extrapolation BC (outflow)
elif bcs[ornt].types[i] == "outflow":
funcs[key][ornt] = outflow_factory(ngh, ornt)
# linear extrapolation BC
elif bcs[ornt].types[i] == "extrap":
funcs[key][ornt] = linear_extrap_factory(nngh[ornt], ngh, ornt, dtype)
# constant, i.e., Dirichlet
elif bcs[ornt].types[i] == "const":
funcs[key][ornt] = constant_bc_factory(
bcs[ornt].values[i], nngh[ornt], ngh, ornt, dtype)
# inflow, i.e., constant non-conservative variables
elif bcs[ornt].types[i] == "inflow":
topo.check()
funcs[key][ornt] = inflow_bc_factory(
bcs[ornt].values[i], nngh[ornt], ngh, ornt, topo, i, dtype)
# this shouldn't happen because pydantic should have catched the error
else:
raise ValueError("{} is not recognized.".format(bcs[ornt].types[i]))
def updater(soln):
for key in ["w", "hu", "hv"]:
for ornt in ["west", "east", "south", "north"]:
soln.q[key] = funcs[key][ornt](soln.q[key])
return soln
# store the functions as an attribute for debug
updater.funcs = funcs
return updater