pyg-nightly 2.6.0.dev20240319__py3-none-any.whl → 2.7.0.dev20250114__py3-none-any.whl

torch_geometric/utils/map.py

@@ -1,6 +1,7 @@
 import warnings
 from typing import Optional, Tuple, Union
 
+import numpy as np
 import torch
 from torch import Tensor
 from torch.utils.dlpack import from_dlpack
@@ -13,7 +14,7 @@ def map_index(
     inclusive: bool = False,
 ) -> Tuple[Tensor, Optional[Tensor]]:
     r"""Maps indices in :obj:`src` to the positional value of their
-    corresponding occurence in :obj:`index`.
+    corresponding occurrence in :obj:`index`.
     Indices must be strictly positive.
 
     Args:
@@ -110,7 +111,12 @@ def map_index(
         result = pd.merge(left_ser, right_ser, how='left', left_on='left_ser',
                           right_index=True)
 
-        out = torch.from_numpy(result['right_ser'].values).to(index.device)
+        out_numpy = result['right_ser'].values
+        if (index.device.type == 'mps'  # MPS does not support `float64`
+                and issubclass(out_numpy.dtype.type, np.floating)):
+            out_numpy = out_numpy.astype(np.float32)
+
+        out = torch.from_numpy(out_numpy).to(index.device)
 
         if out.is_floating_point() and inclusive:
             raise ValueError("Found invalid entries in 'src' that do not have "

torch_geometric/utils/smiles.py

@@ -77,32 +77,18 @@ e_map: Dict[str, List[Any]] = {
 }
 
 
-def from_smiles(smiles: str, with_hydrogen: bool = False,
-                kekulize: bool = False) -> 'torch_geometric.data.Data':
-    r"""Converts a SMILES string to a :class:`torch_geometric.data.Data`
-    instance.
+def from_rdmol(mol: Any) -> 'torch_geometric.data.Data':
+    r"""Converts a :class:`rdkit.Chem.Mol` instance to a
+    :class:`torch_geometric.data.Data` instance.
 
     Args:
-        smiles (str): The SMILES string.
-        with_hydrogen (bool, optional): If set to :obj:`True`, will store
-            hydrogens in the molecule graph. (default: :obj:`False`)
-        kekulize (bool, optional): If set to :obj:`True`, converts aromatic
-            bonds to single/double bonds. (default: :obj:`False`)
+        mol (rdkit.Chem.Mol): The :class:`rdkit` molecule.
     """
-    from rdkit import Chem, RDLogger
+    from rdkit import Chem
 
     from torch_geometric.data import Data
 
-    RDLogger.DisableLog('rdApp.*')
-
-    mol = Chem.MolFromSmiles(smiles)
-
-    if mol is None:
-        mol = Chem.MolFromSmiles('')
-    if with_hydrogen:
-        mol = Chem.AddHs(mol)
-    if kekulize:
-        Chem.Kekulize(mol)
+    assert isinstance(mol, Chem.Mol)
 
     xs: List[List[int]] = []
     for atom in mol.GetAtoms():
