1""" 2This module gathers the most important classes and helper functions used for scripting. 3""" 4import os 5import collections 6 7#################### 8### Monty import ### 9#################### 10from monty.os.path import which 11from monty.termcolor import cprint 12 13####################### 14### Pymatgen import ### 15####################### 16# Tools for unit conversion 17import pymatgen.core.units as units 18FloatWithUnit = units.FloatWithUnit 19ArrayWithUnit = units.ArrayWithUnit 20 21#################### 22### Abipy import ### 23#################### 24from abipy.flowtk import Pseudo, PseudoTable, Mrgscr, Mrgddb, Flow, Work, TaskManager, AbinitBuild, flow_main 25from abipy.core.release import __version__, min_abinit_version 26from abipy.core.globals import enable_notebook, in_notebook, disable_notebook 27from abipy.core import restapi 28from abipy.core.structure import (Lattice, Structure, StructureModifier, dataframes_from_structures, 29 mp_match_structure, mp_search, cod_search) 30from abipy.core.mixins import CubeFile 31from abipy.core.func1d import Function1D 32from abipy.core.kpoints import set_atol_kdiff 33from abipy.abio.robots import Robot 34from abipy.abio.inputs import AbinitInput, MultiDataset, AnaddbInput, OpticInput 35from abipy.abio.abivars import AbinitInputFile 36from abipy.abio.outputs import AbinitLogFile, AbinitOutputFile, OutNcFile, AboRobot 37from abipy.tools.printing import print_dataframe 38from abipy.tools.notebooks import print_source, print_doc 39from abipy.tools.plotting import get_ax_fig_plt, get_axarray_fig_plt, get_ax3d_fig_plt 40from abipy.abio.factories import * 41from abipy.electrons.ebands import (ElectronBands, ElectronBandsPlotter, ElectronDos, ElectronDosPlotter, 42 dataframe_from_ebands, EdosFile) 43from abipy.electrons.gsr import GsrFile, GsrRobot 44from abipy.electrons.eskw import EskwFile 45from abipy.electrons.psps import PspsFile 46from abipy.electrons.gw import SigresFile, SigresRobot 47from abipy.electrons.bse import MdfFile, MdfRobot 48from abipy.electrons.scissors import ScissorsBuilder 49from abipy.electrons.scr import ScrFile 50from abipy.electrons.denpot import (DensityNcFile, VhartreeNcFile, VxcNcFile, VhxcNcFile, PotNcFile, 51 DensityFortranFile, Cut3dDenPotNcFile) 52from abipy.electrons.fatbands import FatBandsFile 53from abipy.electrons.optic import OpticNcFile, OpticRobot 54from abipy.electrons.fold2bloch import Fold2BlochNcfile 55from abipy.dfpt.phonons import (PhbstFile, PhbstRobot, PhononBands, PhononBandsPlotter, PhdosFile, PhononDosPlotter, 56 PhdosReader, phbands_gridplot) 57from abipy.dfpt.ddb import DdbFile, DdbRobot 58from abipy.dfpt.anaddbnc import AnaddbNcFile, AnaddbNcRobot 59from abipy.dfpt.gruneisen import GrunsNcFile 60from abipy.dynamics.hist import HistFile, HistRobot 61from abipy.waves import WfkFile 62from abipy.eph.a2f import A2fFile, A2fRobot 63from abipy.eph.sigeph import SigEPhFile, SigEPhRobot 64from abipy.eph.eph_plotter import EphPlotter 65from abipy.eph.v1sym import V1symFile 66from abipy.eph.gkq import GkqFile, GkqRobot 67from abipy.eph.v1qnu import V1qnuFile 68from abipy.eph.v1qavg import V1qAvgFile 69from abipy.eph.rta import RtaFile, RtaRobot 70from abipy.eph.transportfile import TransportFile 71from abipy.wannier90 import WoutFile, AbiwanFile, AbiwanRobot 72from abipy.electrons.lobster import CoxpFile, ICoxpFile, LobsterDoscarFile, LobsterInput, LobsterAnalyzer 73#from abipy.electrons.abitk import ZinvConvFile, TetraTestFile 74 75# Abinit Documentation. 76from abipy.abio.abivars_db import get_abinit_variables, abinit_help, docvar 77 78 79def _straceback(): 80 """Returns a string with the traceback.""" 81 import traceback 82 return traceback.format_exc() 83 84 85# Abinit text files. Use OrderedDict for nice output in show_abiopen_exc2class. 86ext2file = collections.OrderedDict([ 87 (".abi", AbinitInputFile), 88 (".in", AbinitInputFile), 89 (".abo", AbinitOutputFile), 90 (".out", AbinitOutputFile), 91 (".log", AbinitLogFile), 92 (".cif", Structure), 93 ("POSCAR", Structure), 94 (".cssr", Structure), 95 (".cube", CubeFile), 96 ("anaddb.nc", AnaddbNcFile), 97 ("DEN", DensityFortranFile), 98 (".psp8", Pseudo), 99 (".pspnc", Pseudo), 100 (".fhi", Pseudo), 101 ("JTH.xml", Pseudo), 102 (".wout", WoutFile), 103 # Lobster files. 104 ("COHPCAR.lobster", CoxpFile), 105 ("COOPCAR.lobster", CoxpFile), 106 ("ICOHPLIST.lobster", ICoxpFile), 107 ("DOSCAR.lobster", LobsterDoscarFile), 108 #("ZINVCONV.nc", ZinvConvFile), 109 #("TETRATEST.nc", TetraTestFile), 110 ("EDOS", EdosFile), 111]) 112 113# Abinit files require a special treatment. 114abiext2ncfile = collections.OrderedDict([ 115 ("GSR.nc", GsrFile), 116 ("ESKW.nc", EskwFile), 117 ("DEN.nc", DensityNcFile), 118 ("OUT.nc", OutNcFile), 119 ("VHA.nc", VhartreeNcFile), 120 ("VXC.nc", VxcNcFile), 121 ("VHXC.nc", VhxcNcFile), 122 ("POT.nc", PotNcFile), 123 ("WFK.nc", WfkFile), 124 ("HIST.nc", HistFile), 125 ("PSPS.nc", PspsFile), 126 ("DDB", DdbFile), 127 ("PHBST.nc", PhbstFile), 128 ("PHDOS.nc", PhdosFile), 129 ("SCR.nc", ScrFile), 130 ("SIGRES.nc", SigresFile), 131 ("GRUNS.nc", GrunsNcFile), 132 ("MDF.nc", MdfFile), 133 ("FATBANDS.nc", FatBandsFile), 134 ("FOLD2BLOCH.nc", Fold2BlochNcfile), 135 ("CUT3DDENPOT.nc", Cut3dDenPotNcFile), 136 ("OPTIC.nc", OpticNcFile), 137 ("A2F.nc", A2fFile), 138 ("SIGEPH.nc", SigEPhFile), 139 ("TRANSPORT.nc",TransportFile), 140 ("RTA.nc",RtaFile), 141 ("V1SYM.nc", V1symFile), 142 ("GKQ.nc", GkqFile), 143 ("V1QNU.nc", V1qnuFile), 144 ("V1QAVG.nc", V1qAvgFile), 145 ("ABIWAN.nc", AbiwanFile), 146]) 147 148 149def abiopen_ext2class_table(): 150 """ 151 Print the association table between file extensions and File classes. 152 """ 153 from itertools import chain 154 from tabulate import tabulate 155 table = [] 156 157 for ext, cls in chain(ext2file.items(), abiext2ncfile.items()): 158 table.append((ext, str(cls))) 159 160 return tabulate(table, headers=["Extension", "Class"]) 161 162 163def abifile_subclass_from_filename(filename): 164 """ 165 Returns the appropriate class associated to the given filename. 166 """ 167 if os.path.basename(filename) == Flow.PICKLE_FNAME: 168 return Flow 169 170 from abipy.tools.text import rreplace 171 for ext, cls in ext2file.items(): 172 # This to support gzipped files. 173 if filename.endswith(".gz"): filename = rreplace(filename, ".gz", "", occurrence=1) 174 if filename.endswith(ext): return cls 175 176 ext = filename.split("_")[-1] 177 try: 178 return abiext2ncfile[ext] 179 except KeyError: 180 for ext, cls in abiext2ncfile.items(): 181 if filename.endswith(ext): return cls 182 183 msg = ("No class has been registered for file:\n\t%s\n\nFile extensions supported:\n\n%s" % 184 (filename, abiopen_ext2class_table())) 185 raise ValueError(msg) 186 187 188def dir2abifiles(top, recurse=True): 189 """ 190 Analyze the filesystem starting from directory `top` and 191 return an ordered dictionary mapping the directory name to the list 192 of files supported by ``abiopen`` contained within that directory. 193 If not ``recurse``, children directories are not analyzed. 194 """ 195 dl = collections.defaultdict(list) 196 197 if recurse: 198 for dirpath, dirnames, filenames in os.walk(top): 199 for f in filenames: 200 path = os.path.join(dirpath, f) 201 if not isabifile(path): continue 202 dl[dirpath].append(path) 203 else: 204 for f in os.listdir(top): 205 path = os.path.join(top, f) 206 if not isabifile(path): continue 207 dl[top].append(path) 208 209 return collections.OrderedDict([(k, dl[k]) for k in sorted(dl.keys())]) 210 211 212def isabifile(filepath): 213 """ 214 Return True if `filepath` can be opened with ``abiopen``. 215 """ 216 try: 217 abifile_subclass_from_filename(filepath) 218 return True 219 except ValueError: 220 return False 221 222 223def abiopen(filepath): 224 """ 225 Factory function that opens any file supported by abipy. 226 File type is detected from the extension 227 228 Args: 229 filepath: string with the filename. 230 """ 231 # Handle ~ in filepath. 232 filepath = os.path.expanduser(filepath) 233 234 # Handle zipped files by creating temporary file with correct extension. 235 root, ext = os.path.splitext(filepath) 236 if ext in (".bz2", ".gz", ".z"): 237 from monty.io import zopen 238 with zopen(filepath, "rt") as f: 239 import tempfile 240 _, tmp_path = tempfile.mkstemp(suffix=os.path.basename(root), text=True) 241 cprint("Creating temporary file: %s" % tmp_path, "yellow") 242 with open(tmp_path, "wt") as t: 243 t.write(f.read()) 244 filepath = tmp_path 245 246 if os.path.basename(filepath) == "__AbinitFlow__.pickle": 247 return Flow.pickle_load(filepath) 248 249 # Handle old output files produced by Abinit. 250 import re 251 outnum = re.compile(r".+\.out[\d]+") 252 abonum = re.compile(r".+\.abo[\d]+") 253 if outnum.match(filepath) or abonum.match(filepath): 254 return AbinitOutputFile.from_file(filepath) 255 256 if os.path.basename(filepath) == "log": 257 # Assume Abinit log file. 258 return AbinitLogFile.from_file(filepath) 259 260 cls = abifile_subclass_from_filename(filepath) 261 return cls.from_file(filepath) 262 263 264def display_structure(obj, **kwargs): 265 """ 266 Use Jsmol to display a structure in the jupyter notebook. 267 Requires `nbjsmol` notebook extension installed on the local machine. 268 Install it with `pip install nbjsmol`. See also https://github.com/gmatteo/nbjsmol. 269 270 Args: 271 obj: Structure object or file with a structure or python object with a `structure` attribute. 272 kwargs: Keyword arguments passed to `nbjsmol_display` 273 """ 274 try: 275 from nbjsmol import nbjsmol_display 276 except ImportError as exc: 277 raise ImportError(str(exc) + 278 "\ndisplay structure requires nbjsmol package\n." 279 "Install it with `pip install nbjsmol.`\n" 280 "See also https://github.com/gmatteo/nbjsmol.") 281 282 # Cast to structure, get string with cif data and pass it to nbjsmol. 283 structure = Structure.as_structure(obj) 284 return nbjsmol_display(structure.to(fmt="cif"), ext=".cif", **kwargs) 285 286 287def mjson_load(filepath, **kwargs): 288 """ 289 Read JSON file in MSONable format with MontyDecoder. Return dict with python objects. 290 """ 291 import json 292 from monty.json import MontyDecoder 293 with open(filepath, "rt") as fh: 294 return json.load(fh, cls=MontyDecoder, **kwargs) 295 296 297def mjson_loads(string, **kwargs): 298 """ 299 Read JSON string in MSONable format with MontyDecoder. Return dict with python objects. 300 """ 301 import json 302 from monty.json import MontyDecoder 303 return json.loads(string, cls=MontyDecoder, **kwargs) 304 305 306def mjson_write(d, filepath, **kwargs): 307 """ 308 Write dictionary d to filepath in JSON format using MontyDecoder 309 """ 310 import json 311 from monty.json import MontyEncoder 312 with open(filepath, "wt") as fh: 313 json.dump(d, fh, cls=MontyEncoder, **kwargs) 314 315 316def software_stack(): 317 """ 318 Import all the hard dependencies. Returns ordered dict: package --> string with version info. 319 """ 320 import platform 321 system, node, release, version, machine, processor = platform.uname() 322 # These packages are required 323 import numpy, scipy, netCDF4, pymatgen, apscheduler, pydispatch, yaml 324 325 try: 326 from pymatgen.core import __version__ as pmg_version 327 #from pymatgen.settings import __version__ as pmg_version 328 except AttributeError: 329 pmg_version = pymatgen.__version__ 330 331 d = collections.OrderedDict([ 332 ("system", system), 333 ("python_version", platform.python_version()), 334 ("numpy", numpy.version.version), 335 ("scipy", scipy.version.version), 336 ("netCDF4", netCDF4.__version__), 337 ("apscheduler", apscheduler.version), 338 ("pydispatch", pydispatch.__version__), 339 ("yaml", yaml.__version__), 340 ("pymatgen", pmg_version), 341 ]) 342 343 # Optional but strongly suggested. 344 #try: 345 # import matplotlib 346 # d["matplotlib"] = "%s (backend: %s)" % (matplotlib.__version__, matplotlib.get_backend()) 347 #except ImportError: 348 # pass 349 350 return d 351 352 353def abicheck(verbose=0): 354 """ 355 This function tests if the most important ABINIT executables 356 can be found in $PATH and whether the python modules needed 357 at run-time can be imported. Return string with error messages, empty if success. 358 """ 359 360 err_lines = [] 361 app = err_lines.append 362 363 try: 364 manager = TaskManager.from_user_config() 365 except Exception: 366 manager = None 367 app(_straceback()) 368 369 # Get info on the Abinit build. 370 from abipy.core.testing import cmp_version 371 from abipy.flowtk import PyFlowScheduler 372 373 if manager is not None: 374 cprint("AbiPy Manager:\n%s\n" % str(manager), color="green") 375 build = AbinitBuild(manager=manager) 376 if not build.has_netcdf: app("Abinit executable does not support netcdf") 377 cprint("Abinitbuild:\n%s" % str(build), color="magenta") 378 if verbose: cprint(str(build.info), color="magenta") 379 print() 380 if not cmp_version(build.version, min_abinit_version, op=">="): 381 app("Abipy requires Abinit version >= %s but got %s" % (min_abinit_version, build.version)) 382 383 # Get info on the scheduler. 384 try: 385 scheduler = PyFlowScheduler.from_user_config() 386 cprint("Abipy Scheduler:\n%s\n" % str(scheduler), color="yellow") 387 except Exception as exc: 388 app(_straceback()) 389 390 from tabulate import tabulate 391 try: 392 d = software_stack() 393 cprint("Installed packages:", color="blue") 394 cprint(tabulate(list(d.items()), headers=["Package", "Version"]), color="blue") 395 print() 396 except ImportError: 397 app(_straceback()) 398 399 return "\n".join(err_lines) 400 401 402def install_config_files(workdir=None, force_reinstall=False): 403 """ 404 Install pre-defined configuration files for the TaskManager and the Scheduler 405 in the workdir directory. 406 407 Args: 408 workdir: Directory when configuration files should be produced. Use ~/abinit/abipy/ if None 409 force_reinstall: Allow overwrite pre-existent configuration files. By default, the function 410 raises RuntimeError if configuration files are already present. 411 """ 412 workdir = os.path.join(os.path.expanduser("~"), ".abinit", "abipy") if workdir is None else workdir 413 print("Installing configuration files in directory:", workdir) 414 from monty.os import makedirs_p 415 makedirs_p(workdir) 416 417 scheduler_path = os.path.join(workdir, "scheduler.yaml") 418 scheduler_yaml = """ 419# The launcher will stop submitting jobs when the 420# number of jobs in the queue is >= Max number of jobs 421max_njobs_inqueue: 2 422 423# Maximum number of cores that can be used by the scheduler. 424max_ncores_used: 2 425 426# number of hours to wait. 427#hours: 0 428 429# number of minutes to wait. 430#minutes: 0 431 432# number of seconds to wait. 433seconds: 2 434 435# Send mail to the specified address (accepts string or list of strings). 436# PRO TIP: the scheduler WILL try to send and email after a default time of 4 days. If you 437# comment out the mailto address, this will cause the scheduler to terminate, with 438# potentially nefarious effects on your running jobs. If you do not wish to receive 439# emails, a work around is to set the variable `remindme_s` below to something very 440# large (say, 100 days). 441#mailto: nobody@nowhere.com 442 443# verbosity level (int, default 0) 444#verbose: 0 445 446# The scheduler will shutdown when the number of python exceptions is > max_num_pyexcs 447#max_num_pyexcs: 2 448 449# The scheduler will shutdown when the number of Abinit errors is > max_num_abierrs 450#max_num_abierrs: 0 451 452# The scheduler will shutdow when the total number of tasks launched is > safety_ratio * tot_num_tasks. 453#safety_ratio: 5 454 455# Send an e-mail to mailto every remindme_s seconds. 456#remindme_s: 345600 457""" 458 459 manager_path = os.path.join(workdir, "manager.yaml") 460 manager_yaml = """ 461qadapters: 462 - 463 priority: 1 464 queue: 465 qname: abipy 466 qtype: shell 467 job: 468 mpi_runner: mpirun 469 pre_run: 470 - export OMP_NUM_THREADS=1 471 # IMPORTANT: Change the below line so that the abinit executable is in PATH 472 #- export PATH=$HOME/git_repos/abinit/_build/src/98_main:$PATH 473 #- ulimit -s unlimited; ulimit -n 2048 474 475 limits: 476 min_cores: 1 477 max_cores: 2 478 timelimit: 0:10:0 479 hardware: 480 num_nodes: 1 481 sockets_per_node: 1 482 cores_per_socket: 2 483 mem_per_node: 4 Gb 484""" 485 486 # Write configuration files. 487 if not os.path.isfile(scheduler_path) or force_reinstall: 488 with open(scheduler_path, "wt") as fh: 489 fh.write(scheduler_yaml) 490 print("Scheduler configuration file written to:", scheduler_path) 491 else: 492 raise RuntimeError("Configuration file: %s already exists.\nUse force_reinstall option to overwrite it" % scheduler_path) 493 494 if not os.path.isfile(manager_path) or force_reinstall: 495 with open(manager_path, "wt") as fh: 496 fh.write(manager_yaml) 497 print("Manager configuration file written to:", manager_path) 498 else: 499 raise RuntimeError("Configuration file: %s already exists.\nUse force_reinstall option to overwrite it" % manager_path) 500 501 print(""" 502Configuration files installed successfully. 503Please edit the configuration options according to your installation. 504In particular, edit the `pre_run` section in manager.yml 505so that the abinit executable is in $PATH. 506""") 507 508 return 0 509 510 511def abipy_logo1(): 512 """http://www.text-image.com/convert/pic2ascii.cgi""" 513 return r""" 514 `:- -:` 515 --` .+/` ` `/+. .-. 516 `. :+. /s- `yy .yo -s/ :+. .` 517 ./. +o` /s/ `-::-` `yy.-::-` `:- .:::-` -:` .:` /s/ :s- ./. 518.o. /o: .oo. .oyo++syo. `yyys++oys. -ys -syo++sy+` sy- +y: .oo- oo` `o. 519++ oo. /oo yy- -yy `yy: .yy`-ys .ys` /yo sy- +y: oo/ /o: ++ 520+/ oo` /oo `yy. .yy` yy. `yy`-ys :ys :yo oy/ oy: +o/ :o: /o 521-/ :+. -++` -sy+::+yyy` .sy+::+yy- -ys :yys/::oys. `oyo::/syy: `++- /+. /: 522 -- `// /+- -/++/-// -/++/- `+: :yo:/++/. .:++/:oy: -+/ `+- -- 523 `.` -: :/` :yo +y: `/:` `:. `.` 524 `.. .:. .` `. .:. `.. 525 ... ... 526""" 527 528 529def abipy_logo2(): 530 """http://www.text-image.com/convert/pic2ascii.cgi""" 531 return r""" 532MMMMMMMMMMMMMMMMNhdMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMdhmMMMMMMMMMMMMMMM 533MMMMMMMMMddNMMmoyNMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMNMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMNyomMMmhmMMMMMMMM 534MMMmmMMhomMMMy/hMMMMMMMMMMMMMMMMMMMN::MMMMMMMMMm:oMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMd+yMMMhomMmmMMM 535MmsmMMs+NMMMy+yMMMMMMMMMMMNhyyhmMMMN::mhyyhmMMMNhdMMMMmhyydNMMMdyNMMMMMmyNMMMMMMMMMMMMy+yMMMh+dMmsmM 536m+mMMy+hMMMd++mMMMMMMMMMm+:+ss+:+mMN:::/ss+:+mMd:/MMd/:+so/:oNM/:dMMMMMo:yMMMMMMMMMMMMm++dMMMo+NMN+m 537osMMMo+mMMMy+oMMMMMMMMMM::dMMMMd:/MN::hMMMMd::Nd:/Mm:/NMMMMy:oM/:dMMMMMo:yMMMMMMMMMMMMMo+yMMMy+hMMso 538oyMMMooNMMMyooMMMMMMMMMN::mMMMMm::NM::mMMMMN::Nd:/Mh:/MMMMMh:+Mo:yMMMMM+:yMMMMMMMMMMMMMooyMMMyohMMyo 539dyMMMysmMMMdooNMMMMMMMMMd/:oyys:::NMd/:oyys::dMd:/Mh::/shyo:/mMN+:+yhs/::yMMMMMMMMMMMMNoodMMMysmMMyh 540MddMMNyhMMMMysdMMMMMMMMMMMdyooydssMMMMdysoydMMMNsyMh:+hsosydMMMMMmysosho:yMMMMMMMMMMMMdsyMMMNydMMddM 541MMNmNMMddMMMMhyNMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMh:oMMMMMMMMMMMMMMMMMo:yMMMMMMMMMMMNyhMMMNhmMNmNMM 542MMMMMMMMNmmMMMmhmMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMmNMMMMMMMMMMMMMMMMMNdMMMMMMMMMMMmhmMMNmmMMMMMMMM 543MMMMMMMMMMMMMMMMmmNMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMNmmMMMMMMMMMMMMMMM 544""" 545 546 547def abipy_logo3(): 548 """http://www.text-image.com/convert/pic2ascii.cgi""" 549 return r"""\ 550 `-. `--` 551 -:. `//` `/. :: `+: `-:` 552 --``+: `o+ `.--.` -y/.--.` :: `---` `- .. `o+ `o-`-- 553:/ /o /o- -oo/:+s/ -yy+:/os- ss .oo/:/s+`:y. yo :o: :o. /: 554o- o/ oo` ss /y..y/ ss ss +y` -y::y- yo -o/ .o- -o 555:- // /+. :s+--/sy- +s/--+s: ss oyo:-:so``os:-:oyo -+: -+. -/ 556 -` `/. `/: `-:::-:` `-::-` -- oy-:::. .:::-yo //` :- `- 557 ` ..` `:- :+ /: --` `-` ` 558 `.` ..` 559""" 560