Extending Newton with Custom Kernels and User-Defined Forces

You can extend Newton's physics engine by writing custom Warp kernels that manipulate the simulation state directly, enabling arbitrary force laws and constraints without modifying the core library.

Newton is an open-source physics simulation framework built on Warp – NVIDIA's JIT-compiled GPU/CPU kernel language. Because Newton implements all built-in forces (gravity, springs, contacts) as standard Warp kernels, you can inject your own user-defined forces by writing Python functions decorated with @wp.kernel and launching them against the simulation state. This article demonstrates how to extend Newton with custom kernels, from registering data attributes to integrating forces into the simulation loop.

Understanding Newton's Kernel Architecture

Newton's simulation core delegates all computationally intensive work to Warp kernels – ordinary Python functions that Warp compiles to CUDA or CPU instructions at runtime. The architecture separates concerns into three layers:

• Kernel definitions: Pure Python functions annotated with @wp.kernel that read from and write to State arrays (positions, velocities, forces). These live in files like newton/_src/viewer/kernels.py and newton/_src/solvers/vbd/rigid_vbd_kernels.py.
• Launch helpers: Python code that calls wp.launch(kernel, dim=..., inputs=[...], device=...), centralized in newton/_src/viewer/viewer.py and various *_utils.py modules.
• Force accumulation: Each rigid body maintains a body_f array (spatial force) that kernels update using wp.atomic_add to prevent race conditions during parallel execution.

When the solver advances a time step, it iterates through registered force kernels, accumulates their contributions into body_f, then proceeds to constraint solving. By registering your kernel before this loop, you make your user-defined force part of the standard accumulation phase.

Registering Custom Attributes for Kernel Data

To pass per-body parameters to your kernel (force magnitudes, drag coefficients, magnetic charges), Newton provides a custom attribute system defined in newton/_src/sim/builder.py. This allows you to attach scalar or vector data to bodies that kernels can read during execution.

Register a custom attribute during model construction:

from newton import ModelBuilder
import warp as wp
import numpy as np

builder = ModelBuilder()

# Register a per-body scalar attribute

builder.add_custom_attribute(
    ModelBuilder.CustomAttribute(
        key="my_force",                           # Accessed in kernel as state.custom["my_force"]

        frequency=ModelBuilder.Frequency.BODY,    # Per-body data (vs. per-vertex)

        dtype=wp.float32,
        default=0.0,
    )
)

# Populate values (e.g., sinusoidal force magnitudes)

builder.add_custom_values(
    key="my_force",
    values=[np.sin(t) for t in np.linspace(0, 2*np.pi, builder.body_count)]
)

The data becomes available in the simulation state as state.custom["my_force"], which you pass as an input array to your Warp kernel.

Writing a Custom Warp Kernel for User-Defined Forces

Custom forces are implemented as Warp kernels that read body properties and write to the force accumulator body_f. The kernel signature must accept Warp arrays corresponding to the state data you intend to read or modify.

Create a kernel in newton/_src/viewer/kernels.py or your own script:

import warp as wp

@wp.kernel
def my_custom_force_kernel(
    body_f: wp.array(dtype=wp.spatial_vector),   # Accumulated spatial forces (output)

    body_mass: wp.array(dtype=wp.float32),         # Read-only mass data

    my_force: wp.array(dtype=wp.float32),          # Custom per-body scalar input

):
    i = wp.tid()  # Thread index corresponds to body index

    
    # Guard against out-of-bounds launches

    if i >= body_f.shape[0]:
        return
    
    # Compute force: upward push scaled by custom attribute

    upward = wp.vec3(0.0, my_force[i], 0.0)
    torque = wp.vec3(0.0, 0.0, 0.0)
    
    # Atomic add to avoid race conditions with other force kernels

    wp.atomic_add(body_f, i, wp.spatial_vector(upward, torque))

Critical implementation details:

• Thread indexing: Use wp.tid() to obtain the current thread ID, which maps to the body index in parallel launches.
• Bounds checking: Always verify i < array.shape[0] to handle launch dimensions that may exceed the actual data size.
• Atomic operations: Write forces using wp.atomic_add rather than direct assignment. Multiple kernels (gravity, springs, your custom force) may write to body_f simultaneously.
• Spatial vectors: Forces are stored as wp.spatial_vector, which combines a linear force vector and a torque vector.

