Warp Basics

This example is from NVIDIA’s Accelerated Computing Hub.

import __main__
import numpy as np
import warp as wp
from tqdm.auto import trange

# Warp reads __main__.__loader__, which is missing in some notebook kernels
if not hasattr(__main__, "__loader__"):
    __main__.__loader__ = None

num_particles = 10_000_000
num_steps = 100

mass = 0.1
g = 9.81
b = 0.05

dt = 0.01 * (2 * mass / b)

gravity = wp.vec3([0.0, 0.0, -g])


@wp.kernel
def integrate(positions: wp.array(dtype=wp.vec3), velocities: wp.array(dtype=wp.vec3)):  # type: ignore
    i = wp.tid()
    acceleration = (-b * velocities[i]) / mass + gravity
    velocities[i] += acceleration * dt
    positions[i] += velocities[i] * dt

# Initial positions: random values between -1.0 and 1.0 for x, y, and z
rng = np.random.default_rng(12345)
positions_np = rng.uniform(low=-1.0, high=1.0, size=(num_particles, 3))
positions = wp.array(positions_np, dtype=wp.vec3)

# Initial velocities: random values between -0.5 and 0.5 m/s for vx, vy, and vz
velocities_np = rng.uniform(low=-0.5, high=0.5, size=(num_particles, 3))
velocities = wp.array(velocities_np, dtype=wp.vec3)

print(f"Initial positions:\n{positions}")

for step in trange(num_steps):
    wp.launch(integrate, dim=(num_particles,), inputs=[positions, velocities])

print(f"Final positions:\n{positions}")

Warp 1.18.0 initialized:
   CUDA not enabled in this build
   Devices:
     "cpu"      : "arm"
   Kernel cache:
     ~/Library/Caches/warp/1.18.0
Initial positions:
[[-0.54532796 -0.36648333  0.5947309 ]
 [ 0.35250935 -0.2177809  -0.33437213]
 [ 0.19661751 -0.6265316   0.3455121 ]
 ...
 [-0.55922866 -0.93023884 -0.83259517]
 [-0.39899704  0.97424763  0.94057095]
 [ 0.38531193  0.26505202 -0.60369706]]

  0%|          | 0/100 [00:00<?, ?it/s]

Module __main__ 4f23e15 load on device 'cpu' took 18.01 ms  (cached)

   ⋮ (12 progress updates omitted)
100%|██████████| 100/100 [00:01<00:00, 73.16it/s]

Final positions:
[[ -0.09391715  -1.0830888  -43.828407  ]
 [  0.6937012   -0.6013651  -45.414093  ]
 [  0.5957543    0.08006047 -44.8957    ]
 ...
 [ -0.72598     -0.8991223  -45.298027  ]
 [ -0.32058033   0.70321625 -43.714462  ]
 [ -0.08150801  -0.0732362  -46.081345  ]]