-
Pierre Augier authoredPierre Augier authored
util.py 13.77 KiB
import os
import sys
from itertools import product
from math import sqrt
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
from fluiddyn.util import has_to_be_made as _has_to_be_made
from fluidsim.util import get_dataframe_from_paths, times_start_last_from_path
from scipy import optimize
path_base = os.environ["STRAT_TURB_POLO2022"]
path_base_proj = os.environ["STRAT_TURB_POLO_PROJ2022"]
paths_all = sorted(Path(path_base).glob("simul_folders/ns3d*"))
paths_all_proj = sorted(Path(path_base_proj).glob("simul_folders/ns3d*"))
here = Path(__file__).absolute().parent
tmp_dir = here.parent / "tmp"
tmp_dir.mkdir(exist_ok=True)
def has_to_be_made(name, sources: list):
if not isinstance(name, str):
names = name
return any(has_to_be_made(name, sources) for name in names)
if isinstance(sources, str):
sources = [sources]
sources.append("util.py")
if not any(name.endswith(ext) for ext in (".png", ".tex")):
name += ".png"
return _has_to_be_made(tmp_dir / name, sources, source_dir=here)
height = 3.7
plt.rc("figure", figsize=(1.33 * height, height))
def get_paths(N, Rb, nh, proj=False, ratio_one=False, reverse=False):
str_N = f"_N{N}_"
if not ratio_one:
str_Rb = f"_Rb{Rb:.3g}_"
str_Rb2 = f"_Rb{Rb}_"
str_nh = f"_{nh}x{nh}"
if proj:
paths_tmp = paths_all_proj
else:
paths_tmp = paths_all
paths = [
p
for p in paths_tmp
if str_N in p.name
and (str_Rb in p.name or str_Rb2 in p.name)
and str_nh in p.name
]
else:
str_proj = "_projpoloidal_"
paths_tmp = paths_all_ratio_one
if proj:
paths = [
p
for p in paths_tmp
if str_N in p.name and str_proj in p.name and str_nh in p.name
]
else:
paths = [
p
for p in paths_tmp
if str_N in p.name and str_proj not in p.name and str_nh in p.name
]
paths.sort(key=lambda p: int(p.name.split("x")[1]), reverse=reverse)
return paths[0]
def get_paths_couple(N, Rb, proj=False, ratio_one=False, reverse=False):
str_N = f"_N{N}_"
if not ratio_one:
str_Rb = f"_Rb{Rb:.3g}_"
str_Rb2 = f"_Rb{Rb}_"
if proj:
paths_tmp = paths_all_proj
else:
paths_tmp = paths_all
paths_couple = [
p
for p in paths_tmp
if str_N in p.name and (str_Rb in p.name or str_Rb2 in p.name)
]
else:
str_proj = "_projpoloidal_"
paths_tmp = paths_all_ratio_one
if proj:
paths_couple = [
p for p in paths_tmp if str_N in p.name and str_proj in p.name
]
else:
paths_couple = [
p for p in paths_tmp if str_N in p.name and str_proj not in p.name
]
paths_couple.sort(key=lambda p: int(p.name.split("x")[1]), reverse=reverse)
return paths_couple
def get_path_finer_resol(N, Rb, proj=False, ratio_one=False):
paths_couple = get_paths_couple(N, Rb, proj, ratio_one, reverse=True)
for path in paths_couple:
t_start, t_last = times_start_last_from_path(path)
if t_last > t_start + 1:
return path
def lprod(a, b):
return list(product(a, b))
couples320 = set(
lprod([10, 20, 40], [5, 10, 20, 40, 80, 160])
+ lprod([30], [10, 20, 40])
+ lprod([6.5], [100, 200])
+ lprod([4], [250, 500])
+ lprod([3], [450, 900])
+ lprod([2], [1000, 2000])
+ lprod([0.66], [9000, 18000])
+ [(14.5, 20), (5.2, 150), (2.9, 475), (1.12, 3200), (0.25, 64000)]
)
couples320.add((60, 10))
couples320.add((60, 20))
couples320.add((80, 10))
couples320.add((100, 10))
couples320.add((120, 10))
couples320.remove((40, 160))
# Small Rb
couples320.update(lprod([20], [1, 2]))
couples320.update(lprod([40], [1, 2]))
couples320.update(lprod([80], [0.5, 1]))
N_ratio_one = [10, 20, 50, 80, 120]
has_to_save = "SAVE" in sys.argv
def save_fig(fig, name):
if has_to_save:
print(f"saving file {tmp_dir.name}/{name}")
fig.savefig(tmp_dir / name, dpi=300)
def customize(result, sim):
EKh = result["EKh"]
EKz = result["EKz"]
EK = EKh + EKz
U = sqrt(2 * EK / 3)
nu_2 = sim.params.nu_2
epsK = result["epsK"]
result["name"] = sim.output.name_run
if nu_2 != 0.0:
result["lambda"] = sqrt(U**2 * nu_2 / epsK)
result["Re_lambda"] = U * result["lambda"] / nu_2
result["Rb"] = float(sim.params.short_name_type_run.split("_Rb")[-1])
# else:
# result["lambda"] = INFINITY
# result["Re_lambda"] = INFINITY
# result["Rb"] = INFINITY
result["nx"] = sim.params.oper.nx
result["nz"] = sim.params.oper.nz
result["proj"] = sim.params.projection
def get_customized_dataframe(paths):
df = get_dataframe_from_paths(
paths, tmin="t_last-2", use_cache=1, customize=customize
)
if "nu0.0" in paths[0].name:
columns_old = df.columns.tolist()
# fmt: off
first_columns = [
"N", "nx", "nz", "Fh", "R4", "Gamma",
"lx1", "lx2", "lz1", "lz2", "I_velocity", "I_dissipation"]
# fmt: on
else:
df["Re"] = df.Rb * df.N**2
columns_old = df.columns.tolist()
# fmt: off
first_columns = [
"N", "Rb", "Re", "nx", "nz", "Fh", "R2", "k_max*eta", "epsK2/epsK", "Gamma",
"lx1", "lx2", "lz1", "lz2", "I_velocity", "I_dissipation"]
# fmt: on
columns = first_columns.copy()
for key in columns_old:
if key not in columns:
columns.append(key)
df = df[columns]
return df
def plot(
df,
x,
y,
logx=True,
logy=False,
c=None,
cmap=None,
vmin=None,
vmax=None,
s=None,
ax=None,
):
ax = df.plot.scatter(
x=x,
y=y,
logx=logx,
logy=logy,
c=c,
edgecolors="k",
vmin=vmin,
vmax=vmax,
s=s,
ax=ax,
)
if c is not None:
pc = ax.collections[-1]
if cmap == "seismic":
pc.set_cmap("seismic")
else:
pc.set_cmap("binary")
cbar = plt.colorbar(pc, ax=ax)
if vmin == -1 and vmax == 1:
cbar.set_ticks([-1.0, -0.5, 0.0, 0.5, 1.0])
return ax
N_1couple = 40 # 40
Rb_1couple = 20 # 20
paths_1couple = get_paths_couple(N_1couple, Rb_1couple)
print([p.name for p in paths_1couple])
params_simuls_regimes = {
"D": (40, 2),
"L": (40, 20), # (40, 20),
"O": (10, 80),
"W": (6.5, 200),
"P": (0.66, 18000),
"U": (80, 10),
}
paths_simuls_regimes = {
k: get_path_finer_resol(*params)
for k, params in params_simuls_regimes.items()
}
paths_simuls_regimes = {
k: v for k, v in paths_simuls_regimes.items() if v is not None
}
paths_simuls_regimes_proj = {
k: get_path_finer_resol(*params, proj=True)
for k, params in params_simuls_regimes.items()
}
paths_simuls_regimes_proj = {
k: v for k, v in paths_simuls_regimes_proj.items() if v is not None
}
def formatter_R(v):
if v % 1 == 0 or v >= 100:
return f"{v:.0f}"
else:
return f"{v:.1f}"
def formatter_N(v):
if v % 1 == 0:
return f"{v:.0f}"
elif v < 10:
return f"{v:.2f}"
else:
return f"{v:.1f}"
formatters = {
"N": formatter_N,
"Rb": formatter_R,
"k_max*eta": lambda v: f"{v:.2f}",
"k_max*lambda": lambda v: f"{v:.2f}",
"epsK2/epsK": lambda v: f"{v:.2f}",
"Fh": lambda v: f"{v:.2e}",
"R2": formatter_R,
"R4": lambda v: f"{v:.2e}",
"Re_lambda": formatter_R,
"Re": formatter_R,
}
Fh_limit = 0.13
R2_limit = 10.0
def compute_kf_kb_ko_keta_kd(sim, tmin):
mean_values = sim.output.get_mean_values(tmin=tmin, customize=customize)
# data = sim.output.spectra.load_kzkh_mean(tmin)
# kh = data["kh_spectra"]
# delta_kh = kh[1]
# assert delta_kh == 2 * np.pi / sim.params.oper.Lx
delta_kh = 2 * np.pi / sim.params.oper.Lx
if sim.params.oper.nx == sim.params.oper.nz:
kf = 3.5 * delta_kh
else:
kf = 20 * delta_kh
N = sim.params.N
nu = sim.params.nu_2
Fh = mean_values["Fh"]
R4 = mean_values["R4"]
Uh2 = mean_values["Uh2"]
kb = N / Uh2**0.5
epsK = mean_values["epsK"]
ko = (N**3 / epsK) ** 0.5
kd = (N / nu) ** 0.5
if sim.params.nu_2 != 0.0:
R2 = mean_values["R2"]
keta = ko * R2**0.75
else:
keta = 1e16
ketah = ko * (Fh * R4) ** (3 / 10)
keta = min(keta, ketah)
return kf, kb, ko, keta, kd
def pos_closest_value(input_list, input_value):
arr = np.asarray(input_list)
i = (np.abs(arr - input_value)).argmin()
return i
def compute_omega_emp_vs_kzkh(
N,
spectrum,
kh_spectra,
kz_spectra,
omegas,
):
KH, KZ = np.meshgrid(kh_spectra, kz_spectra)
K = (KH**2 + KZ**2) ** 0.5
K_NOZERO = K.copy()
K_NOZERO[K_NOZERO == 0] = 1e-16
omega_disp = N * KH / K_NOZERO
omega_emp = np.zeros((len(kz_spectra), len(kh_spectra)))
delta_omega_emp = np.zeros((len(kz_spectra), len(kh_spectra)))
omega_norm = np.zeros((len(kz_spectra), len(kh_spectra)))
# we compute omega_emp first
for io in range(len(omegas)):
omega_emp += omegas[io] * spectrum[:, :, io]
omega_norm += spectrum[:, :, io]
omega_norm[omega_norm == 0] = 1e-16
omega_emp = omega_emp / omega_norm
# then we conpute delta_omega_emp
for io in range(len(omegas)):
delta_omega_emp += ((omegas[io] - omega_disp) ** 2) * spectrum[:, :, io]
delta_omega_emp = (np.divide(delta_omega_emp, omega_norm)) ** 0.5
return omega_emp, delta_omega_emp
def spectra_vs_kzomega_slice(
spectrum,
kz_spectra,
omegas,
ikh,
normalize=True,
):
spectrum_kzomega = spectrum[:, ikh, :]
if normalize:
spectrum_normalized = np.zeros(spectrum_kzomega.shape)
for ikz in range(len(kz_spectra)):
norm = sum(spectrum_kzomega[ikz, :])
for io in range(len(omegas)):
spectrum_normalized[ikz, io] = spectrum_kzomega[ikz, io] / norm
return spectrum_normalized
else:
return spectrum_kzomega
def spectra_vs_khomega_slice(
spectrum,
kh_spectra,
omegas,
ikz,
normalize=True,
):
spectrum_khomega = spectrum[ikz, :, :]
if normalize:
spectrum_normalized = np.zeros(spectrum_khomega.shape)
for ikh in range(len(kh_spectra)):
norm = sum(spectrum_khomega[ikh, :])
for io in range(len(omegas)):
spectrum_normalized[ikh, io] = spectrum_khomega[ikh, io] / norm
return spectrum_normalized
else:
return spectrum_khomega
def spectral_fit_kh_kz(
spectrum, kh, kz, kmin, kmax, thetamin=0.0, thetamax=np.pi / 2, plot=False
):
F = spectrum.flatten()
KH = kh.flatten()
KZ = kz.flatten()
K = np.sqrt(KH**2 + KZ**2)
SINT = KH / K
KH = KH[F >= 1e-16]
KZ = KZ[F >= 1e-16]
K = K[F >= 1e-16]
SINT = SINT[F >= 1e-16]
F = F[F >= 1e-16]
F = F[K >= kmin]
KH = KH[K >= kmin]
KZ = KZ[K >= kmin]
SINT = SINT[K >= kmin]
K = K[K >= kmin]
F = F[K <= kmax]
KH = KH[K <= kmax]
KZ = KZ[K <= kmax]
SINT = SINT[K <= kmax]
K = K[K <= kmax]
F = F[SINT <= np.sin(thetamax)]
KZ = KZ[SINT <= np.sin(thetamax)]
K = K[SINT <= np.sin(thetamax)]
KH = KH[SINT <= np.sin(thetamax)]
SINT = SINT[SINT <= np.sin(thetamax)]
F = F[SINT >= np.sin(thetamin)]
KH = KH[SINT >= np.sin(thetamin)]
K = K[SINT >= np.sin(thetamin)]
KZ = KZ[SINT >= np.sin(thetamin)]
SINT = SINT[SINT >= np.sin(thetamin)]
def error(params):
c, ah, az = params
res = np.sum(
(np.log(F) - np.log(c) - ah * np.log(KH) - az * np.log(KZ)) ** 2
)
return res
if plot:
plt.figure()
c = 0.01
for ah in np.linspace(-4, 4, 20):
for az in np.linspace(-4, 4, 20):
plt.scatter(
ah, az, c=np.log(error([c, ah, az])), vmin=-3, vmax=10
)
# print(error([c, ah, az]))
plt.show()
initial_guess = [1e-1, -2, 0]
result = optimize.minimize(error, initial_guess, method="Nelder-Mead")
if result.success:
fitted_params = result.x
print("[c, ah, az] = ", fitted_params)
return fitted_params
else:
raise ValueError(result.message)
def spectral_fit_k_sint(
spectrum, k, sint, kmin, kmax, thetamin=0.0, thetamax=np.pi / 2, plot=False
):
F = spectrum.flatten()
K = k.flatten()
SINT = sint.flatten()
K = K[F >= 1e-16]
SINT = SINT[F >= 1e-16]
F = F[F >= 1e-16]
F = F[K >= kmin]
SINT = SINT[K >= kmin]
K = K[K >= kmin]
F = F[K <= kmax]
SINT = SINT[K <= kmax]
K = K[K <= kmax]
F = F[SINT <= np.sin(thetamax)]
K = K[SINT <= np.sin(thetamax)]
SINT = SINT[SINT <= np.sin(thetamax)]
F = F[SINT >= np.sin(thetamin)]
K = K[SINT >= np.sin(thetamin)]
SINT = SINT[SINT >= np.sin(thetamin)]
def error(params):
c, ak, at = params
res = np.sum(
(np.log(F) - np.log(c) - ak * np.log(K) - at * np.log(SINT)) ** 2
)
return res
if plot:
plt.figure()
c = 0.01
for ak in np.linspace(-10, 10, 20):
for at in np.linspace(-10, 10, 20):
plt.scatter(
ak, at, c=np.log10(error([c, ak, at])), vmin=0, vmax=6
)
# print(error([c, ak, at]))
plt.show()
initial_guess = [1e-1, -2, 0]
result = optimize.minimize(error, initial_guess, method="Nelder-Mead")
if result.success:
fitted_params = result.x
print("[c, ah, az] = ", fitted_params)
return fitted_params
else:
raise ValueError(result.message)