Coverage for /builds/hweiske/ase/ase/calculators/openmx/openmx.py: 27.33%

450 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2024-04-22 11:22 +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.config import cfg 

38from ase.geometry import cell_to_cellpar 

39 

40 

41def parse_omx_version(txt): 

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

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

44 return match.group(1) 

45 

46 

47class OpenMX(FileIOCalculator): 

48 """ 

49 Calculator interface to the OpenMX code. 

50 """ 

51 

52 implemented_properties = [ 

53 'free_energy', # Same value with energy 

54 'energy', 

55 'energies', 

56 'forces', 

57 'stress', 

58 'dipole', 

59 'chemical_potential', 

60 'magmom', 

61 'magmoms', 

62 'eigenvalues'] 

63 

64 default_parameters = OpenMXParameters() 

65 

66 default_pbs = { 

67 'processes': 1, 

68 'walltime': "10:00:00", 

69 'threads': 1, 

70 'nodes': 1 

71 } 

72 

73 default_mpi = { 

74 'processes': 1, 

75 'threads': 1 

76 } 

77 

78 default_output_setting = { 

79 'nohup': True, 

80 'debug': False 

81 } 

82 

83 def __init__(self, restart=None, 

84 ignore_bad_restart_file=FileIOCalculator._deprecated, 

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

86 pbs=None, **kwargs): 

87 

88 # Initialize and put the default parameters. 

89 self.initialize_pbs(pbs) 

90 self.initialize_mpi(mpi) 

91 self.initialize_output_setting(**kwargs) 

92 

93 FileIOCalculator.__init__(self, restart, ignore_bad_restart_file, 

94 label, atoms, command, **kwargs) 

95 

96 def __getitem__(self, key): 

97 """Convenience method to retrieve a parameter as 

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

99 

100 Parameters: 

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

102 """ 

103 return self.parameters[key] 

104 

105 def __setitem__(self, key, value): 

106 self.parameters[key] = value 

107 

108 def initialize_output_setting(self, **kwargs): 

109 output_setting = {} 

110 self.output_setting = dict(self.default_output_setting) 

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

112 if key in self.default_output_setting: 

113 output_setting[key] = value 

114 self.output_setting.update(output_setting) 

115 self.__dict__.update(self.output_setting) 

116 

117 def initialize_pbs(self, pbs): 

118 if pbs: 

119 self.pbs = dict(self.default_pbs) 

120 for key in pbs: 

121 if key not in self.default_pbs: 

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

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

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

125 .format(key, allowed)) 

126 # Put dictionary into python variable 

127 self.pbs.update(pbs) 

128 self.__dict__.update(self.pbs) 

129 else: 

130 self.pbs = None 

131 

132 def initialize_mpi(self, mpi): 

133 if mpi: 

134 self.mpi = dict(self.default_mpi) 

135 for key in mpi: 

136 if key not in self.default_mpi: 

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

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

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

140 .format(key, allowed)) 

141 # Put dictionary into python variable 

142 self.mpi.update(mpi) 

143 self.__dict__.update(self.mpi) 

144 else: 

145 self.mpi = None 

146 

147 def run(self): 

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

149 if self.pbs is not None: 

150 run = self.run_pbs 

151 elif self.mpi is not None: 

152 run = self.run_mpi 

153 else: 

154 run = self.run_openmx 

155 run() 

156 

157 def run_openmx(self): 

158 def isRunning(process=None): 

159 ''' Check mpi is running''' 

160 return process.poll() is None 

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

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

163 olddir = os.getcwd() 

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

165 try: 

166 os.chdir(abs_dir) 

167 if self.command is None: 

168 self.command = 'openmx' 

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

170 command = command % (runfile, outfile) 

171 self.prind(command) 

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

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

174 finally: 

175 os.chdir(olddir) 

176 self.prind("Calculation Finished") 

177 

178 def run_mpi(self): 

179 """ 

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

181 run. 

182 """ 

183 def isRunning(process=None): 

