cosmotool/python/cosmotool/grafic.py
2014-06-04 14:00:54 +02:00

67 lines
2.3 KiB
Python

import struct
import numpy as np
def readGrafic(filename):
with file(filename, mode="rb") as f:
p = struct.unpack("IIIIffffffffI", f.read(4*11 + 2*4))
checkPoint0, Nx, Ny, Nz, delta, _, _, _, scalefac, omega0, omegalambda0, h, checkPoint1 = p
if checkPoint0 != checkPoint1 or checkPoint0 != 4*11:
raise ValueError("Invalid unformatted access")
a = np.empty((Nx,Ny,Nz), dtype=np.float32)
BoxSize = delta * Nx * h
checkPoint = 4*Ny*Nz
for i in xrange(Nx):
checkPoint = struct.unpack("I", f.read(4))[0]
if checkPoint != 4*Ny*Nz:
raise ValueError("Invalid unformatted access")
a[i, :, :] = np.fromfile(f, dtype=np.float32, count=Ny*Nz).reshape((Ny, Nz))
checkPoint = struct.unpack("I", f.read(4))[0]
if checkPoint != 4*Ny*Nz:
raise ValueError("Invalid unformatted access")
return a, BoxSize, scalefac, omega0, omegalambda0, h
def writeGrafic(filename, field, BoxSize, scalefac, **cosmo):
with file(filename, mode="wb") as f:
checkPoint = 4*11
Nx,Ny,Nz = field.shape
delta = BoxSize/Nx/cosmo['h']
bad = 0.0
f.write(struct.pack("IIIIffffffffI", checkPoint,
Nx, Ny, Nz,
delta,
bad, bad, bad,
scalefac,
cosmo['omega_M_0'], cosmo['omega_lambda_0'], 100*cosmo['h'], checkPoint))
checkPoint = 4*Ny*Nx
field = field.reshape(field.shape, order='F')
for i in xrange(Nz):
f.write(struct.pack("I", checkPoint))
f.write(field[i].astype(np.float32).tostring())
f.write(struct.pack("I", checkPoint))
def writeWhitePhase(filename, field):
with file(filename, mode="wb") as f:
Nx,Ny,Nz = field.shape
checkPoint = 4*4
f.write(struct.pack("IIIIII", checkPoint, Nx, Ny, Nz, 0, checkPoint))
checkPoint = struct.pack("I", 4*Ny*Nz)
for i in xrange(Nx):
f.write(checkPoint)
f.write(field[i].astype(np.float32).tostring())
f.write(checkPoint)