Newer
Older
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 fluidsim.util import get_dataframe_from_paths, times_start_last_from_path
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*"))
try:
path_base_ratio_one = os.environ["STRAT_WAVES2022"]
except KeyError:
print("warning: STRAT_WAVES2022 environment variable not set")
else:
paths_all_ratio_one = sorted(
Path(path_base_ratio_one).glob("simul_folders/ns3d*")
)
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
here = Path(__file__).absolute().parent
tmp_dir = here.parent / "tmp"
tmp_dir.mkdir(exist_ok=True)
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=None, 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=None, proj=False, ratio_one=False):
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
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
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])
result["nx"] = sim.params.oper.nx
result["nz"] = sim.params.oper.nz
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
t_start, t_last = sim.output.print_stdout.get_times_start_last()
tmin = t_start + 2
data = sim.output.spectra.load1d_mean(tmin, verbose=False)
kz = data["kz"]
delta_kz = kz[1]
EKx = data["spectra_vx_kz"].sum() * delta_kz
EKy = data["spectra_vy_kz"].sum() * delta_kz
EKz = data["spectra_vz_kz"].sum() * delta_kz
result["EK"] = EKx + EKy + EKz
EKhr = data["spectra_Khr_kz"].sum() * delta_kz
EKhd = data["spectra_Khd_kz"].sum() * delta_kz
EKz = data["spectra_vz_kz"].sum() * delta_kz
result["Epolo"] = EKhd + EKz
result["Etoro"] = EKhr
# Get spatiotemporal spectra
path_run = Path(sim.output.path_run)
paths_spec = sorted(path_run.glob("spatiotemporal/periodogram_[0-9]*.h5"))
if not paths_spec:
return
path_spec = paths_spec[-1]
with h5py.File(path_spec, "r") as f:
kh = f["kh_spectra"][:]
kz = f["kz_spectra"][:]
omegas = f["omegas"][:]
EA = f["spectrum_A"][:]
Epolo = f["spectrum_K"][:] - f["spectrum_Khr"][:]
Eequi = 2 * np.minimum(EA, Epolo)
KH, KZ = np.meshgrid(kh, kz)
K = (KH**2 + KZ**2) ** 0.5
K_NOZERO = K.copy()
K_NOZERO[K_NOZERO == 0] = 1e-16
omega_disp = sim.params.N * KH / K_NOZERO
delta = 0.1
E_waves = np.zeros((len(kz), len(kh)))
for index_omega, omega in enumerate(omegas):
weight = gaussian_weight(omega, omega_disp, delta * omega_disp)
E_waves += weight * Eequi[:, :, index_omega]
dk2_dom = kh[1] * kz[1] * omegas[1]
result["E_waves"] = np.sum(np.nan_to_num(E_waves)) * dk2_dom
result["E_waves_norm"] = np.sum(EA + Epolo) * dk2_dom
def get_customized_dataframe(paths):
df = get_dataframe_from_paths(
paths, tmin="t_start+2", use_cache=1, customize=customize
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,
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]
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),
"O": (10, 80),
"W": (6.5, 200),
"P": (0.66, 18000),
"U": (80, 10),
"S1": (20, 160),
"S2": (40, 40),
"S3": (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
}
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
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
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)
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
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 gaussian_weight(x, mu, sigma):
return np.exp(-0.5 * ((x - mu) / sigma) ** 2)
def compute_E_waves_vs_kh_kz(sim, delta=0.1):
N = sim.params.N
t_start, t_last = sim.output.print_stdout.get_times_start_last()
tmin = t_last - 2.0
# mean_values = sim.output.get_mean_values(tmin=tmin, customize=customize)
path = sim.params.path_run
# Get spatiotemporal spectra
path_spec = sorted(path.glob(f"spatiotemporal/periodogram_[0-9]*.h5"))
assert len(path_spec) == 1, f"Not only 1 periodogram in {path} \n"
path_spec = path_spec[0]
with h5py.File(path_spec, "r") as f:
kh = f["kh_spectra"][:]
kz = f["kz_spectra"][:]
omegas = f["omegas"][:]
EA = f["spectrum_A"][:]
Etoro = f["spectrum_Khr"][:]
KH, KZ = np.meshgrid(kh, kz)
K = (KH**2 + KZ**2) ** 0.5
K_NOZERO = K.copy()
K_NOZERO[K_NOZERO == 0] = 1e-16
omega_disp = N * KH / K_NOZERO
E_waves_omega = np.zeros((len(kz), len(kh), len(omegas)))
E_waves = np.zeros((len(kz), len(kh)))
norm = np.zeros((len(kz), len(kh)))
for io in range(len(omegas)):
omega = omegas[io]
weight = gaussian_weight(omega, omega_disp, delta * omega_disp)
E_waves_omega[:, :, io] = tmp
E_tot += Epolo[:, :, io] + EA[:, :, io] # + Etoro[:,:,io]
return E_waves_omega, E_waves, E_tot, kh, kz, omegas
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
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
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
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)