Add option to set a custom vdw radii when loading structure

AtomPacker/core/io/mmcif/file.py

@@ -9,21 +9,27 @@
 __all__ = ["load_mmcif"]
 
 
+from typing import Dict, Optional
+
 import numpy
 from Bio.PDB.MMCIF2Dict import MMCIF2Dict
 from MDAnalysis import Universe
 
 from ..vdw import _lookup_radii
 
 
-def load_mmcif(filename: str) -> Universe:
+def load_mmcif(filename: str, vdw: Optional[Dict[str, float]] = None) -> Universe:
     """
     Load a PDBx/mmCIF file into the :class:`MDAnalysis.Univese` object.
 
     Parameters
     ----------
     filename : str
         The filename of the structure file.
+    vdw : Dict[str, float]
+        A dictionary containing the van der Waals radii for each atom type,
+        by default None. If None, the radii will be looked up from the
+        `pyKVFinder` package.
 
     Returns
     -------
@@ -51,7 +57,7 @@ def load_mmcif(filename: str) -> Universe:
     # Add radii to topology
     universe.add_TopologyAttr(
         "radii",
-        _lookup_radii(universe.atoms.names),
+        _lookup_radii(universe.atoms.names, vdw),
     )
 
     return universe

AtomPacker/core/io/mol2/file.py

@@ -10,6 +10,7 @@
 
 import warnings
 from string import digits
+from typing import Dict, Optional
 
 from MDAnalysis import Universe
 
@@ -19,14 +20,18 @@
 warnings.filterwarnings("ignore", category=UserWarning)
 
 
-def load_mol2(filename: str) -> Universe:
+def load_mol2(filename: str, vdw: Optional[Dict[str, float]] = None) -> Universe:
     """
     Load a MOL2 file into the :class:`MDAnalysis.Univese` object.
 
     Parameters
     ----------
     filename : str
         The filename of the structure file.
+    vdw : Dict[str, float]
+        A dictionary containing the van der Waals radii for each atom type,
+        by default None. If None, the radii will be looked up from the
+        `pyKVFinder` package.
 
     Returns
     -------
@@ -57,7 +62,7 @@ def load_mol2(filename: str) -> Universe:
     # Add radii to topology
     universe.add_TopologyAttr(
         "radii",
-        _lookup_radii(universe.atoms.elements),
+        _lookup_radii(universe.atoms.elements, vdw),
     )
 
     return universe

AtomPacker/core/io/pdb/file.py

@@ -8,19 +8,25 @@
 
 __all__ = ["load_pdb"]
 
+from typing import Dict, Optional
+
 from MDAnalysis import Universe
 
 from ..vdw import _lookup_radii
 
 
-def load_pdb(filename: str) -> Universe:
+def load_pdb(filename: str, vdw: Optional[Dict[str, float]] = None) -> Universe:
     """
     Load a PDB file into the :class:`MDAnalysis.Univese` object.
 
     Parameters
     ----------
     filename : str
         The filename of the structure file.
+    vdw : Dict[str, float]
+        A dictionary containing the van der Waals radii for each atom type,
+        by default None. If None, the radii will be looked up from the
+        `pyKVFinder` package.
 
     Returns
     -------
@@ -33,7 +39,7 @@ def load_pdb(filename: str) -> Universe:
     # Add radii to topology
     universe.add_TopologyAttr(
         "radii",
-        _lookup_radii(universe.atoms.names),
+        _lookup_radii(universe.atoms.names, vdw),
     )
 
     return universe

AtomPacker/core/io/vdw/file.py

@@ -8,27 +8,43 @@
 
 __all__ = ["_lookup_radii"]
 
+from typing import Dict, Optional
+
 import numpy
 from pyKVFinder import read_vdw
 
 
-def _lookup_radii(names: numpy.ndarray) -> numpy.ndarray:
+def _lookup_radii(
+    names: numpy.ndarray, vdw: Optional[Dict[str, float]]
+) -> numpy.ndarray:
     """
     Look up the van der Waals radii for a list of atom names.
 
     Parameters
     ----------
-    vdw : Dict[str, float]
-        A dictionary containing the van der Waals radii for each atom type.
     names : numpy.ndarray
         An array containing the names of the atoms.
+    vdw : Dict[str, float]
+        A dictionary containing the van der Waals radii for each atom type.
+        If None, the van der Waals radii are looked up from the `pyKVFinder`
+        package.
 
     Returns
     -------
     radii : numpy.ndarray
         An array containing the van der Waals radii for each atom.
