Repository: PlasmaControl/DESC

Issue:
`Transform`s built under jit cant interpolate from real to spectral space
The `fit` method fails because the `method` is `"jitable"`.

```python
    def fit(self, x):
        """Transform from physical domain to spectral using weighted least squares fit.
    
        Parameters
        ----------
        x : ndarray, shape(num_nodes,)
            values in real space at coordinates specified by grid
    
        Returns
        -------
        c : ndarray, shape(num_coeffs,)
            spectral coefficients in basis
    
        """
        if not self.built_pinv:
            raise RuntimeError(
                "Transform must be built with transform.build_pinv() before being used"
            )
    
        if self.method == "direct1":
            Ainv = self.matrices["pinv"]
            c = jnp.matmul(Ainv, x)
        elif self.method == "direct2":
            Ainv = self.matrices["pinvA"]
            Binv = self.matrices["pinvB"]
            yy = jnp.matmul(Ainv, x.reshape((-1, self.grid.num_zeta), order="F"))
            c = jnp.matmul(Binv, yy.T).T.flatten()[self.fft_index]
        elif self.method == "fft":
            Ainv = self.matrices["pinvA"]
            c_fft = jnp.matmul(Ainv, x.reshape((Ainv.shape[1], -1), order="F"))
            c_cplx = jnp.fft.fft(c_fft)
            c_unpad = c_cplx[:, 1 : (c_cplx.shape[1] - self.pad_dim - 1) // 2 + 1]
            c0 = c_cplx[:, :1].real / self.grid.num_zeta
            c2 = c_unpad.real / (self.grid.num_zeta / 2)
            c1 = -c_unpad.imag[:, ::-1] / (self.grid.num_zeta / 2)
            c_diff = jnp.hstack([c1, c0, c2])
            c = c_diff.flatten()[self.fft_index]
>       return c
               ^
E       UnboundLocalError: cannot access local variable 'c' where it is not associated with a value
E       --------------------
E       For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
```

Discussion:

Fix (patch):
diff --git a/desc/transform.py b/desc/transform.py
index 640a3df19..2401ac00c 100644
--- a/desc/transform.py
+++ b/desc/transform.py
@@ -482,7 +482,7 @@ class Transform(IOAble):
                 "Transform must be built with transform.build_pinv() before being used"
             )
 
-        if self.method == "direct1":
+        if self.method in ["direct1", "jitable"]:
             Ainv = self.matrices["pinv"]
             c = jnp.matmul(Ainv, x)
         elif self.method == "direct2":
@@ -532,7 +532,7 @@ class Transform(IOAble):
                 )
             )
 
-        if self.method == "direct1":
+        if self.method in ["direct1", "jitable"]:
             A = self.matrices["direct1"][0][0][0]
             return jnp.matmul(A.T, y)
 
@@ -719,7 +719,10 @@ class Transform(IOAble):
 
     @property
     def method(self):
-        """{``'direct1'``, ``'direct2'``, ``'fft'``}: method of computing transform."""
+        """{``'direct1'``, ``'direct2'``, ``'fft'``, ``'jitable'``}.
+
+        Transform compute method.
+        """
         return self.__dict__.setdefault("_method", "direct1")
 
     @method.setter


Test changes:
diff --git a/tests/test_transform.py b/tests/test_transform.py
index 97c4d1ce6..db9690303 100644
--- a/tests/test_transform.py
+++ b/tests/test_transform.py
@@ -491,8 +491,8 @@ class TestTransform:
         assert t.method == "direct1"
 
     @pytest.mark.unit
-    def test_fit_direct1(self):
-        """Test fitting with direct1 method."""
+    def test_fit_direct1_and_jitable(self):
+        """Test fitting with direct1 and jitable method."""
         basis = FourierZernikeBasis(3, 3, 2, spectral_indexing="ansi")
         grid = ConcentricGrid(3, 3, 2, node_pattern="ocs")
         transform = Transform(grid, basis, method="direct1", build_pinv=True)
@@ -501,6 +501,11 @@ class TestTransform:
         x = transform.transform(c)
         c1 = transform.fit(x)
         np.testing.assert_allclose(c, c1, atol=1e-12)
+        # also test jitable which is the same as direct1
+        transform = Transform(grid, basis, method="jitable", build_pinv=True)
+        x = transform.transform(c)
+        c1 = transform.fit(x)
+        np.testing.assert_allclose(c, c1, atol=1e-12)
 
     @pytest.mark.unit
     def test_fit_direct2(self):