import numpy as np
import requests
import os
from EPWpy.structure.position_atoms import *
from EPWpy.utilities.printing import *
from EPWpy.utilities.display_struct import *
from EPWpy.error_handling import error_handler
from EPWpy.structure.QE_position_atoms import *
from EPWpy.default import default_dicts as dc
mayavi_found = False
try:
from mayavi import mlab
from mayavi.modules.iso_surface import IsoSurface
from tvtk.api import tvtk
mayavi_found = True
except ImportError:
error_handler.error_mayavi()
try:
from mp_api.client import MPRester
except ImportError:
error_handler.error_mp()
[docs]
class Lattice:
"""
A class representing a crystal lattice and its associated data.
This class provides a basic interface for storing and managing lattice-related
information. It can serve as the foundational class for structural manipulation,
symmetry analysis, or visualization routines in materials simulations.
Parameters
----------
data : dict or object
Input data structure containing lattice information. Typically includes
lattice vectors, atomic positions, and element types, but may be any
user-defined format compatible with downstream methods.
Attributes
----------
data : dict or object
The input lattice data, stored for use by other methods.
home : str
The working directory path at the time of object initialization.
Examples
--------
>>> lattice_data = {
... 'lattice_vectors': [[1.0, 0.0, 0.0],
... [0.0, 1.0, 0.0],
... [0.0, 0.0, 1.0]],
... 'atomic_positions': [[0.0, 0.0, 0.0]],
... 'elements': ['Si']
... }
>>> lat = Lattice(lattice_data)
>>> lat.home
'/home/user/project'
"""
def __init__(self, data):
self.data = data
self.home = os.getcwd()
[docs]
def get_atom(self):
"""
Gets atomic attributes from a poscar structure file
"""
if (self.structure != None):
atom_pos,atom_pos2,lattice_vec,mat,materials,natoms,a = main_extract([0.0,0.0],self.structure,T=1)
else:
atom_pos,atom_pos2,lattice_vec,mat,materials,natoms,a = self.get_QE_structure()
return(atom_pos, atom_pos2, lattice_vec, mat, materials, natoms, a)
[docs]
def atom_cart(self):
"""
Gets atomic positions in cartesian co-ordinate
"""
atom_pos,atom_pos2,lattice_vec,mat,materials,natoms,a = main_extract([0.0,0.0],self.structure,T=1)
return(atom_pos2)
[docs]
def get_xyz(self, supercell=None):
if supercell is None:
supercell = []
"""
Writes a xyz file
"""
atom_pos,atom_pos2,lattice_vec,mat,materials,natoms,a = self.get_atom()
#atom_pos2 = self.atom_cart()
if(len(supercell) !=0):
atom_pos2,natoms,mat = get_supercell(atom_pos2,natoms,mat,lattice_vec,a,supercell)
gen_xyz(atom_pos2,natoms,mat)
[docs]
def get_supercell(self, supercell=None, pl_vec=None):
"""
returns the co-ordinates for supercell
"""
if supercell is None:
supercell = [1,1,1]
if pl_vec is None:
pl_vec = []
atom_pos,atom_pos2,lattice_vec,mat,materials,natoms,a = self.get_atom()
#atom_pos2 = self.atom_cart()
if(len(supercell) != 0):
if (len(pl_vec) == 0):
atom_pos2,natoms,mat = get_supercell(atom_pos2,natoms,mat,lattice_vec,a,supercell)
return(atom_pos2,natoms,mat)
else:
atom_pos2,natoms,mat,pl_vec = get_supercell_pl(atom_pos2,natoms,mat,lattice_vec,a,supercell,pl_vec)
return(atom_pos2,natoms,mat,pl_vec)
[docs]
def get_QE_structure(self):
"""
Gets structure from a QE input
"""
atom_pos = self.pw_atomic_positions['atomic_pos']
atom_pos_type = self.pw_atomic_positions['atomic_position_type']
mat = self.pw_atomic_positions['atoms']
materials = self.pw_atomic_species['atomic_species']
natoms = self.pw_atomic_positions['num']
ibrav = self.pw_system['ibrav']
cell_parameters = self.pw_cell_parameters['lattice_vector']
cell_type = self.pw_cell_parameters['cell_type']
try:
celldm ={'1':self.pw_system['celldm(1)'],
'2':self.pw_system['celldm(2)'],
'3':self.pw_system['celldm(3)']}
except KeyError:
celldm ={'1':1.0,
'2':self.pw_system['celldm(2)'],
'3':self.pw_system['celldm(3)']}
atom_pos2,lattice_vec = atom_pos_lat_crystal({'ibrav': ibrav,
'cell_parameters': cell_parameters,
'cell_type': cell_type,
'atom_pos': atom_pos,
'atom_pos_type':atom_pos_type,
'celldm': celldm})
return(atom_pos,atom_pos2,lattice_vec,mat,materials,natoms,1.0)
[docs]
def display_lattice_legacy(self, supercell=None, view=None):
if supercell is None:
supercell = []
if view is None:
view = {}
"""
Display lattice
"""
self.get_xyz(supercell)
return(display_molecule(view))#self.atom_cart())
[docs]
def display_lattice(self, supercell=None, view=None, bond_length=3.5, type_view='mayavi'):
"""
Visualize the atomic lattice of the system.
This method displays the atomic structure of the simulated material
using either **Mayavi** (for interactive 3D visualization) or
**Matplotlib** (for static plotting). It supports supercell expansion
and customizable view parameters.
Parameters
----------
supercell : list or None, optional
Supercell expansion factors along each lattice direction, e.g.,
``[2, 2, 1]`` for a 2×2×1 supercell.
If ``None``, uses the primitive cell.
view : dict or None, optional
Dictionary specifying visualization options such as rotation,
zoom, or color scheme. Example:
``{'azimuth': 45, 'elevation': 30, 'wrap': False}``
If ``None``, defaults are used.
bond_length : float, optional
Maximum interatomic distance (in Å) for drawing bonds (default: ``3.5``).
type_view : str, optional
Visualization backend:
- ``'mayavi'`` → interactive 3D visualization (default)
- ``'matplotlib'`` → static 3D plotting for environments without Mayavi
Returns
-------
object
Visualization object returned by the selected backend:
- ``Mayavi`` scene if available and selected.
- ``Matplotlib`` figure otherwise.
Notes
-----
- Requires Mayavi to be installed for full 3D interactive visualization.
- If Mayavi is not available, falls back automatically to Matplotlib rendering.
- The visualization includes lattice vectors and atomic species.
Example
-------
>>> epw = EPWProperties(file='epw2.out', system='MoS2')
>>> epw.display_lattice(supercell=[2,2,1], type_view='matplotlib')
"""
if supercell is None:
supercell = []
if view is None:
view = {}
"""
Display lattice
"""
if ((mayavi_found) & (type_view == 'mayavi')):
atom_pos2, natoms, mat = self.get_supercell(supercell)
Data= {'positions':atom_pos2,'mat':mat}
return(display_crystal(Data, view = view, bond_length=bond_length))#self.atom_cart())
else:
atom_pos2, natoms, mat = self.get_supercell(supercell)
_,_,lattice_vec,_,materials,natoms,a = self.get_atom()
Data= {'positions':atom_pos2,
'mat':mat,
'natoms':natoms,
'atomic_species':materials,
'lattice_vec':lattice_vec}
return(display_crystal_matplotlib(Data, view = view, bond_length = bond_length))
[docs]
def display_phonons(self, pl_vec, supercell=None, view=None, bond_length=3.5, type_view='mayavi'):
"""
Display the phonon displacement modes of the crystal structure.
This method visualizes phonon modes (atomic vibrations) using the
provided phonon displacement vectors (`pl_vec`). The visualization
can be generated using either Mayavi or Matplotlib, depending on
availability and the `type_view` argument.
Parameters
----------
pl_vec : array-like
Phonon displacement vectors for each atom in the structure.
Typically obtained from a phonon calculation or interpolation step.
Must have shape (N_atoms, 3).
supercell : list[int], optional
Supercell expansion factors in each lattice direction [nx, ny, nz].
If not provided, the primitive cell is used.
view : dict, optional
Visualization parameters, such as camera orientation or zoom level.
Default is an empty dictionary.
bond_length : float, optional
Maximum bond distance (in Ångströms) used to determine which atoms
are connected by bonds in the visualization. Default is 3.5 Å.
type_view : str, optional
Visualization backend to use. Options are:
- `'mayavi'`: use Mayavi for 3D interactive rendering (default)
- `'matplotlib'`: use Matplotlib for static rendering
Returns
-------
object
A visualization object or figure handle created by the selected
visualization backend.
Notes
-----
This method constructs the full crystal structure (including optional
supercell expansion) and overlays atomic displacement vectors
representing phonon modes. The resulting data is passed to an external
visualization routine (`display_crystal_phonon`).
Examples
--------
>>> pl_vec = np.random.randn(4, 3) # Example displacement vectors
>>> model.display_phonons(pl_vec, supercell=[2, 2, 1], type_view='matplotlib')
"""
if supercell is None:
supercell = []
if view is None:
view = {}
atom_pos2, natoms, mat, pl_vec = self.get_supercell(supercell, pl_vec)
_,_,lattice_vec,_,materials,natoms,a = self.get_atom()
Data= {'positions':atom_pos2,
'mat':mat,
'natoms':natoms,
'atomic_species':materials,
'lattice_vec':lattice_vec,
'forces':pl_vec}
return(display_crystal_phonon(Data, view = view, bond_length = bond_length))
[docs]
def get_Data(self, pl_vec, supercell=None, view=None, bond_length=3.5, type_view='mayavi'):
"""
Display phonon modes
"""
if supercell is None:
supercell = []
if view is None:
view = {}
atom_pos2, natoms, mat, pl_vec = self.get_supercell(supercell, pl_vec)
_,_,lattice_vec,_,materials,natoms,a = self.get_atom()
Data= {'positions':atom_pos2,
'mat':mat,
'natoms':natoms,
'atomic_species':materials,
'lattice_vec':lattice_vec,
'forces':pl_vec}
return(Data)
[docs]
def get_poscar(self):
"""
Gets the structure file from materials project
"""
try:
with MPRester("DZYdCD0qMIVqUZQRlojThyabVVnfANXR") as mpr:
docs = mpr.materials.summary.search(material_ids=[f'{self.matid}'], fields=["structure"])
structure = docs[0].structure
# -- Shortcut for a single Materials Project ID:
structure = mpr.get_structure_by_material_id(f'{self.matid}')
#conventional_cell = structure.to_conventional()
#conventional_cell.to(fmt="POSCAR", filename=f'POSCAR_{self.matid}')
prim_cell = structure.to_cell('primitive',symprec=0.1,angle_tolerance=2)
prim_cell.to(fmt="POSCAR", filename=f'POSCAR_{self.matid}')
except NameError:
print(f"Cannot retrieve the structure with mp-api ID = {self.matid}")
print("Install EPWpy with mp-api or specify the structure using celldm")
print("EPWpy will be terminated")
raise SystemExit(1)
[docs]
def get_pseudo(self,name='pseudo'):
"""
Automatically download pseudopotential from pseudodojo
"""
try:
os.mkdir('pseudo')
except OSError:
pass
#print('pseudo folder found')
cwd = self.home
pseudos=[]
resp = []
for atoms in self.atomic_species:
response = self.get_response(atoms)
with open(f'./pseudo/{atoms}_r.{self.pseudo_end}', 'w') as f:
for line in response.text:
f.write(line)
pseudos.append(f'{atoms}_r.{self.pseudo_end}')
resp.append(response.ok)
return(pseudos, str(cwd)+'/pseudo/', resp)
[docs]
def get_ecutwfc(self):
"""
Gets standard ecutwfc
"""
cutoff_array = []
for atoms in self.pw_atomic_species['atomic_species']:
atom_cutoff = dc.ecutwfc_standard[atoms]
cutoff_array.append(int(atom_cutoff))
ecutwfc = max(cutoff_array)
return ecutwfc
[docs]
def get_mass(self):
"""
get mass
"""
mass_array = []
for atoms in self.pw_atomic_species['atomic_species']:
atom_mass = dc.atomic_mass[atoms]
mass_array.append(float(atom_mass))
return mass_array
[docs]
def get_response(self, atoms):
pseudo_type = self.pseudo_typ
end = self.pseudo_end
"""
Get the response from pseudodojo website
"""
for orbital in self.pseudo_orbitals:
url = f'https://raw.githubusercontent.com/PseudoDojo/ONCVPSP-{pseudo_type}/master/{atoms}/{atoms}{orbital}.{end}'#/As/As-d_r.upf
print(url)
response = requests.get(url)
if (response.ok == True):
found = f'ONCVPSP-{pseudo_type}/{atoms}{orbital}.{end}'
break
if(response.ok == False):
print(f'pseudo not found for {atoms} in pseudodojo, please download manually')
print(f'Add tags pseudo_dir:<location of pseudo> and pseudo:[\'each species\'] in EPWpy object')
else:
print(f'pseudo found at pseudodojo : {found}')
return(response)
[docs]
def gen_supercell(self, supercell=None):
"""
returns the co-ordinates for supercell
"""
if supercell is None:
supercell = [1,1,1]
atom_pos,atom_pos2,lattice_vec,mat,materials,natoms,a = self.get_atom()
#atom_pos2 = self.atom_cart()
if(len(supercell) != 0):
atom_pos2, natoms, mat = get_supercell(atom_pos2,natoms,mat,lattice_vec,a,supercell)
lattice_vec_sup = lattice_vec @ np.diag(supercell)
atom_pos_new = np.zeros((np.shape(atom_pos2)), dtype=float)
# Cartesian to Direct
atom_pos_new[:,:] = np.dot(atom_pos2[:,:],np.linalg.inv(lattice_vec_sup))
materials_n, natoms_n = np.unique(mat, return_counts=True)
# Generate POSCAR
gen_POSCAR(lattice_vec_sup,atom_pos_new,natoms_n,materials_n)
# Return everything
return({'atom_pos':atom_pos_new,'lattice_vec':lattice_vec_sup,
'natoms':natoms_n, 'materials':materials_n})
return (None)
@property
def structure_params(self):
"""
Returns structure parameters as a dictionary
Units are Angstrom for variables in Cartesian
"""
atom_pos,atom_pos2,lattice_vec,mat,materials,natoms,a = self.get_atom()
atomic_weight = []
atomic_mass = []
for atm in materials:
atomic_weight.append(dc.atomic_number[atm])
atomic_mass.append(dc.atomic_mass[atm])
return({'atomic_positions_crystal':atom_pos,
'atomic_positions_cartesian':atom_pos2,
'lattice_vector':lattice_vec,
'atomic_species':materials,
'atoms':mat,
'natoms':natoms,
'atomic_weight':atomic_weight,
'atomic_mass':atomic_mass,
'lattice_constant':a})
@property
def structure_params_summary(self):
"""
Print structure as a formatted string
"""
data = self.structure_params
lines = []
lines.append("=== Structure Information ===")
lines.append(f"Lattice constant: {data['lattice_constant']}")
lines.append(f"Number of atoms: {data['natoms']}")
lines.append(f"Atomic species : {', '.join(data['atomic_species'])}")
lines.append(f"Atomic weights : {data['atomic_weight']}")
lines.append(f"Atomic masses : {data['atomic_mass']}\n")
lines.append("--- Lattice Vectors (Å) ---")
for row in data['lattice_vector']:
lines.append(" {:+10.6f} {:+10.6f} {:+10.6f}".format(*row))
lines.append("\n--- Atomic Positions (crystal) ---")
for atom, pos in zip(data['atoms'], data['atomic_positions_crystal']):
lines.append(f" {atom:2s} : {pos[0]:8.3f} {pos[1]:8.3f} {pos[2]:8.3f}")
lines.append("\n--- Atomic Positions (cartesian, Å) ---")
for atom, pos in zip(data['atoms'], data['atomic_positions_cartesian']):
lines.append(f" {atom:2s} : {pos[0]:10.6f} {pos[1]:10.6f} {pos[2]:10.6f}")
formatted_output = "\n".join(lines)
return formatted_output
[docs]
def get_supercell(atom_pos2,natoms,mat,lattice_vec,a,supercell):
"""
Builds supercell of a certain size
"""
Nx = supercell[0]
Ny = supercell[1]
Nz = supercell[2]
matn = []
natoms = []
atom_posn = np.zeros((len(atom_pos2[:,0])*Nx*Ny*Nz,3), dtype = float)
p=0
for t in range(len(atom_pos2[:,0])):
for i in range(Nx):
for j in range(Ny):
for k in range(Nz):
atom_posn[p,:] = atom_pos2[t,:]+a*(i*lattice_vec[0,:]+j*lattice_vec[1,:]+k*lattice_vec[2,:])
matn.append(mat[t])
natoms.append(1)
p +=1
return(atom_posn,natoms,matn)
[docs]
def get_supercell_pl(atom_pos2,natoms,mat,lattice_vec,a,supercell,pl_vec):
"""
Builds supercell of a certain size for the polarization vector
"""
Nx = supercell[0]
Ny = supercell[1]
Nz = supercell[2]
matn = []
natoms = []
atom_posn = np.zeros((len(atom_pos2[:,0])*Nx*Ny*Nz,3), dtype = float)
pl_vecn = np.zeros((len(atom_pos2[:,0])*Nx*Ny*Nz,3), dtype = float)
p=0
for t in range(len(atom_pos2[:,0])):
for i in range(Nx):
for j in range(Ny):
for k in range(Nz):
atom_posn[p,:] = atom_pos2[t,:]+a*(i*lattice_vec[0,:]+j*lattice_vec[1,:]+k*lattice_vec[2,:])
pl_vecn[p,:] = pl_vec[t,:]
matn.append(mat[t])
natoms.append(1)
p +=1
return(atom_posn,natoms,matn,pl_vecn)
