Additional NCFlex features and updates

matscipy/cauchy_born.py

@@ -8,7 +8,7 @@
 from itertools import permutations
 
 from scipy.stats.qmc import LatinHypercube, scale
-
+from matscipy.elasticity import full_3x3_to_Voigt_6_strain, Voigt_6_to_full_3x3_strain
 import ase
 
 
@@ -68,12 +68,12 @@ def predict_gradient(self, grad_phi):
         return grad_phi @ self.model
 
     def save(self):
-        """Saves the regression model to file: CB_model.npy"""
-        np.save('CB_model.npy', self.model)
+        """Saves the regression model to file: CB_model.txt"""
+        np.savetxt('CB_model.txt', self.model)
 
     def load(self):
-        """Loads the regression model from file: CB_model.npy"""
-        self.model = np.load('CB_model.npy')
+        """Loads the regression model from file: CB_model.txt"""
+        self.model = np.loadtxt('CB_model.txt')
 
 
 class CubicCauchyBorn:
@@ -129,7 +129,7 @@ def set_sublattices(self, atoms, A, read_from_atoms=False):
             whether or not to read the sublattices directly
             from the atoms object provided, rather than working them
             out by applying a small strain and measuring forces.
-            The user will have to set the sublattices themselves 
+            The user will have to set the sublattices themselves
             when building the atoms structure
         """
 
@@ -140,6 +140,9 @@ def set_sublattices(self, atoms, A, read_from_atoms=False):
 
         # generate U (right hand stretch tensor) to apply
         if read_from_atoms:
+            if ('lattice1mask' not in atoms.arrays) or (
+                    'lattice2mask' not in atoms.arrays):
+                raise KeyError('Lattice masks not found in atoms object')
             lattice1mask = atoms.arrays['lattice1mask']
             lattice2mask = atoms.arrays['lattice2mask']
         else:
@@ -183,6 +186,10 @@ def set_sublattices(self, atoms, A, read_from_atoms=False):
                 if all(element == lattice1mask[0] for element in lattice1mask):
                     lattice1mask = force_diff_lattice[:, 2] > 0
                     lattice2mask = force_diff_lattice[:, 2] < 0
+                    if all(element == lattice1mask[0]
+                           for element in lattice1mask):
+                        raise RuntimeError(
+                            'No forces detected in any direction when shift applied - cannot assign sublattices')
 
             # set new array in atoms
             try:  # make new arrays if they didn't already exist
@@ -209,10 +216,11 @@ def switch_sublattices(self, atoms):
         ----------
         atoms : ASE atoms object with defined sublattices
         """
-        tmpl1 = self.lattice1mask.copy()
-        tmpl2 = self.lattice2mask.copy()
-        self.lattice1mask = tmpl2.copy()
-        self.lattice2mask = tmpl1.copy()
+        self.lattice1mask, self.lattice2mask = self.lattice2mask.copy(), self.lattice1mask.copy()
+        # tmpl1 = self.lattice1mask.copy()
+        # tmpl2 = self.lattice2mask.copy()
+        # self.lattice1mask = tmpl2.copy()
+        # self.lattice2mask = tmpl1.copy()
         lattice1maskatoms = atoms.arrays['lattice1mask']
         lattice2maskatoms = atoms.arrays['lattice2mask']
         lattice1maskatoms[:] = self.lattice1mask
