diff --git a/essos/surfaces.py b/essos/surfaces.py index 28cdd687..0267bf7a 100644 --- a/essos/surfaces.py +++ b/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``. diff --git a/tests/test_fields.py b/tests/test_fields.py index dfdcccde..956e4328 100644 --- a/tests/test_fields.py +++ b/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) 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)