Fixing backward Euler - using scalar_product for rhs

Author: mikaem

src/jaxfun/galerkin/ChebyshevU.py

@@ -17,7 +17,7 @@ class ChebyshevU(Jacobi):
 
     Implements a Chebyshev basis via the Jacobi formulation with
     alpha = beta = 1/2. Provides several evaluation kernels:
-      * eval_basis_function: Single T_i evaluation (iterative).
+      * eval_basis_function: Single U_i(x) evaluation (iterative).
       * eval_basis_functions: Vectorized generation of all modes < N.
 
     The series expansion (degree N-1):
@@ -226,21 +226,17 @@ def gn(self, n: Symbol | int) -> Expr:
 def matrices(
     test: tuple[ChebyshevU, int], trial: tuple[ChebyshevU, int]
 ) -> sparse.BCOO | None:
-    """Sparse operator matrices between test/trial Chebyshev modes.
+    """Sparse operator matrices between test/trial ChebyshevU modes.
 
     Constructs (possibly rectangular) sparse differentiation / mass-like
     matrices for combinations of test index i and trial index j flags:
 
         (i, j):
           (0,0): Diagonal mass-matrix.
-          (0,1): First derivative.
-          (1,0): Transpose of (0,1).
-          (0,2): Second derivative.
-          (2,0): Transpose of (0,2).
 
     Args:
-        test: Tuple (v, i) with Chebyshev space v and number of derivatives i.
-        trial: Tuple (u, j) with Chebyshev space u and number of derivatives j.
+        test: Tuple (v, i) with ChebyshevU space v and number of derivatives i.
+        trial: Tuple (u, j) with ChebyshevU space u and number of derivatives j.
 
     Returns:
         jax.experimental.sparse.BCOO or None if combination unsupported.

src/jaxfun/galerkin/Fourier.py

@@ -149,7 +149,7 @@ def scalar_product(self, c: Array) -> Array:
         Returns:
             Coefficients scaled by 2π / domain_factor.
         """
-        out = jnp.fft.fft(c, norm="forward") * 2 * jnp.pi / self.domain_factor
+        out = jnp.fft.fft(c, norm="forward") * 2 * jnp.pi / float(self.domain_factor)
         if len(c) > self.N:
             return out[self.wavenumbers()]
         return out

src/jaxfun/integrators/backward_euler.py

@@ -11,8 +11,8 @@
 class BackwardEuler(BaseIntegrator):
     """First-order implicit Euler for linear terms (IMEX for nonlinear terms)."""
 
-    _system_diag: Array | None = nnx.data(None)
-    _system_matrix: Array | None = nnx.data(None)
+    _system_diag: Array | None = None
+    _system_matrix: Array | None = None
 
     def setup(self, dt: float) -> None:
         """Precompute the implicit system matrix for the given step size."""
@@ -33,7 +33,7 @@ def step(self, u_hat: Array, dt: float, N: Padding = None) -> Array:
         if self.linear_forcing is not None:
             rhs = rhs + dt * jnp.asarray(self.linear_forcing)
         if self.has_nonlinear:
-            rhs = rhs + dt * self.apply_mass(self.nonlinear_rhs(u_hat, N))
+            rhs = rhs + dt * self.nonlinear_rhs_scalar_product(u_hat, N)
 
         if self._system_diag is not None:
             return rhs / self._system_diag

src/jaxfun/integrators/base.py

@@ -230,6 +230,17 @@ def nonlinear_rhs(self, uh: Array, N: Padding = None) -> Array:
         assert self._nonlinear_evaluator is not None
         return self.functionspace.forward(self._nonlinear_evaluator(uh, N))
 
+    def nonlinear_rhs_scalar_product(self, uh: Array, N: Padding = None) -> Array:
+        """Return the nonlinear contribution in coefficient space.
+
+        Do *not* apply the mass inverse to complete the forward transformation,
+        because the mass inverse may be required elsewhere.
+        """
+        if not self.has_nonlinear:
+            return jnp.zeros_like(uh)
+        assert self._nonlinear_evaluator is not None
+        return self.functionspace.scalar_product(self._nonlinear_evaluator(uh, N))
+
     def linear_rhs(self, uh: Array) -> Array:
         """Return the linear contribution after applying the inverse mass matrix."""
         rhs = jnp.zeros_like(uh)