Skip to content

Commit

Permalink
Add files via upload
Browse files Browse the repository at this point in the history
  • Loading branch information
@JThornton2020
JThornton2020 authored Jul 5, 2017
1 parent 034c455 commit 086e1be
Showing 1 changed file with 156 additions and 22 deletions.
178 changes: 156 additions & 22 deletions sim_cluster.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@
# Python Package Importing #
# ------------------------------------- #

# Add a as a read in from the restart and crash function for Tyler's code!

# Importing Necessary System Packages
import sys, os, math
import numpy as np
Expand All @@ -11,6 +13,18 @@
from optparse import OptionParser
import glob

# Tyler's imports
import hashlib
import copy
import traceback
import signal

from time import gmtime
from time import mktime
from time import clock

from collections import defaultdict

# Importing cPickle/Pickle
try:
import cPickle as pickle
Expand Down Expand Up @@ -38,38 +52,48 @@
from amuse.couple import multiples

# Import the Tycho Packages
from tycho import create, util, read, write
from tycho import create, util, read, write, encounter_db

# ------------------------------------- #
# Main Production Script #
# ------------------------------------- #

# Tyler uses a start time variable for his naming

global Starting

Starting = mktime(gmtime())

if __name__=="__main__":

start_time = tp.time()

# Creating Command Line Argument Parser
parser = OptionParser()
parser.add_option("-g", "--no-gpu", dest="use_gpu",default=1, type="int", \
help="Disable GPU for computation by setting to a value other than 1.")
parser.add_option("-i", "--gpu-id", dest="gpu_ID", default= -1, type="int", \
help="Select which GPU to use by device ID.")
parser.add_option("-p", "--num-psys", dest="num_psys", default=32, type="int", \
help="Enter the number of planetary systems desired.")
help="Enter the number of planetary systems desired.")
parser.add_option("-s", "--num-stars", dest="num_stars", default=750, type="int", \
help="Enter the number of stars desired.")
help="Enter the number of stars desired.")
parser.add_option("-t", "--timestep", dest="dt", default=0.05, type="float", \
help="Enter the PH4 timestep in N-Body Units.")
help="Enter the PH4 timestep in N-Body Units.")
parser.add_option("-c", "--cluster-name", dest="cluster_name", default=None, type="str", \
help="Enter the name of the Cluster (Defaults to Numerical Naming Scheme).")
help="Enter the name of the Cluster (Defaults to Numerical Naming Scheme).")
parser.add_option("-w", "--w0", dest="w0", default=2.5, type="float", \
help="Enter the w0 parameter for the King's Model.")
help="Enter the w0 parameter for the King's Model.")
parser.add_option("-N", "--num-steps", dest="num_steps", default=1000, type="int", \
help="Enter the total number of time-steps to take.")
help="Enter the total number of time-steps to take.")
parser.add_option("-b", "--IBF", dest="IBF", default = 0.5, type ="float", \
help = "Enter the initial binary fraction.")
help = "Enter the initial binary fraction.")
parser.add_option("-S", "--seed", dest="seed", default = 1234, type="int", \
help = "Enter a random seed for the simulation")
parser.add_option("-R","--restart",dest="restart_file",default="_restart", type="str", \
help = "Enter the name for the restart_file, (Defaults to _restart.hdf5")
parser.add_option("-D", "--database", dest="database", default="cluster_db", type="str")

(options, args) = parser.parse_args()

# Set Commonly Used Python Variables from Options
Expand All @@ -78,10 +102,12 @@
cluster_name = options.cluster_name
restart_file = "Restart/"+cluster_name+options.restart_file
write_file_base = restart_file
database = options.database
crash_base = "CrashSave/"+cluster_name

# Try to Import the Cluster from File or Create a New One
try:
MasterSet, ic_array, converter = read.read_initial_state(options.cluster_name)
MasterSet, ic_array, converter = read.read_initial_state(cluster_name)
# TODO: Create the "stars" and "planets" sets and open up their channels to the master set.
# It should match the style designed below.
read_from_file = True
Expand All @@ -100,8 +126,15 @@
#channel_planets_master = systems_NB.new_channel_to(MasterSet)
read_from_file = False

# Create Initial Conditions Array
initial_conditions = util.store_ic(converter, options)
# Create Initial Conditions Array
initial_conditions = util.store_ic(converter, options)

try:
search = glob.glob("CrashSave/"+cluster_name+"*.hdf5")[-1]
crash_file = search[:-18]
crash = True
except:
crash = False

# Write the Initial State
if not read_from_file:
Expand All @@ -117,7 +150,7 @@
use_gpu = options.use_gpu
gpu_ID = options.gpu_ID

# Setting PH4 as the Top-Level Gravity Code
# Setting PH4 as the Top-Level Gravity Code
if use_gpu == 1:
gravity = ph4(number_of_workers = num_workers, redirection = "none", mode = "gpu")
#try:
Expand Down Expand Up @@ -157,7 +190,7 @@
util.init_smalln()

# Initializing MULTIPLES
if read_from_file:
if read_from_file and crash:
time, multiples_code = read.recover_crash(crash_file, gravity, kep, util.new_smalln)
else:
multiples_code = multiples.Multiples(gravity, util.new_smalln, kep)
Expand All @@ -177,14 +210,101 @@
sys.stdout = f

