1+ import numpy as np
2+ import os
3+
4+ class Cell :
5+ def __init__ (self ):
6+ self .atom = None
7+ self .a = None # Lattice vectors
8+ self .unit = 'Angstrom' # Default unit
9+ self .spin = 0 # Default spin
10+ self .charge = 0 # Default charge
11+ self .lattice_constant = 6.1416 #
12+ self .basis = None
13+ self .pseudo = None
14+ self .orbitals = []
15+ self .pseudo_potentials = {}
16+ self .pseudo_dir = ''
17+ self .orbital_dir = ''
18+ self .basis_type = ''
19+ self .built = False
20+ self ._kspace = None
21+ self .precision = 1e-8 # Default precision
22+ self ._mesh = None
23+ self .ke_cutoff = None
24+ self .rcut = None
25+
26+ @classmethod
27+ def from_file (cls , stru_file ):
28+ cell = cls ()
29+ cell ._parse_stru (stru_file )
30+ cell ._built = True
31+ return cell
32+
33+ def build (self ):
34+ if self .atom is None :
35+ raise ValueError ("Atom information must be set before building." )
36+
37+ if isinstance (self .atom , str ):
38+ if self .atom .endswith ('.xyz' ):
39+ self ._parse_xyz (self .atom )
40+ else :
41+ raise ValueError ("Unsupported file format. Use .xyz files or provide atom list directly." )
42+ elif isinstance (self .atom , list ):
43+ self .atoms = [[atom [0 ], np .array (atom [1 ])] for atom in self .atom ]
44+ else :
45+ raise ValueError ("Unsupported atom format." )
46+
47+ # Automatically set parameters based on precision
48+ self ._set_auto_parameters ()
49+
50+ self ._built = True
51+
52+ def _set_auto_parameters (self ):
53+ if self .a is not None :
54+ self .mesh = [int (np .ceil (np .linalg .norm (v ) / self .precision )) for v in self .a ] # TODO: Check the formula!
55+ else :
56+ self .mesh = [10 , 10 , 10 ] # Default mesh if lattice vectors are not set
57+
58+ self .ke_cutoff = - np .log (self .precision ) * 10 # TODO: Check the formula!
59+ self .rcut = - np .log (self .precision ) * 2 # TODO: Check the formula!
60+
61+ def _parse_stru (self , stru_file ):
62+ self .atoms = []
63+ with open (stru_file , 'r' ) as f :
64+ lines = f .readlines ()
65+ i = 0
66+ while i < len (lines ):
67+ line = lines [i ].strip ()
68+ if 'ATOMIC_SPECIES' in line :
69+ i += 1
70+ while i < len (lines ) and lines [i ].strip ():
71+ parts = lines [i ].split ()
72+ if len (parts ) == 3 :
73+ species , mass , pp_file = parts
74+ pp_file = pp_file .lstrip ('./' )
75+ self .pseudo_potentials [species ] = {
76+ 'mass' : float (mass ),
77+ 'pseudo_file' : pp_file
78+ }
79+ i += 1
80+ elif 'NUMERICAL_ORBITAL' in line :
81+ i += 1
82+ while i < len (lines ) and lines [i ].strip ():
83+ orbital = lines [i ].split ()
84+ self .orbitals .append (orbital )
85+ i += 1
86+ elif 'LATTICE_CONSTANT' in line :
87+ i += 1
88+ self .lattice_constant = float (lines [i ].strip ())
89+ i += 1
90+ elif 'LATTICE_VECTORS' in line :
91+ self .a = np .array ([
92+ list (map (float , lines [i + 1 ].split ())),
93+ list (map (float , lines [i + 2 ].split ())),
94+ list (map (float , lines [i + 3 ].split ()))
95+ ])
96+ i += 4
97+ elif 'ATOMIC_POSITIONS' in line :
98+ i += 3
99+ while i < len (lines ) and lines [i ].strip ():
100+ species = lines [i ].strip ()
101+ i += 2
102+ num_atoms = int (lines [i ].strip ())
103+ i += 1
104+ for _ in range (num_atoms ):
105+ pos = np .array (list (map (float , lines [i ].split ()[:3 ])))
106+ self .atoms .append ([species , pos ])
107+ i += 2
108+ else :
109+ i += 1
110+
111+ def _parse_xyz (self , xyz_file ):
112+ self .atoms = []
113+ with open (xyz_file , 'r' ) as f :
114+ lines = f .readlines ()
115+ num_atoms = int (lines [0 ])
116+ # Skip the comment line
117+ for line in lines [2 :2 + num_atoms ]:
118+ parts = line .split ()
119+ species = parts [0 ]
120+ coords = np .array (list (map (float , parts [1 :4 ])))
121+ self .atoms .append ([species , coords ])
122+
123+ def get_atom_positions (self ):
124+ if not self ._built :
125+ raise RuntimeError ("Cell has not been built. Call build() first." )
126+ return np .array ([atom [1 ] for atom in self .atoms ])
127+
128+ def get_atom_species (self ):
129+ if not self ._built :
130+ raise RuntimeError ("Cell has not been built. Call build() first." )
131+ return [atom [0 ] for atom in self .atoms ]
132+
133+ @property
134+ def unit (self ):
135+ return self ._unit
136+
137+ @unit .setter
138+ def unit (self , value ):
139+ if value .lower () in ['angstrom' , 'a' ]:
140+ self ._unit = 'Angstrom'
141+ elif value .lower () in ['bohr' , 'b' , 'au' ]:
142+ self ._unit = 'Bohr'
143+ else :
144+ raise ValueError ("Unit must be 'Angstrom' or 'Bohr'" )
145+
146+ @property
147+ def lattice_constant (self ):
148+ return self ._lattice_constant
149+
150+ @lattice_constant .setter
151+ def lattice_constant (self , value ):
152+ self ._lattice_constant = value
153+
154+ @property
155+ def precision (self ):
156+ return self ._precision
157+
158+ @precision .setter
159+ def precision (self , value ):
160+ if value <= 0 :
161+ raise ValueError ("Precision must be a positive number" )
162+ self ._precision = value
163+
164+ @property
165+ def kspace (self ):
166+ return self ._kspace
167+
168+ @kspace .setter
169+ def kspace (self , value ):
170+ if value <= 0 :
171+ raise ValueError ("k-space must be a positive number" )
172+ self ._kspace = value
173+
174+ def make_kpts (self , mesh , with_gamma_point = True ):
175+ if self .a is None :
176+ raise ValueError ("Lattice vectors (self.a) must be set before generating k-points." )
177+
178+ kpts = []
179+ for i in range (mesh [0 ]):
180+ for j in range (mesh [1 ]):
181+ for k in range (mesh [2 ]):
182+ if with_gamma_point :
183+ kpt = np .array ([i / mesh [0 ], j / mesh [1 ], k / mesh [2 ]])
184+ else :
185+ kpt = np .array ([(i + 0.5 )/ mesh [0 ], (j + 0.5 )/ mesh [1 ], (k + 0.5 )/ mesh [2 ]])
186+ kpts .append (kpt )
187+
188+ # Convert to cartesian coordinates
189+ recip_lattice = 2 * np .pi * np .linalg .inv (self .a .T )
190+ kpts = np .dot (kpts , recip_lattice )
191+
192+ return np .array (kpts )
0 commit comments