CIC code in python. Moved simu python code to simu.py

python/cosmotool/__init__.py

@@ -1,119 +1,7 @@
 from _cosmotool import *
 from grafic import writeGrafic, writeWhitePhase, readGrafic
 from borg import read_borg_vol
+from cic import cicParticles
+from simu import loadRamsesAll, simpleWriteGadget, SimulationBare
 
 
-class SimulationBare(PySimulationBase):
-
-  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 = [q.copy() for q in s.getPositions()] if s.getPositions() is not None else None
-    self.velocities = [q.copy() for q in s.getVelocities()] if s.getVelocities() is not None else None
-    self.identifiers = s.getIdentifiers().copy() if s.getIdentifiers() is not None else None
-    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,b)]
-        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.getTime()
-    assert self.Hubble == other.getHubble()
-    assert self.boxsize == other.getBoxsize()
-    assert self.Omega_M == other.getOmega_M()
-    assert self.Omega_Lambda == other.getOmega_Lambda()
-    
-    self.positions = _safe_merge(self.positions, other.getPositions())
-    self.velocities = _safe_merge(self.velocities, other.getVelocities())
-    self.identifiers = _safe_merge0(self.identifiers, other.getIdentifiers())
-    
-  def getPositions(self):
-    return self.positions
-    
-  def getVelocities(self):
-    return self.velocities
-    
-  def getIdentifiers(self):
-    return self.identifiers
-
-  def getTime(self):
-    return self.time
-    
-  def getHubble(self):
-    return self.Hubble
-
-  def getBoxsize(self):
-    return self.boxsize
-
-  def getOmega_M(self):
-    return self.Omega_M
-
-  def getOmega_Lambda(self):
-    return self.Omega_Lambda
-
-
-def simpleWriteGadget(filename, positions, boxsize=1.0, Hubble=100, Omega_M=0.30, time=1, velocities=None, identifiers=None):
-
-  s = SimulationBare()
-  
-  s.positions = positions
-  
-  if velocities:
-    s.velocities = velocities
-  else:
-    s.velocities = [np.zeros(positions[0].size,dtype=np.float32)]*3
-    
-  if identifiers:
-    s.identifiers = identifiers
-  else:
-    s.identifiers = np.arange(positions[0].size, dtype=np.int64)
-    
-  s.Hubble = Hubble
-  s.time = time
-  s.Omega_M = Omega_M
-  s.Omega_Lambda = 1-Omega_M
-  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), snapshot_id, cpu_id, **kwargs)
-    if s == None:
-      break
-    if output == None:
-      output = SimulationBare(s)
-    else:
-      output.merge(s)
-      
-    cpu_id += 1
-
-  return output 

python/cosmotool/cic.py

@@ -0,0 +1,20 @@
+import numpy as np
+
+
+def cicParticles(particles, L, N):
+
+    for d in xrange(3):
+        q = particles[d]*N/L
+        i.append(np.floor(q).astype(int))
+        r.append(q-i[-1])
+    
+    density  = np.bincount(shifted(i, (0,0,0)), weights=   r[0]*    r[1]*   r[2], minlength=N*N*N)
+    density += np.bincount(shifted(i, (0,0,1)), weights=   r[0]*    r[1]*(1-r[2]), minlength=N*N*N)
+    density += np.bincount(shifted(i, (0,1,0)), weights=   r[0]*(1-r[1])*   r[2], minlength=N*N*N)
+    density += np.bincount(shifted(i, (0,1,1)), weights=   r[0]*(1-r[1])*(1-r[2]), minlength=N*N*N)
+    density += np.bincount(shifted(i, (1,0,0)), weights=(1-r[0])*    r[1]*   r[2], minlength=N*N*N)
+    density += np.bincount(shifted(i, (1,0,1)), weights=(1-r[0])*    r[1]*(1-r[2]), minlength=N*N*N)
+    density += np.bincount(shifted(i, (1,1,0)), weights=(1-r[0])*(1-r[1])*   r[2], minlength=N*N*N)
+    density += np.bincount(shifted(i, (1,1,1)), weights=(1-r[0])*(1-r[1])*(1-r[2]), minlength=N*N*N)
+
+    return density

python/cosmotool/simu.py

@@ -0,0 +1,116 @@
+from _cosmotool import *
+
+class SimulationBare(PySimulationBase):
+
+  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 = [q.copy() for q in s.getPositions()] if s.getPositions() is not None else None
+    self.velocities = [q.copy() for q in s.getVelocities()] if s.getVelocities() is not None else None
+    self.identifiers = s.getIdentifiers().copy() if s.getIdentifiers() is not None else None
+    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,b)]
+        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.getTime()
+    assert self.Hubble == other.getHubble()
+    assert self.boxsize == other.getBoxsize()
+    assert self.Omega_M == other.getOmega_M()
+    assert self.Omega_Lambda == other.getOmega_Lambda()
+    
+    self.positions = _safe_merge(self.positions, other.getPositions())
+    self.velocities = _safe_merge(self.velocities, other.getVelocities())
+    self.identifiers = _safe_merge0(self.identifiers, other.getIdentifiers())
+    
+  def getPositions(self):
+    return self.positions
+    
+  def getVelocities(self):
+    return self.velocities
+    
+  def getIdentifiers(self):
+    return self.identifiers
+
+  def getTime(self):
+    return self.time
+    
+  def getHubble(self):
+    return self.Hubble
+
+  def getBoxsize(self):
+    return self.boxsize
+
+  def getOmega_M(self):
+    return self.Omega_M
+
+  def getOmega_Lambda(self):
+    return self.Omega_Lambda
+
+
+def simpleWriteGadget(filename, positions, boxsize=1.0, Hubble=100, Omega_M=0.30, time=1, velocities=None, identifiers=None):
+
+  s = SimulationBare()
+  
+  s.positions = positions
+  
+  if velocities:
+    s.velocities = velocities
+  else:
+    s.velocities = [np.zeros(positions[0].size,dtype=np.float32)]*3
+    
+  if identifiers:
+    s.identifiers = identifiers
+  else:
+    s.identifiers = np.arange(positions[0].size, dtype=np.int64)
+    
+  s.Hubble = Hubble
+  s.time = time
+  s.Omega_M = Omega_M
+  s.Omega_Lambda = 1-Omega_M
+  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), snapshot_id, cpu_id, **kwargs)
+    if s == None:
+      break
+    if output == None:
+      output = SimulationBare(s)
+    else:
+      output.merge(s)
+      
+    cpu_id += 1
+
+  return output