Merge pull request #223 from nlesc-nano/plams_ase_fox

ENH: Add support for lattices in PLAMS/ASE/Auto-FOX interconversions

FOX/classes/multi_mol.py

@@ -2124,7 +2124,6 @@ def as_Molecule(self, mol_subset: MolSubset = None,
         at_symbols = self.symbol
 
         # Construct a template molecule and fill it with atoms
-        assert self.atoms is not None
         mol_template = Molecule()
         mol_template.properties = self.properties.copy()
         add_atom = mol_template.add_atom
@@ -2145,31 +2144,41 @@ def as_Molecule(self, mol_subset: MolSubset = None,
                     at2 = mol_template[bond_idx[j]]
                     add_bond(Bond(atom1=at1, atom2=at2, order=order/10.0))
 
+        if self.lattice is None:
+            lattice_iter = ([] for _ in self[m_subset])
+        elif self.lattice.ndim == 2:
+            lattice_iter = (self.lattice.tolist() for _ in self[m_subset])
+        elif self.lattice.ndim == 3:
+            lattice_iter = iter(self.lattice.tolist())
+        else:
+            raise ValueError
+
         # Create copies of the template molecule; update their cartesian coordinates
         ret: List[Molecule] = []
         ret_append = ret.append
-        for i, xyz in enumerate(self[m_subset]):
+        for i, (lat, xyz) in enumerate(zip(lattice_iter, self[m_subset])):
             mol = mol_template.copy()
             mol.from_array(xyz[at_subset])
             mol.properties.frame = i
+            mol.lattice = lat
             ret_append(mol)
 
         return ret
 
     @classmethod
     def from_Molecule(cls: Type[MT], mol_list: Union[Molecule, Sequence[Molecule]],
-                      subset: Container[str] = frozenset({'atoms'})) -> MT:
+                      subset: None | Container[str] = frozenset({'atoms'})) -> MT:
         """Construct a :class:`.MultiMolecule` instance from one or more PLAMS molecules.
 
         Parameters
         ----------
         mol_list : :class:`plams.Molecule <scm.plams.mol.molecule.Molecule>` or :class:`Sequence[plams.Molecule] <collections.abc.Sequence>`
             A PLAMS molecule or list of PLAMS molecules.
-        subset : :class:`Container[str] <collections.abc.Container>`
+        subset : :class:`Container[str] <collections.abc.Container>`, optional
             Transfer a subset of *plams.Molecule* attributes to this instance.
             If :data:`None`, transfer all attributes.
             Accepts one or more of the following values as strings:
-            ``"properties"``, ``"atoms"`` and/or ``"bonds"``.
+            ``"properties"``, ``"atoms"``, ``"lattice"`` and/or ``"bonds"``.
 
         Returns
         -------
@@ -2182,7 +2191,7 @@ def from_Molecule(cls: Type[MT], mol_list: Union[Molecule, Sequence[Molecule]],
             mol_list = (mol_list,)
         else:
             plams_mol = mol_list[0]
-        subset = subset or {'atoms', 'bonds', 'properties'}
+        subset = subset if subset is not None else {'atoms', 'bonds', 'properties', 'lattice'}
 
         # Convert coordinates
         n_mol = len(mol_list)
@@ -2209,6 +2218,10 @@ def from_Molecule(cls: Type[MT], mol_list: Union[Molecule, Sequence[Molecule]],
                                         bond in plams_mol.bonds], dtype=int)
             plams_mol.unset_atoms_id()
 
+        # Convert lattice
+        if 'lattice' in subset:
+            lat = np.array([m.lattice for m in mol_list], dtype=np.float64)
+            kwargs['lattice'] = lat if lat.size != 0 else None
         return cls(coords, **kwargs)
 
     def as_ase(self, mol_subset: MolSubset = None,
@@ -2243,7 +2256,21 @@ def as_ase(self, mol_subset: MolSubset = None,
 
         symbols = self.symbol[j]
         positions_iter = self[i, j]
-        return [Atoms(symbols=symbols, positions=p) for p in positions_iter]
+
+        if "cell" in kwargs:
+            lattice_iter = repeat(kwargs.pop("cell"))
+        else:
+            if self.lattice is None:
+                lattice_iter = repeat(None)
+            elif self.lattice.ndim == 2:
+                lattice_iter = repeat(self.lattice)
+            elif self.lattice.ndim == 3:
+                lattice_iter = iter(self.lattice[i])
+            else:
+                raise ValueError
+
+        return [Atoms(symbols=symbols, positions=p, cell=lat, **kwargs) for lat, p in
+                zip(lattice_iter, positions_iter)]
 
     @classmethod
     def from_ase(cls: Type[MT], mol_list: Union[Atoms, Sequence[Atoms]]) -> MT:
@@ -2271,7 +2298,11 @@ def from_ase(cls: Type[MT], mol_list: Union[Atoms, Sequence[Atoms]]) -> MT:
         coords = [m.positions for m in mol_list]
         _atoms = group_by_values(enumerate(mol_list[0].numbers))
         atoms = {PeriodicTable.get_symbol(k): v for k, v in _atoms.items()}
-        return cls(coords, atoms=atoms)
+
+        lattice = np.array([m.cell for m in mol_list], dtype=np.float64)
+        if not lattice.any():
+            lattice = None
+        return cls(coords, atoms=atoms, lattice=lattice)
 
     @classmethod
     def from_xyz(cls: Type[MT], filename: Union[str, bytes, PathLike],

FOX/classes/multi_mol_magic.py

@@ -235,7 +235,6 @@ def __new__(
         atoms: Optional[IdxMapping] = None,
         bonds: Optional[np.ndarray] = None,
         properties: Optional[Mapping] = None,
-        *,
         atoms_alias: Optional[Mapping[str, slice]] = None,
         lattice: None | npt.ArrayLike = None,
     ) -> MT:
@@ -532,3 +531,16 @@ def __str__(self) -> str:
     def __repr__(self) -> str:
         """Return the canonical string representation of this instance."""
         return f'{self.__class__.__name__}(..., shape={self.shape}, dtype={repr(self.dtype.name)})'
+
+    def __reduce__(self: MT) -> Tuple[MT, Tuple[Any, Any, Any, Any, Any, Any]]:
+        """Helper function for :mod:`pickle`."""
+        cls = type(self)
+        args = (
+            self.view(np.ndarray),
+            self._atoms,
+            self._bonds,
+            self._properties,
+            self._atoms_alias,
+            self._lattice,
+        )
+        return cls, args

tests/test_multi_mol.py

@@ -2,10 +2,13 @@
 
 from __future__ import annotations
 
+import copy
+import pickle
+import weakref
 from os.path import join
 from pathlib import Path
 from itertools import chain, combinations
-from typing import Mapping, Any, Type, Set, Iterable
+from typing import Mapping, Any, Type, Set, Iterable, Callable
 
 import pytest
 import yaml
@@ -370,23 +373,47 @@ def test_from_mass_weighted():
     assertion.allclose(np.abs((mol_new - mol).mean()), 0.0, abs_tol=10**-8)
 
 
-def test_as_molecule():
+@pytest.mark.parametrize(
+    "mol", [MOL[::10], MOL_LATTICE_2D, MOL_LATTICE_3D]
+)
+def test_as_molecule(mol: MultiMolecule):
     """Test :meth:`.MultiMolecule.as_Molecule`."""
-    mol = MOL.copy()
-
     mol_list = mol.as_Molecule()
     mol_array = np.array([i.as_array() for i in mol_list])
     np.testing.assert_allclose(mol_array, mol)
 
+    lattice = np.array([m.lattice for m in mol_list])
 
-def test_from_molecule():
-    """Test :meth:`.MultiMolecule.from_Molecule`."""
-    mol = MOL.copy()
+    assert mol.lattice is None or mol.lattice.ndim in {2, 3}
+    if mol.lattice is None:
+        assertion.eq(lattice.size, 0)
+    elif mol.lattice.ndim == 2:
+        lat_ref = np.full((len(mol), 3, 3), mol.lattice)
+        np.testing.assert_allclose(lattice, lat_ref)
+    elif mol.lattice.ndim == 3:
+        np.testing.assert_allclose(lattice, mol.lattice)
 
+
+@pytest.mark.parametrize(
+    "mol", [MOL[::10], MOL_LATTICE_2D, MOL_LATTICE_3D]
+)
+def test_from_molecule(mol: MultiMolecule):
+    """Test :meth:`.MultiMolecule.from_Molecule`."""
     mol_list = mol.as_Molecule()
-    mol_new = MultiMolecule.from_Molecule(mol_list)
+    mol_new = MultiMolecule.from_Molecule(mol_list, subset=None)
     np.testing.assert_allclose(mol_new, mol)
 
+    assertion.eq(type(mol.lattice), type(mol_new.lattice))
+    assert mol.lattice is None or mol.lattice.ndim in {2, 3}
+
+    if mol.lattice is None:
+        assertion.is_(mol_new.lattice, mol.lattice)
+    elif mol.lattice.ndim == 2:
+        lat_ref = np.full(mol_new.lattice.shape, mol.lattice)
+        np.testing.assert_allclose(mol_new.lattice, lat_ref)
+    elif mol.lattice.ndim == 3:
+        np.testing.assert_allclose(mol_new.lattice, mol.lattice)
+
 
 @delete_finally(join(PATH, 'mol.xyz'))
 def test_as_xyz():
@@ -483,23 +510,76 @@ def test_raises(self, name: str, exc: Type[Exception], dct: Mapping[str, Any]) -
 
 
 @pytest.mark.parametrize(
-    "kwargs",
+    "mol,kwargs",
     [
-        {'mol_subset': np.s_[::10]},
-        {'atom_subset': ['Cd', 'Se']},
-        {'mol_subset': np.s_[10:30], 'atom_subset': ['Cd', 'O']},
-        {'masses': MOL.mass},
-    ]
+        (MOL[::10], {'mol_subset': np.s_[::10]}),
+        (MOL[::10], {'atom_subset': ['Cd', 'Se']}),
+        (MOL[::10], {'mol_subset': np.s_[10:30], 'atom_subset': ['Cd', 'O']}),
+        (MOL[::10], {'masses': MOL[::10].mass}),
+        (MOL_LATTICE_3D, {'mol_subset': np.s_[::10]}),
+        (MOL_LATTICE_3D, {'atom_subset': ['Pb', 'Cs']}),
+        (MOL_LATTICE_3D, {'mol_subset': np.s_[5:15], 'atom_subset': ['Pb', 'Br']}),
+        (MOL_LATTICE_3D, {'masses': MOL_LATTICE_3D.mass}),
+        (MOL_LATTICE_2D, {'mol_subset': np.s_[::10]}),
+        (MOL_LATTICE_2D, {'atom_subset': ['Pb', 'Cs']}),
+        (MOL_LATTICE_2D, {'mol_subset': np.s_[5:15], 'atom_subset': ['Pb', 'Br']}),
+        (MOL_LATTICE_2D, {'masses': MOL_LATTICE_2D.mass}),
+    ],
 )
-def test_ase(kwargs) -> None:
+def test_ase(mol: MultiMolecule, kwargs: Mapping[str, Any]) -> None:
     """Test :meth:`MultiMolecule.as_ase` and :meth:`MultiMolecule.from_ase`."""
-    mol_ref = MOL[:100]
-    ase_mols = mol_ref.as_ase(**kwargs)
-    mol = MultiMolecule.from_ase(ase_mols)
+    ase_mols = mol.as_ase(**kwargs)
+    mol_new = MultiMolecule.from_ase(ase_mols)
 
     i = kwargs.get('mol_subset', np.s_[:])
     j = kwargs.get('atom_subset')
     if j is None:
-        np.testing.assert_allclose(mol, mol_ref[i])
+        np.testing.assert_allclose(mol_new, mol[i])
     else:
-        np.testing.assert_allclose(mol, mol_ref[i].loc[j])
+        np.testing.assert_allclose(mol_new, mol[i].loc[j])
+
+    assert mol.lattice is None or mol.lattice.ndim in {2, 3}
+    if mol.lattice is None:
+        assertion.is_(mol_new.lattice, mol.lattice)
+    elif mol.lattice.ndim == 2:
+        lat_ref = np.full(mol_new.lattice.shape, mol.lattice)
+        np.testing.assert_allclose(mol_new.lattice, lat_ref)
+    elif mol.lattice.ndim == 3:
+        np.testing.assert_allclose(mol_new.lattice, mol.lattice[i])
+
+
+@pytest.mark.parametrize(
+    "func",
+    [
+        lambda m: m.copy(),
+        lambda m: pickle.loads(pickle.dumps(m)),
+        lambda m: weakref.ref(m)(),
+        lambda m: copy.copy(m),
+        lambda m: copy.deepcopy(m),
+    ],
+    ids=["copy", "pickle", "weakref", "__copy__", "__deepcopy__"],
+)
+@pytest.mark.parametrize("mol", [MOL[::10], MOL_LATTICE_2D, MOL_LATTICE_3D])
+def test_copy(func: Callable[[MultiMolecule], MultiMolecule], mol: MultiMolecule) -> None:
+    mol_new = func(mol)
+    assert isinstance(mol_new, MultiMolecule)
+
+    assertion.eq(mol.dtype, mol_new.dtype)
+    np.testing.assert_allclose(mol, mol_new)
+    if mol.lattice is not None and mol_new.lattice is not None:
+        np.testing.assert_allclose(mol.lattice, mol_new.lattice)
+
+    assertion.eq(mol.atoms.keys(), mol_new.atoms.keys())
+    iter1 = ((k, mol.atoms[k], mol_new.atoms[k]) for k in mol.atoms)
+    for k, v1, v2 in iter1:
+        np.testing.assert_array_equal(v1, v2, err_msg=k)
+
+    assertion.eq(mol.atoms_alias.keys(), mol_new.atoms_alias.keys())
+    iter2 = ((k, mol.atoms_alias[k], mol_new.atoms_alias[k]) for k in mol.atoms_alias)
+    for k, (at1, v3), (at2, v4) in iter2:
+        assertion.eq(type(v3), type(v4), message=k)
+        assertion.eq(at1, at2, message=k)
+        if isinstance(v3, np.ndarray):
+            np.testing.assert_array_equal(v3, v4, err_msg=k)
+        else:
+            assertion.eq(v3, v4, message=k)