def create_atoms():
     a = Atoms(
         symbols=["Fe", "Cu", "Ni", "Al"],
         positions=np.random.random((4, 3)),
         cell=np.eye(3),
     )
     b = a.copy()
     b.positions = np.random.random(b.positions.shape)
     c = a.copy()
     c[0] = "Cu"
     d = a + b
     return a, b, c, d
Exemple #2
0
 def test_copy(self):
     pos, cell = generate_fcc_lattice()
     basis = Atoms(symbols='Al', positions=pos, cell=cell)
     basis_copy = basis.copy()
     self.assertEqual(basis, basis_copy)
     basis_copy[:] = "Pt"
     self.assertNotEqual(basis, basis_copy)
class HessianJob(GenericInteractive):
    def __init__(self, project, job_name):
        super(HessianJob, self).__init__(project, job_name)
        self.__version__ = "0.0.1"
        self.__name__ = "HessianJob"
        self._python_only_job = True
        self._force_constants = None
        self._reference_structure = None
        self._energy_pot = 0
        self._forces = np.zeros(3)
        self._pressure = np.zeros((3, 3))
        self._stiffness_tensor = np.zeros((6, 6))
        self._pressure_times_volume = 0
        self._displacements = np.zeros(3)

    @property
    def structure(self):
        return GenericInteractive.structure.fget(self)

    @structure.setter
    def structure(self, structure):
        if self._reference_structure is None:
            self.set_reference_structure(structure)
        GenericInteractive.structure.fset(self, structure)

    def run_static(self):
        run_mode = self.server.run_mode.mode
        self.interactive_open()
        self.run_if_interactive()
        self.interactive_close()
        self.server.run_mode = run_mode

    def set_elastic_moduli(self, bulk_modulus=0, shear_modulus=0):
        self._stiffness_tensor = np.zeros((6, 6))
        self._stiffness_tensor[:3, :3] = bulk_modulus-2*shear_modulus/3
        self._stiffness_tensor[:3, :3] += np.eye(3)*2*shear_modulus
        self._stiffness_tensor[3:, 3:] = np.eye(3)*shear_modulus

    def set_force_constants(self, force_constants):
        if self.structure is None:
            raise ValueError('Set reference structure via set_reference_structure() first')
        n_atom = len(self.structure.positions)
        if len(np.array([force_constants]).flatten()) == 1:
            self._force_constants = force_constants*np.eye(3*n_atom)
        elif np.array(force_constants).shape == (3*n_atom, 3*n_atom):
            self._force_constants = force_constants
        elif np.array(force_constants).shape == (n_atom, n_atom):
            na = np.newaxis
            self._force_constants = (np.array(force_constants)[:, na, :, na]*np.eye(3)[na, :, na, :]).flatten()
        elif len(np.shape(force_constants)) == 4:
            force_shape = np.shape(force_constants)
            if force_shape[2] == 3 and force_shape[3] == 3:
                force_reshape = force_shape[0] * force_shape[2]
                self._force_constants = np.transpose(
                    force_constants,
                    (0, 2, 1, 3)
                ).reshape((force_reshape, force_reshape))
            elif force_shape[1] == 3 and force_shape[3] == 3:
                self._force_constants = np.array(force_constants).reshape(3*n_atom, 3*n_atom)
            else:
                raise AssertionError('force constant shape not recognized')
        else:
            raise AssertionError('force constant shape not recognized')

    def set_reference_structure(self, structure):
        self._reference_structure = structure.copy()
        if self.structure is None:
            self.structure = structure.copy()

    def validate_ready_to_run(self):
        super(HessianJob, self).validate_ready_to_run()
        if self._force_constants is None:
            raise AssertionError('set force constants by set_force_constants before run')
        if self._reference_structure is None:
            raise AssertionError('set reference structure by set_reference_structure before run')

    def interactive_forces_getter(self):
        return self._forces

    def interactive_energy_pot_getter(self):
        return self._energy_pot

    def interactive_energy_tot_getter(self):
        return self.interactive_energy_pot_getter()

    def interactive_positions_getter(self):
        return self.structure.positions.copy()

    def interactive_cells_getter(self):
        return self.structure.cell.copy()

    def interactive_volume_getter(self):
        return self.structure.get_volume()

    def interactive_pressures_getter(self):
        return self._pressure

    def interactive_cells_setter(self, cell):
        if np.sum(self._stiffness_tensor) != 0:
            epsilon = np.einsum(
                'ij,jk->ik',
                self.structure.cell,
                np.linalg.inv(self._reference_structure.cell)
            )-np.eye(3)
            epsilon = (epsilon+epsilon.T)*0.5
            epsilon = np.append(
                epsilon.diagonal(),
                np.roll(epsilon, -1, axis=0).diagonal()
            )
            pressure = -np.einsum('ij,j->i', self._stiffness_tensor, epsilon)
            self._pressure = pressure[3:]*np.roll(np.eye(3), -1, axis=1)
            self._pressure += self._pressure.T+np.eye(3)*pressure[:3]
            self._pressure_times_volume = -np.sum(epsilon*pressure)*self.structure.get_volume()

    def interactive_positions_setter(self, positions):
        displacements = self.structure.get_scaled_positions()
        displacements -= self._reference_structure.get_scaled_positions()
        displacements -= np.rint(displacements)
        self._displacements = np.einsum('ji,ni->nj', self.structure.cell, displacements)

    def calculate_forces(self):
        position_transformed = self._displacements.reshape(
            self._displacements.shape[0] * self._displacements.shape[1])
        forces_transformed = -np.dot(self._force_constants, position_transformed)
        self._forces = forces_transformed.reshape(self._displacements.shape)
        self._energy_pot = -1 / 2 * np.dot(position_transformed, forces_transformed)
        self._energy_pot += self._pressure_times_volume

    def interactive_collect(self):
        super(HessianJob, self).interactive_collect()
        self.interactive_cache["displacements"].append(self._displacements)

    def interactive_initialize_interface(self):
        self.interactive_positions_setter(self.structure.positions)
        self.interactive_cells_setter(self.structure.cell)
        self._interactive_library = True

    def run_if_interactive(self):
        """
        Run the job as Python library and store the result in the HDF5 File.

        Returns:
            int: job ID
        """
        super(HessianJob, self).run_if_interactive()
        self.calculate_forces()
        self.interactive_collect()

    def to_hdf(self, hdf=None, group_name=None):
        """
        Store the ExampleJob object in the HDF5 File

        Args:
            hdf (ProjectHDFio): HDF5 group object - optional
            group_name (str): HDF5 subgroup name - optional
        """
        super(HessianJob, self).to_hdf(hdf=hdf, group_name=group_name)
        with self.project_hdf5.open("input") as hdf5_input:
            if self._force_constants is not None:
                hdf5_input["force_constants"] = self._force_constants
            if self._reference_structure is not None:
                self._reference_structure.to_hdf(hdf5_input)

    def from_hdf(self, hdf=None, group_name=None):
        """
        Restore the ExampleJob object in the HDF5 File

        Args:
            hdf (ProjectHDFio): HDF5 group object - optional
            group_name (str): HDF5 subgroup name - optional
        """
        super(HessianJob, self).from_hdf(hdf=hdf, group_name=group_name)
        with self.project_hdf5.open("input") as hdf5_input:
            if "structure" in hdf5_input.list_groups():
                self._reference_structure = Atoms().from_hdf(hdf5_input)
                self._structure = self._reference_structure.copy()
            if "force_constants" in hdf5_input.list_nodes():
                self._force_constants = hdf5_input["force_constants"]

    def interactive_close(self):
        if self.interactive_is_activated():
            super(HessianJob, self).interactive_close()
            with self.project_hdf5.open("output") as h5:
                if "interactive" in h5.list_groups():
                    for key in h5["interactive"].list_nodes():
                        h5["generic/" + key] = h5["interactive/" + key]
