def protodensity(self, coordinates: torch.Tensor, labels: torch.Tensor,
centers: torch.Tensor):
"""
evaluates protomolecule density at coordinates
:param coordinates:
:param labels:
:param centers:
:return:
"""
dv = distance_vectors(sample_coords=coordinates,
mol_coords=centers,
labels=labels,
device=self.device)
r = distance(dv)
density = torch.zeros_like(r)
p_splitted = self.p.split(1, dim=1)
b_splitted = self.b.split(1, dim=1)
a_splitted = self.a.split(1, dim=1)
for i, (p, b, a) in enumerate(zip(p_splitted, b_splitted, a_splitted)):
p_s = expand_parameter(labels, p)
a_s = expand_parameter(labels, a)
b_s = expand_parameter(labels, b)
k = gen_exponential_kernel(r, a_s, b_s, p_s)
density += k
return density.sum(2)
def forward(self, coordinates: torch.Tensor, coefficients: torch.Tensor,
labels: torch.Tensor, centers: torch.Tensor):
"""
evaluates density at coordinates
:param coordinates:
:param coefficients:
:param labels:
:param centers:
:return:
"""
dv = distance_vectors(sample_coords=coordinates,
mol_coords=centers,
labels=labels,
device=self.device)
r = distance(dv)
density = torch.zeros_like(r)
c_splitted = coefficients.split(1, dim=2)
p_splitted = self.p.split(1, dim=1)
b_splitted = self.b.split(1, dim=1)
a_splitted = self.a.split(1, dim=1)
for i, (c_s, p, b, a) in enumerate(
zip(c_splitted, p_splitted, b_splitted, a_splitted)):
p_s = expand_parameter(labels, p)
a_s = expand_parameter(labels, a)
b_s = expand_parameter(labels, b)
k = gen_exponential_kernel(r, a_s, b_s, p_s)
c_s = c_s.transpose(1, 2)
density += c_s * k
return density.sum(2)
def integrate(self, labels: torch.Tensor):
"""
evaluates functions integrals
:param labels:
:return:
"""
integrals = torch.zeros(labels.size()[0],
labels.size()[1],
self.maxfun,
dtype=self.dtype,
device=self.device)
split_p = self.p.split(1, dim=1)
split_b = self.b.split(1, dim=1)
split_a = self.a.split(1, dim=1)
factorial = lambda n: (n + 1).to(torch.float).lgamma().exp()
for i, (p, b, a) in enumerate(zip(split_p, split_b, split_a)):
p_s = expand_parameter(labels, p).to(self.device)
b_s = expand_parameter(labels, b).to(self.device)
a_s = expand_parameter(labels, a).to(self.device)
f = factorial(2 + p_s)
integrals[:, :, i] = a_s * f * b_s.clamp(
0.1, None).pow(-3 - p_s) * 4 * math.pi
return integrals
def forward(self, l, d):
"""
Returns density
:param l: labels
:param d: distances
:return:
"""
mask = self.mask_input(l)
expanded_params = dict((i, expand_parameter(mask, self.params[i]))
for i in self.params.keys())
return exponential_kernel(
d, **self.filter_args(exponential_kernel, expanded_params))