Launching Custom Kernels in the Simulation Loop

To apply your user-defined force during simulation, launch the kernel before the solver's constraint step. The launch requires the dimension (number of bodies), input arrays from the state, and the target device.

Manual kernel launch from a simulation script:

import warp as wp
import newton

# Assuming 'state' is a newton.State and 'model' is built

device = wp.get_default_device()

# Launch custom force kernel

wp.launch(
    my_custom_force_kernel,
    dim=state.body_f.shape[0],  # One thread per body

    inputs=[
        state.body_f,           # Output: force accumulator

        state.body_mass,        # Input: mass data

        state.custom["my_force"]  # Input: custom attribute

    ],
    device=device,
)

# Proceed with simulation step

sim.step()

The kernel adds its contribution to state.body_f. When sim.step() executes, the solver uses these accumulated forces to update velocities and positions.

Integrating Custom Forces via the Viewer Hook

For interactive simulations, Newton's viewer (newton.ViewerGL) provides an apply_forces() method that runs every frame before the physics step. This is the idiomatic location to inject custom kernels for real-time force application.

Extend the viewer in newton/_src/viewer/viewer.py or subclass it:

class CustomViewer(newton.ViewerGL):
    def apply_forces(self, state: newton.State) -> None:
        # Call built-in forces (picking, wind)

        super().apply_forces(state)
        
        # Inject custom force kernel

        wp.launch(
            my_custom_force_kernel,
            dim=state.body_f.shape[0],
            inputs=[
                state.body_f,
                state.body_mass,
                state.custom["my_force"]
            ],
            device=self.device,
        )

By overriding apply_forces(), your custom kernel executes automatically each simulation step alongside Newton's built-in gravity and contact forces.

Complete Working Example

The following script demonstrates the full pipeline: registering a custom attribute, defining a kernel, and running a simulation with the custom force applied each step.

# example_custom_force.py

import warp as wp
import newton
from newton import ModelBuilder
import numpy as np

# ------------------------------------------------------------------

# 1. Build model with custom attribute

builder = ModelBuilder()
builder.add_body(name="sphere", mass=1.0, inertia=(0.1, 0.1, 0.1))
builder.add_shape_sphere(body="sphere", radius=0.5)

# Register per-body force magnitude

builder.add_custom_attribute(
    ModelBuilder.CustomAttribute(
        key="my_force",
        frequency=ModelBuilder.Frequency.BODY,
        dtype=wp.float32,
        default=0.0,
    )
)

# Set upward force of 5N

builder.add_custom_values(key="my_force", values=[5.0])

model = builder.build()
state = model.state

# ------------------------------------------------------------------

# 2. Define custom Warp kernel

@wp.kernel
def my_custom_force_kernel(
    body_f: wp.array(dtype=wp.spatial_vector),
    my_force: wp.array(dtype=wp.float32),
):
    i = wp.tid()
    if i >= body_f.shape[0]:
        return
    
    upward = wp.vec3(0.0, my_force[i], 0.0)
    torque = wp.vec3(0.0, 0.0, 0.0)
    wp.atomic_add(body_f, i, wp.spatial_vector(upward, torque))

# ------------------------------------------------------------------

# 3. Run simulation with custom force

dt = 0.01
sim = newton.Simulator(model, dt=dt, solver=newton.SolverXPBD())

for step in range(200):
    # Apply user-defined force before solver step

    wp.launch(
        my_custom_force_kernel,
        dim=state.body_f.shape[0],
        inputs=[state.body_f, state.custom["my_force"]],
        device=wp.get_default_device(),
    )
    
    # Advance physics

    sim.step()

# ------------------------------------------------------------------

# 4. Visualize result

viewer = newton.ViewerGL()
viewer.show(state)

This example creates a sphere with a constant 5N upward force, demonstrating how custom attributes flow from the builder through the state to the kernel.

Key Source Files Reference

When implementing custom kernels, consult these specific files in the newton-physics/newton repository:

