Small todos

[dpanici]

• Add the mhdin files for each shot to this repo
• add utility fxn to use PTDATA to grab data for a given shot and time, for all MP and flux loops available, and using the mhdin files (basically replace the code blocks below)
• also place the finite difference wrapper somewhere so it is not just copied from script to script (that should make scripts much shorter)

code block for getting signals

### Get coil currents from ptdata ###
try:
    from ptdata import fetch_ptdata, PtDataFetcher
except Exception as e:
    print("could not import ptdata, got exception")
    print(e)

PCS_SYS_D3 = ":/fusion/projects/codes/pcs/data/ptdata:/fusion/projects/codes/pcs/data/ptdata/uncomp:"
import numpy as np
from scipy.interpolate import interp1d

import sys

shot = int(sys.argv[1])  # full shot numver
time = int(sys.argv[2])  # time is 4 digit time


def compare_to_ptdata_diags_Bp_probes(
    ptnames,
    shot,
    time,
    coilset_in=None,
    coil_grid=None,
    eq=None,
    get_coilset_kwargs={},
):

    shot = round(shot)

    bad_names = []
    interpolable_ptnames = {}
    print(shot)
    for ptname in ptnames:
        # print(ptname)
        try:
            fetcher = PtDataFetcher(ptname, round(shot), sys_d3=PCS_SYS_D3)

            header = fetcher.header
            results = fetcher.fetch()
            interpolable_ptnames[ptname] = interp1d(
                x=results["times"], y=results["data"], kind="nearest"
            )
        except Exception as e:
            print(f"Shot {shot} does not have ptname {ptname}")
            print(e)
            bad_names.append(ptname)
    bad_name_inds = np.asarray(
        [i for i, name in enumerate(ptnames) if name not in bad_names]
    )
    probe_names_minus_bad_names = [name for name in ptnames if name not in bad_names]

    print("#" * 15)
    print(f"Shot {shot}")
    print("#" * 15)
    print(f"t = {time}")

    # poloidal probes
    Bps_measured = [
        interpolable_ptnames[name](time) for name in probe_names_minus_bad_names
    ]
    return Bps_measured, bad_names


import pickle
import f90nml

# get diag signal values and positions
# retrieve the flux loop and B probe info from the EFIT files
mhdin = f90nml.read("../mhdin_for_EFIT_112000_to__156014.dat")
# magnetic probes for poloidal field
efit_mp_names = mhdin["IN3"]["MPNAM2"]
for i in range(len(efit_mp_names)):
    efit_mp_names[i] = efit_mp_names[i].strip()

efit_mp_measurements, bad_names = compare_to_ptdata_diags_Bp_probes(
    efit_mp_names, shot, time
)

fn = "efit_mp_values.p"
with open(fn, "wb") as f:
    efit_mp_names = pickle.dump(efit_mp_measurements, f)

if len(bad_names) > 0:
    for i in range(len(bad_names)):
        bad_names[i] = bad_names[i] + "\n"

with open("mp_probes_to_remove.txt", "w+") as f:
    f.writelines(bad_names)

code block for cleaning signals given mhdin

import os
from desc.magnetic_fields import SplineMagneticField

coils = load("../g0198489.01600_coilset_DESC.h5")
if not os.path.exists("../g0198489.01600_coilset_DESC_splined_field_N500_grid.h5"):
    R = np.linspace(0.5, 3, 100)
    Z = np.linspace(-2, 2, 100)
    phi = np.linspace(0, 2 * np.pi, 2)
    splined_field = SplineMagneticField.from_field(
        coils, R, phi, Z, source_grid=LinearGrid(N=500), chunk_size=100
    )
    splined_field.save("../g0198489.01600_coilset_DESC_splined_field_N500_grid.h5")
else:
    splined_field = load("../g0198489.01600_coilset_DESC_splined_field_N500_grid.h5")
field = splined_field

# get diag signal values and positions
# retrieve the flux loop and B probe info from the EFIT files
mhdin = f90nml.read("../mhdin_197555_and_later.dat")
# magnetic probes for poloidal field
efit_mp_names = mhdin["IN3"]["MPNAM2"]
for i in range(len(efit_mp_names)):
    efit_mp_names[i] = efit_mp_names[i].strip()
efit_mp_coords_rpz = np.vstack(
    [mhdin["IN3"]["XMP2"], np.zeros_like(mhdin["IN3"]["XMP2"]), mhdin["IN3"]["YMP2"]]
).T
efit_mp_angles = np.asarray(mhdin["IN3"]["AMP2"])
dir_phis_Bp_probes = np.zeros_like(efit_mp_angles)
dir_Rs_Bp_probes = np.cos(efit_mp_angles / 180 * np.pi)
dir_Zs_Bp_probes = np.sin(efit_mp_angles / 180 * np.pi)
efit_mp_directions_rpz = np.vstack(
    [dir_Rs_Bp_probes, dir_phis_Bp_probes, dir_Zs_Bp_probes]
).T