Exemple #4
0
class Yaff(AtomisticGenericJob):
    def __init__(self, project, job_name):
        super(Yaff, self).__init__(project, job_name)
        self.__name__ = "Yaff"
        self._executable_activate(enforce=True)
        self.input = YaffInput()
        self.ffatypes = None
        self.ffatype_ids = None
        self.enhanced = None  # should have more generic name e.g. enhanced

    def calc_minimize(self,
                      cell=False,
                      gpos_tol=1e-8,
                      dpos_tol=1e-6,
                      grvecs_tol=1e-8,
                      drvecs_tol=1e-6,
                      max_iter=1000,
                      n_print=5):
        """
        Set up an optimization calculation.

        **Arguments**

        cell (bool): Set True if the cell also has to be optimized
        gpos_tol (float): Convergence criterion for RMS of gradients towards atomic coordinates
        dpos_tol (float): Convergence criterion for RMS of differences of atomic coordinates
        grvecs_tol (float): Convergence criterion for RMS of gradients towards cell parameters
        drvecs_tol (float): Convergence criterion for RMS of differences of cell parameters
        max_iter (int): Maximum number of optimization steps
        n_print (int):  Print frequency
        """
        if cell:
            self.input['jobtype'] = 'opt_cell'
        else:
            self.input['jobtype'] = 'opt'
        self.input['nsteps'] = max_iter
        self.input['h5step'] = n_print
        self.input['gpos_rms'] = gpos_tol
        self.input['dpos_rms'] = dpos_tol
        self.input['grvecs_rms'] = grvecs_tol
        self.input['drvecs_rms'] = drvecs_tol

        super(Yaff, self).calc_minimize(max_iter=max_iter, n_print=n_print)

    def calc_static(self):
        """
        Set up a static force field calculation.
        """

        self.input['jobtype'] = 'sp'
        super(Yaff, self).calc_static()

    def calc_md(self,
                temperature=None,
                pressure=None,
                nsteps=1000,
                time_step=1.0 * femtosecond,
                n_print=5,
                timecon_thermo=100.0 * femtosecond,
                timecon_baro=1000.0 * femtosecond):
        """
        Set an MD calculation within Yaff. Nosé Hoover chain is used by default.

        **Arguments**

        temperature (None/float): Target temperature. If set to None, an NVE calculation is performed.
                                  It is required when the pressure is set
        pressure (None/float): Target pressure. If set to None, an NVE or an NVT calculation is performed.
        nsteps (int): Number of md steps
        time_step (float): Step size between two steps.
        n_print (int):  Print frequency
        timecon_thermo (float): The time associated with the thermostat adjusting the temperature.
        timecon_baro (float): The time associated with the barostat adjusting the temperature.

        """
        self.input['temp'] = temperature
        self.input['press'] = pressure
        self.input['nsteps'] = nsteps
        self.input['timestep'] = time_step
        self.input['h5step'] = n_print
        self.input['timecon_thermo'] = timecon_thermo
        self.input['timecon_baro'] = timecon_baro

        if temperature is None:
            self.input['jobtype'] = 'nve'
        else:
            if pressure is None:
                self.input['jobtype'] = 'nvt'
            else:
                self.input['jobtype'] = 'npt'

        super(Yaff, self).calc_md(temperature=temperature,
                                  pressure=pressure,
                                  n_ionic_steps=nsteps,
                                  time_step=time_step,
                                  n_print=n_print,
                                  temperature_damping_timescale=timecon_thermo,
                                  pressure_damping_timescale=timecon_baro)

    def load_chk(self, fn):
        """
        Load the atom types, atom type ids and structure by reading a .chk file.

        **Arguments**

        fn      the path to the chk file
        """

        system = System.from_file(fn)
        system.set_standard_masses()
        if len(system.pos.shape) != 2:
            raise IOError(
                "Something went wrong, positions in CHK file %s should have Nx3 dimensions"
                % fn)
        if system.cell.rvecs is not None and len(system.cell.rvecs) > 0:
            self.structure = Atoms(
                positions=system.pos.copy() / angstrom,
                numbers=system.numbers,
                masses=system.masses,
                cell=system.cell.rvecs / angstrom,
                pbc=True,
            )
        else:
            self.structure = Atoms(
                positions=system.pos.copy() / angstrom,
                numbers=system.numbers,
                masses=system.masses,
            )
        if system.ffatypes is not None:
            self.ffatypes = system.ffatypes
        if system.ffatype_ids is not None:
            self.ffatype_ids = system.ffatype_ids

    def set_mtd(self,
                ics,
                height,
                sigma,
                pace,
                fn='HILLS',
                fn_colvar='COLVAR',
                stride=10,
                temp=300):
        """
        Setup a Metadynamics run using PLUMED along the internal coordinates
        defined in the ICs argument.

        **Arguments**

        ics     a list of entries defining each internal coordinate. Each
                of these entries should be of the form (kind, [i, j, ...])

                Herein, kind defines the kind of IC as implemented in PLUMED:

                    i.e. distance, angle, torsion, volume, cell, ... see
                    https://www.plumed.org/doc-v2.5/user-doc/html/_colvar.html
                    for more information).

                and [i, j, ...] is a list of atom indices, starting from 0, involved in this
                IC. If no atom indices are required for e.g. volume, provide an empty list.

                An example for a 1D metadynamica using the distance between
                atoms 2 and 4:

                    ics = [('distance', [2,4])]

        height  the height of the Gaussian hills, can be a single value
                (the gaussian hills for each IC have identical height) or
                a list of values, one for each IC defined.

        sigmas  the sigma of the Gaussian hills, can be a single value
                (the gaussian hills for each IC have identical height) or
                a list of values, one for each IC defined.

        pace    the number of steps after which the gaussian hills are
                updated.

        fn      the PLUMED output file for the gaussian hills

        fn_colvar
                the PLUMED output file for logging of collective variables

        stride  the number of steps after which the internal coordinate
                values and bias are printed to the COLVAR output file.

        temp    the system temperature
        """
        for l in ics:
            assert len(l) == 2
            assert isinstance(l[0], str)
            assert isinstance(l[1], list) or isinstance(l[1], tuple)
        ickinds = np.array([ic[0] for ic in ics], dtype='S22')
        icindices = np.array([np.array(ic[1]) + 1
                              for ic in ics])  # plumed starts counting from 1
        if not isinstance(height, list) and not isinstance(height, np.ndarray):
            height = np.array([height])
        if not isinstance(sigma, list) and not isinstance(sigma, np.ndarray):
            sigma = np.array([sigma])
        self.enhanced = {
            'ickinds': ickinds,
            'icindices': icindices,
            'height': height,
            'sigma': sigma,
            'pace': pace,
            'file': fn,
            'file_colvar': fn_colvar,
            'stride': stride,
            'temp': temp
        }

    def set_us(self, ics, kappa, loc, fn_colvar='COLVAR', stride=10, temp=300):
        """
        Setup an Umbrella sampling run using PLUMED along the internal coordinates
        defined in the ICs argument.

        **Arguments**

        ics     a list of entries defining each an internal coordinate. Each
                of these entries should be of the form (kind, [i, j, ...])

                Herein, kind defines the kind of IC as implemented in PLUMED:

                    i.e. distance, angle, torsion, volume, cell, ... see
                    https://www.plumed.org/doc-v2.5/user-doc/html/_colvar.html
                    for more information).

                and [i, j, ...] is a list of atom indices, starting from 0, involved in this
                IC. If no atom indices are required for e.g. volume, provide an empty list.

                An example for a 1D metadynamica using the distance between
                atoms 2 and 4:

                    ics = [('distance', [2,4])]

        kappa   the value of the force constant of the harmonic bias potential,
                can be a single value (the harmonic bias potential for each IC has identical kappa)
                or a list of values, one for each IC defined.

        loc     the location of the umbrella
                (should have a length equal to the number of ICs)

        fn_colvar
                the PLUMED output file for logging of collective variables

        stride  the number of steps after which the internal coordinate
                values and bias are printed to the COLVAR output file.

        temp    the system temperature
        """
        for l in ics:
            assert len(l) == 2
            assert isinstance(l[0], str)
            assert isinstance(l[1], list) or isinstance(l[1], tuple)
        ickinds = np.array([ic[0] for ic in ics], dtype='S22')
        icindices = np.array([np.array(ic[1]) + 1
                              for ic in ics])  # plumed starts counting from 1
        if not isinstance(kappa, list) and not isinstance(kappa, np.ndarray):
            kappa = np.array([kappa])
        if not isinstance(loc, list) and not isinstance(loc, np.ndarray):
            loc = np.array([loc])
        assert len(loc) == len(ics)
        self.enhanced = {
            'ickinds': ickinds,
            'icindices': icindices,
            'kappa': kappa,
            'loc': loc,
            'file_colvar': fn_colvar,
            'stride': stride,
            'temp': temp
        }

    def detect_ffatypes(self,
                        ffatypes=None,
                        ffatype_rules=None,
                        ffatype_level=None):
        """
        Define atom types by explicitely giving them through the
        ffatypes keyword, defining atype rules using the ATSELECT
        language implemented in Yaff (see the Yaff documentation at
        http://molmod.github.io/yaff/ug_atselect.html) or by specifying
        the ffatype_level employing the built-in routine in QuickFF.
        """
        numbers = np.array([
            pt[symbol].number
            for symbol in self.structure.get_chemical_symbols()
        ])
        if self.structure.cell is not None and np.all(
                np.array(self.structure.cell) != np.zeros([3, 3])):
            system = System(numbers,
                            self.structure.positions.copy() * angstrom,
                            rvecs=np.array(self.structure.cell) * angstrom)
        else:
            system = System(numbers,
                            self.structure.positions.copy() * angstrom)
        system.detect_bonds()

        if not sum([
                ffatypes is None, ffatype_rules is None, ffatype_level is None
        ]) == 2:
            raise IOError(
                'Exactly one of ffatypes, ffatype_rules and ffatype_level should be defined'
            )

        if ffatypes is not None:
            system.ffatypes = ffatypes
            system.ffatype_ids = None
            system._init_derived_ffatypes()
        if ffatype_rules is not None:
            system.detect_ffatypes(ffatype_rules)
        if ffatype_level is not None:
            set_ffatypes(system, ffatype_level)

        self.ffatypes = system.ffatypes.copy()
        self.ffatype_ids = system.ffatype_ids.copy()

    def write_input(self):
        input_dict = {
            'jobtype':
            self.input['jobtype'],
            'symbols':
            self.structure.get_chemical_symbols(),
            'numbers':
            np.array([
                pt[symbol].number
                for symbol in self.structure.get_chemical_symbols()
            ]),
            'ffatypes':
            self.ffatypes,
            'ffatype_ids':
            self.ffatype_ids,
            'ffpars':
            self.input['ffpars'],
            'pos':
            self.structure.positions,
            'rcut':
            self.input['rcut'],
            'alpha_scale':
            self.input['alpha_scale'],
            'gcut_scale':
            self.input['gcut_scale'],
            'smooth_ei':
            self.input['smooth_ei'],
            'nsteps':
            self.input['nsteps'],
            'h5step':
            self.input['h5step'],
            'gpos_rms':
            self.input['gpos_rms'],
            'dpos_rms':
            self.input['dpos_rms'],
            'grvecs_rms':
            self.input['grvecs_rms'],
            'drvecs_rms':
            self.input['drvecs_rms'],
            'hessian_eps':
            self.input['hessian_eps'],
            'timestep':
            self.input['timestep'],
            'temp':
            self.input['temp'],
            'press':
            self.input['press'],
            'timecon_thermo':
            self.input['timecon_thermo'],
            'timecon_baro':
            self.input['timecon_baro'],
            'enhanced':
            self.enhanced,
            'cell':
            None
        }
        if self.structure.cell is not None:
            input_dict['cell'] = np.array(self.structure.get_cell())
        write_chk(input_dict, working_directory=self.working_directory)
        write_pars(input_dict=input_dict,
                   working_directory=self.working_directory)
        if self.input['jobtype'] == 'sp':
            write_ysp(input_dict=input_dict,
                      working_directory=self.working_directory)
        elif self.input['jobtype'] == 'opt':
            write_yopt(input_dict=input_dict,
                       working_directory=self.working_directory)
        elif self.input['jobtype'] == 'opt_cell':
            write_yopt_cell(input_dict=input_dict,
                            working_directory=self.working_directory)
        elif self.input['jobtype'] == 'hess':
            write_yhess(input_dict=input_dict,
                        working_directory=self.working_directory)
        elif self.input['jobtype'] == 'nve':
            write_ynve(input_dict=input_dict,
                       working_directory=self.working_directory)
        elif self.input['jobtype'] == 'nvt':
            write_ynvt(input_dict=input_dict,
                       working_directory=self.working_directory)
        elif self.input['jobtype'] == 'npt':
            write_ynpt(input_dict=input_dict,
                       working_directory=self.working_directory)
        else:
            raise IOError('Invalid job type for Yaff job, received %s' %
                          self.input['jobtype'])
        if not self.enhanced is None:
            write_plumed_enhanced(input_dict,
                                  working_directory=self.working_directory)

    def collect_output(self):
        output_dict = collect_output(
            output_file=posixpath.join(self.working_directory, 'output.h5'))
        with self.project_hdf5.open("output") as hdf5_output:
            for k, v in output_dict.items():
                hdf5_output[k] = v

    def to_hdf(self, hdf=None, group_name=None):
        super(Yaff, self).to_hdf(hdf=hdf, group_name=group_name)
        with self.project_hdf5.open("input") as hdf5_input:
            self.structure.to_hdf(hdf5_input)
            self.input.to_hdf(hdf5_input)
            hdf5_input['generic/ffatypes'] = np.asarray(self.ffatypes, 'S22')
            hdf5_input['generic/ffatype_ids'] = self.ffatype_ids
            if not self.enhanced is None:
                grp = hdf5_input.create_group('generic/enhanced')
                for k, v in self.enhanced.items():
                    grp[k] = v

    def from_hdf(self, hdf=None, group_name=None):
        super(Yaff, self).from_hdf(hdf=hdf, group_name=group_name)
        with self.project_hdf5.open("input") as hdf5_input:
            self.input.from_hdf(hdf5_input)
            self.structure = Atoms().from_hdf(hdf5_input)
            self.ffatypes = np.char.decode(
                hdf5_input['generic/ffatypes'])  # decode byte string literals
            self.ffatype_ids = hdf5_input['generic/ffatype_ids']

            if "enhanced" in hdf5_input['generic'].keys():
                self.enhanced = {}
                for key, val in hdf5_input['generic/enhanced'].items():
                    if key == 'ickinds':
                        self.enhanced[key] = np.char.decode(val)
                    else:
                        self.enhanced[key] = val

    def get_structure(self, iteration_step=-1, wrap_atoms=True):
        """
        Overwrite the get_structure routine from AtomisticGenericJob because we want to avoid
        defining a unit cell when one does not exist
        """
        if not (self.structure is not None):
            raise AssertionError()

        positions = self.get("output/generic/positions")
        cells = self.get("output/generic/cells")

        snapshot = self.structure.copy()
        snapshot.positions = positions[iteration_step]
        if cells is not None:
            snapshot.cell = cells[iteration_step]
        indices = self.get("output/generic/indices")
        if indices is not None:
            snapshot.indices = indices[iteration_step]
        if wrap_atoms and cells is not None:
            return snapshot.center()
        else:
            return snapshot

    # Plot functions are deprecated while yaff is no longer in atomic units!
    def plot(self,
             ykey,
             xkey='generic/steps',
             xunit='au',
             yunit='au',
             ref=None,
             linestyle='-',
             rolling_average=False):
        xs = self['output/%s' % xkey] / parse_unit(xunit)
        ys = self['output/%s' % ykey] / parse_unit(yunit)
        if rolling_average:
            ra = np.zeros(len(ys))
            for i, y in enumerate(ys):
                if i == 0:
                    ra[i] = ys[0]
                else:
                    ra[i] = (i * ra[i - 1] + ys[i]) / (i + 1)
            ys = ra.copy()

        self._ref(ys, ref)

        pp.clf()
        pp.plot(xs, ys, linestyle)
        pp.xlabel('%s [%s]' % (xkey, xunit))
        pp.ylabel('%s [%s]' % (ykey, yunit))
        pp.show()

    # Plot functions are deprecated while yaff is no longer in atomic units!
    def plot_multi(self,
                   ykeys,
                   xkey='generic/steps',
                   xunit='au',
                   yunit='au',
                   ref=None,
                   linestyle='-',
                   rolling_average=False):
        # Assume that all ykeys have the same length than the xkey
        xs = self['output/%s' % xkey] / parse_unit(xunit)
        yss = np.array(
            [self['output/%s' % ykey] / parse_unit(yunit) for ykey in ykeys])

        if rolling_average:
            for ys in yss:
                ra = np.zeros(len(ys))
                for i, y in enumerate(ys):
                    if i == 0:
                        ra[i] = ys[0]
                    else:
                        ra[i] = (i * ra[i - 1] + ys[i]) / (i + 1)
                ys = ra.copy()

        if not isinstance(ref, list):
            for ys in yss:
                self._ref(ys, ref)
        else:
            assert len(ref) == len(yss)
            for n in range(len(ref)):
                _ref(yss[n], ref[n])

        pp.clf()
        for n, ys in enumerate(yss):
            pp.plot(xs, ys, linestyle, label=ykeys[n])
        pp.xlabel('%s [%s]' % (xkey, xunit))
        pp.ylabel('[%s]' % (yunit))
        pp.legend()
        pp.show()

    @staticmethod
    def _ref(ys, ref):
        if isinstance(ref, int):
            ys -= ys[ref]
        elif isinstance(ref, float):
            ys -= ref
        elif isinstance(ref, str):
            if ref == 'min':
                ys -= min(ys)
            elif ref == 'max':
                ys -= max(ys)
            elif ref == 'mean':
                ys -= np.mean(ys)

    def log(self):
        with open(posixpath.join(self.working_directory, 'yaff.log')) as f:
            print(f.read())

    def get_yaff_system(self, snapshot=0):
        numbers = np.array([
            pt[symbol].number
            for symbol in self.structure.get_chemical_symbols()
        ])
        if snapshot == 0:
            struct = self.structure
        else:
            struct = self.get_structure(iteration_step=snapshot,
                                        wrap_atoms=False)
        pos = struct.positions.reshape(-1, 3) * angstrom
        cell = struct.cell
        if cell is None:
            system = System(numbers,
                            pos,
                            ffatypes=self.ffatypes,
                            ffatype_ids=self.ffatype_ids)
        else:
            system = System(numbers,
                            pos,
                            rvecs=cell * angstrom,
                            ffatypes=self.ffatypes,
                            ffatype_ids=self.ffatype_ids)
        system.detect_bonds()
        system.set_standard_masses()
        return system

    def get_yaff_ff(self, system=None):
        if system is None:
            system = self.get_yaff_system()
        fn_pars = posixpath.join(self.working_directory, 'pars.txt')
        if not os.path.isfile(fn_pars):
            raise IOError(
                'No pars.txt file find in job working directory. Have you already run the job?'
            )
        ff = ForceField.generate(system,
                                 fn_pars,
                                 rcut=self.input['rcut'],
                                 alpha_scale=self.input['alpha_scale'],
                                 gcut_scale=self.input['gcut_scale'],
                                 smooth_ei=self.input['smooth_ei'])
        return ff

    def mtd_sum_hills_1d(self, fn=None):
        """
        Creates a fes.dat file for plotting the free energy surface after a mtd simulation.

        **Arguments**

        fn      path to the hills file or hills files (comma separated)
        """
        if fn is None:
            fn = posixpath.join(self.working_directory, self.enhanced['file'])
        fn_out = posixpath.join(self.working_directory, 'fes.dat')

        subprocess.check_output(
            "ml load PLUMED/2.5.2-intel-2019a-Python-3.7.2; plumed sum_hills --hills {} --outfile {}"
            .format(fn, fn_out),
            stderr=subprocess.STDOUT,
            universal_newlines=True,
            shell=True)
