def __init__(self, sim, oper=None):
self.sim = sim
self.params = sim.params
if oper is None:
self.oper = sim.oper
else:
self.oper = oper
# creation of the SetOfVariables state_spect and state_phys
self.keys_state_phys = sim.info.solver.classes.State.keys_state_phys
try:
self.keys_computable = sim.info.solver.classes.State.keys_computable
except AttributeError:
self.keys_computable = []
self.state_phys = SetOfVariables(
keys=self.keys_state_phys,
shape_variable=self.oper.shapeX_loc,
dtype=np.float64,
info="state_phys",
)
self.vars_computed = {}
self.it_computed = {}
self.is_initialized = False
def tendencies_nonlin(self, variables=None, old=None):
"""Return a null SetOfVariables object."""
if old is None:
tendencies = SetOfVariables(
like=self.state.state_phys, info="tendencies_nonlin"
)
else:
tendencies = old
tendencies.initialize(value=0.0)
return tendencies
def __init__(self, sim, oper=None):
super().__init__(sim, oper)
self.keys_state_spect = sim.info.solver.classes.State.keys_state_spect
self.state_spect = SetOfVariables(
keys=self.keys_state_spect,
shape_variable=self.oper.shapeK_loc,
dtype=np.complex128,
info="state_spect",
)
def invert_to_get_solution(self, A, b):
"""Solve the linear system :math:`Ax = b`."""
state_phys = self.sim.state.state_phys
arr = spsolve(A, b).reshape(state_phys.shape)
return SetOfVariables(
input_array=arr, keys=state_phys.keys, info=state_phys.info
)
def tendencies_nonlin(self, state=None, old=None):
r"""Compute the nonlinear tendencies.
Parameters
----------
state : :class:`fluidsim.base.setofvariables.SetOfVariables`
optional
Array containing the state. If `state is not None`, the variables
are computed from it, otherwise, they are taken from the global
state of the simulation, `self.state`.
These two possibilities are used during the Runge-Kutta
time-stepping.
Returns
-------
tendencies : :class:`fluidsim.base.setofvariables.SetOfVariables`
An array containing the tendencies for the variables.
Notes
-----
The Lotka-Volterra equations can be written
.. math::
\dot X = \sigma (Y - X),
\dot Y = \rho X - Y - XZ.
\dot Z = X Y - \beta Z.
"""
p = self.params
if state is None:
state = self.state.state_phys
X = state.get_var("X")
Y = state.get_var("Y")
Z = state.get_var("Z")
if old is None:
tendencies = SetOfVariables(like=self.state.state_phys)
else:
tendencies = old
tendencies.set_var("X", p.sigma * (Y - X))
tendencies.set_var("Y", p.rho * X - Y - X * Z)
tendencies.set_var("Z", X * Y - p.beta * Z)
if self.params.forcing.enable:
# TODO: Not implemented, but would be nice to study small perturbations
# cf: Vallis 2nd edition 11.4
tendencies += self.forcing.get_forcing()
return tendencies
def return_statephys_from_statespect(self, state_spect=None):
"""Return the state in physical space as a new object separate from
``self.state_phys``.
"""
state_phys = SetOfVariables(like=self.state_phys)
self.statephys_from_statespect(state_spect, state_phys)
return state_phys
def return_statephys_from_statespect(self, state_spect=None):
"""Return the physical variables computed from the spectral variables."""
ifft = self.oper.ifft
if state_spect is None:
state_spect = self.state_spect
state_phys = SetOfVariables(like=self.state_phys)
for ik in range(self.state_spect.nvar):
state_phys[ik] = ifft(state_spect[ik])
return state_phys
def tendencies_nonlin(self, variables=None, old=None):
r"""Compute the nonlinear tendencies.
This function has to be overridden in a child class.
Returns
-------
tendencies_fft : :class:`fluidsim.base.setofvariables.SetOfVariables`
An array containing only zeros.
"""
if old is None:
tendencies = SetOfVariables(like=self.state.state_spect)
else:
tendencies = old
tendencies.initialize(value=0.0)
return tendencies
def tendencies_nonlin(self, state_spect=None, old=None):
if old is None:
tendencies_fft = SetOfVariables(like=self.state.state_spect,
info="tendencies_nonlin")
else:
tendencies_fft = old
tendencies_fft[:] = 0.0
return tendencies_fft
def tendencies_nonlin(self, state_phys=None, old=None):
"""Compute the "nonlinear" tendencies."""
if old is None:
tendencies = SetOfVariables(
like=self.state.state_phys, info="tendencies", value=0.0
)
else:
tendencies = old
if self.params.forcing.enable:
tendencies += self.forcing.tendencies
return tendencies
def tendencies_nonlin(self, state_spect=None, old=None):
"""Compute the "nonlinear" tendencies."""
oper = self.oper
if state_spect is None:
state_spect = self.state.state_spect
if old is None:
tendencies_fft = SetOfVariables(
like=self.state.state_spect, info="tendencies_nonlin"
)
else:
tendencies_fft = old
w_fft = state_spect.get_var("w_fft")
z_fft = state_spect.get_var("z_fft")
F_fft = tendencies_fft.get_var("w_fft")
tendencies_fft.set_var("z_fft", w_fft)
mamp_zz = oper.monge_ampere_from_fft(z_fft, z_fft)
chi_fft = -oper.invlaplacian_fft(oper.fft2(mamp_zz), order=4)
mamp_zchi = oper.monge_ampere_from_fft(z_fft, chi_fft)
Nw_fft = oper.fft2(mamp_zchi)
lap2z_fft = oper.laplacian_fft(z_fft, order=4)
F_fft[:] = -lap2z_fft + Nw_fft
oper.dealiasing(tendencies_fft)
# ratio = self.test_tendencies_nonlin(
# tendencies_fft, w_fft, z_fft, chi_fft)
# print('ratio:', ratio)
if self.params.forcing.enable:
tendencies_fft += self.forcing.get_forcing()
return tendencies_fft
def tendencies_nonlin(self, state_spect=None, old=None):
oper = self.oper
fft2 = oper.fft2
if state_spect is None:
state_phys = self.state.state_phys
state_spect = self.state.state_spect
else:
state_phys = self.state.return_statephys_from_statespect(state_spect)
ux = state_phys.get_var("ux")
uy = state_phys.get_var("uy")
eta = state_phys.get_var("eta")
rot = state_phys.get_var("rot")
# compute the nonlinear terms for ux, uy and eta
N1x = +rot * uy
N1y = -rot * ux
gradu2_x_fft, gradu2_y_fft = oper.gradfft_from_fft(
fft2(ux ** 2 + uy ** 2) / 2
)
Nx_fft = fft2(N1x) - gradu2_x_fft
Ny_fft = fft2(N1y) - gradu2_y_fft
jx_fft = fft2(eta * ux)
jy_fft = fft2(eta * uy)
Neta_fft = -oper.divfft_from_vecfft(jx_fft, jy_fft)
# self.verify_tendencies(state_spect, state_phys,
# Nx_fft, Ny_fft, Neta_fft)
# compute the nonlinear terms for q, ap and am
(Nq_fft, Np_fft, Nm_fft) = self.oper.qapamfft_from_uxuyetafft(
Nx_fft, Ny_fft, Neta_fft
)
# Nq_fft = self.oper.create_arrayK(value=0)
# Np_fft = self.oper.create_arrayK(value=0)
# Nm_fft = self.oper.create_arrayK(value=0)
oper.dealiasing(Nq_fft, Np_fft, Nm_fft)
if old is None:
tendencies_fft = SetOfVariables(
like=self.state.state_spect, info="tendencies_nonlin"
)
else:
tendencies_fft = old
tendencies_fft.set_var("q_fft", Nq_fft)
tendencies_fft.set_var("ap_fft", Np_fft)
tendencies_fft.set_var("am_fft", Nm_fft)
if self.params.forcing.enable:
tendencies_fft += self.forcing.get_forcing()
return tendencies_fft
def tendencies_nonlin(self, state_spect=None, old=None):
oper = self.oper
fft2 = oper.fft2
ifft2 = oper.ifft2
if state_spect is None:
state_phys = self.state.state_phys
state_spect = self.state.state_spect
else:
state_phys = self.state.return_statephys_from_statespect(
state_spect)
# compute the nonlinear terms for ux, uy and eta
ux = state_phys.get_var("ux")
uy = state_phys.get_var("uy")
eta = state_phys.get_var("eta")
rot_fft = self.state.get_var("rot_fft")
ux_rot_fft, uy_rot_fft = oper.vecfft_from_rotfft(rot_fft)
ux_rot = ifft2(ux_rot_fft)
uy_rot = ifft2(uy_rot_fft)
N1x_fft = fft2(ux_rot * ux)
N2x_fft = fft2(uy_rot * ux)
N1y_fft = fft2(ux_rot * uy)
N2y_fft = fft2(uy_rot * uy)
Nx_fft = -oper.divfft_from_vecfft(N1x_fft, N2x_fft)
Ny_fft = -oper.divfft_from_vecfft(N1y_fft, N2y_fft)
jx_fft = fft2(ux_rot * eta)
jy_fft = fft2(uy_rot * eta)
Neta_fft = -oper.divfft_from_vecfft(jx_fft, jy_fft)
# self.verify_tendencies(state_spect, state_phys, Nx_fft, Ny_fft, Neta_fft)
# compute the nonlinear terms for q, ap and am
(Nq_fft, Np_fft,
Nm_fft) = oper.qapamfft_from_uxuyetafft(Nx_fft, Ny_fft, Neta_fft)
oper.dealiasing(Nq_fft, Np_fft, Nm_fft)
if old is None:
tendencies_fft = SetOfVariables(like=self.state.state_spect,
info="tendencies_nonlin")
else:
tendencies_fft = old
tendencies_fft.set_var("q_fft", Nq_fft)
tendencies_fft.set_var("ap_fft", Np_fft)
tendencies_fft.set_var("am_fft", Nm_fft)
if self.params.forcing.enable:
tendencies_fft += self.forcing.get_forcing()
return tendencies_fft
def tendencies_nonlin(self, state=None, old=None):
r"""Compute the nonlinear tendencies.
Parameters
----------
state : :class:`fluidsim.base.setofvariables.SetOfVariables`
optional
Array containing the state. If `state is not None`, the variables
are computed from it, otherwise, they are taken from the global
state of the simulation, `self.state`.
These two possibilities are used during the Runge-Kutta
time-stepping.
Returns
-------
tendencies : :class:`fluidsim.base.setofvariables.SetOfVariables`
An array containing the tendencies for the variables.
Notes
-----
The Lotka-Volterra equations can be written
.. math::
\dot X = AX - B XY,
\dot Y = -CY + D XY.
"""
p = self.params
if state is None:
state = self.state.state_phys
X = state.get_var("X")
Y = state.get_var("Y")
if old is None:
tendencies = SetOfVariables(like=self.state.state_phys)
else:
tendencies = old
tendencies.set_var("X", p.A * X - p.B * X * Y)
tendencies.set_var("Y", -p.C * Y + p.D * X * Y)
if self.params.forcing.enable:
tendencies += self.forcing.get_forcing()
return tendencies
def return_statephys_from_statespect(self, state_spect=None):
"""Return the state in physical space."""
ifft2 = self.oper.ifft2
if state_spect is None:
state_spect = self.state_spect
ux_fft = state_spect.get_var("ux_fft")
uy_fft = state_spect.get_var("uy_fft")
eta_fft = state_spect.get_var("eta_fft")
state_phys = SetOfVariables(like=self.state_phys)
state_phys.set_var("ux", ifft2(ux_fft))
state_phys.set_var("uy", ifft2(uy_fft))
state_phys.set_var("eta", ifft2(eta_fft))
return state_phys
def return_statephys_from_statespect(self, state_spect=None):
"""Return the state in physical space."""
ifft2 = self.oper.ifft2
if state_spect is None:
state_spect = self.state_spect
q_fft = self.oper.create_arrayK(value=0)
ap_fft = state_spect.get_var("ap_fft")
am_fft = state_spect.get_var("am_fft")
(ux_fft, uy_fft, eta_fft) = self.oper.uxuyetafft_from_qapamfft(
q_fft, ap_fft, am_fft
)
state_phys = SetOfVariables(like=self.state_phys)
state_phys.set_var("ux", ifft2(ux_fft))
state_phys.set_var("uy", ifft2(uy_fft))
state_phys.set_var("eta", ifft2(eta_fft))
return state_phys
def tendencies_nonlin(self, state_spect=None, old=None):
oper = self.oper
ifft_as_arg = oper.ifft_as_arg
ifft_as_arg_destroy = oper.ifft_as_arg_destroy
fft_as_arg = oper.fft_as_arg
if state_spect is None:
spect_get_var = self.state.state_spect.get_var
else:
spect_get_var = state_spect.get_var
vx_fft = spect_get_var("vx_fft")
vy_fft = spect_get_var("vy_fft")
vz_fft = spect_get_var("vz_fft")
b_fft = spect_get_var("b_fft")
omegax_fft, omegay_fft, omegaz_fft = oper.rotfft_from_vecfft(
vx_fft, vy_fft, vz_fft
)
if self.params.f is not None:
self._modif_omegafft_with_f(omegax_fft, omegay_fft, omegaz_fft)
omegax = self.state.fields_tmp[3]
omegay = self.state.fields_tmp[4]
omegaz = self.state.fields_tmp[5]
ifft_as_arg_destroy(omegax_fft, omegax)
ifft_as_arg_destroy(omegay_fft, omegay)
ifft_as_arg_destroy(omegaz_fft, omegaz)
if state_spect is None:
vx = self.state.state_phys.get_var("vx")
vy = self.state.state_phys.get_var("vy")
vz = self.state.state_phys.get_var("vz")
else:
vx = self.state.fields_tmp[0]
vy = self.state.fields_tmp[1]
vz = self.state.fields_tmp[2]
ifft_as_arg(vx_fft, vx)
ifft_as_arg(vy_fft, vy)
ifft_as_arg(vz_fft, vz)
fx, fy, fz = vector_product(vx, vy, vz, omegax, omegay, omegaz)
if old is None:
tendencies_fft = SetOfVariables(
like=self.state.state_spect, info="tendencies_nonlin"
)
else:
tendencies_fft = old
fx_fft = tendencies_fft.get_var("vx_fft")
fy_fft = tendencies_fft.get_var("vy_fft")
fz_fft = tendencies_fft.get_var("vz_fft")
fft_as_arg(fx, fx_fft)
fft_as_arg(fy, fy_fft)
fft_as_arg(fz, fz_fft)
fz_fft += b_fft
oper.project_perpk3d(fx_fft, fy_fft, fz_fft)
if state_spect is None:
b = self.state.state_phys.get_var("b")
else:
b = self.state.fields_tmp[3]
ifft_as_arg(b_fft, b)
fb_fft = (
-oper.div_vb_fft_from_vb(vx, vy, vz, b) - self.params.N ** 2 * vz_fft
)
tendencies_fft.set_var("b_fft", fb_fft)
if self.is_forcing_enabled:
tendencies_fft += self.forcing.get_forcing()
return tendencies_fft
def tendencies_nonlin(self, state_spect=None, old=None):
r"""Compute the nonlinear tendencies of the solver ns2d.strat.
Parameters
----------
state_spect : :class:`fluidsim.base.setofvariables.SetOfVariables`
optional
Array containing the state, i.e. the vorticity, in Fourier
space. If `state_spect`, the variables vorticity and the
velocity are computed from it, otherwise, they are taken
from the global state of the simulation, `self.state`.
These two possibilities are used during the Runge-Kutta
time-stepping.
Returns
-------
tendencies_fft : :class:`fluidsim.base.setofvariables.SetOfVariables`
An array containing the tendencies for the vorticity and the
buoyancy perturbation.
Notes
-----
.. |p| mathmacro:: \partial
The 2D Navier-Stokes equation can be written
.. math:: \p_t \hat\zeta = \hat N(\zeta) - \sigma(k) \hat \zeta,
and
.. math:: \p_t \hat b = \hat N(b) - \sigma(k) \hat b
This function compute the nonlinear terms ("tendencies")
:math:`N(\zeta) = - \mathbf{u}\cdot \mathbf{\nabla} \zeta +
\mathbf{\nabla}\wedge b\mathbf{e_z} = - \mathbf{u}\cdot \mathbf{\nabla}
\zeta + \p_x b` and :math:`N(b) = - \mathbf{u}\cdot \mathbf{\nabla} b +
N^2u_y`. """
oper = self.oper
fft_as_arg = oper.fft_as_arg
ifft_as_arg = oper.ifft_as_arg
if old is None:
tendencies_fft = SetOfVariables(like=self.state.state_spect)
else:
tendencies_fft = old
f_rot_fft = tendencies_fft.get_var("rot_fft")
f_b_fft = tendencies_fft.get_var("b_fft")
if state_spect is None:
rot_fft = self.state.state_spect.get_var("rot_fft")
b_fft = self.state.state_spect.get_var("b_fft")
ux = self.state.state_phys.get_var("ux")
uy = self.state.state_phys.get_var("uy")
else:
rot_fft = state_spect.get_var("rot_fft")
b_fft = state_spect.get_var("b_fft")
ux_fft, uy_fft = oper.vecfft_from_rotfft(rot_fft)
ux = self.state.field_tmp0
uy = self.state.field_tmp1
ifft_as_arg(ux_fft, ux)
ifft_as_arg(uy_fft, uy)
px_rot_fft, py_rot_fft = oper.gradfft_from_fft(rot_fft)
px_b_fft, py_b_fft = oper.gradfft_from_fft(b_fft)
px_rot = self.state.field_tmp2
py_rot = self.state.field_tmp3
px_b = self.state.field_tmp4
py_b = self.state.field_tmp5
ifft_as_arg(px_rot_fft, px_rot)
ifft_as_arg(py_rot_fft, py_rot)
ifft_as_arg(px_b_fft, px_b)
ifft_as_arg(py_b_fft, py_b)
f_rot, f_b = tendencies_nonlin_ns2dstrat(ux, uy, px_rot, py_rot, px_b,
py_b, self.params.N)
fft_as_arg(f_b, f_b_fft)
fft_as_arg(f_rot, f_rot_fft)
f_rot_fft += px_b_fft
oper.dealiasing(tendencies_fft)
# # CHECK NON-LINEAR TRANSFER rot
# T_rot = np.real(f_rot_fft.conj()*rot_fft)
# print(('sum(T_rot) = {0:9.4e} ; sum(abs(T_rot)) = {1:9.4e}').format(
# self.oper.sum_wavenumbers(T_rot),
# self.oper.sum_wavenumbers(abs(T_rot))))
# # CHECK NON-LINEAR TRANSFER b
# T_b = np.real(f_b_fft.conj()*b_fft)
# print(('sum(T_b) = {0:9.4e} ; sum(abs(T_b)) = {1:9.4e}').format(
# self.oper.sum_wavenumbers(T_b),
# self.oper.sum_wavenumbers(abs(T_b))))
if self.params.forcing.enable:
tendencies_fft += self.forcing.get_forcing()
# CHECK ENERGY CONSERVATION
# self.check_energy_conservation(rot_fft, b_fft, f_rot_fft, f_b_fft)
return tendencies_fft
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, in Fourier
space. If `state_spect`, the variables vorticity and the
velocity are computed from it, otherwise, they are taken
from the global state of the simulation, `self.state`.
These two possibilities are used during the Runge-Kutta
time-stepping.
Returns
-------
tendencies_fft : :class:`fluidsim.base.setofvariables.SetOfVariables`
An array containing the tendencies for the vorticity.
Notes
-----
.. |p| mathmacro:: \partial
The 2D Navier-Stokes equation can be written as:
.. math:: \p_t \hat\zeta = \hat N(\zeta) - \sigma(k) \hat \zeta,
This function compute the nonlinear term ("tendencies")
:math:`N(\zeta) = - \mathbf{u}\cdot \mathbf{\nabla} \zeta`.
"""
# the operator and the fast Fourier transform
oper = self.oper
ifft_as_arg = oper.ifft_as_arg
# get or compute rot_fft, ux and uy
if state_spect is None:
rot_fft = self.state.state_spect.get_var("rot_fft")
ux = self.state.state_phys.get_var("ux")
uy = self.state.state_phys.get_var("uy")
else:
rot_fft = state_spect.get_var("rot_fft")
ux_fft, uy_fft = oper.vecfft_from_rotfft(rot_fft)
ux = self.state.field_tmp0
uy = self.state.field_tmp1
ifft_as_arg(ux_fft, ux)
ifft_as_arg(uy_fft, uy)
# "px" like $\partial_x$
px_rot_fft, py_rot_fft = oper.gradfft_from_fft(rot_fft)
px_rot = self.state.field_tmp2
py_rot = self.state.field_tmp3
ifft_as_arg(px_rot_fft, px_rot)
ifft_as_arg(py_rot_fft, py_rot)
Frot = compute_Frot(ux, uy, px_rot, py_rot, self.params.beta)
if old is None:
tendencies_fft = SetOfVariables(like=self.state.state_spect)
else:
tendencies_fft = old
Frot_fft = tendencies_fft.get_var("rot_fft")
oper.fft_as_arg(Frot, Frot_fft)
oper.dealiasing(Frot_fft)
# import numpy as np
# T_rot = np.real(Frot_fft.conj()*rot_fft + Frot_fft*rot_fft.conj())/2.
# print(('sum(T_rot) = {0:9.4e} ; sum(abs(T_rot)) = {1:9.4e}'
# ).format(self.oper.sum_wavenumbers(T_rot),
# self.oper.sum_wavenumbers(abs(T_rot))))
if self.params.forcing.enable:
tendencies_fft += self.forcing.get_forcing()
return tendencies_fft
class StateBase:
"""Contains the state variables and handles the access to fields.
Parameters
----------
sim : child of :class:`fluidsim.base.solvers.base.SimulBase`
oper : Optional[operators]
"""
@staticmethod
def _complete_info_solver(info_solver):
"""Complete the ParamContainer info_solver.
This is a static method!
"""
info_solver.classes.State._set_attribs({
"keys_state_phys": ["X"],
"keys_computable": [],
"keys_phys_needed": ["X"],
})
def __init__(self, sim, oper=None):
self.sim = sim
self.params = sim.params
if oper is None:
self.oper = sim.oper
else:
self.oper = oper
# creation of the SetOfVariables state_spect and state_phys
self.keys_state_phys = sim.info.solver.classes.State.keys_state_phys
try:
self.keys_computable = sim.info.solver.classes.State.keys_computable
except AttributeError:
self.keys_computable = []
self.state_phys = SetOfVariables(
keys=self.keys_state_phys,
shape_variable=self.oper.shapeX_loc,
dtype=np.float64,
info="state_phys",
)
self.vars_computed = {}
self.it_computed = {}
self.is_initialized = False
def compute(self, key):
"""Compute a not stored variable from the stored variables"""
raise ValueError('No method to compute key "' + key + '"')
def clear_computed(self):
"""Clear the stored computed variables."""
self.vars_computed.clear()
def has_vars(self, *keys):
"""Checks if all of the keys are present in the union of
``keys_state_phys`` and ``keys_computable``.
Parameters
----------
keys: str, str ...
Strings indicating state variable names.
Returns
-------
bool
Examples
--------
>>> sim.state.has_vars('ux', 'uy')
>>> sim.state.has_vars('ux')
>>> sim.state.has_vars('ux', 'vx', strict=False)
.. todo::
``strict=True`` can be a Python 3 compatible keywords-only argument
with the function like::
def has_vars(self, *keys, strict=True):
...
if strict:
return keys.issubset(keys_state)
else:
return len(keys.intersection(keys_state)) > 0
When ``True``, checks if all keys form a subset of state keys. When
``False``, checks if the intersection of the keys and the state keys
has atleast one member.
"""
keys_state = set(self.keys_state_phys + self.keys_computable)
keys = set(keys)
return keys.issubset(keys_state)
def get_var(self, key):
"""Get a physical variable (from the storage array or computed).
This is one of the main method of the state classes.
It tries to return the array corresponding to a physical variable. If
it is stored in the main storage array of the state class, it is
directly returned. Otherwise, we try to compute the quantity with the
method :func:`compute`.
It should not be necessary to redefine this method in child class.
"""
if key in self.keys_state_phys:
return self.state_phys.get_var(key)
else:
it = self.sim.time_stepping.it
if key in self.vars_computed and it == self.it_computed[key]:
return self.vars_computed[key]
else:
value = self.compute(key)
self.vars_computed[key] = value
self.it_computed[key] = it
return value
def __call__(self, key):
raise DeprecationWarning("Do not call a state object. "
"Instead, use get_var method.")
def __setitem__(self, key, value):
"""General setter function to set the value of a variable
It should not be necessary to redefine this method in child class.
"""
if key in self.keys_state_phys:
self.state_phys.set_var(key, value)
else:
raise ValueError('key "' + key + '" is not known')
def can_this_key_be_obtained(self, key):
"""To check whether a variable can be obtained.
.. deprecated:: 0.2.0
Use ``has_vars`` method instead.
"""
raise DeprecationWarning("Do not call can_this_key_be_obtained. "
"Instead, use has_vars method.")
def init_statephys_from(self, **kwargs):
"""Initialize `state_phys` from arrays.
Parameters
----------
**kwargs : {key: array, ...}
keys and arrays used for the initialization. The other keys
are set to zero.
Examples
--------
.. code-block:: python
kwargs = {'a': Fa}
init_statespect_from(**kwargs)
init_statespect_from(ux=ux, uy=uy, eta=eta)
"""
self.state_phys[:] = 0.0
for key, value in list(kwargs.items()):
if key not in self.keys_state_phys:
raise ValueError(
f'Do not know how to initialize with key "{key}".')
self.state_phys.set_var(key, value)
def __init__(self, sim):
super().__init__(sim)
params = sim.params
self.forcing_fft = SetOfVariables(
like=sim.state.state_spect, info="forcing_fft", value=0.0
)
if params.forcing.nkmax_forcing < params.forcing.nkmin_forcing:
raise ValueError(
"params.forcing.nkmax_forcing < \n" "params.forcing.nkmin_forcing"
)
self.kmax_forcing = self.oper.deltak * params.forcing.nkmax_forcing
self.kmin_forcing = self.oper.deltak * params.forcing.nkmin_forcing
self.forcing_rate = params.forcing.forcing_rate
if params.forcing.key_forced is not None:
self.key_forced = params.forcing.key_forced
else:
self.key_forced = self._key_forced_default
try:
n = 2 * fftw_grid_size(params.forcing.nkmax_forcing)
except ImportError:
warn("To use smaller forcing arrays: pip install pulp")
i = 0
while 2 * params.forcing.nkmax_forcing > 2 ** i:
i += 1
n = 2 ** i
self._check_forcing_shape([n], sim.oper.shapeX_seq)
try:
angle = self.params.forcing[self.tag].angle
except AttributeError:
pass
else:
if isinstance(angle, str):
if angle.endswith("°"):
angle = radians(float(angle[:-1]))
else:
raise ValueError(
"Angle should be a string with \n"
+ "the degree symbol or a float in radians"
)
self.angle = angle
self.kxmax_forcing = np.sin(angle) * self.kmax_forcing
self.kymax_forcing = np.cos(angle) * self.kmax_forcing
if mpi.rank == 0:
params_coarse = deepcopy(params)
params_coarse.oper.nx = n
# The "+ 1" aims to give some gap between the kxmax and
# the boundary of the oper_coarse.
try:
params_coarse.oper.nx = 2 * fftw_grid_size(
int(self.kxmax_forcing / self.oper.deltakx) + 1
)
except AttributeError:
pass
try:
params_coarse.oper.ny = n
try:
params_coarse.oper.ny = 2 * fftw_grid_size(
int(self.kymax_forcing / self.oper.deltaky) + 1
)
except AttributeError:
pass
except AttributeError:
pass
try:
params_coarse.oper.nz = n
except AttributeError:
pass
params_coarse.oper.type_fft = "sequential"
params_coarse.oper.coef_dealiasing = 1.0
self.oper_coarse = sim.oper.__class__(params=params_coarse)
self.shapeK_loc_coarse = self.oper_coarse.shapeK_loc
self.COND_NO_F = self._compute_cond_no_forcing()
self.nb_forced_modes = (
self.COND_NO_F.size
- np.array(self.COND_NO_F, dtype=np.int32).sum()
)
if not self.nb_forced_modes:
raise ValueError("0 modes forced.")
try:
hasattr(self, "plot_forcing_region")
except NotImplementedError:
pass
else:
mpi.printby0(
"To plot the forcing modes, you can use:\n"
"sim.forcing.forcing_maker.plot_forcing_region()"
)
self.ind_forcing = (
np.logical_not(self.COND_NO_F).flatten().nonzero()[0]
)
self.fstate_coarse = sim.state.__class__(sim, oper=self.oper_coarse)
else:
self.shapeK_loc_coarse = None
if mpi.nb_proc > 1:
self.shapeK_loc_coarse = mpi.comm.bcast(
self.shapeK_loc_coarse, root=0
)
class StatePseudoSpectral(StateBase):
"""Contains the state variables and handles the access to fields.
This is the general class for the pseudo-spectral solvers.
Parameters
----------
sim : child of :class:`fluidsim.base.solvers.base.SimulBase`
oper : Optional[operators]
"""
@staticmethod
def _complete_info_solver(info_solver):
"""Complete the ParamContainer info_solver.
This is a static method!
"""
StateBase._complete_info_solver(info_solver)
info_solver.classes.State.keys_state_phys = ["ux", "uy"]
info_solver.classes.State.keys_phys_needed = ["ux", "uy"]
info_solver.classes.State._set_attribs(
{"keys_state_spect": ["ux_fft", "uy_fft"]})
def __init__(self, sim, oper=None):
super().__init__(sim, oper)
self.keys_state_spect = sim.info.solver.classes.State.keys_state_spect
self.state_spect = SetOfVariables(
keys=self.keys_state_spect,
shape_variable=self.oper.shapeK_loc,
dtype=np.complex128,
info="state_spect",
)
def has_vars(self, *keys):
"""Checks if all of the keys are present in the union of
``keys_state_phys``, ``keys_computable``, and ``keys_state_spect``.
Parameters
----------
keys: str, str ...
Strings indicating state variable names.
Returns
-------
bool
Examples
--------
>>> sim.state.has_vars('ux', 'uy', 'ux_fft')
>>> sim.state.has_vars('rot')
"""
keys_state = set(self.keys_state_phys + self.keys_computable +
self.keys_state_spect)
keys = set(keys)
return keys.issubset(keys_state)
def get_var(self, key):
"""Get a variable (from the storage arrays or computed).
This is one of the main method of the state classes.
It tries to return the array corresponding to a physical variable. If
it is stored in the main storage arrays (`state_phys` and `state_spec`)
of the state class, it is directly returned. Otherwise, we try to
compute the quantity with the method :func:`compute`.
It should not be necessary to redefine this method in child class.
"""
if key in self.keys_state_spect:
return self.state_spect.get_var(key)
elif key in self.keys_state_phys:
return self.state_phys.get_var(key)
else:
it = self.sim.time_stepping.it
if key in self.vars_computed and it == self.it_computed[key]:
return self.vars_computed[key]
else:
value = self.compute(key)
self.vars_computed[key] = value
self.it_computed[key] = it
return value
def __setitem__(self, key, value):
"""General setter function to set the value of a variable
It should not be necessary to redefine this method in child class.
"""
if key in self.keys_state_spect:
self.state_spect.set_var(key, value)
elif key in self.keys_state_phys:
self.state_phys.set_var(key, value)
else:
raise ValueError('key "' + key + '" is not known')
def statespect_from_statephys(self):
"""Compute the spectral variables from the physical variables.
When you implement a new solver, check that this method does the job!
"""
for ik in range(self.state_spect.nvar):
self.oper.fft_as_arg(self.state_phys[ik], self.state_spect[ik])
def statephys_from_statespect(self):
"""Compute the physical variables from the spectral variables.
When you implement a new solver, check that this method does the job!
"""
for ik in range(self.state_spect.nvar):
self.oper.ifft_as_arg(self.state_spect[ik], self.state_phys[ik])
def return_statephys_from_statespect(self, state_spect=None):
"""Return the physical variables computed from the spectral variables."""
ifft = self.oper.ifft
if state_spect is None:
state_spect = self.state_spect
state_phys = SetOfVariables(like=self.state_phys)
for ik in range(self.state_spect.nvar):
state_phys[ik] = ifft(state_spect[ik])
return state_phys
def can_this_key_be_obtained(self, key):
return (key in self.keys_state_phys or key in self.keys_computable
or key in self.keys_state_spect)
def init_statespect_from(self, **kwargs):
"""Initialize `state_spect` from arrays.
Parameters
----------
**kwargs : {key: array, ...}
keys and arrays used for the initialization. The other keys
are set to zero.
Examples
--------
.. code-block:: python
kwargs = {'a_fft': Fa_fft}
init_statespect_from(**kwargs)
ux_fft, uy_fft, eta_fft = oper.uxuyetafft_from_qfft(q_fft)
init_statespect_from(ux_fft=ux_fft, uy_fft=uy_fft, eta_fft=eta_fft)
"""
self.state_spect[:] = 0.0
for key, value in list(kwargs.items()):
if key not in self.keys_state_spect:
raise ValueError(
f'Do not know how to initialize with key "{key}".')
self.state_spect.set_var(key, value)
def tendencies_nonlin(self, state_spect=None, old=None):
oper = self.oper
fft2 = oper.fft2
ifft2 = oper.ifft2
if state_spect is None:
state_phys = self.state.state_phys
state_spect = self.state.state_spect
else:
state_phys = self.state.return_statephys_from_statespect(
state_spect)
ux = state_phys.get_var("ux")
uy = state_phys.get_var("uy")
# eta = state_phys.get_var('eta')
ux_fft = state_spect.get_var("ux_fft")
uy_fft = state_spect.get_var("uy_fft")
eta_fft = state_spect.get_var("eta_fft")
# compute Fx_fft and Fy_fft
rot_fft = oper.rotfft_from_vecfft(ux_fft, uy_fft)
ux_rot_fft, uy_rot_fft = oper.vecfft_from_rotfft(rot_fft)
ux_rot = ifft2(ux_rot_fft)
uy_rot = ifft2(uy_rot_fft)
dxux_fft, dyux_fft = oper.gradfft_from_fft(ux_fft)
dxux = ifft2(dxux_fft)
dyux = ifft2(dyux_fft)
dxuy_fft, dyuy_fft = oper.gradfft_from_fft(uy_fft)
dxuy = ifft2(dxuy_fft)
dyuy = ifft2(dyuy_fft)
FNLx = -ux_rot * dxux - uy_rot * dyux
FNLy = -ux_rot * dxuy - uy_rot * dyuy
if self.params.beta != 0:
f = self.params.f + self.params.beta * oper.YY
else:
f = self.params.f
FCx = +f * uy
FCy = -f * ux
Fgradx_fft, Fgrady_fft = oper.gradfft_from_fft(self.params.c2 *
eta_fft)
Fx_fft = fft2(FCx + FNLx) - Fgradx_fft
Fy_fft = fft2(FCy + FNLy) - Fgrady_fft
# compute Feta_fft
dxeta_fft, dyeta_fft = oper.gradfft_from_fft(eta_fft)
dxeta = ifft2(dxeta_fft)
dyeta = ifft2(dyeta_fft)
div_fft = oper.divfft_from_vecfft(ux_fft, uy_fft)
Feta_fft = -fft2(ux_rot * dxeta + uy_rot * dyeta) - div_fft
if old is None:
tendencies_fft = SetOfVariables(like=self.state.state_spect,
info="tendencies_nonlin")
else:
tendencies_fft = old
tendencies_fft.set_var("ux_fft", Fx_fft)
tendencies_fft.set_var("uy_fft", Fy_fft)
tendencies_fft.set_var("eta_fft", Feta_fft)
oper.dealiasing(tendencies_fft)
if self.params.forcing.enable:
tendencies_fft += self.forcing.get_forcing()
return tendencies_fft
class SpecificForcingPseudoSpectralCoarse(SpecificForcing):
"""Specific forcing for pseudo-spectra solvers"""
tag = "pseudo_spectral"
_key_forced_default = "rot_fft"
@staticmethod
def _check_forcing_shape(shape_forcing, shape):
"""Check if shape of the forcing array exceeds the shape
of the global array.
Parameters
----------
shape_forcing: array-like
A single-element array containing index of largest forcing
wavenumber or a tuple indincating shape of the forcing array.
shape: array-like
A tuple indicating the shape of an array or Operators instance.
"""
if any(np.greater(shape_forcing, shape)):
raise NotImplementedError(
"The resolution is too small for the required forcing: "
"any(np.greater({}, {}))".format(shape_forcing, shape)
)
def __init__(self, sim):
super().__init__(sim)
params = sim.params
self.forcing_fft = SetOfVariables(
like=sim.state.state_spect, info="forcing_fft", value=0.0
)
if params.forcing.nkmax_forcing < params.forcing.nkmin_forcing:
raise ValueError(
"params.forcing.nkmax_forcing < \n" "params.forcing.nkmin_forcing"
)
self.kmax_forcing = self.oper.deltak * params.forcing.nkmax_forcing
self.kmin_forcing = self.oper.deltak * params.forcing.nkmin_forcing
self.forcing_rate = params.forcing.forcing_rate
if params.forcing.key_forced is not None:
self.key_forced = params.forcing.key_forced
else:
self.key_forced = self._key_forced_default
try:
n = 2 * fftw_grid_size(params.forcing.nkmax_forcing)
except ImportError:
warn("To use smaller forcing arrays: pip install pulp")
i = 0
while 2 * params.forcing.nkmax_forcing > 2 ** i:
i += 1
n = 2 ** i
self._check_forcing_shape([n], sim.oper.shapeX_seq)
try:
angle = self.params.forcing[self.tag].angle
except AttributeError:
pass
else:
if isinstance(angle, str):
if angle.endswith("°"):
angle = radians(float(angle[:-1]))
else:
raise ValueError(
"Angle should be a string with \n"
+ "the degree symbol or a float in radians"
)
self.angle = angle
self.kxmax_forcing = np.sin(angle) * self.kmax_forcing
self.kymax_forcing = np.cos(angle) * self.kmax_forcing
if mpi.rank == 0:
params_coarse = deepcopy(params)
params_coarse.oper.nx = n
# The "+ 1" aims to give some gap between the kxmax and
# the boundary of the oper_coarse.
try:
params_coarse.oper.nx = 2 * fftw_grid_size(
int(self.kxmax_forcing / self.oper.deltakx) + 1
)
except AttributeError:
pass
try:
params_coarse.oper.ny = n
try:
params_coarse.oper.ny = 2 * fftw_grid_size(
int(self.kymax_forcing / self.oper.deltaky) + 1
)
except AttributeError:
pass
except AttributeError:
pass
try:
params_coarse.oper.nz = n
except AttributeError:
pass
params_coarse.oper.type_fft = "sequential"
params_coarse.oper.coef_dealiasing = 1.0
self.oper_coarse = sim.oper.__class__(params=params_coarse)
self.shapeK_loc_coarse = self.oper_coarse.shapeK_loc
self.COND_NO_F = self._compute_cond_no_forcing()
self.nb_forced_modes = (
self.COND_NO_F.size
- np.array(self.COND_NO_F, dtype=np.int32).sum()
)
if not self.nb_forced_modes:
raise ValueError("0 modes forced.")
try:
hasattr(self, "plot_forcing_region")
except NotImplementedError:
pass
else:
mpi.printby0(
"To plot the forcing modes, you can use:\n"
"sim.forcing.forcing_maker.plot_forcing_region()"
)
self.ind_forcing = (
np.logical_not(self.COND_NO_F).flatten().nonzero()[0]
)
self.fstate_coarse = sim.state.__class__(sim, oper=self.oper_coarse)
else:
self.shapeK_loc_coarse = None
if mpi.nb_proc > 1:
self.shapeK_loc_coarse = mpi.comm.bcast(
self.shapeK_loc_coarse, root=0
)
def _compute_cond_no_forcing(self):
if hasattr(self.oper_coarse, "K"):
K = self.oper_coarse.K
else:
K = np.sqrt(self.oper_coarse.K2)
COND_NO_F = np.logical_or(K > self.kmax_forcing, K < self.kmin_forcing)
COND_NO_F[self.oper_coarse.shapeK_loc[0] // 2] = True
COND_NO_F[:, self.oper_coarse.shapeK_loc[1] - 1] = True
return COND_NO_F
def put_forcingc_in_forcing(self):
"""Copy data from self.fstate_coarse.state_spect into forcing_fft."""
if mpi.rank == 0:
state_spect = self.fstate_coarse.state_spect
oper_coarse = self.oper_coarse
else:
state_spect = None
oper_coarse = None
self.oper.put_coarse_array_in_array_fft(
state_spect, self.forcing_fft, oper_coarse, self.shapeK_loc_coarse
)
def verify_injection_rate(self):
"""Verify injection rate."""
f_fft = self.forcing_fft.get_var(self.key_forced)
var_fft = self.sim.state.state_spect.get_var(self.key_forced)
deltat = self.sim.time_stepping.deltat
P_forcing1 = np.real(f_fft.conj() * var_fft)
P_forcing2 = abs(f_fft) ** 2
P_forcing2 = deltat / 2 * self.oper.sum_wavenumbers(P_forcing2)
P_forcing1 = self.oper.sum_wavenumbers(P_forcing1)
if mpi.rank == 0:
print(
"P_f = {:9.4e} ; P_1 = {:9.4e}; P_2 = {:9.4e}".format(
P_forcing1 + P_forcing2, P_forcing1, P_forcing2
)
)
def verify_injection_rate_coarse(self, var_fft=None):
"""Verify injection rate."""
if var_fft is None:
var_fft = self.sim.state.state_spect.get_var(self.key_forced)
var_fft = self.oper.coarse_seq_from_fft_loc(
var_fft, self.shapeK_loc_coarse
)
if mpi.rank == 0:
f_fft = self.fstate_coarse.get_var(self.key_forced)
deltat = self.sim.time_stepping.deltat
P_forcing1 = np.real(f_fft.conj() * var_fft)
P_forcing2 = abs(f_fft) ** 2
P_forcing2 = deltat / 2 * self.oper_coarse.sum_wavenumbers(P_forcing2)
P_forcing1 = self.oper_coarse.sum_wavenumbers(P_forcing1)
print(
"P_f = {:9.4e} ; P_1 = {:9.4e}; P_2 = {:9.4e} (coarse)".format(
P_forcing1 + P_forcing2, P_forcing1, P_forcing2
)
)
def compute(self):
"""compute a forcing normalize with a 2nd degree eq."""
if mpi.rank == 0:
obj = self.compute_forcingc_fft_each_time()
if isinstance(obj, dict):
kwargs = obj
else:
kwargs = {self.key_forced: obj}
self.fstate_coarse.init_statespect_from(**kwargs)
self.oper_coarse.dealiasing_setofvar(self.fstate_coarse.state_spect)
self.put_forcingc_in_forcing()
def compute_forcingc_fft_each_time(self):
raise NotImplementedError
def _get_state_from_simul(self, sim_in):
# Warning: this function is for 2d pseudo-spectral solver!
# We have to write something more general.
# It should be done directly in the operators.
if mpi.nb_proc > 1:
raise ValueError(
"BE CARREFUL, THIS WILL BE WRONG !"
" DO NOT USE THIS METHOD WITH MPI"
)
sim = self.sim
sim.time_stepping.t = sim_in.time_stepping.t
if (
sim.params.oper.nx == sim_in.params.oper.nx
and sim.params.oper.ny == sim_in.params.oper.ny
):
state_spect = deepcopy(sim_in.state.state_spect)
else:
# modify resolution
state_spect = SetOfVariables(like=sim.state.state_spect)
keys_state_spect = sim_in.info.solver.classes.State[
"keys_state_spect"
]
for index_key in range(len(keys_state_spect)):
field_fft_seq_in = sim_in.state.state_spect[index_key]
field_fft_seq_new_res = sim.oper.create_arrayK(value=0.0)
[nk0_seq, nk1_seq] = field_fft_seq_new_res.shape
[nk0_seq_in, nk1_seq_in] = field_fft_seq_in.shape
nk0_min = min(nk0_seq, nk0_seq_in)
nk1_min = min(nk1_seq, nk1_seq_in)
# it is a little bit complicate to take into account ky
for ik1 in range(nk1_min):
field_fft_seq_new_res[0, ik1] = field_fft_seq_in[0, ik1]
field_fft_seq_new_res[nk0_min // 2, ik1] = field_fft_seq_in[
nk0_min // 2, ik1
]
for ik0 in range(1, nk0_min // 2):
for ik1 in range(nk1_min):
field_fft_seq_new_res[ik0, ik1] = field_fft_seq_in[
ik0, ik1
]
field_fft_seq_new_res[-ik0, ik1] = field_fft_seq_in[
-ik0, ik1
]
state_spect[index_key] = field_fft_seq_new_res
if sim.output.name_solver == sim_in.output.name_solver:
sim.state.state_spect = state_spect
else: # complicated case... untested solution !
raise NotImplementedError
# state_spect = SetOfVariables('state_spect')
# for k in sim.info.solver.classes.State["keys_state_spect"]:
# if k in sim_in.info.solver.classes.State["keys_state_spect"]:
# sim.state.state_spect[k] = state_spect[k]
# else:
# sim.state.state_spect[k] = self.oper.create_arrayK(value=0.0)
sim.state.statephys_from_statespect()
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
F*x, Fuy = compute_Frot(rot, ux, uy, oper.f_radial)
oper.divrotsh_from_vec(F*x, 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
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 horizontal velocities and
surface displacement scalar in Fourier space. When `state_spect` is
provided, the variables vorticity and the velocities and surface
displacement are computed from it, otherwise, they are taken from
the global state of the simulation, `self.state`.
These two possibilities are used during the Runge-Kutta
time-stepping.
Returns
-------
tendencies_fft : :class:`fluidsim.base.setofvariables.SetOfVariables`
An array containing the tendencies for the velocities
:math:`\mathbf u` and the surface displacement scalar :math:`\eta`.
Notes
-----
.. |p| mathmacro:: \partial
The one-layer shallow water equations can be written as:
.. math:: \partial_t \mathbf u = - \nabla (|\mathbf u|^2/2 + c^2 \eta)
- (\zeta + f) \times \mathbf u
.. math:: \partial_t \eta = - \nabla. ((\eta + H) \mathbf u)
This function computes all terms on the RHS of the equations above,
including the nonlinear term. The linear dissipation term (not shown
above) is computed implicitly, as demonstrated in
:class:`fluidsim.base.time_stepping.pseudo_spect.TimeSteppingPseudoSpectral`.
"""
oper = self.oper
if state_spect is None:
state_phys = self.state.state_phys
else:
state_phys = self.state.return_statephys_from_statespect(
state_spect)
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_fft = SetOfVariables(like=self.state.state_spect,
info="tendencies_nonlin")
else:
tendencies_fft = old
F1x, F1y = compute_Frot(rot, ux, uy, self.params.f)
gradx_fft, grady_fft = oper.gradfft_from_fft(
oper.fft2(compute_pressure(self.params.c2, eta, ux, uy)))
oper.dealiasing(gradx_fft, grady_fft)
Fx_fft = tendencies_fft.get_var("ux_fft")
Fy_fft = tendencies_fft.get_var("uy_fft")
Feta_fft = tendencies_fft.get_var("eta_fft")
Fx_fft[:] = oper.fft2(F1x) - gradx_fft
Fy_fft[:] = oper.fft2(F1y) - grady_fft
Feta_fft[:] = -oper.divfft_from_vecfft(oper.fft2(
(eta + 1) * ux), oper.fft2((eta + 1) * uy))
oper.dealiasing(tendencies_fft)
if self.params.forcing.enable:
tendencies_fft += self.forcing.get_forcing()
return tendencies_fft
