MIT_Data_Processig.py

from pathlib import Path
import numpy as np
from scipy.signal import detrend, get_window
import MITgcmutils as MIT

# -------------------------
# Lectura
# -------------------------
def cargar_datos_mitgcm(outdir_bay: Path, outdir_nobay: Path, use_wildcard: bool = False):
    pattern = 'state.*.t*.nc' if use_wildcard else 'state.0000007200.t*.nc'
    nc_bay   = MIT.mnc_files(str(outdir_bay   / pattern))
    gr_bay   = MIT.mnc_files(str(outdir_bay   / 'grid.t*.nc'))
    nc_nobay = MIT.mnc_files(str(outdir_nobay / pattern))
    time = nc_bay.variables['T'][:]  # segundos
    return {'nc_bay': nc_bay, 'nc_nobay': nc_nobay, 'gr_bay': gr_bay, 'time': time}

# -------------------------
# Estaciones por grilla
# -------------------------
def make_stations_from_grid(gr_bay, j_centro: int = 315, radio: int = 50):
    """
    Genera estaciones (Seiche, A, B, C) a partir de la grilla.
    Devuelve:
      - stations_save: [(i, j, label)]  → para funciones de guardado
      - stations_meta: [{label, i, j, color}] → para metadatos y graficado futuro
    """
    # Dimensiones
    ny, nx = gr_bay.variables['XC'].shape
    i_centro = nx // 2

    # Clamp de índices para evitar out-of-range
    def clamp(v, lo, hi): return max(lo, min(hi, v))
    j_centro = clamp(j_centro, 0, ny-1)
    iS, jS = i_centro, clamp(j_centro + 20, 0, ny-1)
    iA, jA = clamp(i_centro - radio, 0, nx-1), j_centro
    iB, jB = clamp(i_centro + radio, 0, nx-1), j_centro
    iC, jC = i_centro, clamp(j_centro - radio, 0, ny-1)
    # iG, jG = i_centro, clamp(j_centro - 2*radio, 0, ny-1)  # estación extra (Gamma)
    # iO

    stations_save = [
        (iS, jS, 'Seiche'),
        (iA, jA, 'A'),
        (iB, jB, 'B'),
        (iC, jC, 'C'),
    ]
    stations_meta = [
        {"label": "Seiche", "i": int(iS), "j": int(jS), "color": "#9d02d7"},
        {"label": "A",      "i": int(iA), "j": int(jA), "color": "#ea5f94"},
        {"label": "B",      "i": int(iB), "j": int(jB), "color": "#0000ff"},
        {"label": "C",      "i": int(iC), "j": int(jC), "color": "#fa8775"},
    ]
    return stations_save, stations_meta

# -----------------------------------------------------------
# Guardado sin gráficas
# -----------------------------------------------------------
def guardar_series_y_espectros(
    nc_bay,
    nc_nobay,
    time,
    stations,                 # [(i,j,label)] o [(i,j,label,color)]
    tmin_horas: float = 27.0,
    save_path: Path = Path("."),
    save_prefix: str = "resultado",
    extra_meta: dict = None
):
    import json
    save_path.mkdir(parents=True, exist_ok=True)

    # Tiempo recortado
    t = np.array(time, dtype=float)
    tmin_seg = tmin_horas * 3600.0
    idx_start = np.where(t >= tmin_seg)[0][0] if np.any(t >= tmin_seg) else 0
    t_cut = t[idx_start:]
    t_cut_horas = t_cut / 3600.0

    series_npz, espec_npz = {}, {}
    tabla_estaciones = []  # metadatos

    for st in stations:
        # Permite (i,j,label) o (i,j,label,color)
        if len(st) == 3:
            i_st, j_st, label = st
            color = None
        else:
            i_st, j_st, label, color = st

        eta_bay   = np.array(nc_bay.variables['Eta'][:, j_st, i_st], dtype=float)
        eta_nobay = np.array(nc_nobay.variables['Eta'][:, j_st, i_st], dtype=float)

        eta_b_cut = eta_bay[idx_start:]
        eta_n_cut = eta_nobay[idx_start:]
        eta_diff_cut = eta_b_cut - eta_n_cut

        series_npz[f"{label}_t_horas"]  = t_cut_horas
        series_npz[f"{label}_eta_diff"] = eta_diff_cut

        # FFT Welch/Parseval con ventana Hann  [m² Hz⁻¹ — notebook aplica ×3600]
        dt = t_cut[1] - t_cut[0] if len(t_cut) > 1 else 1.0   # segundos
        N   = len(eta_diff_cut)
        win = get_window('hann', N, fftbins=True)
        U_w = float((win**2).sum())
        eta_detrend = detrend(eta_diff_cut) * win
        Y    = np.fft.rfft(eta_detrend)
        freqs = np.fft.rfftfreq(N, d=dt)
        PSD  = (np.abs(Y)**2) / (U_w * (1.0 / dt))   # m² Hz⁻¹
        freq_cph = freqs * 3600.0

        espec_npz[f"{label}_freq_cph"] = freq_cph
        espec_npz[f"{label}_PSD"]      = PSD

        tabla_estaciones.append({"label": str(label), "i": int(i_st), "j": int(j_st), "color": color})

    meta = {
        "tmin_horas": float(tmin_horas),
        "idx_start": int(idx_start),
        "dt_seconds": float(t_cut[1] - t_cut[0]) if len(t_cut) > 1 else None,
        "stations": tabla_estaciones
    }
    if extra_meta:
        meta.update(extra_meta)
    meta_json = json.dumps(meta, ensure_ascii=False)

    series_npz["meta_json"] = np.array(meta_json)
    espec_npz["meta_json"]  = np.array(meta_json)

    out_series    = save_path / f"{save_prefix}_series.npz"
    out_espectros = save_path / f"{save_prefix}_espectros.npz"
    np.savez_compressed(out_series, **series_npz)
    np.savez_compressed(out_espectros, **espec_npz)
    print(f"[OK] Guardados:\n  - {out_series}\n  - {out_espectros}")

