diff options
Diffstat (limited to 'fg21sim/extragalactic/clusters')
| -rw-r--r-- | fg21sim/extragalactic/clusters/mergertree.py | 45 | 
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):      """  | 
