From 368040bf520ef6652b0150a70dfa6a0aeb466c7c Mon Sep 17 00:00:00 2001 From: Corentin Cadiou Date: Wed, 6 Sep 2023 10:14:18 +0100 Subject: [PATCH 01/13] Update pre-commit --- .pre-commit-config.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 508ad450..fbbe36b4 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -5,7 +5,7 @@ ci: repos: - repo: https://github.com/asottile/setup-cfg-fmt - rev: v2.3.0 + rev: v2.4.0 hooks: - id: setup-cfg-fmt - repo: https://github.com/pre-commit/pre-commit-hooks @@ -19,7 +19,7 @@ repos: - id: check-executables-have-shebangs - id: check-yaml - repo: https://github.com/asottile/pyupgrade - rev: v3.6.0 + rev: v3.10.1 hooks: - id: pyupgrade args: [--py38-plus] From caa87dc625529bb91f19d6e07cec94c338921e85 Mon Sep 17 00:00:00 2001 From: Corentin Cadiou Date: Wed, 6 Sep 2023 10:14:54 +0100 Subject: [PATCH 02/13] Add possibility to x-link with yt --- tangos/input_handlers/yt.py | 83 ++++++++++++++++++++++++++++++++++++- 1 file changed, 81 insertions(+), 2 deletions(-) diff --git a/tangos/input_handlers/yt.py b/tangos/input_handlers/yt.py index f9730f16..85a4a647 100644 --- a/tangos/input_handlers/yt.py +++ b/tangos/input_handlers/yt.py @@ -2,6 +2,7 @@ import glob import os +from typing import List, Tuple import numpy as np @@ -46,8 +47,86 @@ def load_object(self, ts_extension, finder_id, finder_offset, object_typetag='ha def load_tracked_region(self, ts_extension, track_data, mode=None): raise NotImplementedError("Tracked regions not implemented for yt") - def match_objects(self, ts1, ts2, halo_min, halo_max, dm_only=False, threshold=0.005, object_typetag='halo'): - raise NotImplementedError("Matching halos still needs to be implemented for yt") + def match_objects( + self, + ts1: str, + ts2: str, + halo_min: int, + halo_max: int, + dm_only: bool=False, + threshold:float =0.005, + object_typetag: str="halo", + output_handler_for_ts2=None, + fuzzy_match_kwa={}, + ) -> Tuple[List[int], List[List[Tuple[int, int]]]]: + if output_handler_for_ts2 is None: + raise NotImplementedError( + "Alternative output_handler_for_ts2 is not implemented for yt." + ) + if fuzzy_match_kwa: + raise NotImplementedError( + "Fuzzy matching is not implemented for yt." + ) + + if halo_min is None: + halo_min = 0 + if halo_max is None: + halo_max = np.inf + + h1 = self._load_halo_cat(ts1, object_typetag) + if output_handler_for_ts2 is None: + h2 = self._load_halo_cat(ts2, object_typetag) + else: + h2 = output_handler_for_ts2._load_halo_cat(ts2, object_typetag) + + # Compute the sets of particle ids in each halo + members2 = np.concatenate([ + h2.halo("halos", i).member_ids + for i in h2.r["particle_identifier"].astype(int) + if halo_min <= i <= halo_max + ]) + + members2halo2 = np.concatenate( + np.repeat(i, len(h2.halo("halos", i).member_ids)) + for i in h2.r["particle_identifier"].astype(int) + if halo_min <= i <= halo_max + ) + + # Compute size of intersection of all sets in h1 with those in h2 + cat, possibilities = [], [] + for ihalo1 in h1.r["particle_identifier"].astype(int): + if not (halo_min <= ihalo1 <= halo_max): + continue + + ids1 = h1.halo("halos", ihalo1).member_ids + mask = np.in1d(ids1, members2) + if mask.sum() == 0: + continue + + # Get the halo ids of the particles in the other snapshot + idhalo2 = members2halo2[mask] + + # Count the number of particles in each halo + idhalo2, counts = np.unique(idhalo2, return_counts=True) + weights = counts / len(ids1) + + # Sort the links by decreasing number of particles + _order = np.argsort(weights)[::-1] + idhalo2 = idhalo2[_order] + weights = weights[_order] + + # Keep only the links with a significant number of particles + mask = weights > threshold + if mask.sum() == 0: + continue + + idhalo2 = idhalo2[mask] + weights = weights[mask] + + cat.append(ihalo1) + possibilities.append(list(zip(idhalo2, counts))) + + return cat, possibilities def enumerate_objects(self, ts_extension, object_typetag="halo", min_halo_particles=config.min_halo_particles): if object_typetag!="halo": From 49f23f6ae2c4bede93efdd2338bb307faa8a3d13 Mon Sep 17 00:00:00 2001 From: Anatole Storck Date: Tue, 19 Sep 2023 17:02:39 +0200 Subject: [PATCH 03/13] support RAMSES data with rockstar halos --- tangos/input_handlers/yt.py | 125 +++++++++++++++++++++++++++--- tangos/tools/crosslink.py | 2 + tangos/util/read_datasets_file.py | 5 +- 3 files changed, 120 insertions(+), 12 deletions(-) diff --git a/tangos/input_handlers/yt.py b/tangos/input_handlers/yt.py index 85a4a647..d910b5bc 100644 --- a/tangos/input_handlers/yt.py +++ b/tangos/input_handlers/yt.py @@ -73,11 +73,11 @@ def match_objects( if halo_max is None: halo_max = np.inf - h1 = self._load_halo_cat(ts1, object_typetag) + h1, _ = self._load_halo_cat(ts1, object_typetag) if output_handler_for_ts2 is None: - h2 = self._load_halo_cat(ts2, object_typetag) + h2, _ = self._load_halo_cat(ts2, object_typetag) else: - h2 = output_handler_for_ts2._load_halo_cat(ts2, object_typetag) + h2, _ = output_handler_for_ts2._load_halo_cat(ts2, object_typetag) # Compute the sets of particle ids in each halo members2 = np.concatenate([ @@ -86,11 +86,11 @@ def match_objects( if halo_min <= i <= halo_max ]) - members2halo2 = np.concatenate( + members2halo2 = np.concatenate([ np.repeat(i, len(h2.halo("halos", i).member_ids)) for i in h2.r["particle_identifier"].astype(int) if halo_min <= i <= halo_max - ) + ]) # Compute size of intersection of all sets in h1 with those in h2 cat, possibilities = [], [] @@ -99,8 +99,12 @@ def match_objects( continue ids1 = h1.halo("halos", ihalo1).member_ids - mask = np.in1d(ids1, members2) + #mask = np.in1d(ids1, members2) + mask = np.in1d(members2, ids1) if mask.sum() == 0: + cat.append( + [] + ) continue # Get the halo ids of the particles in the other snapshot @@ -118,15 +122,25 @@ def match_objects( # Keep only the links with a significant number of particles mask = weights > threshold if mask.sum() == 0: + cat.append( + [] + ) continue idhalo2 = idhalo2[mask] weights = weights[mask] - cat.append(ihalo1) - possibilities.append(list(zip(idhalo2, counts))) + cat.append( + list(zip(idhalo2.value, weights)) + ) + + #mport pdb + #pdb.set_trace() - return cat, possibilities + #import pdb + #pdb.set_trace() + + return cat def enumerate_objects(self, ts_extension, object_typetag="halo", min_halo_particles=config.min_halo_particles): if object_typetag!="halo": @@ -164,6 +178,99 @@ def _load_halo_cat_without_caching(self, ts_extension, snapshot_file): def get_properties(self): return {} +class YtRamsesRockstarInputHandler(YtInputHandler): + patterns = ["output_0????"] + auxiliary_file_patterns = ["halos_*.bin"] + + def load_timestep_without_caching(self, ts_extension, mode=None): + from yt.data_objects.particle_filters import add_particle_filter + if mode is not None: + raise ValueError("Custom load modes are not supported with yt") + f = yt.load(self._extension_to_filename(ts_extension)) + + def Stars(pfilter, data): + filter = data[("all", "particle_family")] == 2 # DM = 1, Stars = 2 + return filter + + add_particle_filter("stars", function=Stars, filtered_type='all', \ + requires=["particle_family"]) + + def AllDarkMatter(pfilter, data): + #filter = np.logical_or(data[("all", "particle_type")] == 4,data[("all", "particle_type")] == 1) # DM = 1, Stars = 2 + filter = data[("all", "particle_family")] == 1 + return filter + + add_particle_filter("dark_matter", function=AllDarkMatter, filtered_type='all', \ + requires=["particle_family"]) + + def MustRefineParticles(pfilter, data): + filter = data[("all", "particle_family")] == 4 + return filter + + add_particle_filter("mrp_dark_matter", function=MustRefineParticles, filtered_type='all', \ + requires=["particle_family"]) + + f.add_particle_filter("stars") + f.add_particle_filter("dark_matter") + f.add_particle_filter("mrp_dark_matter") + + return f + + def _load_halo_cat_without_caching(self, ts_extension, snapshot_file): + # Check whether datasets.txt exists (i.e., if rockstar was run with yt) + if os.path.exists(self._extension_to_filename("datasets.txt")): + fnum = read_datasets(self._extension_to_filename(""),ts_extension) + print (ts_extension) + print (fnum) + else: # otherwise, assume a one-to-one correspondence + overdir = self._extension_to_filename("") + snapfiles = glob.glob(overdir+ts_extension[:2]+len(ts_extension[2:].split('/')[0])*'?') + rockfiles = glob.glob(overdir+"out_*.list") + sortind = np.array([int(rname.split('.')[0].split('_')[-1]) for rname in rockfiles]) + sortord = np.argsort(sortind) + snapfiles.sort() + rockfiles = np.array(rockfiles)[sortord] + timestep_ind = np.argwhere(np.array([s.split('/')[-1] for s in snapfiles])==ts_extension.split('/')[0])[0] + fnum = int(rockfiles[timestep_ind][0].split('.')[0].split('_')[-1]) + cat = yt.frontends.rockstar.RockstarDataset(self._extension_to_filename("halos_"+str(fnum)+".0.bin")) + cat_data = cat.all_data() + # Check whether rockstar was run with Behroozi's distribution or Wise's + if np.any(cat_data["halos","particle_identifier"]<0): + del cat + del cat_data + cat = yt.frontends.rockstar.RockstarDataset(self._extension_to_filename("halos_"+str(fnum)+".0.bin")) + cat.parameters['format_revision'] = 2 # + cat_data = cat.all_data() + return cat, cat_data + + def enumerate_objects(self, ts_extension, object_typetag="halo", min_halo_particles=config.min_halo_particles): + if object_typetag!="halo": + return + if self._can_enumerate_objects_from_statfile(ts_extension, object_typetag): + yield from self._enumerate_objects_from_statfile(ts_extension, object_typetag) + else: + logger.warn("No halo statistics file found for timestep %r", ts_extension) + logger.warn(" => enumerating %ss directly using yt", object_typetag) + + catalogue, catalogue_data = self._load_halo_cat(ts_extension, object_typetag) + num_objects = len(catalogue_data["halos", "virial_radius"]) + + for i in range(num_objects): + obj = self.load_object(ts_extension, int(catalogue_data["halos","particle_identifier"][i]), i, object_typetag) + NDM = len(obj["dark_matter","particle_mass"]) + NGas = 0 # cells + NStar = len(obj["stars","particle_mass"]) + if NDM + NGas + NStar> min_halo_particles: + yield i, int(catalogue_data["halos","particle_identifier"][i]), NDM, NStar, NGas + + def load_object(self, ts_extension, finder_id, finder_offset, object_typetag='halo', mode=None): + f = self.load_timestep(ts_extension, mode) + cat, cat_dat = self._load_halo_cat(ts_extension, object_typetag) + center = cat_dat["halos","particle_position"][cat_dat["halos","particle_identifier"]==finder_id][0] + center+=f.domain_left_edge-cat.domain_left_edge + radius = cat_dat["halos","virial_radius"][cat_dat["halos","particle_identifier"]==finder_id][0] + return f.sphere(center.in_cgs(), radius.in_cgs()) + class YtChangaAHFInputHandler(YtInputHandler): patterns = ["*.00???", "*.00????"] diff --git a/tangos/tools/crosslink.py b/tangos/tools/crosslink.py index 73172f03..54bda2ae 100644 --- a/tangos/tools/crosslink.py +++ b/tangos/tools/crosslink.py @@ -86,6 +86,8 @@ def need_crosslink_ts(self, ts1, ts2, object_typecode=0): def create_db_objects_from_catalog(self, cat, finder_offset_to_halos_1, finder_offset_to_halos_2, same_d_id): items = [] missing_db_object = 0 + #import pdb + #pdb.set_trace() for i, possibilities in enumerate(cat): h1 = finder_offset_to_halos_1.get(i, None) for cat_i, weight in possibilities: diff --git a/tangos/util/read_datasets_file.py b/tangos/util/read_datasets_file.py index 8590e6be..1ac2324c 100644 --- a/tangos/util/read_datasets_file.py +++ b/tangos/util/read_datasets_file.py @@ -5,7 +5,6 @@ def read_datasets(basedir,filename): if os.path.exists(os.path.join(basedir, "datasets.txt")): with open(os.path.join(basedir, "datasets.txt")) as f: for l in f: - if l.split()[0].endswith(filename): + if filename in l.split()[0]: return int(l.split()[1]) - else: - raise AssertionError("Unable to open datasets.txt") + raise AssertionError("Unable to open datasets.txt") From c72c993c9bc55620d967c860c6aa81b676294c4f Mon Sep 17 00:00:00 2001 From: Corentin Cadiou Date: Tue, 19 Sep 2023 19:16:49 +0200 Subject: [PATCH 04/13] Make sure indices are handled correctly tangos is numbering from 0 onwards, while yt is using the rockstar ids in some random but deterministic (?) order. --- tangos/input_handlers/yt.py | 22 ++++++++++------------ 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/tangos/input_handlers/yt.py b/tangos/input_handlers/yt.py index d910b5bc..4389566e 100644 --- a/tangos/input_handlers/yt.py +++ b/tangos/input_handlers/yt.py @@ -87,18 +87,18 @@ def match_objects( ]) members2halo2 = np.concatenate([ - np.repeat(i, len(h2.halo("halos", i).member_ids)) - for i in h2.r["particle_identifier"].astype(int) - if halo_min <= i <= halo_max + np.repeat(itangos, len(h2.halo("halos", irockstar).member_ids)) + for itangos, irockstar in enumerate(h2.r["particle_identifier"].astype(int)) + if halo_min <= itangos <= halo_max ]) # Compute size of intersection of all sets in h1 with those in h2 - cat, possibilities = [], [] - for ihalo1 in h1.r["particle_identifier"].astype(int): - if not (halo_min <= ihalo1 <= halo_max): + cat = [] + for ihalo1_tangos, ihalo1_rockstar in enumerate(h1.r["particle_identifier"].astype(int)): + if not (halo_min <= ihalo1_tangos <= halo_max): continue - ids1 = h1.halo("halos", ihalo1).member_ids + ids1 = h1.halo("halos", ihalo1_rockstar).member_ids #mask = np.in1d(ids1, members2) mask = np.in1d(members2, ids1) if mask.sum() == 0: @@ -130,9 +130,7 @@ def match_objects( idhalo2 = idhalo2[mask] weights = weights[mask] - cat.append( - list(zip(idhalo2.value, weights)) - ) + cat.append(list(zip(idhalo2, weights))) #mport pdb #pdb.set_trace() @@ -232,13 +230,13 @@ def _load_halo_cat_without_caching(self, ts_extension, snapshot_file): rockfiles = np.array(rockfiles)[sortord] timestep_ind = np.argwhere(np.array([s.split('/')[-1] for s in snapfiles])==ts_extension.split('/')[0])[0] fnum = int(rockfiles[timestep_ind][0].split('.')[0].split('_')[-1]) - cat = yt.frontends.rockstar.RockstarDataset(self._extension_to_filename("halos_"+str(fnum)+".0.bin")) + cat = yt.load(self._extension_to_filename("halos_"+str(fnum)+".0.bin")) cat_data = cat.all_data() # Check whether rockstar was run with Behroozi's distribution or Wise's if np.any(cat_data["halos","particle_identifier"]<0): del cat del cat_data - cat = yt.frontends.rockstar.RockstarDataset(self._extension_to_filename("halos_"+str(fnum)+".0.bin")) + cat = yt.load(self._extension_to_filename("halos_"+str(fnum)+".0.bin")) cat.parameters['format_revision'] = 2 # cat_data = cat.all_data() return cat, cat_data From 0a02ccd40eedd6624a3a3c6453c3d52423e45415 Mon Sep 17 00:00:00 2001 From: Corentin Cadiou Date: Tue, 19 Sep 2023 22:04:21 +0200 Subject: [PATCH 05/13] Allow importing properties --- tangos/input_handlers/yt.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/tangos/input_handlers/yt.py b/tangos/input_handlers/yt.py index 4389566e..56c40beb 100644 --- a/tangos/input_handlers/yt.py +++ b/tangos/input_handlers/yt.py @@ -176,6 +176,27 @@ def _load_halo_cat_without_caching(self, ts_extension, snapshot_file): def get_properties(self): return {} + def available_object_property_names_for_timestep(self, ts_extension, object_typetag): + h, _ = self._load_halo_cat(ts_extension, object_typetag) + return [fn for ft, fn in h.field_list if ft == "halos"] + + + def iterate_object_properties_for_timestep(self, ts_extension, object_typetag, property_names): + h, ad = self._load_halo_cat(ts_extension, object_typetag) + + props_with_ftype = [ + ("halos", name) for name in property_names + ] + + ad.get_data(props_with_ftype) + + Nhalo = len(ad["halos", "particle_identifier"]) + data = zip(range(Nhalo), range(Nhalo), *( + ad[_] for _ in props_with_ftype + )) + + return data + class YtRamsesRockstarInputHandler(YtInputHandler): patterns = ["output_0????"] auxiliary_file_patterns = ["halos_*.bin"] From bba70942016b3dfa94a9d4b63bf98e16ddca8053 Mon Sep 17 00:00:00 2001 From: Corentin Cadiou Date: Mon, 25 Sep 2023 13:39:19 +0200 Subject: [PATCH 06/13] Attempt to cache parent dataset --- tangos/input_handlers/yt.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tangos/input_handlers/yt.py b/tangos/input_handlers/yt.py index 56c40beb..16b02bfb 100644 --- a/tangos/input_handlers/yt.py +++ b/tangos/input_handlers/yt.py @@ -271,11 +271,19 @@ def enumerate_objects(self, ts_extension, object_typetag="halo", min_halo_partic logger.warn("No halo statistics file found for timestep %r", ts_extension) logger.warn(" => enumerating %ss directly using yt", object_typetag) - catalogue, catalogue_data = self._load_halo_cat(ts_extension, object_typetag) + _catalogue, catalogue_data = self._load_halo_cat(ts_extension, object_typetag) num_objects = len(catalogue_data["halos", "virial_radius"]) + # Make sure this isn't garbage collected + _f = self.load_timestep(ts_extension) + for i in range(num_objects): - obj = self.load_object(ts_extension, int(catalogue_data["halos","particle_identifier"][i]), i, object_typetag) + obj = self.load_object( + ts_extension, + int(catalogue_data["halos","particle_identifier"][i]), + i, + object_typetag + ) NDM = len(obj["dark_matter","particle_mass"]) NGas = 0 # cells NStar = len(obj["stars","particle_mass"]) From 599b62da92486129821e645fb7d64754660c25b2 Mon Sep 17 00:00:00 2001 From: Anatole Storck Date: Thu, 2 Nov 2023 11:52:12 +0100 Subject: [PATCH 07/13] removed logging/degugging statements --- tangos/input_handlers/yt.py | 6 ++---- tangos/tools/add_simulation.py | 3 ++- tangos/tools/crosslink.py | 2 -- 3 files changed, 4 insertions(+), 7 deletions(-) diff --git a/tangos/input_handlers/yt.py b/tangos/input_handlers/yt.py index a4202982..7c69f119 100644 --- a/tangos/input_handlers/yt.py +++ b/tangos/input_handlers/yt.py @@ -239,8 +239,6 @@ def _load_halo_cat_without_caching(self, ts_extension, snapshot_file): # Check whether datasets.txt exists (i.e., if rockstar was run with yt) if os.path.exists(self._extension_to_filename("datasets.txt")): fnum = read_datasets(self._extension_to_filename(""),ts_extension) - print (ts_extension) - print (fnum) else: # otherwise, assume a one-to-one correspondence overdir = self._extension_to_filename("") snapfiles = glob.glob(overdir+ts_extension[:2]+len(ts_extension[2:].split('/')[0])*'?') @@ -268,8 +266,8 @@ def enumerate_objects(self, ts_extension, object_typetag="halo", min_halo_partic if self._can_enumerate_objects_from_statfile(ts_extension, object_typetag): yield from self._enumerate_objects_from_statfile(ts_extension, object_typetag) else: - logger.warn("No halo statistics file found for timestep %r", ts_extension) - logger.warn(" => enumerating %ss directly using yt", object_typetag) + #logger.warn("No halo statistics file found for timestep %r", ts_extension) + #logger.warn(" => enumerating %ss directly using yt", object_typetag) _catalogue, catalogue_data = self._load_halo_cat(ts_extension, object_typetag) num_objects = len(catalogue_data["halos", "virial_radius"]) diff --git a/tangos/tools/add_simulation.py b/tangos/tools/add_simulation.py index 86af7eb7..da72849e 100644 --- a/tangos/tools/add_simulation.py +++ b/tangos/tools/add_simulation.py @@ -150,7 +150,8 @@ def add_objects_to_timestep(self, ts, create_class=core.halo.Halo): halos.append(h) with pt.ExclusiveLock("db_write_lock"): - logger.info("Add %d %ss to timestep %r", len(halos), create_class.__name__, ts) + if (create_class.__name__ != "Group"): + logger.info("Add %d %ss to timestep %r", len(halos), create_class.__name__, ts) self.session.add_all(halos) self.session.commit() diff --git a/tangos/tools/crosslink.py b/tangos/tools/crosslink.py index e5af6885..c71b8ace 100644 --- a/tangos/tools/crosslink.py +++ b/tangos/tools/crosslink.py @@ -87,8 +87,6 @@ def need_crosslink_ts(self, ts1, ts2, object_typecode=0): def create_db_objects_from_catalog(self, cat, finder_offset_to_halos_1, finder_offset_to_halos_2, same_d_id): items = [] missing_db_object = 0 - #import pdb - #pdb.set_trace() for i, possibilities in enumerate(cat): h1 = finder_offset_to_halos_1.get(i, None) for cat_i, weight in possibilities: From 6d8370cf99938905b13134b75f6652d8e70f2b23 Mon Sep 17 00:00:00 2001 From: Anatole Storck Date: Thu, 2 Nov 2023 12:27:52 +0100 Subject: [PATCH 08/13] Update add_simulation.py --- tangos/tools/add_simulation.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tangos/tools/add_simulation.py b/tangos/tools/add_simulation.py index da72849e..86af7eb7 100644 --- a/tangos/tools/add_simulation.py +++ b/tangos/tools/add_simulation.py @@ -150,8 +150,7 @@ def add_objects_to_timestep(self, ts, create_class=core.halo.Halo): halos.append(h) with pt.ExclusiveLock("db_write_lock"): - if (create_class.__name__ != "Group"): - logger.info("Add %d %ss to timestep %r", len(halos), create_class.__name__, ts) + logger.info("Add %d %ss to timestep %r", len(halos), create_class.__name__, ts) self.session.add_all(halos) self.session.commit() From 3e51429b15233fbf483ee1bfc5de7048b33db634 Mon Sep 17 00:00:00 2001 From: Anatole Storck Date: Fri, 10 Nov 2023 15:51:41 +0100 Subject: [PATCH 09/13] Yield from statfile if it exists --- tangos/input_handlers/yt.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/tangos/input_handlers/yt.py b/tangos/input_handlers/yt.py index 7c69f119..7d305852 100644 --- a/tangos/input_handlers/yt.py +++ b/tangos/input_handlers/yt.py @@ -132,12 +132,6 @@ def match_objects( cat.append(list(zip(idhalo2, weights))) - #mport pdb - #pdb.set_trace() - - #import pdb - #pdb.set_trace() - return cat def enumerate_objects(self, ts_extension, object_typetag="halo", min_halo_particles=config.min_halo_particles): @@ -182,6 +176,11 @@ def available_object_property_names_for_timestep(self, ts_extension, object_type def iterate_object_properties_for_timestep(self, ts_extension, object_typetag, property_names): + try: + yield from super().iterate_object_properties_for_timestep(ts_extension, object_typetag, property_names) + return + except OSError: + pass h, ad = self._load_halo_cat(ts_extension, object_typetag) props_with_ftype = [ @@ -191,11 +190,10 @@ def iterate_object_properties_for_timestep(self, ts_extension, object_typetag, p ad.get_data(props_with_ftype) Nhalo = len(ad["halos", "particle_identifier"]) - data = zip(range(Nhalo), range(Nhalo), *( + yield from zip(range(Nhalo), range(Nhalo), *( ad[_] for _ in props_with_ftype )) - - return data + class YtRamsesRockstarInputHandler(YtInputHandler): patterns = ["output_0????"] From fd30444733c080faf4c84a966242516486001419 Mon Sep 17 00:00:00 2001 From: Anatole Storck Date: Fri, 10 Nov 2023 15:57:04 +0100 Subject: [PATCH 10/13] removed random space --- tangos/input_handlers/yt.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tangos/input_handlers/yt.py b/tangos/input_handlers/yt.py index 7d305852..d3baedbf 100644 --- a/tangos/input_handlers/yt.py +++ b/tangos/input_handlers/yt.py @@ -193,7 +193,7 @@ def iterate_object_properties_for_timestep(self, ts_extension, object_typetag, p yield from zip(range(Nhalo), range(Nhalo), *( ad[_] for _ in props_with_ftype )) - + class YtRamsesRockstarInputHandler(YtInputHandler): patterns = ["output_0????"] From 83d4fbcd8c76b6b0fcc3a8a4f8b020bb57c0e238 Mon Sep 17 00:00:00 2001 From: Anatole Storck Date: Mon, 13 Nov 2023 14:06:37 +0100 Subject: [PATCH 11/13] Apply suggestions from code review Co-authored-by: Corentin Cadiou --- tangos/input_handlers/yt.py | 54 +++++++--------------------------- tangos/tools/add_simulation.py | 2 +- 2 files changed, 12 insertions(+), 44 deletions(-) diff --git a/tangos/input_handlers/yt.py b/tangos/input_handlers/yt.py index d3baedbf..50b028a0 100644 --- a/tangos/input_handlers/yt.py +++ b/tangos/input_handlers/yt.py @@ -102,9 +102,7 @@ def match_objects( #mask = np.in1d(ids1, members2) mask = np.in1d(members2, ids1) if mask.sum() == 0: - cat.append( - [] - ) + cat.append([]) continue # Get the halo ids of the particles in the other snapshot @@ -200,38 +198,9 @@ class YtRamsesRockstarInputHandler(YtInputHandler): auxiliary_file_patterns = ["halos_*.bin"] def load_timestep_without_caching(self, ts_extension, mode=None): - from yt.data_objects.particle_filters import add_particle_filter if mode is not None: raise ValueError("Custom load modes are not supported with yt") - f = yt.load(self._extension_to_filename(ts_extension)) - - def Stars(pfilter, data): - filter = data[("all", "particle_family")] == 2 # DM = 1, Stars = 2 - return filter - - add_particle_filter("stars", function=Stars, filtered_type='all', \ - requires=["particle_family"]) - - def AllDarkMatter(pfilter, data): - #filter = np.logical_or(data[("all", "particle_type")] == 4,data[("all", "particle_type")] == 1) # DM = 1, Stars = 2 - filter = data[("all", "particle_family")] == 1 - return filter - - add_particle_filter("dark_matter", function=AllDarkMatter, filtered_type='all', \ - requires=["particle_family"]) - - def MustRefineParticles(pfilter, data): - filter = data[("all", "particle_family")] == 4 - return filter - - add_particle_filter("mrp_dark_matter", function=MustRefineParticles, filtered_type='all', \ - requires=["particle_family"]) - - f.add_particle_filter("stars") - f.add_particle_filter("dark_matter") - f.add_particle_filter("mrp_dark_matter") - - return f + return yt.load(self._extension_to_filename(ts_extension)) def _load_halo_cat_without_caching(self, ts_extension, snapshot_file): # Check whether datasets.txt exists (i.e., if rockstar was run with yt) @@ -247,13 +216,11 @@ def _load_halo_cat_without_caching(self, ts_extension, snapshot_file): rockfiles = np.array(rockfiles)[sortord] timestep_ind = np.argwhere(np.array([s.split('/')[-1] for s in snapfiles])==ts_extension.split('/')[0])[0] fnum = int(rockfiles[timestep_ind][0].split('.')[0].split('_')[-1]) - cat = yt.load(self._extension_to_filename("halos_"+str(fnum)+".0.bin")) + cat = yt.load(self._extension_to_filename(f"halos_{fnum}.0.bin")) cat_data = cat.all_data() # Check whether rockstar was run with Behroozi's distribution or Wise's if np.any(cat_data["halos","particle_identifier"]<0): - del cat - del cat_data - cat = yt.load(self._extension_to_filename("halos_"+str(fnum)+".0.bin")) + cat = yt.load(self._extension_to_filename(f"halos_{fnum}.0.bin")) cat.parameters['format_revision'] = 2 # cat_data = cat.all_data() return cat, cat_data @@ -280,19 +247,20 @@ def enumerate_objects(self, ts_extension, object_typetag="halo", min_halo_partic i, object_typetag ) - NDM = len(obj["dark_matter","particle_mass"]) + NDM = len(obj["DM", "particle_ones"]) NGas = 0 # cells - NStar = len(obj["stars","particle_mass"]) + NStar = len(obj["star", "particle_ones"]) if NDM + NGas + NStar> min_halo_particles: yield i, int(catalogue_data["halos","particle_identifier"][i]), NDM, NStar, NGas def load_object(self, ts_extension, finder_id, finder_offset, object_typetag='halo', mode=None): f = self.load_timestep(ts_extension, mode) cat, cat_dat = self._load_halo_cat(ts_extension, object_typetag) - center = cat_dat["halos","particle_position"][cat_dat["halos","particle_identifier"]==finder_id][0] - center+=f.domain_left_edge-cat.domain_left_edge - radius = cat_dat["halos","virial_radius"][cat_dat["halos","particle_identifier"]==finder_id][0] - return f.sphere(center.in_cgs(), radius.in_cgs()) + index = np.argwhere(cat_dat["halos", "particle_identifier"] == finder_id)[0, 0] + center = cat_dat["halos","particle_position"][index] + center += f.domain_left_edge - cat.domain_left_edge + radius = cat_dat["halos", "virial_radius"][index] + return f.sphere(center, radius) class YtChangaAHFInputHandler(YtInputHandler): diff --git a/tangos/tools/add_simulation.py b/tangos/tools/add_simulation.py index da72849e..ee44f951 100644 --- a/tangos/tools/add_simulation.py +++ b/tangos/tools/add_simulation.py @@ -150,7 +150,7 @@ def add_objects_to_timestep(self, ts, create_class=core.halo.Halo): halos.append(h) with pt.ExclusiveLock("db_write_lock"): - if (create_class.__name__ != "Group"): + if create_class.__name__ != "Group": logger.info("Add %d %ss to timestep %r", len(halos), create_class.__name__, ts) self.session.add_all(halos) self.session.commit() From c32192a515563dcdc14f1ab2fee78a5b5c4810b3 Mon Sep 17 00:00:00 2001 From: Anatole Storck Date: Mon, 13 Nov 2023 14:07:57 +0100 Subject: [PATCH 12/13] removed commented out logging --- tangos/tools/add_simulation.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tangos/tools/add_simulation.py b/tangos/tools/add_simulation.py index ee44f951..86af7eb7 100644 --- a/tangos/tools/add_simulation.py +++ b/tangos/tools/add_simulation.py @@ -150,8 +150,7 @@ def add_objects_to_timestep(self, ts, create_class=core.halo.Halo): halos.append(h) with pt.ExclusiveLock("db_write_lock"): - if create_class.__name__ != "Group": - logger.info("Add %d %ss to timestep %r", len(halos), create_class.__name__, ts) + logger.info("Add %d %ss to timestep %r", len(halos), create_class.__name__, ts) self.session.add_all(halos) self.session.commit() From 36cf449e1cedf02f3ba6e518f22abc4210b60f68 Mon Sep 17 00:00:00 2001 From: Anatole Storck Date: Mon, 13 Nov 2023 14:09:10 +0100 Subject: [PATCH 13/13] removed more commented out logging --- tangos/input_handlers/yt.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tangos/input_handlers/yt.py b/tangos/input_handlers/yt.py index 50b028a0..2dcc2ea4 100644 --- a/tangos/input_handlers/yt.py +++ b/tangos/input_handlers/yt.py @@ -231,8 +231,8 @@ def enumerate_objects(self, ts_extension, object_typetag="halo", min_halo_partic if self._can_enumerate_objects_from_statfile(ts_extension, object_typetag): yield from self._enumerate_objects_from_statfile(ts_extension, object_typetag) else: - #logger.warn("No halo statistics file found for timestep %r", ts_extension) - #logger.warn(" => enumerating %ss directly using yt", object_typetag) + logger.warn("No halo statistics file found for timestep %r", ts_extension) + logger.warn(" => enumerating %ss directly using yt", object_typetag) _catalogue, catalogue_data = self._load_halo_cat(ts_extension, object_typetag) num_objects = len(catalogue_data["halos", "virial_radius"])