+
+    Raises
+    ------
+    ValueError
+        If not all elements are provided in the `vdw` dictionary.
     """
     # Read van der Waals radii
-    vdw = read_vdw()["GEN"]
-
-    return numpy.vectorize(lambda x: vdw[x.upper()])(names)
+    if vdw is None:
+        vdw = read_vdw()["GEN"]
+        return numpy.vectorize(lambda x: vdw[x.upper()])(names)
+    elif isinstance(vdw, dict):
+        # Check if all elements are provided in vdw
+        if not all([name in vdw for name in names]):
+            raise ValueError("Not all elements are provided in the `vdw` dictionary.")
+        return numpy.vectorize(lambda x: vdw[x])(names)

AtomPacker/core/io/xyz/file.py

@@ -8,19 +8,25 @@
 
 __all__ = ["load_xyz"]
 
+from typing import Dict, Optional
+
 from MDAnalysis import Universe
 
 from ..vdw import _lookup_radii
 
 
-def load_xyz(filename: str) -> Universe:
+def load_xyz(filename: str, vdw: Optional[Dict[str, float]] = None) -> Universe:
     """
     Load a XYZ file into the :class:`MDAnalysis.Univese` object.
 
     Parameters
     ----------
     filename : str
         The filename of the structure file.
+    vdw : Dict[str, float]
+        A dictionary containing the van der Waals radii for each atom type,
+        by default None. If None, the radii will be looked up from the
+        `pyKVFinder` package.
 
     Returns
     -------
@@ -35,17 +41,15 @@ def load_xyz(filename: str) -> Universe:
         universe.add_TopologyAttr(
             "chainIDs", ["X"] * universe.atoms.n_atoms
         )  # chainIDs
-    
+
     # if resnames not provided, set them to 'X'
     if "resnames" not in universe.atoms._SETATTR_WHITELIST:
-        universe.add_TopologyAttr(
-            "resnames", ["UNK"]
-        )  # resnames 
+        universe.add_TopologyAttr("resnames", ["UNK"])  # resnames
 
     # Add radii to topology
     universe.add_TopologyAttr(
         "radii",
-        _lookup_radii(universe.atoms.names),
+        _lookup_radii(universe.atoms.names, vdw),
     )
 
     return universe

AtomPacker/structure/Cage.py

@@ -11,8 +11,8 @@
 """
 
 import itertools
-import os
 import logging
+import os
 import warnings
 from copy import deepcopy
 from typing import Any, Dict, Optional, Tuple
@@ -31,7 +31,7 @@
 from .data import get_lattice_constants
 
 
-class Cage:
+class Cage(object):
     """
     A container for the atoms of a macromolecular structure.
 
@@ -174,7 +174,7 @@ def detect_cavity(
         if self.cluster is not None:
             self.cluster = None
 
-    def load(self, filename: str) -> None:
+    def load(self, filename: str, vdw: Optional[Dict[str, float]] = None) -> None:
         """
         Load a supramolecular cage structure file into the :class:`MDAnalysis
         .Univese` object.
@@ -184,6 +184,10 @@ def load(self, filename: str) -> None:
         filename : str
             The filename of the structure file. The file format is determined
             by the suffix.  Supported formats are: .cif, .mol2, .pdb, .xyz.
+        vdw : Dict[str, float], optional
+            A dictionary containing the van der Waals radii for each atom type,
+            by default None. If None, the van der Waals radii are looked up
+            from the `pyKVFinder` package.
 
         Returns
         -------
@@ -201,13 +205,13 @@ def load(self, filename: str) -> None:
         # Match the suffix to the appropriate file format
         match suffix:
             case ".cif":
-                self.universe = load_mmcif(filename)
+                self.universe = load_mmcif(filename, vdw)
             case ".mol2":
-                self.universe = load_mol2(filename)
+                self.universe = load_mol2(filename, vdw)
             case ".pdb":
-                self.universe = load_pdb(filename)
+                self.universe = load_pdb(filename, vdw)
             case ".xyz":
-                self.universe = load_xyz(filename)
+                self.universe = load_xyz(filename, vdw)
             case _:
                 raise ValueError(
                     f"Unsupported file format: {suffix}. Supported formats \

AtomPacker/structure/data.py

@@ -13,7 +13,7 @@
 
 __all__ = ["get_lattice_constants", "lattice_constants"]
 
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, Tuple
 from warnings import warn
 
 

tests/test_cage.py

@@ -10,7 +10,7 @@
 from AtomPacker import Cage, Cavity, Cluster
 from AtomPacker.core.io.vdw.file import read_vdw
 
-DATADIR = os.path.join(os.path.abspath(os.path.dirname(__file__)),  "data")
+DATADIR = os.path.join(os.path.abspath(os.path.dirname(__file__)), "data")
 
 
 # Define a fixture for different cage file formats
@@ -64,6 +64,18 @@ def test_inexistent_file():
         cage.load("inexistent_file.pdb")
 
 
+def test_custom_vdw(pdb):
+    cage = Cage()
+    cage.load(pdb, vdw={"H": 0.91, "C": 1.66, "Si": 2.1, "O": 1.69})
+    assert isinstance(cage.universe, Universe)
+
+
+def test_custom_vdw_with_missing_atom(pdb):
+    cage = Cage()
+    with pytest.raises(ValueError):
+        cage.load(pdb, vdw={"H": 0.91, "C": 1.66, "Si": 2.1})
+
+
 def test_atomic(pdb):
     cage = Cage()
     cage.load(pdb)