tboudreaux/GridFire
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf(GridFire)

Browse Source 
More preformance improvmnets

1. Switch to mimalloc which gave a roughly 10% improvment accross the
board
2. Use much faster compososition caching
3. Reusing work vector
This commit is contained in:
Emily Boudreaux
2025-12-07 12:34:12 -05:00
parent e48b62f231
commit 8cfa067ad0
23 changed files with 306 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
validation/vv/GridFireValidationSuite.py
View File

@@ -1,7 +1,8 @@
from gridfire.policy import MainSequencePolicy, NetworkPolicy
from gridfire.engine import DynamicEngine, GraphEngine
from gridfire.engine import DynamicEngine, GraphEngine, EngineTypes
from gridfire.type import NetIn
from typing import Dict
from fourdst.composition import Composition
from testsuite import TestSuite
@@ -9,6 +10,12 @@ from utils import init_netIn, init_composition, years_to_seconds
from enum import Enum
EngineNameToType: Dict[str, EngineTypes] = {
"graphengine": EngineTypes.GRAPH_ENGINE,
"multiscalepartitioningengineview": EngineTypes.MULTISCALE_PARTITIONING_ENGINE_VIEW,
"adaptiveengineview": EngineTypes.ADAPTIVE_ENGINE_VIEW
}
class SolarLikeStar_QSE_Suite(TestSuite):
def __init__(self):
initialComposition : Composition = init_composition()
@@ -22,11 +29,11 @@ class SolarLikeStar_QSE_Suite(TestSuite):
notes="Thermodynamically Static, MultiscalePartitioning Engine View"
)
def __call__(self):
def __call__(self, pynucastro_compare: bool = False, pync_engine: str = "AdaptiveEngineView"):
policy : MainSequencePolicy = MainSequencePolicy(self.composition)
engine : DynamicEngine = policy.construct()
netIn : NetIn = init_netIn(self.temperature, self.density, self.tMax, self.composition)
self.evolve(engine, netIn)
self.evolve(engine, netIn, pynucastro_compare = pynucastro_compare, engine_type=EngineNameToType[pync_engine.lower()])
class MetalEnhancedSolarLikeStar_QSE_Suite(TestSuite):
def __init__(self):
@@ -41,7 +48,7 @@ class MetalEnhancedSolarLikeStar_QSE_Suite(TestSuite):
notes="Thermodynamically Static, MultiscalePartitioning Engine View, Z enhanced by 1 dex, temperature reduced to 80% of solar core"
)
def __call__(self):
def __call__(self, pynucastro_compare: bool = False, pync_engine: str = "AdaptiveEngineView"):
policy : MainSequencePolicy = MainSequencePolicy(self.composition)
engine : GraphEngine = policy.construct()
netIn : NetIn = init_netIn(self.temperature, self.density, self.tMax, self.composition)
@@ -59,7 +66,7 @@ class MetalDepletedSolarLikeStar_QSE_Suite(TestSuite):
notes="Thermodynamically Static, MultiscalePartitioning Engine View, Z depleted by 1 dex, temperature increased to 120% of solar core"
)
def __call__(self):
def __call__(self, pynucastro_compare: bool = False, pync_engine: str = "AdaptiveEngineView"):
policy : MainSequencePolicy = MainSequencePolicy(self.composition)
engine : GraphEngine = policy.construct()
netIn : NetIn = init_netIn(self.temperature, self.density, self.tMax, self.composition)
@@ -78,7 +85,7 @@ class SolarLikeStar_No_QSE_Suite(TestSuite):
notes="Thermodynamically Static, No MultiscalePartitioning Engine View"
)
def __call__(self):
def __call__(self, pynucastro_compare: bool = False, pync_engine: str = "AdaptiveEngineView"):
engine : GraphEngine = GraphEngine(self.composition, 3)
netIn : NetIn = init_netIn(self.temperature, self.density, self.tMax, self.composition)
self.evolve(engine, netIn)
@@ -94,9 +101,19 @@ if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="Run some subset of the GridFire validation suite.")
parser.add_argument('--suite', type=str, choices=[suite.name for suite in ValidationSuites], nargs="+", help="The validation suite to run.")
parser.add_argument("--all", action="store_true", help="Run all validation suites.")
parser.add_argument("--pynucastro-compare", action="store_true", help="Generate pynucastro comparison data.")
parser.add_argument("--pync-engine", type=str, choices=["GraphEngine", "MultiscalePartitioningEngineView", "AdaptiveEngineView"], default="AdaptiveEngineView", help="The GridFire engine to use to select the reactions for pyuncastro comparison.")
args = parser.parse_args()
for suite_name in args.suite:
suite = ValidationSuites[suite_name]
instance : TestSuite = suite.value()
instance()
if args.all:
for suite in ValidationSuites:
instance : TestSuite = suite.value()
instance(args.pynucastro_compare, args.pync_engine)
else:
for suite_name in args.suite:
suite = ValidationSuites[suite_name]
instance : TestSuite = suite.value()
instance(args.pynucastro_compare, args.pync_engine)

