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 max cutoff.
+        sweep_val = ceil(np.max(cutoffs) / 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

@@ -8,7 +8,8 @@
 used to train ML models.
 """
 import numpy as np
-from flare.util import element_to_Z, Z_to_element, NumpyEncoder
+from flare.util import element_to_Z, Z_to_element, NumpyEncoder, \
+    get_max_cutoff
 from json import dumps, loads
 
 from typing import List, Union, Any
@@ -60,6 +61,10 @@ 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 compatible with a 3x3x3 supercell of the
+        # structure.
+        self.max_cutoff = get_max_cutoff(self.cell)
+
         # get cell matrices for wrapping coordinates
         self.cell_transpose = self.cell.transpose()
         self.cell_transpose_inverse = np.linalg.inv(self.cell_transpose)
@@ -129,7 +134,7 @@ def get_cell_dot(self):
 
     @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 +158,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 +174,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 +267,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 +288,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 +301,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 +331,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 +360,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 +423,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 +435,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 +458,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 +487,3 @@ def get_unique_species(species: List[Any]) -> (List, List[int]):
     coded_species = np.array(coded_species)
 
     return unique_species, coded_species
-
-