Unit-Aware Type Annotations

Unit-Aware Type Annotations#

brainunit provides a comprehensive type annotation system that lets you express physical-unit constraints directly in Python type hints. This enables:

Self-documenting code — function signatures show expected units
Runtime validation — catch unit mismatches at the function boundary
isinstance support — check if a quantity has the right dimension at runtime
IDE support — type checkers see Annotated[Quantity, ...] metadata

The system is built on PEP 593 (typing.Annotated) and follows the same patterns as astropy’s unit-aware type hints.

import brainunit as u
import jax.numpy as jnp

1. The `Quantity[...]` Subscript Syntax#

The simplest way to annotate unit-aware code is with Quantity[unit]. This produces a typing.Annotated type that records the unit constraint as metadata.

Annotating with a specific Unit#

Pass any Unit object to Quantity[...]:

# Annotate with a specific unit
u.Quantity[u.meter]

Quantity[m]

# Compound units work too
u.Quantity[u.meter / u.second]

Quantity[m / s]

u.Quantity[u.kilogram * u.meter / u.second ** 2]

Quantity[N]

Annotating with a physical type (dimension name)#

Instead of a specific unit, you can annotate with a physical type string. This is more flexible — it accepts any unit of that dimension:

# Annotate with a physical type string
u.Quantity["length"]

Quantity['length']

u.Quantity["speed"]

Quantity['speed']

u.Quantity["voltage"]

Quantity['voltage']

Available physical types#

The following physical type strings are recognized:

Base dimensions	Derived dimensions	Compound dimensions
`"length"`	`"frequency"`	`"speed"` / `"velocity"`
`"mass"`	`"force"`	`"acceleration"`
`"time"`	`"energy"`	`"area"`
`"current"` / `"electric current"`	`"power"`	`"volume"`
`"temperature"`	`"pressure"`	`"density"`
`"substance"` / `"amount of substance"`	`"charge"`	`"momentum"`
`"luminosity"` / `"luminous intensity"`	`"voltage"` / `"electric potential"`	`"angular velocity"`
`"dimensionless"`	`"resistance"`	`"torque"`
	`"capacitance"`
	`"conductance"`
	`"magnetic flux"`
	`"magnetic field"`
	`"inductance"`

2. isinstance Support#

Quantity[...] and PhysicalType(...) both support Python’s isinstance for runtime dimension checking. This lets you check whether a quantity has the right physical type:

# Check with Quantity[unit] — dimension-based matching
x = 2.0 * u.kmeter

print(f"isinstance(x, Quantity[u.meter])  = {isinstance(x, u.Quantity[u.meter])}")   # True (km has length dim)
print(f"isinstance(x, Quantity[u.second]) = {isinstance(x, u.Quantity[u.second])}")  # False (not time)
print(f"isinstance(x, Quantity['length']) = {isinstance(x, u.Quantity['length'])}")   # True
print(f"isinstance(x, Quantity['mass'])   = {isinstance(x, u.Quantity['mass'])}")     # False

isinstance(x, Quantity[u.meter])  = True
isinstance(x, Quantity[u.second]) = False
isinstance(x, Quantity['length']) = True
isinstance(x, Quantity['mass'])   = False

from brainunit.typing import PhysicalType

# PhysicalType also works directly with isinstance
v = 10.0 * u.meter / u.second

print(f"isinstance(v, PhysicalType('speed'))  = {isinstance(v, PhysicalType('speed'))}")   # True
print(f"isinstance(v, PhysicalType('length')) = {isinstance(v, PhysicalType('length'))}")  # False

# Non-Quantity values always return False
print(f"isinstance(42, PhysicalType('length')) = {isinstance(42, PhysicalType('length'))}")  # False

isinstance(v, PhysicalType('speed'))  = True
isinstance(v, PhysicalType('length')) = False
isinstance(42, PhysicalType('length')) = False

3. Using Annotations in Function Signatures#

Annotated functions are self-documenting — the signature tells you exactly what units are expected:

