import copy
from collections import Counter
from itertools import groupby
from pathlib import Path
from typing import List, Any
import numpy as np
import yaml
from gvasp.lib.base_bind import search_image as search_image_bind
from gvasp.common.utils import redefine_frac
from gvasp.common.setting import RootDir
yaml.warnings({'YAMLLoadWarning': False})
[docs]class Lattice(object):
def __init__(self, matrix: np.ndarray):
"""
Initialize of Lattice class
Args:
matrix (np.ndarray): use a (3x3) np.ndarray matrix to initialize the Lattice
"""
self.matrix = matrix
def __eq__(self, other):
return np.all(self.matrix == other.matrix)
def __hash__(self):
return hash(self.volume)
def __repr__(self):
return f"{self.matrix}"
@property
def length(self) -> np.ndarray:
"""
Calculate the lattice length
Returns:
length (np.ndarray): store the length, shape = (3x1)
"""
return np.power(np.sum(np.power(self.matrix, 2), axis=1), 0.5)
@property
def angle(self) -> np.ndarray:
"""
Calculate the lattice angle
Returns:
angle (np.ndarray): store the angle, shape = (3x1)
"""
alpha = np.arccos(np.dot(self.matrix[1], self.matrix[2]) / (self.length[1] * self.length[2])) * 180 / np.pi
beta = np.arccos(np.dot(self.matrix[0], self.matrix[2]) / (self.length[0] * self.length[2])) * 180 / np.pi
gamma = np.arccos(np.dot(self.matrix[0], self.matrix[1]) / (self.length[0] * self.length[1])) * 180 / np.pi
return np.array([alpha, beta, gamma])
@property
def volume(self) -> float:
"""
Calculate the lattice volume
Returns:
volume (float): store the volume
"""
return np.linalg.det(self.matrix)
@property
def inverse(self) -> np.ndarray:
"""
Calculate the inverse matrix of lattice
Returns:
inverse (np.ndarray): store the inverse matrix
"""
return np.linalg.inv(self.matrix)
[docs] @staticmethod
def from_string(string):
"""
Construct a Lattice instance from string
Args:
string (List[str]): three-line <string>
Examples
---------
>>> string
array([[ 7.707464, 0.000000, 0.000000],
-3.853732, 6.674860, 0.000000],
0.000000, 0.000000, 28.319031]])
Returns:
lattice (Lattice): Lattice instance
"""
matrix = np.array([[float(ii) for ii in item.split()[:3]] for item in string])
return Lattice(matrix)
[docs] @staticmethod
def from_POSCAR(name):
"""
Construct a Lattice instance from POSCAR file
Args:
name ([str, Path]): path of POSCAR
Returns:
lattice (Lattice): Lattice instance
"""
with open(name) as f:
cfg = f.readlines()
return Lattice.from_string(cfg[2:5])
@property
def strings(self) -> str:
"""
Transform a Lattice instance to string
Returns:
strings (str): Lattice instance in string format
"""
return "".join([" ".join([f"{ii:>9.6f}" for ii in item]) + "\n" for item in self.matrix])
[docs]class Atom(object):
"""
Atom class represent one atom in periodic solid system
Attributes:
formula (str): chemical formula
number (int): atomic number
period (int): atomic period in element period table
group (int): atomic group in element period table
color (str): atomic color using RGB
order (int): atomic order in <Structure class>, default: 0
frac_coord (np.ndarray): fractional coordinates
cart_coord (np.ndarray): cartesian coordinates
_default_bonds: atomic default bond property {atom: bond-length}
Methods:
atom_type(self) -> property: register atom_type property
set_coord(self, lattice: Lattice) -> object: recalculate frac_coord from cart_coord / recalculate cart_coord
from frac_coord
search_image(atom_i, atom_j) -> np.ndarray: search the nearest image in which have minimum distance between
two atoms
__load_config(cls): load element.yaml file
__initialize_attrs(self): initialize the attributes from the element.yaml
"""
_config_file = Path(f"{RootDir}/element.yaml")
_attributes_yaml = ['number', 'period', 'group', 'color', '_default_bonds']
_load = False
_attrs = None
def __new__(cls, *args, **kwargs):
cls.__load_config()
return object.__new__(cls)
def __init__(self, formula, order: (int, list) = 0, frac_coord=None, cart_coord=None, selective_matrix=None,
constrain=False, spin=0):
self.formula = formula
self.order = order
self.frac_coord = redefine_frac(np.array(frac_coord)) if frac_coord is not None else None
self.cart_coord = np.array(cart_coord) if cart_coord is not None else None
self.selective_matrix = np.array(selective_matrix) if selective_matrix is not None else None
self.constrain = constrain
self.spin = spin
# config atom from `element.yaml`
self.number, self.period, self.group, self.color, self._default_bonds = (None, None, None, None, [])
self.__initialize_attrs()
# atom_type property
self.coordination_number = None
# bonds information
self.bonds = []
def __eq__(self, other):
return self.number == other.number and self.order == other.order
def __lt__(self, other):
return self.number < other.number or self.order < other.order
def __ge__(self, other):
return self.number >= other.number or self.order >= other.order
def __hash__(self):
return hash(self.number) + hash(str(self.order))
def __repr__(self):
return f"(Atom {self.order} : {self.formula} : {self.cart_coord})"
@classmethod
def __load_config(cls):
"""load element.yaml file (private classmethod)"""
if not cls._load:
with open(cls._config_file) as f:
cfg = f.read()
cls._attrs = yaml.safe_load(cfg)
cls._load = True
def __initialize_attrs(self):
"""initialize the attributes from the element.yaml (private method)"""
if isinstance(self.formula, str): # <class Atom>
for key, value in self._attrs[f'Element {self.formula}'].items():
setattr(self, key, value)
elif isinstance(self.formula, list): # <class Atoms>
for attr in self._attributes_yaml:
setattr(self, attr, [self._attrs[f'Element {formula}'][attr] for formula in self.formula])
@property
def atom_type(self):
"""register atom_type property"""
return f"{self.formula}{self.coordination_number}c"
[docs] def set_coord(self, lattice: Lattice):
"""
recalculate frac_coord from cart_coord / recalculate cart_coord from frac_coord
Args:
lattice (Lattice): Lattice instance
"""
assert lattice is not None
if self.cart_coord is not None and self.frac_coord is None:
self.frac_coord = np.dot(self.cart_coord, lattice.inverse)
elif self.frac_coord is not None and self.cart_coord is None:
self.cart_coord = np.dot(self.frac_coord, lattice.matrix)
return self
[docs] @staticmethod
def search_image(atom_i, atom_j) -> np.ndarray:
"""
search the nearest image in which have minimum distance between two atoms
Args:
atom_i (Atom): Atom instance
atom_j (Atom): Atom instance
Returns:
image (np.ndarray): record the transform direction
"""
if not isinstance(atom_i, Atom) or not isinstance(atom_j, Atom):
SystemError("The parameters should be <class Atom>!")
return search_image_bind(atom_i.frac_coord, atom_j.frac_coord)
[docs] @staticmethod
def distance(atom_i, atom_j, lattice: Lattice):
if not isinstance(atom_i, Atom) or not isinstance(atom_j, Atom):
raise TypeError(f"{atom_i} and {atom_j} should be <class Atom>")
image = Atom.search_image(atom_i, atom_j)
atom_j_image = copy.deepcopy(atom_j)
atom_j_image.frac_coord += image
atom_i.cart_coord, atom_j_image.cart_coord = None, None
atom_i.set_coord(lattice)
atom_j_image.set_coord(lattice)
return np.sum((atom_i.cart_coord - atom_j_image.cart_coord) ** 2) ** 0.5
[docs]class Atoms(Atom):
"""
`Atoms class represent atom set in periodic solid system`
@property
formula: chemical formula, <list>
number: atomic number, <list>
period: atomic period in element period table, <list>
group: atomic group in element period table, <list>
color: atomic color using RGB, <list>
order: atomic order in <Structure class>, default: <list(range(len(formula)))>
frac_coord: fractional coordinates, <list>
cart_coord: cartesian coordinates, <list>
_default_bonds: atomic default bond property {atom: bond-length}
count: total atom number in atom set
size: list atom number according to their formula
@func
__initialize_attrs: initialize the attributes from the element.yaml
from_list: construct the <class Atoms> from an Atom list, i.e., [Atom, Atom, Atom] --> Atoms
"""
def __init__(self, *args, **kwargs):
super(Atoms, self).__init__(*args, **kwargs)
self.order = list(range(len(self.formula))) if isinstance(self.order, int) or None in self.order else self.order
self.frac_coord = [None] * len(self.formula) if self.frac_coord is None else self.frac_coord
self.cart_coord = [None] * len(self.formula) if self.cart_coord is None else self.cart_coord
self.selective_matrix = [None] * len(self.formula) if self.selective_matrix is None else self.selective_matrix
self.constrain = [False] * len(self.formula) if self.constrain is False else self.constrain
self.spin = [0] * len(self.formula) if isinstance(self.spin, int) and self.spin == 0 else self.spin
self.__index_list = [] # inner index-list
self._atom_list = [] # static atoms-list
def __len__(self) -> int:
return len(self.formula)
def __repr__(self):
string = ""
for order, formula, cart_coord in zip(self.order, self.formula, self.cart_coord):
string += f"(Atom {order} : {formula} : {cart_coord}) \n"
return string
def __iter__(self):
self.__index_list.append(0)
return self
def __next__(self):
current_loop_index = self.__index_list[-1] # record current-loop-index
if current_loop_index < len(self):
current_loop_index += 1
self.__index_list[-1] = current_loop_index # overwrite the original index
return self[current_loop_index - 1]
else:
del (self.__index_list[-1])
raise StopIteration
def __contains__(self, atom):
if atom in list(self):
return True
else:
return False
def __getitem__(self, index):
return self.atom_list[index]
def __deepcopy__(self, memo=None):
atoms = Atoms(formula=self.formula, order=self.order, frac_coord=self.frac_coord, cart_coord=self.cart_coord,
selective_matrix=self.selective_matrix, constrain=self.constrain, spin=self.spin)
atoms.coordination_number = self.coordination_number
atoms.bonds = self.bonds
atoms._atom_list = []
return atoms
@property
def count(self) -> int:
return len(self)
@property
def elements(self):
return [(key, len(list(group))) for key, group in groupby(self.formula)]
@property
def size(self):
return Counter(self.formula)
@property
def atom_list(self): # atom may update, so we don't use a static list
if not len(self._atom_list):
for index in range(len(self.formula)):
atom = Atom(formula=self.formula[index], order=self.order[index],
frac_coord=self.frac_coord[index], cart_coord=self.cart_coord[index],
selective_matrix=self.selective_matrix[index], constrain=self.constrain[index],
spin=self.spin[index])
# update coordination number && bonds
atom.coordination_number = self.coordination_number[
index] if self.coordination_number is not None else None
atom.bonds = self.bonds[index] if len(self.bonds) != 0 else []
self._atom_list.append(atom)
return self._atom_list
@property
def atom_type(self): # override this property
return [f"{atom.formula}{atom.coordination_number}c" for atom in self]
[docs] def set_coord(self, lattice: Lattice): # overwrite this method
assert lattice is not None
if None not in self.cart_coord and None in self.frac_coord:
self.frac_coord = np.dot(self.cart_coord, lattice.inverse)
elif None not in self.frac_coord and None in self.cart_coord:
self.cart_coord = np.dot(self.frac_coord, lattice.matrix)
# update __atoms_list
for index in range(len(self)):
self.atom_list[index].frac_coord = self.frac_coord[index]
self.atom_list[index].cart_coord = self.cart_coord[index]
return self
[docs] def perturb(self, lattice: Lattice, threshold=0.1, groups=9):
"""
Perturb the atoms' coords
@params:
lattice: Lattice type
threshold: max offset of atom
groups: how many groups you want to get
"""
if None in self.cart_coord:
self.set_coord(lattice=lattice)
# group from z_value
z_max, z_min = np.max(self.cart_coord[:, 2]), np.min(self.cart_coord[:, 2])
group_np = np.linspace(z_min, z_max, groups + 1)
group_index = []
for index in range(groups):
group_index.append(np.where((self.cart_coord[:, 2] >= group_np[index]) &
(self.cart_coord[:, 2] <= group_np[index + 1])))
# exponential decay
decay = -np.linspace(0., (z_max - z_min) / 2, int(groups / 2))
exp_decay = np.exp(decay) * threshold
# construct perturb vector
perturb_vector = np.empty(groups)
perturb_vector[:int(groups / 2)] = exp_decay
perturb_vector[-int(groups / 2):] = exp_decay[::-1]
# construct perturb matrix
perturb_matrix = np.empty((self.count, 3))
for index, group in enumerate(group_index):
perturb_matrix[group] = perturb_vector[index]
random_normal = np.random.normal(loc=0., scale=0.01, size=(self.count, 3))
perturb_item = perturb_matrix + random_normal
# resolve fix_atom
fix_item = np.where(self.selective_matrix == "F", 0, 1)
perturb_item = perturb_item * fix_item
self.cart_coord = self.cart_coord + perturb_item
self.frac_coord = [None] * self.count
self.set_coord(lattice=lattice)
return self
[docs] @staticmethod
def from_list(atoms: list):
formula = [atom.formula for atom in atoms]
order: List[Any] = [atom.order for atom in atoms]
frac_coord = [atom.frac_coord for atom in atoms]
cart_coord = [atom.cart_coord for atom in atoms]
selective_matrix = [atom.selective_matrix for atom in atoms]
constrain = [atom.constrain for atom in atoms]
spin = [atom.spin for atom in atoms]
coordination_number = [atom.coordination_number for atom in atoms]
bonds = [atom.bonds for atom in atoms]
# update coordination number && bonds
new_atoms = Atoms(formula=formula, order=order, frac_coord=frac_coord, cart_coord=cart_coord,
selective_matrix=selective_matrix, constrain=constrain, spin=spin)
new_atoms.coordination_number = coordination_number
new_atoms.bonds = bonds
return new_atoms