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 time 

24import subprocess 

25import re 

26import warnings 

27import numpy as np 

28from ase.geometry import cell_to_cellpar 

29from ase.calculators.calculator import (FileIOCalculator, Calculator, equal, 

30 all_changes, kptdensity2monkhorstpack) 

31from ase.calculators.openmx.parameters import OpenMXParameters 

32from ase.calculators.openmx.default_settings import default_dictionary 

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

34from ase.calculators.openmx.writer import write_openmx 

35 

36 

37def parse_omx_version(txt): 

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

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

40 return match.group(1) 

41 

42 

43class OpenMX(FileIOCalculator): 

44 """ 

45 Calculator interface to the OpenMX code. 

46 """ 

47 

48 implemented_properties = [ 

49 'free_energy', # Same value with energy 

50 'energy', 

51 'energies', 

52 'forces', 

53 'stress', 

54 'dipole', 

55 'chemical_potential', 

56 'magmom', 

57 'magmoms', 

58 'eigenvalues'] 

59 

60 default_parameters = OpenMXParameters() 

61 

62 default_pbs = { 

63 'processes': 1, 

64 'walltime': "10:00:00", 

65 'threads': 1, 

66 'nodes': 1 

67 } 

68 

69 default_mpi = { 

70 'processes': 1, 

71 'threads': 1 

72 } 

73 

74 default_output_setting = { 

75 'nohup': True, 

76 'debug': False 

77 } 

78 

79 def __init__(self, restart=None, 

80 ignore_bad_restart_file=FileIOCalculator._deprecated, 

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

82 pbs=None, **kwargs): 

83 

84 # Initialize and put the default parameters. 

85 self.initialize_pbs(pbs) 

86 self.initialize_mpi(mpi) 

87 self.initialize_output_setting(**kwargs) 

88 

89 FileIOCalculator.__init__(self, restart, ignore_bad_restart_file, 

90 label, atoms, command, **kwargs) 

91 

92 def __getitem__(self, key): 

93 """Convenience method to retrieve a parameter as 

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

95 

96 Parameters: 

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

98 """ 

99 return self.parameters[key] 

100 

101 def __setitem__(self, key, value): 

102 self.parameters[key] = value 

103 

104 def initialize_output_setting(self, **kwargs): 

105 output_setting = {} 

106 self.output_setting = dict(self.default_output_setting) 

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

108 if key in self.default_output_setting: 

109 output_setting[key] = value 

110 self.output_setting.update(output_setting) 

111 self.__dict__.update(self.output_setting) 

112 

113 def initialize_pbs(self, pbs): 

114 if pbs: 

115 self.pbs = dict(self.default_pbs) 

116 for key in pbs: 

117 if key not in self.default_pbs: 

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

119 raise TypeError('Unexpected keyword "{0}" in "pbs" ' 

120 'dictionary. Must be one of: {1}' 

121 .format(key, allowed)) 

122 # Put dictionary into python variable 

123 self.pbs.update(pbs) 

124 self.__dict__.update(self.pbs) 

125 else: 

126 self.pbs = None 

127 

128 def initialize_mpi(self, mpi): 

129 if mpi: 

130 self.mpi = dict(self.default_mpi) 

131 for key in mpi: 

132 if key not in self.default_mpi: 

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

134 raise TypeError('Unexpected keyword "{0}" in "mpi" ' 

135 'dictionary. Must be one of: {1}' 

136 .format(key, allowed)) 

137 # Put dictionary into python variable 

138 self.mpi.update(mpi) 

139 self.__dict__.update(self.mpi) 

140 else: 

141 self.mpi = None 

142 

143 def run(self): 

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

145 if self.pbs is not None: 

146 run = self.run_pbs 

147 elif self.mpi is not None: 

148 run = self.run_mpi 

149 else: 

150 run = self.run_openmx 

151 run() 

152 

153 def run_openmx(self): 

154 def isRunning(process=None): 

155 ''' Check mpi is running''' 

156 return process.poll() is None 

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

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

159 olddir = os.getcwd() 

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

161 try: 

162 os.chdir(abs_dir) 

163 if self.command is None: 

164 self.command = 'openmx' 

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

166 command = command % (runfile, outfile) 

167 self.prind(command) 

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

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