def kinetic_energy(
    m: u.Quantity[u.kilogram],
    v: u.Quantity[u.meter / u.second],
) -> u.Quantity[u.joule]:
    """Calculate kinetic energy: KE = 0.5 * m * v^2."""
    return 0.5 * m * v ** 2


result = kinetic_energy(10.0 * u.kilogram, 3.0 * u.meter / u.second)
print(f"KE = {result}")
print(f"Dimension matches joule? {result.dim == u.joule.dim}")

KE = 45. J
Dimension matches joule? True

With physical type strings (more flexible — allows any unit of that dimension):

def travel_time(
    distance: u.Quantity["length"],
    speed: u.Quantity["speed"],
) -> u.Quantity["time"]:
    """Calculate travel time = distance / speed."""
    return distance / speed


# Works with any length/speed units
t = travel_time(100.0 * u.kmeter, 50.0 * u.kmeter / u.second)
print(f"Time = {t}")

Time = 2. s

4. Pre-built Type Aliases#

The brainunit.typing module provides ready-made aliases for common physical types. These also support isinstance checks.

from brainunit.typing import LENGTH, MASS, TIME, SPEED, ENERGY, VOLTAGE, FORCE


def displacement(v: SPEED, t: TIME) -> LENGTH:
    """Calculate displacement = velocity * time."""
    return v * t


d = displacement(10.0 * u.meter / u.second, 5.0 * u.second)
print(f"Displacement = {d}")

Displacement = 50. m

def newton_second_law(m: MASS, a: u.Quantity["acceleration"]) -> FORCE:
    """F = m * a."""
    return m * a


force = newton_second_law(5.0 * u.kilogram, 9.8 * u.meter / u.second ** 2)
print(f"Force = {force}")
print(f"Dimension matches newton? {force.dim == u.newton.dim}")

Force = 49. N
Dimension matches newton? True

Full list of pre-built aliases#

Alias	Physical Type
`HAS_UNIT`	Any `Quantity` (no constraint)
`DIMENSIONLESS_TYPE`	Dimensionless
`LENGTH`	Length
`MASS`	Mass
`TIME`	Time
`CURRENT`	Electric current
`TEMPERATURE`	Temperature
`SUBSTANCE`	Amount of substance
`LUMINOSITY`	Luminous intensity
`FREQUENCY`	Frequency
`FORCE`	Force
`ENERGY`	Energy
`POWER`	Power
`PRESSURE`	Pressure
`CHARGE`	Electric charge
`VOLTAGE`	Voltage / Electric potential
`RESISTANCE`	Resistance
`CAPACITANCE`	Capacitance
`CONDUCTANCE`	Conductance
`MAGNETIC_FLUX`	Magnetic flux
`MAGNETIC_FIELD`	Magnetic field
`INDUCTANCE`	Inductance
`SPEED`	Speed
`ACCELERATION`	Acceleration
`AREA`	Area
`VOLUME`	Volume
`DENSITY`	Density

5. Core Type Aliases#

For general-purpose typing, brainunit.typing also provides:

QuantityLike — any value that can be converted to a Quantity (numbers, arrays, existing Quantities)
UnitLike — any value interpretable as a Unit (Unit objects, strings like "mV", or None)
DimensionLike — any value interpretable as a Dimension (Dimension objects or strings)

from brainunit.typing import QuantityLike, UnitLike


def make_quantity(value: QuantityLike, unit: UnitLike = None) -> u.Quantity:
    """Create a Quantity from flexible input types."""
    if unit is not None:
        return u.Quantity(value, unit=unit)
    return u.Quantity(value)


# All of these work
print(make_quantity(3.14, u.meter))
print(make_quantity(jnp.array([1.0, 2.0]), "mV"))
print(make_quantity(5.0 * u.second))

3.14 m
[1. 2.] mV
5. s

6. Runtime Validation with `@validate_units`#

The @validate_units decorator checks that Quantity arguments match their annotated units/dimensions at call time. It inspects Quantity[...] metadata and checks every annotated argument on each call.

