Skip to content
Snippets Groups Projects
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)