fix: 🐛 preserve input precision in apply and compress

src/src_method/_sweep.py

@@ -19,18 +19,18 @@
 from time import perf_counter_ns
 from typing import TYPE_CHECKING, NamedTuple
 
-import numpy as np
 import structlog
 from opt_einsum import contract_expression, get_symbol
 
 from ._tensor_train import pad, unpad
-from .utils import to_numpy, truncated_qr
+from .utils import gaussian_sketch, to_numpy, truncated_qr
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
     from types import ModuleType
 
-    from numpy.typing import NDArray
+    import numpy as np
+    from numpy.typing import DTypeLike, NDArray
     from opt_einsum.contract import ContractExpression
 
     from ._tensor_train import TrainKind
@@ -100,7 +100,7 @@ def sweep(
     xp: ModuleType,
     *,
     cutoff: float = 0.0,
-    dtype: type = np.float64,
+    dtype: DTypeLike,
 ) -> list[NDArray]:
     """Contract and compress a stack in ket form with one SRC sweep.
 
@@ -136,7 +136,7 @@ def sweep(
     C = [xp.ones((chi_out,) + (1,) * depth, dtype=dtype)]
     for j in range(n_sites - 1):
         up, down = sites[j][0].shape[2], sites[j][-1].shape[3]
-        omega = xp.asarray(prng.normal(size=(chi_out, up, down))).astype(dtype)
+        omega = gaussian_sketch(prng, (chi_out, up, down), dtype, xp)
         C.append(contract(eqs.ltr, C[j], omega, *sites[j]))
     logger.debug("Left-to-right sweep", seconds=(perf_counter_ns() - tms) * 1e-9)
 

src/src_method/apply.py

@@ -10,15 +10,13 @@
 
 from typing import TYPE_CHECKING
 
-import numpy as np
-
 from ._tensor_train import infer_kind
 from .stack import src
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
-    from numpy.typing import NDArray
+    from numpy.typing import DTypeLike, NDArray
 
 
 def apply(
@@ -27,7 +25,7 @@ def apply(
     chi_out: int,
     *,
     cutoff: float = 0.0,
-    dtype: type = np.float64,
+    dtype: DTypeLike | None = None,
     seed: int | None = None,
     device: str = "cpu",
 ) -> list[NDArray]:
@@ -51,7 +49,8 @@ def apply(
             site during the right-to-left sweep.  The SVD operates on the
             small ``(chi_out, chi_out)`` R factor from QR, so overhead is
             minimal.  Set to 0.0 (default) to keep all bonds at chi_out.
-        dtype: The data type for the computation.
+        dtype: Data type of the random sketches. Defaults to the promoted
+            floating dtype of the inputs; an explicit dtype can promote the result.
         seed: An optional seed for the random number generator.
         device: ``"cpu"`` (default, numpy) or ``"gpu"`` (cupy).  Requires
             the optional ``cupy`` dependency for GPU execution.

src/src_method/compress.py

@@ -10,22 +10,20 @@
 
 from typing import TYPE_CHECKING
 
-import numpy as np
-
 from .stack import src
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
-    from numpy.typing import NDArray
+    from numpy.typing import DTypeLike, NDArray
 
 
 def compress(
     tensor: Sequence[NDArray],
     chi_out: int,
     *,
     cutoff: float = 0.0,
-    dtype: type = np.float64,
+    dtype: DTypeLike | None = None,
     seed: int | None = None,
     device: str = "cpu",
 ) -> list[NDArray]:
@@ -47,7 +45,8 @@ def compress(
             site during the right-to-left sweep.  The SVD operates on the
             small ``(chi_out, chi_out)`` R factor from QR, so overhead is
             minimal.  Set to 0.0 (default) to keep all bonds at chi_out.
-        dtype: The data type for the computation.
+        dtype: Data type of the random sketches. Defaults to the promoted
+            floating dtype of the inputs; an explicit dtype can promote the result.
         seed: An optional seed for the random number generator.
         device: ``"cpu"`` (default, numpy) or ``"gpu"`` (cupy).  Requires
             the optional ``cupy`` dependency for GPU execution.

src/src_method/stack.py

@@ -9,7 +9,6 @@
 
 from typing import TYPE_CHECKING
 
-import numpy as np
 import structlog
 
 from ._sweep import sweep
@@ -19,12 +18,12 @@
     exact_stack,
     normalize_stack,
 )
-from .utils import default_rng, get_xp, setup_logging
+from .utils import default_rng, get_xp, setup_logging, sketch_dtype
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
-    from numpy.typing import NDArray
+    from numpy.typing import DTypeLike, NDArray
 
 setup_logging()
 logger = structlog.get_logger(__name__)
@@ -39,7 +38,7 @@ def src(
     *trains: Sequence[NDArray],
     chi_out: int,
     cutoff: float = 0.0,
-    dtype: type = np.float64,
+    dtype: DTypeLike | None = None,
     seed: int | None = None,
     device: str = "cpu",
 ) -> list[NDArray]:
@@ -75,7 +74,10 @@ def src(
             site during the right-to-left sweep. Set to 0.0 (default) to keep
             all bonds at ``chi_out``. Ignored for two-site stacks, which are
             contracted and truncated to ``chi_out`` exactly.
-        dtype: The data type for the computation.
+        dtype: Data type of the random sketches. Defaults to the promoted
+            floating dtype of the inputs, so single precision stays single and
+            complex inputs get complex Ginibre sketches. An explicit dtype
+            overrides this and can promote the result.
         seed: An optional seed for the random number generator.
         device: ``"cpu"`` (default, numpy) or ``"gpu"`` (cupy). Requires
             the optional ``cupy`` dependency for GPU execution.
@@ -109,6 +111,14 @@ def src(
         output=kind,
         device=xp.__name__,
     )
-    result = sweep(layers, kind, chi_out, prng, xp, cutoff=cutoff, dtype=dtype)
+    result = sweep(
+        layers,
+        kind,
+        chi_out,
+        prng,
+        xp,
+        cutoff=cutoff,
+        dtype=sketch_dtype(dtype, *layers),
+    )
     logger.info("SRC complete.")
     return result

src/src_method/utils/__init__.py

@@ -2,14 +2,16 @@
 
 from __future__ import annotations
 
-from ._backend import default_rng, get_xp, to_numpy
+from ._backend import default_rng, gaussian_sketch, get_xp, sketch_dtype, to_numpy
 from .linalg import truncated_qr
 from .logging_config import setup_logging
 
 __all__ = [
     "default_rng",
+    "gaussian_sketch",
     "get_xp",
     "setup_logging",
+    "sketch_dtype",
     "to_numpy",
     "truncated_qr",
 ]

src/src_method/utils/_backend.py

@@ -13,9 +13,10 @@
 import numpy as np
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
     from types import ModuleType
 
-    from numpy.typing import NDArray
+    from numpy.typing import DTypeLike, NDArray
 
 
 def get_xp(device: str) -> ModuleType:
@@ -60,3 +61,55 @@ def to_numpy(arr: NDArray) -> np.ndarray:
     # cupy.ndarray exposes .get(); fall back to np.asarray for other dispatchers.
     get = getattr(arr, "get", None)
     return get() if callable(get) else np.asarray(arr)
+
+
+def sketch_dtype(dtype: DTypeLike | None, *inputs: Sequence[NDArray]) -> np.dtype:
+    """Resolve the sketch dtype.
+
+    Args:
+        dtype: Explicit sketch dtype, or None to follow the inputs.
+        *inputs: MPS or MPO tensors.
+
+    Returns:
+        The explicit dtype, or the promoted floating dtype of the inputs, so that
+        the sketch never changes the precision or the field (real or complex) of
+        the result. Integer and boolean inputs give ``float64``.
+    """
+    if dtype is not None:
+        return np.dtype(dtype)
+    input_dtype = np.result_type(*(arr.dtype for inpt in inputs for arr in inpt))
+    if input_dtype.kind in "fc":
+        return input_dtype
+    return np.dtype(np.float64)
+
+
+def gaussian_sketch(
+    prng: np.random.Generator,
+    shape: tuple[int, ...],
+    dtype: DTypeLike,
+    xp: ModuleType,
+) -> NDArray:
+    """Draw a Gaussian test tensor of the given dtype on the device of ``xp``.
+
+    Real dtypes get i.i.d. standard normal entries. Complex dtypes get the complex
+    Ginibre ensemble, ``(x + iy) / sqrt(2)`` with ``x, y`` i.i.d. standard normal,
+    whose law is invariant under unitary transformations; a real sketch of a complex
+    operator is not, and loses the guarantees of randomized range finding.
+
+    The draw is always made on the host in double precision and then cast, so a
+    seed gives the same sketch on every device and at every precision.
+
+    Args:
+        prng: Host-side random number generator.
+        shape: Shape of the tensor.
+        dtype: Dtype of the tensor.
+        xp: Array module (``numpy`` or ``cupy``).
+
+    Returns:
+        The sketch as an array of ``xp``.
+    """
+    dtype = np.dtype(dtype)
+    if dtype.kind == "c":
+        re, im = prng.normal(size=(2, *shape)) / np.sqrt(2.0)
+        return xp.asarray(re + 1j * im).astype(dtype)
+    return xp.asarray(prng.normal(size=shape)).astype(dtype)

tests/test_package.py

@@ -8,11 +8,15 @@
 
 """
 
