CeresExpression.kt

src/jvmMain/kotlin/afield/math/optimization/ceres/CeresExpression.kt · 998 lines

@file:afield.runtime.completion.CompletionCategory("Solvers")

package afield.math.optimization.ceres

import kotlin.math.PI

Residuals written in Kotlin, compiled to a little register machine, and differentiated exactly by Ceres.

Every other constraint in this vocabulary is a C++ functor with a matching opcode, which means a new kind of residual costs a native rebuild, a Kotlin change, and a redeployed dylib. This is the way out: build an expression here, CeresExpressionResidual encodes it, and the native EXPR opcode evaluates it. The evaluator is templated on the scalar type, so Ceres runs it on ceres::Jet and gets analytic derivatives — not finite differences — for a residual it has never seen.

// "these two points are exactly 2.5 apart", written from scratch
val program = buildCeresProgram {
    expression(a, b) { (pa, pb) ->
        listOf((pa - pb).length() - 2.5.c)
    }
}

What it costs

A hand-written functor is faster: the VM interprets, allocates a slot vector per evaluation, and Ceres cannot specialise the parameter-block sizes the way [ceres::AutoDiffCostFunction] does with a fixed arity. Reach for a real opcode when a residual is hot and settled; reach for this while you are still finding out what the residual should be, or when it is caller-specific and will never belong in the wrapper.

What to watch

• abs, min and max are not differentiable at their kinks. Ceres will hand you a subgradient there, which is usually what you want for a one-sided barrier and is not what you want in the middle of a smooth objective.
• Nothing checks for division by zero or a square root of a negative. A residual that produces a NaN will make the solve fail rather than corrupt anything, but it will not tell you where.
• Ceres forbids one residual block naming the same point twice, so expression rejects a repeated point rather than letting the native side discover it.

Sharing

The expression is a DAG, not a tree: CeresExpr instances are hash-consed as you build them, so writing d / d evaluates d once. That is why the machine is register-based — with named destination slots, sharing is free.

Free variables

An expression can also read unknowns of its own — quantities Ceres solves for that are not points. CeresProgramBuilder.variable declares one, with an initial value and an optional CeresManifold, and hands back a CeresVariableRef that any number of expressions can name:

// fit a sphere: the centre and the radius are unknowns, not points
val centre = variable("centre", Vec3(0.0, 0.0, 0.0))
val radius = variable("radius", 1.0)
points.forEach { i ->
    expression(listOf(i), listOf(centre, radius)) { (p), (c, r) ->
        listOf(p.distanceTo(c.vec()) - r.scalar)
    }
}

This is the mechanism PLANAR_HESSE hardcodes in C++ — a plane as a unit normal on a SphereManifold plus a scalar offset — made available to a residual written at runtime. It is worth the machinery: on kkC the plane-as-a-free-variable formulation converges in 12 iterations to tangency 3e-08 where the same objective with the plane eliminated inside the residual does not converge at all. Well-conditioned auxiliary unknowns are not overhead; they are frequently what makes a problem tractable.

The solved values come back in SolveResult.variables, in declaration order.

sealed interface CeresExpr

One node of a residual expression: a constant, a single coordinate of one of the residual's points, or an operation over other nodes. Immutable, and hash-consed as you build, so an identical subexpression written twice is one node and is evaluated once.

The mixed-arithmetic operators are declared here and only here, so the expression has to be on the left: expr * 2.0 compiles, 2.0 * expr does not. Write 2.0.c * expr — see Double.c — or reorder. There is no comparison, no boolean and no conditional: a residual is an arithmetic expression, and the one-sided behaviour a barrier needs comes from max and barrierBelow instead.

sealed interface CeresExpr {

operator fun plus(other: CeresExpr): CeresExpr = binary(ExprOp.ADD, this, other)

Sum. The Double overload wraps the number as a constant node first.

operator fun minus(other: CeresExpr): CeresExpr = binary(ExprOp.SUB, this, other)

Difference. The Double overload wraps the number as a constant node first.

operator fun times(other: CeresExpr): CeresExpr = binary(ExprOp.MUL, this, other)

Product. The Double overload wraps the number as a constant node first.

operator fun div(other: CeresExpr): CeresExpr = binary(ExprOp.DIV, this, other)

Quotient, with no check for a zero divisor. The Double overload wraps the number first.

operator fun unaryMinus(): CeresExpr = unary(ExprOp.NEG, this)

Negation, as one instruction rather than a multiply by minus one.

operator fun plus(other: Double): CeresExpr = this + constant(other)

operator fun minus(other: Double): CeresExpr = this - constant(other)

operator fun times(other: Double): CeresExpr = this * constant(other)

operator fun div(other: Double): CeresExpr = this / constant(other)

internal enum class ExprOp(val code: Int, val arity: Int)

internal enum class ExprOp(val code: Int, val arity: Int) {
    CONST(0, 0),
    PARAM(1, 0),
    ADD(2, 2),
    SUB(3, 2),
    MUL(4, 2),
    DIV(5, 2),
    NEG(6, 1),
    SQRT(7, 1),
    ABS(8, 1),
    SIN(9, 1),
    COS(10, 1),
    TAN(11, 1),
    EXP(12, 1),
    LOG(13, 1),
    POW(14, 2),
    MIN(15, 2),
    MAX(16, 2),
    ATAN2(17, 2),
    SQUARE(18, 1),
    RECIP(19, 1),

A component of one of the residual's free variables. See CeresVariableRef.