184 ''' Check mpi is running''' 

185 return process.poll() is None 

186 processes = self.processes 

187 threads = self.threads 

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

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

190 olddir = os.getcwd() 

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

192 try: 

193 os.chdir(abs_dir) 

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

195 self.prind(command) 

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

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

198 finally: 

199 os.chdir(olddir) 

200 self.prind("Calculation Finished") 

201 

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

203 """ 

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

205 Your system should have Scheduler. PBS 

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

207 Super computer user 

208 """ 

209 nodes = self.nodes 

210 processes = self.processes 

211 

212 prefix = self.prefix 

213 olddir = os.getcwd() 

214 try: 

215 os.chdir(self.abs_directory) 

216 except AttributeError: 

217 os.chdir(self.directory) 

218 

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

220 """ 

221 Check submitted job is still Running 

222 """ 

223 def runCmd(exe): 

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

225 stderr=subprocess.STDOUT, 

226 universal_newlines=True) 

227 while True: 

228 line = p.stdout.readline() 

229 if line != '': 

230 # the real code does filtering here 

231 yield line.rstrip() 

232 else: 

233 break 

234 jobs = runCmd('qstat') 

235 columns = None 

236 for line in jobs: 

237 if str(jobNum) in line: 

238 columns = line.split() 

239 self.prind(line) 

240 if columns is not None: 

241 return columns[-2] == status 

242 else: 

243 return False 

244 

245 inputfile = self.label + '.dat' 

246 outfile = self.label + '.log' 

247 

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

249 jobName = prefix 

250 cmd = bashArgs + \ 

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

252 inputfile, outfile) 

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

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

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

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

257 self.prind(wholeCmd) 

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

259 stdout=subprocess.PIPE, universal_newlines=True) 

260 out = out.communicate()[0] 

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

262 

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

264 '\nWaiting for the Queue to start') 

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

266 time.sleep(5) 

267 self.prind('.') 

268 self.prind('Start Calculating') 

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

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

271 

272 os.chdir(olddir) 

273 self.prind('Calculation Finished!') 

274 return jobNum 

275 

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

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

278 

279 The default files generated OpenMX will be deleted IF this 

280 method is called. 

281 

282 """ 

283 self.prind("Cleaning Data") 

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

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

286 files = [ 

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

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

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

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

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

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

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

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

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

296 ] 

297 for f in files: 

298 try: 

299 self.prind("Removing" + f) 

300 os.remove(f) 

301 except OSError: 

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

303 

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

305 system_changes=all_changes): 

306 """ 

307 Capture the RuntimeError from FileIOCalculator.calculate 

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

309 See base FileIOCalculator for documentation. 

310 """ 

311 if self.parameters.data_path is None: 

312 if 'OPENMX_DFT_DATA_PATH' not in cfg: 

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

314 'enviroment variable or specify "data_path" as' 

315 'a keyword argument') 

316 

317 self.prind("Start Calculation") 

318 if properties is None: 

319 properties = self.implemented_properties 

320 try: 

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

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

323 properties=properties, 

324 system_changes=system_changes) 

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

326 self.run() 

327 # self.read_results() 

328 self.version = self.read_version() 

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

330 self.output_atoms = output_atoms 

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

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

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

334 self.results = output_atoms.calc.results 

335 # self.clean() 

336 except RuntimeError as e: 

337 try: 

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

339 lines = fd.readlines() 

340 debug_lines = 10 

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

342 for line in lines[-20:]: 

343 print(line.strip()) 

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

345 raise e 

346 except RuntimeError as e: 

347 raise e 

348 

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

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

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

352 See calculator.py for further details. 

353 

354 Parameters: 

355 - atoms : The Atoms object to write. 

356 - properties : The properties which should be calculated. 

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

358 """ 

359 # Call base calculator. 

360 if atoms is None: 

361 atoms = self.atoms 

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

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

364 properties=properties, system_changes=system_changes) 

365 

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

