Coverage for /builds/kinetik161/ase/ase/calculators/openmx/openmx.py: 27.17%

449 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-12-10 11:04 +0000

1""" 

2 The ASE Calculator for OpenMX <http://www.openmx-square.org> 

3 A Python interface to the software package for nano-scale 

4 material simulations based on density functional theories. 

5 Copyright (C) 2017 Charles Thomas Johnson, Jae Hwan Shim and JaeJun Yu 

6 

7 This program is free software: you can redistribute it and/or modify 

8 it under the terms of the GNU Lesser General Public License as published by 

9 the Free Software Foundation, either version 2.1 of the License, or 

10 (at your option) any later version. 

11 

12 This program is distributed in the hope that it will be useful, 

13 but WITHOUT ANY WARRANTY; without even the implied warranty of 

14 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

15 GNU Lesser General Public License for more details. 

16 

17 You should have received a copy of the GNU Lesser General Public License 

18 along with ASE. If not, see <http://www.gnu.org/licenses/>. 

19 

20""" 

21 

22import os 

23import re 

24import subprocess 

25import time 

26import warnings 

27 

28import numpy as np 

29 

30from ase.calculators.calculator import (Calculator, FileIOCalculator, 

31 all_changes, equal, 

32 kptdensity2monkhorstpack) 

33from ase.calculators.openmx.default_settings import default_dictionary 

34from ase.calculators.openmx.parameters import OpenMXParameters 

35from ase.calculators.openmx.reader import get_file_name, read_openmx 

36from ase.calculators.openmx.writer import write_openmx 

37from ase.geometry import cell_to_cellpar 

38 

39 

40def parse_omx_version(txt): 

41 """Parse version number from stdout header.""" 

42 match = re.search(r'Welcome to OpenMX\s+Ver\.\s+(\S+)', txt, re.M) 

43 return match.group(1) 

44 

45 

46class OpenMX(FileIOCalculator): 

47 """ 

48 Calculator interface to the OpenMX code. 

49 """ 

50 

51 implemented_properties = [ 

52 'free_energy', # Same value with energy 

53 'energy', 

54 'energies', 

55 'forces', 

56 'stress', 

57 'dipole', 

58 'chemical_potential', 

59 'magmom', 

60 'magmoms', 

61 'eigenvalues'] 

62 

63 default_parameters = OpenMXParameters() 

64 

65 default_pbs = { 

66 'processes': 1, 

67 'walltime': "10:00:00", 

68 'threads': 1, 

69 'nodes': 1 

70 } 

71 

72 default_mpi = { 

73 'processes': 1, 

74 'threads': 1 

75 } 

76 

77 default_output_setting = { 

78 'nohup': True, 

79 'debug': False 

80 } 

81 

82 def __init__(self, restart=None, 

83 ignore_bad_restart_file=FileIOCalculator._deprecated, 

84 label='./openmx', atoms=None, command=None, mpi=None, 

85 pbs=None, **kwargs): 

86 

87 # Initialize and put the default parameters. 

88 self.initialize_pbs(pbs) 

89 self.initialize_mpi(mpi) 

90 self.initialize_output_setting(**kwargs) 

91 

92 FileIOCalculator.__init__(self, restart, ignore_bad_restart_file, 

93 label, atoms, command, **kwargs) 

94 

95 def __getitem__(self, key): 

96 """Convenience method to retrieve a parameter as 

97 calculator[key] rather than calculator.parameters[key] 

98 

99 Parameters: 

100 -key : str, the name of the parameters to get. 

101 """ 

102 return self.parameters[key] 

103 

104 def __setitem__(self, key, value): 

105 self.parameters[key] = value 

106 

107 def initialize_output_setting(self, **kwargs): 

108 output_setting = {} 

109 self.output_setting = dict(self.default_output_setting) 

110 for key, value in kwargs.items(): 

111 if key in self.default_output_setting: 

112 output_setting[key] = value 

113 self.output_setting.update(output_setting) 

114 self.__dict__.update(self.output_setting) 

115 

116 def initialize_pbs(self, pbs): 

