{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0ce95f99",
   "metadata": {},
   "source": [
    "# Choose a parameter transform"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1003c194",
   "metadata": {},
   "source": [
    "This how-to is the transform catalog: which transform maps onto which constrained domain,\n",
    "how to compose them, and how to choose. Four related documents cover different ground:\n",
    "\n",
    "| Document | What it gives you |\n",
    "| --- | --- |\n",
    "| [Constrain and regularize parameters](constrain_and_regularize_parameters.ipynb) | The next step — combining a transform with a regularization penalty |\n",
    "| [Parameters, transforms, and regularization](../tutorials/core/05_parameters_transforms_regularization.ipynb) | The guided tour, with three worked models |\n",
    "| [The parameter model](../concepts/the_parameter_model.md) | The *why* behind the design |\n",
    "| [Parameter containers API](../apis/nn/parameters.rst) | Exact signatures |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "de2f27fe",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-27T07:51:06.677717Z",
     "iopub.status.busy": "2026-07-27T07:51:06.677717Z",
     "iopub.status.idle": "2026-07-27T07:51:07.363895Z",
     "shell.execute_reply": "2026-07-27T07:51:07.363895Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0.5.2'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import brainunit as u\n",
    "\n",
    "import brainstate\n",
    "import brainstate.nn as nn\n",
    "\n",
    "brainstate.random.seed(0)\n",
    "brainstate.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99d4e9ea",
   "metadata": {},
   "source": [
    "## 1. What a transform is, and why not clipping\n",
    "\n",
    "A transform is a **bijection** between an unconstrained space — all of $\\mathbb{R}$, where\n",
    "gradient descent is well behaved — and a constrained space such as the positives, an interval,\n",
    "or the probability simplex. The optimizer moves the unconstrained value; the model reads the\n",
    "constrained one.\n",
    "\n",
    "The tempting alternative is to clip after each update. It breaks learning at exactly the place\n",
    "you care about: below the bound the clipped output is constant, so its gradient is zero and the\n",
    "optimizer gets no signal to climb back."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1d9e7234",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-27T07:51:07.364898Z",
     "iopub.status.busy": "2026-07-27T07:51:07.364898Z",
     "iopub.status.idle": "2026-07-27T07:51:07.520733Z",
     "shell.execute_reply": "2026-07-27T07:51:07.520733Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "at theta = -3.0\n",
      "  clip     -> value 0.000000, gradient 0.000000\n",
      "  SoftplusT-> value 0.048587, gradient 0.047426\n"
     ]
    }
   ],
   "source": [
    "theta = jnp.array(-3.0)\n",
    "soft = nn.SoftplusT(lower=0.0)\n",
    "\n",
    "clip_grad = jax.grad(lambda t: jnp.clip(t, 0.0, None).sum())(theta)\n",
    "soft_grad = jax.grad(lambda t: soft.forward(t).sum())(theta)\n",
    "\n",
    "print(f'at theta = {float(theta)}')\n",
    "print(f'  clip     -> value {float(jnp.clip(theta, 0.0, None)):.6f}, gradient {float(clip_grad):.6f}')\n",
    "print(f'  SoftplusT-> value {float(soft.forward(theta)):.6f}, gradient {float(soft_grad):.6f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8ed4181",
   "metadata": {},
   "source": [
    "The clipped parameter is stuck: a zero gradient means the optimizer cannot move it, and any\n",
    "momentum it had is wasted. The transformed parameter has a small but non-zero gradient\n",
    "everywhere, so it can always recover."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57635425",
   "metadata": {},
   "source": [
    "## 2. The catalog, by constrained domain\n",
    "\n",
    "| Constrained domain | Transforms |\n",
    "| --- | --- |\n",
    "| $(\\text{lower}, \\infty)$ — positive quantities | `SoftplusT(lower)`, `ExpT(lower)`, `PositiveT()`, `ReluT(lower_bound=0.0)` |\n",
    "| $(-\\infty, \\text{upper})$ — negative quantities | `NegSoftplusT(upper)`, `NegativeT()` |\n",
    "| $(\\text{lower}, \\text{upper})$ — a bounded interval | `SigmoidT(lower, upper)`, `ScaledSigmoidT(lower, upper, beta=1.0)`, `TanhT(lower, upper)`, `SoftsignT(lower, upper)`, `ClipT(lower, upper)` |\n",
    "| Non-negative, summing to one — probabilities | `SimplexT()` |\n",
    "| Unit $L_2$ norm | `UnitVectorT()` |\n",
    "| Monotonically increasing entries | `OrderedT()` |\n",
    "| Reparameterizations | `AffineT(scale, shift)`, `PowerT(lmbda=0.5)`, `LogT(lower)` |\n",
    "| Composition and masking | `ChainT(*transforms)`, `MaskedT(mask, transform)` |\n",
    "\n",
    "`IdentityT()` is the default: no constraint at all.\n",
    "\n",
    "Note that `TanhT` and `SoftsignT` take explicit `lower` and `upper` bounds — they are bounded\n",
    "interval maps, not fixed $(-1, 1)$ squashers."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f06139cb",
   "metadata": {},
   "source": [
    "### 2.1 Every transform round-trips\n",
    "\n",
    "`inverse()` undoes `forward()`. That is what lets `Param` store an unconstrained value while\n",
    "you read and write constrained ones. The check below runs the whole catalog through a\n",
    "round-trip."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5b89fd1b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-27T07:51:07.521739Z",
     "iopub.status.busy": "2026-07-27T07:51:07.521739Z",
     "iopub.status.idle": "2026-07-27T07:51:09.678286Z",
     "shell.execute_reply": "2026-07-27T07:51:09.678286Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "transform              round-trip   max abs error\n",
      "IdentityT()            True         0.0e+00\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SoftplusT(0.0)         True         3.0e-08\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ExpT(0.0)              True         0.0e+00\n",
      "PositiveT()            True         0.0e+00\n",
      "ReluT()                True         0.0e+00\n",
      "LogT(0.0)              True         0.0e+00\n",
      "NegSoftplusT(0.0)      True         3.0e-08\n",
      "NegativeT()            True         3.0e-08\n",
      "SigmoidT(0, 1)         True         0.0e+00\n",
      "ScaledSigmoidT(0, 1)   True         0.0e+00\n",
      "ClipT(0, 1)            True         0.0e+00\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TanhT(-1, 1)           True         0.0e+00\n",
      "SoftsignT(-1, 1)       True         0.0e+00\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SimplexT()             True         3.0e-08\n",
      "UnitVectorT()          True         0.0e+00\n",
      "OrderedT()             True         0.0e+00\n",
      "AffineT(2, 1)          True         0.0e+00\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PowerT()               True         0.0e+00\n"
     ]
    }
   ],
   "source": [
    "CASES = [\n",
    "    ('IdentityT()', nn.IdentityT(), jnp.array([-1.0, 0.0, 2.0])),\n",
    "    ('SoftplusT(0.0)', nn.SoftplusT(lower=0.0), jnp.array([0.5, 2.0, 8.0])),\n",
    "    ('ExpT(0.0)', nn.ExpT(lower=0.0), jnp.array([0.5, 2.0, 8.0])),\n",
    "    ('PositiveT()', nn.PositiveT(), jnp.array([0.5, 2.0, 8.0])),\n",
    "    ('ReluT()', nn.ReluT(), jnp.array([0.5, 2.0, 8.0])),\n",
    "    ('LogT(0.0)', nn.LogT(lower=0.0), jnp.array([0.5, 2.0, 8.0])),\n",
    "    ('NegSoftplusT(0.0)', nn.NegSoftplusT(upper=0.0), jnp.array([-0.5, -2.0, -8.0])),\n",
    "    ('NegativeT()', nn.NegativeT(), jnp.array([-0.5, -2.0, -8.0])),\n",
    "    ('SigmoidT(0, 1)', nn.SigmoidT(lower=0.0, upper=1.0), jnp.array([0.1, 0.5, 0.9])),\n",
    "    ('ScaledSigmoidT(0, 1)', nn.ScaledSigmoidT(lower=0.0, upper=1.0), jnp.array([0.1, 0.5, 0.9])),\n",
    "    ('ClipT(0, 1)', nn.ClipT(lower=0.0, upper=1.0), jnp.array([0.1, 0.5, 0.9])),\n",
    "    ('TanhT(-1, 1)', nn.TanhT(lower=-1.0, upper=1.0), jnp.array([-0.5, 0.0, 0.5])),\n",
    "    ('SoftsignT(-1, 1)', nn.SoftsignT(lower=-1.0, upper=1.0), jnp.array([-0.5, 0.0, 0.5])),\n",
    "    ('SimplexT()', nn.SimplexT(), jnp.array([0.2, 0.3, 0.5])),\n",
    "    ('UnitVectorT()', nn.UnitVectorT(), jnp.array([0.6, 0.8])),\n",
    "    ('OrderedT()', nn.OrderedT(), jnp.array([-1.0, 0.5, 2.0])),\n",
    "    ('AffineT(2, 1)', nn.AffineT(scale=2.0, shift=1.0), jnp.array([-1.0, 0.0, 3.0])),\n",
    "    ('PowerT()', nn.PowerT(), jnp.array([0.5, 2.0, 8.0])),\n",
    "]\n",
    "\n",
    "print(f'{\"transform\":<22} {\"round-trip\":<12} max abs error')\n",
    "for name, t, x in CASES:\n",
    "    back = u.get_magnitude(t.forward(t.inverse(x)))\n",
    "    err = float(jnp.max(jnp.abs(back - x)))\n",
    "    print(f'{name:<22} {str(bool(err < 1e-4)):<12} {err:.1e}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acb70511",
   "metadata": {},
   "source": [
    "Every transform recovers the original constrained value to within floating-point tolerance.\n",
    "\n",
    "Note the `u.get_magnitude(...)` call. Most transforms return a plain JAX array, but a few —\n",
    "`ExpT` and `LogT` among them — return a dimensionless `brainunit.Quantity`. `get_magnitude`\n",
    "normalizes both to a plain array. See the pitfalls in section 6."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99b8828f",
   "metadata": {},
   "source": [
    "## 3. Composing transforms\n",
    "\n",
    "`ChainT` applies its transforms in order, so you can reach a domain no single transform covers.\n",
    "`MaskedT` applies a transform to only the entries its mask selects, leaving the rest\n",
    "unconstrained."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "357e4f32",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-27T07:51:09.679578Z",
     "iopub.status.busy": "2026-07-27T07:51:09.679578Z",
     "iopub.status.idle": "2026-07-27T07:51:09.716620Z",
     "shell.execute_reply": "2026-07-27T07:51:09.716620Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ChainT round-trip max error: 2.9802322387695312e-08\n",
      "MaskedT keeps the unmasked entry negative: [ 0.49999997 -2.          8.        ]\n"
     ]
    }
   ],
   "source": [
    "chained = nn.ChainT(nn.AffineT(scale=2.0, shift=1.0), nn.SoftplusT(lower=0.0))\n",
    "x = jnp.array([0.5, 2.0, 8.0])\n",
    "print('ChainT round-trip max error:',\n",
    "      float(jnp.max(jnp.abs(u.get_magnitude(chained.forward(chained.inverse(x))) - x))))\n",
    "\n",
    "# Constrain entries 0 and 2 to be positive; leave entry 1 free.\n",
    "masked = nn.MaskedT(jnp.array([1.0, 0.0, 1.0]), nn.SoftplusT(lower=0.0))\n",
    "y = jnp.array([0.5, -2.0, 8.0])\n",
    "print('MaskedT keeps the unmasked entry negative:',\n",
    "      u.get_magnitude(masked.forward(masked.inverse(y))))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e925d2f",
   "metadata": {},
   "source": [
    "## 4. Choosing one\n",
    "\n",
    "| If you need | Use |\n",
    "| --- | --- |\n",
    "| A rate, conductance, or time constant that must stay positive | `SoftplusT(lower)` |\n",
    "| The same, but spanning orders of magnitude | `ExpT(lower)` |\n",
    "| A mixing weight or probability in $[0, 1]$ | `SigmoidT(0.0, 1.0)` |\n",
    "| A categorical distribution over $k$ outcomes | `SimplexT()` |\n",
    "| A direction, with magnitude handled separately | `UnitVectorT()` |\n",
    "| Sorted thresholds or bin edges | `OrderedT()` |\n",
    "| To constrain only some entries of an array | `MaskedT(mask, transform)` |\n",
    "| A domain no single transform covers | `ChainT(...)` |\n",
    "\n",
    "`SoftplusT` is the default choice for positivity: it is gentler near the bound than `ExpT`,\n",
    "which grows exponentially and can turn a moderate unconstrained value into a very large\n",
    "constrained one."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a14ea374",
   "metadata": {},
   "source": [
    "## 5. Attaching a transform to a parameter\n",
    "\n",
    "Pass it as `t=`. From then on `value()` returns the constrained value, and `set_value()`\n",
    "accepts a constrained value and applies the inverse for you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "004edd36",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-27T07:51:09.718627Z",
     "iopub.status.busy": "2026-07-27T07:51:09.718627Z",
     "iopub.status.idle": "2026-07-27T07:51:09.828734Z",
     "shell.execute_reply": "2026-07-27T07:51:09.828734Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "value()      : 5.0\n",
      "stored in val: 3.9815146923065186\n",
      "after a large negative update: 1.0 (still above the 1.0 floor)\n"
     ]
    }
   ],
   "source": [
    "tau = nn.Param(jnp.array(5.0), t=nn.SoftplusT(lower=1.0))\n",
    "print('value()      :', float(tau.value()))\n",
    "print('stored in val:', float(tau.val.value))\n",
    "\n",
    "tau.val.value = jnp.array(-50.0)      # an aggressive optimizer step\n",
    "print('after a large negative update:', float(tau.value()), '(still above the 1.0 floor)')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd9ee2ea",
   "metadata": {},
   "source": [
    "Note what the printout above says about the constructor: `Param(5.0, t=SoftplusT(lower=1.0))`\n",
    "gives `value() == 5.0`. The value you pass in is interpreted in **constrained** space, and the\n",
    "inverse transform is applied once to derive the unconstrained number stored in `val` — here\n",
    "`3.98`. You never have to work out the unconstrained value yourself.\n",
    "\n",
    "`Param` also accepts `precompute=`, a callable applied to the constrained value after the\n",
    "transform. Use it when the model needs a *derived* quantity rather than the constrained value\n",
    "itself — for instance a decay factor computed from a time constant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2c22c9bb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-27T07:51:09.829949Z",
     "iopub.status.busy": "2026-07-27T07:51:09.829949Z",
     "iopub.status.idle": "2026-07-27T07:51:09.860951Z",
     "shell.execute_reply": "2026-07-27T07:51:09.860951Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "without precompute, value() is the time constant : 5.0\n",
      "with precompute, value() is the derived factor   : 0.8187307715415955\n",
      "check exp(-1 / tau)                              : 0.8187307715415955\n"
     ]
    }
   ],
   "source": [
    "plain = nn.Param(jnp.array(5.0), t=nn.SoftplusT(lower=1.0))\n",
    "decay = nn.Param(jnp.array(5.0), t=nn.SoftplusT(lower=1.0),\n",
    "                 precompute=lambda tau: jnp.exp(-1.0 / tau))\n",
    "\n",
    "print('without precompute, value() is the time constant :', float(plain.value()))\n",
    "print('with precompute, value() is the derived factor   :', float(decay.value()))\n",
    "print('check exp(-1 / tau)                              :',\n",
    "      float(jnp.exp(-1.0 / plain.value())))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb6c4d59",
   "metadata": {},
   "source": [
    "## 6. Pitfalls\n",
    "\n",
    "**Saturation.** `SigmoidT` and `TanhT` flatten far from the origin. If the unconstrained value\n",
    "drifts to $\\pm 100$, the gradient is effectively zero and the parameter stops moving. Keep\n",
    "initial values near the middle of the range.\n",
    "\n",
    "**Initializing exactly on the bound.** `SoftplusT(lower=0.0)` cannot represent `0.0` — the\n",
    "inverse diverges. Initialize strictly inside the domain.\n",
    "\n",
    "**Inverting outside the domain.** `inverse()` is defined only on the constrained domain. Calling\n",
    "`SimplexT().inverse()` on a vector that does not sum to one is undefined behaviour, not an error.\n",
    "\n",
    "**Quantity vs. plain array.** `forward()` does not return the same type for every transform:\n",
    "`ExpT` and `LogT` return a dimensionless `brainunit.Quantity`, while `SoftplusT` returns a plain\n",
    "JAX array. Passing a `Quantity` to a raw `jnp` function raises a `TypeError`. Wrap the result in\n",
    "`u.get_magnitude(...)` when you need a plain array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ce1d9dc1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-27T07:51:09.862459Z",
     "iopub.status.busy": "2026-07-27T07:51:09.862459Z",
     "iopub.status.idle": "2026-07-27T07:51:09.907951Z",
     "shell.execute_reply": "2026-07-27T07:51:09.907951Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SoftplusT forward type: ArrayImpl\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ExpT      forward type: Quantity\n",
      "raw jnp on a Quantity -> TypeError\n",
      "with get_magnitude    -> 2.7182817459106445\n"
     ]
    }
   ],
   "source": [
    "print('SoftplusT forward type:', type(nn.SoftplusT(lower=0.0).forward(jnp.array(1.0))).__name__)\n",
    "print('ExpT      forward type:', type(nn.ExpT(lower=0.0).forward(jnp.array(1.0))).__name__)\n",
    "\n",
    "try:\n",
    "    jnp.abs(nn.ExpT(lower=0.0).forward(jnp.array(1.0)))\n",
    "except TypeError as e:\n",
    "    print('raw jnp on a Quantity ->', type(e).__name__)\n",
    "print('with get_magnitude    ->',\n",
    "      float(jnp.abs(u.get_magnitude(nn.ExpT(lower=0.0).forward(jnp.array(1.0))))))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eda20654",
   "metadata": {},
   "source": [
    "## 7. Summary\n",
    "\n",
    "- A transform is a bijection between unconstrained $\\mathbb{R}$ and a constrained domain. The\n",
    "  optimizer works unconstrained; `value()` returns the constrained value.\n",
    "- Clipping zeroes the gradient at the bound; a transform never does.\n",
    "- Pick by constrained domain using the table in section 4. `SoftplusT` is the default for\n",
    "  positivity.\n",
    "- Compose with `ChainT`; restrict to part of an array with `MaskedT`.\n",
    "- Watch for saturation, initialization on a bound, out-of-domain inverses, and the\n",
    "  `Quantity`-vs-array difference between transforms."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}