591 lines
20 KiB
Python
591 lines
20 KiB
Python
#!/usr/bin/env python
|
|
# coding: utf-8
|
|
#
|
|
# Copyright © 2016-2020 - Rennes Physics Institute
|
|
#
|
|
# This file is part of msspec.
|
|
#
|
|
# msspec is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU General Public License as published by
|
|
# the Free Software Foundation, either version 3 of the License, or
|
|
# (at your option) any later version.
|
|
#
|
|
# msspec is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU General Public License for more details.
|
|
#
|
|
# You should have received a copy of the GNU General Public License
|
|
# along with this msspec. If not, see <http://www.gnu.org/licenses/>.
|
|
#
|
|
# Source file : src/msspec/utils.py
|
|
# Last modified: Mon, 27 Sep 2021 17:49:48 +0200
|
|
# Committed by : sylvain tricot <sylvain.tricot@univ-rennes1.fr>
|
|
|
|
|
|
"""
|
|
Module utils
|
|
============
|
|
|
|
|
|
"""
|
|
|
|
|
|
import os
|
|
import re
|
|
|
|
import numpy as np
|
|
import ase.atom
|
|
from ase import Atom
|
|
from ase import Atoms
|
|
|
|
from msspec.iodata import Data
|
|
from msspec.misc import LOGGER
|
|
|
|
|
|
class ForeignPotential(object):
|
|
def __init__(self):
|
|
self.data = Data(title='Foreign Potential')
|
|
# Load exported potentials
|
|
# phagen_data is a dictionary with:
|
|
# self.phagen_data = {
|
|
# 'VintTotal' : <float>,
|
|
# 'VintCoulomb': <float>,
|
|
# 'RHOint' : <float>,
|
|
# 'types': [
|
|
# {
|
|
# 'atom_type': <int>,
|
|
# 'Z' : <int>,
|
|
# 'RWS' : <float>,
|
|
# 'data' : <np.array(..., 4, dtype=float)>
|
|
# },
|
|
# ...
|
|
# ...
|
|
# ...
|
|
# ]
|
|
# }
|
|
self.phagen_data = {'types': []}
|
|
|
|
def write(self, filename, prototypical_atoms):
|
|
LOGGER.debug("Writing Phagen input potential file: {}".format(filename))
|
|
|
|
def DEPRECATEDappend_atom_potential(atom):
|
|
Z = atom.number
|
|
# Find the right type (Z) in the phagen_data
|
|
itypes = []
|
|
for i, t in enumerate(self.phagen_data['types']):
|
|
if t['Z'] == Z:
|
|
itypes.append(i)
|
|
# Check now that we have only one type in the list
|
|
# otherwise we do not know yet how to deal with this.
|
|
assert len(itypes) > 0, "Cannot find the data for atom with Z={}".format(Z)
|
|
assert len(itypes) == 1, "Too many datasets for atom with Z={}".format(Z)
|
|
# So far so good, let's write the block
|
|
t = self.phagen_data['types'][itypes[0]]
|
|
s = "{:<7d}{:<10d}{:1.4f}\n".format(
|
|
t['Z'], len(t['data']), t['RWS'])
|
|
line_fmt = "{:+1.8e} " * 4 + "\n"
|
|
for row in t['data']:
|
|
s += line_fmt.format(*row)
|
|
return s
|
|
|
|
def append_atom_potential(atom):
|
|
line_fmt = "{:+1.8e} " * 4 + "\n"
|
|
atom_type = atom.get('atom_type')
|
|
assert atom_type != None, "Unable get the atom type!"
|
|
for t in self.phagen_data['types']:
|
|
if t['atom_type'] == atom_type:
|
|
s = "{:<7d}{:<10d}{:1.4f}\n".format(
|
|
t['Z'], len(t['data']), t['RWS'])
|
|
for row in t['data']:
|
|
s += line_fmt.format(*row)
|
|
return s
|
|
|
|
content = ""
|
|
# Start by writing the header line
|
|
content += "{:.2f} {:.4f} {:.2f}\n".format(
|
|
self.phagen_data['VintCoulomb'],
|
|
self.phagen_data['RHOint'],
|
|
self.phagen_data['VintTotal'])
|
|
# Then for each atom in the given prototypical cluster,
|
|
# write the data block
|
|
for atom in prototypical_atoms:
|
|
content += append_atom_potential(atom)
|
|
|
|
# Write the content to filename
|
|
try:
|
|
with open(filename, 'r') as fd:
|
|
old_content = fd.read()
|
|
except IOError:
|
|
old_content = ''
|
|
|
|
modified = False
|
|
if content != old_content:
|
|
with open(filename, 'w') as fd:
|
|
fd.write(content)
|
|
modified = True
|
|
|
|
return modified
|
|
|
|
|
|
class SPRKKRPotential(ForeignPotential):
|
|
def __init__(self, atoms, potfile, *exported_files):
|
|
super().__init__()
|
|
self.atoms = atoms
|
|
self.potfile = potfile
|
|
self.load_sprkkr_atom_types()
|
|
for f in exported_files:
|
|
LOGGER.info("Loading file {}...".format(f))
|
|
# get the IT from the filename
|
|
m=re.match('SPRKKR-IT_(?P<IT>\d+)-PHAGEN.*', os.path.basename(f))
|
|
it = int(m.group('IT'))
|
|
|
|
# load the content of the header (2 lines)
|
|
with open(f, 'r') as fd:
|
|
first_line, second_line = [fd.readline().strip()
|
|
for _ in range(2)]
|
|
|
|
# load Coulomb and Total interstitial potential
|
|
pattern = (r'#\s*VMTZ_TOTAL\s*=\s*(?P<TOTAL>.*?)\s+'
|
|
r'VMTZ_Coulomb\s*=\s*(?P<COULOMB>.*?)\s+.*$')
|
|
m = re.match(pattern, first_line)
|
|
self.phagen_data.update(VintCoulomb=float(m.group('COULOMB')),
|
|
VintTotal=float(m.group('TOTAL')),
|
|
RHOint=0.)
|
|
|
|
# load Z, Wigner-Seitz radius from 2nd line of header
|
|
type_data = {}
|
|
_ = re.split(r'\s+', second_line.strip("# "))
|
|
type_data.update(atom_type=int(it), Z=int(_[0]), RWS=float(_[3]))
|
|
|
|
# load the data
|
|
data = np.loadtxt(f, comments='#')
|
|
type_data.update(data=data)
|
|
|
|
self.phagen_data['types'].append(type_data)
|
|
|
|
# store the sprkkr type number in the sprkkr_info property of each
|
|
# atom in the provided self.atoms
|
|
|
|
def load_sprkkr_atom_types(self):
|
|
def read(content, pattern, types):
|
|
# compile the pattern for regex matching
|
|
pat = re.compile(pattern, re.MULTILINE)
|
|
# get the keys and data as list of strings
|
|
keys = re.split(r'\s+',
|
|
pat.search(content).group('KEYS').strip(' \n'))
|
|
txt = pat.search(content).group('DATA').strip('\n').split('\n')
|
|
data = []
|
|
for line in txt:
|
|
# Unpacking values
|
|
values_txt = re.split(r'\s+', line.strip())
|
|
data_dict = {}
|
|
for i, _ in enumerate(values_txt):
|
|
# Type casting
|
|
value = types[i].__call__(_)
|
|
data_dict[keys[i]] = value
|
|
# push to the list
|
|
data.append(data_dict)
|
|
return data
|
|
|
|
# load info in *.pot file
|
|
LOGGER.info("Loading SPRKKR *.pot file {}...".format(self.potfile))
|
|
with open(self.potfile, 'r') as fd:
|
|
content = fd.read()
|
|
|
|
self.sites_data = read(content,
|
|
(r'^\s*SITES\s*\n((.*\n)+?\s*(?P<KEYS>IQ.*)\n'
|
|
r'(?P<DATA>(.*\n)+?))\*+'),
|
|
[int] + [float] * 3)
|
|
self.types_data = read(content,
|
|
(r'^\s*TYPES\s*\n(\s*(?P<KEYS>IT.*)\n'
|
|
r'(?P<DATA>(.*\n)+?))\*+'),
|
|
[int, str] + [int] * 4)
|
|
self.occ_data = read(content,
|
|
(r'^\s*OCCUPATION\s*\n(\s*(?P<KEYS>IQ.*)\n'
|
|
r'(?P<DATA>(.*\n)+?))\*+'),
|
|
[int] * 5 + [float])
|
|
|
|
LOGGER.debug("SITES:")
|
|
for _ in self.sites_data:
|
|
LOGGER.debug(_)
|
|
|
|
LOGGER.debug("TYPES:")
|
|
for _ in self.types_data:
|
|
LOGGER.debug(_)
|
|
|
|
LOGGER.debug("OCCUPATION:")
|
|
for _ in self.occ_data:
|
|
LOGGER.debug(_)
|
|
|
|
for site in self.sites_data:
|
|
IQ = site['IQ']
|
|
# Get the Atom at x, y, z
|
|
x = site['QBAS(X)']
|
|
y = site['QBAS(Y)']
|
|
z = site['QBAS(Z)']
|
|
i = get_atom_index(self.atoms, x, y, z, scaled=True)
|
|
for occupation in self.occ_data:
|
|
if occupation['IQ'] == IQ:
|
|
IT = occupation['ITOQ']
|
|
atom = self.atoms[i]
|
|
atom.set('atom_type', IT)
|
|
LOGGER.debug("Site #{} is type #{}, atom {}".format(IQ, IT, atom))
|
|
|
|
|
|
|
|
class EmptySphere(Atom):
|
|
def __init__(self, *args, **kwargs):
|
|
Atom.__init__(self, *args, **kwargs)
|
|
self.symbol = 'X'
|
|
|
|
|
|
def get_atom_index(atoms, x, y, z, scaled=False):
|
|
""" 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
|
|
:param bool scaled: whether the x,y,z coordinates are scaled
|
|
|
|
:return: the index of the atom as an integer
|
|
:rtype: int
|
|
"""
|
|
# get all distances
|
|
if scaled:
|
|
positions = atoms.get_scaled_positions()
|
|
else:
|
|
positions = atoms.get_positions()
|
|
d = np.linalg.norm(positions - np.array([x, y, z]), axis=1)
|
|
# get the index of the min distance
|
|
i = np.argmin(d)
|
|
# return the index
|
|
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, center=(0, 0, 0)):
|
|
""" 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
|
|
"""
|
|
assert radius >= 0, "Please give a positive radius value"
|
|
radii = np.linalg.norm(atoms.positions - center, axis=1)
|
|
indices = np.where(radii <= radius)[0]
|
|
return atoms[indices]
|
|
|
|
|
|
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}`. # noqa: W605
|
|
:type radius: float
|
|
:param z: The height of the cone in :math:`\mathring{A}`. # noqa: W605
|
|
:type z: float
|
|
:return: A new cluster.
|
|
:rtype: :py:class:`ase.Atoms`
|
|
"""
|
|
new_atoms = atoms.copy()
|
|
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)
|
|
|
|
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 Exception:
|
|
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, shape='spherical'):
|
|
|
|
"""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 = np.inf
|
|
|
|
for atom in cluster:
|
|
if ze - eps < atom.z < ze + eps and (atom.symbol == symbol or
|
|
symbol is None):
|
|
n = np.sqrt(atom.x**2 + atom.y**2)
|
|
if (n < nmin):
|
|
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
|
|
|
|
# the number of planes in the cluster
|
|
p = np.alen(np.unique(np.round(cluster.get_positions()[:, 2], 4)))
|
|
# the symbol of your emitter
|
|
symbol = cluster[np.where(cluster.get_tags() == emitter_tag)[0][0]].symbol
|
|
|
|
assert (diameter != 0 or planes != 0), \
|
|
"At least one of diameter or planes parameter must be use."
|
|
|
|
if diameter == 0:
|
|
# calculate the minimal diameter according to the number of planes
|
|
min_diameter = 1+2*(planes*c/p+1)
|
|
else:
|
|
min_diameter = diameter
|
|
|
|
# number of repetition in each direction
|
|
rep = int(3*min_diameter/min(a, c))
|
|
|
|
# repeat the cluster
|
|
cluster = cluster.repeat((rep, rep, rep))
|
|
|
|
# center the cluster
|
|
center_cluster(cluster)
|
|
|
|
# reset the cell
|
|
cluster.set_cell(cell)
|
|
|
|
# cut the cluster so that we have a centered surface
|
|
cluster = cut_plane(cluster, z=eps)
|
|
|
|
# positions where atoms have the tag of the emitter_tag
|
|
i = np.where(cluster.get_tags() == emitter_tag)
|
|
|
|
# an array of all unique z corresponding to where we have the right
|
|
# atom's tag
|
|
all_ze = np.sort(np.unique(np.round(cluster.get_positions()[:, 2][i], 4)))
|
|
|
|
# an array of all unique z
|
|
all_z = np.sort(np.unique(np.round(cluster.get_positions()[:, 2], 4)))
|
|
|
|
# calculate the number of planes above the emitter's plane
|
|
n = np.where(all_z == all_z.max())[0][0] - np.where(
|
|
all_z == all_ze.max())[0][0]
|
|
|
|
# the height of the emitter's plane
|
|
ze = all_ze.max()
|
|
|
|
# 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()
|
|
|
|
# values of x and y of the emitter
|
|
tx, ty = get_xypos(cluster, ze, symbol)
|
|
|
|
# center the cluster on the emitter
|
|
Atoms.translate(cluster, [-tx, -ty, 0])
|
|
|
|
# calculate where to cut to get the right number of planes above the
|
|
# emitter
|
|
z_cut = all_z[np.where(all_z == all_ze.max())[0][0] + emitter_plane]
|
|
|
|
# translate the surface at z=0
|
|
Atoms.translate(cluster, [0, 0, -z_cut])
|
|
|
|
# cut the planes above those we want to keep
|
|
cluster = cut_plane(cluster, z=eps)
|
|
|
|
radius = diameter/2
|
|
if planes != 0:
|
|
# an array of all unique remaining z
|
|
all_z = np.sort(np.unique(np.round(cluster.get_positions()[:, 2], 4)))
|
|
|
|
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)
|
|
|
|
if shape in ('spherical'):
|
|
# cut a sphere in our cluster with the diameter which is indicate in
|
|
# the parameters
|
|
cluster = cut_sphere(cluster, radius=radius + eps)
|
|
elif shape in ('cylindrical'):
|
|
# cut a sphere in our cluster with the diameter which is indicate in
|
|
# the parameters
|
|
cluster = cut_cylinder(cluster, radius=radius + eps)
|
|
else:
|
|
raise NameError('Unkknown shape specifier: \"{}\"'.format(shape))
|
|
|
|
if planes != 0:
|
|
# calculate where to cut to get the right number of planes
|
|
positions = np.unique(np.round(cluster.get_positions()[:, 2], 4))
|
|
zcut = np.sort(positions)[::-1][planes-1] - eps
|
|
|
|
# cut the right number of planes
|
|
cluster = cut_plane(cluster, z=zcut)
|
|
|
|
# an array of all unique remaining z
|
|
all_z = np.sort(np.unique(np.round(cluster.get_positions()[:, 2], 4)))
|
|
|
|
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
|