@@ -256,13 +264,7 @@ def f_gl(E_vec, calc, unitcell):
             """
             # green lagrange version of function
             # turn E into matrix
-            E = np.zeros([3, 3])
-            E[0, 0] = E_vec[0]
-            E[1, 1] = E_vec[1]
-            E[2, 2] = E_vec[2]
-            E[2, 1], E[1, 2] = E_vec[3], E_vec[3]
-            E[2, 0], E[0, 2] = E_vec[4], E_vec[4]
-            E[0, 1], E[1, 0] = E_vec[5], E_vec[5]
+            E = Voigt_6_to_full_3x3_strain(E_vec)
             # get U^2
             Usqr = 2 * E + np.eye(3)
             # square root matrix to get U
@@ -954,13 +956,7 @@ def eval_eps(eps_vec, grad_f, hess_f):
             return predictions
 
         # turn E into voigt vector
-        E_voigt = np.zeros([np.shape(E)[0], 6])
-        E_voigt[:, 0] = E[:, 0, 0]
-        E_voigt[:, 1] = E[:, 1, 1]
-        E_voigt[:, 2] = E[:, 2, 2]
-        E_voigt[:, 3] = E[:, 1, 2]
-        E_voigt[:, 4] = E[:, 0, 2]
-        E_voigt[:, 5] = E[:, 0, 1]
+        E_voigt = full_3x3_to_Voigt_6_strain(E)
 
         # return the shift vectors
         if method == 'taylor':
@@ -1160,10 +1156,10 @@ def check_for_refit(
             tol = tol * len(atoms[mask])
         else:
             tol = tol * len(atoms)
-        print('tol', tol)
+        # print('tol', tol)
 
         cb_err = self.get_cb_error(atoms, forces=forces, mask=mask)
-        print('cb_err', cb_err)
+        # print('cb_err', cb_err)
         if abs(cb_err) < tol:  # do not refit in this case
             return 0, cb_err / len(atoms)
 
@@ -1219,13 +1215,7 @@ def check_for_refit(
         # print('Es',E)
         # we get back the strain vectors in the lattice frame
         # turn E into voigt vectors
-        E_voigt = np.zeros([np.shape(E)[0], 6])
-        E_voigt[:, 0] = E[:, 0, 0]
-        E_voigt[:, 1] = E[:, 1, 1]
-        E_voigt[:, 2] = E[:, 2, 2]
-        E_voigt[:, 3] = E[:, 1, 2]
-        E_voigt[:, 4] = E[:, 0, 2]
-        E_voigt[:, 5] = E[:, 0, 1]
+        E_voigt = full_3x3_to_Voigt_6_strain(E)
 
         # pass set of Es to evaluate to refit_regression
         self.refit_regression(atoms, E_voigt)
@@ -1336,13 +1326,7 @@ def eval_shift(self, E_vec, unitcell):
         """
         # green lagrange version of function
         # turn E into matrix
-        E = np.zeros([3, 3])
-        E[0, 0] = E_vec[0]
-        E[1, 1] = E_vec[1]
-        E[2, 2] = E_vec[2]
-        E[2, 1], E[1, 2] = E_vec[3], E_vec[3]
-        E[2, 0], E[0, 2] = E_vec[4], E_vec[4]
-        E[0, 1], E[1, 0] = E_vec[5], E_vec[5]
+        E = Voigt_6_to_full_3x3_strain(E_vec)
         # get U^2
         Usqr = 2 * E + np.eye(3)
         # square root matrix
@@ -1407,7 +1391,7 @@ def basis_function_evaluation(self, E_vecs):
 
     def grad_basis_function_evaluation(self, E_vecs, dE_vecs):
         """Function that evaluates the derivative of the polynomial
-        basis functions used in the regression model on a set of 
+        basis functions used in the regression model on a set of
         strain vectors to build the gradient design matrix grad_phi.
 
         Parameters
@@ -1458,43 +1442,29 @@ def evaluate_shift_gradient_regression(self, E, dE):
         -------
         predictions : array_like
             prediction of gradient of the Cauchy-Born shift corrector model
-            at each of the given E points. 
+            at each of the given E points.
         """
 
         # turn E and dE into voigt vectors
-        E_voigt = np.zeros([np.shape(E)[0], 6])
-        E_voigt[:, 0] = E[:, 0, 0]
-        E_voigt[:, 1] = E[:, 1, 1]
-        E_voigt[:, 2] = E[:, 2, 2]
-        E_voigt[:, 3] = E[:, 1, 2]
-        E_voigt[:, 4] = E[:, 0, 2]
-        E_voigt[:, 5] = E[:, 0, 1]
-
-        dE_voigt = np.zeros([np.shape(dE)[0], 6])
-        dE_voigt[:, 0] = dE[:, 0, 0]
-        dE_voigt[:, 1] = dE[:, 1, 1]
-        dE_voigt[:, 2] = dE[:, 2, 2]
-        dE_voigt[:, 3] = dE[:, 1, 2]
-        dE_voigt[:, 4] = dE[:, 0, 2]
-        dE_voigt[:, 5] = dE[:, 0, 1]
+        E_voigt = full_3x3_to_Voigt_6_strain(E)
+        dE_voigt = full_3x3_to_Voigt_6_strain(dE)
 
         # Get the rotated versions of the strain and strain deriv tensors
