Source code for autodepgraph.graph

import logging
import numpy as np
import types
from typing import Dict, Any, List, Optional
from datetime import datetime
import matplotlib.pyplot as plt
from os.path import join, split
import os
import tempfile
import webbrowser
import warnings

import networkx as nx
import autodepgraph
from autodepgraph.visualization import state_cmap
from autodepgraph import visualization as vis

# Used to find functions in modules
from importlib import import_module
try:
    # Only used for finding instrument methods.
    from qcodes.instrument.base import Instrument
except ImportError:
    Instrument = None


[docs]class AutoDepGraph_DAG(nx.DiGraph): """ Attributes: node_states: Allowed states for the nodes matplotlib_edge_properties: Properties passed to networkx plotting of edges matplotlib_label_properties: Properties passed to networkx plotting of labels """ node_states: List[str] = ['good', 'needs calibration', 'bad', 'unknown', 'active'] matplotlib_edge_properties: Dict[str, Any] = {'edge_color': 'k', 'alpha': .8} matplotlib_label_properties: Dict[str, Any] = {'font_color': 'k'} def __init__(self, name, cfg_plot_mode='svg', incoming_graph_data=None, **attr): """ Directed Acyclic Graph used for calibrations. Inherits from a networkx DiGraph. The AutoDepGraph provides several constraints to nodes and edges as well as add the functionality to walk over the graph in order to calibrate a node and visualize the graph in real-time. """ attr['name'] = name self.cfg_plot_mode = cfg_plot_mode self.cfg_plot_mode_args = {'fig': None} super().__init__(incoming_graph_data, **attr) # internal attributes self._DiGraphWindow = None # used for pyqtgraph plotting self._graph_changed_since_plot = True # used for pyqtgraph plotting # counters to count how often functions get called for debugging # purposes self._exec_cnt = 0 self._calib_cnt = 0 self._check_cnt = 0 @property def cfg_svg_filename(self): """ Default location for storing svg based visualizations of the DAG. """ _path_name = split(__file__)[:-1][0] return join(_path_name, 'svg_viewer', 'adg_graph.svg') def fresh_copy(self): return AutoDepGraph_DAG(name=self.name, cfg_plot_mode=self.cfg_plot_mode)
[docs] def add_node(self, node_for_adding, **attr): """ Adds a node to the graph, including starting attributes. attr: name (type) = default_value calibrate_function = 'NotImplementedCalibration' check_functions = 'return_fixed_value' tolerance (float) = 0 timeout (float) = np.inf state (str) = 'unknown' """ # there are set here to ensure these node attributes exist. # setting in this way will most likely interfere # with joining multiple graphs attr['timeout'] = attr.get('timeout', np.inf) attr['calibrate_function'] = attr.get( 'calibrate_function', 'autodepgraph.node_functions.calibration_functions' + '.NotImplementedCalibration') attr['check_function'] = attr.get( 'check_function', 'autodepgraph.node_functions.check_functions' + '.return_fixed_value') # zero default tolerance -> always recalibrate attr['tolerance'] = attr.get('tolerance', 0) super().add_node(node_for_adding, **attr) self.set_node_state(node_for_adding, state=attr.get('state', 'unknown')) # this adds a helper method to the attributes for quick access self._construct_maintenance_methods(nodes=[node_for_adding])
def _construct_maintenance_methods(self, nodes): for n in nodes: self._construct_maintenance_method(node_name=n) def _construct_maintenance_method(self, node_name): node_name_no_space = node_name.replace(' ', '_').replace('-', '_') # This name exists so that text based storing falls back # on the right placeholder function def _construct_maintenance_method(): self.maintain_node(node_name) self.__setattr__('maintain_{}'.format(node_name_no_space), _construct_maintenance_method)
[docs] def add_edge(self, u_of_edge, v_of_edge, **attr): """ Adds an edge that denotes a dependency in the calibration graph. u_of_edge -> v_of_edge denotes that u depends on v. """ # Nodes must already exist to ensure they have the right properties if u_of_edge not in self.nodes(): raise KeyError('{} not in nodes'.format(u_of_edge)) if v_of_edge not in self.nodes(): raise KeyError('{} not in nodes'.format(v_of_edge)) super().add_edge(u_of_edge, v_of_edge, **attr)
def get_node_state(self, node_name): Delta_T = (datetime.now() - self.nodes[node_name]['last_update']).total_seconds() if (Delta_T > self.nodes[node_name]['timeout']): self.nodes[node_name]['state'] = 'unknown' return self.nodes[node_name]['state'] def set_node_state(self, node_name, state, update_monitor=True): if state not in self.node_states: raise IndexError(f'state {state} not in {self.node_states}') self.nodes[node_name]['state'] = state self.nodes[node_name]['last_update'] = datetime.now() if update_monitor: self.update_monitor() def is_manual_node(self, node_name): if isinstance(self.nodes[node_name]['calibrate_function'], (types.MethodType, types.FunctionType)): return False elif 'manual' in self.nodes[node_name]['calibrate_function']: return True elif 'NotImplemented' in self.nodes[node_name]['calibrate_function']: return True else: return False
[docs] def maintain_node(self, node: str, verbose=True) -> str: """ Maintaining a node attempts to go from any state to a good state. any_state -> good Maintain a node performs the following steps: 1. get the state of the dependency nodes. If all are OK or unknown, perform a check on the node itself. If a node is any other state, execute it to move it to a good state 2. perform the "check" experiment on the node itself. This quick check 3. Perform calibration and second round of maintaining dependencies Returns: State of the node after maintaining the node Raises: Exception if the node could not be calibrated """ self._exec_cnt += 1 if verbose: print('Maintaining node "{}".'.format(node)) # 1. Going over the states of all the required nodes and ensure # these are all in a 'Good' state. for req_node_name in self.adj[node]: req_node_state = self.nodes[req_node_name]['state'] if req_node_state in ['good', 'unknown']: continue # assume req_node is in a good state else: # maintaining the node to ensure it is in a good state req_node_state = self.maintain_node(req_node_name, verbose=verbose) if req_node_state == 'bad': raise ValueError('Could not calibrate "{}"'.format( req_node_name)) # 2. Once all required nodes are OK, determine action to be taken state = self.nodes[node]['state'] if state == 'needs calibration': # action is clear, no check required pass else: # determine latest status of node state = self.check_node(node, verbose=verbose) # 3. Take action based on the stae of the node if state == 'needs calibration': cal_succes = self.calibrate_node(node, verbose=verbose) # the calibration can still fail if dependencies that were good # or unknown were bad. In that case all dependencies will # explicitly be executed and calibration will be retried if not cal_succes: state = 'bad' if verbose: print('Initial calibration of "{}" failed, ' 'retrying.'.format(node)) if state == 'bad': # if the state is bad it will execute *all* dependencies. Even # the ones that were updated before. if verbose: print('State of node "{}" is bad, maintaining all required' ' nodes.'.format(node)) for req_node_name in self.adj[node]: req_node_state = self.maintain_node(req_node_name, verbose=verbose) cal_succes = self.calibrate_node(node, verbose=verbose) if not cal_succes: raise ValueError( 'Calibration of "{}" failed.'.format(node)) state = self.nodes[node]['state'] return state
[docs] def check_node(self, node, verbose=False): """ Perform check method on specified node Args: node: Node to check verbose: Verbosity level Returns: Returns node state after the check """ if verbose: print('\tChecking node {}.'.format(node)) self.set_node_state(node, 'active') func = _get_function(self.nodes[node]['check_function']) result = func() if isinstance(result, float): if result < self.nodes[node]['tolerance']: self.set_node_state(node, 'good') if verbose: print('\tNode {} is within tolerance.'.format(node)) else: self.set_node_state(node, 'needs calibration') if verbose: print('\tNode {} needs calibration.'.format(node)) elif result == False: if verbose: print('\tNode {} is in a "bad" state.'.format(node)) # if the function returns False it means the check is broken self.set_node_state(node, 'bad') else: raise ValueError('Expected float or "False", ' 'result is: {}'.format(result)) return self.nodes[node]['state']
[docs] def calibrate_node(self, node: str, verbose: bool = False): """ Calibrate specified node Args: node: Node to calibration verbose: Verbosity level Returns: Returns True if the calibration was succesfull, otherwise False """ if verbose: print('\tCalibrating node {}.'.format(node)) self.set_node_state(node, 'active') func = _get_function(self.nodes[node]['calibrate_function']) try: result = func() except Exception as e: self.set_node_state(node, 'bad') logging.warning(e) return False if result: self.set_node_state(node, 'good') if verbose: print('\tCalibration of node {} successful.'.format(node)) return True else: self.set_node_state(node, 'bad') if verbose: print('\tCalibration of node {} failed.'.format(node)) return False
def set_all_node_states(self, state): for node_dat in self.nodes.values(): node_dat['state'] = state def update_monitor(self): if self.cfg_plot_mode == 'matplotlib': self.update_monitor_mpl() elif self.cfg_plot_mode == 'svg': self.draw_svg() elif self.cfg_plot_mode is None or self.cfg_plot_mode == 'None': return else: raise ValueError('cfg_plot_mode should be in ["matplotlib",' ' "svg", "None" ]')
[docs] def update_monitor_mpl(self): """ Updates a plot using the draw_graph_mpl based on matplotlib. """ fig = self.cfg_plot_mode_args.get('fig', None) if fig is not None: plt.figure(fig) plt.clf() self.draw_mpl(plt.gca()) plt.draw() plt.pause(.01)
def _generate_node_positions(self, node_positions: Optional[dict] = None): nodes = self.nodes() if node_positions is None: node_positions = {} def position_generator(N=10, centre=[0, 5]): """ Generate circle of positions around centre """ idx = 0 while True: phi = 2*np.pi*(idx/N)+np.pi/2 pos = 2.1*np.array([np.cos(phi), np.sin(phi)]) + centre yield pos idx = idx+1 positions = position_generator(len(nodes)) pos = dict([(node, node_positions.get(node, next(positions))) for node in nodes]) return pos def draw_mpl(self, ax=None): if ax is None: f, ax = plt.subplots() ax.axis('off') ax.set_title(self.name) colors_list = [state_cmap[node_dat['state']] for node_dat in self.nodes.values()] node_positions = getattr(self, 'node_positions', None) if node_positions is None: pos = nx.nx_agraph.graphviz_layout(self, prog='dot') else: pos = self._generate_node_positions(node_positions) nx.draw_networkx_nodes(self, pos, ax=ax, node_color=colors_list) nx.draw_networkx_edges(self, pos, ax=ax, arrows=True, **self.matplotlib_edge_properties) nx.draw_networkx_labels( self, pos, ax=ax, **self.matplotlib_label_properties) self._format_mpl_plot(ax) @staticmethod def _format_mpl_plot(ax): """ Method to format the generated matplotlib figure """ ax.set_xticks([]) ax.set_yticks([])
[docs] def draw_svg(self, filename: str = None): """ """ if filename is None: filename = self.cfg_svg_filename self._update_drawing_attrs() vis.draw_graph_svg(self, filename)
[docs] def open_html_viewer(self): """ Open html viewer for the file specified by the svg backend """ template = os.path.join(os.path.split(autodepgraph.__file__)[ 0], 'svg_viewer', 'svg_graph_viewer.html') with open(template, 'rt') as fid: x = fid.read() base, file = os.path.split(self.cfg_svg_filename) x = x.replace('adg_graph.svg', file) tfile = tempfile.mktemp(prefix='svgviewer-', suffix='.html', dir=base) with open(tfile, 'wt') as fid: fid.write(x) webbrowser.open_new_tab(tfile) return tfile
[docs] def set_node_attribute(self, node, attribute, value): """ Set the attribute of the specified node Args: node (str): name of the node attribute (str): attribute to set value (ojbect): value to set """ if attribute in ['state']: raise Exception('please use set_state directly') nx.set_node_attributes(self, {node: {attribute: value}})
[docs] def get_node_attribute(self, node, attribute): """ Return the attribute of the specified node Args: node (str): name of the node attribute (str): attribute to get Returns: value (ojbect): attribute of the object """ if attribute in ['state']: raise Exception('please use get_state directly') return self.nodes[node][attribute]
[docs] def set_node_description(self, node, description): """ Set the node description field Args: node (str): name of the node description (str): description to set """ nx.set_node_attributes(self, {node: {'description': description}})
[docs] def calibration_state(self): """ Return dictionary with current calibration state """ return dict(self.nodes)
def _update_drawing_attrs(self): for node_name, node_attrs in self.nodes(True): state = self.get_node_state(node_name) color = vis.state_cmap[state] shape = 'hexagon' if self.is_manual_node(node_name) else 'ellipse' attr_dict = {'shape': shape, 'style': 'filled', 'color': color, # 'fixedsize':'shape', # 'fixedsize' : b"true", 'fixedsize': "false", 'fillcolor': color} node_attrs.update(attr_dict)
def _construct_maintenance_method(): # This placeholder exists to allow reading and writing graphs in a graph # based format. print('Placeholder function.') print('Call DAG._construct_maintenance_methods() to update methods.') def _get_function(funcStr): if isinstance(funcStr, (types.MethodType, types.FunctionType)): warnings.warn('please set function as a str', DeprecationWarning) f = funcStr elif '.' in funcStr: try: instr_name, method = funcStr.split('.') instr = Instrument.find_instrument(instr_name) f = getattr(instr, method) except Exception as e: f = get_function_from_module(funcStr) else: raise Exception('could not find function %s' % funcStr) return f
[docs]def get_function_from_module(funcStr): """ """ split_idx = funcStr.rfind('.') module_name = funcStr[:split_idx] mod = import_module(module_name) f = getattr(mod, funcStr[(split_idx+1):]) return f
def update_node_state(graph_to_update, graph_to_update_from): for node_name, attrs in graph_to_update_from.nodes(True): if node_name in graph_to_update.nodes(): graph_to_update.nodes[node_name]['state'] = attrs['state'] graph_to_update.nodes[node_name]['last_update'] = attrs['last_update']