from enum import Enum
import numpy as np
from sandalwood import MultivariateTaylorFunction
from .vector_fields import VectorField
[docs]
class Backend(str, Enum):
PYTHON = "python"
MPI = "mpi"
COSY = "cosy"
MPI_COSY = "mpi_cosy"
[docs]
def calculate_b_field(coil_instance, field_points, backend=Backend.PYTHON):
"""
Calculates the magnetic field generated by a coil at a set of
field points using the Biot-Savart law.
Args:
coil_instance (Coil): An instance of a Coil subclass.
field_points (np.ndarray or np.ndarray of
mtf.MultivariateTaylorFunction):
The points (x, y, z) where the magnetic field should be
calculated. Can be a (M, 3) NumPy array of numbers or MTF objects.
backend (Backend or str, optional): The backend to use for the calculation.
Defaults to Backend.PYTHON.
Returns:
VectorField: A VectorField object containing the field points and the
calculated FieldVector objects.
"""
# Normalize backend to Enum
if isinstance(backend, str):
try:
backend = Backend(backend)
except ValueError:
raise ValueError(f"Backend '{backend}' is not a valid option.")
# Input validation for field_points
if not isinstance(field_points, np.ndarray) or (
field_points.ndim == 2 and field_points.shape[1] != 3
):
raise TypeError("field_points must be a NumPy array of shape (N, 3).")
# Check if segments have been generated
if coil_instance.segment_centers is None:
raise RuntimeError("Coil segments have not been generated.")
if coil_instance.use_mtf_for_segments:
element_centers_np = np.array([
[v.x, v.y, v.z] for v in coil_instance.segment_centers
])
element_directions_np = np.array([
[v.x, v.y, v.z] for v in coil_instance.segment_directions
])
else:
element_centers_np = np.array([
c.to_numpy_array() for c in coil_instance.segment_centers
])
element_directions_np = np.array([
d.to_numpy_array() for d in coil_instance.segment_directions
])
if backend == Backend.MPI:
# Tuple unpacking for MPI result
bx, by, bz = mpi_biot_savart(
element_centers=element_centers_np,
element_lengths=coil_instance.segment_lengths,
element_directions=element_directions_np,
field_points=field_points,
backend=Backend.PYTHON,
)
elif backend == Backend.MPI_COSY:
bx, by, bz = mpi_biot_savart(
element_centers=element_centers_np,
element_lengths=coil_instance.segment_lengths,
element_directions=element_directions_np,
field_points=field_points,
backend=Backend.COSY,
)
elif backend in (Backend.PYTHON, Backend.COSY):
bx, by, bz = serial_biot_savart(
element_centers=element_centers_np,
element_lengths=coil_instance.segment_lengths,
element_directions=element_directions_np,
field_points=field_points,
backend=backend,
)
else:
raise ValueError(f"Unknown backend: {backend}")
# Optimized Vector Field Creation (SoA enforced)
if coil_instance.use_mtf_for_segments:
# Optimization: If current is a constant MTF, extract value to avoid
# overhead and potential COSY type errors with empty/zero MTFs.
current_scalar = coil_instance.current
if isinstance(current_scalar, MultivariateTaylorFunction):
try:
# Check if it has only constant term (0 exponents)
# This is a heuristic; valid for simple constants
if current_scalar.exponents.shape[0] == 1 and np.all(
current_scalar.exponents == 0
):
current_scalar = current_scalar.extract_coefficient(
tuple([0] * current_scalar.dimension)
)
except Exception:
pass
# Apply scalar current
bx = bx * current_scalar
by = by * current_scalar
bz = bz * current_scalar
# Apply integration vectorized/map
# Helper to integrate array
def integrate_arr(arr):
# Try vectorized if possible, else list comp (array of objects)
return np.array([v.integrate(4, -1, 1) for v in arr], dtype=object)
bx = integrate_arr(bx)
by = integrate_arr(by)
bz = integrate_arr(bz)
else:
# Fast Path: bx, by, bz are numerical arrays (float)
# Apply scaling and factor 2 (approximation for discrete segments)
scale_factor = coil_instance.current * 2
bx = bx * scale_factor
by = by * scale_factor
bz = bz * scale_factor
# Return using Structure of Arrays (tuple)
return VectorField(vectors=(bx, by, bz), field_points=field_points)
mu_0_4pi = 1e-7 # Define mu_0_4pi if it's not already globally defined
def _python_biot_savart_core(source_points, dl_vectors, field_points, order=None):
"""
Core vectorized Biot-Savart calculation in pure Python.
Returns (Bx, By, Bz) tuple (SoA).
"""
source_points_reshaped = source_points[:, np.newaxis, :]
field_points_reshaped = field_points[np.newaxis, :, :]
r_vectors = field_points_reshaped - source_points_reshaped
r_squared = np.sum(r_vectors**2, axis=2)
# Avoid division by zero at the source point location
# Optimization: Use vectorized masking for numeric arrays
if r_squared.size > 0:
is_mtf = isinstance(r_squared.flat[0], MultivariateTaylorFunction)
if is_mtf:
# Fallback to loops for MTF objects
# Note: With SoA we still have this loop for r_squared check/fix
# unless we implement vectorized ops on MTF arrays.
for i in range(r_squared.shape[0]):
for j in range(r_squared.shape[1]):
val = r_squared[i, j]
const_term = val.get_constant()
if abs(const_term) < 1e-18:
r_squared[i, j] += 1e-18
else:
# Vectorized masking for numeric arrays
mask = np.abs(r_squared) < 1e-18
r_squared[mask] = 1e-18
dl_vectors_reshaped = dl_vectors[:, np.newaxis, :]
cross_products = np.cross(dl_vectors_reshaped, r_vectors, axis=2)
# Calculate 1/r^3 for magnitude scaling.
inv_r_cubed = np.reciprocal(r_squared) * r_squared ** (-0.5)
inv_r_cubed_expanded = np.expand_dims(inv_r_cubed, axis=2)
dB_contributions = (mu_0_4pi * cross_products) * inv_r_cubed_expanded
# Summing for B-field (SoA)
# dB_contributions shape is (N_source, M_field, 3)
Bx = np.sum(dB_contributions[:, :, 0], axis=0)
By = np.sum(dB_contributions[:, :, 1], axis=0)
Bz = np.sum(dB_contributions[:, :, 2], axis=0)
# If an order is specified, truncate
if order is not None:
def truncate_arr(arr):
if arr.size > 0 and isinstance(arr.flat[0], MultivariateTaylorFunction):
return np.array([v.truncate(order) for v in arr.flat]).reshape(
arr.shape
)
return arr
Bx = truncate_arr(Bx)
By = truncate_arr(By)
Bz = truncate_arr(Bz)
return Bx, By, Bz
[docs]
def numpy_biot_savart(
element_centers, element_lengths, element_directions, field_points, order=None
):
"""
NumPy vectorized Biot-Savart calculation.
Wrapper for serial_biot_savart, maintaining return types.
"""
bx, by, bz = serial_biot_savart(
element_centers, element_lengths, element_directions, field_points, order=order
)
# Reconstruct (N, 3) for backward compatibility if needed by external callers?
# docstring says it returns (M, 3).
# We should probably fix callers or reconstruct here.
# The docstring says it returns ndarray (M, 3). So we must stack.
return np.column_stack((bx, by, bz))
[docs]
def mpi_biot_savart(
element_centers,
element_lengths,
element_directions,
field_points,
order=None,
backend=Backend.PYTHON,
):
"""
Parallel Biot-Savart calculation using mpi4py.
Returns tuple (Bx, By, Bz).
"""
try:
from mpi4py import MPI
except ImportError:
raise ImportError("mpi4py is not installed, cannot run in MPI mode.")
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
element_centers = np.array(element_centers)
element_lengths = np.array(element_lengths)
element_directions = np.array(element_directions)
field_points = np.array(field_points)
num_field_points = field_points.shape[0]
chunk_size = num_field_points // size
remainder = num_field_points % size
start_index = rank * chunk_size + min(rank, remainder)
end_index = start_index + chunk_size + (1 if rank < remainder else 0)
local_field_points = field_points[start_index:end_index]
# Returns tuple (bx, by, bz)
local_bx, local_by, local_bz = serial_biot_savart(
element_centers,
element_lengths,
element_directions,
local_field_points,
order=order,
backend=backend,
)
# Gather tuples? list of tuples [(bx0, by0, bz0), (bx1, ...)]
all_chunks = comm.gather((local_bx, local_by, local_bz), root=0)
if rank == 0:
# Unzip and concatenate
# all_chunks is list of tuples
all_bx = [c[0] for c in all_chunks]
all_by = [c[1] for c in all_chunks]
all_bz = [c[2] for c in all_chunks]
bx = np.concatenate(all_bx, axis=0)
by = np.concatenate(all_by, axis=0)
bz = np.concatenate(all_bz, axis=0)
res = (bx, by, bz)
else:
res = (None, None, None)
# Broadcast results to all ranks
res = comm.bcast(res, root=0)
return res
[docs]
def serial_biot_savart(
element_centers,
element_lengths,
element_directions,
field_points,
order=None,
backend=Backend.PYTHON,
):
"""
Serial Biot-Savart calculation.
Returns (Bx, By, Bz) tuple.
"""
element_centers = np.array(element_centers)
element_lengths = np.array(element_lengths)
element_directions = np.array(element_directions)
field_points = np.array(field_points)
if element_centers.ndim != 2 or element_centers.shape[1] != 3:
raise ValueError("element_centers must be a NumPy array of shape (N, 3)")
# ... (skipping generic checks for brevity, they remain same) ...
# We should keep the checks to be safe or rely on caller?
# Let's assume checks are preserved if we don't overwrite them or copy them.
# The tool replaces CONTIGUOUS blocks. We can replace the body.
# Re-implementing checks (replacing function body).
if (
element_lengths.ndim != 1
or element_lengths.shape[0] != element_centers.shape[0]
):
raise ValueError(
"element_lengths must be a NumPy array of shape (N,) and same len"
)
if (
element_directions.ndim != 2
or element_directions.shape[1] != 3
or element_directions.shape[0] != element_centers.shape[0]
):
raise ValueError(
"element_directions must be a NumPy array of shape (N, 3) and same len"
)
if field_points.ndim != 2 or field_points.shape[1] != 3:
raise ValueError("field_points must be a NumPy array of shape (M, 3)")
source_points = element_centers
dl_vectors = 0.5 * element_lengths[:, np.newaxis] * element_directions
if backend is None:
backend = MultivariateTaylorFunction._IMPLEMENTATION
if isinstance(backend, str):
backend = Backend(backend)
if backend == Backend.COSY:
from sandalwood.backends.cosy.cosy_backend import CosyBackend
existing_order = getattr(CosyBackend, "_order", 0)
existing_dim = getattr(CosyBackend, "_dim", 0)
initialized = getattr(CosyBackend, "_initialized", False)
req_order = (
order
if order is not None
else (existing_order if existing_order > 0 else 1)
)
req_dim = max(3, existing_dim)
if not initialized or existing_order < req_order or existing_dim < req_dim:
CosyBackend.initialize(order=req_order, dim=req_dim)
dl_vectors = 0.5 * element_lengths[:, np.newaxis] * element_directions
pos_x = np.ascontiguousarray(field_points[:, 0])
pos_y = np.ascontiguousarray(field_points[:, 1])
pos_z = np.ascontiguousarray(field_points[:, 2])
src_x = np.ascontiguousarray(element_centers[:, 0])
src_y = np.ascontiguousarray(element_centers[:, 1])
src_z = np.ascontiguousarray(element_centers[:, 2])
dl_x = np.ascontiguousarray(dl_vectors[:, 0])
dl_y = np.ascontiguousarray(dl_vectors[:, 1])
dl_z = np.ascontiguousarray(dl_vectors[:, 2])
is_discrete_mode = isinstance(element_centers.flat[0], (float, int, np.number))
if is_discrete_mode:
# Fast Path (Discrete Geometry)
bx, by, bz = CosyBackend.biot_savart_batch(
pos_x, pos_y, pos_z, src_x, src_y, src_z, dl_x, dl_y, dl_z
)
# Apply scale
scale = mu_0_4pi
bx *= scale
by *= scale
bz *= scale
return bx, by, bz
else:
# Parametric Path (MTF/DA)
c_bx, c_by, c_bz = CosyBackend.biot_savart_batch_indices(
pos_x, pos_y, pos_z, src_x, src_y, src_z, dl_x, dl_y, dl_z
)
dim = CosyBackend._dim
mtf_x = MultivariateTaylorFunction.from_cosy_indices(c_bx, dim)
mtf_y = MultivariateTaylorFunction.from_cosy_indices(c_by, dim)
mtf_z = MultivariateTaylorFunction.from_cosy_indices(c_bz, dim)
scale = mu_0_4pi
bx = mtf_x * scale
by = mtf_y * scale
bz = mtf_z * scale
return bx, by, bz
return _python_biot_savart_core(source_points, dl_vectors, field_points, order)