Merge pull request #166 from luigibonati/fix_scatter

Remove torch-scatter dependency (add only torch-scatter)

Author: luigibonati

devtools/conda-envs/test_env.yaml

@@ -29,5 +29,4 @@ dependencies:
   - pip:
     - KDEpy
     - nbmake
-    - torch-scatter
 

docs/notebooks/examples/.gitignore

@@ -0,0 +1 @@
+*.pt

docs/requirements.yaml

@@ -26,7 +26,6 @@ dependencies:
   - ipykernel
   - scikit-learn
   - scipy
-  - torch-scatter
 
   # Pip-only installs
   - pip:

mlcolvar/core/loss/committor_loss.py

@@ -15,7 +15,7 @@
 # =============================================================================
 
 import torch
-from torch_scatter import scatter
+from typing import Optional
 
 # =============================================================================
 # LOSS FUNCTIONS
@@ -82,6 +82,37 @@ def forward(
                                 descriptors_derivatives=self.descriptors_derivatives
                             )
 
+def broadcast(src: torch.Tensor, other: torch.Tensor, dim: int):
+    """Broadcast util, from torch_scatter"""
+    if dim < 0:
+        dim = other.dim() + dim
+    if src.dim() == 1:
+        for _ in range(0, dim):
+            src = src.unsqueeze(0)
+    for _ in range(src.dim(), other.dim()):
+        src = src.unsqueeze(-1)
+    src = src.expand(other.size())
+    return src
+
+def scatter_sum(src: torch.Tensor,
+                index: torch.Tensor,
+                dim: int = -1,
+                out: Optional[torch.Tensor] = None,
+                dim_size: Optional[int] = None) -> torch.Tensor:
+    """Scatter sum function, from torch_scatter module (https://github.com/rusty1s/pytorch_scatter/blob/master/torch_scatter/scatter.py)"""
+    index = broadcast(index, src, dim)
+    if out is None:
+        size = list(src.size())
+        if dim_size is not None:
+            size[dim] = dim_size
+        elif index.numel() == 0:
+            size[dim] = 0
+        else:
+            size[dim] = int(index.max()) + 1
+        out = torch.zeros(size, dtype=src.dtype, device=src.device)
+        return out.scatter_add_(dim, index, src)
+    else:
+        return out.scatter_add_(dim, index, src)
 
 def committor_loss(x: torch.Tensor, 
                   q: torch.Tensor, 
@@ -291,7 +322,7 @@ def _get_scatter_indices(self, batch_ind, atom_ind, dim_ind):
 
         # get the number of elements in each batch 
         # e.g. [17, 18, 18, 18] 
-        batch_elements = scatter(torch.ones_like(batch_ind), batch_ind, reduce='sum')
+        batch_elements = scatter_sum(torch.ones_like(batch_ind), batch_ind)
         batch_elements[0] -= 1 # to make the later indexing consistent
 
         # compute the pointer idxs to the beginning of each batch by summing the number of elements in each batch
@@ -344,7 +375,7 @@ def _compute_square_modulus(self, x : torch.Tensor, indeces : torch.Tensor, n_at
         indeces = indeces.long().to(x.device)
 
         # this sums the elements of x according to the indeces, this way we get the contributions of different descriptors to the same atom
-        out = scatter(x, indeces.long())
+        out = scatter_sum(x, indeces.long())
         # now make the square
         out = out.pow(2)
         # reshape, this needs to have the correct number of atoms as we need to mulply it by the mass vector later

requirements.txt

@@ -3,5 +3,4 @@ torch
 numpy<2
 pandas
 matplotlib
-kdepy
-torch-scatter
+kdepy