+from __future__ import annotations
+
 import numpy as np
 import pytest
 import quimb.tensor as qtn
 
 from src_method import apply, compress
+from src_method.stack import src
+from src_method.utils import gaussian_sketch
 
 # -------------
 # --- Utils ---
@@ -605,3 +609,64 @@ def test_benchmark_src_mpo_mpo(benchmark):
 
     # Still has to be correct
     np.testing.assert_allclose(H1.distance(as_mpo(result_mpo)), 0.0, atol=1e-6)
+
+
+@pytest.mark.parametrize("make_train", [qtn.MPS_rand_state, qtn.MPO_rand])
+@pytest.mark.parametrize("dtype", [np.float32, np.float64, np.complex64, np.complex128])
+def test_compress_precision(make_train, dtype) -> None:
+    tensor = make_train(4, bond_dim=2, dtype=dtype, seed=12)
+
+    result = compress(tensor.arrays, chi_out=2, seed=14)
+
+    assert all(arr.dtype == dtype for arr in result)
+    tolerance = 2e-5 if np.finfo(dtype).bits == 32 else 1e-12
+    np.testing.assert_allclose(
+        type(tensor)(result).to_dense(),
+        tensor.to_dense(),
+        rtol=tolerance,
+        atol=tolerance,
+    )
+
+
+@pytest.mark.parametrize("make_train", [qtn.MPS_rand_state, qtn.MPO_rand])
+@pytest.mark.parametrize("dtype", [np.float32, np.float64, np.complex64, np.complex128])
+def test_apply_precision(make_train, dtype) -> None:
+    tensor = make_train(4, bond_dim=2, dtype=dtype, seed=12)
+    identity = qtn.MPO_identity(4, phys_dim=2, dtype=dtype)
+
+    result = apply(identity.arrays, tensor.arrays, chi_out=2, seed=14)
+
+    assert all(arr.dtype == dtype for arr in result)
+    tolerance = 2e-5 if np.finfo(dtype).bits == 32 else 1e-12
+    np.testing.assert_allclose(
+        type(tensor)(result).to_dense(),
+        tensor.to_dense(),
+        rtol=tolerance,
+        atol=tolerance,
+    )
+
+
+def test_stack_dtype_follows_inputs() -> None:
+    mpo = qtn.MPO_rand(4, bond_dim=2, dtype=np.complex64, seed=1)
+    result = src(mpo.arrays, mpo.arrays, chi_out=4, seed=0)
+    assert all(arr.dtype == np.complex64 for arr in result)
+
+
+@pytest.mark.parametrize("dtype", [np.float32, np.float64, np.complex64, np.complex128])
+def test_gaussian_sketch(dtype) -> None:
+    prng = np.random.default_rng(0)
+    omega = gaussian_sketch(prng, (400, 400), dtype, np)
+
+    assert omega.dtype == dtype
+    assert omega.var() == pytest.approx(1.0, rel=0.05)
+    if np.dtype(dtype).kind == "c":
+        # Circularly symmetric: real and imaginary parts share the variance.
+        assert omega.real.var() == pytest.approx(0.5, rel=0.05)
+        assert omega.imag.var() == pytest.approx(0.5, rel=0.05)
+        assert abs((omega**2).mean()) < 0.02
+
+
+def test_gaussian_sketch_seed_is_precision_independent() -> None:
+    single = gaussian_sketch(np.random.default_rng(3), (5, 5), np.complex64, np)
+    double = gaussian_sketch(np.random.default_rng(3), (5, 5), np.complex128, np)
+    np.testing.assert_allclose(single, double, rtol=1e-6)