367 """ 

368 For a debugging purpose, print the .dat file 

369 """ 

370 if debug is None: 

371 debug = self.debug 

372 if nohup is None: 

373 nohup = self.nohup 

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

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

376 if not nohup: 

377 with open(filename) as fd: 

378 while True: 

379 line = fd.readline() 

380 print(line.strip()) 

381 if not line: 

382 break 

383 

384 def read(self, label): 

385 self.parameters = {} 

386 self.set_label(label) 

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

388 label = label[:-4] 

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

390 self.update_atoms(atoms) 

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

392 self.results = atoms.calc.results 

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

394 self.parameters['label'] = label 

395 

396 def read_version(self, label=None): 

397 version = None 

398 if label is None: 

399 label = self.label 

400 for line in open(get_file_name('.log', label)): 

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

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

403 break 

404 return version 

405 

406 def update_atoms(self, atoms): 

407 self.atoms = atoms.copy() 

408 

409 def set(self, **kwargs): 

410 """Set all parameters. 

411 

412 Parameters: 

413 -kwargs : Dictionary containing the keywords defined in 

414 OpenMXParameters. 

415 """ 

416 

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

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

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

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

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

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

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

424 default_dict = self.default_parameters[key] 

425 for subkey in kwargs[key]: 

426 if subkey not in default_dict: 

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

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

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

430 .format(subkey, key, allowed)) 

431 

432 # Find out what parameter has been changed 

433 changed_parameters = {} 

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

435 oldvalue = self.parameters.get(key) 

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

437 changed_parameters[key] = value 

438 self.parameters[key] = value 

439 

440 # Set the parameters 

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

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

443 self.parameters[key] = value 

444 

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

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

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

448 self.results = {} 

449 

450 value = kwargs.get('energy_cutoff') 

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

452 and value > 0): 

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

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

455 raise ValueError(mess) 

456 

457 atoms = kwargs.get('atoms') 

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

459 self.atoms = atoms.copy() 

460 

461 def set_results(self, results): 

462 # Not Implemented fully 

463 self.results.update(results) 

464 

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

466 # Contruct the command to send to the operating system 

467 abs_dir = os.getcwd() 

468 command = '' 

469 self.prind(self.command) 

470 if self.command is None: 

471 self.command = 'openmx' 

472 # run processes specified by the system variable OPENMX_COMMAND 

473 if processes is None: 

474 command += cfg.get('OPENMX_COMMAND') 

475 if command is None: 

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

477 variable or specify processes as a keyword argument') 

478 else: # run with a specified number of processes 

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

480 if threads is None: 

481 threads_string = '' 

482 command += 'mpirun -np ' + \ 

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

484 ' %s ' + threads_string + ' |tee %s' 

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

486 

487 if runfile is None: 

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

489 if outfile is None: 

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

491 try: 

492 command = command % (runfile, outfile) 

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

494 except TypeError: # in case the OPENMX_COMMAND is incompatible 

495 raise ValueError( 

496 "The 'OPENMX_COMMAND' environment must " + 

497 "be a format string" + 

498 " with four string arguments.\n" + 

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

500 f"Got '{command}'") 

501 return command 

502 

503 def get_stress(self, atoms=None): 

504 if atoms is None: 

505 atoms = self.atoms 

506 

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

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

509 

510 return stress 

511 

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

513 """ 

514 This is band structure function. It is compatible to 

515 ase dft module """ 

516 from ase.dft import band_structure 

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

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

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

520 

521 def get_bz_k_points(self): 

522 kgrid = self['kpts'] 

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

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

525 bz_k_points = [] 

526 n1 = kgrid[0] 

527 n2 = kgrid[1] 

528 n3 = kgrid[2] 

529 for i in range(n1): 

530 for j in range(n2): 

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

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

533 for k in range(n3): 

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

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

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

537 return np.array(bz_k_points) 

538 

539 def get_ibz_k_points(self): 

540 if self['band_kpath'] is None: 

