Calculate coordinate mapping for graph construction. This function handles the high-level logic behind Blockwise graph construction. The output is a tuple containing: The mapping between input and output block coordinates (`coord_maps`), the axes along which to concatenate for each
(
dims,
out_indices,
numblocks,
argpairs,
concatenate,
)
| 821 | |
| 822 | |
| 823 | def _get_coord_mapping( |
| 824 | dims, |
| 825 | out_indices, |
| 826 | numblocks, |
| 827 | argpairs, |
| 828 | concatenate, |
| 829 | ): |
| 830 | """Calculate coordinate mapping for graph construction. |
| 831 | |
| 832 | This function handles the high-level logic behind Blockwise graph |
| 833 | construction. The output is a tuple containing: The mapping between |
| 834 | input and output block coordinates (`coord_maps`), the axes along |
| 835 | which to concatenate for each input (`concat_axes`), and the dummy |
| 836 | indices needed for broadcasting (`dummies`). |
| 837 | |
| 838 | Used by `make_blockwise_graph` and `Blockwise._cull_dependencies`. |
| 839 | |
| 840 | Parameters |
| 841 | ---------- |
| 842 | dims : dict |
| 843 | Mapping between each index specified in `argpairs` and |
| 844 | the number of output blocks for that index. Corresponds |
| 845 | to the Blockwise `dims` attribute. |
| 846 | out_indices : tuple |
| 847 | Corresponds to the Blockwise `output_indices` attribute. |
| 848 | numblocks : dict |
| 849 | Corresponds to the Blockwise `numblocks` attribute. |
| 850 | argpairs : tuple |
| 851 | Corresponds to the Blockwise `indices` attribute. |
| 852 | concatenate : bool |
| 853 | Corresponds to the Blockwise `concatenate` attribute. |
| 854 | """ |
| 855 | |
| 856 | block_names = set() |
| 857 | all_indices = set() |
| 858 | for name, ind in argpairs: |
| 859 | if ind is not None: |
| 860 | block_names.add(name) |
| 861 | for x in ind: |
| 862 | all_indices.add(x) |
| 863 | assert set(numblocks) == block_names, (numblocks, block_names) |
| 864 | |
| 865 | dummy_indices = all_indices - set(out_indices) |
| 866 | |
| 867 | # For each position in the output space, we'll construct a |
| 868 | # "coordinate set" that consists of |
| 869 | # - the output indices |
| 870 | # - the dummy indices |
| 871 | # - the dummy indices, with indices replaced by zeros (for broadcasting), we |
| 872 | # are careful to only emit a single dummy zero when concatenate=True to not |
| 873 | # concatenate the same array with itself several times. |
| 874 | # - a 0 to assist with broadcasting. |
| 875 | |
| 876 | index_pos, zero_pos = {}, {} |
| 877 | for i, ind in enumerate(out_indices): |
| 878 | index_pos[ind] = i |
| 879 | zero_pos[ind] = -1 |
| 880 |
no test coverage detected