Skip to content
Snippets Groups Projects
util.py 15.8 KiB
Newer Older
import os
import sys
from itertools import product
from math import sqrt
from pathlib import Path
vlabarre's avatar
vlabarre committed
import h5py

import matplotlib.pyplot as plt
import numpy as np

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*"))

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*")
    )

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):
    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

    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,
vlabarre's avatar
vlabarre committed
    cmap="inferno",
    vmin=None,
    vmax=None,
vlabarre's avatar
vlabarre committed
    marker="o",
    s=None,
    ax=None,
):
    ax = df.plot.scatter(
        x=x,
        y=y,
        logx=logx,
        logy=logy,
        c=c,
        edgecolors="k",
        vmin=vmin,
        vmax=vmax,
vlabarre's avatar
vlabarre committed
        marker=marker,
        s=s,
        ax=ax,
    )

    if c is not None:
        pc = ax.collections[-1]
vlabarre's avatar
vlabarre committed
        if cmap == "inferno":
            pc.set_cmap("inferno")
        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),
    "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
}


vlabarre's avatar
vlabarre committed
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
vlabarre's avatar
vlabarre committed
    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
Pierre Augier's avatar
Pierre Augier committed
    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)

vlabarre's avatar
vlabarre committed
    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 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):
vlabarre's avatar
vlabarre committed
    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)
vlabarre's avatar
vlabarre committed
    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"][:]
vlabarre's avatar
vlabarre committed
        Epolo = f["spectrum_K"][:] - f["spectrum_Khr"][:]
        Etoro = f["spectrum_Khr"][:]
vlabarre's avatar
vlabarre committed
        Eequi = 2 * np.minimum(EA, Epolo)
vlabarre's avatar
vlabarre committed
    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)))
vlabarre's avatar
vlabarre committed
    E_tot = np.zeros((len(kz), len(kh)))
vlabarre's avatar
vlabarre committed
    for io in range(len(omegas)):
        omega = omegas[io]
        weight = gaussian_weight(omega, omega_disp, delta * omega_disp)
vlabarre's avatar
vlabarre committed
        tmp = weight * Eequi[:, :, io]
        E_waves_omega[:, :, io] = tmp
vlabarre's avatar
vlabarre committed
        E_waves += tmp
        E_tot += Epolo[:, :, io] + EA[:, :, io]  # + Etoro[:,:,io]

vlabarre's avatar
vlabarre committed
    return E_waves_omega, E_waves, E_tot, kh, kz, omegas
Pierre Augier's avatar
Pierre Augier committed

def spectra_vs_kzomega_slice(
Pierre Augier's avatar
Pierre Augier committed
    spectrum,
    kz_spectra,
    omegas,
    ikh,
    normalize=True,
):
    spectrum_kzomega = spectrum[:, ikh, :]
    if normalize:
        spectrum_normalized = np.zeros(spectrum_kzomega.shape)
Pierre Augier's avatar
Pierre Augier committed
        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(
Pierre Augier's avatar
Pierre Augier committed
    spectrum,
    kh_spectra,
    omegas,
    ikz,
    normalize=True,
):
    spectrum_khomega = spectrum[ikz, :, :]
    if normalize:
        spectrum_normalized = np.zeros(spectrum_khomega.shape)
Pierre Augier's avatar
Pierre Augier committed
        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)