#!/usr/bin/env python3

import numpy as np
import pickle
import gc
import  os
import  time
import  shutil
import  logging
import  yaml
import  sys
from    datetime import datetime
import  pandas as pd
import  itertools
import  io
import  torch
from    nicqs.lib import RegularGridInterpolator, fit_angle_dependency
from    pathlib import Path
from    nicqs.lib.visualization import data_to_vtu, write_pvd_file, vector_to_vtu
from    nicqs.lib import ProgressBar, print_nicqs_logo
from    nicqs.lib.materials import MaterialProperties
from    nicqs.lib.debugging import DebuggingGraphs
from    nicqs.lib.utils import utils
from    nicqs.version import check_geometry_version

# Scipy
import  scipy
import  scipy.integrate
from    scipy import integrate
from    scipy.optimize import fsolve
from    scipy.interpolate import interp1d
from    scipy.integrate import solve_ivp

# For the removal of a matplotlib font warning.
import warnings
warnings.filterwarnings("ignore", message="findfont: Font family.*")

# Tim Mulder - 2D Non-Insulated Coil Quench Model
# tim.mulder@cern.ch // tim.mulder@gmail.com

class NI_Coil_Solver:
    # Operational parameters
    T0: list = None  # List of initial temperatures.
    I0: list = None  # List of initial currents.
    unique_regions: list = None # List of unique regions
    conductor_list: list = None # List of regions
    Regions: list = ['SC1']

    Icref = {}

    # Geometric parameters
    Mat: list = ['']  # List of conductor.
    lengths                 = None  # List of equivalent lengths of the path of the solenoids.
    cross_sections          = None  # List of equivalent cross-sectional areas of the solenoids.
    cross_sections_inner    = None
    cross_sections_outer    = None
    volumes: list           = None  # List of volumes of the solenoid.
    masses                  = None  # List of masses of the solenoid.
    radial_resistance       = None
    struct_resistance       = None
    parallel                = None
    serial                  = None
    conductor               = None
    magnet_mass             = np.inf
    turns_per_element       = None

    Index_array = {}

    # Conductor Properties
    Conductor_type: dict                = None
    Conductor_SC_properties: dict       = None
    Conductor_materials                 = None
    Conductor_materials_cross_section   = None
    Conductor_resistor                  = None
    Interpolation_functions             = {}    # Calculated fit of the current sharing between SC and Normal cond.
    x0                                  = None  # List of guesses for method 1, np.ones(len(SC con))
    ref_resistivity                     = None
    R_bypass                            = None
    ref_resistance                      = None
    ref_radial_resistance               = None
    ref_radial_resistance_material      = None
    ref_conductivity                    = None
    Tref_array                          = None
    Bref_array                          = None
    Conductor_conductivity_materials    = None
    Conductor_conductivity_thickness    = None
    Conductor_conductivity_width        = None

    # Circuit properties
    I_circuit_interpolation = None

    # Index of the connections
    Index_radial_connections = None  # For resistance calculation
    Index_inductive_connections = None  # For magnetic field calculation
    Index_struct_connections = None  # Structural conducting connections
    Index_normal_connections = None  # All normal conducting connection
    Index_temperatures = None  # Corresponding temperatures

    # Dictionary with indexes
    index_dict = None # A dictionary with the current indices of the circuits, coils and pancakes.

    # Information for the vtu creation.
    vtu_info = None
    vtu_post_BOOL = True
    vtu_num_elements = 50 # Maximum number vtu files that will be created.
    grid_information: list = []
    seed_points = None

    # Index of the connections per region
    Index_radial_connections_region = None  # For resistance calculation
    Index_inductive_connections_region = None  # For magnetic field calculation
    Index_struct_connections_region = None  # Structural conducting connections
    Index_normal_connections_region = None  # All normal conducting connection
    Index_temperatures_region = None  # Corresponding temperatures

    # b-array indices
    b_V_array_index = None
    b_dIdt_array_index = None

    # Circuits
    circuits = None
    unique_circuits = None
    circuits_list = None
    circuit_data = None
    circuit_interp_dIdt = dict()
    circuit_interp_I = dict()
    signs = None

    # PID ramp
    save_n_pid = 1
    savemod_pid = 1
    output_n = 1
    output_mod = 1
    max_rr = dict()
    min_rr = dict()
    setpoint = dict()
    nominal_current = dict()
    integral = dict()
    Kp = 0.05
    Kd = 1.0
    Ki = 0.00005
    rr = dict()
    rr_remap = dict() # Allows ramp rate remaps based on the ramp rate in another circuit.
    offset = dict()
    pid_bool = dict()
    PID_total_power = False
    ramp_up = True
    ramp_down = False

    # Initial conditions
    initial_conditions = {}
    Index_initial_Iind = {}
    Index_initial_IR = {}
    Index_initial_T = {}

    # Index post processing
    Index_post = None
    Index_post_inductive = None
    Index_post_resistive = None
    Index_post_thermal = None
    element_to_coil_nr = None

    # Cooling
    T_cooling = 10.0

    # Force
    Lt = None # Length * turns

    # Quench detection and protection
    V_detect = 1e99
    T_detect = 1e99
    quench_detected = False
    quench_detected_time = 1e99
    protection_delay = 0.005

    # Follow external ramp rate
    use_external_ramp_rate = False
    external_ramp_rate_path = ''
    external_ramp_rate = None

    # Follow T from file
    # If used, the thermal part will be turned off automatically.
    use_external_temperature = False
    external_temperature_path = ''
    external_temperature = None

    # Change the time constant
    R_multiplier: float | dict = None

    # Change the quench fraction
    quench_fraction = 1.0

    # Thermal part
    T_con = None    # Thermal connections for conduction
    power_dis_matrix = None

    # Solver settings
    t_offset: float                 = 0.0
    solver: str                     = 'vode'    # vode or lsoda
    dt: float                       = 0.001     # Initial step
    t0: list                        = [0.0, 100.0]
    atol: float                     = 1e-8
    rtol: float                     = 1e-10
    min_step: float                 = 1e-99
    max_step: float                 = 1e1
    stop_criterion_E: float         = 5000.0
    stop_criterion_T: float         = 300.0
    sampling_time_dI                = 1000.0
    sampling_time_dT                = 1.0
    max_timestep                    = 1.0
    min_timestep                    = 1.0
    t_offset_bool: bool             = False # If True, it will take the timestamp of the input line and adds this value to the sim time.
    keep_smallest_timestep: bool    = False # Only keeps the smallest timestep used. Normally a user would not use this.
    electrical_part: bool           = True # Full electrical part on or off, should be True.
    thermal_part: bool              = True # Full thermal part on or off
    thermal_conductivity_bool: bool = True #
    thermal_cooling_bool: bool      = False
    stop_on_E                       = False
    stop_on_T                       = False
    T_stop                          = True # Stops the solver ABOVE a certain temperature.
    E_stop                          = True # Stops the solver BELOW a certain stored energy, useful for quench simulation, not for ramp.
    results                         = None # Result of the solver.
    sampling_time_custom            = None
    overwrite_sample_time_custom    = True  # Overwrites the custom sampling time if a quench is detected.
    quench_sampling_time            = 0.001 # When a quench is detected, the sampling time is set to this value
    store_n                         = 0
    save_n                          = 0
    save_mod                        = 10
    start                           = 0.0
    fit_func_name                   = None
    FLAG_USE_STEAM_MAT_LIB          = True
    FLAG_SUMMARY_ONLY               = False
    FLAG_MINIMAL_OUTPUT             = False
    FLAG_NO_OUTPUT                  = False
    FLAG_DAKOTA                     = False
    FLAG_Finished                   = False

    # Magnetic field for material properties
    set_B = None # if not None, the magnetic field will be set to this value.
    B_min = 0.11 # if not None, the minimum magnetic field will be set to this value.

    # yaml files:
    specific_output_path    = None
    geometry_yaml_path      = None
    solver_settings_path    = None
    input_yaml_path         = None
    path_dict               = None

    # Material and Ic dictionaries:
    deg_dep: dict = {}
    SC_Ic_Fit = None
    Normal_rho_FIT = None
    Normal_k_FIT = None
    Normal_Cv_FIT = None
    SC_Cur_Sharing = None
    SC_Ic_Scaling = None
    
    # OPTIMIZATION: Cache for frequently computed values to reduce STEAM material calls
    _material_cache = {}
    _last_material_inputs = {}
    _cached_resistivity = {}

    # Solver input
    inv     = []  # Full inverted matrix
    inv_1   = []  # Part to get the dI/dt
    inv_2   = []  # Part to get the voltage
    matrix_size: int = None
    Mutual  = None  # Mutual inductance matrix (loaded from file)

    # Power supply
    ramp_bool   = True      # Normal ramp behavior, turned off when a quench is detector or if set to False
    PSU_off     = False     # Turns of the PSU
    PSU_off_on_Quench_Detection = False
    PSU_off_time_constant       = 0.002
    PSU_off_time                = 1e99 # Turns the PSU off at this time.
    I_circuit = {} # Dictionary of currents in the circuit

    # Calculated matrices / arrays:
    power_dis_matrix_sum = None # Needed for averaging the right magnetic field on the resisitive nodes.
    Conductor_materials_cross_section_sum = {} # Sum of the cross-sections

    # Paths and other related variables
    mutual_inductance_matrix_path: str = ''  # Path to the mutual inductance matrix, prepared using XXX script.
    Bz_matrix_path: str = ''  # Path to the Bz matrix, prepared using YYY script.
    Br_matrix_path: str = ''  # Path to the Br matrix, prepared using YYY script.
    Bz_center_path: str = ''
    Br_center_path: str = ''
    Bz_peak_path: str = ''
    Br_peak_path: str = ''
    Bz_full_path: str = ''
    Br_full_path: str = ''
    power_dis_matrix_path: str = ''
    thermal_connections_matrix_path: str = ''
    sim_name: str = ''
    tag: str = ''
    delimiter = ','
    results_path = None
    sim_output = None
    names = None

    # Input from the geometry yaml:
    electrical_matrix_path  = None
    Nr_of_nodes             = None
    Nr_of_Connections       = None
    electrical_matrix       = None
    Br_conductor_path       = None
    Bz_conductor_path       = None

    # Debugging and analytics
    max_current: list = []
    min_current: list = []
    nan_detected: list = []
    inf_detected: list = []

    # Capacitor Discharge
    Capacitor_Voltage = 0.0
    Capacitor_Current = 0.0
    Capacitor_Discharge_on_Quench_Detection = False
    Capacitor_Discharge_delay = 0.005 # 5 ms delay

    # Quench Heaters
    Heaters_on_Quench_Detection = False
    Quench_Heater_power = 1000.0
    Quench_Heater_time = 1.0
    Quench_Heater_positions = None

    # Capacitor Discharge
    Cap_discharge = False
    Cap_discharge_time = 999999999.0
    Discharge_R0 = 0.0
    Discharge_L0 = 1e-5
    Discharge_C0 = 100e-3
    Dis_t0 = [0.0, 0.02]
    Dis_t_array = [Dis_t0[0]]
    Dis_List_Coils = None
    Dis_List: list = []
    Dis_List_Resistance = None
    fit_current = None
    fit_dI = None

    # Cooling:
    Thelium: float   = 4.2
    c1: (int, float) = 500.     # For He1 cooling
    c2: (int, float) = 5e4      # For He1 cooling
    c3: (int, float) = 250      # For He1 cooling
    c4: (int, float) = 200      # Not used, transient cooling regime
    c5: (int, float) = 200      # [600] Cooling parameter for superfluid    -> H(T) = c5*(T1^p-T0^p)
    c6: (int, float) = 200      # [250] Cooling parameter for filmboiling   -> H(T) = c6*(T1-T0)
    p: (int, float) = 4.0       # [3.4] Cooling parameter for superfluid    -> H(T) = c5*(T1^p-T0^p)
    h1: (int, float) = 10       # For He1 cooling
    h2: (int, float) = 1e4      # For He1 cooling
    Pmax: (int, float) = 35e3   # Maximum power before going to the film boiling regime
    Tstar: float = 0.           # Calculated
    T_star_1: float = 0.        # Calculated
    T_star_2: float = 0.        # Calculated
    coolant_temperature: float = 4.5
    coolant_h: float = 500.0
    max_cooling: float = 1e99

    max_cooling: bool = False           # Is set to True, where will be a maximum cooling amount, all excess heat won't be cooled away.
    max_cooling_value: float = 0.0      # Maximum cooling value.

    inner_cooling_bool: bool = True
    outer_cooling_bool: bool = True

    thermal_timer = -10.0
    thermal_timer_sig_c = 100
    cryocooler_fit = None

    # Thermal Connections
    T_con_inner = None
    T_con_outer = None
    Thermal_alpha = None
    Thermal_area = None
    Thermal_index = None
    Thermal_resistance = None
    T1_con = None
    T2_con = None
    thermal_element_to_coil_nr = None # From Geometry creator - used for the calculation of the loss per magnet
    thermal_element_to_solenoid_nr = None # From geometry creator - used for the calculation of the loss per solenoid

    # Heater
    heater = True   # If true, the heater can be turned on
    heater_nodes = [0]
    heater_power = 0.0
    heater_time = 0.0
    heater_sig_c = 500.0
    heater_stop = 80.0

    # Torch variables:
    # Some variables are moved to torch specific variables while keeping the normal numpy/list ones.
    lengths_torch = None
    volumes_torch = None
    b0 = None
    B0 = None
    T0_tensor = None

    # Profiler
    profile_bool = True

    # Distributed heating
    R_joints = 0.0

    # GPU
    torch_device = 'cuda' # 'cuda' or 'cpu'
    torch_set_float = 'float64'
    torch_float_type = torch.float64

    # File bools:
    bool_summary = True
    bool_current = True
    bool_temperature = True
    bool_voltage = True
    bool_magnetic_field = True
    bool_force = True
    bool_loss = True

    # Files
    init_path = None
    output_path = None
    file_summary = None
    file_temperature = None
    file_current = None
    file_current_R = None
    file_P_AC = None
    file_Prad = None
    file_Pin = None
    file_P_AC_density = None
    file_Prad_density = None
    file_Pin_density = None
    file_magnetic_field = None
    file_analytics_maxI = None
    file_analytics_minI = None
    file_analytics_nan = None
    file_analytics_inf = None
    file_voltage = None
    file_Fz = None
    file_Fr = None
    file_Fz_density = None
    file_Fr_density = None
    file_Fz_density_peak = None
    file_Fr_density_peak = None
    file_Bz_conductor = None
    file_Br_conductor = None
    file_I_density = None
    file_Power_per_Coil_Element = None
    file_Power_per_Solenoid = None
    file_Fz_per_Coil = None
    file_Fr_per_Coil = None
    file_Fz_per_Pancake = None
    file_Fr_per_Pancake = None
    file_R_Radial_per_Coil = None
    file_R_Radial_per_Pancake = None
    Bz_Conductor = None
    Br_Conductor = None
    Bz_center = None
    Br_center = None
    Bz_peak = None
    Br_peak = None

    # Additional attributes for files:
    columns: dict = {}

    # Result arrays
    t_array: list = []  # Time array
    T_array: list = []  # Temperature array
    I_array: list = []  # Current array
    V_array: list = []  # Voltage array
    IR_array: list = []  # List with current in the resistor
    E_array: list = []  # Stored energy array
    I_circuit_array: list = [] # Circuit current
    V_circuit_array: list = []  # Circuit voltage
    I_circuit_last: dict = {}  # Circuit current
    V_circuit_last: list = []  # Circuit voltage
    Bz_center_array: list = [] # Central Field
    Br_center_array: list = []  # Central Field
    P_AC_array: list = []  # Power dissipation array
    Prad_array: list = []  # Power dissipation array
    Pin_array: list = []  # Input power
    P_AC_density_array: list = []  # Power dissipation array
    Prad_density_array: list = []  # Power dissipation array
    Pin_density_array: list = []  # Input power
    D_array: list = []  # Dissipated energy array
    B_array: list = []  # Total magnetic field
    Bz_array: list = []  # Bz on conductor array
    Br_array: list = []  # Br on conductor array
    Fz_array: list = [] # Fz N
    Fz_density_array: list = []  # Fz N/m3
    Fz_density_peak_array: list = []  # Fz N/m3
    Fr_array: list = []  # Fz N
    Fr_density_array: list = []  # Fr N/m3
    Fr_density_peak_array: list = []  # Fr N/m3
    I_density_array: list = [] # I density / A/mm2
    Capacitor_Voltage_array: list = [] # V
    Capacitor_Current_array: list = [] # A
    Quench_detected_array: list = [] # Quench detection status
    Power_per_Coil_array: list = [] # Sum of the power per coil element
    Power_per_Solenoid_array: list = []
    Fz_per_Coil_array: list = [] # Sum of the power per coil element
    Fr_per_Coil_array: list = [] # Sum of the power per coil element
    Fz_per_Pancake_array: list = [] # Sum of the power per pancake element
    Fr_per_Pancake_array: list = [] # Sum of the power per pancake element
    R_Radial_per_Coil_array: list = [] # Sum of the radial resistance per coil
    R_Radial_per_Pancake_array: list = [] # Sum of the radial resistance per pancake
    R_Radial_array: list = [] # All radial resistances
    R_Conductor_array: list = [] # All conductor resistances.
    Current_Sharing_x_array: list = [] # All values for the current sharing
    PSU_power_array: list = []
    P_cooling_array: list = []
    SC_loss_array: list = []
    Radial_loss_array: list = []
    Field_Angle_array: list = []
    Critical_Ratio_array: list = []
    solver_t_array: list = [] # Solver time
    E_dis: None # Integrated dissipated power

    # Postprocessing and plotting
    Ic_T_fixed = 5.0 # K
    Ic_B_fixed = 40.0 # K
    post_process: dict = {}

    def __init__(self, input_yaml, path_dict = None):
        """
        Initialize the NI Coil Solver with configuration and setup.

        This constructor initializes the solver by loading configuration files,
        checking hardware availability (CUDA), preparing output directories,
        importing electromagnetic matrices, calculating material properties,
        and executing the simulation with automatic post-processing.

        Args:
            input_yaml: Path to the input YAML file containing simulation
                parameters and references to geometry/solver configuration.
            path_dict: Optional dictionary with custom paths to required files.
                If None, paths from the YAML configuration are used.

        Side Effects:
            - Prints NICQS logo to console
            - Creates output directories and log files
            - Loads and parses multiple YAML configuration files
            - Imports large matrices from disk (mutual inductance, B-fields)
            - Executes the full simulation via start_solver()
            - Generates VTU visualization files and post-processing graphs
            - May be interrupted with Ctrl+C, triggering partial post-process
        """
        # Print the NIQCS ascii logo
        print_nicqs_logo()

        # Initialize materials library
        self.materials = MaterialProperties()
        
        # Loads the input yaml file.
        # The input yaml file has input parameters and links to the correct geometry and solver_setting yaml
        self.input_yaml_path = input_yaml
        self.path_dict = path_dict
        self.load_yaml()

        # Checks if CUDA is available:
        self.check_cuda_availability()

        # Prepares the logging and results file.
        self.prepare_files()

        # Initialize debugging graphs library
        self.debug_graphs = DebuggingGraphs(self)

        # Create interpolation functions
        self.create_ramp_rate_interpolation()
        self.create_time_step_interpolation()

        # Makes the required connection matrix for the solver
        self.import_matrix()

        # Makes the material properties fits
        self.calc_init()
        self.load_ramp_rate()

        # Calculating the conductor properties
        self.calculate_geometric_properties()

        # Starts the solver.
        try:
            self.start_solver()
            self.cleaning()
            # Post processing
            self.debug_graphs.post_graph()
            self.results_to_vtu()
        except KeyboardInterrupt:
            print("\nSimulation interrupted by user (Ctrl+C).")
            # choice = input("Do you want to run post-processing anyway? [y/N]: ").strip().lower()
            # if choice in ('y', 'yes'):
            self.cleaning()
            self.results_to_vtu()
            self.debug_graphs.post_graph()

        self.results_post_process() # For comsol export, should be more elegant.

    def check_cuda_availability(self) -> None:
        """
        Check CUDA availability and configure device settings.
        
        This method verifies if CUDA is available on the system and configures
        the torch device accordingly. If CUDA is not available or explicitly 
        set to 'cpu', it falls back to CPU processing. Also configures the
        floating point precision for torch tensors.
        
        Returns:
            None
            
        Side Effects:
            - Sets self.torch_device to 'cpu' or 'cuda'
            - Sets self.torch_float_type based on self.torch_set_float
            - Clears GPU memory cache if CUDA is available
            - Logs device configuration and memory status
        """

        if not torch.cuda.is_available() or self.torch_device=='cpu':
            self.torch_device = 'cpu'
            print('Torch device is set to "cpu" since cuda is not available or it was already set to cpu.')
        else:
            gc.collect()
            torch.cuda.empty_cache()
            print('Memory allocated:', torch.cuda.memory_allocated(device=self.torch_device))
            print('Max memory allocated:', torch.cuda.max_memory_allocated(device=self.torch_device))
            print('Emptied GPU memory cache.')

        if self.torch_set_float == 'float64':
            self.torch_float_type = torch.float64
        elif self.torch_set_float == 'float32':
            self.torch_float_type = torch.float64
        elif self.torch_set_float == 'float16':
            self.torch_float_type = torch.float16
        else:
            logging.info('Unknown torch float type provided, setting it to float64.')
            self.torch_float_type = torch.float64

    def calc_init(self):
        """
        Perform initial calculations and convert NumPy arrays to PyTorch tensors.

        This method prepares the solver by converting key data structures from
        NumPy arrays to PyTorch tensors with the appropriate dtype and device
        (CPU/CUDA). It also performs initial calculations for thermal properties,
        flattens nested lists, and validates conductor configurations.

        Side Effects:
            - Converts multiple arrays to PyTorch tensors on configured device
            - Initializes PID controller parameters if enabled
            - Calculates thermal star temperatures (Tstar, T_star_1, T_star_2)
            - Logs mass and length validation checks
            - May apply resistance multiplier to radial_resistance
        """

        # Creating some often used tensors:
        self.torch_zero = torch.tensor(0.0, dtype=self.torch_float_type, device=self.torch_device)

        # To float type and to the device
        self.Br_Conductor = torch.tensor(self.Br_Conductor, dtype=self.torch_float_type, device=self.torch_device)
        self.Bz_Conductor = torch.tensor(self.Bz_Conductor, dtype=self.torch_float_type, device=self.torch_device)
        self.Br_peak = torch.tensor(self.Br_peak, dtype=self.torch_float_type, device=self.torch_device)
        self.Bz_peak = torch.tensor(self.Bz_peak, dtype=self.torch_float_type, device=self.torch_device)

        # Current sign to torch:
        self.signs = torch.tensor(self.signs, dtype=self.torch_float_type, device=self.torch_device) if self.signs is not None else torch.tensor(np.ones(self.circuits), dtype=self.torch_float_type, device=self.torch_device)

        # Power distribution matrix to torch.
        if self.torch_float_type != torch.float32:
            self.power_dis_matrix = torch.tensor(self.power_dis_matrix, dtype=self.torch_float_type, device=self.torch_device)
        else:
            self.power_dis_matrix = torch.tensor(self.power_dis_matrix, dtype=self.torch_float_type, device=self.torch_device).to_sparse()
        self.power_dis_matrix_sum = torch.sum(self.power_dis_matrix, dim=1) - 1
        self.lengths_torch = torch.tensor(self.lengths, dtype=self.torch_float_type, device=self.torch_device)
        self.volumes_torch = torch.matmul(self.power_dis_matrix, torch.tensor(self.volumes, dtype=self.torch_float_type, device=self.torch_device))
        self.Mutual = torch.tensor(self.Mutual, dtype=self.torch_float_type, device=self.torch_device)

        # Thermal conductivity tensors:
        self.T1_con_tensor = torch.tensor(self.T1_con, dtype=torch.int, device=self.torch_device)
        self.T2_con_tensor = torch.tensor(self.T2_con, dtype=torch.int, device=self.torch_device)

        # Cross-section to tensor:
        self.cross_sections_inner = torch.tensor(self.cross_sections_inner, dtype=self.torch_float_type, device=self.torch_device)
        self.cross_sections_outer = torch.tensor(self.cross_sections_outer, dtype=self.torch_float_type, device=self.torch_device)

        self.I0 = [j for i in self.I0 for j in i]
        self.T0 = [j for i in self.T0 for j in i]
        # self.radial_resistance = torch.tensor([j for i in self.radial_resistance for j in i], dtype=self.torch_float_type, device=self.torch_device)
        self.radial_resistance = torch.tensor(self.radial_resistance, dtype=self.torch_float_type, device=self.torch_device)
        # self.Icref = self.Iref / self.Iop / self.Ic(self.Tref, self.Bref)
        self.cross_sections = np.array(self.cross_sections)
        self.lengths = np.array(self.lengths)
        self.Index_temperatures = np.array(self.Index_temperatures)
        self.masses = np.array(self.masses)
        self.Tstar = ((self.Pmax / self.c5) + self.Thelium ** self.p) ** (1 / self.p) - self.Thelium
        self.T_star_1 = self.h1 / self.c1 + self.Thelium
        self.T_star_2 = (self.h2 / self.c2) ** (1 / 2.5) + self.Thelium
        self.I_circuit = np.zeros(len(self.unique_circuits))
        if self.pid_bool:
            for i in self.unique_circuits:
                self.rr[i] = 0.0
                self.integral[i] = 0.0
                self.offset[i] = 0.0
        else:
            for i in self.unique_circuits:
                self.pid_bool[i] = False

        # Logging important values:
        logging.debug("### CHECK ### Sum of masses: %s kg" % str(np.round(np.sum(self.masses),2)))
        logging.debug("### CHECK ### Coil mass from geometry: %s kg" % str(np.round(np.sum(self.magnet_mass), 2)))
        logging.debug("### CHECK ### Sum of length: %s m" % str(np.round(np.sum(self.lengths), 2)))

        # Needs to be a numpy array:
        self.turns_per_element = np.array(self.turns_per_element)

        for i in self.unique_regions:
            self.Conductor_materials_cross_section_sum[i] = torch.sum(torch.tensor(self.Conductor_materials_cross_section[i], dtype=self.torch_float_type, device=self.torch_device))

        # Manually changing the time constant of the coil after geometry creation.
        if self.R_multiplier is not None:
            if type(self.R_multiplier) in [int, float, np.float64]:
                self.radial_resistance = self.radial_resistance * self.R_multiplier
            elif type(self.R_multiplier) == dict:
                for i in self.R_multiplier.keys():
                    self.radial_resistance[self.Index_initial_IR[i]] = self.radial_resistance[self.Index_initial_IR[i]] * self.R_multiplier[i]

        # Cryocooler data
        c_T = [0, 1,2.5,3,3.5,4,6,8,10,12,14,16,30,50,156.2,227.7,282.1,300]
        parasitic = 1.0 # 1 W of parasitic heat
        c_P = np.array([0,0,0.1,0.47,0.98,1.55,3.3,9.3,12.5,14.4,16,18,24.7,29.5,42,50,51,52]) - parasitic
        c_P = np.where(c_P < 0.0, 0.0, c_P)
        self.cryocooler_fit = interp1d(c_T, c_P, bounds_error=False, fill_value=(c_P[0], c_P[-1]))

        # Capacitor discharge
        if self.Cap_discharge or self.Capacitor_Discharge_on_Quench_Detection:
            logging.info("Making capacitor discharge fit and setting up the index lists.")
            if self.Dis_List_Coils is not None:
                # selecting all nodes
                self.Dis_List = []
                for i in self.Dis_List_Coils:
                    self.Dis_List.append(np.unique(np.array(self.Index_post_thermal)[self.Index_post[i]])) # Index post has connections per coil, post_thermal converts it to all nodes, then I take all unique nodes
                self.Dis_List = list(itertools.chain.from_iterable(self.Dis_List))
                self.Dis_List_Resistance = list(np.array(self.Index_normal_connections)[self.Dis_List])
            logging.info("Finished Making capacitor discharge fit and setting up the index lists.")

    def load_ramp_rate(self):
        """
        Load external ramp rate data from pickle file if configured.
        
        This method loads a pre-computed ramp rate profile from a pickle file
        when use_external_ramp_rate is enabled, allowing the simulation to follow
        an externally defined current ramping schedule.
        
        Args:
            None
            
        Returns:
            None
            
        Side Effects:
            - Sets self.external_ramp_rate to loaded interpolation function
            - Only executes if self.use_external_ramp_rate is True
        """
        if self.use_external_ramp_rate:
            ramp_rate_path = self.specific_output_path + os.sep + self.external_ramp_rate_path
            self.external_ramp_rate = self.load_pkl_file(ramp_rate_path)


    def load_yaml(self):
        """
        Load and parse all YAML configuration files for the simulation.

        This method loads three YAML files in sequence: the input YAML (with
        simulation parameters), the geometry YAML (with coil geometry data),
        and the solver settings YAML (with numerical solver configuration).
        All parameters from these files are set as class attributes.

        Side Effects:
            - Reads input_yaml_path, geometry_yaml_path, solver_settings_path
            - Sets all keys from YAML files as instance attributes via setattr
            - Validates presence of required paths (path_dict, geometry, solver)
            - May use default paths if geometry/solver settings not specified
            - Raises Exception if required paths are missing
        """

        # Reads the contents of the input yaml file
        with open(self.input_yaml_path) as file:
            # Loads the contents of the input yaml.
            contents = yaml.load(file, Loader=yaml.FullLoader)

        # Checks if paths are provided in either the path_dict or in the input yaml.
        if 'path_dict' in contents.keys() and self.path_dict is None:
            setattr(self, 'path_dict', contents['path_dict'])
            contents.pop('path_dict', None) # removing the key from the dictionary as we don't want to overwrite a custom path_dict.

        # Some Exceptions.
        if self.path_dict is None:
            raise Exception("No path dictionary (path_dict) provided.")

        if 'geometry_yaml_path' not in self.path_dict.keys():
            raise Exception("No geometry path provided in the path dictionary.")

        if 'solver_settings_path' not in self.path_dict.keys():
            raise Exception("No solver settings path provided in the path dictionary.")

        if 'specific_output_path' not in self.path_dict.keys():
            raise Exception("No base path provided in the path dictionary.")

        # Loading the paths from the path dictionary.
        for i in self.path_dict.keys():
            setattr(self, i, self.path_dict[i])

        # Exception for if there is no output path specified.
        if self.specific_output_path is None:
            raise Exception("No path to the simulation folder is provided.")

        # Checks if the geometry, and solver settings paths are provided, otherwise they will be loaded from the same folder as the input_yaml file.
        if self.geometry_yaml_path is None:
            self.geometry_yaml_path = Path(self.specific_output_path) / 'geometry.yaml'

        if self.solver_settings_path is None:
            self.solver_settings_path = Path(self.specific_output_path) / 'solver_settings.yaml'

        with open(self.geometry_yaml_path) as file:
            contents_geom = yaml.load(file, Loader=yaml.FullLoader)
        # Sets the values in the class
        for i in contents_geom:
            setattr(self, i, contents_geom[i])

        # Check geometry version compatibility
        if 'geometry_version' not in contents_geom:
            warning_msg = (
                "\n" + "="*70 + "\n" +
                "WARNING: Geometry file missing version information!\n" +
                "="*70 + "\n" +
                "  The geometry file was created with an older version of the code\n" +
                "  that did not include version tracking. This may cause errors.\n" +
                "\n" +
                "  RECOMMENDED ACTION:\n" +
                "  Please regenerate the geometry using nicqs/geometry.py with\n" +
                "  the current version of the code.\n" +
                "="*70 + "\n"
            )
            print(warning_msg)
        else:
            is_current, warning_msg = check_geometry_version(
                contents_geom['geometry_version']
            )
            if not is_current:
                print(warning_msg)

        # Reads the contents of the yaml file
        with open(self.solver_settings_path) as file:
            contents_solv = yaml.load(file, Loader=yaml.FullLoader)
        # Sets the values in the class
        for i in contents_solv:
            setattr(self, i, contents_solv[i])

        # Sets the values in the class
        for i in contents:
            setattr(self, i, contents[i])

    @staticmethod
    def load_pkl_file(pkl_path: str):
        """
        Load a pickle file with error handling.
        
        Args:
            pkl_path (str): Path to the pickle file to load
            
        Returns:
            object or None: The loaded object if successful, None if failed
        """
        try:
            with open(pkl_path, 'rb') as fin:
                data = pickle.load(fin)
        except OSError as e:
            logging.info("Loading of %s has failed: %s" % (pkl_path, str(e)))
            return None
        except Exception as e:
            logging.info("Error loading %s: %s" % (pkl_path, str(e)))
            return None
        else:
            logging.debug("Successfully loaded %s" % pkl_path)
            return data

    def prepare_files(self):
        """
        Prepare output directories, logging, and result files.

        This method sets up the complete output infrastructure for the
        simulation, including creating directories, copying configuration
        files, loading pre-computed matrices (B-fields, power distribution),
        and initializing CSV output files with appropriate column headers.

        Side Effects:
            - Creates results directory and sim-specific subdirectory
            - Starts logger via start_logger()
            - Copies solver script and YAML config files to output directory
            - Loads magnetic field matrices (Br, Bz at conductor, peak, center)
            - Loads power_dis_matrix and thermal_connections_matrix
            - Configures output flags based on FLAG_SUMMARY_ONLY, etc.
            - Defines file paths for all CSV outputs (temperature, current,
              voltage, forces, losses, etc.)
            - Sets up column headers for CSV files
        """

        # Creates the general output folder if it doesn't exist.
        self.results_path = self.specific_output_path + os.sep + "results"

        try:
            os.mkdir(self.results_path)
        except OSError:
            print("Creation of the directory %s failed" % self.results_path)
        else:
            print("Successfully created the directory %s " % self.results_path)

        # Creates the specific simulation output folder.
        if self.sim_name == '':
            self.sim_name = datetime.now().strftime("%Y-%m-%d")

        if self.tag != "":
            self.sim_name = self.sim_name + ' - ' + self.tag

        self.sim_output = utils.uniquify(
            self.results_path + os.sep + 'result' + ' - ' + self.sim_name)

        try:
            os.mkdir(self.sim_output)
        except OSError:
            print("Creation of the directory %s failed" % self.sim_output)
        else:
            print("Successfully created the directory %s " % self.sim_output)

        self.start_logger()

        copied_script_name = time.strftime("%Y-%m-%d_%H%M") + '_' + os.path.basename(__file__)
        shutil.copy(__file__, self.sim_output + os.sep + copied_script_name)
        shutil.copy(self.solver_settings_path, self.sim_output + os.sep + 'solver_settings.yaml')
        shutil.copy(self.input_yaml_path, self.sim_output + os.sep + 'circuit.yaml')
        shutil.copy(self.geometry_yaml_path, self.sim_output + os.sep + 'geometry.yaml')

        ############################################################################

        self.Br_Conductor = self.load_pkl_file(self.Br_conductor_path)
        self.Bz_Conductor = self.load_pkl_file(self.Bz_conductor_path)
        self.Br_peak = self.load_pkl_file(self.Br_peak_path)
        self.Bz_peak = self.load_pkl_file(self.Bz_peak_path)
        self.Bz_center = self.load_pkl_file(self.Bz_center_path)
        self.Br_center = self.load_pkl_file(self.Br_center_path)
        self.power_dis_matrix = self.load_pkl_file(self.power_dis_matrix_path)

        ##########################################################################

        if self.FLAG_SUMMARY_ONLY:
            self.bool_summary = True
            self.bool_current = False
            self.bool_temperature = False
            self.bool_voltage = False
            self.bool_magnetic_field = False
            self.bool_force = False
            self.bool_loss = False
        elif self.FLAG_MINIMAL_OUTPUT:
            self.bool_summary = True
            self.bool_current = True
            self.bool_temperature = True
            self.bool_voltage = False
            self.bool_magnetic_field = False
            self.bool_force = False
            self.bool_loss = False
        elif self.FLAG_NO_OUTPUT:
            self.bool_summary = False
            self.bool_current = False
            self.bool_temperature = True
            self.bool_voltage = False
            self.bool_magnetic_field = False
            self.bool_force = False
            self.bool_loss = False
        else:
            self.bool_summary = True
            self.bool_current = True
            self.bool_temperature = True
            self.bool_voltage = True
            self.bool_magnetic_field = True
            self.bool_force = True
            self.bool_loss = True

        ##########################################################################
        # Output files
        self.file_summary = self.sim_output + os.sep + 'summary.csv'
        self.file_temperature = self.sim_output + os.sep + 'temperature.csv'
        self.file_current = self.sim_output + os.sep + 'current.csv'
        self.file_current_R = self.sim_output + os.sep + 'current_resistor.csv'
        self.file_P_AC = self.sim_output + os.sep + 'P_AC.csv'
        self.file_Prad = self.sim_output + os.sep + 'Prad.csv'
        self.file_Pin = self.sim_output + os.sep + 'Pin.csv'
        self.file_P_AC_density = self.sim_output + os.sep + 'P_AC_density.csv'
        self.file_Prad_density = self.sim_output + os.sep + 'Prad_density.csv'
        self.file_Pin_density = self.sim_output + os.sep + 'Pin_density.csv'
        self.file_magnetic_field = self.sim_output + os.sep + 'B_total.csv'
        self.file_voltage = self.sim_output + os.sep + 'voltage.csv'
        self.file_Fz = self.sim_output + os.sep + 'Fz.csv'
        self.file_Fr = self.sim_output + os.sep + 'Fr.csv'
        self.file_Fz_density = self.sim_output + os.sep + 'Fz_density.csv'
        self.file_Fr_density = self.sim_output + os.sep + 'Fr_density.csv'
        self.file_Fz_density_peak = self.sim_output + os.sep + 'Fz_density_peak.csv'
        self.file_Fr_density_peak = self.sim_output + os.sep + 'Fr_density_peak.csv'
        self.file_Bz_conductor = self.sim_output + os.sep + 'Bz_conductor.csv'
        self.file_Br_conductor = self.sim_output + os.sep + 'Br_conductor.csv'
        self.file_I_density = self.sim_output + os.sep + 'I_density.csv'
        self.file_Power_per_Coil_Element = self.sim_output + os.sep + 'Power_per_Coil.csv'
        self.file_Power_per_Solenoid = self.sim_output + os.sep + 'Power_per_Solenoid.csv'
        self.file_Fz_per_Coil = self.sim_output + os.sep + 'Fz_per_Coil.csv'
        self.file_Fr_per_Coil = self.sim_output + os.sep + 'Fr_per_Coil.csv'
        self.file_Fz_per_Pancake = self.sim_output + os.sep + 'Fz_per_Pancake.csv'
        self.file_Fr_per_Pancake = self.sim_output + os.sep + 'Fr_per_Pancake.csv'
        self.file_R_Radial_per_Coil = self.sim_output + os.sep + 'R_Radial_per_Coil.csv'
        self.file_R_Radial_per_Pancake = self.sim_output + os.sep + 'R_Radial_per_Pancake.csv'

        # Columns
        columns = ['time'] + ['p' + str(i) for i in range(sum([ len(i) for i in self.T0]))]
        columns_2 = ['time'] + ['p' + str(i) for i in range(len(self.Index_inductive_connections))]
        columns_3 = ['time'] + ['Maximum_T'] + ['Average_T'] + ['Stored_Energy'] + ['I_' + str(i) for i in self.unique_circuits] + ['Bz_Center'] + ['Bx_Center'] + ['P_cooling'] + ['P_AC'] + ['P_radial'] + ['V_'+str(i) for i in self.unique_circuits] + ['Max_Fr_Density'] + ['Max_Fz_Density'] + ['Quench_Detected'] + ['Capacitor_Voltage'] + ['Capacitor_Current'] + ['Solver_time'] + ['PSU_power']
        columns_4 = ['time'] + ['p' + str(i) for i in set(self.thermal_element_to_coil_nr)]
        columns_5 = ['time'] + ['p' + str(i) for i in set(self.thermal_element_to_solenoid_nr)]
        columns_6 = ['time'] + ['c' + str(i) for i in self.index_dict['coils'].keys()]
        columns_7 = ['time'] + ['p' + str(i) for i in self.index_dict['pancakes'].keys()]

        if self.names is not None:
            columns_5 = ['time'] + self.names

        # New columns
        self.columns['columns1'] = columns
        self.columns['columns2'] = columns_2
        self.columns['columns3'] = columns_3
        self.columns['columns4'] = columns_4
        self.columns['columns5'] = columns_5
        self.columns['columns6'] = columns_6
        self.columns['columns7'] = columns_7

        if self.bool_summary:
            self.create_csv(self.file_summary, 'columns3')

        if self.bool_temperature:
            self.create_csv(self.file_temperature, 'columns1')

        if self.bool_current:
            self.create_csv(self.file_current, 'columns2')
            self.create_csv(self.file_current_R, 'columns1')
            self.create_csv(self.file_R_Radial_per_Coil, 'columns6')
            self.create_csv(self.file_R_Radial_per_Pancake, 'columns7')

        if self.bool_loss:
            self.create_csv(self.file_Prad, 'columns1')
            self.create_csv(self.file_P_AC, 'columns1')
            self.create_csv(self.file_Pin, 'columns1')
            self.create_csv(self.file_Prad_density, 'columns1')
            self.create_csv(self.file_P_AC_density, 'columns1')
            self.create_csv(self.file_Pin_density, 'columns1')
            self.create_csv(self.file_Power_per_Coil_Element, 'columns4')
            self.create_csv(self.file_Power_per_Solenoid, 'columns5')

        if self.bool_magnetic_field:
            self.create_csv(self.file_magnetic_field, 'columns2')
            self.create_csv(self.file_Bz_conductor, 'columns2')
            self.create_csv(self.file_Br_conductor, 'columns2')

        if self.bool_voltage:
            self.create_csv(self.file_voltage, 'columns1')

        if self.bool_force:
            self.create_csv(self.file_Fz, 'columns2')
            self.create_csv(self.file_Fr, 'columns2')
            self.create_csv(self.file_Fz_density, 'columns2')
            self.create_csv(self.file_Fr_density, 'columns2')
            self.create_csv(self.file_Fz_density_peak, 'columns2')
            self.create_csv(self.file_Fr_density_peak, 'columns2')
            self.create_csv(self.file_Fz_per_Coil, 'columns6')
            self.create_csv(self.file_Fr_per_Coil, 'columns6')
            self.create_csv(self.file_Fz_per_Pancake, 'columns7')
            self.create_csv(self.file_Fr_per_Pancake, 'columns7')

    def start_logger(self):
        """
        Initialize and start the logging system for the simulation.

        Configures Python's logging module to write to both a log file and
        console (stdout). The log file is created in the simulation output
        directory with timestamp and level information for each log entry.

        Side Effects:
            - Creates log file at self.sim_output/logfile
            - Configures logging to INFO level
            - Adds console output handler for logging
            - Writes initial log entries with simulation timestamp
        """
        print('Starting the logger.')

        logger = logging.getLogger()  # root logger
        logger.setLevel(logging.DEBUG)  # capture everything centrally

        # file handler
        file_handler = logging.FileHandler(self.sim_output + os.sep + "logfile", mode="w")
        file_handler.setLevel(logging.DEBUG)
        file_handler.setFormatter(logging.Formatter("%(asctime)-15s %(levelname)-8s %(message)s"))

        # console handler
        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.INFO)  # only INFO+ to console

        logger.addHandler(file_handler)
        logger.addHandler(console_handler)

        logging.info('Logging started after path creation.')
        logging.info('NICQS 2D, Date: ' + str(datetime.now().strftime("%Y-%m-%d %H:%M")))

    def import_matrix(self):
        """
        Load electrical and mutual inductance matrices from disk.

        This method loads pre-computed matrices generated by the geometry
        builder: the electrical matrix (resistance/inductance network) and
        the mutual inductance matrix. It then inverts the electrical matrix
        and splits it into current and voltage parts for the ODE solver.

        Side Effects:
            - Loads electrical_matrix from self.electrical_matrix_path
            - Loads Mutual inductance matrix from self.mutual_inductance_matrix_path
            - Computes full matrix inverse and stores as self.inv
            - Splits inverse into inv_1 (current part) and inv_2 (voltage part)
            - Converts inv_1 and inv_2 to PyTorch tensors on configured device
            - Sets self.matrix_size to the dimension of the electrical matrix
        """
        # Imports the required electrical matrices
        # These matrices are made by the geometry builder.

        # Loading the electrical matrix
        path = self.electrical_matrix_path
        try:
            with open(path, 'rb') as fin:
                self.electrical_matrix = pickle.load(fin)
        except OSError:
            logging.info("Loading of %s has failed." % path)
        else:
            logging.debug("Successfully loaded %s" % path)

        path = self.mutual_inductance_matrix_path
        try:
            with open(path, 'rb') as fin:
                self.Mutual = pickle.load(fin)
        except OSError:
            logging.info("Loading of %s has failed." % path)
        else:
            logging.debug("Successfully loaded %s" % path)

        self.inv = np.linalg.inv(self.electrical_matrix)  # Total inversed matrix (only the current part of it is needed in the ODE solver)
        self.inv_1 = torch.tensor(self.inv[0 : self.Nr_of_Connections, :], dtype=self.torch_float_type, device=self.torch_device)  # The current part
        self.inv_2 = torch.tensor(self.inv[self.Nr_of_Connections:, :], dtype=self.torch_float_type, device=self.torch_device)  # The voltage part
        self.matrix_size = self.inv.shape[0]


    def calculate_geometric_properties(self):
        """
        Calculate geometric properties and initialize material lookup tables.

        This method computes conductor properties (resistivity, conductivity,
        resistance), creates interpolation functions for material properties
        (Ic(T,B), resistivity, thermal conductivity, heat capacity), and
        initializes lookup tables for the power-law E-J relationship. These
        precomputed functions are used extensively during the ODE solution.

        Side Effects:
            - Initializes ref_resistivity, R_bypass, ref_resistance arrays
            - Creates material property dictionaries (SC_Ic_Fit, Normal_rho_FIT,
              Normal_k_FIT, Normal_Cv_FIT, SC_Cur_Sharing, SC_Ic_Scaling)
            - Generates interpolation functions for all materials and conductors
            - Creates power-law lookup tables (LUT) for E-J characteristics
            - Computes reference resistances and conductivities for each connection
            - Logs material property warnings and validation information
        """

        self.ref_resistivity = torch.zeros((self.Nr_of_Connections,), dtype=self.torch_float_type, device=self.torch_device)
        self.R_bypass = np.zeros((self.Nr_of_Connections,))
        self.ref_resistance = torch.zeros((self.Nr_of_Connections,), dtype=self.torch_float_type, device=self.torch_device)
        self.ref_radial_resistance = torch.zeros((self.Nr_of_Connections,), dtype=self.torch_float_type, device=self.torch_device)
        self.ref_conductivity = np.zeros((self.Nr_of_Connections,))

        # Creating the normal and SC dictionaries
        self.SC_Ic_Fit = {}
        self.Normal_rho_FIT = {}
        self.Normal_k_FIT = {}
        self.Normal_Cv_FIT = {}
        self.SC_Cur_Sharing = {}
        self.SC_Ic_Scaling = {}

        # Creating a list of all normal conducting properties
        all_mats = list(np.concatenate([[j for j in self.Conductor_conductivity_materials[i]] for i in self.Conductor_conductivity_materials.keys()])) + [self.Conductor_resistor[i] for i in self.Conductor_resistor.keys()] + list(np.concatenate([[j for j in self.Conductor_materials[i]] for i in self.Conductor_materials.keys()]))
        all_mats = list(set(all_mats))

        # Creating a list of all conductor names
        all_conductors = list(set(list(self.Conductor_type.keys()))) if self.Conductor_type is not None else []

        # Lut minmax
        self.lut_limits = {
            'IIc_min': 1e-9,
            'IIc_max': 1e3,
            'IIc_count': 300,
            'U_min': 1e-15,
            'U_max': 1e3,
            'U_count': 300,
        }

        self.value_limits = {
            'T_min': 1.0,
            'T_max': 100.0,
            'T_count': 100,
            'T_rho_min': 1.0,
            'T_rho_max': 1000,
            'T_rho_count': 200,
            'B_min': 0.001,
            'B_max': 100.0,
            'B_count': 100,
            'Angle_min': 0.0,
            'Angle_max': 360.0,
            'Angle_count': 361,
        }

        # Setting the minimum value of the magnetic field to a certain value.
        if self.B_min is not None:
            self.value_limits['B_min'] = self.B_min

        # Limits
        T = np.linspace(self.value_limits['T_min'], self.value_limits['T_max'], self.value_limits['T_count'])  # Temperature range for the superconductor
        T_rho = np.geomspace(self.value_limits['T_rho_min'], self.value_limits['T_rho_max'], self.value_limits['T_rho_count'])  # Temperature range for the normal conducting parts
        B = np.geomspace(self.value_limits['B_min'], self.value_limits['B_max'], self.value_limits['B_count'])  # Magnetic field range, geomspace, so 100 T is fine
        Degr = np.linspace(self.value_limits['Angle_min'], self.value_limits['Angle_max'], self.value_limits['Angle_count'])  # Angle
        T_torch = torch.tensor(T, dtype=self.torch_float_type, device=self.torch_device)
        T_rho_torch = torch.tensor(T_rho, dtype=self.torch_float_type, device=self.torch_device)
        B_torch = torch.tensor(B, dtype=self.torch_float_type, device=self.torch_device)
        Degr_torch = torch.tensor(Degr, dtype=self.torch_float_type, device=self.torch_device)

        # I/Ic and U for the current sharing lookup table.
        IIc = np.concatenate((
            np.geomspace(self.lut_limits['IIc_min'], 0.4, int(self.lut_limits['IIc_count'] * 0.15), endpoint=False),
            np.linspace(0.4, 5.0, int(self.lut_limits['IIc_count'] * 0.35), endpoint=False),
            np.linspace(5.0, 100.0, int(self.lut_limits['IIc_count'] * 0.25), endpoint=False),
            np.geomspace(100.0, self.lut_limits['IIc_max'], int(self.lut_limits['IIc_count'] * 0.25)),
        ))

        U = np.concatenate((
            np.geomspace(self.lut_limits['U_min'], 1e-4, int(self.lut_limits['U_count'] * 0.2), endpoint=False),
            np.geomspace(1e-4, 100.0, int(self.lut_limits['U_count'] * 0.4), endpoint=False),
            np.linspace(100.0, self.lut_limits['U_max'], int(self.lut_limits['U_count'] * 0.4))
        ))

        # Making them torch tensors.
        IIc_torch = torch.tensor(IIc, dtype=self.torch_float_type, device=self.torch_device)
        U_torch = torch.tensor(U, dtype=self.torch_float_type, device=self.torch_device)

        # Making the resistivity, thermal conductivity and heat capacity interpolation function
        for i in all_mats:
            input_list = np.array([i for i in itertools.product(T_rho, B)])

            R = self.materials.resistivity_selector(i, input_list[:, 0], input_list[:, 1], reshape=True, set_B=self.set_B, B_min=self.B_min)
            k = self.materials.conductivity_selector(i, input_list[:, 0], input_list[:, 1], reshape=True, set_B=self.set_B)
            Cv = self.materials.specific_heat_selector(i, input_list[:, 0], reshape=True)
            R_torch = torch.tensor(R.reshape((len(T_rho), len(B))), dtype=self.torch_float_type, device=self.torch_device)
            k_torch = torch.tensor(k.reshape((len(T_rho), len(B))), dtype=self.torch_float_type, device=self.torch_device)
            Cv_torch = torch.tensor(Cv.reshape((len(T_rho), len(B))), dtype=self.torch_float_type, device=self.torch_device)
            R_fit = RegularGridInterpolator((T_rho_torch, B_torch), R_torch)
            k_fit = RegularGridInterpolator((T_rho_torch, B_torch), k_torch)
            Cv_fit = RegularGridInterpolator((T_rho_torch, B_torch), Cv_torch)
            self.Normal_rho_FIT[i] = R_fit
            self.Normal_k_FIT[i] = k_fit
            self.Normal_Cv_FIT[i] = Cv_fit

        for i in self.unique_regions:
            # Unique regions are for example 'SC1'
            # Pre-calculates a few important values.
            # Checks the fitname and turns angle dependency on or off.

            if self.Conductor_type[i] == 'Superconductor':
                # Calculates the scaling factor: F_scaling = Ic_ref / Ic(T, B, deg)

                # Checks if the reference angle is provided. Otherwise the angle dependency is set to False.
                if 'fit-name' in self.Conductor_SC_properties[i].keys() and 'angle' in self.Conductor_SC_properties[i].keys():
                    self.deg_dep[i] = True if fit_angle_dependency(self.Conductor_SC_properties[i]['fit-name']) else False
                else:
                    self.deg_dep[i] = False

                if self.deg_dep[i]:
                    Tref = self.Conductor_SC_properties['SC1']['Tref']
                    Bref = self.Conductor_SC_properties['SC1']['Bref']
                    Deg = self.Conductor_SC_properties['SC1']['angle']
                    fit_name = self.Conductor_SC_properties[i]['fit-name']
                    input_array = MaterialProperties.create_steam_material_input(
                        fit_name, Temperature=Tref, Magnetic_Field=Bref, Field_Angle=Deg
                    )
                    max_points = input_array.shape[1] if len(input_array.shape) > 1 else input_array.shape[0]
                    steam = self.materials.get_cached_steam_material(fit_name, max_points)
                    Ic = steam.evaluate(input_array)
                    self.SC_Ic_Scaling[i] = torch.tensor(self.Conductor_SC_properties[i]['Ic'] / Ic / self.Conductor_SC_properties[i]['parallel'], dtype=self.torch_float_type, device=self.torch_device)
                else:
                    # Getting the reference Ic with an angle of zero.
                    Tref = self.Conductor_SC_properties['SC1']['Tref']
                    Bref = self.Conductor_SC_properties['SC1']['Bref']
                    fit_name = self.Conductor_SC_properties[i]['fit-name']
                    
                    if fit_angle_dependency(fit_name):
                        input_array = MaterialProperties.create_steam_material_input(
                            fit_name, Temperature=Tref, Magnetic_Field=Bref, Field_Angle=0.0
                        )
                    else:
                        input_array = MaterialProperties.create_steam_material_input(
                            fit_name, Temperature=Tref, Magnetic_Field=Bref
                        )

                    max_points = input_array.shape[1] if len(input_array.shape) > 1 else input_array.shape[0]
                    steam = self.materials.get_cached_steam_material(fit_name, max_points)
                    Ic = steam.evaluate(input_array)
                    self.SC_Ic_Scaling[i] = torch.tensor(self.Conductor_SC_properties[i]['Ic'] / Ic / self.Conductor_SC_properties[i]['parallel'], dtype=self.torch_float_type, device=self.torch_device)

                # Making the Ic interpolation
                fit_name = self.Conductor_SC_properties[i]['fit-name']
                if self.deg_dep[i]:
                    input_list = np.array([j for j in itertools.product(T, B, Degr)])
                    temp = MaterialProperties.create_steam_material_input(
                        fit_name,
                        Temperature=input_list[:, 0],
                        Magnetic_Field=input_list[:, 1],
                        Field_Angle=input_list[:, 2]
                    )
                    max_points = temp.shape[1] if len(temp.shape) > 1 else temp.shape[0]
                    steam = self.materials.get_cached_steam_material(fit_name, max_points)
                    Ic = steam.evaluate(temp) * self.SC_Ic_Scaling[i].cpu().numpy()
                    Ic_torch = torch.tensor(Ic.reshape((len(T), len(B), len(Degr))), dtype=self.torch_float_type, device=self.torch_device)
                    self.SC_Ic_Fit[i] = RegularGridInterpolator((T_torch, B_torch, Degr_torch), Ic_torch)
                else:
                    if fit_angle_dependency(fit_name):
                        input_list = np.array([j for j in itertools.product(T, B, Degr*0.0)])
                        temp = MaterialProperties.create_steam_material_input(
                            fit_name,
                            Temperature=input_list[:, 0],
                            Magnetic_Field=input_list[:, 1],
                            Field_Angle=input_list[:, 2]
                        )
                        max_points = temp.shape[1] if len(temp.shape) > 1 else temp.shape[0]
                        steam = self.materials.get_cached_steam_material(fit_name, max_points)
                        Ic = steam.evaluate(temp) * self.SC_Ic_Scaling[i].cpu().numpy()
                        Ic_torch = torch.tensor(Ic.reshape((len(T), len(B), len(Degr))), dtype=self.torch_float_type, device=self.torch_device)
                        self.SC_Ic_Fit[i] = RegularGridInterpolator((T_torch, B_torch, Degr_torch), Ic_torch)
                    else:
                        input_list = np.array([j for j in itertools.product(T, B)])
                        temp = MaterialProperties.create_steam_material_input(
                            fit_name,
                            Temperature=input_list[:, 0],
                            Magnetic_Field=input_list[:, 1]
                        )
                        max_points = temp.shape[1] if len(temp.shape) > 1 else temp.shape[0]
                        steam = self.materials.get_cached_steam_material(fit_name, max_points)
                        Ic = steam.evaluate(temp) * self.SC_Ic_Scaling[i].cpu().numpy()
                        Ic_torch = torch.tensor(Ic.reshape((len(T), len(B))), dtype=self.torch_float_type, device=self.torch_device)
                        self.SC_Ic_Fit[i] = RegularGridInterpolator((T_torch, B_torch), Ic_torch)

                # Making the x interpolation function
                logging.info('Generating the look-up-tables.')
                input_list = np.array([j for j in itertools.product(IIc, U)])
                x = np.zeros((len(input_list), 1))
                for j in range(len(input_list)):
                    x[j] = fsolve(self.powerlaw, 1.0, args=(input_list[j][0], input_list[j][1], self.Conductor_SC_properties[i]['n-value']), fprime=self.powerlaw_jac, maxfev=1000000, xtol=1e-9)
                x_torch = torch.clip(torch.tensor(x.reshape((len(IIc), len(U))), dtype=self.torch_float_type, device=self.torch_device), min = 0.0, max = 1.0)
                self.SC_Cur_Sharing[i] = RegularGridInterpolator((IIc_torch, U_torch), x_torch)

                # Radial resistances for the coil packs.
                index_norm = self.Index_radial_connections_region[i]
                # T_rad_res = torch.tensor([self.T0[m] for m in [n for n in self.Index_temperatures[index_norm]]], dtype=self.torch_float_type, device=self.torch_device)
                T_rad_res = torch.tensor(np.array(self.Tref_array)[index_norm], dtype=self.torch_float_type, device=self.torch_device)
                B_rad_res = torch.tensor(np.array(self.Bref_array)[index_norm], dtype=self.torch_float_type, device=self.torch_device)

                self.ref_resistivity[index_norm] = self.materials.resistivity_selector_torch(self.Conductor_resistor[i], T=T_rad_res, B=B_rad_res, Normal_rho_FIT=self.Normal_rho_FIT)
                self.ref_resistance[index_norm] = self.radial_resistance[index_norm]

            elif self.Conductor_type[i].lower() == 'normal':
                logging.info("Calculating the reference resistance and resistivity of the normal conducting elements.")

                # First: reference resistance
                index_norm = self.Index_normal_connections_region[i]
                index_ind = self.Index_inductive_connections_region[i]
                T_norm_res = torch.tensor([self.T0[m] for m in [n for n in self.Index_temperatures[index_norm]]], dtype=self.torch_float_type, device=self.torch_device)
                self.ref_resistivity[index_norm] = self.materials.resistivity_selector_torch(self.Conductor_resistor[i], T_norm_res, Normal_rho_FIT=self.Normal_rho_FIT) # Gets the resistivity based on temperature
                self.ref_resistance[index_norm] = self.ref_resistivity[index_norm] * torch.tensor(self.lengths[index_ind] / self.cross_sections[index_ind], dtype=self.torch_float_type, device=self.torch_device) # Calculates the resistance based on temperature and overwrites any other value.

        ## Thermal connections ##
        if self.Thermal_alpha is not None:
            for i in self.Thermal_area.keys():
                for j in self.Thermal_area[i].keys():
                    # Generating a dictionary of torch tensors.
                    self.Thermal_area[i][j] = torch.tensor(self.Thermal_area[i][j], dtype=self.torch_float_type, device=self.torch_device)

            # Thermal alpha =
            self.Thermal_alpha = torch.tensor(self.Thermal_alpha, dtype=self.torch_float_type, device=self.torch_device)
            self.Thermal_resistance = torch.zeros(np.array(self.T0).shape, dtype=self.torch_float_type, device=self.torch_device)
            for i in self.unique_regions:
                T_res = torch.tensor(np.array(self.T0)[self.Thermal_index[i]], dtype=self.torch_float_type, device=self.torch_device)
                thermal_sum = torch.zeros_like(T_res)
                for j in self.Conductor_conductivity_materials[i]:
                    thermal_sum = thermal_sum + self.materials.conductivity_selector_torch(j, T_res, Normal_k_FIT=self.Normal_k_FIT) * self.Thermal_area[i][j]
                self.Thermal_resistance[self.Thermal_index[i]] = self.Thermal_alpha[self.Thermal_index[i]] * 1 / thermal_sum

        # Generates the Ic graphs.
        self.debug_graphs.generate_ic_graph()

        # Generates the Radial Resistance graph.
        self.debug_graphs.generate_radial_resistance_graph()

        # Generate the fit current sharing surfaces:
        self.debug_graphs.generate_current_sharing_surface()

        # Generate Ic surface plots (Ic vs T and B at angle=0):
        self.debug_graphs.generate_ic_surface()

        # Generate a graph of the Ic vs Angle for the superconductors used:
        self.debug_graphs.generate_ic_vs_angle_graph()

        # Generate time constant vs temperature graphs for each circuit:
        self.debug_graphs.generate_time_constant_vs_temperature()

        # Generate time constant vs magnetic field graphs for each circuit:
        self.debug_graphs.generate_time_constant_vs_field()


    def create_initial_conditions(self):
        """
        Set initial conditions for currents and temperatures.

        This method initializes the simulation state by either loading initial
        conditions from saved CSV files (for restart simulations) or using
        prescribed values from the configuration. It sets I0 (initial currents)
        and T0 (initial temperatures) for all elements.

        Side Effects:
            - Modifies self.I0 with initial currents for all connections
            - Modifies self.T0 with initial temperatures for all nodes
            - Sets self.t_offset if loading from file and t_offset_bool is True
            - Reads CSV files (current.csv, current_resistor.csv, temperature.csv)
              if Load_from_file is True for any initial_conditions entry
        """
        for i in self.initial_conditions.keys():
            if self.initial_conditions[i]['Load_from_file']:
                # Loading the current data
                df = pd.read_csv(self.initial_conditions[i]['File_path'] + os.sep + 'current.csv')

                # getting the time offset (only applied to the save files)
                self.t_offset = df.iloc[self.initial_conditions[i]['Line_index']].values[0] if self.t_offset_bool else 0.0

                # All inductive connections
                temp_list = []
                for k in self.Index_initial_Iind.keys():
                    for kk in self.Index_initial_Iind[k]:
                        temp_list.append(kk)
                for j in range(len(df.iloc[self.initial_conditions[i]['Line_index']].values[1::])):
                    self.I0[temp_list[j]] = df.iloc[self.initial_conditions[i]['Line_index']].values[1::][j]

                # All resistive connections
                df = pd.read_csv(self.initial_conditions[i]['File_path'] + os.sep + 'current_resistor.csv')
                temp_list = []
                for k in self.Index_initial_IR.keys():
                    for kk in self.Index_initial_IR[k]:
                        temp_list.append(kk)
                for j in range(len(df.iloc[self.initial_conditions[i]['Line_index']].values[1::])):
                    self.I0[temp_list[j]] = df.iloc[self.initial_conditions[i]['Line_index']].values[1::][j]

                # Temperatures
                df = pd.read_csv(self.initial_conditions[i]['File_path'] + os.sep + 'temperature.csv')
                self.T0 = list(df.iloc[self.initial_conditions[i]['Line_index']].values[1::])
            else:
                for j in self.Index_initial_Iind[i]:
                    self.I0[j] = self.initial_conditions[i]['I_ind']
                for j in self.Index_initial_IR[i]:
                    self.I0[j] = self.initial_conditions[i]['I_R']
                for j in self.Index_initial_T[i]:
                    self.T0[j] = self.initial_conditions[i]['T']

    def load_external_temperature(self):
        """
        Load external temperature profile from pickle file.

        When use_external_temperature is enabled, this method loads a
        pre-computed temperature evolution from disk. The thermal ODE is
        then bypassed during simulation, and temperatures follow the loaded
        profile.

        Side Effects:
            - Loads external_temperature interpolation function from pickle
            - Logs success/failure messages
        """
        try:
            with open(self.specific_output_path + os.sep + self.external_temperature_path, 'rb') as fin:
                self.external_temperature = pickle.load(fin)
        except OSError:
            logging.info("Loading of %s has failed." % self.external_temperature_path)
        else:
            logging.debug("Successfully loaded %s" % self.external_temperature_path)

    def start_solver(self):
        """
        Main solver method that runs the time integration for the NI coil simulation.
        
        This method performs the following major steps:
        1. Initialize simulation parameters and conditions
        2. Set up the ODE solver
        3. Run the main time integration loop
        4. Handle quench detection and protection systems
        5. Log and save results throughout the simulation
        """
        
        # === 1. INITIALIZATION PHASE ===
        logging.info('*** Initializing solver parameters ***')
        
        # Initialize time tracking variables
        t = 0.0  # Current simulation time
        dt = self.dt  # Initial time step
        self.output_n = self.output_mod  # Counter for output frequency
        FIRST_BOOL = True  # Flag to handle first iteration special cases
        n_ode = 0  # ODE integration step counter
        self.start = time.time()  # Wall clock start time for performance tracking
        
        # Set up initial conditions from configuration
        self.create_initial_conditions()
        
        # Handle external temperature input if specified
        if self.use_external_temperature:
            self.load_external_temperature()
            self.T0 = [self.external_temperature(0.0)] * len(self.T0)
        
        # Create the initial state vector: [currents, temperatures, capacitor_voltage, capacitor_current]
        y = np.array(self.I0 + self.T0 + [self.Capacitor_Voltage] + [self.Capacitor_Current])
        
        # Initialize torch tensors for GPU computation
        self._initialize_solver_tensors()
        
        # Create circuit current interpolation for ramping
        self._setup_circuit_interpolation()
        
        # === 2. INITIAL CONDITIONS CALCULATION ===
        logging.info('*** Computing initial field and force distributions ***')
        
        # Calculate initial values for all physical quantities
        initial_results = self._compute_initial_conditions(y)
        E, total_power = initial_results['energy'], initial_results['total_power']
        
        # Store and save the initial timestep results
        self._store_and_save_initial_results(initial_results)
        
        # === 3. ODE SOLVER SETUP ===
        logging.info('*** Setting up ODE solver ***')
        
        # Enable profiler if requested
        if self.profile_bool: 
            profiler.enable()
        
        # Configure the appropriate ODE solver based on user selection
        self._setup_ode_solver(y, t)
        
        # === 4. MAIN TIME INTEGRATION LOOP ===
        logging.info('*** Starting main time integration loop ***')
        
        while t < self.t0[1] and self.T_stop and self.E_stop:
            
            # Advance the ODE solver by one time step
            t, y = self._advance_ode_solver(t, dt, y)
            
            # Extract physical quantities from the state vector
            state_vars = self._extract_state_variables(y, t)
            I, I1, I2, T, V_cap, I_cap = state_vars
            
            # Determine what operations are needed this timestep
            flags = self._determine_timestep_flags()
            SAVE_BOOL, OUTPUT_BOOL, STORE_BOOL = flags
            
            # Convert currents to torch tensors for field calculations
            I_torch = torch.tensor(I, dtype=self.torch_float_type, device=self.torch_device)
            I1_torch = torch.tensor(I1, dtype=self.torch_float_type, device=self.torch_device)
            
            # Compute power dissipation, fields, and other derived quantities
            physics_results = self.solve_this_torch(t, y, 2)
            P_AC, Prad, Pin, B, V, P_cooling, SC_loss, Radial_loss, R_bypass, R_Conductor, IIc, x_array = physics_results
            
            # Calculate stored magnetic energy
            E = torch.sum(0.5 * self.Mutual * torch.outer(I_torch, I_torch)).cpu().numpy()
            
            # Compute detailed magnetic field and force distributions (only when needed for storage)
            if STORE_BOOL or FIRST_BOOL:
                field_results = self._compute_magnetic_fields_and_forces(I1_torch, I1)
                
            # === 5. QUENCH DETECTION AND PROTECTION ===
            now = datetime.now().strftime("%H:%M:%S")
            
            # Check for quench conditions and handle protection systems
            self._handle_quench_detection(t, T, V, now)
            
            # Manage power supply and heater states
            self._manage_power_systems(t, T)
            
            # === 6. POWER DENSITY AND PID CONTROL ===
            if STORE_BOOL or OUTPUT_BOOL or FIRST_BOOL:
                # Calculate power densities for output
                Pin_density, Prad_density, P_AC_density, energy_integrals = self._compute_power_densities(Pin, Prad, P_AC)
                
            # Handle PID control for power-limited ramping
            self._handle_pid_control(Prad, P_AC, dt)
            
            # === 7. SIMULATION TERMINATION CHECKS ===
            # Check if temperature or energy stopping criteria are met
            self._check_termination_criteria(E, T)
            
            # Calculate total electrical power
            total_power = np.sum([np.sum(V[self.circuits_list[i]]) * self.I_circuit_last[i] for i in self.unique_circuits])
            
            # === 8. RESULTS STORAGE AND OUTPUT ===
            # Store detailed results for post-processing (selective based on PID requirements)
            if STORE_BOOL or FIRST_BOOL:
                self._store_timestep_results(t, T, I1, I2, V_cap, I_cap, total_power, 
                                           field_results, physics_results, energy_integrals)
                FIRST_BOOL = False
            else:
                self.save_n_pid += 1
            n_ode += 1
            
            # Save buffered results to disk periodically
            save_asterix = self._handle_result_saving(SAVE_BOOL)
            
            # === 9. ADAPTIVE TIME STEPPING ===
            # Calculate next timestep based on solution dynamics and user constraints
            dt = self._calculate_adaptive_timestep(t)
            
            # === 10. PROGRESS LOGGING ===
            # Output simulation progress information
            if OUTPUT_BOOL:
                self._log_simulation_progress(now, save_asterix, n_ode, t, E, dt, T, I1, B, 
                                            energy_integrals if (STORE_BOOL or OUTPUT_BOOL) else (0, 0), 
                                            P_AC, Prad)
                self.output_n = 1
            else:
                self.output_n += 1
        
        # === 11. SIMULATION COMPLETION ===
        # Mark successful completion
        self.FLAG_Finished = True
        
        # Log total simulation time
        end = time.time()
        now = datetime.now().strftime("%H:%M:%S")
        logging.info(f'{now} || Simulation completed. Total elapsed time: {int(end - self.start)} seconds.')

    def _initialize_solver_tensors(self):
        """
        Initialize torch tensors needed for the solver.
        
        Creates zero-filled tensors on the appropriate device for:
        - b-array for ODE right-hand side
        - Magnetic field storage
        - Temperature tensor storage
        """
        self.b0 = torch.zeros((self.inv.shape[0],), dtype=self.torch_float_type, device=self.torch_device)
        self.B0 = torch.zeros((len(self.I0),), dtype=self.torch_float_type, device=self.torch_device)
        self.T0_tensor = torch.zeros((len(self.T0),), dtype=self.torch_float_type, device=self.torch_device)

    def _setup_circuit_interpolation(self):
        """
        Create interpolation functions for circuit current ramping.
        
        This function generates time-current interpolation based on the specified
        ramp rate profiles, allowing smooth current ramping during simulation.
        """
        # Create time array for circuit current calculation
        t_circuit = np.linspace(self.t0[0], self.t0[1] * 1.01, 100000)
        
        # Integrate ramp rate to get current vs time
        I_circuit = integrate.cumulative_trapezoid(
            [self.ramp_rate(t, 1) for t in t_circuit], 
            x=t_circuit, 
            initial=0
        )
        
        # Create interpolation function for circuit current
        self.I_circuit_interpolation = interp1d(t_circuit, I_circuit, fill_value="extrapolate")

    def _compute_initial_conditions(self, y):
        """
        Compute all physical quantities at t=0 for the initial conditions.
        
        Args:
            y: Initial state vector
            
        Returns:
            dict: Dictionary containing all computed initial quantities
        """
        # Convert initial currents to torch tensor
        I0_tensor = torch.tensor(self.I0, dtype=self.torch_float_type, device=self.torch_device)
        
        # Calculate initial stored magnetic energy
        E = torch.sum(0.5 * self.Mutual * torch.outer(I0_tensor, I0_tensor)).cpu().numpy()
        
        # Compute magnetic field components on conductors
        Bz_conductor = torch.matmul(self.Bz_Conductor, I0_tensor[self.Index_inductive_connections])
        Br_conductor = torch.matmul(self.Br_Conductor, I0_tensor[self.Index_inductive_connections])
        B_tot = torch.sqrt(torch.addcmul(Bz_conductor.square(), Br_conductor, Br_conductor))
        Field_Angle = self._calculate_field_angle(Br_conductor, B_tot)
        
        # Convert to numpy for storage
        Bz_conductor_np = Bz_conductor.cpu().numpy()
        Br_conductor_np = Br_conductor.cpu().numpy()
        B_tot_np = B_tot.cpu().numpy()
        Field_Angle_np = Field_Angle.cpu().numpy()
        
        # Compute peak magnetic fields
        Bz_peak = torch.matmul(self.Bz_peak, I0_tensor[self.Index_inductive_connections]).cpu().numpy()
        Br_peak = torch.matmul(self.Br_peak, I0_tensor[self.Index_inductive_connections]).cpu().numpy()
        
        # Compute central magnetic fields
        Bz_center = np.dot(self.Bz_center, np.array(self.I0)[self.Index_inductive_connections])
        Br_center = np.dot(self.Br_center, np.array(self.I0)[self.Index_inductive_connections])
        
        # Calculate current density
        I_density = (np.array(self.I0)[self.Index_inductive_connections] / 
                    self.cross_sections[self.Index_inductive_connections] * 
                    self.turns_per_element[self.Index_inductive_connections] * 1e-6)
        
        # Compute electromagnetic forces
        force_results = self._compute_electromagnetic_forces(
            Bz_conductor_np, Br_conductor_np, Bz_peak, Br_peak, np.array(self.I0)[self.Index_inductive_connections]
        )
        
        # Calculate power dissipation and other physics
        physics_results = self.solve_this_torch(0.0, y, 2)
        P_AC, Prad, Pin = physics_results[0], physics_results[1], physics_results[2]
        
        # Compute power densities
        Pin_density = Pin / self.cross_sections[self.Index_normal_connections]
        Prad_density = Prad / self.cross_sections[self.Index_normal_connections]
        P_AC_density = P_AC / self.cross_sections[self.Index_normal_connections]
        
        # Calculate total electrical power
        total_power = np.sum([np.sum(physics_results[4][self.circuits_list[i]]) * self.I_circuit_last[i] 
                            for i in self.unique_circuits])
        
        return {
            'energy': E,
            'B_tot': B_tot_np,
            'Bz_conductor': Bz_conductor_np,
            'Br_conductor': Br_conductor_np,
            'Field_Angle': Field_Angle_np,
            'Bz_center': Bz_center,
            'Br_center': Br_center,
            'I_density': I_density,
            'forces': force_results,
            'power_densities': (Pin_density, Prad_density, P_AC_density),
            'physics_results': physics_results,
            'total_power': total_power
        }

    def _compute_electromagnetic_forces(self, Bz_conductor, Br_conductor, Bz_peak, Br_peak, currents):
        """
        Calculate electromagnetic forces and force densities.
        
        Args:
            Bz_conductor: Axial magnetic field on conductors
            Br_conductor: Radial magnetic field on conductors  
            Bz_peak: Peak axial magnetic field
            Br_peak: Peak radial magnetic field
            currents: Current array
            
        Returns:
            dict: Dictionary containing force arrays and densities
        """
        if self.turns_per_element is not None:
            # Calculate Lorentz forces: F = I × B
            # Axial force: Fz = -Br * I * length
            Fz = -Br_conductor * np.array(currents) * self.lengths[self.Index_inductive_connections]
            # Radial force: Fr = Bz * I * length  
            Fr = Bz_conductor * np.array(currents) * self.lengths[self.Index_inductive_connections]
            # Peak forces using peak field values
            Fz_peak = -Br_peak * np.array(currents) * self.lengths[self.Index_inductive_connections]
            Fr_peak = Bz_peak * np.array(currents) * self.lengths[self.Index_inductive_connections]
        else:
            # No turns specified, set forces to zero
            n_elements = len(np.array(currents))
            Fz = np.zeros(n_elements)
            Fr = np.zeros(n_elements)
            Fz_peak = np.zeros(n_elements)
            Fr_peak = np.zeros(n_elements)
        
        # Calculate force densities (force per unit volume, accounting for turns)
        cross_sections = self.cross_sections[self.Index_inductive_connections]
        lengths = self.lengths[self.Index_inductive_connections]
        turns = self.turns_per_element[self.Index_inductive_connections]
        
        Fz_density = Fz / cross_sections / lengths * turns
        Fr_density = Fr / cross_sections / lengths * turns
        Fz_density_peak = Fz_peak / cross_sections / lengths * turns
        Fr_density_peak = Fr_peak / cross_sections / lengths * turns
        
        return {
            'Fz': Fz, 'Fr': Fr, 'Fz_peak': Fz_peak, 'Fr_peak': Fr_peak,
            'Fz_density': Fz_density, 'Fr_density': Fr_density,
            'Fz_density_peak': Fz_density_peak, 'Fr_density_peak': Fr_density_peak
        }

    def _store_and_save_initial_results(self, results):
        """
        Store and save the initial timestep results.
        
        Args:
            results: Dictionary containing all computed initial quantities
        """
        forces = results['forces']
        Pin_density, Prad_density, P_AC_density = results['power_densities']
        physics_results = results['physics_results']
        
        # Store initial results using the main storage function
        self.store_results(
            self.t0[0], self.T0, 
            np.array(self.I0)[self.Index_inductive_connections], 
            np.array(self.I0)[self.Index_normal_connections], 
            physics_results[0], physics_results[1], physics_results[2],  # P_AC, Prad, Pin
            P_AC_density, Prad_density, Pin_density,
            results['B_tot'], results['Bz_center'], results['Br_center'], results['energy'],
            physics_results[4],  # V
            forces['Fz'], forces['Fz_density'], forces['Fz_density_peak'],
            forces['Fr'], forces['Fr_density'], forces['Fr_density_peak'],
            results['Bz_conductor'], results['Br_conductor'], results['I_density'],
            self.Capacitor_Voltage, self.Capacitor_Current, results['total_power'],
            physics_results[5], physics_results[6], physics_results[7],  # P_cooling, SC_loss, Radial_loss
            physics_results[8], physics_results[9],  # R_bypass, R_Conductor
            results['Field_Angle'], physics_results[10], physics_results[11]  # IIc, x_array
        )
        
        # Save initial results to disk
        self.save_results()

    def _setup_ode_solver(self, y, t):
        """
        Configure the ODE solver based on the selected method.
        
        Args:
            y: Initial state vector
            t: Initial time
        """
        solver_method = self.solver.lower()
        
        if solver_method in ['lsoda', 'vode']:
            if solver_method == 'lsoda':
                # LSODA: Automatic switching between stiff and non-stiff methods
                self.results = scipy.integrate.ode(self.solve_this_torch).set_integrator(
                    'lsoda', method='bdf', rtol=self.rtol, atol=self.atol,
                    min_step=self.min_step, max_step=self.max_step, 
                    nsteps=int(1e9), max_order_s=15, max_hnil=99
                )
            elif solver_method == 'vode':
                # VODE: Variable-coefficient ODE solver
                self.results = scipy.integrate.ode(self.solve_this_torch).set_integrator(
                    'vode', method='bdf', rtol=self.rtol, atol=self.atol,
                    min_step=self.min_step, max_step=self.max_step, 
                    nsteps=int(1e5), order=5
                )
            
            # Set the initial conditions
            self.results.set_initial_value(y, t)
            
        elif solver_method == 'diffrax':
            # DiffRax: Modern JAX-based differential equation solver (not implemented)
            print('Setting up the diffrax solver, not yet implemented.')
            
        elif solver_method == 'ivp':
            # solve_ivp: Modern scipy interface (configured per timestep)
            print('Setting up the solve_ivp solver.')

    def _advance_ode_solver(self, t, dt, y=None):
        """
        Advance the ODE solver by one timestep.
        
        Args:
            t: Current time
            dt: Time step size
            y: Initialization array, needed for the ivp solver.
            
        Returns:
            tuple: New time and state vector (t_new, y_new)
        """
        solver_method = self.solver.lower()
        
        if solver_method in ['vode', 'lsoda']:
            # Use the pre-configured ODE solver
            self.results.integrate(self.results.t + dt)
            return self.results.t, self.results.y
            
        elif solver_method == 'ivp':
            # Use solve_ivp with BDF method for stiff problems
            method = 'BDF'
            self.results = solve_ivp(
                self.solve_this_torch, (t, t + dt), y, 
                method=method, max_step=dt/1.0
            )
            return self.results.t[-1], self.results.y[:, -1]
        
        # Default fallback (should not reach here with valid solver names)
        return t + dt, y

    def _extract_state_variables(self, y, t):
        """
        Extract physical quantities from the state vector.
        
        Args:
            y: State vector [currents, temperatures, capacitor_voltage, capacitor_current]
            t: Current time
            
        Returns:
            tuple: Extracted variables (I, I1, I2, T, V_cap, I_cap)
        """
        # Extract current arrays
        I = y[:self.Nr_of_Connections]  # All currents
        I1 = y[self.Index_inductive_connections]  # Inductive currents only
        I2 = y[self.Index_normal_connections]  # Resistive currents only
        
        # Extract temperature array (handle external temperature override)
        if self.use_external_temperature:
            T = np.array([self.external_temperature(t)] * len(self.T0))
        else:
            T = y[self.Nr_of_Connections:-2]
        
        # Extract capacitor state
        V_cap = y[-2]  # Capacitor voltage
        I_cap = y[-1]  # Capacitor current
        
        return I, I1, I2, T, V_cap, I_cap

    def _determine_timestep_flags(self):
        """
        Determine what operations are needed for this timestep.
        
        Returns:
            tuple: Boolean flags (SAVE_BOOL, OUTPUT_BOOL, STORE_BOOL)
        """
        # Check if results should be saved to disk
        SAVE_BOOL = (self.store_n - self.save_n) >= self.save_mod
        
        # Check if progress should be output to console
        OUTPUT_BOOL = self.output_n >= self.output_mod
        
        # Check if detailed results should be stored in memory (for PID control)
        STORE_BOOL = self.save_n_pid >= self.savemod_pid
        
        return SAVE_BOOL, OUTPUT_BOOL, STORE_BOOL

    def _compute_magnetic_fields_and_forces(self, I1_torch, I1):
        """
        Compute detailed magnetic field and force distributions.
        
        Args:
            I1_torch: Inductive currents as torch tensor
            I1: Inductive currents as numpy array
            
        Returns:
            dict: Dictionary containing field and force results
        """
        # Compute magnetic field components
        Bz_conductor = torch.matmul(self.Bz_Conductor, I1_torch)
        Br_conductor = torch.matmul(self.Br_Conductor, I1_torch)
        B_tot = torch.sqrt(torch.addcmul(Bz_conductor.square(), Br_conductor, Br_conductor))
        Field_Angle = self._calculate_field_angle(Br_conductor, B_tot)
        
        # Convert to numpy
        Bz_conductor_np = Bz_conductor.cpu().numpy()
        Br_conductor_np = Br_conductor.cpu().numpy()
        B_tot_np = B_tot.cpu().numpy()
        Field_Angle_np = Field_Angle.cpu().numpy()
        
        # Compute peak fields
        Bz_peak = torch.matmul(self.Bz_peak, I1_torch).cpu().numpy()
        Br_peak = torch.matmul(self.Br_peak, I1_torch).cpu().numpy()
        
        # Compute center fields - return as arrays for compatibility
        Bz_center = np.dot(self.Bz_center, I1)
        Br_center = np.dot(self.Br_center, I1)
        
        # Calculate current density
        I_density = (I1 / self.cross_sections[self.Index_inductive_connections] * 
                    self.turns_per_element[self.Index_inductive_connections] * 1e-6)
        
        # Compute forces
        force_results = self._compute_electromagnetic_forces(
            Bz_conductor_np, Br_conductor_np, Bz_peak, Br_peak, I1
        )
        
        return {
            'B_tot': B_tot_np,
            'Bz_conductor': Bz_conductor_np,
            'Br_conductor': Br_conductor_np,
            'Field_Angle': Field_Angle_np,
            'Bz_center': Bz_center,
            'Br_center': Br_center,
            'I_density': I_density,
            'forces': force_results
        }

    def _handle_quench_detection(self, t, T, V, timestamp):
        """
        Handle quench detection logic and trigger protection systems.
        
        Args:
            t: Current time
            T: Temperature array
            V: Voltage array  
            timestamp: Current timestamp string for logging
        """
        # Calculate voltage per circuit for detection
        V_circuit = [np.sum(V[self.circuits_list[i]]) for i in self.unique_circuits]
        
        # Check quench detection criteria
        voltage_quench = np.max(np.abs(V_circuit)) > self.V_detect
        temperature_quench = np.max(T) > self.T_detect
        
        if (voltage_quench or temperature_quench) and not self.quench_detected:
            # Quench detected for the first time
            self.quench_detected = True
            self.ramp_bool = False  # Stop normal ramping
            self.quench_detected_time = t
            
            logging.critical(f'\n{timestamp} || Quench detected\n')
            
            # Adjust sampling time for higher resolution during quench
            if self.overwrite_sample_time_custom:
                protection_window = (self.protection_delay + self.Capacitor_Discharge_delay) * 2
                self.sampling_time_custom = [t, t + protection_window, self.quench_sampling_time]
            
            # Activate capacitor discharge protection if enabled
            if self.Capacitor_Discharge_on_Quench_Detection:
                self.Cap_discharge = True
                self.Cap_discharge_time = t + self.protection_delay + self.Capacitor_Discharge_delay
            
            # Schedule power supply turn-off if enabled
            if self.PSU_off_on_Quench_Detection:
                self.PSU_off_time = t + self.protection_delay
                self.PSU_off = True
                self.PSU_off_t = self.PSU_off_time + self.PSU_off_time_constant

    def _manage_power_systems(self, t, T):
        """
        Manage power supply and heater system states.
        
        Args:
            t: Current time
            T: Temperature array
        """
        # Handle scheduled power supply shutdown
        if t >= self.PSU_off_time and not self.PSU_off:
            self.ramp_bool = False
            self.PSU_off = True
        
        # Handle heater temperature-based shutdown
        if self.heater and np.max(T) > self.heater_stop:
            self.heater = False

    def _compute_power_densities(self, Pin, Prad, P_AC):
        """
        Calculate power densities and energy integrals.
        
        Args:
            Pin: Input power array
            Prad: Radial power dissipation array
            P_AC: AC power dissipation array
            
        Returns:
            tuple: Power densities and energy integrals
        """
        # Calculate power densities (power per unit cross-sectional area)
        cross_sections = self.cross_sections[self.Index_normal_connections]
        Pin_density = Pin / cross_sections
        Prad_density = Prad / cross_sections
        P_AC_density = P_AC / cross_sections
        
        # Calculate integrated energy dissipation over time
        if len(self.t_array) > 1:
            Psum_AC = np.sum(integrate.trapezoid(np.array(self.P_AC_array).T, self.t_array))
            Psum_rad = np.sum(integrate.trapezoid(np.array(self.Prad_array).T, self.t_array))
        else:
            Psum_AC = 0.0
            Psum_rad = 0.0
        
        return Pin_density, Prad_density, P_AC_density, (Psum_AC, Psum_rad)

    def _handle_pid_control(self, Prad, P_AC, dt):
        """
        Handle PID control for power-limited current ramping.
        
        Args:
            Prad: Radial power dissipation array
            P_AC: AC power dissipation array
            dt: Current timestep
        """

        # Calculate total power for PID control if enabled
        if self.PID_total_power and any(self.pid_bool.values()):
            Power = 0.0
            for i in self.unique_circuits:
                Power += np.sum(Prad[self.circuits_list[i]]) + np.sum(P_AC[self.circuits_list[i]])
        
        # Apply PID control per circuit
        for i in self.unique_circuits:
            if self.pid_bool[i]:
                # Calculate circuit-specific power if not using total power
                if not self.PID_total_power:
                    Power = np.sum(Prad[self.circuits_list[i]]) + np.sum(P_AC[self.circuits_list[i]])
                
                # Apply PID control based on ramping direction
                current_magnitude = abs(self.I_circuit_last[i])
                if self.ramp_up:
                    self.rr[i] = self.PID(Power, i, dt) if current_magnitude < self.nominal_current[i] else 0.0
                elif self.ramp_down:
                    self.rr[i] = -1 * self.PID(Power, i, dt) if current_magnitude > self.nominal_current[i] else 0.0

    def _check_termination_criteria(self, E, T):
        """
        Check if simulation termination criteria are met.
        
        Args:
            E: Stored magnetic energy
            T: Temperature array (current timestep only, not historical)
        """
        # Check temperature stopping criterion
        if hasattr(self, 'T_array') and len(self.T_array) > 0:
            max_temp = np.max(self.T_array[-1])  # Use last stored temperature
            if (max_temp > self.stop_criterion_T) and self.stop_on_T:
                self.T_stop = False
        
        # Check energy stopping criterion
        if (E < self.stop_criterion_E) and self.stop_on_E:
            self.E_stop = False

    def _store_timestep_results(self, t, T, I1, I2, V_cap, I_cap, total_power, field_results, physics_results, energy_integrals):
        """
        Store results for the current timestep.
        
        Args:
            t: Current time
            T: Temperature array
            I1: Inductive currents
            I2: Resistive currents
            V_cap: Capacitor voltage
            I_cap: Capacitor current
            total_power: Total electrical power
            field_results: Magnetic field computation results (if computed)
            physics_results: Physics computation results from solve_this_torch
            energy_integrals: Energy integrals (if computed)
        """
        P_AC, Prad, Pin, B, V, P_cooling, SC_loss, Radial_loss, R_bypass, R_Conductor, IIc, x_array = physics_results
        
        if field_results is not None:
            # Use computed field results
            B_tot = field_results['B_tot']
            Bz_conductor = field_results['Bz_conductor']
            Br_conductor = field_results['Br_conductor']
            Field_Angle = field_results['Field_Angle']
            Bz_center = field_results['Bz_center']
            Br_center = field_results['Br_center']
            I_density = field_results['I_density']
            forces = field_results['forces']
        else:
            # Use placeholder values for non-computed timesteps
            n_inductive = len(I1)
            B_tot = np.zeros(n_inductive)
            Bz_conductor = np.zeros(n_inductive)
            Br_conductor = np.zeros(n_inductive)
            Field_Angle = np.zeros(n_inductive)
            Bz_center = 0.0
            Br_center = 0.0
            I_density = np.zeros(n_inductive)
            forces = {key: np.zeros(n_inductive) for key in ['Fz', 'Fr', 'Fz_peak', 'Fr_peak', 
                     'Fz_density', 'Fr_density', 'Fz_density_peak', 'Fr_density_peak']}
        
        # Calculate power densities
        if energy_integrals is not None:
            cross_sections_normal = self.cross_sections[self.Index_normal_connections]
            Pin_density = Pin / cross_sections_normal
            Prad_density = Prad / cross_sections_normal
            P_AC_density = P_AC / cross_sections_normal
        else:
            # Use placeholder values
            n_normal = len(I2)
            Pin_density = np.zeros(n_normal)
            Prad_density = np.zeros(n_normal)
            P_AC_density = np.zeros(n_normal)
        
        # Calculate stored magnetic energy using full current vector
        # Need to reconstruct full current vector from I1 and I2
        I_full = torch.zeros(len(self.I0), dtype=self.torch_float_type, device=self.torch_device)
        I_full[self.Index_inductive_connections] = torch.tensor(I1, dtype=self.torch_float_type, device=self.torch_device)
        I_full[self.Index_normal_connections] = torch.tensor(I2, dtype=self.torch_float_type, device=self.torch_device)
        E = torch.sum(0.5 * self.Mutual * torch.outer(I_full, I_full)).cpu().numpy()
        
        # Store all results
        self.store_results(
            t, T, I1, I2, P_AC, Prad, Pin, P_AC_density, Prad_density, Pin_density,
            B_tot, Bz_center, Br_center, E,
            V, forces['Fz'], forces['Fz_density'], forces['Fz_density_peak'],
            forces['Fr'], forces['Fr_density'], forces['Fr_density_peak'],
            Bz_conductor, Br_conductor, I_density, V_cap, I_cap, total_power,
            P_cooling, SC_loss, Radial_loss, R_bypass, R_Conductor, 
            Field_Angle, IIc, x_array
        )
        
        # Reset PID counter and mark first timestep as complete
        self.save_n_pid = 1

    def _handle_result_saving(self, SAVE_BOOL):
        """
        Handle periodic saving of results to disk.
        
        Args:
            SAVE_BOOL: Whether results should be saved this timestep
            
        Returns:
            str: Save indicator for logging
        """
        if SAVE_BOOL:
            self.save_results()
            return '*'  # Asterisk indicates save occurred
        else:
            return ''

    def _calculate_adaptive_timestep(self, t):
        """
        Calculate the next timestep based on solution dynamics and constraints.
        
        Args:
            t: Current time
            
        Returns:
            float: Next timestep size
        """
        # Timestep based on current and temperature derivatives
        dts = self.sampling_timestep()
        
        # Timestep based on circuit interpolation
        dtt = np.max([self.time_step_dt(t), 0.01]).item()
        
        # Timestep based on user-defined time points
        if t < np.max(self.unique_sorted_timesteps):
            next_timestep = np.min(self.unique_sorted_timesteps[self.unique_sorted_timesteps > t]).item()
            dtu = next_timestep - t - self.min_timestep
        else:
            dtu = 1e99  # Large value if beyond user time points
        
        # Take the minimum of all constraints, but respect global limits
        dt = np.max([
            np.min([dts, dtt, dtu, self.max_timestep]),
            self.min_timestep
        ])
        
        return dt

    def _log_simulation_progress(self, timestamp, save_indicator, n_ode, t, E, dt, T, I1, B, energy_integrals, P_AC, Prad):
        """
        Log simulation progress information.
        
        Args:
            timestamp: Current time string
            save_indicator: Save indicator character
            n_ode: ODE step counter
            t: Current simulation time
            E: Stored magnetic energy
            dt: Current timestep
            T: Temperature array
            I1: Inductive current array
            B: Magnetic field array
            energy_integrals: Tuple of (AC_energy, radial_energy)
            P_AC: AC power array
            Prad: Radial power array
        """
        Psum_AC, Psum_rad = energy_integrals
        
        output_string = (
            f"{timestamp} |{save_indicator}| {n_ode} || {utils.output_format(t, '{:.3f}')}, "
            f"E = {utils.output_format(E, '{:.2E}')} J, "
            f"E_density = {utils.output_format(E * 1e-3 / self.magnet_mass, '{:.2E}')} kJ/kg, "
            f"dt = {utils.output_format(dt, '{:.3E}')} s, "
            f"max T = {utils.output_format(np.max(T), '{:.2f}')} K, "
            f"Imax = {utils.output_format(np.max(abs(I1)), '{:.1f}')} A, "
            f"Bmax = {utils.output_format(np.max(B), '{:.2f}')} T, "
            f"E_AC = {utils.output_format(np.max(Psum_AC), '{:.2E}')} J, "
            f"Erad = {utils.output_format(np.max(Psum_rad), '{:.2E}')} J, "
            f"P_AC = {utils.output_format(np.sum(P_AC), '{:.3E}')} W, "
            f"Prad = {utils.output_format(np.sum(Prad), '{:.3E}')} W"
        )
        logging.info(output_string)

    def cleaning(self):
        """
        Perform final cleanup operations after simulation completion.

        This method handles post-simulation cleanup including saving any
        remaining buffered results, clearing GPU memory cache, and reporting
        memory usage statistics. Called automatically at simulation end or
        on user interruption.

        Side Effects:
            - Calls save_results() if unsaved data exists (save_n < store_n)
            - Empties CUDA cache if GPU is available
            - Prints GPU memory usage statistics
            - Runs garbage collection
        """

        # Saving final results if needed.
        if self.save_n < self.store_n:
            self.save_results()

        if torch.cuda.is_available():
            gc.collect()
            torch.cuda.empty_cache()
            print('Memory allocated:', torch.cuda.memory_allocated(device=self.torch_device))
            print('Max memory allocated:', torch.cuda.max_memory_allocated(device=self.torch_device))

        if self.profile_bool:
            profiler.disable()
            now = datetime.now().strftime("%H:%M:%S")
            logging.info(str(now) + ' || Storing profiler information.')
            s = io.StringIO()
            stats = pstats.Stats(profiler, stream=s).sort_stats('cumulative')
            stats.print_stats(25)
            profiler.dump_stats(self.sim_output + os.sep + 'profile.pstats')
            with open(self.sim_output + os.sep + 'profiler.txt', 'w+') as f:
                f.write(s.getvalue())
            # subprocess.call(['gprof2dot -f pstats', self.sim_output + os.sep + 'profile.pstats', '| dot -Tpng -o', self.sim_output + os.sep + 'profile.png'])


    def _parse_solver_inputs(self, x):
        """
        Parse and validate solver inputs from state vector.
        OPTIMIZED: Reduces tensor creation overhead and CPU transfers.
        
        Args:
            x (array): State vector containing currents, temperatures, and capacitor states
            
        Returns:
            tuple: Parsed tensors (I_tensor, T_tensor, I_B_tensor, Capacitor_Voltage, Capacitor_Current, T_CPU)
        """
        # CPU numpy variants of T for the material library - avoid copy if possible
        T_CPU = None
        if self.FLAG_USE_STEAM_MAT_LIB:
            if isinstance(x, np.ndarray):
                T_CPU = x[self.Nr_of_Connections:-2]  # Direct slice, no copy
            else:
                T_CPU = np.asarray(x[self.Nr_of_Connections:-2])
            T_CPU = np.clip(T_CPU, self.value_limits['T_rho_min'], self.value_limits['T_rho_max'])

        # Convert to torch tensor only if not already a tensor
        if isinstance(x, torch.Tensor):
            x_tensor = x.to(dtype=self.torch_float_type, device=self.torch_device)
        else:
            x_tensor = torch.tensor(x, dtype=self.torch_float_type, device=self.torch_device)
        
        # Direct indexing without reshape (more efficient)
        I_tensor = x_tensor[:self.Nr_of_Connections]
        T_tensor = x_tensor[self.Nr_of_Connections:-2]
        I_B_tensor = I_tensor[self.Index_inductive_connections]  # Reuse I_tensor
        Capacitor_Voltage = x_tensor[-2]
        Capacitor_Current = x_tensor[-1]

        # Limit temperature range for material property calculations
        if self.thermal_part:
            T_tensor = torch.clamp(T_tensor, min=self.value_limits['T_rho_min'], max=self.value_limits['T_rho_max'])  # clamp is faster than clip
            
        return I_tensor, T_tensor, I_B_tensor, Capacitor_Voltage, Capacitor_Current, T_CPU

    def _process_circuit_ramp(self, b, t, I_tensor, T_tensor, Capacitor_Current):
        """
        Helper function to process circuit ramp rates and power supply behavior.
        OPTIMIZED: Extracted from main solver to improve readability and maintainability.
        
        Args:
            b (torch.Tensor): The b-array for the ODE solver
            t (float): Current time
            I_tensor (torch.Tensor): Current tensor
            T_tensor (torch.Tensor): Temperature tensor
            Capacitor_Current (torch.Tensor): Capacitor current

        Returns:
            b (torch.Tensor): The b-array for the ODE solver
        """
        for i in self.unique_circuits:
            circ_index = self.circuits_list[i]
            sign_list = self.signs[circ_index]
            
            if self.quench_detected and t > self.PSU_off_time and self.PSU_off_on_Quench_Detection:
                # Quench protection
                Cap_cur = torch.zeros_like(T_tensor)
                if self.Dis_List:  # Only process if discharge list exists
                    Cap_cur[self.Dis_List] = Capacitor_Current
                b[circ_index] -= (torch.matmul(self.power_dis_matrix, I_tensor) - Cap_cur)[circ_index] * sign_list / self.PSU_off_time_constant
                
            elif self.PSU_off and not self.quench_detected:
                # Fast discharge based on the current in each circuit
                Cap_cur = torch.zeros_like(T_tensor)
                if self.Dis_List:  # Only process if discharge list exists
                    Cap_cur[self.Dis_List] = Capacitor_Current
                b[circ_index] -= (torch.matmul(self.power_dis_matrix, I_tensor) - Cap_cur)[circ_index] * sign_list / self.PSU_off_time_constant
                
            elif self.ramp_bool and self.pid_bool[i]:
                # Standard PID loop ramp rate
                b[circ_index] += self.rr[i]
                
            elif self.ramp_bool and not self.pid_bool[i] and i in self.rr_remap:
                # Non-standard PID loop ramp rate based on another circuit
                if self.pid_bool[self.rr_remap[i][0]]:
                    b[circ_index] += self.rr[self.rr_remap[i][0]] * self.rr_remap[i][1]
                else:
                    ramp_val = self.ramp_rate(t, self.rr_remap[i][0]) * self.rr_remap[i][1]
                    b[circ_index] += torch.tensor(ramp_val, dtype=self.torch_float_type, device=self.torch_device)
                    
            elif self.ramp_bool and not self.pid_bool[i]:
                # Standard ramp scheme as specified in yaml
                ramp_val = self.ramp_rate(t, i)
                b[circ_index] += torch.tensor(ramp_val, dtype=self.torch_float_type, device=self.torch_device)

            # Apply correct sign to the ramp
            b[circ_index] *= sign_list
        return b

    def _calculate_field_angle(self, Br: torch.Tensor, B_tot: torch.Tensor) -> torch.Tensor:
        """
        Calculate the magnetic field angle from radial and total field components.
        
        The field angle is defined as the angle between the magnetic field vector 
        and the axial direction (z-axis), calculated from the radial component.
        
        Args:
            Br (torch.Tensor): Radial magnetic field component [T]
            B_tot (torch.Tensor): Total magnetic field magnitude [T]
            
        Returns:
            torch.Tensor: Field angle in degrees [0-180°]
            
        Notes:
            - Angle is calculated as arccos(Br / B_tot) * 180 / π
            - NaN values (from division by zero) are replaced with 0.0
            - Result is reshaped to 1D tensor
        """
        # Calculate the angle between field vector and z-axis using radial component
        # arccos(Br/B_tot) gives angle from z-axis in radians
        angle_rad = torch.arccos(Br / B_tot)
        
        # Convert to degrees and take absolute value to ensure positive angles
        angle_deg = torch.abs(angle_rad) * 180.0 / torch.pi
        
        # Handle NaN values (when B_tot = 0) by replacing with 0.0
        angle_clean = torch.nan_to_num(angle_deg, nan=0.0)
        
        # Reshape to 1D tensor as expected by the calling code
        return torch.reshape(angle_clean, (-1,))

    def _compute_magnetic_fields(self, I_B_tensor):
        """
        Helper function to compute magnetic fields efficiently.
        OPTIMIZED: Centralized magnetic field calculations with better memory usage.
        
        Args:
            I_B_tensor (torch.Tensor): Inductive currents tensor
            
        Returns:
            tuple: (B_tot, B_tensor, Bavg, angle)
        """
        # Compute magnetic fields - avoid unnecessary clones and use faster operations
        Bz = torch.matmul(self.Bz_Conductor, I_B_tensor)
        Br = torch.matmul(self.Br_Conductor, I_B_tensor)
        B_tot = torch.sqrt(torch.addcmul(Bz.square(), Br, Br))  # Faster than Bz**2 + Br**2

        # Reuse existing tensors instead of cloning - major memory allocation reduction
        B_tensor = torch.zeros_like(self.B0)
        B_tensor[self.Index_inductive_connections] = B_tot
        
        # Pre-compute this division to avoid repeated calculations
        Bavg = torch.matmul(self.power_dis_matrix, B_tensor) / self.power_dis_matrix_sum

        # Field angle calculation using helper function
        angle = torch.zeros_like(self.B0)
        angle[self.Index_inductive_connections] = self._calculate_field_angle(Br, B_tot)

        angle = torch.clamp(angle, min=self.value_limits['Angle_min'], max=self.value_limits['Angle_max'])
        
        return B_tot, B_tensor, Bavg, angle

    def solve_this_torch(self, t, x, output=1):
        """
        Compute time derivatives for the coupled electro-thermal ODE system.

        This is the core ODE function solved by the integrator. It computes
        dI/dt (current derivatives) from Faraday's law: dI/dt = inv(M) * b,
        where b contains resistive voltages and external ramp rates. It also
        computes dT/dt (temperature derivatives) from the heat equation with
        Joule heating, cooling, and thermal conduction.

        The method handles:
        - Superconductor current sharing and transition resistance
        - Normal conductor resistance (temperature/field dependent)
        - Radial turn-to-turn resistance in NI coils
        - Circuit ramp rates and power supply behavior
        - Capacitor discharge protection
        - Magnetic field calculation and coupling

        Args:
            t: Current simulation time [s]
            x: State vector [I (currents), T (temperatures), V_cap, I_cap]
            output: Return mode. 1=derivatives (for ODE), 2=power/voltages
                (for post-processing and logging)

        Returns:
            If output=1: Array of time derivatives [dI/dt, dT/dt, dV/dt, dI/dt]
            If output=2: Tuple (P_AC, Prad, Pin, B, V, P_cooling, SC_loss,
                Radial_loss, R_bypass, R_Conductor, Field_Angle, IIc, x_array)
        """

        # Parse and validate solver inputs
        I_tensor, T_tensor, I_B_tensor, Capacitor_Voltage, Capacitor_Current, T_CPU = self._parse_solver_inputs(x)

        # Making the b-array filled with zeros.
        b = self.b0.detach().clone()

        # Process circuit ramp rates and power supply behavior
        b = self._process_circuit_ramp(b, t, I_tensor, T_tensor, Capacitor_Current)

        # Compute magnetic fields efficiently
        B_tot, B_tensor, Bavg, angle = self._compute_magnetic_fields(I_B_tensor)

        # Clone zero tensors
        R_Conductor = self.B0.detach().clone()
        R_bypass = self.B0.detach().clone()
        IIc_array = self.B0.detach().clone() if output == 2 else None
        x_array = self.B0.detach().clone() if output == 2 else None

        for i in self.unique_regions:
            if self.Conductor_type[i].lower() == 'superconductor':
                index = self.Index_inductive_connections_region[i]
                I_IND = I_tensor[index]
                T_IND = T_tensor[self.Index_temperatures[index]]
                B_IND = torch.clip(B_tensor[index], min=self.B_min)
                Deg_IND = angle[index] if self.deg_dep[i] else torch.zeros_like(angle[index])

                # Apply angle mirroring if enabled for this region
                # Mirrors angles around 90 deg (e.g., 150->30, 40->140, 90->90)
                if self.Conductor_SC_properties[i].get('mirror_angle', False):
                    Deg_IND = 180.0 - Deg_IND

                if self.FLAG_USE_STEAM_MAT_LIB:
                    # Temperatures of the inductive elements to CPU, clipping for material functions.
                    T_IND_CPU = np.clip(T_IND.cpu().numpy(), self.value_limits['T_rho_min'], self.value_limits['T_rho_max'])

                    # Clipping this only for material function input.
                    B_IND_CPU = np.clip(B_IND.cpu().numpy(), self.value_limits['B_min'], self.value_limits['B_max'])

                    # Field angle, already clipped.
                    Deg_IND_CPU = Deg_IND.cpu().numpy()

                    # Critical Current of superconductor i
                    fit_name = self.Conductor_SC_properties[i]['fit-name']
                    temp = MaterialProperties.create_steam_material_input(
                            fit_name,
                            Temperature=T_IND_CPU,
                            Magnetic_Field=B_IND_CPU,
                            Field_Angle=Deg_IND_CPU
                        )
                    max_points = temp.shape[1] if len(temp.shape) > 1 else temp.shape[0]
                    steam = self.materials.get_cached_steam_material(fit_name, max_points)
                    Ic = steam.evaluate(temp) * self.SC_Ic_Scaling[i].cpu().numpy()
                    Ic = torch.clip(torch.tensor(Ic, dtype=self.torch_float_type, device=self.torch_device), min=1e-99)
                else:
                    T_IND = torch.clamp(T_IND, self.value_limits['T_rho_min'], self.value_limits['T_rho_max'])
                    B_IND = torch.clamp(T_IND, self.value_limits['B_min'], self.value_limits['B_max'])
                    Ic = self.SC_Ic_Fit[i]((T_IND, B_IND, Deg_IND))

                # Resistance of superconductor 'i' per one meter of tape [Ohm/m].
                R = torch.zeros_like(T_IND)
                for l in range(len(self.Conductor_materials[i])):
                    if self.FLAG_USE_STEAM_MAT_LIB:
                        rho = torch.tensor(self.materials.resistivity_selector(self.Conductor_materials[i][l], T_IND_CPU, B_IND_CPU, reshape=True, set_B=self.set_B, B_min=self.B_min), dtype=self.torch_float_type, device=self.torch_device)
                        R += 1 / (rho / self.Conductor_materials_cross_section[i][l])
                    else:
                        R += 1 / (self.materials.resistivity_selector_torch(self.Conductor_materials[i][l], T_IND, B_IND, Normal_rho_FIT=self.Normal_rho_FIT) / self.Conductor_materials_cross_section[i][l])

                # The resistance is calculated here and multiplied by the amount of parallel sections.
                # Because each section is connected in parallel.
                R = (1 / R) * self.Conductor_SC_properties[i]['parallel']

                # For the x interpolation function, the inputs are clipped to assure they will not be outside the LUT range.
                IIc = torch.clip(torch.abs(I_IND / Ic), min=self.lut_limits['IIc_min'], max=self.lut_limits['IIc_max'])
                U = torch.clip(torch.abs(R * I_IND), min=self.lut_limits['U_min'], max=self.lut_limits['U_max'])
                current_sharing_x = self.SC_Cur_Sharing[i]((IIc, U))

                # One newton iteration:
                err = self.powerlaw(current_sharing_x, IIc, U, self.Conductor_SC_properties[i]['n-value'])
                current_sharing_x -= err / self.powerlaw_jac(current_sharing_x, IIc, U, self.Conductor_SC_properties[i]['n-value'])

                # Calculating the resistance.
                R_Conductor[index] = torch.abs((R * (1 - current_sharing_x))) * self.lengths_torch[index]

                if output == 2:
                    # Adding the critical current ratio to an array that will be returned:
                    IIc_array[index] = IIc
                    x_array[index] = current_sharing_x


            # Normal conducting parts, both for the radial resistance and the normal conducting conductors.
            index_norm = self.Index_normal_connections_region[i]
            index_norm2 = self.Index_temperatures[index_norm]
            if self.FLAG_USE_STEAM_MAT_LIB:
                resistivity = self.materials.resistivity_selector(self.Conductor_resistor[i], T_tensor[index_norm2].cpu().numpy(), Bavg[index_norm2].cpu().numpy(), reshape=True)
                resistivity = torch.tensor(resistivity, dtype=self.torch_float_type, device=self.torch_device)
            else:
                resistivity = self.materials.resistivity_selector_torch(self.Conductor_resistor[i], T_tensor[index_norm2], Bavg[index_norm2], Normal_rho_FIT=self.Normal_rho_FIT)
            R_bypass[index_norm] = self.ref_resistance[index_norm] * resistivity / self.ref_resistivity[index_norm]

        # Separated for clarity.
        b[self.b_V_array_index] += -(R_Conductor * I_tensor) # Voltage due to the resistance in the SC
        b[self.b_V_array_index] += -(R_bypass * I_tensor) # Voltage due to radial resistance and current

        # Capacitor discharge: #
        if self.Cap_discharge and t > self.Cap_discharge_time:
            dIdt_cap = (Capacitor_Voltage - Capacitor_Current * (torch.sum(R_bypass[self.Dis_List_Resistance]) + self.Discharge_R0)) / self.Discharge_L0
            dVdt_cap = -Capacitor_Current / self.Discharge_C0
            for i in self.Dis_List:
                b[i] += dIdt_cap
        else:
            dVdt_cap = self.torch_zero.detach().clone()
            dIdt_cap = self.torch_zero.detach().clone()

        if self.thermal_part:
            # Thermal part: only the ohmic heat
            SC_loss = I_tensor ** 2 * R_Conductor
            Radial_loss = I_tensor ** 2 * R_bypass

            Pdis1 = torch.matmul(self.power_dis_matrix, SC_loss)
            Pdis2 = torch.matmul(self.power_dis_matrix, Radial_loss)
            dQ = torch.zeros_like(T_tensor, dtype=self.torch_float_type, device=self.torch_device)
            Pin = torch.zeros_like(T_tensor, dtype=self.torch_float_type, device=self.torch_device)
            P_cooling_sum = torch.tensor(0.0, dtype=self.torch_float_type, device=self.torch_device)
            dQ +=  Pdis1+Pdis2

            if self.heater:
                for ii in self.heater_nodes:
                    dQ[ii] += self.heater_power * utils.sig(t-self.heater_time, self.heater_sig_c)

            if self.quench_detected and self.Heaters_on_Quench_Detection and self.Quench_Heater_positions is not None:
                for ii in self.Quench_Heater_positions:
                    dQ[ii] += self.Quench_Heater_power * utils.sig(t-self.quench_detected_time+self.protection_delay+self.Quench_Heater_time, self.heater_sig_c)

            ## Cooling:
            # Provides global cooling
            # if self.max_cooling:
            #     dQ = np.where(dQ < self.max_cooling_value, 0.0, dQ-self.max_cooling_value)
                #### dQ = np.where(dQ < self.max_cooling_value, -1.0, dQ-self.max_cooling_value)
                #### dQ = np.where((T <= 2.0) & (dQ < 0.0), 0.0, dQ)

            ## Thermal connections in the radial direction.
            if self.Thermal_alpha is not None:
                T1 = T_tensor[self.T1_con]
                T2 = T_tensor[self.T2_con]

                # Updating the thermal resistance.
                for i in self.unique_regions:
                    thermal_sum = torch.zeros_like(T_tensor[self.Thermal_index[i]])
                    for j in self.Conductor_conductivity_materials[i]:
                        if self.FLAG_USE_STEAM_MAT_LIB:
                            thermal_sum_temp = self.materials.conductivity_selector(j, T_CPU[self.Thermal_index[i]], reshape=True) * self.Thermal_area[i][j].cpu().numpy()
                            thermal_sum = thermal_sum + torch.tensor(thermal_sum_temp, dtype=self.torch_float_type, device=self.torch_device)
                        else:
                            thermal_sum = thermal_sum + self.materials.conductivity_selector_torch(j, T_tensor[self.Thermal_index[i]], Normal_k_FIT=self.Normal_k_FIT) * self.Thermal_area[i][j]
                    self.Thermal_resistance[self.Thermal_index[i]] = self.Thermal_alpha[self.Thermal_index[i]] * 1 / thermal_sum

                # Calculating the heat flows.
                dQ_cond = (T1 - T2) / (self.Thermal_resistance[self.T1_con] + self.Thermal_resistance[self.T2_con])

                if self.thermal_conductivity_bool:
                    # Adding and substracting the heat flows from the corresponding nodes.
                    dQ.index_add_(0, self.T2_con_tensor, dQ_cond)
                    dQ.index_add_(0, self.T1_con_tensor, -dQ_cond)

                P_cooling_sum = self.torch_zero.detach().clone() # Set cooling sum value to 0.

                if self.thermal_cooling_bool and self.inner_cooling_bool:
                    # Helium cooling on the pancakes' inner radius
                    T_inner = T_tensor[self.T_con_inner]
                    P_cooling = torch.zeros_like(T_inner)
                    if self.cross_sections_inner is not None:
                        P_cooling += self.coolant_h * self.cross_sections_inner[self.T_con_inner] * (T_inner - self.coolant_temperature)
                    P_cooling_sum += torch.sum(P_cooling)
                    dQ[self.T_con_inner] -= P_cooling

                if self.thermal_cooling_bool and self.outer_cooling_bool:
                    # Helium cooling on the pancakes' inner radius
                    T_outer = T_tensor[self.T_con_outer]
                    P_cooling = torch.zeros_like(T_outer)
                    if self.cross_sections_inner is not None:
                        P_cooling += self.coolant_h * self.cross_sections_outer[self.T_con_outer] * (T_outer - self.coolant_temperature)
                    P_cooling_sum += torch.sum(P_cooling)
                    dQ[self.T_con_outer] -= P_cooling

            C = torch.zeros_like(dQ, dtype=self.torch_float_type, device=self.torch_device)
            for i in self.unique_regions:
                index_norm = self.Index_normal_connections_region[i]
                index_norm2 = self.Index_temperatures[index_norm]
                T_IND = T_tensor[index_norm2]
                if self.FLAG_USE_STEAM_MAT_LIB:
                    T_IND_CPU = T_CPU[index_norm2]
                for l in range(len(self.Conductor_materials[i])):
                    if self.FLAG_USE_STEAM_MAT_LIB:
                        C[index_norm2] += torch.tensor(self.materials.specific_heat_selector(self.Conductor_materials[i][l], T_IND_CPU, reshape=True), dtype=self.torch_float_type, device=self.torch_device) * self.Conductor_materials_cross_section[i][l] / self.Conductor_materials_cross_section_sum[i] * self.volumes_torch[index_norm2]
                    else:
                        C[index_norm2] += self.materials.specific_heat_selector_torch(self.Conductor_materials[i][l], T_IND, Normal_Cv_FIT=self.Normal_Cv_FIT) * self.Conductor_materials_cross_section[i][l] / self.Conductor_materials_cross_section_sum[i] * self.volumes_torch[index_norm2]
            dTdt = dQ / (C * self.quench_fraction) * utils.sig(t - self.thermal_timer, self.thermal_timer_sig_c)
        else:
            dTdt = self.T0_tensor.clone().detach()
            Pdis1 = self.T0_tensor.clone().detach()
            Pdis2 = self.T0_tensor.clone().detach()
            Pin = self.T0_tensor.clone().detach()

        dIdt = torch.matmul(self.inv_1, b)

        # Output 1: only for the solver.
        if output == 1:
            return torch.cat((dIdt,dTdt,torch.unsqueeze(dVdt_cap,0),torch.unsqueeze(dIdt_cap,0))).cpu().numpy()
        if output == 2:
            self.I_circuit_last = dict()
            for i in self.unique_circuits:
                Cap_cur = torch.zeros_like(T_tensor)
                Cap_cur[self.Dis_List] = Capacitor_Current
                self.I_circuit_last[i] = (torch.matmul(self.power_dis_matrix, I_tensor) - Cap_cur)[self.circuits_list[i]][0].cpu().numpy()
            V = torch.matmul(self.inv_2, b)
            self.V_circuit_last = []
            for i in self.unique_circuits:
                circ_index = self.circuits_list[i]
                self.V_circuit_last.append(torch.sum(V[circ_index] * self.signs[circ_index]).cpu().numpy())
            return Pdis1.cpu().numpy(), Pdis2.cpu().numpy(), Pin.cpu().numpy(), B_tensor.cpu().numpy(), V.cpu().numpy(), P_cooling_sum.cpu().numpy(), SC_loss.cpu().numpy(), Radial_loss.cpu().numpy(), R_bypass.cpu().numpy(), R_Conductor.cpu().numpy(), IIc_array.cpu().numpy(), x_array.cpu().numpy()

    def ramp_rate(self, t, circuit = 0):
        """
        Get the ramp rate (dI/dt) for a circuit at a given time.

        Returns the power supply ramp rate from either the external ramp rate
        file (if use_external_ramp_rate is True) or from the circuit_data
        interpolation functions created during initialization.

        Args:
            t: Current time [s]
            circuit: Circuit index, default 0

        Returns:
            float: Ramp rate [A/s] for the specified circuit at time t
        """
        if self.use_external_ramp_rate:
            if circuit == 1:
                rr = self.external_ramp_rate(t)
            else:
                rr = 0.0
        else:
            if circuit == 0:
                rr = 0
            else:
                rr = self.circuit_interp_dIdt[circuit](t)
        return rr

    def store_results(self, t, T, I, IR, P_AC, Prad, Pin, P_AC_density, Prad_density, Pin_density, B, Bz_center, Br_center, E, V, Fz, Fz_density, Fz_density_peak, Fr, Fr_density, Fr_density_peak, Bz_conductor, Br_conductor, I_density, Capacitor_Voltage, Capacitor_Current, total_power, P_cooling, SC_loss, Radial_loss, R_bypass, R_Conductor, Field_Angle, IIc, x_array):  # , E, P, D, Bz, Br):
        """
        Store simulation results in memory arrays for later saving.

        This method appends the current timestep's results to in-memory arrays.
        Results include currents, temperatures, powers, forces, magnetic fields,
        and various derived quantities. Data is stored in memory until
        save_results() is called to write to CSV files.

        Args:
            All simulation quantities at the current timestep (t, T, I, etc.)
            See method signature for complete list.

        Side Effects:
            - Increments self.store_n counter
            - Appends values to numerous *_array attributes
            - Computes and stores derived quantities (forces per coil/pancake,
              power per coil/solenoid, radial resistances, etc.)
        """
        self.store_n += 1
        self.t_array.append(t)
        self.T_array.append(T)
        self.I_array.append(I)
        self.IR_array.append(IR)
        self.P_AC_array.append(P_AC)
        self.Prad_array.append(Prad)
        self.Pin_array.append(Pin)
        self.P_AC_density_array.append(P_AC_density)
        self.Prad_density_array.append(Prad_density)
        self.Pin_density_array.append(Pin_density)
        self.B_array.append(B)
        self.Bz_center_array.append(Bz_center)
        self.Br_center_array.append(Br_center)
        self.E_array.append(E)
        self.I_circuit_array.append([self.I_circuit_last[i] for i in self.unique_circuits])
        self.V_circuit_array.append(self.V_circuit_last.copy())
        self.V_array.append(V)
        self.Fz_array.append(Fz)
        self.Fz_density_array.append(Fz_density)
        self.Fz_density_peak_array.append(Fz_density_peak)
        self.Fr_array.append(Fr)
        self.Fr_density_array.append(Fr_density)
        self.Fr_density_peak_array.append(Fr_density_peak)
        self.I_density_array.append(I_density)
        self.Bz_array.append(Bz_conductor)
        self.Br_array.append(Br_conductor)
        self.Capacitor_Voltage_array.append(Capacitor_Voltage)
        self.Capacitor_Current_array.append(Capacitor_Current)
        self.Quench_detected_array.append(self.quench_detected)
        self.solver_t_array.append(time.time() - self.start)
        self.Power_per_Coil_array.append([np.sum(P_AC[np.array(self.thermal_element_to_coil_nr) == i]) + np.sum(Prad[np.array(self.thermal_element_to_coil_nr) == i]) + np.sum(Pin[np.array(self.thermal_element_to_coil_nr) == i]) for i in set(self.thermal_element_to_coil_nr)])
        self.Power_per_Solenoid_array.append([np.sum(P_AC[np.array(self.thermal_element_to_solenoid_nr) == i]) + np.sum(Prad[np.array(self.thermal_element_to_solenoid_nr) == i]) + np.sum(Pin[np.array(self.thermal_element_to_solenoid_nr) == i]) for i in set(self.thermal_element_to_solenoid_nr)])
        self.PSU_power_array.append(total_power)
        self.P_cooling_array.append(P_cooling)
        self.SC_loss_array.append(SC_loss[self.Index_inductive_connections])
        self.Radial_loss_array.append((Radial_loss[self.Index_radial_connections + self.Index_struct_connections]/self.power_dis_matrix_sum.cpu().numpy())[self.Index_temperatures][self.Index_inductive_connections]) # -> to only the radial loss -> spreading it -> moving it to the inductor elements for plotting. # TODO divide by the number of elements.
        self.Field_Angle_array.append(Field_Angle)
        self.Critical_Ratio_array.append(IIc)
        self.Current_Sharing_x_array.append(x_array)

        # Force per coil and force per pancake from the force arrays:
        self.Fz_per_Coil_array.append([np.sum(Fz[self.index_dict['coils'][i]['I']]) for i in self.index_dict['coils'].keys()])
        self.Fr_per_Coil_array.append([np.sum(Fr[self.index_dict['coils'][i]['I']]) for i in self.index_dict['coils'].keys()])
        self.Fz_per_Pancake_array.append([np.sum(Fz[self.index_dict['pancakes'][i]['I']]) for i in self.index_dict['pancakes'].keys()])
        self.Fr_per_Pancake_array.append([np.sum(Fr[self.index_dict['pancakes'][i]['I']]) for i in self.index_dict['pancakes'].keys()])

        # Radial resistance per coil and per pancake.
        self.R_Radial_array.append(R_bypass[self.Index_normal_connections])
        self.R_Conductor_array.append(R_Conductor[self.Index_inductive_connections])
        self.R_Radial_per_Coil_array.append([np.sum(R_bypass[self.index_dict['coils'][i]['IR']]) for i in self.index_dict['coils'].keys()])
        self.R_Radial_per_Pancake_array.append([np.sum(R_bypass[self.index_dict['pancakes'][i]['IR']]) for i in self.index_dict['pancakes'].keys()])

    def save_results(self):
        """
        Write stored results from memory arrays to CSV files.

        This method writes buffered simulation results to disk, appending data
        to the appropriate CSV files based on configured output flags
        (bool_summary, bool_temperature, etc.). It writes all data from
        save_n to store_n, then updates save_n.

        Side Effects:
            - Appends data to multiple CSV files in self.sim_output directory
            - Updates self.save_n to match self.store_n
            - Logs errors if file writes fail
            - Files written: summary, temperature, current, voltage, magnetic
              field, loss, force (depending on bool_* flags)
        """
        saverange = np.arange(self.save_n, self.store_n)

        if self.bool_summary:
            # Vectorized summary writing - much faster than row-by-row formatting
            try:
                if len(saverange) == 0:
                    return

                # Pre-compute all columns for all timesteps in saverange
                # This avoids millions of individual utils.output_format() calls
                n_rows = len(saverange)

                # Build each column as a numpy array
                time_vals = np.array([self.t_array[j] + self.t_offset
                                     for j in saverange])
                T_max_vals = np.array([np.max(self.T_array[j])
                                      for j in saverange])
                T_mean_vals = np.array([np.mean(self.T_array[j])
                                       for j in saverange])
                E_vals = np.array([self.E_array[j] for j in saverange])

                # I_circuit_array[j] can be an array, need to handle carefully
                I_circuit_vals = np.array([self.I_circuit_array[j]
                                          for j in saverange])

                Bz_center_vals = np.array([self.Bz_center_array[j]
                                          for j in saverange])
                Br_center_vals = np.array([self.Br_center_array[j]
                                          for j in saverange])
                P_cooling_sum_vals = np.array([np.sum(self.P_cooling_array[j])
                                              for j in saverange])
                P_AC_sum_vals = np.array([np.sum(self.P_AC_array[j])
                                         for j in saverange])
                Prad_sum_vals = np.array([np.sum(self.Prad_array[j])
                                         for j in saverange])

                # V_circuit_array[j] can be an array
                V_circuit_vals = np.array([self.V_circuit_array[j]
                                          for j in saverange])

                Fr_max_vals = np.array([np.max(self.Fr_density_array[j])
                                       for j in saverange])
                Fz_max_vals = np.array([np.max(self.Fz_density_array[j])
                                       for j in saverange])
                Quench_vals = np.array([self.Quench_detected_array[j]
                                       for j in saverange])
                Cap_V_vals = np.array([self.Capacitor_Voltage_array[j]
                                      for j in saverange])
                Cap_I_vals = np.array([self.Capacitor_Current_array[j]
                                      for j in saverange])
                solver_t_vals = np.array([self.solver_t_array[j]
                                         for j in saverange])
                PSU_power_vals = np.array([self.PSU_power_array[j]
                                          for j in saverange])

                # Write rows with efficient string formatting
                # Helper function to format values that might be scalars or arrays
                def format_val(val, fmt="{:.5E}"):
                    """Format a value that might be a scalar or numpy array."""
                    # Extract scalar from array if needed
                    if isinstance(val, np.ndarray):
                        if val.size == 1:
                            val = val.item()
                        else:
                            raise ValueError(f"Expected scalar, got array of size {val.size}")
                    return fmt.format(val)

                with open(self.file_summary, 'a+') as f:
                    for idx in range(n_rows):
                        # Build data list using numpy formatted strings
                        data = [
                            str(time_vals[idx]),
                            format_val(T_max_vals[idx], "{:.2f}"),
                            format_val(T_mean_vals[idx], "{:.2f}"),
                            format_val(E_vals[idx])
                        ]

                        # Expand I_circuit array elements
                        I_circ = I_circuit_vals[idx]
                        if np.isscalar(I_circ) or (isinstance(I_circ, np.ndarray) and I_circ.size == 1):
                            data.append(format_val(I_circ))
                        else:
                            data.extend([format_val(val) for val in np.atleast_1d(I_circ)])

                        # Continue with scalar values
                        data.extend([
                            format_val(Bz_center_vals[idx]),
                            format_val(Br_center_vals[idx]),
                            format_val(P_cooling_sum_vals[idx]),
                            format_val(P_AC_sum_vals[idx]),
                            format_val(Prad_sum_vals[idx])
                        ])

                        # Expand V_circuit array elements
                        V_circ = V_circuit_vals[idx]
                        if np.isscalar(V_circ) or (isinstance(V_circ, np.ndarray) and V_circ.size == 1):
                            data.append(format_val(V_circ))
                        else:
                            data.extend([format_val(val) for val in np.atleast_1d(V_circ)])

                        # Final scalar values
                        data.extend([
                            format_val(Fr_max_vals[idx]),
                            format_val(Fz_max_vals[idx]),
                            str(Quench_vals[idx]),
                            format_val(Cap_V_vals[idx]),
                            format_val(Cap_I_vals[idx]),
                            format_val(solver_t_vals[idx]),
                            format_val(PSU_power_vals[idx])
                        ])

                        f.write(self.delimiter.join(data) + '\n')

            except OSError:
                logging.info("Writing on %s has failed." % self.file_temperature)

        if self.bool_temperature:
            self.append_to_csv(self.file_temperature, self.T_array, saverange, output_format="{:0.2f}")

        if self.bool_current:
            self.append_to_csv(self.file_current, self.I_array, saverange)
            self.append_to_csv(self.file_current_R, self.IR_array, saverange)
            self.append_to_csv(self.file_I_density, self.I_density_array, saverange)
            self.append_to_csv(self.file_R_Radial_per_Coil, self.R_Radial_per_Coil_array, saverange)
            self.append_to_csv(self.file_R_Radial_per_Pancake, self.R_Radial_per_Pancake_array, saverange)

        if self.bool_voltage:
            self.append_to_csv(self.file_voltage, self.V_array, saverange)

        if self.bool_magnetic_field:
            self.append_to_csv(self.file_magnetic_field, self.B_array, saverange)
            self.append_to_csv(self.file_Bz_conductor, self.Bz_array, saverange)
            self.append_to_csv(self.file_Br_conductor, self.Br_array, saverange)

        if self.bool_loss:
            self.append_to_csv(self.file_Pin, self.Pin_array, saverange)
            self.append_to_csv(self.file_P_AC, self.P_AC_array, saverange)
            self.append_to_csv(self.file_Prad, self.Prad_array, saverange)
            self.append_to_csv(self.file_Pin_density, self.Pin_density_array, saverange)
            self.append_to_csv(self.file_P_AC_density, self.P_AC_density_array, saverange)
            self.append_to_csv(self.file_Prad_density, self.Prad_density_array, saverange)
            self.append_to_csv(self.file_Power_per_Coil_Element, self.Power_per_Coil_array, saverange)
            self.append_to_csv(self.file_Power_per_Solenoid, self.Power_per_Solenoid_array, saverange)

        if self.bool_force:
            self.append_to_csv(self.file_Fz, self.Fz_array, saverange)
            self.append_to_csv(self.file_Fr, self.Fr_array, saverange)
            self.append_to_csv(self.file_Fz_density, self.Fz_density_array, saverange)
            self.append_to_csv(self.file_Fr_density, self.Fr_density_array, saverange)
            self.append_to_csv(self.file_Fz_density_peak, self.Fz_density_peak_array, saverange)
            self.append_to_csv(self.file_Fr_density_peak, self.Fr_density_peak_array, saverange)
            self.append_to_csv(self.file_Fz_per_Coil, self.Fz_per_Coil_array, saverange)
            self.append_to_csv(self.file_Fr_per_Coil, self.Fr_per_Coil_array, saverange)
            self.append_to_csv(self.file_Fz_per_Pancake, self.Fz_per_Pancake_array, saverange)
            self.append_to_csv(self.file_Fr_per_Pancake, self.Fr_per_Pancake_array, saverange)

        self.save_n = self.store_n

    

    def sampling_timestep(self):
        """
        Calculate adaptive sampling timestep based on solution variation.

        Determines the next data output timestep based on the rate of change
        of temperature and current. The timestep is chosen to keep changes
        below configured thresholds (sampling_time_dT, sampling_time_dI),
        and is bounded by min_timestep and max_timestep.

        Returns:
            float: The calculated sampling timestep [s]

        Side Effects:
            - May update self.max_timestep if keep_smallest_timestep is True
            - Considers sampling_time_custom for special periods (e.g., discharge)
        """
        # Selects an appropriate time step
        dI = self.sampling_time_dI
        dT = self.sampling_time_dT
        dt = self.t_array[-1] - self.t_array[-2]

        # Check maximum temperature based timestep
        if self.thermal_part:
            dTdt = max(abs(self.T_array[-1] - self.T_array[-2]) / dt)
            T_ts = dT / dTdt if dTdt > 0 else self.max_timestep
        else:
            T_ts = self.max_timestep

        # Check maximum current timestep
        dIdt = max(abs(self.I_array[-1] - self.I_array[-2]) / dt)
        I_ts = dI / dIdt if dIdt > 0 else self.max_timestep

        # Taking the smallest of the two:
        sampling_ts = min(T_ts, I_ts)

        # Checking if the timestep is not too big
        if sampling_ts > self.max_timestep:
            sampling_ts = self.max_timestep
        elif sampling_ts < self.min_timestep:
            sampling_ts = self.min_timestep
        if sampling_ts < self.max_timestep and self.keep_smallest_timestep:
            self.max_timestep = sampling_ts

        # Custom - for capacitor discharge
        if self.sampling_time_custom is not None:
            if self.sampling_time_custom[0] <= self.t_array[-1] <= self.sampling_time_custom[1]:
                sampling_ts = self.sampling_time_custom[2] if sampling_ts > self.sampling_time_custom[2] else sampling_ts

        return sampling_ts

    def create_ramp_rate_interpolation(self):
        """
        Create interpolation functions for circuit ramp rates.

        Builds scipy interp1d functions for each circuit's dI/dt profile from
        the circuit_data dictionary. These interpolators are used during
        simulation to determine the power supply ramp rate at any time.

        Side Effects:
            - Populates self.circuit_interp_dIdt with interpolation functions
            - Calls debug_graphs.plot_ramp_rate_interpolation() for visualization
        """

        for i in self.unique_circuits:
            self.circuit_interp_dIdt[i] = interp1d(self.circuit_data[i]['t'], self.circuit_data[i]['dIdt'], bounds_error=False, fill_value=(self.circuit_data[i]['dIdt'][0], self.circuit_data[i]['dIdt'][-1]))
            
        self.debug_graphs.plot_ramp_rate_interpolation()

    def create_time_step_interpolation(self):
        """
        Create interpolation function for adaptive timestep guidance.

        Generates a densified timestep schedule from the unique circuit ramp
        times, creating a smooth interpolation function (time_step_dt) that
        suggests appropriate ODE solver timesteps. Also creates an event
        function for detecting critical time points.

        Side Effects:
            - Sets self.unique_sorted_timesteps (all unique circuit times)
            - Creates self.event (event detection function)
            - Creates self.time_step_dt (timestep interpolation function)
            - Calls debug_graphs.plot_time_step_interpolation() for visualization
        """

        time_steps = [self.circuit_data[i]['t'] for i in self.unique_circuits]
        self.unique_sorted_timesteps = np.sort(list(set(np.concatenate(time_steps))))

        # Creating an event function for the unique sorted timesteps:
        self.event = utils.time_event_factory(self.unique_sorted_timesteps)

        # Densifying the timesteps.
        t_densified = utils.densify_with_endpoint_gaussian(self.unique_sorted_timesteps, total_points=500)
        t_diff = np.diff(t_densified, prepend=[0.01])

        # Making a call-able fit function.
        self.time_step_dt = interp1d(t_densified, t_diff, bounds_error=False, fill_value=(0.01, self.max_timestep))
        self.debug_graphs.plot_time_step_interpolation(t_densified, t_diff)

    def PID(self, Power, i, dt=0):
        """
        Compute PID-controlled ramp rate for power-limited ramping.

        Implements a PI controller (proportional + integral, derivative term
        commented out) to adjust the circuit ramp rate based on the error
        between actual dissipated power and the target setpoint. This enables
        power-limited ramping to avoid excessive heating.

        Args:
            Power: Current total dissipated power [W]
            i: Circuit index
            dt: Timestep for integral calculation [s], default 0

        Returns:
            float: Controlled ramp rate [A/s], clamped to max_rr[i]

        Side Effects:
            - Updates self.integral[i] with integrated error
        """
        # Power is current power.

        # PID calculations
        e = self.setpoint[i] - Power

        P = self.Kp*e
        self.integral[i] += self.Ki*e*dt

        # Testing showed the derivative step in not important in almost all situations.
        # D = self.Kd*(e - e_prev)/dt

        # calculate manipulated variable - MV
        rr = self.offset[i] + P + self.integral[i] # + D

        # update stored data for next iteration
        # e_prev = e

        if rr > self.max_rr[i]: rr = self.max_rr[i]
        return rr

    def HeliumI(self, T):
        """ Function to calculate the heat transfer coefficient between the outer layer and liquid He-I.

        Args:
            T: Temperature of the outermost layer:

        Output: Heat transfer coefficient

        """

        THe = self.Thelium  # Setting the helium temperature
        T = np.where(T < THe, THe, T)  # Assuming nothing can be colder than the temperature of the helium.

        h1 = self.c1 * (T - THe) * (1 - (1 / (1 + np.exp(-(T - self.T_star_1) * 100))))
        h2 = self.c2 * (T - THe) ** 2.5 * (1 / (1 + np.exp(-(T - self.T_star_1) * 100))) * (
                    1 - (1 / (1 + np.exp(-(T - self.T_star_2) * 100))))
        h3 = self.c3 * (1 / (1 + np.exp(-(T - self.T_star_2) * 100)))

        # Transient cooling regime is ignored, see:
        # Source: CUDI: A MODEL FOR CALCULATION OF ELECTRODYNAMIC AND THERMAL BEHAVIOUR OF SUPERCONDUCTING RUTHERFORD CABLES
        # A. P. Verweij

        return h1 + h2 + h3

    def HeliumII(self, T):
        """ Function to calculate the heat transfer coefficient between the outer layer and liquid He-II.

        Args:
            T: Temperature of the outermost layer:

        Output: Heat transfer coefficient

        """
        c5 = self.c5
        c6 = self.c6
        p = self.p
        Thelium = self.Thelium
        Tstar = self.Tstar

        T = np.where((T < (self.Thelium + 1e-3)), self.Thelium + 1e-3, T)
        h_out = ((c5 * (T ** p - Thelium ** p)) * (1 - utils.sig(T - Thelium - Tstar, 100)) + (
                    c6 * (T - Thelium)) * utils.sig(T - Thelium - Tstar, 100)) / (T - Thelium)
        # h_out = ((c5*(T**p-Thelium**p)) * (1-(1/(1+np.exp(-((T-Thelium)-(Tstar))*100)))) + (c6*(T-Thelium)) * (1/(1+np.exp(-((T-Thelium)-(Tstar))*100)  )))/(T-Thelium)
        return np.array(h_out)

    def power_cryocooler(self, T, M):
        """
        Calculate distributed cryocooler cooling power.

        Computes the cooling power provided by a cryocooler, distributing it
        across elements proportionally to their temperature-dependent cooling
        capacity while maintaining the total average cooling power.

        Args:
            T: Array of temperatures [K]
            M: Array of masses [kg]

        Returns:
            Array of cooling powers [W] per element
        """
        CPower = self.cryocooler_fit(T)*M
        Cpower_avg = self.cryocooler_fit(np.mean(T))
        return CPower / np.sum(CPower) * Cpower_avg

    def results_to_vtu(self):
        """
        Convert simulation results to VTU/VTP files for ParaView visualization.

        Creates VTK Unstructured Grid (VTU) files and polydata (VTP) files
        from the simulation results, enabling 3D visualization in ParaView.
        Generates a limited number of output files (vtu_num_elements) evenly
        spaced in time, and creates a PVD collection file for animation.

        Side Effects:
            - Creates 'vtu' subdirectory in sim_output
            - Writes multiple .vtu and .vtp files
            - Writes .pvd collection file for ParaView
            - Only executes if vtu_info is not None and vtu_post_BOOL is True
        """

        if self.vtu_info is not None and self.vtu_post_BOOL:
            logging.info('Post-processing: Starting the .vtu creation.')
            # Creating the output folder:
            vtu_folder = Path(self.sim_output) / 'vtu'
            try:
                os.mkdir(vtu_folder)
            except OSError:
                logging.info("Creation of the directory %s failed" % vtu_folder)
            else:
                logging.debug("Successfully created the directory %s " % vtu_folder)

            # Loading the field grid files.
            try:
                with open(self.Bz_full_path, 'rb') as fin:
                    Bz_full = pickle.load(fin)
            except OSError:
                logging.info("Loading of %s has failed." % (self.Bz_full_path))
            else:
                logging.debug("Successfully loaded %s" % (self.Bz_full_path))

            try:
                with open(self.Br_full_path, 'rb') as fin:
                    Br_full = pickle.load(fin)
            except OSError:
                logging.info("Loading of %s has failed." % (self.Br_full_path))
            else:
                logging.debug("Successfully loaded %s" % (self.Br_full_path))

            # Selecting timesteps
            if len(self.t_array) > self.vtu_num_elements:
                # Making a selection of timesteps.
                idx = np.round(np.linspace(0, len(self.t_array) - 1, self.vtu_num_elements)).astype(int)
            else:
                # Selecting all timesteps if the len of the arrays is smaller than the amount of timesteps specified.
                self.vtu_num_elements = len(self.t_array)
                idx = list(range(self.vtu_num_elements))

            # Initializing some variables
            seed_points = self.seed_points
            filenames = []
            field_filenames = []
            timesteps = []

            # Starting the progressbar:
            prog_bar = ProgressBar(self.vtu_num_elements, 'Vtu\'s created  :', interval=1)
            processed = 0

            for i in range(len(idx)):
                prog_bar.update(processed)
                filename = "vtu_" + str(i) + ".vtu"
                field_filename = "field_" + str(i) + ".vtp"

                data_fields = {
                    'Current': self.I_array[idx[i]],
                    'Current Density': self.I_density_array[idx[i]],
                    'Critial Ratio': self.Critical_Ratio_array[idx[i]][self.Index_inductive_connections],
                    'Current Sharing x': self.Current_Sharing_x_array[idx[i]][self.Index_inductive_connections],
                    'B_tot': self.B_array[idx[i]],
                    'B_z': self.Bz_array[idx[i]],
                    'B_r': self.Br_array[idx[i]],
                    'Field Angle': self.Field_Angle_array[idx[i]],
                    'Fz': self.Fz_array[idx[i]],
                    'Fr': self.Fr_array[idx[i]],
                    'Fz Density': self.Fz_density_array[idx[i]],
                    'Fr Density': self.Fr_density_array[idx[i]],
                    'Azim. Power Dissipation': self.SC_loss_array[idx[i]],
                    'Rad. Power Dissipation': self.Radial_loss_array[idx[i]],
                    'Temperature': np.array(self.T_array[idx[i]])[self.Index_temperatures][self.Index_inductive_connections],
                    'Voltage': np.array(self.V_array[idx[i]])[self.Index_temperatures][self.Index_inductive_connections],
                    'Radial Currents': np.array(self.IR_array[idx[i]])[self.Index_temperatures][self.Index_inductive_connections],
                    'Radial Resistance': np.array(self.R_Radial_array[idx[i]])[self.Index_temperatures][self.Index_inductive_connections],
                    'Conductor Resistance': self.R_Conductor_array[idx[i]],
                    'Element ID': self.Index_inductive_connections,
                    'Node ID': self.Index_temperatures[self.Index_inductive_connections],
                }

                args = {
                    'r_centers': self.vtu_info['inductive']['r_centers'],
                    'z_centers': self.vtu_info['inductive']['z_centers'],
                    'dr': self.vtu_info['inductive']['dr'],
                    'dz': self.vtu_info['inductive']['dz'],
                    'data_fields': data_fields,
                    'n_theta': 72,
                    'filename': vtu_folder / filename
                }

                # Computing the fields on grid.
                Bz = np.dot(Bz_full, self.I_array[idx[i]])
                Br = np.dot(Br_full, self.I_array[idx[i]])

                args_vector = {
                    'r': np.linspace(self.grid_information[0], self.grid_information[1], self.grid_information[4]),
                    'z': np.linspace(self.grid_information[2], self.grid_information[3], self.grid_information[5]),
                    'vec_r': Br,
                    'vec_z': Bz,
                    'n_theta': 11,
                    'filename': vtu_folder / field_filename,
                    'seed_points': seed_points,
                }

                # Saving the data to the vtu.
                data_to_vtu(**args)
                seed_points = vector_to_vtu(**args_vector)

                # Gathering data for the time series pvd file.
                filenames += [filename]
                field_filenames += [field_filename]
                timesteps += [self.t_array[idx[i]]]
                processed += 1

            if len(filenames) > 0:
                # Writing the time series to file.
                write_pvd_file(filenames, timesteps, output_path=vtu_folder / "ramp_series.pvd")

            if len(field_filenames) > 0:
                # Writing the time series to file.
                write_pvd_file(field_filenames, timesteps, output_path=vtu_folder / "field_series.pvd")

    def results_post_process(self):
        """
        Perform optional post-processing for external visualization tools.

        If post_process configuration is provided, this method exports
        simulation results to formats compatible with ParaView and/or COMSOL.
        It converts the homogenized 2D results to 3D cylindrical geometry
        at specified angular positions.

        Side Effects:
            - Calls Results_To_Paraview module if configured
            - May create additional output files for COMSOL/ParaView
            - Validates post_process dictionary structure
            - Only executes if self.post_process is not empty
        """
        if bool(self.post_process):
            from Results_To_Paraview import to_paraview
            logging.info('The post_process started.')
            # Checks if the key types in the post_process dictionary are okay.
            types_check, error = utils.check_dict_keys_types(self.post_process, {'parameters': list, 'angle': (int, float), 'paraview': bool, 'comsol': bool, 'offset_removal': bool, 'npoints': int})

            # If all are okay, post-processing starts.
            if types_check:
                for i in self.post_process['parameters']:
                    logging.info('Post processing %s.' % i)
                    inputs = {
                        'results_path': self.sim_output,
                        'result': i,
                        'angle': self.post_process['angle'],
                        'comsol': self.post_process['comsol'],
                        'paraview': self.post_process['paraview'],
                        'offset_removal': self.post_process['offset_removal'],
                        'npoints': self.post_process['npoints'],
                    }

                    # Results to paraview function
                    r2p = to_paraview(**inputs)
            else:
                logging.info('Post processing type check failed, no post-processing is performed.')
        else:
            logging.info('The post_process dictionary is empty, no post-processing is performed on this dataset.')

    def create_csv(self, csv_path, csv_columns):
        """
        Create a new CSV file with column headers.

        Args:
            csv_path: Full path to the CSV file to create
            csv_columns: Key in self.columns dict specifying column headers

        Side Effects:
            - Creates new CSV file at csv_path with headers
            - Logs success/failure message
        """
        # Creates the csv files and adds the columns.
        try:
            with open(csv_path, 'w+') as f:
                f.write(self.delimiter.join(self.columns[csv_columns]) + '\n')
        except OSError:
            logging.info("Creation of %s has failed." % csv_path)
        else:
            logging.debug("Successfully created %s" % csv_path)

    def append_to_csv(self, csv_path, data_array, saverange, output_format = "{:.5E}"):
        """
        Append data rows to an existing CSV file using vectorized operations.

        Args:
            csv_path: Full path to the CSV file
            data_array: List of data arrays to append
            saverange: Range of indices to write from data_array
            output_format: Format string for numeric output, default "{:.5E}"

        Side Effects:
            - Appends rows to csv_path file
            - Logs error if write fails

        Performance:
            Vectorized implementation using numpy.savetxt is ~10-50x faster
            than the previous row-by-row string formatting approach.
            Eliminates millions of calls to utils.output_format().
        """
        try:
            if len(saverange) == 0:
                return

            # Build 2D array: rows are timesteps, columns are data elements
            # Each data_array[j] is already an array, so stack them vertically
            data_2d = np.array([data_array[j] for j in saverange])

            # Create time column with offset
            time_col = (np.array([self.t_array[j] for j in saverange]) +
                       self.t_offset).reshape(-1, 1)

            # Combine time column with data (time is first column)
            full_data = np.hstack([time_col, data_2d])

            # Convert format string to numpy fmt string
            # "{:.5E}" -> "%.5E", "{:0.2f}" -> "%.2f"
            if output_format == "{:.5E}":
                numpy_fmt = "%.5E"
            elif output_format == "{:0.2f}":
                numpy_fmt = "%.2f"
            else:
                # Parse generic format string: {:.XE} or {:X.Yf}
                # This handles most common cases
                import re
                match = re.search(r'\{:([0-9.]*)([Eef])\}', output_format)
                if match:
                    numpy_fmt = f"%{match.group(1)}{match.group(2).upper()}"
                else:
                    # Fallback: use a reasonable default
                    numpy_fmt = "%.5E"

            # Use numpy.savetxt for vectorized writing
            # This is much faster than manual string formatting and joining
            with open(csv_path, 'ab') as f:  # 'ab' for append in binary mode
                np.savetxt(
                    f,
                    full_data,
                    delimiter=self.delimiter,
                    fmt=numpy_fmt,
                    newline='\n'
                )

        except OSError:
            logging.info("Writing on %s has failed." % csv_path)

    @staticmethod
    def powerlaw(x, IIc, U, n=21.0, E0=1e-4):
        """
        Superconductor power law.
        
        OPTIMIZED: Vectorized computation using torch operations for better performance
        with both numpy arrays and torch tensors.
        
        Args:
            x: Current sharing factor (0-1)
            IIc: Current ratio I/Ic
            U: Voltage parameter
            n: n-value (default 21.0)
            E0: Electric field constant (default 1e-4)
            
        Returns:
            Power law result: E0 * (IIc)^n * x^n - U * (1 - x)
        """
        # Use torch operations for vectorized computation - works with both numpy and torch
        if hasattr(x, 'device'):  # torch tensor
            IIc_pow_n = torch.pow(IIc, n)
            x_pow_n = torch.pow(x, n)
            one_minus_x = 1.0 - x
        else:  # numpy array
            import numpy as np
            IIc_pow_n = np.power(IIc, n)
            x_pow_n = np.power(x, n)
            one_minus_x = 1.0 - x
            
        return E0 * IIc_pow_n * x_pow_n - U * one_minus_x

    @staticmethod
    def powerlaw_jac(x, IIc, U, n=21.0, E0=1e-4):
        """
        Jacobian of superconductor power law.
        
        OPTIMIZED: Vectorized computation and mathematical simplification for better performance.
        Simplified from: n * E0 * IIc * (IIc * x)^(n-1) + U
        To: n * E0 * IIc^n * x^(n-1) + U
        
        Args:
            x: Current sharing factor (0-1)
            IIc: Current ratio I/Ic  
            U: Voltage parameter
            n: n-value (default 21.0)
            E0: Electric field constant (default 1e-4)
            
        Returns:
            Jacobian: n * E0 * IIc^n * x^(n-1) + U
        """
        # Use torch/numpy operations for vectorized computation
        if hasattr(x, 'device'):  # torch tensor
            IIc_pow_n = torch.pow(IIc, n)
            x_pow_n_minus_1 = torch.pow(x, n - 1.0)
        else:  # numpy array
            import numpy as np
            IIc_pow_n = np.power(IIc, n)
            x_pow_n_minus_1 = np.power(x, n - 1.0)
            
        return n * E0 * IIc_pow_n * x_pow_n_minus_1 + U

if __name__ == "__main__":
    import cProfile, pstats
    import json
    profiler = cProfile.Profile()

    if len(sys.argv) == 2:
        # Normal use of the solver.
        NI_sim = NI_Coil_Solver(sys.argv[1])
    elif len(sys.argv) == 3:
        # Called with a dictionary of paths, used for the tests.
        # Load the path_dict from JSON file
        with open(sys.argv[2], 'r') as f:
            path_dict = json.load(f)
        NI_sim = NI_Coil_Solver(sys.argv[1], path_dict)
