Molenspin v0.4 — GPU integration and flexible constraint API

I've been sitting on most of this for three months waiting for the GPU path to stabilize. It finally did, so here's the full release writeup for 0.4.0 and the follow-up 0.4.1 patch that landed last week.

GPU-accelerated integration

The headline feature is optional CUDA/Metal offloading for the RK4 inner loop. For large assemblies — I've been testing with a 48-body planetary system that models a wind turbine gearbox — the GPU path is around 14× faster than the Rust CPU path on my M3 Max, and about 22× faster than the original NumPy baseline from 0.2.

The API change is a single keyword argument:

result = asm.simulate(duration=10.0, dt=1e-4, device="gpu")

On hardware without a compatible GPU it falls back to CPU without raising an error (it logs a warning at INFO level). The GPU path requires PyTorch ≥ 2.2 or the optional molenspin[metal] install for Apple Silicon. I decided against bundling a CUDA runtime — the dependency size was unreasonable for a niche library like this.

Rewritten constraint API

The old Constraint system required users to supply both a residual function and a hand-written Jacobian. Almost nobody wrote correct Jacobians — including me, for the first six months. The new API accepts a residual-only implementation and uses automatic differentiation (via the jax optional dependency, or a fallback finite-difference approximation) to compute the Jacobian internally.

import molenspin as ms

class FixedRatioConstraint(ms.Constraint):
    def __init__(self, shaft_a, shaft_b, ratio):
        self.shaft_a = shaft_a
        self.shaft_b = shaft_b
        self.ratio = ratio

    def residual(self, state):
        omega_a = state[self.shaft_a.state_index]
        omega_b = state[self.shaft_b.state_index]
        return omega_a * self.ratio - omega_b

asm.add_constraint(FixedRatioConstraint(s1, s2, ratio=2.5))

This is a breaking change from 0.3.x — constraints that previously supplied a jacobian() method will still work (the method is called if present), but they no longer need to.

Per-gear backlash

You can now specify backlash per gear pair rather than globally:

g1 = asm.add_gear(teeth=24, driven_by=shaft_in)
g2 = asm.add_gear(teeth=72, meshes_with=g1, backlash=0.003)  # 3 mrad

Backlash introduces a dead-band in the contact constraint. The solver handles it correctly as a complementarity condition — it's not a hack. I wrote up the math in a separate post if you're curious how contact detection interacts with the constraint graph.

Breaking changes in 0.4.0

• SimResult.shaft_histories renamed to SimResult.shafts
• Assembly.add_bevel() removed — use ms.BevelPair() constructor directly
• Constraint jacobian() is now optional (still respected if present)
• Minimum Python version raised to 3.10 (was 3.9)

What's next

I want to add cam-follower support before 0.5. The current model has no way to express a lift curve, which rules out a whole class of mechanisms I care about. Beyond that, there's a long-standing request for a 3D coordinate output alongside the SVG — probably JSON with per-body transform matrices at each timestep.

Full diff on the changelog page.