332 lines
11 KiB
Python
332 lines
11 KiB
Python
|
# -*- 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
|