Skip to content

deepsets.kt

src/jvmMain/kotlin/afield/ai/learning/deepsets.kt · 878 lines

package and imports — 73 lines, from line 1 
@file:afield.runtime.completion.CompletionCategory("Machine learning")

package afield.ai.learning

import afield.math.autodiff.sqrt
import afield.simulation.particles.glsl

import afield.data.collections.get
import afield.data.collections.invoke
import afield.data.collections.r
import afield.data.collections.times
import afield.document.context.Why
import afield.document.graph.minusAssign
import afield.document.graph.plusAssign
import afield.document.properties.compareTo
import afield.document.workspace.file
import afield.geometry.mesh.contains
import afield.geometry.mesh.scaled
import afield.graphics.render.rangeTo
import afield.graphics.text.layout
import afield.math.autodiff.compareTo
import afield.math.autodiff.div
import afield.math.autodiff.dot
import afield.math.autodiff.floor
import afield.math.autodiff.minus
import afield.math.autodiff.norm
import afield.math.autodiff.plus
import afield.math.autodiff.product
import afield.math.autodiff.times
import afield.math.linalg.f
import afield.math.linalg.invoke
import afield.math.linalg.remAssign
import afield.math.linalg.set
import afield.math.linalg.times
import afield.math.units.div
import afield.math.units.minus
import afield.math.units.plus
import afield.math.units.times
import afield.media.image.channel
import afield.media.image.close
import afield.media.image.rotate
import afield.runtime.execution.go
import afield.runtime.lifecycle.once
import afield.simulation.agents.Aggregate
import afield.study.ui.numbers

import afield.audio.synth.*
import afield.data.collections.m
import afield.data.collections.s
import afield.geometry.mesh.gradient
import afield.math.autodiff.AD
import afield.math.autodiff.ADTape
import afield.math.autodiff.adam
import afield.math.autodiff.maxOf
import afield.math.autodiff.mean
import afield.math.autodiff.pow
import afield.math.autodiff.sq
import afield.math.autodiff.sum
import afield.math.autodiff.tanh
import afield.math.control.params
import afield.math.functions.parameterCount
import afield.math.functions.tanh
import afield.math.linalg.d
import afield.math.linalg.i
import afield.math.linalg.x
import afield.math.statistics.random
import afield.simulation.agents.step
import afield.simulation.automata.step
import afield.simulation.particles.MlpActivation
import afield.simulation.particles.MlpWeights
import kotlin.math.max
import kotlin.math.sqrt
import kotlin.random.Random

A learned two-body interaction: the acceleration on a particle as a function of the neighbours around it, fitted by autodiff from examples of any rule that can already be evaluated.

flowMatch learns v(x, t): where a particle should go given only where it is. That can carry a cloud onto a shape, and it is blind to the other particles. An interaction is the other kind of rule:

a_i = phi( v_i, SUM_j psi(x_j - x_i, v_j - v_i), MEAN_j psi(...) )

two small networks and a reduction over neighbours. This is the Deep Sets form, and its aggregation is permutation-invariant over any number of neighbours, so one set of weights works for a particle with three neighbours and a particle with three hundred. Classical boids is one particular choice of psi and phi, so the family contains it.

Why both a sum and a mean

Reynolds separation is a sum: being crowded by twenty neighbours pushes harder than being crowded by two. Alignment and cohesion are means, since they are averages of neighbour velocity and position. A mean-only aggregator cannot represent separation at any width, and a sum-only one makes alignment scale with crowding, so phi is fed both, at a cost of messageWidth extra inputs. DeepSetsWeights.sumScale keeps the summed branch in a sane range.

The feature map is where the physics goes

psi is not fed the raw offset. It is fed interactionFeatures: the offset and velocity difference, plus the unit direction scaled by a small radial basis, u*p and u*p^2, where p is a bounded proximity that rises as neighbours get close.

Separation goes as 1/r over a four-hundred-fold range of r, and a tanh network cannot form the product of a direction and a steep radial profile out of Cartesian components at any practical width. Alignment and cohesion are plain means, and linear in the inputs, so they survive the raw encoding. Measured by ablation against the flocking rule: fed raw Cartesian offsets the model reproduces those two to a relative error of 0.036 and separation only to 0.78, which is barely better than predicting the average.

A radial interaction is a direction times a profile, so handing the network u already multiplied by a couple of profile terms states the form of a two-body law rather than transcribing boids. The profile's coefficients, the velocity dependence, and how the summed and meaned branches combine are all still learned.

Trained by afield.math.autodiff.ADTape, not by Smile

Smile's MLP cannot express this: the gradient has to flow back through a sum over a variable-length neighbour set, and psi is applied once per neighbour with shared weights. autodiff.kt is reverse-mode over a tape, which handles it without knowing anything about networks, since the tape records whatever the forward pass did. Cost is one forward and one backward sweep per batch whatever the neighbour counts are.

