Source code for torx.plotting.plot_magnetic_well_m

"""Plotting routines for 3D magnetic field tracing plots."""
import numpy as np



def plot_3D_field_lines(r_norm, z_norm, phi_array, ax, plot_kwargs={}, axis_set_kwargs={}):
    """Plot traced field-line data on a 3D plot."""
    assert ax.name == "3d"
    # Convert cylindrical data to cartesian
    x = r_norm * np.cos(phi_array)        
    y = r_norm * np.sin(phi_array)         
    z = z_norm
    
    # Plot each field line separately so they are different colors
    for l in range(r_norm.sizes["line"]):
        ax.plot(x.isel({"line":l}),
                y.isel({"line":l}),
                z.isel({"line":l}),
                **plot_kwargs)
        
    ax.set(**axis_set_kwargs)
    
    return ax




def plot_magnetic_well_data(field_line_length, absB, ax, plot_kwargs={}, axis_set_kwargs={}):
    """Plot absolute magnetic field versus the field-line length."""
    # Transpose the data so each field-line is automatically its own color
    ax.plot(field_line_length.T, absB.T, **plot_kwargs)
    
    ax.set(**axis_set_kwargs)
    
    return ax