Exemple #5
0
class TestAtoms(unittest.TestCase):
    @classmethod
    def tearDownClass(cls):
        if sys.version_info[0] >= 3:
            file_location = os.path.dirname(os.path.abspath(__file__))
            if os.path.isfile(
                    os.path.join(file_location,
                                 "../../static/atomistics/test_hdf")):
                os.remove(
                    os.path.join(file_location,
                                 "../../static/atomistics/test_hdf"))

    def setUp(self):
        pass
        self.CO2 = Atoms("CO2",
                         positions=[[0, 0, 0], [0, 0, 1.5], [0, 1.5, 0]])
        C = Atom('C').element
        self.C3 = Atoms([C, C, C], positions=[[0, 0, 0], [0, 0, 2], [0, 2, 0]])
        self.C2 = Atoms(2 * [Atom('C')])

    def test__init__(self):
        pos, cell = generate_fcc_lattice()
        pse = PeriodicTable()
        el = pse.element("Al")
        basis = Atoms()
        self.assertIsInstance(basis, Atoms)
        self.assertIsInstance(basis.info, dict)
        self.assertIsInstance(basis.arrays, dict)
        self.assertIsInstance(basis.adsorbate_info, dict)
        self.assertIsInstance(basis.units, dict)
        self.assertIsInstance(basis.pbc, (bool, list, np.ndarray))
        self.assertIsInstance(basis.indices, np.ndarray)
        self.assertIsNone(basis._internal_positions)
        self.assertIsNone(basis.positions)
        self.assertIsNone(basis.scaled_positions)
        self.assertIsInstance(basis.species, list)
        self.assertIsInstance(basis.elements, np.ndarray)
        self.assertIsNone(basis.cell)
        basis = Atoms(symbols='Al', positions=pos, cell=cell)
        self.assertIsInstance(basis, Atoms)
        self.assertEqual(basis.get_spacegroup()["Number"], 225)
        basis = Atoms(elements='Al', positions=pos, cell=cell)
        self.assertIsInstance(basis, Atoms)
        basis = Atoms(elements=['Al'], positions=pos, cell=cell)
        self.assertIsInstance(basis, Atoms)
        self.assertRaises(ValueError,
                          Atoms,
                          symbols="Pt",
                          elements='Al',
                          positions=pos,
                          cell=cell)
        basis = Atoms(numbers=[13], positions=pos, cell=cell)
        self.assertEqual(basis.get_majority_species()[1], "Al")
        basis = Atoms(species=[el], indices=[0], positions=pos, cell=cell)
        self.assertEqual(basis.get_majority_species()[1], "Al")
        self.assertIsInstance(basis, Atoms)
        self.assertIsInstance(basis.info, dict)
        self.assertIsInstance(basis.arrays, dict)
        self.assertIsInstance(basis.adsorbate_info, dict)
        self.assertIsInstance(basis.units, dict)
        self.assertIsInstance(basis.pbc, (bool, list, np.ndarray))
        self.assertIsInstance(basis.indices, np.ndarray)
        self.assertIsInstance(basis.species, list)
        self.assertIsInstance(basis.cell, np.ndarray)
        self.assertIsInstance(basis._internal_positions, np.ndarray)
        self.assertIsInstance(basis.positions, np.ndarray)
        self.assertIsInstance(basis.scaled_positions, np.ndarray)
        self.assertIsInstance(basis.elements, np.ndarray)

    def test_set_species(self):
        pos, cell = generate_fcc_lattice()
        pse = PeriodicTable()
        el = pse.element("Pt")
        basis = Atoms(symbols='Al', positions=pos, cell=cell)
        self.assertEqual(basis.get_chemical_formula(), "Al")
        basis.set_species([el])
        self.assertEqual(basis.get_chemical_formula(), "Pt")
        self.assertTrue("Al" not in [sp.Abbreviation]
                        for sp in basis._species_to_index_dict.keys())
        self.assertTrue("Pt" in [sp.Abbreviation]
                        for sp in basis._species_to_index_dict.keys())

    def test_new_array(self):
        pos, cell = generate_fcc_lattice()
        basis = Atoms(symbols='Al', positions=pos, cell=cell)
        basis.set_repeat([10, 10, 10])
        spins = np.ones(len(basis))
        basis.new_array(name="spins", a=spins)
        self.assertTrue(np.array_equal(basis.arrays['spins'], spins))

    def test_set_array(self):
        pos, cell = generate_fcc_lattice()
        basis = Atoms(symbols='Al', positions=pos, cell=cell)
        basis.set_repeat([10, 10, 10])
        spins = np.ones(len(basis), dtype=float)
        basis.set_array(name="spins", a=2 * spins, dtype=int)
        self.assertTrue(np.array_equal(basis.arrays['spins'], 2 * spins))

    def test_get_array(self):
        pos, cell = generate_fcc_lattice()
        basis = Atoms(symbols='Al', positions=pos, cell=cell)
        basis.set_repeat([10, 10, 10])
        spins = np.ones(len(basis), dtype=float)
        basis.set_array(name="spins", a=2 * spins, dtype=int)
        self.assertTrue(np.array_equal(basis.arrays['spins'], 2 * spins))
        self.assertTrue(
            np.array_equal(basis.get_array(name="spins"), 2 * spins))

    def test_add_tags(self):
        self.CO2.add_tag(test_tag="a")
        self.assertIsInstance(self.CO2.test_tag, SparseList)
        self.assertEqual(self.CO2.test_tag[0], "a")
        self.assertEqual(self.CO2.test_tag[0], self.CO2.test_tag[2])
        self.assertIsInstance(self.CO2.test_tag.list(), list)
        self.CO2.add_tag(selective_dynamics=[True, True, True])
        self.CO2.selective_dynamics[1] = [True, False, True]
        self.assertEqual(self.CO2.selective_dynamics[1], [True, False, True])
        self.assertIsInstance(self.CO2.selective_dynamics.list(), list)

    def test_get_tags(self):
        self.CO2.add_tag(test_tag="a")
        self.assertIsInstance(self.CO2.test_tag, SparseList)
        self.assertIsInstance(self.CO2.get_tags(), type(dict().keys()))

    def test_get_pbc(self):
        self.assertTrue(np.array_equal(self.CO2.pbc, self.CO2.get_pbc()))
        self.assertEqual(len(self.CO2.get_pbc()), 3)

    def test_set_pbc(self):
        self.CO2.set_pbc(value=[True, True, False])
        self.assertTrue(np.array_equal(self.CO2.pbc, self.CO2.get_pbc()))
        self.assertTrue(np.array_equal([True, True, False],
                                       self.CO2.get_pbc()))
        self.CO2.set_pbc(value=False)
        self.assertTrue(
            np.array_equal([False, False, False], self.CO2.get_pbc()))
        self.assertTrue(np.array_equal(self.CO2.pbc, self.CO2.get_pbc()))

    def test_chemical_element(self):
        self.assertIsInstance(self.CO2.convert_element('C'), ChemicalElement)
        self.assertEqual(len(self.CO2.species), 2)

    def test_copy(self):
        pos, cell = generate_fcc_lattice()
        basis = Atoms(symbols='Al', positions=pos, cell=cell)
        basis_copy = basis.copy()
        self.assertEqual(basis, basis_copy)
        basis_copy[:] = "Pt"
        self.assertNotEqual(basis, basis_copy)

    def test_numbers_to_elements(self):
        num_list = [1, 12, 13, 6]
        self.assertTrue(
            np.array_equal([
                el.Abbreviation
                for el in self.CO2.numbers_to_elements(num_list)
            ], ['H', 'Mg', 'Al', 'C']))

    def test_to_hdf(self):
        if sys.version_info[0] >= 3:
            filename = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                    "../../static/atomistics/test_hdf")
            abs_filename = os.path.abspath(filename)
            hdf_obj = FileHDFio(abs_filename)
            pos, cell = generate_fcc_lattice()
            basis = Atoms(symbols='Al', positions=pos, cell=cell)
            basis.set_repeat([2, 2, 2])
            basis.to_hdf(hdf_obj, "test_structure")
            self.assertTrue(
                np.array_equal(hdf_obj["test_structure/positions"],
                               basis.positions))
            basis_new = Atoms().from_hdf(hdf_obj, "test_structure")
            self.assertEqual(basis, basis_new)

    def test_from_hdf(self):
        if sys.version_info[0] >= 3:
            filename = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                    "../../static/atomistics/test_hdf")
            abs_filename = os.path.abspath(filename)
            hdf_obj = FileHDFio(abs_filename)
            pos, cell = generate_fcc_lattice()
            basis_store = Atoms(symbols='Al', positions=pos, cell=cell)
            basis_store.set_repeat([2, 2, 2])
            basis_store.to_hdf(hdf_obj, "simple_structure")
            basis = Atoms().from_hdf(hdf_obj, group_name="simple_structure")
            self.assertEqual(len(basis), 8)
            self.assertEqual(basis.get_majority_species()[1], "Al")
            self.assertEqual(basis.get_spacegroup()['Number'], 225)

    def create_Fe_bcc(self):
        self.pse = PeriodicTable()
        self.pse.add_element("Fe", "Fe_up", spin="up", pseudo_name='GGA')
        self.pse.add_element("Fe", "Fe_down", spin="down", pseudo_name='GGA')
        Fe_up = self.pse.element("Fe_up")
        Fe_down = self.pse.element("Fe_down")
        self.Fe_bcc = Atoms([Fe_up, Fe_down],
                            scaled_positions=[[0, 0, 0], [0.25, 0.25, 0.25]],
                            cell=np.identity(3))
        self.Fe_bcc.add_tag("group")
        self.Fe_bcc.group[:] = 0

    def test_convert_formula(self):
        self.assertEqual(self.CO2.convert_formula('C'), ['C'])
        self.assertEqual(self.CO2.convert_formula('C3'), ['C', 'C', 'C'])
        self.assertEqual(self.CO2.convert_formula('CO2'), ['C', 'O', 'O'])
        self.assertEqual(self.CO2.convert_formula('CO2Fe'),
                         ['C', 'O', 'O', 'Fe'])
        self.assertEqual(self.CO2.convert_formula('CO2FeF21'),
                         ['C', 'O', 'O', 'Fe', 'F', 'F'])

    def test__getitem__(self):
        self.assertEqual(self.CO2[0].symbol, 'C')
        self.assertEqual(self.C3[2].position.tolist(), [0, 2, 0])
        self.assertTrue((self.C3[1:].positions == np.array([[0, 0, 2],
                                                            [0, 2, 0]])).all())
        short_basis = self.CO2[0]
        self.assertIsInstance(short_basis, Atom)
        short_basis = self.CO2[[0]]
        self.assertIsInstance(short_basis, Atoms)
        self.assertEqual(short_basis.indices[0], 0)
        self.assertEqual(len(short_basis.species), 1)
        short_basis = self.CO2[[2]]
        self.assertIsInstance(short_basis, Atoms)
        self.assertEqual(short_basis.indices[0], 0)
        self.assertEqual(len(short_basis.species), 1)
        basis_Mg = CrystalStructure("Mg",
                                    bravais_basis="fcc",
                                    lattice_constant=4.2)
        basis_O = CrystalStructure("O",
                                   bravais_basis="fcc",
                                   lattice_constant=4.2)
        basis_O.positions += [0., 0., 0.5]
        basis = basis_Mg + basis_O
        basis.center_coordinates_in_unit_cell()
        basis.set_repeat([3, 3, 3])
        mg_indices = basis.select_index("Mg")
        o_indices = basis.select_index("O")
        basis_new = basis[mg_indices] + basis[o_indices]
        self.assertEqual(len(basis_new._tag_list),
                         len(basis[mg_indices]) + len(basis[o_indices]))
        self.assertEqual(basis_new.get_spacegroup()["Number"], 225)

    def test_positions(self):
        self.assertEqual(self.CO2[1:].positions[1:].tolist(),
                         [[0.0, 1.5, 0.0]])
        self.CO2.positions[1][0] = 5.
        self.assertEqual(self.CO2.positions[1].tolist(), [5.0, 0, 1.5])

    def test_set_positions(self):
        pos, cell = generate_fcc_lattice()
        basis = Atoms(symbols='Al', positions=pos, cell=cell)
        basis.set_positions(np.array([[2.5, 2.5, 2.5]]))
        self.assertTrue(np.array_equal(basis.positions, [[2.5, 2.5, 2.5]]))

    def test_set_scaled_positions(self):
        pos, cell = generate_fcc_lattice()
        basis = Atoms(symbols='Al', positions=pos, cell=cell, a=4.2)
        basis.set_scaled_positions(np.array([[0.5, 0.5, 0.5]]))
        self.assertTrue(
            np.array_equal(basis.scaled_positions, [[0.5, 0.5, 0.5]]))
        self.assertTrue(
            np.array_equal(basis.positions,
                           np.dot([[0.5, 0.5, 0.5]], basis.cell)))

    def test_cell(self):
        CO = Atoms("CO",
                   positions=[[0, 0, 0], [0, 0, 2]],
                   cell=[[1, 0, 0], [0, 1, 0], [0, 0, 1]],
                   pbc=[True, True, True])
        self.assertTrue((CO.get_cell() == np.identity(3)).all())
        self.assertTrue((CO.cell == np.identity(3)).all())
        CO.cell[2][2] = 10.
        self.assertTrue(CO.cell[2, 2] == 10.)

    def test_add(self):
        COX = self.C2 + Atom("O", position=[0, 0, -2])
        COX += Atom("O", position=[0, 0, -4])
        COX += COX
        n_objects = len(set(COX.get_species_objects()))
        n_species = len(set(COX.get_chemical_elements()))
        self.assertEqual(n_objects, n_species)

    def test_pbc(self):
        CO = Atoms("CO",
                   positions=[[0, 0, 0], [0, 0, 2]],
                   cell=[[1, 0, 0], [0, 1, 0], [0, 0, 1]],
                   pbc=[True, True, True])
        self.assertTrue((CO.pbc == np.array([True, True, True])).all())
        CO.set_pbc((True, True, False))

    def test_get_masses_DOF(self):
        self.assertEqual(len(self.CO2.get_masses_dof()),
                         len(self.CO2.positions.flatten()))

    def test_get_parent_basis(self):
        periodic_table = PeriodicTable()
        periodic_table.add_element(parent_element="O", new_element="O_up")
        O_up = periodic_table.element("O_up")

        O_basis = Atoms([O_up],
                        cell=10.0 * np.eye(3),
                        scaled_positions=[[0.5, 0.5, 0.5]])
        O_simple = Atoms(["O"],
                         cell=10.0 * np.eye(3),
                         scaled_positions=[[0.5, 0.5, 0.5]])
        O_parent = O_basis.get_parent_basis()
        self.assertNotEqual(O_basis, O_parent)
        self.assertEqual(O_simple, O_parent)
        self.assertEqual(O_parent[0].symbol, "O")
        periodic_table.add_element(parent_element="O", new_element="O_down")
        O_down = periodic_table.element("O_down")
        O_basis = Atoms([O_up, O_down],
                        cell=10.0 * np.eye(3),
                        scaled_positions=[[0.5, 0.5, 0.5], [0, 0, 0]])
        O_simple = Atoms(["O", "O"],
                         cell=10.0 * np.eye(3),
                         scaled_positions=[[0.5, 0.5, 0.5]])
        O_parent = O_basis.get_parent_basis()
        self.assertNotEqual(O_basis, O_parent)
        self.assertEqual(O_simple, O_parent)
        self.assertEqual(O_parent.get_chemical_formula(), "O2")
        self.assertEqual(len(O_basis.species), 2)
        self.assertEqual(len(O_simple.species), 1)
        self.assertEqual(len(O_parent.species), 1)

    def test_profiling(self):
        num = 1000
        C100 = Atoms(num * ["C"], positions=[(0, 0, 0) for _ in range(num)])
        self.assertEqual(len(C100), num)

    def test_Au(self):
        a = 4.05  # Gold lattice constant
        b = a / 2.
        fcc = Atoms(['Au'], cell=[(0, b, b), (b, 0, b), (b, b, 0)], pbc=True)
        # print fcc
        # print "volume: ", fcc.get_volume()

    def test_set_absolute(self):
        a = 4.05  # Gold lattice constant
        b = a / 2.
        positions = np.array([(0.5, 0.4, 0.)])
        fcc = Atoms(symbols=['Au'],
                    scaled_positions=positions,
                    cell=[(0, b, b), (b, 0, b), (b, b, 0)],
                    pbc=True)
        # fcc.set_absolute()
        # print fcc.positions
        # fcc.set_relative()
        self.assertTrue(
            np.linalg.norm(fcc.scaled_positions - positions) < 1e-10)

    def test_repeat(self):
        basis_Mg = CrystalStructure("Mg",
                                    bravais_basis="fcc",
                                    lattice_constant=4.2)
        basis_O = CrystalStructure("O",
                                   bravais_basis="fcc",
                                   lattice_constant=4.2)
        basis_O.scaled_positions += [0., 0., 0.5]
        basis = basis_Mg + basis_O
        basis.center_coordinates_in_unit_cell()
        basis.add_tag(selective_dynamics=[True, True, True])
        basis.selective_dynamics[basis.select_index("O")] = [
            False, False, False
        ]
        len_before = len(basis)
        sel_dyn_before = np.array(basis.selective_dynamics.list())
        self.assertTrue(
            np.alltrue(
                np.logical_not(
                    np.alltrue(sel_dyn_before[basis.select_index("O")],
                               axis=1))))
        self.assertTrue(
            np.alltrue(
                np.alltrue(sel_dyn_before[basis.select_index("Mg")], axis=1)))
        basis.set_repeat([3, 3, 2])
        sel_dyn_after = np.array(basis.selective_dynamics.list())
        len_after = len(basis)
        self.assertEqual(basis.get_spacegroup()["Number"], 225)
        self.assertEqual(len_before * 18, len_after)
        self.assertEqual(len(sel_dyn_before) * 18, len(sel_dyn_after))
        self.assertTrue(
            np.alltrue(
                np.logical_not(
                    np.alltrue(sel_dyn_after[basis.select_index("O")],
                               axis=1))))
        self.assertTrue(
            np.alltrue(
                np.alltrue(sel_dyn_after[basis.select_index("Mg")], axis=1)))

    def test_boundary(self):
        cell = 2.2 * np.identity(3)
        NaCl = Atoms('NaCl',
                     scaled_positions=[(0, 0, 0), (0.5, 0.5, 0.5)],
                     cell=cell)
        NaCl.set_repeat([3, 3, 3])
        # NaCl.plot3d()
        NaCl_bound = NaCl.get_boundary_region(0.2)
        # NaCl_bound.plot3d()

    def test_get_distance(self):
        cell = 2.2 * np.identity(3)
        NaCl = Atoms('NaCl',
                     scaled_positions=[(0, 0, 0), (0.5, 0.5, 0.5)],
                     cell=cell)
        self.assertAlmostEqual(NaCl.get_distance(0, 1), 2.2 * 0.5 * np.sqrt(3))
        self.assertAlmostEqual(NaCl.get_distance(0, [0, 0, 0.5]), 0.5)
        self.assertAlmostEqual(NaCl.get_distance([0, 0, 0], [0, 0, 0.5]), 0.5)

    def test_get_neighbors(self):
        cell = 2.2 * np.identity(3)
        NaCl = Atoms('NaCl',
                     scaled_positions=[(0, 0, 0), (0.5, 0.5, 0.5)],
                     cell=cell)
        # NaCl.repeat([3, 3, 3])
        # NaCl.positions = [(1,1,1)]
        boundary = NaCl.get_boundary_region(3.5)
        extended_cell = NaCl + boundary
        # extended_cell.plot3d()
        nbr_dict = NaCl.get_neighbors(num_neighbors=12, t_vec=True)
        # print nbr_dict.distances
        # print [set(s) for s in nbr_dict.shells]

    def test_center_coordinates(self):
        cell = 2.2 * np.identity(3)
        NaCl = Atoms('NaCl',
                     scaled_positions=[(0, 0, 0), (0.5, 0.5, 0.5)],
                     cell=cell)
        NaCl.set_repeat([3, 3, 3])
        NaCl.positions += [2.2, 2.2, 2.2]
        NaCl.center_coordinates_in_unit_cell(origin=-0.5)
        self.assertTrue(-0.5 < np.min(NaCl.scaled_positions))
        self.assertTrue(np.max(NaCl.scaled_positions < 0.5))
        NaCl.center_coordinates_in_unit_cell(origin=0.)
        self.assertTrue(0 <= np.min(NaCl.positions))
        self.assertTrue(np.max(NaCl.scaled_positions < 1))

    def test_get_shells(self):
        dim = 3
        cell = 2.2 * np.identity(dim)
        Al_sc = Atoms('AlAl',
                      scaled_positions=[(0, 0, 0), (0.5, 0.5, 0.5)],
                      cell=cell)
        Al_sc.set_repeat([3, 3, 3])
        self.assertEqual(np.round(Al_sc.get_shells()[2], 6), 2.2)

    def test_get_shell_matrix(self):
        basis = Atoms('FeFe',
                      scaled_positions=[(0, 0, 0), (0.5, 0.5, 0.5)],
                      cell=np.identity(3))
        output = basis.get_shell_matrix(shell=1)
        self.assertIsInstance(output, np.ndarray)
        self.assertEqual(np.sum(output), 16)
        self.assertTrue(np.all(np.dot(output, output) == np.identity(2) * 64))

    def test_get_distance_matrix(self):
        basis = Atoms('FeFe',
                      scaled_positions=[(0, 0, 0), (0.5, 0.5, 0.5)],
                      cell=np.identity(3))
        output = basis.get_distance_matrix()
        self.assertIsInstance(output, np.ndarray)
        output = np.rint(output * 2 / np.sqrt(3))
        self.assertTrue(np.all(np.dot(output, output) == np.identity(2)))

    def test_cluster_analysis(self):
        import random
        cell = 2.2 * np.identity(3)
        Al_sc = Atoms(elements=['Al', 'Al'],
                      scaled_positions=[(0, 0, 0), (0.5, 0.5, 0.5)],
                      cell=cell)
        Al_sc.set_repeat([4, 4, 4])
        radius = Al_sc.get_shell_radius()
        neighbors = Al_sc.get_neighbors(radius=radius,
                                        num_neighbors=100,
                                        t_vec=False,
                                        exclude_self=True)

        c_Zn = 0.1
        pse = PeriodicTable()
        Zn = pse.element("Zn")
        random.seed(123456)
        for _ in range(1):
            Zn_ind = random.sample(range(len(Al_sc)), int(c_Zn * len(Al_sc)))
            # for i_Zn in Zn_ind:
            #     Al_sc.elements[i_Zn] = Zn

            cluster = Al_sc.cluster_analysis(Zn_ind, neighbors)
            cluster_len = np.sort([len(v) for k, v in cluster.items()])
            # print np.histogram(cluster_len), np.sum(cluster_len), len(Zn_ind)
            # for key, value in cluster.items():
            #     el = pse.Element((key % 100) + 1)
            #     for i_el in value:
            #         Al_sc.elements[i_el] = el
            # Al_sc.plot3d()

    def test_get_bonds(self):
        dim = 3
        cell = 2.62 * np.identity(dim)
        d1, d2 = 0.6, 0.6
        H2O = Atoms('H2O',
                    scaled_positions=[(d1, d2, 0), (d1, -d2, 0), (0, 0, 0)],
                    cell=cell)
        H2O.set_repeat([1, 1, 3])
        # H2O.plot3d(show_bonds=True) #, bond_stretch=2)
        # print H2O.get_bonds(radius=2.)[0]
        # print np.sum(H2O.get_masses())/H2O.get_volume()

    def test_get_symmetry(self):
        cell = 2.2 * np.identity(3)
        Al = Atoms('AlAl', positions=[(0, 0, 0), (0.5, 0.5, 0.5)],
                   cell=cell).repeat(2)
        self.assertEqual(len(set(Al.get_symmetry()['equivalent_atoms'])), 1)
        self.assertEqual(len(Al.get_symmetry()['translations']), 96)
        self.assertEqual(len(Al.get_symmetry()['translations']),
                         len(Al.get_symmetry()['rotations']))

    def _get_voronoi_vertices(self):
        cell = 2.2 * np.identity(3)
        Al = Atoms('AlAl',
                   scaled_positions=[(0, 0, 0), (0.5, 0.5, 0.5)],
                   cell=cell)
        pos, box = Al._get_voronoi_vertices()
        self.assertEqual(len(pos), 14)

    def get_equivalent_voronoi_vertices(self):
        cell = 2.2 * np.identity(3)
        Al = Atoms('AlAl', positions=[(0, 0, 0), (0.5, 0.5, 0.5)],
                   cell=cell).repeat(2)
        pos, box = Al._get_voronoi_vertices()
        self.assertEqual(len(Al), 69)
        self.assertEqual(len(len(Al.get_species_symbols())), 2)
        Al = Atoms('AlAl',
                   scaled_positions=[(0, 0, 0), (0.5, 0.5, 0.5)],
                   cell=cell).repeat(2)
        pos = Al.get_equivalent_voronoi_vertices()
        self.assertEqual(len(pos), 1)

    def test_get_parent_symbols(self):
        self.assertTrue(
            np.array_equal(self.CO2.get_parent_symbols(), ["C", "O", "O"]))
        self.assertTrue(
            np.array_equal(self.CO2.get_parent_symbols(),
                           self.CO2.get_chemical_symbols()))
        cell = np.eye(3) * 10.0
        pse = PeriodicTable()
        pse.add_element("O", "O_up", spin="up")
        o_up = pse.element("O_up")
        basis = Atoms([o_up], scaled_positions=[[0.27, 0.27, 0.27]], cell=cell)
        self.assertTrue(np.array_equal(basis.get_parent_symbols(), ["O"]))
        self.assertFalse(
            np.array_equal(basis.get_parent_symbols(),
                           basis.get_chemical_symbols()))

    def test_get_chemical_symbols(self):
        self.assertTrue(
            np.array_equal(self.CO2.get_chemical_symbols(), ["C", "O", "O"]))
        cell = np.eye(3) * 10.0
        pse = PeriodicTable()
        pse.add_element("O", "O_up", spin="up")
        o_up = pse.element("O_up")
        basis = Atoms([o_up], scaled_positions=[[0.27, 0.27, 0.27]], cell=cell)
        self.assertTrue(np.array_equal(basis.get_chemical_symbols(), ["O_up"]))

    def test_get_symmetry_dataset(self):
        cell = 2.2 * np.identity(3)
        Al_sc = Atoms('AlAl',
                      scaled_positions=[(0, 0, 0), (0.5, 0.5, 0.5)],
                      cell=cell)
        Al_sc.set_repeat([2, 2, 2])
        self.assertEqual(Al_sc.get_symmetry_dataset()['number'], 229)

    def test_get_space_group(self):
        cell = 2.2 * np.identity(3)
        Al_sc = Atoms('AlAl',
                      scaled_positions=[(0, 0, 0), (0.5, 0.5, 0.5)],
                      cell=cell)
        self.assertEqual(Al_sc.get_spacegroup()['InternationalTableSymbol'],
                         'Im-3m')
        self.assertEqual(Al_sc.get_spacegroup()['Number'], 229)
        cell = 4.2 * (0.5 * np.ones((3, 3)) - 0.5 * np.eye(3))
        Al_fcc = Atoms('Al', scaled_positions=[(0, 0, 0)], cell=cell)
        self.assertEqual(Al_fcc.get_spacegroup()['InternationalTableSymbol'],
                         'Fm-3m')
        self.assertEqual(Al_fcc.get_spacegroup()['Number'], 225)
        a = 3.18
        c = 1.623 * a
        cell = np.eye(3)
        cell[0, 0] = a
        cell[2, 2] = c
        cell[1, 0] = -a / 2.
        cell[1, 1] = np.sqrt(3) * a / 2.
        pos = np.array([[0., 0., 0.], [1. / 3., 2. / 3., 1. / 2.]])
        Mg_hcp = Atoms('Mg2', scaled_positions=pos, cell=cell)
        self.assertEqual(Mg_hcp.get_spacegroup()['Number'], 194)
        cell = np.eye(3)
        cell[0, 0] = a
        cell[2, 2] = c
        cell[1, 1] = np.sqrt(3) * a
        pos = np.array([[0., 0., 0.], [0.5, 0.5, 0.], [0.5, 0.16666667, 0.5],
                        [0., 0.66666667, 0.5]])
        Mg_hcp = Atoms('Mg4', scaled_positions=pos, cell=cell)
        self.assertEqual(Mg_hcp.get_spacegroup()['Number'], 194)

    def test_get_primitive_cell(self):
        cell = 2.2 * np.identity(3)
        Al_sc = Atoms('AlFe',
                      scaled_positions=[(0, 0, 0), (0.5, 0.5, 0.5)],
                      cell=cell)
        Al_sc.set_repeat([2, 2, 2])
        primitive_cell = Al_sc.get_primitive_cell()
        self.assertEqual(primitive_cell.get_spacegroup()['Number'], 221)

    def test_get_ir_reciprocal_mesh(self):
        cell = 2.2 * np.identity(3)
        Al_sc = Atoms('AlAl',
                      scaled_positions=[(0, 0, 0), (0.5, 0.5, 0.5)],
                      cell=cell)
        self.assertEqual(len(Al_sc.get_ir_reciprocal_mesh([3, 3, 3])[0]), 27)

    def test_get_number_species_atoms(self):
        self.assertEqual(list(self.CO2.get_number_species_atoms().values()),
                         [1, 2])

    def test_get_chemical_formula(self):
        self.assertEqual(self.CO2.get_chemical_formula(), "CO2")

    def test_get_equivalent_atoms(self):
        cell = 2.2 * np.identity(3)
        Al_sc = Atoms('AlFe',
                      scaled_positions=[(0, 0, 0), (0.5, 0.5, 0.5)],
                      cell=cell)
        Al_sc.set_repeat([2, 2, 2])

    def test_center(self):
        old_pos = self.CO2.positions.copy()
        self.CO2.center(vacuum=5)
        new_array = old_pos + 5 * np.ones(3)
        self.assertTrue(np.array_equal(self.CO2.positions, new_array))

    def test_get_positions(self):
        basis_Mg = CrystalStructure("Mg",
                                    bravais_basis="fcc",
                                    lattice_constant=4.2)
        self.assertTrue(
            np.array_equal(basis_Mg.positions, basis_Mg.get_positions()))

    def test_get_scaled_positions(self):
        basis_Mg = CrystalStructure("Mg",
                                    bravais_basis="fcc",
                                    lattice_constant=4.2)
        self.assertTrue(
            np.array_equal(basis_Mg.scaled_positions,
                           basis_Mg.get_scaled_positions()))

    def test_occupy_lattice(self):
        basis_Mg = CrystalStructure("Mg",
                                    bravais_basis="fcc",
                                    lattice_constant=4.2)
        basis_O = CrystalStructure("O",
                                   bravais_basis="fcc",
                                   lattice_constant=4.2)
        basis_O.scaled_positions += [0., 0., 0.5]
        basis = basis_Mg + basis_O
        basis.center_coordinates_in_unit_cell()
        orig_basis = basis.copy()
        self.assertEqual(basis.get_chemical_formula(), "Mg4O4")
        Mg_indices = basis.select_index("Mg")
        O_indices = basis.select_index("O")
        basis.occupy_lattice(Na=Mg_indices)
        self.assertEqual(basis.get_chemical_formula(), "Na4O4")
        basis.occupy_lattice(Cl=O_indices)
        self.assertEqual(basis.get_chemical_formula(), "Cl4Na4")
        self.assertTrue(np.array_equal(basis.select_index("Na"), Mg_indices))
        self.assertTrue(np.array_equal(basis.select_index("Cl"), O_indices))
        orig_basis.set_repeat([2, 2, 2])
        Mg_indices = orig_basis.select_index("Mg")
        O_indices = orig_basis.select_index("O")
        orig_basis.occupy_lattice(Cl=O_indices, Na=Mg_indices)
        self.assertEqual(orig_basis.get_chemical_formula(), "Cl32Na32")
        orig_basis.occupy_lattice(H=O_indices[0])
        self.assertEqual(orig_basis.get_chemical_formula(), "Cl31HNa32")

    def test_select_index(self):
        self.assertTrue(np.array_equal(self.CO2.select_index("C"), [0]))
        self.assertTrue(np.array_equal(self.CO2.select_index("O"), [1, 2]))

    def test_parent_index(self):
        basis_Mg = CrystalStructure("Mg",
                                    bravais_basis="fcc",
                                    lattice_constant=4.2)
        basis_O = CrystalStructure("O",
                                   bravais_basis="fcc",
                                   lattice_constant=4.2)
        basis_O.positions += [0., 0., 0.5]
        basis = basis_Mg + basis_O
        basis.center_coordinates_in_unit_cell()
        basis.set_repeat([2, 2, 2])
        o_indices = basis.select_index("O")
        pse = PeriodicTable()
        pse.add_element("O", "O_up", spin="up")
        o_up = pse.element("O_up")
        basis[o_indices] = o_up
        self.assertTrue(np.array_equal(o_indices, basis.select_index(o_up)))
        self.assertEqual(len(basis.select_index("O")), 0)
        self.assertTrue(
            np.array_equal(o_indices, basis.select_parent_index("O")))

    def test__eq__(self):
        test_basis = self.CO2.copy()
        self.assertEqual(test_basis, self.CO2)
        test_basis.positions[2] += 0.0
        self.assertEqual(test_basis, self.CO2)
        self.assertNotEqual(self.C2, self.CO2)

    def test__add__(self):
        cell = np.eye(3) * 10.0
        basis_0 = Atoms(["O"], scaled_positions=[[0.5, 0.5, 0.5]], cell=cell)
        basis_1 = Atoms(["H"],
                        scaled_positions=[[0.75, 0.75, 0.75]],
                        cell=cell)
        basis_2 = Atoms(["H"],
                        scaled_positions=[[0.25, 0.25, 0.25]],
                        cell=cell)
        basis_3 = Atoms(["H", "O", "N"],
                        scaled_positions=[[0.35, 0.35, 0.35], [0., 0., 0.],
                                          [0., 0., 0.1]],
                        cell=cell)

        pse = PeriodicTable()
        pse.add_element("O", "O_up", spin="up")
        o_up = pse.element("O_up")
        basis_4 = Atoms([o_up],
                        scaled_positions=[[0.27, 0.27, 0.27]],
                        cell=np.eye(3) * 20.0)
        b = basis_0 + basis_1
        self.assertEqual(b.get_chemical_formula(), "HO")
        b = basis_0 + basis_1 + basis_2
        self.assertEqual(b.get_chemical_formula(), "H2O")
        b += basis_2
        self.assertEqual(b.get_chemical_formula(), "H3O")
        b = basis_0 + basis_1 + basis_2 + basis_3
        self.assertEqual(b.get_chemical_formula(), "H3NO2")
        self.assertTrue(
            np.array_equal(b.scaled_positions[b.select_index("N")],
                           [[0., 0., 0.1]]))
        self.assertTrue(
            np.allclose(
                b.scaled_positions[b.select_index("H")],
                [[0.75, 0.75, 0.75], [0.25, 0.25, 0.25], [0.35, 0.35, 0.35]]))
        self.assertTrue(
            np.allclose(b.scaled_positions[b.select_index("O")],
                        [[0.5, 0.5, 0.5], [0., 0., 0.]]))
        b.set_repeat([2, 2, 2])
        self.assertEqual(b.get_chemical_formula(), "H24N8O16")
        b += basis_4
        self.assertEqual(b.get_chemical_formula(), "H24N8O16O_up")
        self.assertTrue(
            np.allclose(b.scaled_positions[b.select_index(o_up)],
                        [[0.27, 0.27, 0.27]]))
        COX = self.C2 + Atom("O", position=[0, 0, -2])
        COX += Atom("O", position=[0, 0, -4])
        COX += COX
        n_objects = len(set(COX.get_species_objects()))
        n_species = len(set(COX.get_chemical_elements()))
        self.assertEqual(n_objects, n_species)
        self.assertEqual(n_objects, 2)
        self.assertEqual(n_species, 2)
        basis_Mg = CrystalStructure("Mg",
                                    bravais_basis="fcc",
                                    lattice_constant=4.2)
        basis_O = CrystalStructure("O",
                                   bravais_basis="fcc",
                                   lattice_constant=4.2)
        # basis_O.set_relative()
        basis_O.scaled_positions += [0., 0., 0.5]
        basis = basis_Mg + basis_O
        self.assertEqual(len(basis._tag_list),
                         len(basis_Mg._tag_list) + len(basis_O._tag_list))
        basis.center_coordinates_in_unit_cell()
        self.assertEqual(basis.get_spacegroup()["Number"], 225)

    def test__delitem__(self):
        cell = np.eye(3) * 10.0
        basis_0 = Atoms(["O"], scaled_positions=[[0.5, 0.5, 0.5]], cell=cell)
        basis_1 = Atoms(["H"],
                        scaled_positions=[[0.75, 0.75, 0.75]],
                        cell=cell)
        basis_2 = Atoms(["H"],
                        scaled_positions=[[0.25, 0.25, 0.25]],
                        cell=cell)
        basis_3 = Atoms(["H", "O", "N"],
                        scaled_positions=[[0.35, 0.35, 0.35], [0., 0., 0.],
                                          [0., 0., 0.1]],
                        cell=cell)

        pse = PeriodicTable()
        pse.add_element("O", "O_up", spin="up")
        o_up = pse.element("O_up")
        basis_4 = Atoms([o_up],
                        scaled_positions=[[0.27, 0.27, 0.27]],
                        cell=cell)
        b = basis_0 + basis_1 + basis_2 + basis_3 + basis_4
        O_indices = b.select_index("O")
        self.assertEqual(len(b), 7)
        self.assertEqual(len(b.indices), 7)
        self.assertEqual(len(b.species), 4)
        b.__delitem__(O_indices[0])
        self.assertEqual(b.get_chemical_formula(), "H3NOO_up")
        self.assertEqual(len(b), 6)
        self.assertEqual(len(b.indices), 6)
        self.assertEqual(len(b._tag_list), 6)
        self.assertEqual(len(b.species), 4)
        O_indices = b.select_index("O")
        b.__delitem__(O_indices)
        self.assertEqual(b.get_chemical_formula(), "H3NO_up")
        self.assertEqual(len(b), 5)
        self.assertEqual(len(b.indices), 5)
        self.assertEqual(len(b.species), 3)
        self.assertEqual(np.max(b.indices), 2)
        N_indices = b.select_index("N")
        b.__delitem__(N_indices)
        self.assertEqual(b.get_chemical_formula(), "H3O_up")
        self.assertEqual(len(b), 4)
        self.assertEqual(len(b.indices), 4)
        self.assertEqual(len(b.species), 2)
        self.assertEqual(np.max(b.indices), 1)
        O_indices = b.select_index(o_up)
        b.__delitem__(O_indices)
        self.assertEqual(b.get_chemical_formula(), "H3")
        self.assertEqual(len(b), 3)
        self.assertEqual(len(b.indices), 3)
        self.assertEqual(len(b.species), 1)
        self.assertEqual(np.max(b.indices), 0)

    def test__setitem__(self):
        basis = self.CO2.copy()
        basis[0] = 'H'
        basis[1] = 'H'
        self.assertEqual(basis.get_chemical_formula(), "H2O")
        self.assertEqual(len(basis.species), 2)
        self.assertEqual(len(basis.get_species_symbols()), 2)
        basis = self.CO2.copy()
        basis[0] = 'H'
        basis[np.int64(0)] = 'H'
        self.assertEqual(basis.get_chemical_formula(), "HO2")
        self.assertEqual(len(basis.species), 2)
        self.assertEqual(len(basis.get_species_symbols()), 2)
        basis[0] = 'O'
        self.assertEqual(basis.get_chemical_formula(), "O3")
        self.assertEqual(len(basis.species), 1)
        self.assertEqual(len(basis.get_species_symbols()), 1)
        basis = self.CO2.copy()
        basis[[2]] = 'N'
        self.assertEqual(basis.get_chemical_formula(), "CNO")
        self.assertEqual(len(basis.species), 3)
        self.assertEqual(len(basis.get_species_symbols()), 3)
        basis = self.CO2.copy()
        basis[[0]] = 'H'
        basis[np.array([0])] = 'H'
        self.assertEqual(basis.get_chemical_formula(), "HO2")
        self.assertEqual(len(basis.species), 2)
        self.assertEqual(len(basis.get_species_symbols()), 2)

        basis = self.CO2.copy()
        basis[[0]] = 'N'
        self.assertEqual(basis.get_chemical_formula(), "NO2")
        self.assertEqual(len(basis.species), 2)
        self.assertEqual(len(basis.get_species_symbols()), 2)
        basis[[0]] = 'O'
        self.assertEqual(basis.get_chemical_formula(), "O3")
        self.assertEqual(len(basis.species), 1)
        self.assertEqual(len(basis.get_species_symbols()), 1)
        basis[[0, 2]] = 'H'
        self.assertEqual(basis.get_chemical_formula(), "H2O")
        self.assertEqual(len(basis.species), 2)
        self.assertEqual(len(basis.get_species_symbols()), 2)
        pse = PeriodicTable()
        pse.add_element("O", "O_up", spin="up")
        o_up = pse.element("O_up")
        basis[[0, 2]] = o_up
        self.assertEqual(basis.get_chemical_formula(), "OO_up2")
        self.assertEqual(len(basis.species), 2)
        self.assertEqual(len(basis.get_species_symbols()), 2)
        basis[0:3] = "N"
        self.assertEqual(basis.get_chemical_formula(), "N3")
        self.assertEqual(len(basis.species), 1)
        self.assertEqual(len(basis.get_species_symbols()), 1)
        basis[:] = "Ne"
        self.assertEqual(basis.get_chemical_formula(), "Ne3")
        self.assertEqual(len(basis.species), 1)
        self.assertEqual(len(basis.get_species_symbols()), 1)
        basis[-2:] = "H"
        self.assertEqual(basis.get_chemical_formula(), "H2Ne")
        self.assertEqual(len(basis.species), 2)
        self.assertEqual(len(basis.get_species_symbols()), 2)
        basis[0:3] = "O"
        self.assertEqual(basis.get_chemical_formula(), "O3")
        self.assertEqual(len(basis.species), 1)
        self.assertEqual(len(basis.get_species_symbols()), 1)