@@ -142,16 +128,51 @@ def from_smiles(smiles: str, with_hydrogen: bool = False,
         perm = (edge_index[0] * x.size(0) + edge_index[1]).argsort()
         edge_index, edge_attr = edge_index[:, perm], edge_attr[perm]
 
-    return Data(x=x, edge_index=edge_index, edge_attr=edge_attr, smiles=smiles)
+    return Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
 
 
-def to_smiles(data: 'torch_geometric.data.Data',
-              kekulize: bool = False) -> Any:
-    """Converts a :class:`torch_geometric.data.Data` instance to a SMILES
-    string.
+def from_smiles(
+    smiles: str,
+    with_hydrogen: bool = False,
+    kekulize: bool = False,
+) -> 'torch_geometric.data.Data':
+    r"""Converts a SMILES string to a :class:`torch_geometric.data.Data`
+    instance.
 
     Args:
-        data (torch_geometric.data.Data): The molecular graph.
+        smiles (str): The SMILES string.
+        with_hydrogen (bool, optional): If set to :obj:`True`, will store
+            hydrogens in the molecule graph. (default: :obj:`False`)
+        kekulize (bool, optional): If set to :obj:`True`, converts aromatic
+            bonds to single/double bonds. (default: :obj:`False`)
+    """
+    from rdkit import Chem, RDLogger
+
+    RDLogger.DisableLog('rdApp.*')  # type: ignore
+
+    mol = Chem.MolFromSmiles(smiles)
+
+    if mol is None:
+        mol = Chem.MolFromSmiles('')
+    if with_hydrogen:
+        mol = Chem.AddHs(mol)
+    if kekulize:
+        Chem.Kekulize(mol)
+
+    data = from_rdmol(mol)
+    data.smiles = smiles
+    return data
+
+
+def to_rdmol(
+    data: 'torch_geometric.data.Data',
+    kekulize: bool = False,
+) -> Any:
+    """Converts a :class:`torch_geometric.data.Data` instance to a
+    :class:`rdkit.Chem.Mol` instance.
+
+    Args:
+        data (torch_geometric.data.Data): The molecular graph data.
         kekulize (bool, optional): If set to :obj:`True`, converts aromatic
             bonds to single/double bonds. (default: :obj:`False`)
     """
@@ -172,7 +193,7 @@ def to_smiles(data: 'torch_geometric.data.Data',
             data.x[i, 5])])
         atom.SetHybridization(Chem.rdchem.HybridizationType.values[int(
             data.x[i, 6])])
-        atom.SetIsAromatic(int(data.x[i, 7]))
+        atom.SetIsAromatic(bool(data.x[i, 7]))
         mol.AddAtom(atom)
 
     edges = [tuple(i) for i in data.edge_index.t().tolist()]
@@ -207,4 +228,21 @@ def to_smiles(data: 'torch_geometric.data.Data',
     Chem.SanitizeMol(mol)
     Chem.AssignStereochemistry(mol)
 
+    return mol
+
+
+def to_smiles(
+    data: 'torch_geometric.data.Data',
+    kekulize: bool = False,
+) -> str:
+    """Converts a :class:`torch_geometric.data.Data` instance to a SMILES
+    string.
+
+    Args:
+        data (torch_geometric.data.Data): The molecular graph.
+        kekulize (bool, optional): If set to :obj:`True`, converts aromatic
+            bonds to single/double bonds. (default: :obj:`False`)
+    """
+    from rdkit import Chem
+    mol = to_rdmol(data, kekulize=kekulize)
     return Chem.MolToSmiles(mol, isomericSmiles=True)

torch_geometric/utils/sparse.py

@@ -6,6 +6,7 @@ import torch
 from torch import Tensor
 
 import torch_geometric.typing
+from torch_geometric.index import index2ptr, ptr2index
 from torch_geometric.typing import SparseTensor
 from torch_geometric.utils import coalesce, cumsum
 
@@ -197,15 +198,23 @@ def to_torch_coo_tensor(
         # edge_attr = edge_attr.expand(edge_index.size(1))
         edge_attr = torch.ones(edge_index.size(1), device=edge_index.device)
 
-    adj = torch.sparse_coo_tensor(
+    if not torch_geometric.typing.WITH_PT21:
+        adj = torch.sparse_coo_tensor(
+            indices=edge_index,
+            values=edge_attr,
+            size=tuple(size) + edge_attr.size()[1:],
+            device=edge_index.device,
+        )
+        adj = adj._coalesced_(True)
+        return adj
+
+    return torch.sparse_coo_tensor(
         indices=edge_index,
         values=edge_attr,
         size=tuple(size) + edge_attr.size()[1:],
         device=edge_index.device,
+        is_coalesced=True,
     )
-    adj = adj._coalesced_(True)
-
-    return adj
 
 
 def to_torch_csr_tensor(
@@ -483,65 +492,70 @@ def set_sparse_value(adj: Tensor, value: Tensor) -> Tensor:
     raise ValueError(f"Unexpected sparse tensor layout (got '{adj.layout}')")
 
 
-def ptr2index(ptr: Tensor, output_size: Optional[int] = None) -> Tensor:
-    index = torch.arange(ptr.numel() - 1, dtype=ptr.dtype, device=ptr.device)
-    return index.repeat_interleave(ptr.diff(), output_size=output_size)
-
-
-def index2ptr(index: Tensor, size: Optional[int] = None) -> Tensor:
-    if size is None:
-        size = int(index.max()) + 1 if index.numel() > 0 else 0
-
-    return torch._convert_indices_from_coo_to_csr(
-        index, size, out_int32=index.dtype == torch.int32)
-
-
 def cat_coo(tensors: List[Tensor], dim: Union[int, Tuple[int, int]]) -> Tensor:
     assert dim in {0, 1, (0, 1)}
     assert tensors[0].layout == torch.sparse_coo
 
     indices, values = [], []
     num_rows = num_cols = 0
+    is_coalesced = True
 
     if dim == 0:
         for i, tensor in enumerate(tensors):
             if i == 0:
-                indices.append(tensor.indices())
+                indices.append(tensor._indices())
             else:
                 offset = torch.tensor([[num_rows], [0]], device=tensor.device)
-                indices.append(tensor.indices() + offset)
-            values.append(tensor.values())
+                indices.append(tensor._indices() + offset)
+            values.append(tensor._values())
             num_rows += tensor.size(0)
             num_cols = max(num_cols, tensor.size(1))
+            if not tensor.is_coalesced():
+                is_coalesced = False
 
     elif dim == 1:
         for i, tensor in enumerate(tensors):
             if i == 0:
-                indices.append(tensor.indices())
+                indices.append(tensor._indices())
             else:
                 offset = torch.tensor([[0], [num_cols]], device=tensor.device)
                 indices.append(tensor.indices() + offset)
-            values.append(tensor.values())
+            values.append(tensor._values())
             num_rows = max(num_rows, tensor.size(0))
             num_cols += tensor.size(1)
+            is_coalesced = False
 
     else:
         for i, tensor in enumerate(tensors):
             if i == 0:
-                indices.append(tensor.indices())
+                indices.append(tensor._indices())
             else:
                 offset = torch.tensor([[num_rows], [num_cols]],
                                       device=tensor.device)
-                indices.append(tensor.indices() + offset)
-            values.append(tensor.values())
+                indices.append(tensor._indices() + offset)
+            values.append(tensor._values())
             num_rows += tensor.size(0)
             num_cols += tensor.size(1)
+            if not tensor.is_coalesced():
+                is_coalesced = False
+
+    if not torch_geometric.typing.WITH_PT21:
+        out = torch.sparse_coo_tensor(
+            indices=torch.cat(indices, dim=-1),
+            values=torch.cat(values),
+            size=(num_rows, num_cols) + values[-1].size()[1:],
+            device=tensor.device,
+        )
+        if is_coalesced:
+            out = out._coalesced_(True)
+        return out
 
     return torch.sparse_coo_tensor(
         indices=torch.cat(indices, dim=-1),
         values=torch.cat(values),
         size=(num_rows, num_cols) + values[-1].size()[1:],
         device=tensor.device,
+        is_coalesced=True if is_coalesced else None,
     )
 
 

torch_geometric/visualization/graph.py

@@ -132,7 +132,7 @@ def _visualize_graph_via_networkx(
             xy=pos[src],
             xytext=pos[dst],
             arrowprops=dict(
-                arrowstyle="->",
+                arrowstyle="<-",
                 alpha=data['alpha'],
                 shrinkA=sqrt(node_size) / 2.0,
                 shrinkB=sqrt(node_size) / 2.0,
@@ -140,9 +140,8 @@ def _visualize_graph_via_networkx(
             ),
         )
 
-    nodes = nx.draw_networkx_nodes(g, pos, node_size=node_size,
-                                   node_color='white', margins=0.1)
-    nodes.set_edgecolor('black')
+    nx.draw_networkx_nodes(g, pos, node_size=node_size, node_color='white',
+                           margins=0.1, edgecolors='black')
     nx.draw_networkx_labels(g, pos, font_size=10)
 
     if path is not None: