diff --git a/optimistix/_complex.py b/optimistix/_complex.py new file mode 100644 index 00000000..e2411444 --- /dev/null +++ b/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) diff --git a/optimistix/_least_squares.py b/optimistix/_least_squares.py index 6d3a6c99..80d7deb7 100644 --- a/optimistix/_least_squares.py +++ b/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 diff --git a/optimistix/_minimise.py b/optimistix/_minimise.py index 82070fb9..d2229df9 100644 --- a/optimistix/_minimise.py +++ b/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 diff --git a/tests/conftest.py b/tests/conftest.py index 8368845e..dfc72fae 100644 --- a/tests/conftest.py +++ b/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() diff --git a/tests/helpers.py b/tests/helpers.py index 5480952d..b1683e39 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -28,6 +28,33 @@ def tree_allclose(x, y, *, rtol=1e-5, atol=1e-8): return eqx.tree_equal(x, y, typematch=True, rtol=rtol, atol=atol) +def tree_as_dtype(tree, dtype): + def _cast(x): + if eqx.is_array(x) and jnp.issubdtype(x.dtype, jnp.inexact): + if jnp.issubdtype(dtype, jnp.floating) and jnp.issubdtype( + x.dtype, jnp.complexfloating + ): + x = x.real + return x.astype(dtype) + return x + + return jtu.tree_map(_cast, tree) + + +def make_nonreal(tree): + def _add_imaginary_part(x): + if eqx.is_array(x) and jnp.issubdtype(x.dtype, jnp.complexfloating): + return x + jnp.asarray(0.01j, dtype=x.dtype) + return x + + return jtu.tree_map(_add_imaginary_part, tree) + + +def norm_sq(x): + """Pointwise squared norm, for both real and complex arrays.""" + return jnp.real(x * jnp.conj(x)) + + def finite_difference_jvp(fn, primals, tangents, eps=None, **kwargs): assert jax.config.jax_enable_x64 # pyright: ignore out = fn(*primals, **kwargs) @@ -52,6 +79,52 @@ def finite_difference_jvp(fn, primals, tangents, eps=None, **kwargs): return out_ε, tangents_out +def implicit_minimise_jvp(fn, y, args, t_args): + """Compute an argmin JVP from the objective's stationarity equation. + + At a local minimiser ``y(args)``, the gradient satisfies + ``grad_y fn(y, args) = 0``. Differentiating this identity gives + ``Hessian_y(fn) @ dy = -d_args(grad_y(fn))``; this helper constructs both + sides with JAX and solves that linear system for ``dy``. + + Complex PyTrees are first flattened into an explicit R² representation by + concatenating their real and imaginary components. Thus the Hessian is an + ordinary real matrix and the reference does not rely on JAX's convention for + complex gradients or on any complex-to-real Lineax operator. This also makes + the reference independent of the particular iterative minimiser under test. + """ + flat_y, unravel_y = jfu.ravel_pytree(y) + if jnp.issubdtype(flat_y.dtype, jnp.complexfloating): + size = flat_y.size + real_y = jnp.concatenate((flat_y.real, flat_y.imag)) + + def from_real(flat_real_y): + flat_complex_y = jax.lax.complex(flat_real_y[:size], flat_real_y[size:]) + return unravel_y(flat_complex_y) + + else: + real_y = flat_y + + def from_real(flat_real_y): + return unravel_y(flat_real_y) + + def objective(flat_real_y, objective_args): + return fn(from_real(flat_real_y), objective_args) + + grad_fn = jax.grad(objective) + hessian = jax.jacfwd(grad_fn)(real_y, args) + if jtu.tree_leaves(t_args): + _, rhs = jax.jvp( + lambda objective_args: grad_fn(real_y, objective_args), + (args,), + (t_args,), + ) + tangent = jnp.linalg.solve(hessian, -rhs) + else: + tangent = jnp.zeros_like(real_y) + return from_real(tangent) + + # # NOTE: `GN` is shorthand for `gauss_newton`. We want to be sure we test every # branch of `GN=True` and `GN=False` for all of these solvers. @@ -317,13 +390,14 @@ def bowl(tree: PyTree[Array], args: Array): # Trivial quadratic bowl smoke test for convergence. (y, _) = jfu.ravel_pytree(tree) matrix = args - return y.T @ matrix @ y + return jnp.real(y.T.conj() @ matrix @ y) def diagonal_quadratic_bowl(tree: PyTree[Array], args: PyTree[Array]): # A diagonal quadratic bowl smoke test for convergence. - weight_vector = args - return (ω(tree).call(jnp.square) * (0.1 + weight_vector**ω)).ω + weight_vector = jtu.tree_map(jnp.abs, args) + squared_norm = ω(tree).call(norm_sq) + return (squared_norm * (0.1 + weight_vector**ω)).ω def rosenbrock(tree: PyTree[Array], args: Scalar): @@ -338,8 +412,10 @@ def _himmelblau(tree: PyTree[Array], args: PyTree): # Wiki (y, z) = tree const1, const2 = args - term1 = ((ω(y).call(jnp.square) + z**ω - const1) ** 2).ω - term2 = ((y**ω + ω(z).call(jnp.square) - const2) ** 2).ω + term1 = (ω(y).call(jnp.square) + z**ω - const1).ω + term1 = ω(term1).call(norm_sq).ω + term2 = (y**ω + ω(z).call(jnp.square) - const2).ω + term2 = ω(term2).call(norm_sq).ω return (term1**ω + term2**ω).ω @@ -347,8 +423,10 @@ def matyas(tree: PyTree[Array], args: PyTree): # Wiki (y, z) = tree const1, const2 = args - term1 = (const1 * (ω(y).call(jnp.square) + ω(z).call(jnp.square))).ω - term2 = (const2 * y**ω * z**ω).ω + const1 = jnp.abs(const1) + const2 = jnp.abs(const2) + term1 = (const1 * (ω(y).call(norm_sq) + ω(z).call(norm_sq))).ω + term2 = (const2 * (ω(y).call(jnp.conj) * z**ω).call(jnp.real)).ω return (term1**ω - term2**ω).ω @@ -366,9 +444,9 @@ def beale(tree: PyTree[Array], args: PyTree): # Wiki (y, z) = tree const1, const2, const3 = args - term1 = ((const1 - y**ω + y**ω * z**ω) ** 2).ω - term2 = ((const2 - y**ω + y**ω * ω(z).call(jnp.square)) ** 2).ω - term3 = ((const3 - y**ω + y**ω * ω(z).call(lambda x: x**3)) ** 2).ω + term1 = (const1 - y**ω + y**ω * z**ω).call(norm_sq).ω + term2 = (const2 - y**ω + y**ω * ω(z).call(jnp.square)).call(norm_sq).ω + term3 = (const3 - y**ω + y**ω * ω(z).call(lambda x: x**3)).call(norm_sq).ω return (term1**ω + term2**ω + term3**ω).ω @@ -401,7 +479,7 @@ def simple_nn(model_dynamic: PyTree[Array], args: PyTree): model = eqx.combine(model_dynamic, model_static) key = jr.PRNGKey(17) model_key, data_key = jr.split(key, 2) - x = jnp.linspace(0, 1, 100)[..., None] + x = jnp.linspace(0, 1, 100, dtype=data.dtype)[..., None] y = data**2 def loss(model, x, y): @@ -413,7 +491,35 @@ def loss(model, x, y): def square_minus_one(x: Array, args: PyTree): """A simple ||x||^2 - 1 function.""" - return jnp.sum(jnp.square(x)) - 1.0 + return jnp.sum(norm_sq(x)) - 1.0 + + +def complex_quadratic(x: Array, target: Array): + """A smooth C->R quadratic with an argument-dependent minimum.""" + diff = x - target + return jnp.sum(norm_sq(diff)) + + +def complex_to_real_residual(x: Array, target: Array): + """A C->R^2 residual with an argument-dependent root.""" + diff = x - target + return jnp.stack((diff.real, diff.imag)) + + +def holomorphic_residual(x: Array, target: Array): + """A holomorphic C->C residual.""" + return x - target + + +def antiholomorphic_residual(x: Array, target: Array): + """An anti-holomorphic C->C residual.""" + return jnp.conj(x) - target + + +def nonholomorphic_residual(x: Array, target: Array): + """A non-holomorphic C->(C, R) residual with mixed output dtypes.""" + diff = x - target + return diff + 0.25 * jnp.conj(diff), diff.real # @@ -428,7 +534,21 @@ def get_weights(model): return layer1.weight, layer1.bias, layer2.weight, layer2.bias -ffn_init = eqx.nn.MLP(in_size=1, out_size=1, width_size=8, depth=1, key=jr.PRNGKey(17)) +def complex_relu(x): + """ReLU on R, extended componentwise to C.""" + if jnp.issubdtype(x.dtype, jnp.complexfloating): + return jax.lax.complex(jax.nn.relu(x.real), jax.nn.relu(x.imag)) + return jax.nn.relu(x) + + +ffn_init = eqx.nn.MLP( + in_size=1, + out_size=1, + width_size=8, + depth=1, + activation=complex_relu, + key=jr.PRNGKey(17), +) weight1 = jnp.array( [ [3.39958394], @@ -479,6 +599,16 @@ def get_weights(model): [jr.normal(key, leaf.shape, leaf.dtype) ** 2 for leaf in leaves] ) +diagonal_bowl_init_complex = ( + {"a": (0.05 + 0.01j) * jnp.ones((2, 3, 3), dtype=jnp.complex128)}, + ((0.01 + 0.05j) * jnp.ones(2, dtype=jnp.complex128)), +) +leaves_complex, treedef_complex = jtu.tree_flatten(diagonal_bowl_init_complex) +key = jr.PRNGKey(17) +diagonal_bowl_args_complex = treedef.unflatten( + [jr.normal(key, leaf.shape, leaf.real.dtype) ** 2 for leaf in leaves_complex] +) + # neural net args ffn_data = jnp.linspace(0, 1, 100)[..., None] ffn_args = (ffn_static, ffn_data) @@ -490,6 +620,36 @@ def get_weights(model): diagonal_bowl_init, diagonal_bowl_args, ), + ( + diagonal_quadratic_bowl, + jnp.array(0.0), + diagonal_bowl_init_complex, + diagonal_bowl_args_complex, + ), + ( + complex_to_real_residual, + jnp.array(0.0), + jnp.array(1.0 + 1.1j, dtype=jnp.complex128), + jnp.array(-0.5 + 0.25j, dtype=jnp.complex128), + ), + ( + holomorphic_residual, + jnp.array(0.0), + jnp.array(1.0 + 1.1j, dtype=jnp.complex128), + jnp.array(-0.5 + 0.25j, dtype=jnp.complex128), + ), + ( + antiholomorphic_residual, + jnp.array(0.0), + jnp.array(1.0 + 1.1j, dtype=jnp.complex128), + jnp.array(-0.5 + 0.25j, dtype=jnp.complex128), + ), + ( + nonholomorphic_residual, + jnp.array(0.0), + jnp.array(1.0 + 1.1j, dtype=jnp.complex128), + jnp.array(-0.5 + 0.25j, dtype=jnp.complex128), + ), ( rosenbrock, jnp.array(0.0), @@ -559,6 +719,18 @@ def get_weights(model): ), # Problems with initial value of 0 (square_minus_one, jnp.array(-1.0), jnp.array(1.0), None), + ( + square_minus_one, + jnp.array(-1.0), + jnp.array(1.0 + 1.1j, dtype=jnp.complex128), + None, + ), + ( + complex_quadratic, + jnp.array(0.0), + jnp.array(1.0 + 1.1j, dtype=jnp.complex128), + jnp.array(-0.5 + 0.25j, dtype=jnp.complex128), + ), ) # ROOT FIND/FIXED POINT PROBLEMS @@ -887,20 +1059,23 @@ def forward_only_ode(k, args): dy = lambda t, y, k: -k * y def solve(_k): - return dfx.diffeqsolve( - dfx.ODETerm(dy), - dfx.Tsit5(), - 0.0, - 10.0, - 0.1, - 10.0, - args=_k, - adjoint=dfx.ForwardMode(), - ) - - data = jnp.asarray(solve(jnp.array(0.5)).ys) # seems to make type checkers happy + with jax.numpy_dtype_promotion("standard"): + return dfx.diffeqsolve( + dfx.ODETerm(dy), + dfx.Tsit5(), + 0.0, + 10.0, + 0.1, + jnp.asarray(10.0, dtype=_k.dtype), + args=_k, + adjoint=dfx.ForwardMode(), + ) + + reference_k = jnp.asarray(0.5, dtype=k.dtype) + data = jnp.asarray(solve(reference_k).ys) # seems to make type checkers happy fit = jnp.asarray(solve(k).ys) - return jnp.sum((data - fit) ** 2) + diff = data - fit + return jnp.sum(norm_sq(diff)) forward_only_fn_init_options_expected = ( diff --git a/tests/test_least_squares.py b/tests/test_least_squares.py index 36b710c8..fa877a77 100644 --- a/tests/test_least_squares.py +++ b/tests/test_least_squares.py @@ -11,13 +11,20 @@ import pytest from .helpers import ( + antiholomorphic_residual, + complex_to_real_residual, diagonal_quadratic_bowl, finite_difference_jvp, + holomorphic_residual, least_squares_fn_minima_init_args, least_squares_optimisers, + make_nonreal, + nonholomorphic_residual, + norm_sq, rosenbrock, simple_nn, tree_allclose, + tree_as_dtype, ) @@ -26,8 +33,17 @@ @pytest.mark.parametrize("solver", least_squares_optimisers) @pytest.mark.parametrize("_fn, minimum, init, args", least_squares_fn_minima_init_args) -def test_least_squares(solver, _fn, minimum, init, args): +@pytest.mark.parametrize("dtype", [jnp.float64, jnp.complex128]) +def test_least_squares(solver, _fn, minimum, init, args, dtype): + init = make_nonreal(tree_as_dtype(init, dtype)) + args = tree_as_dtype(args, dtype) atol = rtol = 1e-4 + if dtype == jnp.complex128 and isinstance(solver, optx.NelderMead): + atol = rtol = 1e-3 + if dtype == jnp.complex128 and isinstance(solver, optx.OptaxMinimiser): + max_steps = 100_000 + else: + max_steps = 10_000 has_aux = random.choice([True, False]) if has_aux: fn = lambda x, args: (_fn(x, args), smoke_aux) @@ -40,7 +56,13 @@ def test_least_squares(solver, _fn, minimum, init, args): context = contextlib.nullcontext() with context: optx_argmin = optx.least_squares( - fn, solver, init, has_aux=has_aux, args=args, max_steps=10_000, throw=False + fn, + solver, + init, + has_aux=has_aux, + args=args, + max_steps=max_steps, + throw=False, ).value out = fn(optx_argmin, args) if has_aux: @@ -48,14 +70,17 @@ def test_least_squares(solver, _fn, minimum, init, args): else: residual = out optx_min = jtu.tree_reduce( - lambda x, y: x + y, jtu.tree_map(lambda x: jnp.sum(x**2), residual) + lambda x, y: x + y, jtu.tree_map(lambda x: jnp.sum(norm_sq(x)), residual) ) assert tree_allclose(optx_min, minimum, atol=atol, rtol=rtol) @pytest.mark.parametrize("solver", least_squares_optimisers) @pytest.mark.parametrize("_fn, minimum, init, args", least_squares_fn_minima_init_args) -def test_least_squares_jvp(getkey, solver, _fn, minimum, init, args): +@pytest.mark.parametrize("dtype", [jnp.float64, jnp.complex128]) +def test_least_squares_jvp(getkey, solver, _fn, minimum, init, args, dtype): + init = make_nonreal(tree_as_dtype(init, dtype)) + args = tree_as_dtype(args, dtype) if _fn in (simple_nn, diagonal_quadratic_bowl): # These are ridiculously finickity to get references values for the derivatives return @@ -67,8 +92,10 @@ def test_least_squares_jvp(getkey, solver, _fn, minimum, init, args): fn = _fn dynamic_args, static_args = eqx.partition(args, eqx.is_array) - t_init = jtu.tree_map(lambda x: jr.normal(getkey(), x.shape), init) - t_dynamic_args = jtu.tree_map(lambda x: jr.normal(getkey(), x.shape), dynamic_args) + t_init = jtu.tree_map(lambda x: jr.normal(getkey(), x.shape, dtype=x.dtype), init) + t_dynamic_args = jtu.tree_map( + lambda x: jr.normal(getkey(), x.shape, dtype=x.dtype), dynamic_args + ) def least_squares(x, dynamic_args, *, adjoint): args = eqx.combine(dynamic_args, static_args) @@ -98,12 +125,22 @@ def least_squares(x, dynamic_args, *, adjoint): (t_init, t_dynamic_args), adjoint=otd, ) - if _fn is rosenbrock: + if _fn in ( + complex_to_real_residual, + holomorphic_residual, + nonholomorphic_residual, + rosenbrock, + ): # Finite difference does a bad job on this one, but we can figure it out # analytically. assert isinstance(args, jax.Array) expected_out = jtu.tree_map(lambda x: jnp.full_like(x, args), init) t_expected_out = jtu.tree_map(lambda x: jnp.full_like(x, t_dynamic_args), init) + elif _fn is antiholomorphic_residual: + expected_out = jtu.tree_map(lambda x: jnp.full_like(x, jnp.conj(args)), init) + t_expected_out = jtu.tree_map( + lambda x: jnp.full_like(x, jnp.conj(t_dynamic_args)), init + ) else: expected_out, t_expected_out = finite_difference_jvp( least_squares, @@ -129,7 +166,8 @@ def least_squares(x, dynamic_args, *, adjoint): # assert tree_allclose(t_out2, t_expected_out, atol=atol, rtol=rtol) -def test_gauss_newton_jacrev(): +@pytest.mark.parametrize("dtype", [jnp.float64, jnp.complex128]) +def test_gauss_newton_jacrev(dtype): @jax.custom_vjp def f(y, _): return dict(bar=y["foo"] ** 2) @@ -143,9 +181,11 @@ def f_bwd(sign, g): f.defvjp(f_fwd, f_bwd) solver = optx.LevenbergMarquardt(rtol=1e-8, atol=1e-8) - y0 = dict(foo=jnp.arange(3.0)) + y0 = dict(foo=make_nonreal(tree_as_dtype(jnp.arange(3.0), dtype))) out = optx.least_squares(f, solver, y0, options=dict(jac="bwd"), max_steps=512) - assert tree_allclose(out.value, dict(foo=jnp.zeros(3)), rtol=1e-3, atol=1e-2) + assert tree_allclose( + out.value, dict(foo=jnp.zeros(3, dtype=dtype)), rtol=1e-3, atol=1e-2 + ) with pytest.raises(TypeError, match="forward-mode autodiff"): optx.least_squares(f, solver, y0, options=dict(jac="fwd"), max_steps=512) @@ -253,13 +293,11 @@ def compute2(y2): assert tree_allclose(grad_dot2, true_grad_dot2) - # TODO: figure out what is going on here, complex numbers don't seem to be behaving. - pytest.skip() assert tree_allclose(grad_dot1, true_grad_dot1) # For context, the complex dot product between two scalars is # `(a + bi)^bar . (c + di) = ac + bd + i(ad - bc)` # The real dot product is # `(a, b) . (c, d) = ac + bd` - # In general we expect the real part of the complex dot product to agree with the - # real dot product. - assert tree_allclose(grad_dot1.real, grad_dot2) + # JAX represents the real gradient `(a, b)` as the complex number `a - bi`, so its + # R^2 directional derivative is instead the real part of the bilinear product. + assert tree_allclose(jnp.sum(grad1 * z).real, grad_dot2) diff --git a/tests/test_minimise.py b/tests/test_minimise.py index bc171bfa..475eebb7 100644 --- a/tests/test_minimise.py +++ b/tests/test_minimise.py @@ -12,15 +12,14 @@ import pytest from .helpers import ( - beale, - bowl, - finite_difference_jvp, forward_only_fn_init_options_expected, golden_search_fn_y0_options_expected, - matyas, + implicit_minimise_jvp, + make_nonreal, minimisation_fn_minima_init_args, minimisers, tree_allclose, + tree_as_dtype, ) @@ -32,12 +31,19 @@ ) @pytest.mark.parametrize("solver", minimisers) @pytest.mark.parametrize("_fn, minimum, init, args", minimisation_fn_minima_init_args) -def test_minimise(solver, _fn, minimum, init, args, options): +@pytest.mark.parametrize("dtype", [jnp.float64, jnp.complex128]) +def test_minimise(solver, _fn, minimum, init, args, options, dtype): + init = make_nonreal(tree_as_dtype(init, dtype)) + args = tree_as_dtype(args, dtype) if isinstance(solver, optx.GradientDescent): max_steps = 100_000 else: max_steps = 10_000 atol = rtol = 1e-4 + if dtype == jnp.complex128 and isinstance(solver, optx.NelderMead): + # The R² representation doubles the dimension of the simplex problem. + # Nelder--Mead is particularly inaccurate in these larger dimensions. + atol = rtol = 1e-2 has_aux = random.choice([True, False]) if has_aux: fn = lambda x, args: (_fn(x, args), smoke_aux) @@ -67,13 +73,18 @@ def test_minimise(solver, _fn, minimum, init, args, options): ) @pytest.mark.parametrize("solver", minimisers) @pytest.mark.parametrize("_fn, minimum, init, args", minimisation_fn_minima_init_args) -def test_minimise_jvp(getkey, solver, _fn, minimum, init, args, options): +@pytest.mark.parametrize("dtype", [jnp.float64, jnp.complex128]) +def test_minimise_jvp(getkey, solver, _fn, minimum, init, args, options, dtype): + init = make_nonreal(tree_as_dtype(init, dtype)) + args = tree_as_dtype(args, dtype) if isinstance(solver, (optx.GradientDescent, optx.NonlinearCG)): max_steps = 100_000 atol = rtol = 1e-2 else: max_steps = 10_000 atol = rtol = 1e-3 + if dtype == jnp.complex128 and isinstance(solver, optx.NelderMead): + atol = rtol = 1e-2 has_aux = random.choice([True, False]) if has_aux: fn = lambda x, args: (_fn(x, args), smoke_aux) @@ -81,8 +92,10 @@ def test_minimise_jvp(getkey, solver, _fn, minimum, init, args, options): fn = _fn dynamic_args, static_args = eqx.partition(args, eqx.is_array) - t_init = jtu.tree_map(lambda x: jr.normal(getkey(), x.shape), init) - t_dynamic_args = jtu.tree_map(lambda x: jr.normal(getkey(), x.shape), dynamic_args) + t_init = jtu.tree_map(lambda x: jr.normal(getkey(), x.shape, dtype=x.dtype), init) + t_dynamic_args = jtu.tree_map( + lambda x: jr.normal(getkey(), x.shape, dtype=x.dtype), dynamic_args + ) def minimise(x, dynamic_args, *, adjoint): args = eqx.combine(dynamic_args, static_args) @@ -107,26 +120,6 @@ def minimise(x, dynamic_args, *, adjoint): out, t_out = eqx.filter_jit(ft.partial(eqx.filter_jvp, minimise))( (init, dynamic_args), (t_init, t_dynamic_args), adjoint=otd ) - if _fn is bowl: - # Finite difference is very inaccurate on this problem. - expected_out = t_expected_out = jtu.tree_map(jnp.zeros_like, init) - elif _fn in (beale, matyas): - if isinstance(solver, optx.NonlinearCG): - eps = 1e-3 - atol = rtol = 1e-2 # finite difference does a really bad job on this one - else: - eps = 1e-4 - expected_out, t_expected_out = finite_difference_jvp( - minimise, - (init, dynamic_args), - (t_init, t_dynamic_args), - adjoint=otd, - eps=eps, - ) - else: - expected_out, t_expected_out = finite_difference_jvp( - minimise, (init, dynamic_args), (t_init, t_dynamic_args), adjoint=otd - ) # TODO(kidger): reinstate once we can do jvp-of-custom_vjp. Right now this errors # because of the line searches used internally. # @@ -137,10 +130,15 @@ def minimise(x, dynamic_args, *, adjoint): # out2, t_out2 = eqx.filter_jvp( # minimise, (init, dynamic_args), (t_init, t_dynamic_args), adjoint=dto, # ) - assert tree_allclose(out, expected_out, atol=atol, rtol=rtol) - if not isinstance(solver, optx.NelderMead): + if isinstance(solver, optx.NelderMead): # Nelder-Mead does such a bad job that the finite-difference gradients are # noticeably different. + assert tree_allclose(_fn(out, args), minimum, atol=atol, rtol=rtol) + else: + t_expected_out = implicit_minimise_jvp( + _fn, out, args, eqx.combine(t_dynamic_args, static_args) + ) + assert tree_allclose(_fn(out, args), minimum, atol=atol, rtol=rtol) assert tree_allclose(t_out, t_expected_out, atol=atol, rtol=rtol) # assert tree_allclose(expected_out2, expected_out, atol=atol, rtol=rtol) # assert tree_allclose(out2, expected_out, atol=atol, rtol=rtol) @@ -148,18 +146,21 @@ def minimise(x, dynamic_args, *, adjoint): # assert tree_allclose(t_out2, t_expected_out, atol=atol, rtol=rtol) +@pytest.mark.parametrize("dtype", [jnp.float64, jnp.complex128]) @pytest.mark.parametrize( "method", [optx.polak_ribiere, optx.fletcher_reeves, optx.hestenes_stiefel, optx.dai_yuan], ) -def test_nonlinear_cg_methods(method): +def test_nonlinear_cg_methods(method, dtype): solver = optx.NonlinearCG(rtol=1e-10, atol=1e-10, method=method) def f(y, _): - A = jnp.array([[2.0, -1.0], [-1.0, 3.0]]) - b = jnp.array([-100.0, 5.0]) - c = jnp.array(100.0) - return jnp.einsum("ij,i,j", A, y, y) + jnp.dot(b, y) + c + A = jnp.array([[2.0, -1.0], [-1.0, 3.0]], dtype=dtype) + b = jnp.array([-100.0, 5.0], dtype=dtype) + c = jnp.array(100.0, dtype=dtype) + quadratic = jnp.einsum("ij,i,j", A, jnp.conj(y), y).real + linear = jnp.vdot(b, y).real + return quadratic + linear + c.real # Analytic minimum: # 0 = df/dyk @@ -168,9 +169,11 @@ def f(y, _): # => y = -0.5 A^{-1} b # = [[-0.3, 0.1], [0.1, 0.2]] [-100, 5] # = [29.5, 9] - y0 = jnp.array([2.0, 3.0]) + y0 = make_nonreal(jnp.array([2.0, 3.0], dtype=dtype)) sol = optx.minimise(f, solver, y0, max_steps=500) - assert tree_allclose(sol.value, jnp.array([29.5, 9.0]), rtol=1e-5, atol=1e-5) + assert tree_allclose( + sol.value, jnp.array([29.5, 9.0], dtype=dtype), rtol=1e-5, atol=1e-5 + ) def test_optax_recompilation(): @@ -204,12 +207,20 @@ def f(x, _): @pytest.mark.parametrize( "fn, y0, options, expected", forward_only_fn_init_options_expected ) -def test_forward_minimisation(fn, y0, options, expected, solver): +@pytest.mark.parametrize("dtype", [jnp.float64, jnp.complex128]) +def test_forward_minimisation(fn, y0, options, expected, solver, dtype): + y0 = make_nonreal(tree_as_dtype(y0, dtype)) + expected = tree_as_dtype(expected, dtype) if isinstance(solver, optx.OptaxMinimiser): # No support for forward option return else: - # Many steps because gradient descent takes ridiculously long - sol = optx.minimise(fn, solver, y0, options=options, max_steps=2**10) + if dtype == jnp.complex128: + context = pytest.warns(match="Complex dtype support in Diffrax") + else: + context = contextlib.nullcontext() + with context: + # Many steps because gradient descent takes ridiculously long + sol = optx.minimise(fn, solver, y0, options=options, max_steps=2**10) assert sol.result == optx.RESULTS.successful assert tree_allclose(sol.value, expected, atol=1e-4, rtol=1e-4) diff --git a/tests/test_root_find.py b/tests/test_root_find.py index 11a4ddc6..7f6dd36c 100644 --- a/tests/test_root_find.py +++ b/tests/test_root_find.py @@ -12,8 +12,10 @@ from .helpers import ( finite_difference_jvp, fixed_point_fn_init_args, + make_nonreal, PiggybackAdjoint, tree_allclose, + tree_as_dtype, ) @@ -31,7 +33,10 @@ @pytest.mark.parametrize("solver", _root_finders) @pytest.mark.parametrize("_fn, init, args", fixed_point_fn_init_args) -def test_root_find(solver, _fn, init, args): +@pytest.mark.parametrize("dtype", [jnp.float64, jnp.complex128]) +def test_root_find(solver, _fn, init, args, dtype): + init = make_nonreal(tree_as_dtype(init, dtype)) + args = tree_as_dtype(args, dtype) atol = rtol = 1e-5 has_aux = random.choice([True, False]) @@ -43,9 +48,14 @@ def root_find_problem(y, args): fn = lambda x, args: (root_find_problem(x, args), smoke_aux) else: fn = root_find_problem - optx_root = optx.root_find( - fn, solver, init, has_aux=has_aux, args=args, max_steps=10_000, throw=False - ).value + if dtype == jnp.complex128: + context = pytest.warns(match="Complex support in Optimistix is a work in") + else: + context = contextlib.nullcontext() + with context: + optx_root = optx.root_find( + fn, solver, init, has_aux=has_aux, args=args, max_steps=10_000, throw=False + ).value out = fn(optx_root, args) if has_aux: fn_val, _ = out diff --git a/tests/test_solve.py b/tests/test_solve.py index ae268b53..38012973 100644 --- a/tests/test_solve.py +++ b/tests/test_solve.py @@ -2,30 +2,39 @@ import jax import jax.numpy as jnp import optimistix as optx +import pytest +from .helpers import make_nonreal, norm_sq -def test_minimise(): + +@pytest.mark.parametrize("dtype", [jnp.float64, jnp.complex128]) +def test_minimise(dtype): @jax.grad def f(offset): def fn(x, _): - return x**2 + offset + with jax.numpy_dtype_promotion("standard"): + return norm_sq(x) + offset solver = optx.GradientDescent(learning_rate=0.1, rtol=0.1, atol=0.1) - return optx.minimise(fn, solver, 0.0).value + y0 = make_nonreal(jnp.array(0.0, dtype=dtype)) + return norm_sq(optx.minimise(fn, solver, y0).value) - f(0.0) + f(jnp.array(0.0)) -def test_least_squares(): +@pytest.mark.parametrize("dtype", [jnp.float64, jnp.complex128]) +def test_least_squares(dtype): @jax.grad def f(offset): def fn(x, _): - return x + offset + with jax.numpy_dtype_promotion("standard"): + return jnp.conj(x) + offset solver = optx.Dogleg(rtol=0.1, atol=0.1) - return optx.least_squares(fn, solver, 0.0).value + y0 = make_nonreal(jnp.array(0.0, dtype=dtype)) + return norm_sq(optx.least_squares(fn, solver, y0).value) - f(0.0) + f(jnp.array(0.0)) def test_root_find(): @@ -52,14 +61,15 @@ def fn(x, _): f(0.0) -def test_forward_mode(): +@pytest.mark.parametrize("dtype", [jnp.float64, jnp.complex128]) +def test_forward_mode(dtype): def f(y, _): return eqxi.nondifferentiable_backward(y) optx.least_squares( f, optx.LevenbergMarquardt(rtol=1e-4, atol=1e-4), - jnp.arange(3.0), + make_nonreal(jnp.arange(3.0, dtype=dtype)), options=dict(jac="fwd"), )