NumPy: How to Split an Array Vertically (row-wise)

Introduction

In the realm of scientific computing with Python, NumPy is a cornerstone library that offers comprehensive support for working with arrays. One common array operation is splitting, which allows you to divide an array into several smaller arrays. This tutorial will guide you through splitting a NumPy array vertically (or row-wise), including a range of code examples from basic to advanced use cases.

Understanding Vertical Split

To split an array vertically means to divide it along its rows. Imagine a matrix or a 2D array; a vertical split will cut it across the horizontal axis, resulting in sub-arrays that each contain some number of full rows from the original array. NumPy provides a specific function to perform this operation, which is vsplit().

Before getting started, ensure you have NumPy installed:

pip install numpy

Basic Vertical Splitting

Let’s take a look at a basic example. Assume we have a 2D array with an even number of rows:

import numpy as np

# Create an example 2D array
array = np.array([[1, 2], [3, 4], [5, 6], [7, 8]])
print("Original array:")
print(array)
 

You can now divide this array into two smaller 2D arrays with np.vsplit():

# Split the array into 2 equally sized sub-arrays
arrays = np.vsplit(array, 2)
print("Splitted arrays:")
for a in arrays:
    print(a)

This will output:

Original array:
[[1 2]
 [3 4]
 [5 6]
 [7 8]]

Splitted arrays:
[[1 2]
 [3 4]]
[[5 6]
 [7 8]]

Splitting Into Unequal Sub-arrays

What if you need to split your array into uneven chunks? This is also possible with vsplit(). Suppose you want the first sub-array to have one row and the second to have the rest:

# Split the array with an array section index
arrays = np.vsplit(array, [1])
print("Splitted arrays:")
for a in arrays:
    print(a)

The output now will be:

Original array:
[[1 2]
 [3 4]
 [5 6]
 [7 8]]

Splitted arrays:
[[1 2]]
[[3 4]
 [5 6]
 [7 8]]

Vertical Splitting with array_split()

But there’s more! If you want to split your array into a specific number of sub-arrays that may not result in equal divisions, you can use np.array_split(). Slightly different from vsplit(), this function allows for an uneven split while always trying to balance the sub-arrays sizes:

# Split the array into 3 sub-arrays using array_split()
arrays = np.array_split(array, 3)
print("Splitted arrays:")
for a in arrays:
    print(a)

The output demonstrates an uneven split into three parts:

Original array:
[[1 2]
 [3 4]
 [5 6]
 [7 8]]

Splitted arrays:
[[1 2]
 [3 4]]
[[5 6]]
[[7 8]]

Advanced Splitting Scenarios

In more sophisticated use cases, such as working with multi-dimensional arrays, you might need to handle vertical splitting differently. For an array with more than two dimensions, you would use dsplit() for a depth-wise split, which corresponds to splitting along the third axis.

Example: 3D Array Splitting

import numpy as np

# Create an example 3D array
array_3d = np.arange(24).reshape(2, 3, 4)
print("Original 3D array:")
print(array_3d)

# Performing a vertical split on the 3D array
arrays_3d = np.dsplit(array_3d, 2)
print("Splitted 3D arrays:")
for a in arrays_3d:
    print(a)

This creates:

Original 3D array:
[[[ 0  1  2  3]
  [ 4  5  6  7]
  [ 8  9 10 11]]

 [[12 13 14 15]
  [16 17 18 19]
  [20 21 22 23]]]

Splitted 3D arrays:
[[[ 0  1]
  [ 4  5]
  [ 8  9]]

 [[12 13]
  [16 17]
  [20 21]]]

[[[ 2  3]
  [ 6  7]
  [10 11]]

 [[14 15]
  [18 19]
  [22 23]]]

Error Handling and Debugging

It’s essential to note that attempting to perform a vertical split on an array with fewer rows than the specified sections will raise an error. Additionally, you’ll want to ensure that the dimensions of your array support the kind of split you’re attempting, or you’ll also encounter errors.

Dealing with Errors

try:
    # Attempting to split into more parts than there are rows
    np.vsplit(array, 5)
except ValueError as e:
    print("Error:", e)

This will produce an error such as:

Error: array split does not result in an equal division

Conclusion

Mastering array splitting in NumPy can significantly facilitate data manipulation and processing tasks in Python. With functions such as vsplit() and array_split(), you have the flexibility to structure your data as needed for various scientific computing and data analysis operations.