import numpy as np import matplotlib.pyplot as plt from scipy.integrate import solve_ivp from scipy.interpolate import CubicSpline, NearestNDInterpolator from scipy.spatial import cKDTree import warnings # ------------------------------------------------------------------- # 1. Define a synthetic 2D vector field (magnetic field) # ------------------------------------------------------------------- def magnetic_field_2d(x, y): """ Example: Dipole-like field with a quadrupole perturbation. Bx = 2*x*y / (x^2+y^2)^2 (dipole component) By = (y^2 - x^2) / (x^2+y^2)^2 This is a standard 2D dipole field (in-plane). """ r2 = x*x + y*y if r2 < 1e-6: return 0.0, 0.0 r4 = r2*r2 Bx = 2*x*y / r4 By = (y*y - x*x) / r4 return Bx, By def field_unit_tangent(x, y): Bx, By = magnetic_field_2d(x, y) Bmag = np.hypot(Bx, By) if Bmag < 1e-8: return 0.0, 0.0 return Bx/Bmag, By/Bmag # ------------------------------------------------------------------- # 2. Trace a field line from a seed point using arclength parameter s # dr/ds = unit tangent (B/|B|) # ------------------------------------------------------------------- def trace_field_line(seed, s_max=5.0, step=0.05): def ode(s, r): x, y = r tx, ty = field_unit_tangent(x, y) return [tx, ty] sol = solve_ivp(ode, [0, s_max], seed, t_eval=np.arange(0, s_max+step, step), method='RK45', rtol=1e-6, atol=1e-9) points = sol.y.T # (N, 2) return points # ------------------------------------------------------------------- # 3. Build cubic spline representation for a field line (x(s), y(s)) # ------------------------------------------------------------------- class FieldLineSpline: def __init__(self, points): self.points = np.asarray(points) # Remove duplicate consecutive points diff = np.linalg.norm(np.diff(self.points, axis=0), axis=1) mask = np.concatenate(([True], diff > 1e-12)) self.points = self.points[mask] # Compute cumulative arclength diffs = np.diff(self.points, axis=0) seg_lengths = np.linalg.norm(diffs, axis=1) self.s = np.concatenate(([0], np.cumsum(seg_lengths))) # Fit cubic splines for x(s) and y(s) self.spline_x = CubicSpline(self.s, self.points[:, 0], bc_type='natural') self.spline_y = CubicSpline(self.s, self.points[:, 1], bc_type='natural') def position(self, s): return np.array([self.spline_x(s), self.spline_y(s)]) def tangent(self, s): dx = self.spline_x.derivative()(s) dy = self.spline_y.derivative()(s) norm = np.hypot(dx, dy) if norm < 1e-8: return np.array([0.0, 0.0]) return np.array([dx/norm, dy/norm]) # unit tangent def get_coefficients(self): coeffs_x = self.spline_x.c coeffs_y = self.spline_y.c knots = self.spline_x.x segs = [] for i in range(len(knots)-1): segs.append({ 's_start': knots[i], 's_end': knots[i+1], 'cx': coeffs_x[:, i], 'cy': coeffs_y[:, i] }) return segs # ------------------------------------------------------------------- # 4. Build a set of field lines covering the region # ------------------------------------------------------------------- def generate_field_lines(seeds, s_max=4.0, step=0.05): lines = [] for seed in seeds: points = trace_field_line(seed, s_max, step) if len(points) > 5: lines.append(FieldLineSpline(points)) return lines # ------------------------------------------------------------------- # 5. Reconstruct field from splines (at any point) # Nearest field line method: find closest point on any spline, # then use that spline's tangent. # ------------------------------------------------------------------- class FieldReconstructor: def __init__(self, field_lines): self.lines = field_lines # Build a KD-tree of all sampled points from all lines self.all_points = [] self.line_indices = [] self.s_values = [] for i, line in enumerate(self.lines): pts = line.points s_vals = line.s for j, pt in enumerate(pts): self.all_points.append(pt) self.line_indices.append(i) self.s_values.append(s_vals[j]) self.all_points = np.array(self.all_points) self.kdtree = cKDTree(self.all_points) def reconstruct_field(self, x, y): """Return unit tangent vector at (x,y) from nearest field line point.""" query_point = np.array([x, y]) dist, idx = self.kdtree.query(query_point) line_idx = self.line_indices[idx] s_val = self.s_values[idx] # Use the spline's tangent at that s tangent = self.lines[line_idx].tangent(s_val) return tangent # ------------------------------------------------------------------- # 6. Verification: compare original field vs spline-reconstructed field # ------------------------------------------------------------------- def verify_accuracy(field_lines, reconstructor, grid_points, title="Verification"): """ grid_points: list of (x,y) points to evaluate (can be a mesh). Returns: mean angular error (degrees), mean magnitude ratio error (if magnitude reconstructed) Here we only reconstruct direction, so we compare angle. """ errors_deg = [] for (x,y) in grid_points: # Original field direction tx_orig, ty_orig = field_unit_tangent(x, y) # Reconstructed direction trec = reconstructor.reconstruct_field(x, y) # Compute angle difference dot = tx_orig*trec[0] + ty_orig*trec[1] dot = np.clip(dot, -1.0, 1.0) angle = np.arccos(dot) * 180.0 / np.pi errors_deg.append(angle) mean_err = np.mean(errors_deg) std_err = np.std(errors_deg) return mean_err, std_err, errors_deg # ------------------------------------------------------------------- # 7. Main: generate field lines, reconstruct, measure loss # ------------------------------------------------------------------- if __name__ == "__main__": # Define seeds: points on a circle around the origin radii = [1.0, 1.5, 2.0, 2.5, 3.0] angles = np.linspace(0, 2*np.pi, 120, endpoint=False) seeds = [] for r in radii: for theta in angles: seeds.append([r*np.cos(theta), r*np.sin(theta)]) # Generate field lines (arclength up to 4.0) print("Tracing field lines...") lines = generate_field_lines(seeds, s_max=4.0, step=0.05) print(f"Generated {len(lines)} field lines.") # Build reconstructor reconstructor = FieldReconstructor(lines) # Create test grid (points not necessarily on field lines) x_grid = np.linspace(-3.5, 3.5, 30) y_grid = np.linspace(-3.5, 3.5, 30) test_points = [] for x in x_grid: for y in y_grid: if np.hypot(x,y) > 0.5 and np.hypot(x,y) < 3.8: # avoid singularity and boundary test_points.append((x,y)) # Verify accuracy mean_err, std_err, errors = verify_accuracy(lines, reconstructor, test_points) print(f"\n--- Verification Results ---") print(f"Number of test points: {len(test_points)}") print(f"Mean angular error: {mean_err:.2f}° ± {std_err:.2f}°") # Loss metric: fraction of points with error > 10° loss = np.mean(np.array(errors) > 10.0) * 100 print(f"Loss (error > 10°): {loss:.1f}%") # Visualize error map x_vals = [p[0] for p in test_points] y_vals = [p[1] for p in test_points] E_vals = np.array(errors) plt.figure(figsize=(10,8)) plt.scatter(x_vals, y_vals, c=E_vals, cmap='hot', s=10) plt.colorbar(label='Angular error (degrees)') # Plot some field lines for line in lines[:10]: pts = line.points plt.plot(pts[:,0], pts[:,1], 'b-', linewidth=0.5, alpha=0.5) plt.title(f"Field Reconstruction Error (mean = {mean_err:.1f}°)") plt.xlabel('x'); plt.ylabel('y') plt.axis('equal') plt.tight_layout() plt.show() # Also compute error along the field lines themselves (should be near zero) print("\n--- Self-consistency check (on field lines) ---") on_line_points = [] for line in lines: # sample points along the spline (not just original traced points) s_vals = np.linspace(line.s[0], line.s[-1], 20) for s in s_vals: x,y = line.position(s) on_line_points.append((x,y)) mean_err_on, _, _ = verify_accuracy(lines, reconstructor, on_line_points) print(f"Mean angular error on field lines: {mean_err_on:.2f}°") # Print spline coefficients for first line as example coeffs = lines[0].get_coefficients() print("\nExample cubic spline coefficients (first segment of first line):") print(f" x(s) = {coeffs[0]['cx'][0]:.4f}*(s-si)^3 + {coeffs[0]['cx'][1]:.4f}*(s-si)^2 + {coeffs[0]['cx'][2]:.4f}*(s-si) + {coeffs[0]['cx'][3]:.4f}") print(f" y(s) = {coeffs[0]['cy'][0]:.4f}*(s-si)^3 + {coeffs[0]['cy'][1]:.4f}*(s-si)^2 + {coeffs[0]['cy'][2]:.4f}*(s-si) + {coeffs[0]['cy'][3]:.4f}")