code: def heterog_laplacian(
      x: FourierSeries, c: FourierSeries, *, params=None
  ) -> FourierSeries:
      r"""Computes the position-varying laplacian using algorithm 4 of
      [[Johnson, 2011]](https://math.mit.edu/~stevenj/fft-deriv.pdf).
  
      Args:
        x (FourierSeries): Input field
        c (FourierSeries): Coefficient field
  
      Returns:
        The Laplacian of the field
      """
      assert x.dims == 1  # Laplacian only defined for scalar fields
  
      k_vec = params["k_vec"]
      u = x.params[..., 0]
      v = c.params[..., 0]
  
      def single_coordinate(axis, u):
          u = jnp.moveaxis(u, axis, -1)
          U = x._ffts[0](u, axis=-1)
          U_prime = 1j * k_vec[axis] * U
  
          # Handle Nyquist frequency
          N_on_L = 1 / u.domain.dx[axis]
          N = u.domain.N[axis]
          if N % 2 == 0:
              U_prime = U_prime.at[..., N // 2].set(
                  U[..., N // 2] * N_on_L * jnp.pi * 1j
              )
  
          # Multiply with heterogeneous field
          u_prime = x._ffts[1](U_prime, axis=-1, n=u.shape[-1])
          c_uprime = v * u_prime
  
          # Get the second derivative
          V = x._ffts[0](c_uprime, axis=-1)
          V_prime = 1j * k_vec[axis] * V
  
          # Handle Nyquist frequency
          if N % 2 == 0:
              V_prime = V_prime.at[..., N // 2].set(
                  V[..., N // 2] * N_on_L * jnp.pi * 1j
              )
  
          # Return to space
          ddu = x._ffts[1](V_prime, axis=-1, n=u.shape[-1])
  
          return jnp.moveaxis(ddu, -1, axis)
  
      new_params = jnp.sum(
          jnp.stack(
              [single_coordinate(i, u) for i in range(x.domain.ndim)],
              axis=-1,
          ),
          axis=-1,
          keepdims=True,
      )
      return FourierSeries(new_params, x.domain)
  
imports: from typing import Union
  import jax
  import numpy
  from jax import numpy
  from jaxdf.conv import fd_coefficients_fornberg, reflection_conv
  from jaxdf.core import operator
  from jaxdf.discretization import Continuous, FiniteDifferences, FourierSeries
label_desc: Use image node_id 20 for calling a global function with example usage: heterog_laplacian(x, c) and returns: FourierSeries, jnp