import numpy as np
import types
from ytree.data_structures.fields import FieldContainer
_selection_types = ("forest", "tree", "prog")
[docs]
class NodeContainer:
"""
A container to hold a persistent collection of unrelated TreeNodes.
This is a convenience object to hold a series of TreeNodes
and provide simple field access similar to querying fields
for an entire arbor. Field access for the container is cached
for fast second-time access.
Parameters
----------
nodes : Iterable of TreeNodes
A list, array, or generator of TreeNode objects.
Examples
--------
>>> import ytree
>>> a = ytree.load("tiny_ctrees/locations.dat")
>>> first = a[0]
>>> last = a[-1]
>>> container = a.container([first, last])
>>> for node in container:
>>> print (node["mass"])
1.0231655e+12 Msun
2.0633094e+11 Msun
>>> print (container["mass"])
[1.0231655e+12 2.0633094e+11] Msun
>>> fn = a.save_arbor(trees=container)
>>> a_new = ytree.load(fn)
>>> print (a_new["mass"])
[1.0231655e+12 2.0633094e+11] Msun
>>> import ytree
>>> a = ytree.load("tiny_ctrees/locations.dat")
>>> a_slice = a.container(a[::8])
>>> print (a.size, a_slice.size)
32 4
>>> print (a_slice["mass"])
[1.0231655e+12 6.7683457e+11 2.9294966e+11 2.8014389e+11] Msun
>>> my_tree = a_slice[0]
>>> print (my_tree)
TreeNode[1457223360]
>>> import ytree
>>> a = ytree.load("tiny_ctrees/locations.dat")
>>> my_tree = a[0]
>>> leaves = my_tree.get_leaf_nodes()
>>> leaf_container = a.container(leaves)
>>> print (leaf_container["mass"])
[2.6503598e+08 4.2388490e+08 5.6964032e+08 5.5640288e+08 3.5769786e+08
1.9870504e+08 6.3597126e+08] Msun
>>> print (leaf_container.nodes)
[TreeNode[1169524360], TreeNode[758740170], TreeNode[278238650],
TreeNode[142383462], TreeNode[38885675], TreeNode[19602159],
TreeNode[5057761]]
"""
[docs]
def __init__(self, nodes, arbor=None):
self._nodes = nodes
self.field_data = FieldContainer(arbor)
@property
def nodes(self):
if isinstance(self._nodes, types.GeneratorType):
self._nodes = list(self._nodes)
return self._nodes
def __len__(self):
return len(self.nodes)
@property
def size(self):
return self.__len__()
def __iter__(self):
for node in self.nodes:
yield node
def __getitem__(self, key):
if isinstance(key, str):
if key in _selection_types:
raise SyntaxError("Argument must be a field or integer.")
if key not in self.field_data:
self.field_data[key] = self.field_data.arbor.arr(
[tree[key] for tree in self.nodes]
)
return self.field_data[key]
if isinstance(key, (int, np.integer, slice)):
return self.nodes[key]
else:
raise ValueError(f"Unrecognized argument type: {key} ({type(key)}).")