Add QSS (quasi steady state) analysis in chemistry_reduce.py

Author: jannisteunissen

tools/chemistry_reduce.py

@@ -30,6 +30,10 @@ def load_data(simname):
     n_species = len(species)
     n_reactions = len(reactions)
 
+    # Load amounts of species in the total simulation volume
+    amounts = np.loadtxt(simname + '_amounts.txt')
+    amounts = amounts[:, 1:] # Drop time column
+
     assert stoich_matrix.shape == (n_species, n_reactions), \
         f"Stoichiometric matrix shape {stoich_matrix.shape} != " \
         f"({n_species}, {n_reactions})"
@@ -45,6 +49,7 @@ def load_data(simname):
         "species": species,
         "reactions": reactions,
         "stoich_matrix": stoich_matrix,
+        "amounts": amounts
     }
 
 
@@ -156,7 +161,7 @@ def kept_reactions_from_species(stoich_matrix, kept_species_set):
     return kept
 
 
-def drg_reduce(data, target_species, threshold=0.01):
+def drg_reduce(data, target_species, threshold=0.01, verbose=False):
     """
     Apply the DRG method to reduce the mechanism.
 
@@ -172,16 +177,13 @@ def drg_reduce(data, target_species, threshold=0.01):
 
     Returns
     -------
-    kept_species : list of str
-        Species retained after reduction.
-    kept_reactions : list of str
-        Reactions retained after reduction.
-    kept_species_idx : list of int
-        Indices of retained species.
-    kept_reactions_idx : list of int
-        Indices of retained reactions.
-    r_AB : array
-        The DRG coefficient matrix.
+    result : dict
+        Dictionary with keys:
+        - kept_species : list of str
+        - kept_reactions : list of str
+        - kept_species_idx : list of int
+        - kept_reactions_idx : list of int
+        - r_AB : array, the DRG coefficient matrix
     """
     species = data["species"]
     reactions = data["reactions"]
@@ -205,9 +207,7 @@ def drg_reduce(data, target_species, threshold=0.01):
     if not target_indices:
         raise ValueError("No valid target species provided!")
 
-    print("\nComputing DRG coefficients")
     r_AB = compute_drg_coefficients(stoich_matrix, rates)
-    print(f"DRG coefficient matrix computed (max={r_AB.max():.4f})")
 
     kept_species_set = bfs_reachable(r_AB, target_indices, threshold)
     kept_species_idx = sorted(kept_species_set)
@@ -217,13 +217,20 @@ def drg_reduce(data, target_species, threshold=0.01):
                                                      kept_species_set)
     kept_reactions = [reactions[i] for i in kept_reactions_idx]
 
-    print(f"\nDRG Reduction Results (threshold={threshold}):")
-    print(f"  Species: {n_species} -> {len(kept_species)} "
-          f"(removed {n_species - len(kept_species)})")
-    print(f"  Reactions: {n_reactions} -> {len(kept_reactions)} "
-          f"(removed {n_reactions - len(kept_reactions)})")
+    if verbose:
+        print(f"\nDRG Reduction Results (threshold={threshold}):")
+        print(f"  Species: {n_species} -> {len(kept_species)} "
+              f"(removed {n_species - len(kept_species)})")
+        print(f"  Reactions: {n_reactions} -> {len(kept_reactions)} "
+              f"(removed {n_reactions - len(kept_reactions)})")
 
-    return kept_species, kept_reactions, kept_species_idx, kept_reactions_idx, r_AB
+    return {
+        "kept_species": kept_species,
+        "kept_reactions": kept_reactions,
+        "kept_species_idx": kept_species_idx,
+        "kept_reactions_idx": kept_reactions_idx,
+        "r_AB": r_AB,
+    }
 
 
 def sweep_thresholds(data, target_species, thresholds=None):
@@ -288,6 +295,135 @@ def sweep_thresholds(data, target_species, thresholds=None):
     return results
 
 
