# -*- encoding: utf-8 -*- # vim: set fdm=indent ts=4 sw=4 sts=4 et tw=80 ai cc=+0 mouse=a nu : # """ Module utils ============ """ import numpy as np from ase import Atoms, Atom from ase.visualize import view class MsSpecAtoms(Atoms): def __init__(self, *args, **kwargs): Atoms.__init__(self, *args, **kwargs) self.__absorber_index = None def set_absorber(self, index): self.__absorber_index = index def get_absorber(self): return self.__absorber_index class EmptySphere(Atom): def __init__(self, *args, **kwargs): Atom.__init__(self, *args, **kwargs) self.symbol = 'X' def get_atom_index(atoms, x, y, z): """ Return the index of the atom that is the closest to the coordiantes given as parameters. :param ase.Atoms atoms: an ASE Atoms object :param float x: the x position in angstroms :param float y: the y position in angstroms :param float z: the z position in angstroms :return: the index of the atom as an integer :rtype: int """ # get all distances d = np.linalg.norm(atoms.get_positions() - np.array([x, y, z]), axis = 1) # get the index of the min distance i = np.argmin(d) # return the index and the corresponding distance return i def center_cluster(atoms, invert=False): """ Centers an Atoms object by translating it so the origin is roughly at the center of the cluster. The function supposes that the cluster is wrapped inside the unit cell, with the origin being at the corner of the cell. It is used in combination with the cut functions, which work only if the origin is at the center of the cluster :param ase.Atoms atoms: an ASE Atoms object :param bool invert: if True, performs the opposite translation (uncentering the cluster) """ for i, cell_vector in enumerate(atoms.get_cell()): if invert: atoms.translate(0.5*cell_vector) else: atoms.translate(-0.5*cell_vector) def cut_sphere(atoms, radius): assert radius >= 0, "Please give a positive radius value" radii = np.linalg.norm(atoms.positions, axis=1) indices = np.where(radii <= radius)[0] return atoms[indices] def _cut_sphere(atoms, radius=None): """ Removes all the atoms of an Atoms object outside a sphere with a given radius :param ase.Atoms atoms: an ASE Atoms object :param float radius: the radius of the sphere :return: The modified atom cluster :rtype: ase.Atoms """ if radius is None: raise ValueError("radius not set") new_atoms = atoms.copy() del_list = [] for index, position in enumerate(new_atoms.positions): if np.linalg.norm(position) > radius: del_list.append(index) del_list.reverse() for index in del_list: del new_atoms[index] return new_atoms def cut_cylinder(atoms, axis="z", radius=None): """ Removes all the atoms of an Atoms object outside a cylinder with a given axis and radius :param ase.Atoms atoms: an ASE Atoms object :param str axis: string "x", "y", or "z". The axis of the cylinder, "z" by default :param float radius: the radius of the cylinder :return: The modified atom cluster :rtype: ase.Atoms """ if radius is None: raise ValueError("radius not set") new_atoms = atoms.copy() dims = {"x": 0, "y": 1, "z": 2} if axis in dims: axis = dims[axis] else: raise ValueError("axis not valid, must be 'x','y', or 'z'") del_list = [] for index, position in enumerate(new_atoms.positions): # calculating the distance of the atom to the given axis r = 0 for dim in range(3): if dim != axis: r = r + position[dim]**2 r = np.sqrt(r) if r > radius: del_list.append(index) del_list.reverse() for index in del_list: del new_atoms[index] return new_atoms def cut_cone(atoms, radius, z = 0): """Shapes the cluster as a cone. Keeps all the atoms of the input Atoms object inside a cone of based radius *radius* and of height *z*. :param atoms: The cluster to modify. :type atoms: :py:class:`ase.Atoms` :param radius: The base cone radius in :math:`\mathring{A}`. :type radius: float :param z: The height of the cone in :math:`\mathring{A}`. :type z: float :return: A new cluster. :rtype: :py:class:`ase.Atoms` """ new_atoms = atoms.copy() origin = np.array((0, 0, 0)) max_theta = np.arctan(radius/(-z)) u = np.array((0, 0, -z)) normu = np.linalg.norm(u) new_atoms.translate(u) indices = [] for i in range(len(new_atoms)): v = new_atoms[i].position normv = np.linalg.norm(v) _ = np.dot(u, v)/normu/normv if _ == 0: print(v) theta = np.arccos(_) if theta <= max_theta: indices.append(i) new_atoms = new_atoms[indices] new_atoms.translate(-u) # pylint: disable=invalid-unary-operand-type return new_atoms def cut_plane(atoms, x=None, y=None, z=None): """ Removes the atoms whose coordinates are higher (or lower, for a negative cutoff value) than the coordinates given for every dimension. For example, .. code-block:: python cut_plane(atoms, x=[-5,5], y=3.6, z=0) #every atom whose x-coordinate is higher than 5 or lower than -5, and/or #y-coordinate is higher than 3.6, and/or z-coordinate is higher than 0 #is deleted. :param ase.Atoms atoms: an ASE Atoms object :param int x: x cutoff value :param int y: y cutoff value :param int z: z cutoff value :return: The modified atom cluster :rtype: ase.Atoms """ dim_names = ('x', 'y', 'z') dim_values = [x, y, z] for i, (name, value) in enumerate(zip(dim_names, dim_values)): assert isinstance(value, (int, float, list, tuple, type(None))), "Wrong type" if isinstance(value, (tuple, list)): assert len(value) == 2 and np.all([isinstance(el, (int, float, type(None))) for el in value]), \ "Wrong length" else: try: if value >= 0: dim_values[i] = [-np.inf, value] else: dim_values[i] = [value, np.inf] except: dim_values[i] = [value, value] if dim_values[i][0] is None: dim_values[i][0] = -np.inf if dim_values[i][1] is None: dim_values[i][1] = np.inf dim_values = np.array(dim_values) def constraint(coordinates): return np.all(np.logical_and(coordinates >= dim_values[:,0], coordinates <= dim_values[:,1])) indices = np.where(list(map(constraint, atoms.positions)))[0] return atoms[indices] def hemispherical_cluster(cluster, emitter_tag=0, emitter_plane=0, diameter=0, planes=0): """Creates and returns a cluster based on an Atoms object and some parameters. :param cluster: the Atoms object used to create the cluster :type cluster: Atoms object :param emitter_tag: the tag of your emitter :type emitter_tag: integer :param diameter: the diameter of your cluster in Angströms :type diameter: float :param planes: the number of planes of your cluster :type planes: integer :param emitter_plane: the plane where your emitter will be starting by 0 for the first plane :type emitter_plane: integer See :ref:`hemispherical_cluster_faq` for more informations. """ def get_xypos(cluster, ze, symbol=None): nmin = None for atom in cluster: if ze - eps < atom.z < ze + eps and (atom.symbol == symbol or symbol == None): n = np.sqrt(atom.x**2 + atom.y**2) if (n < nmin) or (nmin is None): nmin = n iatom = atom.index pos = cluster.get_positions()[iatom] tx, ty = pos[0], pos[1] return tx, ty cell = cluster.get_cell() eps = 0.01 # a useful small value c = cell[:, 2].max() # a lattice parameter a = cell[:, 0].max() # a lattice parameter p = np.alen(np.unique(np.round(cluster.get_positions()[:, 2], 4))) # the number of planes in the cluster symbol = cluster[np.where(cluster.get_tags() == emitter_tag)[0][0]].symbol # the symbol of your emitter assert (diameter != 0 or planes != 0), "At least one of diameter or planes parameter must be use." if diameter == 0: l = 1+2*(planes*c/p+1) # calculate the minimal diameter according to the number of planes else: l = diameter rep = int(2*l/min(a,c)) # number of repetition in each direction cluster = cluster.repeat((rep, rep, rep)) # repeat the cluster center_cluster(cluster) # center the cluster cluster.set_cell(cell) # reset the cell cluster = cut_plane(cluster, z=eps) # cut the cluster so that we have a centered surface i = np.where(cluster.get_tags() == emitter_tag) # positions where atoms have the tag of the emitter_tag all_ze = np.sort(np.unique(np.round(cluster.get_positions()[:, 2][i], 4))) # an array of all unique z corresponding to where we have the right atom's tag all_z = np.sort(np.unique(np.round(cluster.get_positions()[:, 2], 4))) # an array of all unique z n = np.where(all_z == all_z.max())[0][0] - np.where(all_z == all_ze.max())[0][0] # calculate the number of planes above the emitter's plane ze = all_ze.max() # the height of the emitter's plane # if the number of planes above the emitter's plane is smaller than it must be, recalculate n and ze while n < emitter_plane: all_ze = all_ze[:-1] n = np.where(all_z == all_z.max())[0][0] - np.where(all_z == all_ze.max())[0][0] ze = all_ze.max() tx, ty = get_xypos(cluster, ze, symbol) # values of x and y of the emitter Atoms.translate(cluster, [-tx, -ty, 0]) # center the cluster on the emitter z_cut = all_z[np.where(all_z == all_ze.max())[0][0] + emitter_plane] # calculate where to cut to get the right number of planes above the emitter Atoms.translate(cluster, [0, 0, -z_cut]) # translate the surface at z=0 cluster = cut_plane(cluster, z=eps) # cut the planes above those we want to keep radius = diameter/2 if planes!=0: all_z = np.sort(np.unique(np.round(cluster.get_positions()[:, 2], 4))) # an array of all unique remaining z zplan = all_z[-planes] xplan, yplan = get_xypos(cluster, zplan) radius = np.sqrt(xplan**2 + yplan**2 + zplan**2) if diameter!=0: assert (radius <= diameter/2), "The number of planes is too high compared to the diameter." radius = max(radius, diameter/2) cluster = cut_sphere(cluster, radius=radius + eps) # cut a sphere in our cluster with the diameter which is indicate in the parameters if planes!=0: zcut = np.sort(np.unique(np.round(cluster.get_positions()[:, 2], 4)))[::-1][planes-1] - eps # calculate where to cut to get the right number of planes cluster = cut_plane(cluster, z=zcut) # cut the right number of planes all_z = np.sort(np.unique(np.round(cluster.get_positions()[:, 2], 4))) # an array of all unique remaining z assert emitter_plane < np.alen(all_z), "There are not enough existing plans." ze = all_z[- emitter_plane - 1] # the z-coordinate of the emitter Atoms.translate(cluster, [0, 0, -ze]) # put the emitter in (0,0,0) return cluster