    VAR(20, 0),
}

internal data class ExprConst(val value: Double) : CeresExpr

init

    init {
        require(value.isFinite()) { "expression constants must be finite" }
    }
}

internal data class ExprParam(val block: Int, val component: Int) : CeresExpr

Component component (0=x, 1=y, 2=z) of the expression's block-th point.

init

    init {
        require(block >= 0) { "parameter block index must be non-negative" }
        require(component in 0..2) { "component must be 0, 1 or 2" }
    }
}

internal data class ExprVar(val block: Int, val component: Int) : CeresExpr

Component component of the residual's block-th free variable.

block is the position in the list of variables handed to CeresProgramBuilder.expression — not the program-wide declaration index. A residual block only receives the parameter blocks it names, so the VM addresses them locally; CeresVariableRef holds the program-wide index and the encoder writes that.

init

    init {
        require(block >= 0) { "variable index must be non-negative" }
        require(component >= 0) { "variable component must be non-negative" }
    }
}

internal data class ExprNode(val op: ExprOp, val a: CeresExpr, val b: CeresExpr?) : CeresExpr

private fun binary(op: ExprOp, a: CeresExpr, b: CeresExpr): CeresExpr = ExprNode(op, a, b)
private fun unary(op: ExprOp, a: CeresExpr): CeresExpr = ExprNode(op, a, null)

fun constant(value: Double): CeresExpr = ExprConst(value)

A literal. 2.5.c reads better than constant(2.5) inside an expression.

Constants are pooled by value when the expression is compiled, so repeating one is free. Must be finite: a NaN or an infinity here would only surface as a failed solve with no indication of where.

val Double.c: CeresExpr get() = ExprConst(this)

A literal, as a property. Also the fix for 2.0 * expr, which does not compile — write 2.0.c * expr.

val Int.c: CeresExpr get() = ExprConst(toDouble())

A literal, as a property.

fun sqrt(x: CeresExpr): CeresExpr = unary(ExprOp.SQRT, x)

Square root. No guard: the native evaluator will happily take the root of a negative and produce a NaN, which fails the solve without saying where. Prefer CeresExprVec.length, which keeps a small epsilon under the root precisely so the derivative at zero stays finite.

fun square(x: CeresExpr): CeresExpr = unary(ExprOp.SQUARE, x)

x * x as one instruction, and one node — cheaper and clearer than multiplying an expression by itself.

fun reciprocal(x: CeresExpr): CeresExpr = unary(ExprOp.RECIP, x)

1 / x as one instruction. Nothing checks for zero; see sqrt for what a NaN costs you.

fun exp(x: CeresExpr): CeresExpr = unary(ExprOp.EXP, x)

e to the power x, as one instruction. The inverse of ln.

fun ln(x: CeresExpr): CeresExpr = unary(ExprOp.LOG, x)

Natural log. Undefined at and below zero, and unguarded.

fun sin(x: CeresExpr): CeresExpr = unary(ExprOp.SIN, x)

Sine, in radians — as everything angular here is. CERES_PI is on hand for the literal.

fun cos(x: CeresExpr): CeresExpr = unary(ExprOp.COS, x)

Cosine, in radians.

fun tan(x: CeresExpr): CeresExpr = unary(ExprOp.TAN, x)

Tangent, in radians. Unbounded at the odd multiples of pi/2, and unguarded.

fun abs(x: CeresExpr): CeresExpr = unary(ExprOp.ABS, x)

Not differentiable at zero; Ceres will use a subgradient.

fun min(a: CeresExpr, b: CeresExpr): CeresExpr = binary(ExprOp.MIN, a, b)

Not differentiable where the arguments cross.

fun max(a: CeresExpr, b: CeresExpr): CeresExpr = binary(ExprOp.MAX, a, b)

Not differentiable where the arguments cross. The building block of a one-sided barrier.

fun pow(base: CeresExpr, exponent: CeresExpr): CeresExpr = binary(ExprOp.POW, base, exponent)

base to the power exponent, both expressions — so the exponent may itself be solved for. For a fixed integer power prefer square or a product; this is the general and slowest case.

fun atan2(y: CeresExpr, x: CeresExpr): CeresExpr = binary(ExprOp.ATAN2, y, x)

Two-argument arctangent, in radians, y first. The only way to get a signed angle out of this vocabulary.

fun barrierBelow(x: CeresExpr, floor: CeresExpr): CeresExpr = max(0.0.c, floor - x)

max(0, floor - x): zero while x >= floor, and a penalty below it.

Contributing exactly nothing while satisfied is the useful part — a barrier can be added to a tuned objective without changing its answer where it was already healthy. The kink at the floor is non-differentiable and Ceres takes a subgradient, which is what you want here.

fun barrierAbove(x: CeresExpr, ceiling: CeresExpr): CeresExpr = max(0.0.c, x - ceiling)

max(0, x - ceiling): zero while x <= ceiling, and a penalty above it. The mirror of barrierBelow.

data class CeresExprVec(val x: CeresExpr, val y: CeresExpr, val z: CeresExpr)

One of the expression's points, as three scalar expressions — the vector algebra you actually want to write a residual in.

Nothing here is a new kind of node: every operator expands to scalar arithmetic on the three components as you write it, so (a - b).lengthSquared() is an expression tree, not a call the VM knows about. Sharing still saves the work — a difference used twice compiles once.

The scalar operators are not quite symmetric: vec * 2.0 works because times has a Double overload, vec / 2.0 does not, because div only takes a CeresExpr. Write vec / 2.0.c.

data class CeresExprVec(val x: CeresExpr, val y: CeresExpr, val z: CeresExpr) {

operator fun plus(o: CeresExprVec) = CeresExprVec(x + o.x, y + o.y, z + o.z)

Componentwise sum: three scalar additions, written once.

operator fun minus(o: CeresExprVec) = CeresExprVec(x - o.x, y - o.y, z - o.z)

Componentwise difference, which is how an edge or an offset is written.

operator fun times(k: CeresExpr) = CeresExprVec(x * k, y * k, z * k)

Scale all three components. The Double overload goes through Double.c.

operator fun times(k: Double) = times(k.c)

operator fun div(k: CeresExpr) = CeresExprVec(x / k, y / k, z / k)

Divide all three components. There is no Double overload, so write vec / 2.0.c.

operator fun unaryMinus() = CeresExprVec(-x, -y, -z)

Negate all three components, reversing the direction.

infix fun dot(o: CeresExprVec): CeresExpr = x * o.x + y * o.y + z * o.z

Dot product — five instructions, and the basis of every angle and projection residual here.

infix fun cross(o: CeresExprVec)

Cross product. Three residuals whose magnitude is an area (or, between unit vectors, a sine), so it makes a good scale-free "these are parallel" test without any square root.

    infix fun cross(o: CeresExprVec) = CeresExprVec(
        y * o.z - z * o.y,
        z * o.x - x * o.z,
        x * o.y - y * o.x,
    )

fun lengthSquared(): CeresExpr = this dot this

Squared length. Prefer this to length() when you only need to compare.

fun length(epsilon: Double = 1e-20): CeresExpr = sqrt(lengthSquared() + epsilon.c)

Length, with epsilon added under the root so the derivative stays finite at the origin. A plain sqrt(0) has an infinite derivative and will poison the Jacobian.

fun distanceTo(o: CeresExprVec, epsilon: Double = 1e-20): CeresExpr = (this - o).length(epsilon)

Distance between two points, epsilon-guarded like length. The most-reached-for residual in the vocabulary.

companion object

    companion object {

val ORIGIN = CeresExprVec(0.0.c, 0.0.c, 0.0.c)

The world origin as constants — for a residual measured from a fixed point rather than between two points.

fun of(v: Vec3) = CeresExprVec(v.x.c, v.y.c, v.z.c)

A fixed position, frozen into the expression as three constants.

The vector is read when the expression is built, so a later change to it has no effect and the point is not something the solver can move. That is the only way to reach outside the residual's own parameter blocks: an expression cannot read a point it was not bound to, and cannot declare a free variable of its own.

const val CERES_PI: Double = PI

Handy when an expression wants a literal angle; every trigonometric op here is in radians.

Free variables

enum class CeresManifold(val code: Int)

The space a free variable lives in — which is to say, how Ceres is allowed to step it.

A manifold is not decoration. PLANAR_HESSE's whole advantage over PLANAR is that its plane normal is on a SPHERE: three raw components, two effective parameters, and no freedom to change length. Left as a plain 3-vector with a "keep it unit" residual, the same variable would carry a scale direction that the objective barely constrains, and a trust-region minimiser will find it.

enum class CeresManifold(val code: Int) {

Euclidean: every component free, tangent space the same size as the block.

    NONE(0),

Unit norm (ceres::SphereManifold). Needs at least two components.

    SPHERE(1),

A rotation as a unit quaternion, w first (ceres::QuaternionManifold). Exactly four components, three effective parameters.

    QUATERNION(2),
}

data class CeresVariableDecl

An unknown Ceres solves for that is not a point: dimension scalars with an initial value and a manifold.

Declare one with CeresProgramBuilder.variable. The name never crosses the ABI — the native side is handed dimensions and initial values, and this side reconstructs names when it reads the solved values back — so it costs nothing and is worth choosing well; it is what CompiledCeresExpression.disassemble and SolveResult.variables label things with.

data class CeresVariableDecl(
    /** A label for this side only; it never crosses the ABI. Blank is rejected. */
    val name: String,
    /** The starting value, one entry per component, and what fixes [dimension]. */
    val initial: List<Double>,
    /** The geometry the solver steps the block on. Euclidean unless you say otherwise. */
    val manifold: CeresManifold = CeresManifold.NONE,
    /**
     * Per-component box constraints, `component to (lower .. upper)`. A component may appear once.
     *
     * These are Ceres' own `SetParameterLowerBound` / `SetParameterUpperBound`, and they are the right
     * tool for "this radius is positive" — a one-sided barrier residual can only *penalise* an
     * excursion after the minimiser has already taken it, and it changes the objective you wrote down
     * while doing so.
     *
     * Only on [CeresManifold.NONE]. A bound is a statement about an ambient coordinate, and a
     * manifold's `Plus` is free to step straight through one, so the pair does not mean what it
     * looks like it means. Rejected rather than silently allowed.
     */
    val bounds: Map<Int, ClosedFloatingPointRange<Double>> = emptyMap(),
) {

val dimension: Int get() = initial.size

How many scalars the block holds, which is initial.size and is capped by the wrapper.

init

    init {
        require(name.isNotBlank()) { "a free variable needs a name" }
        require(initial.isNotEmpty()) { "free variable '$name' needs at least one component" }
        require(initial.size <= MAX_DIMENSION) {
            "free variable '$name' has ${initial.size} components; the limit is $MAX_DIMENSION"
        }
        require(initial.all { it.isFinite() }) { "free variable '$name' has a non-finite initial value" }
        when (manifold) {
            CeresManifold.NONE -> Unit
            CeresManifold.SPHERE -> {
                require(initial.size >= 2) {
                    "free variable '$name' is on a sphere manifold, which needs at least two components"
                }
                require(initial.sumOf { it * it } > 1e-20) {
                    "free variable '$name' is on a sphere manifold and cannot start at the origin"
                }
            }

            CeresManifold.QUATERNION -> {
                require(initial.size == 4) {
                    "free variable '$name' is a quaternion, which needs exactly four components (w first)"
                }
                require(initial.sumOf { it * it } > 1e-20) {
                    "free variable '$name' is a quaternion and cannot start at zero"
                }
            }
        }
        if (bounds.isNotEmpty()) {
            require(manifold == CeresManifold.NONE) {
                "free variable '$name' is on a ${manifold.name.lowercase()} manifold, which steps in the " +
                    "tangent space and so can walk straight through a bound on an ambient coordinate"
            }
            bounds.forEach { (component, range) ->
                require(component in 0 until dimension) {
                    "free variable '$name' has $dimension component(s); bound given for $component"
                }
                require(range.start.isFinite() || range.endInclusive.isFinite()) {
                    "the bound on '$name'[$component] constrains nothing; leave it out instead"
                }
                require(range.start <= range.endInclusive) {
                    "the bound on '$name'[$component] is empty: ${range.start} .. ${range.endInclusive}"
                }
                require(initial[component] in range) {
                    "free variable '$name'[$component] starts at ${initial[component]}, outside its " +
                        "bound ${range.start} .. ${range.endInclusive}; Ceres requires a feasible start"
                }
            }
        }
    }

internal fun boundOf(component: Int): Pair<Double, Double>?

The bound on component, widened to the doubles the ABI carries.

    internal fun boundOf(component: Int): Pair<Double, Double>? = bounds[component]?.let {
        val lower = if (it.start.isFinite()) it.start else -Double.MAX_VALUE
        val upper = if (it.endInclusive.isFinite()) it.endInclusive else Double.MAX_VALUE
        lower to upper
    }

companion object

    companion object {

const val MAX_DIMENSION: Int = 64

A bound on a single block's size. Nothing here needs a hundred-component unknown, and the native side allocates from this number, so it is checked rather than trusted.

fun scalar

A free scalar, optionally box-constrained. atLeast = 0.0 is the usual reason to reach for this: a radius, a scale or a duration that has no meaning below zero.

        fun scalar(
            name: String,
            initial: Double,
            atLeast: Double = Double.NEGATIVE_INFINITY,
            atMost: Double = Double.POSITIVE_INFINITY,
        ): CeresVariableDecl = CeresVariableDecl(
            name,
            listOf(initial),
            bounds = if (atLeast.isFinite() || atMost.isFinite()) mapOf(0 to atLeast..atMost) else emptyMap(),
        )

fun vector(name: String, initial: Vec3): CeresVariableDecl

        fun vector(name: String, initial: Vec3): CeresVariableDecl =
            CeresVariableDecl(name, listOf(initial.x, initial.y, initial.z))

fun vector(name: String, vararg initial: Double): CeresVariableDecl

        fun vector(name: String, vararg initial: Double): CeresVariableDecl =
            CeresVariableDecl(name, initial.toList())

fun unitVector(name: String, initial: Vec3): CeresVariableDecl

A direction on the unit sphere. The seed is normalised for you; it only has to be non-zero.

        fun unitVector(name: String, initial: Vec3): CeresVariableDecl {
            val length = kotlin.math.sqrt(initial.x * initial.x + initial.y * initial.y + initial.z * initial.z)
            require(length > 1e-10) { "unit-vector variable '$name' needs a non-zero seed direction" }
            return CeresVariableDecl(
                name,
                listOf(initial.x / length, initial.y / length, initial.z / length),
                CeresManifold.SPHERE,
            )
        }

fun quaternion(name: String, w: Double, x: Double, y: Double, z: Double): CeresVariableDecl

A rotation, w first. Normalised for you.

        fun quaternion(name: String, w: Double, x: Double, y: Double, z: Double): CeresVariableDecl {
            val length = kotlin.math.sqrt(w * w + x * x + y * y + z * z)
            require(length > 1e-10) { "quaternion variable '$name' needs a non-zero seed" }
            return CeresVariableDecl(
                name,
                listOf(w / length, x / length, y / length, z / length),
                CeresManifold.QUATERNION,
            )
        }

fun rotation(name: String): CeresVariableDecl = quaternion(name, 1.0, 0.0, 0.0, 0.0)

The identity rotation.

data class CeresVariableRef(val index: Int, val declaration: CeresVariableDecl)

A declared variable, as the program refers to it.

Variables are declared at program scope rather than per residual, so several residual blocks can share one — which is the difference between "each face gets its own plane" and "these fifty observations all see the same camera". index is the program-wide declaration order, which is also the order SolveResult.variables reports the solved values in.

init

    init {
        require(index >= 0) { "variable index must be non-negative" }
    }

val name: String get() = declaration.name

The declared name, forwarded from declaration.

val dimension: Int get() = declaration.dimension

The declared component count, forwarded from declaration.

class CeresExprVariable internal constructor

One of an expression's free variables, bound to its position in that residual's variable list.

Distinct from CeresVariableRef: that one knows where the variable sits in the program, this one knows where it sits in this residual's parameter blocks, and only the second is what the VM addresses.

class CeresExprVariable internal constructor(
    internal val block: Int,
    /** What was declared at program scope: the name, the initial value and the manifold. */
    val declaration: CeresVariableDecl,
) {

val name: String get() = declaration.name

The declared name, forwarded from declaration.

val dimension: Int get() = declaration.dimension

How many components get will accept, forwarded from declaration.

operator fun get(component: Int): CeresExpr

Component component of this variable.

    operator fun get(component: Int): CeresExpr {
        require(component in 0 until dimension) {
            "variable '$name' has $dimension component(s); asked for $component"
        }
        return ExprVar(block, component)
    }

val scalar: CeresExpr

The variable itself, when it is a single scalar.

    val scalar: CeresExpr
        get() {
            require(dimension == 1) { "variable '$name' has $dimension components, not one" }
            return get(0)
        }

fun vec(): CeresExprVec

The first three components, as a vector.

    fun vec(): CeresExprVec {
        require(dimension >= 3) { "variable '$name' has $dimension components; vec() needs three" }
        return CeresExprVec(get(0), get(1), get(2))
    }

fun rotate(v: CeresExprVec): CeresExprVec

A four-component variable read as a rotation, applied to v — the quaternion sandwich written out, so no trigonometry and no normalisation appear in the residual.

Pair it with CeresManifold.QUATERNION and Ceres keeps the block on the unit sphere itself; with CeresManifold.NONE the rotation is only a rotation while the quaternion happens to be unit, which is a thing to remember rather than a thing to rely on.

    fun rotate(v: CeresExprVec): CeresExprVec {
        require(dimension == 4) { "variable '$name' has $dimension components; rotate() needs four" }
        val w = get(0)
        val u = CeresExprVec(get(1), get(2), get(3))
        // v + 2 * u x (u x v + w * v), the standard expansion of q v q*.
        val t = (u cross v) + v * w
        return v + (u cross t) * 2.0
    }
}

data class CeresExpressionResidual

A residual block whose value is computed by expressions, one residual each, over points.

Build one with CeresProgramBuilder.expression rather than directly — that hands you the point expressions already bound.

One residual block over all of points, so its Jacobian row is dense in them: prefer several small expressions to one that names half the cloud. A repeated point is rejected here rather than at the native side, because Ceres forbids a residual block naming the same parameter block twice.

data class CeresExpressionResidual(
    /** The cloud points the expressions may read, distinct, and `ALL_POINTS` is not one of them. */
    val points: List<PointId>,
    /** One expression per residual in the block. At least one is required. */
    val expressions: List<CeresExpr>,
    /** Square-rooted into the functor like every other weight here, so the cost scales by `weight`. */
    override val weight: Double = 1.0,
    /**
     * The free variables these expressions read, in the order the expressions' [CeresExprVariable]
     * handles were bound. Declared at program scope, so a variable named here may also be named by
     * other residuals.
     */
    val variables: List<CeresVariableRef> = emptyList(),
) : GeometryConstraint {

init

    init {
        // No points is legal as long as there are free variables. A fit — "what sphere passes through
        // this fixed data" — has no point unknowns at all: the data are constants in the expression
        // and every unknown is a free variable. Requiring a point would have forced such a residual to
        // name one it does not read, which is the zero-Jacobian-column defect this layer refuses
        // elsewhere.
        require(points.isNotEmpty() || variables.isNotEmpty()) {
            "an expression residual needs at least one point or one free variable"
        }
        requirePointList(points, 0, allowAll = false)
        require(points.distinct().size == points.size) {
            "Ceres forbids one residual block naming the same point twice; points were $points"
        }
        require(variables.map { it.index }.distinct().size == variables.size) {
            "Ceres forbids one residual block naming the same parameter block twice; variables were " +
                variables.map { it.name }
        }
        require(expressions.isNotEmpty()) { "an expression residual needs at least one residual" }
        requireWeight(weight)
        // Compiling here is what validates slots, components and variable dimensions, and it is
        // cheap. An expression that cannot compile should fail where it was written.
        compile()
    }

override fun encodeLegacy(out: MutableList<Double>)

Append the Expr opcode and the whole compiled program (point and variable indices, residual slots, slot count, constants and instructions) to the bytecode being built.

    override fun encodeLegacy(out: MutableList<Double>) {
        val compiled = compile()
        out.op(CeresOpcode.Expr)
        out += points.size.toDouble()
        points.forEach { out += it.value.toDouble() }
        out += variables.size.toDouble()
        variables.forEach { out += it.index.toDouble() }
        out += compiled.residualSlots.size.toDouble()
        compiled.residualSlots.forEach { out += it.toDouble() }
        out += compiled.slotCount.toDouble()
        out += compiled.constants.size.toDouble()
        compiled.constants.forEach { out += it }
        out += (compiled.code.size / 4).toDouble()
        compiled.code.forEach { out += it.toDouble() }
        out += weight
    }

fun compile(): CompiledCeresExpression

The compiled program, exposed so a caller can see what their expression became.

Compiles afresh on every call and is called again by encodeLegacy, so this is for inspection — pair it with CompiledCeresExpression.disassemble. Throws if an expression reads a point index beyond the ones bound.

    fun compile(): CompiledCeresExpression =
        CeresExpressionCompiler(points.size, variables.map { it.declaration }).compile(expressions)
}

data class CompiledCeresExpression

The register-machine program an expression compiles to: straight-line code, no branches, no stack.

Every instruction writes a slot it alone owns and reads only slots already written, so evaluation is a single forward pass and shared subexpressions cost nothing extra. That is the whole argument for a register machine over a stack machine here — sharing is structural rather than something the interpreter has to juggle.

data class CompiledCeresExpression(
    /** Four ints per instruction: `op, dst, a, b`. */
    val code: List<Int>,
    /** The constant pool the `CONST` instructions index, deduplicated by the compiler. */
    val constants: List<Double>,
    /** Which slot each residual reads, in the order the expressions were given. */
    val residualSlots: List<Int>,
    /** Slots the evaluator must allocate — one per instruction, since each writes a fresh one. */
    val slotCount: Int,
    /** The free variables this residual reads, in local (parameter-block) order. */
    val variables: List<CeresVariableDecl> = emptyList(),
) {

val instructionCount: Int get() = code.size / 4

code``.size / 4: the cost of one evaluation, and what disassemble prints a line each for.

fun disassemble(): String

A readable disassembly. The point of an expression VM is that you can look at the program.

    fun disassemble(): String = buildString {
        variables.forEachIndexed { at, variable ->
            append("var v%-2d = %s[%d]".format(at, variable.name, variable.dimension))
            if (variable.manifold != CeresManifold.NONE) {
                append(" on ").append(variable.manifold.name.lowercase())
            }
            variable.bounds.keys.sorted().forEach { component ->
                val (lower, upper) = requireNotNull(variable.boundOf(component))
                append(
                    "  [%d] in %s..%s".format(
                        component,
                        if (lower <= -Double.MAX_VALUE) "-inf" else "%.6g".format(lower),
                        if (upper >= Double.MAX_VALUE) "+inf" else "%.6g".format(upper),
                    ),
                )
            }
            appendLine(
                variable.initial.joinToString(", ", prefix = "  seed [", postfix = "]") { "%.6g".format(it) },
            )
        }
        for (i in 0 until instructionCount) {
            val op = ExprOp.entries.first { it.code == code[4 * i] }
            val dst = code[4 * i + 1]
            val a = code[4 * i + 2]
            val b = code[4 * i + 3]
            append("s%-3d = ".format(dst))
            when (op) {
                ExprOp.CONST -> append("${constants[a]}")
                ExprOp.PARAM -> append("p$a.${"xyz"[b]}")
                ExprOp.VAR -> append("v$a[$b]")
                else -> when (op.arity) {
                    1 -> append("${op.name.lowercase()}(s$a)")
                    else -> append("${op.name.lowercase()}(s$a, s$b)")
                }
            }
            appendLine()
        }
        residualSlots.forEachIndexed { at, slot -> appendLine("r$at  = s$slot") }
    }
}

A residual family built out of free variables: the Hesse plane

data class CeresHessePlane(val normal: Vec3, val offset: Double)

A plane as { x : x . normal == offset }, with normal a unit vector.

fun hessePlaneOf(points: List<Vec3>): CeresHessePlane

The total-least-squares plane through points, in Hesse normal form: the centroid, and the smallest eigenvector of the scatter matrix about it.

This is the closed form of the fit, and it is also exactly the seed the PLANAR_HESSE opcode computes in C++ from the pose the problem was compiled with. An expression has to be handed the pose and seed its own plane here — which is not a shortcoming: the caller picks the seed explicitly and can print it, where the opcode's is invisible.

The normal's sign is canonicalised (largest-magnitude component positive) so that the same points always give the same seed. n and -n with d and -d are the same plane and negate every residual that reads them, which leaves the objective and the step identical — but a seed that flips run to run makes results harder to compare than they need to be.

fun hessePlaneOf(points: List<Vec3>): CeresHessePlane {
    require(points.size >= 3) { "a plane fit needs at least three points" }
    val cx = points.sumOf { it.x } / points.size
    val cy = points.sumOf { it.y } / points.size
    val cz = points.sumOf { it.z } / points.size

    val scatter = DoubleArray(9)
    points.forEach { p ->
        val v = doubleArrayOf(p.x - cx, p.y - cy, p.z - cz)
        for (i in 0..2) for (j in 0..2) scatter[3 * i + j] += v[i] * v[j]
    }

    val (values, vectors) = symmetricEigen3(scatter)
    var smallest = 0
    for (i in 1..2) if (values[i] < values[smallest]) smallest = i
    var nx = vectors[smallest]
    var ny = vectors[3 + smallest]
    var nz = vectors[6 + smallest]
    val length = kotlin.math.sqrt(nx * nx + ny * ny + nz * nz)
    if (length < 1e-12) {
        // A perfectly degenerate scatter matrix — every point coincident. Any normal will do, and
        // saying so beats dividing by zero.
        return CeresHessePlane(Vec3(1.0, 0.0, 0.0), cx)
    }
    nx /= length
    ny /= length
    nz /= length
    val dominant = listOf(nx, ny, nz).withIndex().maxBy { kotlin.math.abs(it.value) }
    if (dominant.value < 0.0) {
        nx = -nx
        ny = -ny
        nz = -nz
    }
    return CeresHessePlane(Vec3(nx, ny, nz), nx * cx + ny * cy + nz * cz)
}

internal fun symmetricEigen3(matrix: DoubleArray): Pair<DoubleArray, DoubleArray>

Eigenvalues and eigenvectors of a symmetric 3x3 matrix, by cyclic Jacobi rotations.

Returns the eigenvalues and a row-major matrix whose columns are the corresponding unit eigenvectors. Jacobi rather than anything cleverer because 3x3 is small, the iteration is unconditionally stable on a symmetric matrix, and it agrees with Eigen::SelfAdjointEigenSolver to machine precision — which matters, since the C++ side of this comparison uses exactly that.

internal fun symmetricEigen3(matrix: DoubleArray): Pair<DoubleArray, DoubleArray> {
    require(matrix.size == 9) { "expected a row-major 3x3 matrix" }
    val a = matrix.copyOf()
    val v = doubleArrayOf(1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0)

    repeat(64) {
        // Largest off-diagonal magnitude; stop when it is negligible against the diagonal.
        val offDiagonal = kotlin.math.abs(a[1]) + kotlin.math.abs(a[2]) + kotlin.math.abs(a[5])
        val scale = kotlin.math.abs(a[0]) + kotlin.math.abs(a[4]) + kotlin.math.abs(a[8])
        if (offDiagonal <= 1e-18 * kotlin.math.max(scale, 1e-300)) return values3(a) to v

        for ((p, q) in listOf(0 to 1, 0 to 2, 1 to 2)) {
            val apq = a[3 * p + q]
            if (kotlin.math.abs(apq) < 1e-300) continue
            val tau = (a[3 * q + q] - a[3 * p + p]) / (2.0 * apq)
            val t = (if (tau >= 0.0) 1.0 else -1.0) / (kotlin.math.abs(tau) + kotlin.math.sqrt(1.0 + tau * tau))
            val c = 1.0 / kotlin.math.sqrt(1.0 + t * t)
            val s = t * c

            // A <- J^T A J and V <- V J, written out rather than looped: three-by-three, and the
            // explicit form is the one that is obviously the similarity transform it claims to be.
            val j = DoubleArray(9)
            j[0] = 1.0; j[4] = 1.0; j[8] = 1.0
            j[3 * p + p] = c
            j[3 * p + q] = s
            j[3 * q + p] = -s
            j[3 * q + q] = c
            val jt = DoubleArray(9) { idx -> j[3 * (idx % 3) + idx / 3] }
            val next = multiply3(multiply3(jt, a), j)
            next.copyInto(a)
            multiply3(v, j).copyInto(v)
        }
    }
    return values3(a) to v
}

private fun values3(a: DoubleArray) = doubleArrayOf(a[0], a[4], a[8])

private fun multiply3(a: DoubleArray, b: DoubleArray) = DoubleArray(9) { idx ->
    val row = idx / 3
    val column = idx % 3
    a[3 * row] * b[column] + a[3 * row + 1] * b[3 + column] + a[3 * row + 2] * b[6 + column]
}

fun CeresProgramBuilder.planarHesseExpression

PLANAR_HESSE, written as an expression: this face's corners against a free unit normal and a free scalar offset.

The opcode and this compile to different things — a C++ template against an interpreted register program — but they are the same objective on the same parameterisation, and the recorded evidence for that parameterisation is what makes free variables worth having at all: on kkC the plane as a free variable converges in 12 iterations to tangency 3e-08, where the same objective with the plane eliminated inside the residual (PLANAR_FIT) does not converge.

pose is the flat xyz array the problem will be compiled with — the seed the opcode reads for itself and an expression has to be given.

fun CeresProgramBuilder.planarHesseExpression(
    points: List<Int>,
    pose: DoubleArray,
    weight: Double = 1.0,
    name: String = "plane",
): CeresProgramBuilder {
    require(points.size >= 3) { "a planar face needs at least three points" }
    require(points.distinct().size == points.size) {
        "a Hesse-plane expression names each corner once; face $points repeats one"
    }
    points.forEach {
        require(3 * it + 2 < pose.size) { "point $it is outside the pose of ${pose.size / 3} points" }
    }
    val seed = hessePlaneOf(points.map { Vec3(pose[3 * it], pose[3 * it + 1], pose[3 * it + 2]) })
    val normal = unitVariable("$name.n", seed.normal)
    val offset = variable("$name.d", seed.offset)
    return expression(points, listOf(normal, offset), weight) { corners, (n, d) ->
        val direction = n.vec()
        corners.map { (it dot direction) - d.scalar }
    }
}

Flattens an expression DAG into register-machine instructions.

Nodes are memoised by structural identity, so a subexpression written twice — or reached twice through the DAG — is computed once and read from its slot thereafter. Constants are pooled the same way.

private class CeresExpressionCompiler(
    private val paramCount: Int,
    private val variables: List<CeresVariableDecl>,
) {
    private val code = mutableListOf<Int>()
    private val constants = mutableListOf<Double>()
    private val constantSlots = mutableMapOf<Double, Int>()
    private val slots = mutableMapOf<CeresExpr, Int>()
    private val variablesRead = mutableSetOf<Int>()
    private var nextSlot = 0

    fun compile(expressions: List<CeresExpr>): CompiledCeresExpression {
        val residualSlots = expressions.map { emit(it) }
        // A variable named by a residual but never read gives Ceres a Jacobian column of exact zeros
        // — a rank deficiency the minimiser has to regularise its way around, for a parameter that
        // provably cannot affect the answer. Almost always a typo in the expression.
        variables.indices.forEach { block ->
            require(block in variablesRead) {
                "variable '${variables[block].name}' is named by this expression but never read by it"
            }
        }
        return CompiledCeresExpression(
            code.toList(),
            constants.toList(),
            residualSlots,
            nextSlot,
            variables,
        )
    }

    private fun emit(expr: CeresExpr): Int = slots.getOrPut(expr) {
        when (expr) {
            is ExprConst -> {
                val index = constantSlots.getOrPut(expr.value) {
                    constants += expr.value
                    constants.size - 1
                }
                instruction(ExprOp.CONST, index, 0)
            }

            is ExprParam -> {
                require(expr.block < paramCount) {
                    "expression reads point ${expr.block} but only $paramCount were bound"
                }
                instruction(ExprOp.PARAM, expr.block, expr.component)
            }

            is ExprVar -> {
                require(expr.block < variables.size) {
                    "expression reads free variable ${expr.block} but only ${variables.size} were named"
                }
                val declaration = variables[expr.block]
                require(expr.component < declaration.dimension) {
                    "expression reads component ${expr.component} of '${declaration.name}', " +
                        "which has ${declaration.dimension}"
                }
                variablesRead += expr.block
                instruction(ExprOp.VAR, expr.block, expr.component)
            }

            is ExprNode -> {
                // Operands first: a slot must be written before it is read.
                val a = emit(expr.a)
                val b = expr.b?.let { emit(it) } ?: 0
                instruction(expr.op, a, b)
            }
        }
    }

    private fun instruction(op: ExprOp, a: Int, b: Int): Int {
        val dst = nextSlot++
        code += op.code
        code += dst
        code += a
        code += b
        return dst
    }
}

internal fun CeresExpr.variableBlocksRead(): Set<Int>

The free-variable blocks this reads, as ExprVar.block indices. A caller that builds expressions over a program-wide numbering of scalar variables (the spreadsheet does: variable i is ExprVar(i, 0)) uses this with remapVariableBlocks to hand each residual only the blocks it reads, which is what CeresExpressionResidual and Ceres require.

internal fun CeresExpr.variableBlocksRead(): Set<Int> {
    val out = linkedSetOf<Int>()
    fun walk(e: CeresExpr) {
        when (e) {
            is ExprVar -> out += e.block
            is ExprNode -> { walk(e.a); e.b?.let(::walk) }
            is ExprConst, is ExprParam -> Unit
        }
    }
    walk(this)
    return out
}

internal fun CeresExpr.remapVariableBlocks(replace: (block: Int, component: Int) -> CeresExpr): CeresExpr = when (this) {

this with every ExprVar(block, c) replaced by replace``(block, c); see variableBlocksRead.

internal fun CeresExpr.remapVariableBlocks(replace: (block: Int, component: Int) -> CeresExpr): CeresExpr = when (this) {
    is ExprVar -> replace(block, component)
    is ExprNode -> ExprNode(op, a.remapVariableBlocks(replace), b?.remapVariableBlocks(replace))
    is ExprConst, is ExprParam -> this
}

internal fun ceresVariableComponent(block: Int, component: Int = 0): CeresExpr

A variable reference for remapVariableBlocks: component component of the residual's block-th variable.

internal fun ceresVariableComponent(block: Int, component: Int = 0): CeresExpr = ExprVar(block, component)

internal fun CeresExpr.evaluateOnJvm(variable: (block: Int, component: Int) -> Double): Double = when (this) {

Evaluate this on the JVM, with variable supplying free-variable components. The same arithmetic the native VM runs, for reading a residual back at a solution without a round trip. Point parameters are not supported (a sheet has none).

internal fun CeresExpr.evaluateOnJvm(variable: (block: Int, component: Int) -> Double): Double = when (this) {
    is ExprConst -> value
    is ExprVar -> variable(block, component)
    is ExprParam -> error("evaluateOnJvm does not bind point parameters")
    is ExprNode -> {
        val x = a.evaluateOnJvm(variable)
        val y = b?.evaluateOnJvm(variable) ?: 0.0
        when (op) {
            ExprOp.ADD -> x + y
            ExprOp.SUB -> x - y
            ExprOp.MUL -> x * y
            ExprOp.DIV -> x / y
            ExprOp.NEG -> -x
            ExprOp.SQRT -> kotlin.math.sqrt(x)
            ExprOp.ABS -> kotlin.math.abs(x)
            ExprOp.SIN -> kotlin.math.sin(x)
            ExprOp.COS -> kotlin.math.cos(x)
            ExprOp.TAN -> kotlin.math.tan(x)
            ExprOp.EXP -> kotlin.math.exp(x)
            ExprOp.LOG -> kotlin.math.ln(x)
            ExprOp.POW -> Math.pow(x, y)
            ExprOp.MIN -> kotlin.math.min(x, y)
            ExprOp.MAX -> kotlin.math.max(x, y)
            ExprOp.ATAN2 -> kotlin.math.atan2(x, y)
            ExprOp.SQUARE -> x * x
            ExprOp.RECIP -> 1.0 / x
            ExprOp.CONST, ExprOp.PARAM, ExprOp.VAR -> error("leaf opcode $op in a node")
        }
    }
}