From c443fc2ce3eedf4abf9654acc55b2479c43a2895 Mon Sep 17 00:00:00 2001 From: EnricoTrizio Date: Wed, 8 May 2024 11:32:51 +0200 Subject: [PATCH] Added multiple descriptors wrapper class --- mlcolvar/core/transform/__init__.py | 2 +- .../core/transform/descriptors/__init__.py | 5 +- .../descriptors/multiple_descriptors.py | 92 +++++++++++++++++++ ...test_core_transform_multipledescriptors.py | 4 + 4 files changed, 100 insertions(+), 3 deletions(-) create mode 100644 mlcolvar/core/transform/descriptors/multiple_descriptors.py create mode 100644 mlcolvar/tests/test_core_transform_multipledescriptors.py diff --git a/mlcolvar/core/transform/__init__.py b/mlcolvar/core/transform/__init__.py index 53c19b2f..80e88304 100644 --- a/mlcolvar/core/transform/__init__.py +++ b/mlcolvar/core/transform/__init__.py @@ -1,4 +1,4 @@ -__all__ = ["Transform","Normalization","Statistics","SwitchingFunctions","PairwiseDistances","EigsAdjMat","ContinuousHistogram","Inverse",'TorsionalAngle'] +__all__ = ["Transform","Normalization","Statistics","SwitchingFunctions","MultipleDescriptors","PairwiseDistances","EigsAdjMat","ContinuousHistogram","Inverse",'TorsionalAngle'] from .transform import * from .utils import * diff --git a/mlcolvar/core/transform/descriptors/__init__.py b/mlcolvar/core/transform/descriptors/__init__.py index 2bc10bb7..588cce97 100644 --- a/mlcolvar/core/transform/descriptors/__init__.py +++ b/mlcolvar/core/transform/descriptors/__init__.py @@ -1,6 +1,7 @@ -__all__ = ["CoordinationNumbers", "EigsAdjMat", "PairwiseDistances", "TorsionalAngle"] +__all__ = ["MultipleDescriptors", "CoordinationNumbers", "EigsAdjMat", "PairwiseDistances", "TorsionalAngle"] from .coordination_numbers import * from .eigs_adjacency_matrix import * from .pairwise_distances import * -from .torsional_angle import * \ No newline at end of file +from .torsional_angle import * +from .multiple_descriptors import * \ No newline at end of file diff --git a/mlcolvar/core/transform/descriptors/multiple_descriptors.py b/mlcolvar/core/transform/descriptors/multiple_descriptors.py new file mode 100644 index 00000000..89dc0cb5 --- /dev/null +++ b/mlcolvar/core/transform/descriptors/multiple_descriptors.py @@ -0,0 +1,92 @@ +import torch + +__all__ = ["MultipleDescriptors"] + +class MultipleDescriptors(torch.nn.Module): + """Wrapper class to combine multiple descriptor transform objects acting on the same set of atomic positions in a single module""" + def __init__(self, + descriptors_list: list, + n_atoms: int, + ): + """_summary_ + + Parameters + ---------- + descriptors_list : list + List of descriptor transform objects to be combined + n_atoms : int + Number of atoms in the system + """ + super().__init__() + self.in_features = n_atoms * 3 + self.descriptors_list = descriptors_list + + self.out_features = 0 + for d in self.descriptors_list: + self.out_features += d.out_features + + def forward(self, pos): + for i,d in enumerate(self.descriptors_list): + if i == 0: + out = d(pos) + else: + aux = d(pos) + out = torch.concatenate((out, aux), 1) + return out + +def test_multipledescriptors(): + from .torsional_angle import TorsionalAngle + from .pairwise_distances import PairwiseDistances + + # check using torsional angles and distances in alanine + pos = torch.Tensor([[[ 0.3887, -0.4169, -0.1212], + [ 0.4264, -0.4374, -0.0983], + [ 0.4574, -0.4136, -0.0931], + [ 0.4273, -0.4797, -0.0871], + [ 0.4684, 0.4965, -0.0692], + [ 0.4478, 0.4571, -0.0441], + [-0.4933, 0.4869, -0.1026], + [-0.4840, 0.4488, -0.1116], + [-0.4748, -0.4781, -0.1232], + [-0.4407, -0.4781, -0.1569]], + [[ 0.3910, -0.4103, -0.1189], + [ 0.4334, -0.4329, -0.1020], + [ 0.4682, -0.4145, -0.1013], + [ 0.4322, -0.4739, -0.0867], + [ 0.4669, -0.4992, -0.0666], + [ 0.4448, 0.4670, -0.0375], + [-0.4975, 0.4844, -0.0981], + [-0.4849, 0.4466, -0.0991], + [-0.4818, -0.4870, -0.1291], + [-0.4490, -0.4933, -0.1668]]]) + pos.requires_grad = True + cell = torch.Tensor([3.0233, 3.0233, 3.0233]) + + # model 1 and 2 for torsional angles, model 3 for distances + model_1 = TorsionalAngle([1,3,4,6], 10, ['angle'], False, cell, False) + model_2 = TorsionalAngle([3,4,6,8], 10, ['angle'], False, cell, False) + model_3 = PairwiseDistances(n_atoms = 10, + PBC = True, + cell = cell, + scaled_coords = False, + slicing_pairs=[[0, 1], [0, 2]]) + # compute single references + angle_1 = model_1(pos) + angle_2 = model_2(pos) + distances = model_3(pos) + + # stack torsional angles + model_tot = MultipleDescriptors([model_1, model_2], n_atoms = 10) + out = model_tot(pos) + out.sum().backward() + for i in range(len(pos)): + assert(torch.allclose(out[i, 0], angle_1[i])) + assert(torch.allclose(out[i, 1], angle_2[i])) + + # stack torsional angle and two distances + model_tot = MultipleDescriptors([model_1, model_3], n_atoms = 10) + out = model_tot(pos) + out.sum().backward() + for i in range(len(pos)): + assert(torch.allclose(out[i, 0], angle_1[i])) + assert(torch.allclose(out[i, 1:], distances[i])) \ No newline at end of file diff --git a/mlcolvar/tests/test_core_transform_multipledescriptors.py b/mlcolvar/tests/test_core_transform_multipledescriptors.py new file mode 100644 index 00000000..6c3e92bf --- /dev/null +++ b/mlcolvar/tests/test_core_transform_multipledescriptors.py @@ -0,0 +1,4 @@ +from mlcolvar.core.transform.descriptors.multiple_descriptors import test_multipledescriptors + +if __name__ == "__main__": + test_multipledescriptors() \ No newline at end of file