Support units of measurement in PyMC

[drbenvincent]

Many statistical models (especially in scientific, engineering, and health applications) are built around physical quantities that carry units (e.g., meters, seconds, kilograms). Currently, PyMC treats all variables as unitless, which may lead to misinterpretations or errors when combining data with different units or when interpreting model parameters.

Adding optional support for units would:

• Improve model readability and transparency
• Enable automated unit-checking to catch errors (e.g., adding kg to m)
• Facilitate interpretation of parameter estimates and priors

This could be implemented via integration with existing Python libraries like pint

In general, the idea would be to optionally specify the units. Initially it might be required to specify the units of all data and parameters and pymc could help with unit checking to avoid errors.

This would be useful as we could check unit consistency, but also make errors less likely (e.g. expressing slope priors in the wrong units).

It could be very fun to explore unit inference. For example, if you specify the units of data but not the parameters, if we are regressing weight ~ age where weight is in kg and age is in years, the model could infer that the slope is in units of kg/year and the intercept is in units of kg.

I'll leave it there - this is an initial proposal which is intended to spark discussion.

[ricardoV94]

Code examples?

[drbenvincent]

Something vaguely like this?

import pint
import pymc as pm

ureg = pint.UnitRegistry()

age = pm.Data("age", [10, 20, 30, 40] * ureg.year)
weight = pm.Data("weight", [50, 60, 65, 67] * ureg.kg)

with pm.Model() as model:
    intercept = pm.Normal("intercept", mu=1.2 * ureg.kg, sigma=0.2 * ureg.kg)
    beta = pm.Normal("beta", mu=0.5 * ureg.kg / ureg.year, sigma=0.1 * ureg.kg / ureg.year) 
    mu = pm.Deterministic("mu", intercept + beta * age)  # PyMC infers mu should be in kg
    pm.Normal("obs", mu=mu, sigma=0.5 * ureg.kg, observed=weight)

• PyMC does unit checks and throws errors if there are incompatibilities
• PyMC optionally infers units of any nodes where units are not provided, or throws an error if it is not possible, asking for units of more variables.
• In theory I guess you could allow the intercept's mu to be provided in kg and the sigma in another weight unit and auto-convert, but perhaps emit a warning.
• Units would be incorporated in the idata