117 if pbs: 

118 self.pbs = dict(self.default_pbs) 

119 for key in pbs: 

120 if key not in self.default_pbs: 

121 allowed = ', '.join(list(self.default_pbs.keys())) 

122 raise TypeError('Unexpected keyword "{}" in "pbs" ' 

123 'dictionary. Must be one of: {}' 

124 .format(key, allowed)) 

125 # Put dictionary into python variable 

126 self.pbs.update(pbs) 

127 self.__dict__.update(self.pbs) 

128 else: 

129 self.pbs = None 

130 

131 def initialize_mpi(self, mpi): 

132 if mpi: 

133 self.mpi = dict(self.default_mpi) 

134 for key in mpi: 

135 if key not in self.default_mpi: 

136 allowed = ', '.join(list(self.default_mpi.keys())) 

137 raise TypeError('Unexpected keyword "{}" in "mpi" ' 

138 'dictionary. Must be one of: {}' 

139 .format(key, allowed)) 

140 # Put dictionary into python variable 

141 self.mpi.update(mpi) 

142 self.__dict__.update(self.mpi) 

143 else: 

144 self.mpi = None 

145 

146 def run(self): 

147 '''Check Which Running method we r going to use and run it''' 

148 if self.pbs is not None: 

149 run = self.run_pbs 

150 elif self.mpi is not None: 

151 run = self.run_mpi 

152 else: 

153 run = self.run_openmx 

154 run() 

155 

156 def run_openmx(self): 

157 def isRunning(process=None): 

158 ''' Check mpi is running''' 

159 return process.poll() is None 

160 runfile = get_file_name('.dat', self.label, absolute_directory=False) 

161 outfile = get_file_name('.log', self.label) 

162 olddir = os.getcwd() 

163 abs_dir = os.path.join(olddir, self.directory) 

164 try: 

165 os.chdir(abs_dir) 

166 if self.command is None: 

167 self.command = 'openmx' 

168 command = self.command + ' %s > %s' 

169 command = command % (runfile, outfile) 

170 self.prind(command) 

171 p = subprocess.Popen(command, shell=True, universal_newlines=True) 

172 self.print_file(file=outfile, running=isRunning, process=p) 

173 finally: 

174 os.chdir(olddir) 

175 self.prind("Calculation Finished") 

176 

177 def run_mpi(self): 

178 """ 

179 Run openmx using MPI method. If keyword `mpi` is declared, it will 

180 run. 

181 """ 

182 def isRunning(process=None): 

183 ''' Check mpi is running''' 

184 return process.poll() is None 

185 processes = self.processes 

186 threads = self.threads 

187 runfile = get_file_name('.dat', self.label, absolute_directory=False) 

188 outfile = get_file_name('.log', self.label) 

189 olddir = os.getcwd() 

190 abs_dir = os.path.join(olddir, self.directory) 

191 try: 

192 os.chdir(abs_dir) 

193 command = self.get_command(processes, threads, runfile, outfile) 

194 self.prind(command) 

195 p = subprocess.Popen(command, shell=True, universal_newlines=True) 

196 self.print_file(file=outfile, running=isRunning, process=p) 

197 finally: 

198 os.chdir(olddir) 

199 self.prind("Calculation Finished") 

200 

201 def run_pbs(self, prefix='test'): 

202 """ 

203 Execute the OpenMX using Plane Batch System. In order to use this, 

204 Your system should have Scheduler. PBS 

205 Basically, it does qsub. and wait until qstat signal shows c 

206 Super computer user 

207 """ 

208 nodes = self.nodes 

209 processes = self.processes 

210 

211 prefix = self.prefix 

212 olddir = os.getcwd() 

213 try: 

214 os.chdir(self.abs_directory) 

215 except AttributeError: 

216 os.chdir(self.directory) 

217 

218 def isRunning(jobNum=None, status='Q', qstat='qstat'): 

219 """ 

220 Check submitted job is still Running 

221 """ 

222 def runCmd(exe): 

