def test_get_strain_state_dict(self):
     strain_inds = [(0,), (1,), (2,), (1, 3), (1, 2, 3)]
     vecs = {}
     strain_states = []
     for strain_ind in strain_inds:
         ss = np.zeros(6)
         np.put(ss, strain_ind, 1)
         strain_states.append(tuple(ss))
         vec = np.zeros((4, 6))
         rand_values = np.random.uniform(0.1, 1, 4)
         for i in strain_ind:
             vec[:, i] = rand_values
         vecs[strain_ind] = vec
     all_strains = [Strain.from_voigt(v).zeroed() for vec in vecs.values()
                    for v in vec]
     random.shuffle(all_strains)
     all_stresses = [Stress.from_voigt(np.random.random(6)).zeroed()
                     for s in all_strains]
     strain_dict = {k.tostring():v for k,v in zip(all_strains, all_stresses)}
     ss_dict = get_strain_state_dict(all_strains, all_stresses, add_eq=False)
     # Check length of ss_dict
     self.assertEqual(len(strain_inds), len(ss_dict))
     # Check sets of strain states are correct
     self.assertEqual(set(strain_states), set(ss_dict.keys()))
     for strain_state, data in ss_dict.items():
         # Check correspondence of strains/stresses
         for strain, stress in zip(data["strains"], data["stresses"]):
             self.assertArrayAlmostEqual(Stress.from_voigt(stress), 
                                         strain_dict[Strain.from_voigt(strain).tostring()])

 def test_find_eq_stress(self):
     random_strains = [Strain.from_voigt(s) for s in np.random.uniform(0.1, 1, (20, 6))]
     random_stresses = [Strain.from_voigt(s) for s in np.random.uniform(0.1, 1, (20, 6))]
     with warnings.catch_warnings(record=True):
         no_eq = find_eq_stress(random_strains, random_stresses)
         self.assertArrayAlmostEqual(no_eq, np.zeros((3,3)))
     random_strains[12] = Strain.from_voigt(np.zeros(6))
     eq_stress = find_eq_stress(random_strains, random_stresses)
     self.assertArrayAlmostEqual(random_stresses[12], eq_stress)

    def setUp(self):
        self.norm_str = Strain.from_deformation([[1.02, 0, 0], [0, 1, 0], [0, 0, 1]])
        self.ind_str = Strain.from_deformation([[1, 0.02, 0], [0, 1, 0], [0, 0, 1]])

        self.non_ind_str = Strain.from_deformation([[1, 0.02, 0.02], [0, 1, 0], [0, 0, 1]])

        with warnings.catch_warnings(record=True):
            warnings.simplefilter("always")
            self.no_dfm = Strain([[0.0, 0.01, 0.0], [0.01, 0.0002, 0.0], [0.0, 0.0, 0.0]])

 def test_new(self):
     test_strain = Strain([[0., 0.01, 0.],
                           [0.01, 0.0002, 0.],
                           [0., 0., 0.]])
     self.assertArrayAlmostEqual(
         test_strain,
         test_strain.get_deformation_matrix().green_lagrange_strain)
     self.assertRaises(ValueError, Strain, [[0.1, 0.1, 0],
                                            [0, 0, 0],
                                            [0, 0, 0]])

 def test_from_index_amount(self):
     # From voigt index
     test = Strain.from_index_amount(2, 0.01)
     should_be = np.zeros((3, 3))
     should_be[2, 2] = 0.01
     self.assertArrayAlmostEqual(test, should_be)
     # from full-tensor index
     test = Strain.from_index_amount((1, 2), 0.01)
     should_be = np.zeros((3, 3))
     should_be[1, 2] = should_be[2, 1] = 0.01
     self.assertArrayAlmostEqual(test, should_be)

 def test_get_compliance_expansion(self):
     ce_exp = self.exp_cu.get_compliance_expansion()
     et_comp = ElasticTensorExpansion(ce_exp)
     strain_orig = Strain.from_voigt([0.01, 0, 0, 0, 0, 0])
     stress = self.exp_cu.calculate_stress(strain_orig)
     strain_revert = et_comp.calculate_stress(stress)
     self.assertArrayAlmostEqual(strain_orig, strain_revert, decimal=4)

 def test_get_strain_from_stress(self):
     strain = Strain.from_voigt([0.05, 0, 0, 0, 0, 0])
     stress3 = self.exp_cu.calculate_stress(strain)
     strain_revert3 = self.exp_cu.get_strain_from_stress(stress3)
     self.assertArrayAlmostEqual(strain, strain_revert3, decimal=2)
     # fourth order
     stress4 = self.exp_cu_4.calculate_stress(strain)
     strain_revert4 = self.exp_cu_4.get_strain_from_stress(stress4)
     self.assertArrayAlmostEqual(strain, strain_revert4, decimal=2)

 def test_find_eq_stress(self):
     test_strains = deepcopy(self.strains)
     test_stresses = deepcopy(self.pk_stresses)
     with warnings.catch_warnings(record=True):
         no_eq = find_eq_stress(test_strains, test_stresses)
         self.assertArrayAlmostEqual(no_eq, np.zeros((3,3)))
     test_strains[3] = Strain.from_voigt(np.zeros(6))
     eq_stress = find_eq_stress(test_strains, test_stresses)
     self.assertArrayAlmostEqual(test_stresses[3], eq_stress)

 def test_from_strain_stress_list(self):
     strain_list = [Strain.from_deformation(def_matrix)
                    for def_matrix in self.def_stress_dict['deformations']]
     stress_list = [stress for stress in self.def_stress_dict['stresses']]
     with warnings.catch_warnings(record = True):
         et_fl = -0.1*ElasticTensor.from_strain_stress_list(strain_list, stress_list)
         self.assertArrayAlmostEqual(et_fl.round(2),
                                     [[59.29, 24.36, 22.46, 0, 0, 0],
                                      [28.06, 56.91, 22.46, 0, 0, 0],
                                      [28.06, 25.98, 54.67, 0, 0, 0],
                                      [0, 0, 0, 26.35, 0, 0],
                                      [0, 0, 0, 0, 26.35, 0],
                                      [0, 0, 0, 0, 0, 26.35]])

    def calculate_stress(self, strain):
        """
        Calculate's a given elastic tensor's contribution to the
        stress using Einstein summation

        Args:
            strain (3x3 array-like): matrix corresponding to strain
        """
        strain = np.array(strain)
        if strain.shape == (6,):
            strain = Strain.from_voigt(strain)
        assert strain.shape == (3, 3), "Strain must be 3x3 or voigt-notation"
        stress_matrix = self.einsum_sequence([strain]*(self.order - 1)) \
                / factorial(self.order - 1)
        return Stress(stress_matrix)

    def run_task(self, fw_spec):
        v, _ = get_vasprun_outcar(self.get("calc_dir", "."), parse_dos=False, parse_eigen=False)
        stress = v.ionic_steps[-1]['stress']
        defo = self['deformation']
        d_ind = np.nonzero(defo - np.eye(3))
        delta = Decimal((defo - np.eye(3))[d_ind][0])
        # Shorthand is d_X_V, X is voigt index, V is value
        dtype = "_".join(["d", str(reverse_voigt_map[d_ind][0]),
                          "{:.0e}".format(delta)])
        strain = Strain.from_deformation(defo)
        defo_dict = {'deformation_matrix': defo,
                     'strain': strain.tolist(),
                     'stress': stress}

        return FWAction(mod_spec=[{'_set': {
            'deformation_tasks->{}'.format(dtype): defo_dict}}])

    def calculate_stress(self, strain):
        """
        Calculate's a given elastic tensor's contribution to the
        stress using Einstein summation

        Args:
            strain (3x3 array-like): matrix corresponding to strain
        """
        strain = np.array(strain)
        if strain.shape == (6,):
            strain = Strain.from_voigt(strain)
        assert strain.shape == (3, 3), "Strain must be 3x3 or voigt-notation"
        lc = string.ascii_lowercase[:self.rank-2]
        lc_pairs = map(''.join, zip(*[iter(lc)]*2))
        einsum_string = "ij" + lc + ',' + ','.join(lc_pairs) + "->ij"
        einsum_args = [self] + [strain] * (self.order - 1)
        stress_matrix = np.einsum(einsum_string, *einsum_args) \
            / factorial(self.order - 1)
        return Stress(stress_matrix)

    def test_energy_density(self):

        film_elac = ElasticTensor.from_voigt([
            [324.32,  187.3,   170.92,    0.,      0.,      0.],
            [187.3,   324.32,  170.92,    0.,      0.,      0.],
            [170.92,  170.92,  408.41,    0.,      0.,      0.],
            [0.,      0.,      0.,    150.73,    0.,      0.],
            [0.,      0.,      0.,      0.,    150.73,    0.],
            [0.,      0.,      0.,      0.,      0.,    238.74]])

        dfm = Deformation([[ -9.86004855e-01,2.27539582e-01,-4.64426035e-17],
                           [ -2.47802121e-01,-9.91208483e-01,-7.58675185e-17],
                           [ -6.12323400e-17,-6.12323400e-17,1.00000000e+00]])

        self.assertAlmostEqual(film_elac.energy_density(dfm.green_lagrange_strain),
            0.00125664672793)

        film_elac.energy_density(Strain.from_deformation([[ 0.99774738,  0.11520994, -0.        ],
                                        [-0.11520994,  0.99774738,  0.        ],
                                        [-0.,         -0.,          1.,        ]]))