from brainunit.typing import validate_units


@validate_units
def ohms_law(
    V: u.Quantity[u.volt],
    R: u.Quantity[u.ohm],
) -> u.Quantity[u.amp]:
    """Calculate current: I = V / R."""
    return V / R


# Correct units — works fine
current = ohms_law(12.0 * u.volt, 4.0 * u.ohm)
print(f"Current = {current}")

Current = 3. A

# Compatible units also work (millivolt is still a voltage)
current2 = ohms_law(500.0 * u.mvolt, 100.0 * u.ohm)
print(f"Current = {current2}")

Current = 5. mA

# Wrong units — raises an error!
try:
    ohms_law(12.0 * u.meter, 4.0 * u.ohm)  # meter is not a voltage
except Exception as e:
    print(f"Caught: {type(e).__name__}: {e}")

Caught: DimensionMismatchError: Argument 'V' of 'ohms_law' expected dimension compatible with V, got m.

# Non-Quantity argument — also caught
try:
    ohms_law(12.0, 4.0 * u.ohm)  # plain float is not a Quantity
except TypeError as e:
    print(f"Caught: {e}")

Caught: Argument 'V' of 'ohms_law' expected a Quantity, got float.

Strict mode#

By default, @validate_units only checks dimensional compatibility. Use strict=True to require exact unit match (same scale):

@validate_units(strict=True)
def precise_voltage(V: u.Quantity[u.volt]) -> u.Quantity[u.volt]:
    """Only accepts values in volts, not millivolts or kilovolts."""
    return V


# Exact unit match — OK
print(precise_voltage(5.0 * u.volt))

# Different scale — rejected in strict mode
try:
    precise_voltage(500.0 * u.mvolt)
except Exception as e:
    print(f"Strict mode caught: {type(e).__name__}: {e}")

5. V
Strict mode caught: UnitMismatchError: Argument 'V' of 'precise_voltage' expected unit V, got mV.

Physical type validation#

@validate_units also works with physical type string annotations:

@validate_units
def momentum(
    m: u.Quantity["mass"],
    v: u.Quantity["speed"],
) -> u.Quantity["momentum"]:
    """Calculate momentum: p = m * v."""
    return m * v


p = momentum(2.0 * u.kilogram, 5.0 * u.meter / u.second)
print(f"Momentum = {p}")

# Wrong dimension
try:
    momentum(2.0 * u.kilogram, 5.0 * u.meter)  # length, not speed
except Exception as e:
    print(f"Caught: {type(e).__name__}: {e}")

Momentum = 10. kg * m / s
Caught: DimensionMismatchError: Argument 'v' of 'momentum' expected physical type 'speed', got dimension m.

7. The `PhysicalType` Class#

PhysicalType is the annotation marker that records dimension constraints. You usually don’t need to create it directly — Quantity["speed"] does it for you — but it’s available for advanced use:

from brainunit.typing import PhysicalType

pt = PhysicalType("speed")
print(f"Physical type: {pt}")
print(f"Dimension: {pt.dimension}")
print(f"Same as meter/second dim? {pt.dimension == (u.meter / u.second).dim}")

Physical type: PhysicalType('speed')
Dimension: m s^-1
Same as meter/second dim? True

# PhysicalType objects are hashable and comparable
print(PhysicalType("length") == PhysicalType("length"))  # True
print(PhysicalType("length") == PhysicalType("mass"))    # False

# Case-insensitive
print(PhysicalType("Length") == PhysicalType("length"))   # True

True
False
True

8. Composing with `typing.Optional` and `typing.Union`#

Unit-annotated Quantities work with standard typing constructs:

from typing import Optional


def apply_force(
    mass: u.Quantity[u.kilogram],
    acceleration: u.Quantity[u.meter / u.second ** 2],
    friction: Optional[u.Quantity[u.newton]] = None,
) -> u.Quantity[u.newton]:
    """Calculate net force, optionally subtracting friction."""
    force = mass * acceleration
    if friction is not None:
        force = force - friction
    return force