223 p = subprocess.Popen(exe, stdout=subprocess.PIPE, 

224 stderr=subprocess.STDOUT, 

225 universal_newlines=True) 

226 while True: 

227 line = p.stdout.readline() 

228 if line != '': 

229 # the real code does filtering here 

230 yield line.rstrip() 

231 else: 

232 break 

233 jobs = runCmd('qstat') 

234 columns = None 

235 for line in jobs: 

236 if str(jobNum) in line: 

237 columns = line.split() 

238 self.prind(line) 

239 if columns is not None: 

240 return columns[-2] == status 

241 else: 

242 return False 

243 

244 inputfile = self.label + '.dat' 

245 outfile = self.label + '.log' 

246 

247 bashArgs = "#!/bin/bash \n cd $PBS_O_WORKDIR\n" 

248 jobName = prefix 

249 cmd = bashArgs + \ 

250 'mpirun -hostfile $PBS_NODEFILE openmx {} > {}'.format( 

251 inputfile, outfile) 

252 echoArgs = ["echo", f"$' {cmd}'"] 

253 qsubArgs = ["qsub", "-N", jobName, "-l", "nodes=%d:ppn=%d" % 

254 (nodes, processes), "-l", "walltime=" + self.walltime] 

255 wholeCmd = " ".join(echoArgs) + " | " + " ".join(qsubArgs) 

256 self.prind(wholeCmd) 

257 out = subprocess.Popen(wholeCmd, shell=True, 

258 stdout=subprocess.PIPE, universal_newlines=True) 

259 out = out.communicate()[0] 

260 jobNum = int(re.match(r'(\d+)', out.split()[0]).group(1)) 

261 

262 self.prind('Queue number is ' + str(jobNum) + 

263 '\nWaiting for the Queue to start') 

264 while isRunning(jobNum, status='Q'): 

265 time.sleep(5) 

266 self.prind('.') 

267 self.prind('Start Calculating') 

268 self.print_file(file=outfile, running=isRunning, 

269 jobNum=jobNum, status='R', qstat='qstat') 

270 

271 os.chdir(olddir) 

272 self.prind('Calculation Finished!') 

273 return jobNum 

274 

275 def clean(self, prefix='test', queue_num=None): 

276 """Method which cleans up after a calculation. 

277 

278 The default files generated OpenMX will be deleted IF this 

279 method is called. 

280 

281 """ 

282 self.prind("Cleaning Data") 

283 fileName = get_file_name('', self.label) 

284 pbs_Name = get_file_name('', self.label) 

285 files = [ 

286 # prefix+'.out',#prefix+'.dat',#prefix+'.BAND*', 

287 fileName + '.cif', fileName + '.dden.cube', fileName + \ 

288 '.ene', fileName + '.md', fileName + '.md2', 

289 fileName + '.tden.cube', fileName + '.sden.cube', fileName + \ 

290 '.v0.cube', fileName + '.v1.cube', 

291 fileName + '.vhart.cube', fileName + '.den0.cube', fileName + \ 

292 '.bulk.xyz', fileName + '.den1.cube', 

293 fileName + '.xyz', pbs_Name + '.o' + \ 

294 str(queue_num), pbs_Name + '.e' + str(queue_num) 

295 ] 

296 for f in files: 

297 try: 

298 self.prind("Removing" + f) 

299 os.remove(f) 

300 except OSError: 

301 self.prind("There is no such file named " + f) 

302 

303 def calculate(self, atoms=None, properties=None, 

304 system_changes=all_changes): 

305 """ 

306 Capture the RuntimeError from FileIOCalculator.calculate 

307 and add a little debug information from the OpenMX output. 

308 See base FileIOCalculator for documentation. 

309 """ 

310 if self.parameters.data_path is None: 

311 if 'OPENMX_DFT_DATA_PATH' not in os.environ: 

312 warnings.warn('Please either set OPENMX_DFT_DATA_PATH as an' 

313 'enviroment variable or specify "data_path" as' 

314 'a keyword argument') 

315 

316 self.prind("Start Calculation") 

317 if properties is None: 

318 properties = self.implemented_properties 

319 try: 

