Skip to content

Commit

Permalink
Improved formatting of classes and functions
Browse files Browse the repository at this point in the history
  • Loading branch information
@EnricoTrizio
EnricoTrizio committed May 8, 2024
1 parent c443fc2 commit aee5c97
Show file tree
Hide file tree
Showing 8 changed files with 147 additions and 188 deletions.
45 changes: 19 additions & 26 deletions mlcolvar/core/transform/descriptors/coordination_numbers.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,23 +14,23 @@ class CoordinationNumbers(Transform):
"""

def __init__(self,
group_A : list,
group_B : list,
cutoff : float,
n_atoms : int,
PBC : bool,
cell : Union[float, list],
mode : str,
scaled_coords : bool = False,
group_A: list,
group_B: list,
cutoff: float,
n_atoms: int,
PBC: bool,
cell: Union[float, list],
mode: str,
scaled_coords: bool = False,
switching_function = None) -> torch.Tensor:
"""Initialize a coordination number object between two groups of atoms A and B.
Parameters
----------
group_A : list
Zero-based indeces of group A atoms
Zero-based indices of group A atoms
group_B : list
Zero-based indeces of group B atoms
Zero-based indices of group B atoms
cutoff : float
Cutoff radius for coordination number evaluation
n_atoms : int
Expand Down Expand Up @@ -61,15 +61,11 @@ def __init__(self,
self.group_B = group_B
self._group_B_size = len(group_B)
self._reordering = np.concatenate((self.group_A, self.group_B))


self.cutoff = cutoff

self.n_atoms = n_atoms
self.PBC = PBC
self.cell = cell
self.scaled_coords = scaled_coords

self.mode = mode
self.switching_function = switching_function

Expand All @@ -82,8 +78,6 @@ def compute_coordination_number(self, pos):
PBC=self.PBC,
cell=self.cell,
scaled_coords=self.scaled_coords)
# batch_size = dist.shape[0]
# device = pos.device

