diff --git a/src/src_method/apply.py b/src/src_method/apply.py index 4ea0e65..7b1fad5 100644 --- a/src/src_method/apply.py +++ b/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. diff --git a/src/src_method/compress.py b/src/src_method/compress.py index af8eacf..cc6cca7 100644 --- a/src/src_method/compress.py +++ b/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. diff --git a/src/src_method/utils/_backend.py b/src/src_method/utils/_backend.py index 7a09be9..1ef6103 100644 --- a/src/src_method/utils/_backend.py +++ b/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) diff --git a/tests/test_package.py b/tests/test_package.py index 23ed5ee..3375057 100644 --- a/tests/test_package.py +++ b/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, + )