Split an array into multiple sub-arrays. Please refer to the ``split`` documentation. The only difference between these functions is that ``array_split`` allows `indices_or_sections` to be an integer that does *not* equally divide the axis. For an array of length l that should
(ary, indices_or_sections, axis=0)
| 730 | |
| 731 | @array_function_dispatch(_array_split_dispatcher) |
| 732 | def array_split(ary, indices_or_sections, axis=0): |
| 733 | """ |
| 734 | Split an array into multiple sub-arrays. |
| 735 | |
| 736 | Please refer to the ``split`` documentation. The only difference |
| 737 | between these functions is that ``array_split`` allows |
| 738 | `indices_or_sections` to be an integer that does *not* equally |
| 739 | divide the axis. For an array of length l that should be split |
| 740 | into n sections, it returns l % n sub-arrays of size l//n + 1 |
| 741 | and the rest of size l//n. |
| 742 | |
| 743 | See Also |
| 744 | -------- |
| 745 | split : Split array into multiple sub-arrays of equal size. |
| 746 | |
| 747 | Examples |
| 748 | -------- |
| 749 | >>> x = np.arange(8.0) |
| 750 | >>> np.array_split(x, 3) |
| 751 | [array([0., 1., 2.]), array([3., 4., 5.]), array([6., 7.])] |
| 752 | |
| 753 | >>> x = np.arange(9) |
| 754 | >>> np.array_split(x, 4) |
| 755 | [array([0, 1, 2]), array([3, 4]), array([5, 6]), array([7, 8])] |
| 756 | |
| 757 | """ |
| 758 | try: |
| 759 | Ntotal = ary.shape[axis] |
| 760 | except AttributeError: |
| 761 | Ntotal = len(ary) |
| 762 | try: |
| 763 | # handle array case. |
| 764 | Nsections = len(indices_or_sections) + 1 |
| 765 | div_points = [0] + list(indices_or_sections) + [Ntotal] |
| 766 | except TypeError: |
| 767 | # indices_or_sections is a scalar, not an array. |
| 768 | Nsections = int(indices_or_sections) |
| 769 | if Nsections <= 0: |
| 770 | raise ValueError('number sections must be larger than 0.') from None |
| 771 | Neach_section, extras = divmod(Ntotal, Nsections) |
| 772 | section_sizes = ([0] + |
| 773 | extras * [Neach_section+1] + |
| 774 | (Nsections-extras) * [Neach_section]) |
| 775 | div_points = _nx.array(section_sizes, dtype=_nx.intp).cumsum() |
| 776 | |
| 777 | sub_arys = [] |
| 778 | sary = _nx.swapaxes(ary, axis, 0) |
| 779 | for i in range(Nsections): |
| 780 | st = div_points[i] |
| 781 | end = div_points[i + 1] |
| 782 | sub_arys.append(_nx.swapaxes(sary[st:end], axis, 0)) |
| 783 | |
| 784 | return sub_arys |
| 785 | |
| 786 | |
| 787 | def _split_dispatcher(ary, indices_or_sections, axis=None): |