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()