warp.map#

warp.map( func, *inputs, out=None, return_kernel=False, block_dim=256, device=None, )[source]#

Map a function over the elements of one or more arrays.

You can use a Warp function, a regular Python function, or a lambda expression to map it to a set of arrays.

a = wp.array([1, 2, 3], dtype=wp.float32)
b = wp.array([4, 5, 6], dtype=wp.float32)
c = wp.array([7, 8, 9], dtype=wp.float32)
result = wp.map(lambda x, y, z: x + 2.0 * y - z, a, b, c)
print(result)

[2. 4. 6.]

Clamp values in an array in place:

xs = wp.array([-1.0, 0.0, 1.0], dtype=wp.float32)
wp.map(wp.clamp, xs, -0.5, 0.5, out=xs)
print(xs)

[-0.5  0.   0.5]

Note that only one of the inputs must be a Warp array. For example, it is possible vectorize the function warp.transform_point() over a collection of points with a given input transform as follows:

tf = wp.transform((1.0, 2.0, 3.0), wp.quat_rpy(0.2, -0.6, 0.1))
points = wp.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=wp.vec3)
transformed = wp.map(wp.transform_point, tf, points)

Besides regular Warp arrays, other array types, such as the warp.indexedarray, are supported as well:

arr = wp.array(data=np.arange(10, dtype=np.float32))
indices = wp.array([1, 3, 5, 7, 9], dtype=int)
iarr = wp.indexedarray1d(arr, [indices])
out = wp.map(lambda x: x * 10.0, iarr)
print(out)

[10. 30. 50. 70. 90.]

If multiple arrays are provided, the NumPy broadcasting rules are applied to determine the shape of the output array. Two shapes are compatible when: starting from the trailing dimension,

the two dimensions are equal, or
one of the dimensions is 1.

For example, given arrays of shapes (3, 1, 4) and (5, 4), the broadcasted shape is (3, 5, 4).

If no array(s) are provided to the out argument, the output array(s) are created automatically. The data type(s) of the output array(s) are determined by the type of the return value(s) of the function. The requires_grad flag for an automatically created output array is set to True if any of the input arrays have it set to True and the respective output array’s dtype is a type that supports differentiation.

Parameters:

func (Callable | Function) – The function to map over the arrays.
*inputs (Array[DType] | Any) – The input arrays or values to pass to the function.
out (Array[DType] | list[Array[DType]] | None) – Optional output array(s) to store the result(s). If None, the output array(s) will be created automatically.
return_kernel (bool) – If True, only return the generated kernel without performing the mapping operation.
block_dim (int) – The number of threads per block for the kernel launch.
device (Device | str | None) – The device on which to run the kernel.

Returns:

The resulting array(s) of the mapping. If return_kernel is True, only returns the kernel used for mapping.

Return type:

array | list[array] | Kernel