Skip to content

GitLab

Commit 3af8c49e authored by Gao, Xiang's avatar Gao, Xiang Committed by Farhad Ramezanghorbani
Browse files

Use @ syntax (#371)

parent ffb075e6
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
torchani/aev.py
View file @ 3af8c49e
...@@ -149,7 +149,7 @@ def neighbor_pairs(padding_mask, coordinates, cell, shifts, cutoff): ...@@ -149,7 +149,7 @@ def neighbor_pairs(padding_mask, coordinates, cell, shifts, cutoff):
shifts_all = torch.cat([shifts_center, shifts_outide]) shifts_all = torch.cat([shifts_center, shifts_outide])
p1_all = torch.cat([p1_center, p1]) p1_all = torch.cat([p1_center, p1])
p2_all = torch.cat([p2_center, p2]) p2_all = torch.cat([p2_center, p2])
shift_values = torch.mm(shifts_all.to(cell.dtype), cell) shift_values = shifts_all.to(cell.dtype) @ cell
# step 5, compute distances, and find all pairs within cutoff # step 5, compute distances, and find all pairs within cutoff
distances = (coordinates.index_select(1, p1_all) - coordinates.index_select(1, p2_all) + shift_values).norm(2, -1) distances = (coordinates.index_select(1, p1_all) - coordinates.index_select(1, p2_all) + shift_values).norm(2, -1)
...@@ -241,7 +241,7 @@ def compute_aev(species, coordinates, cell, shifts, triu_index, constants, sizes ...@@ -241,7 +241,7 @@ def compute_aev(species, coordinates, cell, shifts, triu_index, constants, sizes
coordinates = coordinates.flatten(0, 1) coordinates = coordinates.flatten(0, 1)
species1 = species[atom_index1] species1 = species[atom_index1]
species2 = species[atom_index2] species2 = species[atom_index2]
shift_values = torch.mm(shifts.to(cell.dtype), cell) shift_values = shifts.to(cell.dtype) @ cell
vec = coordinates.index_select(0, atom_index1) - coordinates.index_select(0, atom_index2) + shift_values vec = coordinates.index_select(0, atom_index1) - coordinates.index_select(0, atom_index2) + shift_values
distances = vec.norm(2, -1) distances = vec.norm(2, -1)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment