Support C->R case

optimistix/_complex.py

@@ -0,0 +1,59 @@
+from typing import Any, Generic
+
+import equinox as eqx
+import jax.lax as lax
+import jax.numpy as jnp
+import jax.tree_util as jtu
+from jaxtyping import Array, PyTree
+
+from ._custom_types import Args, Aux, Fn, Out, Y
+from ._solution import Solution
+
+
+class _RealPair(eqx.Module):
+    real: Array
+    imag: Array
+
+
+class _FromRealFn(eqx.Module, Generic[Y, Out, Aux]):
+    fn: Fn[Y, Out, Aux]
+    convert_output: bool = eqx.field(static=True)
+
+    def __call__(self, y: PyTree, args: Args) -> tuple[Out, Aux]:
+        out, aux = self.fn(_real_to_complex(y), args)
+        if self.convert_output:
+            out = _complex_to_real(out)
+        return out, aux
+
+
+def _is_complex_leaf(x: Any) -> bool:
+    return hasattr(x, "dtype") and jnp.issubdtype(x.dtype, jnp.complexfloating)
+
+
+def _has_complex(x: PyTree) -> bool:
+    return any(_is_complex_leaf(leaf) for leaf in jtu.tree_leaves(x))
+
+
+def _complex_to_real(x: PyTree) -> PyTree:
+    def _to_real_pair(leaf):
+        if _is_complex_leaf(leaf):
+            return _RealPair(leaf.real, leaf.imag)
+        return leaf
+
+    return jtu.tree_map(_to_real_pair, x)
+
+
+def _real_to_complex(x: PyTree) -> PyTree:
+    def _to_complex(leaf):
+        if isinstance(leaf, _RealPair):
+            return lax.complex(leaf.real, leaf.imag)
+        return leaf
+
+    return jtu.tree_map(
+        _to_complex, x, is_leaf=lambda leaf: isinstance(leaf, _RealPair)
+    )
+
+
+def _restore_solution(solution: Solution) -> Solution:
+    value = _real_to_complex(solution.value)
+    return eqx.tree_at(lambda sol: sol.value, solution, value)

optimistix/_least_squares.py

@@ -6,6 +6,12 @@
 from jaxtyping import PyTree, Scalar
 
 from ._adjoint import AbstractAdjoint, ImplicitAdjoint
+from ._complex import (
+    _complex_to_real,
+    _FromRealFn,
+    _has_complex,
+    _restore_solution,
+)
 from ._custom_types import Args, Aux, Fn, MaybeAuxFn, Out, SolverState, Y
 from ._iterate import AbstractIterativeSolver, iterative_solve
 from ._minimise import AbstractMinimiser, minimise
@@ -38,7 +44,10 @@ class _ToMinimiseFn(eqx.Module, Generic[Y, Out, Aux]):
 
     def __call__(self, y: Y, args: Args) -> tuple[Scalar, Aux]:
         residual, aux = self.residual_fn(y, args)
-        return 0.5 * sum_squares(residual), aux
+        # A residual pytree may mix real and complex leaves. Their squared norms are
+        # all real, but reducing them can still encounter mixed dtypes before `.real`.
+        with jax.numpy_dtype_promotion("standard"):
+            return 0.5 * sum_squares(residual), aux
 
 
 @eqx.filter_jit
@@ -60,6 +69,9 @@ def least_squares(
 
     Given a nonlinear function `fn(y, args)` which returns a pytree of residuals,
     this returns the solution to $\min_y \sum_i \textrm{fn}(y, \textrm{args})_i^2$.
+    When the inputs are complex-valued, complex inputs and residuals are represented
+    internally by their real and imaginary parts. Thus the Jacobian and linear solves
+    are performed over the reals.
 
     **Arguments:**
 
@@ -114,12 +126,16 @@ def least_squares(
         )
     else:
         y0 = jtu.tree_map(inexact_asarray, y0)
+        complex_to_real = _has_complex(y0)
+        if complex_to_real:
+            y0 = _complex_to_real(y0)
+            fn = _FromRealFn(fn, convert_output=True)
         fn = eqx.filter_closure_convert(fn, y0, args)  # pyright: ignore
         fn = cast(Fn[Y, Out, Aux], fn)
         f_struct, aux_struct = fn.out_struct  # pyright: ignore[reportFunctionMemberAccess]
         if options is None:
             options = {}
-        return iterative_solve(
+        solution = iterative_solve(
             fn,
             solver,
             y0,
@@ -133,3 +149,6 @@ def least_squares(
             aux_struct=aux_struct,
             rewrite_fn=_rewrite_fn,
         )
+        if complex_to_real:
+            solution = _restore_solution(solution)
+        return solution

optimistix/_minimise.py

@@ -7,6 +7,12 @@
 from jaxtyping import PyTree, Scalar
 
 from ._adjoint import AbstractAdjoint, ImplicitAdjoint
+from ._complex import (
+    _complex_to_real,
+    _FromRealFn,
+    _has_complex,
+    _restore_solution,
+)
 from ._custom_types import Aux, Fn, MaybeAuxFn, SolverState, Y
 from ._iterate import AbstractIterativeSolver, iterative_solve
 from ._misc import inexact_asarray, NoneAux, OutAsArray
@@ -52,6 +58,8 @@ def minimise(
     """Minimise a function.
 
     This minimises a nonlinear function `fn(y, args)` which returns a scalar value.
+    Complex-valued inputs are represented internally by their real and imaginary
+    parts, so that differentiation and linear solves are performed over the reals.
 
     **Arguments:**
 
@@ -88,6 +96,10 @@ def minimise(
     if not has_aux:
         fn = NoneAux(fn)  # pyright: ignore
     fn = OutAsArray(fn)
+    complex_to_real = _has_complex(y0)
+    if complex_to_real:
+        y0 = _complex_to_real(y0)
+        fn = _FromRealFn(fn, convert_output=False)
     fn = eqx.filter_closure_convert(fn, y0, args)  # pyright: ignore
     fn = cast(Fn[Y, Scalar, Aux], fn)
     f_struct, aux_struct = fn.out_struct  # pyright: ignore[reportFunctionMemberAccess]
@@ -103,7 +115,7 @@ def minimise(
             "minimisation function must output a single floating-point scalar."
         )
 
-    return iterative_solve(
+    solution = iterative_solve(
         fn,
         solver,
         y0,
@@ -117,3 +129,6 @@ def minimise(
         f_struct=f_struct,
         rewrite_fn=_rewrite_fn,
     )
+    if complex_to_real:
+        solution = _restore_solution(solution)
+    return solution

tests/conftest.py

@@ -1,3 +1,4 @@
+import equinox as eqx
 import equinox.internal as eqxi
 import jax
 import pytest
@@ -8,6 +9,14 @@
 jax.config.update("jax_numpy_dtype_promotion", "strict")
 
 
+@pytest.fixture(scope="module", autouse=True)
+def clear_caches_between_modules():
+    """Bound memory from compiled executables in the full parametrised test suite."""
+    yield
+    eqx.clear_caches()
+    jax.clear_caches()
+
+
 @pytest.fixture
 def getkey():
     return eqxi.GetKey()