From 8bec0f7b1dcda64eda4c0fcc1087129db0e3037a Mon Sep 17 00:00:00 2001 From: Rogerio Jorge Date: Fri, 23 Oct 2026 23:02:16 -0500 Subject: [PATCH 1/3] Fix boundary modes, vacuum flux export, and surface transform fit --- essos/fields.py | 13 ++ essos/surfaces.py | 246 ++++++++++++++++++++++--------- tests/test_fields.py | 12 +- tests/test_surface_flux.py | 80 ++++++++++ tests/test_surface_vmec_input.py | 94 ++++++++++++ 5 files changed, 368 insertions(+), 77 deletions(-) create mode 100644 tests/test_surface_flux.py create mode 100644 tests/test_surface_vmec_input.py diff --git a/essos/fields.py b/essos/fields.py index 7cb66a1b..83193d04 100644 --- a/essos/fields.py +++ b/essos/fields.py @@ -145,6 +145,19 @@ def B(self, points): return jnp.mean(dB_sum, axis=0) + @jit + def A(self, points): + """Cartesian filament vector potential [T m] at a point [m]. + + Uses the same periodic coil quadrature and currents as :meth:`B`, + so ``curl(A) = B`` and shape/current derivatives remain traceable. + """ + distance = jnp.linalg.norm(jnp.asarray(points) - self.coils.gamma, axis=-1) + tangent_over_distance = self.coils.gamma_dash / distance[..., None] + potential = jnp.einsum('i,ijk->jk', self.coils.currents * 1e-7, + tangent_over_distance) + return jnp.mean(potential, axis=0) + @jit def b_cyl(self, R, phi, Z): """Return ``(B_R, B_phi, B_Z)`` on broadcast cylindrical arrays. diff --git a/essos/surfaces.py b/essos/surfaces.py index 527669d8..191962f6 100644 --- a/essos/surfaces.py +++ b/essos/surfaces.py @@ -20,17 +20,20 @@ def _cacheable(*values): @jit def toroidal_flux(surface, field, idx=0) -> jnp.ndarray: - curve = surface.gamma[idx] - dl = jnp.roll(curve, -1, axis=0) - curve - A_vals = vmap(field.A)(curve) - Adl = jnp.sum(A_vals * dl, axis=1) - tf = jnp.sum(Adl) - #curve = surface.gamma[idx] - #dl = surface.gammadash_theta[idx] - #A_vals = vmap(field.A)(curve) - #Adl = jnp.sum(A_vals * dl, axis=1)/surface.ntheta - #tf = jnp.sum(Adl) - return tf + """Signed magnetic flux through the poloidal loop at toroidal grid index idx. + + Integrate ``A . dx`` using the analytic theta tangent and periodic + quadrature. For ``R = R0 + a cos(theta), Z = b sin(theta)`` the loop + normal points in the negative toroidal direction: positive B_phi gives + negative flux. The vector potential must be single-valued on the loop. + """ + count = surface.ntheta - int(surface.close) + if count < 2: + raise ValueError('Toroidal flux requires at least two distinct theta points') + curve = surface.gamma[idx, :count] + tangent = surface.gammadash_theta[idx, :count] + return 2 * jnp.pi * jnp.mean(jnp.sum(vmap(field.A)(curve) * tangent, axis=-1)) + @jit def poloidal_flux(surface, field, idx=0) -> jnp.ndarray: @@ -89,44 +92,71 @@ def B_contravariant_theta_phi_on_surface(surface, field): # The final index selects theta (0) or phi (1). return jnp.stack((B_contra_theta, B_contra_phi), axis=-1) -def iota_approx_on_surface(surface, field, regularization=1e-10): - """ - Rotational transform (iota) on the surface. - Only meaningful when B.n ~ 0 (the surface is a flux surface). +def iota_approx_on_surface(surface, field, regularization=1e-10, *, + lambda_mpol=None, lambda_ntor=None, return_diagnostics=False): + r"""Fit ``theta* = theta + lambda`` and a constant rotational transform. + + This is a surface PDE fit, not a field-line measurement. It requires a flux + surface and nonzero B^phi. Resolve lambda independently of boundary shape: + even a circular boundary can need high lambda harmonics. Defaults preserve + the boundary's Fourier basis; supply both cutoffs to test convergence. + The fit does not enforce invertibility of theta*. Optional diagnostics give + the RMS PDE residual and minimum ``1 + d(lambda)/d(theta)``; a small residual + alone does not establish a valid straight-field-line coordinate system. """ - - # --- 1. Local field-line slope on the surface ------------------------------------- - # A field line tangent to the surface moves as dtheta/dphi = B^theta / B^phi. - # This slope depends on how theta is parametrized, so it is NOT iota yet. - B_contra = B_contravariant_theta_phi_on_surface(surface, field) - slope = ( B_contra[..., 0] / B_contra[..., 1] ).ravel() # shape (nphi*ntheta,) - - # --- 2. Basis for lambda, reusing the surface's own Fourier modes ------------------ - # Straight-field-line angle: theta* = theta + lambda(theta, phi), where iota is constant. - # Stellarator symmetry -> lambda = sum_mn l_mn sin(m theta - xn phi). - # Drop mode (0,0): sin(0) = 0 has no effect. - xm = surface.xm[1:] - xn = surface.xn[1:] - angles = surface.angles[1:].reshape(xm.size, -1).T # shape (npoints, nmodes) - - # --- 3. Linear system: one equation per grid point --------------------------------- - # Along a field line, d(theta*)/dphi = iota -> slope*(1 + dlambda/dtheta) + dlambda/dphi = iota - # With dlambda/dtheta = sum l m cos(.), dlambda/dphi = -sum l xn cos(.): - # iota - sum_mn l_mn (m*slope - xn) cos(angle) = slope - # Unknowns: x = [iota, l_mn]. Linear in x. - lambda_columns = -( xm[None, :] * slope[:, None] - xn[None, :] ) * jnp.cos(angles) - A = jnp.concatenate([ jnp.ones((slope.size, 1)), lambda_columns ], axis=1) - - # --- 4. Least-squares solution (normal equations) ---------------------------------- - # Small system (1 + nmodes unknowns), cheap and differentiable with JAX. - # The tiny regularization only avoids a singular matrix. - AtA = A.T @ A + regularization * jnp.eye(A.shape[1]) - Ats = A.T @ slope - solution = jnp.linalg.solve(AtA, Ats) - - # --- 5. iota is the first unknown --------------------------------------------------- - iota = solution[0] - return iota + if (lambda_mpol is None) != (lambda_ntor is None): + raise ValueError("Specify both lambda_mpol and lambda_ntor, or neither") + if lambda_mpol is None: + xm, xn = surface.xm[1:], surface.xn[1:] + else: + if lambda_mpol < 0 or lambda_ntor < 0: + raise ValueError("Lambda Fourier cutoffs must be nonnegative") + modes = [(m, n * surface.nfp) for m in range(lambda_mpol + 1) + for n in range(-lambda_ntor, lambda_ntor + 1) + if m > 0 or n > 0] + xm = jnp.asarray([m for m, _ in modes]) + xn = jnp.asarray([n for _, n in modes]) + angles = (xm[:, None, None] * surface.theta2d + - xn[:, None, None] * surface.phi2d).reshape(xm.size, surface.nphi * surface.ntheta).T + contra = B_contravariant_theta_phi_on_surface(surface, field) + slope = (contra[..., 0] / contra[..., 1]).ravel() + derivative = xm[None, :] * slope[:, None] - xn[None, :] + columns = -derivative * jnp.cos(angles) + theta_derivative_columns = xm[None, :] * jnp.cos(angles) + if surface.rs is not None or surface.zc is not None: + # Cosine lambda modes are needed when stellarator symmetry is absent. + columns = jnp.concatenate((columns, derivative * jnp.sin(angles)), axis=1) + theta_derivative_columns = jnp.concatenate( + (theta_derivative_columns, -xm[None, :] * jnp.sin(angles)), axis=1) + A = jnp.concatenate((jnp.ones((slope.size, 1)), columns), axis=1) + + # Apply the same endpoint quadrature as surface integrals. Normalize it so + # regularization has the same meaning at every grid resolution. + weights = jnp.ones((surface.nphi, surface.ntheta)) + if surface.close: + weights = weights.at[:, -1].set(0.0) + if surface.range_torus == "full torus": + weights = weights.at[-1, :].set(0.0) + else: + weights = weights.at[jnp.array([0, -1]), :].multiply(0.5) + weights = weights.ravel() / jnp.sum(weights) + sqrt_weights = jnp.sqrt(weights) + # Ridge only the coordinate correction, leaving a constant iota unbiased. + # QR avoids squaring the condition number as the normal equations did. + ridge = jnp.sqrt(regularization) * jnp.eye(A.shape[1]).at[0, 0].set(0.0) + design = jnp.concatenate((sqrt_weights[:, None] * A, ridge), axis=0) + rhs = jnp.concatenate((sqrt_weights * slope, jnp.zeros(A.shape[1]))) + Q, R = jnp.linalg.qr(design, mode="reduced") + solution = jnp.linalg.solve(R, Q.T @ rhs) + if return_diagnostics: + mean_squared_residual = jnp.sum(weights * (A @ solution - slope)**2) + angle_jacobian = 1 + theta_derivative_columns @ solution[1:] + return {"iota": solution[0], + "mean_squared_residual": mean_squared_residual, + "rms_residual": jnp.sqrt(mean_squared_residual), + "angle_jacobian": angle_jacobian.reshape(surface.nphi, surface.ntheta), + "minimum_angle_jacobian": jnp.min(angle_jacobian)} + return solution[0] @jit def BdotN(surface, field): @@ -138,6 +168,26 @@ def BdotN(surface, field): def BdotN_over_B(surface, field, **kwargs): return BdotN(surface, field) / jnp.linalg.norm(B_on_surface(surface, field), axis=2) +def _surface_mean(surface, values): + """Average of ``values`` (on the ``(nphi, ntheta)`` grid) over the whole torus. + + Theta is periodic, so a duplicated ``close=True`` end column is dropped. + Over a full torus phi is treated the same way; over a half period with + ``close=True`` both ends are sampled and the trapezoid rule is used. + A half period with ``close=False`` is a midpoint grid. Surfaces without + ``close``/``range_torus`` (e.g. VMEX's) are taken as open full-torus grids. + Integrals over the torus are this mean times ``(2 pi)**2``, because the + normal is taken with respect to theta and phi in radians and a + half-period integrand repeats over the ``2 nfp`` half periods. + """ + close = getattr(surface, "close", False) + if close: + values = values[:, :-1] + if getattr(surface, "range_torus", "full torus") == "full torus" or not close: + return jnp.mean(values[:-1] if close else values) + weights = jnp.ones(values.shape[0]).at[jnp.array([0, -1])].set(0.5) + return jnp.sum(weights @ values) / (jnp.sum(weights) * values.shape[1]) + @jit def _squared_flux_local(surface, field): return 0.5 * jnp.mean(BdotN(surface, field)**2 / jnp.sum(B_on_surface(surface, field)**2, axis=2) @@ -305,32 +355,68 @@ def _compute_scaling(xm, xn, scaling_type, scaling_factor): @classmethod - def from_input_file(cls, file, ntheta=30, nphi=30, close=True, range_torus='full torus'): - """Read indexed boundary coefficients; truncate outside MPOL/NTOR as VMEC does.""" + def from_input_file(cls, file, ntheta=30, nphi=30, close=True, range_torus='full torus', + *, mpol=None, ntor=None, max_mode=None): + """Read a boundary independently of the VMEC solver's MPOL and NTOR. + + By default the surface contains the nonzero boundary modes (including + asymmetric partners when LASYM is true). Explicit ``mpol`` and ``ntor`` + are maximum surface mode numbers: smaller limits truncate the boundary, + while larger limits pad it with zeros. ``max_mode`` sets both limits. + """ from f90nml import Parser nml = Parser().read(file)['indata'] - nfp, mpol, ntor = int(nml.get('nfp', 1)), int(nml.get('mpol', 6)) - 1, int(nml.get('ntor', 0)) - if mpol < 0 or ntor < 0: - raise ValueError('Require MPOL >= 1 and NTOR >= 0') - def coefficients(name): - matrix = np.zeros((mpol + 1, 2 * ntor + 1)) - if name not in nml: - return jnp.asarray(matrix.T) - n0, m0 = nml.start_index[name] - for i, row in enumerate(nml[name]): - for j, value in enumerate(row or []): - m, n = m0 + i, n0 + j - if value is not None and 0 <= m <= mpol and abs(n) <= ntor: + def mode_limit(value, name): + if value is not None and (isinstance(value, bool) or + not isinstance(value, (int, np.integer)) or value < 0): + raise ValueError(f'{name} must be a non-negative integer') + return int(value) if value is not None else None + + mpol, ntor = mode_limit(mpol, 'mpol'), mode_limit(ntor, 'ntor') + max_mode = mode_limit(max_mode, 'max_mode') + if max_mode is not None: + if any(value is not None and value != max_mode for value in (mpol, ntor)): + raise ValueError('max_mode conflicts with mpol or ntor') + mpol = ntor = max_mode + + # f90nml reverses Fortran's axes: rows vary m, entries within a row + # vary n. Each family can have its own lower bounds and ragged rows. + names = ['rbc', 'zbs'] + if nml.get('lasym', False): + names += [name for name in ('rbs', 'zbc') if name in nml] + tables = {} + for name in names: + table = {} + if name in nml: + n0, m0 = nml.start_index[name] + for i, row in enumerate(nml[name]): + for j, value in enumerate(row or []): + m, n = m0 + i, n0 + j + if value is None or m < 0: + continue + # The m=0 negative-n modes duplicate positive-n modes. sign = -1 if m == 0 and n < 0 and name in ('rbs', 'zbs') else 1 n = abs(n) if m == 0 else n - matrix[m, n + ntor] += sign * value - return jnp.asarray(matrix.T) - partners = {name: coefficients(name) for name in ('rbs', 'zbc') - if nml.get('lasym', False) and name in nml} - return surfacerzfourier_from_boundary(coefficients('rbc'), coefficients('zbs'), nfp, + table[m, n] = table.get((m, n), 0.0) + sign * value + tables[name] = table + support = [(m, n) for table in tables.values() for (m, n), value in table.items() if value != 0] + if mpol is None: + mpol = max((m for m, n in support), default=0) + if ntor is None: + ntor = max((abs(n) for m, n in support), default=0) + + def coefficients(name): + matrix = np.zeros((2 * ntor + 1, mpol + 1)) + for (m, n), value in tables[name].items(): + if m <= mpol and abs(n) <= ntor: + matrix[n + ntor, m] = value + return jnp.asarray(matrix) + + partners = {name: coefficients(name) for name in ('rbs', 'zbc') if name in tables} + return surfacerzfourier_from_boundary(coefficients('rbc'), coefficients('zbs'), int(nml.get('nfp', 1)), ntheta=ntheta, nphi=nphi, close=close, range_torus=range_torus, _cls=cls, **partners) - + @classmethod def from_vmec(cls, vmec, s=1, ntheta=30, nphi=30, close=True, range_torus='full torus', offset=0.): """Flux surface ``s`` of a VMEC field; ``offset`` [m] grows its m = 1, n = 0 modes, e.g. a wall that far outside the LCFS.""" @@ -534,6 +620,9 @@ def _compute_meshgrid(self): div, end_val = self.nfp, 0.5 quadpoints_theta = jnp.linspace(0, 2 * jnp.pi, num=self.ntheta, endpoint=self.close) quadpoints_phi = jnp.linspace(0, 2 * jnp.pi * end_val / div, num=self.nphi, endpoint=self.close) + if self.range_torus == "half period" and not self.close: + # Midpoint quadrature for integrands repeated by stellarator symmetry. + quadpoints_phi = quadpoints_phi + jnp.pi * end_val / div / self.nphi theta2d, phi2d = jnp.meshgrid(quadpoints_theta, quadpoints_phi) return quadpoints_theta, quadpoints_phi, theta2d, phi2d @@ -887,12 +976,27 @@ def to_vtk(self, filename, extra_data=None, field=None): pointData = {**pointData, **extra_data} gridToVTK(str(filename), x, y, z, pointData=pointData) - def to_vmec(self, filename): - """Write a VMEC boundary namelist, including any asymmetric partners.""" + def to_vmec(self, filename, *, mpol=None, ntor=None, + ftol_array=(1e-12,), ns_array=(51,), niter_array=(3500,)): + """Write a vacuum VMEC input, including any asymmetric boundary partners. + + Solver MPOL/NTOR default to the surface mode maxima plus two. They can + be overridden independently of the boundary representation; MPOL is a + count (VMEC retains m < MPOL), while NTOR is a maximum mode number. + The default zero-current constraint lets VMEC compute vacuum iota. + """ families = {'RBC': self.rc, 'ZBS': self.zs, 'RBS': self.rs, 'ZBC': self.zc} asym = self.rs is not None or self.zc is not None + solver_mpol = self.mpol + 2 if mpol is None else mpol + solver_ntor = self.ntor + 2 if ntor is None else ntor + if (not isinstance(solver_mpol, (int, np.integer)) or solver_mpol <= self.mpol or + not isinstance(solver_ntor, (int, np.integer)) or solver_ntor < self.ntor): + raise ValueError('Solver MPOL/NTOR must retain all surface modes') lines = ['&INDATA', f'LASYM = .{str(asym).upper()}.', f'NFP = {self.nfp}', - f'MPOL = {self.mpol + 1}', f'NTOR = {self.ntor}'] + f'MPOL = {solver_mpol}', f'NTOR = {solver_ntor}', + 'NCURR = 1', 'CURTOR = 0.0', 'AC = 0.0', 'AM = 0.0'] + for name, values in (('FTOL_ARRAY', ftol_array), ('NS_ARRAY', ns_array), ('NITER_ARRAY', niter_array)): + lines.append(f'{name} = ' + ', '.join(str(value) for value in np.atleast_1d(values))) for m, xn, i in zip(self.xm, self.xn, range(len(self.xm))): if int(xn) % self.nfp: raise ValueError('Toroidal mode numbers must be multiples of nfp') diff --git a/tests/test_fields.py b/tests/test_fields.py index ffb25bd5..0e2bb709 100644 --- a/tests/test_fields.py +++ b/tests/test_fields.py @@ -477,7 +477,7 @@ def test_asymmetric_surface_sparse_input_coefficients(tmp_path): 'ZBS(0,1)=.8, ZBS(-1,0)=-.1, ZBS(1,0)=.2, ' 'RBS(-1,1)=.2, RBS(-1,0)=.4, RBS(1,0)=.5, ' 'ZBC(0,0)=.1, ZBC(-1,0)=.2, ZBC(1,0)=.1 /') - surface = SurfaceRZFourier.from_input_file(filename, ntheta=5, nphi=7, close=False) + surface = SurfaceRZFourier.from_input_file(filename, ntheta=5, nphi=7, close=False, mpol=1, ntor=1) np.testing.assert_allclose(surface.rc, [3, .3, 0, 0, .7], atol=1e-14) np.testing.assert_allclose(surface.zs, [0, .3, 0, .8, 0], atol=1e-14) np.testing.assert_allclose(surface.rs, [0, .1, .2, 0, 0], atol=1e-14) @@ -499,11 +499,11 @@ class DerivedSurface(SurfaceRZFourier): assert (symmetric.nfp, symmetric.mpol, symmetric.ntor) == (1, 1, 0) filename.write_text('&INDATA RBC(0,0)=3 /') default = SurfaceRZFourier.from_input_file(filename) - assert (default.nfp, default.mpol, default.ntor) == (1, 5, 0) - np.testing.assert_array_equal(default.zs, np.zeros(6)) + assert (default.nfp, default.mpol, default.ntor) == (1, 0, 0) + np.testing.assert_array_equal(default.zs, np.zeros(1)) filename.write_text('&INDATA MPOL=0, NTOR=0 /') - with pytest.raises(ValueError, match='MPOL >= 1'): - SurfaceRZFourier.from_input_file(filename) + with pytest.raises(ValueError, match='mpol must be a non-negative integer'): + SurfaceRZFourier.from_input_file(filename, mpol=-1) @pytest.mark.parametrize('phase', [0.0, 0.37]) @@ -541,7 +541,7 @@ def test_surface_vmec_export_preserves_multimode_geometry(tmp_path, asymmetric): surface.to_vmec(filename) text = filename.read_text() assert f'LASYM = .{str(asymmetric).upper()}.' in text - assert 'MPOL = 3' in text and 'NTOR = 1' in text and 'RBC(1,1)' in text + assert 'MPOL = 4' in text and 'NTOR = 3' in text and 'RBC(1,1)' in text rebuilt = SurfaceRZFourier.from_input_file(filename, ntheta=8, nphi=9, close=False) for name in ('rc', 'zs', 'rs', 'zc'): value = getattr(surface, name) diff --git a/tests/test_surface_flux.py b/tests/test_surface_flux.py new file mode 100644 index 00000000..23123ae7 --- /dev/null +++ b/tests/test_surface_flux.py @@ -0,0 +1,80 @@ +"""Signed toroidal flux for VMEC export and its coil vector potential.""" +import equinox as eqx +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from essos.surfaces import SurfaceRZFourier, toroidal_flux + + +class ToroidalField(eqx.Module): + strength: float = 2.0 + + def A(self, xyz): + # curl(-B R e_z) = B e_phi. + return jnp.array([0., 0., -self.strength * jnp.hypot(xyz[0], xyz[1])]) + + +def ellipse(a=.7, b=.4, count=12, close=False, range_torus='full torus'): + return SurfaceRZFourier(jnp.array([3., a]), jnp.array([0., b]), + nfp=3, mpol=1, ntor=0, ntheta=count, nphi=count, + close=close, range_torus=range_torus) + + +@pytest.mark.parametrize('close', [False, True]) +@pytest.mark.parametrize('count', [5, 9, 17]) +@pytest.mark.parametrize('range_torus', ['full torus', 'half period']) +def test_toroidal_flux_analytic_area_and_sign(close, count, range_torus): + surface = ellipse(count=count, close=close, range_torus=range_torus) + for idx in (0, 2): + np.testing.assert_allclose(toroidal_flux(surface, ToroidalField(), idx), + -2 * np.pi * .7 * .4, atol=2e-14) + + +@pytest.mark.parametrize('close,count', [(False, 1), (True, 2)]) +def test_toroidal_flux_requires_distinct_points(close, count): + with pytest.raises(ValueError, match='distinct theta'): + toroidal_flux(ellipse(count=count, close=close), ToroidalField()) + + +@pytest.mark.parametrize('close', [False, True]) +def test_toroidal_flux_shape_gradient(close): + def flux(a): + return toroidal_flux(ellipse(a=a, close=close), ToroidalField()) + derivative = jax.jit(jax.grad(flux))(.7) + finite_difference = (flux(.7 + 1e-5) - flux(.7 - 1e-5)) / 2e-5 + np.testing.assert_allclose(derivative, -2 * np.pi * .4, rtol=1e-12) + np.testing.assert_allclose(derivative, finite_difference, rtol=1e-9) + + +def test_biot_savart_vector_potential_curl_matches_field(): + from essos.coils import Coils, CreateEquallySpacedCurves + from essos.fields import BiotSavart + curves = CreateEquallySpacedCurves(2, 2, 3., 1., n_segments=40, nfp=2, stellsym=True) + field = BiotSavart(Coils(curves, jnp.array([8e4, -12e4]))) + for xyz in (jnp.array([.6, .2, -.3]), jnp.array([2., .4, .1])): + jacobian = jax.jacfwd(field.A)(xyz) + curl = jnp.array([jacobian[2, 1] - jacobian[1, 2], + jacobian[0, 2] - jacobian[2, 0], + jacobian[1, 0] - jacobian[0, 1]]) + np.testing.assert_allclose(curl, field.B(xyz), rtol=2e-13, atol=2e-15) + + +def test_coil_toroidal_flux_matches_independent_disk_integral(): + from essos.coils import Coils, Curves + from essos.fields import BiotSavart + # A vertical circular coil surrounds the small elliptical poloidal loop. + coefficients = jnp.zeros((1, 3, 3)).at[0, 0, 0].set(3.) + coefficients = coefficients.at[0, 0, 2].set(1.).at[0, 2, 1].set(1.) + field = BiotSavart(Coils(Curves(coefficients, n_segments=64, stellsym=False), jnp.array([1e5]))) + surface = ellipse(a=.3, b=.2, count=32) + nodes, weights = np.polynomial.legendre.leggauss(20) + rho = jnp.asarray((nodes + 1) / 2)[:, None] + theta = jnp.linspace(0, 2 * jnp.pi, 48, endpoint=False)[None, :] + x, z = 3 + .3 * rho * jnp.cos(theta), .2 * rho * jnp.sin(theta) + points = jnp.stack((x, jnp.zeros_like(x), z), axis=-1) + by = jax.vmap(field.B)(points.reshape(-1, 3))[:, 1].reshape(x.shape) + # Increasing theta gives a disk normal of -e_phi at phi=0. + disk_flux = jnp.sum(-by * (.3 * .2 * rho) * jnp.asarray(weights / 2)[:, None]) * (2 * jnp.pi / 48) + np.testing.assert_allclose(toroidal_flux(surface, field), disk_flux, rtol=1e-12, atol=1e-14) diff --git a/tests/test_surface_vmec_input.py b/tests/test_surface_vmec_input.py new file mode 100644 index 00000000..fb1407d6 --- /dev/null +++ b/tests/test_surface_vmec_input.py @@ -0,0 +1,94 @@ +"""Boundary mode support and equilibrium solver controls are independent.""" +import numpy as np +import pytest +from f90nml import read + +from essos.surfaces import SurfaceRZFourier + + +def coefficient(surface, family, m, n): + mask = (np.asarray(surface.xm) == m) & (np.asarray(surface.xn) == n * surface.nfp) + return np.asarray(getattr(surface, family))[mask].item() + + +@pytest.mark.parametrize('solver_resolution', ['MPOL=2, NTOR=0,', 'MPOL=20, NTOR=16,', '']) +def test_input_infers_nonzero_boundary_support(tmp_path, solver_resolution): + filename = tmp_path / 'input.boundary' + filename.write_text('&INDATA NFP=3, ' + solver_resolution + + 'RBC(0,0)=3, RBC(0,1)=.4, RBC(-2,3)=.03, ' + 'RBC(-9,8)=0, ZBS(0,1)=.5, ZBS(4,2)=-.02, ZBS(11,10)=0 /') + surface = SurfaceRZFourier.from_input_file(filename, ntheta=7, nphi=9, close=False) + assert (surface.mpol, surface.ntor, surface.nfp) == (3, 4, 3) + assert coefficient(surface, 'rc', 3, -2) == .03 + assert coefficient(surface, 'zs', 2, 4) == -.02 + assert coefficient(surface, 'rc', 0, 0) == 3 + # Off-diagonal signed modes catch both swapped indices and sign mistakes. + theta, phi = surface.theta2d, surface.phi2d + radius = 3 + .4 * np.cos(theta) + .03 * np.cos(3 * theta + 6 * phi) + height = .5 * np.sin(theta) - .02 * np.sin(2 * theta - 12 * phi) + np.testing.assert_allclose(surface.gamma[..., 0], radius * np.cos(phi), atol=1e-14) + np.testing.assert_allclose(surface.gamma[..., 2], height, atol=1e-14) + + +def test_independent_input_limits_pad_and_truncate(tmp_path): + filename = tmp_path / 'input.limits' + filename.write_text('&INDATA MPOL=9, NTOR=9, RBC(0,0)=3, RBC(-2,3)=.03, ZBS(1,1)=.5 /') + inferred = SurfaceRZFourier.from_input_file(filename) + assert (inferred.mpol, inferred.ntor) == (3, 2) + truncate_m = SurfaceRZFourier.from_input_file(filename, mpol=1) + assert (truncate_m.mpol, truncate_m.ntor) == (1, 2) + assert np.count_nonzero(truncate_m.rc) == 1 + truncate_n = SurfaceRZFourier.from_input_file(filename, ntor=1) + assert (truncate_n.mpol, truncate_n.ntor) == (3, 1) + assert np.count_nonzero(truncate_n.rc) == 1 + padded = SurfaceRZFourier.from_input_file(filename, mpol=5, ntor=4) + assert (padded.mpol, padded.ntor) == (5, 4) + np.testing.assert_allclose(padded.gamma, inferred.gamma, atol=1e-14) + limited = SurfaceRZFourier.from_input_file(filename, max_mode=1) + assert (limited.mpol, limited.ntor) == (1, 1) + for kwargs in ({'mpol': -1}, {'ntor': 1.5}, {'max_mode': True}, {'max_mode': 2, 'ntor': 1}): + with pytest.raises(ValueError): + SurfaceRZFourier.from_input_file(filename, **kwargs) + + +def test_asymmetric_support_and_negative_n_folding(tmp_path): + filename = tmp_path / 'input.asymmetric' + filename.write_text('&INDATA LASYM=.true., RBC(0,0)=3, RBC(-6,0)=.2, RBC(6,0)=-.2, ' + 'ZBS(0,1)=.5, RBS(-2,3)=.03, RBS(-1,0)=.1, RBS(1,0)=.4, ' + 'ZBC(4,2)=.02, ZBC(-1,0)=.1, ZBC(1,0)=.2 /') + surface = SurfaceRZFourier.from_input_file(filename) + assert (surface.mpol, surface.ntor) == (3, 4) + assert coefficient(surface, 'rs', 3, -2) == .03 + assert coefficient(surface, 'rs', 0, 1) == pytest.approx(.3) + assert coefficient(surface, 'zc', 2, 4) == .02 + assert coefficient(surface, 'zc', 0, 1) == pytest.approx(.3) + filename.write_text('&INDATA RBC(8,7)=0, ZBS(-10,9)=0 /') + empty = SurfaceRZFourier.from_input_file(filename) + assert (empty.mpol, empty.ntor) == (0, 0) + np.testing.assert_array_equal(empty.rc, [0.]) + + +def test_vmec_export_vacuum_solver_defaults_and_overrides(tmp_path): + filename = tmp_path / 'input.surface' + filename.write_text('&INDATA RBC(0,0)=3, RBC(0,1)=.4, RBC(-2,3)=.03, ZBS(0,1)=.5 /') + surface = SurfaceRZFourier.from_input_file(filename) + surface.to_vmec(filename) + nml = read(filename)['indata'] + assert (nml['mpol'], nml['ntor'], nml['ncurr']) == (5, 4, 1) + assert nml['ftol_array'] == 1e-12 + assert nml['ns_array'] == 51 + assert nml['niter_array'] == 3500 + assert nml['curtor'] == nml['ac'] == nml['am'] == 0 + rebuilt = SurfaceRZFourier.from_input_file(filename) + assert (rebuilt.mpol, rebuilt.ntor) == (surface.mpol, surface.ntor) + np.testing.assert_allclose(rebuilt.gamma, surface.gamma, atol=1e-14) + surface.to_vmec(filename, mpol=7, ntor=5, ftol_array=[1e-8, 1e-13], ns_array=[11, 51], niter_array=[100, 4000]) + nml = read(filename)['indata'] + assert (nml['mpol'], nml['ntor']) == (7, 5) + assert nml['ftol_array'] == [1e-8, 1e-13] + assert nml['ns_array'] == [11, 51] + assert nml['niter_array'] == [100, 4000] + with pytest.raises(ValueError, match='retain all surface modes'): + surface.to_vmec(filename, mpol=3) + with pytest.raises(ValueError, match='retain all surface modes'): + surface.to_vmec(filename, ntor=1) From 5c89bd01c05e233f311b01fc068b34178e6ec27f Mon Sep 17 00:00:00 2001 From: Rogerio Jorge Date: Fri, 23 Oct 2026 23:02:16 -0500 Subject: [PATCH 2/3] Correct single-stage QS and transform objectives --- essos/objective_functions.py | 108 +++++++---- tests/test_single_stage_objectives.py | 247 ++++++++++++++++++++++++++ 2 files changed, 316 insertions(+), 39 deletions(-) create mode 100644 tests/test_single_stage_objectives.py diff --git a/essos/objective_functions.py b/essos/objective_functions.py index 69829d6f..25267805 100644 --- a/essos/objective_functions.py +++ b/essos/objective_functions.py @@ -7,7 +7,7 @@ from functools import partial from essos.dynamics import Tracing from essos.fields import BiotSavart,BiotSavart_from_gamma -from essos.surfaces import BdotN_over_B, B_contravariant_theta_phi_on_surface, iota_approx_on_surface +from essos.surfaces import BdotN_over_B, B_contravariant_theta_phi_on_surface, iota_approx_on_surface, _surface_mean from essos.coils import Curves, Coils from essos.constants import mu_0 from essos.coil_perturbation import perturb_curves, perturb_curves_systematic, perturb_curves_statistic @@ -187,6 +187,12 @@ def loss_BdotN_sum(field,surface): def loss_BdotN_mean(field,surface): return jnp.mean(jnp.abs(BdotN_over_B(surface, field))) + +def loss_BdotN_mean_squared(field, surface): + """Mean squared normal field, with periodic endpoint quadrature.""" + return _surface_mean(surface, jnp.square(BdotN_over_B(surface, field))) + + @partial(jit, static_argnames=['target_tol']) def loss_BdotN_constraint(field,surface,target_tol=1.e-6): bdotn_over_b = BdotN_over_B(surface, field) @@ -196,55 +202,75 @@ def loss_BdotN_constraint(field,surface,target_tol=1.e-6): ########################### QUASI-SYMMETRY LOSS ########################### gftd13@gmail.com -def quasi_symmetry_residual_on_surface( field , surface ): - """ - Return a pointwise quasi-symmetry residual on the surface. - Shape is usually (nphi, ntheta) or flattened to (npoints,). +def quasi_symmetry_residual_on_surface(field, surface): + r"""Dimensionless surface triple-product residual, shape ``(nphi, ntheta)``. + + The numerator is ``(n x grad|B|).grad(B.grad|B|)``. On a flux surface, + replacing grad(psi) by its unit direction preserves the QS zero condition + (Rodriguez, Paul & Bhattacharjee, JPP 88, 905880109, 2022, section 2.3). + It does not select a helicity or prove QS when the supplied surface is not + a flux surface, so this objective must accompany a normal-field constraint. + + Multiplication by ``L_B**3 / |B|**3``, with + ``L_B = sqrt(2)|B| / ||grad B||_F``, makes the residual dimensionless. + Evaluate the cancelled expression directly. A relative gradient floor of + 1e-14 times the RMS field divided by the RMS surface extent keeps constant + fields differentiable while preserving field-amplitude and length scaling. + A magnetic null or a collapsed surface is outside this diagnostic's domain. """ - # 1. get surface points - surf_xyz = surface.gamma.reshape(-1, 3) - - # 2. field and gradient of |B| - # BBvec_xyz = jax.vmap(field.B)(surf_xyz) - BBmod_xyz = jax.vmap(field.AbsB)(surf_xyz) - gradB_xyz = jax.vmap(field.dAbsB_by_dX)(surf_xyz) - - # 3. surface normal - unitnormal_xyz = surface.unitnormal.reshape(-1, 3) - - # 4. compute L_B and normalize the quasi-symmetry residual: L_B = sqrt(2) * |B| / ||grad(B)||_F - LB_xyz = jax.vmap(field.L_gradB)(surf_xyz) - norm_factor = LB_xyz**3 / BBmod_xyz**3 - - # 5. quasi-symmetry condition - # 5.1. Compute grad(B · grad|B|) = grad(B · grad|B|) - grad_B_dot_gradB_xyz = jax.vmap(jax.grad(lambda x: jnp.dot(field.B(x), field.dAbsB_by_dX(x))))(surf_xyz) - - # 5.2. Compute (n x grad|B|) · grad(B · grad|B|) and normalized quasi-symmetry residual - QS_residual_xyz = norm_factor * jnp.sum(jnp.cross(unitnormal_xyz, gradB_xyz) * grad_B_dot_gradB_xyz, axis=1) - - return QS_residual_xyz + points = surface.gamma.reshape(-1, 3) + normals = surface.unitnormal.reshape(-1, 3) + grad_magnitude = jax.vmap(field.dAbsB_by_dX)(points) + grad_parallel_derivative = jax.vmap(jax.grad( + lambda x: jnp.dot(field.B(x), field.dAbsB_by_dX(x))))(points) + triple_product = jnp.sum(jnp.cross(normals, grad_magnitude) + * grad_parallel_derivative, axis=-1) + + jacobian = jax.vmap(field.dB_by_dX)(points) + gradient_squared = jnp.sum(jacobian**2, axis=(-2, -1)) + field_squared = jnp.mean(jax.vmap(field.AbsB)(points)**2) + extent_squared = jnp.mean(jnp.sum((points - jnp.mean(points, axis=0))**2, axis=-1)) + denominator = (gradient_squared + 1e-28 * field_squared / extent_squared)**1.5 + residual = 2 * jnp.sqrt(2.0) * triple_product / denominator + return residual.reshape(surface.gamma.shape[:-1]) def loss_quasi_symmetry(field, surface): + """Mean squared normalized triple product, independent of grid-point count. + + This replaces the former sum of absolute residuals. Existing optimization + weights for that objective need to be retuned. """ - Scalar objective: smaller means closer to quasi-symmetry. - """ - QS_residual_xyz = quasi_symmetry_residual_on_surface(field, surface) - QS_residual = jnp.sum(jnp.abs(QS_residual_xyz)) - return QS_residual + residual = quasi_symmetry_residual_on_surface(field, surface) + return _surface_mean(surface, residual**2) -def loss_iota_approx(field, surface, iota_target): +def loss_iota_approx(field, surface, iota_target, *, lambda_mpol=None, lambda_ntor=None): """ Penalize the distance between the surface rotational transform and a target. iota is dimensionless, so no normalization is needed. """ # Rotational transform from the straight-field-line fit on the surface. - iota = iota_approx_on_surface(surface, field) + iota = iota_approx_on_surface(surface, field, lambda_mpol=lambda_mpol, lambda_ntor=lambda_ntor) # Quadratic penalty: smooth, with gradient -> 0 only at the target. return jnp.square(iota - iota_target) + +def loss_straight_field_line_residual(field, surface, *, lambda_mpol=None, + lambda_ntor=None, min_jacobian=0.1): + r"""Penalize the PDE residual and folds of the fitted straight-field angle. + + A target iota alone permits a large, noninvertible coordinate correction. + The squared PDE residual controls fit error, while a squared hinge below + ``min_jacobian`` enforces positive ``1 + d(lambda)/d(theta)`` on the sampled + grid. Use with normal-field and target-iota objectives, and independently + check both diagnostics on a finer grid. This is not a field-line estimate. + """ + diagnostics = iota_approx_on_surface(surface, field, lambda_mpol=lambda_mpol, + lambda_ntor=lambda_ntor, return_diagnostics=True) + deficit = jnp.maximum(min_jacobian - diagnostics["angle_jacobian"], 0.0) + return diagnostics["mean_squared_residual"] + _surface_mean(surface, deficit**2) + def loss_poloidal_slope_min(field, surface, slope_min, B_phi_reference, phi_floor_fraction): """Penalize insufficient positive poloidal slope and small B^phi.""" # Components are ordered (B^theta, B^phi). @@ -407,10 +433,14 @@ def loss_coil_curvature(coils, max_coil_curvature=0): return jnp.mean(pointwise_curvature_loss*jnp.linalg.norm(coils.gamma_dash, axis=-1), axis=1) # Adapter for optimization losses that use "field" as their dependency. -@partial(jit, static_argnames=["max_coil_curvature"]) -def loss_coil_curvature_from_field(field, max_coil_curvature=0): - coil_losses = loss_coil_curvature( field.coils , max_coil_curvature=max_coil_curvature ) - return jnp.mean(coil_losses) +@partial(jit, static_argnames=["max_coil_curvature", "n_segments"]) +def loss_coil_curvature_from_field(field, max_coil_curvature=0, *, n_segments=None): + """Curvature penalty; Fourier coils can use independent curvature sampling.""" + geometry = field.coils + if n_segments is not None: + geometry = field.coils.curves.copy() + geometry.n_segments = n_segments + return jnp.mean(loss_coil_curvature(geometry, max_coil_curvature=max_coil_curvature)) def compute_candidates(coils, min_separation): centers = coils.curves.curves[:, :, 0] diff --git a/tests/test_single_stage_objectives.py b/tests/test_single_stage_objectives.py new file mode 100644 index 00000000..038c7cce --- /dev/null +++ b/tests/test_single_stage_objectives.py @@ -0,0 +1,247 @@ +"""Manufactured fields and derivatives for the active single-stage objectives.""" +import equinox as eqx +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from essos.coils import Coils, Curves +from essos.fields import BiotSavart, MagneticField +from essos.losses import custom_loss +from essos.objective_functions import ( + loss_BdotN_mean_squared, loss_coil_curvature_from_field, loss_coil_length_max, + loss_iota_approx, loss_quasi_symmetry, loss_straight_field_line_residual, + quasi_symmetry_residual_on_surface, +) +from essos.surfaces import (BdotN_over_B, B_contravariant_theta_phi_on_surface, + SurfaceRZFourier, iota_approx_on_surface) + + +def torus(a=.4, *, ntheta=24, nphi=12, close=False): + rc, zs = jnp.array([2., a]), jnp.array([0., a]) + return SurfaceRZFourier(rc, zs, 1, 1, 0, ntheta=ntheta, nphi=nphi, close=close) + + +class PolynomialField(MagneticField, eqx.Module): + """Curl-free, divergence-free Cartesian polynomial field.""" + coefficients: jax.Array + + def __init__(self, coefficients): + self.coefficients = coefficients + + def B(self, x): + a, b, c = self.coefficients + return jnp.array([1 + a*x[0] + b*x[1], + 1 + b*x[0] - a*x[1] + c*x[2], 2 + c*x[1]]) + + +class ToroidalField(MagneticField, eqx.Module): + amplitude: jax.Array + + def __init__(self, amplitude): + self.amplitude = amplitude + + def B(self, x): + return self.amplitude * jnp.array([-x[1], x[0], 0.]) / (x[0]**2+x[1]**2) + + +class StraightAngleField(eqx.Module): + iota: jax.Array + correction: jax.Array + harmonic: int = eqx.field(static=True, default=2) + amplitude: float = 1. + + def B(self, x): + R = jnp.hypot(x[0], x[1]) + er = jnp.array([x[0]/R, x[1]/R, 0.]) + etheta = -x[2]*er + jnp.array([0., 0., R-2.]) + ephi = jnp.array([-x[1], x[0], 0.]) + theta = jnp.arctan2(x[2], R-2.) + dlam = jnp.cos(self.harmonic*theta) * self.harmonic * self.correction + return self.amplitude * (ephi + self.iota/(1+dlam)*etheta) + + +def directional_check(fun, x, *, rtol=2e-5, atol=1e-8, step=2e-5): + direction = jnp.cos(jnp.arange(x.size)+.31).reshape(x.shape) + direction /= jnp.linalg.norm(direction) + value, gradient = jax.jit(jax.value_and_grad(fun))(x) + exact = jnp.vdot(gradient, direction) + finite = (fun(x+step*direction)-fun(x-step*direction))/(2*step) + assert value.shape == () and gradient.shape == x.shape + assert np.isfinite(value) and np.isfinite(gradient).all() + np.testing.assert_allclose(exact, finite, rtol=rtol, atol=atol) + + +def test_qs_matches_direct_triple_product_and_field_scaling(): + surface = torus(ntheta=9, nphi=7) + coefficients = jnp.array([.13, .21, -.17]) + field = PolynomialField(coefficients) + residual = quasi_symmetry_residual_on_surface(field, surface) + points = surface.gamma.reshape(-1, 3) + gradb = jax.vmap(jax.grad(field.AbsB))(points) + gradparallel = jax.vmap(jax.grad(lambda x: field.B(x) @ jax.grad(field.AbsB)(x)))(points) + triple = jnp.sum(jnp.cross(surface.unitnormal.reshape(-1, 3), gradb)*gradparallel, axis=-1) + norm = jax.vmap(lambda x: jnp.linalg.norm(jax.jacfwd(field.B)(x)))(points) + expected = 2*jnp.sqrt(2.) * triple/norm**3 + np.testing.assert_allclose(residual.ravel(), expected, rtol=1e-12, atol=1e-14) + # Scaling all currents by a negative factor reverses the signed residual; + # the objective is unchanged, also at very small absolute field strength. + class Scaled(PolynomialField): + scale: float + def __init__(self, coefficients, scale): + self.coefficients = coefficients + self.scale = scale + def B(self, x): + return self.scale * super().B(x) + for scale in (3., -2., 1e-18): + actual = quasi_symmetry_residual_on_surface(Scaled(coefficients, scale), surface) + np.testing.assert_allclose(actual, np.sign(scale)*residual, rtol=2e-12, atol=1e-13) + directional_check(lambda c: loss_quasi_symmetry(PolynomialField(c), surface), coefficients) + + +def test_qs_axisymmetry_and_uniform_field_have_finite_gradients(): + surface = torus(ntheta=9, nphi=7) + assert float(loss_quasi_symmetry(ToroidalField(jnp.array(1.)), surface)) < 1e-25 + zero = jnp.zeros(3) + assert float(loss_quasi_symmetry(PolynomialField(zero), surface)) == 0 + gradient = jax.grad(lambda c: loss_quasi_symmetry(PolynomialField(c), surface))(zero) + assert np.isfinite(gradient).all() + np.testing.assert_allclose(gradient, 0.) + + +@pytest.mark.parametrize('objective', [loss_quasi_symmetry, loss_BdotN_mean_squared]) +def test_grid_averaging_drops_duplicated_endpoints(objective): + field = PolynomialField(jnp.array([.13, .21, -.17])) + values = [objective(field, torus(ntheta=nt, nphi=np_, close=close)) + for nt, np_, close in [(24, 24, False), (25, 25, True), (48, 48, False)]] + np.testing.assert_allclose(values, values[0], rtol=1e-6) + + +def test_iota_needs_coordinate_resolution_independent_of_shape(): + surface = torus(ntheta=64, nphi=8) + field = StraightAngleField(jnp.array(.4), jnp.array(.3)) + # theta* = theta + .3 sin(2 theta), whose Jacobian stays positive. + # The shape has only m=1. Its default lambda basis gives 0.5, not 0.4. + np.testing.assert_allclose(iota_approx_on_surface(surface, field), .5, rtol=1e-8) + for m in (2, 3, 4): + diagnostics = iota_approx_on_surface(surface, field, lambda_mpol=m, lambda_ntor=0, + return_diagnostics=True) + np.testing.assert_allclose(diagnostics['iota'], .4, atol=1e-9) + assert float(diagnostics['rms_residual']) < 1e-9 + np.testing.assert_allclose(diagnostics['minimum_angle_jacobian'], .4, atol=1e-8) + directional_check(lambda x: iota_approx_on_surface( + surface, StraightAngleField(x[0], x[1]), lambda_mpol=3, lambda_ntor=0), jnp.array([.4, .3])) + assert float(loss_iota_approx(field, surface, .4, lambda_mpol=3, lambda_ntor=0)) < 1e-18 + + +def test_iota_is_independent_of_field_sign_and_closed_grid(): + for close in (False, True): + surface = torus(ntheta=65 if close else 64, nphi=9 if close else 8, close=close) + for amplitude in (1., -3., 1e-12): + field = StraightAngleField(jnp.array(.4), jnp.array(.3), amplitude=amplitude) + np.testing.assert_allclose(iota_approx_on_surface(surface, field, lambda_mpol=3, lambda_ntor=0), + .4, atol=1e-9) + + +def test_actual_biot_savart_coil_and_surface_directional_derivatives(): + # Poloidal loops give a toroidal field with nonzero B^phi everywhere on + # this grid, as required by the straight-field-angle fit. A lone loop + # normal to the radial direction instead has B^phi zeros and reversals. + curve_dofs = jnp.array([[[2., 0., .8], [0., 0., 0.], [0., .8, 0.]]]) + curves = Curves(curve_dofs, n_segments=12, nfp=4, stellsym=False) + field = BiotSavart(Coils(curves, jnp.array([1e5]))) + surface = torus(ntheta=6, nphi=5) + bphi = B_contravariant_theta_phi_on_surface(surface, field)[..., 1] + assert jnp.all(bphi < 0.) + assert jnp.min(jnp.abs(bphi)) > .05*jnp.max(jnp.abs(bphi)) + qs = custom_loss(loss_quasi_symmetry, 'field', 'surface') + geometry = custom_loss(lambda f: loss_coil_length_max(f, 3.) + + loss_coil_curvature_from_field(f, 1.), 'field') + normal = custom_loss(loss_BdotN_mean_squared, 'field', 'surface') + iota = custom_loss(loss_iota_approx, 'field', 'surface', iota_target=.4, + lambda_mpol=2, lambda_ntor=1) + validity = custom_loss(loss_straight_field_line_residual, 'field', 'surface', + lambda_mpol=2, lambda_ntor=1) + combined = qs + .03*geometry + .02*normal + .001*iota + .001*validity + combined.dependencies = {'field': field, 'surface': surface} + x = combined.starting_dofs + directional_check(combined, x, rtol=2e-4, step=1e-6) + np.testing.assert_allclose(combined.grad(x), jax.grad(combined)(x), rtol=2e-10, atol=1e-10) + # Independently verify B.n/|B| current scaling, including sign reversal. + reference = BdotN_over_B(surface, field) + for scale in (-3., 1e-12): + scaled = BiotSavart(Coils(curves, jnp.array([1e5*scale]))) + np.testing.assert_allclose(BdotN_over_B(surface, scaled), np.sign(scale)*reference, rtol=2e-12, atol=1e-14) + np.testing.assert_allclose(loss_BdotN_mean_squared(scaled, surface), + loss_BdotN_mean_squared(field, surface), rtol=2e-12) + + +@pytest.mark.parametrize('active', [False, True]) +def test_individual_coil_penalties_and_gradients_on_circle(active): + def circle_field(radius): + coefficients = jnp.zeros((1, 3, 3)).at[0, 0, 2].set(radius).at[0, 1, 1].set(radius) + return BiotSavart(Coils(Curves(coefficients, n_segments=16, stellsym=False), jnp.array([1e5]))) + radius = .7 + length_target, curvature_target = (3., 1.) if active else (10., 10.) + def penalties(x): + field = circle_field(x[0]) + return jnp.stack((loss_coil_length_max(field, length_target), + loss_coil_curvature_from_field(field, curvature_target))) + expected = np.array([max(2*np.pi*radius-length_target, 0.), + max(1/radius-curvature_target, 0.)**2*2*np.pi*radius]) + np.testing.assert_allclose(penalties(jnp.array([radius])), expected, rtol=1e-12) + jacobian = jax.jacrev(penalties)(jnp.array([radius]))[:, 0] + finite = (penalties(jnp.array([radius+1e-6]))-penalties(jnp.array([radius-1e-6])))/(2e-6) + np.testing.assert_allclose(jacobian, finite, rtol=1e-7, atol=1e-9) + if not active: + np.testing.assert_allclose(jacobian, 0.) + + +def test_straight_angle_validity_loss_accepts_an_invertible_manufactured_field(): + surface = torus(ntheta=24, nphi=6) + def fun(x): + field = StraightAngleField(x[0], x[1]) + return loss_straight_field_line_residual(field, surface, lambda_mpol=3, lambda_ntor=0) + x = jnp.array([.4, .2]) + value, gradient = jax.value_and_grad(fun)(x) + assert float(value) < 1e-18 + assert np.isfinite(gradient).all() + np.testing.assert_allclose(gradient, 0., atol=1e-15) + + +def test_straight_angle_validity_loss_detects_a_folded_fit_and_differentiates(): + # A smooth generic field with substantial normal component cannot in + # general be assigned a valid straight-field angle on this chosen torus. + # Least squares can still report an iota, while its coordinate folds. + surface = torus(ntheta=16, nphi=12) + coefficients = jnp.array([.13, .21, -.17]) + field = PolynomialField(coefficients) + diagnostics = iota_approx_on_surface(surface, field, lambda_mpol=3, lambda_ntor=1, + return_diagnostics=True) + assert float(diagnostics['minimum_angle_jacobian']) < 0. + value = loss_straight_field_line_residual(field, surface, lambda_mpol=3, lambda_ntor=1) + assert float(value) > float(diagnostics['mean_squared_residual']) + directional_check(lambda c: loss_straight_field_line_residual( + PolynomialField(c), surface, lambda_mpol=3, lambda_ntor=1), coefficients, + rtol=5e-5, atol=1e-7, step=1e-6) + + +def test_curvature_sampling_is_independent_of_field_quadrature(): + coefficients = jnp.zeros((1, 3, 5)).at[0, 0, 2].set(2.).at[0, 1, 1].set(.7) + coefficients = coefficients.at[0, 2, 4].set(.2) + + def penalty(c, field_segments, curvature_segments=None): + field = BiotSavart(Coils(Curves(c, n_segments=field_segments, nfp=2, + stellsym=True), jnp.array([1.]), currents_scale=1.)) + value = loss_coil_curvature_from_field(field, 1., n_segments=curvature_segments) + assert field.coils.n_segments == field_segments + return value + + resolved = lambda c: penalty(c, 8, 256) + reference = lambda c: penalty(c, 256) + for a, b in zip(jax.value_and_grad(resolved)(coefficients), + jax.value_and_grad(reference)(coefficients)): + np.testing.assert_allclose(a, b, rtol=1e-12, atol=1e-12) + assert not np.isclose(float(resolved(coefficients)), float(penalty(coefficients, 8))) + directional_check(resolved, coefficients) + From 015b5045b6ed6d7694bea4796a4abaceaf0351ec Mon Sep 17 00:00:00 2001 From: Rogerio Jorge Date: Fri, 23 Oct 2026 23:02:16 -0500 Subject: [PATCH 3/3] Tune the existing single-stage example with global controls --- .../optimize_coils_and_surface_v2.py | 175 +++++++++++------- examples/input_files/coils_QA_mode1.json | 158 ++++++++++++++++ examples/input_files/input.QA_mode1 | 18 ++ 3 files changed, 283 insertions(+), 68 deletions(-) create mode 100644 examples/input_files/coils_QA_mode1.json create mode 100644 examples/input_files/input.QA_mode1 diff --git a/examples/coil_optimization/optimize_coils_and_surface_v2.py b/examples/coil_optimization/optimize_coils_and_surface_v2.py index 26d18fb4..4f228aa5 100644 --- a/examples/coil_optimization/optimize_coils_and_surface_v2.py +++ b/examples/coil_optimization/optimize_coils_and_surface_v2.py @@ -1,6 +1,6 @@ import os from time import time -# import jax # gftd13@gmail.ocm +import jax import jax.numpy as jnp import matplotlib.pyplot as plt @@ -11,12 +11,12 @@ from essos.coils import Coils, CreateEquallySpacedCurves from essos.fields import BiotSavart from essos.surfaces import ( SurfaceRZFourier, BdotN_over_B, B_on_surface, B_contravariant_theta_phi_on_surface, - iota_approx_on_surface, SquaredFlux ) + iota_approx_on_surface, SquaredFlux, toroidal_flux ) from essos.losses import custom_loss -from essos.objective_functions import ( loss_BdotN_mean, loss_coil_curvature_from_field, +from essos.objective_functions import ( loss_BdotN_mean_squared, loss_coil_curvature_from_field, loss_coil_length_max, loss_cross_sectional_area_mean, loss_cross_sectional_area_uniformity, loss_surface_normal_displacement, - loss_quasi_symmetry, loss_iota_approx, loss_surface_curvature_section, + loss_quasi_symmetry, loss_iota_approx, loss_straight_field_line_residual, loss_surface_curvature_section, loss_surface_non_axisymmetric_amplitude, loss_poloidal_slope_min, loss_surface_poloidal_derivative, quasi_symmetry_residual_on_surface, loss_major_radius ) @@ -40,11 +40,12 @@ # Exporting results and the folder name EXPORT = True -output_filepath = os.path.join(os.path.dirname(__file__), "output_coils_surface_opt/test_QA_iota0p4_QSsum10_it00") +output_filepath = os.path.join(os.path.dirname(__file__), "output_coils_surface_opt/mode1") os.makedirs(output_filepath, exist_ok=True) # Make the path # True -> run a fixed-boundary VMEX equilibrium on the exported surface RUN_VMEX = True +VMEC_MPOL = 9; VMEC_NTOR = 9 # solver resolution; None uses boundary maxima + 2 # True -> skip optimization, load saved dofs from output_filepath LOAD_RESULT = False @@ -54,21 +55,25 @@ # True -> start from an already optimized surface and coils, at a higher surface resolution (select the folder). RESTART_FROM = False restart_filepath = os.path.join(os.path.dirname(__file__), "output_coils_surface_opt/test_QA_iota0p4_it00") -# SURFACE_MPOL = 2; SURFACE_NTOR = 3 # maximum m and n of the surface (ESSOS convention) +SURFACE_MPOL = 1; SURFACE_NTOR = 1 # maximum boundary m and |n| +# Continuation: set RESTART_FROM, restart_filepath, output_filepath, and surface modes. +# At mode 3 use QS_WEIGHT=1e4 and IOTA_MPOL=IOTA_NTOR=8. # PATH TO THE VMEC INPUT FILE (wout file) FOR THE SURFACE ---------------------------- input_filepath = os.path.join(os.path.dirname(__file__), "..", "input_files") # vmec_input = os.path.join(input_filepath, 'wout_LandremanPaul2021_QA_reactorScale_lowres.nc') -vmec_input = os.path.join(input_filepath, 'input.toroidal_surface_02') +vmec_input = os.path.join(input_filepath, 'input.QA_mode1') +# Fixed-boundary fitted coils; set None to use equally spaced coils. +INITIAL_COILS = os.path.join(input_filepath, 'coils_QA_mode1.json') # SURFACE PARAMETERS ----------------------------------------------------------------- -ntheta = 22; nphi = 22 ; Npoints = ntheta * nphi +ntheta = 32; nphi = 32 ; Npoints = ntheta * nphi # COILS PARAMETERS ------------------------------------------------------------------- # The coils are initialized as equally spaced curves around a torus. N_COILS = 4; FOURIER_ORDER = 4; LARGE_R = 10; SMALL_R = 5.6; NFP = 2; -N_SEGMENTS = 30; STELLSYM = True +N_SEGMENTS = 64; STELLSYM = True COIL_CURRENT = 1. # Amperes (optimization does not depend on current magnitude) # CONTROL PARAMETERS FOR THE OPTIMIZATION ------------------------------------------- @@ -80,15 +85,16 @@ MAJOR_RADIUS_WEIGHT = 0 # The target will be calculated from the initial surface # Coils ~~~~~~~~~~~~~~~~~~~~ -LENGTH_WEIGHT = 10.; LENGTH_TARGET = 40.; +LENGTH_WEIGHT = 10.; LENGTH_TARGET = 50.; CURVATURE_WEIGHT = 500.; CURVATURE_TARGET = 2. +CURVATURE_SEGMENTS = 256 # Field ~~~~~~~~~~~~~~~~~~~~ -NORMAL_FIELD_WEIGHT = 1e4; -QS_WEIGHT = 1 +NORMAL_FIELD_WEIGHT = 1e6; +QS_WEIGHT = 1e3 # Surface ~~~~~~~~~~~~~~~~~~~~ -CROSS_SECTIONAL_AREA_WEIGHT = 1e5 +CROSS_SECTIONAL_AREA_WEIGHT = 1e6 AREA_UNIFORMITY_WEIGHT = 0. # 1e3 AREA_UNIFORMITY_TOLERANCE = 0.05 # |A(phi)/ - 1| allowed without penalty (~ Amax/Amin <= 1.1) @@ -107,8 +113,10 @@ A3D_WEIGHT = 0. # Rotational transform ~~~~~~~~~~~~~~~~~~~~ -IOTA_APPROX_WEIGHT = 1e4 +IOTA_APPROX_WEIGHT = 1e6 IOTA_APPROX_TARGET = 0.41 +IOTA_MPOL = 6; IOTA_NTOR = 6 +IOTA_FIT_WEIGHT = 1e6 # reject inaccurate or folded straight-angle fits # dtheta/dphi slope penalty SLOPE_MIN = 0.05 @@ -120,7 +128,8 @@ SCALE_SURFACE = 1. # Numerical method parameters -------------------------------------------------------- -MAXITER=250 +MAXITER=1500 +MAXCOR=50 # ==================================================================================== @@ -165,21 +174,22 @@ def diagnostic_gradient(loss_total, dofs, field, surface, scale_field, scale_sur # ------------------------------------------------------------------------------------ # Define the initial surface, coils, and their corresponding field. -if not RESTART_FROM: - # Initial torus from vmec_input and equally spaced coils. - init_curves = CreateEquallySpacedCurves(N_COILS, FOURIER_ORDER, LARGE_R, SMALL_R, n_segments=N_SEGMENTS, nfp=NFP, stellsym=STELLSYM) - init_coils = Coils(curves=init_curves, currents=[COIL_CURRENT]*N_COILS) - surface_init = SurfaceRZFourier.from_input_file(vmec_input, ntheta=ntheta, nphi=nphi, close=True, range_torus='half period') -else: - # The optimized coils of the previous run are the initial coils of this one. +if LOAD_RESULT: + init_coils = Coils.from_json(os.path.join(output_filepath, "init_coils_coils_surface.json")) + surface_file = os.path.join(output_filepath, "input.initial_surface") +elif RESTART_FROM: init_coils = Coils.from_json(os.path.join(restart_filepath, "opt_coils_coils_surface.json")) - - # The previous optimized boundary is read with VMEX and padded with zero modes up to the new resolution - # (VMEC's MPOL counts poloidal modes, so MPOL = mpol + 1). It is saved as the initial surface of this run. - restart_input_file = os.path.join(output_filepath, "input.restart_surface") - vj.VmecInput.from_file(os.path.join(restart_filepath, "input.coils_surface")).change_resolution( - mpol=SURFACE_MPOL + 1, ntor=SURFACE_NTOR).to_indata(restart_input_file) - surface_init = SurfaceRZFourier.from_input_file(restart_input_file, ntheta=ntheta, nphi=nphi, close=True, range_torus='half period') + surface_file = os.path.join(restart_filepath, "input.coils_surface") +else: + init_curves = CreateEquallySpacedCurves(N_COILS, FOURIER_ORDER, LARGE_R, SMALL_R, n_segments=N_SEGMENTS, nfp=NFP, stellsym=STELLSYM) + init_coils = (Coils.from_json(INITIAL_COILS) if INITIAL_COILS else + Coils(curves=init_curves, currents=[COIL_CURRENT]*N_COILS)) + surface_file = vmec_input +init_coils.n_segments = N_SEGMENTS +surface_init = SurfaceRZFourier.from_input_file(surface_file, ntheta=ntheta, nphi=nphi, close=False, + range_torus='half period', mpol=SURFACE_MPOL, ntor=SURFACE_NTOR) +if surface_init.nfp != init_coils.nfp: + raise ValueError("Boundary and coils must have the same NFP") field_init = BiotSavart(init_coils) @@ -189,7 +199,7 @@ def diagnostic_gradient(loss_total, dofs, field, surface, scale_field, scale_sur print(f"\nRESTART: initial coils and surface from '{os.path.basename(restart_filepath)}', " f"surface resolution raised to mpol = {SURFACE_MPOL}, ntor = {SURFACE_NTOR}\n") else: - print(f"\nNEW RUN: initial surface from '{os.path.basename(vmec_input)}' and equally spaced coils\n") + print(f"\nNEW RUN: initial surface from '{os.path.basename(vmec_input)}' and coils from {INITIAL_COILS or 'equally spaced initialization'}\n") # ------------------------------------------------------------------------------------ # Initialize reference values for the surface geometry. These will be used to compute the surface normal displacement loss. @@ -201,9 +211,11 @@ def diagnostic_gradient(loss_total, dofs, field, surface, scale_field, scale_sur kappa_reference = float(jnp.max(surface_init.curvature_section_by_phi())) -MAJOR_RADIUS_TARGET = float(surface_init.rc[0]) - -CROSS_SECTIONAL_AREA_TARGET = float(jnp.mean(surface_init.area_section_by_phi())) +# Retain the original geometry targets when increasing boundary resolution. +surface_reference = SurfaceRZFourier.from_input_file(vmec_input, ntheta=ntheta, nphi=nphi, + close=False, range_torus='half period') +MAJOR_RADIUS_TARGET = float(surface_reference.rc[0]) +CROSS_SECTIONAL_AREA_TARGET = float(jnp.mean(surface_reference.area_section_by_phi())) LENGTH_SCALE_SURFACE = float(jnp.sqrt(CROSS_SECTIONAL_AREA_TARGET / jnp.pi)) kappa_max = ALPHA_KAPPA * kappa_reference @@ -217,9 +229,9 @@ def diagnostic_gradient(loss_total, dofs, field, surface, scale_field, scale_sur # ==================================================================================== # ==================================================================================== -L_normal_field = custom_loss( loss_BdotN_mean, "field", "surface" ) +L_normal_field = custom_loss( loss_BdotN_mean_squared, "field", "surface" ) L_length_max = custom_loss( loss_coil_length_max , "field", max_coil_length=LENGTH_TARGET ) -L_curvature = custom_loss( loss_coil_curvature_from_field , "field" , max_coil_curvature=CURVATURE_TARGET ) +L_curvature = custom_loss( loss_coil_curvature_from_field , "field" , max_coil_curvature=CURVATURE_TARGET, n_segments=CURVATURE_SEGMENTS ) L_quasi_symmetry = custom_loss(loss_quasi_symmetry, "field", "surface") L_cross_sectional_area = custom_loss( loss_cross_sectional_area_mean, "surface", target_area=CROSS_SECTIONAL_AREA_TARGET ) L_area_uniformity = custom_loss( loss_cross_sectional_area_uniformity, "surface", tolerance=AREA_UNIFORMITY_TOLERANCE ) @@ -231,8 +243,10 @@ def diagnostic_gradient(loss_total, dofs, field, surface, scale_field, scale_sur L_surface_non_axisymmetric_amplitude = custom_loss( loss_surface_non_axisymmetric_amplitude, "surface", A3D_min=A3D_MIN ) L_poloidal_slope_min = custom_loss( loss_poloidal_slope_min, "field", "surface", slope_min=SLOPE_MIN, B_phi_reference=B_phi_reference, phi_floor_fraction=PHI_FLOOR_FRACTION ) -L_iota_approx = custom_loss( loss_iota_approx, "field", "surface", iota_target=IOTA_APPROX_TARGET ) -L_major_radius = MAJOR_RADIUS_WEIGHT * custom_loss( loss_major_radius, "surface", target_major_radius=MAJOR_RADIUS_TARGET ) +L_iota_approx = custom_loss( loss_iota_approx, "field", "surface", iota_target=IOTA_APPROX_TARGET, lambda_mpol=IOTA_MPOL, lambda_ntor=IOTA_NTOR ) +L_iota_fit = custom_loss(loss_straight_field_line_residual, "field", "surface", + lambda_mpol=IOTA_MPOL, lambda_ntor=IOTA_NTOR) +L_major_radius = custom_loss( loss_major_radius, "surface", target_major_radius=MAJOR_RADIUS_TARGET ) # ==================================================================================== # ==================================================================================== @@ -247,6 +261,7 @@ def diagnostic_gradient(loss_total, dofs, field, surface, scale_field, scale_sur (AREA_UNIFORMITY_WEIGHT, L_area_uniformity), (NORMAL_DISPLACEMENT_WEIGHT, L_surface_normal_displacement), (IOTA_APPROX_WEIGHT, L_iota_approx), + (IOTA_FIT_WEIGHT, L_iota_fit), (A3D_WEIGHT, L_surface_non_axisymmetric_amplitude), (POLOIDAL_SLOPE_WEIGHT, L_poloidal_slope_min), (QQ_WEIGHT, L_surface_poloidal_derivative), @@ -320,7 +335,7 @@ def loss_callback(intermediate_result): print(f"{0:>6}{cost_previous:>19.4e}{'':>18}{'':>15}{optimality_init:>15.2e}") # SciPy provides the parameters and objective value after this iteration. - dofs_current = jnp.asarray(intermediate_result.x) + dofs_current = jnp.asarray(intermediate_result.x) * x_scale_optimization cost_current = float(intermediate_result.fun) # Compare this accepted point with the previous accepted point. @@ -355,8 +370,14 @@ def loss_callback(intermediate_result): res = SimpleNamespace(x=x_opt, nfev=0) else: - res = minimize(L_total, L_total.starting_dofs, jac=L_total.grad, method="L-BFGS-B", - callback=loss_callback, options={'maxiter': MAXITER}) + value_and_grad = jax.jit(jax.value_and_grad(L_total)) + def scaled_value_and_grad(x): + value, gradient = value_and_grad(jnp.asarray(x) * x_scale_optimization) + return float(value), np.asarray(gradient * x_scale_optimization) + res = minimize(scaled_value_and_grad, np.asarray(L_total.starting_dofs / x_scale_optimization), + jac=True, method="L-BFGS-B", callback=loss_callback, + options={'maxiter': MAXITER, 'maxcor': MAXCOR, 'maxls': 40, 'ftol': 1e-12, 'gtol': 1e-8}) + res.x = np.asarray(res.x * x_scale_optimization) jnp.save(os.path.join(output_filepath, "res_x.npy"), res.x) # res = minimize(L_total, L_total.starting_dofs, jac = L_total.grad, method = "L-BFGS-B", callback=loss_callback , @@ -427,8 +448,8 @@ def loss_callback(intermediate_result): if NORMAL_FIELD_WEIGHT != 0.0: print("\n---------------------------------------------------------------------------") print("Normal-field weighted losses:") - print("NORMAL-FIELD LOSS (INITIAL, WEIGHTED):", NORMAL_FIELD_WEIGHT * loss_BdotN_mean(field_init, surface_init) ) - print("NORMAL-FIELD LOSS (OPTIMIZED, WEIGHTED):",NORMAL_FIELD_WEIGHT * loss_BdotN_mean(field_opt, surface_opt) ) + print("NORMAL-FIELD LOSS (INITIAL, WEIGHTED):", NORMAL_FIELD_WEIGHT * loss_BdotN_mean_squared(field_init, surface_init) ) + print("NORMAL-FIELD LOSS (OPTIMIZED, WEIGHTED):",NORMAL_FIELD_WEIGHT * loss_BdotN_mean_squared(field_opt, surface_opt) ) if QS_WEIGHT != 0.0: print("\n---------------------------------------------------------------------------") @@ -439,8 +460,8 @@ def loss_callback(intermediate_result): if IOTA_APPROX_WEIGHT != 0.0: print("\n---------------------------------------------------------------------------") print("Rotational transform weighted losses:") - print("IOTA LOSS (INITIAL, WEIGHTED):", IOTA_APPROX_WEIGHT * loss_iota_approx(field_init, surface_init,iota_target=IOTA_APPROX_TARGET) ) - print("IOTA LOSS (OPTIMIZED, WEIGHTED):", IOTA_APPROX_WEIGHT * loss_iota_approx(field_opt, surface_opt,iota_target=IOTA_APPROX_TARGET) ) + print("IOTA LOSS (INITIAL, WEIGHTED):", IOTA_APPROX_WEIGHT * loss_iota_approx(field_init, surface_init,iota_target=IOTA_APPROX_TARGET, lambda_mpol=IOTA_MPOL, lambda_ntor=IOTA_NTOR) ) + print("IOTA LOSS (OPTIMIZED, WEIGHTED):", IOTA_APPROX_WEIGHT * loss_iota_approx(field_opt, surface_opt,iota_target=IOTA_APPROX_TARGET, lambda_mpol=IOTA_MPOL, lambda_ntor=IOTA_NTOR) ) if A3D_WEIGHT !=0.0: print("\n---------------------------------------------------------------------------") @@ -500,8 +521,8 @@ def loss_callback(intermediate_result): if CURVATURE_WEIGHT != 0.0: print("\n---------------------------------------------------------------------------") print("Coil-curvature residuals and losses:") - print("COIL CURVATURE LOSS (INITIAL, WEIGHTED):", CURVATURE_WEIGHT * loss_coil_curvature_from_field(field_init, max_coil_curvature=CURVATURE_TARGET)) - print("COIL CURVATURE LOSS (OPTIMIZED, WEIGHTED):", CURVATURE_WEIGHT * loss_coil_curvature_from_field(field_opt, max_coil_curvature=CURVATURE_TARGET)) + print("COIL CURVATURE LOSS (INITIAL, WEIGHTED):", CURVATURE_WEIGHT * loss_coil_curvature_from_field(field_init, max_coil_curvature=CURVATURE_TARGET, n_segments=CURVATURE_SEGMENTS)) + print("COIL CURVATURE LOSS (OPTIMIZED, WEIGHTED):", CURVATURE_WEIGHT * loss_coil_curvature_from_field(field_opt, max_coil_curvature=CURVATURE_TARGET, n_segments=CURVATURE_SEGMENTS)) @@ -518,8 +539,8 @@ def loss_callback(intermediate_result): print("max abs residual (initial):",jnp.max(jnp.abs(B_dot_n_over_B_init)) ) print("max abs residual (optimized):",jnp.max(jnp.abs(B_dot_n_over_B_opt)) ) if NORMAL_FIELD_WEIGHT != 0.0: - print("Normal-field loss (initial):",loss_BdotN_mean(field_init, surface_init) ) - print("Normal-field loss (optimized):",loss_BdotN_mean(field_opt, surface_opt) ) + print("Normal-field loss (initial):",loss_BdotN_mean_squared(field_init, surface_init) ) + print("Normal-field loss (optimized):",loss_BdotN_mean_squared(field_opt, surface_opt) ) print("\n---------------------------------------------------------------------------") @@ -543,8 +564,11 @@ def loss_callback(intermediate_result): print("\n---------------------------------------------------------------------------") print("Rotational transform:") print("iota (target):", IOTA_APPROX_TARGET) -print("iota (initial):", iota_approx_on_surface(surface_init, field_init)) -print("iota (optimized):", iota_approx_on_surface(surface_opt, field_opt)) +print("iota (initial):", iota_approx_on_surface(surface_init, field_init, lambda_mpol=IOTA_MPOL, lambda_ntor=IOTA_NTOR)) +print("iota (optimized):", iota_approx_on_surface(surface_opt, field_opt, lambda_mpol=IOTA_MPOL, lambda_ntor=IOTA_NTOR)) +fit = iota_approx_on_surface(surface_opt, field_opt, lambda_mpol=IOTA_MPOL, lambda_ntor=IOTA_NTOR, return_diagnostics=True) +print("Straight-angle fit RMS and minimum Jacobian:", fit["rms_residual"], fit["minimum_angle_jacobian"]) +print("RMS B.n/B:", jnp.sqrt(jnp.mean(B_dot_n_over_B_opt**2))) print("\n---------------------------------------------------------------------------") print("3D modes amplitude residuals and losses:") @@ -663,8 +687,8 @@ def loss_callback(intermediate_result): if CURVATURE_WEIGHT != 0.0: - print("Curvature loss (initial):", loss_coil_curvature_from_field(field_init, max_coil_curvature=CURVATURE_TARGET)) - print("Curvature loss (optimized):", loss_coil_curvature_from_field(field_opt, max_coil_curvature=CURVATURE_TARGET)) + print("Curvature loss (initial):", loss_coil_curvature_from_field(field_init, max_coil_curvature=CURVATURE_TARGET, n_segments=CURVATURE_SEGMENTS)) + print("Curvature loss (optimized):", loss_coil_curvature_from_field(field_opt, max_coil_curvature=CURVATURE_TARGET, n_segments=CURVATURE_SEGMENTS)) print("\n---------------------------------------------------------------------------") @@ -776,8 +800,10 @@ def section_area(surface, phi): if EXPORT: """ Save the coils to a json file """ - init_coils.to_json(os.path.join(output_filepath, "init_coils_coils_surface.json")) - opt_coils.to_json(os.path.join(output_filepath, "opt_coils_coils_surface.json")) + if not LOAD_RESULT: + init_coils.to_json(os.path.join(output_filepath, "init_coils_coils_surface.json")) + surface_init.to_vmec(os.path.join(output_filepath, "input.initial_surface")) + opt_coils.to_json(os.path.join(output_filepath, "opt_coils_coils_surface.json")) """ Save results in vtk format to analyze in Paraview """ surface_init.to_vtk(os.path.join(output_filepath, "init_surface_coils_surface"), field=field_init) @@ -786,11 +812,21 @@ def section_area(surface, phi): opt_coils.to_vtk(os.path.join(output_filepath, "opt_coils_coils_surface")) """ Save the optimized boundary as a VMEC input file for VMEX """ - # to_vmec writes only the boundary. VMEX reads it back, sets the run controls, and rewrites a complete input: + # Match the vacuum equilibrium flux to the optimized coil field: # NCURR = 1 holds the toroidal current at zero, so iota is computed (default 0 prescribes iota = 0). vmec_input_file = os.path.join(output_filepath, "input.coils_surface") - surface_opt.to_vmec(vmec_input_file) - vmec_inp = replace(vj.VmecInput.from_file(vmec_input_file), ncurr=1, niter_array=[20000]) + surface_opt.to_vmec(vmec_input_file, mpol=VMEC_MPOL, ntor=VMEC_NTOR) + flux_surface = SurfaceRZFourier.from_input_file(vmec_input_file, ntheta=128, nphi=64, + close=False, range_torus="full torus") + signed_area = jnp.sum(jnp.linalg.norm(flux_surface.gamma[..., :2], axis=-1) + * flux_surface.gammadash_theta[..., 2], axis=1) + fluxes = jnp.array([-jnp.sign(signed_area[k]) * toroidal_flux(flux_surface, field_opt, k) + for k in range(flux_surface.nphi)]) + phiedge = float(jnp.mean(fluxes)) + if not np.isfinite(phiedge) or phiedge == 0 or bool(jnp.any(signed_area == 0)): + raise ValueError("VMEX requires finite nonzero toroidal flux and valid boundary sections") + print("Toroidal flux (Wb), relative spread:", phiedge, float(jnp.ptp(fluxes) / abs(phiedge))) + vmec_inp = replace(vj.VmecInput.from_file(vmec_input_file), ncurr=1, phiedge=phiedge) vmec_inp.to_indata(vmec_input_file) """ Fixed-boundary VMEX equilibrium of the optimized surface """ @@ -800,6 +836,8 @@ def section_area(surface, phi): os.makedirs(vmex_output_filepath, exist_ok=True) vmec_result = vj.solve_multigrid(vmec_inp, verbose=False) + if not bool(vmec_result.converged): + raise RuntimeError("VMEX did not converge; increase solver resolution or iteration limit") # The equilibrium is saved as a wout file, and VMEX's diagnostic figures are written next to it. wout = vj.wout_from_result(vmec_inp, vmec_result) @@ -830,26 +868,27 @@ def section_area(surface, phi): # The input parameters are grouped as in the INPUT DATA section and written as "NAME = value" lines, # so the file can later be read back to set up a run. RUN_PARAMETERS = { - "Run control": ["EXPORT", "RUN_VMEX", "LOAD_RESULT", "MAXITER"], + "Run control": ["EXPORT", "RUN_VMEX", "VMEC_MPOL", "VMEC_NTOR", "LOAD_RESULT", "MAXITER", "MAXCOR"], "Continuation": ["RESTART_FROM", "restart_filepath", "SURFACE_MPOL", "SURFACE_NTOR"], "Initial surface": ["ntheta", "nphi"], "Coils": ["N_COILS", "FOURIER_ORDER", "LARGE_R", "SMALL_R", "NFP", "N_SEGMENTS", - "STELLSYM", "COIL_CURRENT"], + "STELLSYM", "COIL_CURRENT", "INITIAL_COILS"], "Geometry losses": ["MAJOR_RADIUS_WEIGHT"], - "Coil losses": ["LENGTH_WEIGHT", "LENGTH_TARGET", "CURVATURE_WEIGHT", "CURVATURE_TARGET"], + "Coil losses": ["LENGTH_WEIGHT", "LENGTH_TARGET", "CURVATURE_WEIGHT", "CURVATURE_TARGET", "CURVATURE_SEGMENTS"], "Field losses": ["NORMAL_FIELD_WEIGHT", "QS_WEIGHT"], "Surface losses": ["CROSS_SECTIONAL_AREA_WEIGHT", "AREA_UNIFORMITY_WEIGHT", "AREA_UNIFORMITY_TOLERANCE", "NORMAL_DISPLACEMENT_WEIGHT", "QQ_WEIGHT", "ALPHA_QQ", "KAPPA_WEIGHT", "ALPHA_KAPPA", "A3D_MIN", "A3D_WEIGHT"], - "Rotational transform": ["IOTA_APPROX_WEIGHT", "IOTA_APPROX_TARGET", + "Rotational transform": ["IOTA_APPROX_WEIGHT", "IOTA_APPROX_TARGET", "IOTA_MPOL", "IOTA_NTOR", "IOTA_FIT_WEIGHT", "SLOPE_MIN", "PHI_FLOOR_FRACTION", "POLOIDAL_SLOPE_WEIGHT"], "x_scale": ["SCALE_FIELD", "SCALE_SURFACE"], } -with open(os.path.join(output_filepath, "essos_surf_coils.in"), "w") as f: - f.write(f"# ESSOS single-stage run parameters ({datetime.datetime.now():%Y-%m-%d %H:%M})\n") - f.write(f"VMEC_INPUT = {os.path.basename(vmec_input)!r}\n") - for section, names in RUN_PARAMETERS.items(): - f.write(f"\n# {section}\n") - for name in names: - f.write(f"{name} = {globals()[name]!r}\n") \ No newline at end of file +if not LOAD_RESULT: + with open(os.path.join(output_filepath, "essos_surf_coils.in"), "w") as f: + f.write(f"# ESSOS single-stage run parameters ({datetime.datetime.now():%Y-%m-%d %H:%M})\n") + f.write(f"VMEC_INPUT = {os.path.basename(vmec_input)!r}\n") + for section, names in RUN_PARAMETERS.items(): + f.write(f"\n# {section}\n") + for name in names: + f.write(f"{name} = {globals()[name]!r}\n") \ No newline at end of file diff --git a/examples/input_files/coils_QA_mode1.json b/examples/input_files/coils_QA_mode1.json new file mode 100644 index 00000000..d1db3d84 --- /dev/null +++ b/examples/input_files/coils_QA_mode1.json @@ -0,0 +1,158 @@ +{ + "nfp": 2, + "stellsym": true, + "order": 4, + "n_segments": 64, + "dofs_curves_raw": [ + [ + [ + 12.165997732130961, + 0.2833595929830598, + 4.925694711727081, + 0.6754372679030735, + 0.5166146512195183, + 0.2499255201007313, + 0.05678029712653307, + -0.04652912417676418, + 0.058713669195272725 + ], + [ + 2.1824744792379214, + 1.5949954731305909, + 1.389334681317779, + -1.4527630469653598, + -0.6396748696588712, + -0.4871525454739271, + -0.33799192488029123, + -0.06027803092857291, + -0.035167898357700206 + ], + [ + 0.5671825831262702, + -8.072417667591626, + -0.506157171799504, + -0.8238726730338162, + -0.02329543138210656, + -0.33991299294274024, + -0.15981852457772242, + 0.006863201665878661, + 0.02702352598214967 + ] + ], + [ + [ + 9.652476664244022, + 0.5300098985847638, + 3.9091096747804133, + 0.8763910387311621, + 1.8936842244320284, + 0.12983244886713716, + 0.8086279554379214, + -0.03562821763267742, + 0.03677849977535758 + ], + [ + 5.790240142975378, + 1.44744030029469, + 4.066609589130031, + -0.006614083273075611, + -0.7034851729150142, + 0.2066013801977743, + -0.41901084342964784, + 0.04690731771770946, + -0.13706096986310304 + ], + [ + 1.5696756417034128, + -7.146067747508914, + -1.327940080769105, + -0.7865381169560787, + 0.08149759328415468, + -0.10951424342304536, + -0.19314245730997087, + -0.028770775543661514, + -0.0017305552221522284 + ] + ], + [ + [ + 5.681048343221239, + 1.1276086563910093, + 2.360800573440978, + -0.811216592258916, + 2.3957929312971262, + -0.9008754733038813, + 0.39246976814386403, + -0.09831425158943331, + -0.02508279715523907 + ], + [ + 7.673927782039761, + 0.7243150775704081, + 5.896652481547694, + 0.8500505785804349, + 0.6334400422331147, + -0.043129308252194354, + 0.004485019027339625, + 0.06497936538903008, + -0.07633959578895576 + ], + [ + 1.9467568866829998, + -5.331586445016807, + -2.14775475334169, + -1.0456563920165056, + 0.6788104276285729, + -0.09572827973222013, + -0.03058045242175043, + 0.11467418931047524, + -0.16839558414869574 + ] + ], + [ + [ + 1.3386194084863965, + 1.5450072061448048, + 1.2311332745944106, + -1.986382546041568, + 0.4848077701465219, + -0.11606250449236098, + 0.40359995330283427, + -0.385428237753316, + 0.09991450919292977 + ], + [ + 8.271292661387895, + 0.4725332671707741, + 6.6193422088127045, + 0.8042197629208929, + 0.8926730257965785, + -0.22604755943082253, + 0.08903608450043489, + -0.14560195284505642, + 0.3073119983715116 + ], + [ + 0.9989411749195406, + -3.984581009311121, + -0.6559060515497677, + -2.3196453755499506, + 0.7801313051819453, + 0.2467268133021224, + 0.09382108403949402, + -0.48047048679860926, + -0.19251467388380278 + ] + ] + ], + "scaling_type": 2, + "scaling_factor": 0.0, + "scale_fixed": 1.0, + "dofs_currents_raw": [ + 1.1560515202831294, + 1.1606695044806632, + 1.2106895248995644, + 1.7419860562443727 + ], + "currents_scale": 1.0 +} \ No newline at end of file diff --git a/examples/input_files/input.QA_mode1 b/examples/input_files/input.QA_mode1 new file mode 100644 index 00000000..ebaa7cfa --- /dev/null +++ b/examples/input_files/input.QA_mode1 @@ -0,0 +1,18 @@ +&INDATA +LASYM = .FALSE. +NFP = 2 +MPOL = 3 +NTOR = 3 +NCURR = 1 +CURTOR = 0.0 +AC = 0.0 +AM = 0.0 +FTOL_ARRAY = 1e-12 +NS_ARRAY = 51 +NITER_ARRAY = 3500 +RBC(0,0) = 1.000000000000000e+01, ZBS(0,0) = 0.000000000000000e+00 +RBC(1,0) = 1.635736469508825e+00, ZBS(1,0) = -1.253658594587032e+00 +RBC(-1,1) = 4.397572638097199e-01, ZBS(-1,1) = 3.995062357599209e-01 +RBC(0,1) = 1.973645774447548e+00, ZBS(0,1) = 2.562798679477748e+00 +RBC(1,1) = -1.295463562917905e+00, ZBS(1,1) = 9.523560437661636e-01 +/