def set_sd_flags(poscar_input=None, n_layers=2, top=True, bottom=True,
                 poscar_output='POSCAR2'):
    """
    set the relaxation flags for top and bottom layers of interface.
    The upper and lower bounds of the z coordinate are determined
    based on the slab.
    Args:
         poscar_input: input poscar file name
         n_layers: number of layers to be relaxed
         top: whether n_layers from top are be relaxed
         bottom: whether n_layers from bottom are be relaxed
         poscar_output: output poscar file name
    Returns:
         None
         writes the modified poscar file
    """
    poscar1 = Poscar.from_file(poscar_input)
    sd_flags = np.zeros_like(poscar1.structure.frac_coords)
    z_coords = poscar1.structure.frac_coords[:, 2]
    z_lower_bound, z_upper_bound = None, None
    if bottom:
        z_lower_bound = np.unique(z_coords)[n_layers - 1]
        sd_flags[np.where(z_coords <= z_lower_bound)] = np.ones((1, 3))
    if top:
        z_upper_bound = np.unique(z_coords)[-n_layers]
        sd_flags[np.where(z_coords >= z_upper_bound)] = np.ones((1, 3))
    poscar2 = Poscar(poscar1.structure, selective_dynamics=sd_flags.tolist())
    poscar2.write_file(filename=poscar_output)

    def test_velocities(self):
        si = 14
        coords = list()
        coords.append([0, 0, 0])
        coords.append([0.75, 0.5, 0.75])

        # Silicon structure for testing.
        latt = [[3.8401979337, 0.00, 0.00],
                [1.9200989668, 3.3257101909, 0.00],
                [0.00, -2.2171384943, 3.1355090603]]
        struct = Structure(latt, [si, si], coords)
        poscar = Poscar(struct)
        poscar.set_temperature(900)

        v = np.array(poscar.velocities)

        for x in np.sum(v, axis=0):
            self.assertAlmostEqual(x, 0, 7)

        temperature = struct[0].specie.atomic_mass.to("kg") * \
                      np.sum(v ** 2) / (3 * const.k) * 1e10
        self.assertAlmostEqual(temperature, 900, 4,
                               'Temperature instantiated incorrectly')

        poscar.set_temperature(700)
        v = np.array(poscar.velocities)
        for x in np.sum(v, axis=0):
            self.assertAlmostEqual(
                x, 0, 7, 'Velocities initialized with a net momentum')

        temperature = struct[0].specie.atomic_mass.to("kg") * \
                      np.sum(v ** 2) / (3 * const.k) * 1e10
        self.assertAlmostEqual(temperature, 700, 4,
                               'Temperature instantiated incorrectly')

    def test_significant_figures(self):
        si = 14
        coords = list()
        coords.append([0, 0, 0])
        coords.append([0.75, 0.5, 0.75])

        # Silicon structure for testing.
        latt = [[3.8401979337, 0.00, 0.00],
                [1.9200989668, 3.3257101909, 0.00],
                [0.00, -2.2171384943, 3.1355090603]]
        struct = Structure(latt, [si, si], coords)
        poscar = Poscar(struct)
        expected_str = '''Si2
1.0
3.84 0.00 0.00
1.92 3.33 0.00
0.00 -2.22 3.14
Si
2
direct
0.00 0.00 0.00 Si
0.75 0.50 0.75 Si
'''

        actual_str = poscar.get_string(significant_figures=2)
        self.assertEqual(actual_str, expected_str, "Wrong POSCAR output!")

 def setup(self):
     """
     setup static jobs for all the calibrate objects
     copies CONTCAR to POSCAR
     sets NSW = 0
     """
     for cal in self.cal_objs:
         for i, jdir in enumerate(cal.old_job_dir_list):
             job_dir = self.job_dir + os.sep \
                 + jdir.replace(os.sep, '_').replace('.', '_') \
                 + os.sep + 'STATIC'
             logger.info('setting up job in {}'.format(job_dir))
             cal.incar = Incar.from_file(jdir + os.sep + 'INCAR')
             cal.incar['EDIFF'] = '1E-6'
             cal.incar['NSW'] = 0
             cal.potcar = Potcar.from_file(jdir + os.sep + 'POTCAR')
             cal.kpoints = Kpoints.from_file(jdir + os.sep + 'KPOINTS')
             contcar_file = jdir + os.sep + 'CONTCAR'
             if os.path.isfile(contcar_file):
                 logger.info('setting poscar file from {}'
                             .format(contcar_file))
                 cal.poscar = Poscar.from_file(contcar_file)
                 cal.add_job(job_dir=job_dir)
             else:
                 logger.critical("""CONTCAR doesnt exist.
                 Setting up job using input set in the old
                 calibration directory""")
                 cal.poscar = Poscar.from_file(jdir + os.sep + 'POSCAR')
                 cal.add_job(job_dir=job_dir)

 def setUp(self):
     p = Poscar.from_file(os.path.join(test_dir, 'POSCAR.Li2O'),
                          check_for_POTCAR=False)
     s1 = p.structure
     p = Poscar.from_file(os.path.join(test_dir, 'CONTCAR.Li2O'),
                          check_for_POTCAR=False)
     s2 = p.structure
     self.analyzer = RelaxationAnalyzer(s1, s2)

 def test_write(self):
     filepath = os.path.join(test_dir, "POSCAR")
     poscar = Poscar.from_file(filepath)
     tempfname = "POSCAR.testing"
     poscar.write_file(tempfname)
     p = Poscar.from_file(tempfname)
     self.assertArrayAlmostEqual(poscar.structure.lattice.abc, p.structure.lattice.abc, 5)
     os.remove(tempfname)

 def test_write(self):
     filepath = self.TEST_FILES_DIR / 'POSCAR'
     poscar = Poscar.from_file(filepath)
     tempfname = Path("POSCAR.testing")
     poscar.write_file(tempfname)
     p = Poscar.from_file(tempfname)
     self.assertArrayAlmostEqual(poscar.structure.lattice.abc,
                                 p.structure.lattice.abc, 5)
     tempfname.unlink()

