from __future__ import division
import sys,os
if sys.version_info.major>=3:
from io import StringIO
else:
from StringIO import StringIO
from .nbody import NbodySnapshot
from .. import extern as ext
from .settings import LTSettings
import numpy as np
import astropy.units as u
from astropy.cosmology import w0waCDM,LambdaCDM
############################################################
################Gadget2Settings class#######################
############################################################
[docs]class Gadget2Settings(LTSettings):
"""
Class handler of the tunable settings in a Gadget2 run
"""
file_names = ["InitCondFile","OutputDir","EnergyFile","InfoFile","TimingsFile","CpuFile","RestartFile","SnapshotFileBase","OutputListFilename"]
cpu_timings = ["TimeLimitCPU","ResubmitOn","ResubmitCommand"]
code_options = ["ICFormat","SnapFormat","ComovingIntegrationOn","TypeOfTimestepCriterion","OutputListOn","PeriodicBoundariesOn"]
characteristics_of_run = ["TimeMax"]
output_frequency = ["TimeBetSnapshot","TimeOfFirstSnapshot","CpuTimeBetRestartFile","TimeBetStatistics","NumFilesPerSnapshot","NumFilesWrittenInParallel"]
accuracy_time_integration = ["ErrTolIntAccuracy","MaxRMSDisplacementFac","CourantFac","MaxSizeTimestep","MinSizeTimestep"]
tree_algorithm = ["ErrTolTheta","TypeOfOpeningCriterion","ErrTolForceAcc","TreeDomainUpdateFrequency"]
sph = ["DesNumNgb","MaxNumNgbDeviation","ArtBulkViscConst","InitGasTemp","MinGasTemp"]
memory_allocation = ["PartAllocFactor","TreeAllocFactor","BufferSize"]
system_of_units = ["UnitLength_in_cm","UnitMass_in_g","UnitVelocity_in_cm_per_s","GravityConstantInternal"]
softening = ["MinGasHsmlFractional","SofteningGas","SofteningHalo","SofteningDisk","SofteningBulge","SofteningStars","SofteningBndry","SofteningGasMaxPhys","SofteningHaloMaxPhys","SofteningDiskMaxPhys","SofteningBulgeMaxPhys","SofteningStarsMaxPhys","SofteningBndryMaxPhys"]
def __init__(self,**kwargs):
#Default outputs
self.OutputScaleFactor = np.array([0.2463687286034,0.25331915596808,0.2603792960781,0.26755062217624,0.27483471475984,0.282233268286,0.28974809827109,0.29738114881152,0.30513450055509,0.31301037915502,0.32101116424033,0.32913939894726,0.3373978000372,0.34578926866836,0.35431690185511,0.36298400467612,0.37179410329025,0.38075095882651,0.38985858222059,0.39912125007794,0.40854352165158,0.4181302570317,0.42788663665433,0.43781818224776,0.4479307793476,0.45823070152614,0.46872463649679,0.47941971427274,0.49032353757851,0.50144421473605,0.51279039527177,0.52437130852031,0.53619680553253,0.54827740463258,0.56062434101017,0.57324962078195,0.58616608000982,0.59938744922579,0.61292842408364,0.62680474283878,0.64103327145057,0.65563209720868,0.67062063190864,0.68601972574487,0.70185179325524,0.71814095284413,0.73491318163551,0.75219648767026,0.77002110176951,0.78841969174729,0.80742760208188,0.82708312265889,0.84742779079644,0.86850673147378,0.89036904153293,0.91306822464037,0.9366626850185,0.96121628943351,0.98679900871668,1.0])
#File names
self.InitCondFile = "gadget_ic"
self.OutputDir = "snapshots"
self.EnergyFile = "energy.txt"
self.InfoFile = "info.txt"
self.TimingsFile = "timings.txt"
self.CpuFile = "cpu.txt"
self.RestartFile = "restart"
self.SnapshotFileBase = "snapshot"
self.OutputListFilename = "outputs.txt"
#CPU Timings
self.TimeLimitCPU = 1.0*u.day
self.ResubmitOn = 0
self.ResubmitCommand = "my-scriptfile"
#Code options
self.ICFormat = 1
self.SnapFormat = 1
self.ComovingIntegrationOn = 1
self.TypeOfTimestepCriterion = 0
self.OutputListOn = 1
self.PeriodicBoundariesOn = 1
#Caracteristics of run
self.TimeMax = 1.0
#Output frequency
self.TimeBetSnapshot = 0.5
self.TimeOfFirstSnapshot = 0
self.CpuTimeBetRestartFile = 12.5*u.hour
self.TimeBetStatistics = 0.05
self.NumFilesPerSnapshot = 16
self.NumFilesWrittenInParallel = 8
#Accuracy of time integration
self.ErrTolIntAccuracy = 0.025
self.MaxRMSDisplacementFac = 0.2
self.CourantFac = 0.15
self.MaxSizeTimestep = 0.02
self.MinSizeTimestep = 0.0
#Tree algorithm, force accuracy, domain update frequency
self.ErrTolTheta = 0.45
self.TypeOfOpeningCriterion = 1
self.ErrTolForceAcc = 0.005
self.TreeDomainUpdateFrequency = 0.025
#Further parameters of SPH
self.DesNumNgb = 33
self.MaxNumNgbDeviation = 2
self.ArtBulkViscConst = 0.8
self.InitGasTemp = 1000.0
self.MinGasTemp = 50.0
#Memory allocation
self.PartAllocFactor = 1.3
self.TreeAllocFactor = 0.7
self.BufferSize = 20*u.Mbyte
#System of units
self.UnitLength_in_cm = 3.085678e21 # ; 1.0 kpc
self.UnitMass_in_g = 1.989e43 #; 1.0e10 solar masses
self.UnitVelocity_in_cm_per_s = 1.0e5 # ; 1 km/sec
self.GravityConstantInternal = 0
#Softening lengths
self.MinGasHsmlFractional = 0.25
self.SofteningGas = 0
self.SofteningHalo = 9.000000
self.SofteningDisk = 0
self.SofteningBulge = 0
self.SofteningStars = 0
self.SofteningBndry = 0
self.SofteningGasMaxPhys = 0
self.SofteningHaloMaxPhys = 9.000000
self.SofteningDiskMaxPhys = 0
self.SofteningBulgeMaxPhys = 0
self.SofteningStarsMaxPhys = 0
self.SofteningBndryMaxPhys = 0
#Allow for kwargs override
for key in kwargs.keys():
setattr(self,key,kwargs[key])
def sections(self):
return [ "file_names","cpu_timings","code_options","characteristics_of_run","output_frequency","accuracy_time_integration","tree_algorithm","sph","memory_allocation","system_of_units","softening" ]
def showSection(self,section):
if section not in self.sections():
raise ValueError("Parameter file does not admit a section named {0}".format(section))
for option in getattr(self,section):
print("{0} = {1}".format(option,getattr(self,option)))
def show(self):
for section in self.sections():
print(section+":\n")
self.showSection(section)
print("\n")
[docs] def writeSection(self,section):
"""
Writes the corresponding section of the Gadget2 parameter file
"""
output = StringIO()
#Write preamble
output.write("% {0}\n\n".format(section))
#Cycle through options
for option in getattr(self,section):
#Read the corresponding value
value = getattr(self,option)
#Convert units as necessary
if type(value)==u.quantity.Quantity:
if value.unit.physical_type=="time":
value = value.to(u.s).value
elif value.unit.physical_type=="speed":
value = value.to(u.cm/u.s).value
elif "byte" in value.unit.to_string():
value = value.to(u.Mbyte).value
#Write the line
output.write("{0} {1}\n".format(option,value))
#Finish
output.write("\n\n")
output.seek(0)
return output.read()
@classmethod
[docs] def default(cls):
"""
Generate default settings
"""
return cls()
############################################################
################Gadget2Header class#########################
############################################################
class Gadget2Header(dict):
"""
Class handler of a Gadget2 snapshot header
"""
def __init__(self,HeaderDict=dict()):
super(Gadget2Header,self).__init__()
for key in HeaderDict.keys():
self[key] = HeaderDict[key]
def __repr__(self):
keys = self.keys()
keys.sort()
return "\n".join([ "{0} : {1}".format(key,self[key]) for key in keys ])
def __add__(self,rhs):
assert isinstance(rhs,Gadget2Header),"addition not defined if rhs is not a Gadget2Header!"
#Check that it makes sense to add the snapshots (cosmological parameters, box size, time and redshift must agree)
fields_to_match = ["Ode0","Om0","h","w0","wa","box_size","endianness","flag_cooling","flag_feedback","flag_sfr","num_files"]
fields_to_match += ["num_particles_total","num_particles_total_gas","num_particles_total_side","num_particles_total_with_mass","redshift","scale_factor"]
for field in fields_to_match:
assert self[field] == rhs[field],"{0} fields do not match!".format(field)
assert np.all(self["masses"]==rhs["masses"])
assert np.all(self["num_particles_total_of_type"]==rhs["num_particles_total_of_type"])
#Construct the header of the merged snapshot
merged_header = self.copy()
merged_header["files"] += rhs["files"]
merged_header["num_particles_file"] += rhs["num_particles_file"]
merged_header["num_particles_file_gas"] += rhs["num_particles_file_gas"]
merged_header["num_particles_file_of_type"] += rhs["num_particles_file_of_type"]
merged_header["num_particles_file_with_mass"] += rhs["num_particles_file_with_mass"]
return merged_header
##############################################################
#################Gadget2Snapshot class######################
##############################################################
[docs]class Gadget2Snapshot(NbodySnapshot):
"""
A class that handles Gadget2 snapshots, mainly I/O from the binary format and spatial information statistics.Inherits from the abstract NbodySnapshot
"""
###############################################################################################
#########################Abstract methods implementation#######################################
###############################################################################################
@classmethod
def buildFilename(cls,root,pool,**kwargs):
if pool is not None:
return root+".{0}".format(pool.rank)
else:
return root
@classmethod
def int2root(cls,name,n):
return name + "_{0:03d}".format(n)
############################################################################################
def getHeader(self):
header = Gadget2Header(ext._gadget2.getHeader(self.fp))
header["files"] = [self.fp.name]
header["w0"] = -1.0
header["wa"] = 0.0
header["comoving_distance"] = LambdaCDM(H0=header["h"]*100,Om0=header["Om0"],Ode0=header["Ode0"]).comoving_distance(header["redshift"]).to(u.kpc).value * header["h"]
return header
############################################################################################
def setLimits(self):
self._first = None
self._last = None
############################################################################################
[docs] def getPositions(self,first=None,last=None,save=True):
"""
Reads in the particles positions (read in of a subset is allowed): when first and last are specified, the numpy array convention is followed (i.e. getPositions(first=a,last=b)=getPositions()[a:b])
:param first: first particle in the file to be read, if None 0 is assumed
:type first: int. or None
:param last: last particle in the file to be read, if None the total number of particles is assumed
:type last: int. or None
:param save: if True saves the particles positions as attribute
:type save: bool.
:returns: numpy array with the particle positions
"""
#The file must not be closed
assert not self.fp.closed
#Particles do not have structure
self.weights = None
self.virial_radius = None
self.concentration = None
numPart = self._header["num_particles_file"]
#Calculate the offset from the beginning of the file: 4 bytes (endianness) + 256 bytes (header) + 8 bytes (void)
offset = 4 + 256 + 8
#If first is specified, offset the file pointer by that amount
if first is not None:
assert first>=0
offset += 4 * 3 * first
numPart -= first
if last is not None:
if first is not None:
assert last>=first and last<=self._header["num_particles_file"]
numPart = last - first
else:
assert last<=self._header["num_particles_file"]
numPart = last
#Read in the particles positions and return the corresponding array
try:
positions = (ext._gadget2.getPosVel(self.fp,offset,numPart) * self.kpc_over_h).to(self.Mpc_over_h)
except AttributeError:
positions = ext._gadget2.getPosVel(self.fp,offset,numPart) * u.kpc
if save:
self.positions = positions
return self.positions
#Return
return positions
############################################################################################
[docs] def getVelocities(self,first=None,last=None,save=True):
"""
Reads in the particles velocities (read in of a subset is allowed): when first and last are specified, the numpy array convention is followed (i.e. getVelocities(first=a,last=b)=getVelocities()[a:b])
:param first: first particle in the file to be read, if None 0 is assumed
:type first: int. or None
:param last: last particle in the file to be read, if None the total number of particles is assumed
:type last: int. or None
:param save: if True saves the particles velocities as attrubute
:type save: bool.
:returns: numpy array with the particle velocities
"""
assert not self.fp.closed
numPart = self._header["num_particles_file"]
#Calculate the offset from the beginning of the file: 4 bytes (endianness) + 256 bytes (header) + 8 bytes (void)
offset = 4 + 256 + 8
#Skip all the particle positions
offset += 4 * 3 * numPart
#Skip other 8 void bytes
offset += 8
#If first is specified, offset the file pointer by that amount
if first is not None:
assert first>=0
offset += 4 * 3 * first
numPart -= first
if last is not None:
if first is not None:
assert last>=first and last<=self._header["num_particles_file"]
numPart = last - first
else:
assert last<=self._header["num_particles_file"]
numPart = last
#Read in the particles positions and return the corresponding array
velocities = ext._gadget2.getPosVel(self.fp,offset,numPart)
#Scale units
velocities *= self._velocity_unit
velocities *= u.cm / u.s
if save:
self.velocities = velocities
return self.velocities
#Return
return velocities
############################################################################################
[docs] def getID(self,first=None,last=None,save=True):
"""
Reads in the particles IDs, 4 byte ints, (read in of a subset is allowed): when first and last are specified, the numpy array convention is followed (i.e. getID(first=a,last=b)=getID()[a:b])
:param first: first particle in the file to be read, if None 0 is assumed
:type first: int. or None
:param last: last particle in the file to be read, if None the total number of particles is assumed
:type last: int. or None
:param save: if True saves the particles IDs as attribute
:type save: bool.
:returns: numpy array with the particle IDs
"""
assert not self.fp.closed
numPart = self._header["num_particles_file"]
#Calculate the offset from the beginning of the file: 4 bytes (endianness) + 256 bytes (header) + 8 bytes (void)
offset = 4 + 256 + 8
#Skip all the particle positions
offset += 4 * 3 * numPart
#Skip other 8 void bytes
offset += 8
#Skip all the particle velocities
offset += 4 * 3 * numPart
#Skip other 8 void bytes
offset += 8
#If first is specified, offset the file pointer by that amount
if first is not None:
assert first>=0
offset += 4 * first
numPart -= first
if last is not None:
if first is not None:
assert last>=first and last<=self._header["num_particles_file"]
numPart = last - first
else:
assert last<=self._header["num_particles_file"]
numPart = last
#Read in the particles positions and return the corresponding array
ids = ext._gadget2.getID(self.fp,offset,numPart)
if save:
self.id = ids
return self.id
#Return
return ids
############################################################################################
[docs] def write(self,filename,files=1):
"""
Writes particles information (positions, velocities, etc...) to a properly formatter Gadget snapshot
:param filename: name of the file to which to write the snapshot
:type filename: str.
:param files: number of files on which to split the writing of the snapshot (useful if the number of particles is large); if > 1 the extension ".n" is appended to the filename
:type files: int.
"""
#Sanity checks
assert hasattr(self,"positions"),"Positions must be specified!!"
assert self.positions.shape[1]==3
if not hasattr(self,"_header"):
self.setHeaderInfo()
#Build a bare header based on the available info (need to convert units back to the Gadget ones)
_header_bare = self._header.copy()
_header_bare["box_size"] = _header_bare["box_size"].to(self.kpc_over_h).value
_header_bare["masses"] = _header_bare["masses"].to(u.g).value * _header_bare["h"] / self._mass_unit
_header_bare["num_particles_file_of_type"] = _header_bare["num_particles_file_of_type"].astype(np.int32)
_header_bare["num_particles_total_of_type"] = _header_bare["num_particles_total_of_type"].astype(np.int32)
_header_bare["comoving_distance"] = _header_bare["comoving_distance"].to(self.Mpc_over_h).value * 1.0e3
#Convert units for positions and velocities
_positions_converted = self.positions.to(self.kpc_over_h).value.astype(np.float32)
if hasattr(self,"velocities"):
assert self.positions.shape==self.velocities.shape
_velocities_converted = (self.velocities.to(u.cm/u.s).value / self._velocity_unit).astype(np.float32)
writeVel = 1
else:
_velocities_converted = np.zeros((1,3),dtype=np.float32)
writeVel = 0
#Check if we want to split on multiple files (only DM particles supported so far for this feature)
if files>1:
#Number of files the snapshot is split into
_header_bare["num_files"] = files
#Update the header with the file names
self.header["files"] = [ "{0}.{1}".format(filename,n) for n in range(files) ]
#Distribute particles among files
particles_per_file = _header_bare["num_particles_total"] // files
#Write each file
for n in range(files - 1):
#Update header
_header_bare["num_particles_file"] = particles_per_file
#TODO all particles are DM, fix distribution in the future
_header_bare["num_particles_file_of_type"] = np.array([0,particles_per_file,0,0,0,0],dtype=np.int32)
#Write it!
filename_with_extension = "{0}.{1}".format(filename,n)
ext._gadget2.write(_header_bare,_positions_converted[n*particles_per_file:(n+1)*particles_per_file],_velocities_converted[n*particles_per_file:(n+1)*particles_per_file],n*particles_per_file+1,filename_with_extension,writeVel)
#The last file might have a different number of particles
particles_last_file = len(_positions_converted[particles_per_file*(files-1):])
#Update header
_header_bare["num_particles_file"] = particles_last_file
#TODO all particles are DM, fix distribution in the future
_header_bare["num_particles_file_of_type"] = np.array([0,particles_last_file,0,0,0,0],dtype=np.int32)
#Write it!
filename_with_extension = "{0}.{1}".format(filename,files-1)
ext._gadget2.write(_header_bare,_positions_converted[particles_per_file*(files-1):],_velocities_converted[particles_per_file*(files-1):],(files-1)*particles_per_file+1,filename_with_extension,writeVel)
else:
#Update the num_files key only if not present already
if "num_files" not in _header_bare.keys():
_header_bare["num_files"] = 1
#Update the header with the file names
self.header["files"] = [ filename ]
#Write it!!
ext._gadget2.write(_header_bare,_positions_converted,_velocities_converted,1,filename,writeVel)
############################################################################################
###########################Extra methods####################################################
############################################################################################
[docs] def writeParameterFile(self,filename,settings):
"""
Writes a Gadget2 parameter file to evolve the current snapshot using Gadget2
:param filename: name of the file to which to write the parameters
:type filename: str.
:param settings: tunable settings of Gadget2 (see Gadget2 manual)
:type settings: Gadget2Settings
"""
#Create output directory if not existent already
outputdir = settings.OutputDir
if not(os.path.isdir(outputdir)):
os.mkdir(outputdir)
#Update the settings according to the physical units of the current snapshot
settings.UnitLength_in_cm = self._length_unit
settings.UnitMass_in_g = self._mass_unit
settings.UnitVelocity_in_cm_per_s = self._velocity_unit
#Set the appropriate name for the initial condition file
if "files" in self.header.keys():
suffix = self.header["files"][0].split(".")[-1]
settings.InitCondFile = os.path.abspath(self.header["files"][0].rstrip(".{0}".format(suffix)))
#Write the options
with open(filename,"w") as paramfile:
#Filenames section
paramfile.write(settings.writeSection("file_names"))
#CPU time limit section
paramfile.write(settings.writeSection("cpu_timings"))
#Code options section
paramfile.write(settings.writeSection("code_options"))
#Initial scale factor time
paramfile.write("{0} {1}\n".format("TimeBegin",self.header["scale_factor"]))
#Characteristics of run section
paramfile.write(settings.writeSection("characteristics_of_run"))
#Cosmological parameters
paramfile.write("{0} {1}\n".format("Omega0",self.header["Om0"]))
paramfile.write("{0} {1}\n".format("OmegaLambda",self.header["Ode0"]))
paramfile.write("{0} {1}\n".format("OmegaBaryon",0.046))
paramfile.write("{0} {1}\n".format("HubbleParam",self.header["h"]))
paramfile.write("{0} {1}\n".format("BoxSize",self.header["box_size"].to(self.kpc_over_h).value))
paramfile.write("{0} {1}\n".format("w0",self.header["w0"]))
paramfile.write("{0} {1}\n\n".format("wa",self.header["wa"]))
#Output frequency section
paramfile.write(settings.writeSection("output_frequency"))
#Accuracy of time integration section
paramfile.write(settings.writeSection("accuracy_time_integration"))
#Tree algorithm section
paramfile.write(settings.writeSection("tree_algorithm"))
#SPH section
paramfile.write(settings.writeSection("sph"))
#Memory allocation section
paramfile.write(settings.writeSection("memory_allocation"))
#System of units section
paramfile.write(settings.writeSection("system_of_units"))
#Softening lengths section
paramfile.write(settings.writeSection("softening"))
##############################################################
#################Gadget2SnapshotDE class########################
##############################################################
[docs]class Gadget2SnapshotDE(Gadget2Snapshot):
"""
A class that handles Gadget2 snapshots, mainly I/O from the binary format and spatial information statistics.Inherits from Gadget2Snapshot; assumes that the header includes Dark Energy information
"""
def getHeader(self):
header = Gadget2Header(ext._gadget2.getHeader(self.fp))
header["files"] = [self.fp.name]
return header
##################################################################
#################Gadget2SnapshotPipe class########################
##################################################################
[docs]class Gadget2SnapshotPipe(Gadget2SnapshotDE):
"""
Read in the particle positions when calling the constructor, without calling fseek
"""
def __init__(self,*args,**kwargs):
#Call parent constructor
super(Gadget2SnapshotPipe,self).__init__(*args,**kwargs)
#Read in the positions
npart = self.header["num_particles_file"]
self.fp.read(8)
try:
self.positions = (np.fromstring(self.fp.read(4*3*npart),dtype=np.float32).reshape(npart,3) * self.kpc_over_h).to(self.Mpc_over_h)
except AttributeError:
pass
#Read the rest
self.fp.read()
#Particles do not have structure
self.weights = None
self.virial_radius = None
self.concentration = None