Repository: f0uriest/quadax

Issue:
Complex numbers support
Thanks a lot for your excellent library! 👏

When integrating a complex-valued function, for example just just changing the README example to `fun = lambda t: t * jnp.log(1 + t) * 1j`, I get a `TypeError`:
```python
TypeError: true_fun and false_fun output must have identical types, got
('ShapedArray(complex64[])', 'ShapedArray(float32[], weak_type=True)', 'DIFFERENT ShapedArray(complex64[]) vs. ShapedArray(float32[])', 'DIFFERENT ShapedArray(complex64[]) vs. ShapedArray(float32[])').
```

This happens even when the norm is real-valued (which is [the default](https://github.com/f0uriest/quadax/blob/main/quadax/adaptive.py#L411C51-L411C89) anyway I think).

Hints from issue comments:
@f0uriest: Thanks, this was due to some type mismatches when shortcutting a zero length interval. Should be fixed by #20 

Solution patch:
diff --git a/README.rst b/README.rst
index 300ab73..1ff3729 100644
--- a/README.rst
+++ b/README.rst
@@ -11,7 +11,7 @@ quadax is a library for numerical quadrature and integration using JAX.
 - ``vmap``-able, ``jit``-able, differentiable.
 - Scalar or vector valued integrands.
 - Finite or infinite domains with discontinuities or singularities within the domain of integration.
-- Globally adaptive Gauss-Konrod and Clenshaw-Curtis quadrature for smooth integrands (similar to ``scipy.integrate.quad``)
+- Globally adaptive Gauss-Kronrod and Clenshaw-Curtis quadrature for smooth integrands (similar to ``scipy.integrate.quad``)
 - Adaptive tanh-sinh quadrature for singular or near singular integrands.
 - Quadrature from sampled values using trapezoidal and Simpsons methods.
 
diff --git a/docs/api.rst b/docs/api.rst
index 69768d4..2685242 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -11,7 +11,7 @@ Adaptive integration of a callable function or method
     :toctree: _api/
     :recursive:
 
-    quadgk    -- General purpose integration using Gauss-Konrod scheme
+    quadgk    -- General purpose integration using Gauss-Kronrod scheme
     quadcc    -- General purpose integration using Clenshaw-Curtis scheme
     quadts    -- General purpose integration using tanh-sinh (aka double exponential) scheme
     romberg   -- Adaptive trapezoidal integration with Richardson extrapolation
@@ -27,7 +27,7 @@ Quadrature Rules
     :template: class.rst
 
     AbstractQuadratureRule -- Abstract base class for all quadrature rules
-    GaussKronrodRule       -- Fixed order integration over finite interval using Gauss-Konrod scheme
+    GaussKronrodRule       -- Fixed order integration over finite interval using Gauss-Kronrod scheme
     ClenshawCurtisRule     -- Fixed order integration over finite interval using Clenshaw-Curtis scheme
     TanhSinhRule           -- Fixed order integration over finite interval using tanh-sinh (aka double exponential) scheme
 
diff --git a/quadax/adaptive.py b/quadax/adaptive.py
index c929698..a374a9f 100644
--- a/quadax/adaptive.py
+++ b/quadax/adaptive.py
@@ -39,7 +39,7 @@ def quadgk(
     order=21,
     norm=jnp.inf,
 ):
-    """Global adaptive quadrature using Gauss-Konrod rule.
+    """Global adaptive quadrature using Gauss-Kronrod rule.
 
     Integrate fun from `interval[0]` to `interval[-1]` using a h-adaptive scheme with
     error estimate. Breakpoints can be specified in `interval` where integration
@@ -429,12 +429,14 @@ def adaptive_quadrature(
     state = {}
     state["neval"] = 0  # number of evaluations of local quadrature rule
     state["ninter"] = len(interval) - 1  # current number of intervals
-    state["r_arr"] = jnp.zeros((max_ninter, *shape))  # local results from each interval
+    state["r_arr"] = jnp.zeros(
+        (max_ninter, *shape), f.dtype
+    )  # local results from each interval
     state["e_arr"] = jnp.zeros(max_ninter)  # local error est. from each interval
     state["a_arr"] = jnp.zeros(max_ninter)  # start of each interval
     state["b_arr"] = jnp.zeros(max_ninter)  # end of each interval
     state["s_arr"] = jnp.zeros(
-        (max_ninter, *shape)
+        (max_ninter, *shape), f.dtype
     )  # global est. of I from n intervals
     state["a_arr"] = state["a_arr"].at[: state["ninter"]].set(interval[:-1])
     state["b_arr"] = state["b_arr"].at[: state["ninter"]].set(interval[1:])
@@ -442,7 +444,7 @@ def adaptive_quadrature(
     state["roundoff2"] = 0  # for keeping track of roundoff errors
     state["status"] = 0  # error flag
     state["err_bnd"] = 0.0  # error bound we're trying to reach
-    state["area"] = jnp.zeros(shape)  # current best estimate for I
+    state["area"] = jnp.zeros(shape, f.dtype)  # current best estimate for I
     state["err_sum"] = 0.0  # current estimate for error in I
 
     def init_body(i, state_):
diff --git a/quadax/fixed_order.py b/quadax/fixed_order.py
index be7a8a3..b98babc 100644
--- a/quadax/fixed_order.py
+++ b/quadax/fixed_order.py
@@ -118,7 +118,7 @@ def integrate(
         def truefun():
             f = jax.eval_shape(vfun, jnp.array(0.0))
             z = jnp.zeros(f.shape, f.dtype)
-            return z, 0.0, z, z
+            return z, self.norm(z), jnp.abs(z), jnp.abs(z)
 
         def falsefun():
 


Test patch:
diff --git a/tests/test_adaptive.py b/tests/test_adaptive.py
index 4ef0acc..f0b74c9 100644
--- a/tests/test_adaptive.py
+++ b/tests/test_adaptive.py
@@ -87,11 +87,17 @@
         "interval": [0, 1, 2],
         "val": -4,
     },
+    # problem 16 - complex function
+    {
+        "fun": lambda t: t * jnp.log(1 + t) * 1j,
+        "interval": [0, 1],
+        "val": 0.25j,
+    },
 ]
 
 
 class TestQuadGK:
-    """Tests for Gauss-Konrod quadrature."""
+    """Tests for Gauss-Kronrod quadrature."""
 
     def _base(self, i, tol, fudge=1, **kwargs):
         prob = example_problems[i]
@@ -210,6 +216,12 @@ def test_prob15(self):
         self._base(14, 1e-8)
         self._base(14, 1e-12)
 
+    def test_prob16(self):
+        """Test for example problem #16."""
+        self._base(16, 1e-4)
+        self._base(16, 1e-8)
+        self._base(16, 1e-12)
+
 
 class TestQuadCC:
     """Tests for Clenshaw-Curtis quadrature."""
@@ -331,6 +343,12 @@ def test_prob15(self):
         self._base(14, 1e-8)
         self._base(14, 1e-12)
 
+    def test_prob16(self):
+        """Test for example problem #16."""
+        self._base(16, 1e-4)
+        self._base(16, 1e-8)
+        self._base(16, 1e-12)
+
 
 class TestQuadTS:
     """Tests for adaptive tanh-sinh quadrature."""
@@ -452,6 +470,12 @@ def test_prob15(self):
         self._base(14, 1e-8)
         self._base(14, 1e-12)
 
+    def test_prob16(self):
+        """Test for example problem #16."""
+        self._base(16, 1e-4)
+        self._base(16, 1e-8)
+        self._base(16, 1e-12)
+
 
 class TestRombergTS:
     """Tests for tanh-sinh quadrature with adaptive refinement."""
@@ -567,6 +591,12 @@ def test_prob15(self):
         self._base(14, 1e-8)
         self._base(14, 1e-12)
 
+    def test_prob16(self):
+        """Test for example problem #16."""
+        self._base(16, 1e-4)
+        self._base(16, 1e-8)
+        self._base(16, 1e-12)
+
 
 class TestRomberg:
     """Tests for Romberg's method (only for well behaved integrands)."""
@@ -670,3 +700,9 @@ def test_prob15(self):
         self._base(14, 1e-4)
         self._base(14, 1e-8)
         self._base(14, 1e-12)
+
+    def test_prob16(self):
+        """Test for example problem #16."""
+        self._base(16, 1e-4)
+        self._base(16, 1e-8)
+        self._base(16, 1e-12)