def _db_import_export(target_session, from_session, *sims): external_id_to_internal_halo = {} translated_halolink_ids = [] if len(sims)==0: sims = [x.id for x in all_simulations(from_session)] for sim in sims: ext_sim = get_simulation(sim, from_session) sim = Simulation(ext_sim.basename) target_session.add(sim) logger.info("Transferring simulation %s", ext_sim) halos_this_ts = [] for p_ext in ext_sim.properties: dic = get_or_create_dictionary_item( target_session, p_ext.name.text) p = SimulationProperty(sim, dic, p_ext.data) halos_this_ts.append(p) for tk_ext in ext_sim.trackers: tk = TrackData(sim, tk_ext.halo_number) tk.particles = tk_ext.particles tk.use_iord = tk_ext.use_iord halos_this_ts.append(tk) target_session.add_all(halos_this_ts) for ts_ext in ext_sim.timesteps: logger.info("Transferring timestep %s",ts_ext) ts = TimeStep(sim, ts_ext.extension) ts.redshift = ts_ext.redshift ts.time_gyr = ts_ext.time_gyr ts.available = True target_session.add(ts) halos_this_ts = [] logger.info("Transferring objects for %s", ts_ext) for h_ext in ts_ext.objects: h = SimulationObjectBase(ts, h_ext.halo_number, h_ext.finder_id, h_ext.finder_offset, h_ext.NDM, h_ext.NStar, h_ext.NGas, h_ext.object_typecode) h.external_id = h_ext.id halos_this_ts.append(h) target_session.add_all(halos_this_ts) target_session.commit() for h in halos_this_ts: assert h.id is not None and h.id > 0 external_id_to_internal_halo[h.external_id] = h properties_this_ts = [] logger.info("Transferring object properties for %s", ts_ext) for h_ext in ts_ext.objects: h_new = external_id_to_internal_halo[h_ext.id] for p_ext in h_ext.properties: dic = get_or_create_dictionary_item( target_session, p_ext.name.text) dat = p_ext.data_raw if dat is not None: p = HaloProperty(h_new, dic, dat) target_session.add(p) target_session.commit() for ts_ext in ext_sim.timesteps: logger.info("Transferring halolinks for timestep %s", ts_ext) sys.stdout.flush() _translate_halolinks( target_session, ts_ext.links_from, external_id_to_internal_halo, translated_halolink_ids) _translate_halolinks( target_session, ts_ext.links_to, external_id_to_internal_halo, translated_halolink_ids) target_session.commit() logger.info("Done")
def add_tracker(halo, size=None): import pynbody try: halo = get_halo(halo) except: sim = get_simulation(halo) print("Adding tracker for isolated run", sim) halo = None if halo is not None: # check we can open the tracker file hfile = halo.load() hfile.physical_units() use_iord = True try: hfile.dm['iord'] except: use_iord = False # get the centre cen = halo.get('SSC', None) if cen is None: cen = pynbody.analysis.halo.shrink_sphere_center(hfile.dm) hfile.ancestor.dm['pos'] -= cen if size is None: size = '500 pc' size = pynbody.units.Unit(size) try: size.in_units("kpc") X = hfile.ancestor.dm[pynbody.filt.Sphere(size)] except: size.in_units("kpc km s^-1") X = hfile.ancestor.dm[pynbody.filt.LowPass("j2", size ** 2)] size = str(size.in_units("kpc km s^-1")) + "_kks" if len(X) < 2: print("Alert! Track data is too short") import pdb pdb.set_trace() # setup the tracker tx = TrackData(halo.timestep.simulation) print("Tracker halo ID is", tx.halo_number) print("Length of track data is", len(X)) tx.select(X, use_iord) ts_trigger = halo.timestep else: f = sim.timesteps[0].load() tx = TrackData(sim) if tx.halo_number != 1: print("Already have a tracker for this simulation -> abort") return print("Tracker halo ID is", tx.halo_number) tx.particles = np.array( np.arange(0, len(f.dm) - 1, 1), dtype=int) tx.use_iord = False ts_trigger = None core.get_default_session().add(tx) tx.create_objects(first_timestep=ts_trigger) if halo is not None: targ = halo.timestep.halos.filter_by( object_typecode=1, halo_number=tx.halo_number).first() size = str(size).replace(" ", "") halo["dm_central_" + size] = targ print("Copying SSC...", end=' ') sys.stdout.flush() while halo is not None: try: targ['SSC'] = halo['SSC'] except KeyError: pass halo = halo.next targ = targ.next print("done") core.get_default_session().commit()