Finds the contraction for a given set of input and output sets. Parameters ---------- positions : iterable Integer positions of terms used in the contraction. input_sets : list List of sets that represent the lhs side of the einsum subscript output_set : set
(positions, input_sets, output_set)
| 83 | |
| 84 | |
| 85 | def _find_contraction(positions, input_sets, output_set): |
| 86 | """ |
| 87 | Finds the contraction for a given set of input and output sets. |
| 88 | |
| 89 | Parameters |
| 90 | ---------- |
| 91 | positions : iterable |
| 92 | Integer positions of terms used in the contraction. |
| 93 | input_sets : list |
| 94 | List of sets that represent the lhs side of the einsum subscript |
| 95 | output_set : set |
| 96 | Set that represents the rhs side of the overall einsum subscript |
| 97 | |
| 98 | Returns |
| 99 | ------- |
| 100 | new_result : set |
| 101 | The indices of the resulting contraction |
| 102 | remaining : list |
| 103 | List of sets that have not been contracted, the new set is appended to |
| 104 | the end of this list |
| 105 | idx_removed : set |
| 106 | Indices removed from the entire contraction |
| 107 | idx_contraction : set |
| 108 | The indices used in the current contraction |
| 109 | |
| 110 | Examples |
| 111 | -------- |
| 112 | |
| 113 | # A simple dot product test case |
| 114 | >>> pos = (0, 1) |
| 115 | >>> isets = [set('ab'), set('bc')] |
| 116 | >>> oset = set('ac') |
| 117 | >>> _find_contraction(pos, isets, oset) |
| 118 | ({'a', 'c'}, [{'a', 'c'}], {'b'}, {'a', 'b', 'c'}) |
| 119 | |
| 120 | # A more complex case with additional terms in the contraction |
| 121 | >>> pos = (0, 2) |
| 122 | >>> isets = [set('abd'), set('ac'), set('bdc')] |
| 123 | >>> oset = set('ac') |
| 124 | >>> _find_contraction(pos, isets, oset) |
| 125 | ({'a', 'c'}, [{'a', 'c'}, {'a', 'c'}], {'b', 'd'}, {'a', 'b', 'c', 'd'}) |
| 126 | """ |
| 127 | |
| 128 | idx_contract = set() |
| 129 | idx_remain = output_set.copy() |
| 130 | remaining = [] |
| 131 | for ind, value in enumerate(input_sets): |
| 132 | if ind in positions: |
| 133 | idx_contract |= value |
| 134 | else: |
| 135 | remaining.append(value) |
| 136 | idx_remain |= value |
| 137 | |
| 138 | new_result = idx_remain & idx_contract |
| 139 | idx_removed = (idx_contract - new_result) |
| 140 | remaining.append(new_result) |
| 141 | |
| 142 | return (new_result, remaining, idx_removed, idx_contract) |
no test coverage detected