import numpy as np
from adios2 import FileReader
from scipy.interpolate import RectBivariateSpline, CubicSpline
from matplotlib.tri import LinearTriInterpolator
from xgc_analysis.plane_data import PlaneData
from xgc_analysis.plane import Plane
[docs]
class MagneticField:
def __init__(self, plane_instance, data_dir="."):
"""
Initialize the MagneticField from two BP files:
- xgc.equil.bp: contains equilibrium quantities.
- xgc.bfield.bp: contains the magnetic field data on the planar mesh.
Args:
plane_instance: an instance of the Plane class (from plane.py). It must have:
- n_n: the number of vertices on the plane.
- get_triangulation_data(n_int): a method returning an object with attribute 'triObj'.
- Attributes 'cnv_to_surf' and 'cnv_from_surf' that are used by the flux‐surface averaging
functions in PlaneData.
"""
# ---------------------------
# Read equilibrium data from xgc.equil.bp
# ---------------------------
with FileReader(data_dir+"/xgc.equil.bp") as reader:
self.bp_sign = float(reader.read("bp_sign").item())
self.bt_sign = float(reader.read("bt_sign").item())
self.eq_I_array = reader.read("eq_I") # shape: (129,)
self.eq_axis_b = float(reader.read("eq_axis_b").item())
self.eq_axis_r = float(reader.read("eq_axis_r").item())
self.eq_axis_z = float(reader.read("eq_axis_z").item())
self.eq_max_r = float(reader.read("eq_max_r").item())
self.eq_max_z = float(reader.read("eq_max_z").item())
self.eq_min_r = float(reader.read("eq_min_r").item())
self.eq_min_z = float(reader.read("eq_min_z").item())
self.eq_mpsi = int(reader.read("eq_mpsi").item())
self.eq_mr = int(reader.read("eq_mr").item())
self.eq_mz = int(reader.read("eq_mz").item())
self.eq_psi_grid = reader.read("eq_psi_grid") # length: (129,)
# Skip eq_psi_norm.
self.eq_psi_rz = reader.read("eq_psi_rz") # shape: (129, 129)
self.eq_x_psi = float(reader.read("eq_x_psi").item())
self.eq_x_r = float(reader.read("eq_x_r").item())
self.eq_x_z = float(reader.read("eq_x_z").item())
# Set up a bicubic spline interpolator for eq_psi_rz.
# Create uniform grids in R and Z based on eq_min_r, eq_max_r (and similarly for Z) using eq_mr, eq_mz.
R_grid = np.linspace(self.eq_min_r, self.eq_max_r, self.eq_mr)
Z_grid = np.linspace(self.eq_min_z, self.eq_max_z, self.eq_mz)
# Note that we have to transpose eq_psi_rz here. It's inner dimension (axis=1 in C-ordering) is R
# and the outer dimension is Z. But if we want the first argument to the spline interpolation
# to be R and the second Z coordinates, the order must be opposite, hence the transpose op.
self.eq_psi_rz_spline = RectBivariateSpline(R_grid, Z_grid, self.eq_psi_rz.T, kx=3, ky=3)
# Set up a cubic spline interpolator for eq_I.
self.eq_I_spline = CubicSpline(self.eq_psi_grid, self.eq_I_array)
# ---------------------------
# Read magnetic field data from xgc.bfield.bp
# ---------------------------
with FileReader(data_dir+"/xgc.bfield.bp") as reader:
bfield_array = reader.read("bfield") # shape: (3, n_n)
jpar_bg = reader.read("jpar_bg") # shape: (n_n,)
self.jpar_bg_fs_avg = reader.read("jpar_bg_fs_avg") # shape: (149,)
n_n_file = int(reader.read("n_n").item())
psi_array = reader.read("psi") # shape: (n_n,)
# Check that the number of vertices in the file matches the Plane instance.
if n_n_file != plane_instance.n_n:
raise ValueError("Mismatch in number of vertices: n_n in xgc.bfield.bp does not equal plane_instance.n_n")
# Use the PlaneData class for bfield and psi.
# bfield is a vector field with 3 components. The file stores bfield as shape (3, n_n);
# we store it as (n_n, 3).
self.bfield = PlaneData(plane_instance, n_components=3, data_array=bfield_array.T)
# psi is a scalar field.
self.psi_pd = PlaneData(plane_instance, n_components=1, data_array=psi_array)
# jpar_bg is also a scalar field.
self.jpar_bg_pd = PlaneData(plane_instance, n_components=1, data_array=jpar_bg)
self.jpar_bg_pd.data = jpar_bg
# ---------------------------
# Set up a triangle interpolator for ψ on the planar mesh.
# ---------------------------
tri_data = plane_instance.get_triangulation_data(n_int=400) # n_int can be adjusted as needed.
self.psi_tri_interp = LinearTriInterpolator(tri_data.triObj, self.psi_pd.data)
[docs]
def compute_background_field(self, R, Z):
"""
Return the magnetic-field components (B_R, B_Z, B_t) at the given
cylindrical coordinates.
Parameters
----------
R, Z : float or ndarray
Cylindrical coordinates. Scalars and arbitrary-shaped NumPy
arrays are both accepted; broadcasting **is** supported as long
as R and Z have compatible shapes.
Returns
-------
B_R, B_Z, B_t : float or ndarray
Radial, vertical (poloidal) and toroidal components. Their
type/shape matches the broadcasted shape of *R* and *Z*.
"""
import numpy as np
# --- convert inputs to arrays without copying when possible -----
R_arr = np.asanyarray(R)
Z_arr = np.asanyarray(Z)
# --- equilibrium poloidal flux and its derivatives --------------
psi_val = self.eq_psi_rz_spline.ev(R_arr, Z_arr) # ψ(R,Z)
dpsi_dR = self.eq_psi_rz_spline.ev(R_arr, Z_arr, dx=1, dy=0) # ∂ψ/∂R
dpsi_dZ = self.eq_psi_rz_spline.ev(R_arr, Z_arr, dx=0, dy=1) # ∂ψ/∂Z
# --- cylindrical-component formulas -----------------------------
invR = 1.0 / R_arr
B_R = -self.bp_sign * invR * dpsi_dZ
B_Z = self.bp_sign * invR * dpsi_dR
# --- toroidal field from I(ψ) spline ----------------------------
I_val = self.eq_I_spline(psi_val) # I(ψ)
B_t = self.bt_sign * I_val * invR
# --- if the caller passed scalars, unwrap to Python float -------
if np.isscalar(R) and np.isscalar(Z):
return float(B_R), float(B_Z), float(B_t)
return B_R, B_Z, B_t