aboutsummaryrefslogtreecommitdiffstats
path: root/fg21sim
diff options
context:
space:
mode:
authorAaron LI <aly@aaronly.me>2018-01-01 16:24:04 +0800
committerAaron LI <aly@aaronly.me>2018-01-01 16:26:20 +0800
commit151dafecd3ae38c62d0fa2e10ff2160b7b17ff5c (patch)
tree8dd33fa1102d938b9e7c18d983fc67018f4204c8 /fg21sim
parent97b76931c6a16dd4e99ccb254b14eb0e0840efcc (diff)
downloadfg21sim-151dafecd3ae38c62d0fa2e10ff2160b7b17ff5c.tar.bz2
clusters/mergertree: implement itermain(), lmain, imain(idx)
Diffstat (limited to 'fg21sim')
-rw-r--r--fg21sim/extragalactic/clusters/mergertree.py45
1 files changed, 45 insertions, 0 deletions
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):
"""