r"""Shallow water 2 layers solver (:mod:`fluidsim_ocean.sw2l.solver`)
====================================================================
.. autoclass:: InfoSolverSW2L
:members:
:private-members:
.. autoclass:: Simul
:members:
:private-members:
.. todo::
Write a decent unittest file for this solver...
"""
import pickle
import numpy as np
from scipy.linalg import eig
from fluiddyn.util import mpi
from fluidsim.base.setofvariables import SetOfVariables
from fluidsim.base.solvers.pseudo_spect import (
SimulBasePseudoSpectral,
InfoSolverPseudoSpectral,
)
#from fluidsim.solvers.sw1l.util_pythran import compute_Frot
from fluidsim.solvers.sw1l.solver import compute_Frot
from ..util_pythran import gm_plus_kin_energy
[docs]class InfoSolverSW2L(InfoSolverPseudoSpectral):
"""Information about the solver SW2L."""
[docs] def _init_root(self):
"""The simulation object is instantiated with classes defined in this
function.
The super function `InfoSolverPseudoSpectral._init_root` is
called first. A few classes defined by this function are retained as it
is:
- :class:`fluidsim.solvers.sw1l.OperatorsPseudoSpectralSW1L`
- :class:`fluidsim.base.preprocess.pseudo_spect.PreprocessPseudoSpectral`
Solver-specific first-level classes are added:
- :class:`fluidsim_ocean.solvers.sw2l.solver.Simul`
- :class:`fluidsim_ocean.solvers.sw2l.state.StateSW2L`
- :class:`fluidsim_ocean.time_stepping.TimeSteppingPseudoSpectralSWnL`
- :class:`fluidsim_ocean.solvers.sw2l.init_fields.InitFieldsSW2L`
- :class:`fluidsim_ocean.solvers.sw2l.output.OutputSW2L`
- :class:`fluidsim_ocean.solvers.sw2l.forcing.ForcingSW2L`
"""
super(InfoSolverSW2L, self)._init_root()
package = "fluidsim_ocean.sw2l"
self.module_name = package + ".solver"
self.class_name = "Simul"
self.short_name = "sw2l"
classes = self.classes
classes.Operators.module_name = "fluidsim.solvers.sw1l.operators"
classes.Operators.class_name = "OperatorsPseudoSpectralSW1L"
classes.State.module_name = package + ".state"
classes.State.class_name = "StateSW2L"
classes.TimeStepping.module_name = "fluidsim_ocean.time_stepping"
classes.TimeStepping.class_name = "TimeSteppingPseudoSpectralSWnL"
classes.InitFields.module_name = package + ".init_fields"
classes.InitFields.class_name = "InitFieldsSW2L"
classes.Output.module_name = package + ".output"
classes.Output.class_name = "OutputSW2L"
classes.Forcing.module_name = package + ".forcing"
classes.Forcing.class_name = "ForcingSW2L"
[docs]class Simul(SimulBasePseudoSpectral):
"""A solver of the shallow-water 2 layers equations (SW2L).
.. inheritance-diagram:: Simul
.. todo::
I guest the topography (bathymetry) can be handle in this class. We will
need specific parameters about the topography.
"""
InfoSolver = InfoSolverSW2L
[docs] @staticmethod
def _complete_params_with_default(params):
r"""This static method is used to complete the *params* container.
Notes
-----
- :math:`f` is the Coriolis frequency [1/s].
- densities, list/array of layer densities :math:`\rho` [kg/m^3]
- :math:`g` gravity [m/s]
- :math:`H` background layer thicknesses [m]
- :math:`drag_coef` bottom friction coefficient [1/s]
"""
SimulBasePseudoSpectral._complete_params_with_default(params)
_H = [1000.0, 3000.0]
attribs = {"f": 1.0e-4, "densities": [1034.0, 1037.0], "g": 9.81, "H": _H, "drag_coef":0.}
params._set_attribs(attribs)
# todo: unused parameters...
params._set_child("topography", {"kind": None})
def __init__(self, params):
self.compute_vertical_modes(params)
super(Simul, self).__init__(params)
# if mpi.rank == 0:
# self.output.print_stdout(
# "c2 = {0:6.5g} ; f = {1:6.5g} ; kd2 = {2:6.5g}".format(
# params.c2, params.f, params.kd2
# )
# )
# derive other parameters
_H = self.params.H
_topo = 0.0 # should include real topo
self.Z = [sum(_H[::-1][:i]) + _topo for i in range(len(_H) + 1)][::-1]
[docs] def compute_vertical_modes(self, params):
""" Compute vertical mode structure
"""
# job performed by mpi process 0 and scattered
if mpi.rank == 0:
H = params.H
rho = params.densities
g = params.g
nb_layers = len(params.H)
A = np.zeros((nb_layers, nb_layers))
#
A[0, 0] = 1.0 + rho[0] / (rho[1] - rho[0])
A[0, 1] = -rho[1] / (rho[1] - rho[0])
#
for n in range(1, nb_layers - 1):
A[n, n - 1] = -rho[n - 1] / (rho[n] - rho[n - 1])
A[n, n] = rho[n] / (rho[n] - rho[n - 1]) + rho[n] / (
rho[n + 1] - rho[n]
)
A[n, n + 1] = -rho[n + 1] / (rho[n + 1] - rho[n])
#
n = nb_layers - 1
A[n, n - 1] = -rho[n - 1] / (rho[n] - rho[n - 1])
A[n, n] = rho[n] / (rho[n] - rho[n - 1])
#
B = np.diag(H)
#
w, v = eig(A, B)
# sort by eigenvalues
isort = np.argsort(np.real(w))
v = v[:, isort]
w = np.real(w[isort])
# fix normalization
sc = (B.dot(v ** 2)).sum(axis=0) / np.sum(H)
self.phi_vmode = v.dot(np.diag(1.0 / np.sqrt(sc)))
self.phih_vmode = B.dot(self.phi_vmode.dot(np.diag(w))) # h modes
self.iphi_vmode = B.dot(self.phi_vmode) / np.sum(H)
self.c_vmode = np.sqrt(g / w)
else:
# dummy assignment on other mpi processes
self.phi_vmode = 0.0
self.phih_vmode = 0.0
self.iphi_vmode = 0.0
self.c_vmode = 0.0
# broadcast result
if mpi.nb_proc > 1:
self.phi_vmode = mpi.comm.bcast(self.phi_vmode, root=0)
self.phih_vmode = mpi.comm.bcast(self.phih_vmode, root=0)
self.iphi_vmode = mpi.comm.bcast(self.iphi_vmode, root=0)
self.c_vmode = mpi.comm.bcast(self.c_vmode, root=0)
# # store vertical modes in params
if mpi.rank == 0 :
attribs = {"phi_vmode": self.phi_vmode,
"phih_vmode": self.phih_vmode,
"iphi_vmode": self.iphi_vmode,
"c_vmode": self.c_vmode,}
output = open("vmodes.pkl", 'wb')
pickle.dump(attribs, output)
output.close()
# params._set_attribs(attribs)
# print("phi",type(self.phi_vmode.tolist()),self.phi_vmode.tolist())
# print("H",type(H))
[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 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_n` and layer thicknesses :math:`h_n`.
Notes
-----
.. |uu| mathmacro:: \mathbf{u}
.. |p| mathmacro:: \partial
.. |ez| mathmacro:: \mathbf{e_z}
.. |bnabla| mathmacro:: \boldsymbol{\nabla}
.. |D| mathmacro:: \mbox{D}
.. |xx| mathmacro:: \mathbf{x}
In this function, we have to compute the right hand sides of these
equations:
.. math::
\p_t \uu_n(x,y) &= -\bnabla \Big \{ g m_n + \frac{1}{2} |\uu_n|^2 \Big \}
- (\zeta_n + f) \ez\times\uu_n,
\p_t h_n(x,y) &= - \nabla \cdot (\mathbf{u}_n h_n)
The gradient of the Montgomery potential is related to interface vertical
positions according to:
.. math::
\bnabla m_0 = \bnabla z_0 = \bnabla (T + h_0 + h_1),
.. math::
\bnabla m_1 = \frac{\rho_0}{\rho_1} \bnabla m_0
+ \big(1 - \frac{\rho_0}{\rho_1}\big) \bnabla (T + h_1).
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)
ux0 = state_phys.get_var("ux0")
uy0 = state_phys.get_var("uy0")
h0 = state_phys.get_var("h0")
rot0 = state_phys.get_var("rot0")
ux1 = state_phys.get_var("ux1")
uy1 = state_phys.get_var("uy1")
h1 = state_phys.get_var("h1")
rot1 = state_phys.get_var("rot1")
# Parameters for bottom friction
drag_coef = self.params.drag_coef
H = self.params.H
ux0_t0 = 0.
uy0_t0 = 0.
h0_t0 = H[0]
ux1_t0 = 0.
uy1_t0 = 0.
h1_t0 = H[1]
if old is None:
tendencies_fft = SetOfVariables(
like=self.state.state_spect, info="tendencies_nonlin"
)
else:
tendencies_fft = old
Fx0_fft = tendencies_fft.get_var("ux0_fft")
Fy0_fft = tendencies_fft.get_var("uy0_fft")
Fh0_fft = tendencies_fft.get_var("h0_fft")
Fx1_fft = tendencies_fft.get_var("ux1_fft")
Fy1_fft = tendencies_fft.get_var("uy1_fft")
Fh1_fft = tendencies_fft.get_var("h1_fft")
fft = oper.fft
f = self.params.f
# compute normalized Montgomery potentials
# for now ``topography = 0``
m0 = h0 + h1
rho0, rho1 = self.params.densities
ratio_rhos = rho0 / rho1
m1 = ratio_rhos * m0 + (1 - ratio_rhos) * h1
# compute nonlinear tendencies for layer 0
Fx, Fy = compute_Frot(rot0, ux0, uy0, f)
gradx_fft, grady_fft = oper.gradfft_from_fft(
fft(gm_plus_kin_energy(self.params.g, m0, ux0, uy0))
)
# Compute bottom friction for layer 0
bt_ux = -drag_coef * (ux0 - ux0_t0)
bt_uy = -drag_coef * (uy0 - uy0_t0)
bt_h = -drag_coef * (h0 - h0_t0)
# oper.dealiasing(gradx_fft, grady_fft)
Fx0_fft[:] = fft(Fx + bt_ux) - gradx_fft
Fy0_fft[:] = fft(Fy + bt_uy) - grady_fft
Fh0_fft[:] = -oper.divfft_from_vecfft(fft(h0 * ux0), fft(h0 * uy0)) + fft(bt_h)
# compute nonlinear tendencies for layer 1
Fx, Fy = compute_Frot(rot1, ux1, uy1, f)
gradx_fft, grady_fft = oper.gradfft_from_fft(
fft(gm_plus_kin_energy(self.params.g, m1, ux1, uy1))
)
# Compute bottom friction for layer 1
bt_ux = -drag_coef * (ux1 - ux1_t0)
bt_uy = -drag_coef * (uy1 - uy1_t0)
bt_h = -drag_coef * (h1 - h1_t0)
# oper.dealiasing(gradx_fft, grady_fft)
Fx1_fft[:] = fft(Fx + bt_ux) - gradx_fft
Fy1_fft[:] = fft(Fy + bt_uy) - grady_fft
Fh1_fft[:] = -oper.divfft_from_vecfft(fft(h1 * ux1), fft(h1 * uy1)) + fft(bt_h)
oper.dealiasing(tendencies_fft)
if self.params.forcing.enable:
tendencies_fft += self.forcing.get_forcing()
return tendencies_fft
if __name__ == "__main__":
import numpy as np
import fluiddyn as fld
params = Simul.create_default_params()
params.short_name_type_run = "test"
nh = 32 * 8
Lh = 2 * np.pi
params.oper.nx = nh
params.oper.ny = nh
params.oper.Lx = Lh
params.oper.Ly = Lh
dx = Lh / nh
delta_x = params.oper.Lx / params.oper.nx
params.nu_8 = 1e-4
params.time_stepping.t_end = 1.0
params.time_stepping.USE_CFL = False
params.time_stepping.deltat0 = 0.1
params.init_fields.type = "noise"
params.forcing.enable = False
params.forcing.type = "waves"
params.output.periods_print.print_stdout = 0.25
# params.output.periods_save.phys_fields = 0.5
# params.output.periods_save.spectra = 0.5
# params.output.periods_save.spect_energy_budg = 0.5
# params.output.periods_save.increments = 0.5
# params.output.periods_save.pdf = 0.5
# params.output.periods_save.time_signals_fft = True
params.output.periods_plot.phys_fields = 1.0
params.output.phys_fields.field_to_plot = "h0"
sim = Simul(params)
# sim.output.phys_fields.plot(field="h0")
sim.time_stepping.start()
# sim.output.phys_fields.plot(field="h0")
fld.show()