def run_aconvasp_command(command, structure):
    """
    Helper function for calling aconvasp with different arguments
    """
    poscar = Poscar(structure)
    p = subprocess.Popen(command, stdout=subprocess.PIPE,
                         stdin=subprocess.PIPE,
                         stderr=subprocess.PIPE)
    output = p.communicate(input=poscar.get_string())
    return output

    def test_init(self):
        filepath = os.path.join(test_dir, 'POSCAR')
        poscar = Poscar.from_file(filepath,check_for_POTCAR=False)
        comp = poscar.structure.composition
        self.assertEqual(comp, Composition("Fe4P4O16"))

        # Vasp 4 type with symbols at the end.
        poscar_string = """Test1
1.0
3.840198 0.000000 0.000000
1.920099 3.325710 0.000000
0.000000 -2.217138 3.135509
1 1
direct
0.000000 0.000000 0.000000 Si
0.750000 0.500000 0.750000 F
"""
        poscar = Poscar.from_string(poscar_string)
        self.assertEqual(poscar.structure.composition, Composition("SiF"))

        poscar_string = ""
        self.assertRaises(ValueError, Poscar.from_string, poscar_string)

        # Vasp 4 tyle file with default names, i.e. no element symbol found.
        poscar_string = """Test2
1.0
3.840198 0.000000 0.000000
1.920099 3.325710 0.000000
0.000000 -2.217138 3.135509
1 1
direct
0.000000 0.000000 0.000000
0.750000 0.500000 0.750000
"""
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            poscar = Poscar.from_string(poscar_string)
        self.assertEqual(poscar.structure.composition, Composition("HHe"))
        # Vasp 4 tyle file with default names, i.e. no element symbol found.
        poscar_string = """Test3
1.0
3.840198 0.000000 0.000000
1.920099 3.325710 0.000000
0.000000 -2.217138 3.135509
1 1
Selective dynamics
direct
0.000000 0.000000 0.000000 T T T Si
0.750000 0.500000 0.750000 F F F O
"""
        poscar = Poscar.from_string(poscar_string)
        self.assertEqual(poscar.selective_dynamics, [[True, True, True],
                                                     [False, False, False]])
        self.selective_poscar = poscar

 def plot_images(self, outfile):
     """
     Generates a POSCAR with the calculated diffusion path with respect to the first endpoint.
     :param outfile: Output file for the POSCAR
     """
     sum_struct = self.__images[0].sites
     for image in self.__images:
         for site_i in self.__relax_sites:
             sum_struct.append(PeriodicSite(image.sites[site_i].specie, image.sites[site_i].frac_coords,
                                            self.__images[0].lattice, to_unit_cell=True, coords_are_cartesian=False))
     sum_struct = Structure.from_sites(sum_struct, validate_proximity=False)
     p = Poscar(sum_struct)
     p.write_file(outfile)

    def test_write_MD_poscar(self):
        # Parsing from an MD type run with velocities and predictor corrector data
        # And writing a new POSCAR from the new structure
        p = Poscar.from_file(os.path.join(test_dir, "CONTCAR.MD"), check_for_POTCAR=False)

        tempfname = "POSCAR.testing"
        p.write_file(tempfname)
        p3 = Poscar.from_file(tempfname)

        self.assertArrayAlmostEqual(p.structure.lattice.abc, p3.structure.lattice.abc, 5)
        self.assertArrayAlmostEqual(p.velocities, p3.velocities, 5)
        self.assertArrayAlmostEqual(p.predictor_corrector, p3.predictor_corrector, 5)
        self.assertEqual(p.predictor_corrector_preamble, p3.predictor_corrector_preamble)
        os.remove(tempfname)

    def test_apply_transformation(self):
        enum_trans = EnumerateStructureTransformation(refine_structure=True)
        p = Poscar.from_file(os.path.join(test_dir, 'POSCAR.LiFePO4'),
                             check_for_POTCAR=False)
        struct = p.structure
        expected_ans = [1, 3, 1]
        for i, frac in enumerate([0.25, 0.5, 0.75]):
            trans = SubstitutionTransformation({'Fe': {'Fe': frac}})
            s = trans.apply_transformation(struct)
            oxitrans = OxidationStateDecorationTransformation(
                {'Li': 1, 'Fe': 2, 'P': 5, 'O': -2})
            s = oxitrans.apply_transformation(s)
            alls = enum_trans.apply_transformation(s, 100)
            self.assertEqual(len(alls), expected_ans[i])
            self.assertIsInstance(trans.apply_transformation(s), Structure)
            for s in alls:
                self.assertIn("energy", s)

        #make sure it works for non-oxidation state decorated structure
        trans = SubstitutionTransformation({'Fe': {'Fe': 0.5}})
        s = trans.apply_transformation(struct)
        alls = enum_trans.apply_transformation(s, 100)
        self.assertEqual(len(alls), 3)
        self.assertIsInstance(trans.apply_transformation(s), Structure)
        for s in alls:
            self.assertNotIn("energy", s)

 def test_get_structure(self):
     if not aio.ase_loaded:
         raise SkipTest("ASE not present. Skipping...")
     p = Poscar.from_file(os.path.join(test_dir, 'POSCAR'))
     atoms = aio.AseAtomsAdaptor.get_atoms(p.structure)
     self.assertEqual(aio.AseAtomsAdaptor.get_structure(atoms).formula,
                      "Fe4 P4 O16")

