Computes all possible pair contractions, sieves the results based on ``memory_limit`` and returns the lowest cost path. This algorithm scales factorial with respect to the elements in the list ``input_sets``. Parameters ---------- input_sets : list List of sets that
(input_sets, output_set, idx_dict, memory_limit)
| 143 | |
| 144 | |
| 145 | def _optimal_path(input_sets, output_set, idx_dict, memory_limit): |
| 146 | """ |
| 147 | Computes all possible pair contractions, sieves the results based |
| 148 | on ``memory_limit`` and returns the lowest cost path. This algorithm |
| 149 | scales factorial with respect to the elements in the list ``input_sets``. |
| 150 | |
| 151 | Parameters |
| 152 | ---------- |
| 153 | input_sets : list |
| 154 | List of sets that represent the lhs side of the einsum subscript |
| 155 | output_set : set |
| 156 | Set that represents the rhs side of the overall einsum subscript |
| 157 | idx_dict : dictionary |
| 158 | Dictionary of index sizes |
| 159 | memory_limit : int |
| 160 | The maximum number of elements in a temporary array |
| 161 | |
| 162 | Returns |
| 163 | ------- |
| 164 | path : list |
| 165 | The optimal contraction order within the memory limit constraint. |
| 166 | |
| 167 | Examples |
| 168 | -------- |
| 169 | >>> isets = [set('abd'), set('ac'), set('bdc')] |
| 170 | >>> oset = set() |
| 171 | >>> idx_sizes = {'a': 1, 'b':2, 'c':3, 'd':4} |
| 172 | >>> _optimal_path(isets, oset, idx_sizes, 5000) |
| 173 | [(0, 2), (0, 1)] |
| 174 | """ |
| 175 | |
| 176 | full_results = [(0, [], input_sets)] |
| 177 | for iteration in range(len(input_sets) - 1): |
| 178 | iter_results = [] |
| 179 | |
| 180 | # Compute all unique pairs |
| 181 | for curr in full_results: |
| 182 | cost, positions, remaining = curr |
| 183 | for con in itertools.combinations(range(len(input_sets) - iteration), 2): |
| 184 | |
| 185 | # Find the contraction |
| 186 | cont = _find_contraction(con, remaining, output_set) |
| 187 | new_result, new_input_sets, idx_removed, idx_contract = cont |
| 188 | |
| 189 | # Sieve the results based on memory_limit |
| 190 | new_size = _compute_size_by_dict(new_result, idx_dict) |
| 191 | if new_size > memory_limit: |
| 192 | continue |
| 193 | |
| 194 | # Build (total_cost, positions, indices_remaining) |
| 195 | total_cost = cost + _flop_count(idx_contract, idx_removed, len(con), idx_dict) |
| 196 | new_pos = positions + [con] |
| 197 | iter_results.append((total_cost, new_pos, new_input_sets)) |
| 198 | |
| 199 | # Update combinatorial list, if we did not find anything return best |
| 200 | # path + remaining contractions |
| 201 | if iter_results: |
| 202 | full_results = iter_results |
no test coverage detected