Source code for abipy.abio.robots

# coding: utf-8
"""
This module defines the Robot BaseClass. Robots operates on multiple files and provide helper
functions to plot the data e.g. convergence studies and to build pandas dataframes from the output files.
"""
import sys
import os
import inspect
import itertools
import numpy as np

from collections import OrderedDict, deque
from functools import wraps
from monty.string import is_string, list_strings
from monty.termcolor import cprint
from abipy.core.mixins import NotebookWriter
from abipy.tools.numtools import sort_and_groupby
from abipy.tools import duck
from abipy.tools.plotting import (plot_xy_with_hue, add_fig_kwargs, get_ax_fig_plt, get_axarray_fig_plt,
    rotate_ticklabels, set_visible)


[docs]class Robot(NotebookWriter): """ This is the base class from which all Robot subclasses should derive. A Robot supports the `with` context manager: Usage example: .. code-block:: python with Robot([("label1", "file1"), (label2, "file2")]) as robot: # Do something with robot. files are automatically closed when we exit. for label, abifile in self.items(): print(label) """ # filepaths are relative to `start`. None for asbolute paths. This flag is set in trim_paths start = None # Used in iter_lineopt to generate matplotlib linestyles. _LINE_COLORS = ["b", "r", "g", "m", "y", "k", "c"] _LINE_STYLES = ["-", ":", "--", "-.",] _LINE_WIDTHS = [2, ] def __init__(self, *args): """ Args: args is a list of tuples (label, filepath) """ self._abifiles, self._do_close = OrderedDict(), OrderedDict() self._exceptions = deque(maxlen=100) for label, abifile in args: self.add_file(label, abifile)
[docs] @classmethod def get_supported_extensions(self): """List of strings with extensions supported by Robot subclasses.""" # This is needed to have all subclasses. from abipy.abilab import Robot return sorted([cls.EXT for cls in Robot.__subclasses__()])
[docs] @classmethod def class_for_ext(cls, ext): """Return the Robot subclass associated to the given extension.""" for subcls in cls.__subclasses__(): if subcls.EXT in (ext, ext.upper()): return subcls # anaddb.nc does not follow the extension rule... if ext.lower() == "anaddb": from abipy.dfpt.anaddbnc import AnaddbNcRobot as subcls return subcls raise ValueError("Cannot find Robot subclass associated to extension %s\n" % ext + "The list of supported extensions (case insensitive) is:\n%s" % str(cls.get_supported_extensions()))
[docs] @classmethod def from_dir(cls, top, walk=True, abspath=False): """ This class method builds a robot by scanning all files located within directory `top`. This method should be invoked with a concrete robot class, for example: robot = GsrRobot.from_dir(".") Args: top (str): Root directory walk: if True, directories inside `top` are included as well. abspath: True if paths in index should be absolute. Default: Relative to `top`. """ new = cls(*cls._open_files_in_dir(top, walk)) if not abspath: new.trim_paths(start=top) return new
[docs] @classmethod def from_dirs(cls, dirpaths, walk=True, abspath=False): """ Similar to `from_dir` but accepts a list of directories instead of a single directory. Args: walk: if True, directories inside `top` are included as well. abspath: True if paths in index should be absolute. Default: Relative to `top`. """ items = [] for top in list_strings(dirpaths): items.extend(cls._open_files_in_dir(top, walk)) new = cls(*items) if not abspath: new.trim_paths(start=os.getcwd()) return new
[docs] @classmethod def from_dir_glob(cls, pattern, walk=True, abspath=False): """ This class method builds a robot by scanning all files located within the directories matching `pattern` as implemented by glob.glob This method should be invoked with a concrete robot class, for example: robot = GsrRobot.from_dir_glob("flow_dir/w*/outdata/") Args: pattern: Pattern string walk: if True, directories inside `top` are included as well. abspath: True if paths in index should be absolute. Default: Relative to getcwd(). """ import glob items = [] for top in filter(os.path.isdir, glob.iglob(pattern)): items += cls._open_files_in_dir(top, walk=walk) new = cls(*items) if not abspath: new.trim_paths(start=os.getcwd()) return new
@classmethod def _open_files_in_dir(cls, top, walk): """Open files in directory tree starting from `top`. Return list of Abinit files.""" if not os.path.isdir(top): raise ValueError("%s: no such directory" % str(top)) from abipy.abilab import abiopen items = [] if walk: for dirpath, dirnames, filenames in os.walk(top): filenames = sorted([f for f in filenames if cls.class_handles_filename(f)]) for f in filenames: abifile = abiopen(os.path.join(dirpath, f)) if abifile is not None: items.append((abifile.filepath, abifile)) else: filenames = [f for f in os.listdir(top) if cls.class_handles_filename(f)] for f in filenames: abifile = abiopen(os.path.join(top, f)) if abifile is not None: items.append((abifile.filepath, abifile)) return items
[docs] @classmethod def class_handles_filename(cls, filename): """True if robot class handles filename.""" # Special treatment of AnaddbNcRobot if cls.EXT == "anaddb" and os.path.basename(filename).lower() == "anaddb.nc": return True return (filename.endswith("_" + cls.EXT + ".nc") or filename.endswith("." + cls.EXT)) # This for .abo
[docs] @classmethod def from_files(cls, filenames, labels=None, abspath=False): """ Build a Robot from a list of `filenames`. if labels is None, labels are automatically generated from absolute paths. Args: abspath: True if paths in index should be absolute. Default: Relative to `top`. """ filenames = list_strings(filenames) from abipy.abilab import abiopen filenames = [f for f in filenames if cls.class_handles_filename(f)] items = [] for i, f in enumerate(filenames): try: abifile = abiopen(f) except Exception as exc: cprint("Exception while opening file: `%s`" % str(f), "red") cprint(exc, "red") abifile = None if abifile is not None: label = abifile.filepath if labels is None else labels[i] items.append((label, abifile)) new = cls(*items) if labels is None and not abspath: new.trim_paths(start=None) return new
[docs] @classmethod def from_flow(cls, flow, outdirs="all", nids=None, ext=None, task_class=None): """ Build a robot from a |Flow| object. Args: flow: |Flow| object outdirs: String used to select/ignore the files in the output directory of flow, works and tasks outdirs="work" selects only the outdir of the Works, outdirs="flow+task" selects the outdir of the Flow and the outdirs of the tasks outdirs="-work" excludes the outdir of the Works. Cannot use ``+`` and ``-`` flags in the same string. Default: `all` that is equivalent to "flow+work+task" nids: List of node identifiers used to select particular nodes. Not used if None ext: File extension associated to the robot. Mainly used if method is invoked with the BaseClass task_class: Task class or string with the class name used to select the tasks in the flow. None implies no filtering. Usage example: .. code-block:: python with abilab.GsrRobot.from_flow(flow) as robot: print(robot) # That is equivalent to: with Robot.from_flow(flow, ext="GSR") as robot: print(robot) Returns: ``Robot`` subclass. """ robot = cls() if ext is None else cls.class_for_ext(ext)() all_opts = ("flow", "work", "task") if outdirs == "all": tokens = all_opts elif "+" in outdirs: assert "-" not in outdirs tokens = outdirs.split("+") elif "-" in outdirs: assert "+" not in outdirs tokens = [s for s in all if s not in outdirs.split("-")] else: tokens = list_strings(outdirs) if not all(t in all_opts for t in tokens): raise ValueError("Wrong outdirs string %s" % outdirs) if "flow" in tokens: robot.add_extfile_of_node(flow, nids=nids, task_class=task_class) if "work" in tokens: for work in flow: robot.add_extfile_of_node(work, nids=nids, task_class=task_class) if "task" in tokens: for task in flow.iflat_tasks(): robot.add_extfile_of_node(task, nids=nids, task_class=task_class) return robot
[docs] def add_extfile_of_node(self, node, nids=None, task_class=None): """ Add the file produced by this node to the robot. Args: node: |Flow| or |Work| or |Task| object. nids: List of node identifiers used to select particular nodes. Not used if None task_class: Task class or string with class name used to select the tasks in the flow. None implies no filtering. """ if nids and node.node_id not in nids: return filepath = node.outdir.has_abiext(self.EXT) if not filepath: # Look in run.abi directory. filepath = node.wdir.has_abiext(self.EXT) # This to ignore DDB.nc files (only text DDB are supported) if filepath and filepath.endswith("_DDB.nc"): return if filepath: try: label = os.path.relpath(filepath) except OSError: # current working directory may not be defined! label = filepath # Filter by task_class (class or string with class name) if task_class is not None and not node.isinstance(task_class): return None self.add_file(label, filepath)
[docs] def scan_dir(self, top, walk=True): """ Scan directory tree starting from ``top``. Add files to the robot instance. Args: top (str): Root directory walk: if True, directories inside ``top`` are included as well. Return: Number of files found. """ count = 0 for filepath, abifile in self.__class__._open_files_in_dir(top, walk): count += 1 self.add_file(filepath, abifile) return count
[docs] def add_file(self, label, abifile, filter_abifile=None): """ Add a file to the robot with the given label. Args: label: String used to identify the file (must be unique, ax exceptions is raised if label is already present. abifile: Specify the file to be added. Accepts strings (filepath) or abipy file-like objects. filter_abifile: Function that receives an ``abifile`` object and returns True if the file should be added to the plotter. """ if is_string(abifile): from abipy.abilab import abiopen abifile = abiopen(abifile) if filter_abifile is not None and not filter_abifile(abifile): abifile.close() return # Open file here --> have to close it. self._do_close[abifile.filepath] = True if label in self._abifiles: raise ValueError("label %s is already present!" % label) self._abifiles[label] = abifile
#def pop_filepath(self, filepath): # """ # Remove the file with the given `filepath` and close it. # """ # if label, abifile in self._abifiles.items(): # if abifile.filepath != filepath: continue # self._abifiles.pop(label) # if self._do_close.pop(abifile.filepath, False): # try: # abifile.close() # except Exception as exc: # print("Exception while closing: ", abifile.filepath) # print(exc)
[docs] def iter_lineopt(self): """Generates matplotlib linestyles.""" for o in itertools.product( self._LINE_WIDTHS, self._LINE_STYLES, self._LINE_COLORS): yield {"linewidth": o[0], "linestyle": o[1], "color": o[2]}
[docs] @staticmethod def ordered_intersection(list_1, list_2): """Return ordered intersection of two lists. Items must be hashable.""" set_2 = frozenset(list_2) return [x for x in list_1 if x in set_2]
#def _get_ointersection_i(self, iattrname): # if len(self.abifiles) == 0: return [] # values = list(range(getattr(self.abifiles[0], iattrname))) # if len(self.abifiles) == 1: return values # for abifile in self.abifiles[1:]: # values = self.ordered_intersection(values, range(getattr(abifile, iattrname))) # return values @staticmethod def _to_relpaths(paths): """Convert a list of absolute paths to relative paths.""" root = os.getcwd() return [os.path.relpath(p, root) for p in paths]
[docs] def remove(self): """Close the file handle, remove the file from disk for each file in the robot.""" for abifile in self.abifiles: abifile.remove()
[docs] def pop_label(self, label): """ Remove file with the given ``label`` and close it. """ if label in self._abifiles: abifile = self._abifiles.pop(label) if self._do_close.pop(abifile.filepath, False): try: abifile.close() except Exception as exc: print("Exception while closing: ", abifile.filepath) print(exc)
[docs] def change_labels(self, new_labels, dryrun=False): """ Change labels of the files. Args: new_labels: List of strings (same length as self.abifiles) dryrun: True to activate dryrun mode. Return: mapping new_label --> old_label. """ if len(new_labels) != len(self): raise ValueError("Robot has %d files while len(new_labels) = %d" % (len(new_labels), len(self))) old_labels = list(self._abifiles.keys()) if not dryrun: old_abifiles, self._abifiles = self._abifiles, OrderedDict() new2old = OrderedDict() for old, new in zip(old_labels, new_labels): new2old[new] = old if not dryrun: self._abifiles[new] = old_abifiles[old] else: print("old [%s] --> new [%s]" % (old, new)) return new2old
[docs] def remap_labels(self, function, dryrun=False): """ Change labels of the files by executing ``function`` Args: function: Callable object e.g. lambda function. The output of function(abifile) is used as new label. Note that the function shall not return duplicated labels when applied to self.abifiles. dryrun: True to activate dryrun mode. Return: mapping new_label --> old_label. """ new_labels = [function(afile) for afile in self.abifiles] # Labels must be unique and hashable. if len(set(new_labels)) != len(new_labels): raise ValueError("Duplicated labels are not allowed. Change input function.\nnew_labels %s" % str(new_labels)) return self.change_labels(new_labels, dryrun=dryrun)
[docs] def trim_paths(self, start=None): """ Replace absolute filepaths in the robot with relative paths wrt to ``start`` directory. If start is None, os.getcwd() is used. Set ``self.start`` attribute, return ``self.start``. """ self.start = os.getcwd() if start is None else start old_paths = list(self._abifiles.keys()) old_new_paths = [(p, os.path.relpath(os.path.abspath(p), start=self.start)) for p in old_paths] old_abifiles = self._abifiles self._abifiles = OrderedDict() for old, new in old_new_paths: self._abifiles[new] = old_abifiles[old] return self.start
@property def exceptions(self): """List of exceptions.""" return self._exceptions def __len__(self): return len(self._abifiles) #def __iter__(self): # return iter(self._abifiles) #def __contains__(self, item): # return item in self._abifiles def __getitem__(self, key): # self[key] return self._abifiles.__getitem__(key) def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): """Activated at the end of the with statement.""" self.close()
[docs] def keys(self): return self._abifiles.keys()
[docs] def items(self): return self._abifiles.items()
@property def labels(self): """ List of strings used to create labels in matplotlib figures when plotting results taked from multiple files. By default, labels is initialized with the path of the files in the robot. Use change_labels to change the list. """ return list(self._abifiles.keys())
[docs] def get_label_files_str(self): """Return string with [label, filepath].""" from tabulate import tabulate return tabulate([(label, abifile.relpath) for label, abifile in self.items()], headers=["Label", "Relpath"]) + "\n"
[docs] def show_files(self, stream=sys.stdout): """Show label --> file path""" stream.write(self.get_label_files_str())
def __repr__(self): """Invoked by repr.""" return self.get_label_files_str() def __str__(self): """Invoked by str.""" return self.to_string()
[docs] def to_string(self, verbose=0): """String representation.""" lines = ["%s with %d files in memory:\n" % (self.__class__.__name__, len(self.abifiles))] app = lines.append for i, f in enumerate(self.abifiles): app(f.to_string(verbose=verbose)) app("\n") return "\n".join(lines)
def _repr_html_(self): """Integration with jupyter_ notebooks.""" return '<ol start="0">\n{}\n</ol>'.format("\n".join("<li>%s</li>" % label for label, abifile in self.items())) @property def abifiles(self): """List of netcdf files.""" return list(self._abifiles.values())
[docs] def has_different_structures(self, rtol=1e-05, atol=1e-08): """ Check if structures are equivalent, return string with info about differences (if any). """ if len(self) <= 1: return "" formulas = set([af.structure.composition.formula for af in self.abifiles]) if len(formulas) != 1: return "Found structures with different full formulas: %s" % str(formulas) lines = [] s0 = self.abifiles[0].structure for abifile in self.abifiles[1:]: s1 = abifile.structure if not np.allclose(s0.lattice.matrix, s1.lattice.matrix, rtol=rtol, atol=atol): lines.append("Structures have different lattice:") if not np.allclose(s0.frac_coords, s1.frac_coords, rtol=rtol, atol=atol): lines.append("Structures have different atomic positions:") return "\n".join(lines)
#def apply(self, func_or_string, args=(), **kwargs): # """ # Applies function to all ``abifiles`` available in the robot. # Args: # func_or_string: If callable, the output of func_or_string(abifile, ...) is used. # If string, the output of getattr(abifile, func_or_string)(...) # args (tuple): Positional arguments to pass to function in addition to the array/series # kwargs: Additional keyword arguments will be passed as keywords to the function # Return: List of results # """ # if callable(func_or_string): # return [func_or_string(abifile, *args, *kwargs) for abifile in self.abifiles] # else: # return [duck.getattrd(abifile, func_or_string)(*args, **kwargs) for abifile in self.abifiles]
[docs] def is_sortable(self, aname, raise_exc=False): """ Return True if ``aname`` is an attribute of the netcdf file If raise_exc is True, AttributeError with an explicit message is raised. """ try: obj = None try: # abiifile.foo.bar? obj = duck.getattrd(self.abifiles[0], aname) except AttributeError: # abifile.params[aname] ? if hasattr(self.abifiles[0], "params") and aname in self.abifiles[0].params: obj = self.abifiles[0].params[aname] # Let's try to convert obj to scalar. float(obj) return True except Exception: if not raise_exc: return False attrs = [] for key, obj in inspect.getmembers(self.abifiles[0]): # Ignores anything starting with underscore if key.startswith('_') or callable(obj) or hasattr(obj, "__len__"): continue attrs.append(key) # Add entries in params. if hasattr(self.abifiles[0], "params") and hasattr(self.abifiles[0].params, "keys"): attrs.extend(self.abifiles[0].params.keys()) raise AttributeError("""\ `%s` object has no attribute `%s`. Choose among: %s Note that this list is automatically generated. Not all entries are sortable (Please select number-like quantities)""" % (self.__class__.__name__, aname, str(attrs)))
def _sortby_labelfile_list(self, labelfile_list, func_or_string, reverse=False, unpack=False): """ Return: list of (label, abifile, param) tuples where param is obtained via ``func_or_string``. or labels, abifiles, params if ``unpack`` """ if not func_or_string: # Catch None or empty items = [(label, abifile, label) for (label, abifile) in labelfile_list] if not unpack: return items else: return [t[0] for t in items], [t[1] for t in items], [t[2] for t in items] elif callable(func_or_string): items = [(label, abifile, func_or_string(abifile)) for (label, abifile) in labelfile_list] else: # Assume string and attribute with the same name. # try in abifile.params if not hasattrd(abifile, func_or_string) self.is_sortable(func_or_string, raise_exc=True) if duck.hasattrd(self.abifiles[0], func_or_string): items = [(label, abifile, duck.getattrd(abifile, func_or_string)) for (label, abifile) in labelfile_list] else: items = [(label, abifile, abifile.params[func_or_string]) for (label, abifile) in labelfile_list] items = sorted(items, key=lambda t: t[2], reverse=reverse) if not unpack: return items else: return [t[0] for t in items], [t[1] for t in items], [t[2] for t in items]
[docs] def sortby(self, func_or_string, reverse=False, unpack=False): """ Sort files in the robot by ``func_or_string``. Args: func_or_string: Either None, string, callable defining the quantity to be used for sorting. If string, it's assumed that the abifile has an attribute with the same name and getattr is invoked. If callable, the output of func_or_string(abifile) is used. If None, no sorting is performed. reverse: If set to True, then the list elements are sorted as if each comparison were reversed. unpack: Return (labels, abifiles, params) if True Return: list of (label, abifile, param) tuples where param is obtained via ``func_or_string``. or labels, abifiles, params if ``unpack`` """ labelfile_list = list(self.items()) return self._sortby_labelfile_list(labelfile_list, func_or_string, reverse=reverse, unpack=unpack)
[docs] def group_and_sortby(self, hue, func_or_string): """ Group files by ``hue`` and, inside each group` sort items by ``func_or_string``. Args: hue: Variable that define subsets of the data, which will be drawn on separate lines. Accepts callable or string If string, it's assumed that the abifile has an attribute with the same name and getattr is invoked. Dot notation is also supported e.g. hue="structure.formula" --> abifile.structure.formula If callable, the output of hue(abifile) is used. func_or_string: Either None, string, callable defining the quantity to be used for sorting. If string, it's assumed that the abifile has an attribute with the same name and getattr is invoked. If callable, the output of func_or_string(abifile) is used. If None, no sorting is performed. Return: List of :class:`HueGroup` instance. """ # Group by hue. # This is the section in which we support: callable, abifile.attr.name syntax or abifile.params["key"] items = list(self.items()) if callable(hue): key = lambda t: hue(t[1]) else: # Assume string. if duck.hasattrd(self.abifiles[0], hue): key = lambda t: duck.getattrd(t[1], hue) else: # Try in abifile.params if hasattr(self.abifiles[0], "params") and hue in self.abifiles[0].params: key = lambda t: t[1].params[hue] else: raise TypeError("""\ Cannot interpret hue argument of type `%s` and value `%s`. Expecting callable or attribute name or key in abifile.params""" % (type(hue), str(hue))) groups = [] for hvalue, labelfile_list in sort_and_groupby(items, key=key): # Use func_or_string to sort each group labels, abifiles, xvalues = self._sortby_labelfile_list(labelfile_list, func_or_string, unpack=True) groups.append(HueGroup(hvalue, xvalues, abifiles, labels)) return groups
[docs] def close(self): """ Close all files that have been opened by the Robot. """ for abifile in self.abifiles: if self._do_close.pop(abifile.filepath, False): try: abifile.close() except Exception as exc: print("Exception while closing: ", abifile.filepath) print(exc)
#@classmethod #def open(cls, obj, nids=None, **kwargs): # """ # Flexible constructor. obj can be a :class:`Flow` or a string with the directory containing the Flow. # `nids` is an optional list of :class:`Node` identifiers used to filter the set of :class:`Task` in the Flow. # """ # has_dirpath = False # if is_string(obj): # try: # from abipy.flowtk import Flow # obj = Flow.pickle_load(obj) # except: # has_dirpath = True # if not has_dirpath: # # We have a Flow. smeth is the name of the Task method used to open the file. # items = [] # smeth = "open_" + cls.EXT.lower() # for task in obj.iflat_tasks(nids=nids): #, status=obj.S_OK): # open_method = getattr(task, smeth, None) # if open_method is None: continue # abifile = open_method() # if abifile is not None: items.append((task.pos_str, abifile)) # return cls(*items) # else: # # directory --> search for files with the appropriate extension and open it with abiopen. # if nids is not None: raise ValueError("nids cannot be used when obj is a directory.") # return cls.from_dir(obj) #def get_attributes(self, attr_name, obj=None, retdict=False): # od = OrderedDict() # for label, abifile in self.items(): # obj = abifile if obj is None else getattr(abifile, obj) # od[label] = getattr(obj, attr_name) # if retdict: # return od # else: # return list(od.values()) def _exec_funcs(self, funcs, arg): """ Execute list of callable functions. Each function receives arg as argument. """ if not isinstance(funcs, (list, tuple)): funcs = [funcs] d = {} for func in funcs: try: key, value = func(arg) d[key] = value except Exception as exc: cprint("Exception: %s" % str(exc), "red") self._exceptions.append(str(exc)) return d
[docs] @staticmethod def sortby_label(sortby, param): """Return the label to be used when files are sorted with ``sortby``.""" return "%s %s" % (sortby, param) if not (callable(sortby) or sortby is None) else str(param)
[docs] def get_structure_dataframes(self, abspath=False, filter_abifile=None, **kwargs): """ Wrap dataframes_from_structures function. Args: abspath: True if paths in index should be absolute. Default: Relative to getcwd(). filter_abifile: Function that receives an ``abifile`` object and returns True if the file should be added to the plotter. """ from abipy.core.structure import dataframes_from_structures if "index" not in kwargs: index = list(self._abifiles.keys()) if not abspath: index = self._to_relpaths(index) kwargs["index"] = index abifiles = self.abifiles if filter_abifile is not None else list(filter(filter_abifile, self.abifiles)) return dataframes_from_structures(struct_objects=abifiles, **kwargs)
[docs] def get_lattice_dataframe(self, **kwargs): """Return |pandas-DataFrame| with lattice parameters.""" dfs = self.get_structure_dataframes(**kwargs) return dfs.lattice
[docs] def get_coords_dataframe(self, **kwargs): """Return |pandas-DataFrame| with atomic positions.""" dfs = self.get_structure_dataframes(**kwargs) return dfs.coords
[docs] def get_params_dataframe(self, abspath=False): """ Return |pandas-DataFrame| with the most important parameters. that are usually subject to convergence studies. Args: abspath: True if paths in index should be absolute. Default: Relative to `top`. """ rows, row_names = [], [] for label, abifile in self.items(): if not hasattr(abifile, "params"): import warnings warnings.warn("%s does not have `params` attribute" % type(abifile)) break rows.append(abifile.params) row_names.append(label) row_names = row_names if abspath else self._to_relpaths(row_names) import pandas as pd return pd.DataFrame(rows, index=row_names, columns=list(rows[0].keys()))
############################################## # Helper functions to plot pandas dataframes # ############################################## @staticmethod @wraps(plot_xy_with_hue) def plot_xy_with_hue(*args, **kwargs): return plot_xy_with_hue(*args, **kwargs) @staticmethod def _get_label(func_or_string): """ Return label associated to ``func_or_string``. If callable, docstring __doc__ is used. """ if func_or_string is None: return "" elif callable(func_or_string): if getattr(func_or_string, "__doc__", ""): return func_or_string.__doc__.strip() else: return func_or_string.__name__ else: return str(func_or_string)
[docs] @add_fig_kwargs def plot_convergence(self, item, sortby=None, hue=None, ax=None, fontsize=8, **kwargs): """ Plot the convergence of ``item`` wrt the ``sortby`` parameter. Values can optionally be grouped by ``hue``. Args: item: Define the quantity to plot. Accepts callable or string If string, it's assumed that the abifile has an attribute with the same name and `getattr` is invoked. Dot notation is also supported e.g. hue="structure.formula" --> abifile.structure.formula If callable, the output of item(abifile) is used. sortby: Define the convergence parameter, sort files and produce plot labels. Can be None, string or function. If None, no sorting is performed. If string and not empty it's assumed that the abifile has an attribute with the same name and `getattr` is invoked. If callable, the output of sortby(abifile) is used. hue: Variable that define subsets of the data, which will be drawn on separate lines. Accepts callable or string If string, it's assumed that the abifile has an attribute with the same name and getattr is invoked. If callable, the output of hue(abifile) is used. ax: |matplotlib-Axes| or None if a new figure should be created. fontsize: legend and label fontsize. kwargs: keyword arguments passed to matplotlib plot method. Returns: |matplotlib-Figure| Example: robot.plot_convergence("energy") robot.plot_convergence("energy", sortby="nkpt") robot.plot_convergence("pressure", sortby="nkpt", hue="tsmear") """ ax, fig, plt = get_ax_fig_plt(ax=ax) if "marker" not in kwargs: kwargs["marker"] = "o" def get_yvalues(abifiles): if callable(item): return [float(item(a)) for a in abifiles] else: return [float(getattr(a, item)) for a in abifiles] if hue is None: labels, abifiles, params = self.sortby(sortby, unpack=True) yvals = get_yvalues(abifiles) #print("params", params, "\nyvals", yvals) ax.plot(params, yvals, **kwargs) else: groups = self.group_and_sortby(hue, sortby) for g in groups: yvals = get_yvalues(g.abifiles) label = "%s: %s" % (self._get_label(hue), g.hvalue) ax.plot(g.xvalues, yvals, label=label, **kwargs) ax.grid(True) ax.set_xlabel("%s" % self._get_label(sortby)) if sortby is None: rotate_ticklabels(ax, 15) ax.set_ylabel("%s" % self._get_label(item)) if hue is not None: ax.legend(loc="best", fontsize=fontsize, shadow=True) return fig
[docs] @add_fig_kwargs def plot_convergence_items(self, items, sortby=None, hue=None, fontsize=6, **kwargs): """ Plot the convergence of a list of ``items`` wrt to the ``sortby`` parameter. Values can optionally be grouped by ``hue``. Args: items: List of attributes (or callables) to be analyzed. sortby: Define the convergence parameter, sort files and produce plot labels. Can be None, string or function. If None, no sorting is performed. If string and not empty it's assumed that the abifile has an attribute with the same name and `getattr` is invoked. If callable, the output of sortby(abifile) is used. hue: Variable that define subsets of the data, which will be drawn on separate lines. Accepts callable or string If string, it's assumed that the abifile has an attribute with the same name and getattr is invoked. Dot notation is also supported e.g. hue="structure.formula" --> abifile.structure.formula If callable, the output of hue(abifile) is used. fontsize: legend and label fontsize. kwargs: keyword arguments are passed to ax.plot Returns: |matplotlib-Figure| """ # Note: in principle one could call plot_convergence inside a loop but # this one is faster as sorting is done only once. # Build grid plot. nrows, ncols = len(items), 1 ax_list, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols, sharex=True, sharey=False, squeeze=False) ax_list = ax_list.ravel() # Sort and group files if hue. if hue is None: labels, ncfiles, params = self.sortby(sortby, unpack=True) else: groups = self.group_and_sortby(hue, sortby) marker = kwargs.pop("marker", "o") for i, (ax, item) in enumerate(zip(ax_list, items)): if hue is None: # Extract data. if callable(item): yvals = [float(item(gsr)) for gsr in self.abifiles] else: yvals = [duck.getattrd(gsr, item) for gsr in self.abifiles] if not is_string(params[0]): ax.plot(params, yvals, marker=marker, **kwargs) else: # Must handle list of strings in a different way. xn = range(len(params)) ax.plot(xn, yvals, marker=marker, **kwargs) ax.set_xticks(xn) ax.set_xticklabels(params, fontsize=fontsize) else: for g in groups: # Extract data. if callable(item): yvals = [float(item(gsr)) for gsr in g.abifiles] else: yvals = [duck.getattrd(gsr, item) for gsr in g.abifiles] label = "%s: %s" % (self._get_label(hue), g.hvalue) ax.plot(g.xvalues, yvals, label=label, marker=marker, **kwargs) ax.grid(True) ax.set_ylabel(self._get_label(item)) if i == len(items) - 1: ax.set_xlabel("%s" % self._get_label(sortby)) if sortby is None: rotate_ticklabels(ax, 15) if i == 0 and hue is not None: ax.legend(loc="best", fontsize=fontsize, shadow=True) return fig
[docs] @add_fig_kwargs def plot_lattice_convergence(self, what_list=None, sortby=None, hue=None, fontsize=8, **kwargs): """ Plot the convergence of the lattice parameters (a, b, c, alpha, beta, gamma). wrt the``sortby`` parameter. Values can optionally be grouped by ``hue``. Args: what_list: List of strings with the quantities to plot e.g. ["a", "alpha", "beta"]. None means all. item: Define the quantity to plot. Accepts callable or string If string, it's assumed that the abifile has an attribute with the same name and `getattr` is invoked. If callable, the output of item(abifile) is used. sortby: Define the convergence parameter, sort files and produce plot labels. Can be None, string or function. If None, no sorting is performed. If string and not empty it's assumed that the abifile has an attribute with the same name and `getattr` is invoked. If callable, the output of sortby(abifile) is used. hue: Variable that define subsets of the data, which will be drawn on separate lines. Accepts callable or string If string, it's assumed that the abifile has an attribute with the same name and getattr is invoked. Dot notation is also supported e.g. hue="structure.formula" --> abifile.structure.formula If callable, the output of hue(abifile) is used. ax: |matplotlib-Axes| or None if a new figure should be created. fontsize: legend and label fontsize. Returns: |matplotlib-Figure| Example: robot.plot_lattice_convergence() robot.plot_lattice_convergence(sortby="nkpt") robot.plot_lattice_convergence(sortby="nkpt", hue="tsmear") """ if not self.abifiles: return None # The majority of AbiPy files have a structure object # whereas Hist.nc defines final_structure. Use geattr and key to extract structure object. key = "structure" if not hasattr(self.abifiles[0], "structure"): if hasattr(self.abifiles[0], "final_structure"): key = "final_structure" else: raise TypeError("Don't know how to extract structure from %s" % type(self.abifiles[0])) # Define callbacks. docstrings will be used as ylabels. def a(afile): "a (Ang)" return getattr(afile, key).lattice.a def b(afile): "b (Ang)" return getattr(afile, key).lattice.b def c(afile): "c (Ang)" return getattr(afile, key).lattice.c def volume(afile): r"$V$" return getattr(afile, key).lattice.volume def alpha(afile): r"$\alpha$" return getattr(afile, key).lattice.alpha def beta(afile): r"$\beta$" return getattr(afile, key).lattice.beta def gamma(afile): r"$\gamma$" return getattr(afile, key).lattice.gamma items = [a, b, c, volume, alpha, beta, gamma] if what_list is not None: locs = locals() items = [locs[what] for what in list_strings(what_list)] # Build plot grid. nrows, ncols = len(items), 1 ax_list, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols, sharex=True, sharey=False, squeeze=False) marker = kwargs.pop("marker", "o") for i, (ax, item) in enumerate(zip(ax_list.ravel(), items)): self.plot_convergence(item, sortby=sortby, hue=hue, ax=ax, fontsize=fontsize, marker=marker, show=False) if i != 0: set_visible(ax, False, "legend") if i != len(items) - 1: set_visible(ax, False, "xlabel") return fig
[docs] def get_baserobot_code_cells(self, title=None): """ Return list of jupyter_ cells with calls to methods provided by the base class. """ # Try not pollute namespace with lots of variables. nbformat, nbv = self.get_nbformat_nbv() title = "## Code to compare multiple Structure objects" if title is None else str(title) return [ nbv.new_markdown_cell(title), nbv.new_code_cell("robot.get_lattice_dataframe()"), nbv.new_code_cell("""# robot.plot_lattice_convergence(sortby="nkpt", hue="tsmear")"""), nbv.new_code_cell("#robot.get_coords_dataframe()"), ]
[docs]class HueGroup(object): """ This small object is used by ``group_and_sortby`` to store information abouth the group. """ def __init__(self, hvalue, xvalues, abifiles, labels): """ Args: hvalue: Hue value. xvalues: abifiles are sorted by ``func_or_string`` and these are the values associated to ``abifiles``. abifiles: List of file with this hue value. labels: List of labels associated to ``abifiles``. """ self.hvalue = hvalue self.abifiles = abifiles self.labels = labels self.xvalues = xvalues assert len(abifiles) == len(labels) assert len(abifiles) == len(xvalues) def __len__(self): return len(self.abifiles) def __iter__(self): """Iterate over (label, abifile, xvalue).""" return zip(self.labels, self.abifiles, self.xvalues)
#@lazy_property #def pretty_hvalue(self): # """Return pretty string with hvalue.""" # if duck.is_intlike(self.hvalue): # return "%d" % self.havalue # else: # try: # return "%.3f" % self.hvalue # except: # return str(self.hvalue)