1import numpy as np
2from scipy.sparse import csr_matrix
3from openfermion.third_party.representability._dualbasis import \
4    DualBasisElement, DualBasis
5
6
7class TMap(object):
8
9    def __init__(self, tensors):
10        """
11        provide a map of tensor name to tensors
12
13        :param tensors: list of tensors
14        :return: TMap object
15        """
16        self.tensors = tensors
17        self._map = dict(map(lambda x: (x.name, x), self.tensors))
18
19    def __getitem__(self, i):
20        return self._map[i]
21
22    def __iter__(self):
23        for tt in self.tensors:
24            yield tt
25
26
27class MultiTensor(object):
28
29    def __init__(self, tensors, dual_basis=DualBasis()):
30        """
31        A collection of tensor objects with maps from name to tensor
32
33        Args:
34            tensors: a dictionary or tuple of tensors and their associated call
35                     name.
36            DualBasisElement dual_basis: the set of linear operators restricting
37                                         the feasible elements of the linear
38                                         space generated by vectorizing the
39                                         tensors.
40        """
41        if not isinstance(tensors, list):
42            raise TypeError("MultiTensor accepts a list")
43
44        # this preserves the order the user passes with the tensors
45        self.tensors = TMap(tensors)
46
47        # since all the tensors are indexed from zero...I need to know their
48        # numbering offset when combined with everything.
49        self.off_set_map = self.make_offset_dict(self.tensors)
50
51        # An iterable object that provides access to the dual basis elements
52        self.dual_basis = dual_basis
53        self.vec_dim = sum([vec.size for vec in self.tensors])
54
55    @staticmethod
56    def make_offset_dict(tensors):
57        if not isinstance(tensors, TMap):
58            raise TypeError("Tensors must be a TMap")
59
60        tally = 0
61        offsets = {}
62        # this is why ordered dicts are great.  Remembering orderings
63        for tensor_value in tensors:
64            offsets[tensor_value.name] = tally
65            tally += tensor_value.size
66
67        return offsets
68
69    def vectorize_tensors(self):
70        """
71        vectorize the tensors
72
73        :returns: a vectorized form of the tensors
74        """
75        vec = np.empty((0, 1))
76        for tensor in self.tensors:
77            vec = np.vstack((vec, tensor.vectorize()))
78        return vec
79
80    def add_dual_elements(self, dual_element):
81        """
82        Add  a dual element to the dual basis
83        """
84        if not isinstance(dual_element, DualBasisElement):
85            raise TypeError(
86                "dual_element variable needs to be a DualBasisElement type")
87
88        # we should extend TMap to add
89        self.dual_basis.elements.extend(dual_element)
90
91    def synthesize_dual_basis(self):
92        """
93        from the list of maps create a m x n sparse matrix for Ax=b
94
95        where x is the vectorized form of all the tensors.  This would be
96        the very last step usually.
97
98        :returns: sparse matrix
99        """
100        # go throught the dual basis list and synthesize each element
101        dual_row_indices = []
102        dual_col_indices = []
103        dual_data_values = []
104
105        # this forms the b-vector of ax + b = c
106        bias_data_values = []
107        # this forms the c-vector of ax + b = c
108        inner_prod_data_values = []
109        for index, dual_element in enumerate(self.dual_basis):
110            dcol, dval = self.synthesize_element(dual_element)
111            dual_row_indices.extend([index] * len(dcol))
112            dual_col_indices.extend(dcol)
113            dual_data_values.extend(dval)
114            inner_prod_data_values.append(float(dual_element.dual_scalar))
115            bias_data_values.append(dual_element.constant_bias)
116        sparse_dual_operator = csr_matrix(
117            (dual_data_values, (dual_row_indices, dual_col_indices)),
118            [index + 1, self.vec_dim])
119
120        sparse_bias_vector = csr_matrix(
121            (bias_data_values, (range(index + 1), [0] * (index + 1))),
122            [index + 1, 1])
123
124        sparse_innerp_vector = csr_matrix(
125            (inner_prod_data_values, (range(index + 1), [0] * (index + 1))),
126            [index + 1, 1])
127
128        return sparse_dual_operator, sparse_bias_vector, sparse_innerp_vector
129
130    def synthesize_element(self, element):
131        """
132        Generate the row index and column index for an element
133
134        :param DualBasisElement element: element of the dual basis to vectorize
135        """
136        col_idx = []
137        data_vals = []
138        for tlabel, velement, coeff in element:
139            col_idx.append(self.off_set_map[tlabel] +
140                           self.tensors[tlabel].index_vectorized(*velement))
141            data_vals.append(coeff)
142
143        return col_idx, data_vals
144