import numpy as np
from scipy.interpolate import RectBivariateSpline, CubicSpline
from matplotlib.tri import LinearTriInterpolator
from .plane_data import PlaneData
from .catalog import has_static_product, read_static_variables
[docs]
class MagneticField:
def __init__(self, plane_instance, data_dir=".", catalog=None, source_reader=None):
"""
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.
data_dir: retained for compatibility with existing call sites.
catalog: SimulationCatalog used for static product reads. Direct
local BP fallbacks are disabled.
source_reader: optional backend for catalog static reads.
"""
# ---------------------------
# Read equilibrium data from xgc.equil.bp
# ---------------------------
equil_vars = (
"bp_sign",
"bt_sign",
"eq_I",
"eq_axis_b",
"eq_axis_r",
"eq_axis_z",
"eq_max_r",
"eq_max_z",
"eq_min_r",
"eq_min_z",
"eq_mpsi",
"eq_mr",
"eq_mz",
"eq_psi_grid",
"eq_psi_rz",
"eq_x_psi",
"eq_x_r",
"eq_x_z",
)
if has_static_product(catalog, "xgc.equil.bp", equil_vars):
equil_values = read_static_variables(
catalog,
"xgc.equil.bp",
equil_vars,
source_reader=source_reader,
)
else:
raise RuntimeError(
"MagneticField requires catalog product 'xgc.equil.bp' with equilibrium variables; "
"direct local BP fallback is disabled."
)
self.bp_sign = float(np.asarray(equil_values["bp_sign"]).item())
self.bt_sign = float(np.asarray(equil_values["bt_sign"]).item())
self.eq_I_array = np.asarray(equil_values["eq_I"]) # shape: (129,)
self.eq_axis_b = float(np.asarray(equil_values["eq_axis_b"]).item())
self.eq_axis_r = float(np.asarray(equil_values["eq_axis_r"]).item())
self.eq_axis_z = float(np.asarray(equil_values["eq_axis_z"]).item())
self.eq_max_r = float(np.asarray(equil_values["eq_max_r"]).item())
self.eq_max_z = float(np.asarray(equil_values["eq_max_z"]).item())
self.eq_min_r = float(np.asarray(equil_values["eq_min_r"]).item())
self.eq_min_z = float(np.asarray(equil_values["eq_min_z"]).item())
self.eq_mpsi = int(np.asarray(equil_values["eq_mpsi"]).item())
self.eq_mr = int(np.asarray(equil_values["eq_mr"]).item())
self.eq_mz = int(np.asarray(equil_values["eq_mz"]).item())
self.eq_psi_grid = np.asarray(equil_values["eq_psi_grid"]) # length: (129,)
# Skip eq_psi_norm.
self.eq_psi_rz = np.asarray(equil_values["eq_psi_rz"]) # shape: (129, 129)
self.eq_x_psi = float(np.asarray(equil_values["eq_x_psi"]).item())
self.eq_x_r = float(np.asarray(equil_values["eq_x_r"]).item())
self.eq_x_z = float(np.asarray(equil_values["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
# ---------------------------
bfield_vars = ("bfield", "psi")
if has_static_product(catalog, "xgc.bfield.bp", bfield_vars):
bfield_values = read_static_variables(
catalog,
"xgc.bfield.bp",
("bfield", "n_n", "psi"),
source_reader=source_reader,
missing="skip",
)
bfield_array = np.asarray(bfield_values["bfield"]) # shape: (3, n_n)
n_n_file = int(np.asarray(bfield_values["n_n"]).item()) if "n_n" in bfield_values else plane_instance.n_n
psi_array = np.asarray(bfield_values["psi"]) # shape: (n_n,)
else:
raise RuntimeError(
"MagneticField requires catalog product 'xgc.bfield.bp' with 'bfield' and 'psi'; "
"direct local BP/Stream fallback is disabled."
)
# 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 background data may live in xgc.current_drive.bp (newer) or
# xgc.bfield.bp (older). If unavailable in both, default to zeros.
# ---------------------------------------------------------------
jpar_bg, jpar_bg_fs_avg = self._load_jpar_background(
data_dir=data_dir,
n_n=plane_instance.n_n,
nsurf=plane_instance.nsurf,
catalog=catalog,
source_reader=source_reader,
)
# 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
self.jpar_bg_fs_avg = jpar_bg_fs_avg
# ---------------------------
# 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)
def _try_read_jpar_catalog(self, catalog, product_key, source_reader=None):
"""
Try reading jpar background variables from a catalog product.
Parameters
----------
catalog : xgc_analysis.catalog.SimulationCatalog or None
Catalog to inspect.
product_key : str
Candidate product key.
source_reader : callable or None, optional
Optional backend for catalog static reads.
"""
if not has_static_product(catalog, product_key, ("jpar_bg", "jpar_bg_fs_avg")):
return None
try:
values = read_static_variables(
catalog,
product_key,
("jpar_bg", "jpar_bg_fs_avg"),
source_reader=source_reader,
)
return np.asarray(values["jpar_bg"]), np.asarray(values["jpar_bg_fs_avg"])
except Exception:
return None
def _load_jpar_background(self, data_dir, n_n, nsurf, catalog=None, source_reader=None):
"""
Load ``jpar_bg`` and ``jpar_bg_fs_avg`` with catalog lookup order:
1) xgc.current_drive.bp
2) xgc.bfield.bp
3) zeros if not found in catalog metadata
"""
for fname in ("xgc.current_drive.bp", "xgc.bfield.bp"):
result = self._try_read_jpar_catalog(catalog, fname, source_reader=source_reader)
if result is None:
continue
jpar_bg, jpar_bg_fs_avg = result
# Basic shape validation / normalization
jpar_bg = np.squeeze(jpar_bg)
jpar_bg_fs_avg = np.squeeze(jpar_bg_fs_avg)
if jpar_bg.shape != (n_n,):
# If shape is incompatible, ignore and continue fallback chain.
continue
if jpar_bg_fs_avg.ndim != 1:
continue
return jpar_bg, jpar_bg_fs_avg
# Default to zeros if unavailable.
return np.zeros(n_n, dtype=float), np.zeros(nsurf, dtype=float)
[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