57
validation/vv/testsuite.py
View File

@@ -2,8 +2,11 @@ from abc import ABC, abstractmethod
import fourdst.atomic
import scipy.integrate
import gridfire
from fourdst.composition import Composition
from gridfire.engine import DynamicEngine, GraphEngine
from gridfire.engine import DynamicEngine, GraphEngine, AdaptiveEngineView, MultiscalePartitioningEngineView
from gridfire.engine import EngineTypes
from gridfire.policy import MainSequencePolicy
from gridfire.type import NetIn, NetOut
from gridfire.exceptions import GridFireError
from gridfire.solver import CVODESolverStrategy
@@ -21,6 +24,12 @@ import numpy as np
import json
import time
EngineTypeLookup : Dict[EngineTypes, Any] = {
EngineTypes.ADAPTIVE_ENGINE_VIEW: AdaptiveEngineView,
EngineTypes.MULTISCALE_PARTITIONING_ENGINE_VIEW: MultiscalePartitioningEngineView,
EngineTypes.GRAPH_ENGINE: GraphEngine
}
def load_network_module(filepath):
module_name = os.path.basename(filepath).replace(".py", "")
if module_name in sys.modules: # clear any existing module with the same name
@@ -103,12 +112,16 @@ class TestSuite(ABC):
self.composition : Composition = composition
self.notes : str = notes
def evolve_pynucastro(self, engine: GraphEngine):
def evolve_pynucastro(self, engine: DynamicEngine):
print("Evolution complete. Now building equivalent pynucastro network...")
# Build equivalent pynucastro network for comparison
reaclib_library : pyna.ReacLibLibrary = pyna.ReacLibLibrary()
rate_names = [r.id().replace("e+","").replace("e-","").replace(", ", ",") for r in engine.getNetworkReactions()]
with open(f"{self.name}_rate_names_pynuc.txt", "w") as f:
for r_name in rate_names:
f.write(f"{r_name}\n")
goodRates : List[pyna.rates.reaclib_rate.ReacLibRate] = []
missingRates = []
@@ -156,7 +169,23 @@ class TestSuite(ABC):
atol=1e-8
)
endTime = time.time()
initial_duration = endTime - startTime
print("Pynucastro integration complete. Writing results to JSON...")
print("Running pynucastro a second time to account for any JIT compilation overhead...")
startTime = time.time()
sol = scipy.integrate.solve_ivp(
net.rhs,
[0, self.tMax],
Y0,
args=(self.density, self.temperature),
method="BDF",
jac=net.jacobian,
rtol=1e-5,
atol=1e-8
)
endTime = time.time()
final_duration = endTime - startTime
print(f"Pynucastro second integration complete. Initial run time: {initial_duration: .4f} s, Second run time: {final_duration: .4f} s")
data: List[Dict[str, Union[float, Dict[str, float]]]] = []
@@ -182,7 +211,8 @@ class TestSuite(ABC):
"Temperature": self.temperature,
"Density": self.density,
"tMax": self.tMax,
"ElapsedTime": endTime - startTime,
"RunTime0": initial_duration,
"RunTime1": final_duration,
"DateCreated": datetime.now().isoformat()
},
"Steps": data
@@ -191,7 +221,7 @@ class TestSuite(ABC):
with open(f"GridFireValidationSuite_{self.name}_pynucastro.json", "w") as f:
json.dump(pynucastro_json, f, indent=4)
def evolve(self, engine: GraphEngine, netIn: NetIn, pynucastro_compare: bool = True):
def evolve(self, engine: DynamicEngine, netIn: NetIn, pynucastro_compare: bool = True, engine_type: EngineTypes | None = None):
solver : CVODESolverStrategy = CVODESolverStrategy(engine)
stepLogger : StepLogger = StepLogger()
@@ -232,10 +262,23 @@ class TestSuite(ABC):
)
if pynucastro_compare:
self.evolve_pynucastro(engine)
if engine_type is not None:
if engine_type == EngineTypes.ADAPTIVE_ENGINE_VIEW:
print("Pynucastro comparison using AdaptiveEngineView...")
self.evolve_pynucastro(engine)
elif engine_type == EngineTypes.MULTISCALE_PARTITIONING_ENGINE_VIEW:
print("Pynucastro comparison using MultiscalePartitioningEngineView...")
graphEngine : GraphEngine = GraphEngine(self.composition, depth=3)
multiScaleEngine : MultiscalePartitioningEngineView = MultiscalePartitioningEngineView(graphEngine)
self.evolve_pynucastro(multiScaleEngine)
elif engine_type == EngineTypes.GRAPH_ENGINE:
print("Pynucastro comparison using GraphEngine...")
graphEngine : GraphEngine = GraphEngine(self.composition, depth=3)
self.evolve_pynucastro(graphEngine)
else:
print(f"Pynucastro comparison not implemented for engine type: {engine_type}")
@abstractmethod
def __call__(self):
def __call__(self, pynucastro_compare: bool = False, pync_engine: str = "AdaptiveEngineView"):
pass