From 89857f546f6cd39697d02b321ff9697740de5a0b Mon Sep 17 00:00:00 2001 From: Jonathan Schilling Date: Wed, 17 Jul 2024 09:41:31 +0200 Subject: [PATCH] fix mpol inconsistency between VMEC and SurfaceRZFourier --- src/simsopt/geo/surfacerzfourier.py | 39 +++++++++++++++++------------ src/simsopt/mhd/vmec.py | 15 ++++++++--- 2 files changed, 34 insertions(+), 20 deletions(-) diff --git a/src/simsopt/geo/surfacerzfourier.py b/src/simsopt/geo/surfacerzfourier.py index f66a5c304..3eb3116af 100644 --- a/src/simsopt/geo/surfacerzfourier.py +++ b/src/simsopt/geo/surfacerzfourier.py @@ -35,6 +35,10 @@ class SurfaceRZFourier(sopp.SurfaceRZFourier, Surface): and the same for :math:`z(\theta, \phi)`. + Note that the `mpol` in this is different than the `mpol` in VMEC: + VMEC uses poloidal mode numbers `m=0, 1, ..., (mpol-1)`, + but here we use `m=0, 1, ..., mpol`. + Here, :math:`(r,\phi, z)` are standard cylindrical coordinates, and theta is any poloidal angle. @@ -190,7 +194,10 @@ def from_wout(cls, filename: str, s: float = 1.0, interp = interp1d(s_full_grid, zmnc, kind=interp_kind, axis=0) zbc = interp(s) - mpol = int(np.max(xm)) + # VMEC only considers poloidal modes up to `m=(mpol-1)`, + # but `SurfaceRZFourier` includes `m=mpol`. + # Here, `xm` comes from VMEC and `mpol` goes into `SurfaceRZFourier`. + mpol = int(np.max(xm)) - 1 ntor = int(np.max(np.abs(xn)) / nfp) ntheta = kwargs.pop("ntheta", None) @@ -640,31 +647,31 @@ def write_nml(self, filename: str): def fourier_transform_scalar(self, scalar, mpol=None, ntor=None, normalization=None, **kwargs): r""" Compute the Fourier components of a scalar on the surface. The scalar - is evaluated at the quadrature points on the surface. - The Fourier uses the conventions of the FourierRZSurface series, + is evaluated at the quadrature points on the surface. + The Fourier uses the conventions of the FourierRZSurface series, with `npol` going from `-ntor` to `ntor` and `mpol` from 0 to `mpol` - i.e.: + i.e.: :math:`f(\theta, \phi) = \Sum_{m=0}^{mpol} \Sum_{n=-npol}^{npol} A^{mn}_s \sin(m\theta - n*Nfp*\phi)\\ + A^{mn}_c \cos(m\theta - n*Nfp*\phi)` Where the cosine series is only evaluated if the surface is not stellarator - symmetric (if the scalar does not adhere to the symmetry of the surface, + symmetric (if the scalar does not adhere to the symmetry of the surface, request the cosine series by setting the kwarg stellsym=False) - By default, the poloidal and toroidal resolution are the same as those of the surface, but different quantities can be specified in the kwargs. + By default, the poloidal and toroidal resolution are the same as those of the surface, but different quantities can be specified in the kwargs. *Arguments*: - scalar: 2D array of shape (numquadpoints_phi, numquadpoints_theta). - mpol: maximum poloidal mode number of the transform, if None, the mpol attribute of the surface is used. - - ntor: maximum toroidal mode number of the transform if None, + - ntor: maximum toroidal mode number of the transform if None, the ntor attribute of the surface is used. *Optional keyword arguments*: - - normalization: Fourier transform normalization. Can be: + - normalization: Fourier transform normalization. Can be: None: forward and back transform are not normalized float: forward transform is divided by this number - - stellsym: boolean to override the stellsym attribute + - stellsym: boolean to override the stellsym attribute of the surface if you want to force the calculation of the cosine series *Returns*: - A_mns: 2D array of shape (mpol+1, 2*ntor+1) containing the sine coefficients - - A_mnc: 2D array of shape (mpol+1, 2*ntor+1) containing the cosine coefficients + - A_mnc: 2D array of shape (mpol+1, 2*ntor+1) containing the cosine coefficients (these are zero if the surface is stellarator symmetric) """ assert scalar.shape[0] == self.quadpoints_phi.size, "scalar must be evaluated at the quadrature points on the surface.\n the scalar you passed in has shape {}".format(scalar.shape) @@ -684,9 +691,9 @@ def fourier_transform_scalar(self, scalar, mpol=None, ntor=None, normalization=N factor = 2.0 / (ntheta_grid * nphi_grid) phi2d, theta2d = np.meshgrid(2 * np.pi * self.quadpoints_phi, - 2 * np.pi * self.quadpoints_theta, + 2 * np.pi * self.quadpoints_theta, indexing="ij") - + for m in range(mpol + 1): nmin = -ntor if m == 0: nmin = 1 @@ -701,7 +708,7 @@ def fourier_transform_scalar(self, scalar, mpol=None, ntor=None, normalization=N if not stellsym: cosangle = np.cos(angle) A_mnc[m, n + ntor] = np.sum(scalar * cosangle * factor2) - + if not stellsym: A_mnc[0, ntor] = np.sum(scalar) / (ntheta_grid * nphi_grid) if normalization is not None: @@ -723,7 +730,7 @@ def inverse_fourier_transform_scalar(self, A_mns, A_mnc, normalization=None, **k symmetric. *Arguments*: - A_mns: 2D array of shape (mpol+1, 2*ntor+1) containing the sine coefficients - - A_mnc: 2D array of shape (mpol+1, 2*ntor+1) containing the cosine coefficients + - A_mnc: 2D array of shape (mpol+1, 2*ntor+1) containing the cosine coefficients (these are zero if the surface is stellarator symmetric) *Optional keyword arguments*: - normalization: Fourier transform normalization. Can be: @@ -752,10 +759,10 @@ def inverse_fourier_transform_scalar(self, A_mns, A_mnc, normalization=None, **k if not stellsym: cosangle = np.cos(angle) scalars = scalars + A_mnc[m, n + ntor] * cosangle - + if not stellsym: scalars = scalars + A_mnc[0, ntor] - if normalization is not None: + if normalization is not None: if not isinstance(normalization, float): raise ValueError("normalization must be a float") scalars = scalars * normalization diff --git a/src/simsopt/mhd/vmec.py b/src/simsopt/mhd/vmec.py index 507c8c9a7..85061905a 100644 --- a/src/simsopt/mhd/vmec.py +++ b/src/simsopt/mhd/vmec.py @@ -344,9 +344,13 @@ def __init__(self, # object, but the mpol/ntor values of either the vmec object # or the boundary surface object can be changed independently # by the user. + # `vi.mpol` comes from VMEC, but VMEC only uses `m=0, 1, ..., (mpol-1)`, + # but `SurfaceRZFourier` includes `m=mpol`, so we have to + # initialize `SurfaceRZFourier` with the highest poloidal mode + # that VMEC uses, which is `vi.mpol-1`. self._boundary = SurfaceRZFourier.from_nphi_ntheta(nfp=vi.nfp, stellsym=not vi.lasym, - mpol=vi.mpol, + mpol=vi.mpol - 1, ntor=vi.ntor, ntheta=ntheta, nphi=nphi, @@ -354,7 +358,8 @@ def __init__(self, self.free_boundary = bool(vi.lfreeb) # Transfer boundary shape data from fortran to the ParameterArray: - for m in range(vi.mpol + 1): + # The highest mode number relevant to VMEC in its input is `vi.mpol-1`. + for m in range(vi.mpol): for n in range(-vi.ntor, vi.ntor + 1): self._boundary.rc[m, n + vi.ntor] = vi.rbc[101 + n, m] self._boundary.zs[m, n + vi.ntor] = vi.zbs[101 + n, m] @@ -523,7 +528,9 @@ def set_indata(self): mpol_capped = np.min([boundary_RZFourier.mpol, 101]) ntor_capped = np.min([boundary_RZFourier.ntor, 101]) # Transfer boundary shape data from the surface object to VMEC: - for m in range(mpol_capped + 1): + # The highest mode number that VMEC can use is `m=100` + # if the max resolution is `mpol=101`. + for m in range(mpol_capped): for n in range(-ntor_capped, ntor_capped + 1): vi.rbc[101 + n, m] = boundary_RZFourier.get_rc(m, n) vi.zbs[101 + n, m] = boundary_RZFourier.get_zs(m, n) @@ -756,7 +763,7 @@ def run(self): os.remove(filename) except FileNotFoundError: logger.debug(f"Tried to delete the file {filename} but it was not found") - + self.files_to_delete = [] # Record the latest output file to delete if we run again: