Jon/feature/auto sweep

flare/env.py

@@ -2,7 +2,7 @@
 environment of an atom. :class:`AtomicEnvironment` objects are inputs to the
 2-, 3-, and 2+3-body kernels."""
 import numpy as np
-from math import sqrt
+from math import sqrt, ceil
 from numba import njit
 from flare.struc import Structure
 
@@ -21,12 +21,16 @@ class AtomicEnvironment:
     :type cutoffs: np.ndarray
     """
 
-    def __init__(self, structure: Structure, atom: int, cutoffs, sweep=1):
+    def __init__(self, structure: Structure, atom: int, cutoffs):
         self.structure = structure
         self.positions = structure.wrapped_positions
         self.cell = structure.cell
         self.species = structure.coded_species
-        self.sweep_array = np.arange(-sweep, sweep+1, 1)
+
+        # Set the sweep array based on the 2-body cutoff.
+        sweep_val = ceil(cutoffs[0] / structure.max_cutoff)
+        self.sweep_val = sweep_val
+        self.sweep_array = np.arange(-sweep_val, sweep_val + 1, 1)
 
         self.atom = atom
         self.ctype = structure.coded_species[atom]
@@ -107,7 +111,7 @@ def from_dict(dictionary):
             cutoffs = dictionary['cutoffs']
         else:
             cutoffs = []
-            for cutoff in ['cutoff_2','cutoff_3','cutoff_mb']:
+            for cutoff in ['cutoff_2', 'cutoff_3', 'cutoff_mb']:
                 if dictionary.get(cutoff):
                     cutoffs.append(dictionary[cutoff])
 
@@ -212,10 +216,12 @@ class to allow for njit acceleration with Numba.
 
 
 @njit
-def get_2_body_arrays_ind(positions, atom: int, cell, cutoff_2: float, species):
-    """Returns distances, coordinates, species of atoms, and indexes of neighbors
-    in the 2-body local environment. This method is implemented outside
-    the AtomicEnvironment class to allow for njit acceleration with Numba.
+def get_2_body_arrays_ind(positions, atom: int, cell, cutoff_2: float,
+                          species: np.ndarray):
+    """Returns distances, coordinates, species of atoms, and indexes of
+        neighbors in the 2-body local environment. This method is implemented
+        outside the AtomicEnvironment class to allow for njit acceleration
+        with Numba.
 
     :param positions: Positions of atoms in the structure.
     :type positions: np.ndarray
@@ -377,10 +383,11 @@ def get_3_body_arrays(bond_array_2, bond_positions_2, cutoff_3: float):
 @njit
 def get_m_body_arrays(positions, atom: int, cell, cutoff_mb: float, species,
                       sweep: np.ndarray):
-    """Returns distances, and species of atoms in the many-body
-    local environment, and returns distances and numbers of neighbours for atoms in the one
-    many-body local environment. This method is implemented outside the AtomicEnvironment
-    class to allow for njit acceleration with Numba.
+    """Returns distances, and species of atoms in the many-body local
+        environment, and returns distances and numbers of neighbours for atoms
+        in the one many-body local environment. This method is implemented
+        outside the AtomicEnvironment class to allow for njit acceleration
+        with Numba.
 
     :param positions: Positions of atoms in the structure.
     :type positions: np.ndarray
@@ -393,7 +400,8 @@ class to allow for njit acceleration with Numba.
     :type cutoff_mb: float
     :param species: Numpy array of species represented by their atomic numbers.
     :type species: np.ndarray
-    :param indexes: Boolean indicating whether indexes of neighbours are returned
+    :param indexes: Boolean indicating whether indexes of neighbours are
+        returned
     :type indexes: boolean
     :return: Tuple of arrays describing pairs of atoms in the 2-body local
      environment.
@@ -412,11 +420,13 @@ class to allow for njit acceleration with Numba.
 
      num_neighs_mb: number of neighbours of each atom in the local environment
 
-     etypes_mb_array: species of neighbours of each atom in the local environment
+     etypes_mb_array: species of neighbours of each atom in the local
+        environment
 
     :rtype: np.ndarray, np.ndarray, np.ndarray, np.ndarray
     """
-    # TODO: this can be probably improved using stored arrays, redundant calls to get_2_body_arrays
+    # TODO: this can be probably improved using stored arrays, redundant calls
+    #  to get_2_body_arrays
     # Get distances, positions, species and indexes of neighbouring atoms
     bond_array_mb, __, etypes, bond_inds = get_2_body_arrays_ind(
         positions, atom, cell, cutoff_mb, species)
@@ -443,7 +453,8 @@ class to allow for njit acceleration with Numba.
         neigh_dists_mb[i, :num_neighs_mb[i]] = neighbouring_dists[i]
         etypes_mb_array[i, :num_neighs_mb[i]] = neighbouring_etypes[i]
 
-    return bond_array_mb, neigh_dists_mb, num_neighs_mb, etypes_mb_array, etypes
+    return bond_array_mb, neigh_dists_mb, num_neighs_mb, etypes_mb_array, \
+        etypes
 
 
 if __name__ == '__main__':

flare/gp.py

@@ -344,8 +344,7 @@ def check_instantiation(self):
             self.hyps_mask = None
 
     def update_db(self, struc: Structure, forces: List,
-                  custom_range: List[int] = (), energy: float = None,
-                  sweep: int = 1):
+                  custom_range: List[int] = (), energy: float = None):
         """Given a structure and forces, add local environments from the
         structure to the training set of the GP. If energy is given, add the
         entire structure to the training set.
@@ -370,7 +369,7 @@ def update_db(self, struc: Structure, forces: List,
         if forces is not None:
             for atom in update_indices:
                 env_curr = \
-                    AtomicEnvironment(struc, atom, self.cutoffs, sweep=sweep)
+                    AtomicEnvironment(struc, atom, self.cutoffs)
                 forces_curr = np.array(forces[atom])
 
                 self.training_data.append(env_curr)
@@ -386,7 +385,7 @@ def update_db(self, struc: Structure, forces: List,
             structure_list = []  # Populate with all environments of the struc
             for atom in range(noa):
                 env_curr = \
-                    AtomicEnvironment(struc, atom, self.cutoffs, sweep=sweep)
+                    AtomicEnvironment(struc, atom, self.cutoffs)
                 structure_list.append(env_curr)
 
             self.energy_labels.append(energy)

flare/struc.py

@@ -60,6 +60,9 @@ def __init__(self, cell: 'ndarray', species: Union[List[str], List[int]],
         self.vec2 = self.cell[1, :]
         self.vec3 = self.cell[2, :]
 
+        # Compute the max cutoff for sweep = 1.
+        self.max_cutoff = self.get_max_cutoff()
+
         # get cell matrices for wrapping coordinates
         self.cell_transpose = self.cell.transpose()
         self.cell_transpose_inverse = np.linalg.inv(self.cell_transpose)
@@ -127,9 +130,49 @@ def get_cell_dot(self):
 
         return cell_dot
 
+    def get_max_cutoff(self) -> float:
+        """Compute the maximum cutoff compatible with a 3x3x3 supercell of the
+            structure.
+
+        Returns:
+            float: maximum cutoff
+        """
+        # Retrieve the lattice vectors.
+        a_vec = self.cell[0]
+        b_vec = self.cell[1]
+        c_vec = self.cell[2]
+
+        # Compute dot products and norms of lattice vectors.
+        a_dot_b = np.dot(a_vec, b_vec)
+        a_dot_c = np.dot(a_vec, c_vec)
+        b_dot_c = np.dot(b_vec, c_vec)
+
+        a_norm = np.linalg.norm(a_vec)
+        b_norm = np.linalg.norm(b_vec)
+        c_norm = np.linalg.norm(c_vec)
+
+        # Compute the six independent altitudes of the cell faces.
+        # The smallest is the maximum atomic environment cutoff that can be
+        # used with sweep=1.
+        max_candidates = np.zeros(6)
+        max_candidates[0] = \
+            a_norm * np.sqrt(1 - (a_dot_b / (a_norm * b_norm))**2)
+        max_candidates[1] = \
+            b_norm * np.sqrt(1 - (a_dot_b / (a_norm * b_norm))**2)
+        max_candidates[2] = \
+            a_norm * np.sqrt(1 - (a_dot_c / (a_norm * c_norm))**2)
+        max_candidates[3] = \
+            c_norm * np.sqrt(1 - (a_dot_c / (a_norm * c_norm))**2)
+        max_candidates[4] = \
+            b_norm * np.sqrt(1 - (b_dot_c / (b_norm * c_norm))**2)
+        max_candidates[5] = \
+            c_norm * np.sqrt(1 - (b_dot_c / (b_norm * c_norm))**2)
+
+        return np.min(max_candidates)
+
     @staticmethod
     def raw_to_relative(positions: 'ndarray', cell_transpose: 'ndarray',
-                        cell_dot_inverse: 'ndarray')-> 'ndarray':
+                        cell_dot_inverse: 'ndarray') -> 'ndarray':
         """Convert Cartesian coordinates to relative (fractional) coordinates,
         expressed in terms of the cell vectors set in self.cell.
 
@@ -153,7 +196,7 @@ def raw_to_relative(positions: 'ndarray', cell_transpose: 'ndarray',
     @staticmethod
     def relative_to_raw(relative_positions: 'ndarray',
                         cell_transpose_inverse: 'ndarray',
-                        cell_dot: 'ndarray')-> 'ndarray':
+                        cell_dot: 'ndarray') -> 'ndarray':
         """Convert fractional coordinates to raw (Cartesian) coordinates.
 
         :param relative_positions: fractional coordinates.
@@ -169,16 +212,16 @@ def relative_to_raw(relative_positions: 'ndarray',
         return np.matmul(np.matmul(relative_positions, cell_dot),
                          cell_transpose_inverse)
 
-    def wrap_positions(self, in_place: bool = True)-> 'ndarray':
+    def wrap_positions(self, in_place: bool = True) -> 'ndarray':
         """
         Convenience function which folds atoms outside of the unit cell back
         into the unit cell. in_place flag controls if the wrapped positions
         are set in the class.
 
-        :param in_place: If true, set the current structure 
-		positions to be the wrapped positions.
+        :param in_place: If true, set the current structure positions to be
+            the wrapped positions.
         :return: Cartesian coordinates of positions all in unit cell
-	:rtype: np.ndarray
+        :rtype: np.ndarray
         """
         rel_pos = \
             self.raw_to_relative(self.positions, self.cell_transpose,
@@ -262,7 +305,7 @@ def as_str(self) -> str:
         return dumps(self.as_dict(), cls=NumpyEncoder)
 
     @staticmethod
-    def from_dict(dictionary: dict)-> 'flare.struc.Structure':
+    def from_dict(dictionary: dict) -> 'flare.struc.Structure':
         """
         Assembles a Structure object from a dictionary parameterizing one.
 
@@ -283,7 +326,7 @@ def from_dict(dictionary: dict)-> 'flare.struc.Structure':
         return struc
 
     @staticmethod
-    def from_ase_atoms(atoms: 'ase.Atoms')-> 'flare.struc.Structure':
+    def from_ase_atoms(atoms: 'ase.Atoms') -> 'flare.struc.Structure':
         """
         From an ASE Atoms object, return a FLARE structure
 
@@ -296,12 +339,10 @@ def from_ase_atoms(atoms: 'ase.Atoms')-> 'flare.struc.Structure':
                           species=atoms.get_chemical_symbols())
         return struc
 
-    def to_ase_atoms(self)-> 'ase.Atoms':
+    def to_ase_atoms(self) -> 'ase.Atoms':
         from ase import Atoms
-        return Atoms(self.species_labels,
-                     positions=self.positions,
-                     cell=self.cell, 
-                     pbc=True)
+        return Atoms(self.species_labels, positions=self.positions,
+                     cell=self.cell, pbc=True)
 
     def to_pmg_structure(self):
         """
@@ -328,7 +369,7 @@ def to_pmg_structure(self):
                                   )
 
     @staticmethod
-    def from_pmg_structure(structure: 'pymatgen Structure')-> \
+    def from_pmg_structure(structure: 'pymatgen Structure') -> \
             'flare Structure':
         """
         Returns Pymatgen structure as FLARE structure.
@@ -357,11 +398,9 @@ def from_pmg_structure(structure: 'pymatgen Structure')-> \
 
         return new_struc
 
-    def to_xyz(self, extended_xyz: bool = True,
-                     print_stds: bool = False,
-                     print_forces : bool = False,
-                     print_max_stds: bool = False,
-                     write_file: str = '')->str:
+    def to_xyz(self, extended_xyz: bool = True, print_stds: bool = False,
+               print_forces: bool = False, print_max_stds: bool = False,
+               write_file: str = '') -> str:
         """
         Convenience function which turns a structure into an extended .xyz
         file; useful for further input into visualization programs like VESTA
@@ -422,7 +461,8 @@ def to_xyz(self, extended_xyz: bool = True,
 
     @staticmethod
     def from_file(file_name, format='') -> Union['flare.struc.Structure',
-                                                List['flare.struc.Structure']]:
+                                                 List['flare.struc.Structure']
+                                                 ]:
         """
         Load a FLARE structure from a file or a series of FLARE structures
         :param file_name:
@@ -433,7 +473,7 @@ def from_file(file_name, format='') -> Union['flare.struc.Structure',
         try:
             with open(file_name, 'r') as _:
                 pass
-        except:
+        except FileNotFoundError:
             raise FileNotFoundError
 
         if 'xyz' in file_name or 'xyz' in format.lower():
@@ -456,16 +496,16 @@ def from_file(file_name, format='') -> Union['flare.struc.Structure',
             else:
                 return structures
 
-        is_poscar = 'POSCAR' in file_name or 'CONTCAR' in file_name or 'vasp' in format.lower()
+        is_poscar = 'POSCAR' in file_name or 'CONTCAR' in file_name \
+            or 'vasp' in format.lower()
         if is_poscar and _pmg_present:
             pmg_structure = pmgvaspio.Poscar.from_file(file_name).structure
             return Structure.from_pmg_structure(pmg_structure)
         elif is_poscar and not _pmg_present:
-            raise ImportError("Pymatgen not imported; " \
+            raise ImportError("Pymatgen not imported; "
                               "functionality requires pymatgen.")
 
 
-
 def get_unique_species(species: List[Any]) -> (List, List[int]):
     """
     Returns a list of the unique species passed in, and a list of
@@ -485,5 +525,3 @@ def get_unique_species(species: List[Any]) -> (List, List[int]):
     coded_species = np.array(coded_species)
 
     return unique_species, coded_species
-
-

tests/test_env.py

@@ -11,13 +11,13 @@ def test_species_count(cutoff):
     species = [1, 2, 3]
     positions = np.array([[0, 0, 0], [0.5, 0.5, 0.5], [0.1, 0.1, 0.1]])
     struc_test = Structure(cell, species, positions)
-    env_test = AtomicEnvironment(structure=struc_test,
-                                 atom=0,
+    env_test = AtomicEnvironment(structure=struc_test, atom=0,
                                  cutoffs=np.array([1, 1]))
     assert (len(struc_test.positions) == len(struc_test.coded_species))
     assert (len(env_test.bond_array_2) == len(env_test.etypes))
     assert (isinstance(env_test.etypes[0], np.int8))
 
+
 @pytest.mark.parametrize('cutoff', cutoff_list)
 def test_env_methods(cutoff):
     cell = np.eye(3)
@@ -39,3 +39,43 @@ def test_env_methods(cutoff):
     assert np.array_equal(remade_env.bond_array_2, env_test.bond_array_2)
     assert np.array_equal(remade_env.bond_array_3, env_test.bond_array_3)
     assert np.array_equal(remade_env.bond_array_mb, env_test.bond_array_mb)
+
+
+def test_auto_sweep():
+    """Test that the number of neighbors inside the local environment is
+        correctly computed."""
+
+    # Make an arbitrary non-cubic structure.
+    cell = np.array([[1.3, 0.5, 0.8],
+                     [-1.2, 1, 0.73],
+                     [-0.8, 0.1, 0.9]])
+    positions = np.array([[1.2, 0.7, 2.3],
+                          [3.1, 2.5, 8.9],
+                          [-1.8, -5.8, 3.0],
+                          [0.2, 1.1, 2.1],
+                          [3.2, 1.1, 3.3]])
+    species = np.array([1, 2, 3, 4, 5])
+    arbitrary_structure = Structure(cell, species, positions)
+
+    # Construct an environment.
+    cutoffs = np.array([4., 3.])
+    arbitrary_environment = \
+        AtomicEnvironment(arbitrary_structure, 0, cutoffs)
+
+    # Count the neighbors.
+    n_neighbors_1 = len(arbitrary_environment.etypes)
+
+    # Reduce the sweep value, and check that neighbors are missing.
+    sweep_val = arbitrary_environment.sweep_val
+    arbitrary_environment.sweep_array = \
+        np.arange(-sweep_val + 1, sweep_val, 1)
+    arbitrary_environment.compute_env()
+    n_neighbors_2 = len(arbitrary_environment.etypes)
+    assert(n_neighbors_1 > n_neighbors_2)
+
+    # Increase the sweep value, and check that the count is the same.
+    arbitrary_environment.sweep_array = \
+        np.arange(-sweep_val - 1, sweep_val + 2, 1)
+    arbitrary_environment.compute_env()
+    n_neighbors_3 = len(arbitrary_environment.etypes)
+    assert(n_neighbors_1 == n_neighbors_3)