From 161b6cad1643ee4f3a8f00881f91d89c272e1450 Mon Sep 17 00:00:00 2001 From: Johnnie Gray Date: Thu, 8 Feb 2024 14:47:57 -0800 Subject: [PATCH 1/2] randn: allow dist="rademacher", including TN builders --- docs/changelog.md | 8 +++++++ quimb/gen/rand.py | 14 +++++++------ quimb/tensor/tensor_builder.py | 38 +++++++++++++++++++++++++++------- 3 files changed, 47 insertions(+), 13 deletions(-) diff --git a/docs/changelog.md b/docs/changelog.md index 40d3044a..927d7119 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -2,6 +2,14 @@ Release notes for `quimb`. +(whats-new-1-7-3)= +## v1.7.3 (unreleased) + +**Enhancements:** + +- [qu.randn](quimb.randn): support `dist="rademacher"`. +- support `dist` and other `randn` options in various TN builders. + (whats-new-1-7-2)= ## v1.7.2 (2024-01-30) diff --git a/quimb/gen/rand.py b/quimb/gen/rand.py index ad998a2f..501aeac9 100644 --- a/quimb/gen/rand.py +++ b/quimb/gen/rand.py @@ -159,14 +159,12 @@ def randn( dtype : {'complex128', 'float64', 'complex64' 'float32'}, optional The data-type of the output array. scale : float, optional - The width of the distribution (standard deviation if - ``dist='normal'``). + A multiplicative scale for the random numbers. loc : float, optional - The location of the distribution (lower limit if - ``dist='uniform'``). + An additive location for the random numbers. num_threads : int, optional How many threads to use. If ``None``, decide automatically. - dist : {'normal', 'uniform', 'exp'}, optional + dist : {'normal', 'uniform', 'rademacher', 'exp'}, optional Type of random number to generate. """ if seed is not None: @@ -178,6 +176,10 @@ def randn( else: d = prod(shape) + if dist == "rademacher": + # not parallelized for now + return rand_rademacher(shape, scale=scale, loc=loc, dtype=dtype) + if num_threads is None: # only multi-thread for big ``d`` if d <= 32768: @@ -289,7 +291,7 @@ def rand_rademacher(shape, scale=1, loc=0.0, dtype=float): else: raise TypeError( - f"dtype {dtype} not understood - should be float or " "complex." + f"dtype {dtype} not understood - should be float or complex." ) x = _choice(entries, shape) diff --git a/quimb/tensor/tensor_builder.py b/quimb/tensor/tensor_builder.py index 0ddac100..cad00bd8 100644 --- a/quimb/tensor/tensor_builder.py +++ b/quimb/tensor/tensor_builder.py @@ -71,7 +71,15 @@ def fill_fn(shape): @random_seed_fn def rand_tensor( - shape, inds, tags=None, dtype="float64", left_inds=None, **randn_opts + shape, + inds, + tags=None, + dtype="float64", + dist="normal", + scale=1.0, + loc=0.0, + left_inds=None, + **randn_opts, ): """Generate a random tensor with specified shape and inds. @@ -85,6 +93,12 @@ def rand_tensor( Labels to tag this tensor with. dtype : {'float64', 'complex128', 'float32', 'complex64'}, optional The underlying data type. + dist : {'normal', 'uniform', 'rademacher', 'exp'}, optional + Type of random number to generate, defaults to 'normal'. + scale : float, optional + A multiplier for the random numbers. + loc : float, optional + An offset for the random numbers. left_inds : sequence of str, optional Which, if any, indices to group as 'left' indices of an effective matrix. This can be useful, for example, when automatically applying @@ -95,7 +109,9 @@ def rand_tensor( ------- Tensor """ - data = randn(shape, dtype=dtype, **randn_opts) + data = randn( + shape, dtype=dtype, dist=dist, scale=scale, loc=loc, **randn_opts + ) return Tensor(data=data, inds=inds, tags=tags, left_inds=left_inds) @@ -1210,8 +1226,8 @@ def TN2D_rand( String specifier for naming convention of row tags. y_tag_id : str, optional String specifier for naming convention of column tags. - dist : str, optional - The distribution to sample from. + dist : {'normal', 'uniform', 'rademacher', 'exp'}, optional + Type of random number to generate, defaults to 'normal'. loc : float, optional The 'location' of the distribution, its meaning depends on ``dist``. scale : float, optional @@ -3327,6 +3343,7 @@ def MPS_rand_state( normalize=True, cyclic=False, dtype="float64", + dist="normal", trans_invar=False, **mps_opts, ): @@ -3360,7 +3377,11 @@ def MPS_rand_state( "boundary conditions." ) array = sensibly_scale( - randn(shape=(bond_dim, bond_dim, phys_dim), dtype=dtype) + randn( + shape=(bond_dim, bond_dim, phys_dim), + dtype=dtype, + dist=dist, + ) ) def fill_fn(shape): @@ -3369,7 +3390,7 @@ def fill_fn(shape): else: def fill_fn(shape): - return sensibly_scale(randn(shape, dtype=dtype)) + return sensibly_scale(randn(shape, dtype=dtype, dist=dist)) mps = MatrixProductState.from_fill_fn( fill_fn, @@ -3749,6 +3770,7 @@ def MPO_rand( cyclic=False, herm=False, dtype="float64", + dist="normal", **mpo_opts, ): """Generate a random matrix product state. @@ -3768,6 +3790,8 @@ def MPO_rand( open boundary conditions. dtype : {float, complex} or numpy dtype, optional Data type of the tensor network. + dist : {'normal', 'uniform', 'rademacher', 'exp'}, optional + Type of random number to generate, defaults to 'normal'. herm : bool, optional Whether to make the matrix hermitian (or symmetric if real) or not. mpo_opts @@ -3782,7 +3806,7 @@ def MPO_rand( ] def gen_data(shape): - data = randn(shape, dtype=dtype) + data = randn(shape, dtype=dtype, dist=dist) if not herm: return data From b507abc4536445eff2975cafc73dd35f03f7c839 Mon Sep 17 00:00:00 2001 From: Johnnie Gray Date: Thu, 8 Feb 2024 14:53:39 -0800 Subject: [PATCH 2/2] svd_truncated: restore scipy fallback --- docs/changelog.md | 5 +++++ quimb/tensor/decomp.py | 24 ++++++++++++++++++++++-- 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/docs/changelog.md b/docs/changelog.md index 927d7119..2cf48671 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -10,6 +10,11 @@ Release notes for `quimb`. - [qu.randn](quimb.randn): support `dist="rademacher"`. - support `dist` and other `randn` options in various TN builders. +**Bug fixes:** + +- restore fallback (to `scipy.linalg.svd` with driver='gesvd') behavior for + truncated SVD with numpy backend. + (whats-new-1-7-2)= ## v1.7.2 (2024-01-30) diff --git a/quimb/tensor/decomp.py b/quimb/tensor/decomp.py index 3eb1f6ae..b5fcb574 100644 --- a/quimb/tensor/decomp.py +++ b/quimb/tensor/decomp.py @@ -3,10 +3,12 @@ import functools import operator +import warnings import numpy as np -import scipy.sparse.linalg as spla +import scipy.linalg as scla import scipy.linalg.interpolative as sli +import scipy.sparse.linalg as spla from autoray import ( astype, backend_like, @@ -313,7 +315,6 @@ def _trim_and_renorm_svd_result_numba( return U, None, VH -@svd_truncated.register("numpy") @njit # pragma: no cover def svd_truncated_numba( x, cutoff=-1.0, cutoff_mode=4, max_bond=-1, absorb=0, renorm=0 @@ -325,6 +326,25 @@ def svd_truncated_numba( ) +@svd_truncated.register("numpy") +def svd_truncated_numpy( + x, cutoff=-1.0, cutoff_mode=4, max_bond=-1, absorb=0, renorm=0 +): + """Numpy version of ``svd_truncated``, trying the accelerated version + first, then falling back to the more stable scipy version. + """ + try: + return svd_truncated_numba( + x, cutoff, cutoff_mode, max_bond, absorb, renorm + ) + except np.linalg.LinAlgError as e: # pragma: no cover + warnings.warn(f"Got: {e}, falling back to scipy gesvd driver.") + U, s, VH = scla.svd(x, full_matrices=False, lapack_driver="gesvd") + return _trim_and_renorm_svd_result_numba( + U, s, VH, cutoff, cutoff_mode, max_bond, absorb, renorm + ) + + @svd_truncated.register("autoray.lazy") @lazy.core.lazy_cache("svd_truncated") def svd_truncated_lazy(