HelmholtzCoil.py

import numpy as np
import matplotlib.pyplot as plt
from scipy.special import ellipk, ellipe

def single_loop_B(rho, z, R, I, mu0):
    alpha2 = (R + rho)**2 + z**2
    beta2 = (R - rho)**2 + z**2
    # avoid k2=1 exactly or k2 close to 1 leading to large errors
    k2 = 1 - beta2 / alpha2
    k2 = np.clip(k2, 0, 1-1e-10) 
    
    K = ellipk(k2)
    E = ellipe(k2)
    denom = np.sqrt(alpha2)
    
    factor = mu0 * I / (2 * np.pi)
    
    Bz = factor / denom * (K + (R**2 - rho**2 - z**2) / beta2 * E)
    
    # For Brho, handle rho=0 singularity
    # Using a slightly more robust masking approach for the division by rho
    Brho = np.zeros_like(rho)
    # Mask for non-zero rho and avoid the exact wire location singularity
    mask = (rho > 1e-10) & (beta2 > 1e-10)

    if np.any(mask):
        Brho[mask] = factor * (z[mask] / (rho[mask] * denom[mask])) * \
                     (-K[mask] + (R**2 + rho[mask]**2 + z[mask]**2) / beta2[mask] * E[mask])
    
    return Brho, Bz

def helmholtz_B(rho, z, R, d, I, mu0):
    Br1, Bz1 = single_loop_B(rho, z + d/2, R, I, mu0)
    Br2, Bz2 = single_loop_B(rho, z - d/2, R, I, mu0)
    return Br1 + Br2, Bz1 + Bz2

# Sample parameters
R = 1.0  # radius
d = R    # separation for Helmholtz
I = 1.0  # current
mu0 = 1.0  # normalized units

# Grid for full side view
# Increased grid resolution slightly for smoother color contours
ny = 200
nz = 200
# Zoom out slightly to show field decay
y = np.linspace(-2.0*R, 2.0*R, ny)
z = np.linspace(-2.0*R, 2.0*R, nz)
Z, Y = np.meshgrid(z, y)
rho = np.abs(Y)

Brho_abs, Bz = helmholtz_B(rho, Z, R, d, I, mu0)
By = np.sign(Y) * Brho_abs

# --- NEW CALCULATION ---
# Calculate Magnetic Field Magnitude |B|
B_mag = np.sqrt(Brho_abs**2 + Bz**2)

# Plot
fig, ax = plt.subplots(figsize=(10, 8))

# --- NEW PLOTTING SECTION: COLOR CONTOUR ---
# We use contourf to show the magnitude.
# Note: The field goes to infinity exactly at the wires. To prevent the 
# colormap from being washed out by these singularities, we define levels
# focusing on the region of interest (0 to about 2x the center field).
# The theoretical center field is approx 0.7155 in these units.
levels = np.linspace(0, 1.5, 100) 
# Use 'extend="max"' to color everything above 1.5 the brightest color
cp = ax.contourf(Z, Y, B_mag, levels=levels, cmap='plasma', extend='max')

# Add a colorbar showing the scale
cbar = fig.colorbar(cp, ax=ax)
cbar.set_label('Magnetic Field Magnitude $|B|$ (normalized)')
# -------------------------------------------

# Streamplot on top
# Changed color to black/dark grey for better contrast against the 'plasma' map
ax.streamplot(z, y, Bz, By, density=1.5, color='black', linewidth=0.8, arrowsize=1)

ax.set_xlabel('z (axial direction)')
ax.set_ylabel('y (radial direction)')
ax.set_title('Helmholtz Coil: Field Lines and Magnitude')
ax.set_aspect('equal')
ax.set_xlim(z.min(), z.max())
ax.set_ylim(y.min(), y.max())

# Indicate coil planes
ax.axvline(-d/2, color='white', linestyle='--', alpha=0.5)
ax.axvline(d/2, color='white', linestyle='--', alpha=0.5)

# Coil cross-sections (wire positions)
# Changed marker color to cyan for visibility
for pos in [-d/2, d/2]:
    ax.plot(pos, R, 'co', markersize=6, markeredgecolor='k')
    ax.plot(pos, -R, 'co', markersize=6, markeredgecolor='k')

plt.tight_layout()
plt.show()