Return the indices to access the main diagonal of an n-dimensional array. See `diag_indices` for full details. Parameters ---------- arr : array, at least 2-D See Also -------- diag_indices Notes ----- .. versionadded:: 1.4.0 Examples -------
(arr)
| 991 | |
| 992 | @array_function_dispatch(_diag_indices_from) |
| 993 | def diag_indices_from(arr): |
| 994 | """ |
| 995 | Return the indices to access the main diagonal of an n-dimensional array. |
| 996 | |
| 997 | See `diag_indices` for full details. |
| 998 | |
| 999 | Parameters |
| 1000 | ---------- |
| 1001 | arr : array, at least 2-D |
| 1002 | |
| 1003 | See Also |
| 1004 | -------- |
| 1005 | diag_indices |
| 1006 | |
| 1007 | Notes |
| 1008 | ----- |
| 1009 | .. versionadded:: 1.4.0 |
| 1010 | |
| 1011 | Examples |
| 1012 | -------- |
| 1013 | |
| 1014 | Create a 4 by 4 array. |
| 1015 | |
| 1016 | >>> a = np.arange(16).reshape(4, 4) |
| 1017 | >>> a |
| 1018 | array([[ 0, 1, 2, 3], |
| 1019 | [ 4, 5, 6, 7], |
| 1020 | [ 8, 9, 10, 11], |
| 1021 | [12, 13, 14, 15]]) |
| 1022 | |
| 1023 | Get the indices of the diagonal elements. |
| 1024 | |
| 1025 | >>> di = np.diag_indices_from(a) |
| 1026 | >>> di |
| 1027 | (array([0, 1, 2, 3]), array([0, 1, 2, 3])) |
| 1028 | |
| 1029 | >>> a[di] |
| 1030 | array([ 0, 5, 10, 15]) |
| 1031 | |
| 1032 | This is simply syntactic sugar for diag_indices. |
| 1033 | |
| 1034 | >>> np.diag_indices(a.shape[0]) |
| 1035 | (array([0, 1, 2, 3]), array([0, 1, 2, 3])) |
| 1036 | |
| 1037 | """ |
| 1038 | |
| 1039 | if not arr.ndim >= 2: |
| 1040 | raise ValueError("input array must be at least 2-d") |
| 1041 | # For more than d=2, the strided formula is only valid for arrays with |
| 1042 | # all dimensions equal, so we check first. |
| 1043 | if not np.all(diff(arr.shape) == 0): |
| 1044 | raise ValueError("All dimensions of input must be of equal length") |
| 1045 | |
| 1046 | return diag_indices(arr.shape[0], arr.ndim) |