+def detect_qss_species(data, threshold=0.1, min_fraction=0.9,
+                       exclude_species=None, verbose=False):
+    """
+    Detect species satisfying the quasi-steady-state approximation.
+
+    A species is flagged as QSS if, for at least `min_fraction` of all
+    time steps, its net production/consumption imbalance is below
+    `threshold`.
+
+    Parameters
+    ----------
+    data : dict
+        Data loaded by load_data().
+    threshold : float
+        Maximum allowed relative imbalance delta_A for QSS.
+    min_fraction : float
+        Fraction of time steps that must satisfy the criterion
+        (avoids flagging species that are only transiently in steady state).
+    exclude_species : list of str, optional
+        Species that should never be marked as QSS (e.g., targets,
+        major species, electrons).
+
+    Returns
+    -------
+    qss_species : list of str
+        Species names identified as quasi-steady-state.
+    delta : array, shape (n_times, n_species)
+        The imbalance metric at each time step for each species.
+    """
+    if exclude_species is None:
+        exclude_species = []
+
+    stoich = data["stoich_matrix"]   # (n_species, n_reactions)
+    rates = data["rates"]            # (n_times, n_reactions)
+    species = data["species"]
+
+    n_times, n_reactions = rates.shape
+    n_species = len(species)
+
+    # Contribution of each reaction to each species: nu_{A,i} * R_i(t)
+    # shape: (n_times, n_species, n_reactions)
+    omega = stoich[np.newaxis, :, :] * rates[:, np.newaxis, :]
+
+    # Production: positive contributions; Consumption: negative contributions
+    production = np.where(omega > 0, omega, 0.0).sum(axis=2)   # (n_times, n_species)
+    consumption = np.where(omega < 0, -omega, 0.0).sum(axis=2) # (n_times, n_species)
+
+    # Net rate and imbalance
+    net = np.abs(production - consumption)
+    scale = np.maximum(production, consumption)
+    scale = np.maximum(scale, ATOL)  # avoid division by zero
+
+    delta = net / scale  # (n_times, n_species)
+
+    # A species is QSS if delta < threshold for >= min_fraction of time steps
+    fraction_steady = (delta < threshold).sum(axis=0) / n_times  # (n_species,)
+
+    qss_species = []
+
+    if verbose:
+        print(f"\n{'Species':>20} | {'Frac steady':>12} | {'Mean delta':>12} | "
+              f"{'QSS?':>5}")
+        print("-" * 60)
+
+    for i, sp in enumerate(species):
+        is_qss = (fraction_steady[i] >= min_fraction
+                  and sp not in exclude_species)
+        if is_qss:
+            qss_species.append(sp)
+
+        # Print species that are close to QSS for diagnostics
+        if verbose and (fraction_steady[i] > 0.5 or is_qss):
+            print(f"{sp:>20} | {fraction_steady[i]:>11.3f} | "
+                  f"{delta[:, i].mean():>12.4e} | "
+                  f"{'YES' if is_qss else 'no':>5}")
+
+    if verbose:
+        print(f"\nIdentified {len(qss_species)} QSS species (threshold="
+              f"{threshold}, min_fraction={min_fraction})")
+
+    return qss_species, delta
+
+
+def drg_with_qss(data, target_species, threshold=0.1, qss_threshold=0.1,
+                 qss_min_fraction=0.9, verbose=False):
+    """
+    Two-stage reduction:
+    1. DRG to remove unimportant species
+    2. QSS detection among the kept species to identify algebraic species
+    """
+    # Stage 1: DRG graph-based removal
+    drg_result = drg_reduce(data, target_species, threshold)
+    kept_species = drg_result["kept_species"]
+    kept_reactions = drg_result["kept_reactions"]
+    r_AB = drg_result["r_AB"]
+
+    # Stage 2: Among kept species, find QSS candidates
+    qss_candidates, delta = detect_qss_species(
+        data,
+        threshold=qss_threshold,
+        min_fraction=qss_min_fraction,
+        exclude_species=target_species,  # never eliminate targets
+        verbose=verbose
+    )
+
+    # Only consider QSS species that survived DRG
+    qss_in_kept = [sp for sp in qss_candidates if sp in kept_species]
+
+    # Kept species that are not qss
+    ode_species = [sp for sp in kept_species if sp not in qss_in_kept]
+
+    if verbose:
+        print(f"  Original:    {len(data['species'])} species, "
+              f"{len(data['reactions'])} reactions")
+        print(f"  After DRG:   {len(kept_species)} species, "
+              f"{len(kept_reactions)} reactions")
+        print(f"  QSS species: {len(qss_in_kept)} species (removed by QSS)")
+        print(f"  ODE species: {len(ode_species)} species (kept after QSS)")
+
+    return {
+        "kept_species": kept_species,
+        "kept_reactions": kept_reactions,
+        "qss_species": qss_in_kept,
+        "ode_species": ode_species,
+        "r_AB": r_AB,
+        "delta": delta,
+    }
+
+
 def main():
     import argparse
 