320 Calculator.calculate(self, atoms, properties, system_changes) 

321 self.write_input(atoms=self.atoms, parameters=self.parameters, 

322 properties=properties, 

323 system_changes=system_changes) 

324 self.print_input(debug=self.debug, nohup=self.nohup) 

325 self.run() 

326 # self.read_results() 

327 self.version = self.read_version() 

328 output_atoms = read_openmx(filename=self.label, debug=self.debug) 

329 self.output_atoms = output_atoms 

330 # XXX The parameters are supposedly inputs, so it is dangerous 

331 # to update them from the outputs. --askhl 

332 self.parameters.update(output_atoms.calc.parameters) 

333 self.results = output_atoms.calc.results 

334 # self.clean() 

335 except RuntimeError as e: 

336 try: 

337 with open(get_file_name('.log')) as fd: 

338 lines = fd.readlines() 

339 debug_lines = 10 

340 print('##### %d last lines of the OpenMX output' % debug_lines) 

341 for line in lines[-20:]: 

342 print(line.strip()) 

343 print('##### end of openMX output') 

344 raise e 

345 except RuntimeError as e: 

346 raise e 

347 

348 def write_input(self, atoms=None, parameters=None, 

349 properties=[], system_changes=[]): 

350 """Write input (dat)-file. 

351 See calculator.py for further details. 

352 

353 Parameters: 

354 - atoms : The Atoms object to write. 

355 - properties : The properties which should be calculated. 

356 - system_changes : List of properties changed since last run. 

357 """ 

358 # Call base calculator. 

359 if atoms is None: 

360 atoms = self.atoms 

361 FileIOCalculator.write_input(self, atoms, properties, system_changes) 

362 write_openmx(label=self.label, atoms=atoms, parameters=self.parameters, 

363 properties=properties, system_changes=system_changes) 

364 

365 def print_input(self, debug=None, nohup=None): 

366 """ 

367 For a debugging purpose, print the .dat file 

368 """ 

369 if debug is None: 

370 debug = self.debug 

371 if nohup is None: 

372 nohup = self.nohup 

373 self.prind('Reading input file' + self.label) 

374 filename = get_file_name('.dat', self.label) 

375 if not nohup: 

376 with open(filename) as fd: 

377 while True: 

378 line = fd.readline() 

379 print(line.strip()) 

380 if not line: 

381 break 

382 

383 def read(self, label): 

384 self.parameters = {} 

385 self.set_label(label) 

386 if label[-5:] in ['.dat', '.out', '.log']: 

387 label = label[:-4] 

388 atoms = read_openmx(filename=label, debug=self.debug) 

389 self.update_atoms(atoms) 

390 self.parameters.update(atoms.calc.parameters) 

391 self.results = atoms.calc.results 

392 self.parameters['restart'] = self.label 

393 self.parameters['label'] = label 

394 

395 def read_version(self, label=None): 

396 version = None 

397 if label is None: 

398 label = self.label 

399 for line in open(get_file_name('.out', label)): 

400 if line.find('Ver.') != -1: 

401 version = line.split()[-1] 

402 break 

403 return version 

404 

405 def update_atoms(self, atoms): 

406 self.atoms = atoms.copy() 

407 

408 def set(self, **kwargs): 

409 """Set all parameters. 

410 

411 Parameters: 

412 -kwargs : Dictionary containing the keywords defined in 

413 OpenMXParameters. 

414 """ 

415 

416 for key, value in kwargs.items(): 

417 if key not in self.default_parameters.keys(): 

418 raise KeyError(f'Unkown keyword "{key}" and value "{value}".') 

419 if key == 'xc' and value not in self.default_parameters.allowed_xc: 

420 raise KeyError(f'Given xc "{value}" is not allowed') 

421 if key in ['dat_arguments'] and isinstance(value, dict): 

422 # For values that are dictionaries, verify subkeys, too. 

423 default_dict = self.default_parameters[key] 

424 for subkey in kwargs[key]: 

425 if subkey not in default_dict: 

426 allowed = ', '.join(list(default_dict.keys())) 

427 raise TypeError('Unknown subkeyword "{}" of keyword ' 

428 '"{}". Must be one of: {}' 

429 .format(subkey, key, allowed)) 

430 

431 # Find out what parameter has been changed 

432 changed_parameters = {} 

433 for key, value in kwargs.items(): 

434 oldvalue = self.parameters.get(key) 

435 if key not in self.parameters or not equal(value, oldvalue): 

436 changed_parameters[key] = value 

437 self.parameters[key] = value 

438 

439 # Set the parameters 

440 for key, value in kwargs.items(): 

441 # print(' Setting the %s as %s'%(key, value)) 

442 self.parameters[key] = value 

443 

444 # If Changed Parameter is Critical, we have to reset the results 

445 for key, value in changed_parameters.items(): 

446 if key in ['xc', 'kpts', 'energy_cutoff']: 

447 self.results = {} 

448 

449 value = kwargs.get('energy_cutoff') 

450 if value is not None and not (isinstance(value, (float, int)) 

451 and value > 0): 

452 mess = "'{}' must be a positive number(in eV), \ 

453 got '{}'".format('energy_cutoff', value) 

454 raise ValueError(mess) 

455 

456 atoms = kwargs.get('atoms') 

457 if atoms is not None and self.atoms is None: 

458 self.atoms = atoms.copy() 

459 

460 def set_results(self, results): 

461 # Not Implemented fully 

462 self.results.update(results) 

463 

464 def get_command(self, processes, threads, runfile=None, outfile=None): 

465 # Contruct the command to send to the operating system 

466 abs_dir = os.getcwd() 

467 command = '' 

468 self.prind(self.command) 

469 if self.command is None: 

470 self.command = 'openmx' 

471 # run processes specified by the system variable OPENMX_COMMAND 

472 if processes is None: 

473 command += os.environ.get('OPENMX_COMMAND') 

474 if command is None: 

