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