# ------------------------------------------------------------------------
# Envoltura por escenario (siempre usa estaciones derivadas de la grilla)
# ------------------------------------------------------------------------
def guardar_series_y_espectros_escenario(
    scenario_label: str,
    outdir_bay: Path,
    outdir_nobay: Path,
    tmin_horas: float = 27.0,
    save_dir: Path = Path("./npz_out"),
    use_wildcard: bool = False,
    j_centro: int = 315,
    radio: int = 50
):
    data = cargar_datos_mitgcm(outdir_bay, outdir_nobay, use_wildcard=use_wildcard)

    # Estaciones consistentes para este escenario:
    stations_save, stations_meta = make_stations_from_grid(
        data["gr_bay"], j_centro=j_centro, radio=radio
    )

    prefix = f"{scenario_label}"
    extra_meta = {
        "scenario": scenario_label,
        "paths": {"bay": str(outdir_bay), "nobay": str(outdir_nobay)},
        "station_colors": stations_meta
    }

    # Si quieres persistir también los colores dentro de las tuplas, añade color:
    stations_with_color = [(i, j, lab, meta["color"]) 
                           for (i, j, lab), meta in zip(stations_save, stations_meta)]

    guardar_series_y_espectros(
        nc_bay     = data["nc_bay"],
        nc_nobay   = data["nc_nobay"],
        time       = data["time"],
        stations   = stations_with_color,      # se guardan colores en meta
        tmin_horas = tmin_horas,
        save_path  = save_dir,
        save_prefix= prefix,
        extra_meta = extra_meta
    )

# --------------------------------------------------------------
# Batch para varios escenarios (todos con las mismas reglas)
# --------------------------------------------------------------
def procesar_varios_escenarios(
    scenarios_dict: dict,            # {'realistic': Path(...), ...}
    subdir_bay: str = "run_expand",
    subdir_nobay: str = "run_nobay",
    tmin_horas: float = 27.0,
    save_dir: Path = Path("./npz_out"),
    use_wildcard: bool = False,
    j_centro: int = 315,
    radio: int = 50
):
    for scen_label, base_path in scenarios_dict.items():
        outdir_bay   = Path(base_path) / subdir_bay
        outdir_nobay = Path(base_path) / subdir_nobay
        print(f"\n=== Escenario: {scen_label} ===")
        guardar_series_y_espectros_escenario(
            scenario_label = scen_label,
            outdir_bay     = outdir_bay,
            outdir_nobay   = outdir_nobay,
            tmin_horas     = tmin_horas,
            save_dir       = save_dir,
            use_wildcard   = use_wildcard,
            j_centro       = j_centro,
            radio          = radio
        )
    print("\n[FIN] Procesamiento de todos los escenarios.")
