Emily Boudreaux
bb1d6bbb24
Python bindings have now been brought back up to feature pairity with C++. Further, stubs have been added for all python features so that code completion will work
96 lines
3.3 KiB
Python
96 lines
3.3 KiB
Python
from fourdst.composition import Composition
|
|
from gridfire.type import NetIn
|
|
from gridfire.policy import MainSequencePolicy
|
|
from gridfire.solver import CVODESolverStrategy
|
|
from enum import Enum
|
|
from typing import Dict, Union, SupportsFloat
|
|
import json
|
|
import dicttoxml
|
|
|
|
def init_composition() -> Composition:
|
|
Y = [7.0262E-01, 9.7479E-06, 6.8955E-02, 2.5000E-04, 7.8554E-05, 6.0144E-04, 8.1031E-05, 2.1513E-05]
|
|
S = ["H-1", "He-3", "He-4", "C-12", "N-14", "O-16", "Ne-20", "Mg-24"]
|
|
return Composition(S, Y)
|
|
|
|
def init_netIn(temp: float, rho: float, time: float, comp: Composition) -> NetIn:
|
|
netIn = NetIn()
|
|
netIn.temperature = temp
|
|
netIn.density = rho
|
|
netIn.tMax = time
|
|
netIn.dt0 = 1e-12
|
|
netIn.composition = comp
|
|
return netIn
|
|
|
|
class StepData(Enum):
|
|
TIME = 0
|
|
DT = 1
|
|
COMP = 2
|
|
CONTRIB = 3
|
|
|
|
|
|
class StepLogger:
|
|
def __init__(self):
|
|
self.num_steps: int = 0
|
|
self.step_data: Dict[int, Dict[StepData, Union[SupportsFloat, Dict[str, SupportsFloat]]]] = {}
|
|
|
|
def log_step(self, context):
|
|
engine = context.engine
|
|
self.step_data[self.num_steps] = {}
|
|
self.step_data[self.num_steps][StepData.TIME] = context.t
|
|
self.step_data[self.num_steps][StepData.DT] = context.dt
|
|
comp_data: Dict[str, SupportsFloat] = {}
|
|
for species in engine.getNetworkSpecies():
|
|
sid = engine.getSpeciesIndex(species)
|
|
comp_data[species.name()] = context.state[sid]
|
|
self.step_data[self.num_steps][StepData.COMP] = comp_data
|
|
self.num_steps += 1
|
|
|
|
def to_json (self, filename: str):
|
|
serializable_data = {
|
|
stepNum: {
|
|
StepData.TIME.name: step[StepData.TIME],
|
|
StepData.DT.name: step[StepData.DT],
|
|
StepData.COMP.name: step[StepData.COMP],
|
|
}
|
|
for stepNum, step in self.step_data.items()
|
|
}
|
|
with open(filename, 'w') as f:
|
|
json.dump(serializable_data, f, indent=4)
|
|
|
|
def to_xml(self, filename: str):
|
|
serializable_data = {
|
|
stepNum: {
|
|
StepData.TIME.name: step[StepData.TIME],
|
|
StepData.DT.name: step[StepData.DT],
|
|
StepData.COMP.name: step[StepData.COMP],
|
|
}
|
|
for stepNum, step in self.step_data.items()
|
|
}
|
|
xml_data = dicttoxml.dicttoxml(serializable_data, custom_root='StepLog', attr_type=False)
|
|
with open(filename, 'wb') as f:
|
|
f.write(xml_data)
|
|
|
|
def main(temp: float, rho: float, time: float):
|
|
comp = init_composition()
|
|
netIn = init_netIn(temp, rho, time, comp)
|
|
|
|
policy = MainSequencePolicy(comp)
|
|
engine = policy.construct()
|
|
|
|
solver = CVODESolverStrategy(engine)
|
|
|
|
step_logger = StepLogger()
|
|
solver.set_callback(lambda context: step_logger.log_step(context))
|
|
|
|
solver.evaluate(netIn, False)
|
|
step_logger.to_xml("log_data.xml")
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
parser = argparse.ArgumentParser(description="Simple python example of GridFire usage")
|
|
parser.add_argument("-t", "--temp", type=float, help="Temperature in K", default=1.5e7)
|
|
parser.add_argument("-r", "--rho", type=float, help="Density in g/cm^3", default=1.5e2)
|
|
parser.add_argument("--tMax", type=float, help="Time in s", default=3.1536 * 1e17)
|
|
args = parser.parse_args()
|
|
main(args.temp, args.rho, args.tMax)
|