# from the linked mhd.in stuff

efit_probe_dict = {}
mhdin = f90nml.read("../mhdin_197555_and_later.dat")
efit_probe_names = mhdin["IN3"]["MPNAM2"]
for i in range(len(efit_probe_names)):
    efit_probe_names[i] = efit_probe_names[i].strip()

# flux loops for poloidal flux
efit_flux_loop_names = mhdin["IN3"]["LPNAME"]
for i in range(len(efit_flux_loop_names)):
    efit_flux_loop_names[i] = efit_flux_loop_names[i].strip()
efit_flux_loop_coords_rz = np.vstack([mhdin["IN3"]["RSI"], mhdin["IN3"]["ZSI"]]).T

flux_loops = CoilSet(
    [
        FourierRZCoil(current=0.0, R_n=r, Z_n=z, sym=False)
        for r, z in efit_flux_loop_coords_rz
    ],
    check_intersection=False,
)

# efit_flux_loop_dict = {}
# for name, efit_loop_data in zip(efit_flux_loop_names, efit_data_dict["flux_loops"]):
#     efit_flux_loop_dict[name] = efit_loop_data

import pickle

fn = "efit_flux_loop_values.p"
with open(fn, "rb") as f:
    efit_flux_loop_values_with_plasma = pickle.load(f)

# convert from Wb/rad to Wb
efit_flux_loop_values_with_plasma_in_Wb_shifted = (
    (np.asarray(efit_flux_loop_values_with_plasma)) * 2 * np.pi
)
# all except first (PSF1A) are measured relative to PSF1A, so shift back to get absolute measurement
efit_flux_loop_values_with_plasma_in_Wb_shifted[
    1:
] += efit_flux_loop_values_with_plasma_in_Wb_shifted[0]

import pickle

fn = "efit_mp_values.p"
with open(fn, "rb") as f:
    efit_mp_values_with_plasma = pickle.load(f)

# DSL6U157 does not exist so remove that one, I know this from when I fetched the ptname and it failed to find this for 198489
with open("mp_probes_to_remove.txt", "r") as f:
    mp_names_to_remove = f.readlines()

mp_inds_remove = []
for mpname in mp_names_to_remove:
    mp_inds_remove.append(efit_mp_names.index(mpname.strip()))

# DSL6U157 does not exist so remove that one, I know this from when I fetched the ptname and it failed to find this for 198489
with open("mp_probes_to_remove.txt", "r") as f:
    mp_names_to_remove = f.readlines()

mp_inds_remove = []
for mpname in mp_names_to_remove:
    mp_inds_remove.append(efit_mp_names.index(mpname.strip()))
print(f"{mp_inds_remove=}")

print(f"efit_mp_values_with_plasma size {len(efit_mp_values_with_plasma)}")
print(f"efit mp names length: {len(efit_mp_names)}")

# further clean by removing ones ==0, and ones very large
mp_additionally_inds_remove_due_to_data = []
if np.where(np.array(efit_mp_values_with_plasma) < -0.3)[0].size > 0:
    mp_additionally_inds_remove_due_to_data += [
        ind for ind in np.where(np.array(efit_mp_values_with_plasma) < -0.3)[0]
    ]
if np.where(np.isclose(efit_mp_values_with_plasma, 0, atol=1e-6))[0].size > 0:
    mp_additionally_inds_remove_due_to_data += [
        ind for ind in np.where(np.isclose(efit_mp_values_with_plasma, 0, atol=1e-3))[0]
    ]


flux_loop_inds_remove = np.where(
    np.logical_or(
        np.isclose(efit_flux_loop_values_with_plasma_in_Wb_shifted, 0, atol=1e-1),
        np.array(efit_flux_loop_values_with_plasma_in_Wb_shifted) < -2,
    )
)[0]


with open("rogowski_plasma_current_value.p", "rb") as f:
    measured_plasma_current = pickle.load(f)


if len(mp_inds_remove):
    # mp_inds_remove.sort()
    # mp_inds_remove.reverse()
    # for ind_remove in mp_inds_remove:
    mp_meas_coords = np.delete(efit_mp_coords_rpz, mp_inds_remove, axis=0)
    mp_meas_dirs = np.delete(efit_mp_directions_rpz, mp_inds_remove, axis=0)
if len(mp_additionally_inds_remove_due_to_data):
    mp_meas_coords = np.delete(
        mp_meas_coords, mp_additionally_inds_remove_due_to_data, axis=0
    )
    mp_meas_dirs = np.delete(
        mp_meas_dirs, mp_additionally_inds_remove_due_to_data, axis=0
    )
    # also remove those values from the efit stuff
    efit_mp_values_with_plasma = np.delete(
        efit_mp_values_with_plasma, mp_additionally_inds_remove_due_to_data
    )