aboutsummaryrefslogtreecommitdiffstats
path: root/fg21sim/extragalactic
diff options
context:
space:
mode:
Diffstat (limited to 'fg21sim/extragalactic')
-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):
"""