Source code for treefarm.ancestry_filter

"""
AncestryFilter functions



"""

#-----------------------------------------------------------------------------
# Copyright (c) ytree development team. All rights reserved.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
#-----------------------------------------------------------------------------

import numpy as np

from yt.utilities.operator_registry import \
    OperatorRegistry

ancestry_filter_registry = OperatorRegistry()

[docs]def add_ancestry_filter(name, function): """ Add an ancestry filter function to the registry. """ ancestry_filter_registry[name] = AncestryFilter(function)
[docs]class AncestryFilter(object): r""" An AncestryFilter takes a halo and a list of ancestors and returns a filtered list of filtered list of ancestors. For example, a filter could return only the most massive ancestor. Parameters ---------- halo: halo data container Data container of the descendent halo. ancestors : list of halo data containers List of data containers for the ancestor halos. The function should return a list of data containers. """
[docs] def __init__(self, function, args=None, kwargs=None): self.function = function self.args = args if self.args is None: self.args = [] self.kwargs = kwargs if self.kwargs is None: self.kwargs = {}
def __call__(self, halo, ancestors): return self.function(halo, ancestors, *self.args, **self.kwargs)
[docs]def most_massive(halo, ancestors): r""" Return only the most massive ancestor. Parameters ---------- halo: halo data container Data container of the descendent halo. ancestors : list of halo data containers List of data containers for the ancestor halos. Returns ------- filtered_ancestors : list of halo data containers List containing data container of most massive halo. """ if len(ancestors) == 0: return [] i_max = np.argmax([my_halo.mass for my_halo in ancestors]) return [ancestors[i_max]]
add_ancestry_filter("most_massive", most_massive)