tree.network

  1from __future__ import annotations
  2
  3from typing import List, Iterable
  4from varname.core import varname
  5
  6class _SliceMixin:
  7    """
  8    **Private mix‑in** that equips a container with NumPy‑style multi‑axis
  9    slicing and broadcasting.
 10
 11    It is mixed into :class:`Tree`, :class:`Network` and :class:`_Mapper`
 12    so they can all share one robust implementation of ``__getitem__``.
 13
 14    --------------------
 15    Supported selectors
 16    --------------------
 17
 18    ╔═══════════════════════╦══════════════════════════════════════════════════╗
 19    ║ Selector type         ║ Behaviour                                        ║
 20    ╠═══════════════════════╬══════════════════════════════════════════════════╣
 21    ║ ``int``               ║ Pick a single element, *collapse* the axis.      ║
 22    ║ ``Tree`` instance     ║ Same as ``int`` but selects the exact object.    ║
 23    ║ ``str``               ║ • If it matches a child’s ``name`` → pick it.    ║
 24    ║                       ║ • Otherwise treat as attribute broadcast.        ║
 25    ║ ``slice``             ║ Keep the axis; later keys decide whether to      ║
 26    ║                       ║ collapse (mimics NumPy’s “slice‑then‑scalar”).   ║
 27    ║ ``list`` / ``_Mapper``║ Advanced (a.k.a. *fancy*) indexing; keeps axis.  ║
 28    ╚═══════════════════════╩══════════════════════════════════════════════════╝
 29
 30    Attribute broadcasting lets you ask, for example::
 31
 32        tree[:, 'mass']          # list of each child.mass
 33        tree[['a', 'b'], :, 'color']
 34
 35    ------------------
 36    Implementation
 37    ------------------
 38    The heavy lifting is done in the private
 39    :py:meth:`_SliceMixin.__array_slice` helper; the public entry point
 40    is ``__getitem__`` which forwards to it.
 41
 42    The class is **deliberately prefixed with an underscore** because
 43    user code should never inherit from it directly; consume it only via
 44    the public classes that already depend on it.
 45    """
 46
 47    # ‑‑ internal implementation pulled from your working array_slice() ‑‑
 48    def __array_slice(self, *keys):
 49        if len(keys) == 1 and isinstance(keys[0], tuple):
 50            keys = keys[0]
 51        if not keys:
 52            return _Mapper(self)
 53
 54        key, *rest = keys
 55        children = list(self)
 56
 57        # normalise first selector ------------------------------------------------
 58        if isinstance(key, Tree):
 59            selectors, grouping = [key], False
 60        elif isinstance(key, slice):
 61            selectors, grouping = children[key], True
 62        elif isinstance(key, (list, _Mapper)):
 63            selectors, grouping = list(key), True
 64        else:
 65            selectors, grouping = [key], False
 66
 67        resolved = []
 68        name_map = {c.name: c for c in children}
 69
 70        for sel in selectors:
 71            if sel is None:
 72                continue
 73            if isinstance(sel, Tree):
 74                resolved.append(sel)
 75            elif isinstance(sel, int):
 76                resolved.append(children[sel])
 77            elif isinstance(sel, slice):
 78                resolved.append(children[sel])
 79            elif isinstance(sel, _Mapper):
 80                resolved.append(sel)
 81            elif isinstance(sel, str):
 82                if sel in name_map:
 83                    resolved.append(name_map[sel])
 84                else:
 85                    resolved.append(
 86                        _Mapper([getattr(ch, sel, None) for ch in children])
 87                    )
 88            else:
 89                raise ValueError(f"Unsupported key type: {type(sel)}")
 90
 91        # recurse ----------------------------------------------------------------
 92        if rest:
 93            resolved = [item.__array_slice(*rest) for item in resolved]
 94
 95            collapse = False
 96            if not grouping:
 97                collapse = True                       # scalar pick
 98            elif isinstance(key, slice):
 99                if not isinstance(rest[0], slice):
