# Extending Newton with Custom Kernels and User-Defined Forces

> Extend Newton physics engine with custom Warp kernels and user-defined forces. Manipulate simulation state directly for custom force laws and constraints without core library changes.

- Repository: [Newton Physics/newton](https://github.com/newton-physics/newton)
- Tags: how-to-guide
- Published: 2026-03-19

---

**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`](https://github.com/newton-physics/newton/blob/main/newton/_src/viewer/kernels.py) and [`newton/_src/solvers/vbd/rigid_vbd_kernels.py`](https://github.com/newton-physics/newton/blob/main/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`](https://github.com/newton-physics/newton/blob/main/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`](https://github.com/newton-physics/newton/blob/main/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:

```python
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`](https://github.com/newton-physics/newton/blob/main/newton/_src/viewer/kernels.py) or your own script:

```python
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:

```python
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`](https://github.com/newton-physics/newton/blob/main/newton/_src/viewer/viewer.py) or subclass it:

```python
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.

```python

# 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](https://github.com/newton-physics/newton) repository:

- **[`newton/_src/viewer/kernels.py`](https://github.com/newton-physics/newton/blob/main/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`](https://github.com/newton-physics/newton/blob/main/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`](https://github.com/newton-physics/newton/blob/main/newton/_src/solvers/xpbd/kernels.py)** – XPBD-specific kernel signatures for constraint and force processing.
- **[`newton/_src/sim/builder.py`](https://github.com/newton-physics/newton/blob/main/newton/_src/sim/builder.py)** – Definition of `ModelBuilder.CustomAttribute` and APIs for registering per-body or per-vertex data.
- **[`newton/_src/viewer/viewer.py`](https://github.com/newton-physics/newton/blob/main/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`](https://github.com/newton-physics/newton/blob/main/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`](https://github.com/newton-physics/newton/blob/main/newton/_src/viewer/kernels.py) and [`newton/_src/solvers/vbd/rigid_vbd_kernels.py`](https://github.com/newton-physics/newton/blob/main/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`](https://github.com/newton-physics/newton/blob/main/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`](https://github.com/newton-physics/newton/blob/main/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`](https://github.com/newton-physics/newton/blob/main/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.