# --------------------------------------------------------------
# Ejemplo de uso:
if __name__ == "__main__":
    scenarios = {
        'realistic': base_ivan / 'BayIW_Rectan',
        'linear'   : base_linear / 'BayIW_Rectan_linear',
        '2layer'   : base_ivan / 'BayIW_Rectan_2layer',
    }
    procesar_varios_escenarios(
        scenarios_dict = scenarios,
        subdir_bay     = "run_expand",
        subdir_nobay   = "run_nobay",
        tmin_horas     = 27.0,
        save_dir       = Path("./npz_out"),
        use_wildcard   = False,
        j_centro       = 315,
        radio          = 50
    )
# --------------------------------------------------------------
# Fin del script

# --------------------------------------------------------------
# AGREGADO: FUNCION POR ESCENARIO PARA GRAFICAR
# --------------------------------------------------------------
# - Lectura de npz y graficado
import json
import numpy as np
import matplotlib.pyplot as plt

def _cargar_npz_con_meta(series_npz_path, espectros_npz_path):
    """Carga .npz de series y espectros y devuelve (Dseries, Despectros, meta)."""
    dseries = np.load(series_npz_path, allow_pickle=False)
    despec  = np.load(espectros_npz_path, allow_pickle=False)

    # meta_json guardado como array de texto
    meta = json.loads(str(dseries["meta_json"]))

    return dseries, despec, meta

def graficar_estaciones_en_escenario(
    series_npz_path: str,
    espectros_npz_path: str,
    estaciones_orden: list = None,   # e.g. ["Seiche","A","B","C"]; si None, usa lo que hay en meta
    seiche_freq_cph: float = 0.303356,
    coriolis_f: float = 0.04,   # cph for lat ~ 41.5°N
    xlim_cph=(0, 0.8),
    ylim_psd=(1e-13, 1e-7),
    titulo_suffix: str = None        # para agregar extra al título
):
    """
    Lee un par de archivos .npz de un ESCENARIO (series + espectros) y
    grafica: (1) series temporales ΔEta de todas las estaciones del escenario
             (2) espectros PSD de todas las estaciones del escenario.
    """
    dseries, despec, meta = _cargar_npz_con_meta(series_npz_path, espectros_npz_path)
    scenario = meta.get("scenario", "escenario")

    # Orden de estaciones
    stations_meta = meta.get("stations", [])
    if estaciones_orden is None:
        estaciones_orden = [st["label"] for st in stations_meta]

    # Mapa label->color (si existe)
    color_map = {st["label"]: (st.get("color") or None) for st in stations_meta}

    # ---------- Figura 1: series temporales ----------
    plt.figure(figsize=(7,5))
    for lab in estaciones_orden:
        t_h = dseries[f"{lab}_t_horas"]
        eta = dseries[f"{lab}_eta_diff"]
        plt.plot(t_h, eta, label=f"Estación {lab}", color=color_map.get(lab, None))

    plt.xlabel("Tiempo (horas)")
    plt.ylabel("Δ Eta (m)")
    title = f"ΔEta (bahía - sin bahía) • {scenario}"
    if titulo_suffix:
        title += f"\n{titulo_suffix}"
    plt.title(title)
    plt.grid(True)
    plt.legend()
    plt.tight_layout()
    plt.show()

    # ---------- Figura 2: espectros ----------
    plt.figure(figsize=(7,5))
    plt.yscale("log", base=10)
    for lab in estaciones_orden:
        freq_cph = despec[f"{lab}_freq_cph"]
        PSD      = despec[f"{lab}_PSD"] * 3600.0   # Hz⁻¹ → cph⁻¹
        plt.plot(freq_cph, PSD, label=f"Estación {lab}", color=color_map.get(lab, None))

    plt.axvline(seiche_freq_cph, color='k', linestyle='--', linewidth=0.7, label='f seiche')
    plt.axvline(coriolis_f, color='gray', linestyle='--', linewidth=0.7, label='f Coriolis')
    plt.xlim(*xlim_cph)
    plt.ylim(*ylim_psd)
    plt.xlabel("Frecuencia (ciclos/hora)")
    plt.ylabel("PSD [m²/cph]")
    plt.title(f"Espectro ΔEta • {scenario}")
    plt.grid(True, which="both", ls=':')
    plt.legend()
    plt.tight_layout()
    plt.show()

