Add 4th order derivative transforms

desc/basis.py

@@ -608,7 +608,7 @@
         For L>0, the indexing scheme defines order of the basis functions:
 
         ``'ansi'``: ANSI indexing fills in the pyramid with triangles of
-        decreasing size, ending in a triagle shape. For L == M,
+        decreasing size, ending in a triangle shape. For L == M,
         the traditional ANSI pyramid indexing is recovered. For L>M, adds rows
         to the bottom of the pyramid, increasing L while keeping M constant,
         giving a "house" shape.
@@ -651,7 +651,7 @@
             For L>0, the indexing scheme defines order of the basis functions:
 
             ``'ansi'``: ANSI indexing fills in the pyramid with triangles of
-            decreasing size, ending in a triagle shape. For L == M,
+            decreasing size, ending in a triangle shape. For L == M,
             the traditional ANSI pyramid indexing is recovered. For L>M, adds rows
             to the bottom of the pyramid, increasing L while keeping M constant,
             giving a "house" shape.
@@ -830,7 +830,7 @@
         For L>0, the indexing scheme defines order of the basis functions:
 
         ``'ansi'``: ANSI indexing fills in the pyramid with triangles of
-        decreasing size, ending in a triagle shape. For L == M,
+        decreasing size, ending in a triangle shape. For L == M,
         the traditional ANSI pyramid indexing is recovered. For L>M, adds rows
         to the bottom of the pyramid, increasing L while keeping M constant,
         giving a "house" shape.
@@ -875,7 +875,7 @@
             For L>0, the indexing scheme defines order of the basis functions:
 
             ``'ansi'``: ANSI indexing fills in the pyramid with triangles of
-            decreasing size, ending in a triagle shape. For L == M,
+            decreasing size, ending in a triangle shape. For L == M,
             the traditional ANSI pyramid indexing is recovered. For L>M, adds rows
             to the bottom of the pyramid, increasing L while keeping M constant,
             giving a "house" shape.
@@ -1252,37 +1252,49 @@
     beta = 0
     n = (l - m) // 2
     s = (-1) ** n
+    jacobi_arg = 1 - 2 * r**2
     if dr == 0:
-        out = r**m * _jacobi(n, alpha, beta, 1 - 2 * r**2, 0)
+        out = r**m * _jacobi(n, alpha, beta, jacobi_arg, 0)
     elif dr == 1:
-        f = _jacobi(n, alpha, beta, 1 - 2 * r**2, 0)
-        df = _jacobi(n, alpha, beta, 1 - 2 * r**2, 1)
+        f = _jacobi(n, alpha, beta, jacobi_arg, 0)
+        df = _jacobi(n, alpha, beta, jacobi_arg, 1)
         out = m * r ** jnp.maximum(m - 1, 0) * f - 4 * r ** (m + 1) * df
     elif dr == 2:
-        f = _jacobi(n, alpha, beta, 1 - 2 * r**2, 0)
-        df = _jacobi(n, alpha, beta, 1 - 2 * r**2, 1)
-        d2f = _jacobi(n, alpha, beta, 1 - 2 * r**2, 2)
+        f = _jacobi(n, alpha, beta, jacobi_arg, 0)
+        df = _jacobi(n, alpha, beta, jacobi_arg, 1)
+        d2f = _jacobi(n, alpha, beta, jacobi_arg, 2)
         out = (
-            m * (m - 1) * r ** jnp.maximum((m - 2), 0) * f
-            - 2 * 4 * m * r**m * df
-            + r**m * (16 * r**2 * d2f - 4 * df)
+            (m - 1) * m * r ** jnp.maximum(m - 2, 0) * f
+            - 4 * (2 * m + 1) * r**m * df
+            + 16 * r ** (m + 2) * d2f
         )
     elif dr == 3:
-        f = _jacobi(n, alpha, beta, 1 - 2 * r**2, 0)
-        df = _jacobi(n, alpha, beta, 1 - 2 * r**2, 1)
-        d2f = _jacobi(n, alpha, beta, 1 - 2 * r**2, 2)
-        d3f = _jacobi(n, alpha, beta, 1 - 2 * r**2, 3)
+        f = _jacobi(n, alpha, beta, jacobi_arg, 0)
+        df = _jacobi(n, alpha, beta, jacobi_arg, 1)
+        d2f = _jacobi(n, alpha, beta, jacobi_arg, 2)
+        d3f = _jacobi(n, alpha, beta, jacobi_arg, 3)
         out = (
             (m - 2) * (m - 1) * m * r ** jnp.maximum(m - 3, 0) * f
-            - 12 * (m - 1) * m * r ** jnp.maximum(m - 1, 0) * df
-            + 48 * r ** (m + 1) * d2f
+            - 12 * m**2 * r ** jnp.maximum(m - 1, 0) * df
+            + 48 * (m + 1) * r ** (m + 1) * d2f
             - 64 * r ** (m + 3) * d3f
-            + 48 * m * r ** (m + 1) * d2f
-            - 12 * m * r ** jnp.maximum(m - 1, 0) * df
+        )
+    elif dr == 4:
+        f = _jacobi(n, alpha, beta, jacobi_arg, 0)
+        df = _jacobi(n, alpha, beta, jacobi_arg, 1)
+        d2f = _jacobi(n, alpha, beta, jacobi_arg, 2)
+        d3f = _jacobi(n, alpha, beta, jacobi_arg, 3)
+        d4f = _jacobi(n, alpha, beta, jacobi_arg, 4)
+        out = (
+            (m - 3) * (m - 2) * (m - 1) * m * r ** jnp.maximum(m - 4, 0) * f
+            - 8 * m * (2 * m**2 - 3 * m + 1) * r ** jnp.maximum(m - 2, 0) * df
+            + 48 * (2 * m**2 + 2 * m + 1) * r**m * d2f
+            - 128 * (2 * m + 3) * r ** (m + 2) * d3f
+            + 256 * r ** (m + 4) * d4f
         )
     else:
         raise NotImplementedError(
-            "Analytic radial derivatives of zernike polynomials for order>3 "
+            "Analytic radial derivatives of zernike polynomials for order>4 "
             + "have not been implemented"
         )
     return s * jnp.where((l - m) % 2 == 0, out, 0)

