fix: 🐛 preserve input precision in apply and compress

src/src_method/apply.py

@@ -13,7 +13,6 @@
 from time import perf_counter_ns
 from typing import TYPE_CHECKING
 
-import numpy as np
 import structlog
 from opt_einsum import contract
 
@@ -30,12 +29,14 @@
     to_numpy,
     truncated_qr,
 )
+from .utils._backend import sketch_dtype
 
 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
 
 # Set up logger
 setup_logging()
@@ -61,7 +62,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]:
@@ -85,7 +86,10 @@ 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 Gaussian random sketch.
+            Defaults to the real floating type matching the
+            dtype of the inputs.
+            An explicit override 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.
@@ -124,6 +128,7 @@ def apply(
         check_exact_supported(len(left_tensor))
         logger.warning(LOG_WARN_SMALL)
         return exact_apply(left_tensor, right_tensor, chi_out, right_kind)
+    dtype = sketch_dtype(dtype, left_tensor, right_tensor)
     if right_kind == "mps":
         return _src_mpo_mps(
             left_tensor, right_tensor, chi_out, prng, xp, cutoff=cutoff, dtype=dtype
@@ -146,7 +151,7 @@ def _src_mpo_mps(
     xp: ModuleType,
     *,
     cutoff: float = 0.0,
-    dtype: type = np.float64,
+    dtype: DTypeLike,
 ) -> list[NDArray]:
     """Computes the compressed product |η> ≈ H|ψ> using the SRC method.
 
@@ -157,7 +162,7 @@ def _src_mpo_mps(
         prng: A numpy / cupy random number generator.
         xp: Array module (``numpy`` or ``cupy``).
         cutoff: Relative singular-value cutoff for adaptive bond truncation.
-        dtype: The data type for the computation.
+        dtype: Data type of the Gaussian random sketch.
 
     Returns:
         The site arrays of the compressed MPS |η> in right-canonical form.
@@ -259,7 +264,7 @@ def _src_mpo_mpo(
     xp: ModuleType,
     *,
     cutoff: float = 0.0,
-    dtype: type = np.float64,
+    dtype: DTypeLike,
 ) -> list[NDArray]:
     """Computes the compressed product H_new ≈ H1 @ H2 using the SRC method.
 
@@ -270,7 +275,7 @@ def _src_mpo_mpo(
         prng: A numpy / cupy random number generator.
         xp: Array module (``numpy`` or ``cupy``).
         cutoff: Relative singular-value cutoff for adaptive bond truncation.
-        dtype: The data type for the computation.
+        dtype: Data type of the Gaussian random sketch.
 
     Returns:
         The site arrays of the compressed MPO in right-canonical form.

src/src_method/compress.py

@@ -13,7 +13,6 @@
 from time import perf_counter_ns
 from typing import TYPE_CHECKING
 
-import numpy as np
 import structlog
 from opt_einsum import contract
 
@@ -30,12 +29,14 @@
     to_numpy,
     truncated_qr,
 )
+from .utils._backend import sketch_dtype
 
 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
 
 # Set up logger
 setup_logging()
@@ -60,7 +61,7 @@ def compress(
     chi_out: int,
     *,
     cutoff: float = 0.0,
-    dtype: type = np.float64,
+    dtype: DTypeLike | None = None,
     seed: int | None = None,
     device: str = "cpu",
 ) -> list[NDArray]:
@@ -82,7 +83,10 @@ 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 Gaussian random sketch.
+            Defaults to the real floating type matching the
+            dtype of the inputs.
+            An explicit override 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.
@@ -112,6 +116,7 @@ def compress(
         check_exact_supported(len(tensor))
         logger.warning(LOG_WARN_SMALL)
         return exact_compress(tensor, chi_out, kind)
+    dtype = sketch_dtype(dtype, tensor)
     if kind == "mps":
         return _src_mps(tensor, chi_out, prng, xp, cutoff=cutoff, dtype=dtype)
     return _src_mpo(tensor, chi_out, prng, xp, cutoff=cutoff, dtype=dtype)
@@ -129,7 +134,7 @@ def _src_mpo(
     xp: ModuleType,
     *,
     cutoff: float = 0.0,
-    dtype: type = np.float64,
+    dtype: DTypeLike,
 ) -> list[NDArray]:
     """Compress an MPO using the SRC method.
 
@@ -139,7 +144,7 @@ def _src_mpo(
         prng: A numpy / cupy random number generator instance.
         xp: Array module (``numpy`` or ``cupy``).
         cutoff: Relative singular-value cutoff for adaptive bond truncation.
-        dtype: The data type for the computation.
+        dtype: Data type of the Gaussian random sketch.
 
     Returns:
         The site arrays of the compressed MPO.
@@ -233,7 +238,7 @@ def _src_mps(
     xp: ModuleType,
     *,
     cutoff: float = 0.0,
-    dtype: type = np.float64,
+    dtype: DTypeLike,
 ) -> list[NDArray]:
     """Compress an MPS using the SRC method.
 
@@ -243,7 +248,7 @@ def _src_mps(
         prng: A numpy / cupy random number generator instance.
         xp: Array module (``numpy`` or ``cupy``).
         cutoff: Relative singular-value cutoff for adaptive bond truncation.
-        dtype: The data type for the computation.
+        dtype: Data type of the Gaussian random sketch.
 
     Returns:
         The site arrays of the compressed MPS.

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,21 @@ 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 real floating counterpart of the inputs dtype.
+    """
+    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 np.finfo(input_dtype).dtype
+    return np.dtype(np.float64)

tests/test_package.py

@@ -8,6 +8,8 @@
 
 """
 
+from __future__ import annotations
+
 import numpy as np
 import pytest
 import quimb.tensor as qtn
@@ -605,3 +607,38 @@ 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,
+    )