diff --git a/src/src_method/_tensor_train.py b/src/src_method/_tensor_train.py index 8496aea..eb0b36a 100644 --- a/src/src_method/_tensor_train.py +++ b/src/src_method/_tensor_train.py @@ -21,6 +21,7 @@ import numpy as np from opt_einsum import contract +from ._validation import validate_open_boundary from .utils import to_numpy if TYPE_CHECKING: @@ -210,8 +211,9 @@ def normalize_stack( the contracted train. Raises: - ValueError: If the stack is empty, if the trains differ in length, or if - adjacent trains have mismatched physical dimensions. + ValueError: If the stack is empty, if a train is not open-boundary, if the + trains differ in length, or if adjacent trains have mismatched physical + dimensions. TypeError: If a train has an unrecognised layout or an MPS sits anywhere other than at one end of the stack. """ @@ -220,6 +222,10 @@ def normalize_stack( raise ValueError(msg) kinds: list[TrainKind | None] = [infer_kind(train) for train in trains] _check_roles(kinds) + for i, (train, kind) in enumerate(zip(trains, kinds)): + validate_open_boundary( + train, _MPS_BOUNDARY_NDIM if kind == "mps" else _MPO_BOUNDARY_NDIM, i + ) sizes = [len(train) for train in trains] if len(set(sizes)) > 1: msg = f"All tensor trains must have the same number of sites, got {sizes}." @@ -237,7 +243,8 @@ def _check_roles(kinds: Sequence[TrainKind | None]) -> None: if None in kinds: msg = ( f"Unsupported tensor network layout for train {kinds.index(None)}: " - "expected an MPS or MPO given as a list of per-site arrays." + "expected an MPS or MPO given as a list of per-site arrays " + "(periodic boundary conditions are not supported)." ) raise TypeError(msg) if len(kinds) > 1 and ("mps" in kinds[1:-1] or kinds[0] == kinds[-1] == "mps"): diff --git a/src/src_method/_validation.py b/src/src_method/_validation.py new file mode 100644 index 0000000..5fbf893 --- /dev/null +++ b/src/src_method/_validation.py @@ -0,0 +1,88 @@ +"""Argument and layout checks shared by the public entry points. + +A malformed call should fail at once with a message naming the culprit, rather than +as an opaque shape error deep inside the sweep. +""" + +from __future__ import annotations + +import numbers +from typing import TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from collections.abc import Sequence + + from numpy.typing import NDArray + +__all__ = ["validate_chi_out", "validate_cutoff", "validate_open_boundary"] + + +def validate_chi_out(chi_out: int) -> None: + """Reject a non-integer or non-positive output bond dimension. + + Any `numbers.Integral` is accepted, so NumPy integers pass; `bool` does not. + + Args: + chi_out: The requested output bond dimension. + + Raises: + TypeError: If `chi_out` is not an integer. + ValueError: If `chi_out` is not positive. + """ + if isinstance(chi_out, bool) or not isinstance(chi_out, numbers.Integral): + msg = f"chi_out must be an integer, got {chi_out!r}." + raise TypeError(msg) + if chi_out < 1: + msg = f"chi_out must be positive, got {chi_out!r}." + raise ValueError(msg) + + +def validate_cutoff(cutoff: float) -> None: + """Reject a cutoff outside ``[0.0, 1.0)``. + + A cutoff of 1 or more would discard every singular value, and a negative one + would silently mean "no truncation". + + Args: + cutoff: The relative singular-value cutoff. + + Raises: + ValueError: If `cutoff` is not a real number in ``[0.0, 1.0)``. + """ + if not isinstance(cutoff, numbers.Real) or not 0.0 <= cutoff < 1.0: + msg = f"cutoff must be in [0.0, 1.0), got {cutoff!r}." + raise ValueError(msg) + + +def validate_open_boundary( + train: Sequence[NDArray], boundary_ndim: int, index: int +) -> None: + """Check that a train has the layout of an open-boundary MPS or MPO. + + An open-boundary train has rank ``boundary_ndim`` at its two ends and one more + (the second bond) at every interior site. A periodic train has the interior rank + everywhere; read from its first site it is misclassified as a train of the + next-lower rank, which this check unmasks. + + A two-site train has no interior site, so a periodic one cannot be told apart + from an open one of the same rank and is not caught. + + Args: + train: The site arrays of the train. + boundary_ndim: The rank of the boundary sites, as inferred from the first + site. + index: The position of the train in the stack, for the error message. + + Raises: + ValueError: If the ranks do not follow the open-boundary pattern. + """ + expected = [boundary_ndim, *[boundary_ndim + 1] * (len(train) - 2), boundary_ndim] + ranks = [np.ndim(site) for site in train] + if len(train) > 1 and ranks != expected: + msg = ( + f"Train {index} has site ranks {ranks}, expected {expected}: only " + "open boundary conditions are supported." + ) + raise ValueError(msg) diff --git a/src/src_method/apply.py b/src/src_method/apply.py index 4e076f1..534c00d 100644 --- a/src/src_method/apply.py +++ b/src/src_method/apply.py @@ -59,10 +59,13 @@ def apply( The site arrays of the compressed tensor network (MPS or MPO). Raises: - TypeError: If the combination of input tensor types is unsupported. - ValueError: If the two trains differ in length or in the physical - dimensions they join, if a sub-three-site train is not exactly two - sites, or if ``device`` is not recognised. + TypeError: If ``chi_out`` is not an integer or the combination of input + tensor types is unsupported. + ValueError: If ``chi_out`` is not positive, if ``cutoff`` is not + in ``[0.0, 1.0)``, if a train is not open-boundary, if the two trains + differ in length or in the physical dimensions they join, if a + sub-three-site train is not exactly two sites, or if ``device`` is not + recognised. ImportError: If ``device="gpu"`` but cupy is not installed. """ left_kind = infer_kind(left_tensor) diff --git a/src/src_method/compress.py b/src/src_method/compress.py index f9bc7a4..bdff465 100644 --- a/src/src_method/compress.py +++ b/src/src_method/compress.py @@ -55,9 +55,12 @@ def compress( The site arrays of the compressed tensor network (MPS or MPO). Raises: - TypeError: If the input tensor type is unsupported. - ValueError: If a sub-three-site train is not exactly two sites, or if - ``device`` is not recognised. + TypeError: If ``chi_out`` is not an integer or the input tensor type is + unsupported. + ValueError: If ``chi_out`` is not positive, if ``cutoff`` is not + in ``[0.0, 1.0)``, if the train is not open-boundary, if a + sub-three-site train is not exactly two sites, or if ``device`` is not + recognised. ImportError: If ``device="gpu"`` but cupy is not installed. """ return src( diff --git a/src/src_method/stack.py b/src/src_method/stack.py index c8a8d54..adb661b 100644 --- a/src/src_method/stack.py +++ b/src/src_method/stack.py @@ -18,6 +18,7 @@ exact_stack, normalize_stack, ) +from ._validation import validate_chi_out, validate_cutoff from .utils import default_rng, get_xp, setup_logging, sketch_dtype if TYPE_CHECKING: @@ -87,13 +88,17 @@ def src( form, as numpy arrays (host-side, whatever the ``device``). Raises: - TypeError: If a train has an unrecognised layout or an MPS sits anywhere - other than at one end of the stack. - ValueError: If the stack is empty, if the trains differ in length or in - the physical dimensions they join, if a sub-three-site stack is not - exactly two sites, or if ``device`` is not recognised. + TypeError: If ``chi_out`` is not an integer, if a train has an unrecognised + layout or an MPS sits anywhere other than at one end of the stack. + ValueError: If ``chi_out`` is not positive, if ``cutoff`` is not in + ``[0.0, 1.0)``, if a train is not open-boundary, if the stack is empty, + if the trains differ in length or in the physical dimensions they join, + if a sub-three-site stack is not exactly two sites, or if ``device`` is + not recognised. ImportError: If ``device="gpu"`` but cupy is not installed. """ + validate_chi_out(chi_out) + validate_cutoff(cutoff) xp = get_xp(device) prng = default_rng(seed) layers, kind = normalize_stack(trains) diff --git a/tests/test_package.py b/tests/test_package.py index 1239052..3a9e326 100644 --- a/tests/test_package.py +++ b/tests/test_package.py @@ -573,13 +573,111 @@ def test_apply_unsupported_types(n_sites, phys_dim, chi_out, array_type): def test_compress_unsupported_type(n_sites, chi_out, array_type): """Tests that compress raises TypeError for unsupported tensor layouts.""" - # Rank-4 boundary tensors are neither an MPS nor an MPO (e.g. a PEPS row) + # Rank-4 boundary tensors are neither an MPS nor an MPO (e.g. a PEPS row, or a + # periodic MPO); the message points at the periodic case. peps_like = [np.zeros((2, 2, 2, 2), dtype=array_type)] * n_sites - with pytest.raises(TypeError): + with pytest.raises(TypeError, match="periodic"): compress(peps_like, chi_out=chi_out, dtype=array_type) +def test_periodic_boundary_mps_raises(n_sites, chi_out, array_type): + """A periodic MPS has rank 3 everywhere, so its first site looks like an MPO + boundary; the interior rank gives it away. + """ + periodic_mps = [np.zeros((2, 2, 2), dtype=array_type)] * n_sites + + with pytest.raises(ValueError, match="open boundary"): + compress(periodic_mps, chi_out=chi_out, dtype=array_type) + + ok_mpo = qtn.MPO_identity(n_sites, phys_dim=2, dtype=array_type).arrays + with pytest.raises(ValueError, match="open boundary"): + apply(periodic_mps, ok_mpo, chi_out=chi_out, dtype=array_type) + + +def test_inconsistent_site_ranks_raise(n_sites, chi_out, array_type): + """Both ends can be rank 3 and still disagree with an interior site.""" + mpo = list(qtn.MPO_identity(n_sites, phys_dim=2, dtype=array_type).arrays) + mpo[-2] = np.zeros((mpo[-2].shape[0], mpo[-2].shape[1], 2), dtype=array_type) + + with pytest.raises(ValueError, match="open boundary"): + compress(mpo, chi_out=chi_out, dtype=array_type) + + ok_mpo = qtn.MPO_identity(n_sites, phys_dim=2, dtype=array_type).arrays + with pytest.raises(ValueError, match="open boundary"): + apply(ok_mpo, mpo, chi_out=chi_out, dtype=array_type) + + +def test_mismatched_boundary_ranks_raise(chi_out, array_type): + """The two ends of a train must have the same rank, also for two sites.""" + two_site = [ + np.zeros((2, 2, 2), dtype=array_type), + np.zeros((2, 2), dtype=array_type), + ] + with pytest.raises(ValueError, match="open boundary"): + compress(two_site, chi_out=chi_out, dtype=array_type) + + longer = [ + np.zeros((2, 2, 2), dtype=array_type), + np.zeros((2, 2, 2, 2), dtype=array_type), + np.zeros((2, 2, 2), dtype=array_type), + np.zeros((2, 2, 2, 2), dtype=array_type), + ] + with pytest.raises(ValueError, match="open boundary"): + compress(longer, chi_out=chi_out, dtype=array_type) + + +# ----------------------------------------------- +# --- Test chi_out / cutoff argument validation -- +# ----------------------------------------------- + + +@pytest.mark.parametrize( + ("chi_out", "error"), + [(0, ValueError), (-3, ValueError), (2.5, TypeError), (True, TypeError)], +) +def test_invalid_chi_out_raises(n_sites, phys_dim, chi_out, error, array_type): + """Non-positive chi_out raises ValueError, a non-integer one TypeError.""" + H = qtn.MPO_identity(n_sites, phys_dim=phys_dim, dtype=array_type) + psi = qtn.MPS_rand_state(n_sites, bond_dim=4, phys_dim=phys_dim, dtype=array_type) + + with pytest.raises(error, match="chi_out"): + apply(H.arrays, psi.arrays, chi_out=chi_out, dtype=array_type) + with pytest.raises(error, match="chi_out"): + compress(H.arrays, chi_out=chi_out, dtype=array_type) + + +@pytest.mark.parametrize("cutoff", [2.0, -1.0]) +def test_invalid_cutoff_raises(n_sites, phys_dim, chi_out, cutoff, array_type): + """cutoff outside [0.0, 1.0) must raise ValueError, not silently misbehave.""" + H = qtn.MPO_identity(n_sites, phys_dim=phys_dim, dtype=array_type) + psi = qtn.MPS_rand_state( + n_sites, bond_dim=chi_out, phys_dim=phys_dim, dtype=array_type + ) + + with pytest.raises(ValueError, match="cutoff"): + apply(H.arrays, psi.arrays, chi_out=chi_out, cutoff=cutoff, dtype=array_type) + with pytest.raises(ValueError, match="cutoff"): + compress(H.arrays, chi_out=chi_out, cutoff=cutoff, dtype=array_type) + + +def test_numpy_integer_chi_out_accepted(n_sites, phys_dim, chi_out, array_type): + """chi_out as a numpy integer (np.int64 etc.) must be accepted, not rejected. + + validate_chi_out deliberately checks `numbers.Integral` rather than + `isinstance(chi_out, int)` so numpy integer types aren't rejected. + """ + H = qtn.MPO_identity(n_sites, phys_dim=phys_dim, dtype=array_type) + psi = qtn.MPS_rand_state( + n_sites, bond_dim=chi_out, phys_dim=phys_dim, dtype=array_type + ) + + psi_compress = as_mps( + apply(H.arrays, psi.arrays, chi_out=np.int64(chi_out), dtype=array_type) + ) + np.testing.assert_allclose(psi.distance(psi_compress), 0.0, atol=1e-6) + + # ----------------------------------------- # --- Check for performance regressions --- # -----------------------------------------