"""Spatial means regions
========================
.. autoclass:: SpatialMeansRegions
:members:
:private-members:
"""
from pathlib import Path
import numbers
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from fluiddyn.util import mpi
from fluidfft.fft3d.operators import vector_product
from fluidsim.base.output.base import SpecificOutput
from . import SimulExtender
[docs]
class SpatialMeansRegions(SimulExtender, SpecificOutput):
r"""Specific output for the MILESTONE simulations
Notes
-----
.. |p| mathmacro:: \partial
.. |vv| mathmacro:: \textbf{v}
.. |Fh| mathmacro:: \textbf{F}_h
.. |ez| mathmacro:: \textbf{e}_\textbf{z}
.. |bnabla| mathmacro:: \boldsymbol{\nabla}
.. |bomega| mathmacro:: \boldsymbol{\omega}
This output is used with the solver
:mod:`fluidsim.solvers.ns3d.strat.solver`, which solves the incompressible
Navier-Stokes equations (possibly with hyper-viscosity) under the
Boussinesq approximation with a constant Brunt-Vaisala frequency:
.. math::
\p_t \vv + \vv \cdot \bnabla \vv = - \bnabla p + b \ez + D_\vv + \Fh,
\p_t b + \vv \cdot \bnabla b = - N^2 v_z + D_b,
where :math:`\vv` is the non-divergent velocity (:math:`\bnabla \cdot \vv =
0`), :math:`p` is the pressure, :math:`b` is the buoyancy, :math:`N` is the
(constant) Brunt-Vaisala frequency and :math:`D_\vv` and :math:`D_b` are
dissipative terms.
The total dynamic pressure :math:`P = p + |\vv|^2/2` can be computed from
:math:`\vv` and :math:`b` as:
.. math::
\bnabla^2 P = \bnabla \cdot (\vv \times \bomega) + \bnabla \cdot \Fh + \p_z b.
We consider the energy budget
.. math::
\p_t |\vv|^2/2 + \bnabla \cdot (\vv P) = \vv \cdot D_\vv + \vv \cdot \Fh + v_z b
\p_t e_A + \bnabla \cdot (\vv e_A) = - v_z b + b D_b / N^2,
with :math:`e_A = b^2/(2N^2)`.
We take the average of these equations over regions delimited by surfaces
along the y and z axis. Some terms written as a divergence can be rewritten
as fluxes through 2 surfaces at :math:`x_{min}` and :math:`x_{max}`:
.. math::
\int_V \frac{dV}{V} \bnabla \cdot (\vv P) =
\frac{S}{V} \left( \langle v_x P \rangle_{S_{min}} - \langle v_x P \rangle_{S_{max}} \right).
"""
_tag = "spatial_means_regions"
_module_name = "fluidsim.extend_simul.spatial_means_regions_milestone"
def __init__(self, output):
params = output.sim.params
params_cls = params.output.spatial_means_regions
self.xmin_given = params_cls.xmin
self.xmax_given = params_cls.xmax
if isinstance(self.xmin_given, numbers.Number):
self.xmin_given = [self.xmin_given]
if isinstance(self.xmax_given, numbers.Number):
self.xmax_given = [self.xmax_given]
self.nb_regions = len(self.xmin_given)
if len(self.xmax_given) != self.nb_regions:
raise ValueError("len(self.xmax_given) != len(self.xmin_given)")
oper = output.sim.oper
Lx = params.oper.Lx
if not params.ONLY_COARSE_OPER:
x_seq = oper.x_seq
else:
x_seq = Lx / params.oper.nx * np.arange(params.oper.nx)
self.info_regions = []
_, _, ix_seq_start = oper.seq_indices_first_X
nx_loc = oper.shapeX_loc[2]
for xmin, xmax in zip(self.xmin_given, self.xmax_given):
xmin, xmax = Lx * xmin, Lx * xmax
ixmin = np.argmin(abs(x_seq - xmin))
xmin = x_seq[ixmin]
ixmin_loc = ixmin - ix_seq_start
if ixmin_loc < 0 or ixmin_loc > nx_loc - 1:
# this limit is not in this process
ixmin_loc = None
ixmax = np.argmin(abs(x_seq - xmax))
xmax = x_seq[ixmax]
ixmax_loc = ixmax - ix_seq_start
ixstop_loc = ixmax_loc + 1
# special case when region == whole numerical domain
if ixmax == params.oper.nx - 1 and ixmin == 0:
xmax += x_seq[1]
ixmax += 1
ixmax_surf = ixmax % params.oper.nx
ixmax_surf_loc = ixmax_surf - ix_seq_start
if ixmax_surf_loc < 0 or ixmax_surf_loc > nx_loc - 1:
# this limit is not in this process
ixmax_surf_loc = None
self.info_regions.append(
(xmin, xmax, ixmin, ixmax, ixmin_loc, ixmax_surf_loc, ixstop_loc)
)
super().__init__(
output, period_save=params.output.periods_save.spatial_means_regions
)
if self.period_save == 0:
return
self.masks = []
self.nb_points = []
self.nb_points_xmin = []
self.nb_points_xmax = []
for info_region in self.info_regions:
(ixmin_loc, ixmax_surf_loc, ixstop_loc) = info_region[4:]
mask_loc = np.zeros(shape=oper.shapeX_loc, dtype=np.int8)
mask_loc[:, :, ixmin_loc:ixstop_loc] = 1
self.masks.append(mask_loc)
nb_points = mask_loc.sum()
if mpi.nb_proc > 1:
nb_points = mpi.comm.allreduce(nb_points, op=mpi.MPI.SUM)
self.nb_points.append(nb_points)
def compute_nb_points_surface(ixsurface_loc):
if ixsurface_loc is not None:
nz_loc = oper.shapeX_loc[0]
ny_loc = oper.shapeX_loc[1]
nb_points_xsurface_loc = ny_loc * nz_loc
else:
nb_points_xsurface_loc = np.int8(0)
if mpi.nb_proc > 1:
nb_points_xsurface = mpi.comm.allreduce(
nb_points_xsurface_loc, op=mpi.MPI.SUM
)
else:
nb_points_xsurface = nb_points_xsurface_loc
return nb_points_xsurface
self.nb_points_xmin.append(compute_nb_points_surface(ixmin_loc))
self.nb_points_xmax.append(compute_nb_points_surface(ixmax_surf_loc))
self._save_one_time()
def _init_path_files(self):
self.path_dir = Path(self.output.path_run) / self._tag
self.paths = [
self.path_dir / f"data{iregion}.csv"
for iregion in range(self.nb_regions)
]
def _init_files(self, arrays_1st_time=None):
if mpi.rank == 0:
self.path_dir.mkdir(exist_ok=True)
for path, info_region in zip(self.paths, self.info_regions):
xmin, xmax = info_region[:2]
if not path.exists():
with open(path, "w") as file:
file.write(
f"# xmin = {xmin} ; xmax = {xmax}\n"
"time,EK,EKz,EA,epsK,epsA,conv_K2A,PK,PA,"
"flux_Pnl_xmin,flux_Pnl_xmax,flux_v2_xmin,flux_v2_xmax,"
"flux_P_dz_b_xmin,flux_P_dz_b_xmax,"
"flux_Pforcing_xmin,flux_Pforcing_xmax,"
"flux_A_xmin,flux_A_xmax\n"
)
else:
with open(path, "r") as file:
words = file.readline().split()
xmin_file = float(words[3])
xmax_file = float(words[7])
if xmin_file != xmin or xmax_file != xmax:
raise ValueError(
"xmin_file != xmin or xmax_file != xmax\n"
f"{xmin_file} != {xmin} or {xmax_file} != {xmax}"
)
[docs]
@classmethod
def complete_params_with_default(cls, params):
params.output.periods_save._set_attrib(cls._tag, 0)
params.output._set_child(
cls._tag, attribs={"xmin": [0.1, 0.4, 0.7], "xmax": [0.3, 0.6, 0.9]}
)
[docs]
@classmethod
def get_modif_info_solver(cls):
"""Create a function to modify ``info_solver``.
Note that this function is called when the object ``info_solver`` has
not yet been created (and cannot yet be modified)! This is why one
needs to create a function that will be called later to modify
``info_solver``.
"""
def modif_info_solver(info_solver):
info_solver.classes.Output.classes._set_child(
"SpatialMeansRegions",
attribs={
"module_name": cls._module_name,
"class_name": cls.__name__,
},
)
return modif_info_solver
[docs]
def _online_save(self):
if self._has_to_online_save():
self._save_one_time()
def _compute_means_regions(self, field):
results = []
for nb_points, mask in zip(self.nb_points, self.masks):
result = (mask * field).sum()
if mpi.nb_proc > 1:
result = mpi.comm.allreduce(result, op=mpi.MPI.SUM)
results.append(result / nb_points)
return np.array(results)
[docs]
def _compute_fluxes_regions(self, field, vx):
"""Compute for each region the fluxes over xmin and xmax surface"""
fluxes = []
for index_region, info_region in enumerate(self.info_regions):
xmin, xmax = info_region[:2]
length_region = xmax - xmin
ixmin_loc, ixmax_surf_loc = info_region[4:-1]
def compute_flux(nb_points_surfaces, ixsurface_loc):
nb_points_surface = nb_points_surfaces[index_region]
if ixsurface_loc is not None:
field_xsurface = field[:, :, ixsurface_loc]
vx_xsurface = vx[:, :, ixsurface_loc]
sum_xsurface = (vx_xsurface * field_xsurface).sum()
else:
sum_xsurface = 0.0
if mpi.nb_proc > 1:
sum_xsurface = mpi.comm.allreduce(
sum_xsurface, op=mpi.MPI.SUM
)
return sum_xsurface / nb_points_surface / length_region
flux_xmin = compute_flux(self.nb_points_xmin, ixmin_loc)
flux_xmax = -compute_flux(self.nb_points_xmax, ixmax_surf_loc)
fluxes.append((flux_xmin, flux_xmax))
return np.array(fluxes)
def _save_one_time(self):
tsim = self.sim.time_stepping.t
self.t_last_save = tsim
state = self.sim.state
oper = self.sim.oper
ifft = oper.ifft
fft = oper.fft
ifft_as_arg_destroy = oper.ifft_as_arg_destroy
get_var = state.state_phys.get_var
b = get_var("b")
vx = get_var("vx")
vy = get_var("vy")
vz = get_var("vz")
N2b2 = 0.5 / self.sim.params.N**2 * b * b
EAs = self._compute_means_regions(N2b2)
vh2 = 0.5 * (vx * vx + vy * vy)
vz2 = 0.5 * vz * vz
v2_over_2 = vh2 + vz2
EKhs = self._compute_means_regions(vh2)
EKzs = self._compute_means_regions(vz2)
del vh2, vz2
EKs = EKhs + EKzs
conv_K2A = -self._compute_means_regions(b * vz)
get_var = state.state_spect.get_var
b_fft = get_var("b_fft")
vx_fft = get_var("vx_fft")
vy_fft = get_var("vy_fft")
vz_fft = get_var("vz_fft")
f_d, _ = self.sim.compute_freq_diss()
f_d_vx = ifft(f_d * vx_fft)
f_d_vy = ifft(f_d * vy_fft)
f_d_vz = ifft(f_d * vz_fft)
epsKs = self._compute_means_regions(
vx * f_d_vx + vy * f_d_vy + vz * f_d_vz
)
del f_d_vx, f_d_vy, f_d_vz
epsAs = [
epsA / self.sim.params.N**2
for epsA in self._compute_means_regions(b * ifft(f_d * b_fft))
]
if self.sim.params.forcing.enable:
deltat = self.sim.time_stepping.deltat
forcing_fft = self.sim.forcing.get_forcing()
fx_fft = forcing_fft.get_var("vx_fft")
fy_fft = forcing_fft.get_var("vy_fft")
assert np.allclose(forcing_fft.get_var("vz_fft"), 0)
fx = ifft(fx_fft)
fy = ifft(fy_fft)
PKs = self._compute_means_regions(
vx * fx + vy * fy + (abs(fx) ** 2 + abs(fy) ** 2) * deltat / 2
)
else:
PKs = np.zeros(3)
# Compute spatial fluxes
# Need to compute the pressure P = v^2/2 + p
omegax_fft, omegay_fft, omegaz_fft = oper.rotfft_from_vecfft(
vx_fft, vy_fft, vz_fft
)
ifft_as_arg_destroy = oper.ifft_as_arg_destroy
omegax = state.fields_tmp[3]
omegay = state.fields_tmp[4]
omegaz = 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)
del omegax_fft, omegay_fft, omegaz_fft
f_nl_x, f_nl_y, f_nl_z = vector_product(
vx, vy, vz, omegax, omegay, omegaz
)
del omegax, omegay, omegaz
f_nl_x_fft = fft(f_nl_x)
f_nl_y_fft = fft(f_nl_y)
f_nl_z_fft = fft(f_nl_z)
f_nl_x_fft *= 1j * oper.Kx
dx_f_nl_x_fft = f_nl_x_fft
del f_nl_x_fft
f_nl_y_fft *= 1j * oper.Ky
dy_f_nl_y_fft = f_nl_y_fft
del f_nl_y_fft
f_nl_z_fft *= 1j * oper.Kz
dz_f_nl_z_fft = f_nl_z_fft
del f_nl_z_fft
P_nl_fft = -(dx_f_nl_x_fft + dy_f_nl_y_fft + dz_f_nl_z_fft) / oper.K2_not0
del dx_f_nl_x_fft, dy_f_nl_y_fft, dz_f_nl_z_fft
P_nl = ifft(P_nl_fft)
del P_nl_fft
if self.sim.params.forcing.enable:
fz_fft = forcing_fft.get_var("vz_fft")
P_forcing = -ifft(
oper.divfft_from_vecfft(fx_fft, fy_fft, fz_fft) / oper.K2_not0
)
P_dz_b = -ifft(1j * oper.Kz * b_fft / oper.K2_not0)
fluxes_P_nl = self._compute_fluxes_regions(P_nl, vx)
fluxes_v2 = self._compute_fluxes_regions(v2_over_2, vx)
fluxes_P_forcing = self._compute_fluxes_regions(P_forcing, vx)
fluxes_P_dz_b = self._compute_fluxes_regions(P_dz_b, vx)
fluxes_A = self._compute_fluxes_regions(N2b2, vx)
if mpi.rank > 0:
return
for i, path in enumerate(self.paths):
flux_Pnl_xmin, flux_Pnl_xmax = fluxes_P_nl[i]
flux_v2_xmin, flux_v2_xmax = fluxes_v2[i]
flux_Pforcing_xmin, flux_Pforcing_xmax = fluxes_P_forcing[i]
flux_P_dz_b_xmin, flux_P_dz_b_xmax = fluxes_P_dz_b[i]
flux_A_xmin, flux_A_xmax = fluxes_A[i]
with open(path, "a") as file:
file.write(
f"{tsim},{EKs[i]},{EKzs[i]},{EAs[i]},"
f"{epsKs[i]},{epsAs[i]},{conv_K2A[i]},{PKs[i]},0.0,"
f"{flux_Pnl_xmin},{flux_Pnl_xmax},{flux_v2_xmin},{flux_v2_xmax},"
f"{flux_P_dz_b_xmin},{flux_P_dz_b_xmax},"
f"{flux_Pforcing_xmin},{flux_Pforcing_xmax},"
f"{flux_A_xmin},{flux_A_xmax}\n"
)
def load(self, iregion=0):
df = pd.read_csv(self.paths[iregion], skiprows=1)
return df
[docs]
def print_keys(self):
"""Print the keys associated with the computed quantities"""
df = pd.read_csv(self.paths[0], skiprows=1, nrows=1)
print(df.columns)
[docs]
def plot(self, keys="EK", iregion=0):
"""Plot some quantities for a given region"""
df = self.load(iregion)
ax = df.plot(x="time", y=keys)
xmin, xmax = self.info_regions[iregion][:2]
ax.set_title(self.output.summary_simul + f"\nxmin={xmin}, xmax={xmax}")
return ax
[docs]
def plot_budget(
self, iregion=0, decompose_fluxes=False, plot_conversion=False
):
"""Plot the energy budget for a given region"""
df = self.load(iregion)
times = df["time"]
E_tot = df["EK"] + df["EA"]
dt_E_tot = np.gradient(E_tot, times)
P_tot = df["PK"] + df["PA"]
eps = df["epsK"] + df["epsA"]
fig, ax = plt.subplots()
ax.plot(times, dt_E_tot, "k--", label="$d_t E$", linewidth=2)
ax.plot(times, P_tot, label="Forcing")
ax.plot(times, -eps, label="Viscosity")
kinds = "flux_Pnl flux_Pforcing flux_A flux_P_dz_b".split()
flux_tot = np.zeros_like(times)
for kind in kinds:
flux_kind = df[kind + "_xmin"] + df[kind + "_xmax"]
flux_tot += flux_kind.values
ax.plot(times, flux_tot, label="Surface fluxes")
ax.plot(times, P_tot - eps + flux_tot, label="All terms", linewidth=2)
if decompose_fluxes:
for kind in kinds:
flux_kind = df[kind + "_xmin"] + df[kind + "_xmax"]
ax.plot(times, flux_kind, ":", label=kind)
if plot_conversion:
ax.plot(times, df["conv_K2A"], "-.", label=r"$C_{K\rightarrow A}$")
xmin, xmax = self.info_regions[iregion][:2]
ax.set_title(
self.output.summary_simul + f"\nxmin={xmin:.3f}, xmax={xmax:.3f}"
)
ax.set_xlabel("time")
fig.legend()
fig.tight_layout()
return ax