#! /usr/bin/env python
# Copyright (C) 2013 ETH Zurich, Institute of Astronomy
# System imports
from __future__ import print_function, division, absolute_import, unicode_literals
import os
import ConfigParser
import time as time_mod
import random
import string
import numpy as np
from scipy import interpolate
from os import path
import logging
import datetime
import sys
# External modules
import PyCosmo
import shutil
import tempfile
# cpod imports
CPOD_CONFIG_NAME = "cpod.ini"
CPOD_MUSIC_SECTION = "music"
CPOD_GADGET2_SECTION = "gadget2"
CPOD_ROCKSTAR_SECTION = "rockstar"
CPOD_EXEC_OPTION = "exec"
CPOD_WORKSPACE_SECTION = "workspace"
CPOD_PREFIX_OPTION = "local_base_path"
REMOTE_BASE_PATH_OPTION = "remote_base_path"
MUSIC_SEED_SECTION = "random"
MUSIC_SEED_OPTION = "seed[8]"
OUTPUT_DATA_FOLDER = "data"
CONFIG_FILE_NAME='ics_MUSIC.conf'
DEFAULT_ROCKSTAR_CONFIG_NAME = 'rockstar_confg.cfg'
[docs]class SimManage(object):
"""
This class has been created to house functions that are
used to manage and manipulate N-body simulations that
used for the mass mapping projects.
"""
DELTA_CONFIG_FILE_NAME = "deltaConfig.conf"
def __init__(self, cpodConfPath):
"""
Initialises an instance of SimManage.
"""
self.cpodConfPath = cpodConfPath
if not os.path.exists(cpodConfPath):
raise Exception("Config not found in: '%s'" % cpodConfPath)
self.parser = ConfigParser.ConfigParser()
self.parser.read(cpodConfPath)
print('Welcome to the cPod simulation manager.\n')
self.music_exec = self.parser.get(CPOD_MUSIC_SECTION, CPOD_EXEC_OPTION)
print('The path to the MUSIC excutable has been set to: %s' % self.music_exec)
self.gadget2_exec = self.parser.get(CPOD_GADGET2_SECTION, CPOD_EXEC_OPTION)
print('The path to the gadget2_exec excutable has been set to: %s' % self.gadget2_exec)
self.rockstar_exec = self.parser.get(CPOD_ROCKSTAR_SECTION, CPOD_EXEC_OPTION)
print('The path to the rockstar_exec excutable has been set to: %s' % self.rockstar_exec)
self.remote_base_path = self.parser.get(CPOD_WORKSPACE_SECTION, REMOTE_BASE_PATH_OPTION)
#path_temp = path.dirname(__file__)+'/templates/pycosmo_template.ini'
path_templates = path.join(path.dirname(__file__), 'templates/')
self.templates_dir = path_templates
#sys.path.append(self.templates_dir)
print('The path to the templates folder has been set to: %s' % self.templates_dir)
self.music_param = None
self.gadget_param = None
# self.rockstar_param = None
self.output_path = None
[docs] def setup(self, deltaConfig=None, seed=None):
"""
Sets up the instance by creating the output directory structure
:param root_output_dir: the root of the desired output directory
:param deltaConfig: (optional) the delta config to use
:param seed: (optional) the seed to be used
"""
self.root_output_dir = self.parser.get(CPOD_WORKSPACE_SECTION, CPOD_PREFIX_OPTION)
if len(self.root_output_dir)==0:
print("local_base_path is empty using a temp directory")
self.root_output_dir = tempfile.mkdtemp()
self.deltaConfig = deltaConfig
self.seed = seed
logging.info("Delta config is:" + str(deltaConfig))
logging.info("seed in is:" + str(seed))
logging.info("Setting up dir structure")
output_data_path = os.path.join(self.root_output_dir, OUTPUT_DATA_FOLDER)
logging.info("Output directory is: %s" % output_data_path)
self.output_path = self._prepare_output(output_data_path=output_data_path)
# A1 - setup for MUSIC
music_outfile = self.mk_music_config(output_path=self.output_path, configin=deltaConfig, random_seed=seed)
# A2 - setup for gadget2
gadget_outdir, gadget_param_file, _ = self.mk_gadget2_config(self.output_path)
# A3 - setup for Rockstar
num_snaps = 50
rockstar_outfile = self.mk_rockstar_config(dir_in=self.output_path, dir_out=self.remote_base_path,
filename="snapshot_<snap>", num_snaps=num_snaps, startingsnap=0)
deltaConfigPath = os.path.join(self.output_path, self.DELTA_CONFIG_FILE_NAME)
if deltaConfig is not None:
shutil.copy(deltaConfig, deltaConfigPath)
else:
open(deltaConfigPath, 'a').close()
with open(self.output_path + 'setup_summary.dat', 'w') as f:
f.write(self.output_path + music_outfile + '\n')
f.write(self.output_path + gadget_param_file + '\n')
f.write(rockstar_outfile + '\n')
self.music_param = os.path.join(self.output_path, music_outfile)
self.gadget_param = os.path.join(self.output_path, gadget_param_file)
self.rockstar_param = rockstar_outfile
def _prepare_output(self, output_data_path=None, dir_new=True):
"""
Sets up the output directory for further processing
:param output_data_path: (optional) directory for storing data
:param dir_new: (optional) if True (default) a new directory will be created inside output_data_path
"""
if output_data_path is None:
output_data_path = os.path.normpath(path.dirname(__file__)+'/../../data/')+'/'
print('output_data_path set to:\n')
print(output_data_path)
if dir_new is True: # if true, need to
today = datetime.date.today()
time_now = time_mod.localtime()
datestring = str(today.year).zfill(4)+str(today.month).zfill(2)+str(today.day).zfill(2)
timestring = str(time_now[3]).zfill(2) + str(time_now[4]).zfill(2) + str(time_now[5]).zfill(2)
random_str = random.choice(string.letters)
dir_temp = 'Sim_3D_'+datestring+'_'+timestring+'_'+random_str
output_path = path.join(output_data_path, dir_temp)+ "/"
os.makedirs(output_path)
else:
output_path = output_data_path
shutil.copy(self.cpodConfPath, os.path.join(output_path, CPOD_CONFIG_NAME))
return output_path
[docs] def mk_music_config(self, output_path, outfile=CONFIG_FILE_NAME, configin=None, template_conf_file=None,
random_seed=None, ic_in_configdir=True):
"""
This function is designed to automatically generate the configuration file need to run MUSIC (IC generator)
:param output_path: path to the directory where the music configuration should be created
:param outfile: (optional) name of the output config file
:param configin: (optional) name of a file that contains entries to overwrite default values in template
:param template_conf_file: (optional) template config file to be used
:param random_seed: (optional) random seed to use
:param ic_in_configdir: (optional) if True (default) the output directory for ic files will be same as the config files
:return: the name of the MUSIC config file
"""
config = ConfigParser.ConfigParser()
config.optionxform=str
if template_conf_file is None:
#setting up the default template
template_conf_file = self.templates_dir+'ics_MUSIC_template.conf'
print("Template config file is set to: '%s'" % template_conf_file)
config.read(template_conf_file) # reading the base template
if configin is not None:
config.read(configin) # overwrite some of the entries using input file
if random_seed is not None:
config.set(MUSIC_SEED_SECTION, MUSIC_SEED_OPTION, random_seed)
# if ic_in_configdir is True:
# config.set('output','filename',path.join(output_path, config.get('output','filename')))
#TODO make sure name is always the same for proper housekeeping!
with open(output_path+outfile,'w') as cfgfile:
config.write(cfgfile) ###!!!! NEED TO THINK ABOUT ONLY OUTPUTING WHAT IS IN THE TEMPLATE
cfgfile.close()
return outfile
[docs] def mk_gadget2_config(self,dir_in,ic_datafile='ics_gadget.dat',ic_confgfile='ics_MUSIC.conf',
list_a_file='outfile_a.dat',list_a_new=True,len_fac_out = 0.5,zmax_out = 2.0):
"""
This function is designed to automatically generate a gadget parameter file.
This input for this is a directory that contains both the ic_files and the config files
used to generate that ic data.
:param dir_in: directory containing ic data and ic config files
:param ic_datafile: name of the ic data file
:param ic_confgfile: name of the ic config file
:param list_a_file: file with the output a times
:param list_a_new: if True (default) the list_a_file is created
:param len_fac_out: ???
:param zmax_out: ???
:return: outdir,gadget_param_file,num_a_out
"""
#outdir = dir_in+'3D_boxes/'
outdir = '3D_boxes/'
print(dir_in)
print(ic_datafile)
ic_file_temp = ic_datafile
output_a_list = list_a_file
#ic_file_temp = dir_in+ic_datafile
#output_a_list = dir_in+list_a_file
gadget_param_file = 'gadget_input.param'
if not os.path.exists(path.join(dir_in, outdir)):
os.makedirs(dir_in+outdir)
ic_config = ConfigParser.ConfigParser()
ic_config.optionxform=str
ic_config.read(dir_in+ic_confgfile)
num_a_out = None
if list_a_new is True:
num_a_out = self.mk_list_a(dir_in+output_a_list, ic_config, len_fac = len_fac_out, zmax = zmax_out)
hubbleparam = float(ic_config.get('cosmology','H0'))/100.
timebegin = 1./(1.+float(ic_config.get('setup','zstart')))
timemax = 1.0
gadget_template = path.join(self.templates_dir, 'gadget_template.param')
with open(gadget_template,'r') as fparam_temp:
with open(dir_in+gadget_param_file,'w') as f:
softlen = float(ic_config.get('setup','boxlength'))/float(ic_config.get('random','cubesize'))/30.
f.write('% Variable set by Mass Mapping routines \n')
f.write('\n')
f.write('InitCondFile %s \n' %ic_file_temp)
f.write('OutputDir %s \n' %outdir)
f.write('OutputListFilename %s \n' %output_a_list)
f.write('TimeBegin %s \n' %timebegin)
f.write('TimeMax %s \n' %timemax)
f.write('Omega0 %s \n' %ic_config.get('cosmology','Omega_m'))
f.write('OmegaLambda %s \n' %ic_config.get('cosmology','Omega_L'))
f.write('OmegaBaryon %s \n' %ic_config.get('cosmology','Omega_b'))
f.write('HubbleParam %s \n' %hubbleparam)
f.write('BoxSize %s \n' %ic_config.get('setup','boxlength'))
f.write('SofteningHalo %s \n' %softlen)
f.write('SofteningHaloMaxPhys %s \n' %softlen)
f.write('\n')
f.write(fparam_temp.read())
return outdir,gadget_param_file,num_a_out
[docs] def mk_list_a(self,file_out,ic_config,len_fac = 0.5,zmax = 2.0):
"""
Function for making a list of a value that will be used output results of the
N-body simulations.
:param file_out: the output file where the a list will be stored
:param ic_configin: config instance used to generate initial conditions
:param len_fac: ???
:param zmax: ???
:return: num_a_out
"""
# pycosmo_template = self.templates_dir+'pycosmo_template.ini'
# print("PyCosmo template is set to: '%s'" % pycosmo_template)
amin = 1.0/(1.0+zmax)
#ic_config = ConfigParser.ConfigParser()
#ic_config.read(ic_config_file)
hub_param = str(float(ic_config.get('cosmology','H0'))/100.)
box_len = float(ic_config.get('setup','boxlength'))/float(hub_param) #units in Mpc
cosmo = PyCosmo.Cosmo(paramfile='cpod.templates.pycosmo_template_for_cpod_test')
# cosmo_config = PyCosmo.read_param(pycosmo_template)
# cosmo_config.set('cosmo_params','omega_m_0',ic_config.get('cosmology','Omega_m'))
# cosmo_config.set('cosmo_params','omega_l_0',ic_config.get('cosmology','Omega_L'))
# cosmo_config.set('cosmo_params','omega_b_0',ic_config.get('cosmology','Omega_b'))
# cosmo_config.set('cosmo_params','h',hub_param)
# cosmo_config.set('cosmo_params','n',ic_config.get('cosmology','nspec'))
cosmo.set(omega_m = float(ic_config.get('cosmology','Omega_m')))
cosmo.set(omega_l_in = float(ic_config.get('cosmology','Omega_L')))
cosmo.set(omega_b = float(ic_config.get('cosmology','Omega_b')))
cosmo.set(h = float(hub_param))
cosmo.set(n = float(ic_config.get('cosmology','nspec')))
dist_max = cosmo.background.dist_rad_a(amin) #units in Mpc
num_out = np.floor(dist_max/box_len/len_fac)
dist_int = np.linspace(0.,box_len*len_fac*num_out)
a_arr = np.linspace(1.0,amin,500)
dist_arr = cosmo.background.dist_rad_a(a_arr) #units in Mpc
f_int = interpolate.interp1d(dist_arr,a_arr)
a_out = f_int(dist_int)
np.savetxt(file_out,a_out[::-1])
num_a_out = len(a_out)
return num_a_out
[docs] def mk_rockstar_config(self,dir_in='temp',dir_out='temp',rockstar_config_file=DEFAULT_ROCKSTAR_CONFIG_NAME,filename="snapshot_<snap>",
num_snaps=50,startingsnap=0):
"""
function for making the config files need to run rock star
:param dir_in:
:param dir_out:
:param rockstar_config_file:
:param filename:
:param num_snaps:
:param startingsnap:
"""
head, tail = os.path.split(dir_in)
simulationName = os.path.basename(head)
dir_out_sim = os.path.join(dir_out,simulationName)
dir_out_3Dbox = os.path.join(dir_out_sim,'3D_boxes/')
dir_out_halo = os.path.join(dir_out_sim,'Halodata')
print('\n\n\n\n\ndir_out: ')
print(dir_in)
print(dir_out)
print(dir_out_sim)
print(dir_out_halo)
print(simulationName)
# if not os.path.exists(dir_out_sim):
# print('blub')
# os.mkdir(dir_out_sim)
#
# if not os.path.exists(dir_out_halo):
# os.mkdir(dir_out_halo)
with open(os.path.join(dir_in,rockstar_config_file),'w') as f:
f.write('#Rockstar Halo Finder\n')
f.write('#Parallel config file for multi-cpu, multi-snapshot halo finding\n')
f.write('#automatically generated using python package for cpod.\n')
f.write('\n')
f.write('# example of how to run: start Rockstar server as\n')
f.write('# ./rockstar -c parallel.cfg\n')
f.write('#Then launch the reading/analysis tasks with:\n')
f.write('# ./rockstar -c auto-rockstar.cfg\n')
f.write('\n')
f.write('FILE_FORMAT = "%s" \n' %"GADGET2")
f.write('\n')
f.write('INBASE = "%s"\n' %dir_out_3Dbox)
f.write('OUTBASE="%s"\n' %dir_out_halo)
f.write('\n')
f.write('FILENAME="%s"\n' %filename)
f.write('NUM_SNAPS=%i\n' %num_snaps)
f.write('STARTING_SNAP=%i\n' %startingsnap)
f.write('NUM_BLOCKS=1\n')
f.write('\n')
f.write('SCALE_NOW = 1\n')
f.write('GADGET_LENGTH_CONVERSION = 1\n')
f.write('GADGET_MASS_CONVERSION = 1e+10\n')
f.write('FORCE_RES = 0.001 #Force resolution of simulation, in Mpc/h\n')
f.write('\n')
# f.write('OUTPUT_FORMAT="%s"\n' %'BOTH')
# f.write('DELETE_BINARY_OUTPUT_AFTER_FINISHED=0\n')
# f.write('RESCALE_PARTICLE_MASS=0\n')
# f.write('\n')
#This specifies how many CPUs you want to analyze the particles:
#NUM_WRITERS = 8
f.write('PARALLEL_IO=1\n')
f.write('NUM_READERS=1\n')
f.write('NUM_WRITERS=16\n')
f.write('FORK_READERS_FROM_WRITERS=1\n')
f.write('FORK_PROCESSORS_PER_MACHINE=16\n')
# f.write('FULL_PARTICLE_CHUNKS = 16\n')
return os.path.join(dir_in,rockstar_config_file)
#FILE_FORMAT = "GADGET2"
[docs] def cpod_quicksetup(self):
"""
This function has been written to quickly set up the environment for
runing: MUSIC, Gadget2 and rockstar
"""
# A1 - setup for MUSIC
outdir = None
configin = None
outdir_main = self._prepare_output(output_data_path=outdir)
outfile = self.mk_music_config(configin=configin,output_path=outdir_main)
# A2 - setup for gadget2
outdir_3Dboxes,gadget_param_file,num_a_out = self.mk_gadget2_config(outdir_main)
# A3 - setup for Rockstar
#dir_in = outdir
#dir_out = outdir_temp
num_snaps=num_a_out #!!!NEED to AUTOMATE THIS!!
file_out = self.mk_rockstar_config(dir_in=outdir_main,dir_out=outdir_main,
filename="snapshot_<snap>",num_snaps=num_snaps,startingsnap=0) # changed dir_in from outdir_3Dboxes to outdir_main, but don't know what this quicksetup function is there for, so no clue whether that was a stupid idea...
f = open(outdir_main+'setup_summary.dat','w')
f.write(outdir_main+outfile +'\n')
f.write(outdir_main+gadget_param_file+'\n')
f.write(file_out+'\n')
f.close()
ff = open(outdir_main+'time_stamp.dat','w')
ff.write('Directory and config files created by the function cpod_quicksetup \n')
ff.write('')
ff.write(str(datetime.datetime.now()))
ff.close
return outdir_main