# we can apply the switching cutoff with the switching function
contributions = apply_cutoff(x=dist,
Expand Down Expand Up @@ -111,6 +105,7 @@ def forward(self, pos):

def test_coordination_number():
from mlcolvar.core.transform.tools.switching_functions import SwitchingFunctions

# simple example based on calixarene water coordination numbers
pos = torch.Tensor([[[-0.410219, -0.680065, -2.016121],
[-0.164329, -0.630426, -2.120843],
Expand Down Expand Up @@ -138,17 +133,15 @@ def test_coordination_number():
[-0.528576, -0.202031, -1.534733]]])

cell = 4.0273098

pos.requires_grad = True

n_atoms = 12
cutoff=0.25

switching_function=SwitchingFunctions(in_features=n_atoms*3, name='Rational', cutoff=cutoff, options={'n': 2, 'm' : 6, 'eps' : 1e0})

model = CoordinationNumbers(group_A = [0, 1],
group_B = [2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
cutoff= cutoff,
model = CoordinationNumbers(group_A=[0, 1],
group_B=[2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
cutoff=cutoff,
n_atoms=n_atoms,
PBC=True,
cell=cell,
Expand All @@ -159,7 +152,7 @@ def test_coordination_number():
out = model(pos)
out.sum().backward()

# we shift by hand the 0,1 atoms with 2,3
# we swap by hand the 0,1 atoms with 2,3
pos = torch.Tensor([[[-0.250341, -0.392700, -1.534535],
[-0.277187, -0.615506, -1.335904],
[-0.410219, -0.680065, -2.016121],
Expand Down Expand Up @@ -188,9 +181,9 @@ def test_coordination_number():
pos.requires_grad = True
switching_function=SwitchingFunctions(in_features=n_atoms*3, name='Rational', cutoff=cutoff, options={'n': 2, 'm' : 6, 'eps' : 1e0})

model = CoordinationNumbers(group_A = [2, 3],
group_B = [0, 1, 4, 5, 6, 7, 8, 9, 10, 11],
cutoff= cutoff,
model = CoordinationNumbers(group_A=[2, 3],
group_B=[0, 1, 4, 5, 6, 7, 8, 9, 10, 11],
cutoff=cutoff,
n_atoms=n_atoms,
PBC=True,
cell=cell,
Expand Down
48 changes: 24 additions & 24 deletions mlcolvar/core/transform/descriptors/eigs_adjacency_matrix.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,12 +13,12 @@ class EigsAdjMat(Transform):
"""

def __init__(self,
mode : str,
cutoff : float,
n_atoms : int,
mode: str,
cutoff: float,
n_atoms: int,
PBC: bool,
cell: Union[float, list],
scaled_coords : bool = False,
scaled_coords: bool = False,
switching_function = None) -> torch.Tensor:
"""Initialize an eigenvalues of an adjacency matrix object.
Expand Down Expand Up @@ -59,13 +59,13 @@ def __init__(self,

def compute_adjacency_matrix(self, pos):
pos = compute_adjacency_matrix(pos=pos,
mode = self.mode,
cutoff = self.cutoff,
n_atoms = self.n_atoms,
PBC = self.PBC,
cell = self.cell,
scaled_coords = self.scaled_coords,
switching_function=self.switching_function)
mode=self.mode,
cutoff=self.cutoff,
n_atoms=self.n_atoms,
PBC=self.PBC,
cell=self.cell,
scaled_coords=self.scaled_coords,
switching_function=self.switching_function)
return pos

def get_eigenvalues(self, x):
Expand All @@ -87,28 +87,28 @@ def test_eigs_of_adj_matrix():
[1., 1.1, 1.] ] ]
)
pos.requires_grad = True

cell = torch.Tensor([1., 2., 1.])

cutoff = 1.8
switching_function=SwitchingFunctions(in_features=n_atoms*3, name='Fermi', cutoff=cutoff, options={'q':0.05})

model = EigsAdjMat(mode = 'continuous',
cutoff = cutoff,
n_atoms = n_atoms,
PBC = True,
cell = cell,
scaled_coords = False,
model = EigsAdjMat(mode='continuous',
cutoff=cutoff,
n_atoms=n_atoms,
PBC=True,
cell=cell,
scaled_coords=False,
switching_function=switching_function)
out = model(pos)
out.sum().backward()

pos = torch.einsum('bij,j->bij', pos, 1/cell)
model = EigsAdjMat(mode = 'continuous',
cutoff = cutoff,
n_atoms = n_atoms,
PBC = True,
cell = cell,
scaled_coords = True,
model = EigsAdjMat(mode='continuous',
cutoff=cutoff,
n_atoms=n_atoms,
PBC=True,
cell=cell,
scaled_coords=True,
switching_function=switching_function)
out = model(pos)
assert(out.shape[-1] == model.out_features)
Expand Down
15 changes: 6 additions & 9 deletions mlcolvar/core/transform/descriptors/multiple_descriptors.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,28 +63,25 @@ def test_multipledescriptors():
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]])
model_1 = TorsionalAngle(indices=[1,3,4,6], n_atoms=10, mode=['angle'], PBC=False, cell=cell, scaled_coords=False)
model_2 = TorsionalAngle(indices=[3,4,6,8], n_atoms=10, mode=['angle'], PBC=False, cell=cell, scaled_coords=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)
model_tot = MultipleDescriptors(descriptors_list=[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)
model_tot = MultipleDescriptors(descriptors_list=[model_1, model_3], n_atoms=10)
out = model_tot(pos)
out.sum().backward()
for i in range(len(pos)):
Expand Down
30 changes: 10 additions & 20 deletions mlcolvar/core/transform/descriptors/pairwise_distances.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,11 +13,11 @@ class PairwiseDistances(Transform):
"""

def __init__(self,
n_atoms : int,
n_atoms: int,
PBC: bool,
cell: Union[float, list],
scaled_coords : bool = False,
slicing_pairs : list = None) -> torch.Tensor:
scaled_coords: bool = False,
slicing_pairs: list = None) -> torch.Tensor:
"""Initialize a pairwise distances matrix object.
Parameters
Expand All @@ -31,7 +31,7 @@ def __init__(self,
scaled_coords : bool
Switch for coordinates scaled on cell's vectors use, by default False
slicing_pairs : list
Indeces of the subset of distances to be returned, by default None
indices of the subset of distances to be returned, by default None
Returns
-------
Expand Down Expand Up @@ -93,32 +93,22 @@ def test_pairwise_distances():
0.2976, 0.3748, 0.4262, 0.4821, 0.5043, 0.6376, 0.1447, 0.2449, 0.2454,
0.2705, 0.3597, 0.4833, 0.1528, 0.1502, 0.2370, 0.2408, 0.3805, 0.2472,
0.3243, 0.3159, 0.4527, 0.1270, 0.1301, 0.2440, 0.2273, 0.2819, 0.1482]])

# cell = torch.Tensor([1., 2., 1.])

model = PairwiseDistances(n_atoms = 10,
PBC = True,
cell = cell,
scaled_coords = False)
# PBC no scaled coords
model = PairwiseDistances(n_atoms=10, PBC=True, cell=cell, scaled_coords=False)
out = model(pos_abs)
assert(out.reshape(pos_abs.shape[0], -1).shape[-1] == model.out_features)
assert(torch.allclose(out, ref_distances, atol=1e-3))
out.sum().backward()


model = PairwiseDistances(n_atoms = 10,
PBC = True,
cell = cell,
scaled_coords = False,
slicing_pairs=[[0, 1], [0, 2]])
# PBC no scaled coords slicing
model = PairwiseDistances(n_atoms=10, PBC=True, cell=cell, scaled_coords=False, slicing_pairs=[[0, 1], [0, 2]])
out = model(pos_abs)
assert(torch.allclose(out, ref_distances[:, [0, 1]], atol=1e-3))
out.sum().backward()

model = PairwiseDistances(n_atoms = 10,
PBC = True,
cell = cell,
scaled_coords = True)
# PBC and scaled coords
model = PairwiseDistances(n_atoms=10, PBC=True, cell=cell, scaled_coords=True)
out = model(pos_scaled)
assert(out.reshape(pos_scaled.shape[0], -1).shape[-1] == model.out_features)
assert(torch.allclose(out, ref_distances, atol=1e-3))
Expand Down
Loading

0 comments on commit aee5c97

Please sign in to comment.