clusters/mergertree: implement itermain(), lmain, imain(idx)

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):
     """