541 return self.get_bz_k_points() 

542 else: 

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

544 

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

546 """Convert band_kpath <-> kpts""" 

547 if kpts is None: 

548 kpts = [] 

549 band_kpath = np.array(band_kpath) 

550 band_nkpath = len(band_kpath) 

551 for i, kpath in enumerate(band_kpath): 

552 end = False 

553 nband = int(kpath[0]) 

554 if band_nkpath == i: 

555 end = True 

556 nband += 1 

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

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

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

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

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

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

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

564 elif band_kpath is None: 

565 band_kpath = [] 

566 points = np.asarray(kpts) 

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

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

569 N = len(points) 

570 indices = [0] 

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

572 indices.append(N - 1) 

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

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

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

576 'kpts': 20, 

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

578 return band_kpath 

579 

580 def get_lattice_type(self): 

581 cellpar = cell_to_cellpar(self.atoms.cell) 

582 abc = cellpar[:3] 

583 angles = cellpar[3:] 

584 min_lv = min(abc) 

585 if np.ptp(abc) < 0.01 * min_lv: 

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

587 return 'cubic' 

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

589 return 'fcc' 

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

591 return 'bcc' 

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

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

594 return 'tetragonal' 

595 else: 

596 return 'orthorhombic' 

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

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

599 return 'hexagonal' 

600 else: 

601 return 'not special' 

602 

603 def get_number_of_spins(self): 

604 try: 

605 magmoms = self.atoms.get_initial_magnetic_moments() 

606 if self['scf_spinpolarization'] is None: 

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

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

609 return 1 

610 else: 

611 return 2 

612 else: 

613 raise NotImplementedError 

614 else: 

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

616 return 2 

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

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

619 is not None: 

620 return 1 

621 except KeyError: 

622 return 1 

623 

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

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

626 self.calculate(self.atoms) 

627 if kpt is None and spin is None: 

628 return self.results['eigenvalues'] 

629 else: 

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

631 

632 def get_fermi_level(self): 

633 try: 

634 fermi_level = self.results['chemical_potential'] 

635 except KeyError: 

636 self.calculate() 

637 fermi_level = self.results['chemical_potential'] 

638 return fermi_level 

639 

640 def get_number_of_bands(self): 

641 pag = self.parameters.get 

642 dfd = default_dictionary 

643 if 'number_of_bands' not in self.results: 

644 n = 0 

645 for atom in self.atoms: 

646 sym = atom.symbol 

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

648 d = 1 

649 for orbital in orbitals: 

650 n += d * orbital 

651 d += 2 

652 self.results['number_of_bands'] = n 

653 return self.results['number_of_bands'] 

654 

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

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

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

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

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

660 return dx, dy, dz 

661 

662 def dk(self, dirG): 

663 dx, dy, dz = dirG 

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

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

666 

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

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

669 # only expected to work for no spin polarization 

670 try: 

671 self['bloch_overlaps'] 

672 except KeyError: 

673 self.read_bloch_overlaps() 

674 dirG = tuple(dirG) 

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

676 nr3 = nx * ny * nz 

677 if kpoint is None and nextkpoint is None: 

678 return {kpoint: self['bloch_overlaps' 

679 ][kpoint][dirG][:nbands, :nbands 

680 ] for kpoint in range(nr3)} 

681 if kpoint is None: 

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

683 if nextkpoint is None: 

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

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

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

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

688 

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

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

691 Otherwise, it just ignored''' 

692 if debug is None: 

693 debug = self.debug 

694 if debug: 

695 print(line) 

696 

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

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

699 prev_position = 0 

700 last_position = 0 

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

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

703 time.sleep(5) 

704 with open(file) as fd: 

705 while running(**args): 

706 fd.seek(last_position) 

707 new_data = fd.read() 

708 prev_position = fd.tell() 

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

710 if prev_position != last_position: 

711 if not self.nohup: 

712 print(new_data) 

713 last_position = prev_position 

714 time.sleep(1)