Python – Return NumPy Array From Function

by

Do you need to create a function that returns a NumPy array but you don’t know how? No worries, in sixty seconds, you’ll know! Go! 🚀

A Python function can return any object such as a NumPy Array. To return an array, first create the array object within the function body, assign it to a variable arr, and return it to the caller of the function using the keyword operation “return arr“.

👉 Recommended Tutorial: How to Initialize a NumPy Array? 6 Easy Ways

Create and Return 1D Array

For example, the following code creates a function create_array() of numbers 0, 1, 2, …, 9 using the np.arange() function and returns the array to the caller of the function:

import numpy as np


def create_array():
    ''' Function to return array '''
    return np.arange(10)

numbers = create_array()
print(numbers)
# [0 1 2 3 4 5 6 7 8 9]

The np.arange([start,] stop[, step]) function creates a new NumPy array with evenly-spaced integers between start (inclusive) and stop (exclusive).

The step size defines the difference between subsequent values. For example, np.arange(1, 6, 2) creates the NumPy array [1, 3, 5].

To better understand the function, have a look at this video:

I also created this figure to demonstrate how NumPy’s arange() function works on three examples:

In the code example, we used np.arange(10) with default start=0 and step=1 only specifying the stop=10 argument.

If you need an even deeper understanding, I’d recommend you check out our full guide on the Finxter blog.

👉 Recommended Tutorial: NumPy Arange Function — A Helpful Illustrated Guide

Create and Return 2D NumPy Array

You can also create a 2D (or multi-dimensional) array in a Python function by first creating a 2D or (xD) nested list and converting the nested list to a NumPy array by passing it into the np.array() function.

The following code snippet uses nested list comprehension to create a 2D NumPy array following a more complicated creation pattern:

import numpy as np


def create_array(a,b):
    ''' Function to return array '''
    lst = [[(i+j)**2 for i in range(a)] for j in range(b)]
    return np.array(lst)


arr = create_array(4,3)
print(arr)

Output:

[[ 0  1  4  9]
 [ 1  4  9 16]
 [ 4  9 16 25]]

I definitely recommend reading the following tutorial to understand nested list comprehension in Python:

👉 Recommended Tutorial: Nested List Comprehension in Python

More Ways

There are many other ways to return an array in Python.

For example, you can use either of those methods inside the function body to create and initialize a NumPy array:

To get a quick overview what to put into the function and how these methods work, I’d recommend you check out our full tutorial.

👉 Recommended Tutorial: How to Initialize a NumPy Array? 6 Easy Ways

Related Tutorials

Programmer Humor

Q: How do you tell an introverted computer scientist from an extroverted computer scientist?

A: An extroverted computer scientist looks at your shoes when he talks to you.