diff --git a/src/src_method/_sweep.py b/src/src_method/_sweep.py index d7f6ddd..af8a89f 100644 --- a/src/src_method/_sweep.py +++ b/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) diff --git a/src/src_method/apply.py b/src/src_method/apply.py index c148bd1..4e076f1 100644 --- a/src/src_method/apply.py +++ b/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. diff --git a/src/src_method/compress.py b/src/src_method/compress.py index 6011f9b..f9bc7a4 100644 --- a/src/src_method/compress.py +++ b/src/src_method/compress.py @@ -10,14 +10,12 @@ 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( @@ -25,7 +23,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]: @@ -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. diff --git a/src/src_method/stack.py b/src/src_method/stack.py index cdd2048..c8a8d54 100644 --- a/src/src_method/stack.py +++ b/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 diff --git a/src/src_method/utils/__init__.py b/src/src_method/utils/__init__.py index 50a9ba4..e695123 100644 --- a/src/src_method/utils/__init__.py +++ b/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", ] diff --git a/src/src_method/utils/_backend.py b/src/src_method/utils/_backend.py index 7a09be9..a0f49ec 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,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) diff --git a/tests/test_package.py b/tests/test_package.py index 23ed5ee..1239052 100644 --- a/tests/test_package.py +++ b/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) diff --git a/tests/test_stack.py b/tests/test_stack.py index cd4962c..800b45d 100644 --- a/tests/test_stack.py +++ b/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):