# Without friction
print(apply_force(5.0 * u.kilogram, 2.0 * u.meter / u.second ** 2))

# With friction
print(apply_force(5.0 * u.kilogram, 2.0 * u.meter / u.second ** 2, 3.0 * u.newton))

10. N
7. N

9. Introspecting Annotations#

Unit annotations can be inspected programmatically via the _metadata attribute:

# Introspect the metadata stored in Quantity[...] types
length_type = u.Quantity[u.meter]
speed_type = u.Quantity["speed"]

print(f"Quantity[u.meter] metadata: {length_type._metadata}")
print(f"Quantity['speed'] metadata: {speed_type._metadata}")

Quantity[u.meter] metadata: m
Quantity['speed'] metadata: PhysicalType('speed')

10. Comparison: `Quantity[unit]` vs `@check_dims` vs `@check_units`#

brainunit provides three approaches to unit validation. Here’s when to use each:

Feature	`Quantity[unit]` + `@validate_units`	`@check_dims`	`@check_units`
Style	PEP 593 type hints	Decorator kwargs	Decorator kwargs
isinstance	Yes	No	No
Validates	Input arguments	Input + return	Input + return
Granularity	Unit or dimension	Dimension only	Unit (exact)
Best for	New code, libraries	Legacy code, quick checks	Strict unit matching

# Approach 1: Type annotations (recommended for new code)
@validate_units
def approach1(V: u.Quantity[u.volt], I: u.Quantity[u.amp]) -> u.Quantity[u.watt]:
    return V * I


# Approach 2: check_dims decorator
@u.check_dims(V=u.volt.dim, I=u.amp.dim, result=u.watt.dim)
def approach2(V, I):
    return V * I


# Approach 3: check_units decorator
@u.check_units(V=u.volt, I=u.amp, result=u.watt)
def approach3(V, I):
    return V * I


# All three produce the same result
V, I = 12.0 * u.volt, 2.0 * u.amp
print(f"Approach 1: {approach1(V, I)}")
print(f"Approach 2: {approach2(V, I)}")
print(f"Approach 3: {approach3(V, I)}")

Approach 1: 24. W
Approach 2: 24. W
Approach 3: 24. W

11. Best Practices#

Use physical type strings for public APIs#

Physical type strings ("length", "speed") are more flexible than specific units. They accept any unit of the right dimension:

# Preferred: accepts meter, kilometer, mile, etc.
def distance(x: Quantity["length"]) -> Quantity["length"]: ...

# More restrictive: conceptually expects meters
def distance(x: Quantity[u.meter]) -> Quantity[u.meter]: ...

Use specific units for internal functions#

When a function’s implementation depends on a specific unit (e.g., for numerical constants), annotate with that unit:

@validate_units(strict=True)
def _internal_calc(voltage_mV: Quantity[u.mvolt]) -> Quantity[u.mvolt]:
    # Implementation uses millivolt-specific constants
    return voltage_mV * 0.5

Use `isinstance` for runtime branching#

def process(quantity):
    if isinstance(quantity, Quantity["length"]):
        return quantity.in_unit(u.meter)
    elif isinstance(quantity, Quantity["time"]):
        return quantity.in_unit(u.second)
    else:
        raise TypeError(f"Unexpected dimension: {quantity.dim}")

Combine with `@validate_units` for safety-critical code#

@validate_units
def drug_dosage(
    concentration: Quantity["density"],
    volume: Quantity["volume"],
    patient_mass: Quantity["mass"],
) -> Quantity["dimensionless"]:
    return (concentration * volume) / patient_mass

Use `QuantityLike` and `UnitLike` for flexible inputs#

from brainunit.typing import QuantityLike, UnitLike

def convert(value: QuantityLike, target: UnitLike) -> Quantity:
    q = Quantity(value)
    if target is not None:
        return q.in_unit(target)
    return q