2020-03-31 15:39:59 +02:00
|
|
|
import json
|
|
|
|
from dataclasses import dataclass
|
|
|
|
from pathlib import Path
|
|
|
|
from typing import Dict
|
|
|
|
|
|
|
|
from rebound import Particle
|
|
|
|
|
|
|
|
|
|
|
|
class ExtraData:
|
|
|
|
|
|
|
|
def __init__(self):
|
|
|
|
self.tree = CollisionTree()
|
|
|
|
self.pdata: Dict[int, ParticleData] = {}
|
|
|
|
self.meta = Meta()
|
2020-12-27 13:45:07 +01:00
|
|
|
self.energy = EnergyConservation()
|
2020-03-31 15:39:59 +02:00
|
|
|
|
|
|
|
def save(self, filename: Path):
|
|
|
|
pdata = {}
|
|
|
|
for k, v in self.pdata.items():
|
|
|
|
pdata[k] = v.__dict__
|
|
|
|
|
|
|
|
with filename.open("w") as f:
|
|
|
|
json.dump({
|
|
|
|
"meta": self.meta.save(),
|
|
|
|
"pdata": pdata,
|
2020-12-27 13:40:11 +01:00
|
|
|
"tree": self.tree.save(),
|
|
|
|
"dEs": self.energy._dEs
|
2020-03-31 15:39:59 +02:00
|
|
|
}, f, indent=2)
|
|
|
|
|
|
|
|
@classmethod
|
|
|
|
def load(cls, filename: Path):
|
|
|
|
with filename.open() as f:
|
|
|
|
data = json.load(f)
|
|
|
|
self = cls()
|
|
|
|
self.meta = Meta(**data["meta"])
|
|
|
|
|
|
|
|
self.tree.load(data["tree"])
|
2020-12-27 13:40:11 +01:00
|
|
|
# self.tree._dEs = data["dEs"]
|
2020-03-31 15:39:59 +02:00
|
|
|
|
|
|
|
for k, v in data["pdata"].items():
|
2020-12-27 13:40:11 +01:00
|
|
|
self.pdata[int(k)] = ParticleData(**v)
|
2020-03-31 15:39:59 +02:00
|
|
|
|
|
|
|
return self
|
|
|
|
|
|
|
|
|
|
|
|
@dataclass
|
|
|
|
class ParticleData:
|
|
|
|
water_mass_fraction: float
|
|
|
|
|
|
|
|
|
|
|
|
@dataclass
|
|
|
|
class Meta:
|
|
|
|
tmax: float = None
|
|
|
|
savesteps: int = None
|
|
|
|
max_n: int = None
|
|
|
|
walltime: int = None # seconds
|
|
|
|
cputime: int = None # seconds
|
|
|
|
current_time: float = None
|
|
|
|
current_steps: float = None
|
|
|
|
|
|
|
|
def save(self):
|
|
|
|
return self.__dict__
|
|
|
|
|
|
|
|
|
|
|
|
class CollisionTree:
|
|
|
|
def __init__(self):
|
|
|
|
self._tree = {}
|
|
|
|
|
2020-03-31 19:12:27 +02:00
|
|
|
def add(self, source1: Particle, source2: Particle, to: Particle, metadata: Dict):
|
|
|
|
self._tree[to.hash.value] = {"parents": [source1.hash.value, source2.hash.value], "meta": metadata}
|
2020-03-31 15:39:59 +02:00
|
|
|
|
|
|
|
def save(self):
|
|
|
|
return self._tree
|
|
|
|
|
|
|
|
def load(self, tree):
|
|
|
|
self._tree = tree
|
2020-12-27 13:40:11 +01:00
|
|
|
|
|
|
|
|
2020-12-27 13:45:07 +01:00
|
|
|
class EnergyConservation:
|
2020-12-27 13:40:11 +01:00
|
|
|
def __init__(self):
|
|
|
|
self._dEs = []
|
|
|
|
self.initial_energy = None
|
|
|
|
|
|
|
|
def set_initial_energy(self, initial_energy: float) -> None:
|
|
|
|
self.initial_energy = initial_energy
|
|
|
|
|
|
|
|
def add_energy_value(self, energy: float) -> None:
|
|
|
|
dE = abs((energy - self.initial_energy) / self.initial_energy)
|
|
|
|
self._dEs.append(dE)
|