Allow shared transforms between R, Z, lambda

desc/compute/utils.py

@@ -332,15 +332,28 @@ def get_transforms(keys, obj, grid, jitable=False, **kwargs):
     derivs = get_derivs(keys, obj, has_axis=grid.axis.size)
     transforms = {"grid": grid}
     for c in derivs.keys():
-        if hasattr(obj, c + "_basis"):
-            transforms[c] = Transform(
-                grid,
-                getattr(obj, c + "_basis"),
-                derivs=derivs[c],
-                build=True,
-                method=method,
-            )
-        elif c == "B":
+        if hasattr(obj, c + "_basis"):  # regular stuff like R, Z, lambda etc.
+            basis = getattr(obj, c + "_basis")
+            # first check if we already have a transform with a compatible basis
+            for transform in transforms.values():
+                if basis.equiv(getattr(transform, "basis", None)):
+                    ders = np.unique(
+                        np.vstack([derivs[c], transform.derivatives]), axis=0
+                    ).astype(int)
+                    # don't build until we know all the derivs we need
+                    transform.change_derivatives(ders, build=False)
+                    c_transform = transform
+                    break
+            else:  # if we didn't exit the loop early
+                c_transform = Transform(
+                    grid,
+                    basis,
+                    derivs=derivs[c],
+                    build=False,
+                    method=method,
+                )
+            transforms[c] = c_transform
+        elif c == "B":  # for fitting Boozer harmonics
             transforms["B"] = Transform(
                 grid,
                 DoubleFourierSeries(
@@ -350,11 +363,11 @@ def get_transforms(keys, obj, grid, jitable=False, **kwargs):
                     sym=obj.R_basis.sym,
                 ),
                 derivs=derivs["B"],
-                build=True,
+                build=False,
                 build_pinv=True,
                 method=method,
             )
-        elif c == "w":
+        elif c == "w":  # for fitting Boozer toroidal stream function
             transforms["w"] = Transform(
                 grid,
                 DoubleFourierSeries(
@@ -364,13 +377,18 @@ def get_transforms(keys, obj, grid, jitable=False, **kwargs):
                     sym=obj.Z_basis.sym,
                 ),
                 derivs=derivs["w"],
-                build=True,
+                build=False,
                 build_pinv=True,
                 method=method,
             )
-        elif c not in transforms:
+        elif c not in transforms:  # possible other stuff lumped in with transforms
             transforms[c] = getattr(obj, c)
 
+    # now build them
+    for t in transforms.values():
+        if hasattr(t, "build"):
+            t.build()
+
     return transforms
 
 

desc/transform.py

@@ -82,7 +82,6 @@ def __init__(
         self._method = method
         # assign according to logic in setter function
         self.method = method
-        self._matrices = self._get_matrices()
         if build:
             self.build()
         if build_pinv:
@@ -134,13 +133,14 @@ def _sort_derivatives(self):
 
     def _get_matrices(self):
         """Get matrices to compute all derivatives."""
-        n = np.amax(self.derivatives) + 1
+        n = 4  # hardcode max derivative order for now,
         matrices = {
             "direct1": {
-                i: {j: {k: {} for k in range(n)} for j in range(n)} for i in range(n)
+                i: {j: {k: {} for k in range(n + 1)} for j in range(n + 1)}
+                for i in range(n + 1)
             },
-            "fft": {i: {j: {} for j in range(n)} for i in range(n)},
-            "direct2": {i: {} for i in range(n)},
+            "fft": {i: {j: {} for j in range(n + 1)} for i in range(n + 1)},
+            "direct2": {i: {} for i in range(n + 1)},
         }
         return matrices
 
@@ -389,7 +389,7 @@ def build(self):
         if self.method in ["fft", "direct2"]:
             temp_d = np.hstack(
                 [self.derivatives[:, :2], np.zeros((len(self.derivatives), 1))]
-            )
+            ).astype(int)
             temp_modes = np.hstack([self.lm_modes, np.zeros((self.num_lm_modes, 1))])
             for d in temp_d:
                 self.matrices["fft"][d[0]][d[1]] = self.basis.evaluate(
@@ -398,7 +398,7 @@ def build(self):
         if self.method == "direct2":
             temp_d = np.hstack(
                 [np.zeros((len(self.derivatives), 2)), self.derivatives[:, 2:]]
-            )
+            ).astype(int)
             temp_modes = np.hstack(
                 [np.zeros((self.num_n_modes, 2)), self.n_modes[:, np.newaxis]]
             )