diff --git a/distreqx/bijectors/__init__.py b/distreqx/bijectors/__init__.py index 7805142..ca146ee 100644 --- a/distreqx/bijectors/__init__.py +++ b/distreqx/bijectors/__init__.py @@ -8,6 +8,7 @@ from ._chain import Chain as Chain from ._diag_linear import DiagLinear as DiagLinear from ._linear import AbstractLinearBijector as AbstractLinearBijector +from ._r2_to_complex import R2ToComplex as R2ToComplex from ._scalar_affine import ScalarAffine as ScalarAffine from ._shift import Shift as Shift from ._sigmoid import Sigmoid as Sigmoid diff --git a/distreqx/bijectors/_r2_to_complex.py b/distreqx/bijectors/_r2_to_complex.py new file mode 100644 index 0000000..febf79a --- /dev/null +++ b/distreqx/bijectors/_r2_to_complex.py @@ -0,0 +1,46 @@ +import jax.numpy as jnp +from jaxtyping import Array + +from ._bijector import ( + AbstractBijector, + AbstractForwardInverseBijector, + AbstractFwdLogDetJacBijector, + AbstractInvLogDetJacBijector, +) + + +class R2ToComplex( + AbstractForwardInverseBijector, + AbstractInvLogDetJacBijector, + AbstractFwdLogDetJacBijector, +): + """Maps a real array of shape (..., 2) to a complex array of shape (...).""" + + _is_constant_jacobian: bool = True + _is_constant_log_det: bool = True + + def forward_and_log_det(self, x: Array) -> tuple[Array, Array]: + """Computes y = x[..., 0] + 1j * x[..., 1] and log|det J(f)(x)| = 0.0.""" + if x.shape[-1] != 2: + raise ValueError( + f"Expected the last dimension of x to be 2 for R2 coordinates, " + f"but got {x.shape[-1]}." + ) + + y = x[..., 0] + 1j * x[..., 1] + return y, jnp.zeros((), dtype=x.dtype) + + def inverse_and_log_det(self, y: Array) -> tuple[Array, Array]: + """Computes x = [Re(y), Im(y)] and log|det J(f^{-1})(y)| = 0.0.""" + if not jnp.iscomplexobj(y): + raise ValueError( + f"Expected input to inverse to be a complex array, got {y.dtype}." + ) + + x = jnp.stack([jnp.real(y), jnp.imag(y)], axis=-1) + # The log-det uses the real dtype so type-checkers and JAX don't complain + return x, jnp.zeros((), dtype=jnp.real(y).dtype) + + def same_as(self, other: AbstractBijector) -> bool: + """Returns True if this bijector is guaranteed to be the same as `other`.""" + return type(other) is R2ToComplex diff --git a/docs/api/bijectors/r2_to_complex.md b/docs/api/bijectors/r2_to_complex.md new file mode 100644 index 0000000..e480bd9 --- /dev/null +++ b/docs/api/bijectors/r2_to_complex.md @@ -0,0 +1,6 @@ +# R2ToComplex Bijector + +::: distreqx.bijectors.R2ToComplex + options: + members: false +--- diff --git a/mkdocs.yml b/mkdocs.yml index aa5cf4e..e571f48 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -124,6 +124,7 @@ nav: - 'api/bijectors/shift.md' - 'api/bijectors/sigmoid.md' - 'api/bijectors/tanh.md' + - 'api/bijectors/r2_to_complex.md' - 'api/bijectors/triangular_linear.md' - 'api/bijectors/_bijector.md' - Utilities: diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..8f407ab --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,5 @@ +import jax + +# Must be set before any JAX arrays are initialized. Living here (rather than in +# individual test files) guarantees it runs before any test module import. +jax.config.update("jax_enable_x64", True) diff --git a/tests/r2_to_complex_test.py b/tests/r2_to_complex_test.py new file mode 100644 index 0000000..bad460c --- /dev/null +++ b/tests/r2_to_complex_test.py @@ -0,0 +1,112 @@ +from unittest import TestCase + +import equinox as eqx +import jax +import jax.numpy as jnp +import numpy as np +from parameterized import parameterized # type: ignore + +from distreqx.bijectors import Chain, R2ToComplex, ScalarAffine + + +class R2ToComplexTest(TestCase): + def setUp(self): + self.bij = R2ToComplex() + + def assertion_fn(self, rtol=1e-5): + return lambda x, y: np.testing.assert_allclose(x, y, rtol=rtol) + + def test_invalid_shapes_and_types(self): + with self.assertRaisesRegex(ValueError, "last dimension"): + self.bij.forward_and_log_det(jnp.array([1.0, 2.0, 3.0])) + + with self.assertRaisesRegex(ValueError, "complex array"): + self.bij.inverse_and_log_det(jnp.array([1.0, 2.0])) + + @parameterized.expand( + [ + ("float32", jnp.float32, jnp.complex64), + ("float64", jnp.float64, jnp.complex128), + ] + ) + def test_forward_and_log_det(self, name, real_dtype, comp_dtype): + x = jnp.array([[1.0, 2.0], [3.0, 4.0]], dtype=real_dtype) + y, log_det = self.bij.forward_and_log_det(x) + + expected_y = jnp.array([1.0 + 2.0j, 3.0 + 4.0j], dtype=comp_dtype) + self.assertion_fn()(y, expected_y) + self.assertEqual(y.shape, (2,)) + + self.assertEqual(log_det, 0.0) + self.assertEqual(log_det.dtype, real_dtype) + + @parameterized.expand( + [ + ("float32", jnp.float32, jnp.complex64), + ("float64", jnp.float64, jnp.complex128), + ] + ) + def test_inverse_and_log_det(self, name, real_dtype, comp_dtype): + y = jnp.array([1.0 + 2.0j, 3.0 + 4.0j], dtype=comp_dtype) + x, log_det = self.bij.inverse_and_log_det(y) + + expected_x = jnp.array([[1.0, 2.0], [3.0, 4.0]], dtype=real_dtype) + self.assertion_fn()(x, expected_x) + self.assertEqual(x.shape, (2, 2)) + + self.assertEqual(log_det, 0.0) + self.assertEqual(log_det.dtype, real_dtype) + + def test_jittable(self): + @eqx.filter_jit + def f(bij, x): + return bij.forward_and_log_det(x) + + x = jnp.array([1.0, 2.0]) + y, log_det = f(self.bij, x) + self.assertIsInstance(y, jax.Array) + self.assertIsInstance(log_det, jax.Array) + + def test_same_as(self): + same_bij = R2ToComplex() + self.assertTrue(self.bij.same_as(same_bij)) + + @parameterized.expand( + [ + ("float32", jnp.float32, jnp.complex64), + ("float64", jnp.float64, jnp.complex128), + ] + ) + def test_forward_inverse_round_trip(self, name, real_dtype, comp_dtype): + # forward then inverse should recover the original real input exactly, + # confirming the mapping is a genuine bijection (not just type punning). + x = jnp.array([[1.0, -2.5], [0.0, 3.0], [-4.0, -4.0]], dtype=real_dtype) + y, _ = self.bij.forward_and_log_det(x) + x_rec, _ = self.bij.inverse_and_log_det(y) + self.assertion_fn()(x_rec, x) + self.assertEqual(x_rec.dtype, real_dtype) + + # inverse then forward should recover the original complex input. + z = jnp.array([1.0 - 2.5j, 3.0 + 0.5j], dtype=comp_dtype) + x2, _ = self.bij.inverse_and_log_det(z) + z_rec, _ = self.bij.forward_and_log_det(x2) + self.assertion_fn()(z_rec, z) + self.assertEqual(z_rec.dtype, comp_dtype) + + def test_chaining(self): + # Chain a real-valued affine transform ahead of R2ToComplex, to check + # R2ToComplex composes correctly as part of a larger bijector pipeline. + affine = ScalarAffine(shift=jnp.array(1.0), scale=jnp.array(2.0)) + chain = Chain([R2ToComplex(), affine]) + x = jnp.array([[1.0, 2.0], [-3.0, 4.5]]) + + y, log_det = chain.forward_and_log_det(x) + + affine_out = affine.scale * x + affine.shift + expected_y = affine_out[..., 0] + 1j * affine_out[..., 1] + self.assertion_fn()(y, expected_y) + self.assertion_fn()(log_det, affine.forward_log_det_jacobian(x)) + + x_rec, inv_log_det = chain.inverse_and_log_det(y) + self.assertion_fn()(x_rec, x) + self.assertion_fn()(inv_log_det, -affine.forward_log_det_jacobian(x))