class StrainTest(PymatgenTest):
    def setUp(self):
        self.norm_str = Strain.from_deformation([[1.02, 0, 0],
                                                 [0, 1, 0],
                                                 [0, 0, 1]])
        self.ind_str = Strain.from_deformation([[1, 0.02, 0],
                                                [0, 1, 0],
                                                [0, 0, 1]])

        self.non_ind_str = Strain.from_deformation([[1, 0.02, 0.02],
                                                    [0, 1, 0],
                                                    [0, 0, 1]])

        with warnings.catch_warnings(record=True):
            warnings.simplefilter("always")
            self.no_dfm = Strain([[0., 0.01, 0.],
                                  [0.01, 0.0002, 0.],
                                  [0., 0., 0.]])

    def test_new(self):
        test_strain = Strain([[0., 0.01, 0.],
                              [0.01, 0.0002, 0.],
                              [0., 0., 0.]])
        self.assertArrayAlmostEqual(
            test_strain,
            test_strain.get_deformation_matrix().green_lagrange_strain)
        self.assertRaises(ValueError, Strain, [[0.1, 0.1, 0],
                                               [0, 0, 0],
                                               [0, 0, 0]])

    def test_from_deformation(self):
        self.assertArrayAlmostEqual(self.norm_str,
                                    [[0.0202, 0, 0],
                                     [0, 0, 0],
                                     [0, 0, 0]])
        self.assertArrayAlmostEqual(self.ind_str,
                                    [[0., 0.01, 0.],
                                     [0.01, 0.0002, 0.],
                                     [0., 0., 0.]])
        self.assertArrayAlmostEqual(self.non_ind_str,
                                    [[0., 0.01, 0.01],
                                     [0.01, 0.0002, 0.0002],
                                     [0.01, 0.0002, 0.0002]])

    def test_from_index_amount(self):
        # From voigt index
        test = Strain.from_index_amount(2, 0.01)
        should_be = np.zeros((3, 3))
        should_be[2, 2] = 0.01
        self.assertArrayAlmostEqual(test, should_be)
        # from full-tensor index
        test = Strain.from_index_amount((1, 2), 0.01)
        should_be = np.zeros((3, 3))
        should_be[1, 2] = should_be[2, 1] = 0.01
        self.assertArrayAlmostEqual(test, should_be)

    def test_properties(self):
        # deformation matrix
        self.assertArrayAlmostEqual(self.ind_str.get_deformation_matrix(),
                                    [[1, 0.02, 0],
                                     [0, 1, 0],
                                     [0, 0, 1]])
        symm_dfm = Strain(self.no_dfm).get_deformation_matrix(shape="symmetric")
        self.assertArrayAlmostEqual(symm_dfm, [[0.99995,0.0099995, 0],
                                               [0.0099995,1.00015, 0],
                                               [0, 0, 1]])
        self.assertArrayAlmostEqual(self.no_dfm.get_deformation_matrix(),
                                    [[1, 0.02, 0],
                                     [0, 1, 0],
                                     [0, 0, 1]])

        # voigt
        self.assertArrayAlmostEqual(self.non_ind_str.voigt,
                                    [0, 0.0002, 0.0002, 0.0004, 0.02, 0.02])

    def test_convert_strain_to_deformation(self):
        strain = Tensor(np.random.random((3, 3))).symmetrized
        while not (np.linalg.eigvals(strain) > 0).all():
            strain = Tensor(np.random.random((3, 3))).symmetrized
        upper = convert_strain_to_deformation(strain, shape="upper")
        symm = convert_strain_to_deformation(strain, shape="symmetric")
        self.assertArrayAlmostEqual(np.triu(upper), upper)
        self.assertTrue(Tensor(symm).is_symmetric())
        for defo in upper, symm:
            self.assertArrayAlmostEqual(defo.green_lagrange_strain, strain)

