{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# Unit-Aware Type Annotations\n\nbrainunit provides a comprehensive type annotation system that lets you express physical-unit constraints directly in Python type hints. This enables:\n\n- **Self-documenting code** — function signatures show expected units\n- **Runtime validation** — catch unit mismatches at the function boundary\n- **isinstance support** — check if a quantity has the right dimension at runtime\n- **IDE support** — type checkers see `Annotated[Quantity, ...]` metadata\n\nThe system is built on [PEP 593](https://peps.python.org/pep-0593/) (`typing.Annotated`) and follows the same patterns as [astropy's unit-aware type hints](https://docs.astropy.org/en/stable/units/type_hints.html)."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "import brainunit as u\n",
    "import jax.numpy as jnp"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. The `Quantity[...]` Subscript Syntax\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Annotating with a specific Unit\n",
    "\n",
    "Pass any `Unit` object to `Quantity[...]`:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Annotate with a specific unit\n",
    "u.Quantity[u.meter]"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Compound units work too\n",
    "u.Quantity[u.meter / u.second]"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "u.Quantity[u.kilogram * u.meter / u.second ** 2]"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Annotating with a physical type (dimension name)\n",
    "\n",
    "Instead of a specific unit, you can annotate with a physical type string. This is more flexible — it accepts any unit of that dimension:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Annotate with a physical type string\n",
    "u.Quantity[\"length\"]"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "u.Quantity[\"speed\"]"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "u.Quantity[\"voltage\"]"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Available physical types\n",
    "\n",
    "The following physical type strings are recognized:\n",
    "\n",
    "| **Base dimensions** | **Derived dimensions** | **Compound dimensions** |\n",
    "|---|---|---|\n",
    "| `\"length\"` | `\"frequency\"` | `\"speed\"` / `\"velocity\"` |\n",
    "| `\"mass\"` | `\"force\"` | `\"acceleration\"` |\n",
    "| `\"time\"` | `\"energy\"` | `\"area\"` |\n",
    "| `\"current\"` / `\"electric current\"` | `\"power\"` | `\"volume\"` |\n",
    "| `\"temperature\"` | `\"pressure\"` | `\"density\"` |\n",
    "| `\"substance\"` / `\"amount of substance\"` | `\"charge\"` | `\"momentum\"` |\n",
    "| `\"luminosity\"` / `\"luminous intensity\"` | `\"voltage\"` / `\"electric potential\"` | `\"angular velocity\"` |\n",
    "| `\"dimensionless\"` | `\"resistance\"` | `\"torque\"` |\n",
    "| | `\"capacitance\"` | |\n",
    "| | `\"conductance\"` | |\n",
    "| | `\"magnetic flux\"` | |\n",
    "| | `\"magnetic field\"` | |\n",
    "| | `\"inductance\"` | |"
   ]
  },
  {
   "cell_type": "markdown",
   "source": "## 2. isinstance Support\n\n`Quantity[...]` and `PhysicalType(...)` both support Python's `isinstance` for runtime dimension checking. This lets you check whether a quantity has the right physical type:",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "source": "# Check with Quantity[unit] — dimension-based matching\nx = 2.0 * u.kmeter\n\nprint(f\"isinstance(x, Quantity[u.meter])  = {isinstance(x, u.Quantity[u.meter])}\")   # True (km has length dim)\nprint(f\"isinstance(x, Quantity[u.second]) = {isinstance(x, u.Quantity[u.second])}\")  # False (not time)\nprint(f\"isinstance(x, Quantity['length']) = {isinstance(x, u.Quantity['length'])}\")   # True\nprint(f\"isinstance(x, Quantity['mass'])   = {isinstance(x, u.Quantity['mass'])}\")     # False",
   "metadata": {},
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "source": "from brainunit.typing import PhysicalType\n\n# PhysicalType also works directly with isinstance\nv = 10.0 * u.meter / u.second\n\nprint(f\"isinstance(v, PhysicalType('speed'))  = {isinstance(v, PhysicalType('speed'))}\")   # True\nprint(f\"isinstance(v, PhysicalType('length')) = {isinstance(v, PhysicalType('length'))}\")  # False\n\n# Non-Quantity values always return False\nprint(f\"isinstance(42, PhysicalType('length')) = {isinstance(42, PhysicalType('length'))}\")  # False",
   "metadata": {},
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 3. Using Annotations in Function Signatures\n\nAnnotated functions are self-documenting — the signature tells you exactly what units are expected:"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "def kinetic_energy(\n",
    "    m: u.Quantity[u.kilogram],\n",
    "    v: u.Quantity[u.meter / u.second],\n",
    ") -> u.Quantity[u.joule]:\n",
    "    \"\"\"Calculate kinetic energy: KE = 0.5 * m * v^2.\"\"\"\n",
    "    return 0.5 * m * v ** 2\n",
    "\n",
    "\n",
    "result = kinetic_energy(10.0 * u.kilogram, 3.0 * u.meter / u.second)\n",
    "print(f\"KE = {result}\")\n",
    "print(f\"Dimension matches joule? {result.dim == u.joule.dim}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With physical type strings (more flexible — allows any unit of that dimension):"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "def travel_time(\n",
    "    distance: u.Quantity[\"length\"],\n",
    "    speed: u.Quantity[\"speed\"],\n",
    ") -> u.Quantity[\"time\"]:\n",
    "    \"\"\"Calculate travel time = distance / speed.\"\"\"\n",
    "    return distance / speed\n",
    "\n",
    "\n",
    "# Works with any length/speed units\n",
    "t = travel_time(100.0 * u.kmeter, 50.0 * u.kmeter / u.second)\n",
    "print(f\"Time = {t}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 4. Pre-built Type Aliases\n\nThe `brainunit.typing` module provides ready-made aliases for common physical types. These also support `isinstance` checks."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from brainunit.typing import LENGTH, MASS, TIME, SPEED, ENERGY, VOLTAGE, FORCE\n",
    "\n",
    "\n",
    "def displacement(v: SPEED, t: TIME) -> LENGTH:\n",
    "    \"\"\"Calculate displacement = velocity * time.\"\"\"\n",
    "    return v * t\n",
    "\n",
    "\n",
    "d = displacement(10.0 * u.meter / u.second, 5.0 * u.second)\n",
    "print(f\"Displacement = {d}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "def newton_second_law(m: MASS, a: u.Quantity[\"acceleration\"]) -> FORCE:\n",
    "    \"\"\"F = m * a.\"\"\"\n",
    "    return m * a\n",
    "\n",
    "\n",
    "force = newton_second_law(5.0 * u.kilogram, 9.8 * u.meter / u.second ** 2)\n",
    "print(f\"Force = {force}\")\n",
    "print(f\"Dimension matches newton? {force.dim == u.newton.dim}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Full list of pre-built aliases\n",
    "\n",
    "| Alias | Physical Type |\n",
    "|---|---|\n",
    "| `HAS_UNIT` | Any `Quantity` (no constraint) |\n",
    "| `DIMENSIONLESS_TYPE` | Dimensionless |\n",
    "| `LENGTH` | Length |\n",
    "| `MASS` | Mass |\n",
    "| `TIME` | Time |\n",
    "| `CURRENT` | Electric current |\n",
    "| `TEMPERATURE` | Temperature |\n",
    "| `SUBSTANCE` | Amount of substance |\n",
    "| `LUMINOSITY` | Luminous intensity |\n",
    "| `FREQUENCY` | Frequency |\n",
    "| `FORCE` | Force |\n",
    "| `ENERGY` | Energy |\n",
    "| `POWER` | Power |\n",
    "| `PRESSURE` | Pressure |\n",
    "| `CHARGE` | Electric charge |\n",
    "| `VOLTAGE` | Voltage / Electric potential |\n",
    "| `RESISTANCE` | Resistance |\n",
    "| `CAPACITANCE` | Capacitance |\n",
    "| `CONDUCTANCE` | Conductance |\n",
    "| `MAGNETIC_FLUX` | Magnetic flux |\n",
    "| `MAGNETIC_FIELD` | Magnetic field |\n",
    "| `INDUCTANCE` | Inductance |\n",
    "| `SPEED` | Speed |\n",
    "| `ACCELERATION` | Acceleration |\n",
    "| `AREA` | Area |\n",
    "| `VOLUME` | Volume |\n",
    "| `DENSITY` | Density |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 5. Core Type Aliases\n\nFor general-purpose typing, `brainunit.typing` also provides:\n\n- **`QuantityLike`** — any value that can be converted to a `Quantity` (numbers, arrays, existing Quantities)\n- **`UnitLike`** — any value interpretable as a `Unit` (Unit objects, strings like `\"mV\"`, or `None`)\n- **`DimensionLike`** — any value interpretable as a `Dimension` (Dimension objects or strings)"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from brainunit.typing import QuantityLike, UnitLike\n",
    "\n",
    "\n",
    "def make_quantity(value: QuantityLike, unit: UnitLike = None) -> u.Quantity:\n",
    "    \"\"\"Create a Quantity from flexible input types.\"\"\"\n",
    "    if unit is not None:\n",
    "        return u.Quantity(value, unit=unit)\n",
    "    return u.Quantity(value)\n",
    "\n",
    "\n",
    "# All of these work\n",
    "print(make_quantity(3.14, u.meter))\n",
    "print(make_quantity(jnp.array([1.0, 2.0]), \"mV\"))\n",
    "print(make_quantity(5.0 * u.second))"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 6. Runtime Validation with `@validate_units`\n\nThe `@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."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from brainunit.typing import validate_units\n",
    "\n",
    "\n",
    "@validate_units\n",
    "def ohms_law(\n",
    "    V: u.Quantity[u.volt],\n",
    "    R: u.Quantity[u.ohm],\n",
    ") -> u.Quantity[u.amp]:\n",
    "    \"\"\"Calculate current: I = V / R.\"\"\"\n",
    "    return V / R\n",
    "\n",
    "\n",
    "# Correct units — works fine\n",
    "current = ohms_law(12.0 * u.volt, 4.0 * u.ohm)\n",
    "print(f\"Current = {current}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Compatible units also work (millivolt is still a voltage)\n",
    "current2 = ohms_law(500.0 * u.mvolt, 100.0 * u.ohm)\n",
    "print(f\"Current = {current2}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Wrong units — raises an error!\n",
    "try:\n",
    "    ohms_law(12.0 * u.meter, 4.0 * u.ohm)  # meter is not a voltage\n",
    "except Exception as e:\n",
    "    print(f\"Caught: {type(e).__name__}: {e}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Non-Quantity argument — also caught\n",
    "try:\n",
    "    ohms_law(12.0, 4.0 * u.ohm)  # plain float is not a Quantity\n",
    "except TypeError as e:\n",
    "    print(f\"Caught: {e}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Strict mode\n",
    "\n",
    "By default, `@validate_units` only checks dimensional compatibility. Use `strict=True` to require exact unit match (same scale):"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "@validate_units(strict=True)\n",
    "def precise_voltage(V: u.Quantity[u.volt]) -> u.Quantity[u.volt]:\n",
    "    \"\"\"Only accepts values in volts, not millivolts or kilovolts.\"\"\"\n",
    "    return V\n",
    "\n",
    "\n",
    "# Exact unit match — OK\n",
    "print(precise_voltage(5.0 * u.volt))\n",
    "\n",
    "# Different scale — rejected in strict mode\n",
    "try:\n",
    "    precise_voltage(500.0 * u.mvolt)\n",
    "except Exception as e:\n",
    "    print(f\"Strict mode caught: {type(e).__name__}: {e}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Physical type validation\n",
    "\n",
    "`@validate_units` also works with physical type string annotations:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "@validate_units\n",
    "def momentum(\n",
    "    m: u.Quantity[\"mass\"],\n",
    "    v: u.Quantity[\"speed\"],\n",
    ") -> u.Quantity[\"momentum\"]:\n",
    "    \"\"\"Calculate momentum: p = m * v.\"\"\"\n",
    "    return m * v\n",
    "\n",
    "\n",
    "p = momentum(2.0 * u.kilogram, 5.0 * u.meter / u.second)\n",
    "print(f\"Momentum = {p}\")\n",
    "\n",
    "# Wrong dimension\n",
    "try:\n",
    "    momentum(2.0 * u.kilogram, 5.0 * u.meter)  # length, not speed\n",
    "except Exception as e:\n",
    "    print(f\"Caught: {type(e).__name__}: {e}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 7. The `PhysicalType` Class\n\n`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:"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from brainunit.typing import PhysicalType\n",
    "\n",
    "pt = PhysicalType(\"speed\")\n",
    "print(f\"Physical type: {pt}\")\n",
    "print(f\"Dimension: {pt.dimension}\")\n",
    "print(f\"Same as meter/second dim? {pt.dimension == (u.meter / u.second).dim}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# PhysicalType objects are hashable and comparable\n",
    "print(PhysicalType(\"length\") == PhysicalType(\"length\"))  # True\n",
    "print(PhysicalType(\"length\") == PhysicalType(\"mass\"))    # False\n",
    "\n",
    "# Case-insensitive\n",
    "print(PhysicalType(\"Length\") == PhysicalType(\"length\"))   # True"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 8. Composing with `typing.Optional` and `typing.Union`\n\nUnit-annotated Quantities work with standard `typing` constructs:"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from typing import Optional\n",
    "\n",
    "\n",
    "def apply_force(\n",
    "    mass: u.Quantity[u.kilogram],\n",
    "    acceleration: u.Quantity[u.meter / u.second ** 2],\n",
    "    friction: Optional[u.Quantity[u.newton]] = None,\n",
    ") -> u.Quantity[u.newton]:\n",
    "    \"\"\"Calculate net force, optionally subtracting friction.\"\"\"\n",
    "    force = mass * acceleration\n",
    "    if friction is not None:\n",
    "        force = force - friction\n",
    "    return force\n",
    "\n",
    "\n",
    "# Without friction\n",
    "print(apply_force(5.0 * u.kilogram, 2.0 * u.meter / u.second ** 2))\n",
    "\n",
    "# With friction\n",
    "print(apply_force(5.0 * u.kilogram, 2.0 * u.meter / u.second ** 2, 3.0 * u.newton))"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 9. Introspecting Annotations\n\nUnit annotations can be inspected programmatically via the `_metadata` attribute:"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": "# Introspect the metadata stored in Quantity[...] types\nlength_type = u.Quantity[u.meter]\nspeed_type = u.Quantity[\"speed\"]\n\nprint(f\"Quantity[u.meter] metadata: {length_type._metadata}\")\nprint(f\"Quantity['speed'] metadata: {speed_type._metadata}\")",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 10. Comparison: `Quantity[unit]` vs `@check_dims` vs `@check_units`\n\nbrainunit provides three approaches to unit validation. Here's when to use each:\n\n| Feature | `Quantity[unit]` + `@validate_units` | `@check_dims` | `@check_units` |\n|---|---|---|---|\n| **Style** | PEP 593 type hints | Decorator kwargs | Decorator kwargs |\n| **isinstance** | Yes | No | No |\n| **Validates** | Input arguments | Input + return | Input + return |\n| **Granularity** | Unit or dimension | Dimension only | Unit (exact) |\n| **Best for** | New code, libraries | Legacy code, quick checks | Strict unit matching |"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Approach 1: Type annotations (recommended for new code)\n",
    "@validate_units\n",
    "def approach1(V: u.Quantity[u.volt], I: u.Quantity[u.amp]) -> u.Quantity[u.watt]:\n",
    "    return V * I\n",
    "\n",
    "\n",
    "# Approach 2: check_dims decorator\n",
    "@u.check_dims(V=u.volt.dim, I=u.amp.dim, result=u.watt.dim)\n",
    "def approach2(V, I):\n",
    "    return V * I\n",
    "\n",
    "\n",
    "# Approach 3: check_units decorator\n",
    "@u.check_units(V=u.volt, I=u.amp, result=u.watt)\n",
    "def approach3(V, I):\n",
    "    return V * I\n",
    "\n",
    "\n",
    "# All three produce the same result\n",
    "V, I = 12.0 * u.volt, 2.0 * u.amp\n",
    "print(f\"Approach 1: {approach1(V, I)}\")\n",
    "print(f\"Approach 2: {approach2(V, I)}\")\n",
    "print(f\"Approach 3: {approach3(V, I)}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 11. Best Practices\n\n### Use physical type strings for public APIs\n\nPhysical type strings (`\"length\"`, `\"speed\"`) are more flexible than specific units. They accept any unit of the right dimension:\n\n```python\n# Preferred: accepts meter, kilometer, mile, etc.\ndef distance(x: Quantity[\"length\"]) -> Quantity[\"length\"]: ...\n\n# More restrictive: conceptually expects meters\ndef distance(x: Quantity[u.meter]) -> Quantity[u.meter]: ...\n```\n\n### Use specific units for internal functions\n\nWhen a function's implementation depends on a specific unit (e.g., for numerical constants), annotate with that unit:\n\n```python\n@validate_units(strict=True)\ndef _internal_calc(voltage_mV: Quantity[u.mvolt]) -> Quantity[u.mvolt]:\n    # Implementation uses millivolt-specific constants\n    return voltage_mV * 0.5\n```\n\n### Use `isinstance` for runtime branching\n\n```python\ndef process(quantity):\n    if isinstance(quantity, Quantity[\"length\"]):\n        return quantity.in_unit(u.meter)\n    elif isinstance(quantity, Quantity[\"time\"]):\n        return quantity.in_unit(u.second)\n    else:\n        raise TypeError(f\"Unexpected dimension: {quantity.dim}\")\n```\n\n### Combine with `@validate_units` for safety-critical code\n\n```python\n@validate_units\ndef drug_dosage(\n    concentration: Quantity[\"density\"],\n    volume: Quantity[\"volume\"],\n    patient_mass: Quantity[\"mass\"],\n) -> Quantity[\"dimensionless\"]:\n    return (concentration * volume) / patient_mass\n```\n\n### Use `QuantityLike` and `UnitLike` for flexible inputs\n\n```python\nfrom brainunit.typing import QuantityLike, UnitLike\n\ndef convert(value: QuantityLike, target: UnitLike) -> Quantity:\n    q = Quantity(value)\n    if target is not None:\n        return q.in_unit(target)\n    return q\n```"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}