step_index = 0
# Dummy Variable so it doesn't restart at step 0
flag = 0

# This is where I added Tyler's stuff

# So far this puts one empty list, encounters, into the dictionary
# The particle list is made in the EncounterHandler class
# The particle list will contain all the information about the particles involved in an encounter

encounterInformation = defaultdict(list)

for star in MasterSet:
if star.type == "star":
name = str(star.id)
encounterInformation[name] = []



class EncounterHandler(object):
def __init__(self):
self.run_id = hashlib.sha256(str(Starting)).hexdigest()

# I need to change this from writing into Tyler's database files to writing the infomation into structured arrays

def handle_encounter_v2(self, kepler, conv, time, star1, star2):

encounters1 = []
encounters2 = []
particles1 = []
particles2 = []
name1 = str(star1.id)
name2 = str(star2.id)

M,a,e,r,E,t = multiples.get_component_binary_elements(star1, star2, kepler)
peri = abs(conv.to_si(a * (1.0 - e)).value_in(units.AU))
apo = abs(conv.to_si(a * (1.0 + e)).value_in(units.AU))

r_au = conv.to_si(r).value_in(units.AU)
real_time = conv.to_si(time).value_in(units.Myr)

# OrbitParams returns the peri and apo and convets all of the parameters to SI
orbit = encounter_db.OrbitParams.from_nbody_params(M, a, e, r, E, time, conv)

# from_particle returns the id, mass, radius, position, and velocity in SI
star1_params = encounter_db.EncounterBody.from_particle(star1, conv)
star2_params = encounter_db.EncounterBody.from_particle(star2, conv)

'''
Encounter returns:
[<Encounter None @ t=1.25509037737
peri=1.60014590665e+15 AU, r_init=3.44462289279e+15 m, ecc=3.32880142942
Body 82: <Body 82: mass=1.57566086421e+30 kg>
Body 1: <Body 1: mass=1.50990530471e+29 kg>
'''

encounter = encounter_db.Encounter([star1_params, star2_params], orbit, conv.to_si(time))

# Here I need to add all of this information into the structured arrays

# -------- #
# FLAG #
# -------- #

# First I put everything into the two particles list.

particles1.append(encounter)
particles1.append(star1_params)
particles1.append(star2_params)

particles2.append(encounter)
particles2.append(star2_params)
particles2.append(star1_params)

encounters1.append(particles1)
encounters2.append(particles2)

encounterInformation[name1].append(encounters1)
encounterInformation[name2].append(encounters2)

return True

# cluster_params = encounter_db.ClusterParameters(num_stars, end_time)
# db_writer = encounter_db.EncounterDbWriter(database, cluster_params)
encounters = EncounterHandler()

# Begin Evolving the Cluster
while time < end_time:
sys.stdout.flush()
time += delta_t
multiples_code.evolve_model(time)

def encounter_callback(time, s1, s2):
return encounters.handle_encounter_v2(kep, converter, time, s1, s2)


multiples_code.evolve_model(time, callback=encounter_callback)
#multiples_code.evolve_model(time)
gravity.synchronize_model()

# Copy values from the module to the set in memory.
Expand All @@ -203,9 +323,15 @@
# multi_systems = tree.get_tree_subset().copy_to_new_particles()
# CoMSet.add_particle(multi_systems)
write.write_time_step(MasterSet, converter, time, cluster_name)


# Write out a crash file every 50 steps
if step_index%50 == 0:
step = str(time.number)
crash_file = crash_base+"_t"+step
write.write_crash_save(time, MasterSet, gravity, multiples_code, crash_file)

# Write out the restart file and restart from it every 10 time steps
if step_index%10 == 0 and flag == 1:
if step_index%10 == 0:
step = str(time.number)
write_file=write_file_base+step
write.write_state_to_file(time, MasterSet, gravity, multiples_code, write_file)
Expand Down Expand Up @@ -254,11 +380,8 @@
print '\n [UPDATE] Reset at %s!' %(tp.strftime("%Y/%m/%d-%H:%M:%S", tp.gmtime()))
print '-------------'
sys.stdout.flush()
flag = 0

step_index += 1
if step_index%10 == 0:
flag = 1

# Log that a Step was Taken
print '-------------'
Expand All @@ -269,7 +392,18 @@
# Log that the simulation Ended & Switch to Terminal Output
print '[UPDATE] Run Finished at %s! \n' %(tp.strftime("%Y/%m/%d-%H:%M:%S", tp.gmtime()))
sys.stdout = orig_stdout
# f.close()
f.close()

# Tyler's finalize code should produce the json files and cluster file we need added here
# def finish(time, end_time, ex=None):
# db_writer.finalize(time, end_time, wall_time=time.time()-start_time, ex=ex)

# finish(time.number, end_time.number)

# Pickle the encounter information dictionary
encounter_file = open("Encounters/"+cluster_name+"_encounters.pkl", "wb")
pickle.dump(encounterInformation, encounter_file)
encounter_file.close()

# Alerts the Terminal User that the Run has Ended!
print '[UPDATE] Run Finished at %s! \n' %(tp.strftime("%Y/%m/%d-%H:%M:%S", tp.gmtime()))
Expand Down

0 comments on commit 086e1be

Please sign in to comment.