tree.linalg

 1from __future__ import annotations
 2from varname.core import varname
 3from tree.network import Network
 4from numpy import zeros as linalg_zeros
 5from numpy import transpose as linalg_transpose
 6
 7class MatrixTree(Network):
 8    """
 9    Represent the MatrixTree as a Matrix
10
11    Examples
12    --------
13    
14    Consider a complicated tree:
15    >>> A = MatrixTree()
16    >>> B = MatrixTree()
17    >>> tree = MatrixTree(A, B)
18    >>> tree.add(μ=0)
19    >>> A.hop(B, t=1)
20
21    The matrix representation of the tree is given by
22    >>> tree.matrix
23    array([[0., 1.],
24           [0., 0.]])
25    """
26
27    def __init__(self, *child: MatrixTree) -> None:
28        """
29        A subclass of MatrixTree.
30        """
31        super().__init__(*child)
32
33        try:
34            self.__dict__["__name__"] = varname()
35        except:
36            self.__dict__["__name__"]
37
38    @property
39    def matrix(self):
40        """
41        Construct the matrix representation of the tree
42
43        >>> A = MatrixTree()
44        >>> B = MatrixTree()
45        >>> matrix = MatrixTree(A, B)
46        >>> matrix.add(μ=0.6)
47        >>> A.hop(B, t=1)
48        >>> matrix.matrix
49        array([[0.6, 1. ],
50               [0. , 0.6]])
51
52        >>> A.t = 3
53        >>> matrix.matrix
54        array([[0.6, 3. ],
55               [0. , 0.6]])
56
57        >>> A.t *= 3
58        >>> matrix.matrix
59        array([[0.6, 9. ],
60               [0. , 0.6]])
61        """
62        if not hasattr(self, "_matrix"):
63            self.__set_indices__()
64            self.__dict__["_matrix"] = linalg_zeros([self.n_leaves, self.n_leaves])
65
66        diagonals, offdiagonals = self._diagonal, self._offdiagonal
67
68        for key, items in diagonals.items():
69            for node, item in items.items():
70                value = item["value"]
71                try:
72                    previous_value = self._diagonal[key][node]["previous_value"] 
73                    value -= previous_value
74                except:
75                    pass
76
77                if not value == 0:
78                    indices = node.indices
79                    slice = (indices, indices)
80                    self._matrix[*slice] += value
81                    self._diagonal[key][node]["previous_value"] = item["value"]
82
83        for key, items in offdiagonals.items():
84            for node, neighbours in items.items():
85                for neighbour, item in neighbours.items():
86                    value = item["value"]
87                    try:
88                        previous_value = self._offdiagonal[key][node][neighbour]["previous_value"]
89                        value -= previous_value
90                    except:
91                        pass
92
93                    if not value == 0:
94                        slice = (node.indices, neighbour.indices)
95                        self._matrix[*slice] += value
96                        self._offdiagonal[key][node][neighbour]["previous_value"] = item["value"]
97
98        return self._matrix