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 __init__(self, num_radial=30, max_radius=2):
super().__init__()
sp = rescaled_act.Softplus(beta=5)
RadialModel = partial(CosineBasisModel, max_radius=max_radius, number_of_basis=num_radial, h=100, L=2, act=sp)
K = partial(Kernel, RadialModel=RadialModel)
self.conv = Convolution(K, [(1, 0)], [(1, 1)])
def create_model(args, atomrefs, means, stddevs, properties, avg_n_atoms):
ssp = rescaled_act.ShiftedSoftplus(beta=args.beta)
kernel_conv = create_kernel_conv(
cutoff=args.rad_maxr,
n_bases=args.rad_nb,
n_neurons=args.rad_h,
n_layers=args.rad_L,
act=ssp,
radial_model=args.radial_model
)
sp = rescaled_act.Softplus(beta=args.beta)
if args.res:
net = ResNetwork(
kernel_conv=kernel_conv,
embed=args.embed,
l0=args.l0,
l1=args.l1,
l2=args.l2,
l3=args.l3,
L=args.L,
scalar_act=sp,
gate_act=rescaled_act.sigmoid,
avg_n_atoms=avg_n_atoms
)
else:
net = Network(
kernel_conv=kernel_conv,
embed=args.embed,
l0=args.l0,
l1=args.l1,
l2=args.l2,
l3=args.l3,
L=args.L,
scalar_act=sp,
gate_act=rescaled_act.sigmoid,
avg_n_atoms=avg_n_atoms
)
ident = torch.nn.Identity()
if args.mlp_out:
outnet = OutputMLPNetwork(
kernel_conv=kernel_conv,
previous_Rs=net.Rs[-1],
l0=args.outnet_l0,
l1=args.outnet_l1,
l2=args.outnet_l2,
l3=args.outnet_l3,
L=args.outnet_L,
scalar_act=sp,
gate_act=rescaled_act.sigmoid,
mlp_h=args.outnet_neurons,
mlp_L=args.outnet_layers,
avg_n_atoms=avg_n_atoms
)
else:
outnet = OutputScalarNetwork(
kernel_conv=kernel_conv,
previous_Rs=net.Rs[-1],
scalar_act=ident,
avg_n_atoms=avg_n_atoms
)
output_modules = [
spk.atomistic.Atomwise(
property=prop,
mean=means[prop],
stddev=stddevs[prop],
atomref=atomrefs[prop],
outnet=outnet,
# aggregation_mode='sum'
) for prop in properties
]
model = spk.AtomisticModel(net, output_modules)
return model