def select_ids_jug(self, data_filters={}):
'''Save a set of all ids that match a set of data filters to a file.
Args:
data_filters (dict of dicts): The data filters to apply.
'''
print(
"########################################################################"
)
print("Selecting IDs")
print(
"########################################################################"
)
sys.stdout.flush()
selected_ids = self.get_selected_ids_jug(data_filters)
selected_ids_formatted = jug.Task(self.format_selected_ids,
selected_ids)
jug.Task(self.save_selected_ids, selected_ids_formatted, data_filters)
jug.barrier()
print(
"########################################################################"
)
print("Done selecting IDs!")
sys.stdout.flush()
def find_galaxies_for_particle_tracks_jug(self):
'''Main function when using jug'''
self.read_data()
ptrack_gal_ids = self.get_galaxy_identification_loop_jug()
jug.Task(self.write_galaxy_identifications, ptrack_gal_ids)
jug.barrier()
def save_particle_tracks_jug(self):
'''Loop over all redshifts, get the data, and save the particle tracks.
'''
print("#" * 80)
print("Starting Tracking!")
print("#" * 80)
# Get the target ids
self.get_target_ids()
tracked_data_snapshots = self.get_tracked_data_jug()
formatted_data = jug.Task(self.format_tracked_data,
tracked_data_snapshots)
# Write particle data to the file
jug.Task(self.write_tracked_data, formatted_data)
jug.barrier()
from jug import barrier, Task
import math
def double(x):
# this tests an important regression:
# using __import__ for the jugfile with barrier() would make this code **not** work
val = math.sqrt(2.)*math.sqrt(2.)
return x*val
two = Task(double,1)
barrier()
four = Task(double, two)
def make_call(f, arg):
return f(arg)
eight = Task(make_call, double, four)
debug=False #
#------------#
if debug:
filenames = ['/group_workspaces/jasmin/hiresgw/mj07/xjanpa.pj19910301.u.nc',\
'/group_workspaces/jasmin/hiresgw/mj07/xjanpa.pj19910311.u.nc',\
'/group_workspaces/jasmin/hiresgw/mj07/xjanpa.pj19910321.u.nc',\
'/group_workspaces/jasmin/hiresgw/mj07/xjanpa.pj19910401.u.nc',\
'/group_workspaces/jasmin/hiresgw/mj07/xjanpa.pj19910411.u.nc',\
'/group_workspaces/jasmin/hiresgw/mj07/xjanpa.pj19910421.u.nc']
else: #glob bit
filenames = glob('/group_workspaces/jasmin/hiresgw/mj07/xjanpa.pi*')
filenames.sort()
map(filesplit,filenames)
jug.barrier()
if debug:
filenames = ['/group_workspaces/jasmin/hiresgw/mj07/xjanpa.pj19910301.u.nc',\
'/group_workspaces/jasmin/hiresgw/mj07/xjanpa.pj19910401.u.nc']
else: #glob month bit
filenames = glob('/group_workspaces/jasmin/hiresgw/mj07/xjanpa.pi??????01.v.nc')
filenames.sort()
map(monreduce,filenames)
jug.barrier()
map(finalreduce,[False])
from jug import barrier, TaskGenerator
@TaskGenerator
def double(x):
return 2 * x
s = 1
for i in range(2048):
s = double(s)
barrier()
s2 = s
for i in range(2048):
s2 = double(s2)
def twice(x):
x2 = double(x)
barrier()
return double(value(x2))
def run_linefinder_jug(
tag,
out_dir = None,
sim_data_dir = None,
halo_data_dir = None,
main_mt_halo_id = None,
sim_name = None,
galdef = None,
selector_data_filters = {},
selector_kwargs = {},
sampler_kwargs = {},
tracker_kwargs = {},
gal_linker_kwargs = {},
classifier_kwargs = {},
visualization_kwargs = {},
run_id_selecting = True,
run_id_sampling = True,
run_tracking = True,
run_galaxy_linking = True,
run_classifying = True,
run_visualization = True,
):
'''Main function for running linefinder.
Args:
tag (str):
Filename identifier for data products.
out_dir (str):
Output directory to store the data in.
sim_data_dir (str):
Directory the simulation data is stored in.
halo_data_dir (str):
Directory the halo data (e.g. AHF output) is stored in.
Halo data is necessary for linking particles to galaxies.
main_mt_halo_id (int):
Halo ID for the main merger tree halo that's being tracked.
If not provided defaults to 0 (or whatever value is cataloged for
the sim name).
sim_name (str):
Name of the simulation this is being run for.
If provided then linefinder will automatically choose the location
of the simulation and halo data, according to the linefinder.config
file. The sim_data_dir or halo_data_dir arguments directly
overwrites this.
galdef (str):
Which set of parameters to use for the galaxy_linking and
classification steps? Defaults to the parameters in
linefinder.config
selector_data_filters (dict):
Data filters to pass to select.IDSelector.select_ids()
selector_kwargs (dict):
Arguments to use when selecting what particles to track.
Arguments will be passed to select.IDSelector
sampler_kwargs (dict):
Arguments to use when selecting what particles to track.
Arguments will be passed to select.IDSampler
tracker_kwargs (dict):
Arguments to use when tracking particles.
Arguments will be passedts to pass to track.ParticleTracker
gal_linker_kwargs (dict):
Arguments to use when associating particles with galaxies.
Arguments will be passed to galaxy_link.ParticleTrackGalaxyLinker
classifier_kwargs (dict):
Arguments to use when classifying particles.
Arguments will be passed to classify.Classifier
visualization_kwargs (dict):
Arguments to use when visualizing the data.
Arguments will be passed to visualize.export_to_firefly
run_id_selecting (bool):
If True, then run routines for selecting particles.
run_id_sampling (bool):
If True, then run routines for sampling from the full list of
selected particles.
run_tracking (bool):
If True, then run routines for tracking particles.
run_galaxy_linking (bool):
If True, then run routines for associating particles with galaxies.
run_classifying (bool):
If True, then run routines for classifying particles.
'''
# Expand data dirs, if possible
if out_dir is not None:
out_dir = os.path.expandvars( out_dir )
if sim_data_dir is not None:
sim_data_dir = os.path.expandvars( sim_data_dir )
if halo_data_dir is not None:
halo_data_dir = os.path.expandvars( halo_data_dir )
# Set up for auto-retrieval, if chosen
if sim_name is not None:
file_manager = file_management.FileManager()
if out_dir is None:
out_dir = file_manager.get_linefinder_dir( sim_name )
# Setup for galaxy definitions, if chosen
if galdef is not None:
galdef_dict = linefinder_config.GALAXY_DEFINITIONS[galdef]
# Setup jugdata
jugdir_tail = '{}.jugdata'.format( tag )
jug.set_jugdir( os.path.join( out_dir, jugdir_tail ) )
print( "Starting jug thread..." )
# These are kwargs that could be used at any stage of running linefinder.
general_kwargs = {
'out_dir': out_dir,
'tag': tag,
}
# Run the ID Selecting
if run_id_selecting:
# Update arguments
selector_kwargs = utilities.merge_two_dicts(
selector_kwargs, general_kwargs )
# Check if the snapshot kwargs exist, and if not, create them
if 'snapshot_kwargs' not in list( selector_kwargs.keys() ):
selector_kwargs['snapshot_kwargs'] = {}
# Add in sim data dir if given
if sim_data_dir is not None:
selector_kwargs['snapshot_kwargs']['sdir'] = sim_data_dir
# Add in halo data dir if given
if halo_data_dir is not None:
selector_kwargs['snapshot_kwargs']['halo_data_dir'] = halo_data_dir
# Use sim name to find defaults
if sim_name is not None:
snapshot_kwargs = selector_kwargs['snapshot_kwargs']
if 'sdir' not in snapshot_kwargs:
snapshot_kwargs['sdir'] = file_manager.get_sim_dir( sim_name )
if 'halo_data_dir' not in snapshot_kwargs:
snapshot_kwargs['halo_data_dir'] = file_manager.get_halo_dir( sim_name )
if 'main_halo_id' not in snapshot_kwargs:
snapshot_kwargs['main_halo_id'] = linefinder_config.MAIN_MT_HALO_ID[sim_name]
selector_kwargs['snapshot_kwargs'] = snapshot_kwargs
id_selector = select.IDSelector( **selector_kwargs )
id_selector.select_ids_jug( selector_data_filters )
# Run the ID Sampling
if run_id_sampling:
# Update arguments
sampler_kwargs = utilities.merge_two_dicts(
sampler_kwargs, general_kwargs )
# Check if the snapshot kwargs exist, and if not, create them
if 'snapshot_kwargs' not in list( sampler_kwargs.keys() ):
sampler_kwargs['snapshot_kwargs'] = {}
# Add in sim data dir if given
if sim_data_dir is not None:
sampler_kwargs['snapshot_kwargs']['sdir'] = sim_data_dir
# Add in halo data dir if given
if halo_data_dir is not None:
sampler_kwargs['snapshot_kwargs']['halo_data_dir'] = halo_data_dir
# Use sim name to find defaults
if sim_name is not None:
snapshot_kwargs = sampler_kwargs['snapshot_kwargs']
if 'sdir' not in snapshot_kwargs:
snapshot_kwargs['sdir'] = file_manager.get_sim_dir( sim_name )
if 'halo_data_dir' not in snapshot_kwargs:
snapshot_kwargs['halo_data_dir'] = file_manager.get_halo_dir( sim_name )
if 'main_halo_id' not in snapshot_kwargs:
snapshot_kwargs['main_halo_id'] = linefinder_config.MAIN_MT_HALO_ID[sim_name]
sampler_kwargs['snapshot_kwargs'] = snapshot_kwargs
id_sampler = select.IDSampler( **sampler_kwargs )
jug.Task( id_sampler.sample_ids )
jug.barrier()
# Run the Particle Tracking
if run_tracking:
# Update arguments
tracker_kwargs = utilities.merge_two_dicts(
tracker_kwargs, general_kwargs )
# Add in sim data dir if given
if sim_data_dir is not None:
tracker_kwargs['sdir'] = sim_data_dir
# Choose the sdir automatically, if possible
if 'sdir' not in tracker_kwargs:
# Try and load the default values if using the file manager.
if sim_name is not None:
tracker_kwargs['sdir'] = file_manager.get_sim_dir( sim_name )
# Try to use the sdir passed to the selector kwargs
elif 'snapshot_kwargs' in selector_kwargs:
if 'sdir' in 'snapshot_kwargs':
tracker_kwargs['sdir'] = \
selector_kwargs['snapshot_kwargs']['sdir']
particle_tracker = track.ParticleTracker( **tracker_kwargs )
particle_tracker.save_particle_tracks_jug()
# Run the Galaxy Finding
if run_galaxy_linking:
# Update arguments
gal_linker_kwargs = utilities.merge_two_dicts(
gal_linker_kwargs, general_kwargs )
# Add in halo data dir if given
if halo_data_dir is not None:
gal_linker_kwargs['halo_data_dir'] = halo_data_dir
if sim_name is not None:
if 'halo_data_dir' not in gal_linker_kwargs:
gal_linker_kwargs['halo_data_dir'] = file_manager.get_halo_dir( sim_name )
if 'main_mt_halo_id' not in gal_linker_kwargs:
gal_linker_kwargs['main_mt_halo_id'] = linefinder_config.MAIN_MT_HALO_ID[sim_name]
# Default to halo 0 if MT halo ID not given
if 'main_mt_halo_id' not in gal_linker_kwargs:
gal_linker_kwargs['main_mt_halo_id'] = 0
if galdef is not None:
for key in [ 'galaxy_cut', 'length_scale', 'mt_length_scale' ]:
gal_linker_kwargs[key] = galdef_dict[key]
particle_track_gal_linker = galaxy_link.ParticleTrackGalaxyLinker(
**gal_linker_kwargs
)
particle_track_gal_linker.find_galaxies_for_particle_tracks_jug()
# Run the Classification
if run_classifying:
# Update arguments
classifier_kwargs = utilities.merge_two_dicts(
classifier_kwargs, general_kwargs )
# Add in halo data dir if given
if halo_data_dir is not None:
classifier_kwargs['halo_data_dir'] = halo_data_dir
if sim_name is not None:
if 'halo_data_dir' not in classifier_kwargs:
classifier_kwargs['halo_data_dir'] = file_manager.get_halo_dir( sim_name )
if galdef is not None:
for key in [ 't_pro', 't_m', ]:
classifier_kwargs[key] = galdef_dict[key]
classifier = classify.Classifier( **classifier_kwargs )
jug.Task( classifier.classify_particles )
# Run Visualizing
if run_visualization:
# Add in halo data dir if given
if halo_data_dir is not None:
visualization_kwargs['halo_data_dir'] = halo_data_dir
if sim_name is not None:
if 'halo_data_dir' not in visualization_kwargs:
visualization_kwargs['halo_data_dir'] = file_manager.get_halo_dir( sim_name )
if 'main_mt_halo_id' not in visualization_kwargs:
visualization_kwargs['main_halo_id'] = linefinder_config.MAIN_MT_HALO_ID[sim_name]
jug.Task(
visualize.export_to_firefly,
tag = tag,
data_dir = out_dir,
**visualization_kwargs
)
# Make a file indicating that the visualization completed.
f = os.path.join( out_dir, 'visualized_{}'.format(tag ) )
open(f, 'a').close()
def twice(x):
x2 = double(x)
barrier()
return double(value(x2))
