test(sandbox): some work on sandbox tests

validation/pynucastro/GridFireEquiv/evolve.py

@@ -38,13 +38,17 @@
 #     f.write("t,h1,h2,he3,he4,c12,n14,o16,ne20,mg24\n")
 #     for (t,h1,h2,he3,he4,c12,n14,o16,ne20,mg24) in zip(sol.t, sol.y[network.jp, :], sol.y[network.jd, :], sol.y[network.jhe3, :], sol.y[network.jhe4, :], sol.y[network.jc12, :], sol.y[network.jn14, :], sol.y[network.jo16, :], sol.y[network.jne20, :], sol.y[network.jmg24, :]):
 #         f.write(f"{t},{h1},{h2},{he3},{he4},{c12},{n14},{o16},{ne20},{mg24}\n")
+import sys
 
 from gridfire.engine import GraphEngine, MultiscalePartitioningEngineView, AdaptiveEngineView
 from gridfire.solver import CVODESolverStrategy
 from gridfire.type import NetIn
 
 from fourdst.composition import Composition
-from fourdst.atomic import species
+from fourdst.atomic import species, Species
+
+from datetime import datetime
+from typing import List, Dict, Set, Tuple
 
 symbols : list[str] = ["H-1", "He-3", "He-4", "C-12", "N-14", "O-16", "Ne-20", "Mg-24"]
 X : list[float]     = [0.708, 2.94e-5, 0.276, 0.003, 0.0011, 9.62e-3, 1.62e-3, 5.16e-4]
@@ -61,7 +65,7 @@ netIn = NetIn()
 netIn.composition = comp
 netIn.temperature = 1.5e7
 netIn.density = 1.6e2
-netIn.tMax = 3.14e16
+netIn.tMax = 3e17
 netIn.dt0 = 1e-12
 
 baseEngine = GraphEngine(netIn.composition, 2)
@@ -73,35 +77,46 @@ adaptiveEngine = AdaptiveEngineView(qseEngine)
 
 solver = CVODESolverStrategy(adaptiveEngine)
 
-data = []
+data: List[Tuple[float, Dict[str, Tuple[float, float]]]] = []
 def callback(context):
-    H1Index = context.engine.getSpeciesIndex(species["H-1"])
-    H2Index = context.engine.getSpeciesIndex(species["H-2"])
-    He3Index = context.engine.getSpeciesIndex(species["He-3"])
-    He4Index = context.engine.getSpeciesIndex(species["He-4"])
-    C12Index = context.engine.getSpeciesIndex(species["C-12"])
-    N14Index = context.engine.getSpeciesIndex(species["N-14"])
-    O16Index = context.engine.getSpeciesIndex(species["O-16"])
-    Ne20Index = context.engine.getSpeciesIndex(species["Mg-24"])
-    Mg24Index = context.engine.getSpeciesIndex(species["Mg-24"])
-    data.append([context.t,
-                 context.state[H1Index],
-                 context.state[H2Index],
-                 context.state[He3Index],
-                 context.state[He4Index],
-                 context.state[C12Index],
-                 context.state[N14Index],
-                 context.state[O16Index],
-                 context.state[Ne20Index],
-                 context.state[Mg24Index]
-            ])
-
+    engine = context.engine
+    abundances: Dict[str, Tuple[float, float]] = {}
+    for species in engine.getNetworkSpecies():
+        sid = engine.getSpeciesIndex(species)
+        abundances[species.name()] = (species.mass(), context.state[sid])
+    data.append((context.t,abundances))
 
 solver.set_callback(callback)
-results = solver.evaluate(netIn, False)
 
-# with open("gridfire_results.csv", 'w') as f:
-#     f.write("t,h1,h2,he3,he4,c12,n14,o16,ne20,mg24\n")
-#     for row in data:
-#         rowStr = ','.join([str(x) for x in row])
-#         f.write(f"{rowStr}\n")
+try:
+    results = solver.evaluate(netIn, False)
+    print(f"H-1 final molar abundance: {results.composition.getMolarAbundance("H-1"):0.3f}")
+except:
+    print("Warning: solver did not converge", file=sys.stderr)
+
+
+uniqueSpecies: Set[Tuple[str, float]] = set()
+for t, timestep in data:
+    for (name, (mass, abundance)) in timestep.items():
+        uniqueSpecies.add((name, mass))
+sortedUniqueSpecies = list(sorted(uniqueSpecies, key=lambda e: e[1]))
+
+with open(f"gridfire_results_{datetime.now().strftime("%H-%M-%d_%m_%Y")}.csv", 'w') as f:
+    f.write('t,')
+    for i, (species,mass) in enumerate(sortedUniqueSpecies):
+        f.write(f"{species}")
+        if i < len(sortedUniqueSpecies)-1:
+            f.write(",")
+    f.write("\n")
+
+    for t, timestep in data:
+        f.write(f"{t},")
+        for i, (species, mass) in enumerate(sortedUniqueSpecies):
+            if timestep.get(species, None) is not None:
+                f.write(f"{timestep.get(species)[1]}")
+            else:
+                f.write(f"")
+            if i < len(sortedUniqueSpecies) - 1:
+                f.write(",")
+        f.write("\n")
+