1# coding: utf-8 2""" 3This module defines the Robot BaseClass. Robots operates on multiple files and provide helper 4functions to plot the data e.g. convergence studies and to build pandas dataframes from the output files. 5""" 6import sys 7import os 8import inspect 9import itertools 10import numpy as np 11 12from collections import OrderedDict, deque 13from functools import wraps 14from monty.string import is_string, list_strings 15from monty.termcolor import cprint 16from abipy.core.mixins import NotebookWriter 17from abipy.tools.numtools import sort_and_groupby 18from abipy.tools import duck 19from abipy.tools.plotting import (plot_xy_with_hue, add_fig_kwargs, get_ax_fig_plt, get_axarray_fig_plt, 20 rotate_ticklabels, set_visible) 21 22 23class Robot(NotebookWriter): 24 """ 25 This is the base class from which all Robot subclasses should derive. 26 A Robot supports the `with` context manager: 27 28 Usage example: 29 30 .. code-block:: python 31 32 with Robot([("label1", "file1"), (label2, "file2")]) as robot: 33 # Do something with robot. files are automatically closed when we exit. 34 for label, abifile in self.items(): 35 print(label) 36 """ 37 # filepaths are relative to `start`. None for asbolute paths. This flag is set in trim_paths 38 start = None 39 40 # Used in iter_lineopt to generate matplotlib linestyles. 41 _LINE_COLORS = ["b", "r", "g", "m", "y", "k", "c"] 42 _LINE_STYLES = ["-", ":", "--", "-.",] 43 _LINE_WIDTHS = [2, ] 44 45 def __init__(self, *args): 46 """ 47 Args: 48 args is a list of tuples (label, filepath) 49 """ 50 self._abifiles, self._do_close = OrderedDict(), OrderedDict() 51 self._exceptions = deque(maxlen=100) 52 53 for label, abifile in args: 54 self.add_file(label, abifile) 55 56 @classmethod 57 def get_supported_extensions(self): 58 """List of strings with extensions supported by Robot subclasses.""" 59 # This is needed to have all subclasses. 60 from abipy.abilab import Robot 61 return sorted([cls.EXT for cls in Robot.__subclasses__()]) 62 63 @classmethod 64 def class_for_ext(cls, ext): 65 """Return the Robot subclass associated to the given extension.""" 66 for subcls in cls.__subclasses__(): 67 if subcls.EXT in (ext, ext.upper()): 68 return subcls 69 70 # anaddb.nc does not follow the extension rule... 71 if ext.lower() == "anaddb": 72 from abipy.dfpt.anaddbnc import AnaddbNcRobot as subcls 73 return subcls 74 75 raise ValueError("Cannot find Robot subclass associated to extension %s\n" % ext + 76 "The list of supported extensions (case insensitive) is:\n%s" % 77 str(cls.get_supported_extensions())) 78 79 @classmethod 80 def from_dir(cls, top, walk=True, abspath=False): 81 """ 82 This class method builds a robot by scanning all files located within directory `top`. 83 This method should be invoked with a concrete robot class, for example: 84 85 robot = GsrRobot.from_dir(".") 86 87 Args: 88 top (str): Root directory 89 walk: if True, directories inside `top` are included as well. 90 abspath: True if paths in index should be absolute. Default: Relative to `top`. 91 """ 92 new = cls(*cls._open_files_in_dir(top, walk)) 93 if not abspath: new.trim_paths(start=top) 94 return new 95 96 @classmethod 97 def from_dirs(cls, dirpaths, walk=True, abspath=False): 98 """ 99 Similar to `from_dir` but accepts a list of directories instead of a single directory. 100 101 Args: 102 walk: if True, directories inside `top` are included as well. 103 abspath: True if paths in index should be absolute. Default: Relative to `top`. 104 """ 105 items = [] 106 for top in list_strings(dirpaths): 107 items.extend(cls._open_files_in_dir(top, walk)) 108 new = cls(*items) 109 if not abspath: new.trim_paths(start=os.getcwd()) 110 return new 111 112 @classmethod 113 def from_dir_glob(cls, pattern, walk=True, abspath=False): 114 """ 115 This class method builds a robot by scanning all files located within the directories 116 matching `pattern` as implemented by glob.glob 117 This method should be invoked with a concrete robot class, for example: 118 119 robot = GsrRobot.from_dir_glob("flow_dir/w*/outdata/") 120 121 Args: 122 pattern: Pattern string 123 walk: if True, directories inside `top` are included as well. 124 abspath: True if paths in index should be absolute. Default: Relative to getcwd(). 125 """ 126 import glob 127 items = [] 128 for top in filter(os.path.isdir, glob.iglob(pattern)): 129 items += cls._open_files_in_dir(top, walk=walk) 130 new = cls(*items) 131 if not abspath: new.trim_paths(start=os.getcwd()) 132 return new 133 134 @classmethod 135 def _open_files_in_dir(cls, top, walk): 136 """Open files in directory tree starting from `top`. Return list of Abinit files.""" 137 if not os.path.isdir(top): 138 raise ValueError("%s: no such directory" % str(top)) 139 from abipy.abilab import abiopen 140 items = [] 141 if walk: 142 for dirpath, dirnames, filenames in os.walk(top): 143 filenames = sorted([f for f in filenames if cls.class_handles_filename(f)]) 144 for f in filenames: 145 abifile = abiopen(os.path.join(dirpath, f)) 146 if abifile is not None: items.append((abifile.filepath, abifile)) 147 else: 148 filenames = [f for f in os.listdir(top) if cls.class_handles_filename(f)] 149 for f in filenames: 150 abifile = abiopen(os.path.join(top, f)) 151 if abifile is not None: items.append((abifile.filepath, abifile)) 152 153 return items 154 155 @classmethod 156 def class_handles_filename(cls, filename): 157 """True if robot class handles filename.""" 158 # Special treatment of AnaddbNcRobot 159 if cls.EXT == "anaddb" and os.path.basename(filename).lower() == "anaddb.nc": 160 return True 161 162 return (filename.endswith("_" + cls.EXT + ".nc") or 163 filename.endswith("." + cls.EXT)) # This for .abo 164 165 @classmethod 166 def from_files(cls, filenames, labels=None, abspath=False): 167 """ 168 Build a Robot from a list of `filenames`. 169 if labels is None, labels are automatically generated from absolute paths. 170 171 Args: 172 abspath: True if paths in index should be absolute. Default: Relative to `top`. 173 """ 174 filenames = list_strings(filenames) 175 from abipy.abilab import abiopen 176 filenames = [f for f in filenames if cls.class_handles_filename(f)] 177 items = [] 178 for i, f in enumerate(filenames): 179 try: 180 abifile = abiopen(f) 181 except Exception as exc: 182 cprint("Exception while opening file: `%s`" % str(f), "red") 183 cprint(exc, "red") 184 abifile = None 185 186 if abifile is not None: 187 label = abifile.filepath if labels is None else labels[i] 188 items.append((label, abifile)) 189 190 new = cls(*items) 191 if labels is None and not abspath: new.trim_paths(start=None) 192 return new 193 194 @classmethod 195 def from_flow(cls, flow, outdirs="all", nids=None, ext=None, task_class=None): 196 """ 197 Build a robot from a |Flow| object. 198 199 Args: 200 flow: |Flow| object 201 outdirs: String used to select/ignore the files in the output directory of flow, works and tasks 202 outdirs="work" selects only the outdir of the Works, 203 outdirs="flow+task" selects the outdir of the Flow and the outdirs of the tasks 204 outdirs="-work" excludes the outdir of the Works. 205 Cannot use ``+`` and ``-`` flags in the same string. 206 Default: `all` that is equivalent to "flow+work+task" 207 nids: List of node identifiers used to select particular nodes. Not used if None 208 ext: File extension associated to the robot. Mainly used if method is invoked with the BaseClass 209 task_class: Task class or string with the class name used to select the tasks in the flow. 210 None implies no filtering. 211 212 Usage example: 213 214 .. code-block:: python 215 216 with abilab.GsrRobot.from_flow(flow) as robot: 217 print(robot) 218 219 # That is equivalent to: 220 with Robot.from_flow(flow, ext="GSR") as robot: 221 print(robot) 222 223 Returns: 224 ``Robot`` subclass. 225 """ 226 robot = cls() if ext is None else cls.class_for_ext(ext)() 227 all_opts = ("flow", "work", "task") 228 229 if outdirs == "all": 230 tokens = all_opts 231 elif "+" in outdirs: 232 assert "-" not in outdirs 233 tokens = outdirs.split("+") 234 elif "-" in outdirs: 235 assert "+" not in outdirs 236 tokens = [s for s in all if s not in outdirs.split("-")] 237 else: 238 tokens = list_strings(outdirs) 239 240 if not all(t in all_opts for t in tokens): 241 raise ValueError("Wrong outdirs string %s" % outdirs) 242 243 if "flow" in tokens: 244 robot.add_extfile_of_node(flow, nids=nids, task_class=task_class) 245 246 if "work" in tokens: 247 for work in flow: 248 robot.add_extfile_of_node(work, nids=nids, task_class=task_class) 249 250 if "task" in tokens: 251 for task in flow.iflat_tasks(): 252 robot.add_extfile_of_node(task, nids=nids, task_class=task_class) 253 254 return robot 255 256 def add_extfile_of_node(self, node, nids=None, task_class=None): 257 """ 258 Add the file produced by this node to the robot. 259 260 Args: 261 node: |Flow| or |Work| or |Task| object. 262 nids: List of node identifiers used to select particular nodes. Not used if None 263 task_class: Task class or string with class name used to select the tasks in the flow. 264 None implies no filtering. 265 """ 266 if nids and node.node_id not in nids: return 267 filepath = node.outdir.has_abiext(self.EXT) 268 if not filepath: 269 # Look in run.abi directory. 270 filepath = node.wdir.has_abiext(self.EXT) 271 272 # This to ignore DDB.nc files (only text DDB are supported) 273 if filepath and filepath.endswith("_DDB.nc"): 274 return 275 276 if filepath: 277 try: 278 label = os.path.relpath(filepath) 279 except OSError: 280 # current working directory may not be defined! 281 label = filepath 282 283 # Filter by task_class (class or string with class name) 284 if task_class is not None and not node.isinstance(task_class): 285 return None 286 287 self.add_file(label, filepath) 288 289 def scan_dir(self, top, walk=True): 290 """ 291 Scan directory tree starting from ``top``. Add files to the robot instance. 292 293 Args: 294 top (str): Root directory 295 walk: if True, directories inside ``top`` are included as well. 296 297 Return: 298 Number of files found. 299 """ 300 count = 0 301 for filepath, abifile in self.__class__._open_files_in_dir(top, walk): 302 count += 1 303 self.add_file(filepath, abifile) 304 305 return count 306 307 def add_file(self, label, abifile, filter_abifile=None): 308 """ 309 Add a file to the robot with the given label. 310 311 Args: 312 label: String used to identify the file (must be unique, ax exceptions is 313 raised if label is already present. 314 abifile: Specify the file to be added. Accepts strings (filepath) or abipy file-like objects. 315 filter_abifile: Function that receives an ``abifile`` object and returns 316 True if the file should be added to the plotter. 317 """ 318 if is_string(abifile): 319 from abipy.abilab import abiopen 320 abifile = abiopen(abifile) 321 if filter_abifile is not None and not filter_abifile(abifile): 322 abifile.close() 323 return 324 325 # Open file here --> have to close it. 326 self._do_close[abifile.filepath] = True 327 328 if label in self._abifiles: 329 raise ValueError("label %s is already present!" % label) 330 331 self._abifiles[label] = abifile 332 333 #def pop_filepath(self, filepath): 334 # """ 335 # Remove the file with the given `filepath` and close it. 336 # """ 337 # if label, abifile in self._abifiles.items(): 338 # if abifile.filepath != filepath: continue 339 # self._abifiles.pop(label) 340 # if self._do_close.pop(abifile.filepath, False): 341 # try: 342 # abifile.close() 343 # except Exception as exc: 344 # print("Exception while closing: ", abifile.filepath) 345 # print(exc) 346 347 def iter_lineopt(self): 348 """Generates matplotlib linestyles.""" 349 for o in itertools.product( self._LINE_WIDTHS, self._LINE_STYLES, self._LINE_COLORS): 350 yield {"linewidth": o[0], "linestyle": o[1], "color": o[2]} 351 352 @staticmethod 353 def ordered_intersection(list_1, list_2): 354 """Return ordered intersection of two lists. Items must be hashable.""" 355 set_2 = frozenset(list_2) 356 return [x for x in list_1 if x in set_2] 357 358 #def _get_ointersection_i(self, iattrname): 359 # if len(self.abifiles) == 0: return [] 360 # values = list(range(getattr(self.abifiles[0], iattrname))) 361 # if len(self.abifiles) == 1: return values 362 # for abifile in self.abifiles[1:]: 363 # values = self.ordered_intersection(values, range(getattr(abifile, iattrname))) 364 # return values 365 366 @staticmethod 367 def _to_relpaths(paths): 368 """Convert a list of absolute paths to relative paths.""" 369 root = os.getcwd() 370 return [os.path.relpath(p, root) for p in paths] 371 372 def pop_label(self, label): 373 """ 374 Remove file with the given ``label`` and close it. 375 """ 376 if label in self._abifiles: 377 abifile = self._abifiles.pop(label) 378 if self._do_close.pop(abifile.filepath, False): 379 try: 380 abifile.close() 381 except Exception as exc: 382 print("Exception while closing: ", abifile.filepath) 383 print(exc) 384 385 def change_labels(self, new_labels, dryrun=False): 386 """ 387 Change labels of the files. 388 389 Args: 390 new_labels: List of strings (same length as self.abifiles) 391 dryrun: True to activate dryrun mode. 392 393 Return: 394 mapping new_label --> old_label. 395 """ 396 if len(new_labels) != len(self): 397 raise ValueError("Robot has %d files while len(new_labels) = %d" % (len(new_labels), len(self))) 398 399 old_labels = list(self._abifiles.keys()) 400 if not dryrun: 401 old_abifiles, self._abifiles = self._abifiles, OrderedDict() 402 new2old = OrderedDict() 403 for old, new in zip(old_labels, new_labels): 404 new2old[new] = old 405 if not dryrun: 406 self._abifiles[new] = old_abifiles[old] 407 else: 408 print("old [%s] --> new [%s]" % (old, new)) 409 410 return new2old 411 412 def remap_labels(self, function, dryrun=False): 413 """ 414 Change labels of the files by executing ``function`` 415 416 Args: 417 function: Callable object e.g. lambda function. The output of function(abifile) is used as 418 new label. Note that the function shall not return duplicated labels when applied to self.abifiles. 419 dryrun: True to activate dryrun mode. 420 421 Return: 422 mapping new_label --> old_label. 423 """ 424 new_labels = [function(afile) for afile in self.abifiles] 425 # Labels must be unique and hashable. 426 if len(set(new_labels)) != len(new_labels): 427 raise ValueError("Duplicated labels are not allowed. Change input function.\nnew_labels %s" % str(new_labels)) 428 429 return self.change_labels(new_labels, dryrun=dryrun) 430 431 def trim_paths(self, start=None): 432 """ 433 Replace absolute filepaths in the robot with relative paths wrt to ``start`` directory. 434 If start is None, os.getcwd() is used. Set ``self.start`` attribute, return ``self.start``. 435 """ 436 self.start = os.getcwd() if start is None else start 437 old_paths = list(self._abifiles.keys()) 438 old_new_paths = [(p, os.path.relpath(os.path.abspath(p), start=self.start)) for p in old_paths] 439 440 old_abifiles = self._abifiles 441 self._abifiles = OrderedDict() 442 for old, new in old_new_paths: 443 self._abifiles[new] = old_abifiles[old] 444 445 return self.start 446 447 @property 448 def exceptions(self): 449 """List of exceptions.""" 450 return self._exceptions 451 452 def __len__(self): 453 return len(self._abifiles) 454 455 #def __iter__(self): 456 # return iter(self._abifiles) 457 458 #def __contains__(self, item): 459 # return item in self._abifiles 460 461 def __getitem__(self, key): 462 # self[key] 463 return self._abifiles.__getitem__(key) 464 465 def __enter__(self): 466 return self 467 468 def __exit__(self, exc_type, exc_val, exc_tb): 469 """Activated at the end of the with statement.""" 470 self.close() 471 472 def keys(self): 473 return self._abifiles.keys() 474 475 def items(self): 476 return self._abifiles.items() 477 478 @property 479 def labels(self): 480 """ 481 List of strings used to create labels in matplotlib figures when plotting results 482 taked from multiple files. By default, labels is initialized with the path of the files in the robot. 483 Use change_labels to change the list. 484 """ 485 return list(self._abifiles.keys()) 486 487 def get_label_files_str(self): 488 """Return string with [label, filepath].""" 489 from tabulate import tabulate 490 return tabulate([(label, abifile.relpath) for label, abifile in self.items()], headers=["Label", "Relpath"]) + "\n" 491 492 def show_files(self, stream=sys.stdout): 493 """Show label --> file path""" 494 stream.write(self.get_label_files_str()) 495 496 def __repr__(self): 497 """Invoked by repr.""" 498 return self.get_label_files_str() 499 500 def __str__(self): 501 """Invoked by str.""" 502 return self.to_string() 503 504 def to_string(self, verbose=0): 505 """String representation.""" 506 lines = ["%s with %d files in memory:\n" % (self.__class__.__name__, len(self.abifiles))] 507 app = lines.append 508 for i, f in enumerate(self.abifiles): 509 app(f.to_string(verbose=verbose)) 510 app("\n") 511 512 return "\n".join(lines) 513 514 def _repr_html_(self): 515 """Integration with jupyter_ notebooks.""" 516 return '<ol start="0">\n{}\n</ol>'.format("\n".join("<li>%s</li>" % label for label, abifile in self.items())) 517 518 @property 519 def abifiles(self): 520 """List of netcdf files.""" 521 return list(self._abifiles.values()) 522 523 def has_different_structures(self, rtol=1e-05, atol=1e-08): 524 """ 525 Check if structures are equivalent, 526 return string with info about differences (if any). 527 """ 528 if len(self) <= 1: return "" 529 formulas = set([af.structure.composition.formula for af in self.abifiles]) 530 if len(formulas) != 1: 531 return "Found structures with different full formulas: %s" % str(formulas) 532 533 lines = [] 534 s0 = self.abifiles[0].structure 535 for abifile in self.abifiles[1:]: 536 s1 = abifile.structure 537 if not np.allclose(s0.lattice.matrix, s1.lattice.matrix, rtol=rtol, atol=atol): 538 lines.append("Structures have different lattice:") 539 if not np.allclose(s0.frac_coords, s1.frac_coords, rtol=rtol, atol=atol): 540 lines.append("Structures have different atomic positions:") 541 542 return "\n".join(lines) 543 544 #def apply(self, func_or_string, args=(), **kwargs): 545 # """ 546 # Applies function to all ``abifiles`` available in the robot. 547 548 # Args: 549 # func_or_string: If callable, the output of func_or_string(abifile, ...) is used. 550 # If string, the output of getattr(abifile, func_or_string)(...) 551 # args (tuple): Positional arguments to pass to function in addition to the array/series 552 # kwargs: Additional keyword arguments will be passed as keywords to the function 553 554 # Return: List of results 555 # """ 556 # if callable(func_or_string): 557 # return [func_or_string(abifile, *args, *kwargs) for abifile in self.abifiles] 558 # else: 559 # return [duck.getattrd(abifile, func_or_string)(*args, **kwargs) for abifile in self.abifiles] 560 561 def is_sortable(self, aname, raise_exc=False): 562 """ 563 Return True if ``aname`` is an attribute of the netcdf file 564 If raise_exc is True, AttributeError with an explicit message is raised. 565 """ 566 try: 567 obj = None 568 try: 569 # abiifile.foo.bar? 570 obj = duck.getattrd(self.abifiles[0], aname) 571 except AttributeError: 572 # abifile.params[aname] ? 573 if hasattr(self.abifiles[0], "params") and aname in self.abifiles[0].params: 574 obj = self.abifiles[0].params[aname] 575 576 # Let's try to convert obj to scalar. 577 float(obj) 578 return True 579 580 except Exception: 581 if not raise_exc: return False 582 attrs = [] 583 for key, obj in inspect.getmembers(self.abifiles[0]): 584 # Ignores anything starting with underscore 585 if key.startswith('_') or callable(obj) or hasattr(obj, "__len__"): continue 586 attrs.append(key) 587 588 # Add entries in params. 589 if hasattr(self.abifiles[0], "params") and hasattr(self.abifiles[0].params, "keys"): 590 attrs.extend(self.abifiles[0].params.keys()) 591 592 raise AttributeError("""\ 593`%s` object has no attribute `%s`. Choose among: 594 595 %s 596 597Note that this list is automatically generated. 598Not all entries are sortable (Please select number-like quantities)""" % (self.__class__.__name__, aname, str(attrs))) 599 600 def _sortby_labelfile_list(self, labelfile_list, func_or_string, reverse=False, unpack=False): 601 """ 602 Return: list of (label, abifile, param) tuples where param is obtained via ``func_or_string``. 603 or labels, abifiles, params if ``unpack`` 604 """ 605 if not func_or_string: 606 # Catch None or empty 607 items = [(label, abifile, label) for (label, abifile) in labelfile_list] 608 if not unpack: 609 return items 610 else: 611 return [t[0] for t in items], [t[1] for t in items], [t[2] for t in items] 612 613 elif callable(func_or_string): 614 items = [(label, abifile, func_or_string(abifile)) for (label, abifile) in labelfile_list] 615 616 else: 617 # Assume string and attribute with the same name. 618 # try in abifile.params if not hasattrd(abifile, func_or_string) 619 self.is_sortable(func_or_string, raise_exc=True) 620 if duck.hasattrd(self.abifiles[0], func_or_string): 621 items = [(label, abifile, duck.getattrd(abifile, func_or_string)) for (label, abifile) in labelfile_list] 622 else: 623 items = [(label, abifile, abifile.params[func_or_string]) for (label, abifile) in labelfile_list] 624 625 items = sorted(items, key=lambda t: t[2], reverse=reverse) 626 if not unpack: 627 return items 628 else: 629 return [t[0] for t in items], [t[1] for t in items], [t[2] for t in items] 630 631 def sortby(self, func_or_string, reverse=False, unpack=False): 632 """ 633 Sort files in the robot by ``func_or_string``. 634 635 Args: 636 func_or_string: Either None, string, callable defining the quantity to be used for sorting. 637 If string, it's assumed that the abifile has an attribute with the same name and getattr is invoked. 638 If callable, the output of func_or_string(abifile) is used. 639 If None, no sorting is performed. 640 reverse: If set to True, then the list elements are sorted as if each comparison were reversed. 641 unpack: Return (labels, abifiles, params) if True 642 643 Return: list of (label, abifile, param) tuples where param is obtained via ``func_or_string``. 644 or labels, abifiles, params if ``unpack`` 645 """ 646 labelfile_list = list(self.items()) 647 return self._sortby_labelfile_list(labelfile_list, func_or_string, reverse=reverse, unpack=unpack) 648 649 def group_and_sortby(self, hue, func_or_string): 650 """ 651 Group files by ``hue`` and, inside each group` sort items by ``func_or_string``. 652 653 Args: 654 hue: Variable that define subsets of the data, which will be drawn on separate lines. 655 Accepts callable or string 656 If string, it's assumed that the abifile has an attribute with the same name and getattr is invoked. 657 Dot notation is also supported e.g. hue="structure.formula" --> abifile.structure.formula 658 If callable, the output of hue(abifile) is used. 659 func_or_string: Either None, string, callable defining the quantity to be used for sorting. 660 If string, it's assumed that the abifile has an attribute with the same name and getattr is invoked. 661 If callable, the output of func_or_string(abifile) is used. 662 If None, no sorting is performed. 663 664 Return: List of :class:`HueGroup` instance. 665 """ 666 # Group by hue. 667 # This is the section in which we support: callable, abifile.attr.name syntax or abifile.params["key"] 668 items = list(self.items()) 669 670 if callable(hue): 671 key = lambda t: hue(t[1]) 672 else: 673 # Assume string. 674 if duck.hasattrd(self.abifiles[0], hue): 675 key = lambda t: duck.getattrd(t[1], hue) 676 else: 677 # Try in abifile.params 678 if hasattr(self.abifiles[0], "params") and hue in self.abifiles[0].params: 679 key = lambda t: t[1].params[hue] 680 else: 681 raise TypeError("""\ 682Cannot interpret hue argument of type `%s` and value `%s`. 683Expecting callable or attribute name or key in abifile.params""" % (type(hue), str(hue))) 684 685 groups = [] 686 for hvalue, labelfile_list in sort_and_groupby(items, key=key): 687 # Use func_or_string to sort each group 688 labels, abifiles, xvalues = self._sortby_labelfile_list(labelfile_list, func_or_string, unpack=True) 689 groups.append(HueGroup(hvalue, xvalues, abifiles, labels)) 690 691 return groups 692 693 def close(self): 694 """ 695 Close all files that have been opened by the Robot. 696 """ 697 for abifile in self.abifiles: 698 if self._do_close.pop(abifile.filepath, False): 699 try: 700 abifile.close() 701 except Exception as exc: 702 print("Exception while closing: ", abifile.filepath) 703 print(exc) 704 705 #@classmethod 706 #def open(cls, obj, nids=None, **kwargs): 707 # """ 708 # Flexible constructor. obj can be a :class:`Flow` or a string with the directory containing the Flow. 709 # `nids` is an optional list of :class:`Node` identifiers used to filter the set of :class:`Task` in the Flow. 710 # """ 711 # has_dirpath = False 712 # if is_string(obj): 713 # try: 714 # from abipy.flowtk import Flow 715 # obj = Flow.pickle_load(obj) 716 # except: 717 # has_dirpath = True 718 719 # if not has_dirpath: 720 # # We have a Flow. smeth is the name of the Task method used to open the file. 721 # items = [] 722 # smeth = "open_" + cls.EXT.lower() 723 # for task in obj.iflat_tasks(nids=nids): #, status=obj.S_OK): 724 # open_method = getattr(task, smeth, None) 725 # if open_method is None: continue 726 # abifile = open_method() 727 # if abifile is not None: items.append((task.pos_str, abifile)) 728 # return cls(*items) 729 730 # else: 731 # # directory --> search for files with the appropriate extension and open it with abiopen. 732 # if nids is not None: raise ValueError("nids cannot be used when obj is a directory.") 733 # return cls.from_dir(obj) 734 735 #def get_attributes(self, attr_name, obj=None, retdict=False): 736 # od = OrderedDict() 737 # for label, abifile in self.items(): 738 # obj = abifile if obj is None else getattr(abifile, obj) 739 # od[label] = getattr(obj, attr_name) 740 # if retdict: 741 # return od 742 # else: 743 # return list(od.values()) 744 745 def _exec_funcs(self, funcs, arg): 746 """ 747 Execute list of callable functions. Each function receives arg as argument. 748 """ 749 if not isinstance(funcs, (list, tuple)): funcs = [funcs] 750 d = {} 751 for func in funcs: 752 try: 753 key, value = func(arg) 754 d[key] = value 755 except Exception as exc: 756 cprint("Exception: %s" % str(exc), "red") 757 self._exceptions.append(str(exc)) 758 return d 759 760 @staticmethod 761 def sortby_label(sortby, param): 762 """Return the label to be used when files are sorted with ``sortby``.""" 763 return "%s %s" % (sortby, param) if not (callable(sortby) or sortby is None) else str(param) 764 765 def get_structure_dataframes(self, abspath=False, filter_abifile=None, **kwargs): 766 """ 767 Wrap dataframes_from_structures function. 768 769 Args: 770 abspath: True if paths in index should be absolute. Default: Relative to getcwd(). 771 filter_abifile: Function that receives an ``abifile`` object and returns 772 True if the file should be added to the plotter. 773 """ 774 from abipy.core.structure import dataframes_from_structures 775 if "index" not in kwargs: 776 index = list(self._abifiles.keys()) 777 if not abspath: index = self._to_relpaths(index) 778 kwargs["index"] = index 779 780 abifiles = self.abifiles if filter_abifile is not None else list(filter(filter_abifile, self.abifiles)) 781 return dataframes_from_structures(struct_objects=abifiles, **kwargs) 782 783 def get_lattice_dataframe(self, **kwargs): 784 """Return |pandas-DataFrame| with lattice parameters.""" 785 dfs = self.get_structure_dataframes(**kwargs) 786 return dfs.lattice 787 788 def get_coords_dataframe(self, **kwargs): 789 """Return |pandas-DataFrame| with atomic positions.""" 790 dfs = self.get_structure_dataframes(**kwargs) 791 return dfs.coords 792 793 def get_params_dataframe(self, abspath=False): 794 """ 795 Return |pandas-DataFrame| with the most important parameters. 796 that are usually subject to convergence studies. 797 798 Args: 799 abspath: True if paths in index should be absolute. Default: Relative to `top`. 800 """ 801 rows, row_names = [], [] 802 for label, abifile in self.items(): 803 if not hasattr(abifile, "params"): 804 import warnings 805 warnings.warn("%s does not have `params` attribute" % type(abifile)) 806 break 807 rows.append(abifile.params) 808 row_names.append(label) 809 810 row_names = row_names if abspath else self._to_relpaths(row_names) 811 import pandas as pd 812 return pd.DataFrame(rows, index=row_names, columns=list(rows[0].keys())) 813 814 ############################################## 815 # Helper functions to plot pandas dataframes # 816 ############################################## 817 818 @staticmethod 819 @wraps(plot_xy_with_hue) 820 def plot_xy_with_hue(*args, **kwargs): 821 return plot_xy_with_hue(*args, **kwargs) 822 823 @staticmethod 824 def _get_label(func_or_string): 825 """ 826 Return label associated to ``func_or_string``. 827 If callable, docstring __doc__ is used. 828 """ 829 if func_or_string is None: 830 return "" 831 elif callable(func_or_string): 832 if getattr(func_or_string, "__doc__", ""): 833 return func_or_string.__doc__.strip() 834 else: 835 return func_or_string.__name__ 836 else: 837 return str(func_or_string) 838 839 @add_fig_kwargs 840 def plot_convergence(self, item, sortby=None, hue=None, ax=None, fontsize=8, **kwargs): 841 """ 842 Plot the convergence of ``item`` wrt the ``sortby`` parameter. 843 Values can optionally be grouped by ``hue``. 844 845 Args: 846 item: Define the quantity to plot. Accepts callable or string 847 If string, it's assumed that the abifile has an attribute with the same name and `getattr` is invoked. 848 Dot notation is also supported e.g. hue="structure.formula" --> abifile.structure.formula 849 If callable, the output of item(abifile) is used. 850 sortby: Define the convergence parameter, sort files and produce plot labels. 851 Can be None, string or function. If None, no sorting is performed. 852 If string and not empty it's assumed that the abifile has an attribute 853 with the same name and `getattr` is invoked. 854 If callable, the output of sortby(abifile) is used. 855 hue: Variable that define subsets of the data, which will be drawn on separate lines. 856 Accepts callable or string 857 If string, it's assumed that the abifile has an attribute with the same name and getattr is invoked. 858 If callable, the output of hue(abifile) is used. 859 ax: |matplotlib-Axes| or None if a new figure should be created. 860 fontsize: legend and label fontsize. 861 kwargs: keyword arguments passed to matplotlib plot method. 862 863 Returns: |matplotlib-Figure| 864 865 Example: 866 867 robot.plot_convergence("energy") 868 869 robot.plot_convergence("energy", sortby="nkpt") 870 871 robot.plot_convergence("pressure", sortby="nkpt", hue="tsmear") 872 """ 873 ax, fig, plt = get_ax_fig_plt(ax=ax) 874 if "marker" not in kwargs: 875 kwargs["marker"] = "o" 876 877 def get_yvalues(abifiles): 878 if callable(item): 879 return [float(item(a)) for a in abifiles] 880 else: 881 return [float(getattr(a, item)) for a in abifiles] 882 883 if hue is None: 884 labels, abifiles, params = self.sortby(sortby, unpack=True) 885 yvals = get_yvalues(abifiles) 886 #print("params", params, "\nyvals", yvals) 887 ax.plot(params, yvals, **kwargs) 888 else: 889 groups = self.group_and_sortby(hue, sortby) 890 for g in groups: 891 yvals = get_yvalues(g.abifiles) 892 label = "%s: %s" % (self._get_label(hue), g.hvalue) 893 ax.plot(g.xvalues, yvals, label=label, **kwargs) 894 895 ax.grid(True) 896 ax.set_xlabel("%s" % self._get_label(sortby)) 897 if sortby is None: rotate_ticklabels(ax, 15) 898 ax.set_ylabel("%s" % self._get_label(item)) 899 900 if hue is not None: 901 ax.legend(loc="best", fontsize=fontsize, shadow=True) 902 903 return fig 904 905 @add_fig_kwargs 906 def plot_convergence_items(self, items, sortby=None, hue=None, fontsize=6, **kwargs): 907 """ 908 Plot the convergence of a list of ``items`` wrt to the ``sortby`` parameter. 909 Values can optionally be grouped by ``hue``. 910 911 Args: 912 items: List of attributes (or callables) to be analyzed. 913 sortby: Define the convergence parameter, sort files and produce plot labels. 914 Can be None, string or function. If None, no sorting is performed. 915 If string and not empty it's assumed that the abifile has an attribute 916 with the same name and `getattr` is invoked. 917 If callable, the output of sortby(abifile) is used. 918 hue: Variable that define subsets of the data, which will be drawn on separate lines. 919 Accepts callable or string 920 If string, it's assumed that the abifile has an attribute with the same name and getattr is invoked. 921 Dot notation is also supported e.g. hue="structure.formula" --> abifile.structure.formula 922 If callable, the output of hue(abifile) is used. 923 fontsize: legend and label fontsize. 924 kwargs: keyword arguments are passed to ax.plot 925 926 Returns: |matplotlib-Figure| 927 """ 928 # Note: in principle one could call plot_convergence inside a loop but 929 # this one is faster as sorting is done only once. 930 931 # Build grid plot. 932 nrows, ncols = len(items), 1 933 ax_list, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols, 934 sharex=True, sharey=False, squeeze=False) 935 ax_list = ax_list.ravel() 936 937 # Sort and group files if hue. 938 if hue is None: 939 labels, ncfiles, params = self.sortby(sortby, unpack=True) 940 else: 941 groups = self.group_and_sortby(hue, sortby) 942 943 marker = kwargs.pop("marker", "o") 944 for i, (ax, item) in enumerate(zip(ax_list, items)): 945 if hue is None: 946 # Extract data. 947 if callable(item): 948 yvals = [float(item(gsr)) for gsr in self.abifiles] 949 else: 950 yvals = [duck.getattrd(gsr, item) for gsr in self.abifiles] 951 952 if not is_string(params[0]): 953 ax.plot(params, yvals, marker=marker, **kwargs) 954 else: 955 # Must handle list of strings in a different way. 956 xn = range(len(params)) 957 ax.plot(xn, yvals, marker=marker, **kwargs) 958 ax.set_xticks(xn) 959 ax.set_xticklabels(params, fontsize=fontsize) 960 else: 961 for g in groups: 962 # Extract data. 963 if callable(item): 964 yvals = [float(item(gsr)) for gsr in g.abifiles] 965 else: 966 yvals = [duck.getattrd(gsr, item) for gsr in g.abifiles] 967 label = "%s: %s" % (self._get_label(hue), g.hvalue) 968 ax.plot(g.xvalues, yvals, label=label, marker=marker, **kwargs) 969 970 ax.grid(True) 971 ax.set_ylabel(self._get_label(item)) 972 if i == len(items) - 1: 973 ax.set_xlabel("%s" % self._get_label(sortby)) 974 if sortby is None: rotate_ticklabels(ax, 15) 975 if i == 0 and hue is not None: 976 ax.legend(loc="best", fontsize=fontsize, shadow=True) 977 978 return fig 979 980 @add_fig_kwargs 981 def plot_lattice_convergence(self, what_list=None, sortby=None, hue=None, fontsize=8, **kwargs): 982 """ 983 Plot the convergence of the lattice parameters (a, b, c, alpha, beta, gamma). 984 wrt the``sortby`` parameter. Values can optionally be grouped by ``hue``. 985 986 Args: 987 what_list: List of strings with the quantities to plot e.g. ["a", "alpha", "beta"]. 988 None means all. 989 item: Define the quantity to plot. Accepts callable or string 990 If string, it's assumed that the abifile has an attribute 991 with the same name and `getattr` is invoked. 992 If callable, the output of item(abifile) is used. 993 sortby: Define the convergence parameter, sort files and produce plot labels. 994 Can be None, string or function. 995 If None, no sorting is performed. 996 If string and not empty it's assumed that the abifile has an attribute 997 with the same name and `getattr` is invoked. 998 If callable, the output of sortby(abifile) is used. 999 hue: Variable that define subsets of the data, which will be drawn on separate lines. 1000 Accepts callable or string 1001 If string, it's assumed that the abifile has an attribute with the same name and getattr is invoked. 1002 Dot notation is also supported e.g. hue="structure.formula" --> abifile.structure.formula 1003 If callable, the output of hue(abifile) is used. 1004 ax: |matplotlib-Axes| or None if a new figure should be created. 1005 fontsize: legend and label fontsize. 1006 1007 Returns: |matplotlib-Figure| 1008 1009 Example: 1010 1011 robot.plot_lattice_convergence() 1012 1013 robot.plot_lattice_convergence(sortby="nkpt") 1014 1015 robot.plot_lattice_convergence(sortby="nkpt", hue="tsmear") 1016 """ 1017 if not self.abifiles: return None 1018 1019 # The majority of AbiPy files have a structure object 1020 # whereas Hist.nc defines final_structure. Use geattr and key to extract structure object. 1021 key = "structure" 1022 if not hasattr(self.abifiles[0], "structure"): 1023 if hasattr(self.abifiles[0], "final_structure"): 1024 key = "final_structure" 1025 else: 1026 raise TypeError("Don't know how to extract structure from %s" % type(self.abifiles[0])) 1027 1028 # Define callbacks. docstrings will be used as ylabels. 1029 def a(afile): 1030 "a (Ang)" 1031 return getattr(afile, key).lattice.a 1032 def b(afile): 1033 "b (Ang)" 1034 return getattr(afile, key).lattice.b 1035 def c(afile): 1036 "c (Ang)" 1037 return getattr(afile, key).lattice.c 1038 def volume(afile): 1039 r"$V$" 1040 return getattr(afile, key).lattice.volume 1041 def alpha(afile): 1042 r"$\alpha$" 1043 return getattr(afile, key).lattice.alpha 1044 def beta(afile): 1045 r"$\beta$" 1046 return getattr(afile, key).lattice.beta 1047 def gamma(afile): 1048 r"$\gamma$" 1049 return getattr(afile, key).lattice.gamma 1050 1051 items = [a, b, c, volume, alpha, beta, gamma] 1052 if what_list is not None: 1053 locs = locals() 1054 items = [locs[what] for what in list_strings(what_list)] 1055 1056 # Build plot grid. 1057 nrows, ncols = len(items), 1 1058 ax_list, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols, 1059 sharex=True, sharey=False, squeeze=False) 1060 1061 marker = kwargs.pop("marker", "o") 1062 for i, (ax, item) in enumerate(zip(ax_list.ravel(), items)): 1063 self.plot_convergence(item, sortby=sortby, hue=hue, ax=ax, fontsize=fontsize, 1064 marker=marker, show=False) 1065 if i != 0: 1066 set_visible(ax, False, "legend") 1067 if i != len(items) - 1: 1068 set_visible(ax, False, "xlabel") 1069 1070 return fig 1071 1072 def get_baserobot_code_cells(self, title=None): 1073 """ 1074 Return list of jupyter_ cells with calls to methods provided by the base class. 1075 """ 1076 # Try not pollute namespace with lots of variables. 1077 nbformat, nbv = self.get_nbformat_nbv() 1078 title = "## Code to compare multiple Structure objects" if title is None else str(title) 1079 return [ 1080 nbv.new_markdown_cell(title), 1081 nbv.new_code_cell("robot.get_lattice_dataframe()"), 1082 nbv.new_code_cell("""# robot.plot_lattice_convergence(sortby="nkpt", hue="tsmear")"""), 1083 nbv.new_code_cell("#robot.get_coords_dataframe()"), 1084 ] 1085 1086 1087class HueGroup(object): 1088 """ 1089 This small object is used by ``group_and_sortby`` to store information abouth the group. 1090 """ 1091 1092 def __init__(self, hvalue, xvalues, abifiles, labels): 1093 """ 1094 Args: 1095 hvalue: Hue value. 1096 xvalues: abifiles are sorted by ``func_or_string`` and these are the values 1097 associated to ``abifiles``. 1098 abifiles: List of file with this hue value. 1099 labels: List of labels associated to ``abifiles``. 1100 """ 1101 self.hvalue = hvalue 1102 self.abifiles = abifiles 1103 self.labels = labels 1104 self.xvalues = xvalues 1105 assert len(abifiles) == len(labels) 1106 assert len(abifiles) == len(xvalues) 1107 1108 def __len__(self): 1109 return len(self.abifiles) 1110 1111 def __iter__(self): 1112 """Iterate over (label, abifile, xvalue).""" 1113 return zip(self.labels, self.abifiles, self.xvalues) 1114 1115 #@lazy_property 1116 #def pretty_hvalue(self): 1117 # """Return pretty string with hvalue.""" 1118 # if duck.is_intlike(self.hvalue): 1119 # return "%d" % self.havalue 1120 # else: 1121 # try: 1122 # return "%.3f" % self.hvalue 1123 # except: 1124 # return str(self.hvalue) 1125