desc/transform.py

@@ -58,8 +58,8 @@
         self._rcond = rcond if rcond is not None else "auto"
 
         if (
-            not np.all(self.grid.nodes[:, 2] == 0)
-            and not (self.grid.NFP == self.basis.NFP)
+            np.any(self.grid.nodes[:, 2] != 0)
+            and self.grid.NFP != self.basis.NFP
             and grid.node_pattern != "custom"
         ):
             warnings.warn(
@@ -71,29 +71,19 @@
                 )
             )
 
-        self._derivatives = self._get_derivatives(derivs)
-        self._sort_derivatives()
-        self._method = method
-
         self._built = False
         self._built_pinv = False
-        self._set_up()
+        self._derivatives = self._get_derivatives(derivs)
+        self._sort_derivatives()
+        # assign according to logic in setter function
+        self.method = method
+        # assign according to logic in property function
+        self._matrices = self.matrices
         if build:
             self.build()
         if build_pinv:
             self.build_pinv()
 
-    def _set_up(self):
-
-        self.method = self._method
-        self._matrices = {
-            "direct1": {
-                i: {j: {k: {} for k in range(4)} for j in range(4)} for i in range(4)
-            },
-            "fft": {i: {j: {} for j in range(4)} for i in range(4)},
-            "direct2": {i: {} for i in range(4)},
-        }
-
     def _get_derivatives(self, derivs):
         """Get array of derivatives needed for calculating objective function.
 
@@ -471,41 +461,36 @@
             return np.zeros(self.grid.num_nodes)
 
         if self.method == "direct1":
-            A = self.matrices["direct1"][dr][dt][dz]
+            A = self.matrices["direct1"].get(dr, {}).get(dt, {}).get(dz, {})
             if isinstance(A, dict):
                 raise ValueError(
                     colored("Derivative orders are out of initialized bounds", "red")
                 )
-            return jnp.matmul(A, c)
+            return A @ c
 
         elif self.method == "direct2":
-            A = self.matrices["fft"][dr][dt]
-            B = self.matrices["direct2"][dz]
-
+            A = self.matrices["fft"].get(dr, {}).get(dt, {})
+            B = self.matrices["direct2"].get(dz, {})
             if isinstance(A, dict) or isinstance(B, dict):
                 raise ValueError(
                     colored("Derivative orders are out of initialized bounds", "red")
                 )
             c_mtrx = jnp.zeros((self.num_lm_modes * self.num_n_modes,))
             c_mtrx = put(c_mtrx, self.fft_index, c).reshape((-1, self.num_n_modes))
-
-            cc = jnp.matmul(A, c_mtrx)
-            return jnp.matmul(cc, B.T).flatten(order="F")
+            cc = A @ c_mtrx
+            return (cc @ B.T).flatten(order="F")
 
         elif self.method == "fft":
-            A = self.matrices["fft"][dr][dt]
+            A = self.matrices["fft"].get(dr, {}).get(dt, {})
             if isinstance(A, dict):
                 raise ValueError(
                     colored("Derivative orders are out of initialized bounds", "red")
                 )
-
             # reshape coefficients
             c_mtrx = jnp.zeros((self.num_lm_modes * self.num_n_modes,))
             c_mtrx = put(c_mtrx, self.fft_index, c).reshape((-1, self.num_n_modes))
-
             # differentiate
             c_diff = c_mtrx[:, :: (-1) ** dz] * self.dk**dz * (-1) ** (dz > 1)
-
             # re-format in complex notation
             c_real = jnp.pad(
                 (self.num_z_nodes / 2)
@@ -520,10 +505,9 @@
                     jnp.fliplr(jnp.conj(c_real)),
                 )
             )
-
             # transform coefficients
             c_fft = jnp.real(jnp.fft.ifft(c_cplx))
-            return jnp.matmul(A, c_fft).flatten(order="F")
+            return (A @ c_fft).flatten(order="F")
 
     def fit(self, x):
         """Transform from physical domain to spectral using weighted least squares fit.
@@ -733,24 +717,25 @@
             # if we actually added derivatives and didn't build them, then its not built
             self._built = False
         if build:
-            # we don't update self._built here because it is still built from before
+            # we don't update self._built here because it is still built from before,
             # but it still might have unbuilt matrices from new derivatives
             self.build()
 
     @property
     def matrices(self):
         """dict: transform matrices such that x=A*c."""
-        return self.__dict__.setdefault(
-            "_matrices",
-            {
+        if not hasattr(self, "_matrices"):
+            # to allow computing of highest order derivative
+            n = np.amax(self.derivatives) + 1
+            self._matrices = {
                 "direct1": {
-                    i: {j: {k: {} for k in range(4)} for j in range(4)}
-                    for i in range(4)
+                    i: {j: {k: {} for k in range(n)} for j in range(n)}
+                    for i in range(n)
                 },
-                "fft": {i: {j: {} for j in range(4)} for i in range(4)},
-                "direct2": {i: {} for i in range(4)},
-            },
-        )
+                "fft": {i: {j: {} for j in range(n)} for i in range(n)},
+                "direct2": {i: {} for i in range(n)},
+            }
+        return self._matrices
 
     @property
     def num_nodes(self):