def test_qr_thin_square_fat(backend, shape):
if backend == 'sparse':
pytest.xfail("Sparse doesn't support linear algebra yet...")
x = gen_rand(shape, backend)
Q, R = ar.do('linalg.qr', x)
xn, Qn, Rn = map(ar.to_numpy, (x, Q, R))
assert ar.do('allclose', xn, Qn @ Rn)
def _get_exp_stab_gate(self, coo, tau):
'''
Returns
-------
`exp(tau * multiplier * stabilizer)`: qarray
Expm() of stabilizer centered on empty face at `coo`.
Params
-------
`coo`: tuple (x,y)
(Irrelevant) location of the empty face that
the stabilizer corresponds to. Just a label.
`tau`: float
Imaginary time for the exp(tau * gate)
'''
key = (coo, tau)
if key not in self._exp_stab_gates:
where, gate = self._stab_gates[coo]
el, ev = do('linalg.eigh', gate)
expgate = ev @ do('diag', do('exp', el * tau)) @ dag(ev)
self._exp_stab_gates[key] = (where, expgate)
return self._exp_stab_gates[key]
def _expm_cached(self, x, y):
cache = self._op_cache['expm']
key = (id(x), y)
if key not in cache:
el, ev = do('linalg.eigh', x)
cache[key] = ev @ do('diag', do('exp', el * y)) @ dag(ev)
return cache[key]
def norm_fro(x):
if isinstance(x, numpy.ndarray):
return norm_fro_dense(x.reshape(-1))
try:
return do('linalg.norm', reshape(x, [-1]), 2)
except AttributeError:
return do('sum', do('multiply', do('conj', x), x)) ** 0.5
def test_linalg_svd_square(backend):
if backend == 'sparse':
pytest.xfail("Sparse doesn't support linear algebra yet...")
x = gen_rand((5, 4), backend)
U, s, V = ar.do('linalg.svd', x)
assert (ar.infer_backend(x) == ar.infer_backend(U) == ar.infer_backend(s)
== ar.infer_backend(V) == backend)
y = U @ ar.do('diag', s, like=x) @ V
diff = ar.do('sum', abs(y - x))
assert ar.to_numpy(diff) < 1e-8
def fsim_param_gen(params):
theta, phi = params
a = do('cos', theta)
b = -1j * do('sin', theta)
c = do('exp', -1j * phi)
data = [[[[1, 0], [0, 0]], [[0, a], [b, 0]]],
[[[0, b], [a, 0]], [[0, 0], [0, c]]]]
return do('array', data, like=params)
def get_expm_gate(self, edge, t):
'''Local term for `edge`, matrix-exponentiated
by `t`.
'''
# return self._expm_cached(self.get_gate(edge), x)
key = (edge, t)
if key not in self._exp_gates:
gate = self.get_gate(edge)
el, ev = do('linalg.eigh', gate)
self._exp_gates[key] = ev @ do('diag', do('exp', el * t)) @ dag(ev)
return self._exp_gates[key]
def modified_gram_schmidt(X):
Q = []
for j in range(0, X.shape[0]):
q = X[j, :]
for i in range(0, j):
rij = ar.do('tensordot', ar.do('conj', Q[i]), q, 1)
q = q - rij * Q[i]
rjj = ar.do('linalg.norm', q, 2)
Q.append(q / rjj)
return ar.do('stack', Q, axis=0, like=X)
def test_count_nonzero(backend, array_dtype):
if backend == 'mars':
import mars
if mars._version.version_info < (0, 4, 0, ''):
pytest.xfail('mars count_nonzero bug fixed in version 0.4.')
if array_dtype == 'int':
x = ar.do('array', [0, 1, 2, 0, 3], like=backend)
elif array_dtype == 'float':
x = ar.do('array', [0., 1., 2., 0., 3.], like=backend)
elif array_dtype == 'bool':
x = ar.do('array', [False, True, True, False, True], like=backend)
nz = ar.do('count_nonzero', x)
assert ar.to_numpy(nz) == 3
def _unitize_exp(x):
r"""Perform isometrization using the using anti-symmetric matrix
exponentiation.
.. math::
U_A = \exp{A - A^\dagger}
If ``x`` is rectangular it is completed with zeros first.
"""
m, n = x.shape
d = max(m, n)
x = do('pad', x, [[0, d - m], [0, d - n]], 'constant', constant_values=0.0)
expx = do('linalg.expm', x - dag(x))
return expx[:m, :n]
def test_translator_random_normal(backend):
from autoray import numpy as anp
x = anp.random.normal(100.0, 0.1, size=(4, 5), like=backend)
if backend == 'sparse':
assert (x.data > 90.0).all()
assert (x.data < 110.0).all()
return
assert (ar.to_numpy(x) > 90.0).all()
assert (ar.to_numpy(x) < 110.0).all()
if backend == 'tensorflow':
x32 = ar.do('random.normal',
100.0,
0.1,
dtype='float32',
size=(4, 5),
like=backend)
assert x32.dtype == 'float32'
assert (ar.to_numpy(x32) > 90.0).all()
assert (ar.to_numpy(x32) < 110.0).all()
# test default single scalar
x = anp.random.normal(loc=1500, scale=10, like=backend)
assert 1000 <= ar.to_numpy(x) < 2000
def recursively_stack_chunks(loc, rem):
if not rem:
return chunks[loc]
return do('stack',
[recursively_stack_chunks(loc + (d,), rem[1:])
for d in range(self.size_dict[rem[0]])],
axis=output_pos[rem[0]] - len(loc), like=s)
def test_mgs(backend):
if backend == 'sparse':
pytest.xfail("Sparse doesn't support linear algebra yet...")
x = gen_rand((3, 5), backend)
Ux = modified_gram_schmidt(x)
y = ar.do('sum', Ux @ ar.dag(Ux))
assert ar.to_numpy(y) == pytest.approx(3)
def state_energy(psi, hterms, vterms, **opts):
he = psi.compute_local_expectation(hterms, normalized=True, **opts)
ve = psi.compute_local_expectation(vterms, normalized=True, **opts)
return autoray.do('real', (he + ve))
def test_complex_creation(backend, real_dtype):
if backend == 'torch':
pytest.xfail("Pytorch doesn't support complex numbers yet...")
if (backend == 'sparse') and (real_dtype == 'float32'):
pytest.xfail("Bug in sparse where single precision isn't maintained "
"after scalar multiplication.")
x = ar.do(
'complex',
ar.astype(ar.do('random.normal', size=(3, 4), like=backend),
real_dtype),
ar.astype(ar.do('random.normal', size=(3, 4), like=backend),
real_dtype))
assert ar.get_dtype_name(x) == {
'float32': 'complex64',
'float64': 'complex128'
}[real_dtype]
def tf_qr(x):
U, s, VH = autoray.do('linalg.svd', x)
dtype = autoray.get_dtype_name(U)
if 'complex' in dtype:
s = autoray.astype(s, dtype)
Q = U
R = autoray.reshape(s, (-1, 1)) * VH
return Q, R
def test_register_function(backend):
x = ar.do('ones', shape=(2, 3), like=backend)
def direct_fn(x):
return 1
# first test we can provide the function directly
ar.register_function(backend, 'test_register', direct_fn)
assert ar.do('test_register', x) == 1
def wrap_fn(fn):
def new_fn(*args, **kwargs):
res = fn(*args, **kwargs)
return res + 1
return new_fn
# then check we can wrap the old (previous) function
ar.register_function(backend, 'test_register', wrap_fn, wrap=True)
assert ar.do('test_register', x) == 2
def test_dtype_specials(backend, creation, dtype):
import numpy as np
x = ar.do(creation, shape=(2, 3), like=backend)
if backend == 'torch' and 'complex' in dtype:
pytest.xfail("Pytorch doesn't support complex numbers yet...")
x = ar.astype(x, dtype)
assert ar.get_dtype_name(x) == dtype
x = ar.to_numpy(x)
assert isinstance(x, np.ndarray)
assert ar.get_dtype_name(x) == dtype
def _unitize_qr(x):
"""Perform isometrization using the QR decomposition.
"""
fat = x.shape[0] < x.shape[1]
if fat:
x = transpose(x)
Q = do('linalg.qr', x)[0]
if fat:
Q = transpose(Q)
return Q
def state_energy(psi: beeky.QubitEncodeVector, hterms: dict, vterms: dict,
**opts):
'''Energy <psi|H|psi>, summing contribns from
horiz/vertical terms in Hamiltonian.
'''
# TODO: compute row/col envs first?
he = psi.compute_local_expectation(hterms, normalized=True, **opts)
ve = psi.compute_local_expectation(vterms, normalized=True, **opts)
return autoray.do('real', (he + ve))
def test_triu(backend):
x = gen_rand((4, 4), backend)
xl = ar.do('triu', x)
xln = ar.to_numpy(xl)
assert xln[1, 0] == 0.0
if backend != 'sparse':
# this won't work for sparse because density < 1
assert (xln > 0.0).sum() == 10
xl = ar.do('triu', x, k=-1)
xln = ar.to_numpy(xl)
if backend != 'sparse':
# this won't work for sparse because density < 1
assert xln[1, 0] != 0.0
assert xln[2, 0] == 0.0
if backend != 'sparse':
# this won't work for sparse because density < 1
assert (xln > 0.0).sum() == 13
if backend == 'tensorflow':
with pytest.raises(ValueError):
ar.do('triu', x, 1)
def _unitize_modified_gram_schmidt(A):
"""Perform isometrization explicitly using the modified Gram Schmidt
procedure.
"""
m, n = A.shape
thin = m > n
if thin:
A = do('transpose', A)
Q = []
for j in range(0, min(m, n)):
q = A[j, :]
for i in range(0, j):
rij = do('tensordot', do('conj', Q[i]), q, 1)
q = q - rij * Q[i]
Q.append(q / do('linalg.norm', q, 2))
Q = do('stack', Q, axis=0, like=A)
if thin:
Q = do('transpose', Q)
return Q
def gen_rand(shape, backend, dtype='float64'):
if backend == 'jax':
from jax import random as jrandom
global JAX_RANDOM_KEY
if JAX_RANDOM_KEY is None:
JAX_RANDOM_KEY = jrandom.PRNGKey(42)
JAX_RANDOM_KEY, subkey = jrandom.split(JAX_RANDOM_KEY)
return jrandom.uniform(subkey, shape=shape, dtype=dtype)
elif backend == 'sparse':
return ar.do('random.uniform',
size=shape,
like=backend,
density=0.5,
format='coo',
fill_value=0)
x = ar.do('random.uniform', size=shape, like=backend)
x = ar.astype(x, ar.to_backend_dtype(dtype, backend))
assert ar.get_dtype_name(x) == dtype
return x
def u3_gate_param_gen(params):
theta, phi, lamda = params
c2 = do('cos', theta / 2)
s2 = do('sin', theta / 2)
el = do('exp', 1.j * lamda)
ep = do('exp', 1.j * phi)
elp = do('exp', 1.j * (lamda + phi))
data = [[c2, -el * s2], [ep * s2, elp * c2]]
return do('array', data, like=params)
def _svd(x, cutoff=-1.0, cutoff_mode=3, max_bond=-1, absorb=0):
if isinstance(x, np.ndarray):
return _svd_numpy(x, cutoff, cutoff_mode, max_bond, absorb)
U, s, VH = do('linalg.svd', x)
if cutoff > 0.0:
if cutoff_mode == 1:
n_chi = do('count_nonzero', s > cutoff)
elif cutoff_mode == 2:
n_chi = do('count_nonzero', s > cutoff * s[0])
elif cutoff_mode in (3, 4):
s2 = s * s
cs2 = do('cumsum', s2, 0)
tot = cs2[-1]
if cutoff_mode == 3:
n_chi = do('count_nonzero', (tot - cs2) > cutoff) + 1
else:
n_chi = do('count_nonzero', cs2 < (1 - cutoff) * tot) + 1
n_chi = max(n_chi, 1)
if max_bond > 0:
n_chi = min(n_chi, max_bond)
if n_chi < s.shape[0]:
s = s[:n_chi]
U = U[..., :n_chi]
VH = VH[:n_chi, ...]
if cutoff_mode in (3, 4):
norm = (tot / cs2[n_chi - 1])**0.5
s *= norm
elif max_bond > 0:
s = s[:max_bond]
U = U[..., :max_bond]
VH = VH[:max_bond, ...]
if absorb == -1:
U = U * reshape(s, (1, -1))
elif absorb == 1:
VH = VH * reshape(s, (-1, 1))
else:
s **= 0.5
U = U * reshape(s, (1, -1))
VH = VH * reshape(s, (-1, 1))
return U, VH
def ry_gate_param_gen(params):
phi, = params
c = do('cos', phi / 2)
s = do('sin', phi / 2)
data = [[c, -s], [s, c]]
return do('array', data, like=params)
def constant(x, backend=_DEFAULT_BACKEND):
if backend == 'jax' and isinstance(x, qarray):
x = x.A
return do('array', x, like=backend)
def _lq(x):
if isinstance(x, np.ndarray):
return _lq_numba(x)
Q, L = do('linalg.qr', do('transpose', x))
return do('transpose', L), do('transpose', Q)
def _qr(x):
if isinstance(x, np.ndarray):
return _qr_numba(x)
return do('linalg.qr', x)
def rz_gate_param_gen(params):
phi, = params
c = do('cos', phi / 2)
s = -1j * do('sin', phi / 2)
data = [[c + s, 0], [0, c - s]]
return do('array', data, like=params)
