From 151dafecd3ae38c62d0fa2e10ff2160b7b17ff5c Mon Sep 17 00:00:00 2001 From: Aaron LI Date: Mon, 1 Jan 2018 16:24:04 +0800 Subject: clusters/mergertree: implement itermain(), lmain, imain(idx) --- fg21sim/extragalactic/clusters/mergertree.py | 45 ++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) (limited to 'fg21sim/extragalactic/clusters') diff --git a/fg21sim/extragalactic/clusters/mergertree.py b/fg21sim/extragalactic/clusters/mergertree.py index 312fb92..265de0d 100644 --- a/fg21sim/extragalactic/clusters/mergertree.py +++ b/fg21sim/extragalactic/clusters/mergertree.py @@ -44,6 +44,51 @@ class MergerTree: self.main = main self.sub = sub + def itermain(self): + """ + Iterate by tracing the main cluster. + """ + maintree = self + subtree = None + while maintree: + main = maintree.data + if subtree is not None: + sub = subtree.data + else: + sub = {} + yield (main, sub) + subtree = maintree.sub + maintree = maintree.main + + @property + def lmain(self): + """ + Return the length (i.e., number of events) along the main cluster. + """ + return len(list(self.itermain())) + + def imain(self, idx): + """ + Trace the main cluster to locate the idx-th (0-based) event, and + return its data (both the main and sub clusters). + + Parameters + ---------- + idx : int + The event index (0-based) along the main cluster to be retrieved. + + Returns + ------- + (main, sub) : dict + The data dictionaries of the main and sub clusters. The ``sub`` + dictionary may be ``None`` if it is an accretion event. + """ + for main, sub in self.itermain(): + if idx == 0: + return (main, sub) + idx -= 1 + raise IndexError("index out of range: %d" % idx) + def save_mtree(mtree, outfile, clobber=False): """ -- cgit v1.2.2