def get_wf_elastic_constant(structure, strain_states=None, stencils=None,
                            db_file=None,
                            conventional=False, order=2, vasp_input_set=None,
                            analysis=True,
                            sym_reduce=False, tag='elastic',
                            copy_vasp_outputs=False, **kwargs):
    """
    Returns a workflow to calculate elastic constants.

    Firework 1 : write vasp input set for structural relaxation,
                 run vasp,
                 pass run location,
                 database insertion.

    Firework 2 - number of total deformations: Static runs on the deformed structures

    last Firework : Analyze Stress/Strain data and fit the elastic tensor

    Args:
        structure (Structure): input structure to be optimized and run.
        strain_states (list of Voigt-notation strains): list of ratios of nonzero elements
            of Voigt-notation strain, e. g. [(1, 0, 0, 0, 0, 0), (0, 1, 0, 0, 0, 0), etc.].
        stencils (list of floats, or list of list of floats): values of strain to multiply
            by for each strain state, i. e. stencil for the perturbation along the strain
            state direction, e. g. [-0.01, -0.005, 0.005, 0.01].  If a list of lists,
            stencils must correspond to each strain state provided.
        db_file (str): path to file containing the database credentials.
        conventional (bool): flag to convert input structure to conventional structure,
            defaults to False.
        order (int): order of the tensor expansion to be determined.  Defaults to 2 and
            currently supports up to 3.
        vasp_input_set (VaspInputSet): vasp input set to be used.  Defaults to static
            set with ionic relaxation parameters set.  Take care if replacing this,
            default ensures that ionic relaxation is done and that stress is calculated
            for each vasp run.
        analysis (bool): flag to indicate whether analysis task should be added
            and stresses and strains passed to that task
        sym_reduce (bool): Whether or not to apply symmetry reductions
        tag (str):
        copy_vasp_outputs (bool): whether or not to copy previous vasp outputs.
        kwargs (keyword arguments): additional kwargs to be passed to get_wf_deformations

    Returns:
        Workflow
    """
    # Convert to conventional if specified
    if conventional:
        structure = SpacegroupAnalyzer(
            structure).get_conventional_standard_structure()

    uis_elastic = {"IBRION": 2, "NSW": 99, "ISIF": 2, "ISTART": 1,
                   "PREC": "High"}
    vis = vasp_input_set or MPStaticSet(structure,
                                        user_incar_settings=uis_elastic)
    strains = []
    if strain_states is None:
        strain_states = get_default_strain_states(order)
    if stencils is None:
        stencils = [np.linspace(-0.01, 0.01, 5 + (order - 2) * 2)] * len(
            strain_states)
    if np.array(stencils).ndim == 1:
        stencils = [stencils] * len(strain_states)
    for state, stencil in zip(strain_states, stencils):
        strains.extend(
            [Strain.from_voigt(s * np.array(state)) for s in stencil])

    # Remove zero strains
    strains = [strain for strain in strains if not (abs(strain) < 1e-10).all()]
    vstrains = [strain.voigt for strain in strains]
    if np.linalg.matrix_rank(vstrains) < 6:
        # TODO: check for sufficiency of input for nth order
        raise ValueError("Strain list is insufficient to fit an elastic tensor")

    deformations = [s.get_deformation_matrix() for s in strains]

    if sym_reduce:
        # Note this casts deformations to a TensorMapping
        # with unique deformations as keys to symmops
        deformations = symmetry_reduce(deformations, structure)

    wf_elastic = get_wf_deformations(structure, deformations, tag=tag,
                                     db_file=db_file,
                                     vasp_input_set=vis,
                                     copy_vasp_outputs=copy_vasp_outputs,
                                     **kwargs)
    if analysis:
        defo_fws_and_tasks = get_fws_and_tasks(wf_elastic,
                                               fw_name_constraint="deformation",
                                               task_name_constraint="Transmuted")
        for idx_fw, idx_t in defo_fws_and_tasks:
            defo = \
            wf_elastic.fws[idx_fw].tasks[idx_t]['transformation_params'][0][
                'deformation']
            pass_dict = {
                'strain': Deformation(defo).green_lagrange_strain.tolist(),
                'stress': '>>output.ionic_steps.-1.stress',
                'deformation_matrix': defo}
            if sym_reduce:
                pass_dict.update({'symmops': deformations[defo]})

#.........这里部分代码省略.........