tests/test_stack.py

@@ -98,7 +98,7 @@ def test_sweep_is_exact_without_truncation(spec, rng):
     stack = make_stack(spec, rng)
     kind = "mps" if stack[-1][0].ndim == 2 else "mpo"
 
-    out = sweep(stack, kind, CHI_EXACT, np.random.default_rng(0), np)
+    out = sweep(stack, kind, CHI_EXACT, np.random.default_rng(0), np, dtype=np.float64)
 
     assert rel_error(out, dense_stack(*stack)) < 1e-10
 
@@ -108,7 +108,9 @@ def test_sweep_rectangular_physical_legs(rng):
     B = random_mpo([3, 2, 2, 1], rng, up=3, down=2)
     psi = random_mps([2, 2, 3, 2], rng, phys=2)
 
-    out = sweep([A, B, psi], "mps", CHI_EXACT, np.random.default_rng(0), np)
+    out = sweep(
+        [A, B, psi], "mps", CHI_EXACT, np.random.default_rng(0), np, dtype=np.float64
+    )
 
     assert rel_error(out, dense_stack(A, B, psi)) < 1e-10
 
@@ -117,7 +119,7 @@ def test_sweep_does_not_mutate_inputs(rng):
     stack = make_stack("A B psi", rng)
     before = [[t.copy() for t in train] for train in stack]
 
-    sweep(stack, "mps", 4, np.random.default_rng(0), np)
+    sweep(stack, "mps", 4, np.random.default_rng(0), np, dtype=np.float64)
 
     for train, saved in zip(stack, before):
         for t, s in zip(train, saved):