def __init__(self):
matrix_pairs = {}
rng = np.random.RandomState(SEED)
for s in SIZES:
for n_block in [1, 2, 5]:
arrays = [rng.standard_normal((s, s)) for _ in range(n_block)]
arrays = [arr @ arr.T for arr in arrays]
matrix_pairs[(s, n_block)] = (
matrices.PositiveDefiniteBlockDiagonalMatrix(
matrices.DensePositiveDefiniteMatrix(arr)
for arr in arrays), sla.block_diag(*arrays))
if AUTOGRAD_AVAILABLE:
@primitive
def block_diag(blocks):
return sla.block_diag(*blocks)
def vjp_block_diag(ans, blocks):
blocks = tuple(blocks)
def vjp(g):
i, j = 0, 0
vjp_blocks = []
for block in blocks:
j += block.shape[0]
vjp_blocks.append(g[i:j, i:j])
i = j
return tuple(vjp_blocks)
return vjp
defvjp(block_diag, vjp_block_diag)
def get_param(matrix):
return tuple(block.array for block in matrix._blocks)
def param_func(param, matrix):
return block_diag(param)
else:
param_func, get_param = None, None
super().__init__(matrix_pairs, get_param, param_func, rng)
def add_grad(spline_lst):
funcgrad_lst = []
for spline in spline_lst:
@primitive
def f(ws, yaw=0):
assert np.all(yaw == 0), "Gradients of yaw not implemented"
return spline(ws)
def f_vjp(ans, ws, yaw=0):
def gr(g):
return g * spline.derivative()(ws)
return gr
defvjp(f, f_vjp)
funcgrad_lst.append(f)
return funcgrad_lst
def test_scalar2multi_scalar():
def fxy(x):
return x**2 + 1, 2 * x + 1
def f(x):
fx, fy = fxy(x)
return fx + fy
x = 3.
ref = 8
npt.assert_equal(cs(f)(x), ref)
npt.assert_almost_equal(fd(f)(x), ref, 5)
npt.assert_equal(autograd(f)(x), ref)
pf = primitive(f)
defvjp(pf, lambda ans, x: lambda g: g * (2 * x + 2))
npt.assert_array_equal(autograd(pf, False)(x), ref)
pf = primitive(fxy)
defvjp(pf, lambda ans, x: lambda g: (g[0] * 2 * x, g[1] * 2))
npt.assert_array_equal(autograd(f, False)(x), ref)
def get_sparse_product(z_mat):
"""
Return an autograd-compatible function that calculates
``z_mat @ a`` and ``z_mat.T @ a`` when ``z_mat`` is a sparse matrix.
Parameters
------------
z_mat: A 2d matrix
The matrix by which to multiply. The matrix can be dense, but the only
reason to use ``get_sparse_product`` is with a sparse matrix since
dense matrix multiplication is supported natively by ``autograd``.
Returns
-----------
z_mult:
A function such that ``z_mult(b) = z_mat @ b``.
zt_mult:
A function such that ``zt_mult(b) = z_mat.T @ b``.
Unlike standard sparse matrix multiplication, ``z_mult`` and ``zt_mult``
can be used with ``autograd``.
"""
if z_mat.ndim != 2:
raise ValueError(
'get_sparse_product can only be used with 2d arrays.')
def check_b(b):
b = np.atleast_1d(b)
if (b.ndim > 2):
raise ValueError('The argument must be at most two dimensional.')
return b
@primitive
def z_mult(b):
return z_mat @ check_b(b)
@primitive
def zt_mult(b):
return z_mat.T @ check_b(b)
def z_mult_jvp(g, ans, b):
return z_mult(g) # z_mat @ g
defjvp(z_mult, z_mult_jvp)
def z_mult_vjp(ans, b):
def vjp(g):
return zt_mult(g) # z_mat.T @ g
return vjp
defvjp(z_mult, z_mult_vjp)
def zt_mult_jvp(g, ans, b):
return zt_mult(g) # z_mat.T @ g
defjvp(zt_mult, zt_mult_jvp)
def zt_mult_vjp(ans, b):
def vjp(g):
return z_mult(g) # (z_mat.T).T @ g
return vjp
defvjp(zt_mult, zt_mult_vjp)
return z_mult, zt_mult
result[groups[n], :] += x[n, :]
else:
result[groups[n]] += x[n]
return result
def _ungroup(v, groups):
if v.ndim > 1:
return v[groups, :]
else:
return v[groups]
def grouped_sum_vjp(ans, x, groups, num_groups=None):
def vjp(v):
return _ungroup(v, groups)
return vjp
defvjp(grouped_sum, grouped_sum_vjp)
def grouped_sum_jvp(v, ans, x, groups, num_groups=None):
return grouped_sum(v, groups, num_groups=num_groups)
defjvp(grouped_sum, grouped_sum_jvp)
@primitive
def replace(x_sub, x, inds):
"""Differentiably replace elements of `x[inds]` with `x_sub` by making a copy.
Parameters
------------
x_sub: numpy.ndarray[D]
A numeric array with replacement values.
from __future__ import absolute_import
import numpy as onp
from . import numpy_wrapper as anp
from .numpy_boxes import ArrayBox
from autograd.tracer import primitive
from autograd.core import defvjp
# ----- Binary ufuncs -----
defvjp(anp.add, lambda g, ans, x, y: unbroadcast(x, g),
lambda g, ans, x, y: unbroadcast(y, g))
defvjp(anp.multiply, lambda g, ans, x, y: unbroadcast(x, y * g),
lambda g, ans, x, y: unbroadcast(y, x * g))
defvjp(anp.subtract, lambda g, ans, x, y: unbroadcast(x, g),
lambda g, ans, x, y: unbroadcast(y, -g))
defvjp(anp.divide, lambda g, ans, x, y: unbroadcast(x, g / y),
lambda g, ans, x, y: unbroadcast(y, -g * x / y**2))
defvjp(
anp.power,
lambda g, ans, x, y: unbroadcast(x, g * y * x**anp.where(y, y - 1, 1.)),
lambda g, ans, x, y: unbroadcast(y,
g * anp.log(replace_zero(x, 1.)) * x**y))
def replace_zero(x, val):
return anp.where(x, x, val)
def unbroadcast(target, g, broadcast_idx=0):
while anp.ndim(g) > anp.ndim(target):
g = anp.sum(g, axis=broadcast_idx)
def asarray(x, dtype=None, order=None):
if isinstance(x, (ArrayBox)):
return x
# if any([isinstance(_x, (ArrayBox, SequenceBox)) for _x in np.atleast_1d(x).flatten()]):
# return x
return np.asarray(x, dtype, order)
# replace asarray to support autograd
anp.asarray = asarray
# replace dsqrt to avoid divide by zero if x=0
eps = 2 * np.finfo(float).eps**2
defvjp(anp.sqrt, lambda ans, x: lambda g: g * 0.5 * np.where(x == 0, eps, x)**
-0.5) # @UndefinedVariable
@contextmanager
def use_autograd_in(modules=["py_wake."]):
def get_dict(m):
if isinstance(m, dict):
return [m]
if isinstance(m, str):
return [
v.__dict__ for k, v in sys.modules.items() if k.startswith(m)
and k != __name__ and getattr(v, 'np', None) == np
]
if inspect.ismodule(m):
return [m.__dict__]
def vjp_block_diag(ans, blocks):
blocks = tuple(blocks)
def vjp(g):
i, j = 0, 0
vjp_blocks = []
for block in blocks:
j += block.shape[0]
vjp_blocks.append(g[i:j, i:j])
i = j
return tuple(vjp_blocks)
return vjp
defvjp(block_diag, vjp_block_diag)
class TestPositiveDefiniteBlockDiagonalMatrix(
InvertibleBlockMatrixTests,
DifferentiableMatrixTests,
ExplicitShapePositiveDefiniteMatrixTests,
):
@pytest.fixture
def matrix_pair(self, rng, size, n_block):
arrays = [rng.standard_normal((size, size)) for _ in range(n_block)]
arrays = [arr @ arr.T for arr in arrays]
return (
matrices.PositiveDefiniteBlockDiagonalMatrix(
matrices.DensePositiveDefiniteMatrix(arr) for arr in arrays
),
for k2 in range(k1 + 1):
mat[k1, k2] = vec[_sym_index(k1, k2)]
return mat
# Because we cannot use autograd with array assignment, define the
# vector jacobian product and jacobian vector products of
# _unvectorize_ld_matrix.
def _unvectorize_ld_matrix_vjp(g):
assert g.shape[0] == g.shape[1]
return _vectorize_ld_matrix(g)
defvjp(_unvectorize_ld_matrix,
lambda ans, vec: lambda g: _unvectorize_ld_matrix_vjp(g))
def _unvectorize_ld_matrix_jvp(g):
return _unvectorize_ld_matrix(g)
defjvp(_unvectorize_ld_matrix, lambda g, ans, x: _unvectorize_ld_matrix_jvp(g))
def _exp_matrix_diagonal(mat):
assert mat.shape[0] == mat.shape[1]
# NB: make_diagonal() is only defined in the autograd version of numpy
mat_exp_diag = np.make_diagonal(np.exp(np.diag(mat)),
offset=0,
axis1=-1,
def test_vector2multi_vector():
def fxy(x):
return x**2 + 1, 2 * x + 1
def f0(x):
return fxy(x)[0]
def fsum(x):
fx, fy = fxy(x)
return fx + fy
x = np.array([1., 2, 3])
ref0 = [2, 4, 6]
refsum = [4, 6, 8]
npt.assert_equal(cs(f0)(x), ref0)
npt.assert_almost_equal(fd(f0)(x), ref0, 5)
npt.assert_equal(autograd(f0)(x), ref0)
pf0 = primitive(f0)
defvjp(pf0, lambda ans, x: lambda g: g * (2 * x))
npt.assert_array_equal(autograd(pf0, False)(x), ref0)
npt.assert_equal(cs(fsum)(x), refsum)
npt.assert_almost_equal(fd(fsum)(x), refsum, 5)
npt.assert_equal(autograd(fsum)(x), refsum)
pfsum = primitive(fsum)
defvjp(pfsum, lambda ans, x: lambda g: g * (2 * x + 2))
npt.assert_array_equal(autograd(pfsum, False)(x), refsum)
pfxy = primitive(fxy)
def dfxy(x):
return 2 * x, np.full(x.shape, 2)
def gsum(x):
fx, fy = pfxy(x)
return fx + fy
def g0(x):
return pfxy(x)[0]
pgsum = primitive(gsum)
pg0 = primitive(g0)
defvjp(pgsum, lambda ans, x: lambda g: g * np.sum(dfxy(x), 0))
defvjp(pg0, lambda ans, x: lambda g: g * dfxy(x)[0])
npt.assert_array_equal(autograd(pgsum, False)(x), refsum)
npt.assert_array_equal(autograd(pg0, False)(x), ref0)
defvjp(pfxy, lambda ans, x: lambda g: dfxy(x)[0])
def h0(x):
return pfxy(x)[0]
npt.assert_array_equal(autograd(h0, False)(x), ref0)
defvjp(pfxy, lambda ans, x: lambda g: np.sum(g * np.asarray(dfxy(x)), 0))
def hsum(x):
fx, fy = pfxy(x)
return fx + fy
npt.assert_array_equal(autograd(hsum, False)(x), refsum)