def __init__(self,
Rs_in,
Rs_out,
Rs_mid1,
Rs_mid2,
groups,
convolution,
linear=Linear,
scalar_activation=swish,
gate_activation=sigmoid):
super().__init__()
act_in = GatedBlock(groups * Rs_mid1, scalar_activation,
gate_activation)
self.lin_in = linear(Rs_in, act_in.Rs_in)
self.act_in = act_in
act_mid = GatedBlock(Rs_mid2, scalar_activation, gate_activation)
self.conv = convolution(Rs_mid1, act_mid.Rs_in)
self.act_mid = act_mid
act_out = GatedBlock(Rs_out, scalar_activation, gate_activation)
self.lin_out = linear(groups * Rs_mid2, act_out.Rs_in)
self.act_out = act_out
self.groups = groups
def make_layer(Rs_in, Rs_out):
if feature_product:
tr1 = rs.TransposeToMulL(Rs_in)
lts = LearnableTensorSquare(tr1.Rs_out,
list(range(lmax + 1)),
allow_change_output=True)
tr2 = torch.nn.Flatten(2)
Rs = tr1.mul * lts.Rs_out
act = GatedBlock(Rs_out, swish, sigmoid)
conv = convolution(K(Rs, act.Rs_in))
return torch.nn.ModuleList(
[torch.nn.Sequential(tr1, lts, tr2), conv, act])
else:
act = GatedBlock(Rs_out, swish, sigmoid)
conv = convolution(K(Rs_in, act.Rs_in))
return torch.nn.ModuleList([conv, act])
def test1(self):
with torch_default_dtype(torch.float64):
Rs_in = [(3, 0), (3, 1), (2, 0), (1, 2)]
Rs_out = [(3, 0), (3, 1), (1, 2), (3, 0)]
f = GatedBlock(Rs_out, rescaled_act.Softplus(beta=5),
rescaled_act.sigmoid)
c = Convolution(Kernel(Rs_in, f.Rs_in, ConstantRadialModel))
abc = torch.randn(3)
D_in = o3.direct_sum(
*
[o3.irr_repr(l, *abc) for mul, l in Rs_in for _ in range(mul)])
D_out = o3.direct_sum(*[
o3.irr_repr(l, *abc) for mul, l in Rs_out for _ in range(mul)
])
x = torch.randn(1, 5, sum(mul * (2 * l + 1) for mul, l in Rs_in))
geo = torch.randn(1, 5, 3)
rx = torch.einsum("ij,zaj->zai", (D_in, x))
rgeo = geo @ o3.rot(*abc).t()
y = f(c(x, geo), dim=2)
ry = torch.einsum("ij,zaj->zai", (D_out, y))
self.assertLess((f(c(rx, rgeo)) - ry).norm(), 1e-10 * ry.norm())
def make_layer(Rs_in, Rs_out):
act = GatedBlock(Rs_out, swish, sigmoid)
conv = Convolution(Rs_in,
act.Rs_in,
size,
lmax=lmax,
fuzzy_pixels=True,
padding=size // 2)
return torch.nn.Sequential(conv, act)
def __init__(self,
Rs_in,
Rs_out,
Rs_mid1,
Rs_mid2,
groups,
convolution,
linear=Linear,
scalar_activation=swish,
gate_activation=sigmoid,
final_nonlinearity=True):
"""
:param Rs_in:
:param Rs_out:
:param Rs_mid1:
:param Rs_mid2:
:param convolution: convolution operation that takes Rs_mid1, Rs_mid2
e.g. convolution = lambda Rs_in, Rs_out: Convolution(Kernel(Rs_in, Rs_out, ConstantRadialModel))
:param linear:
:param scalar_activation:
:param gated_activation:
:param final_nonlinearity:
"""
super().__init__()
act_in = GatedBlock(groups * Rs_mid1, scalar_activation,
gate_activation)
self.lin_in = linear(Rs_in, act_in.Rs_in)
self.act_in = act_in
act_mid = GatedBlock(Rs_mid2, scalar_activation, gate_activation)
self.conv = convolution(Rs_mid1, act_mid.Rs_in)
self.act_mid = act_mid
if final_nonlinearity:
act_out = GatedBlock(Rs_out, scalar_activation, gate_activation)
self.lin_out = linear(groups * Rs_mid2, act_out.Rs_in)
self.act_out = act_out
else:
self.lin_out = linear(groups * Rs_mid2, Rs_out)
self.act_out = None
self.groups = groups
def make_layer(Rs_in, Rs_out):
if feature_product:
tp = TensorSquare(Rs_in,
selection_rule=partial(o3.selection_rule,
lmax=lmax))
lin = Linear(tp.Rs_out, Rs_in)
act = GatedBlock(Rs_out, swish, sigmoid)
conv = convolution(K, Rs_in, act.Rs_in)
if feature_product:
return torch.nn.ModuleList([tp, lin, conv, act])
return torch.nn.ModuleList([conv, act])
def __init__(self, repr, Operation):
"""
:param repr: multiplicities
:param Operation: class of signature (Rs_out)
"""
super().__init__()
self.repr = repr
self.z_conv = Operation([(mul, l) for l, mul in enumerate(repr)], [(sum(repr), 0)])
self.z_act = ScalarActivation([(sum(repr), torch.sigmoid)], bias=False)
self.h_tilde = GatedBlock(Operation, repr, tanh, tanh)
def __init__(self, num_classes):
super().__init__()
representations = [(1, ), (2, 2, 2, 1), (4, 4, 4, 4), (6, 4, 4, 0),
(64, )]
representations = [[(mul, l) for l, mul in enumerate(rs)]
for rs in representations]
R = partial(CosineBasisModel,
max_radius=3.0,
number_of_basis=3,
h=100,
L=50,
act=relu)
K = partial(Kernel, RadialModel=R)
C = partial(Convolution, K)
self.firstlayers = torch.nn.ModuleList([
GatedBlock(partial(C, Rs_in), Rs_out, relu, sigmoid)
for Rs_in, Rs_out in zip(representations, representations[1:])
])
self.lastlayers = torch.nn.Sequential(AvgSpacial(),
torch.nn.Linear(64, num_classes))
def make_layer(Rs_in, Rs_out):
act = GatedBlock(Rs_out, relu, sigmoid)
conv = Convolution(K, Rs_in, act.Rs_in)
return torch.nn.ModuleList([conv, act])
R = partial(CosineBasisModel,
max_radius=3.0,
number_of_basis=3,
h=100,
L=3,
act=swish)
K = partial(Kernel, RadialModel=R)
C = partial(Convolution, K)
Rs_in1 = [(2, 0), (1, 1)]
Rs_in2 = [(1, 1)]
Rs_out1 = [(3, 0), (1, 1)]
Rs_out2 = [(3, 1)]
act1 = GatedBlock(Rs_out1, swish, sigmoid) # (Rs1, Rs2) --> Rs1
mix1 = Mixer(C, [Rs_in1, Rs_in2], act1.Rs_in)
act2 = GatedBlock(Rs_out2, swish, sigmoid) # (Rs1, Rs2) --> Rs2
mix2 = Mixer(C, [Rs_in1, Rs_in2], act2.Rs_in)
# First type of atom (eg. Hydrogen)
# There is 2 atoms with their features and positions
fea1 = torch.randn(10, 2, rs.dim(Rs_in1))
geo1 = torch.randn(10, 2, 3)
# Second type of atom (eg. Oxygen)
# There is 3 atoms with their features and positions
fea2 = torch.randn(10, 3, rs.dim(Rs_in2))
geo2 = torch.randn(10, 3, 3)
# The layer is splited in two parts for the two output types of atoms
def make_layer(Rs_in, Rs_out):
act = GatedBlock(Rs_out, relu, sigmoid)
conv = PeriodicConvolution(K, Rs_in, act.Rs_in, max_radius=3.8)
return torch.nn.ModuleList([conv, act])