def test_symmetric_mixed_tensor_sum(G1, G2):
N = 5
rep = T(2)(G1) * T(1)(G2) + 2 * T(0)(G1) * T(2)(G2) + T(1)(G1) + T(1)(G2)
P = rep.equivariant_projector()
v = np.random.rand(rep.size())
v = P @ v
samples = {G1: G1.samples(N), G2: G2.samples(N)}
gv = (vmap(rep.rho_dense)(samples) * v).sum(-1)
err = rel_error(gv, v + jnp.zeros_like(gv))
assert err < 1e-5, f"Symmetric vector fails err {err:.3e} with G={G1}x{G2}"
def test_symmetric_mixed_products(G1, G2):
N = 5
rep1 = (T(0) + 2 * T(1) + T(2))(G1)
rep2 = (T(0) + T(1))(G2)
rep = rep2 * rep1.T
P = rep.equivariant_projector()
v = np.random.rand(rep.size())
v = P @ v
W = v.reshape((rep2.size(), rep1.size()))
x = np.random.rand(N, rep1.size())
g1s = G1.samples(N)
g2s = G2.samples(N)
ring = vmap(rep1.rho_dense)(g1s)
routg = vmap(rep2.rho_dense)(g2s)
gx = (ring @ x[..., None])[..., 0]
Wgx = gx @ W.T
gWx = (routg @ (x @ W.T)[..., None])[..., 0]
equiv_err = rel_error(Wgx, gWx)
assert equiv_err < 1e-5, f"Equivariant gWx=Wgx fails err {equiv_err:.3e} with G={G1}x{G2}"
samples = {G1: g1s, G2: g2s}
gv = (vmap(rep.rho_dense)(samples) * v).sum(-1)
err = rel_error(gv, v + jnp.zeros_like(gv))
assert err < 1e-5, f"Symmetric vector fails err {err:.3e} with G={G1}x{G2}"
def test_bespoke_representations():
class ProductSubRep(Rep):
def __init__(self, G, subgroup_id, size):
""" Produces the representation of the subgroup of G = G1 x G2
with the index subgroup_id in {0,1} specifying G1 or G2.
Also requires specifying the size of the representation given by G1.d or G2.d """
self.G = G
self.index = subgroup_id
self._size = size
def __str__(self):
return "V_" + str(self.G).split('x')[self.index]
def size(self):
return self._size
def rho(self, M):
# Given that M is a LazyKron object, we can just get the argument
return M.Ms[self.index]
def drho(self, A):
return A.Ms[self.index]
def __call__(self, G):
# adding this will probably not be necessary in a future release,
# necessary now because rep is __call__ed in nn.EMLP constructor
assert self.G == G
return self
G1, G2 = SO(3), S(5)
G = G1 * G2
VSO3 = ProductSubRep(G, 0, G1.d)
VS5 = ProductSubRep(G, 1, G2.d)
Vin = VS5 + V(G)
Vout = VSO3
str(Vin >> Vout)
model = emlp.nn.EMLP(Vin, Vout, group=G)
input_point = np.random.randn(Vin.size()) * 10
from emlp.reps.linear_operators import LazyKron
lazy_G_sample = LazyKron([G1.sample(), G2.sample()])
out1 = model(Vin.rho(lazy_G_sample) @ input_point)
out2 = Vout.rho(lazy_G_sample) @ model(input_point)
assert rel_error(
out1, out2) < 1e-4, "EMLP equivariance fails on bespoke productsubrep"
def test_equivariant_matrix(G1, G2):
N = 5
repin = T(2)(G2) + 3 * T(0)(G1) + T(1)(G2) + 2 * T(2)(G1) * T(1)(G2)
repout = (T(1)(G1) + T(2)(G1) * T(0)(G2) + T(1)(G1) * T(1)(G2) + T(0)(G1) +
T(2)(G1) * T(1)(G2))
repW = repout * repin.T
P = repW.equivariant_projector()
W = np.random.rand(repout.size(), repin.size())
W = (P @ W.reshape(-1)).reshape(*W.shape)
x = np.random.rand(N, repin.size())
samples = {G1: G1.samples(N), G2: G2.samples(N)}
ring = vmap(repin.rho_dense)(samples)
routg = vmap(repout.rho_dense)(samples)
gx = (ring @ x[..., None])[..., 0]
Wgx = gx @ W.T
#print(g.shape,([email protected]).shape)
gWx = (routg @ (x @ W.T)[..., None])[..., 0]
equiv_err = rel_error(Wgx, gWx)
assert equiv_err < 1e-5, f"Equivariant gWx=Wgx fails err {equiv_err:.3e} with G={G1}x{G2}"