Auto stash before rebase of "idt"

t3/main.py

@@ -158,7 +158,7 @@ def __init__(self,
                  clean_dir: bool = False,
                  ):
 
-        self.sa_dict = None
+        self.sa_dict, self.sa_dict_idt = None, None
         self.sa_observables = list()
         self.t0 = datetime.datetime.now()  # initialize the timer as datetime object
 
@@ -313,9 +313,12 @@ def execute(self):
                                                         sa_rtol=self.t3['sensitivity']['rtol'],
                                                         )
                     simulate_adapter.simulate()
-                    self.sa_dict = simulate_adapter.get_sa_coefficients()
-                    save_yaml_file(path=self.paths['SA dict'], content=self.sa_dict)
-                if self.t3['sensitivity']['global_observables'] == 'IDT':
+                    self.sa_dict = simulate_adapter.get_sa_coefficients(
+                        top_SA_species=self.t3['sensitivity']['top_SA_species'],
+                        top_SA_reactions=self.t3['sensitivity']['top_SA_reactions'],
+                        max_workers=self.t3['sensitivity']['max_workers'],
+                        save_yaml=True)
+                if self.t3['sensitivity']['global_observables'].lower() == 'idt':
                     simulate_adapter = simulate_factory(simulate_method='CanteraIDT',
                                                         t3=self.t3,
                                                         rmg=self.rmg,
@@ -327,8 +330,11 @@ def execute(self):
                                                         sa_rtol=self.t3['sensitivity']['rtol'],
                                                         )
                     simulate_adapter.simulate()
-                    self.sa_dict_idt = simulate_adapter.get_sa_coefficients()
-                    save_yaml_file(path=self.paths['SA idt dict'], content=self.sa_dict_idt)
+                    self.sa_dict_idt = simulate_adapter.get_sa_coefficients(
+                        top_SA_species=self.t3['sensitivity']['top_SA_species'],
+                        top_SA_reactions=self.t3['sensitivity']['top_SA_reactions'],
+                        max_workers=self.t3['sensitivity']['max_workers'],
+                        save_yaml=True)
 
             additional_calcs_required = self.determine_species_and_reactions_to_calculate()
 