The weights come out as MlpWeights

Both halves are packaged in the same flat layout the flow-matching work already uses, so MlpWeights.glsl emits them for the GPU unchanged and the parity machinery in four.neural-field-parity.smoke applies to them without modification.

class DeepSetsWeights

A trained interaction in the plain Cartesian form: psi and phi, the width of the message between them, and the five scales that normalise a neighbourhood into the networks and the answer back out to world units. EquivariantWeights is the same model constrained so that rotating a neighbourhood rotates the answer with it.

class DeepSetsWeights(
    /** Applied per neighbour. Inputs `(dx, dy, dvx, dvy)` normalised; outputs [messageWidth]. */
    val psi: MlpWeights,
    /** Applied once. Inputs `(vx, vy, sum…, mean…)` normalised; outputs a normalised acceleration. */
    val phi: MlpWeights,
    /** How wide a single neighbour's message is. */
    val messageWidth: Int,
    /** Interaction radius, in world units. Offsets are divided by this. */
    val radius: Float,
    /** Velocity differences are divided by this. */
    val velocityScale: Float,
    /** The output is multiplied by this to get a world-units acceleration. */
    val accelScale: Float,
    /** The summed branch is divided by this, so a typical crowd lands near 1. */
    val sumScale: Float,
    /** Closest approach used by the proximity feature — the same floor the rule itself uses. */
    val minSeparation: Float,
) {

init

    init {
        require(psi.inputs == PSI_INPUTS) {
            "psi takes the $PSI_INPUTS features interactionFeatures() produces, not ${psi.inputs}"
        }
        require(psi.outputs == messageWidth) { "psi must emit $messageWidth, not ${psi.outputs}" }
        require(phi.inputs == 2 + 2 * messageWidth) {
            "phi takes velocity plus a summed and a meaned message (${2 + 2 * messageWidth}), " +
                "this one takes ${phi.inputs}"
        }
        require(phi.outputs == 2) { "phi must emit a 2D acceleration, not ${phi.outputs}" }
    }

val parameterCount: Int get() = psi.parameterCount + phi.parameterCount

How many weights and biases in total.

fun evaluate(vx: Float, vy: Float, neighbours: List<FloatArray>): Pair<Float, Float>

The acceleration on a particle with velocity (vx, vy) given its neighbours, in fp32 and in the same order a shader would accumulate it — the reference the GPU is held to.

neighbours supplies (dx, dy, dvx, dvy) per neighbour, already in world units.

    fun evaluate(vx: Float, vy: Float, neighbours: List<FloatArray>): Pair<Float, Float> {
        val sum = FloatArray(messageWidth)
        val out = FloatArray(messageWidth)
        for (n in neighbours) {
            psi.evaluate(interactionFeatures(n[0], n[1], n[2], n[3], radius, velocityScale, minSeparation), out)
            for (k in 0 until messageWidth) sum[k] += out[k]
        }
        val count = neighbours.size
        val input = FloatArray(phi.inputs)
        input[0] = vx / velocityScale
        input[1] = vy / velocityScale
        for (k in 0 until messageWidth) {
            input[2 + k] = sum[k] / sumScale
            input[2 + messageWidth + k] = if (count > 0) sum[k] / count else 0f
        }
        val a = phi.evaluate(input)
        return a[0] * accelScale to a[1] * accelScale
    }

override fun toString(): String

Both network shapes and the parameter total. The GLSL emitter writes this line into the shader as a comment, so a generated source names the weights it was generated from.

    override fun toString(): String =
        "DeepSetsWeights(psi ${psi}, phi ${phi}, $parameterCount parameters)"

companion object

    companion object {

const val PSI_INPUTS = 8

How many features interactionFeatures produces, and therefore psi's input width.

fun interactionFeatures

The per-neighbour features psi reads — see the discussion on DeepSetsWeights.

[dx/R, dy/R, dvx/V, dvy/V, ux*p, uy*p, ux*p^2, uy*p^2], where u is the unit direction to the neighbour and p = minSeparation / max(r, minSeparation) is a proximity in (0, 1].

The last four are the ones that matter. A 1/r law is exactly u * p / minSeparation for any r past the floor, so it is linear in ux*p — which is what turns separation from the term the network could not fit into the one it fits first.

fun interactionFeatures(
    dx: Float,
    dy: Float,
    dvx: Float,
    dvy: Float,
    radius: Float,
    velocityScale: Float,
    minSeparation: Float,
    into: FloatArray = FloatArray(DeepSetsWeights.PSI_INPUTS),
): FloatArray {
    val distance = sqrt((dx * dx + dy * dy).toDouble()).toFloat()
    val soft = maxOf(distance, minSeparation)
    // At distance zero the direction is undefined; zero is the only answer that does not invent a
    // force out of a rounding error.
    val ux = if (distance > 0f) dx / distance else 0f
    val uy = if (distance > 0f) dy / distance else 0f
    val p = minSeparation / soft
    into[0] = dx / radius
    into[1] = dy / radius
    into[2] = dvx / velocityScale
    into[3] = dvy / velocityScale
    into[4] = ux * p
    into[5] = uy * p
    into[6] = ux * p * p
    into[7] = uy * p * p
    return into
}

class InteractionExample

One training example: a neighbourhood and the acceleration the rule being imitated produces.

class InteractionExample(
    /** The centre particle's velocity, in world units. */
    val velocity: FloatArray,
    /** `(dx, dy, dvx, dvy)` per neighbour, in world units. */
    val neighbours: List<FloatArray>,
    /** The target acceleration, in world units. */
    val target: FloatArray,
)

data class DeepSetsOptions

The knobs trainDeepSets takes.

data class DeepSetsOptions(
    /** Adam steps to take. Each one calls `examples` for a fresh batch. */
    val steps: Int = 4000,
    /**
     * How many examples a batch is expected to hold. [trainDeepSets] divides its loss by the size
     * of the list `examples` actually returned and only checks this is positive; [trainEquivariant]
     * divides by this number, so a batch of a different size there rescales the gradient.
     */
    val batch: Int = 16,
    /**
     * How many numbers `psi` emits per neighbour, and so how wide the summed and the meaned branch
     * `phi` reads each are. [trainEquivariant] reads it as a count of vector channels instead.
     */
    val messageWidth: Int = 6,
    /** Hidden layer widths of the per-neighbour network. `tanh` between them, linear out. */
    val psiHidden: IntArray = intArrayOf(16, 16),
    /** Hidden layer widths of the combining network. */
    val phiHidden: IntArray = intArrayOf(16, 16),
    /** Adam's rate at the first step. */
    val learningRate: Double = 3e-3,
    /** Decayed linearly to this by the last step. Same reasoning as `FlowMatchOptions`. */
    val finalLearningRate: Double = 2e-4,
    /** Seeds the Xavier initialisation, and nothing else: `examples` draws its own randomness. */
    val seed: Long = 1L,
    /** Called with the step number and that step's loss, or null to train silently. */
    val onProgress: ((Int, Double) -> Unit)? = null,
    /** How often [onProgress] fires, in steps. The last step reports only if it lands on one. */
    val progressEvery: Int = 200,
)

fun trainDeepSets

Fits a DeepSetsWeights to whatever rule examples embodies, by reverse-mode autodiff.

examples is called with a step index and must return a fresh batch; drawing new neighbourhoods every step rather than looping over a fixed set is what stops the networks memorising particular configurations, and costs nothing when the examples are generated from a rule you already have.

Plain MSE on the acceleration. Weighting each example by its own target magnitude was tried and removed: it changes nothing, because the regime where it would matter — a settled flock, where separation and cohesion cancel twentyfold — is a regime the examples do not contain. Sample the neighbourhoods from a running flock instead; that is worth a factor of ten in absolute force error where reweighting is worth nothing.

Adam, because plain SGD on a sum-over-neighbours objective has wildly different gradient scales between the psi weights (which see every neighbour) and the phi weights (which see one vector), and a single learning rate suits neither.

body — 81 lines, from line 305 
fun trainDeepSets(
    radius: Float,
    velocityScale: Float,
    accelScale: Float,
    sumScale: Float,
    minSeparation: Float,
    options: DeepSetsOptions = DeepSetsOptions(),
    examples: (Int) -> List<InteractionExample>,
): DeepSetsWeights {
    require(options.steps > 0) { "trainDeepSets needs at least one step" }
    require(options.batch > 0) { "trainDeepSets needs a batch of at least one" }

    val random = Random(options.seed)
    val psiShape = MlpShape(DeepSetsWeights.PSI_INPUTS, options.psiHidden, options.messageWidth)
    val phiShape = MlpShape(2 + 2 * options.messageWidth, options.phiHidden, 2)

    val psiParams = psiShape.initialise(random)
    val phiParams = phiShape.initialise(random)
    val adam = Adam(psiParams.size + phiParams.size)

    var lastLoss = 0.0
    for (step in 0 until options.steps) {
        val tape = ADTape()
        val psiVars = tape.variables(psiParams)
        val phiVars = tape.variables(phiParams)

        var loss = tape.constant(0.0)
        val batch = examples(step)
        for (example in batch) {
            val sum = MutableList(options.messageWidth) { tape.constant(0.0) }
            val features = FloatArray(DeepSetsWeights.PSI_INPUTS)
            for (n in example.neighbours) {
                // Features are data, not parameters, so they go on the tape as constants: no
                // gradient flows into them and they cost one node each.
                interactionFeatures(n[0], n[1], n[2], n[3], radius, velocityScale, minSeparation, features)
                val out = psiShape.forward(features.map { tape.constant(it.toDouble()) }, psiVars)
                for (k in 0 until options.messageWidth) sum[k] = sum[k] + out[k]
            }
            val count = example.neighbours.size
            val input = ArrayList<AD>(phiShape.inputs)
            input += tape.constant(example.velocity[0].toDouble() / velocityScale)
            input += tape.constant(example.velocity[1].toDouble() / velocityScale)
            for (k in 0 until options.messageWidth) input += sum[k] / sumScale.toDouble()
            for (k in 0 until options.messageWidth) {
                input += if (count > 0) sum[k] / count.toDouble() else tape.constant(0.0)
            }
            val predicted = phiShape.forward(input, phiVars)
            for (d in 0 until 2) {
                val residual = predicted[d] - (example.target[d].toDouble() / accelScale)
                loss = loss + residual.sq()
            }
        }
        loss = loss / (batch.size * 2).toDouble()
        lastLoss = loss.value

        val gradients = tape.backward(loss)
        val rate = options.learningRate +
            (options.finalLearningRate - options.learningRate) * step / (options.steps - 1).coerceAtLeast(1)
        adam.step(rate) { i ->
            if (i < psiParams.size) gradients[psiVars[i].index]
            else gradients[phiVars[i - psiParams.size].index]
        }.let { delta ->
            for (i in psiParams.indices) psiParams[i] -= delta[i]
            for (i in phiParams.indices) phiParams[i] -= delta[psiParams.size + i]
        }

        val onProgress = options.onProgress
        if (onProgress != null && (step + 1) % options.progressEvery == 0) onProgress(step + 1, lastLoss)
    }

    return DeepSetsWeights(
        psi = psiShape.toMlpWeights(psiParams),
        phi = phiShape.toMlpWeights(phiParams),
        messageWidth = options.messageWidth,
        radius = radius,
        velocityScale = velocityScale,
        accelScale = accelScale,
        sumScale = sumScale,
        minSeparation = minSeparation,
    )
}

The layer shape of one small network, and the two things a trainer needs from it: a forward pass over AD and a conversion into the flat MlpWeights layout a shader reads.

tanh throughout the hidden layers, and a linear output. tanh rather than the leaky that flowMatch prefers because this is a bounded regression onto a bounded target, and the argument for leaky there — that a velocity field has to keep growing away from the data without bound — does not apply to an interaction with a finite radius.

source — 47 lines, from line 396 
private class MlpShape(val inputs: Int, val hidden: IntArray, val outputs: Int) {
    val widths: IntArray = hidden + outputs
    val inputWidths: IntArray = intArrayOf(inputs) + hidden
    val parameterCount: Int = MlpWeights.floatCount(inputs, widths)

    /** Xavier, so the initial activations neither saturate `tanh` nor vanish through the depth. */
    fun initialise(random: Random): DoubleArray {
        val p = DoubleArray(parameterCount)
        var at = 0
        for (layer in widths.indices) {
            val n = widths[layer]
            val below = inputWidths[layer]
            val limit = sqrt(6.0 / (below + n))
            for (o in 0 until n) for (i in 0 until below) p[at++] = (random.nextDouble() * 2 - 1) * limit
            for (o in 0 until n) p[at++] = 0.0
        }
        return p
    }

    fun forward(x: List<AD>, params: List<AD>): List<AD> {
        var source = x
        var at = 0
        for (layer in widths.indices) {
            val n = widths[layer]
            val below = inputWidths[layer]
            val weightBase = at
            val biasBase = at + n * below
            val destination = ArrayList<AD>(n)
            for (o in 0 until n) {
                var sum = params[biasBase + o]
                for (i in 0 until below) sum = sum + params[weightBase + o * below + i] * source[i]
                destination += if (layer == widths.size - 1) sum else tanh(sum)
            }
            at = biasBase + n
            source = destination
        }
        return source
    }

    fun toMlpWeights(params: DoubleArray) = MlpWeights(
        inputs = inputs,
        outputs = outputs,
        hidden = hidden.copyOf(),
        activation = MlpActivation.TANH,
        flat = FloatArray(params.size) { params[it].toFloat() },
    )
}

Adam, in the fifteen lines it actually is.

private class Adam(size: Int, val beta1: Double = 0.9, val beta2: Double = 0.999, val epsilon: Double = 1e-8) {
    private val m = DoubleArray(size)
    private val v = DoubleArray(size)
    private val delta = DoubleArray(size)
    private var t = 0

    inline fun step(rate: Double, gradient: (Int) -> Double): DoubleArray {
        t++
        val correction1 = 1.0 - Math.pow(beta1, t.toDouble())
        val correction2 = 1.0 - Math.pow(beta2, t.toDouble())
        for (i in delta.indices) {
            val g = gradient(i)
            m[i] = beta1 * m[i] + (1 - beta1) * g
            v[i] = beta2 * v[i] + (1 - beta2) * g * g
            delta[i] = rate * (m[i] / correction1) / (sqrt(v[i] / correction2) + epsilon)
        }
        return delta
    }
}

fun DeepSetsWeights.glsl(prefix: String = "learned"): String

GLSL for the whole learned interaction: both networks, the feature map, and the combination.

Emits, for prefix = "learned":

  • learnedPsiWeights / learnedPhiWeights buffer-reference types and the two evaluation functions, straight from MlpWeights.glsl — the same emitter the flow-matching work uses, so there is one implementation of "an MLP as a shader" and one parity test covering it.
  • learnedFeatures(vec2 d, vec2 dv, out float f[8]), which must agree with interactionFeatures exactly.
  • learnedCombine(...), the phi half: velocity, the summed message and the meaned message.

The aggregation loop itself is not emitted, because it belongs to the caller's neighbour walk — the consumer already has a loop over the 3x3 block, and the message sum has to accumulate inside it.

Place immediately after #version: this carries #extension directives.

body — 56 lines, from line 483 
fun DeepSetsWeights.glsl(prefix: String = "learned"): String = buildString {
    // Captured before the builder starts, because inside `buildString` the implicit receiver is
    // the StringBuilder: writing `$this` in the template below silently interpolates everything
    // appended so far, which emits both networks a second time inside what looks like a comment
    // and fails as "block name cannot be redefined" forty lines later.
    val description = this@glsl.toString()
    append(psi.glsl("${prefix}Psi"))
    append("\n")
    append(phi.glsl("${prefix}Phi"))
    append("\n")
    append(
        """
        // $description
        const int ${prefix}MessageWidth = $messageWidth;
        const float ${prefix}Radius = ${glslFloat(radius)};
        const float ${prefix}VelocityScale = ${glslFloat(velocityScale)};
        const float ${prefix}AccelScale = ${glslFloat(accelScale)};
        const float ${prefix}SumScale = ${glslFloat(sumScale)};
        const float ${prefix}MinSeparation = ${glslFloat(minSeparation)};

        // Must agree with interactionFeatures() in deepsets.kt, term for term.
        void ${prefix}Features(vec2 d, vec2 dv, out float f[${DeepSetsWeights.PSI_INPUTS}]) {
            float distance = length(d);
            float soft = max(distance, ${prefix}MinSeparation);
            // Zero at zero distance: the direction is undefined there and any other answer
            // invents a force out of a rounding error.
            vec2 u = distance > 0.0 ? d / distance : vec2(0.0);
            float p = ${prefix}MinSeparation / soft;
            f[0] = d.x / ${prefix}Radius;
            f[1] = d.y / ${prefix}Radius;
            f[2] = dv.x / ${prefix}VelocityScale;
            f[3] = dv.y / ${prefix}VelocityScale;
            f[4] = u.x * p;
            f[5] = u.y * p;
            f[6] = u.x * p * p;
            f[7] = u.y * p * p;
        }

        vec2 ${prefix}Combine(
            ${prefix}PhiWeights pw, vec2 velocity, float messageSum[${prefix}MessageWidth], float n
        ) {
            float x[${phi.inputs}];
            x[0] = velocity.x / ${prefix}VelocityScale;
            x[1] = velocity.y / ${prefix}VelocityScale;
            for (int k = 0; k < ${prefix}MessageWidth; ++k) {
                x[2 + k] = messageSum[k] / ${prefix}SumScale;
                x[2 + ${prefix}MessageWidth + k] = n > 0.0 ? messageSum[k] / n : 0.0;
            }
            float y[2];
            ${prefix}Phi(pw, x, y);
            return vec2(y[0], y[1]) * ${prefix}AccelScale;
        }
        """.trimIndent()
    )
    append("\n")
}

A float as GLSL source, always with a decimal point so it cannot be read as an int.

private fun glslFloat(v: Float): String {
    val s = v.toString()
    return if ('.' in s || 'e' in s || 'E' in s) s else "$s.0"
}

The equivariant form

class EquivariantWeights

A learned interaction that is rotationally equivariant by construction — a(R x) = R a(x) holds exactly, not approximately.

DeepSetsWeights learns a map from Cartesian components to Cartesian components, and nothing in it forces that map to commute with rotation. Measured on the flocking rule: on a perfectly symmetric ring of neighbours, where the true force is exactly zero by symmetry, the plain form produces a force of 0.14 at twelve neighbours and 0.50 at forty-eight — growing with crowding, because it rides on the aggregated message — and rotating a configuration by ninety degrees changes its answer by 18% of a typical force. A fixed-direction residual that scales with neighbour count is a steady drift, and in the learned_boids scene it is what carried the learned flock out of its containment while the written flock stayed put.

The construction that removes it is the standard one and costs nothing:

  • psi sees only invariants of a neighbour — r, the proximity p, |dv|, and u . dv — and emits scalars: a pair (a_k, b_k) per channel.
  • the per-neighbour message is a vector, a_k * u + b_k * dv, built from the neighbour's own direction and velocity difference.
  • those are summed and meaned into channels aggregated vectors.
  • phi sees only invariants of those — their lengths, and |v| — and emits scalars again, which weight the aggregated vectors into the output.

Rotate every input and the invariants are unchanged, so every scalar the two networks emit is unchanged, so the output rotates with the inputs. Exactly. And for a symmetric neighbourhood the direction vectors cancel before any network sees them, so the answer is identically zero.

Boids is still inside the family: separation is a_k ∝ p against u, cohesion is a_k ∝ r against u, alignment is b_k = 1 against dv. Nothing is given up.

class EquivariantWeights(
    /** Invariants to `2 * channels` scalars: `a_k` then `b_k`. */
    val psi: MlpWeights,
    /** Aggregate invariants to `2 * channels + 1` scalars: the sum weights, mean weights, and `v`. */
    val phi: MlpWeights,
    /** How many vector channels are aggregated. */
    val channels: Int,
    /** Interaction radius, in world units. Neighbour distances are divided by this. */
    val radius: Float,
    /** Divides velocity differences, and the centre particle's own speed, into the invariants. */
    val velocityScale: Float,
    /** The output is multiplied by this to get a world-units acceleration. */
    val accelScale: Float,
    /** The summed channels are divided by this, so a typical crowd lands near 1. */
    val sumScale: Float,
    /** Closest approach used by the proximity invariant, the same floor the rule itself uses. */
    val minSeparation: Float,
) {

init

    init {
        require(psi.inputs == PSI_INVARIANTS) { "psi takes $PSI_INVARIANTS invariants, not ${psi.inputs}" }
        require(psi.outputs == 2 * channels) { "psi emits 2 x $channels scalars, not ${psi.outputs}" }
        require(phi.inputs == 2 * channels + 1) { "phi takes ${2 * channels + 1} invariants, not ${phi.inputs}" }
        require(phi.outputs == 2 * channels + 1) { "phi emits ${2 * channels + 1} scalars, not ${phi.outputs}" }
    }

val parameterCount: Int get() = psi.parameterCount + phi.parameterCount

How many weights and biases in total.

fun evaluate(vx: Float, vy: Float, neighbours: List<FloatArray>): Pair<Float, Float>

The acceleration, in world units. Equivariant to fp32.

    fun evaluate(vx: Float, vy: Float, neighbours: List<FloatArray>): Pair<Float, Float> {
        val sumX = FloatArray(channels)
        val sumY = FloatArray(channels)
        val features = FloatArray(PSI_INVARIANTS)
        val coefficients = FloatArray(2 * channels)

        for (n in neighbours) {
            val distance = sqrt((n[0] * n[0] + n[1] * n[1]).toDouble()).toFloat()
            val ux = if (distance > 0f) n[0] / distance else 0f
            val uy = if (distance > 0f) n[1] / distance else 0f
            invariantFeatures(n[0], n[1], n[2], n[3], radius, velocityScale, minSeparation, features)
            psi.evaluate(features, coefficients)
            for (k in 0 until channels) {
                val a = coefficients[k]
                val b = coefficients[channels + k]
                sumX[k] += a * ux + b * n[2] / velocityScale
                sumY[k] += a * uy + b * n[3] / velocityScale
            }
        }

        val count = neighbours.size
        val input = FloatArray(phi.inputs)
        for (k in 0 until channels) {
            input[k] = sqrt((sumX[k] * sumX[k] + sumY[k] * sumY[k]).toDouble()).toFloat() / sumScale
            input[channels + k] = if (count > 0) input[k] * sumScale / count else 0f
        }
        input[2 * channels] = sqrt((vx * vx + vy * vy).toDouble()).toFloat() / velocityScale
        val w = phi.evaluate(input)

        var ax = 0f
        var ay = 0f
        for (k in 0 until channels) {
            val mean = if (count > 0) 1f / count else 0f
            ax += w[k] * sumX[k] / sumScale + w[channels + k] * sumX[k] * mean
            ay += w[k] * sumY[k] / sumScale + w[channels + k] * sumY[k] * mean
        }
        ax += w[2 * channels] * vx / velocityScale
        ay += w[2 * channels] * vy / velocityScale
        return ax * accelScale to ay * accelScale
    }

override fun toString(): String =

Both network shapes, the channel count and the parameter total. The GLSL emitter writes this line into the shader as a comment.

        "EquivariantWeights(psi $psi, phi $phi, $channels channels, $parameterCount parameters)"

companion object

    companion object {

const val PSI_INVARIANTS = 5

[r/R, p, p^2, |dv|/V, (u . dv)/V] — see invariantFeatures.

fun invariantFeatures

The rotation-invariant description of one neighbour: [r/R, p, p^2, |dv|/V, (u . dv)/V].

Every entry is a length or a dot product, so a rotation of the configuration leaves all five unchanged — which is the whole mechanism. p = minSeparation / max(r, minSeparation) is the same bounded proximity interactionFeatures uses, and it is here for the same reason: a 1/r law is linear in it.

fun invariantFeatures(
    dx: Float,
    dy: Float,
    dvx: Float,
    dvy: Float,
    radius: Float,
    velocityScale: Float,
    minSeparation: Float,
    into: FloatArray = FloatArray(EquivariantWeights.PSI_INVARIANTS),
): FloatArray {
    val distance = sqrt((dx * dx + dy * dy).toDouble()).toFloat()
    val soft = maxOf(distance, minSeparation)
    val ux = if (distance > 0f) dx / distance else 0f
    val uy = if (distance > 0f) dy / distance else 0f
    val p = minSeparation / soft
    into[0] = distance / radius
    into[1] = p
    into[2] = p * p
    into[3] = sqrt((dvx * dvx + dvy * dvy).toDouble()).toFloat() / velocityScale
    into[4] = (ux * dvx + uy * dvy) / velocityScale
    return into
}

fun trainEquivariant

Fits an EquivariantWeights to whatever rule examples embodies.

Same trainer as trainDeepSets — Adam over a reverse-mode tape — but the forward pass builds vectors out of scalar coefficients, so the gradient flows back through the aggregation into coefficients that never see an absolute direction.

No rotation augmentation is needed or useful here. Augmentation is how you ask a non-equivariant model to be approximately equivariant; measured on the plain form it takes the symmetric-ring artefact from 0.50 to 0.13, which is a good return on one line and still not zero. This is zero.

body — 98 lines, from line 706 
fun trainEquivariant(
    radius: Float,
    velocityScale: Float,
    accelScale: Float,
    sumScale: Float,
    minSeparation: Float,
    options: DeepSetsOptions = DeepSetsOptions(),
    examples: (Int) -> List<InteractionExample>,
): EquivariantWeights {
    require(options.steps > 0) { "trainEquivariant needs at least one step" }
    val channels = options.messageWidth
    val random = Random(options.seed)
    val psiShape = MlpShape(EquivariantWeights.PSI_INVARIANTS, options.psiHidden, 2 * channels)
    val phiShape = MlpShape(2 * channels + 1, options.phiHidden, 2 * channels + 1)

    val psiParams = psiShape.initialise(random)
    val phiParams = phiShape.initialise(random)
    val adam = Adam(psiParams.size + phiParams.size)

    val features = FloatArray(EquivariantWeights.PSI_INVARIANTS)
    for (step in 0 until options.steps) {
        val tape = ADTape()
        val psiVars = tape.variables(psiParams)
        val phiVars = tape.variables(phiParams)
        var loss = tape.constant(0.0)

        for (example in examples(step)) {
            val sumX = MutableList(channels) { tape.constant(0.0) }
            val sumY = MutableList(channels) { tape.constant(0.0) }
            for (n in example.neighbours) {
                val distance = sqrt((n[0] * n[0] + n[1] * n[1]).toDouble())
                val ux = if (distance > 0) n[0] / distance else 0.0
                val uy = if (distance > 0) n[1] / distance else 0.0
                invariantFeatures(n[0], n[1], n[2], n[3], radius, velocityScale, minSeparation, features)
                val c = psiShape.forward(features.map { tape.constant(it.toDouble()) }, psiVars)
                for (k in 0 until channels) {
                    sumX[k] = sumX[k] + c[k] * ux + c[channels + k] * (n[2].toDouble() / velocityScale)
                    sumY[k] = sumY[k] + c[k] * uy + c[channels + k] * (n[3].toDouble() / velocityScale)
                }
            }

            val count = example.neighbours.size
            val input = ArrayList<AD>(phiShape.inputs)
            for (k in 0 until channels) {
                // A length is invariant; a component is not. This is the only place the aggregate
                // is allowed to reach the scalar network, and it has to be through a norm.
                input += sqrt(sumX[k].sq() + sumY[k].sq() + 1e-12) / sumScale.toDouble()
            }
            for (k in 0 until channels) {
                input += if (count > 0) {
                    sqrt(sumX[k].sq() + sumY[k].sq() + 1e-12) / count.toDouble()
                } else tape.constant(0.0)
            }
            input += tape.constant(
                sqrt((example.velocity[0] * example.velocity[0] +
                    example.velocity[1] * example.velocity[1]).toDouble()) / velocityScale
            )
            val w = phiShape.forward(input, phiVars)

            var ax = tape.constant(0.0)
            var ay = tape.constant(0.0)
            val mean = if (count > 0) 1.0 / count else 0.0
            for (k in 0 until channels) {
                ax = ax + w[k] * sumX[k] / sumScale.toDouble() + w[channels + k] * sumX[k] * mean
                ay = ay + w[k] * sumY[k] / sumScale.toDouble() + w[channels + k] * sumY[k] * mean
            }
            ax = ax + w[2 * channels] * (example.velocity[0].toDouble() / velocityScale)
            ay = ay + w[2 * channels] * (example.velocity[1].toDouble() / velocityScale)

            loss = loss + (ax - example.target[0].toDouble() / accelScale).sq() +
                (ay - example.target[1].toDouble() / accelScale).sq()
        }

        loss = loss / (options.batch * 2).toDouble()
        val gradients = tape.backward(loss)
        val rate = options.learningRate +
            (options.finalLearningRate - options.learningRate) * step / (options.steps - 1).coerceAtLeast(1)
        val delta = adam.step(rate) { i ->
            if (i < psiParams.size) gradients[psiVars[i].index] else gradients[phiVars[i - psiParams.size].index]
        }
        for (i in psiParams.indices) psiParams[i] -= delta[i]
        for (i in phiParams.indices) phiParams[i] -= delta[psiParams.size + i]

        val onProgress = options.onProgress
        if (onProgress != null && (step + 1) % options.progressEvery == 0) onProgress(step + 1, loss.value)
    }

    return EquivariantWeights(
        psi = psiShape.toMlpWeights(psiParams),
        phi = phiShape.toMlpWeights(phiParams),
        channels = channels,
        radius = radius,
        velocityScale = velocityScale,
        accelScale = accelScale,
        sumScale = sumScale,
        minSeparation = minSeparation,
    )
}

fun EquivariantWeights.glsl(prefix: String = "learned"): String

GLSL for the equivariant interaction: both networks, the invariant features, and the combination — the shader half of EquivariantWeights.

The aggregation loop is the caller's, as with DeepSetsWeights.glsl, but here what it accumulates is channels vec2s rather than scalars:

float f[5];  learnedInvariants(d, dv, f);
float c[2 * K]; learnedPsi(pc.psiWeights, f, c);
for (int k = 0; k < K; ++k) messageSum[k] += c[k] * u + c[K + k] * dv / V;

Place immediately after #version: this carries #extension directives.

body — 58 lines, from line 820 
fun EquivariantWeights.glsl(prefix: String = "learned"): String = buildString {
    // Captured before the builder: inside `buildString` the implicit receiver is the
    // StringBuilder, so `$this` in the template below would paste everything appended so far.
    val description = this@glsl.toString()
    append(psi.glsl("${prefix}Psi"))
    append("\n")
    append(phi.glsl("${prefix}Phi"))
    append("\n")
    append(
        """
        // $description
        const int ${prefix}Channels = $channels;
        const float ${prefix}Radius = ${glslFloat(radius)};
        const float ${prefix}VelocityScale = ${glslFloat(velocityScale)};
        const float ${prefix}AccelScale = ${glslFloat(accelScale)};
        const float ${prefix}SumScale = ${glslFloat(sumScale)};
        const float ${prefix}MinSeparation = ${glslFloat(minSeparation)};

        // Must agree with invariantFeatures() in deepsets.kt. Every entry is a length or a dot
        // product, which is the whole mechanism: rotate the configuration and these do not move.
        void ${prefix}Invariants(vec2 d, vec2 dv, out float f[${EquivariantWeights.PSI_INVARIANTS}]) {
            float distance = length(d);
            float soft = max(distance, ${prefix}MinSeparation);
            vec2 u = distance > 0.0 ? d / distance : vec2(0.0);
            float p = ${prefix}MinSeparation / soft;
            f[0] = distance / ${prefix}Radius;
            f[1] = p;
            f[2] = p * p;
            f[3] = length(dv) / ${prefix}VelocityScale;
            f[4] = dot(u, dv) / ${prefix}VelocityScale;
        }

        vec2 ${prefix}Combine(
            ${prefix}PhiWeights pw, vec2 velocity, vec2 messageSum[${prefix}Channels], float n
        ) {
            float x[${phi.inputs}];
            for (int k = 0; k < ${prefix}Channels; ++k) {
                x[k] = length(messageSum[k]) / ${prefix}SumScale;
                x[${prefix}Channels + k] = n > 0.0 ? length(messageSum[k]) / n : 0.0;
            }
            x[2 * ${prefix}Channels] = length(velocity) / ${prefix}VelocityScale;

            float w[${phi.outputs}];
            ${prefix}Phi(pw, x, w);

            vec2 a = vec2(0.0);
            float mean = n > 0.0 ? 1.0 / n : 0.0;
            for (int k = 0; k < ${prefix}Channels; ++k) {
                a += w[k] * messageSum[k] / ${prefix}SumScale
                   + w[${prefix}Channels + k] * messageSum[k] * mean;
            }
            a += w[2 * ${prefix}Channels] * velocity / ${prefix}VelocityScale;
            return a * ${prefix}AccelScale;
        }
        """.trimIndent()
    )
    append("\n")
}