Exemple #6
0
class TestVaspStructure(unittest.TestCase):
    """
    Testing routines in the vasp/structure module.
    """
    def setUp(self):
        self.file_location = os.path.dirname(os.path.abspath(__file__))
        poscar_directory = os.path.join(
            self.file_location, "../static/vasp_test_files/poscar_samples")
        file_list = os.listdir(poscar_directory)
        self.file_list = [
            posixpath.join(poscar_directory, f) for f in file_list
        ]
        atom_numbers = np.random.randint(low=1, high=99, size=(1, 3)).flatten()
        cell = 10.0 * np.eye(3)
        pos = 0.5 * np.ones((3, 3)) - 0.5 * np.eye(3)
        self.structure = Atoms(numbers=atom_numbers, cell=cell, positions=pos)
        self.assertIsInstance(self.structure, Atoms)
        self.structure.repeat([2, 2, 2])
        self.element_list = self.structure.get_chemical_elements()

    def test_atoms_from_string(self):
        for poscar_file in self.file_list:
            with open(poscar_file, "r") as f:
                lines = f.readlines()
                atoms = atoms_from_string(string=lines)
                self.assertIsInstance(atoms, Atoms)

    def test_read_atoms(self):
        for f in self.file_list:
            atoms = read_atoms(filename=f)
            self.assertIsInstance(atoms, Atoms)
            if f.split("/")[-1] == "POSCAR_1":
                self.assertEqual(len(atoms), 744)
                self.assertEqual(len(atoms.select_index("H")), 432)
                self.assertEqual(len(atoms.select_index("O")), 216)
                self.assertEqual(len(atoms.select_index("Mg")), 96)
            if f.split("/")[-1] == "POSCAR_scaled":
                self.assertEqual(len(atoms), 256)
                self.assertEqual(len(atoms.select_index("Cu")), 256)
                cell = np.eye(3) * 4. * 3.63
                self.assertTrue(np.array_equal(atoms.cell, cell))
                self.assertEqual(atoms.get_spacegroup()["Number"], 225)
            if f.split("/")[-1] == "POSCAR_volume_scaled":
                self.assertEqual(len(atoms), 256)
                self.assertEqual(len(atoms.select_index("Cu")), 256)
                cell = np.eye(3) * 4. * 3.63
                self.assertTrue(np.array_equal(atoms.cell, cell))
                self.assertEqual(atoms.get_spacegroup()["Number"], 225)
            if f.split("/")[-1] == "POSCAR_random":
                self.assertEqual(len(atoms), 33)
                self.assertEqual(len(atoms.selective_dynamics), 33)
                self.assertEqual(len(atoms.select_index("Zn")), 1)
                self.assertFalse(
                    np.array_equal(atoms.selective_dynamics[0],
                                   [True, True, True]))
                self.assertTrue(
                    np.array_equal(atoms.selective_dynamics[0],
                                   [False, False, False]))
                self.assertTrue(
                    np.array_equal(atoms.selective_dynamics[-5],
                                   [True, True, True]))
                self.assertTrue(
                    np.array_equal(atoms.selective_dynamics[-4],
                                   [False, False, False]))

    def test_write_poscar(self):
        write_poscar(structure=self.structure,
                     filename=posixpath.join(self.file_location,
                                             "POSCAR_test"))
        test_atoms = read_atoms(
            posixpath.join(self.file_location, "POSCAR_test"))
        self.assertEqual(self.structure.get_chemical_formula(),
                         test_atoms.get_chemical_formula())
        struct = self.structure.copy()
        struct.add_tag(selective_dynamics=[True, True, True])
        write_poscar(structure=struct,
                     filename=posixpath.join(self.file_location,
                                             "POSCAR_test"))
        test_atoms = read_atoms(
            posixpath.join(self.file_location, "POSCAR_test"))
        truth_array = np.empty_like(struct.positions, dtype=bool)
        truth_array[:] = [True, True, True]
        self.assertTrue(
            np.array_equal(np.array(test_atoms.selective_dynamics),
                           truth_array))
        os.remove(posixpath.join(self.file_location, "POSCAR_test"))

    def test_vasp_sorter(self):
        write_poscar(structure=self.structure,
                     filename=posixpath.join(self.file_location,
                                             "POSCAR_test"))
        test_atoms = read_atoms(
            posixpath.join(self.file_location, "POSCAR_test"))
        vasp_order = vasp_sorter(self.structure)
        self.assertEqual(len(self.structure), len(test_atoms))
        self.assertEqual(self.structure[vasp_order], test_atoms)
        os.remove(posixpath.join(self.file_location, "POSCAR_test"))

    def tearDown(self):
        pass