"""SW1L equations over a sphere (:mod:`fluidsim.solvers.sphere.sw1l`)
=====================================================================
.. autoclass:: InfoSolverSphereSW1L
:members:
:private-members:
.. autoclass:: SimulSphereSW1L
:members:
:private-members:
"""
from transonic import boost
from fluidsim.base.sphericalharmo.solver import (
InfoSolverSphericalHarmo,
SimulSphericalHarmo,
)
from fluidsim.base.setofvariables import SetOfVariables
Af = "float64[:, :]"
@boost
def compute_Frot(rot: Af, ux: Af, uy: Af, f_radial: Af):
"""Compute cross-product of absolute potential vorticity with velocity."""
rot_abs = rot + f_radial
F1x = rot_abs * uy
F1y = -rot_abs * ux
return F1x, F1y
[docs]
class InfoSolverSphereSW1L(InfoSolverSphericalHarmo):
"""Contain the information on the ``sphere.sw1l`` solver."""
[docs]
def _init_root(self):
"""Init. `self` by writting the information on the solver."""
super()._init_root()
here = "fluidsim.solvers.sphere.sw1l"
self.module_name = here + ".solver"
self.class_name = "SimulSphereSW1L"
self.short_name = "sphere.sw1l"
self.classes.State.module_name = here + ".state"
self.classes.State.class_name = "StateSphericalHarmoSW1L"
self.classes.Output.module_name = here + ".output"
self.classes.Output.class_name = "Output"
[docs]
class SimulSphereSW1L(SimulSphericalHarmo):
"""Spherical-harmonics solver for shallow water equation."""
InfoSolver = InfoSolverSphereSW1L
[docs]
@staticmethod
def _complete_params_with_default(params):
"""Complete the `params` container (static method)."""
SimulSphericalHarmo._complete_params_with_default(params)
attribs = {"c2": 20}
params._set_attribs(attribs)
[docs]
def tendencies_nonlin(self, state_spect=None, old=None):
r"""Compute the nonlinear tendencies.
Parameters
----------
state_spect : :class:`fluidsim.base.setofvariables.SetOfVariables`
optional
Array containing the state, i.e. the vorticity, the divergence
and displacement in the spectral space. If ``state_spect`` is
provided, the variables in the physical space are computed from
it, otherwise, they are taken from the global state of the
simulation, ``self.state.state_phys``.
These two possibilities are used during the time-stepping.
old : :class:`fluidsim.base.setofvariables.SetOfVariables`
optional
Array containing the previous ``tendencies_sh``. This array is
reused to save memory and improve performance.
Returns
-------
tendencies_sh : :class:`fluidsim.base.setofvariables.SetOfVariables`
An array containing the tendencies for the vorticity.
Notes
-----
.. |p| mathmacro:: \partial
The 1-layer shallow water equations are solved in the
vector-invariant form (see section 2.2.6 Vallis 2nd edition).
.. math::
\p_t \hat\zeta = \hat N_\zeta - \sigma(k) \hat \zeta,
\p_t \hat\delta = \hat N_\delta - \sigma(k) \hat \delta,
\p_t \hat\eta = \hat N_\eta - \sigma(k) \hat \eta,
This function computes the nonlinear term ("tendencies"). The algorithm is
as follows,
- Compute :math:`N_u` and :math:`N_v`, the tendencies for the velocities.
- Take divergence and curl of the above to obtain
:math:`N_\zeta, N_\delta`.
- Subtract laplacian of total energy K.E. + hydrostatic pressure from
:math:`N_\delta`.
- Compute :math:`N_\eta = -\nabla.((\eta + 1)\mathbf{u})`
"""
# the spherical harmonics operator
oper = self.oper
if state_spect is None:
state_spect = self.state.state_spect
state_phys = self.state.state_phys
else:
state_phys = self.state.return_statephys_from_statespect(state_spect)
# get or compute rot_sh, ux and uy
ux = state_phys.get_var("ux")
uy = state_phys.get_var("uy")
eta = state_phys.get_var("eta")
rot = state_phys.get_var("rot")
if old is None:
tendencies_sh = SetOfVariables(like=self.state.state_spect)
else:
tendencies_sh = old
Frot_sh = tendencies_sh.get_var("rot_sh")
Fdiv_sh = tendencies_sh.get_var("div_sh")
Feta_sh = tendencies_sh.get_var("eta_sh")
c2 = self.params.c2
# Absolute vorticity x Velocity
Fux, Fuy = compute_Frot(rot, ux, uy, oper.f_radial)
oper.divrotsh_from_vec(Fux, Fuy, Fdiv_sh, Frot_sh)
# Subtract laplacian of K.E. + hydrostatic pressure term from
# divergence tendency
Fdiv_sh += oper.laplacian_sh(
oper.sht(0.5 * (ux**2 + uy**2) + c2 * eta), negative=True
)
# Calculate Feta_sh = \nabla.(hu) = \nabla.((1 + \eta)u)
oper.divrotsh_from_vec(-(eta + 1) * ux, -(eta + 1) * uy, Feta_sh)
# oper.dealiasing(tendencies_sh)
# def check_conservation(Fvar_sh, var_sh, var_str):
# import numpy as np
# T = np.real(Fvar_sh.conj() * var_sh + Fvar_sh * var_sh.conj()) / 2.0
# print(
# ("sum(T_{2}) = {0:9.4e} ; sum(abs(T_{2})) = {1:9.4e}").format(
# self.oper.sum_wavenumbers(T),
# self.oper.sum_wavenumbers(abs(T)),
# var_str
# ),
# end =" | "
# )
# rot_sh = state_spect.get_var("rot_sh")
# div_sh = state_spect.get_var("div_sh")
# check_conservation(Frot_sh, rot_sh, "rot")
# check_conservation(Fdiv_sh, div_sh, "div")
# print()
if self.params.forcing.enable:
tendencies_sh += self.forcing.get_forcing()
return tendencies_sh
Simul = SimulSphereSW1L
if __name__ == "__main__":
params = Simul.create_default_params()
params.short_name_type_run = "test"
params.output.sub_directory = "test"
params.nu_2 = 1e-2
params.oper.radius = 10.0
params.oper.omega = 0.0
params.forcing.enable = False
params.init_fields.type = "noise"
params.time_stepping.deltat0 = 1e-3
params.time_stepping.USE_CFL = False
params.time_stepping.t_end = 10.0
# params.time_stepping.deltat0 = 0.1
params.output.periods_save.phys_fields = 0.25
sim = Simul(params)
sim.time_stepping.start()