100                    collapse = True                   # slice → non‑slice
101
102            if collapse:
103                flat = []
104                for item in resolved:
105                    flat.extend(item if isinstance(item, (list, _Mapper)) else [item])
106                resolved = flat
107
108        return _Mapper(resolved)
109
110    # public entry point --------------------------------------------------------
111    def __getitem__(self, *keys):
112        return self.__array_slice(*keys)
113
114class Tree(_SliceMixin, list):
115    """
116    Efficient and dynamic tree structure with cascading attributes. More complex structure than a binary tree, 
117    and more flexible, resembling a network.
118
119    .. note:: Computations: efficient and dynamic 
120        - [x] `parent` attribute is set to child nodes as they are added.
121        - [x] All other attributes are computed *recursively*, with *vectorisation*.
122        - [x] Attributes are set in a cascade down the descendants.
123        - [x] Attributes are accessed via Python's own hashing routine.
124
125    Examples
126    --------
127    
128    Build a tree branch with two `leaves`
129    >>> leaf = Tree()
130    >>> branch = Tree(leaf, leaf)
131    >>> print(branch)
132    branch: ['leaf', 'leaf']
133    >>> len(branch)
134    2
135
136    Attributes cascade down the `descendants`
137    >>> branch.color = 'green'
138    >>> leaf.color
139    'green'
140    >>> leaf.color = 'crimson'
141    >>> leaf.color
142    'crimson'
143
144    We can climb the tree with the `parent` attribute
145    >>> leaf.parent.color
146    'green'
147
148    Consider a more complicated tree:
149    >>> trunk = Tree(branch, branch, leaf)
150    >>> tree = Tree(trunk)
151
152    We can access the `ancestors`': 
153    >>> print(leaf.ancestors)
154    ancestors: ['leaf', 'trunk', 'tree']
155
156    Note the different ways to access the `root`.
157    >>> leaf.root == leaf.ancestors[-1] 
158    True
159
160    >>> print(tree.descendants)
161    descendants: ['tree', 'trunk', 'branch', 'leaf', 'leaf', 'branch', 'leaf', 'leaf', 'leaf']
162
163    Extracting variables from the leaves:
164    >>> tree.a = 1.0
165    >>> print(tree.a)
166    1.0
167    >>> print(tree.leaves.a)
168    [1.0, 1.0, 1.0, 1.0, 1.0]
169    >>> print(extrinsic(tree.leaves.a))
170    5.0
171    >>> print(intrinsic(tree.leaves.a))
172    1.0
173    """
174
175    def __init__(self, *child: Tree) -> None:
176        """
177        A subclass of list.
178        """
179        super().__init__(_child for _child in child)
180
181        self.__dict__["parent"] = None
182        """
183        Parent attribute set by parent node.
184        """
185        
186        try:
187            self.__dict__["__name__"] = varname()
188        except:
189            self.__dict__["__name__"] = str(list(self))
190
191        self.__setparent__(*child)
192
193    def __eq__(self, other):
194        return (id(self) == id(other))
195
196    def __hash__(self):
197        return id(self)
198
199    @property
200    def name(self):
201        return self.__name__
202
203    def __str__(self):
204        if len(self)>0:
205            return self.__name__ + ': ' + str([child.__name__ for child in self])
206        else:
207            return self.__name__
208
209    @property
210    def root(self) -> Tree:
211        """
212        The last ancestor.
213        """
214        return self.ancestors[-1]
215    
216    @property
217    def __ancestors(self) -> Tree:
218        """
219        This node prepended to its parent's ancestors.
220
221        .. note:: The first ancestor is self.
222        """
223        if self.parent == None:
224            return [self]
225        else:
226            return [self] + self.parent.__ancestors
227
228    @property
229    def ancestors(self) -> "Tree":
230        """
231        This node prepended to its parent’s ancestors.
232        The first element is *self*.
233        """
234        # 1. raw list (self … root)
235        chain = self.__ancestors                     # type: List[Tree]
236
237        # 2. remember current parent pointers
238        old_parent = {node: node.parent for node in chain}
239
240        # 3. create a *display* object without breaking the real graph
241        ancestors = Tree(*chain)                     # pretty wrapper
242        ancestors.__dict__["__name__"] = "ancestors"
243
244        # 4. restore parents exactly as they were
245        for node, par in old_parent.items():
246            node.__dict__["parent"] = par
247
248        return ancestors
249
250    @property
251    def descendants(self) -> Tree:
252        """
253
254        .. note:: The first descendant is self.
255
256        Lists all descendant nodes
257        """
258        descendants = Tree(self)
259        for child in self:
260            descendants += (child.descendants)
261        return descendants
262
263    @property
264    def generations(self) -> List[Tree]:
265        """
266        Descendants ordered by generation.
267
268        """
269        descendants = self.descendants
270        index = [descendant.height for descendant in descendants]
271        generations = [Tree() for _ in range(max(index)+1)]
272        for index, descendant in zip(index, descendants):
273            generations[index].append(descendant)
274        for i, generation in enumerate(generations):
275            generation.__dict__["__name__"] = "generation " + str(i)
276        generations = Tree(*generations)
277        return generations
278
279    @property
280    def leaves(self) -> List[Tree]:
281        """
282        Elements of the tree without descendants.
283        """
284        leaves = _Mapper([descendant for child in self for descendant in child.leaves] if len(self)>0 else [self])
285        return leaves
286
287    @property
288    def n_leaves(self):
289        return len(self.leaves)
290
291    @property
292    def size(self) -> int:
293        """
294        Number of `descendants` of current node, including itself.
295        """
296        return len(self.descendants)
297
298    @property
299    def depth(self) -> int:
300        """
301        Number of ancestors of current node excluding itself.
302        """
303        return len(self.ancestors) - 1
304
305    @property
306    def height(self) -> int:
307        """
308        Number of generations of descendants.
309        """
310        if len(self)>0:
311            return 1 + max([child.height for child in self])
312        else:
313            return 0
314
315    def insert(self, index, child):
316        """Insert child before index and link to it as its parent."""
317        super().insert(index, child)
318        self.__setparent__(*child)
319        # try:
320        #     self.__dict__[child.__name__] = child
321        # except:
322        #     pass
323
324    def append(self, child: Tree):
325        """Append child and link to it as its parent."""
326        super().append(child)
327        self.__setparent__(*child)
328        # try:
329        #     self.__dict__[child.__name__] = child
330        # except:
331        #     pass
332
333    def extend(self, children: Iterable[Tree]):
334        """Extend children with an Iterable of Trees, linking 
335        to the children as a parent.
336        """
337        self.__setparent__(*children)
338        super().extend(children)
339        # for child in children:
340        #     try:
341        #         self.__dict__[child.__name__] = child
342        #     except:
343        #         pass
344    
345    def __getattr__(self, name: str):
346        try:
347            return self.__dict__[f"{name}"]
348        except KeyError:
349            raise AttributeError(name + f" not defined!")
350
351    def __setattr__(self, name, value):
352        """set attributes down the lineage"""
353        self.__dict__[f"{name}"] = value
354        [item.__setattr__(name, value) for item in self]
355
356    def __setparent__(self, *children):
357        """
358        Set childs' parent
359        """
360        def f(child):
361            child.__dict__["parent"] = self
362            child.__setparent__(*child)
363        list(map(f, children))
364
365class Graph(Tree):
366    """
367    Add `add` and `hop` functions to Tree to turn the leaves into a Graph
368
369    Examples
370    --------
371    
372    Consider a complicated graph:
373    >>> A = Graph()
374    >>> B = Graph()
375    >>> graph = Graph(A, B)
376
377    Add onsite term:
378    >>> graph.add(μ=0.2)
379
380    Add hopping terms:
381    >>> A.hop(B, t=1)
382    >>> B.hop(A, t=-1)
383
384    Parameter describing diagonal elements are given by the following set
385    >>> A.diagonal
386    {'μ'}
387
388    Parameter describing off-diagonal elements are given by the following
389    dictionary, where the dictionary value is the Tree node specified by `hop`:
390    >>> A.offdiagonal
391    {'t': {[]}}
392
393    The parameter itself is accessed by
394    >>> A.μ
395    0.2
396
397    >>> graph.leaves.t
398    [1, -1]
399    >>> graph.leaves.t *= 2
400    >>> graph.leaves.t
401    [2, -2]
402    """
403    _diagonal = {}
404    _offdiagonal = {}
405
406    def __init__(self, *child: Tree) -> None:
407        """
408        A subclass of Tree.
409        """
410        super().__init__(*child)
411        
412        try:
413            self.__dict__["__name__"] = varname()
414        except:
415            self.__dict__["__name__"]
416
417    def __reset__(self):
418        self._diagonal = {}
419        self._offdiagonal = {}
420
421    @property
422    def diagonal(self):
423        _diagonal = set()
424        for key, items in self._diagonal.items():
425            for node, item in items.items():
426                if self in node.descendants:
427                    _diagonal.add(key)
428        return _diagonal
429
430    @property
431    def offdiagonal(self):
432        _offdiagonal = dict()
433        for key, items in self._offdiagonal.items():
434            for node, neighbours in items.items():
435                for neighbour, items in neighbours.items():
436                    if self in node.descendants:
437                        try:
438                            _offdiagonal[key].add(neighbour)
439                        except:
440                            _offdiagonal[key] = set([neighbour])
441        return _offdiagonal
442
443    def __setattr__(self, name, value, update = True):
444        """
445        Set attributes down the lineage, and add attribute to list of
446        attributes to be updates
447        """
448        if update:
449            self.__cache__(name, value)
450
451        self.__dict__[f"{name}"] = value
452        
453        [item.__setattr__(name, value, update = False) for item in self]
454
455    def __cache__(self, key, value):
456
457        try:
458            self._diagonal[key][self]["value"] = value
459        except:
460            pass
461
462        try:
463            neighbours = self._offdiagonal[key][self]
464            for neighbour, items in neighbours.items():
465                self._offdiagonal[key][self][neighbour]["value"] = value
466        except:
467            pass
468
469    def add(self, **kwarg):
470        """Add vertex term as a parameter described by a key-value pair"""
471        
472        if len(kwarg)==0:
473            raise ValueError("Requires a kwarg!")
474
475        if len(kwarg)>1:
476            raise ValueError("Accepts only one kwarg!")
477
478        key, value = list(kwarg.items())[0]
479
480        if not key in self._diagonal:
481            self._diagonal[key] = {}
482
483        if not self in self._diagonal[key]:
484            self._diagonal[key][self] = {}
485
486        self.__setattr__(key, value)
487
488    def hop(self, neighbour: Tree, **kwarg):
489        """Add directed edge as a neighbour and a parameter (described by a
490        key-value pair)
491        """
492
493        if len(kwarg)==0:
494            raise ValueError("Requires a kwarg!")
495
496        if len(kwarg)>1:
497            raise ValueError("Accepts only one kwarg!")
498
499        if self.n_leaves!=neighbour.n_leaves:
500            raise ValueError("Self and neighbour have different number of leaves. Hopping is not well-defined!")
501
502        key, value = list(kwarg.items())[0]
503
504        if not key in self._offdiagonal:
505            self._offdiagonal[key] = {}
506
507        if not self in self._offdiagonal[key]:
508            self._offdiagonal[key][self] = {}
509
510        if not neighbour in self._offdiagonal[key][self]:
511            self._offdiagonal[key][self][neighbour] = {}
512
513        self.__setattr__(key, value)
514
515class _Mapper(_SliceMixin, list):
516    """
517    **Private list‑like wrapper** returned by all slicing operations and by
518    :pyattr:`Tree.leaves`.
519
520    It behaves exactly like a Python ``list`` *plus* a few tricks:
521
522    * **Vectorised attribute access / assignment**
523
524    >>> tree = Tree(Tree(Tree()),Tree(Tree()))
525    >>> branches = tree[:, 0]
526    >>> print(branches)
527    [[], []]
528    >>> branches.color = "green"      # propagated to every branch
529    >>> print(branches.color)
530    ['green', 'green']
531
532    * **Simple numeric broadcasting**
533
534      ``__iadd__``, ``__imul__`` … modify numbers (or NumPy / torch arrays)
535      in place across the collection.
536
537    * **Recursive slicing**
538
539      Because ``_Mapper`` also inherits :class:`_SliceMixin`, you can keep
540      chaining subscripts indefinitely::
541
542          leaf_masses = tree[:, :, :, 'mass']
543
544    The class is private because library users never need to create it
545    manually; it is only surfaced as a return value.
546    """
547
548    def __init__(self, *items) -> None:
549        """
550        A subclass of list.
551        """
552        if type(items)==tuple and len(items)==1:
553            items = items[0]
554        if type(items) == list:
555            items = [_Mapper(item) if type(item)==list else item for item in items]
556        super().__init__(items)
557
558    def __getattr__(self, name: str):
559        return _Mapper([getattr(item, name) if type(item)!=type(None) else None for item in self])
560
561    def __setattr__(self, name, value):
562        """Set attributes down the lineage"""
563        if type(value) is list:
564            if len(value) == len(self):
565                [setattr(self[i], name, value[i]) for i in range(len(self)) if self[i] is not None]
566        else:
567            [setattr(item, name, value) for item in self if item is not None]
568
569    def __imul__(self, scaler):
570       self = [n*scaler for n in self]
571       return self
572
573    def __itruediv__(self, scaler):
574       self = [n/scaler for n in self]
575       return self
576
577    def __iadd__(self, scaler):
578       self = [n+scaler for n in self]
579       return self
580
581    def __isub__(self, scaler):
582       self = [n-scaler for n in self]
583       return self
584
585    @property
586    def names(self):
587        """Returns name of each element"""
588        return self.name
589
590def extrinsic(items):
591    r"""
592    $$\sum_{i\in \text{items}} i$$
593
594    If items is a numpy/torch object, optimisations are employed.
595    """
596    try:
597        return items.sum()
598    except:
599        return sum(items)
600
601def intrinsic(items):
602    r"""
603    $$\frac{1}{\\#\text{items}}\sum_{i\in \text{items}} i$$
604
605    If items is a numpy/torch object, optimisations are employed.
606    """
607    try:
608        return items.mean()
609    except:
610        return extrinsic(items) / len(items)
611
612class Network(Graph):
613    """
614    Access Network nodes via a powerful slicing method
615
616    Examples
617    --------
618
619    >>> up = Network()
620    >>> dn = Network()
621    >>> branch = Network(up, dn)
622    >>> trunk_1 = Network(branch, branch)
623    >>> trunk_2 = Network(branch, branch)
624    >>> tree = Network(trunk_1, trunk_2)
625
626    >>> items = tree[['trunk_1', 'trunk_2']]
627    >>> items
628    [[[[], []], [[], []]], [[[], []], [[], []]]]
629    >>> items.__name__
630    ['trunk_1', 'trunk_2']
631    >>> items.root.name
632    [['tree'], ['tree']]
633
634    >>> items = tree[[trunk_1, trunk_2], :]
635    >>> items.name
636    [['branch', 'branch'], ['branch', 'branch']]
637
638    >>> items = tree[[trunk_1, trunk_2], :, up]
639    >>> items.name
640    [['up', 'up'], ['up', 'up']]
641    >>> items = tree[trunk_2, :, up]
642    >>> items.name
643    ['up', 'up']
644
645    >>> items = tree[:, :, up]
646    >>> items.name
647    [['up', 'up'], ['up', 'up']]
648
649    >>> items = tree[:, :]
650    >>> items.name
651    [['branch', 'branch'], ['branch', 'branch']]
652
653    >>> items = tree[:]
654    >>> items.name
655    ['trunk_1', 'trunk_2']
656
657    >>> tree[['trunk_2', 'trunk_2']].a = 1.2
658    >>> up.a
659    1.2
660    >>> tree.leaves.a
661    [1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2, 1.2]
662
663    Many ways to slice:
664    >>> tree[['trunk_1', 'trunk_2']].names
665    ['trunk_1', 'trunk_2']
666    >>> tree[['trunk_1', 'trunk_2'], 0].names
667    [['branch'], ['branch']]
668    >>> tree[['trunk_1', 'trunk_2'], 0,0].names
669    [['up'], ['up']]
670    >>> tree[:, :, 1].names
671    [['dn', 'dn'], ['dn', 'dn']]
672    >>> tree[:, :, :].names
673    [[['up', 'dn'], ['up', 'dn']], [['up', 'dn'], ['up', 'dn']]]
674    >>> tree[['trunk_1', trunk_2], :, :].names
675    [[['up', 'dn'], ['up', 'dn']], [['up', 'dn'], ['up', 'dn']]]
676    >>> tree['trunk_1', :].names
677    ['branch', 'branch']
678    >>> tree[('trunk_1')].names
679    ['trunk_1']
680    >>> tree[['trunk_1']].names
681    ['trunk_1']
682    >>> tree[(trunk_1)].names
683    ['trunk_1']
684    >>> tree['trunk_1'].names
685    ['trunk_1']
686    >>> tree[trunk_1].names
687    ['trunk_1']
688    >>> tree[1].names
689    ['trunk_2']
690    >>> tree[:].names
691    ['trunk_1', 'trunk_2']
692    >>> tree[0:2].names
693    ['trunk_1', 'trunk_2']
694    """
695    def __init__(self, *child: Network) -> None:
696        """
697        A subclass of Tree.
698        """
699        super().__init__(*child)
700
701        try:
702            self.__dict__["__name__"] = varname()
703        except:
704            self.__dict__["__name__"]
705
706        for child in child:
707            self.__dict__[child.__name__] = child
708
709    def __set_indices__(self):
710        """
711        Consider:
712        >>> A = Network()
713        >>> B = Network()
714        >>> tree = Network(A, B)
715
716        Set the indices
717        >>> tree.__set_indices__()
718
719        The indices are then given by
720        >>> A.index
721        0
722        >>> B.index
723        1
724
725        """
726        self.initialised = True
727        def f(l, i): l.__dict__["index"] = i
728        [f(leaf, i) for i, leaf in enumerate(self.leaves)]
729
730    @property
731    def indices(self):
732        if not hasattr(self, "_indices"):
733            self.__dict__["_indices"] = [leaf.index for leaf in self.leaves]
734        return self._indices