Merge pull request #174 from alongd/flux-mol-images

Flux diagrams with molecule images + main-workflow wiring

Author: alongd

.github/workflows/ci.yml

@@ -43,7 +43,7 @@ jobs:
         uses: actions/checkout@v6
         with:
           repository: ReactionMechanismGenerator/ARC
-          ref: linear_ts
+          ref: main
           path: ARC
 
       # 3. Checkout RMG Database

.gitignore

@@ -149,3 +149,6 @@ htmlcov/*
 #VSCode files
 .vscode/*
 .vscode/launch.json
+
+# Local-only agent working files (never commit)
+docs/superpowers/

docs/input_reference.md

@@ -19,6 +19,16 @@ t3:
   options:
     flux_adapter: RMG  # optional, can use any implemented simulation adapter, default: 'RMG'
     profiles_adapter: RMG  # optional, can use any implemented simulation adapter, default: 'RMG'
+    generate_flux_diagrams: true  # optional, whether the main T3 workflow generates flux diagrams
+                                  # after the sensitivity-analysis step of each iteration,
+                                  # default: True
+    flux_diagrams_with_images: true  # optional, whether to render each flux-diagram node as a
+                                     # molecule structure image (True) instead of a plain text
+                                     # label (False), default: True
+    flux_diagram_reactors: null  # optional, which reactor(s) to draw a flux diagram for:
+                                 # None -> the first reactor only; a 1-based reactor number or a
+                                 # list of numbers -> those reactors; 'all' -> every reactor,
+                                 # default: None
     collision_violators_thermo: false  # optional, whether to calculate thermo of species participating
                                       # in collision violating reactions, default: False
     collision_violators_rates: false  # optional, whether to calculate rates of core collision violating

t3/main.py

@@ -19,8 +19,11 @@
 import os
 import re
 import shutil
+import traceback
 from collections import deque
 
+import cantera as ct
+
 from arc.common import (get_number_with_ordinal_indicator,
                         get_ordinal_indicator,
                         key_by_val,
@@ -37,6 +40,7 @@
 from t3.common import (DATA_BASE_PATH,
                        PROJECTS_BASE_PATH,
                        VALID_CHARS,
+                       convert_termination_time_to_seconds,
                        delete_root_rmg_log,
                        get_species_by_label,
                        sa_dict_from_yaml,
@@ -49,6 +53,8 @@
 from t3.runners.rmg_runner import rmg_runner, run_arkane_job
 from t3.schema import InputBase
 from t3.simulate.factory import simulate_factory
+from t3.utils.fix_cantera import fix_cantera
+from t3.utils.flux import generate_flux
 from t3.utils.libraries import append_to_rmg_libraries
 from t3.utils.writer import write_pdep_network_file, write_rmg_input_file
 from t3.utils.cantera_parser import load_cantera_yaml_file
@@ -293,6 +299,8 @@ def execute(self):
                             content=comparison,
                         )
 
+            self._generate_flux_diagrams()
+
             additional_calcs_required = self.determine_species_and_reactions_to_calculate()
 
             # ARC
@@ -321,6 +329,7 @@ def execute(self):
                              f'------------------------------------------------------\n')
             self.set_paths()
             self.run_rmg(restart_rmg=False)
+            self._generate_flux_diagrams()
 
         self.logger.log_species_summary(species_dict=self.species)
         self.logger.log_reactions_summary(reactions_dict=self.reactions)
@@ -355,6 +364,7 @@ def set_paths(self,
             'chem annotated': os.path.join(iteration_path, 'RMG', 'chemkin', 'chem_annotated.inp'),
             'species dict': os.path.join(iteration_path, 'RMG', 'chemkin', 'species_dictionary.txt'),
             'figs': os.path.join(iteration_path, 'Figures'),
+            'flux diagrams': os.path.join(iteration_path, 'flux'),
             'SA': os.path.join(iteration_path, 'SA'),
             'SA coefficients': os.path.join(iteration_path, 'SA', 'sa_coefficients.yml'),
             'SA dict': os.path.join(iteration_path, 'SA', 'sa.yaml'),
@@ -1551,6 +1561,189 @@ def add_reaction(self,
                 self.reactions[rxn_key].reasons.append(reason)
         return None
 
+    def _select_flux_reactors(self) -> list[int]:
+        """Resolve flux_diagram_reactors to 0-based reactor indices (None -> first reactor)."""
+        n = len(self.rmg['reactors'])
+        selection = self.t3['options']['flux_diagram_reactors']
+        if selection is None:
+            return [0] if n else []
+        if selection == 'all':
+            return list(range(n))
+        numbers = selection if isinstance(selection, list) else [selection]
+        indices = []
+        for num in numbers:
+            idx = num - 1  # option is 1-based
+            if 0 <= idx < n and idx not in indices:
+                indices.append(idx)
+            elif not (0 <= idx < n):
+                self.logger.warning(f'flux_diagram_reactors: reactor {num} out of range '
+                                    f'(have {n}); skipping.')
+        return indices
+
+    def _flux_observables(self) -> list[str]:
+        """Base-label observables for flux, independent of whether SA ran (order-preserving dedupe)."""
+        labels = [s['label'] for s in self.rmg['species']
+                  if s.get('observable') or s.get('SA_observable')] or list(self.sa_observables)
+        seen, deduped = set(), []
+        for label in labels:
+            if label not in seen:
+                seen.add(label)
+                deduped.append(label)
+        return deduped
+
+    def _cantera_name_map(self, model) -> tuple[dict, set, set]:
+        """Return (base_map, ambiguous_bases, full_names).
+
+        base_map: base label ('H') -> full cantera name ('H(3)').
+        ambiguous_bases: base labels that map to >1 species (unusable as a base match).
+        full_names: the set of exact cantera species names (an exact full-name match always wins).
+        """
+        base_map, ambiguous, full_names = dict(), set(), set()
+        for i in range(model.n_species):
+            name = model.species()[i].name
+            full_names.add(name)
+            base = name.split('(')[0]
+            if base in base_map and base_map[base] != name:
+                ambiguous.add(base)
+            base_map[base] = name
+        return base_map, ambiguous, full_names
+
+    def _resolve_species_name(self, label: str, base_map: dict, ambiguous: set,
+                              full_names: set) -> str | None:
+        """Resolve an rmg label to a full cantera name; exact full-name wins over base match."""
+        if label in full_names:
+            return label
+        if label in base_map and label not in ambiguous:
+            return base_map[label]
+        return None
+
+    def _flux_reactor_type(self) -> tuple[str, bool] | None:
+        """Map the configured simulate adapter to (flux reactor_type, energy) or None if unsupported.
+
+        generate_flux supports only 'BatchP' and 'JSR'. PFR / constant-UV adapters have no flux
+        equivalent, so they are unsupported (skip+warn).
+        """
+        adapter = ((self.t3.get('sensitivity') or {}).get('adapter') or '')
+        supported = {'CanteraJSR': ('JSR', False),
+                     'CanteraConstantTP': ('BatchP', False),
+                     'CanteraConstantHP': ('BatchP', True),   # adiabatic const-P -> energy on
+                     'RMGConstantTP': ('BatchP', False),
+                     '': ('BatchP', False)}                    # no adapter: gas batch const T P default
+        unsupported = {'CanteraConstantUV', 'CanteraPFR', 'CanteraPFRTProfile'}
+        if adapter in supported:
+            if adapter == '':
+                self.logger.info('No sensitivity adapter set; assuming a BatchP flux reactor.')
+            return supported[adapter]
+        if adapter in unsupported:
+            self.logger.warning(f"Flux diagrams are not supported for simulate adapter "
+                                f"'{adapter}'; skipping flux diagrams.")
+            return None
+        self.logger.warning(f"Unknown simulate adapter '{adapter}' — cannot determine the flux "
+                            f"reactor type; skipping flux diagrams.")
+        return None
+
+    def _flux_conditions_from_reactor(self, reactor: dict, base_map: dict, ambiguous: set,
+                                      full_names: set, reactor_type: str, energy: bool) -> dict | None:
+        """Build generate_flux conditions from one RMG reactor, or None if the reactor is unusable.
+
+        A reactor is skipped entirely (returns None) if it is liquid/ranged, lacks a termination
+        time, or if ANY species with positive concentration cannot be unambiguously mapped to a
+        cantera species — dropping a real reactant would silently distort the flux graph.
+        """
+        if reactor.get('P') is None or reactor['type'] != 'gas batch constant T P':
+            self.logger.warning(f"Flux-diagram conditions are derived only from 'gas batch constant "
+                                f"T P' reactor entries; skipping reactor entry of type "
+                                f"'{reactor.get('type')}'.")
+            return None
+        if isinstance(reactor['T'], list) or isinstance(reactor['P'], list):
+            self.logger.warning('Flux diagrams do not support ranged T/P reactors; skipping.')
+            return None
+        if reactor.get('termination_time') is None:
+            self.logger.warning('Flux diagrams require a reactor termination_time; skipping.')
+            return None
+        composition = dict()
+        for spc in self.rmg['species']:
+            conc = spc['concentration']
+            if isinstance(conc, (list, tuple)):
+                self.logger.warning('Flux diagrams do not support ranged concentrations; skipping reactor.')
+                return None
+            if conc <= 0:
+                continue
+            name = self._resolve_species_name(spc['label'], base_map, ambiguous, full_names)
+            if name is None:
+                self.logger.warning(f"Flux diagrams: reactant '{spc['label']}' (concentration "
+                                    f"{conc}) could not be mapped to a cantera species; "
+                                    f"skipping this reactor to avoid a distorted flux graph.")
+                return None
+            composition[name] = conc
+        if not composition:
+            self.logger.warning('Flux diagrams: no positive-concentration species for this reactor; skipping.')
+            return None
+        t_final = convert_termination_time_to_seconds(reactor['termination_time'])
+        return {'reactor_type': reactor_type, 'energy': energy,
+                'T': reactor['T'], 'P': reactor['P'],
+                'times': [f * t_final for f in (0.1, 0.5, 1.0)],  # sample sub-terminal + terminal
+                'composition': composition}
+
+    def _generate_flux_diagrams(self) -> None:
+        """Generate flux diagrams for the selected reactor(s). Never raises."""
+        try:
+            if not self.t3['options']['generate_flux_diagrams']:
+                return
+            model_path = self.paths['cantera annotated']
+            if not os.path.isfile(model_path):
+                self.logger.warning(f'Skipping flux diagrams: cantera model not found at {model_path}.')
+                return
+            indices = self._select_flux_reactors()
+            if not indices:
+                return
+            reactor_kind = self._flux_reactor_type()
+            if reactor_kind is None:
+                return                                    # unsupported adapter, already warned
+            reactor_type, energy = reactor_kind
+            observable_labels = self._flux_observables()
+            if not observable_labels:
+                self.logger.info('Skipping flux diagrams: no observables identified.')
+                return
+            fix_cantera(model_path=model_path)          # fix before loading (generate_flux expects fixed)
+            model = ct.Solution(model_path)
+            base_map, ambiguous, full_names = self._cantera_name_map(model)
+            observables = [name for name in
+                           (self._resolve_species_name(o, base_map, ambiguous, full_names)
+                            for o in observable_labels) if name is not None]
+            if not observables:
+                self.logger.warning(f'Skipping flux diagrams: none of the observables '
+                                    f'{observable_labels} map to cantera species.')
+                return
+            draw_images = self.t3['options']['flux_diagrams_with_images']
+            species_dict_path = self.paths['species dict'] if draw_images else None
+            # Use a per-reactor subfolder whenever the user explicitly selected reactor(s) or when
+            # more than one reactor is drawn; the generic folder is reserved for the default case
+            # (no explicit selection -> the single first reactor), preserving reactor identity.
+            explicit = self.t3['options']['flux_diagram_reactors'] is not None
+            for idx in indices:
+                folder = self.paths['flux diagrams'] if (not explicit and len(indices) == 1) else \
+                    os.path.join(self.paths['flux diagrams'], f'reactor_{idx + 1}')
+                try:
+                    conditions = self._flux_conditions_from_reactor(
+                        self.rmg['reactors'][idx], base_map, ambiguous, full_names,
+                        reactor_type, energy)
+                    if conditions is None:
+                        continue
+                    generate_flux(model_path=model_path, folder_path=folder,
+                                  observables=observables,
+                                  draw_molecule_images=draw_images,
+                                  species_dictionary_path=species_dict_path,
+                                  logger=self.logger, fix_cantera_model=False,
+                                  **conditions)
+                except Exception as e:
+                    self.logger.warning(f'Could not generate flux diagram for reactor {idx + 1} '
+                                        f'(target folder {folder}): {e.__class__.__name__}: {e}')
+                    self.logger.debug(traceback.format_exc())
+        except Exception as e:
+            self.logger.warning(f'Flux-diagram generation failed: {e.__class__.__name__}: {e}')
+            self.logger.debug(traceback.format_exc())
+
     def dump_sa_coefficients(self):
         """
         Save the SA coefficients dictionary to a YAML file for user evaluation

t3/schema.py

@@ -5,7 +5,7 @@
 
 import os
 from enum import Enum
-from typing import Annotated
+from typing import Annotated, Literal
 
 from pydantic import BaseModel, Field, ValidationInfo, field_serializer, field_validator, model_validator
 
@@ -32,6 +32,9 @@ class T3Options(BaseModel):
     A class for validating input.T3.options arguments
     """
     flux_adapter: Annotated[str, Field(max_length=255)] = 'RMG'
+    generate_flux_diagrams: bool = True
+    flux_diagrams_with_images: bool = True
+    flux_diagram_reactors: int | list[int] | Literal['all'] | None = None
     profiles_adapter: Annotated[str, Field(max_length=255)] = 'RMG'
     collision_violators_thermo: bool = False
     collision_violators_rates: bool = False
@@ -66,6 +69,19 @@ def enforce_collision_thermo(self) -> T3Options:
             self.collision_violators_thermo = True
         return self
 
+    @field_validator('flux_diagram_reactors', mode='before')
+    @classmethod
+    def check_flux_diagram_reactors(cls, value):
+        """flux_diagram_reactors: 1-based reactor number(s) or 'all'."""
+        if value is None or value == 'all':
+            return value
+        nums = value if isinstance(value, list) else [value]
+        for n in nums:
+            if isinstance(n, bool) or not isinstance(n, int) or n < 1:
+                raise ValueError(f'flux_diagram_reactors must be positive 1-based int(s) or '
+                                 f'"all"; got {value!r}')
+        return value
+
     @field_validator('library_name')
     @classmethod
     def check_library_name(cls, value):