axiom

A model of your own · step 2 of 5

The mean, as an expression

How do you write a formula axiom can differentiate and check?

Not as a Python function. A mean is built out of typed nodes — Data for a column, Param for something to be fitted, Const for a literal, and Add, Mul, Div, Pow, Apply to combine them.

It is more verbose than lambda x, a, b, k: ... and it buys three things a lambda cannot: the gradient every backend needs, without you writing it; a content hash, so the model that produced a number is identifiable later; and the dimension check in the next step.

step 2
from axiom.core import (
    Add, Apply, Const, D, Data, Div, Mul, Param, Prior,
    dimension, dimensionless,
)

spend = Data(name="spend", dimension=D.currency)

alpha = Param(name="alpha", dimension=D.outcome,
              prior=Prior(family="normal", hyper={"mu": 40.0, "sigma": 10.0}))
beta = Param(name="beta", dimension=D.outcome,
             prior=Prior(family="normal", hyper={"mu": 0.0, "sigma": 10.0}))
k = Param(name="k", dimension=D.currency,
          prior=Prior(family="lognormal", hyper={"mu": 5.5, "sigma": 1.0}))

one = Const(value=1.0, dimension=dimensionless())
minus = Const(value=-1.0, dimension=dimensionless())

# 1 - exp(-spend / k)   -- dimensionless, because spend/k is
saturation = Add(terms=(
    one,
    Mul(factors=(minus, Apply(fn="exp", arg=Mul(factors=(
        minus, Div(numerator=spend, denominator=k),
    ))))),
))
# alpha + beta * saturation
mean = Add(terms=(alpha, Mul(factors=(beta, saturation))))

print(f"the mean's dimension is  {dimension(mean)}")
print(f"the saturation term is   {dimension(saturation)} (as it must be)")
ratio = Div(numerator=spend, denominator=k)
print(f"spend / k is             {dimension(ratio)}")
what it printed
the mean's dimension is  Y
the saturation term is   1 (as it must be)
spend / k is             1

Every line above was captured by running this code at build time. The steps share one namespace, so the code on this page is the code you would type after the steps before it.