def set_periodic_checkpoint(sim, period):
"""
Set up periodic checkpoints of the simulation
The checkpoints are saved in openPMD format, in the directory
`./checkpoints`, with one subdirectory per process.
All the field and particle information of each processor is saved.
NB: Checkpoints are registered among the list of diagnostics
`diags` of the Simulation object `sim`.
Parameters
----------
sim: a Simulation object
The simulation that is to be saved in checkpoints
period: integer
The number of PIC iteration between each checkpoint.
"""
# Only processor 0 creates a directory where checkpoints will be stored
# Make sure that all processors wait until this directory is created
# (Use the global MPI communicator instead of the `BoundaryCommunicator`
# so that this still works in the case `use_all_ranks=False`)
if comm.rank == 0:
if os.path.exists('./checkpoints') is False:
os.mkdir('./checkpoints')
comm.Barrier()
# Choose the name of the directory: one directory per processor
write_dir = 'checkpoints/proc%d/' % comm.rank
# Register a periodic FieldDiagnostic in the diagnostics of the simulation
sim.diags.append(FieldDiagnostic(period, sim.fld, write_dir=write_dir))
# Register a periodic ParticleDiagnostic, which contains all
# the particles which are present in the simulation
particle_dict = {}
for i in range(len(sim.ptcl)):
particle_dict['species %d' % i] = sim.ptcl[i]
sim.diags.append(
ParticleDiagnostic(period, particle_dict, write_dir=write_dir))
def test_synchronize(self):
""" Make sure that we can make the overlap spaces accurate. """
for case in self.cases:
space.initialize_space(case['shape'])
data = np.random.randn(*case['shape']).astype(case['dtype'])
cpu_data = np.empty_like(data)
comm.Allreduce(data, cpu_data)
g = Grid(case['dtype'])
self.assertRaises(TypeError, g.synchronize) # No overlap.
# Test with-overlap cases as well.
for k in range(1, 4):
g = Grid(case['dtype'], x_overlap=k)
# Overwrite entire grid
data = np.random.randn(*case['shape']).astype(case['dtype'])
cpu_data = np.empty_like(data)
comm.Allreduce(data, cpu_data)
cpu_raw_bad = get_cpu_raw(cpu_data, k)
cpu_raw_bad[:k, :, :] += 1 # Mess up padding areas.
cpu_raw_bad[-k:, :, :] += 1
drv.memcpy_htod(g.data.ptr, cpu_raw_bad)
# Prove that the data is not synchronized at this time.
cpu_raw = get_cpu_raw(cpu_data, k)
xx = case['shape'][0]
gd = g._get_raw()
self.assertTrue((gd[:k, :, :] != cpu_raw[:k, :, :]).all())
self.assertTrue((gd[-k:, :, :] != cpu_raw[-k:, :, :]).all())
g.synchronize() # Synchronize the overlapping data.
# Make sure that the overlap data is accurate.
gd = g._get_raw()
self.assertTrue((gd[:k, :, :] == cpu_raw[:k, :, :]).all())
self.assertTrue((gd[-k:, :, :] == cpu_raw[-k:, :, :]).all())
comm.Barrier() # Wait for other mpi nodes to finish.
if isinstance(x, collections.Iterable):
return [a for i in x for a in flatten(i)]
else:
return [x]
rank = CW.Get_rank()
size = CW.Get_size()
two_col_width = 7.20472 #inches
single_col_width = 3.50394 #inches
page_height = 10.62472
font_size = 10
sys.stdout.flush()
CW.Barrier()
pickle_file = sys.argv[1]
true_birth_con_pickle = sys.argv[2]
plot_gradient = False
read_pickle = bool(sys.argv[3])
baseline_yr = float(sys.argv[4])
#plot_key = sys.argv[2]
plot_keys = ['System_semimajor'] #, 'System_ecc', 'System_energies']
sys.stdout.flush()
CW.Barrier()
#check_sub_sys = [[21, 22], [23,97], [23,24], [23,97], [72,76], [78,235]]
jump_time_start = []
jump_time_end = []
def __init__(self, array_or_dtype, x_overlap=0):
""" Create a spatial grid on the GPU(s).
Input variables
array_or_dtype -- can either be a numpy array of the same shape as
the global space, or a numpy dtype. If a valid array is passed,
it will be loaded on to the GPU. If a dtype is passed, then
an array of zeros, of that dtype will be loaded onto the GPU.
Optional variables
x_overlap -- the number of adjacent cells in either the negative or
positive x-direction that need to simultaneously be accessed along
with the current cell. Must be a non-negative integer. Default
value is 0.
"""
shape = get_space_info()['shape'] # Get the shape of the space.
xr = get_space_info()['x_range'] # Get the local x_range.
all_x_ranges = get_space_info()[
'all_x_ranges'] # Get the local x_range.
local_shape = (xr[1] - xr[0], shape[1], shape[2])
self._set_gce_type('grid') # Set the gce type to grid.
# Make sure overlap option is valid.
if type(x_overlap) is not int:
raise TypeError('x_overlap must be an integer.')
elif x_overlap < 0:
raise TypeError('x_overlap must be a non-negative integer.')
if comm.rank == 0:
# Process the array_or_dtype input variable.
if type(array_or_dtype) is np.ndarray: # Input is an array.
array = array_or_dtype
# Make sure the array is of the correct shape.
if array.shape != shape:
raise TypeError(
'Shape of array does not match shape of space.')
# Make sure the array is of a valid datatype.
self._get_dtype(array.dtype.type)
elif type(array_or_dtype) is type: # Input is a datatype.
self._get_dtype(array_or_dtype) # Validate the dtype.
array = np.zeros(shape,
dtype=self.dtype) # Make a zeros array.
else: # Invalid input.
raise TypeError(
'Input variable must be a numpy array or dtype')
# Prepare array to be scattered.
array = [array[r[0]:r[1], :, :] for r in all_x_ranges]
else:
array = None
array = comm.scatter(array)
self._get_dtype(array.dtype.type)
# # Narrow down the array to local x_range.
# array = array[xr[0]:xr[1],:,:]
# Add padding to array, if needed.
self._xlap = x_overlap
if self._xlap is not 0:
padding = np.empty((self._xlap, ) + shape[1:3], dtype=array.dtype)
array = np.concatenate((padding, array, padding), axis=0)
self.to_gpu(array) # Load onto device.
# Determine information needed for synchronization.
if self._xlap is not 0:
# Calculates the pointer to the x offset in a grid.
ptr_dx = lambda x_pos: self.data.ptr + self.data.dtype.itemsize * \
x_pos * shape[1] * shape[2]
# Pointers to different sections of the grid that are relevant
# for synchronization.
self._sync_ptrs = { 'forw_src': ptr_dx(xr[1]-xr[0]), \
'back_dest': ptr_dx(0), \
'back_src': ptr_dx(self._xlap), \
'forw_dest': ptr_dx(xr[1]-xr[0] + self._xlap)}
# Buffers used during synchronization.
self._sync_buffers = [drv.pagelocked_empty( \
(self._xlap, shape[1], shape[2]), \
self.dtype) for k in range(4)]
# Streams used during synchronization.
self._sync_streams = [drv.Stream() for k in range(4)]
# Used to identify neighboring MPI nodes with whom to synchronize.
self._sync_adj = get_space_info()['mpi_adj']
# Offset in bytes to the true start of the grid.
# This is used to "hide" overlap areas from the kernel.
self._xlap_offset = self.data.dtype.itemsize * \
self._xlap * shape[1] * shape[2]
self.synchronize() # Synchronize the grid.
comm.Barrier(
) # Wait for all grids to synchronize before proceeding.
def main():
rank = CW.Get_rank()
size = CW.Get_size()
args = parse_inputs()
n_orb = int(args.no_orbits)
n_systems = int(args.no_systems)
q_min = 0.05
my_orb = bo.random_orbits(n_orb=n_orb)
US_group_vel = 10.
UCL_group_vel = 4.
#Madsen, 2002 gives the STANDARD ERROR of the US and UCL velcs to be 1.3 and 1.9km/s
US_group_std = 1.3 * args.group_velocity_sigma #From Preibisch et al., 2008
UCL_group_std = 1.3 * args.group_velocity_sigma
standard_std = {'F': 1.08, 'G': 0.63, 'K': 1.43, 'M': 2.27} # 2.0
astrophysical_std = args.astrophysical_std #Astrophysical radial velocity uncertainty
Object = []
Region = []
IR_excess = []
Temp_sptype = []
Pref_template = []
Obs_info = []
all_bayes = [[], []]
RV_standard_info = {}
sys.stdout.flush()
CW.Barrier()
#Read in RV standard list
header = 0
with open('/home/100/rlk100/RV_standard_list.csv', 'rU') as f:
reader = csv.reader(f)
for row in reader:
if header != 0:
RV_standard_info[row[0]] = (float(row[5]), float(row[6]),
float(row[7]))
else:
header = 1
f.close()
sys.stdout.flush()
CW.Barrier()
print("Reading in current spreadsheet", args.input_file)
header = 0
reshape_len = -1
with open(args.input_file, 'rU') as f:
reader = csv.reader(f)
for row in reader:
if header != 0:
if 'U4' in row[0]:
row[0] = 'UCAC4' + row[0].split('U4')[-1]
Object.append(row[0])
Region.append(row[1])
IR_excess.append(row[5])
Pref_template.append(row[15]) #row[18])
Temp_sptype.append(row[16]) #row[19])
if len(row) > 17:
Obs = np.array(row[17:])
Obs = np.delete(Obs, np.where(Obs == ''))
if reshape_len == -1:
for ob in Obs:
reshape_len = reshape_len + 1
if '/' in ob and ob != Obs[0]:
break
#if len(Obs) > 5:
# Obs = np.reshape(Obs, (len(Obs)/5, 5))
Obs = np.reshape(Obs,
(len(Obs) / reshape_len, reshape_len))
for ind_obs in Obs:
if '/' in ind_obs[0]:
new_format = '20' + ind_obs[0].split('/')[
-1] + '-' + ind_obs[0].split('/')[-2] + '-' + (
"%02d" % int(ind_obs[0].split('/')[-3]))
ind_obs[0] = new_format
else:
Obs = np.array([])
Obs_info.append(Obs)
if header == 0:
header = 1
f.close()
del header
sys.stdout.flush()
CW.Barrier()
Obj_bayes = np.nan * np.zeros(len(Object))
#Read in currently calculated Bayes Factors:
if args.restart_calc != 'False':
print("Reading in calulated Bayes factors")
header = 0
with open(args.bayes_file, 'rU') as f:
reader = csv.reader(f)
for row in reader:
if header != 0:
ind = Object.index(row[0])
Obj_bayes[ind] = float(row[2])
if row[1] == 'US':
all_bayes[0].append(float(row[2]))
else:
all_bayes[1].append(float(row[2]))
del ind
else:
header = 1
f.close()
del header
sys.stdout.flush()
CW.Barrier()
if args.restart_calc != 'False' and rank == 0:
print("Creating new bayes file")
f = open(args.bayes_file, 'w')
f.write('Object,Region,Bayes_factor\n')
f.close()
sys.stdout.flush()
CW.Barrier()
inds = list(range(len(Object)))
skip_inds = np.where(np.array(IR_excess) == 'NN')[0]
for skit in skip_inds:
inds.remove(skit)
skip_inds = np.where(np.array(Pref_template) == '')[0]
for skit in skip_inds:
inds.remove(skit)
del skip_inds
del IR_excess
rit = 0
sys.stdout.flush()
CW.Barrier()
for obj in inds:
Pref_template_name = Pref_template[obj].split('_')[0]
if np.isnan(Obj_bayes[obj]) and rank == rit:
print("Doing object:", Object[obj], "on rank:", rank)
likelihoods = []
single_likelihoods = []
#Produces masses within +/- 10% of the mass of the template.
#!!! Mike suggests a single mass.
M_1 = (np.random.random(n_systems) *
(RV_standard_info[Pref_template_name][1] -
RV_standard_info[Pref_template_name][0])
) + RV_standard_info[Pref_template_name][0]
#Generates mass ratios with minium mass ratio of q_min (default 0.01?, should this be dependant on the primary mass? Because sometimes low mass ratios could give very low mass companions i.e. BD mass...)
#!!! Mike suggests 0.05 due to brown dwarf desert.
q = (np.random.random(n_systems) * (1 - q_min)) + q_min
#from Primary masses and mass ratios, secondary masses can get calculated
M_2 = M_1 * q
#Get dates of the observations of the object
jds = Obs_info[obj][:, 1].astype(np.float)
#get observed data, and add in the error in the standards in quadrature.
#This relates to the spectrograph stability
#There is also an astrophysical error due to these objects being rapid rotators etc.
RV_standard_err = standard_std[Temp_sptype[obj][0]]
err = np.sqrt(Obs_info[obj][:, 3].astype(float)**2 +
RV_standard_err**2 + astrophysical_std**2)
observed_rv = Obs_info[obj][:, 2].astype(float)
#IN A LOOP iterate over random orbits:
for orb in range(n_orb):
#FIXME: Figure out which velocity to use!
if Region[obj] == 'US':
if args.group_velocity == 'True':
v_group = np.random.normal(
US_group_vel,
np.sqrt(US_group_std**2 + RV_standard_err**2),
n_systems)
else:
v_group = np.random.normal(
np.mean(observed_rv),
np.sqrt(US_group_std**2 + RV_standard_err**2),
n_systems)
else:
if args.group_velocity == 'True':
v_group = np.random.normal(
UCL_group_vel,
np.sqrt(UCL_group_std**2 + RV_standard_err**2),
n_systems)
else:
v_group = np.random.normal(
np.mean(observed_rv),
np.sqrt(UCL_group_std**2 + RV_standard_err**2),
n_systems)
#generate orbit?
#!!! Find just one set of orbital parameters at at a time, and
#scale the RVS. OR if you really want you can compute a, i etc
#yourself and plug these into my_orb, but some RV scalign is still needed.
rho, theta, normalised_vr = bo.binary_orbit(my_orb,
jds,
plot_orbit_no=orb)
for system in range(n_systems):
actual_vr = bo.scale_rv(normalised_vr,
my_orb['P'][orb],
M_1[system],
M_2[system],
my_orb['i'][orb],
group_velocity=v_group[system])
this_likelihood = bo.calc_likelihood(
actual_vr, observed_rv, err)
likelihoods.append(this_likelihood)
#THEN CALCULATE PROBABILITY OF BEING A SINGLE STAR
single_likelihoods.append(
bo.calc_likelihood(v_group[system], observed_rv, err))
del actual_vr
del this_likelihood
del v_group
del M_1
del q
del M_2
del jds
del RV_standard_err
del err
del observed_rv
#THEN CALCULATE BAYES FACTOR
bayes_factor = np.mean(likelihoods) / np.mean(single_likelihoods)
print(("Bayes Factor: {0:5.2f} for ".format(bayes_factor) +
Object[obj]), "on rank", rank, "with SpT", Temp_sptype[obj])
del likelihoods
del single_likelihoods
if Region[obj] == 'US':
send_data = [0.0, float(obj), bayes_factor, Temp_sptype[obj]]
#print "Sending data:", send_data, "from rank:", rank
if rank == 0:
bayes_update = send_data
else:
CW.send(send_data, dest=0, tag=rank)
else:
send_data = [1.0, float(obj), bayes_factor, Temp_sptype[obj]]
#print "Sending data:", send_data, "from rank:", rank
if rank == 0:
bayes_update = send_data
else:
CW.send(send_data, dest=0, tag=rank)
del send_data
if rank == 0:
all_bayes[int(bayes_update[0])].append(bayes_update[2])
Obj_bayes[int(bayes_update[1])] = bayes_update[2]
print("Updated Bayes factors retrieved from rank 0 for object",
Object[int(bayes_update[1])])
f = open(args.bayes_file, 'a')
write_string = Object[int(bayes_update[1])] + ',' + Region[int(
bayes_update[1])] + ',' + str(bayes_update[2]) + ',' + str(
bayes_update[3]) + '\n'
f.write(write_string)
f.close()
del bayes_update
del write_string
rit = rit + 1
if rit == size:
sys.stdout.flush()
CW.Barrier()
rit = 0
if rank == 0:
print("UPDATING CALCULATED BAYES VALUES")
for orit in range(1, size):
bayes_update = CW.recv(source=orit, tag=orit)
all_bayes[int(bayes_update[0])].append(bayes_update[2])
Obj_bayes[int(bayes_update[1])] = bayes_update[2]
print("Updated Bayes factors retrieved from rank", orit,
"for object", Object[int(bayes_update[1])])
f = open(args.bayes_file, 'a')
write_string = Object[int(
bayes_update[1])] + ',' + Region[int(
bayes_update[1])] + ',' + str(
bayes_update[2]) + ',' + str(
bayes_update[3]) + '\n'
f.write(write_string)
f.close()
del bayes_update
del write_string
sys.stdout.flush()
CW.Barrier()
sys.stdout.flush()
CW.Barrier()
if rank == 0:
print("UPDATING CALCULATED BAYES VALUES")
for orit in range(1, size):
bayes_update = CW.recv(source=orit, tag=orit)
all_bayes[int(bayes_update[0])].append(bayes_update[2])
Obj_bayes[int(bayes_update[1])] = bayes_update[2]
print("Updated Bayes factors retrieved from rank", orit,
"for object", Object[int(bayes_update[1])])
f = open(args.bayes_file, 'a')
write_string = Object[int(bayes_update[1])] + ',' + Region[int(
bayes_update[1])] + ',' + str(bayes_update[2]) + ',' + str(
bayes_update[3]) + '\n'
f.write(write_string)
f.close()
del bayes_update
del write_string
sys.stdout.flush()
CW.Barrier()
print("Finished Calculating bayes factors!")
def restart_from_checkpoint(sim, iteration=None):
"""
Fills the Simulation object `sim` with data saved in a checkpoint.
More precisely, the following data from `sim` is overwritten:
- Current time and iteration number of the simulation
- Position of the boundaries of the simulation box
- Values of the field arrays
- Size and values of the particle arrays
Any other information (e.g. diagnostics of the simulation, presence of a
moving window, presence of a laser antenna, etc.) need to be set by hand.
For this reason, a successful restart will often require to modify the
original input script that produced the checkpoint, rather than to start
a new input script from scratch.
NB: This function should always be called *before* the initialization
of the moving window, since the moving window infers the position of
particle injection from the existing particle data.
Parameters
----------
sim: a Simulation object
The Simulation object into which the checkpoint should be loaded
iteration: integer (optional)
The iteration number of the checkpoint from which to restart
If None, the latest checkpoint available will be used.
"""
# Import openPMD-viewer
try:
from opmd_viewer import OpenPMDTimeSeries
except ImportError:
raise ImportError(
'The package `opmd_viewer` is required to restart from checkpoints.'
'\nPlease install it from https://github.com/openPMD/openPMD-viewer'
)
# Verify that the restart is valid (only for the first processor)
# (Use the global MPI communicator instead of the `BoundaryCommunicator`,
# so that this also works for `use_all_ranks=False`)
if comm.rank == 0:
check_restart(sim, iteration)
comm.Barrier()
# Choose the name of the directory from which to restart:
# one directory per processor
checkpoint_dir = 'checkpoints/proc%d/hdf5' % comm.rank
ts = OpenPMDTimeSeries(checkpoint_dir)
# Select the iteration, and its index
if iteration is None:
iteration = ts.iterations[-1]
i_iteration = ts.iterations.index(iteration)
# Modify parameters of the simulation
sim.iteration = iteration
sim.time = ts.t[i_iteration]
# Load the particles
# Loop through the different species
for i in range(len(sim.ptcl)):
name = 'species %d' % i
load_species(sim.ptcl[i], name, ts, iteration, sim.comm)
# Load the fields
# Loop through the different modes
for m in range(sim.fld.Nm):
# Load the fields E and B
for fieldtype in ['E', 'B', 'J']:
for coord in ['r', 't', 'z']:
load_fields(sim.fld.interp[m], fieldtype, coord, ts, iteration)
def execute(cfg, *args, **kwargs):
# Parse keyword arguments.
post_sync_grids = kwargs.get('post_sync', None)
# Parse the inputs.
gpu_params = []
for k in range(len(params)):
if params[k]['gce_type'] is 'number':
gpu_params.append(params[k]['dtype'](args[k]))
elif params[k]['gce_type'] is 'const': # Load Const.
gpu_params.append(args[k].data.ptr)
# Const no longer actually "const" in cuda code.
# d_ptr, size_in_bytes = my_get_global(params[k]['name'])
# drv.memcpy_dtod(d_ptr, args[k].data.gpudata, size_in_bytes)
elif params[k]['gce_type'] is 'grid':
if args[k]._xlap is 0:
gpu_params.append(args[k].data.ptr)
else:
gpu_params.append(args[k].data.ptr + \
args[k]._xlap_offset)
elif params[k]['gce_type'] is 'out':
args[k].data.fill(args[k].dtype(0)) # Initialize the Out.
gpu_params.append(args[k].data.ptr)
else:
raise TypeError('Invalid input type.')
# See if we need to synchronize grids after kernel execution.
if post_sync_grids is None:
sync_pad = 0
else:
sync_pad = max([g._xlap for g in post_sync_grids])
start2.record(stream)
comm.Barrier()
start.record(stream)
# Execute kernel in padded regions first.
execute_range(x_start, x_start + sync_pad, gpu_params, cfg, stream)
execute_range(x_end - sync_pad, x_end, gpu_params, cfg, stream)
pad_done.record(stream) # Just for timing purposes.
stream.synchronize() # Wait for execution to finish.
# Begin kernel execution in remaining "core" region.
execute_range(x_start + sync_pad, x_end - sync_pad, gpu_params,
cfg, stream)
comp_done.record(stream) # Timing only.
# While core kernel is executing, perform synchronization.
if post_sync_grids is not None: # Synchronization needed.
for grid in post_sync_grids:
grid.synchronize_start() # Start synchronization.
# Keep on checking until everything is done.
while not (all([grid.synchronize_isdone() \
for grid in post_sync_grids]) and \
stream.is_done()):
pass
else: # Nothing to synchronize.
stream.synchronize() # Just wait for execution to finish.
sync_done.record() # Timing.
# Obtain the result for all Outs.
batch_reduce(*[args[k] for k in range(len(params)) \
if params[k]['gce_type'] is 'out'])
all_done.record() # Timing.
all_done.synchronize()
return comp_done.time_since(
start) # Return time needed to execute the function.
if rank == 0:
solO = opticalpath(np.loadtxt("dbr.txt")[:, 1])
sol = neighb(solO, n, Neighb)
wyn = solO
cel_wyn = goal(wyn, w)
mkdir(path)
else:
sol = np.array([])
sol = mpi.bcast(sol, root=0)
j = 0
start = timeit.default_timer()
stop = start
end = 500
mpi.Barrier()
while (j <= end) & ((stop - start) < time_limit):
best_sol = np.array([])
good_sol = np.array([])
m = 0
e = 0
cel = np.array([])
period = n // size
full_range = np.linspace(0, size, size + 1) * period
if full_range[-1] < n:
full_range[-1] = n
tempsol = sol[int(full_range[rank]):int(full_range[rank + 1])]
cel = goal(tempsol, w)
def main():
rank = CW.Get_rank()
size = CW.Get_size()
args = parse_inputs()
prev_args = args
print("Starting mosaic_mod_script on rank", rank)
# Read in directories:
input_file = args.input_file
#save_dir = args.save_directory
#if os.path.exists(save_dir) == False:
# os.makedirs(save_dir)
# Read in input file
print("Reading in input mosaic file on rank", rank)
positions = []
paths = []
args_dict = []
with open(input_file, 'rU') as mosaic_file:
reader = csv.reader(mosaic_file)
for row in reader:
if row[0] == 'Grid_inputs:':
glr = float(row[1])
grl = float(row[2])
glw = float(row[3])
ghspace = float(row[4])
elif row[0][0] != '#':
positions.append((int(row[0]), int(row[1])))
paths.append(row[2])
dict = ""
for col in row[3:]:
dict = dict + col
if col != row[-1]:
dict = dict + ','
dict = ast.literal_eval(dict)
args_temp = argparse.Namespace(**vars(args))
for key in list(dict.keys()):
if key in args:
exec("args_temp."+ key + " = " + "str(dict[key])")
args_dict.append(args_temp)
del args_temp
args = prev_args
import pdb
pdb.set_trace()
positions = np.array(positions)
c = define_constants()
mym.set_global_font_size(args.text_font)
files = []
simfo = []
X = []
Y = []
X_vel = []
Y_vel = []
sim_files = []
L = None
for pit in range(len(paths)):
fs = get_files(paths[pit], args_dict[pit])
files.append(fs)
#print "paths =", paths
#print "fs =", fs
#print "args_dict =", args_dict
sfo = sim_info(paths[pit], fs[-1], args_dict[pit])
simfo.append(sfo)
if args_dict[pit].yt_proj == False:
x, y, x_vel, y_vel, cl = mym.initialise_grid(files[pit][-1], zoom_times=args_dict[pit].zoom_times)
X.append(x)
Y.append(y)
X_vel.append(x_vel)
Y_vel.append(y_vel)
else:
x = np.linspace(sfo['xmin'], sfo['xmax'], sfo['dimension'])
y = np.linspace(sfo['ymin'], sfo['ymax'], sfo['dimension'])
x, y = np.meshgrid(x, y)
annotate_space = (simfo[pit]['xmax'] - simfo[pit]['xmin'])/31.
x_ind = []
y_ind = []
counter = 0
while counter < 31:
val = annotate_space*counter + annotate_space/2. + simfo[pit]['xmin']
x_ind.append(int(val))
y_ind.append(int(val))
counter = counter + 1
x_vel, y_vel = np.meshgrid(x_ind, y_ind)
if args_dict[pit].projection_orientation != None:
y_val = 1./np.tan(np.deg2rad(float(args_dict[pit].projection_orientation)))
if np.isinf(y_val):
y_val = 0.0
L = [1.0, y_val, 0.0]
else:
if has_particles == False or len(dd['particle_posx']) == 1:
L = [0.0, 1.0, 0.0]
else:
pos_vec = [np.diff(dd['particle_posx'].value)[0], np.diff(dd['particle_posy'].value)[0]]
L = [-1*pos_vec[-1], pos_vec[0]]
L.append(0.0)
if L[0] > 0.0:
L = [-1.0*L[0], -1.0*L[1], 0.0]
print("SET PROJECTION ORIENTATION L=", L)
L = np.array(L)
X.append(x)
Y.append(y)
X_vel.append(x_vel)
Y_vel.append(y_vel)
if rank == 0:
print("shape of x, y", np.shape(x), np.shape(y))
if args_dict[pit].yt_proj == False and args_dict[pit].image_center != 0:
sim_fs = sorted(glob.glob(paths[pit] + 'WIND_hdf5_plt_cnt*'))
elif args_dict[pit].yt_proj != False and args_dict[pit].image_center != 0:
sim_fs = files
else:
sim_fs = []
sim_files.append(sim_fs)
#myf.set_normal(L)
#print "SET PROJECTION ORIENTATION L=", myf.get_normal()
# Initialise Grid and build lists
if args.plot_time != None:
m_times = [args.plot_time]
else:
m_times = mym.generate_frame_times(files[0], args.time_step, presink_frames=args.presink_frames, end_time=args.end_time)
no_frames = len(m_times)
m_times = m_times[args.start_frame:]
sys.stdout.flush()
CW.Barrier()
usable_files = []
usable_sim_files = []
for pit in range(len(paths)):
usable_fs = mym.find_files(m_times, files[pit])
usable_files.append(usable_fs)
if args_dict[pit].image_center != 0 and args_dict[pit].yt_proj == False:
usable_sfs = mym.find_files(m_times, sim_files[pit])
usable_sim_files.append(usable_fs)
del sim_files[pit]
else:
usable_sim_files.append([])
sys.stdout.flush()
CW.Barrier()
frames = list(range(args.start_frame, no_frames))
sink_form_time = []
for pit in range(len(paths)):
sink_form = mym.find_sink_formation_time(files[pit])
print("sink_form_time", sink_form_time)
sink_form_time.append(sink_form)
del files
# Define colourbar bounds
cbar_max = args.colourbar_max
cbar_min = args.colourbar_min
if L is None:
if args.axis == 'xy':
L = [0.0, 0.0, 1.0]
else:
L = [1.0, 0.0, 0.0]
L = np.array(L)
if args.axis == 'xy':
y_int = 1
else:
y_int = 2
sys.stdout.flush()
CW.Barrier()
rit = args.working_rank
for frame_val in range(len(frames)):
if rank == rit:
time_val = m_times[frame_val]
plt.clf()
columns = np.max(positions[:,0])
rows = np.max(positions[:,1])
width = float(columns)*(14.5/3.)
height = float(rows)*(17./4.)
fig =plt.figure(figsize=(width, height))
gs_left = gridspec.GridSpec(rows, columns-1)
gs_right = gridspec.GridSpec(rows, 1)
gs_left.update(right=glr, wspace=glw, hspace=ghspace)
gs_right.update(left=grl, hspace=ghspace)
axes_dict = {}
counter = 1
for pit in range(len(paths)):
try:
title_parts = args_dict[pit].title
except:
title_parts = args_dict[pit]['title']
title = ''
for part in title_parts:
if part != title_parts[-1]:
title = title + part + ' '
else:
title = title + part
ax_label = 'ax' + str(counter)
yit = np.where(positions[:,1] == positions[pit][1])[0][0]
if positions[pit][0] == 1 and positions[pit][1] == 1:
if columns > 1:
axes_dict.update({ax_label:fig.add_subplot(gs_left[0,0])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
else:
axes_dict.update({ax_label:fig.add_subplot(gs_right[0,0])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
elif positions[pit][0] != columns:
if args.share_x and args.share_y:
if yit >= len(axes_dict):
axes_dict.update({ax_label:fig.add_subplot(gs_left[positions[pit][1]-1,positions[pit][0]-1], sharex=axes_dict['ax1'])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
else:
axes_dict.update({ax_label:fig.add_subplot(gs_left[positions[pit][1]-1,positions[pit][0]-1], sharex=axes_dict['ax1'], sharey=axes_dict[list(axes_dict.keys())[yit]])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
elif args.share_x:
axes_dict.update({ax_label:fig.add_subplot(gs_left[positions[it][1]-1,positions[pit][0]-1], sharex=axes_dict['ax1'])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
elif args.share_y and positions[pit][0]!=1:
yit = np.where(positions[:,1] == positions[pit][1])[0][0]
axes_dict.update({ax_label:fig.add_subplot(gs_left[positions[pit][1]-1,positions[pit][0]-1], sharey=axes_dict[list(axes_dict.keys())[yit]])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
elif args.share_y:
axes_dict.update({ax_label:fig.add_subplot(gs_left[positions[pit][1]-1,positions[pit][0]-1])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
else:
axes_dict.update({ax_label:fig.add_subplot(gs_left[positions[pit][1]-1,positions[pit][0]-1])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
else:
if args.share_x and args.share_y:
yit = np.where(positions[:,1] == positions[pit][1])[0][0]
axes_dict.update({ax_label:fig.add_subplot(gs_right[positions[pit][1]-1,0], sharex=axes_dict['ax1'], sharey=axes_dict[list(axes_dict.keys())[yit]])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
elif args.share_x:
axes_dict.update({ax_label:fig.add_subplot(gs_right[positions[pit][1]-1,0], sharex=axes_dict['ax1'])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
elif args.share_y:
yit = np.where(positions[:,1] == positions[pit][1])[0][0]
axes_dict.update({ax_label:fig.add_subplot(gs_right[positions[pit][1]-1,0], sharey=axes_dict[list(axes_dict.keys())[yit]])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
else:
axes_dict.update({ax_label:fig.add_subplot(gs_right[positions[pit][1]-1,0])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
counter = counter + 1
axes_dict[ax_label].set(adjustable='box-forced', aspect='equal')
if args.yt_proj and args.plot_time==None and os.path.isfile(paths[pit] + "movie_frame_" + ("%06d" % frames[frame_val]) + ".pkl"):
pickle_file = paths[pit] + "movie_frame_" + ("%06d" % frames[frame_val]) + ".pkl"
print("USING PICKLED FILE:", pickle_file)
file = open(pickle_file, 'r')
#weight_fieldstuff = pickle.load(file)
X[pit], Y[pit], image, magx, magy, X_vel[pit], Y_vel[pit], velx, vely, part_info, args_dict[pit], simfo[pit] = pickle.load(file)
#file_time = stuff[17]
file.close()
else:
time_val = m_times[frame_val]
print("FILE =", usable_files[pit][frame_val])
has_particles = has_sinks(usable_files[pit][frame_val])
if has_particles:
part_info = mym.get_particle_data(usable_files[pit][frame_val], args_dict[pit].axis, proj_or=L)
else:
part_info = {}
center_vel = [0.0, 0.0, 0.0]
if args.image_center != 0 and has_particles:
original_positions = [X[pit], Y[pit], X_vel[pit], y_vel[pit]]
x_pos = np.round(part_info['particle_position'][0][args.image_center - 1]/cl)*cl
y_pos = np.round(part_info['particle_position'][1][args.image_center - 1]/cl)*cl
pos = np.array([part_info['particle_position'][0][args.image_center - 1], part_info['particle_position'][1][args.image_center - 1]])
X[pit] = X[pit] + x_pos
Y[pit] = Y[pit] + y_pos
X_vel[pit] = X_vel[pit] + x_pos
Y_vel[pit] = Y_vel[pit] + y_pos
if args.yt_proj == False:
sim_file = usable_sim_files[frame_val][:-12] + 'part' + usable_sim_files[frame_val][-5:]
else:
sim_file = part_file
if len(part_info['particle_mass']) == 1:
part_ind = 0
else:
min_dist = 1000.0
for part in range(len(part_info['particle_mass'])):
f = h5py.File(sim_file, 'r')
temp_pos = np.array([f[list(f.keys())[11]][part][13]/c['au'], f[list(f.keys())[11]][part][13+y_int]/c['au']])
f.close()
dist = np.sqrt(np.abs(np.diff((temp_pos - pos)**2)))[0]
if dist < min_dist:
min_dist = dist
part_ind = part
f = h5py.File(sim_file, 'r')
center_vel = [f[list(f.keys())[11]][part_ind][18], f[list(f.keys())[11]][part_ind][19], f[list(f.keys())[11]][part_ind][20]]
f.close()
xabel, yabel, xlim, ylim = image_properties(X[pit], Y[pit], args_dict[pit], simfo[pit])
if args_dict[pit].axis == 'xy':
center_vel=center_vel[:2]
else:
center_vel=center_vel[::2]
if args_dict[pit].ax_lim != None:
if has_particles and args_dict[pit].image_center != 0:
xlim = [-1*args_dict[pit].ax_lim + part_info['particle_position'][0][args_dict[pit].image_center - 1], args_dict[pit].ax_lim + part_info['particle_position'][0][args_dict[pit].image_center - 1]]
ylim = [-1*args_dict[pit].ax_lim + part_info['particle_position'][1][args_dict[pit].image_center - 1], args_dict[pit].ax_lim + part_info['particle_position'][1][args_dict[pit].image_center - 1]]
else:
xlim = [-1*args_dict[pit].ax_lim, args_dict[pit].ax_lim]
ylim = [-1*args_dict[pit].ax_lim, args_dict[pit].ax_lim]
if args.yt_proj == False:
f = h5py.File(usable_files[pit][frame_val], 'r')
image = get_image_arrays(f, simfo[pit]['field'], simfo[pit], args_dict[pit], X[pit], Y[pit])
magx = get_image_arrays(f, 'mag'+args.axis[0]+'_'+simfo[pit]['movie_file_type']+'_'+args.axis, simfo[pit], args_dict[pit], X[pit], Y[pit])
magy = get_image_arrays(f, 'mag'+args.axis[1]+'_'+simfo[pit]['movie_file_type']+'_'+args.axis, simfo[pit], args_dict[pit], X[pit], Y[pit])
x_pos_min = int(np.round(np.min(X[pit]) - simfo[pit]['xmin_full'])/simfo[pit]['cell_length'])
y_pos_min = int(np.round(np.min(Y[pit]) - simfo[pit]['xmin_full'])/simfo[pit]['cell_length'])
if np.shape(f['vel'+args.axis[0]+'_'+simfo[pit]['movie_file_type']+'_'+args.axis]) == (2048, 2048):
velocity_data = [f['vel'+args.axis[0]+'_'+simfo[pit]['movie_file_type']+'_'+args.axis], f['vel'+args.axis[1]+'_'+simfo[pit]['movie_file_type']+'_'+args.axis]]
elif args.axis == 'xy':
velocity_data = [f['vel'+args.axis[0]+'_'+simfo[pit]['movie_file_type']+'_'+args.axis][:,:,0], f['vel'+args.axis[1]+'_'+simfo[pit]['movie_file_type']+'_'+args.axis][:,:,0]]
else:
velocity_data = [f['vel'+args.axis[0]+'_'+simfo[pit]['movie_file_type']+'_'+args.axis][:,0,:], f['vel'+args.axis[1]+'_'+simfo[pit]['movie_file_type']+'_'+args.axis][:,0,:]]
velx, vely = mym.get_quiver_arrays(y_pos_min, x_pos_min, X[pit], velocity_data[0], velocity_data[1], center_vel=center_vel)
else:
if args_dict[pit].image_center == 0 or has_particles == False:
center_pos = np.array([0.0, 0.0, 0.0])
else:
dd = f.all_data()
center_pos = np.array([dd['particle_posx'][args.image_center-1].in_units('AU'), dd['particle_posy'][args.image_center-1].in_units('AU'), dd['particle_posz'][args.image_center-1].in_units('AU')])
x_width = (xlim[1] -xlim[0])
y_width = (ylim[1] -ylim[0])
thickness = yt.YTArray(args.slice_thickness, 'AU')
proj = yt.OffAxisProjectionPlot(f, L, [simfo[pit]['field'], 'cell_mass', 'velz_mw', 'magz_mw', 'Projected_Magnetic_Field_mw', 'Projected_Velocity_mw'], center=(center_pos, 'AU'), width=(x_width, 'AU'), depth=(args.slice_thickness, 'AU'))
image = (proj.frb.data[simfo[pit]['field']]/thickness.in_units('cm')).value
velx_full = (proj.frb.data[('gas', 'Projected_Velocity_mw')].in_units('g*cm**2/s')/thickness.in_units('cm')).value
vely_full = (proj.frb.data[('gas', 'velz_mw')].in_units('g*cm**2/s')/thickness.in_units('cm')).value
magx = (proj.frb.data[('gas', 'Projected_Magnetic_Field_mw')].in_units('g*gauss*cm')/thickness.in_units('cm')).value
magy = (proj.frb.data[('gas', 'magz_mw')].in_units('g*gauss*cm')/thickness.in_units('cm')).value
mass = (proj.frb.data[('gas', 'cell_mass')].in_units('cm*g')/thickness.in_units('cm')).value
velx_full = velx_full/mass
vely_full = vely_full/mass
magx = magx/mass
magy = magy/mass
del mass
velx, vely = mym.get_quiver_arrays(0.0, 0.0, X[pit], velx_full, vely_full, center_vel=center_vel)
del velx_full
del vely_full
if len(frames) == 1:
if rank == 0:
pickle_file = paths[pit] + "movie_frame_" + ("%06d" % frames[frame_val]) + ".pkl"
file = open(pickle_file, 'w+')
pickle.dump((X[pit], Y[pit], image, magx, magy, X_vel[pit], Y_vel[pit], velx, vely, xlim, ylim, has_particles, part_info, simfo[pit], time_val,xabel, yabel), file)
file.close()
print("Created Pickle:", pickle_file, "for file:", usable_files[pit][frame_val])
else:
pickle_file = paths[pit] + "movie_frame_" + ("%06d" % frames[frame_val]) + ".pkl"
file = open(pickle_file, 'w+')
pickle.dump((X[pit], Y[pit], image, magx, magy, X_vel[pit], Y_vel[pit], velx, vely, xlim, ylim, has_particles, part_info, simfo[pit], time_val,xabel, yabel), file)
file.close()
print("Created Pickle:", pickle_file, "for file:", usable_files[pit][frame_val])
f.close()
plot = axes_dict[ax_label].pcolormesh(X[pit], Y[pit], image, cmap=plt.cm.gist_heat, norm=LogNorm(vmin=cbar_min, vmax=cbar_max), rasterized=True)
plt.gca().set_aspect('equal')
if frame_val > 0 or time_val > -1.0:
axes_dict[ax_label].streamplot(X[pit], Y[pit], magx, magy, density=4, linewidth=0.25, arrowstyle='-', minlength=0.5)
else:
axes_dict[ax_label].streamplot(X[pit], Y[pit], magx, magy, density=4, linewidth=0.25, minlength=0.5)
xlim = args_dict[pit]['xlim']
ylim = args_dict[pit]['ylim']
mym.my_own_quiver_function(axes_dict[ax_label], X_vel[pit], Y_vel[pit], velx, vely, plot_velocity_legend=bool(args_dict[pit]['annotate_velocity']), limits=[xlim, ylim], standard_vel=args.standard_vel)
if args_dict[pit]['has_particles']:
if args.annotate_particles_mass == True:
mym.annotate_particles(axes_dict[ax_label], part_info['particle_position'], part_info['accretion_rad'], limits=[xlim, ylim], annotate_field=part_info['particle_mass'])
else:
mym.annotate_particles(axes_dict[ax_label], part_info['particle_position'], part_info['accretion_rad'], limits=[xlim, ylim], annotate_field=None)
if args.plot_lref == True:
r_acc = np.round(part_info['accretion_rad'])
axes_dict[ax_label].annotate('$r_{acc}$='+str(r_acc)+'AU', xy=(0.98*simfo[pit]['xmax'], 0.93*simfo[pit]['ymax']), va="center", ha="right", color='w', fontsize=args_dict[pit].text_font)
if args.annotate_time == "True" and pit == 0:
print("ANNONTATING TIME:", str(int(time_val))+'yr')
time_text = axes_dict[ax_label].text((xlim[0]+0.01*(xlim[1]-xlim[0])), (ylim[1]-0.03*(ylim[1]-ylim[0])), '$t$='+str(int(time_val))+'yr', va="center", ha="left", color='w', fontsize=args.text_font)
time_text.set_path_effects([path_effects.Stroke(linewidth=3, foreground='black'), path_effects.Normal()])
#ax.annotate('$t$='+str(int(time_val))+'yr', xy=(xlim[0]+0.01*(xlim[1]-xlim[0]), ylim[1]-0.03*(ylim[1]-ylim[0])), va="center", ha="left", color='w', fontsize=args.text_font)
title_text = axes_dict[ax_label].text((np.mean(xlim)), (ylim[1]-0.03*(ylim[1]-ylim[0])), title, va="center", ha="center", color='w', fontsize=(args.text_font+2))
title_text.set_path_effects([path_effects.Stroke(linewidth=3, foreground='black'), path_effects.Normal()])
if positions[pit][0] == columns:
cbar = plt.colorbar(plot, pad=0.0, ax=axes_dict[ax_label])
cbar.set_label('Density (gcm$^{-3}$)', rotation=270, labelpad=14, size=args.text_font)
axes_dict[ax_label].set_xlabel(args_dict[pit]['xabel'], labelpad=-1, fontsize=args.text_font)
if positions[pit][0] == 1:
axes_dict[ax_label].set_ylabel(args_dict[pit]['yabel'], labelpad=-20, fontsize=args.text_font)
axes_dict[ax_label].set_xlim(xlim)
axes_dict[ax_label].set_ylim(ylim)
for line in axes_dict[ax_label].xaxis.get_ticklines():
line.set_color('white')
for line in axes_dict[ax_label].yaxis.get_ticklines():
line.set_color('white')
plt.tick_params(axis='both', which='major', labelsize=16)
for line in axes_dict[ax_label].xaxis.get_ticklines():
line.set_color('white')
for line in axes_dict[ax_label].yaxis.get_ticklines():
line.set_color('white')
if positions[pit][0] != 1:
yticklabels = axes_dict[ax_label].get_yticklabels()
plt.setp(yticklabels, visible=False)
if positions[pit][0] == 1:
axes_dict[ax_label].tick_params(axis='y', which='major', labelsize=args.text_font)
if positions[pit][1] == rows:
axes_dict[ax_label].tick_params(axis='x', which='major', labelsize=args.text_font)
if positions[pit][0] != 1:
xticklabels = axes_dict[ax_label].get_xticklabels()
plt.setp(xticklabels[0], visible=False)
if len(usable_files[pit]) > 1:
if args.output_filename == None:
import pdb
pdb.set_trace()
file_name = save_dir + "movie_frame_" + ("%06d" % frames[frame_val])
else:
file_name = args.output_filename + "_" + str(int(time_val))
else:
if args.output_filename != None:
file_name = args.output_filename
else:
import pdb
pdb.set_trace()
file_name = save_dir + "time_" + str(args.plot_time)
plt.savefig(file_name + ".eps", format='eps', bbox_inches='tight')
#plt.savefig(file_name + ".pdf", format='pdf', bbox_inches='tight')
#plt.savefig(file_name + ".jpg", format='jpeg', bbox_inches='tight')
call(['convert', '-antialias', '-quality', '100', '-density', '200', '-resize', '100%', '-flatten', file_name+'.eps', file_name+'.jpg'])
os.remove(file_name + '.eps')
del image
del magx
del magy
del velx
del vely
if args.image_center != 0 and has_particles:
X[pit], Y[pit], X_vel[pit], Y_vel[pit] = original_positions
print('Created frame', (frames[frame_val]), 'of', str(frames[-1]), 'on rank', rank, 'at time of', str(time_val), 'to save_dir:', file_name + '.eps')
rit = rit +1
if rit == size:
rit = 0
print("completed making movie frames on rank", rank)
rank = CW.Get_rank()
size = CW.Get_size()
if rank == 0:
print("size =", size)
args = parse_inputs()
#Define relevant directories
input_dir = sys.argv[1]
save_dir = sys.argv[2]
global_pickle = sys.argv[3]
if os.path.exists(save_dir) == False and rank == 0:
os.makedirs(save_dir)
sys.stdout.flush()
CW.Barrier()
#Set some plot variables independant on data files
#File files
sink_files = sorted(glob.glob(input_dir + "output*/*.dat"))
files = sorted(glob.glob(input_dir + "*/info*.txt"))
rm_files = []
for info_name in files:
sink_file = info_name.split('info')[0] + 'stars_output.dat'
if sink_file not in sink_files:
rm_files.append(info_name)
for rm_file in rm_files:
files.remove(rm_file)
del sink_files
gc.collect()
from distributed import Client
from mpi4py.MPI import COMM_WORLD as world
from dask_mpi import initialize, send_close_signal
# Split our MPI world into two pieces, one consisting just of
# the old rank 3 process and the other with everything else
new_comm_assignment = 1 if world.rank == 3 else 0
comm = world.Split(new_comm_assignment)
if world.rank != 3:
# run tests with rest of comm
is_client = initialize(comm=comm, exit=False)
if is_client:
with Client() as c:
c.submit(lambda x: x + 1, 10).result() == 11
c.submit(lambda x: x + 1, 20).result() == 21
send_close_signal()
# check that our original comm is intact
world.Barrier()
x = 100 if world.rank == 0 else 200
x = world.bcast(x)
assert x == 100
