Improve explicit solvation (duartegroup#81)

* Timing test fix

* Refactor solvent

* Update changelog

* Add readable explicit solvation

* Add explicit solvent declaration

* Fix LGTM alerts

* Move constructor and add docstrings

* Add optimised solvent structures

* Improve get_solvent and enable to_explicit

* Add tests

* Add tests

* JS changes

.gitignore

@@ -14,6 +14,8 @@ coverage.xml
 
 tests/__pycache__/
 
+tests/*.xyz
+
 autode/conformers/__pycache__/
 
 autode/transition_states/__pycache__/
@@ -31,8 +33,6 @@ tests/htmlcov/
 
 *.pyc
 
-*.xyz
-
 htmlcov/
 
 tests/.autode_calculations

autode/reactions/reaction.py

@@ -72,7 +72,7 @@ def __init__(self,
             self._init_from_molecules(molecules=args)
 
         self.type = reaction_types.classify(self.reacs, self.prods)
-        self.solvent = get_solvent(solvent_name=solvent_name)
+        self.solvent = get_solvent(solvent_name, kind='implicit')
         self.temp = float(temp)
 
         self._check_solvent()

autode/solvent/__init__.py

@@ -1,4 +1,8 @@
-from autode.solvent.solvents import ImplicitSolvent
+from autode.solvent.solvents import Solvent, ImplicitSolvent, get_solvent
+from autode.solvent.explicit_solvent import ExplicitSolvent
 
 
-__all__ = ['ImplicitSolvent']
+__all__ = ['get_solvent',
+           'Solvent',
+           'ImplicitSolvent',
+           'ExplicitSolvent']

autode/solvent/explicit_solvent.py

@@ -1,95 +1,246 @@
-from copy import deepcopy
-from math import ceil
 import numpy as np
+from typing import Optional
+from scipy.spatial import distance_matrix
+from autode.geom import get_points_on_sphere, get_rot_mat_euler
 from autode.log import logger
+from autode.atoms import AtomCollection
+from autode.solvent.solvents import Solvent
 
 
-def add_solvent_molecules(species, n_qm_solvent_mols, n_solvent_mols):
-    """Add a specific number of solvent molecules around a solute"""
-    # Initialise a new random seed and make a copy of the species' atoms. RandomState is thread safe
-    rand = np.random.RandomState()
-
-    logger.info(f'Adding solvent molecules around {species.name}')
-
-    solvent_n_atoms = species.solvent_mol.n_atoms
-    total_n_solvent_atoms = n_solvent_mols * solvent_n_atoms
-
-    centre_species(species.solvent_mol)
-    solvent_coords = species.solvent_mol.coordinates
-    solvent_size = np.linalg.norm(np.max(solvent_coords, axis=0) - np.min(solvent_coords, axis=0))
-    solvent_size = 1 if solvent_size < 1 else solvent_size
-
-    centre_species(species)
-    solute_coords = species.coordinates
-    radius = np.linalg.norm(np.max(solute_coords, axis=0) - np.min(solute_coords, axis=0))
-    radius = 2 if radius < 2 else radius
-
-    solvent_area = (0.9*solvent_size) ** 2 * np.pi
-
-    i = 1
-    all_solvent_atoms = []
-    while len(all_solvent_atoms) <= total_n_solvent_atoms:
-        add_solvent_on_sphere(species, all_solvent_atoms, radius, solvent_area, i, rand)
-        i += 1
-
-    # TODO make this nicer?
-    # Only take closest solvent molecules
-    distances = []
-    for i in range(int(len(all_solvent_atoms)/solvent_n_atoms)):
-        solvent_mol_atoms = all_solvent_atoms[i*solvent_n_atoms:(i+1)*solvent_n_atoms]
-        solvent_mol_coords = [atom.coord for atom in solvent_mol_atoms]
-        distances.append(np.linalg.norm(np.average(solvent_mol_coords, axis=0)))
-
-    species.qm_solvent_atoms = []
-    species.mm_solvent_atoms = []
-    sorted_distances = sorted(distances)
-    for i in range(n_solvent_mols):
-        original_index = distances.index(sorted_distances[i])
-        solvent_mol_atoms = all_solvent_atoms[original_index*solvent_n_atoms:(original_index+1)*solvent_n_atoms]
-        if i < n_qm_solvent_mols:
-            species.qm_solvent_atoms += solvent_mol_atoms
-        else:
-            species.mm_solvent_atoms += solvent_mol_atoms
-
-    return None
-
-
-def centre_species(species):
-    """Translates a species so its centre is at (0,0,0)"""
-    species_coords = species.coordinates
-    species_centre = np.average(species_coords, axis=0)
-    species.translate(-species_centre)
-
-
-def add_solvent_on_sphere(species, solvent_atoms, radius, solvent_mol_area, radius_mult, rand):
-    """Packs solvent molecules semi-evenly on a sphere around the solvent molecule"""
-    rad_to_use = (radius * radius_mult * 0.8) + 0.4
-    fit_on_sphere = ceil((4 * np.pi * rad_to_use**2) / solvent_mol_area)
-    d = fit_on_sphere**(4/5)
-    m_theta = ceil(d/np.pi)
-    total_circum = 0
-    for m in range(0, m_theta):
-        total_circum += 2 * np.pi * np.sin(np.pi * (m+0.5)/m_theta)
-    for m in range(0, m_theta):
-        theta = np.pi * (m+0.5)/m_theta
-        circum = 2 * np.pi * np.sin(theta)
-        n_on_ring = int(round(circum * fit_on_sphere / total_circum))
-        for n in range(0, n_on_ring):
-            if m % 2 == 0:
-                phi = (2 * np.pi * n/n_on_ring) + 0.7*np.pi*(rand.rand()-0.5)/(n_on_ring)
-            else:
-                phi = (2 * np.pi * (n+0.5)/n_on_ring) + 0.7*np.pi*(rand.rand()-0.5)/(n_on_ring)
-            # Add a little bit of randomness to the positioning
-            rand_theta = theta + 0.35*np.pi*(rand.rand()-0.5)/(m_theta-1)
-            rand_add = 0.4*radius * (rand.rand()-0.5)
-            x = (rad_to_use + rand_add) * np.sin(rand_theta) * np.cos(phi)
-            y = (rad_to_use + rand_add) * np.sin(rand_theta) * np.sin(phi)
-            z = (rad_to_use + rand_add) * np.cos(rand_theta)
-            position = [x, y, z]
-            species.solvent_mol.rotate(axis=rand.uniform(-1.0, 1.0, 3), theta=2*np.pi*rand.rand())
-            for atom in species.solvent_mol.atoms:
-                new_atom = deepcopy(atom)
-                new_atom.translate(position)
-                solvent_atoms.append(new_atom)
-
-    return None
+class _RandomPointGenerator:
+    r"""
+    Generator for points (unit vectors) in solvent shells. e.g. where if x is
+    a solute molecule the vectors in the different shells::
+
+            ------
+                   `
+            ---      `
+                \     |
+            x    |    |
+
+    """
+
+    def __init__(self, random_state):
+        """
+        Point generator
+
+        Arguments:
+            random_state (numpy.random.mtrand.RandomState):
+        """
+        self.random_state = random_state
+        self._sphere_n = 1
+        self._points = []
+
+    def random_point(self) -> np.ndarray:
+        """
+        Generate a random point in a solvent shell. Will return points on
+        the surface of the solvent shell (self._sphere_n) and increment the
+        solvent shell when there are none left
+
+        Returns:
+            (np.ndarray): Point on the 3D sphere
+        """
+
+        if len(self._points) == 0:
+            # Surface area of the sphere scales r^2, so square solvent shell
+            self._points = get_points_on_sphere(n_points=self._sphere_n**2 * 10)
+            self._sphere_n += 1
+
+        idx = self.random_state.randint(0, len(self._points))
+        return self._points.pop(idx)
+
+
+class ExplicitSolvent(AtomCollection, Solvent):
+    """Explicit solvation """
+
+    def __init__(self,
+                 solvent: 'autode.species.species.Species',
+                 num:     int,
+                 solute:  Optional['autode.species.species.Species'] = None,
+                 **kwargs):
+        """
+        Explicit solvent. Initial construction attempts to generate a
+        reasonable distribution around the (unmodified) solute. Only supports
+        unicomponent uncharged solvents.
+
+        ----------------------------------------------------------------------
+        Arguments:
+
+            solvent (autode.species.species.Species): Solvent molecule (copied)
+
+            num (int): Number of solvent molecules to add
+
+
+        Keyword Arguments:
+
+            solute (autode.species.species.Species | None): Solute which this
+                   solvent surrounds. If None then no translation to the
+                   explicit solvent molecules will be applied
+
+            aliases (list(str)): List of aliases of this solvent
+        """
+        if num <= 0:
+            raise ValueError('Must solvate with at least a single solvent '
+                             f'molecule. Had {num}')
+
+        solvent_atoms = sum((solvent.atoms.copy() for _ in range(num)), None)
+        AtomCollection.__init__(self, atoms=solvent_atoms)
+        Solvent.__init__(self,
+                         name=solvent.name,
+                         smiles=None,
+                         aliases=kwargs.get('aliases', None))
+
+        self.solvent_n_atoms = solvent.n_atoms
+        # TODO: Something better than this hardcoded value
+        self.solvent_radius = solvent.radius.to('ang') + 2.0
+
+        if solute is not None:
+            self.randomise_around(solute)
+
+    def __eq__(self, other) -> bool:
+        """Equality between two explicit solvent environments"""
+
+        if isinstance(other, ExplicitSolvent) and self.n_atoms == other.n_atoms:
+            return all(o_at.label == at.label
+                       for o_at, at in zip(other.atoms, self.atoms))
+
+        return False
+
+    @property
+    def is_implicit(self) -> bool:
+        """Is this solvent implicit?
+
+        Returns:
+            (bool): False
+        """
+        return False
+
+    @property
+    def n_solvent_molecules(self) -> int:
+        """Number of solvent molecules comprising this explicit solvent
+        cluster
+
+        Returns:
+            (int): n
+        """
+        return self.n_atoms // self.solvent_n_atoms
+
+    def solvent_atom_idxs(self, i: int) -> np.ndarray:
+        """
+        Atom indexes of an particular solvent molecule
+
+        Returns:
+            (np.ndarray): Atom indexes
+        """
+        if i < 0 or i >= self.n_solvent_molecules:
+            raise ValueError(f'Cannot find the indexes for the {i}th solvent '
+                             f'only had {self.n_solvent_molecules}.')
+
+        first_idx = i * self.solvent_n_atoms
+        last_idx = first_idx + self.solvent_n_atoms
+
+        return np.array(range(first_idx, last_idx), dtype=int)
+
+    @staticmethod
+    def _too_close_to_solute(solvent_coords: np.ndarray,
+                             solute_coords:  np.ndarray,
+                             solute_radius:  float) -> bool:
+        """
+        Are a set of solvent coordinates too close to the solute? (for a
+        particular solute radius)
+
+        Arguments:
+            solvent_coords (np.ndarray):  Shape = (N, 3)
+
+            solute_coords (np.ndarray):  Shape = (M, 3)
+
+            solute_radius (float): Radius (Å)
+        """
+        min_dist = np.min(distance_matrix(solute_coords, solvent_coords))
+        return min_dist < solute_radius
+
+    def _too_close_to_solvent(self,
+                              coords:       np.ndarray,
+                              solvent_idxs: np.ndarray,
+                              max_idx:      int) -> bool:
+        """
+        Are a set of solvent coordinates too close to the solvent molecules
+        that have already been translated?
+
+        Arguments:
+            coords (np.ndarray):  Shape = (N, 3) Coordinates of all
+                                  the solvent molecules
+
+            solvent_idxs (np.ndarray): Integer array of atom indexes of a
+                                       particular solvent molecule
+
+            max_idx (int): Indexes up to which the repulsion should be
+                                calculated. NOT INCLUSIVE
+        """
+        if max_idx == 0:
+            return False
+
+        min_dist = np.min(distance_matrix(coords[solvent_idxs],
+                                          coords[:max_idx * self.solvent_n_atoms]))
+
+        return min_dist < self.solvent_radius
+
+    def randomise_around(self,
+                         solute: 'autode.species.species.Species') -> None:
+        r"""
+        Randomise the positions of the solvent molecules around the solute,
+        for example using a methane solute and water solvent::
+
+
+                               H2O
+                        H20
+                 H2o            H2O
+                                        H2O
+               H2O      CH4      H2O
+
+                     H2O      H2O
+
+
+        where the solvent molecules are roughly packed in shells around the
+        solute.
+
+        Arguments:
+            solute (autode.species.species.Species):
+        """
+        logger.info(f'Randomising {self.n_solvent_molecules} solvent '
+                    f'molecules around {solute}')
+
+        coords = self.coordinates
+
+        # ----------------- Properties of the solute molecule -----------------
+        m_radius = solute.radius.to('ang') + 1.0       # Assume some exterior H
+        m_origin = np.average(solute.coordinates, axis=0)
+        m_coords = solute.coordinates - m_origin
+        # ---------------------------------------------------------------------
+
+        rand = np.random.RandomState()
+        pg = _RandomPointGenerator(random_state=rand)
+
+        for i in range(self.n_solvent_molecules):
+
+            idxs = self.solvent_atom_idxs(i)
+            coords[idxs] -= np.average(coords[idxs], axis=0)  # -> origin
+
+            # Apply a random rotation to the solvent molecule
+            rand_rot_mat = get_rot_mat_euler(axis=rand.uniform(-1.0, 1.0, size=3),
+                                             theta=rand.uniform(-np.pi, np.pi))
+
+            coords[idxs] = np.dot(coords[idxs], rand_rot_mat.T)
+
+            # Select a random vector along which this solvent molecule is to be
+            # translated until there is not any close contacts
+            vec = 0.1 * pg.random_point()
+
+            while (self._too_close_to_solute(coords[idxs], m_coords, m_radius)
+                   or self._too_close_to_solvent(coords, idxs, i)):
+
+                coords[idxs] += vec
+
+        # Finally, translate to be centred around the solute's origin
+        self.coordinates = coords + m_origin
+        return None

autode/solvent/lib/1,1,1-trichloroethane.xyz

@@ -0,0 +1,10 @@
+8
+Generated by autodE on: 2021-10-13. E = -1457.581688 Ha
+C    -0.80105   0.07876  -0.26127
+C     0.62710  -0.06102   0.20243
+Cl    1.55876  -1.01879  -0.97368
+Cl    1.38067   1.54339   0.36475
+Cl    0.67057  -0.87901   1.78439
+H    -0.82393   0.62117  -1.21575
+H    -1.37371   0.63526   0.49245
+H    -1.23850  -0.91976  -0.39323

autode/solvent/lib/1,1,2-trichloroethane.xyz

@@ -0,0 +1,10 @@
+8
+Generated by autodE on: 2021-10-13. E = -1457.583711 Ha
+Cl    1.85690   1.10608   0.04843
+C     0.72828  -0.24441  -0.17148
+C    -0.67450   0.14602   0.23946
+Cl   -1.33504   1.40761  -0.81177
+Cl   -1.68683  -1.31513   0.20075
+H     0.73791  -0.55403  -1.22588
+H     1.06785  -1.08007   0.45448
+H    -0.69457   0.53402   1.26600