def read_poscar(i_path, l_get_sorted_symbols=False):
    poscar = Poscar.from_file("{}".format(i_path))
    struct = poscar.structure
    if l_get_sorted_symbols:
        return struct, poscar.site_symbols
    else:
        return struct

 def test_from_md_run(self):
     #Parsing from an MD type run with velocities
     p = Poscar.from_file(os.path.join(test_dir, "CONTCAR.MD"),
                          check_for_POTCAR=False)
     self.assertAlmostEqual(np.sum(np.array(p.velocities)), 0.0065417961324)
     self.assertEqual(p.predictor_corrector[0][0], 1)
     self.assertEqual(p.predictor_corrector[1][0], 2)

    def test_init_from_structure(self):
        filepath = os.path.join(test_dir, 'POSCAR')
        poscar = Poscar.from_file(filepath)
        struct = poscar.structure
        xyz = XYZ(struct)
        ans = """24
Fe4 P4 O16
Fe 2.277347 4.550379 2.260125
Fe 2.928536 1.516793 4.639870
Fe 7.483231 4.550379 0.119620
Fe 8.134420 1.516793 2.499364
P 0.985089 1.516793 1.990624
P 4.220794 4.550379 4.370369
P 6.190973 1.516793 0.389120
P 9.426677 4.550379 2.768865
O 0.451582 4.550379 3.365614
O 1.006219 1.516793 3.528306
O 1.725331 0.279529 1.358282
O 1.725331 2.754057 1.358282
O 3.480552 3.313115 3.738027
O 3.480552 5.787643 3.738027
O 4.199665 4.550379 1.148562
O 4.754301 1.516793 0.985870
O 5.657466 4.550379 3.773620
O 6.212102 1.516793 3.610928
O 6.931215 0.279529 1.021463
O 6.931215 2.754057 1.021463
O 8.686436 3.313115 3.401208
O 8.686436 5.787643 3.401208
O 9.405548 4.550379 1.231183
O 9.960184 1.516793 1.393875"""
        self.assertEqual(str(xyz), ans)

    def run_task(self, fw_spec):
        chgcar_start = False
        # read the VaspInput from the previous run

        poscar = Poscar.from_file(zpath('POSCAR'))
        incar = Incar.from_file(zpath('INCAR'))

        # figure out what GGA+U values to use and override them
        # LDAU values to use
        mpvis = MPVaspInputSet()
        ggau_incar = mpvis.get_incar(poscar.structure).as_dict()
        incar_updates = {k: ggau_incar[k] for k in ggau_incar.keys() if 'LDAU' in k}

        for k in ggau_incar:
            # update any parameters not set explicitly in previous INCAR
            if k not in incar and k in ggau_incar:
                incar_updates[k] = ggau_incar[k]

        incar.update(incar_updates)  # override the +U keys


        # start from the CHGCAR of previous run
        if os.path.exists('CHGCAR'):
            incar['ICHARG'] = 1
            chgcar_start = True

        # write back the new INCAR to the current directory
        incar.write_file('INCAR')
        return FWAction(stored_data={'chgcar_start': chgcar_start})

    def test_apply_transformation(self):
        trans = MagOrderingTransformation({"Fe": 5})
        p = Poscar.from_file(os.path.join(test_dir, 'POSCAR.LiFePO4'),
                             check_for_POTCAR=False)
        s = p.structure
        alls = trans.apply_transformation(s, 10)
        self.assertEqual(len(alls), 3)
        f = SpacegroupAnalyzer(alls[0]["structure"], 0.1)
        self.assertEqual(f.get_space_group_number(), 31)

        model = IsingModel(5, 5)
        trans = MagOrderingTransformation({"Fe": 5},
                                          energy_model=model)
        alls2 = trans.apply_transformation(s, 10)
        # Ising model with +J penalizes similar neighbor magmom.
        self.assertNotEqual(alls[0]["structure"], alls2[0]["structure"])
        self.assertEqual(alls[0]["structure"], alls2[2]["structure"])

        s = self.get_structure('Li2O')
        # Li2O doesn't have magnetism of course, but this is to test the
        # enumeration.
        trans = MagOrderingTransformation({"Li+": 1}, max_cell_size=3)
        alls = trans.apply_transformation(s, 100)
        # TODO: check this is correct, unclear what len(alls) should be
        self.assertEqual(len(alls), 12)

        trans = MagOrderingTransformation({"Ni": 5})
        alls = trans.apply_transformation(self.NiO.get_primitive_structure(),
                                          return_ranked_list=10)
        self.assertEqual(self.NiO_AFM_111.lattice, alls[0]["structure"].lattice)
        self.assertEqual(self.NiO_AFM_001.lattice, alls[1]["structure"].lattice)

    def setUp(self):
        if "VASP_PSP_DIR" not in os.environ:
            os.environ["VASP_PSP_DIR"] = test_dir
        filepath = os.path.join(test_dir, 'POSCAR')
        poscar = Poscar.from_file(filepath)
        self.struct = poscar.structure

        self.mitparamset = MITVaspInputSet()
        self.mitparamset_unsorted = MITVaspInputSet(sort_structure=False)
        self.mithseparamset = MITHSEVaspInputSet()
        self.paramset = MPVaspInputSet()
        self.userparamset = MPVaspInputSet(
            user_incar_settings={'MAGMOM': {"Fe": 10, "S": -5, "Mn3+": 100}}
        )
        self.mitggaparam = MITGGAVaspInputSet()
        self.mpstaticparamset = MPStaticVaspInputSet()
        self.mpnscfparamsetu = MPNonSCFVaspInputSet(
            {"NBANDS": 50}, mode="Uniform")
        self.mpnscfparamsetl = MPNonSCFVaspInputSet(
            {"NBANDS": 60}, mode="Line")
        self.mphseparamset = MPHSEVaspInputSet()
        self.mpbshseparamsetl = MPBSHSEVaspInputSet(mode="Line")
        self.mpbshseparamsetu = MPBSHSEVaspInputSet(
            mode="Uniform", added_kpoints=[[0.5, 0.5, 0.0]])
        self.mpdielparamset = MPStaticDielectricDFPTVaspInputSet()

    def test_static_constructors(self):
        kpoints = Kpoints.gamma_automatic([3, 3, 3], [0, 0, 0])
        self.assertEqual(kpoints.style, Kpoints.supported_modes.Gamma)
        self.assertEqual(kpoints.kpts, [[3, 3, 3]])
        kpoints = Kpoints.monkhorst_automatic([2, 2, 2], [0, 0, 0])
        self.assertEqual(kpoints.style, Kpoints.supported_modes.Monkhorst)
        self.assertEqual(kpoints.kpts, [[2, 2, 2]])
        kpoints = Kpoints.automatic(100)
        self.assertEqual(kpoints.style, Kpoints.supported_modes.Automatic)
        self.assertEqual(kpoints.kpts, [[100]])
        filepath = os.path.join(test_dir, "POSCAR")
        poscar = Poscar.from_file(filepath)
        kpoints = Kpoints.automatic_density(poscar.structure, 500)
        self.assertEqual(kpoints.kpts, [[2, 4, 4]])
        self.assertEqual(kpoints.style, Kpoints.supported_modes.Monkhorst)
        kpoints = Kpoints.automatic_density(poscar.structure, 500, True)
        self.assertEqual(kpoints.style, Kpoints.supported_modes.Gamma)
        kpoints = Kpoints.automatic_density_by_vol(poscar.structure, 1000)
        self.assertEqual(kpoints.kpts, [[6, 11, 13]])
        self.assertEqual(kpoints.style, Kpoints.supported_modes.Gamma)

        s = poscar.structure
        s.make_supercell(3)
        kpoints = Kpoints.automatic_density(s, 500)
        self.assertEqual(kpoints.kpts, [[1, 1, 1]])
        self.assertEqual(kpoints.style, Kpoints.supported_modes.Gamma)