Separate boundary modes from VMEC solver resolution

essos/surfaces.py

@@ -259,32 +259,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."""
@@ -809,12 +845,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 boundary maxima plus two. Since MPOL
+        is a count, this retains m through boundary m + 1 and |n| through
+        boundary |n| + 2. Override these independently and check convergence.
+        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')
@@ -823,7 +874,7 @@ def to_vmec(self, filename):
                                    for name, table in families.items() if table is not None))
         with open(filename, 'w') as f:
             f.write('\n'.join(lines) + '\n/\n')
-            
+
     def mean_cross_sectional_area(self):
         """Mean over phi of the constant-phi cross-section area ``int R dZ/dtheta dtheta``.
 

tests/test_fields.py

@@ -595,7 +595,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)
@@ -617,11 +617,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])
@@ -659,7 +659,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)

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)