def _add_two_properties_different_ranges(): for i in pt.distributed(list(range(1, 10))): with pt.ExclusiveLock('insert', 0.05): tangos.get_halo(i)['my_test_property_2'] = i tangos.core.get_default_session().commit() for i in pt.distributed(list(range(1, 8))): with pt.ExclusiveLock('insert', 0.05): tangos.get_halo(i)['my_test_property_3'] = i tangos.core.get_default_session().commit()
def insert_list(property_list): from tangos import parallel_tasks as pt if pt.backend != None: with pt.ExclusiveLock("insert_list"): _insert_list_unlocked(property_list) else: _insert_list_unlocked(property_list)
def _test_not_run_twice(): import time # For this test we want a staggered start time.sleep(pt.backend.rank() * 0.05) for i in pt.distributed(list(range(3))): with pt.ExclusiveLock("lock"): tangos.get_halo(1)['test_count'] += 1 tangos.get_default_session().commit()
def _test_synchronize_db_creator(): rank = pt.backend.rank() import tangos.parallel_tasks.database # hack: MultiProcessing backend forks so has already "synced" the current creator. tangos.core.creator._current_creator = None pt.database.synchronize_creator_object(tangos.core.get_default_session()) with pt.ExclusiveLock('insert', 0.05): tangos.get_halo(rank)['db_creator_test_property'] = 1.0 tangos.core.get_default_session().commit()
def test_function(): lock = pt.ExclusiveLock("hello") print("Hello from rank", pt.backend.rank()) for i in pt.distributed(range(10)): with lock: print("Task", i) time.sleep(0.1) if pt.backend.rank() == 1: print() print("OK")
def add_missing_trackdata_and_BH_objects(timestep, this_step_bh_iords, existing_bhobj_iords, session): with parallel_tasks.ExclusiveLock("bh"): track, track_nums = db.tracking.get_trackers(timestep.simulation) with session.no_autoflush: tracker_to_add = collect_bh_trackers(this_step_bh_iords, timestep.simulation, track_nums) halo_to_add = generate_missing_bh_objects(this_step_bh_iords, timestep, existing_bhobj_iords) session.add_all(tracker_to_add) session.add_all(halo_to_add) session.commit() logger.info("Committed %d new trackdata and %d new BH objects for %r", len(tracker_to_add), len(halo_to_add), timestep)
def test_function(): lock = pt.ExclusiveLock("hello") rank = pt.backend.rank() for test_len, test_type in zip((10, 10000, 10000), (np.int32, np.float64, np.float32)) : data_to_send = np.arange(test_len, dtype=test_type) if rank == 1: pt.backend.send_numpy_array(data_to_send, 2) elif rank == 2: received = pt.backend.receive_numpy_array(1) assert(received.dtype==data_to_send.dtype) assert(received.shape==data_to_send.shape) assert all(received==data_to_send) if rank==2: print() print("OK")
def crosslink_ts(self, ts1, ts2, halo_min=0, halo_max=None, dmonly=False, threshold=config.default_linking_threshold, object_typecode=0): """Link the halos of two timesteps together :type ts1 tangos.core.TimeStep :type ts2 tangos.core.TimeStep""" logger.info("Gathering halo information for %r and %r", ts1, ts2) halos1 = self.make_finder_id_to_halo_map(ts1, object_typecode) halos2 = self.make_finder_id_to_halo_map(ts2, object_typecode) with parallel_tasks.ExclusiveLock("create_db_objects_from_catalog"): same_d_id = core.dictionary.get_or_create_dictionary_item( self.session, "ptcls_in_common") self.session.commit() output_handler_1 = ts1.simulation.get_output_handler() output_handler_2 = ts2.simulation.get_output_handler() if type(output_handler_1).match_objects != type( output_handler_2).match_objects: logger.error( "Timesteps %r and %r cannot be crosslinked; they are using incompatible file readers", ts1, ts2) return # keep the files alive throughout (so they are not garbage-collected after the first match_objects): snap1 = ts1.load() snap2 = ts2.load() try: cat = output_handler_1.match_objects( ts1.extension, ts2.extension, halo_min, halo_max, dmonly, threshold, core.halo.Halo.object_typetag_from_code(object_typecode), output_handler_for_ts2=output_handler_2) back_cat = output_handler_2.match_objects( ts2.extension, ts1.extension, halo_min, halo_max, dmonly, threshold, core.halo.Halo.object_typetag_from_code(object_typecode), output_handler_for_ts2=output_handler_1) except Exception as e: if isinstance(e, KeyboardInterrupt): raise logger.exception( "Exception during attempt to crosslink timesteps %r and %r", ts1, ts2) return with self.session.no_autoflush: logger.info("Gathering links for %r and %r", ts1, ts2) items = self.create_db_objects_from_catalog( cat, halos1, halos2, same_d_id) logger.info("Identified %d links between %r and %r", len(items), ts1, ts2) items_back = self.create_db_objects_from_catalog( back_cat, halos2, halos1, same_d_id) logger.info("Identified %d links between %r and %r", len(items_back), ts2, ts1) with parallel_tasks.ExclusiveLock("create_db_objects_from_catalog"): logger.info("Preparing to commit links for %r and %r", ts1, ts2) self.session.add_all(items) self.session.add_all(items_back) self.session.commit() logger.info("Finished committing total of %d links for %r and %r", len(items) + len(items_back), ts1, ts2)
def _add_property(): for i in pt.distributed(list(range(1, 10))): with pt.ExclusiveLock('insert', 0.05): tangos.get_halo(i)['my_test_property'] = i tangos.core.get_default_session().commit()
def generate_halolinks(session, fname, pairs): for ts1, ts2 in parallel_tasks.distributed(pairs): bh_log = None if BlackHolesLog.can_load(ts2.filename): bh_log = BlackHolesLog(ts2.filename) elif ShortenedOrbitLog.can_load(ts2.filename): bh_log = ShortenedOrbitLog(ts2.filename) if bh_log is None: logger.error("Warning! No orbit file found!") links = [] mergers_links = [] bh_map = {} logger.info("Gathering BH tracking information for steps %r and %r", ts1, ts2) with parallel_tasks.ExclusiveLock("bh"): dict_obj = db.core.get_or_create_dictionary_item( session, "tracker") dict_obj_next = db.core.get_or_create_dictionary_item( session, "BH_merger_next") dict_obj_prev = db.core.get_or_create_dictionary_item( session, "BH_merger_prev") track_links_n, idf_n, idt_n = db.tracking.get_tracker_links( session, dict_obj_next) bh_objects_1, nums1, id1 = get_bh_objs_numbers_and_dbids(ts1) bh_objects_2, nums2, id2 = get_bh_objs_numbers_and_dbids(ts2) tracker_links, idf, idt = db.tracking.get_tracker_links( session, dict_obj) idf_n = np.array(idf_n) idt_n = np.array(idt_n) if len(nums1) == 0 or len(nums2) == 0: logger.info("No BHs found in either step %r or %r... moving on", ts1, ts2) continue logger.info("Generating BH tracker links between steps %r and %r", ts1, ts2) o1 = np.where(np.in1d(nums1, nums2))[0] o2 = np.where(np.in1d(nums2, nums1))[0] if len(o1) == 0 or len(o2) == 0: continue with session.no_autoflush: for ii, jj in zip(o1, o2): if nums1[ii] != nums2[jj]: raise RuntimeError("BH iords are mismatched") exists = np.where((idf == id1[ii]) & (idt == id2[jj]))[0] if len(exists) == 0: links.append( tangos.core.halo_data.HaloLink(bh_objects_1[ii], bh_objects_2[jj], dict_obj, 1.0)) links.append( tangos.core.halo_data.HaloLink(bh_objects_2[jj], bh_objects_1[ii], dict_obj, 1.0)) logger.info("Generated %d tracker links between steps %r and %r", len(links), ts1, ts2) logger.info("Generating BH Merger information for steps %r and %r", ts1, ts2) for l in open(fname[0]): l_split = l.split() t = float(l_split[6]) bh_dest_id = int(l_split[0]) bh_src_id = int(l_split[1]) ratio = float(l_split[4]) # ratios in merger file are ambiguous (since major progenitor may be "source" rather than "destination") # re-establish using the log file: try: ratio = bh_log.determine_merger_ratio(bh_src_id, bh_dest_id) except (ValueError, AttributeError) as e: logger.debug( "Could not calculate merger ratio for %d->%d from the BH log; assuming the .BHmergers-asserted value is accurate", bh_src_id, bh_dest_id) if t > ts1.time_gyr and t <= ts2.time_gyr: bh_map[bh_src_id] = (bh_dest_id, ratio) resolve_multiple_mergers(bh_map) logger.info("Gathering BH merger links for steps %r and %r", ts1, ts2) with session.no_autoflush: for src, (dest, ratio) in bh_map.items(): if src not in nums1 or dest not in nums2: logger.warning( "Can't link BH %r -> %r; missing BH objects in database", src, dest) continue bh_src_before = bh_objects_1[nums1.index(src)] bh_dest_after = bh_objects_2[nums2.index(dest)] if ((idf_n == bh_src_before.id) & (idt_n == bh_dest_after.id)).sum() == 0: mergers_links.append( tangos.core.halo_data.HaloLink(bh_src_before, bh_dest_after, dict_obj_next, 1.0)) mergers_links.append( tangos.core.halo_data.HaloLink(bh_dest_after, bh_src_before, dict_obj_prev, ratio)) logger.info("Generated %d BH merger links for steps %r and %r", len(mergers_links), ts1, ts2) with parallel_tasks.ExclusiveLock("bh"): logger.info("Committing total %d BH links for steps %r and %r", len(mergers_links) + len(links), ts1, ts2) session.add_all(links) session.add_all(mergers_links) session.commit() logger.info("Finished committing BH links for steps %r and %r", ts1, ts2)
def assign_bh_to_halos(bh_halo_assignment, bh_iord, timestep, linkname, hostname=None): session = Session.object_session(timestep) linkname_dict_id = tangos.core.dictionary.get_or_create_dictionary_item( session, linkname) if hostname is not None: host_dict_id = tangos.core.dictionary.get_or_create_dictionary_item( session, hostname) else: host_dict_id = None logger.info("Gathering %s links for step %r", linkname, timestep) links, link_id_from, link_id_to = db.tracking.get_tracker_links( session, linkname_dict_id) halos = timestep.halos.filter_by(object_typecode=0).all() halo_nums = [h.halo_number for h in halos] halo_catind = [h.finder_offset for h in halos] halo_ids = np.array([h.id for h in halos]) logger.info("Gathering bh halo information for %r", timestep) with parallel_tasks.lock.SharedLock("bh"): bh_database_object, existing_bh_nums, bhobj_ids = get_bh_objs_numbers_and_dbids( timestep) bh_links = [] with session.no_autoflush: for bhi, haloi in zip(bh_iord, bh_halo_assignment): haloi = int(haloi) bhi = int(bhi) if haloi not in halo_catind: logger.warning( "Skipping BH in halo %d as no corresponding halo found in the database", haloi) continue if bhi not in existing_bh_nums: logger.warning("Can't find the database object for BH %d", bhi) print(bhi) print(existing_bh_nums) continue bh_index_in_list = existing_bh_nums.index(bhi) halo_index_in_list = halo_catind.index(haloi) bh_obj = bh_database_object[bh_index_in_list] halo_obj = halos[halo_index_in_list] num_existing_links = ( (link_id_from == halo_ids[halo_index_in_list]) & (link_id_to == bhobj_ids[bh_index_in_list])).sum() if num_existing_links == 0: bh_links.append( tangos.core.halo_data.HaloLink(halo_obj, bh_obj, linkname_dict_id)) if host_dict_id is not None: bh_links.append( tangos.core.halo_data.HaloLink(bh_obj, halo_obj, host_dict_id)) logger.info("Committing %d %s links for step %r...", len(bh_links), linkname, timestep) with parallel_tasks.ExclusiveLock("bh"): session.add_all(bh_links) session.commit() logger.info("...done")
def _get_linkname_dictionaryitem(self): with parallel_tasks.ExclusiveLock("create_db_objects_from_catalog"): same_d_id = core.dictionary.get_or_create_dictionary_item( self.session, "ptcls_in_common") self.session.commit() return same_d_id