"""State for the SW1L solver (:mod:`fluidsim.solvers.sw1l.state`)
=================================================================
.. autoclass:: StateSW1L
:members:
:private-members:
"""
import numpy as np
from fluidsim.base.setofvariables import SetOfVariables
from fluidsim.base.state import StatePseudoSpectral
from fluiddyn.util import mpi
[docs]
class StateSW1L(StatePseudoSpectral):
"""State for the solver sw1l.
Contains the variables corresponding to the state and handles the
access to other fields for the solver SW1L.
.. inheritance-diagram:: StateSW1L
"""
[docs]
@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_spect": ["ux_fft", "uy_fft", "eta_fft"],
"keys_state_phys": ["ux", "uy", "eta", "rot"],
"keys_computable": [],
"keys_phys_needed": ["ux", "uy", "eta"],
"keys_linear_eigenmodes": ["q_fft", "a_fft", "d_fft"],
}
)
[docs]
def compute(self, key, SAVE_IN_DICT=True, RAISE_ERROR=True):
"""Compute and return a variable."""
it = self.sim.time_stepping.it
if key in self.vars_computed and it == self.it_computed[key]:
return self.vars_computed[key]
if key == "Jx":
ux = self.state_phys.get_var("ux")
eta = self.state_phys.get_var("eta")
h = 1 + eta
result = h * ux
elif key == "Jy":
uy = self.state_phys.get_var("uy")
eta = self.state_phys.get_var("eta")
h = 1 + eta
result = h * uy
elif key == "Jx_fft":
Jx = self.compute("Jx")
result = self.oper.fft2(Jx)
elif key == "Jy_fft":
Jy = self.compute("Jy")
result = self.oper.fft2(Jy)
elif key == "rot_fft":
ux_fft = self.state_spect.get_var("ux_fft")
uy_fft = self.state_spect.get_var("uy_fft")
result = self.oper.rotfft_from_vecfft(ux_fft, uy_fft)
elif key == "div_fft":
ux_fft = self.state_spect.get_var("ux_fft")
uy_fft = self.state_spect.get_var("uy_fft")
result = self.oper.divfft_from_vecfft(ux_fft, uy_fft)
elif key == "div":
div_fft = self.compute("div_fft")
result = self.oper.ifft2(div_fft)
elif key == "q":
rot = self.state_phys.get_var("rot")
eta = self.state_phys.get_var("eta")
result = rot - self.params.f * eta
elif key == "q_fft":
ux_fft = self.state_spect.get_var("ux_fft")
uy_fft = self.state_spect.get_var("uy_fft")
eta_fft = self.state_spect.get_var("eta_fft")
rot_fft = self.oper.rotfft_from_vecfft(ux_fft, uy_fft)
result = rot_fft - self.params.f * eta_fft
elif key == "a_fft":
ux_fft = self.state_spect.get_var("ux_fft")
uy_fft = self.state_spect.get_var("uy_fft")
eta_fft = self.state_spect.get_var("eta_fft")
result = self.oper.afft_from_uxuyetafft(ux_fft, uy_fft, eta_fft)
elif key == "h":
eta = self.state_phys.get_var("eta")
result = 1 + eta
elif key == "Floc":
h = self.compute("h")
ux = self.state_phys.get_var("ux")
uy = self.state_phys.get_var("uy")
result = np.sqrt((ux**2 + uy**2) / (self.sim.params.c2 * h))
else:
to_print = 'Do not know how to compute "' + key + '".'
if RAISE_ERROR:
raise ValueError(to_print)
else:
if mpi.rank == 0:
print(to_print + "\nreturn an array of zeros.")
result = self.oper.create_arrayX(value=0.0)
if SAVE_IN_DICT:
self.vars_computed[key] = result
self.it_computed[key] = it
return result
[docs]
def statespect_from_statephys(self):
"""Compute the state in Fourier space."""
ux = self.state_phys.get_var("ux")
uy = self.state_phys.get_var("uy")
eta = self.state_phys.get_var("eta")
ux_fft = self.state_spect.get_var("ux_fft")
uy_fft = self.state_spect.get_var("uy_fft")
eta_fft = self.state_spect.get_var("eta_fft")
self.oper.fft_as_arg(ux, ux_fft)
self.oper.fft_as_arg(uy, uy_fft)
self.oper.fft_as_arg(eta, eta_fft)
[docs]
def statephys_from_statespect(self, state_spect=None, state_phys=None):
"""Compute the state in physical space."""
ifft_as_arg = self.oper.ifft_as_arg
if state_spect is None:
state_spect = self.state_spect
if state_phys is None:
state_phys = self.state_phys
ux_fft = state_spect.get_var("ux_fft")
uy_fft = state_spect.get_var("uy_fft")
eta_fft = state_spect.get_var("eta_fft")
rot_fft = self.oper.rotfft_from_vecfft(ux_fft, uy_fft)
ux = state_phys.get_var("ux")
uy = state_phys.get_var("uy")
eta = state_phys.get_var("eta")
rot = state_phys.get_var("rot")
ifft_as_arg(ux_fft, ux)
ifft_as_arg(uy_fft, uy)
ifft_as_arg(eta_fft, eta)
ifft_as_arg(rot_fft, rot)
[docs]
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
[docs]
def init_statespect_from(self, **kwargs):
"""Initializes *state_spect* using arrays provided as keyword
arguments.
"""
if len(kwargs) == 1:
key_fft, value = list(kwargs.items())[0]
try:
init_from_keyfft = self.__getattribute__(
"init_from_" + key_fft.replace("_", "")
)
init_from_keyfft(value)
except AttributeError:
super().init_statespect_from(**kwargs)
elif len(kwargs) == 2:
if "q_fft" in kwargs and "a_fft" in kwargs:
self.init_from_qafft(**kwargs)
else:
super().init_statespect_from(**kwargs)
[docs]
def init_from_etafft(self, eta_fft):
r"""Initialize from :math:`\hat \eta` and set velocities to zero."""
state_spect = self.state_spect
state_spect.set_var("ux_fft", np.zeros_like(eta_fft))
state_spect.set_var("uy_fft", np.zeros_like(eta_fft))
state_spect.set_var("eta_fft", eta_fft)
self.oper.dealiasing(state_spect)
self.statephys_from_statespect()
[docs]
def init_from_uxuyetafft(self, ux_fft, uy_fft, eta_fft):
"""Self explanatory."""
state_spect = self.state_spect
state_spect.set_var("ux_fft", ux_fft)
state_spect.set_var("uy_fft", uy_fft)
state_spect.set_var("eta_fft", eta_fft)
self.oper.dealiasing(state_spect)
self.statephys_from_statespect()
[docs]
def init_from_rotuxuyfft(self, rot, ux_fft, uy_fft):
"""Initializes from velocities and discards vorticity."""
self.init_from_uxuyfft(ux_fft, uy_fft)
[docs]
def init_from_rotfft(self, rot_fft):
"""Initializes with rotational velocities computed from vorticity."""
ux_fft, uy_fft = self.oper.vecfft_from_rotfft(rot_fft)
self.init_from_uxuyfft(ux_fft, uy_fft)
[docs]
def init_from_qfft(self, q_fft):
"""Initialize from potential vorticity."""
rot_fft = self.oper.rotfft_from_qfft(q_fft)
ux_fft, uy_fft = self.oper.vecfft_from_rotfft(rot_fft)
eta_fft = self.oper.etafft_from_qfft(q_fft)
self.init_from_uxuyetafft(ux_fft, uy_fft, eta_fft)
[docs]
def init_from_afft(self, a_fft):
"""Initialize from ageostrophic variable."""
ux_fft, uy_fft, eta_fft = self.oper.uxuyetafft_from_afft(a_fft)
self.init_from_uxuyetafft(ux_fft, uy_fft, eta_fft)
[docs]
def init_from_qafft(self, q_fft, a_fft):
"""Initialize from potential vorticity and ageostrophic variables."""
rot_fft = self.oper.rotfft_from_qfft(q_fft)
uxq_fft, uyq_fft = self.oper.vecfft_from_rotfft(rot_fft)
etaq_fft = self.oper.etafft_from_qfft(q_fft)
uxa_fft, uya_fft, etaa_fft = self.oper.uxuyetafft_from_afft(a_fft)
self.init_from_uxuyetafft(
uxq_fft + uxa_fft, uyq_fft + uxa_fft, etaq_fft + etaa_fft
)
[docs]
def init_from_uxuyfft(self, ux_fft, uy_fft):
"""Initialize from velocities and adjust surface displacement by solving
a Poisson equation, assuming no mean flow.
"""
oper = self.oper
ifft2 = oper.ifft2
oper.projection_perp(ux_fft, uy_fft)
oper.dealiasing(ux_fft, uy_fft)
ux = ifft2(ux_fft)
uy = ifft2(uy_fft)
rot_fft = oper.rotfft_from_vecfft(ux_fft, uy_fft)
rot = ifft2(rot_fft)
eta_fft = self._etafft_no_div(ux, uy, rot)
eta = ifft2(eta_fft)
state_spect = self.state_spect
state_spect.set_var("ux_fft", ux_fft)
state_spect.set_var("uy_fft", uy_fft)
state_spect.set_var("eta_fft", eta_fft)
state_phys = self.state_phys
state_phys.set_var("rot", rot)
state_phys.set_var("ux", ux)
state_phys.set_var("uy", uy)
state_phys.set_var("eta", eta)
def _etafft_no_div(self, ux, uy, rot):
K2_not0 = self.oper.K2_not0
rot_abs = rot + self.params.f
tempx_fft = -self.oper.fft2(rot_abs * uy)
tempy_fft = self.oper.fft2(rot_abs * ux)
uu2_fft = self.oper.fft2(ux**2 + uy**2)
eta_fft = (
1.0j * self.oper.KX * tempx_fft / K2_not0
+ 1.0j * self.oper.KY * tempy_fft / K2_not0
- uu2_fft / 2.0
) / self.params.c2
if mpi.rank == 0:
eta_fft[0, 0] = 0.0
self.oper.dealiasing(eta_fft)
return eta_fft