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

450 statements  

« prev     ^ index     » next       coverage.py v7.5.3, created at 2024-08-05 14:37 +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', 

289 fileName + '.dden.cube', 

290 fileName + '.ene', 

291 fileName + '.md', 

292 fileName + '.md2', 

293 fileName + '.tden.cube', 

294 fileName + '.sden.cube', 

295 fileName + '.v0.cube', 

296 fileName + '.v1.cube', 

297 fileName + '.vhart.cube', 

298 fileName + '.den0.cube', 

299 fileName + '.bulk.xyz', 

300 fileName + '.den1.cube', 

301 fileName + '.xyz', 

302 pbs_Name + '.o' + str(queue_num), 

303 pbs_Name + '.e' + str(queue_num) 

304 ] 

305 for f in files: 

306 try: 

307 self.prind("Removing" + f) 

308 os.remove(f) 

309 except OSError: 

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

311 

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

313 system_changes=all_changes): 

314 """ 

315 Capture the RuntimeError from FileIOCalculator.calculate 

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

317 See base FileIOCalculator for documentation. 

318 """ 

319 if self.parameters.data_path is None: 

320 if 'OPENMX_DFT_DATA_PATH' not in cfg: 

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

322 'enviroment variable or specify "data_path" as' 

323 'a keyword argument') 

324 

325 self.prind("Start Calculation") 

326 if properties is None: 

327 properties = self.implemented_properties 

328 try: 

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

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

331 properties=properties, 

332 system_changes=system_changes) 

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

334 self.run() 

335 # self.read_results() 

336 self.version = self.read_version() 

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

338 self.output_atoms = output_atoms 

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

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

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

342 self.results = output_atoms.calc.results 

343 # self.clean() 

344 except RuntimeError as e: 

345 try: 

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

347 lines = fd.readlines() 

348 debug_lines = 10 

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

350 for line in lines[-20:]: 

351 print(line.strip()) 

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

353 raise e 

354 except RuntimeError as e: 

355 raise e 

356 

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

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

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

360 See calculator.py for further details. 

361 

362 Parameters: 

363 - atoms : The Atoms object to write. 

364 - properties : The properties which should be calculated. 

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

366 """ 

367 # Call base calculator. 

368 if atoms is None: 

369 atoms = self.atoms 

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

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

372 properties=properties, system_changes=system_changes) 

373 

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

375 """ 

376 For a debugging purpose, print the .dat file 

377 """ 

378 if debug is None: 

379 debug = self.debug 

380 if nohup is None: 

381 nohup = self.nohup 

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

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

384 if not nohup: 

385 with open(filename) as fd: 

386 while True: 

387 line = fd.readline() 

388 print(line.strip()) 

389 if not line: 

390 break 

391 

392 def read(self, label): 

393 self.parameters = {} 

394 self.set_label(label) 

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

396 label = label[:-4] 

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

398 self.update_atoms(atoms) 

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

400 self.results = atoms.calc.results 

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

402 self.parameters['label'] = label 

403 

404 def read_version(self, label=None): 

405 version = None 

406 if label is None: 

407 label = self.label 

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

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

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

411 break 

412 return version 

413 

414 def update_atoms(self, atoms): 

415 self.atoms = atoms.copy() 

416 

417 def set(self, **kwargs): 

418 """Set all parameters. 

419 

420 Parameters: 

421 -kwargs : Dictionary containing the keywords defined in 

422 OpenMXParameters. 

423 """ 

424 

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

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

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

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

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

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

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

432 default_dict = self.default_parameters[key] 

433 for subkey in kwargs[key]: 

434 if subkey not in default_dict: 

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

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

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

438 .format(subkey, key, allowed)) 

439 

440 # Find out what parameter has been changed 

441 changed_parameters = {} 

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

443 oldvalue = self.parameters.get(key) 

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

445 changed_parameters[key] = value 

446 self.parameters[key] = value 

447 

448 # Set the parameters 

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

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

451 self.parameters[key] = value 

452 

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

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

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

456 self.results = {} 

457 

458 value = kwargs.get('energy_cutoff') 

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

460 and value > 0): 

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

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

463 raise ValueError(mess) 

464 

465 atoms = kwargs.get('atoms') 

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

467 self.atoms = atoms.copy() 

468 

469 def set_results(self, results): 

470 # Not Implemented fully 

471 self.results.update(results) 

472 

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

474 # Contruct the command to send to the operating system 

475 abs_dir = os.getcwd() 

476 command = '' 

477 self.prind(self.command) 

478 if self.command is None: 

479 self.command = 'openmx' 

480 # run processes specified by the system variable OPENMX_COMMAND 

481 if processes is None: 

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

483 if command is None: 

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

485 variable or specify processes as a keyword argument') 

486 else: # run with a specified number of processes 

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

488 if threads is None: 

489 threads_string = '' 

490 command += 'mpirun -np ' + \ 

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

492 ' %s ' + threads_string + ' |tee %s' 

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

494 

495 if runfile is None: 

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

497 if outfile is None: 

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

499 try: 

500 command = command % (runfile, outfile) 

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

502 except TypeError: # in case the OPENMX_COMMAND is incompatible 

503 raise ValueError( 

504 "The 'OPENMX_COMMAND' environment must " + 

505 "be a format string" + 

506 " with four string arguments.\n" + 

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

508 f"Got '{command}'") 

509 return command 

510 

511 def get_stress(self, atoms=None): 

512 if atoms is None: 

513 atoms = self.atoms 

514 

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

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

517 

518 return stress 

519 

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

521 """ 

522 This is band structure function. It is compatible to 

523 ase dft module """ 

524 from ase.dft import band_structure 

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

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

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

528 

529 def get_bz_k_points(self): 

530 kgrid = self['kpts'] 

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

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

533 bz_k_points = [] 

534 n1 = kgrid[0] 

535 n2 = kgrid[1] 

536 n3 = kgrid[2] 

537 for i in range(n1): 

538 for j in range(n2): 

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

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

541 for k in range(n3): 

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

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

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

545 return np.array(bz_k_points) 

546 

547 def get_ibz_k_points(self): 

548 if self['band_kpath'] is None: 

549 return self.get_bz_k_points() 

550 else: 

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

552 

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

554 """Convert band_kpath <-> kpts""" 

555 if kpts is None: 

556 kpts = [] 

557 band_kpath = np.array(band_kpath) 

558 band_nkpath = len(band_kpath) 

559 for i, kpath in enumerate(band_kpath): 

560 end = False 

561 nband = int(kpath[0]) 

562 if band_nkpath == i: 

563 end = True 

564 nband += 1 

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

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

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

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

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

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

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

572 elif band_kpath is None: 

573 band_kpath = [] 

574 points = np.asarray(kpts) 

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

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

577 N = len(points) 

578 indices = [0] 

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

580 indices.append(N - 1) 

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

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

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

584 'kpts': 20, 

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

586 return band_kpath 

587 

588 def get_lattice_type(self): 

589 cellpar = cell_to_cellpar(self.atoms.cell) 

590 abc = cellpar[:3] 

591 angles = cellpar[3:] 

592 min_lv = min(abc) 

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

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

595 return 'cubic' 

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

597 return 'fcc' 

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

599 return 'bcc' 

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

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

602 return 'tetragonal' 

603 else: 

604 return 'orthorhombic' 

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

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

607 return 'hexagonal' 

608 else: 

609 return 'not special' 

610 

611 def get_number_of_spins(self): 

612 try: 

613 magmoms = self.atoms.get_initial_magnetic_moments() 

614 if self['scf_spinpolarization'] is None: 

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

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

617 return 1 

618 else: 

619 return 2 

620 else: 

621 raise NotImplementedError 

622 else: 

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

624 return 2 

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

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

627 is not None: 

628 return 1 

629 except KeyError: 

630 return 1 

631 

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

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

634 self.calculate(self.atoms) 

635 if kpt is None and spin is None: 

636 return self.results['eigenvalues'] 

637 else: 

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

639 

640 def get_fermi_level(self): 

641 try: 

642 fermi_level = self.results['chemical_potential'] 

643 except KeyError: 

644 self.calculate() 

645 fermi_level = self.results['chemical_potential'] 

646 return fermi_level 

647 

648 def get_number_of_bands(self): 

649 pag = self.parameters.get 

650 dfd = default_dictionary 

651 if 'number_of_bands' not in self.results: 

652 n = 0 

653 for atom in self.atoms: 

654 sym = atom.symbol 

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

656 d = 1 

657 for orbital in orbitals: 

658 n += d * orbital 

659 d += 2 

660 self.results['number_of_bands'] = n 

661 return self.results['number_of_bands'] 

662 

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

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

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

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

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

668 return dx, dy, dz 

669 

670 def dk(self, dirG): 

671 dx, dy, dz = dirG 

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

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

674 

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

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

677 # only expected to work for no spin polarization 

678 try: 

679 self['bloch_overlaps'] 

680 except KeyError: 

681 self.read_bloch_overlaps() 

682 dirG = tuple(dirG) 

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

684 nr3 = nx * ny * nz 

685 if kpoint is None and nextkpoint is None: 

686 return {kpoint: self['bloch_overlaps' 

687 ][kpoint][dirG][:nbands, :nbands 

688 ] for kpoint in range(nr3)} 

689 if kpoint is None: 

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

691 if nextkpoint is None: 

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

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

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

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

696 

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

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

699 Otherwise, it just ignored''' 

700 if debug is None: 

701 debug = self.debug 

702 if debug: 

703 print(line) 

704 

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

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

707 prev_position = 0 

708 last_position = 0 

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

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

711 time.sleep(5) 

712 with open(file) as fd: 

713 while running(**args): 

714 fd.seek(last_position) 

715 new_data = fd.read() 

716 prev_position = fd.tell() 

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

718 if prev_position != last_position: 

719 if not self.nohup: 

720 print(new_data) 

721 last_position = prev_position 

722 time.sleep(1)