-        epsx = E_voigt
-        epsy = self.permutation(E_voigt, 1)
-        epsz = self.permutation(E_voigt, 2)
-
-        depsx = dE_voigt
-        depsy = self.permutation(dE_voigt, 1)
-        depsz = self.permutation(dE_voigt, 2)
-        predictions = np.zeros([np.shape(epsx)[0], 3])
-
-        # Predict the shift gradients using the fitted regression model
-        predictions[:, 0] = self.RM.predict_gradient(
-            self.grad_basis_function_evaluation(epsx, depsx))
-        predictions[:, 1] = self.RM.predict_gradient(
-            self.grad_basis_function_evaluation(epsy, depsy))
-        predictions[:, 2] = self.RM.predict_gradient(
-            self.grad_basis_function_evaluation(epsz, depsz))
+        eps = [
+            E_voigt, self.permutation(
+                E_voigt, 1), self.permutation(
+                E_voigt, 2)]
+
+        deps = [
+            dE_voigt, self.permutation(
+                dE_voigt, 1), self.permutation(
+                dE_voigt, 2)]
+        predictions = np.zeros([np.shape(eps[0])[0], 3])
+
+        for i in range(3):
+            # Predict the shift gradients using the fitted regression model
+            predictions[:, i] = self.RM.predict_gradient(
+                self.grad_basis_function_evaluation(eps[i], deps[i]))
         return predictions
 
     def get_shift_gradients(
@@ -1510,7 +1480,7 @@ def get_shift_gradients(
         """Get the deformation gradient tensor field or Green-Lagrange strain
         tensor field for a system of atoms in the lab frame, depending on which
         function is provided. From these, get the gradients of the Cauchy-Born shift corrector field
-        field in the lattice frame. The supplied F_func or E_func needs to take 
+        field in the lattice frame. The supplied F_func or E_func needs to take
         parameter de, used for finite differences to evaluate the derivative of the
         given field
 
@@ -1532,7 +1502,7 @@ def get_shift_gradients(
             for the atoms in form of a numpy array shape [natoms,ndims,ndims],
             where F is specified in cartesian coordinates. The returned
             tensor field must be in the lab frame. Optional, but one of
-            F_func or E_func must be provided. The function must accept de as a 
+            F_func or E_func must be provided. The function must accept de as a
             keyword argument for calculating the gradient of the field by finite
             differences.
         E_func : Function
@@ -1544,8 +1514,8 @@ def get_shift_gradients(
             field for the atoms in form of a numpy array shape
             [natoms,ndims,ndims], where E is specified in cartesian
             coordinates, and the returned tensor field is in the lab frame.
-            Optional, but one of F_func or E_func must be provided. The 
-            function must accept de as a keyword argument for calculating 
+            Optional, but one of F_func or E_func must be provided. The
+            function must accept de as a keyword argument for calculating
             the gradient of the field by finite differences.
         coordinates : string
             Specifies the coordinate system that the function
@@ -1570,7 +1540,7 @@ def get_shift_gradients(
             coordinates=coordinates, de=de, *args, **kwargs)
 
         # print(E_higher,E_lower)
-        dE = (E_higher-E_lower)/(2*de)
+        dE = (E_higher - E_lower) / (2 * de)
         # print(dE)
         natoms = len(atoms)
         # get the cauchy born shifts unrotated
@@ -1584,17 +1554,17 @@ def get_shift_gradients(
             # dshift_2[i, :] = np.transpose(A) @ dshift_2_no_rr[i, :]
 
         # need to adjust gradients for different lattices
-        dshifts[self.lattice1mask] = -0.5*(dshifts[self.lattice1mask])
-        dshifts[self.lattice2mask] = 0.5*(dshifts[self.lattice2mask])
+        dshifts[self.lattice1mask] = -0.5 * (dshifts[self.lattice1mask])
+        dshifts[self.lattice2mask] = 0.5 * (dshifts[self.lattice2mask])
 
         return dshifts
 
     def save_regression_model(self):
-        """Saves the regression model to file: CB_model.npy"""
+        """Saves the regression model to file: CB_model.txt"""
         self.RM.save()
 
     def load_regression_model(self):
-        """Loads the regression model from file: CB_model.npy"""
+        """Loads the regression model from file: CB_model.txt"""
         self.RM = RegressionModel()
         self.RM.load()