475 warnings.warn('Either specify OPENMX_COMMAND as an environment\ 

476 variable or specify processes as a keyword argument') 

477 else: # run with a specified number of processes 

478 threads_string = ' -nt ' + str(threads) 

479 if threads is None: 

480 threads_string = '' 

481 command += 'mpirun -np ' + \ 

482 str(processes) + ' ' + self.command + \ 

483 ' %s ' + threads_string + ' |tee %s' 

484 # str(processes) + ' openmx %s' + threads_string + ' > %s' 

485 

486 if runfile is None: 

487 runfile = os.path.join(abs_dir, f'{self.prefix} .dat') 

488 if outfile is None: 

489 outfile = os.path.join(abs_dir, f'{self.prefix} .log') 

490 try: 

491 command = command % (runfile, outfile) 

492 # command += '" > ./%s &' % outfile # outputs 

493 except TypeError: # in case the OPENMX_COMMAND is incompatible 

494 raise ValueError( 

495 "The 'OPENMX_COMMAND' environment must " + 

496 "be a format string" + 

497 " with four string arguments.\n" + 

498 "Example : 'mpirun -np 4 openmx ./%s -nt 2 > ./%s'.\n" + 

499 f"Got '{command}'") 

500 return command 

501 

502 def get_stress(self, atoms=None): 

503 if atoms is None: 

504 atoms = self.atoms 

505 

506 # Note: Stress is only supported from OpenMX 3.8+. 

507 stress = self.get_property('stress', atoms) 

508 

509 return stress 

510 

511 def get_band_structure(self, atoms=None, calc=None): 

512 """ 

513 This is band structure function. It is compatible to 

514 ase dft module """ 

515 from ase.dft import band_structure 

516 if isinstance(self['kpts'], tuple): 

517 self['kpts'] = self.get_kpoints(band_kpath=self['band_kpath']) 

518 return band_structure.get_band_structure(self.atoms, self, ) 

519 

520 def get_bz_k_points(self): 

521 kgrid = self['kpts'] 

522 if type(kgrid) in [int, float]: 

523 kgrid = kptdensity2monkhorstpack(self.atoms, kgrid, False) 

524 bz_k_points = [] 

525 n1 = kgrid[0] 

526 n2 = kgrid[1] 

527 n3 = kgrid[2] 

528 for i in range(n1): 

529 for j in range(n2): 

530 # Monkhorst Pack Grid [H.J. Monkhorst and J.D. Pack, 

531 # Phys. Rev. B 13, 5188 (1976)] 

532 for k in range(n3): 

533 bz_k_points.append((0.5 * float(2 * i - n1 + 1) / n1, 

534 0.5 * float(2 * j - n2 + 1) / n2, 

535 0.5 * float(2 * k - n3 + 1) / n3)) 

536 return np.array(bz_k_points) 

537 

538 def get_ibz_k_points(self): 

539 if self['band_kpath'] is None: 

540 return self.get_bz_k_points() 

541 else: 

542 return self.get_kpoints(band_kpath=self['band_kpath']) 

543 

544 def get_kpoints(self, kpts=None, symbols=None, band_kpath=None, eps=1e-5): 

545 """Convert band_kpath <-> kpts""" 

546 if kpts is None: 

547 kpts = [] 

548 band_kpath = np.array(band_kpath) 

549 band_nkpath = len(band_kpath) 

550 for i, kpath in enumerate(band_kpath): 

551 end = False 

552 nband = int(kpath[0]) 

553 if band_nkpath == i: 

554 end = True 

555 nband += 1 

556 ini = np.array(kpath[1:4], dtype=float) 

557 fin = np.array(kpath[4:7], dtype=float) 

558 x = np.linspace(ini[0], fin[0], nband, endpoint=end) 

559 y = np.linspace(ini[1], fin[1], nband, endpoint=end) 

560 z = np.linspace(ini[2], fin[2], nband, endpoint=end) 

561 kpts.extend(np.array([x, y, z]).T) 

562 return np.array(kpts, dtype=float) 

563 elif band_kpath is None: 

564 band_kpath = [] 

565 points = np.asarray(kpts) 

566 diffs = points[1:] - points[:-1] 

567 kinks = abs(diffs[1:] - diffs[:-1]).sum(1) > eps 

568 N = len(points) 

569 indices = [0] 

570 indices.extend(np.arange(1, N - 1)[kinks]) 

571 indices.append(N - 1) 

572 for start, end, s_sym, e_sym in zip(indices[1:], indices[:-1], 

573 symbols[1:], symbols[:-1]): 

574 band_kpath.append({'start_point': start, 'end_point': end, 

575 'kpts': 20, 

576 'path_symbols': (s_sym, e_sym)}) 

577 return band_kpath 

578 

579 def get_lattice_type(self): 

580 cellpar = cell_to_cellpar(self.atoms.cell) 

581 abc = cellpar[:3] 

582 angles = cellpar[3:] 

583 min_lv = min(abc) 

584 if abc.ptp() < 0.01 * min_lv: 

585 if abs(angles - 90).max() < 1: 

586 return 'cubic' 

587 elif abs(angles - 60).max() < 1: 

588 return 'fcc' 

589 elif abs(angles - np.arccos(-1 / 3.) * 180 / np.pi).max < 1: 

590 return 'bcc' 

591 elif abs(angles - 90).max() < 1: 

592 if abs(abc[0] - abc[1]).min() < 0.01 * min_lv: 

593 return 'tetragonal' 

594 else: 

595 return 'orthorhombic' 

596 elif abs(abc[0] - abc[1]) < 0.01 * min_lv and \ 

597 abs(angles[2] - 120) < 1 and abs(angles[:2] - 90).max() < 1: 

598 return 'hexagonal' 

599 else: 

600 return 'not special' 

601 

602 def get_number_of_spins(self): 

603 try: 

604 magmoms = self.atoms.get_initial_magnetic_moments() 

605 if self['scf_spinpolarization'] is None: 

606 if isinstance(magmoms[0], float): 

607 if abs(magmoms).max() < 0.1: 

608 return 1 

609 else: 

610 return 2 

611 else: 

612 raise NotImplementedError 

613 else: 

614 if self['scf_spinpolarization'] == 'on': 

615 return 2 

616 elif self['scf_spinpolarization'] == 'nc' or \ 

617 np.any(self['initial_magnetic_moments_euler_angles']) \ 

618 is not None: 

619 return 1 

620 except KeyError: 

621 return 1 

622 

623 def get_eigenvalues(self, kpt=None, spin=None): 

624 if self.results.get('eigenvalues') is None: 

625 self.calculate(self.atoms) 

626 if kpt is None and spin is None: 

627 return self.results['eigenvalues'] 

628 else: 

629 return self.results['eigenvalues'][spin, kpt, :] 

630 

631 def get_fermi_level(self): 

632 try: 

633 fermi_level = self.results['chemical_potential'] 

634 except KeyError: 

635 self.calculate() 

636 fermi_level = self.results['chemical_potential'] 

637 return fermi_level 

638 

639 def get_number_of_bands(self): 

640 pag = self.parameters.get 

641 dfd = default_dictionary 

642 if 'number_of_bands' not in self.results: 

643 n = 0 

644 for atom in self.atoms: 

645 sym = atom.symbol 

646 orbitals = pag('dft_data_dict', dfd)[sym]['orbitals used'] 

647 d = 1 

648 for orbital in orbitals: 

649 n += d * orbital 

650 d += 2 

651 self.results['number_of_bands'] = n 

652 return self.results['number_of_bands'] 

653 

654 def dirG(self, dk, bzone=(0, 0, 0)): 

655 nx, ny, nz = self['wannier_kpts'] 

656 dx = dk // (ny * nz) + bzone[0] * nx 

657 dy = (dk // nz) % ny + bzone[1] * ny 

658 dz = dk % nz + bzone[2] * nz 

659 return dx, dy, dz 

660 

661 def dk(self, dirG): 

662 dx, dy, dz = dirG 

663 nx, ny, nz = self['wannier_kpts'] 

664 return ny * nz * (dx % nx) + nz * (dy % ny) + dz % nz 

665 

666 def get_wannier_localization_matrix(self, nbands, dirG, nextkpoint=None, 

667 kpoint=None, spin=0, G_I=(0, 0, 0)): 

668 # only expected to work for no spin polarization 

669 try: 

670 self['bloch_overlaps'] 

671 except KeyError: 

672 self.read_bloch_overlaps() 

673 dirG = tuple(dirG) 

674 nx, ny, nz = self['wannier_kpts'] 

675 nr3 = nx * ny * nz 

676 if kpoint is None and nextkpoint is None: 

677 return {kpoint: self['bloch_overlaps' 

678 ][kpoint][dirG][:nbands, :nbands 

679 ] for kpoint in range(nr3)} 

680 if kpoint is None: 

681 kpoint = (nextkpoint - self.dk(dirG)) % nr3 

682 if nextkpoint is None: 

683 nextkpoint = (kpoint + self.dk(dirG)) % nr3 

684 if dirG not in self['bloch_overlaps'][kpoint].keys(): 

685 return np.zeros((nbands, nbands), complex) 

686 return self['bloch_overlaps'][kpoint][dirG][:nbands, :nbands] 

687 

688 def prind(self, line, debug=None): 

689 ''' Print the value if debugging mode is on. 

690 Otherwise, it just ignored''' 

691 if debug is None: 

692 debug = self.debug 

693 if debug: 

694 print(line) 

695 

696 def print_file(self, file=None, running=None, **args): 

697 ''' Print the file while calculation is running''' 

698 prev_position = 0 

699 last_position = 0 

700 while not os.path.isfile(file): 

701 self.prind(f'Waiting for {file} to come out') 

702 time.sleep(5) 

703 with open(file) as fd: 

704 while running(**args): 

705 fd.seek(last_position) 

706 new_data = fd.read() 

707 prev_position = fd.tell() 

708 # self.prind('pos', prev_position != last_position) 

709 if prev_position != last_position: 

710 if not self.nohup: 

711 print(new_data) 

712 last_position = prev_position 

713 time.sleep(1)