diff --git a/CMakeLists.txt b/CMakeLists.txt index f4fe29b90..68858e9ed 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -56,11 +56,11 @@ else() unset(COMPILER_SUPPORTS_MARCH_NATIVE CACHE) CHECK_CXX_COMPILER_FLAG(-march=native COMPILER_SUPPORTS_MARCH_NATIVE) if(COMPILER_SUPPORTS_MARCH_NATIVE) - set(CMAKE_CXX_FLAGS "-O3 -march=native -mfma -ffp-contract=fast") - elseif(${CMAKE_HOST_SYSTEM_PROCESSOR} STREQUAL "arm64") - set(CMAKE_CXX_FLAGS "-O3 -mcpu=apple-a14 -mfma -ffp-contract=fast") + set(CMAKE_CXX_FLAGS "-O3 -march=native -ffp-contract=fast") + elseif(${CMAKE_HOST_SYSTEM_PROCESSOR} STREQUAL "arm64") + set(CMAKE_CXX_FLAGS "-O3 -mcpu=apple-a14 -ffp-contract=fast") else() - set(CMAKE_CXX_FLAGS "-O3 -ffp-contract=fast") + set(CMAKE_CXX_FLAGS "-O3 -ffp-contract=fast") endif() endif() # Ensure all code used from Eigen does not have LGPL license: @@ -129,7 +129,7 @@ pybind11_add_module(${PROJECT_NAME} src/simsoptpp/integral_BdotN.cpp src/simsoptpp/psc.cpp src/simsoptpp/dipole_field.cpp src/simsoptpp/permanent_magnet_optimization.cpp - src/simsoptpp/dommaschk.cpp src/simsoptpp/reiman.cpp src/simsoptpp/tracing.cpp + src/simsoptpp/dommaschk.cpp src/simsoptpp/reiman.cpp src/simsoptpp/tracing.cpp src/simsoptpp/magneticfield_biotsavart.cpp src/simsoptpp/python_boozermagneticfield.cpp src/simsoptpp/boozerradialinterpolant.cpp ) diff --git a/src/simsopt/field/normal_field.py b/src/simsopt/field/normal_field.py index 8c14d3398..68f09cd3e 100644 --- a/src/simsopt/field/normal_field.py +++ b/src/simsopt/field/normal_field.py @@ -3,6 +3,7 @@ import numpy as np from .._core.optimizable import DOFs, Optimizable +from simsopt.geo import SurfaceRZFourier logger = logging.getLogger(__name__) @@ -20,6 +21,15 @@ class NormalField(Optimizable): ``NormalField`` represents the magnetic field normal to a toroidal surface, for example the computational boundary of SPEC free-boundary. + It consists a surface (the computational boundary), and a set of Fourier harmonics that describe + the normal component of an externally provided field. + + The Fourier harmonics are the degees of freedom, the computational boundary is kept fixed. + The Fourier harmonics are stored in the SPEC convention, + i.e. it is the + fourier components of B.(\partial\vec{r}/ \partial\theta \times \partial\vec{r}/ \partial\zeta) on the surface with a 1/(2\pi)^2 + Fourier normalization factor. + Args: nfp: The number of field period stellsym: Whether (=True) or not (=False) stellarator symmetry is enforced. @@ -37,7 +47,7 @@ class NormalField(Optimizable): """ def __init__(self, nfp=1, stellsym=True, mpol=1, ntor=0, - vns=None, vnc=None): + vns=None, vnc=None, surface=None): self.nfp = nfp self.stellsym = stellsym @@ -47,39 +57,44 @@ def __init__(self, nfp=1, stellsym=True, mpol=1, ntor=0, if vns is None: vns = np.zeros((self.mpol + 1, 2 * self.ntor + 1)) - if vnc is None and not stellsym: + if not self.stellsym and vnc is None: vnc = np.zeros((self.mpol + 1, 2 * self.ntor + 1)) + + if surface is None: + surface = SurfaceRZFourier(nfp=nfp, stellsym=stellsym, mpol=mpol, ntor=ntor) + surface.fix_all() + self.surface = surface if self.stellsym: self.ndof = self.ntor + self.mpol * (2 * self.ntor + 1) else: self.ndof = 2 * (self.ntor + self.mpol * (2 * self.ntor + 1)) + 1 - - # Pack in a single array - dofs = np.zeros((self.ndof,)) - - # Populate dofs array - vns_shape = vns.shape - input_mpol = int(vns_shape[0]-1) - input_ntor = int((vns_shape[1]-1)/2) - for mm in range(0, self.mpol+1): - for nn in range(-self.ntor, self.ntor+1): - if mm == 0 and nn < 0: continue - if mm > input_mpol: continue - if nn > input_ntor: continue - - if not (mm == 0 and nn == 0): - ii = self.get_index_in_dofs(mm, nn, even=False) - dofs[ii] = vns[mm, input_ntor+nn] - - if not self.stellsym: - ii = self.get_index_in_dofs(mm, nn, even=True) - dofs[ii] = vnc[mm, input_ntor+nn] + + self._vns = vns + self._vnc = vnc + + dofs = self.get_dofs() Optimizable.__init__( self, x0=dofs, names=self._make_names()) + + @property + def vns(self): + return self._vns + + @vns.setter + def vns(self, value): + raise AttributeError('Change Vns using set_vns() or set_vns_asarray()') + + @property + def vnc(self): + return self._vnc + + @vnc.setter + def vnc(self, value): + raise AttributeError('Change Vnc using set_vnc() or set_vnc_asarray()') @classmethod def from_spec(cls, filename): @@ -102,6 +117,15 @@ def from_spec(cls, filename): vnc = None else: vnc = np.asarray(ph['vnc']) + mpol = ph['Mpol'] + ntor = ph['Ntor'] + surface = SurfaceRZFourier(nfp=ph['nfp'], stellsym=bool(ph['istellsym']), mpol=mpol, ntor=ntor) + old_ntor = np.array(ph['rbc']).shape[1]//2 + surface.rc[:] = np.array(ph['rbc'])[:mpol+1, old_ntor-ntor:old_ntor+ntor+1] + surface.zs[:] = np.array(ph['zbs'])[:mpol+1, old_ntor-ntor:old_ntor+ntor+1] + if not ph['istellsym']: + surface.zc[:] = np.array(ph['zbc'])[:mpol+1, old_ntor-ntor:old_ntor+ntor+1] + surface.rs[:] = np.array(ph['rbs'])[:mpol+1, old_ntor-ntor:old_ntor+ntor+1] normal_field = cls( nfp=ph['nfp'], @@ -109,11 +133,91 @@ def from_spec(cls, filename): mpol=ph['Mpol'], ntor=ph['Ntor'], vns=vns, - vnc=vnc + vnc=vnc, + surface=surface ) return normal_field + @classmethod + def from_spec_object(cls, spec): + """ + Initialize using the simsopt SPEC object's attributes + """ + if not spec.freebound: + raise ValueError('The given SPEC object is not free-boundary') + + surface = spec.computational_boundary + + # Grab all the attributes from the SPEC object into an input dictionary + input_dict = {'nfp': spec.nfp, + 'stellsym': spec.stellsym, + 'mpol': spec.mpol, + 'ntor': spec.ntor, + 'vns': spec.array_translator(spec.inputlist.vns).as_simsopt, + 'surface': surface} + if not spec.stellsym: + input_dict.append({'vnc': spec.array_translator(spec.inputlist.vnc).as_simsopt}) + + normal_field = cls(**input_dict) + + return normal_field + + def get_dofs(self): + """ + get DOFs from vns and vnc + """ + # Pack in a single array + dofs = np.zeros((self.ndof,)) + + # Populate dofs array + vns_shape = self.vns.shape + input_mpol = int(vns_shape[0]-1) + input_ntor = int((vns_shape[1]-1)/2) + for mm in range(0, self.mpol+1): + for nn in range(-self.ntor, self.ntor+1): + if mm == 0 and nn < 0: continue + if mm > input_mpol: continue + if nn > input_ntor: continue + + if not (mm == 0 and nn == 0): + ii = self.get_index_in_dofs(mm, nn, even=False) + dofs[ii] = self.vns[mm, input_ntor+nn] + + if not self.stellsym: + ii = self.get_index_in_dofs(mm, nn, even=True) + dofs[ii] = self.vnc[mm, input_ntor+nn] + return dofs + + def get_index_in_array(self, m, n, mpol=None, ntor=None, even=False): + """ + Returns position of mode (m,n) in vns or vnc array + + Args: + - m: poloidal mode number + - n: toroidal mode number (normalized by Nfp) + - mpol: resolution of dofs array. If None (by default), use self.mpol + - ntor: resolution of dofs array. If None (by default), use self.ntor + - even: set to True to get vnc. Default is False + """ + + if mpol is None: + mpol = self.mpol + if ntor is None: + ntor = self.ntor + + if m < 0 or m > mpol: + raise ValueError('m out of bound') + if abs(n) > ntor: + raise ValueError('n out of bound') + if m == 0 and n < 0: + raise ValueError('n has to be positive if m==0') + if not even and m == 0 and n == 0: + raise ValueError('m=n=0 not supported for odd series') + + return [m, n+ntor] + + def get_index_in_dofs(self, m, n, mpol=None, ntor=None, even=False): """ Returns position of mode (m,n) in dofs array @@ -161,8 +265,10 @@ def get_vns(self, m, n): def set_vns(self, m, n, value): self.check_mn(m, n) - ii = self.get_index_in_dofs(m, n) - self.local_full_x[ii] = value + i,j = self.get_index_in_array(m, n) + self._vns[i,j] = value + dofs = self.get_dofs() + self.local_full_x = dofs def get_vnc(self, m, n): self.check_mn(m, n) @@ -174,11 +280,13 @@ def get_vnc(self, m, n): def set_vnc(self, m, n, value): self.check_mn(m, n) - ii = self.get_index_in_dofs(m, n, even=True) + i,j = self.get_index_in_array(m, n) if self.stellsym: raise ValueError('Stellarator symmetric has no vnc') else: - self.local_full_x[ii] = value + self._vnc[i,j] = value + dofs = self.get_dofs() + self.local_full_x = dofs def check_mn(self, m, n): if m < 0 or m > self.mpol: @@ -193,10 +301,10 @@ def _make_names(self): Form a list of names of the ``vns``, ``vnc`` """ if self.stellsym: - names = self._make_names_helper(False) + names = self._make_names_helper(even=False) else: - names = np.append(self._make_names_helper(False), - self._make_names_helper(True)) + names = np.append(self._make_names_helper(even=False), + self._make_names_helper(even=True)) return names @@ -287,3 +395,122 @@ def fixed_range(self, mmin, mmax, nmin, nmax, fixed=True): fn(f'vns({m},{n})') if not self.stellsym: fn(f'vnc({m},{n})') + + def get_vns_asarray(self, mpol=None, ntor=None): + """ + Return the vns as a single array + """ + if mpol is None: + mpol = self.mpol + elif mpol > self.mpol: + raise ValueError('mpol out of bound') + + if ntor is None: + ntor = self.ntor + elif ntor > self.ntor: + raise ValueError('ntor out of bound') + + vns = self.vns + + return vns[0:mpol, self.ntor-ntor:self.ntor+ntor+1] + + def get_vnc_asarray(self, mpol=None, ntor=None): + """ + Return the vnc as a single array + """ + if mpol is None: + mpol = self.mpol + elif mpol > self.mpol: + raise ValueError('mpol out of bound') + + if ntor is None: + ntor = self.ntor + elif ntor > self.ntor: + raise ValueError('ntor out of bound') + + vnc = self.vns + + return vnc[0:mpol, self.ntor-ntor:self.ntor+ntor+1] + + def get_vns_vnc_asarray(self, mpol=None, ntor=None): + """ + Return the vns and vnc as two arrays single array + """ + if mpol is None: + mpol = self.mpol + elif mpol > self.mpol: + raise ValueError('mpol out of bound') + + if ntor is None: + ntor = self.ntor + elif ntor > self.ntor: + raise ValueError('ntor out of bound') + + vns = self.get_vns_asarray(mpol, ntor) + vnc = self.get_vnc_asarray(mpol, ntor) + return vns, vnc + + def set_vns_asarray(self, vns, mpol=None, ntor=None): + """ + Set the vns from a single array + """ + if mpol is None: + mpol = self.mpol + elif mpol > self.mpol: + raise ValueError('mpol out of bound') + + if ntor is None: + ntor = self.ntor + elif ntor > self.ntor: + raise ValueError('ntor out of bound') + + self._vns = vns[0:mpol, self.ntor-ntor:self.ntor+ntor+1] + dofs = self.get_dofs() + self.local_full_x = dofs + + def set_vnc_asarray(self, vnc, mpol=None, ntor=None): + """ + Set the vnc from a single array + """ + if mpol is None: + mpol = self.mpol + elif mpol > self.mpol: + raise ValueError('mpol out of bound') + + if ntor is None: + ntor = self.ntor + elif ntor > self.ntor: + raise ValueError('ntor out of bound') + + self._vnc = vnc[0:mpol, self.ntor-ntor:self.ntor+ntor+1] + dofs = self.get_dofs() + self.local_full_x = dofs + + def set_vns_vnc_asarray(self, vns, vnc, mpol=None, ntor=None): + """ + Set the vns and vnc from two single arrays + """ + if mpol is None: + mpol = self.mpol + elif mpol > self.mpol: + raise ValueError('mpol out of bound') + + if ntor is None: + ntor = self.ntor + elif ntor > self.ntor: + raise ValueError('ntor out of bound') + + self.set_vns_asarray(vns, mpol, ntor) + self.set_vnc_asarray(vnc, mpol, ntor) + + def get_real_space_field(self): + """ + Fourier transform the field and get the real-space values of the normal component of the externally + provided field on the computational boundary. + The returned array will be of size specified by the surfaces' quadpoints and located on the quadpoints. + """ + vns, vnc = self.get_vns_vnc_asarray(mpol=self.mpol, ntor=self.ntor) + BdotN_unnormalized = self.surface.inverse_fourier_transform_scalar(vns, vnc, normalization=(2*np.pi)**2, stellsym=self.stellsym) + normal_field_real_space = -1 * BdotN_unnormalized / np.linalg.norm(self.surface.normal(), axis=-1) + return normal_field_real_space + diff --git a/src/simsopt/geo/surface.py b/src/simsopt/geo/surface.py index 9da6fba97..081aa5633 100644 --- a/src/simsopt/geo/surface.py +++ b/src/simsopt/geo/surface.py @@ -75,7 +75,8 @@ def from_nphi_ntheta(cls, nphi=61, ntheta=62, range="full torus", nfp=1, nphi, ntheta, nfp=nfp, range=range) return cls(quadpoints_phi=quadpoints_phi, quadpoints_theta=quadpoints_theta, nfp=nfp, **kwargs) - + + @staticmethod def get_quadpoints(nphi=None, ntheta=None, range=None, @@ -109,7 +110,8 @@ def get_quadpoints(nphi=None, """ return (Surface.get_phi_quadpoints(nphi=nphi, range=range, nfp=nfp), Surface.get_theta_quadpoints(ntheta=ntheta)) - + + @staticmethod def get_theta_quadpoints(ntheta=None): r""" Sets the theta grid points for Surface subclasses. @@ -124,7 +126,8 @@ def get_theta_quadpoints(ntheta=None): if ntheta is None: ntheta = 62 return list(np.linspace(0.0, 1.0, ntheta, endpoint=False)) - + + @staticmethod def get_phi_quadpoints(nphi=None, range=None, nfp=1): r""" Sets the phi grid points for Surface subclasses. @@ -840,6 +843,7 @@ def interpolate_on_arclength_grid(self, function, theta_evaluate): return function_interpolated + def signed_distance_from_surface(xyz, surface): """ Compute the signed distances from points ``xyz`` to a surface. The sign is diff --git a/src/simsopt/geo/surfacerzfourier.py b/src/simsopt/geo/surfacerzfourier.py index 391da4f47..f66a5c304 100644 --- a/src/simsopt/geo/surfacerzfourier.py +++ b/src/simsopt/geo/surfacerzfourier.py @@ -637,6 +637,130 @@ def write_nml(self, filename: str): with open(filename, 'w') as f: f.write(self.get_nml()) + 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, + with `npol` going from `-ntor` to `ntor` and `mpol` from 0 to `mpol` + 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, + 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. + *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, + the ntor attribute of the surface is used. + *Optional keyword arguments*: + - 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 + 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 + (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) + assert scalar.shape[1] == self.quadpoints_theta.size, "scalar must be evaluated at the quadrature points on the surface.\n the scalar you passed in has shape {}".format(scalar.shape) + stellsym = kwargs.pop('stellsym', self.stellsym) + if mpol is None: + try: mpol = self.mpol + except AttributeError: raise ValueError("mpol must be specified") + if ntor is None: + try: ntor = self.ntor + except AttributeError: raise ValueError("ntor must be specified") + A_mns = np.zeros((int(mpol + 1), int(2 * ntor + 1))) # sine coefficients + A_mnc = np.zeros((int(mpol + 1), int(2 * ntor + 1))) # cosine coefficients + ntheta_grid = len(self.quadpoints_theta) + nphi_grid = len(self.quadpoints_phi) + + factor = 2.0 / (ntheta_grid * nphi_grid) + + phi2d, theta2d = np.meshgrid(2 * np.pi * self.quadpoints_phi, + 2 * np.pi * self.quadpoints_theta, + indexing="ij") + + for m in range(mpol + 1): + nmin = -ntor + if m == 0: nmin = 1 + for n in range(nmin, ntor+1): + angle = m * theta2d - n * self.nfp * phi2d + sinangle = np.sin(angle) + factor2 = factor + # The next 2 lines ensure inverse Fourier transform(Fourier transform) = identity + if np.mod(ntheta_grid, 2) == 0 and m == (ntheta_grid/2): factor2 = factor2 / 2 + if np.mod(nphi_grid, 2) == 0 and abs(n) == (nphi_grid/2): factor2 = factor2 / 2 + A_mns[m, n + ntor] = np.sum(scalar * sinangle * factor2) + 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: + if isinstance(normalization, float): + raise ValueError("normalization must be a float") + A_mns = A_mns / normalization + A_mnc = A_mnc / normalization + + return A_mns, A_mnc + + def inverse_fourier_transform_scalar(self, A_mns, A_mnc, normalization=None, **kwargs): + r""" + Compute the inverse Fourier transform of a scalar on the surface, specified by the Fourier coefficients. The quantity must be + is evaluated at the quadrature points on the surface. The Fourier + transform is defined as + :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. + *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 + (these are zero if the surface is stellarator symmetric) + *Optional keyword arguments*: + - normalization: Fourier transform normalization. Can be: + None: forward and back transform are not normalized + float: inverse transform is multiplied by this number + - stellsym: boolean to override the stellsym attribute of the surface + """ + mpol = A_mns.shape[0] - 1 + ntor = int((A_mns.shape[1] - 1) / 2) + stellsym = kwargs.pop('stellsym', self.stellsym) + ntheta_grid = len(self.quadpoints_theta) + nphi_grid = len(self.quadpoints_phi) + + phi2d, theta2d = np.meshgrid(2 * np.pi * self.quadpoints_phi, + 2 * np.pi * self.quadpoints_theta, + indexing="ij") + + scalars = np.zeros((nphi_grid, ntheta_grid)) + for m in range(mpol + 1): + nmin = -ntor + if m == 0: nmin = 1 + for n in range(nmin, ntor+1): + angle = m * theta2d - n * self.nfp * phi2d + sinangle = np.sin(angle) + scalars = scalars + A_mns[m, n + ntor] * sinangle + 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 not isinstance(normalization, float): + raise ValueError("normalization must be a float") + scalars = scalars * normalization + return scalars + def make_rotating_ellipse(self, major_radius, minor_radius, elongation, torsion=0): """ Set the surface shape to be a rotating ellipse with the given diff --git a/src/simsopt/mhd/boozer.py b/src/simsopt/mhd/boozer.py index b7344c873..5fd869642 100644 --- a/src/simsopt/mhd/boozer.py +++ b/src/simsopt/mhd/boozer.py @@ -33,7 +33,6 @@ __all__ = ['Boozer', 'Quasisymmetry'] - class Boozer(Optimizable): """ This class handles the transformation to Boozer coordinates. @@ -124,10 +123,6 @@ def run(self): if isinstance(self.equil, Vmec): #partake in parallel VMEC job self.equil.run() - #skedaddle if you are not proc0 of your group - if (self.mpi is not None) and (not self.mpi.proc0_groups): - logger.info("This proc is skipping the rest of boozer.run since it is not a group leader.") - return wout = self.equil.wout # Shorthand @@ -178,8 +173,8 @@ def run(self): self.bx.mnmax = wout.mnmax self.bx.xm = wout.xm self.bx.xn = wout.xn - logger.info('mnmax:', wout.mnmax, ' len(xm):', len(wout.xm), ' len(xn):', len(wout.xn)) - logger.info('mnmax_nyq:', wout.mnmax_nyq, ' len(xm_nyq):', len(wout.xm_nyq), ' len(xn_nyq):', len(wout.xn_nyq)) + logger.info(f'mnmax: {wout.mnmax} len(xm): {len(wout.xm)} len(xn): {len(wout.xn)}') + logger.info(f'mnmax_nyq: {wout.mnmax_nyq} len(xm_nyq): {len(wout.xm_nyq)} len(xn_nyq): {len(wout.xn_nyq)}') assert len(wout.xm) == wout.mnmax assert len(wout.xn) == wout.mnmax assert len(self.bx.xm) == self.bx.mnmax diff --git a/src/simsopt/mhd/defaults_freebound.sp b/src/simsopt/mhd/defaults_freebound.sp new file mode 100644 index 000000000..b7fa02917 --- /dev/null +++ b/src/simsopt/mhd/defaults_freebound.sp @@ -0,0 +1,210 @@ +&physicslist + Igeometry = 3 + Istellsym = 1 + Lfreebound = 1 + phiedge = 2.000000000000000E+00 + curtor = 0.000000000000000E+00 + curpol = 4.000000000000000E+01 + gamma = 0.000000000000000E+00 + Nfp = 5 + Nvol = 1 + Mpol = 3 + Ntor = 3 + Lrad = 6 2 + tflux = 1.000000000000000E+00 6.218040116397460E+00 + pflux = 0.000000000000000E+00 3.816883073688957E+00 + helicity = 6.443366715051990E-02 1.116525236095663E+00 + pscale = 1.000000000000000E+02 + Ladiabatic = 0 + pressure = 0.000000000000000E+00 0.000000000000000E+00 + adiabatic = 0.000000000000000E+00 0.000000000000000E+00 + mu = 0.000000000000000E+00 0.000000000000000E+00 + Ivolume = 0.000000000000000E+00 0.000000000000000E+00 + Isurf = 0.000000000000000E+00 0.000000000000000E+00 + Lconstraint = 3 + pl = 0 0 0 + ql = 0 0 0 + pr = 0 0 0 + qr = 0 0 0 + iota = 0.000000000000000E+00 2.809417939338480E-01 3.050000000000000E-01 + lp = 0 0 0 + lq = 0 0 0 + rp = 0 0 0 + rq = 0 0 0 + oita = 0.000000000000000E+00 2.809417939338480E-01 3.050000000000000E-01 + mupftol = 1.000000000000000E-12 + mupfits = 128 + Lreflect = 0 + rpol = 1.000000000000000E+00 + rtor = 1.000000000000000E+00 + Rac = 9.645091076834294E+00 -1.273000100134322E-01 7.227734995841448E-04 2.901752615147871E-05 + Zas = 0.000000000000000E+00 -1.201876846458274E-01 1.210937745192939E-03 8.101291120903119E-05 + Ras = 0.000000000000000E+00 0.000000000000000E+00 0.000000000000000E+00 0.000000000000000E+00 + Zac = 0.000000000000000E+00 0.000000000000000E+00 0.000000000000000E+00 0.000000000000000E+00 +Rbc(0,0) = 9.704610460216003E+00 Zbs(0,0) = 0.000000000000000E+00 Rbs(0,0) = 0.000000000000000E+00 Zbc(0,0) = 0.000000000000000E+00 +Rbc(1,0) = -1.045848541227035E-01 Zbs(1,0) = -1.034983900846667E-01 Rbs(1,0) = 0.000000000000000E+00 Zbc(1,0) = 0.000000000000000E+00 +Rbc(2,0) = 5.496924456881049E-04 Zbs(2,0) = 8.105754421674023E-04 Rbs(2,0) = 0.000000000000000E+00 Zbc(2,0) = 0.000000000000000E+00 +Rbc(3,0) = -2.605852916338383E-05 Zbs(3,0) = 2.595333996994953E-06 Rbs(3,0) = 0.000000000000000E+00 Zbc(3,0) = 0.000000000000000E+00 +Rbc(-3,1) = -1.195463633068680E-05 Zbs(-3,1) = 2.841919864618188E-05 Rbs(-3,1) = 0.000000000000000E+00 Zbc(-3,1) = 0.000000000000000E+00 +Rbc(-2,1) = -3.394450906670391E-04 Zbs(-2,1) = 3.596274021066687E-04 Rbs(-2,1) = 0.000000000000000E+00 Zbc(-2,1) = 0.000000000000000E+00 +Rbc(-1,1) = -4.931651888369649E-03 Zbs(-1,1) = 6.963549886130800E-03 Rbs(-1,1) = 0.000000000000000E+00 Zbc(-1,1) = 0.000000000000000E+00 +Rbc(0,1) = 9.918678713491879E-01 Zbs(0,1) = -9.772386678651404E-01 Rbs(0,1) = 0.000000000000000E+00 Zbc(0,1) = 0.000000000000000E+00 +Rbc(1,1) = 8.640027952833292E-02 Zbs(1,1) = 9.241036581289261E-02 Rbs(1,1) = 0.000000000000000E+00 Zbc(1,1) = 0.000000000000000E+00 +Rbc(2,1) = 1.964242130996640E-03 Zbs(2,1) = 2.153208475198715E-03 Rbs(2,1) = 0.000000000000000E+00 Zbc(2,1) = 0.000000000000000E+00 +Rbc(3,1) = 7.998671801673742E-05 Zbs(3,1) = 1.161062030463278E-04 Rbs(3,1) = 0.000000000000000E+00 Zbc(3,1) = 0.000000000000000E+00 +Rbc(-3,2) = -2.762152103889883E-06 Zbs(-3,2) = 4.166379371076043E-06 Rbs(-3,2) = 0.000000000000000E+00 Zbc(-3,2) = 0.000000000000000E+00 +Rbc(-2,2) = -1.726514233457863E-05 Zbs(-2,2) = 1.983462874964590E-05 Rbs(-2,2) = 0.000000000000000E+00 Zbc(-2,2) = 0.000000000000000E+00 +Rbc(-1,2) = -6.283360944862572E-04 Zbs(-1,2) = 7.379659477847013E-04 Rbs(-1,2) = 0.000000000000000E+00 Zbc(-1,2) = 0.000000000000000E+00 +Rbc(0,2) = -2.706714444329093E-02 Zbs(0,2) = 2.833946442940345E-02 Rbs(0,2) = 0.000000000000000E+00 Zbc(0,2) = 0.000000000000000E+00 +Rbc(1,2) = -1.663136766639212E-02 Zbs(1,2) = -2.612324295966568E-03 Rbs(1,2) = 0.000000000000000E+00 Zbc(1,2) = 0.000000000000000E+00 +Rbc(2,2) = -4.482002741556905E-03 Zbs(2,2) = -5.025208765858588E-03 Rbs(2,2) = 0.000000000000000E+00 Zbc(2,2) = 0.000000000000000E+00 +Rbc(3,2) = 6.068006959746730E-05 Zbs(3,2) = -1.030827848856039E-05 Rbs(3,2) = 0.000000000000000E+00 Zbc(3,2) = 0.000000000000000E+00 +Rbc(-3,3) = 4.496337386676320E-06 Zbs(-3,3) = -2.281597748539077E-06 Rbs(-3,3) = 0.000000000000000E+00 Zbc(-3,3) = 0.000000000000000E+00 +Rbc(-2,3) = -1.615850297435840E-05 Zbs(-2,3) = -3.695652852901507E-05 Rbs(-2,3) = 0.000000000000000E+00 Zbc(-2,3) = 0.000000000000000E+00 +Rbc(-1,3) = 1.492859671801447E-04 Zbs(-1,3) = -2.388623508824467E-04 Rbs(-1,3) = 0.000000000000000E+00 Zbc(-1,3) = 0.000000000000000E+00 +Rbc(0,3) = 3.539712137494496E-03 Zbs(0,3) = -4.998079213999107E-03 Rbs(0,3) = 0.000000000000000E+00 Zbc(0,3) = 0.000000000000000E+00 +Rbc(1,3) = -3.845595274173447E-02 Zbs(1,3) = -2.589386621604443E-02 Rbs(1,3) = 0.000000000000000E+00 Zbc(1,3) = 0.000000000000000E+00 +Rbc(2,3) = -4.164306359043766E-02 Zbs(2,3) = -4.284945955446148E-02 Rbs(2,3) = 0.000000000000000E+00 Zbc(2,3) = 0.000000000000000E+00 +Rbc(3,3) = -1.242620719798954E-03 Zbs(3,3) = -1.354757246649995E-03 Rbs(3,3) = 0.000000000000000E+00 Zbc(3,3) = 0.000000000000000E+00 +Rwc(0,0) = 1.000000000000000E+01 Zws(0,0) = 0.000000000000000E+00 Rws(0,0) = 0.000000000000000E+00 Zwc(0,0) = 0.000000000000000E+00 +Rwc(1,0) = 0.000000000000000E+00 Zws(1,0) = 0.000000000000000E+00 Rws(1,0) = 0.000000000000000E+00 Zwc(1,0) = 0.000000000000000E+00 +Rwc(2,0) = 0.000000000000000E+00 Zws(2,0) = 0.000000000000000E+00 Rws(2,0) = 0.000000000000000E+00 Zwc(2,0) = 0.000000000000000E+00 +Rwc(3,0) = 0.000000000000000E+00 Zws(3,0) = 0.000000000000000E+00 Rws(3,0) = 0.000000000000000E+00 Zwc(3,0) = 0.000000000000000E+00 +Rwc(-3,1) = 0.000000000000000E+00 Zws(-3,1) = 0.000000000000000E+00 Rws(-3,1) = 0.000000000000000E+00 Zwc(-3,1) = 0.000000000000000E+00 +Rwc(-2,1) = 0.000000000000000E+00 Zws(-2,1) = 0.000000000000000E+00 Rws(-2,1) = 0.000000000000000E+00 Zwc(-2,1) = 0.000000000000000E+00 +Rwc(-1,1) = 0.000000000000000E+00 Zws(-1,1) = 0.000000000000000E+00 Rws(-1,1) = 0.000000000000000E+00 Zwc(-1,1) = 0.000000000000000E+00 +Rwc(0,1) = 3.000000000000000E+00 Zws(0,1) = -3.000000000000000E+00 Rws(0,1) = 0.000000000000000E+00 Zwc(0,1) = 0.000000000000000E+00 +Rwc(1,1) = 1.250000000000000E+00 Zws(1,1) = 1.250000000000000E+00 Rws(1,1) = 0.000000000000000E+00 Zwc(1,1) = 0.000000000000000E+00 +Rwc(2,1) = 0.000000000000000E+00 Zws(2,1) = 0.000000000000000E+00 Rws(2,1) = 0.000000000000000E+00 Zwc(2,1) = 0.000000000000000E+00 +Rwc(3,1) = 0.000000000000000E+00 Zws(3,1) = 0.000000000000000E+00 Rws(3,1) = 0.000000000000000E+00 Zwc(3,1) = 0.000000000000000E+00 +Rwc(-3,2) = 0.000000000000000E+00 Zws(-3,2) = 0.000000000000000E+00 Rws(-3,2) = 0.000000000000000E+00 Zwc(-3,2) = 0.000000000000000E+00 +Rwc(-2,2) = 0.000000000000000E+00 Zws(-2,2) = 0.000000000000000E+00 Rws(-2,2) = 0.000000000000000E+00 Zwc(-2,2) = 0.000000000000000E+00 +Rwc(-1,2) = 0.000000000000000E+00 Zws(-1,2) = 0.000000000000000E+00 Rws(-1,2) = 0.000000000000000E+00 Zwc(-1,2) = 0.000000000000000E+00 +Rwc(0,2) = 0.000000000000000E+00 Zws(0,2) = 0.000000000000000E+00 Rws(0,2) = 0.000000000000000E+00 Zwc(0,2) = 0.000000000000000E+00 +Rwc(1,2) = 0.000000000000000E+00 Zws(1,2) = 0.000000000000000E+00 Rws(1,2) = 0.000000000000000E+00 Zwc(1,2) = 0.000000000000000E+00 +Rwc(2,2) = 0.000000000000000E+00 Zws(2,2) = 0.000000000000000E+00 Rws(2,2) = 0.000000000000000E+00 Zwc(2,2) = 0.000000000000000E+00 +Rwc(3,2) = 0.000000000000000E+00 Zws(3,2) = 0.000000000000000E+00 Rws(3,2) = 0.000000000000000E+00 Zwc(3,2) = 0.000000000000000E+00 +Rwc(-3,3) = 0.000000000000000E+00 Zws(-3,3) = 0.000000000000000E+00 Rws(-3,3) = 0.000000000000000E+00 Zwc(-3,3) = 0.000000000000000E+00 +Rwc(-2,3) = 0.000000000000000E+00 Zws(-2,3) = 0.000000000000000E+00 Rws(-2,3) = 0.000000000000000E+00 Zwc(-2,3) = 0.000000000000000E+00 +Rwc(-1,3) = 0.000000000000000E+00 Zws(-1,3) = 0.000000000000000E+00 Rws(-1,3) = 0.000000000000000E+00 Zwc(-1,3) = 0.000000000000000E+00 +Rwc(0,3) = 0.000000000000000E+00 Zws(0,3) = 0.000000000000000E+00 Rws(0,3) = 0.000000000000000E+00 Zwc(0,3) = 0.000000000000000E+00 +Rwc(1,3) = 0.000000000000000E+00 Zws(1,3) = 0.000000000000000E+00 Rws(1,3) = 0.000000000000000E+00 Zwc(1,3) = 0.000000000000000E+00 +Rwc(2,3) = 0.000000000000000E+00 Zws(2,3) = 0.000000000000000E+00 Rws(2,3) = 0.000000000000000E+00 Zwc(2,3) = 0.000000000000000E+00 +Rwc(3,3) = 0.000000000000000E+00 Zws(3,3) = 0.000000000000000E+00 Rws(3,3) = 0.000000000000000E+00 Zwc(3,3) = 0.000000000000000E+00 +Vns(0,0) = 0.000000000000000E+00 Bns(0,0) = 0.000000000000000E+00 Vnc(0,0) = 0.000000000000000E+00 Bnc(0,0) = 0.000000000000000E+00 +Vns(1,0) = 0.000000000000000E+00 Bns(1,0) = -2.098314067166154E-02 Vnc(1,0) = 0.000000000000000E+00 Bnc(1,0) = 0.000000000000000E+00 +Vns(2,0) = 0.000000000000000E+00 Bns(2,0) = -8.308195552233343E-04 Vnc(2,0) = 0.000000000000000E+00 Bnc(2,0) = 0.000000000000000E+00 +Vns(3,0) = 0.000000000000000E+00 Bns(3,0) = -1.000049567085383E-05 Vnc(3,0) = 0.000000000000000E+00 Bnc(3,0) = 0.000000000000000E+00 +Vns(-3,1) = 0.000000000000000E+00 Bns(-3,1) = 1.025035236011660E-05 Vnc(-3,1) = 0.000000000000000E+00 Bnc(-3,1) = 0.000000000000000E+00 +Vns(-2,1) = 0.000000000000000E+00 Bns(-2,1) = 6.972200102040577E-05 Vnc(-2,1) = 0.000000000000000E+00 Bnc(-2,1) = 0.000000000000000E+00 +Vns(-1,1) = 0.000000000000000E+00 Bns(-1,1) = -8.857899561787283E-04 Vnc(-1,1) = 0.000000000000000E+00 Bnc(-1,1) = 0.000000000000000E+00 +Vns(0,1) = 0.000000000000000E+00 Bns(0,1) = 1.268794075038921E-02 Vnc(0,1) = 0.000000000000000E+00 Bnc(0,1) = 0.000000000000000E+00 +Vns(1,1) = 0.000000000000000E+00 Bns(1,1) = 5.235427842411160E-03 Vnc(1,1) = 0.000000000000000E+00 Bnc(1,1) = 0.000000000000000E+00 +Vns(2,1) = 0.000000000000000E+00 Bns(2,1) = -1.115762206996049E-03 Vnc(2,1) = 0.000000000000000E+00 Bnc(2,1) = 0.000000000000000E+00 +Vns(3,1) = 0.000000000000000E+00 Bns(3,1) = -4.424638432480133E-06 Vnc(3,1) = 0.000000000000000E+00 Bnc(3,1) = 0.000000000000000E+00 +Vns(-3,2) = 0.000000000000000E+00 Bns(-3,2) = -2.420544046309371E-06 Vnc(-3,2) = 0.000000000000000E+00 Bnc(-3,2) = 0.000000000000000E+00 +Vns(-2,2) = 0.000000000000000E+00 Bns(-2,2) = -1.127639127865856E-05 Vnc(-2,2) = 0.000000000000000E+00 Bnc(-2,2) = 0.000000000000000E+00 +Vns(-1,2) = 0.000000000000000E+00 Bns(-1,2) = -6.680730994641753E-04 Vnc(-1,2) = 0.000000000000000E+00 Bnc(-1,2) = 0.000000000000000E+00 +Vns(0,2) = 0.000000000000000E+00 Bns(0,2) = -2.105169513847248E-02 Vnc(0,2) = 0.000000000000000E+00 Bnc(0,2) = 0.000000000000000E+00 +Vns(1,2) = 0.000000000000000E+00 Bns(1,2) = -3.458317607829524E-02 Vnc(1,2) = 0.000000000000000E+00 Bnc(1,2) = 0.000000000000000E+00 +Vns(2,2) = 0.000000000000000E+00 Bns(2,2) = 7.025733332835698E-03 Vnc(2,2) = 0.000000000000000E+00 Bnc(2,2) = 0.000000000000000E+00 +Vns(3,2) = 0.000000000000000E+00 Bns(3,2) = 1.097429331303196E-03 Vnc(3,2) = 0.000000000000000E+00 Bnc(3,2) = 0.000000000000000E+00 +Vns(-3,3) = 0.000000000000000E+00 Bns(-3,3) = 2.012582723897679E-06 Vnc(-3,3) = 0.000000000000000E+00 Bnc(-3,3) = 0.000000000000000E+00 +Vns(-2,3) = 0.000000000000000E+00 Bns(-2,3) = -1.738204711861772E-06 Vnc(-2,3) = 0.000000000000000E+00 Bnc(-2,3) = 0.000000000000000E+00 +Vns(-1,3) = 0.000000000000000E+00 Bns(-1,3) = 5.534246820790295E-05 Vnc(-1,3) = 0.000000000000000E+00 Bnc(-1,3) = 0.000000000000000E+00 +Vns(0,3) = 0.000000000000000E+00 Bns(0,3) = 1.008641123981896E-03 Vnc(0,3) = 0.000000000000000E+00 Bnc(0,3) = 0.000000000000000E+00 +Vns(1,3) = 0.000000000000000E+00 Bns(1,3) = -8.786509679612788E-03 Vnc(1,3) = 0.000000000000000E+00 Bnc(1,3) = 0.000000000000000E+00 +Vns(2,3) = 0.000000000000000E+00 Bns(2,3) = -1.503823116287205E-02 Vnc(2,3) = 0.000000000000000E+00 Bnc(2,3) = 0.000000000000000E+00 +Vns(3,3) = 0.000000000000000E+00 Bns(3,3) = -9.987559625167259E-04 Vnc(3,3) = 0.000000000000000E+00 Bnc(3,3) = 0.000000000000000E+00 +/ +&numericlist + Linitialize = 0 + LautoinitBn = 0 + Lzerovac = 0 + Ndiscrete = 2 + Nquad = -1 + iMpol = -4 + iNtor = -4 + Lsparse = 0 + Lsvdiota = 0 + imethod = 3 + iorder = 2 + iprecon = 1 + iotatol = -1.000000000000000E+00 + Lextrap = 0 + Mregular = -1 + Lrzaxis = 2 + Ntoraxis = 3 +/ +&locallist + LBeltrami = 4 + Linitgues = 1 + Lmatsolver = 3 + NiterGMRES = 200 + LGMRESprec = 1 + epsGMRES = 1.000000000000000E-14 + epsILU = 1.000000000000000E-12 +/ +&globallist + Lfindzero = 2 + escale = 0.000000000000000E+00 + opsilon = 1.000000000000000E+00 + pcondense = 4.000000000000000E+00 + epsilon = 1.000000000000000E+00 + wpoloidal = 1.000000000000000E+00 + upsilon = 1.000000000000000E+00 + forcetol = 1.000000000000000E-12 + c05xmax = 1.000000000000000E-06 + c05xtol = 1.000000000000000E-12 + c05factor = 1.000000000000000E-04 + LreadGF = F + mfreeits = 1 + gBntol = 1.000000000000000E-04 + gBnbld = 1.000000000000000E-01 + vcasingeps = 1.000000000000000E-10 + vcasingtol = 1.000000000000000E-04 + vcasingits = 8 + vcasingper = 1 +/ +&diagnosticslist + odetol = 1.000000000000000E-07 + nPpts = -1 + Ppts = 0.000000000000000E+00 + nPtrj = 1 1 + LHevalues = F + LHevectors = F + LHmatrix = F + Lperturbed = 0 + dpp = -1 + dqq = -1 + dRZ = 1.000000000000000E-05 + Lcheck = 0 + Ltiming = F +/ +&screenlist +/ + 0 0 9.704610460216003E+00 0.000000000000000E+00 0.000000000000000E+00 0.000000000000000E+00 + 0 1 -1.045848541227035E-01 -1.034983900846667E-01 0.000000000000000E+00 0.000000000000000E+00 + 0 2 5.496924456881049E-04 8.105754421674023E-04 0.000000000000000E+00 0.000000000000000E+00 + 0 3 -2.605852916338383E-05 2.595333996994953E-06 0.000000000000000E+00 0.000000000000000E+00 + 1 -3 -1.195463633068680E-05 2.841919864618188E-05 0.000000000000000E+00 0.000000000000000E+00 + 1 -2 -3.394450906670391E-04 3.596274021066687E-04 0.000000000000000E+00 0.000000000000000E+00 + 1 -1 -4.931651888369649E-03 6.963549886130800E-03 0.000000000000000E+00 0.000000000000000E+00 + 1 0 9.918678713491879E-01 -9.772386678651404E-01 0.000000000000000E+00 0.000000000000000E+00 + 1 1 8.640027952833292E-02 9.241036581289261E-02 0.000000000000000E+00 0.000000000000000E+00 + 1 2 1.964242130996640E-03 2.153208475198715E-03 0.000000000000000E+00 0.000000000000000E+00 + 1 3 7.998671801673742E-05 1.161062030463278E-04 0.000000000000000E+00 0.000000000000000E+00 + 2 -3 -2.762152103889883E-06 4.166379371076043E-06 0.000000000000000E+00 0.000000000000000E+00 + 2 -2 -1.726514233457863E-05 1.983462874964590E-05 0.000000000000000E+00 0.000000000000000E+00 + 2 -1 -6.283360944862572E-04 7.379659477847013E-04 0.000000000000000E+00 0.000000000000000E+00 + 2 0 -2.706714444329093E-02 2.833946442940345E-02 0.000000000000000E+00 0.000000000000000E+00 + 2 1 -1.663136766639212E-02 -2.612324295966568E-03 0.000000000000000E+00 0.000000000000000E+00 + 2 2 -4.482002741556905E-03 -5.025208765858588E-03 0.000000000000000E+00 0.000000000000000E+00 + 2 3 6.068006959746730E-05 -1.030827848856039E-05 0.000000000000000E+00 0.000000000000000E+00 + 3 -3 4.496337386676320E-06 -2.281597748539077E-06 0.000000000000000E+00 0.000000000000000E+00 + 3 -2 -1.615850297435840E-05 -3.695652852901507E-05 0.000000000000000E+00 0.000000000000000E+00 + 3 -1 1.492859671801447E-04 -2.388623508824467E-04 0.000000000000000E+00 0.000000000000000E+00 + 3 0 3.539712137494496E-03 -4.998079213999107E-03 0.000000000000000E+00 0.000000000000000E+00 + 3 1 -3.845595274173447E-02 -2.589386621604443E-02 0.000000000000000E+00 0.000000000000000E+00 + 3 2 -4.164306359043766E-02 -4.284945955446148E-02 0.000000000000000E+00 0.000000000000000E+00 + 3 3 -1.242620719798954E-03 -1.354757246649995E-03 0.000000000000000E+00 0.000000000000000E+00 diff --git a/src/simsopt/mhd/spec.py b/src/simsopt/mhd/spec.py index 05e304a52..8a9655efc 100644 --- a/src/simsopt/mhd/spec.py +++ b/src/simsopt/mhd/spec.py @@ -1,4 +1,3 @@ -# coding: utf-8 # Copyright (c) HiddenSymmetries Development Team. # Distributed under the terms of the MIT License @@ -9,8 +8,10 @@ import copy import logging import os.path +import shutil, os import traceback from typing import Optional +from scipy.constants import mu_0 import numpy as np @@ -80,7 +81,7 @@ class Spec(Optimizable): mpi: A :obj:`simsopt.util.mpi.MpiPartition` instance, from which the worker groups will be used for SPEC calculations. If ``None``, each MPI process will run SPEC independently. - verbose: Whether to print SPEC output to stdout. + verbose: Whether to print SPEC output to stdout (and also a few other statements) keep_all_files: If ``False``, all output files will be deleted except for the first and most recent ones from worker group 0. If ``True``, all output files will be kept. @@ -102,26 +103,24 @@ def __init__(self, if py_spec is None: raise RuntimeError( "Using Spec requires py_spec to be installed.") + if tolerance <= 0: + raise ValueError( + 'tolerance should be greater than zero' + ) + self.tolerance = tolerance self.lib = spec # For the most commonly accessed fortran modules, provide a - # shorthand so ".lib" is not needed: - modules = [ - "inputlist", - "allglobal", - ] - for key in modules: - setattr(self, key, getattr(spec, key)) + # shorthand: + setattr(self, "inputlist", getattr(spec, "inputlist")) + setattr(self, "allglobal", getattr(spec, "allglobal")) self.verbose = verbose # mute screen output if necessary - # TODO: relies on /dev/null being accessible (Windows!) + # TODO: mute SPEC in windows, currently relies on /dev/null being accessible if not self.verbose: self.lib.fileunits.mute(1) - # python wrapper does not need to write files along the run - #self.lib.allglobal.skip_write = True - # If mpi is not specified, use a single worker group: if mpi is None: self.mpi = MpiPartition(ngroups=1) @@ -134,77 +133,53 @@ def __init__(self, # Read default input file, which should be in the same # directory as this file: filename = os.path.join(os.path.dirname(__file__), 'defaults.sp') - logger.info( - f"Initializing a SPEC object from defaults in {filename}") + if self.mpi.proc0_groups: + logger.info( + f"Initializing a SPEC object from defaults in {filename}") else: if not filename.endswith('.sp'): filename = f"{filename}.sp" - logger.info(f"Initializing a SPEC object from file: {filename}") - - if tolerance <= 0: - raise ValueError( - 'tolerance should be greater than zero' - ) - self.tolerance = tolerance - - self.init(filename) + if self.mpi.proc0_groups: + logger.info(f"Group {self.mpi.group}: Initializing a SPEC object from file: {filename}") + + #If spec has run before, clear the f90wrap array caches. + # See https://github.com/hiddenSymmetries/simsopt/pull/431 + # This addresses the issue https://github.com/hiddenSymmetries/simsopt/issues/357 + if spec.allglobal._arrays: + self._clear_f90wrap_array_caches() + + # Initialize the FORTRAN state with values in the input file: + self._init_fortran_state(filename) si = spec.inputlist # Shorthand - # Read number of (plasma) volumes + # Copy useful and commonly used values from SPEC inputlist to + self.stellsym = bool(si.istellsym) + self.freebound = bool(si.lfreebound) + self.nfp = si.nfp + self.mpol = si.mpol + self.ntor = si.ntor self.nvol = si.nvol - - # Read number of (plasma+vacuum) volumes - if si.lfreebound: - self.mvol = self.nvol + 1 - else: - self.mvol = self.nvol + if self.freebound: + self.mvol = si.nvol + 1 + else: + self.mvol = si.nvol # Store initial guess data # The initial guess is a collection of SurfaceRZFourier instances, # stored in a list of size Mvol-1 (the number of inner interfaces) - nmodes = self.allglobal.num_modes - stellsym = bool(si.istellsym) - if nmodes > 0 and self.nvol > 1: - self.initial_guess = [ - SurfaceRZFourier(nfp=si.nfp, stellsym=stellsym, mpol=si.mpol, ntor=si.ntor) for n in range(0, self.mvol-1) - ] - for imode in range(0, nmodes): - mm = self.allglobal.mmrzrz[imode] - nn = self.allglobal.nnrzrz[imode] - if mm > si.mpol: - continue - if abs(nn) > si.ntor: - continue - - # Populate SurfaceRZFourier instances, except plasma boundary - for lvol in range(0, self.nvol-1): - self.initial_guess[lvol].set_rc(mm, nn, self.allglobal.allrzrz[0, lvol, imode]) - self.initial_guess[lvol].set_zs(mm, nn, self.allglobal.allrzrz[1, lvol, imode]) - - if not si.istellsym: - self.initial_guess[lvol].set_rs(mm, nn, self.allglobal.allrzrz[2, lvol, imode]) - self.initial_guess[lvol].set_zc(mm, nn, self.allglobal.allrzrz[3, lvol, imode]) - - if si.lfreebound: # Populate plasma boundary as well - self.initial_guess[self.nvol-1].set_rc(mm, nn, si.rbc[si.mntor+nn, si.mmpol+mm]) - self.initial_guess[self.nvol-1].set_zs(mm, nn, si.zbs[si.mntor+nn, si.mmpol+mm]) - - if not si.istellsym: - self.initial_guess[self.nvol-1].set_rs(mm, nn, si.rbs[si.mntor+nn, si.mmpol+mm]) - self.initial_guess[self.nvol-1].set_zc(mm, nn, si.zbc[si.mntor+nn, si.mmpol+mm]) - + read_initial_guess = self.inputlist.linitialize == 0 and self.nvol > 1 + if read_initial_guess: + self.initial_guess = self._read_initial_guess() # In general, initial guess is NOT a degree of freedom for the # optimization - we thus fix them. - for lvol in range(0, self.mvol-1): - self.initial_guess[lvol].fix_all() - + for surface in self.initial_guess: + surface.fix_all() else: # There is no initial guess - in this case, we let SPEC handle # the construction of the initial guess. This generally means # that the geometry of the inner interfaces will be constructed # by interpolation between the plasma (or computational) boundary # and the magnetic axis - self.initial_guess = None # Store axis data @@ -220,31 +195,28 @@ def __init__(self, self.files_to_delete = [] # Create a surface object for the boundary: - logger.debug(f"In __init__, si.istellsym={si.istellsym} stellsym={stellsym}") - self._boundary = SurfaceRZFourier(nfp=si.nfp, - stellsym=stellsym, - mpol=si.mpol, - ntor=si.ntor) - - # Transfer the boundary shape from fortran to the boundary - # surface object: - for m in range(si.mpol + 1): - for n in range(-si.ntor, si.ntor + 1): - self._boundary.rc[m, - n + si.ntor] = si.rbc[n + si.mntor, - m + si.mmpol] - self._boundary.zs[m, - n + si.ntor] = si.zbs[n + si.mntor, - m + si.mmpol] - if not stellsym: - self._boundary.rs[m, - n + si.ntor] = si.rbs[n + si.mntor, - m + si.mmpol] - self._boundary.zc[m, - n + si.ntor] = si.zbc[n + si.mntor, - m + si.mmpol] + self._boundary = SurfaceRZFourier(nfp=self.nfp, stellsym=self.stellsym, + mpol=self.mpol, ntor=self.ntor) + self._boundary.rc[:] = self.array_translator(si.rbc, style='spec').as_simsopt + self._boundary.zs[:] = self.array_translator(si.zbs, style='spec').as_simsopt + if not self.freebound: + self._boundary.rs[:] = self.array_translator(si.rbs, style='spec').as_simsopt + self._boundary.zc[:] = self.array_translator(si.zbc, style='spec').as_simsopt self._boundary.local_full_x = self._boundary.get_dofs() + # If the equilibrium is freeboundary, we need to read the computational + # boundary as well. Otherwise set the outermost boundary as the + # computational boundary. + if self.freebound: + self._computational_boundary = SurfaceRZFourier(nfp=self.nfp, stellsym=self.stellsym, + mpol=self.mpol, ntor=self.ntor) + self._computational_boundary.rc[:] = self.array_translator(si.rwc, style='spec').as_simsopt + self._computational_boundary.zs[:] = self.array_translator(si.zws, style='spec').as_simsopt + if not self.stellsym: + self._computational_boundary.rs[:] = self.array_translator(si.rws, style='spec').as_simsopt + self._computational_boundary.zc[:] = self.array_translator(si.zwc, style='spec').as_simsopt + self._computational_boundary.local_full_x = self._computational_boundary.get_dofs() + self.need_to_run_code = True self.counter = -1 @@ -262,24 +234,132 @@ def __init__(self, # Define normal field - these are the Vns, Vnc harmonics. Can be used as # dofs in an optimization - if si.lfreebound: - self.normal_field = NormalField.from_spec(filename) + if self.freebound: + vns = self.array_translator(si.vns) + vnc = self.array_translator(si.vnc) + self._normal_field = NormalField(nfp=self.nfp, stellsym=self.stellsym, + mpol=self.mpol, ntor=self.ntor, + vns=vns.as_simsopt, vnc=vnc.as_simsopt, + surface=self._computational_boundary) else: - self.normal_field: Optional[NormalField] = None + self._normal_field: Optional[NormalField] = None # By default, all dofs owned by SPEC directly, as opposed to # dofs owned by the boundary surface object, are fixed. x0 = self.get_dofs() fixed = np.full(len(x0), True) names = ['phiedge', 'curtor'] - if si.lfreebound: - depends_on = [self.normal_field] + if self.freebound: + depends_on = [self._normal_field] else: depends_on = [self._boundary] super().__init__(x0=x0, fixed=fixed, names=names, depends_on=depends_on, external_dof_setter=Spec.set_dofs) + + @classmethod + def default_freeboundary(cls, copy_to_pwd, verbose=True): + """ + Create a default freeboundary SPEC object + Args: + copy_to_pwd: boolean, if True, the default input file will be copied to the current working directory. Has to be set True as free-boundary SPEC can only handle files in the current working directory. + verbose: boolean, if True, print statements will be printed + """ + filename = 'defaults_freebound.sp' + if verbose: + print(f'Copying {os.path.join(os.path.dirname(__file__), filename)} to {os.getcwd()}/{filename}') + shutil.copy(os.path.join(os.path.dirname(__file__), filename), os.getcwd()) + return cls(filename=filename, verbose=verbose) + + @classmethod + def default_fixedboundary(cls): + """ + Create a default fixedboundary SPEC object (same as calling class empty) + """ + return cls() + + def _read_initial_guess(self): + """ + Read initial guesses from SPEC and return a list of surfaceRZFourier objects. + """ + nmodes = self.allglobal.num_modes + interfaces = [] # initialize list + n_guess_surfs = self.nvol if self.freebound else self.nvol - 1 # if freeboundary, plasma boundary is also a guess surface + for lvol in range(0, n_guess_surfs): # loop over volumes + # the fourier comonents of the initial guess are stored in the allrzrz array + thissurf_rc = self.allglobal.allrzrz[0, lvol, :nmodes] + thissurf_zs = self.allglobal.allrzrz[1, lvol, :nmodes] + if not self.stellsym: + thissurf_rs = self.allglobal.allrzrz[2, lvol, :nmodes] + thissurf_zc = self.allglobal.allrzrz[3, lvol, :nmodes] + thissurf = SurfaceRZFourier(nfp=self.nfp, stellsym=self.stellsym, + mpol=self.mpol, ntor=self.ntor) + # the mode number convention is different in SPEC. Best to use the + # in-built array mmrzrz to get the numbers for each fourier component: + for imode in range(0, nmodes): + mm = self.allglobal.mmrzrz[imode] # helper arrays to get index for each mode + nn = self.allglobal.nnrzrz[imode] + # The guess surface could have more modes than we are running SPEC with, skip if case + if mm > self.mpol or abs(nn) > self.ntor: + continue + thissurf.set_rc(mm, nn, thissurf_rc[imode]) + thissurf.set_zs(mm, nn, thissurf_zs[imode]) + if not self.stellsym: + thissurf.set_rs(mm, nn, thissurf_rs[imode]) + thissurf.set_zc(mm, nn, thissurf_zc[imode]) + interfaces.append(thissurf) + return interfaces + + def _set_spec_initial_guess(self): + """ + Set initial guesses in SPEC from a list of surfaceRZFourier objects. + """ + # Set all modes to zero + spec.allglobal.mmrzrz[:] = 0 + spec.allglobal.nnrzrz[:] = 0 + spec.allglobal.allrzrz[:] = 0 + + # transform to SurfaceRZFourier if necessary + initial_guess = [s.to_RZFourier() for s in self.initial_guess] + + # Loop on modes + imn = -1 # counter + for mm in range(0, self.mpol+1): + for nn in range(-self.ntor, self.ntor+1): + if mm == 0 and nn < 0: + continue + + imn += 1 + + spec.allglobal.mmrzrz[imn] = mm + spec.allglobal.nnrzrz[imn] = nn + + # Populate inner plasma boundaries + for lvol in range(0, self.nvol-1): + spec.allglobal.allrzrz[0, lvol, imn] = initial_guess[lvol].get_rc(mm, nn) + spec.allglobal.allrzrz[1, lvol, imn] = initial_guess[lvol].get_zs(mm, nn) + + if not self.stellsym: + spec.allglobal.allrzrz[2, lvol, imn] = initial_guess[lvol].get_rs(mm, nn) + spec.allglobal.allrzrz[3, lvol, imn] = initial_guess[lvol].get_zc(mm, nn) + spec.allglobal.num_modes = imn + 1 + # possibly cleaner way to do this (tested to work Chris Smiet 11/9/2023): +# n_indices, m_values = np.indices([2*si.ntor+1, si.mpol+1]) #get indices for array +# n_values = n_indices - si.ntor # offset indices to correspond with mode numbers +# index_mask = ~np.logical_and(m_values==0, n_values < 0) # twiddle ~ NOT, pick elements +# numel = np.sum(index_mask) # number of trues in mask, i.e. chosen elements +# spec.allglobal.mmrzrz[:numel] = m_values[index_mask] # skip elements, make 1 +# spec.allglobal.nnrzrz[:numel] = n_values[index_mask] # skip m==0 n<0 elements +# +# initial_guess = [s.to_RZFourier() for s in self.initial_guess] +# for lvol, surface in enumerate(initial_guess): +# # EXPLAIN: surface.rc is m,n indexed, transpose, apply above mask which unravels to 1d +# spec.allglobal.allrzrz[0, lvol, :numel] = surface.rc.transpose()[index_mask] +# spec.allglobal.allrzrz[1, lvol, :numel] = surface.zs.transpose()[index_mask] +# if not self.stellsym: +# spec.allglobal.allrzrz[2, lvol, :numel] = surface.rs.transpose()[index_mask] +# spec.allglobal.allrzrz[3, lvol, :numel] = surface.zc.transpose()[index_mask] @property def boundary(self): @@ -291,6 +371,56 @@ def boundary(self): """ return self._boundary + @boundary.setter + def boundary(self, boundary): + """ + Setter for the geometry of the plasma boundary + """ + if not self.freebound: + if self._boundary is not boundary: + self.remove_parent(self._boundary) + self._boundary = boundary + self.append_parent(boundary) + return + else: # in freeboundary case, plasma boundary is is not a parent + self._boundary = boundary + + @property + def normal_field(self): + """ + Getter for the normal field + """ + return self._normal_field + + @normal_field.setter + def normal_field(self, normal_field): + """ + Setter for the normal field + """ + if not self.freebound: + raise ValueError('Normal field can only be set in freeboundary case') + if not isinstance(normal_field, NormalField): + raise ValueError('Input should be a NormalField or CoilNormalField') + if self._normal_field is not normal_field: + self.remove_parent(self._normal_field) + if self._computational_boundary is not normal_field.surface: + normal_field.surface = self._computational_boundary + self._normal_field = normal_field + self.append_parent(normal_field) + return + + @property + def computational_boundary(self): + """ + Getter for the computational boundary. + Same as the plasma boundary in the non-free-boundary case, + gives the surface on which Vns and Vnc are defined in the free-boundary case. + + Returns: + SurfaceRZFourier instance representing the plasma boundary + """ + return self._computational_boundary + @property def pressure_profile(self): """ @@ -617,6 +747,39 @@ def helicity_profile(self, helicity_profile): self.append_parent(helicity_profile) self.need_to_run_code = True + def activate_profile(self, longname): + """ + Take a profile from the inputlist, and make it optimizable in + simsopt. + Args: + longname: string, either + - 'pressure' + - 'volume_current' + - 'interface_current' + - 'iota' + - 'oita' + - 'mu' + - 'pflux' + - 'tflux' + - 'helicity' + """ + profile_dict = { + 'pressure': {'specname': 'pressure', 'cumulative': False, 'length': self.nvol}, + 'volume_current': {'specname': 'ivolume', 'cumulative': True, 'length': self.nvol}, + 'interface_current': {'specname': 'isurf', 'cumulative': False, 'length': self.nvol-1}, + 'helicity': {'specname': 'helicity', 'cumulative': False, 'length': self.mvol}, + 'iota': {'specname': 'iota', 'cumulative': False, 'length': self.mvol}, + 'oita': {'specname': 'oita', 'cumulative': False, 'length': self.mvol}, + 'mu': {'specname': 'mu', 'cumulative': False, 'length': self.mvol}, + 'pflux': {'specname': 'pflux', 'cumulative': True, 'length': self.mvol}, + 'tflux': {'specname': 'tflux', 'cumulative': True, 'length': self.mvol} + } + + profile_data = self.inputlist.__getattribute__(profile_dict[longname]['specname'])[0:profile_dict[longname]['length']] + profile = ProfileSpec(profile_data, cumulative=profile_dict[longname]['cumulative'], psi_edge=self.inputlist.phiedge) + profile.unfix_all() + self.__setattr__(longname + '_profile', profile) + def set_profile(self, longname, lvol, value): """ This function is used to set the pressure, currents, iota, oita, @@ -678,17 +841,6 @@ def get_profile(self, longname, lvol): return profile.f(lvol) - @boundary.setter - def boundary(self, boundary): - """ - Setter for the geometry of the plasma boundary - """ - - if self._boundary is not boundary: - self.remove_parent(self._boundary) - self._boundary = boundary - self.append_parent(boundary) - def recompute_bell(self, parent=None): self.need_to_run_code = True @@ -714,9 +866,29 @@ def set_dofs(self, x): ] for p in profiles: if p is not None: - p.phiedge = x[0] + p.psi_edge = x[0] - def init(self, filename: str): + @property + def poloidal_current_amperes(self): + """ + return the total poloidal current in Amperes, + i.e. the current that must be carried by the coils that link the plasma + poloidally + """ + return self.inputlist.curpol/(mu_0) + + def _clear_f90wrap_array_caches(self): + """ + Clear the f90wrap array caches. This is necessary when a new file is read after SPEC has run before. + + See https://github.com/hiddenSymmetries/simsopt/pull/431 + + This function is for addressing the issue https://github.com/hiddenSymmetries/simsopt/issues/357 + """ + spec.allglobal._arrays = {} + spec.inputlist._arrays = {} + + def _init_fortran_state(self, filename: str): """ Initialize SPEC fortran state from an input file. @@ -742,6 +914,11 @@ def init(self, filename: str): def run(self, update_guess: bool = True): """ Run SPEC, if needed. + The FORTRAN state has been created in the __init__ method (there can be only one for one python kernel), and the values in the FORTRAN memory are updated by python. + + The most important values are in inputlist and allglobal. + + Note: Since all calls to fortran memory access the same memory, there is no difference between spec.inputlist and self.inputlist (or even my_other_spec2.inputlist). Args: - update_guess: boolean. If True, initial guess will be updated with @@ -751,11 +928,17 @@ def run(self, update_guess: bool = True): if not self.need_to_run_code: logger.info("run() called but no need to re-run SPEC.") return - logger.info("Preparing to run SPEC.") + if self.mpi.proc0_groups: + logger.info(f"Group {self.mpi.group}: Preparing to run SPEC.") self.counter += 1 si = self.inputlist # Shorthand + # if freeboundary, store plasma-caused boundary field to re-set if run does not converge + if self.freebound: + initial_bns = np.copy(si.bns) + initial_bnc = np.copy(si.bnc) + # Check that number of volumes in internal memory is consistent with # the input file if self.nvol != si.nvol: @@ -770,21 +953,19 @@ def run(self, update_guess: bool = True): boundary_RZFourier = self.boundary.to_RZFourier() # Transfer boundary data to fortran: - si.rbc[:, :] = 0.0 - si.zbs[:, :] = 0.0 + si.rbc[:, :] = self.array_translator(boundary_RZFourier.rc, style='simsopt').as_spec + si.zbs[:, :] = self.array_translator(boundary_RZFourier.zs, style='simsopt').as_spec si.rbs[:, :] = 0.0 si.zbc[:, :] = 0.0 - mpol_capped = np.min([boundary_RZFourier.mpol, si.mmpol]) - ntor_capped = np.min([boundary_RZFourier.ntor, si.mntor]) - stellsym = bool(si.istellsym) - logger.debug(f"In run, si.istellsym = {si.istellsym} stellsym = {stellsym}") - for m in range(mpol_capped + 1): - for n in range(-ntor_capped, ntor_capped + 1): - si.rbc[n + si.mntor, m + si.mmpol] = boundary_RZFourier.get_rc(m, n) - si.zbs[n + si.mntor, m + si.mmpol] = boundary_RZFourier.get_zs(m, n) - if not stellsym: - si.rbs[n + si.mntor, m + si.mmpol] = boundary_RZFourier.get_rs(m, n) - si.zbc[n + si.mntor, m + si.mmpol] = boundary_RZFourier.get_zc(m, n) + if not self.stellsym: + si.rbs[:, :] = self.array_translator(boundary_RZFourier.rs, style='simsopt').as_spec + si.zbc[:, :] = self.array_translator(boundary_RZFourier.zc, style='simsopt').as_spec + + # transfer normal field to fortran: + if self.freebound: + si.vns[:, :] = self.array_translator(self.normal_field.get_vns_asarray(), style='simsopt').as_spec + if not self.stellsym: + si.vnc[:, :] = self.array_translator(self.normal_field.get_vnc_asarray(), style='simsopt').as_spec # Set the coordinate axis using the lrzaxis=2 feature: si.lrzaxis = 2 @@ -798,46 +979,17 @@ def run(self, update_guess: bool = True): si.zac[0:mn] = self.axis['zac'] # Set initial guess - if self.initial_guess is not None: - # Set all modes to zero - spec.allglobal.mmrzrz[:] = 0 - spec.allglobal.nnrzrz[:] = 0 - spec.allglobal.allrzrz[:] = 0 - - # transform to SurfaceRZFourier if necessary - initial_guess = [s.to_RZFourier() for s in self.initial_guess] - - # Loop on modes - imn = -1 # counter - for mm in range(0, si.mpol+1): - for nn in range(-si.ntor, si.ntor+1): - if mm == 0 and nn < 0: - continue - - imn += 1 - - spec.allglobal.mmrzrz[imn] = mm - spec.allglobal.nnrzrz[imn] = nn - - # Populate inner plasma boundaries - for lvol in range(0, self.nvol-1): - spec.allglobal.allrzrz[0, lvol, imn] = initial_guess[lvol].get_rc(mm, nn) - spec.allglobal.allrzrz[1, lvol, imn] = initial_guess[lvol].get_zs(mm, nn) - - if not si.istellsym: - spec.allglobal.allrzrz[2, lvol, imn] = initial_guess[lvol].get_rs(mm, nn) - spec.allglobal.allrzrz[3, lvol, imn] = initial_guess[lvol].get_zc(mm, nn) - - # Populate plasma boundary - if si.lfreebound: - si.rbc[si.mntor+nn, si.mmpol+mm] = initial_guess[self.nvol-1].get_rc(mm, nn) - si.zbs[si.mntor+nn, si.mmpol+mm] = initial_guess[self.nvol-1].get_zs(mm, nn) - - if not si.istellsym: - si.rbs[si.mntor+nn, si.mmpol+mm] = initial_guess[self.nvol-1].get_rs(mm, nn) - si.zbc[si.mntor+nn, si.mmpol+mm] = initial_guess[self.nvol-1].get_zc(mm, nn) - - spec.allglobal.num_modes = imn + 1 + if self.initial_guess is not None: # note: self.initial_guess is None for workers!! only leaders read and do_stuff with initial_guess + self._set_spec_initial_guess() # workers get the info through a broadcast. this line fails if workers get a guess set + + # write the boundary which is a guess in freeboundary + if self.freebound: + boundaryguess = self.initial_guess[-1].to_RZFourier() + si.rbc[:] = self.array_translator(boundaryguess.rc, style='simsopt').as_spec + si.zbs[:] = self.array_translator(boundaryguess.zs, style='simsopt').as_spec + if not self.stellsym: + si.rbs[:] = self.array_translator(boundaryguess.rs, style='simsopt').as_spec + si.zbc[:] = self.array_translator(boundaryguess.zc, style='simsopt').as_spec # Set profiles from dofs if self.pressure_profile is not None: @@ -936,6 +1088,7 @@ def run(self, update_guess: bool = True): logger.debug('About to call check_inputs') spec.allglobal.check_inputs() logger.debug('About to call broadcast_inputs') + self.mpi.comm_groups.Barrier() # Barrier to ensure all groups are ready to broadcast. spec.allglobal.broadcast_inputs() logger.debug('About to call preset') spec.preset() @@ -968,8 +1121,8 @@ def run(self, update_guess: bool = True): if self.verbose: traceback.print_exc() raise ObjectiveFailure("SPEC did not run successfully.") - - logger.info("SPEC run complete.") + if self.mpi.proc0_groups: + logger.info(f"{filename}: SPEC run complete.") # Barrier so workers do not try to read the .h5 file before it # is finished: self.mpi.comm_groups.Barrier() @@ -983,16 +1136,35 @@ def run(self, update_guess: bool = True): raise ObjectiveFailure( "Unable to read results following SPEC execution") - logger.info("Successfully loaded SPEC results.") + if self.mpi.proc0_groups: + logger.info("Successfully loaded SPEC results.") self.need_to_run_code = False # Deal with unconverged equilibria - these are excluded by # the optimizer, and the objective function is set to a large number if self.results.output.ForceErr > self.tolerance: + # re-set the boundary field guess before failing, so the next run does not start unconverged. + if self.freebound: + logger.info(f"run {filename}: Force balance unconverged, re-setting boundary field guess and raising ObjectiveFailure") + si.bns, si.bnc = initial_bns, initial_bnc raise ObjectiveFailure( 'SPEC could not find force balance' ) - + + # check for picard convergence once SPEC is updated + if self.freebound: + try: + if self.results.output.BnsErr > self.inputlist.gbntol: + if self.mpi.proc0_groups: + logger.info(f"run {filename}: Picard iteration unconverged, re-setting boundary field guess and raising ObjectiveFailure") + si.bns, si.bnc = initial_bns, initial_bnc + raise ObjectiveFailure( + 'Picard iteration failed to reach convergence, increase mfreeits or gbntol' + ) + except AttributeError: + logger.info('Picard convergence not checked, please update SPEC version 3.22 or higher') + print("Picard convergence not checked, please update SPEC version 3.22 or higher") + # Save geometry as initial guess for next iterations if update_guess: new_guess = None @@ -1002,7 +1174,7 @@ def run(self, update_guess: bool = True): ] for ii, (mm, nn) in enumerate(zip(self.results.output.im, self.results.output.in_)): - nnorm = (nn / si.nfp).astype('int') + nnorm = (nn / si.nfp).astype('int') # results.output.in_ is fourier number for 1 field period for reasons... for lvol in range(0, self.mvol-1): new_guess[lvol].set_rc(mm, nnorm, self.results.output.Rbc[lvol+1, ii]) new_guess[lvol].set_zs(mm, nnorm, self.results.output.Zbs[lvol+1, ii]) @@ -1016,8 +1188,9 @@ def run(self, update_guess: bool = True): axis['zas'] = self.results.output.Zbs[0, 0:si.ntor+1] self.axis = copy.copy(axis) - # Enforce SPEC to use initial guess - self.initial_guess = new_guess + # Enforce SPEC to use initial guess, but only group leaders have the necessary arrays + if self.mpi.proc0_groups: + self.initial_guess = new_guess # set this here? or make whole above monster proc0-exclusive self.inputlist.linitialize = 0 # Group leaders handle deletion of files: @@ -1027,6 +1200,8 @@ def run(self, update_guess: bool = True): if (not self.keep_all_files) and (self.mpi.group > 0): os.remove(filename + '.sp.h5') os.remove(filename + '.sp.end') + if os.path.exists('.' + filename + '.sp.DF'): + os.remove('.' + filename + '.sp.DF') # Delete the previous output file, if desired: for file_to_delete in self.files_to_delete: @@ -1038,6 +1213,8 @@ def run(self, update_guess: bool = True): self.counter > 0) and (not self.keep_all_files): self.files_to_delete.append(filename + '.sp.h5') self.files_to_delete.append(filename + '.sp.end') + if self.mpi.proc0_groups: + logger.info(f"file {filename} completed run sucessfully") def volume(self): """ @@ -1052,6 +1229,98 @@ def iota(self): """ self.run() return self.results.transform.fiota[1, 0] + + def array_translator(self, array=None, style='spec'): + """ + Returns a SpecFourierArray object to help transforming between + arrays using SPEC conventions and simsopt conventions. + + create from an array of either style by setting style='spec' or + style='simsopt' upon creation. + + The created class will have properties: + - as_spec: returns the SPEC style array + - as_simsopt: returns the simsopt style array + and setters to set the arrays from either convention after creation. + + SPEC arrays are zero-padded and indexed with Fortran conventions, + in python that means the m=0,n=0 component is at [nmtor,mmpol]. + + use: + vnc = myspec.array_translator(myspec.inputlist.rbc) + mysurf.rc = vnc.as_simsopt + """ + if style == 'spec': + return Spec.SpecFourierArray(self, array) + elif style == 'simsopt': + obj = Spec.SpecFourierArray(self) + obj.as_simsopt = array # setter function + return obj + else: + raise ValueError(f'array style:{style} not supported, options: [simsopt, spec]') + + # Inner class of Spec + class SpecFourierArray: + """ + Helper class to make index-jangling easier and hide away + strange SPEC array sizing and indexing conventions. + + SPEC stores Fourier series in a very large zero-padded + array with Fortran variable indexing and 'n,m' indexing + convention. + In python this means the m=0,n=0 component is at [nmtor,mmpol]. + + + """ + def __init__(self, outer_spec, array=None): + self._outer_spec = outer_spec + self._mmpol = outer_spec.inputlist.mmpol + self._mntor = outer_spec.inputlist.mntor + if array is None: + array = np.zeros((2*self._mntor+1, 2*self._mmpol+1)) + self._array = np.array(array) + + @property + def as_spec(self): + return self._array + + @as_spec.setter + def as_spec(self, array): + if array.shape != [2*self.mntor+1, 2*self.mmpol+1]: + raise ValueError('Array size is not consistent witn mmpol and mntor') + self._array = array + + @property + def as_simsopt(self, ntor=None, mpol=None): + """ + get a simsopt style 'n,m' intexed non-padded array + from the SpecFourierArray. + """ + if ntor is None: + ntor = self._outer_spec.ntor + if mpol is None: + mpol = self._outer_spec.mpol + n_start = self._mntor - ntor + n_end = self._mntor + ntor + 1 + m_start = self._mmpol + m_end = self._mmpol + mpol + 1 + return self._array[n_start:n_end, m_start:m_end].transpose() # switch n and m + + @as_simsopt.setter + def as_simsopt(self, array, ntor=None, mpol=None): + """ + set the SpecFourierArray from a simsopt style 'm,n' intexed non-padded array + """ + if mpol is None: + mpol = array.shape[0]-1 + if ntor is None: + ntor = int((array.shape[1]-1)/2) + n_start = self._mntor - ntor + n_end = self._mntor + ntor + 1 + m_start = self._mmpol + m_end = self._mmpol + mpol + 1 + self._array = np.zeros((2*self._mntor+1, 2*self._mmpol+1)) + self._array[n_start:n_end, m_start:m_end] = array.transpose() class Residue(Optimizable): @@ -1069,7 +1338,7 @@ class Residue(Optimizable): rtol: the relative tolerance of the integrator """ - def __init__(self, spec, pp, qq, vol=1, theta=0, s_guess=None, s_min=-1.0, + def __init__(self, spec, pp, qq, vol=1, theta=0., s_guess=None, s_min=-1.0, s_max=1.0, rtol=1e-9): # if not spec_found: if spec is None: @@ -1113,7 +1382,7 @@ def J(self): self.spec.run() if not self.mpi.proc0_groups: - logger.info( + logger.debug( "This proc is skipping Residue.J() since it is not a group leader.") return 0.0 else: @@ -1133,5 +1402,6 @@ def J(self): if self.fixed_point is None: raise ObjectiveFailure("Residue calculation failed") + logger.info(f"group {self.mpi.group} found residue {self.fixed_point.GreenesResidue} for {self.pp}/{self.qq} in {self.spec.allglobal.ext}") return self.fixed_point.GreenesResidue diff --git a/tests/field/test_normal_field.py b/tests/field/test_normal_field.py index ea639b563..c17ee816b 100644 --- a/tests/field/test_normal_field.py +++ b/tests/field/test_normal_field.py @@ -93,9 +93,13 @@ def test_getter_setter(self): normal_field.set_vns(m=1, n=-1, value=1) self.assertEqual(normal_field.get_vns(m=1, n=-1), 1) + i, j = normal_field.get_index_in_array(m=1, n=-1) + self.assertEqual(normal_field.vns[i, j], 1) normal_field.set_vnc(m=2, n=1, value=0.5) self.assertEqual(normal_field.get_vnc(m=2, n=1), 0.5) + i, j = normal_field.get_index_in_array(m=2, n=1) + self.assertEqual(normal_field.vnc[i, j], 0.5) def test_check_mn(self): """ @@ -208,3 +212,138 @@ def test_fixed_range(self): self.assertTrue(normal_field.is_fixed(ii)) else: self.assertTrue(normal_field.is_free(ii)) + + @unittest.skipIf(py_spec is None, "py_spec not found") + def test_from_spec_object(self): + """ + test classmethod to instantiate from an existing SPEC object + """ + from simsopt.mhd import Spec + # Init from SPEC input file + filename = os.path.join(TEST_DIR, 'Dommaschk.sp') + spec = Spec(filename) + normal_field = NormalField.from_spec_object(spec) + normal_field_2 = NormalField.from_spec(filename) + + self.assertAlmostEqual(normal_field.get_vns(m=3, n=-1), 1.71466651E-03) + self.assertAlmostEqual(normal_field.get_vns(m=5, n=1), -3.56991494E-05) + + self.assertTrue(np.allclose(normal_field.local_full_x, normal_field_2.local_full_x)) + + @unittest.skipIf(py_spec is None, "py_spec not found") + def test_fail_for_fixedb(self): + """ + test if instantiation fails if spec is not freeboundary + """ + from simsopt.mhd import Spec + # Init from SPEC input file + spec = Spec() + with self.assertRaises(ValueError): + NormalField.from_spec_object(spec) + + def test_wrong_index(self): + """ + test if accessing a wrong m or n raises an error + """ + normal_field = NormalField(stellsym=False, mpol=3, ntor=2) + with self.assertRaises(ValueError): + normal_field.get_index_in_array(m=4, n=1) + with self.assertRaises(ValueError): + normal_field.get_index_in_array(m=4, n=1, even=True) + with self.assertRaises(ValueError): + normal_field.get_index_in_array(m=1, n=3) + with self.assertRaises(ValueError): + normal_field.get_index_in_array(m=1, n=3, even=True) + with self.assertRaises(ValueError): + normal_field.get_index_in_array(m=0, n=-1) + with self.assertRaises(ValueError): + normal_field.get_index_in_array(m=0, n=-1, even=True) + with self.assertRaises(ValueError): + normal_field.get_index_in_array(m=3, n=-3) + with self.assertRaises(ValueError): + normal_field.get_index_in_array(m=3, n=-3, even=True) + with self.assertRaises(ValueError): + normal_field.get_index_in_array(m=4, n=-3) + with self.assertRaises(ValueError): + normal_field.get_index_in_array(m=4, n=-3, even=True) + self.assertIsInstance(normal_field.get_index_in_array(m=0, n=0, even=True), list) + normal_field2 = NormalField(stellsym=True, mpol=3, ntor=2) + with self.assertRaises(ValueError): + normal_field2.get_index_in_array(m=0, n=0) + + def test_asarray_getter_setter_raises(self): + """ + test that if wrong bounds are given, an error is raised + """ + normal_field = NormalField(stellsym=False, mpol=3, ntor=2) + with self.assertRaises(ValueError): + normal_field.get_vns_asarray(mpol=4, ntor=5) + with self.assertRaises(ValueError): + normal_field.get_vns_asarray(mpol=3, ntor=4) + with self.assertRaises(ValueError): + normal_field.get_vnc_asarray(mpol=4, ntor=5) + with self.assertRaises(ValueError): + normal_field.get_vnc_asarray(mpol=3, ntor=4) + with self.assertRaises(ValueError): + normal_field.get_vns_vnc_asarray(mpol=4, ntor=5) + with self.assertRaises(ValueError): + normal_field.get_vns_vnc_asarray(mpol=3, ntor=4) + with self.assertRaises(ValueError): + normal_field.set_vns_asarray(np.zeros((3, 7)), mpol=2, ntor=3) + with self.assertRaises(ValueError): + normal_field.set_vnc_asarray(np.zeros((4, 7)), mpol=2, ntor=3) + with self.assertRaises(ValueError): + normal_field.set_vns_vnc_asarray(np.zeros((4, 7)), np.zeros((4, 7)), mpol=2, ntor=3) + + + @unittest.skipIf(py_spec is None, "py_spec not found") + def test_get_set_vns_vnc_asarray(self): + """ + test the array-wise getter and setter functions for the + Vns and Vnc arrays + """ + filename = os.path.join(TEST_DIR, 'M16N08.sp') + normal_field = NormalField.from_spec(filename) + vns = normal_field.get_vns_asarray() + vnc = normal_field.get_vnc_asarray() + vns2, vnc2 = normal_field.get_vns_vnc_asarray() + self.assertTrue(np.allclose(vns, vns2)) + self.assertTrue(np.allclose(vnc, vnc2)) + vns3 = np.copy(vnc) + vnc3 = np.copy(vnc) + i, j = normal_field.get_index_in_array(m=3, n=-1) + vns3[i, j] = 0.5 + normal_field.set_vns_asarray(vns3) + self.assertEqual(normal_field.get_vns(m=3, n=-1), 0.5) + dof_index = normal_field.get_index_in_dofs(m=3, n=-1) + self.assertEqual(normal_field.local_full_x[dof_index], 0.5) + i, j = normal_field.get_index_in_array(m=2, n=1) + vnc3[i, j] = 1.5 + normal_field.set_vnc_asarray(vnc3) + self.assertEqual(normal_field.get_vnc(m=2, n=1), 1.5) + dof_index = normal_field.get_index_in_dofs(m=2, n=1, even=True) + self.assertEqual(normal_field.local_full_x[dof_index], 1.5) + normal_field.set_vns_vnc_asarray(vns2, vnc2) + + @unittest.skipIf(py_spec is None, "py_spec not found") + def test_get_real_space_field(self): + """ + test the conversion to real space + """ + filename = os.path.join(TEST_DIR, 'Dommaschk.sp') + + normal_field = NormalField.from_spec(filename) + real_space_field = normal_field.get_real_space_field() + self.assertTrue(real_space_field is not None) + self.assertEqual(real_space_field.shape, (normal_field.surface.quadpoints_phi.size, normal_field.surface.quadpoints_theta.size)) + + def test_vns_vnc_setter(self): + normal_field = NormalField() + with self.assertRaises(AttributeError): + normal_field.vns = 1 + with self.assertRaises(AttributeError): + normal_field.vnc = 1 + + +if __name__ == '__main__': + unittest.main() \ No newline at end of file diff --git a/tests/geo/test_surface_rzfourier.py b/tests/geo/test_surface_rzfourier.py index 538cf17d2..fce6fec00 100755 --- a/tests/geo/test_surface_rzfourier.py +++ b/tests/geo/test_surface_rzfourier.py @@ -732,6 +732,44 @@ def test_make_rotating_ellipse_iota(self): np.testing.assert_array_less(0, eq.wout.iotaf) np.testing.assert_allclose(eq.mean_iota(), 0.4291137962772453, rtol=1e-6) + def test_fourier_transform_scalar(self): + """ + Test the Fourier transform of a field on a surface. + """ + s = SurfaceRZFourier(mpol=4, ntor=5) + s.rc[0, 0] = 1.3 + s.rc[1, 0] = 0.4 + s.zs[1, 0] = 0.2 + + # Create the grid of quadpoints: + phi2d, theta2d = np.meshgrid(2 * np.pi * s.quadpoints_phi, + 2 * np.pi * s.quadpoints_theta, + indexing='ij') + + # create a test field where only Fourier elements [m=2, n=3] + # and [m=4,n=5] are nonzero: + field = 0.8 * np.sin(2*theta2d - 3*s.nfp*phi2d) + 0.2*np.sin(4*theta2d - 5*s.nfp*phi2d)+ 0.7*np.cos(3*theta2d - 3*s.nfp*phi2d) + + # Transform the field to Fourier space: + ft_sines, ft_cosines = s.fourier_transform_scalar(field, stellsym=False) + self.assertAlmostEqual(ft_sines[2, 3+s.ntor], 0.8) + self.assertAlmostEqual(ft_sines[4, 5+s.ntor], 0.2) + self.assertAlmostEqual(ft_cosines[3, 3+s.ntor], 0.7) + + # Test that all other elements are close to zero + sines_mask = np.ones_like(ft_sines, dtype=bool) + cosines_mask = np.copy(sines_mask) + sines_mask[2, 3 + s.ntor] = False + sines_mask[4, 5 + s.ntor] = False + cosines_mask[3, 3 + s.ntor] = False + self.assertEqual(np.all(np.abs(ft_sines[sines_mask]) < 1e-10), True) + self.assertEqual(np.all(np.abs(ft_cosines[cosines_mask]) < 1e-10), True) + + # Transform back to real space: + field2 = s.inverse_fourier_transform_scalar(ft_sines, ft_cosines, stellsym=False, normalization=1/2*np.pi**2) + + # Check that the result is the same as the original field: + np.testing.assert_allclose(field/2*np.pi**2, field2) class SurfaceRZPseudospectralTests(unittest.TestCase): def test_names(self): diff --git a/tests/mhd/test_spec.py b/tests/mhd/test_spec.py index 0ab5fb536..179a7e2af 100755 --- a/tests/mhd/test_spec.py +++ b/tests/mhd/test_spec.py @@ -2,8 +2,11 @@ import os import unittest + import numpy as np from monty.tempfile import ScratchDir +from scipy.constants import mu_0 +from simsopt._core.util import ObjectiveFailure try: import spec as spec_mod @@ -22,6 +25,7 @@ from simsopt.geo import SurfaceGarabedian from simsopt.mhd import ProfileSpec +from simsopt.field import NormalField from simsopt.objectives import LeastSquaresProblem from simsopt.solve import least_squares_serial_solve @@ -91,6 +95,27 @@ def test_init_freeboundary_nonstellsym(self): self.assertAlmostEqual(s.normal_field.get_vns(3, -1), -1.269776831212886e-04, places) self.assertAlmostEqual(s.normal_field.get_vnc(1, 0), 1.924871538367248e-04, places) self.assertAlmostEqual(s.normal_field.get_vnc(1, -2), 4.070523669489626e-04, places) + + def test_normal_field_setter(self): + """ + Try creating a Spec instance from a freeboundary file that is also + non-stellarator symmetric. + Check value of normal field + """ + + filename = os.path.join(TEST_DIR, 'M16N08.sp') + + with ScratchDir("."): + s = Spec(filename) + surface = s.boundary + old_normal = s.normal_field + new_normal = NormalField(s.nfp, stellsym=s.stellsym, mpol=s.mpol, ntor=s.ntor, surface=surface) + s.normal_field = new_normal + self.assertAlmostEqual(s.normal_field.get_vns(0,1), 0) # set to zeros + self.assertIs(s.normal_field.surface, s._computational_boundary) # normal field surface is set to spec computational boundary. + self.assertIsNot(old_normal, new_normal) + + def test_init_freeboundary(self): """ @@ -196,6 +221,93 @@ def test_set_profile_cumulative(self): else: self.assertEqual(s.get_profile('volume_current', lvol), 1) + def test_activate_profiles(self): + """ + test activate all profiles and confirm that DOFs are + correctly added + """ + profiles = ['pressure', + 'volume_current', + 'interface_current', + 'iota', + 'oita', + 'mu', + 'pflux', + 'tflux', + 'helicity'] + + for profile in profiles: + with ScratchDir("."): + s = Spec.default_freeboundary(copy_to_pwd=True) + startdofs = len(s.x) + self.assertIsNone(s.__getattribute__(profile+'_profile')) + s.activate_profile(profile) + self.assertIsNotNone(s.__getattribute__(profile+'_profile')) + if profile in ['tflux', 'pflux', 'mu', 'oita', 'iota', 'helicity']: + # test that the 'mvol' length profiles have been increased by two + self.assertEqual(len(s.x), startdofs+2) + elif profile in ['volume_current', 'pressure']: + # test that the 'nvol' length profiles have been increased by one + self.assertEqual(len(s.x), startdofs+1) + elif profile in ['interface_current']: + # test that the surface current profile has not been increased + self.assertEqual(len(s.x), startdofs) + else: + raise ValueError(f"Profile {profile} not recognized") + + + def test_freeboundary_default(self): + """ + test the default freeboundary file, and re-set if Picard + iteration fails + """ + with ScratchDir("."): + s = Spec.default_freeboundary(copy_to_pwd=True) + startingboundary = np.copy(s.inputlist.bns) + # run sucessfully as default should be converged + s.run() + # ask to generate guess, not enough freeboundary iterations + s.inputlist.linitialize = 2 + s.inputlist.lautoinitbn = 1 + s.inputlist.mfreeits = 2 + s.recompute_bell() + try: + s.run() + except ObjectiveFailure: + self.assertTrue(np.all(s.inputlist.bns == startingboundary)) + else: + raise ValueError("ObjectiveFailure not raised") + # give enough freeboundary iterations, let ir run successfully + s.inputlist.mfreeits = 10 + s.recompute_bell() + s.run() + + def test_array_translator(self): + """ + Test the array_translator method to convert a SPEC style array to simsopt style array. + """ + spec = Spec() + array = spec.inputlist.rbc + translator = spec.array_translator(array) + translator2 = spec.array_translator(array, style='spec') + translator3 = spec.array_translator(translator.as_simsopt, style='simsopt') + + self.assertTrue(np.all(array == translator.as_spec)) + self.assertTrue(np.all(array == translator2.as_spec)) + self.assertTrue(np.all(array == translator3.as_spec)) + self.assertTrue(np.all(translator.as_simsopt == translator2.as_simsopt)) + self.assertTrue(np.all(translator2.as_simsopt == translator3.as_simsopt)) + # test that the shape of the array is correct: + self.assertEqual(translator2.as_simsopt.shape, (spec.inputlist.ntor+1, (2*spec.inputlist.mpol)+1)) + + def test_poloidal_current_amperes(self): + """ + Test the poloidal current in amperes + """ + filename = os.path.join(TEST_DIR, 'RotatingEllipse_Nvol8.sp') + s = Spec(filename) + self.assertAlmostEqual(s.poloidal_current_amperes, s.inputlist.curpol/mu_0) + def test_integrated_stellopt_scenarios_1dof(self): """ This script implements the "1DOF_circularCrossSection_varyR0_targetVolume"