Understanding ndarray.compress() method in NumPy (5 examples)

Updated: February 26, 2024 By: Guest Contributor Post a comment

Introduction

NumPy is a fundamental package for scientific computing in Python. It offers a powerful N-dimensional array object, which is a multi-dimensional container of items of the same type and size. One of the useful methods provided by NumPy’s ndarray is compress(), which returns a new array containing elements selected from an input array along a specified axis, where a boolean condition is true. In this tutorial, we will delve into the compress() method through five progressively advanced examples.

Syntax & Parameters of compress() Method

Before diving into examples, it’s crucial to understand the syntax of the compress() method:

ndarray.compress(condition, axis=None, out=None)
  • condition: A 1-D array of boolean values that indicates which elements to select.
  • axis: The axis along which to perform the selection. If None, the input array is flattened before the selection.
  • out: An optional output array to place the result in.

Example 1: Basic Usage

Let’s start with the simplest example, selecting elements based on a condition.

import numpy as np

# Sample array
arr = np.array([1, 2, 3, 4, 5])
# Condition: Elements greater than 3
condition = arr > 3

# Using compress
result = arr.compress(condition)
print(result)

Output:

[4 5]

This example demonstrates the basic usage of compress(), where we filtered the array elements greater than 3.

Example 2: Using Axis Parameter

For multidimensional arrays, the axis parameter becomes particularly useful. Here, we’ll select elements from a 2D array.

import numpy as np

# Creating a 2D array
arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
# Condition: Elements greater than 4
condition = arr > 4

# Using compress along axis 0
result = arr.compress(condition, axis=0)
print(result)

Note: In this example, the condition is applied column-wise because the condition array is implicitly flattened when applying along a specific axis. Therefore, the output might not be as expected. To apply a condition row-wise or column-wise, consider using boolean indexing instead.

Example 3: Compress with Flattened Input

When you do not specify an axis, compress() treats the input array as if it were flattened. This example demonstrates that behavior.

import numpy as np

# Creating a 2D array
arr = np.array([[1, 2], [3, 4]])
# Condition: Elements > 2
condition = np.array([True, False, True, False])

# Compress without specifying the axis
flattened_result = arr.compress(condition)
print(flattened_result)

Output:

[1 3]

This example highlights how compress() behaves differently when no axis is specified, effectively flattening the input array before applying the condition.

Example 4: Using out Parameter

The out parameter allows you to place the result of compress() into a pre-allocated array. This can be useful for optimizing memory usage in scenarios where the size of the output is known in advance.

import numpy as np

# Sample array
arr = np.array([1, 2, 3, 4, 5])
# Condition: Elements greater than 2
condition = arr > 2
# Pre-allocated output array
out_array = np.empty(3, dtype=int)

# Using compress with the out parameter
arr.compress(condition, out=out_array)
print(out_array)

Output:

[3 4 5]

This example demonstrates the use of the out parameter, showcasing how to direct the result of compress() into an existing array.

Example 5: Advanced Filtering with compress()

For more complex conditionals, you can combine compress() with logical operators to perform advanced filtering. This example shows how to filter array elements that meet multiple conditions.

import numpy as np

# Sample array
arr = np.array([1, 2, 3, 4, 5])
# Condition: Elements greater than 2 and divisible by 2
condition = np.logical_and(arr > 2, arr % 2 == 0)

# Applying compress with complex condition
result = arr.compress(condition)
print(result)

Output:

[4]

This advanced example combines multiple conditions to demonstrate the versatility of compress() for complex array filtering tasks.

Conclusion

The compress() method in NumPy provides a powerful tool for selective element extraction based on conditions, allowing for efficient and flexible array manipulations. Through these examples, we’ve explored the method’s utility from basic to advanced scenarios, showcasing its capabilities in handling diverse data filtering needs.