Poincare Boundary Condition

desc/basis.py

@@ -9,7 +9,7 @@
 
 from desc.backend import custom_jvp, fori_loop, gammaln, jit, jnp, sign
 from desc.io import IOAble
-from desc.utils import flatten_list
+from desc.utils import copy_coeffs, flatten_list
 
 __all__ = [
     "PowerSeries",
@@ -1726,3 +1726,182 @@ def zernike_norm(l, m):
 
     """
     return np.sqrt((2 * (l + 1)) / (np.pi * (1 + int(m == 0))))
+
+
+def get_basis_poincare(X_lmn_3D, basis_3D, zeta=0):
+    """Convert 3D FourierZernikeBasis to 2D ZernikePolynomial at zeta=0.
+
+    Takes a 3D FourierZernike basis and its coefficients X_lmn_3D and evaluates
+    the coefficients at the zeta=0 or Pi cross-section, returning a 2D ZernikePolynomial
+    basis and its coefficients X_lmn_2D. The output dimension will not be the same,
+    only toroidal mode will be N=0.
+
+    Parameters
+    ----------
+    X_lmn_3D : array, size [basis_3D.num_modes,1]
+        The Fourier-Zernike basis coefficients of the quantity X, that you wish to find
+        the 2D ZernikePolynomial basis corresponding to the quantity's value at the
+        zeta=0 cross-section
+    basis_3D : FourierZernikeBasis
+        The Fourier-Zernike basis corresponding to the coefficients passed.
+
+    Returns
+    -------
+    X_lmn_2D : array, size [basis_2D.num_modes,1]
+        The ZernikePolynomial basis coefficients of the quantity X, such that their
+        evaluation is the same as the input 3D basis when evaluated at zeta=0.
+    basis_2D : FourierZernikeBasis
+        The ZernikePolynomial basis corresponding to the coefficients output. The
+        radial resolution L and poloidal resolution M are equal to the max radial
+        and poloidal resolutions of the 3D basis passed as an input.
+    """
+    # Add up all the X_lm(n>=0) modes
+    # so that the quantity at the zeta=0 surface is described with just lm modes
+    # and get rid of the toroidal modes
+    modes_2D = []
+    X_lmn_2D = (
+        []
+    )  # these are corresponding to the 2D modes of the ZernikePolynomial Basis
+    modes_3D = basis_3D.modes
+    for i, mode in enumerate(modes_3D):
+        if zeta == 0:
+            if mode[-1] < 0:
+                pass  # we do not need the sin(zeta) modes as they = 0 at zeta=0
+            else:
+                if (mode[0], mode[1], 0) not in modes_2D:
+                    modes_2D.append((mode[0], mode[1], 0))
+                    l = mode[0]
+                    m = mode[1]
+
+                    inds = np.where(
+                        np.logical_and(
+                            (modes_3D[:, :2] == [l, m]).all(axis=1),
+                            modes_3D[:, 2] >= 0,
+                        )
+                    )[0]
+
+                    SUM = np.sum(X_lmn_3D[inds])
+                    X_lmn_2D.append(SUM)
+        elif zeta == np.pi:
+            if mode[-1] < 0:
+                pass  # we do not need the sin(zeta) modes as they = 0 at zeta=pi
+            else:
+                if (mode[0], mode[1], 0) not in modes_2D:
+                    modes_2D.append((mode[0], mode[1], 0))
+                    l = mode[0]
+                    m = mode[1]
+
+                    inds = np.where(
+                        np.logical_and(
+                            (modes_3D[:, :2] == [l, m]).all(axis=1),
+                            modes_3D[:, 2] % 2 == 0,
+                        )
+                    )[0]
+
+                    inds1 = np.where(
+                        np.logical_and(
+                            (modes_3D[:, :2] == [l, m]).all(axis=1),
+                            modes_3D[:, 2] % 2 == 1,
+                        )
+                    )[0]
+
+                    SUM = np.sum(X_lmn_3D[inds]) - np.sum(X_lmn_3D[inds1])
+                    X_lmn_2D.append(SUM)
+
+    X_lmn_2D = np.asarray(X_lmn_2D)
+    modes_2D = np.asarray(modes_2D)
+
+    Lmax = np.max(abs(modes_2D[:, 0]))
+    Mmax = np.max(abs(modes_2D[:, 1]))
+    LM_max = max(Lmax, Mmax)
+    basis_2D = ZernikePolynomial(
+        L=LM_max,
+        M=LM_max,
+        spectral_indexing=basis_3D._spectral_indexing,
+        sym=basis_3D.sym,
+    )
+    X_lmn_2D = copy_coeffs(X_lmn_2D, modes_2D, basis_2D.modes)
+    return X_lmn_2D, basis_2D
+
+
+def FourierZernike_to_FourierZernike_no_N_modes(X_lmn_3D, basis_3D, zeta=0):
+    """Converts 3D FourierZernikeBasis to the same basis but with no toroidal modes.
+
+    Takes a 3D FourierZernike basis and its coefficients X_lmn_3D and evaluates the
+    coefficients at the zeta=0 cross-section, returning the same 3D FourierZernike
+    basis but with zero toroidal dependence s.t. the zeta=0 XS is the same as the input
+    so X_lmn_no_N = 0 for N != 0. Input and output array sizes are the same.
+
+    Parameters
+    ----------
+    X_lmn_3D : array, size [basis_3D.num_modes,1]
+        The Fourier-Zernike basis coefficients of the quantity X, that you wish to
+        find the 2D ZernikePolynomial basis
+        corresponding to the quantity's value at the zeta=0 cross-section
+    basis_3D : FourierZernikeBasis
+        The Fourier-Zernike basis corresponding to the coefficients passed.
+
+    Returns
+    -------
+    X_lmn_no_N : array, size [basis_3D.num_modes,1]
+        The FourerZernike basis coefficients of the quantity X, such that their
+        evaluation is the same as the input 3D basis when evaluated at zeta=0.
+        The coefficients with nonzero toroidal modenumbers are all zero.
+    basis_3D : FourierZernikeBasis
+        The Fourier-Zernike basis corresponding to the coefficients passed.
+    """
+    if not (zeta == 0 or zeta == np.pi):
+        raise ValueError("Unvalid zeta value! Only 0 or pi is supported right now!")
+
+    # Add up all the X_lm(n>=0) modes
+    # so that the quantity at the zeta=0 surface is described with just lm modes
+    # and set any mode with nonzero toroidal mode numebr = 0
+    X_lmn_no_N = np.zeros_like(X_lmn_3D)
+    modes_no_N = []
+    modes_3D = basis_3D.modes
+    for i, mode in enumerate(modes_3D):
+        if zeta == 0:
+            if mode[-1] < 0:
+                pass  # we do not need the sin(zeta) modes as they = 0 at zeta=0
+            else:
+                if (mode[0], mode[1], 0) not in modes_no_N:
+                    modes_no_N.append((mode[0], mode[1], 0))
+                    l = mode[0]
+                    m = mode[1]
+
+                    inds = np.where(
+                        np.logical_and(
+                            (modes_3D[:, :2] == [l, m]).all(axis=1),
+                            modes_3D[:, 2] >= 0,
+                        )
+                    )[0]
+
+                    SUM = np.sum(X_lmn_3D[inds])
+                    X_lmn_no_N[i] = SUM
+        elif zeta == np.pi:
+            if mode[-1] < 0:
+                pass  # we do not need the sin(zeta) modes as they = 0 at zeta=pi
+            else:
+                if (mode[0], mode[1], 0) not in modes_no_N:
+                    modes_no_N.append((mode[0], mode[1], 0))
+                    l = mode[0]
+                    m = mode[1]
+
+                    inds = np.where(
+                        np.logical_and(
+                            (modes_3D[:, :2] == [l, m]).all(axis=1),
+                            modes_3D[:, 2] % 2 == 0,
+                        )
+                    )[0]
+
+                    inds1 = np.where(
+                        np.logical_and(
+                            (modes_3D[:, :2] == [l, m]).all(axis=1),
+                            modes_3D[:, 2] % 2 == 1,
+                        )
+                    )[0]
+
+                    SUM = np.sum(X_lmn_3D[inds]) - np.sum(X_lmn_3D[inds1])
+                    X_lmn_no_N[i] = SUM
+
+    return X_lmn_no_N, basis_3D

desc/compute/utils.py

@@ -1367,6 +1367,7 @@ def body(i, mins):
     "Za_n",
     "Rb_lmn",
     "Zb_lmn",
+    "Lb_lmn",
 )
 # map from profile name to equilibrium parameter name
 profile_names = {

desc/equilibrium/equilibrium.py

@@ -11,14 +11,15 @@
 from termcolor import colored
 
 from desc.backend import jnp
-from desc.basis import FourierZernikeBasis, fourier, zernike_radial
+from desc.basis import FourierZernikeBasis, fourier, get_basis_poincare, zernike_radial
 from desc.compat import ensure_positive_jacobian
 from desc.compute import compute as compute_fun
 from desc.compute import data_index
 from desc.compute.utils import get_data_deps, get_params, get_profiles, get_transforms
 from desc.geometry import (
     FourierRZCurve,
     FourierRZToroidalSurface,
+    PoincareSurface,
     Surface,
     ZernikeRZToroidalSection,
 )
@@ -143,7 +144,6 @@ class Equilibrium(IOAble, Optimizable):
         "_atomic_number",
         "_anisotropy",
         "_spectral_indexing",
-        "_bdry_mode",
         "_L_grid",
         "_M_grid",
         "_N_grid",
@@ -222,7 +222,7 @@ def __init__(
         self._Z_sym = "sin" if self.sym else False
 
         # surface
-        self._surface, self._bdry_mode = parse_surface(
+        self._surface = parse_surface(
             surface, self.NFP, self.sym, self.spectral_indexing
         )
 
@@ -423,6 +423,7 @@ def _sort_args(self, args):
             "Za_n",
             "Rb_lmn",
             "Zb_lmn",
+            "Lb_lmn",
         )
         assert sorted(args) == sorted(arg_order)
         return [arg for arg in arg_order if arg in args]
@@ -744,6 +745,48 @@ def get_axis(self):
         axis = FourierRZCurve(R_n, Z_n, modes_R, modes_Z, NFP=self.NFP, sym=self.sym)
         return axis
 
+    def set_poincare_equilibrium(self, zeta=0):
+        """Sets the equilibrium for solving Poincare BC from an existing equilibrium.
+
+        Parameters
+        ----------
+        self : Equilibrium
+            Some equilibrium to be used for creating Poincare equilibrium
+        zeta : float (optional)
+            Zeta angle at which the Poincare section will be fixed
+            Only 0 and Pi is supported for now
+
+        Returns
+        -------
+        eq_poincare : Equilibrium
+            Separate Equilibrium object to be used for Poincare BC problem
+        """
+        surf = self.get_surface_at(zeta=zeta / self.NFP)
+        Lb_lmn, Lb_basis = get_basis_poincare(self.L_lmn, self.L_basis, zeta)
+        surface = PoincareSurface(
+            surface=surf, L_lmn=Lb_lmn, modes_L=Lb_basis.modes, zeta=zeta
+        )
+
+        eq_poincare = Equilibrium(
+            surface=surface,
+            pressure=self.pressure,
+            iota=self.iota,
+            Psi=self.Psi,  # flux (in Webers) within the last closed flux surface
+            NFP=self.NFP,  # number of field periods
+            L=self.L,  # radial spectral resolution
+            M=self.M,  # poloidal spectral resolution
+            N=self.N,  # toroidal spectral resolution
+            L_grid=self.L_grid,  # real space radial resolution, slightly oversampled
+            M_grid=self.M_grid,  # real space poloidal resolution, slightly oversampled
+            N_grid=self.N_grid,  # real space toroidal resolution
+            sym=True,  # explicitly enforce stellarator symmetry
+            spectral_indexing=self._spectral_indexing,
+        )
+
+        eq_poincare.change_resolution(self.L, self.M, self.N)
+        eq_poincare.axis = eq_poincare.get_axis()
+        return eq_poincare
+
     def compute(
         self,
         names,
@@ -1165,11 +1208,6 @@ def sym(self):
         """bool: Whether this equilibrium is stellarator symmetric."""
         return self._sym
 
-    @property
-    def bdry_mode(self):
-        """str: Method for specifying boundary condition."""
-        return self._bdry_mode
-
     @optimizable_parameter
     @property
     def Psi(self):
@@ -1293,6 +1331,31 @@ def Zb_lmn(self):
     def Zb_lmn(self, Zb_lmn):
         self.surface.Z_lmn = Zb_lmn
 
+    @optimizable_parameter
+    @property
+    def Lb_lmn(self):
+        """ndarray: Spectral coefficients of Lambda at the boundary."""
+        if isinstance(self.surface, PoincareSurface):
+            return self.surface.L_lmn
+        else:
+            return np.empty(0)
+
+    @Lb_lmn.setter
+    def Lb_lmn(self, Lb_lmn):
+        self.surface.L_lmn = Lb_lmn
+
+    @property
+    def Lb_basis(self):
+        """ndarray: Spectral coefficients of Lambda at the boundary."""
+        if isinstance(self.surface, PoincareSurface):
+            return self.surface.L_basis
+        else:
+            return np.empty(0)
+
+    @Lb_basis.setter
+    def Lb_basis(self, Lb_basis):
+        self.surface.L_basis = Lb_basis
+
     @optimizable_parameter
     @property
     def Ra_n(self):
@@ -1785,11 +1848,6 @@ def solve(
                     "yellow",
                 )
             )
-        if self.bdry_mode == "poincare":
-            raise NotImplementedError(
-                "Solving equilibrium with poincare XS as BC is not supported yet "
-                + "on master branch."
-            )
 
         things, result = optimizer.optimize(
             self,