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