# --------------------------------------------------------------
# AGREGADO: FUNCION PARA COMPARAR ESCENARIOS
# --------------------------------------------------------------
import json
import numpy as np
import matplotlib.pyplot as plt

def _auto_scenario_colors(escenarios_npz: dict, scenario_colors: dict | None):
    """
    Devuelve un dict {scenario: color} usando 'scenario_colors' si se pasa,
    o asignando colores de forma determinística desde Tab10 en orden de llave.
    """
    keys = list(escenarios_npz.keys())
    if scenario_colors:
        # normaliza a solo los escenarios presentes, manteniendo orden
        return {k: scenario_colors.get(k) for k in keys}
    cmap = plt.get_cmap("tab10")
    return {k: cmap(i % 10) for i, k in enumerate(keys)}

def comparar_escenarios_en_estacion(
    escenarios_npz: dict,
    # dict: { "linear": {"series": ".../linear_series.npz", "espectros":".../linear_espectros.npz"}, ... }
    estacion_label: str = "A",
    seiche_freq_cph: float = 0.303356,
    coriolis_f: float = 0.04,   # cph for lat ~ 41.5°N
    xlim_cph=(0, 0.8),
    ylim_psd=(1e-13, 1e-7),
    scenario_colors: dict | None = None,   # p.ej. {"realistic":"tab:green","linear":"tab:orange","2layer":"tab:purple"}
    titulo_suffix: str | None = None
):
    """
    Grafica (1) series ΔEta y (2) espectros PSD de UNA estación a través de
    varios escenarios, asignando un color único a cada escenario.
    """
    # Colores por escenario (constantes para ambas figuras)
    scen_color = _auto_scenario_colors(escenarios_npz, scenario_colors)

    series_data = []
    espec_data  = []

    for scen, paths in escenarios_npz.items():
        dser = np.load(paths["series"], allow_pickle=False)
        desp = np.load(paths["espectros"], allow_pickle=False)
        meta = json.loads(str(dser["meta_json"]))
        scenario_name = meta.get("scenario", scen)

        # extrae series
        try:
            t_h = dser[f"{estacion_label}_t_horas"]
            eta = dser[f"{estacion_label}_eta_diff"]
        except KeyError:
            raise KeyError(f"El archivo {paths['series']} no contiene claves para la estación '{estacion_label}'.")

        # extrae espectro
        try:
            f_cph = desp[f"{estacion_label}_freq_cph"]
            PSD   = desp[f"{estacion_label}_PSD"]
        except KeyError:
            raise KeyError(f"El archivo {paths['espectros']} no contiene claves para la estación '{estacion_label}'.")

        series_data.append((scenario_name, t_h, eta))
        espec_data.append((scenario_name, f_cph, PSD * 3600.0))   # Hz⁻¹ → cph⁻¹

    # ---------- Figura 1: ΔEta vs tiempo ----------
    plt.figure(figsize=(7,5))
    for scen_name, t_h, eta in series_data:
        plt.plot(t_h, eta, label=scen_name, color=scen_color.get(scen_name))

    plt.xlabel("Tiempo (horas)")
    plt.ylabel("Δ Eta (m)")
    title = f"Comparación de escenarios • Estación {estacion_label}"
    if titulo_suffix:
        title += f"\n{titulo_suffix}"
    plt.title(title)
    plt.grid(True)
    plt.legend(title="Escenario")
    plt.tight_layout()
    plt.show()

    # ---------- Figura 2: PSD ----------
    plt.figure(figsize=(7,5))
    plt.yscale("log", base=10)
    for scen_name, f_cph, PSD in espec_data:
        plt.plot(f_cph, PSD, label=scen_name, color=scen_color.get(scen_name))

    plt.axvline(seiche_freq_cph, color='k', linestyle='--', linewidth=0.7, label='f seiche')
    plt.axvline(coriolis_f, color='b', linestyle='--', linewidth=0.7, label='f Coriolis')
    plt.xlim(*xlim_cph)
    plt.ylim(*ylim_psd)
    plt.xlabel("Frecuencia (ciclos/hora)")
    plt.ylabel("PSD [m²/cph]")
    plt.title(f"Comparación de espectros • Estación {estacion_label}")
    plt.grid(True, which="both", ls=':')
    plt.legend(title="Escenario")
    plt.tight_layout()
    plt.show()

# Final del script
# --------------------------------------------------------------