• newton/_src/viewer/kernels.py – Reference implementations for interactive forces (picking, wind) showing the pattern for reading state and writing to body_f.
• newton/_src/solvers/vbd/rigid_vbd_kernels.py – Low-level force accumulation examples (apply_joint_forces, apply_body_forces) demonstrating wp.atomic_add on spatial vectors.
• newton/_src/solvers/xpbd/kernels.py – XPBD-specific kernel signatures for constraint and force processing.
• newton/_src/sim/builder.py – Definition of ModelBuilder.CustomAttribute and APIs for registering per-body or per-vertex data.
• newton/_src/viewer/viewer.py – The apply_forces() method that serves as the integration point for viewer-driven force injection.
• newton/tests/test_custom_attributes.py – Unit tests verifying custom attribute registration and kernel access patterns.

Performance and Safety Best Practices

When extending Newton with custom kernels, adhere to these guidelines to ensure correct, high-performance simulations:

• Use atomic operations for force accumulation. Always write to body_f via wp.atomic_add (or wp.atomic_sub) rather than direct assignment. Multiple kernels may write concurrently, and atomic operations prevent race conditions.
• Guard against out-of-bounds access. Warp launches operate on a fixed dimension; always check if i >= array.shape[0]: return at the start of your kernel to handle cases where the launch size exceeds the actual data count.
• Keep kernels data-parallel. Avoid Python control flow or loops inside kernels. Warp compiles kernels for GPU execution; use vectorized operations and wp.tid() for parallelism.
• Match data types precisely. Ensure your kernel signatures use wp.float32, wp.spatial_vector, or other Warp types that exactly match the state arrays passed from Python.
• Reuse device arrays. When possible, pre-allocate auxiliary arrays and reuse them across frames rather than allocating inside the simulation loop, minimizing CPU-GPU transfer overhead.

Summary

Extending Newton with custom kernels and user-defined forces leverages the framework's Warp foundation to inject arbitrary physics into the simulation pipeline. The essential workflow involves:

• Registering custom attributes via ModelBuilder.add_custom_attribute() to pass per-body parameters to kernels.
• Writing Warp kernels annotated with @wp.kernel that read state data and write forces to body_f using wp.atomic_add.
• Launching kernels manually before sim.step() or integrating them into the viewer’s apply_forces() hook for automatic per-frame execution.
• Consulting reference implementations in newton/_src/viewer/kernels.py and newton/_src/solvers/vbd/rigid_vbd_kernels.py for correct patterns.

This architecture allows you to implement complex force laws—from magnetic fields to custom drag models—without forking the Newton repository.

Frequently Asked Questions

How do I pass per-body parameters to a custom kernel in Newton?

Use the custom attribute system defined in newton/_src/sim/builder.py. Call builder.add_custom_attribute() with a ModelBuilder.CustomAttribute specifying the key, frequency (BODY or VERTEX), and data type. Populate values with builder.add_custom_values(), then access the array in your kernel via state.custom["your_key"].

Why must I use wp.atomic_add instead of direct assignment when writing forces?

Newton runs multiple force kernels concurrently during each simulation step. If two kernels write to body_f simultaneously, direct assignment creates a race condition where one write overwrites the other. wp.atomic_add ensures thread-safe accumulation by serializing additions to the same memory address, preserving contributions from all kernels.

Where is the best place to inject a custom force for interactive simulations?

Override the apply_forces() method in newton/_src/viewer/viewer.py or subclass ViewerGL. This method executes every frame before the physics solver, making it the natural integration point for forces that must update continuously (wind, user interaction, or procedural forces). Call super().apply_forces(state) first to preserve built-in picking and wind forces, then launch your custom kernel.

Can I create kernels dynamically at runtime rather than pre-defining them?

Yes. Newton demonstrates this pattern in newton/_src/solvers/featherstone/solver_featherstone.py (around line 150), where the solver generates a custom kernel on-the-fly to evaluate the system matrix. You can use Python's dynamic code generation or Warp's just-in-time compilation features to create specialized kernels based on runtime parameters, though for most use cases pre-defined kernels with custom attribute inputs offer better readability and debugging.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}