Update graph building

Author: calvinp0

arc/plotter.py

@@ -51,6 +51,7 @@
                                    )
 from arc.species.perceive import perceive_molecule_from_xyz
 from arc.species.species import ARCSpecies, rmg_mol_to_dict_repr
+from arc.provenance.nodes import _enum_val, NodeType, EdgeType, DecisionKind
 
 
 PRETTY_UNITS = {'(s^-1)': r' (s$^-1$)',
@@ -73,15 +74,132 @@ def _wrap_graph_label(text: str, width: int = 24) -> str:
                      for line in (textwrap.wrap(part, width=width) or ['']))
 
 
+def render_provenance_graph(prov_graph, run_label: str = 'ARC run') -> 'graphviz.Digraph':
+    """
+    Render a :class:`ProvenanceGraph` as a Graphviz directed graph.
+
+    Node styling by type:
+      - **species**: box / aliceblue
+      - **calculation**: box / color by status (honeydew=done, mistyrose=errored, white=pending)
+      - **data**: note / cornsilk
+      - **decision**: diamond / color by kind (lavender, moccasin, mistyrose)
+
+    Edge styling by type:
+      - ``selected_by``: solid green
+      - ``rejected_by``: dashed red
+      - ``troubleshot_by``: dashed orange
+      - ``retried_as`` / ``fine_of``: dotted gray
+      - others: solid black
+
+    Args:
+        prov_graph: A :class:`ProvenanceGraph` instance.
+        run_label (str): Label for the root run node.
+
+    Returns:
+        graphviz.Digraph: The rendered graph object.
+    """
+    gv = graphviz.Digraph(
+        name='arc_provenance',
+        comment=f'ARC provenance for {run_label}',
+        graph_attr={'rankdir': 'LR', 'splines': 'true', 'overlap': 'false'},
+        node_attr={'shape': 'box', 'style': 'rounded,filled', 'fillcolor': 'white', 'fontname': 'Helvetica'},
+        edge_attr={'fontname': 'Helvetica'},
+    )
+
+    # Node styling lookup
+    _calc_colors = {'done': 'honeydew', 'errored': 'mistyrose', 'pending': 'white'}
+    _decision_colors = {
+        'ts_guess_selection': 'lavender',
+        'ts_guess_selection_failed': 'mistyrose',
+        'job_troubleshooting': 'moccasin',
+        'conformer_selection': 'lavender',
+        'ts_guess_clustering': 'lavender',
+        'ts_method_spawning': 'lavender',
+        'ts_validation_freq': 'lightyellow',
+        'ts_validation_nmd': 'lightyellow',
+        'ts_validation_irc': 'lightyellow',
+        'ts_switch': 'mistyrose',
+    }
+
+    # Edge styling lookup
+    _edge_styles = {
+        'selected_by': {'color': 'green3', 'style': 'solid'},
+        'rejected_by': {'color': 'red', 'style': 'dashed'},
+        'troubleshot_by': {'color': 'orange', 'style': 'dashed'},
+        'triggered_by': {'color': 'gray40', 'style': 'solid'},
+        'retried_as': {'color': 'gray60', 'style': 'dotted'},
+        'fine_of': {'color': 'gray60', 'style': 'dotted'},
+        'spawned_by': {'color': 'blue', 'style': 'solid'},
+    }
+
+    for node in prov_graph.nodes.values():
+        nid = _sanitize_graphviz_id(node.node_id)
+        ntype = node.node_type
+
+        if ntype == 'species':
+            lbl = node.label or node.node_id
+            is_ts = (node.metadata or {}).get('is_ts', False)
+            if is_ts:
+                lbl += '\nTS'
+            gv.node(nid, _wrap_graph_label(lbl), shape='box', fillcolor='aliceblue')
+
+        elif ntype == 'calculation':
+            parts = [getattr(node, 'job_type', '') or '', getattr(node, 'job_name', '') or '']
+            if getattr(node, 'job_adapter', None):
+                parts.append(node.job_adapter)
+            if getattr(node, 'level', None):
+                parts.append(node.level)
+            lbl = '\n'.join(p for p in parts if p)
+            status = getattr(node, 'status', 'pending') or 'pending'
+            fillcolor = _calc_colors.get(status, 'white')
+            gv.node(nid, _wrap_graph_label(lbl), shape='box', fillcolor=fillcolor)
+
+        elif ntype == 'data':
+            dk = getattr(node, 'data_kind', '') or ''
+            val = getattr(node, 'value', None)
+            lbl = dk
+            if val is not None and not isinstance(val, (list, dict)):
+                lbl += f'\n{val}'
+            gv.node(nid, _wrap_graph_label(lbl), shape='note', fillcolor='cornsilk')
+
+        elif ntype == 'decision':
+            dk = getattr(node, 'decision_kind', '') or ''
+            outcome = getattr(node, 'outcome', '') or ''
+            lbl = dk.replace('_', ' ')
+            if outcome:
+                lbl += f'\n{outcome}'
+            fillcolor = _decision_colors.get(dk, 'lavender')
+            gv.node(nid, _wrap_graph_label(lbl, width=28), shape='diamond', fillcolor=fillcolor)
+
+        else:
+            gv.node(nid, _wrap_graph_label(node.node_id))
+
+    for edge in prov_graph.edges:
+        src = _sanitize_graphviz_id(edge.source_id)
+        tgt = _sanitize_graphviz_id(edge.target_id)
+        etype = edge.edge_type
+        style_attrs = _edge_styles.get(etype, {})
+        label = etype.replace('_', ' ') if etype not in ('belongs_to', 'input_of', 'output_of') else ''
+        gv.edge(src, tgt, label=label, **style_attrs)
+
+    return gv
+
+
 def save_provenance_artifacts(project_directory: str,
                               provenance: dict,
+                              graph=None,
                               ) -> dict:
     """
     Save provenance YAML and render Graphviz artifacts for an ARC run.
 
+    When a ``graph`` (:class:`ProvenanceGraph`) is provided, the Graphviz
+    visualization is built from the graph's typed nodes and edges rather
+    than the flat event list, producing richer diagrams.
+
     Args:
         project_directory (str): The ARC project directory.
         provenance (dict): A provenance dictionary with an ``events`` list.
+        graph: Optional ProvenanceGraph instance for graph-based rendering.
 
     Returns:
         dict: Paths to generated artifacts.
@@ -99,6 +217,23 @@ def save_provenance_artifacts(project_directory: str,
         save_yaml_file(path=yml_path, content=provenance)
         return {'yml': yml_path, 'dot': None, 'svg': None}
 
+    # Prefer graph-based rendering when a ProvenanceGraph is available.
+    if graph is not None and len(graph) > 0:
+        gv_graph = render_provenance_graph(graph, run_label=run_label)
+        with open(dot_path, 'w') as f:
+            f.write(gv_graph.source)
+        try:
+            svg_data = gv_graph.pipe(format='svg')
+        except (graphviz.ExecutableNotFound, graphviz.CalledProcessError):
+            logger.warning('Could not render ARC provenance SVG because Graphviz is not available on this system.')
+        else:
+            with open(svg_path, 'wb') as f:
+                f.write(svg_data)
+        provenance['updated_at'] = datetime.datetime.now().isoformat(timespec='seconds')
+        save_yaml_file(path=yml_path, content=provenance)
+        return {'yml': yml_path, 'dot': dot_path, 'svg': svg_path if os.path.isfile(svg_path) else None}
+
+    # Fallback: event-based rendering (legacy path).
     graph = graphviz.Digraph(
         name='arc_provenance',
         comment=f'ARC provenance for {run_label}',

arc/plotter_test.py

@@ -9,6 +9,11 @@
 import shutil
 import unittest
 
+try:
+    import graphviz
+except ImportError:
+    graphviz = None
+
 import arc.plotter as plotter
 from arc.common import ARC_PATH, ARC_TESTING_PATH, read_yaml_file, safe_copy_file
 from arc.species.converter import str_to_xyz
@@ -300,6 +305,57 @@ def test_save_provenance_artifacts(self):
             # Troubleshoot follow-up connects from the decision diamond, not the species node.
             self.assertIn('decision_7 -> job_spc1_freq_a3', dot)
 
+    def test_render_provenance_graph(self):
+        """Test Graphviz rendering from a ProvenanceGraph object."""
+        from arc.provenance import (ProvenanceGraph, DecisionKind, DataKind, EdgeType)
+        g = ProvenanceGraph(project='render_test')
+        sid = g.add_species_node(label='ethanol')
+        cid = g.add_calculation_node(label='ethanol', job_name='opt_a1',
+                                     job_type='opt', job_adapter='gaussian',
+                                     level='b3lyp/6-31g(d)', status='done')
+        did = g.add_data_node(label='ethanol', data_kind=DataKind.energy, value=-79.5)
+        dec = g.add_decision_node(label='ethanol',
+                                  decision_kind=DecisionKind.conformer_selection,
+                                  outcome='Selected conformer #0')
+        g.add_edge(sid, cid, EdgeType.input_of)
+        g.add_edge(cid, did, EdgeType.output_of)
+        g.add_edge(did, dec, EdgeType.selected_by)
+
+        if graphviz is not None:
+            gv = plotter.render_provenance_graph(g, run_label='render_test')
+            dot_source = gv.source
+            self.assertIn('ethanol', dot_source)
+            self.assertIn('opt', dot_source)
+            self.assertIn('energy', dot_source)
+            self.assertIn('conformer selection', dot_source)
+            self.assertIn('honeydew', dot_source)  # done calc
+            self.assertIn('cornsilk', dot_source)  # data node
+            self.assertIn('diamond', dot_source)    # decision node
+            self.assertIn('green3', dot_source)     # selected_by edge
+
+    def test_save_provenance_artifacts_with_graph(self):
+        """Test that save_provenance_artifacts prefers graph-based rendering when a graph is provided."""
+        from arc.provenance import (ProvenanceGraph, DecisionKind, EdgeType)
+        project = 'arc_project_for_testing_delete_after_usage'
+        project_directory = os.path.join(ARC_PATH, 'Projects', project)
+        g = ProvenanceGraph(project=project)
+        sid = g.add_species_node(label='spc1')
+        cid = g.add_calculation_node(label='spc1', job_name='opt_a1',
+                                     job_type='opt', status='done')
+        g.add_edge(sid, cid, EdgeType.input_of)
+        provenance = {'project': project, 'events': []}
+        paths = plotter.save_provenance_artifacts(
+            project_directory=project_directory,
+            provenance=provenance,
+            graph=g,
+        )
+        self.assertTrue(os.path.isfile(paths['yml']))
+        if paths['dot'] is not None:
+            with open(paths['dot'], 'r') as f:
+                dot = f.read()
+            # Graph-based rendering uses node IDs like species_1 not event-based species_spc1.
+            self.assertIn('species_1', dot)
+            self.assertIn('honeydew', dot)
 
     @classmethod
     def tearDownClass(cls):

arc/provenance/__init__.py

@@ -0,0 +1,38 @@
+"""
+ARC provenance subpackage — directed acyclic graph for computational provenance.
+
+Tracks the full chain of inputs, calculations, decisions, and outputs that
+produce ARC's results.  Inspired by AiiDA's DAG model but adapted for ARC's
+branching decision trees (TS guess evaluation, conformer selection,
+troubleshooting loops).
+
+Submodules:
+    - ``nodes``: Node types, edge types, and their data classes.
+    - ``graph``: ProvenanceGraph container with query and serialization.
+"""
+
+from arc.provenance.graph import ProvenanceGraph
+from arc.provenance.nodes import (
+    CalculationNode,
+    DataKind,
+    DataNode,
+    DecisionKind,
+    DecisionNode,
+    EdgeType,
+    NodeType,
+    ProvenanceEdge,
+    ProvenanceNode,
+)
+
+__all__ = [
+    'ProvenanceGraph',
+    'ProvenanceNode',
+    'CalculationNode',
+    'DataNode',
+    'DecisionNode',
+    'ProvenanceEdge',
+    'NodeType',
+    'DataKind',
+    'DecisionKind',
+    'EdgeType',
+]