@@ -395,6 +401,7 @@ def set_paths(self,
             'SA input': os.path.join(iteration_path, 'SA', 'input.py'),
             'SA dict': os.path.join(iteration_path, 'SA', 'sa.yaml'),
             'SA IDT dict': os.path.join(iteration_path, 'SA', 'sa_idt.yaml'),
+            'SA IDT dict top X': os.path.join(iteration_path, 'SA', 'sa_idt_top_x.yaml'),
             'PDep SA': os.path.join(iteration_path, 'PDep_SA'),
             'ARC': os.path.join(iteration_path, 'ARC'),
             'ARC input': os.path.join(iteration_path, 'ARC', 'input.yml'),
@@ -691,6 +698,7 @@ def determine_species_and_reactions_to_calculate(self) -> bool:
             bool: Whether additional calculations are required.
         """
         species_keys, reaction_keys, coll_vio_spc_keys, coll_vio_rxn_keys = list(), list(), list(), list()
+        rxn_idt_keys = None
 
         self.rmg_species, self.rmg_reactions = self.load_species_and_reactions_from_chemkin_file()
         self.logger.info(f'This RMG model has {len(self.rmg_species)} species '
@@ -726,6 +734,9 @@ def determine_species_and_reactions_to_calculate(self) -> bool:
             # 1.2. SA
             if sa_observables_exist:
                 species_keys.extend(self.determine_species_based_on_sa())
+            if self.t3['sensitivity']['global_observables'].lower() == 'idt':
+                species_idt_keys, rxn_idt_keys = self.determine_params_based_on_sa_idt()
+                species_keys.extend(species_idt_keys)
             # 1.3. collision violators
             if self.t3['options']['collision_violators_thermo']:
                 species_keys.extend(coll_vio_spc_keys)
@@ -734,6 +745,10 @@ def determine_species_and_reactions_to_calculate(self) -> bool:
         # 2.1. SA
         if sa_observables_exist:
             reaction_keys.extend(self.determine_reactions_based_on_sa())
+        if self.t3['sensitivity']['global_observables'].lower() == 'idt':
+            if rxn_idt_keys is None:
+                species_idt_keys, rxn_idt_keys = self.determine_params_based_on_sa_idt()
+            reaction_keys.extend(rxn_idt_keys)
         # 2.2. collision violators
         if self.t3['options']['collision_violators_rates']:
             reaction_keys.extend(coll_vio_rxn_keys)
@@ -812,6 +827,58 @@ def determine_species_based_on_sa(self) -> List[int]:
 
         return species_keys
 
+    def determine_params_based_on_sa_idt(self) -> Tuple[List[int], List[int]]:
+        """
+        Determine species or reactions to calculate based on sensitivity analysis of IDT.
+
+        Returns:
+            List[int]: Entries are T3 species indices of species determined to be calculated based on IDT SA.
+        """
+        visited_species, species_keys = list(), list()
+        visited_rxns, rxn_keys = list(), list()
+        if self.sa_dict_idt is None:
+            self.logger.error(f"T3's sa_dict_idt was None. Please check that the input file contains a proper "
+                              f"'sensitivity' block, that 'IDT' was defined in the global_observables list, "
+                              f"and/or that SA was run successfully.\n"
+                              f"Not performing refinement based on IDT sensitivity analysis!")
+            return species_keys, rxn_keys
+        for token in ['thermo', 'kinetics']:
+            for r, reactor_idt_data in self.sa_dict_idt[token]['IDT'].items():
+                for phi, phi_data in reactor_idt_data.items():
+                    for p, p_data in phi_data.items():
+                        for t, idt_data in p_data.items():
+                            for index, sa in idt_data.items():
+                                if token == 'thermo':
+                                    if index not in visited_species:
+                                        visited_species.append(index)
+                                        species = self.get_species_by_index(index)
+                                        if self.species_requires_refinement(species=species):
+                                            reason = f'(i {self.iteration}) IDT is sensitive to this species.'
+                                            key = self.add_species(species=species, reasons=reason)
+                                            if key is not None:
+                                                species_keys.append(key)
+                                elif token == 'kinetics':
+                                    if index not in visited_rxns:
+                                        visited_rxns.append(index)
+                                        reaction = self.get_reaction_by_index(index)
+                                        if self.reaction_requires_refinement(reaction):
+                                            reason = f'(i {self.iteration}) IDT is sensitive to this reaction.'
+                                            key = self.add_reaction(reaction=reaction, reasons=reason)
+                                            if key is not None:
+                                                rxn_keys.append(key)
+                                        for spc in reaction.reactants + reaction.products:
+                                            spc_index = self.get_species_key(spc)
+                                            if spc_index not in visited_species:
+                                                visited_species.append(spc_index)
+                                                if self.species_requires_refinement(species=spc):
+                                                    reason = f'(i {self.iteration}) IDT is sensitive to a reaction ' \
+                                                             f'({reaction}) in which this species participates.'
+                                                    key = self.add_species(species=spc, reasons=reason)
+                                                    if key is not None:
+                                                        species_keys.append(key)
+        return species_keys, rxn_keys
+
+
     def determine_reactions_based_on_sa(self) -> List[int]:
         """
         Determine reaction rate coefficients to calculate based on sensitivity analysis.
@@ -1178,6 +1245,21 @@ def get_species_key(self,
                 return key
         return None
 
+    def get_species_by_index(self, index) -> Optional[Species]:
+        """
+        Get a species object by its T3 index.
+
+        Args:
+            index (int): The species index.
+
+        Returns:
+            Optional[Species]: The species object if it exists, ``None`` if it does not.
+        """
+        for key, species_dict in self.species.items():
+            if key == index:
+                return species_dict['object']
+        return None
+
     def get_reaction_key(self,
                          reaction: Optional[Reaction] = None,
                          label: Optional[str] = None,
@@ -1206,6 +1288,21 @@ def get_reaction_key(self,
                 return key
         return None
 
+    def get_reaction_by_index(self, index) -> Optional[Reaction]:
+        """
+        Get a reaction object by its T3 index.
+
+        Args:
+            index (int): The reaction index.
+
+        Returns:
+            Optional[Reaction]: The reaction object if it exists, ``None`` if it does not.
+        """
+        for key, reaction_dict in self.reactions.items():
+            if key == index:
+                return reaction_dict['object']
+        return None
+
     def load_species_and_reactions_from_chemkin_file(self) -> Tuple[List[Species], List[Reaction]]:
         """
         Load RMG Species and Reaction objects from the annotated Chemkin file.
@@ -1251,7 +1348,7 @@ def add_species(self,
                     ) -> Optional[int]:
         """
         Add a species to self.species and to self.qm['species'].
-        If the species already exists in self.species, only the reasons to compute will be appended.
+        If the species already exists in self.species, only the reasons to compute it will be appended.
 
         Args:
             species (Species): The species to consider.

t3/simulate/adapter.py

@@ -5,6 +5,7 @@
 """
 
 from abc import ABC, abstractmethod
+from typing import Optional
 
 
 class SimulateAdapter(ABC):
@@ -27,7 +28,12 @@ def simulate(self):
         pass
 
     @abstractmethod
-    def get_sa_coefficients(self):
+    def get_sa_coefficients(self,
+                            top_SA_species: int = 10,
+                            top_SA_reactions: int = 10,
+                            max_workers: int = 24,
+                            save_yaml: bool = True,
+                            ) -> Optional[dict]:
         """
         Obtain the sensitivity analysis coefficients.
 

t3/simulate/cantera_constantHP.py

@@ -382,12 +382,23 @@ def simulate(self):
 
             self.all_data.append((time, condition_data, reaction_sensitivity_data, thermodynamic_sensitivity_data))
 
-    def get_sa_coefficients(self):
+    def get_sa_coefficients(self,
+                            top_SA_species: int = 10,
+                            top_SA_reactions: int = 10,
+                            max_workers: int = 24,
+                            save_yaml: bool = True,
+                            ) -> Optional[dict]:
         """
         Obtain the SA coefficients.
 
+        Args:
+            top_SA_species (int, optional): The number of top sensitive species to return.
+            top_SA_reactions (int, optional): The number of top sensitive reactions to return.
+            max_workers (int, optional): The maximal number of workers to use for parallel processing.
+            save_yaml (bool, optional): Save the SA dictionary to a YAML file.
+
         Returns:
-             sa_dict (dict): a SA dictionary, whose structure is given in the docstring for T3/t3/main.py
+             sa_dict (Optional[dict]): a SA dictionary, whose structure is given in the docstring for T3/t3/main.py
         """
         sa_dict = {'kinetics': dict(), 'thermo': dict(), 'time': list()}
 
@@ -411,7 +422,7 @@ def get_sa_coefficients(self):
                     sa_dict['thermo'][observable_label] = dict()
                 parameter = get_parameter_from_header(spc)
                 sa_dict['thermo'][observable_label][parameter] = spc.data
-
+        # todo: add save yaml to all adapters, use top SA species and reactions, change usage in main
         return sa_dict
 
     def get_idt_by_T(self):

t3/simulate/cantera_constantTP.py

@@ -381,12 +381,23 @@ def simulate(self):
 
             self.all_data.append((time, condition_data, reaction_sensitivity_data, thermodynamic_sensitivity_data))
 
-    def get_sa_coefficients(self):
+    def get_sa_coefficients(self,
+                            top_SA_species: int = 10,
+                            top_SA_reactions: int = 10,
+                            max_workers: int = 24,
+                            save_yaml: bool = True,
+                            ) -> Optional[dict]:
         """
         Obtain the SA coefficients.
 
+        Args:
+            top_SA_species (int, optional): The number of top sensitive species to return.
+            top_SA_reactions (int, optional): The number of top sensitive reactions to return.
+            max_workers (int, optional): The maximal number of workers to use for parallel processing.
+            save_yaml (bool, optional): Save the SA dictionary to a YAML file.
+
         Returns:
-             sa_dict (dict): a SA dictionary, whose structure is given in the docstring for T3/t3/main.py
+             sa_dict (Optional[dict]): a SA dictionary, whose structure is given in the docstring for T3/t3/main.py
         """
         sa_dict = {'kinetics': dict(), 'thermo': dict(), 'time': list()}
 

t3/simulate/cantera_constantUV.py

@@ -381,12 +381,23 @@ def simulate(self):
 
             self.all_data.append((time, condition_data, reaction_sensitivity_data, thermodynamic_sensitivity_data))
 
-    def get_sa_coefficients(self):
+    def get_sa_coefficients(self,
+                            top_SA_species: int = 10,
+                            top_SA_reactions: int = 10,
+                            max_workers: int = 24,
+                            save_yaml: bool = True,
+                            ) -> Optional[dict]:
         """
         Obtain the SA coefficients.
 
+        Args:
+            top_SA_species (int, optional): The number of top sensitive species to return.
+            top_SA_reactions (int, optional): The number of top sensitive reactions to return.
+            max_workers (int, optional): The maximal number of workers to use for parallel processing.
+            save_yaml (bool, optional): Save the SA dictionary to a YAML file.
+
         Returns:
-             sa_dict (dict): a SA dictionary, whose structure is given in the docstring for T3/t3/main.py
+             sa_dict (Optional[dict]): a SA dictionary, whose structure is given in the docstring for T3/t3/main.py
         """
         sa_dict = {'kinetics': dict(), 'thermo': dict(), 'time': list()}
 

t3/simulate/rmg_constantTP.py

@@ -18,6 +18,8 @@
 from rmgpy.tools.loader import load_rmg_py_job
 from rmgpy.tools.plot import plot_sensitivity
 
+from arc.common import read_yaml_file, save_yaml_file
+
 from t3.common import get_chem_to_rmg_rxn_index_map, get_species_by_label, get_values_within_range, \
     get_observable_label_from_header, get_parameter_from_header, time_lapse
 from t3.simulate.adapter import SimulateAdapter
@@ -224,10 +226,21 @@ def simulate(self):
 
         self.logger.info(f'Simulation via RMG completed, execution time: {time_lapse(tic)}')
 
-    def get_sa_coefficients(self) -> Optional[dict]:
+    def get_sa_coefficients(self,
+                            top_SA_species: int = 10,
+                            top_SA_reactions: int = 10,
+                            max_workers: int = 24,
+                            save_yaml: bool = True,
+                            ) -> Optional[dict]:
         """
         Obtain the SA coefficients.
 
+        Args:
+            top_SA_species (int, optional): The number of top sensitive species to return.
+            top_SA_reactions (int, optional): The number of top sensitive reactions to return.
+            max_workers (int, optional): The maximal number of workers to use for parallel processing.
+            save_yaml (bool, optional): Save the SA dictionary to a YAML file.
+
         Returns:
              sa_dict (Optional[dict]): An SA dictionary, structure is given in the docstring for T3/t3/main.py
         """
@@ -265,6 +278,8 @@ def get_sa_coefficients(self) -> Optional[dict]:
                         parameter = chem_to_rmg_rxn_index_map[int(parameter)] \
                             if all(c.isdigit() for c in parameter) else parameter
                     sa_dict[sa_type][observable_label][parameter] = df[header].values
+        if save_yaml:
+            save_yaml_file(path=self.paths['SA dict'], content=sa_dict)
         return sa_dict
 
     def generate_rmg_reactors_for_simulation(self) -> dict: