diff --git a/python/_cosmotool.pyx b/python/_cosmotool.pyx index 7bbba85..706285c 100644 --- a/python/_cosmotool.pyx +++ b/python/_cosmotool.pyx @@ -301,3 +301,5 @@ def loadRamses(str basepath, int snapshot_id, int cpu_id, bool doublePrecision = return None return _PySimulationAdaptor(wrap_simudata(data, flags)) + +def loadAllRamses(str basepath, int diff --git a/python/cosmotool/__init__.py b/python/cosmotool/__init__.py index 162b5e0..de5d9fa 100644 --- a/python/cosmotool/__init__.py +++ b/python/cosmotool/__init__.py @@ -5,9 +5,53 @@ from borg import read_borg_vol class SimulationBare(PySimulationBase): - def __init__(self): - pass - + def __init__(self, *args): + if len(args) == 0: + return + + if not isinstance(args[0], PySimulationBase): + raise TypeError("Simulation object to mirror must be a PySimulationBase") + s = args[0] + + self.positions = s.getPositions() + self.velocities = s.getVelocities() + self.identifiers = s.getIdentifiers() + self.boxsize = s.getBoxsize() + self.time = s.getTime() + self.Hubble = s.getHubble() + self.Omega_M = s.getOmega_M() + self.Omega_Lambda = s.getOmega_Lambda() + + + def merge(self, other): + + def _safe_merge(a, b): + if b: + if a: + a = [np.append(q, r) for q,r in zip(a,other.positions)] + else: + a = b + return a + + def _safe_merge0(a, b): + if b: + if a: + a = np.append(a, b) + else: + a = b + return a + + + assert self.time == other.time + assert self.Hubble == other.Hubble + assert self.boxsize == other.boxsize + assert self.Omega_M == other.Omega_M + assert self.Omega_Lambda == other.Omega_Lambda + + self.positions = _safe_merge(self.positions, other.positions) + self.velocities = _safe_merge(self.velocities, other.velocities) + self.identifiers = _safe_merge0(self.idenfiers, other.identifiers) + def getPositions(self): return self.positions @@ -56,3 +100,18 @@ def simpleWriteGadget(filename, positions, boxsize=1.0, Hubble=100, Omega_M=0.30 s.boxsize = boxsize writeGadget(filename, s) + +def loadRamsesAll(basepath, snapshot_id, **kwargs): + + cpu_id = 0 + output = None + while True: + s = loadRamses("%s/output_%05d" % (basepath,snapshot_id), cpu_id, **kwargs) + if s == None: + break + if output == None: + output = SimulationBare(s) + else: + output.merge(s) + + cpu_id += 1