database_maker_echo.py

"""
database_maker_echo.py — Build a per-shot ECHO command/target database.

For each shot in SHOT_LIST, fetches from MDSplus (D3D tree, PTDATA):
  - ECHO command angles   PTDATA EOSBANG{N}   poloidal mirror angle commanded
                                               by the ECHO optimiser
  - Duty cycles           PTDATA EOSBDC{P}    per gyrotron power supply
  - ECHO targets          PTDATA EOXTARGET    target deposition profile,
                                               reshaped to (n_time, N_RHO)
  - Gyro failure signal   PTDATA GYTFLTTO{N}  per-gyrotron fault flag
                                               (active when > FAULT_THRESHOLD)

Saves the combined {shot: {...}} dict to echo_command_database.pkl in this
directory.

Usage:
  python database_maker_echo.py
"""

import os
import pickle

import numpy as np

import mds_utils

ATLAS = mds_utils.ATLAS
OUT_DIR = os.path.dirname(os.path.abspath(__file__))
OUT_PATH = os.path.join(OUT_DIR, "echo_command_database.pkl")

# Active ECHO-steered gyrotrons and their PCS G-index
# (matches mds_utils._GYROTRON_G_INDEX)
GYROS = {"LEIA": 4, "R2D2": 5, "YODA": 8, "NASA": 9, "HAN": 11}

# G-index -> power supply number (matches mds_utils.fetch_ech_power fallback map)
GYRO_POWER_SUPPLY = {4: 1, 5: 2, 8: 2, 9: 4, 11: 3} # NOT NECESSARILY TRUE. CHECK ECHRES.

# Same shot list used in ECH_OPT/plot_state_vector.ipynb to build all_shots_data
SHOT_LIST = (list(np.arange(205822, 205839))
             + list(np.arange(205894, 205912))
             + list(np.arange(206527, 206543)))

N_RHO = 101
FAULT_THRESHOLD = 0.1


def _fetch_ptdata(conn, ptname, shot):
    """Return (data, time_ms) for a PTDATA point, or (None, None) on failure."""
    try:
        d = np.array(conn.get(f"PTDATA('{ptname}', {shot})").data(), dtype=float)
        t = np.array(conn.get(f"DIM_OF(PTDATA('{ptname}', {shot}))").data(), dtype=float)
        if d.size == 0:
            return None, None
        return d, mds_utils.ms_time(t)
    except Exception:
        return None, None


def fetch_command_angles(conn, shot):
    """PTDATA EOSBANG{N} — poloidal mirror angle commanded by the ECHO optimiser."""
    out = {}
    for gyro, g_idx in GYROS.items():
        d, t = _fetch_ptdata(conn, f"EOSBANG{g_idx}", shot)
        out[gyro] = {"pol_deg": d, "time_ms": t}
    return out


def fetch_duty_cycles(conn, shot):
    """PTDATA EOSBDC{P} — per-power-supply duty cycle."""
    out = {}
    for ps in sorted(set(GYRO_POWER_SUPPLY.values())):
        d, t = _fetch_ptdata(conn, f"EOSBDC{ps}", shot)
        out[ps] = {"duty_cycle": d, "time_ms": t}
    return out


def fetch_target(conn, shot):
    """PTDATA EOXTARGET — ECHO target deposition profile, reshaped to (n_time, N_RHO)."""
    d, t = _fetch_ptdata(conn, "EOXTARGET", shot)
    if d is None or d.size < N_RHO:
        return {"profile": None, "time_ms": None}
    n_time = d.size // N_RHO
    profile = d[: n_time * N_RHO].reshape(n_time, N_RHO)
    time_ms = t[:n_time] if t is not None else None
    return {"profile": profile, "time_ms": time_ms}


def fetch_gyro_failure(conn, shot):
    """PTDATA GYTFLTTO{N} — per-gyrotron failure flag (active when > FAULT_THRESHOLD)."""
    out = {}
    for gyro, g_idx in GYROS.items():
        d, t = _fetch_ptdata(conn, f"GYTFLTTO{g_idx}", shot)
        fault = (d > FAULT_THRESHOLD) if d is not None else None
        out[gyro] = {"fault": fault, "time_ms": t}
    return out


def fetch_shot(shot):
    if not mds_utils.MDS_AVAILABLE:
        raise RuntimeError("MDSplus unavailable.")
    conn = mds_utils.mds.Connection(ATLAS)
    conn.openTree("D3D", shot)
    try:
        return {
            "command_angles": fetch_command_angles(conn, shot),
            "duty_cycle": fetch_duty_cycles(conn, shot),
            "target": fetch_target(conn, shot),
            "gyro_failure": fetch_gyro_failure(conn, shot),
        }
    finally:
        conn.closeAllTrees()


def main():
    database = {}
    for shot in SHOT_LIST:
        shot = int(shot)
        print(f"Processing shot {shot}...", end=" ")
        try:
            database[shot] = fetch_shot(shot)
            print("done")
        except Exception as e:
            print(f"FAILED: {type(e).__name__}: {e}")

    print(f"\nProcessed {len(database)}/{len(SHOT_LIST)} shots successfully")
    failed = sorted(set(int(s) for s in SHOT_LIST) - set(database.keys()))
    if failed:
        print(f"Failed shots: {failed}")

    with open(OUT_PATH, "wb") as f:
        pickle.dump(database, f, protocol=pickle.HIGHEST_PROTOCOL)
    print(f"Saved -> {OUT_PATH}")


if __name__ == "__main__":
    main()