@@ -304,16 +440,55 @@ def main():
                    help="DRG threshold")
     p.add_argument("-s", "--sweep", action="store_true",
                    help="Sweep over multiple thresholds")
+    p.add_argument("-q", "--qss", action="store_true",
+                   help="Enable QSS detection after DRG reduction")
+    p.add_argument("--qss-threshold", type=float, default=0.1,
+                   help="QSS imbalance threshold")
+    p.add_argument("--qss-min-fraction", type=float, default=0.9,
+                   help="Min. fraction of time steps satisfying QSS criterion")
+    p.add_argument("-v", "--verbose", action="store_true",
+                   help="Enable verbose output")
 
     args = p.parse_args()
 
     data = load_data(args.simname)
 
     if args.sweep:
         sweep_thresholds(data, args.target)
+    elif args.qss:
+        result = drg_with_qss(data, args.target,
+                              threshold=args.threshold,
+                              qss_threshold=args.qss_threshold,
+                              verbose=args.verbose)
+
+        print(f"\nODE species ({len(result['ode_species'])}):")
+        for sp in result["ode_species"]:
+            print(f"  {sp}")
+
+        print(f"\nQSS species ({len(result['qss_species'])}):")
+        for sp in result["qss_species"]:
+            print(f"  {sp}")
+
+        removed_species = sorted(
+            set(data["species"]) - set(result["kept_species"]))
+        print(f"\nRemoved species ({len(removed_species)}):")
+        for sp in removed_species:
+            print(f"  {sp}")
+
+        print(f"\nKept reactions ({len(result['kept_reactions'])}):")
+        for rx in result["kept_reactions"]:
+            print(f"  {rx}")
+
+        removed_reactions = sorted(
+            set(data["reactions"]) - set(result["kept_reactions"]))
+        print(f"\nRemoved reactions ({len(removed_reactions)}):")
+        for rx in removed_reactions:
+            print(f"  {rx}")
     else:
-        kept_species, kept_reactions, kept_species_idx, kept_reactions_idx, r_AB = \
-            drg_reduce(data, args.target, args.threshold)
+        result = drg_reduce(data, args.target, args.threshold,
+                            verbose=args.verbose)
+        kept_species = result["kept_species"]
+        kept_reactions = result["kept_reactions"]
 
         print(f"\nKept species ({len(kept_species)}):")
         for sp in kept_species:
@@ -328,7 +503,8 @@ def main():
         for rx in kept_reactions:
             print(f"  {rx}")
 
-        removed_reactions = sorted(set(data["reactions"]) - set(kept_reactions))
+        removed_reactions = sorted(
+            set(data["reactions"]) - set(kept_reactions))
         print(f"\nRemoved reactions ({len(removed_reactions)}):")
         for rx in removed_reactions:
             print(f"  {rx}")