170 finally: 

171 os.chdir(olddir) 

172 self.prind("Calculation Finished") 

173 

174 def run_mpi(self): 

175 """ 

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

177 run. 

178 """ 

179 def isRunning(process=None): 

180 ''' Check mpi is running''' 

181 return process.poll() is None 

182 processes = self.processes 

183 threads = self.threads 

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

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

186 olddir = os.getcwd() 

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

188 try: 

189 os.chdir(abs_dir) 

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

191 self.prind(command) 

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

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

194 finally: 

195 os.chdir(olddir) 

196 self.prind("Calculation Finished") 

197 

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

199 """ 

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

201 Your system should have Scheduler. PBS 

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

203 Super computer user 

204 """ 

205 nodes = self.nodes 

206 processes = self.processes 

207 

208 prefix = self.prefix 

209 olddir = os.getcwd() 

210 try: 

211 os.chdir(self.abs_directory) 

212 except AttributeError: 

213 os.chdir(self.directory) 

214 

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

216 """ 

217 Check submitted job is still Running 

218 """ 

219 def runCmd(exe): 

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

221 stderr=subprocess.STDOUT, 

222 universal_newlines=True) 

223 while True: 

224 line = p.stdout.readline() 

225 if line != '': 

226 # the real code does filtering here 

227 yield line.rstrip() 

228 else: 

229 break 

230 jobs = runCmd('qstat') 

231 columns = None 

232 for line in jobs: 

233 if str(jobNum) in line: 

234 columns = line.split() 

235 self.prind(line) 

236 if columns is not None: 

237 return columns[-2] == status 

238 else: 

239 return False 

240 

241 inputfile = self.label + '.dat' 

242 outfile = self.label + '.log' 

243 

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

245 jobName = prefix 

246 cmd = bashArgs + \ 

247 'mpirun -hostfile $PBS_NODEFILE openmx %s > %s' % ( 

248 inputfile, outfile) 

249 echoArgs = ["echo", "$' %s'" % cmd] 

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

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

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

253 self.prind(wholeCmd) 

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

255 stdout=subprocess.PIPE, universal_newlines=True) 

256 out = out.communicate()[0] 

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

258 

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

260 '\nWaiting for the Queue to start') 

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

262 time.sleep(5) 

263 self.prind('.') 

264 self.prind('Start Calculating') 

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

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

267 

268 os.chdir(olddir) 

269 self.prind('Calculation Finished!') 

270 return jobNum 

271 

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

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

274 

275 The default files generated OpenMX will be deleted IF this 

276 method is called. 

277 

278 """ 

279 self.prind("Cleaning Data") 

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

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

282 files = [ 

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

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

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

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

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

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

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

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

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

292 ] 

293 for f in files: 

294 try: 

295 self.prind("Removing" + f) 

296 os.remove(f) 

297 except OSError: 

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

299 

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

301 system_changes=all_changes): 

302 """ 

303 Capture the RuntimeError from FileIOCalculator.calculate 

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

305 See base FileIOCalculator for documentation. 

306 """ 

307 if self.parameters.data_path is None: 

308 if 'OPENMX_DFT_DATA_PATH' not in os.environ: 

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

310 'enviroment variable or specify "data_path" as' 

311 'a keyword argument') 

312 

313 self.prind("Start Calculation") 

314 if properties is None: 

315 properties = self.implemented_properties 

316 try: 

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

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

319 properties=properties, 

320 system_changes=system_changes) 

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

322 self.run() 

323 # self.read_results() 

324 self.version = self.read_version() 

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

326 self.output_atoms = output_atoms 

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

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

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

330 self.results = output_atoms.calc.results 

331 # self.clean() 

332 except RuntimeError as e: 

333 try: 

334 with open(get_file_name('.log'), 'r') as fd: 

335 lines = fd.readlines() 

336 debug_lines = 10 

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

338 for line in lines[-20:]: 

339 print(line.strip()) 

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

341 raise e 

342 except RuntimeError as e: 

343 raise e 

344 

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

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

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

348 See calculator.py for further details. 

349 

350 Parameters: 

351 - atoms : The Atoms object to write. 

352 - properties : The properties which should be calculated. 

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

354 """ 

355 # Call base calculator. 

356 if atoms is None: 

357 atoms = self.atoms 

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

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

360 properties=properties, system_changes=system_changes) 

361 

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

363 """ 

364 For a debugging purpose, print the .dat file 

365 """ 

366 if debug is None: 

367 debug = self.debug 

368 if nohup is None: 

369 nohup = self.nohup 

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

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

372 if not nohup: 

373 with open(filename, 'r') as fd: 

374 while True: 

375 line = fd.readline() 

376 print(line.strip()) 

377 if not line: 

378 break 

379 

380 def read(self, label): 

381 self.parameters = {} 

382 self.set_label(label) 

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

384 label = label[:-4] 

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

386 self.update_atoms(atoms) 

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

388 self.results = atoms.calc.results 

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

390 self.parameters['label'] = label 

391 

392 def read_version(self, label=None): 

393 version = None 

394 if label is None: 

395 label = self.label 

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

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

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

399 break 

400 return version 

401 

402 def update_atoms(self, atoms): 

403 self.atoms = atoms.copy() 

404 

405 def set(self, **kwargs): 

406 """Set all parameters. 

407 

408 Parameters: 

409 -kwargs : Dictionary containing the keywords defined in 

410 OpenMXParameters. 

411 """ 

412 

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

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

415 raise KeyError('Unkown keyword "%s" and value "%s".' % 

416 (key, value)) 

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

418 raise KeyError('Given xc "%s" is not allowed' % value) 

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

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

421 default_dict = self.default_parameters[key] 

422 for subkey in kwargs[key]: 

423 if subkey not in default_dict: 

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

425 raise TypeError('Unknown subkeyword "{0}" of keyword ' 

426 '"{1}". Must be one of: {2}' 

427 .format(subkey, key, allowed)) 

428 

429 # Find out what parameter has been changed 

430 changed_parameters = {} 

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

432 oldvalue = self.parameters.get(key) 

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

434 changed_parameters[key] = value 

435 self.parameters[key] = value 

436 

437 # Set the parameters 

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

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

440 self.parameters[key] = value 

441 

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

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

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

445 self.results = {} 

446 

447 value = kwargs.get('energy_cutoff') 

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

449 and value > 0): 

450 mess = "'%s' must be a positive number(in eV), \ 

451 got '%s'" % ('energy_cutoff', value) 

452 raise ValueError(mess) 

453 

454 atoms = kwargs.get('atoms') 

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

456 self.atoms = atoms.copy() 

457 

458 def set_results(self, results): 

459 # Not Implemented fully 

460 self.results.update(results) 

461 

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

463 # Contruct the command to send to the operating system 

464 abs_dir = os.getcwd() 

465 command = '' 

466 self.prind(self.command) 

467 if self.command is None: 

468 self.command = 'openmx' 

469 # run processes specified by the system variable OPENMX_COMMAND 

470 if processes is None: 

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

472 if command is None: 

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

474 variable or specify processes as a keyword argument') 

475 else: # run with a specified number of processes 

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

477 if threads is None: 

478 threads_string = '' 

479 command += 'mpirun -np ' + \ 

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

481 ' %s ' + threads_string + ' |tee %s' 

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

483 

484 if runfile is None: 

485 runfile = abs_dir + '/' + self.prefix + '.dat' 

486 if outfile is None: 

487 outfile = abs_dir + '/' + self.prefix + '.log' 

488 try: 

489 command = command % (runfile, outfile) 

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

491 except TypeError: # in case the OPENMX_COMMAND is incompatible 

492 raise ValueError( 

493 "The 'OPENMX_COMMAND' environment must " + 

494 "be a format string" + 

495 " with four string arguments.\n" + 

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

497 "Got '%s'" % command) 

498 return command 

499 

500 def get_stress(self, atoms=None): 

501 if atoms is None: 

502 atoms = self.atoms 

503 

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

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

506 

507 return stress 

508 

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

510 """ 

511 This is band structure function. It is compatible to 

512 ase dft module """ 

513 from ase.dft import band_structure 

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

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

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

517 

518 def get_bz_k_points(self): 

519 kgrid = self['kpts'] 

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

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

522 bz_k_points = [] 

523 n1 = kgrid[0] 

524 n2 = kgrid[1] 

525 n3 = kgrid[2] 

526 for i in range(n1): 

527 for j in range(n2): 

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

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

530 for k in range(n3): 

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

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

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

534 return np.array(bz_k_points) 

535 

536 def get_ibz_k_points(self): 

537 if self['band_kpath'] is None: 

538 return self.get_bz_k_points() 

539 else: 

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

541 

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

543 """Convert band_kpath <-> kpts""" 

544 if kpts is None: 

545 kpts = [] 

546 band_kpath = np.array(band_kpath) 

547 band_nkpath = len(band_kpath) 

548 for i, kpath in enumerate(band_kpath): 

549 end = False 

550 nband = int(kpath[0]) 

551 if band_nkpath == i: 

552 end = True 

553 nband += 1 

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

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

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

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

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

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

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

561 elif band_kpath is None: 

562 band_kpath = [] 

563 points = np.asarray(kpts) 

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

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

566 N = len(points) 

567 indices = [0] 

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

569 indices.append(N - 1) 

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

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

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

573 'kpts': 20, 

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

575 return band_kpath 

576 

577 def get_lattice_type(self): 

578 cellpar = cell_to_cellpar(self.atoms.cell) 

579 abc = cellpar[:3] 

580 angles = cellpar[3:] 

581 min_lv = min(abc) 

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

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

584 return 'cubic' 

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

586 return 'fcc' 

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

588 return 'bcc' 

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

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

591 return 'tetragonal' 

592 else: 

593 return 'orthorhombic' 

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

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

596 return 'hexagonal' 

597 else: 

598 return 'not special' 

599 

600 def get_number_of_spins(self): 

601 try: 

602 magmoms = self.atoms.get_initial_magnetic_moments() 

603 if self['scf_spinpolarization'] is None: 

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

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

606 return 1 

607 else: 

608 return 2 

609 else: 

610 raise NotImplementedError 

611 else: 

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

613 return 2 

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

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

616 is not None: 

617 return 1 

618 except KeyError: 

619 return 1 

620 

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

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

623 self.calculate(self.atoms) 

624 if kpt is None and spin is None: 

625 return self.results['eigenvalues'] 

626 else: 

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

628 

629 def get_fermi_level(self): 

630 try: 

631 fermi_level = self.results['chemical_potential'] 

632 except KeyError: 

633 self.calculate() 

634 fermi_level = self.results['chemical_potential'] 

635 return fermi_level 

636 

637 def get_number_of_bands(self): 

638 pag = self.parameters.get 

639 dfd = default_dictionary 

640 if 'number_of_bands' not in self.results: 

641 n = 0 

642 for atom in self.atoms: 

643 sym = atom.symbol 

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

645 d = 1 

646 for orbital in orbitals: 

647 n += d * orbital 

648 d += 2 

649 self.results['number_of_bands'] = n 

650 return self.results['number_of_bands'] 

651 

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

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

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

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

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

657 return dx, dy, dz 

658 

659 def dk(self, dirG): 

660 dx, dy, dz = dirG 

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

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

663 

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

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

666 # only expected to work for no spin polarization 

667 try: 

668 self['bloch_overlaps'] 

669 except KeyError: 

670 self.read_bloch_overlaps() 

671 dirG = tuple(dirG) 

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

673 nr3 = nx * ny * nz 

674 if kpoint is None and nextkpoint is None: 

675 return {kpoint: self['bloch_overlaps' 

676 ][kpoint][dirG][:nbands, :nbands 

677 ] for kpoint in range(nr3)} 

678 if kpoint is None: 

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

680 if nextkpoint is None: 

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

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

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

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

685 

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

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

688 Otherwise, it just ignored''' 

689 if debug is None: 

690 debug = self.debug 

691 if debug: 

692 print(line) 

693 

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

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

696 prev_position = 0 

697 last_position = 0 

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

699 self.prind('Waiting for %s to come out' % file) 

700 time.sleep(5) 

701 with open(file, 'r') as fd: 

702 while running(**args): 

703 fd.seek(last_position) 

704 new_data = fd.read() 

705 prev_position = fd.tell() 

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

707 if prev_position != last_position: 

708 if not self.nohup: 

709 print(new_data) 

710 last_position = prev_position 

711 time.sleep(1)