py4u blog

Mastering `numpy.compress()` in Python: A Detailed Guide

NumPy is the cornerstone of numerical computing in Python, enabling efficient manipulation of multi-dimensional arrays. Among its utility functions, numpy.compress() stands out for selecting elements from an array based on a boolean condition, with support for axis-specific operations. This function is invaluable for tasks like data filtering, feature selection, and advanced array manipulation.

2026-10

Table of Contents#

Syntax and Parameters of numpy.compress()#

The syntax of numpy.compress() is:

numpy.compress(condition, a, axis=None, out=None)

Parameters:#

  1. condition: A 1-D array of booleans (or integers, where non-zero values are treated as True). This array determines which elements of a to keep.
  2. a: The input array (numpy.ndarray) from which elements are selected.
  3. axis (optional): The axis along which to apply the condition. If None (default), the array is flattened before selection. For an n-dimensional array, axis can be 0 (rows), 1 (columns), ..., n-1.
  4. out (optional): A pre-allocated array to store the result. Must have the same shape as the expected output (derived from condition and axis).

How numpy.compress() Works#

numpy.compress() selects elements from a where condition is True (or non-zero). Here’s how it handles different scenarios:

  • If axis=None: The array a is flattened (converted to 1D), and condition must have length equal to the number of elements in a. The output is a 1D array of selected elements.
  • If axis is specified: The condition is applied along that axis. For example, if axis=0 (rows in a 2D array), condition must have length equal to the number of rows. The output retains the original array’s dimensions, except the size along axis is reduced to the number of True values in condition.

Example Usage#

1D Array Example#

Filter elements from a 1D array using a boolean condition:

import numpy as np
 
# Input array
arr = np.array([10, 20, 30, 40, 50])
# Condition: select elements at indices 0, 2, 4 (0-based)
condition = np.array([True, False, True, False, True])
 
# Apply compress
compressed = np.compress(condition, arr)
print("Compressed 1D array:", compressed)
# Output: [10 30 50]

2D Array with Axis=0 (Rows)#

Select rows from a 2D array where condition is True:

arr_2d = np.array([[1, 2],
                   [3, 4],
                   [5, 6]])
# Condition: select rows 0 and 2 (0-based)
condition = np.array([True, False, True])
 
# Apply compress along axis=0 (rows)
compressed_axis0 = np.compress(condition, arr_2d, axis=0)
print("Compressed along axis=0:")
print(compressed_axis0)
# Output:
# [[1 2]
#  [5 6]]

2D Array with Axis=1 (Columns)#

Select columns from a 2D array where condition is True:

# Condition: select column 0 (0-based)
condition_axis1 = np.array([True, False])
 
# Apply compress along axis=1 (columns)
compressed_axis1 = np.compress(condition_axis1, arr_2d, axis=1)
print("Compressed along axis=1:")
print(compressed_axis1)
# Output:
# [[1]
#  [3]
#  [5]]

Using the out Parameter#

For large arrays, reusing memory with out avoids creating new arrays:

# Pre-allocate output array (shape matches compressed_axis0)
out_array = np.empty((2, 2), dtype=int)
 
# Apply compress with out
np.compress(condition, arr_2d, axis=0, out=out_array)
print("Output using 'out' parameter:")
print(out_array)
# Output:
# [[1 2]
#  [5 6]]

Common Practices and Best Practices#

When to Use numpy.compress()#

  • Axis-Specific Selection: Use compress when you need to select rows/columns (or higher-dimensional axes) while preserving array structure.
  • Memory Efficiency: Use out to reuse memory for large arrays.
  • Alternative to Boolean Indexing: For 1D arrays, boolean indexing (e.g., arr[condition]) is simpler, but compress is better for axis-aware selection in multi-dimensional arrays.

Best Practices#

  1. Validate condition Length: Ensure condition has length matching the axis size (if axis is set) or the total number of elements (if axis=None).
  2. Specify axis Clearly: Avoid unexpected flattening by explicitly setting axis.
  3. Handle Empty Outputs: If all condition values are False, the output will be an empty array. Check for this to avoid errors.
  4. Use out for Large Arrays: Reuse memory with out to reduce overhead.

Use Cases#

  1. Data Filtering: Select rows from a dataset where a column value meets a criterion (e.g., condition = data[:, 0] > 10 for a 2D dataset).
  2. Feature Selection: In machine learning, select features (columns) where a statistical test (e.g., correlation with target) is significant.
  3. Image Processing: Apply a mask (boolean array) to an image (3D array) to select pixels along a specific axis (e.g., color channels).
  4. Time Series: Select time steps (rows) where a metric (e.g., sales > threshold) is met.

Limitations and Considerations#

  • condition Must Be 1D: If you have a 2D boolean mask, flatten it (e.g., condition = mask.ravel()) or use boolean indexing.
  • Axis Validity: axis must be a valid integer (0 to ndim(a)-1).
  • Performance: For extremely large arrays, compress may be slower than boolean indexing. Profile and optimize as needed.

Conclusion#

numpy.compress() is a versatile tool for selecting array elements based on a condition, with support for axis-specific operations. By mastering its syntax, parameters, and best practices, you can efficiently filter data, select features, and manipulate multi-dimensional arrays in numerical computing tasks.

References#

This guide equips you to leverage numpy.compress() for efficient array manipulation. Experiment with examples and integrate it into your data workflows!