docs/examples/gyrus_overview.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gyrus Overview (`gyrus==0.1.0`)\n",
    "\n",
    "Recursively generate large-scale Nengo models using NumPy semantics.\n",
    "\n",
    "## Cheat Sheet (API Summary)\n",
    "\n",
    "**Operator methods** consume an operator as their first argument, and produce some other\n",
    "operator (which typically depends on the given operator as input). They are also\n",
    "vectorized to work element-wise across \"folds\" (i.e., N-D arrays of operators),\n",
    "analogous to \"ufuncs\" in NumPy. Each method can be called either as a function, or as a\n",
    "method on the operator, in the same way that `np.sum(arr)` is equivalent to `arr.sum()`:\n",
    " - `apply` - Operator that applies a function to each output ideally using a Node.\n",
    " - `configure` - Operator that applies configuration settings to all downstream\n",
    "operators.\n",
    " - `convolve` - Operator that approximates the circular convolution using Product\n",
    "networks.\n",
    " - `filter` - Operator that applies a synaptic filter to each output.\n",
    " - `decode` - Operator that approximates a function of each output using an Ensemble.\n",
    " - `integrate` - Operator that solves $\\dot{x} = u + \\text{integrand}(x)$.\n",
    " - `transform` - Operator that applies a single transform to every output.\n",
    " - `lti` - Operator that solves $\\dot{x} = A.\\text{state}(x) + B.u$. where $A$, $B$ =\n",
    "system.\n",
    " - `multiply` - Operator that approximates an element-wise product using a Product\n",
    "network.\n",
    " - `neurons` - Operator that linearly projects to a layer of neurons.\n",
    " - `slice` - Operator that slices each output using Nengo's object slicing.\n",
    " - `unbundle` - Operator that splits each output into a Fold of one-dimensional outputs.\n",
    "\n",
    "**Operator attributes** can be accessed from any operator or fold:\n",
    " - `input_ops` - Tuple of operators that the operator depends on as inputs.\n",
    " - `label` - Returns the label of the operator's output as it appears in Nengo GUI.\n",
    " - `ndim` - Returns the dimensionality of the underlying NumPy array (or `0`).\n",
    " - `shape` - Returns the shape of the underlying NumPy array (or `()`).\n",
    " - `size_out` - Returns the `size_out` of each element of the underlying NumPy array.\n",
    "\n",
    "**Fold methods** manipulate folds, which are a special kind of operator equivalent to an\n",
    "`np.ndarray` where each element of the array is another operator. These methods are like\n",
    "operator methods, but consume and produce folds:\n",
    " - `bundle` - Operator that joins all of the outputs along a given axis into a single\n",
    "output.\n",
    " - `integrate_fold` - Operator that solves $\\dot{x} = u + \\text{integrand}(x)$ where x\n",
    "is a Fold.\n",
    " - `reduce_transform` - Operator that sums together transforms applied to each output\n",
    "along an axis.\n",
    "\n",
    "**NumPy array functions** are like Fold methods, but in actuality are NumPy array\n",
    "functions applied to the fold's underlying array:\n",
    " - Functional programming routines: `np.apply_along_axis`.\n",
    " - Math routines: `np.dot`, `np.mean`, `np.outer`, `np.prod`, `np.sum`.\n",
    " - Changing array shape: `np.reshape`, `np.ravel`.\n",
    " - Transpose-like operations: `np.moveaxis` `np.rollaxis`, `np.swapaxes`,\n",
    "`np.transpose`.\n",
    " - Changing number of dimensions: `np.atleast_1d`, `np.atleast_2d`, `np.atleast_3d`,\n",
    "`np.broadcast_to`, `np.broadcast_arrays`, `np.expand_dims`, `np.squeeze`.\n",
    " - Joining arrays: `np.concatenate`, `np.stack`, `np.block`, `np.vstack`, `np.hstack`,\n",
    "`np.dstack`, `np.column_stack`.\n",
    " - Splitting arrays: `np.split`, `np.array_split`, `np.dsplit`, `np.hsplit`,\n",
    "`np.vsplit`.\n",
    " - Tiling arrays: `np.repeat`, `np.tile`.\n",
    " - Rearranging elements: `np.flip`, `np.fliplr`, `np.flipud`, `np.reshape`, `np.roll`,\n",
    "`np.rot90`.\n",
    "\n",
    "**NumPy array methods** are like NumPy array functions, but are only accessible as\n",
    "methods on a given fold:\n",
    " - `flatten` - Equivalent to `np.ndarray.flatten`.\n",
    " - `T` - Equivalent to `np.ndarray.T`.\n",
    " - `__getitem__` - Equivalent to `ndarray.__getitem__`.\n",
    "\n",
    "**Fold attributes** can be accessed from any fold:\n",
    " - `array` - Returns the Fold as a read-only NumPy array.\n",
    "\n",
    "**Operator input functions** create roots of the operator graph (i.e., operators without\n",
    "any input operators):\n",
    " - `broadcast_scalar` - Operator that creates a scalar Node and transforms it to match\n",
    "some shape.\n",
    " - `stimuli` - Operator that supplies input data given Nengo objects or Node outputs.\n",
    "(Vectorized version of `stimulus`.)\n",
    " - `stimulus` - Operator that supplies input data given a Nengo object or a Node output.\n",
    "\n",
    "**Top-level functions** that don't fit under any other category:\n",
    " - `asoperator` - Returns x as a single Operator if possible, otherwise as a Fold.\n",
    " - `fold` - Return a fold from an iterable of operators or folds.\n",
    "\n",
    "**Extensions** allow users to easily define their own Gyrus operators and register them\n",
    "as array functions and ufuncs:\n",
    " - `register_method` - Register a function as a method that takes self as its first\n",
    "argument.\n",
    " - `register_ufunc` - Register a function as a NumPy ufunc handler for this type.\n",
    " - `vectorize` - Dynamically creates a vectorized operator type with the given\n",
    "implementation.\n",
    "\n",
    "**Python dunder operators** overload the behaviour of common Python operations:\n",
    " - add (`+`)\n",
    " - subtract (`-`)\n",
    " - multiply (`*`)\n",
    " - true divide (`/`)\n",
    " - power (`**`)\n",
    " - matmul (`@`)\n",
    " - slicing (`[...]`) at the Nengo Node level.\n",
    "\n",
    "**NumPy ufuncs** approximate the behaviour of common NumPy ufuncs using the `decode`\n",
    "operator:\n",
    " - `cos`\n",
    " - `sin`\n",
    " - `tanh`\n",
    " - `square`\n",
    "\n",
    "**Operator outputs** are methods attached to each operator that allow the user to build\n",
    "and simulate the Nengo model that produces the output(s) of the given operator:\n",
    " - `make` - Generates operator's graph within the current Nengo network context.\n",
    "(Interoperable with existing Nengo networks.)\n",
    " - `probe` - Returns a function that consumes a simulator and produces probe data.\n",
    "(Typically used in conjunction with `make`.)\n",
    " - `run` - Generates a Nengo network, then probes, simulates, and returns its data.\n",
    "(Convenience wrapper around `make` and `probe`.)\n",
    "\n",
    "**Optional operators methods** require installation with the extra, `[optional]`, and\n",
    "give access to the following operator methods:\n",
    " - `layer` - Operator that applies a `nengo_dl.Layer` to each output.\n",
    " - `tensor_node` - Operator that applies a `nengo_dl.TensorNode` to each output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import nengo\n",
    "import nengo_dl\n",
    "from nengo_gui.jupyter import InlineGUI\n",
    "\n",
    "import numpy as np\n",
    "from scipy.special import legendre\n",
    "import seaborn as sns\n",
    "\n",
    "import gyrus"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example #1 - Communication Channel\n",
    "\n",
    "The obligatory example of encoding a scalar into a spiking ensemble and decoding out a\n",
    "filtered version of the identity function.\n",
    "\n",
    "Note: `decode` accepts a `function` parameter which defaults to the identity, and an\n",
    "`n_neurons` parameter which currently defaults to 100."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@gyrus.stimulus\n",
    "def stim(t):\n",
    "    return np.sin(2 * np.pi * t)\n",
    "\n",
    "\n",
    "out = stim.decode().filter(5e-3).run(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.title(\"Example #1 - Communication Channel\")\n",
    "plt.plot(out)\n",
    "plt.xlabel(\"Time-step\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The fact that operators are vectorized comes with some pleasantly surprising\n",
    "consequences, such as the ability to take a signal and filter it with a number of\n",
    "distinct synapses to produce a parallel fold:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "synapses = [0.005, 0.05, 0.1]\n",
    "out = np.asarray(stim.decode().filter(synapses).run(1))\n",
    "out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.title(\"Example #1 - Communication Channel (Vectorized Filtering)\")\n",
    "for i, y_i in enumerate(out):\n",
    "    plt.plot(y_i, label=fr\"$\\tau$ = {synapses[i]}s\")\n",
    "plt.xlabel(\"Time-step\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example #2 - Ensemble Array\n",
    "\n",
    "Ensemble arrays are 'trivial' in Gyrus – just add another dimension to the underlying\n",
    "NumPy array. This naturally generalizes to higher-dimensional ndarrays of arbitrary\n",
    "shape (a.k.a., tensors), which serves to overcome Nengo's limitation of representing\n",
    "only scalars or vectors, while still being just Nengo under the hood.\n",
    "\n",
    "Note: The `bundle(axis=0)` method changes the output shape from `(16, 1000, 1)` to\n",
    "`(1000, 16)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d = 16\n",
    "stim = gyrus.stimuli(np.linspace(-1, 1, d))\n",
    "assert stim.shape == (d,)\n",
    "\n",
    "y = stim.decode().filter(1e-2).bundle()\n",
    "assert y.size_out == d\n",
    "\n",
    "out = np.asarray(y.run(1))\n",
    "out.shape  # shape: (time, size_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.title(f\"Example #2 - Ensemble Array (d={d})\")\n",
    "plt.plot(out)\n",
    "plt.xlabel(\"Time-step\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example #3 - Nonlinear Functions\n",
    "\n",
    "Numeric operations are overloaded to automatically use neurons for each nonlinear\n",
    "computation (e.g., `a ** 2` and `(...) * b`) and transforms for each linear computation\n",
    "(e.g., `1 - ...` and `2 * b`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_functions = [\n",
    "    lambda t: 2 * t - 1,\n",
    "    lambda t: 1 - 2 * t,\n",
    "]\n",
    "\n",
    "\n",
    "def f(a, b, tau=1e-2):  # like a subnetwork\n",
    "    return ((1 - a ** 2).filter(tau) * (2 * b)).filter(tau)\n",
    "\n",
    "\n",
    "x = gyrus.stimuli(input_functions)\n",
    "y = f(*x)\n",
    "out = np.asarray(y.run(1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.title(\"Example #3 - Nonlinear Functions\")\n",
    "plt.plot(out)\n",
    "plt.xlabel(\"Time-step\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can make it more accurate by configuring the number of neurons used by all operators\n",
    "downstream of `a` and/or `b`. Configuration follows natural precedence rules (first\n",
    "left-to-right, then upstream)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = f(*x.configure(n_neurons=1000))\n",
    "out = np.asarray(y.run(1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.title(\"Example #3 - Nonlinear Functions (More Neurons)\")\n",
    "plt.plot(out)\n",
    "plt.xlabel(\"Time-step\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example #4 - Building a Gyrus operator into a Nengo network\n",
    "\n",
    "We now show how to make a Gyrus operator explicitly inside of our own `nengo.Network()`\n",
    "definition using the `make()` method, and then use the `stimulus` function to create\n",
    "Node inputs to the Gyrus operator.\n",
    "\n",
    "In general, this pattern enables full interoperability with Nengo: `stimulus(node)` can\n",
    "be used to create Gyrus inputs from Nengo nodes, and the output(s) of `make()` are\n",
    "Nengo nodes that can be connected from just like any other Nengo object. This gives\n",
    "users the ability to interface networks generated using Gyrus with existing Nengo\n",
    "networks, and therefore a way to dig down into certain implementation details at the\n",
    "level of Nengo and reuse existing Nengo code.\n",
    "\n",
    "For sake of example, we remake the same network from example #3, but get the ideal\n",
    "output (i.e., don't approximate the nonlinearities with neurons). To do so, we use\n",
    "Nengo's `model.config` to set `neuron_type=nengo.Direct()` for all ensembles. An\n",
    "equivalent approach would be to invoke `.configure(neuron_type=nengo.Direct())` on\n",
    "either of the two `Stimulus` operators.\n",
    "\n",
    "Note: `p = gyrus.probe((...).make())` and then `p(sim)` is an equivalent way to do the\n",
    "probing, but has the added feature of being able to recursively probe entire folds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with nengo.Network() as model:\n",
    "    model.config[nengo.Ensemble].neuron_type = nengo.Direct()\n",
    "\n",
    "    a = nengo.Node(input_functions[0])\n",
    "    b = nengo.Node(input_functions[1])\n",
    "\n",
    "    out = f(gyrus.stimulus(a), gyrus.stimulus(b)).make()\n",
    "\n",
    "    p = nengo.Probe(out)\n",
    "\n",
    "with nengo.Simulator(model) as sim:\n",
    "    sim.run(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.title(\"Example #4 - Building a Gyrus operator into a Nengo network\")\n",
    "plt.plot(sim.data[p])\n",
    "plt.xlabel(\"Time-step\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also use the Nengo GUI to visualize the generated Nengo model. The \"Decode\"\n",
    "ensemble belongs to the `a ** 2` operator and the \"Multiply\" subnetwork implements the\n",
    "product with `b`. Nodes in between are responsible for linear transformations, and an\n",
    "additional node is automatically created to supply the constant `1` inside of `f`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "InlineGUI(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example #5 - More complicated array functions\n",
    "\n",
    "This example computes the mean-squared power across a vector of $d$ shifted Legendre\n",
    "polynomials provided as input. This also demonstrates how to probe multiple operators\n",
    "simultaneously using a fold."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d = 16\n",
    "input_functions = [lambda t, p=legendre(i): p(2 * t - 1) for i in range(d)]\n",
    "\n",
    "stim = gyrus.stimuli(input_functions)\n",
    "y_hat = (stim ** 2).mean()\n",
    "y_ideal = stim.apply(np.square).mean()\n",
    "\n",
    "out_hat, out_ideal = gyrus.fold([y_hat, y_ideal]).filter(5e-3).run(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.title(\"Example #5 - More complicated array functions\")\n",
    "plt.plot(out_hat, alpha=0.8, label=\"Spiking\")\n",
    "plt.plot(out_ideal, label=\"Ideal\")\n",
    "plt.xlabel(\"Time-step\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example #6 - Building a Nengo network into a Gyrus operator\n",
    "\n",
    "Suppose we want to define our own Gyrus operator using Nengo code, such that it is\n",
    "appropriately vectorized across folds. We can do so using the `@gyrus.vectorize`\n",
    "decorator. Inside of the decorated implementation, we are free to use whatever Nengo\n",
    "code we'd like, and can even `make()` other Gyrus operators and interface them with the\n",
    "Nengo code (see example #4). The implementation is expected to return a single Nengo\n",
    "node.\n",
    "\n",
    "For sake of example, we'll take an existing Nengo subnetwork\n",
    "(`nengo.networks.oscillator`) and turn it into a Gyrus operator. We'll also register the\n",
    "operator as a method on all existing operators using the `@gyrus.register_method`\n",
    "decorator.\n",
    "\n",
    "To demonstrate the benefits of doing this, we show that it becomes trivial to\n",
    "instantiate three different oscillators with three separate kick inputs, simply by\n",
    "expanding the shape of the input. We can even specify three different frequencies, and\n",
    "they will be applied element-wise alongside each respective kick. In general, all Gyrus\n",
    "operators are vectorized and follow the same NumPy semantics as `@np.vectorize` (in\n",
    "fact, that is what is used under the hood).\n",
    "\n",
    "Note: A more accurate oscillator implementation is supported by the `lti` method (see\n",
    "final example #9)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@gyrus.register_method(\"oscillate\")\n",
    "@gyrus.vectorize(\"Oscillate\")\n",
    "def oscillate(node, *, n_neurons, recurrent_tau, frequency, **kwargs):\n",
    "    oscillator = nengo.networks.Oscillator(\n",
    "        n_neurons=n_neurons,\n",
    "        recurrent_tau=recurrent_tau,\n",
    "        frequency=frequency,\n",
    "    )\n",
    "    nengo.Connection(node, oscillator.input, synapse=None)\n",
    "    return oscillator.output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "strength = np.asarray([0.5, 1, 4])\n",
    "hz = np.asarray([0.5, 0.75, 1.0])\n",
    "\n",
    "kicks = gyrus.stimuli([lambda t, r=r: r if t <= 0.05 else 0 for r in strength])\n",
    "x = kicks.transform([[1], [0]]).oscillate(\n",
    "    n_neurons=250,\n",
    "    recurrent_tau=0.1,\n",
    "    frequency=2 * np.pi * hz,\n",
    ")\n",
    "out = np.asarray(x.filter(1e-2).run(1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(5, 5))\n",
    "plt.title(\"Example #6 - Building a Nengo network into a Gyrus operator\")\n",
    "for i, xy in enumerate(out):\n",
    "    plt.plot(*xy.T, label=f\"{hz[i]} Hz\")\n",
    "plt.axis(\"equal\")\n",
    "plt.xlim(-1, 1)\n",
    "plt.ylim(-1, 1)\n",
    "plt.xlabel(\"$x(t)$\")\n",
    "plt.ylabel(\"$y(t)$\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `gyrus.vectorize` decorator can also be used by directly passing it in an\n",
    "implementation (rather than a name followed by several optional arguments). This\n",
    "convenience allows it to be used like a Keras Lambda layer, but for embedding arbitrary\n",
    "Nengo code into the Gyrus model. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def multiply_by_two(x):\n",
    "    y = nengo.Node(size_in=x.size_out)\n",
    "    nengo.Connection(x, y, transform=2, synapse=None)\n",
    "    return y\n",
    "\n",
    "\n",
    "x = gyrus.stimulus(np.ones(3))\n",
    "y = gyrus.vectorize(multiply_by_two)(x)\n",
    "y.run(1, dt=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example #7 - Extending Gyrus to work with existing NumPy functions\n",
    "\n",
    "Suppose we have some NumPy function that uses `np.linalg.norm`, and we wish to convert\n",
    "that function into a Nengo network without changing that piece of code. We'll first see\n",
    "that we get a `TypeError` \"no implementation found\" when we try to use that function on\n",
    "a Gyrus operator:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [],
   "source": [
    "np.linalg.norm(gyrus.stimuli(0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To register an implementation, use the `gyrus.register_method` decorator. If one wishes\n",
    "to override the entire behaviour of `norm`, then one can define:\n",
    "```python\n",
    "@gyrus.register_method(\"norm\")\n",
    "def custom_norm(x, ord=None, axis=None, keepdims=False):\n",
    "    ...\n",
    "```\n",
    "and return a Gyrus operator (or array of operators). This will customize the\n",
    "implementation of `np.linalg.norm` given an array of operators `x`.\n",
    "\n",
    "But there is also a much simpler solution in this case; just defer to the underlying\n",
    "`np.linalg.norm` implementation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gyrus.register_method(\"norm\")(np.linalg.norm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This causes NumPy to invoke its regular routine of squaring, summing, and square-rooting\n",
    "the results (with appropriate handling of axes), where each element being squared,\n",
    "summed, and square-rooted is a single Gyrus operator (i.e., an element of the array).\n",
    "Now, if we try again, we'll get a `TypeError` \"all returned `NotImplemented`\" -- but\n",
    "this time because it can't find an implementation for `np.sqrt`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [],
   "source": [
    "np.linalg.norm(gyrus.stimuli(0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To solve this final problem, we'll need to add an implementation for `np.sqrt`. We do so\n",
    "using a `@gyrus.register_ufunc(np.sqrt)` decorator which tells NumPy how to handle this\n",
    "particular ufunc. We'll also use the `@gyrus.vectorize` decorator so that we can provide\n",
    "the implementation for a single element using Nengo and have it automatically vectorized\n",
    "across folds (see example #6).\n",
    "\n",
    "This particular implementation will use an ensemble to decode `np.sqrt(x)` when `x >=\n",
    "0`, otherwise some `undefined` value for `x < 0` (default `0`). This implementation\n",
    "could be improved somewhat by specializing the Ensemble's parameters, but this serves as\n",
    "a straightforward example.\n",
    "\n",
    "To make this example more complete, we'll also register a number of Ensemble parameters\n",
    "as `configurable` which allows them to be configured via the `configure` operator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gyrus.nengo_helpers import get_params\n",
    "\n",
    "configurable = get_params(nengo.Ensemble) - {\"dimensions\"} | {\"seed\", \"undefined\"}\n",
    "print(\"Configurable parameters:\", configurable)\n",
    "\n",
    "\n",
    "@gyrus.register_ufunc(np.sqrt)\n",
    "@gyrus.vectorize(\"Sqrt\", configurable=configurable)\n",
    "def custom_sqrt(node, *, n_neurons=100, undefined=0, **ens_kwargs):\n",
    "    \"\"\"Straightforward implementation of the sqrt nonlinearity.\"\"\"\n",
    "    x = nengo.Ensemble(n_neurons, dimensions=node.size_out, **ens_kwargs)\n",
    "    nengo.Connection(node, x, synapse=None)\n",
    "\n",
    "    def valid_sqrt(x):\n",
    "        \"\"\"Returns sqrt(x) for x >= 0, otherwise undefined.\"\"\"\n",
    "        y = np.full_like(x, undefined)\n",
    "        y[x >= 0] = np.sqrt(x[x >= 0])\n",
    "        return y\n",
    "\n",
    "    out = nengo.Node(size_in=node.size_out)\n",
    "    nengo.Connection(x, out, function=valid_sqrt, synapse=None)\n",
    "    return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: We could also decorate with `@gyrus.register_method(\"sqrt\")` to attach this as a\n",
    "`.sqrt()` method to every operator (including folds).\n",
    "\n",
    "Let's first try this out to see how it looks. To get the ideal output as well, we\n",
    "configure the same operator in two different ways (this works because operators are\n",
    "_pure_ functions)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stim = gyrus.stimuli(lambda t: 2 * t - 1)\n",
    "x = gyrus.fold(\n",
    "    [\n",
    "        stim.configure(n_neurons=250),\n",
    "        stim.configure(neuron_type=nengo.Direct()),\n",
    "    ]\n",
    ")\n",
    "out_hat, out_ideal = np.sqrt(x).filter(1e-2).run(1)\n",
    "\n",
    "plt.figure()\n",
    "plt.title(\"Example #7 - Extending Gyrus (np.sqrt)\")\n",
    "plt.plot(out_hat, alpha=0.8, label=\"Spiking\")\n",
    "plt.plot(out_ideal, label=\"Ideal\")\n",
    "plt.xlabel(\"Time-step\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now `np.linalg.norm` finally works!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_functions = [\n",
    "    lambda t: t * np.sin(2 * np.pi * t),\n",
    "    lambda t: t * np.cos(2 * np.pi * t),\n",
    "]\n",
    "stim = gyrus.stimuli(input_functions)\n",
    "y_hat = np.linalg.norm(stim)\n",
    "y_ideal = stim.bundle().apply(np.linalg.norm)\n",
    "out_hat, out_ideal = gyrus.fold([y_hat, y_ideal]).filter(nengo.Alpha(1e-2)).run(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.title(\"Example #7 - Extending Gyrus (np.linalg.norm)\")\n",
    "plt.plot(out_hat, alpha=0.8, label=\"Spiking\")\n",
    "plt.plot(out_ideal, label=\"Ideal\")\n",
    "plt.xlabel(\"Time-step\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: Once `np.sqrt` is defined we can also do things like:\n",
    "```python\n",
    "y_hat = nengo.utils.numpy.rms(stim)\n",
    "y_ideal = stim.bundle().apply(nengo.utils.numpy.rms)\n",
    "```\n",
    "to implement other existing functions (e.g., `nengo.utils.numpy.rms`) using spikes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example #8 - Integrating Ordinary Differential Equations\n",
    "\n",
    "Gyrus also supports an `integrate` operator that can be used to integrate arbitrary\n",
    "differential equations produced by other operators. Here is a simple example launching a\n",
    "number of parallel integrators in parallel to with cosine inputs (the integral of cosine\n",
    "is sine).\n",
    "\n",
    "More generally, an `integrand` function can be supplied to integrate any function of the\n",
    "integral, `x`, using Gyrus operators (see example #9)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(gyrus.integrate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d = 32  # number of integrators\n",
    "dt = 1e-3  # time-step of simulation\n",
    "\n",
    "\n",
    "def f(t, hz):  # derivative of sinusoid with given frequency\n",
    "    return np.cos(2 * np.pi * hz * (t - dt)) * (2 * np.pi * hz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u = gyrus.stimuli([lambda t, hz=hz: f(t, hz) for hz in np.linspace(1, 2, d)])\n",
    "x = u.integrate().decode().filter(5e-3)\n",
    "out = np.asarray(x.run(1, dt=dt))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "colors = sns.color_palette(\"cubehelix\", n_colors=d)\n",
    "\n",
    "plt.figure()\n",
    "plt.title(\"Example #8 - Integrating Ordinary Differential Equations\")\n",
    "for x_i, color in zip(out, colors):\n",
    "    plt.plot(x_i.squeeze(axis=1), color=color, alpha=0.7)\n",
    "plt.xlabel(\"Time-step\")\n",
    "sns.despine(offset=10)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example #9 - Improved Oscillator & NengoDL\n",
    "\n",
    "The `lti` method in Gyrus builds on top of the `integrate` method to solve differential\n",
    "equations of the form $\\dot{x} = A.\\text{state}(x) + B.u$. where $A$, $B$ is the desired\n",
    "system. It does so using zero-order hold (ZOH)—assuming $\\text{state}(x) = x$—and\n",
    "compensates for the simulation time-step. Nonlinear dynamics are supported by specifying\n",
    "a nonlinear `state(x)` function for the integrand.\n",
    "\n",
    "Interestingly, the `lti` method is essentially just one line, and thus exists mostly as\n",
    "a convenience wrapper:\n",
    "```python\n",
    "return u.transform(Bmap).integrate(integrand=lambda x: state(x).transform(Amap))\n",
    "```\n",
    "\n",
    "For improved accuracy, we also use `unbundle()` and `bundle()` to unbundle the 2D\n",
    "representation into two 1D representations, and then bundle them back together.\n",
    "\n",
    "In addition, we use the NengoDL simulator to run this on the GPU, by passing\n",
    "`simulator=nengo_dl.Simulator` to the `run` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def oscillator(hertz):\n",
    "    \"\"\"Linear (A, B) system for an oscillator at given frequency.\"\"\"\n",
    "    radians = 2 * np.pi * hertz\n",
    "    return [[0, -radians], [radians, 0]], [1, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kick = gyrus.stimuli(lambda t: 1 / dt if t <= dt else 0)\n",
    "\n",
    "\n",
    "def state(x):\n",
    "    x = x.configure(neuron_type=nengo.SpikingRectifiedLinear())\n",
    "    return x.unbundle().decode().bundle()\n",
    "\n",
    "\n",
    "x = kick.lti(oscillator(hertz=5), state=state, dt=dt)\n",
    "out = x.run(1, dt=dt, simulator=nengo_dl.Simulator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.title(\"Example #9 - Improved Oscillator & NengoDL\")\n",
    "plt.plot(out)\n",
    "plt.xlabel(\"Time-step\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's also super simple to turn this into a controlled oscillator. Just multiply the\n",
    "`state` by a `control` signal. One subtle caveat is the ZOH discretization assumes the\n",
    "dynamics are linear, and so the time-step is no longer ideally accounted for when\n",
    "`control != 1`. The most accurate approach currently known is a bit more involved:\n",
    "https://gl.appliedbrainresearch.com/arvoelke/scratchpad/blob/master/benchmarks/oscillator.py#L37"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "control = gyrus.stimuli(\n",
    "    nengo.processes.Piecewise({0: 1, 0.2: -1, 0.4: 0, 0.6: -0.5, 0.8: 0.8})\n",
    ")\n",
    "\n",
    "\n",
    "def state(x):\n",
    "    x = x.configure(n_neurons=500, neuron_type=nengo.LIF())\n",
    "    return (x.unbundle() * control).bundle()\n",
    "\n",
    "\n",
    "x = kick.lti(oscillator(hertz=5), state=state, dt=dt)\n",
    "out_x, out_control = gyrus.fold([x, control]).run(\n",
    "    1, dt=dt, simulator=nengo_dl.Simulator\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.title(\"Example #9 - Improved Oscillator & NengoDL (Controlled)\")\n",
    "plt.plot(out_x, label=\"$x(t)$\")\n",
    "plt.plot(out_control, linestyle=\"--\", label=\"Control\")\n",
    "plt.xlabel(\"Time-step\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There's still plenty more that can be done with the Gyrus API that is not shown here.\n",
    "Check out the rest of the examples in this directory!"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}