"""Horizontal means
===================
Provides:
.. autoclass:: HorizontalMeans
:members:
:private-members:
:noindex:
:undoc-members:
"""
import numpy as np
import h5py
import matplotlib.pyplot as plt
from fluiddyn.util import mpi
from fluidsim.extend_simul import SimulExtender, extend_simul_class
from .base import SpecificOutput
__all__ = ["extend_simul_class", "HorizontalMeans"]
[docs]
class HorizontalMeans(SpecificOutput, SimulExtender):
"""Horizontal means as functions of the z coordinate
Examples
--------
This is an output class to compute/save/load/plot horizontally averaged
quantities. Since this is useful only for particular simulations, it is a
``SimulExtender``, meaning that it should be used this way:
.. code-block:: python
from fluidsim.solvers.ns3d.solver import Simul
from fluidsim.base.output.horiz_means import HorizontalMeans, extend_simul_class
Simul = extend_simul_class(Simul, HorizontalMeans)
params = Simul.create_default_params()
...
params.output.periods_save.horiz_means = 0.1
sim = Simul(params)
Then, during or after the simulation:
.. code-block:: python
sim.output.horiz_means.plot(tmin=10)
data = sim.output.horiz_means.load(tmin=10)
"""
_tag = "horiz_means"
_module_name = "fluidsim.base.output.horiz_means"
_name_file = _tag + ".h5"
[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(
cls._tag,
attribs={
"module_name": cls._module_name,
"class_name": cls.__name__,
},
)
return modif_info_solver
[docs]
@classmethod
def complete_params_with_default(cls, params):
params.output.periods_save._set_attrib(cls._tag, 0)
[docs]
def _compute_hmean(self, arr3d):
if mpi.nb_proc == 1:
return np.mean(arr3d, axis=(1, 2))
# here (MPI case) it gets a bit more complicated!
sum_local = np.sum(arr3d, axis=(1, 2))
if mpi.rank == 0:
sum_global = np.zeros(self.nz)
# MPI messages and summation
for rank in range(mpi.nb_proc):
if mpi.rank == 0:
nz_loc = self.nzs_local[rank]
if rank == 0 and mpi.rank == 0:
data = sum_local
else:
if mpi.rank == 0:
data = np.empty(nz_loc)
mpi.comm.Recv(data, source=rank, tag=42 * rank)
elif mpi.rank == rank:
mpi.comm.Send(sum_local, dest=0, tag=42 * rank)
if mpi.rank == 0:
iz_start = self.izs_start[rank]
sum_global[iz_start : iz_start + nz_loc] += data
if mpi.rank == 0:
return sum_global / self.nh
[docs]
def _hmean_to_3d_local(self, profile):
if mpi.nb_proc != 1:
profile = mpi.comm.bcast(profile, root=0)
profile = profile[self.iz_start : self.iz_start + self.nz_local]
return np.broadcast_to(
profile[:, np.newaxis, np.newaxis], self.shapeX_loc
)
def __init__(self, output):
self.output = output
sim = output.sim
params = sim.params
self.shapeX_loc = sim.oper.shapeX_loc
if mpi.nb_proc > 1:
self.iz_start, _, _ = sim.oper.oper_fft.get_seq_indices_first_X()
self.izs_start = mpi.comm.gather(self.iz_start, root=0)
nh_local = self.shapeX_loc[1] * self.shapeX_loc[2]
self.nhs_local = mpi.comm.gather(nh_local, root=0)
self.nz_local = self.shapeX_loc[0]
self.nzs_local = mpi.comm.gather(self.nz_local, root=0)
if mpi.rank == 0:
assert len(self.izs_start) == mpi.nb_proc
assert len(self.nhs_local) == mpi.nb_proc
assert len(self.nzs_local) == mpi.nb_proc
print(self.izs_start, self.nhs_local, self.nzs_local)
self.nz = params.oper.nz
self.nh = params.oper.nx * params.oper.ny
dz = params.oper.Lz / self.nz
z = dz * np.arange(self.nz)
super().__init__(
output,
period_save=params.output.periods_save.horiz_means,
arrays_1st_time={"z": z},
)
[docs]
def compute(self):
data = {}
get_var = self.sim.state.state_phys.get_var
vx = get_var("vx")
vy = get_var("vy")
vz = get_var("vz")
def _extend_data_with_hmean(arr3d, name: str):
result = data[name] = self._compute_hmean(arr3d)
return result
vx_mean = _extend_data_with_hmean(vx, "vx")
vy_mean = _extend_data_with_hmean(vy, "vy")
vz_mean = _extend_data_with_hmean(vz, "vz")
vx_mean3d = self._hmean_to_3d_local(vx_mean)
vy_mean3d = self._hmean_to_3d_local(vy_mean)
vz_mean3d = self._hmean_to_3d_local(vz_mean)
vxp = vx - vx_mean3d
vyp = vy - vy_mean3d
vzp = vz - vz_mean3d
_extend_data_with_hmean(vxp**2, "vxp_vxp")
_extend_data_with_hmean(vyp**2, "vyp_vyp")
_extend_data_with_hmean(vzp**2, "vzp_vzp")
_extend_data_with_hmean(vyp * vxp, "vyp_vxp")
_extend_data_with_hmean(vzp * vxp, "vzp_vxp")
_extend_data_with_hmean(vzp * vyp, "vzp_vyp")
return data
[docs]
def load(self, tmin=None, tmax=None, verbose=False):
with h5py.File(self.path_file, "r") as file:
dset_times = file["times"]
times = dset_times[...]
nt = len(times)
z = file["z"][...]
if tmin is None:
imin_plot = 0
else:
imin_plot = np.argmin(abs(times - tmin))
if tmax is None:
imax_plot = nt - 1
else:
imax_plot = np.argmin(abs(times - tmax))
tmin = times[imin_plot]
tmax = times[imax_plot]
if verbose:
print(
"compute time average with\n"
+ (
f"tmin = {tmin:8.6g} ; tmax = {tmax:8.6g}"
f"imin = {imin_plot:8d} ; imax = {imax_plot:8d}"
)
)
data = {"z": z}
for key in list(file.keys()):
if key.startswith("v"):
dset_key = file[key]
profile = dset_key[imin_plot : imax_plot + 1].mean(0)
data[key] = profile
return data
[docs]
def plot(self, tmin=None, tmax=None, verbose=False):
data = self.load(tmin, tmax, verbose)
z = data.pop("z")
fig, (ax_left, ax_right) = plt.subplots(ncols=2, sharey=True)
for letter in "xyz":
arr = data.pop("v" + letter)
ax_left.plot(arr, z, label=f"$<v_{letter}>_{{xy}}$")
arr = data.pop(f"v{letter}p_v{letter}p")
ax_right.plot(arr, z, label=f"$<v_{letter}' v_{letter}'>_{{xy}}$")
ax_left.legend()
ax_left.set_ylabel("$z$")
for l0, l1 in ("yx", "zx", "zy"):
name = f"v{l0}p_v{l1}p"
ax_right.plot(
data[name], z, "--", label=f"$<v_{l0}' v_{l1}'>_{{xy}}$"
)
ax_right.legend()
fig.tight_layout()