{ "cells": [ { "cell_type": "code", "execution_count": null, "id": "d4270b45", "metadata": { "tags": [ "remove-input" ] }, "outputs": [], "source": [ "%config InlineBackend.figure_formats = ['svg']\n", "\n", "import numpy as np\n", "import matplotlib.pyplot as plt" ] }, { "cell_type": "markdown", "id": "476bd7bb", "metadata": {}, "source": [ "(extending)=\n", "\n", "# Extending flatspin\n", "\n", "There are many ways to extend the functionality of flatspin.\n", "Here we discuss two of the most common use cases, namely custom geometries and custom encoders." ] }, { "cell_type": "markdown", "id": "be7d1f05", "metadata": {}, "source": [ "## Custom geometries\n", "\n", "There are two ways to extend flatspin with your own custom geometries:\n", "1. Provide a set of spin positions and angles to {class}`CustomSpinIce `\n", "2. Extend {class}`SpinIce ` and create a parameterized geometry" ] }, { "cell_type": "markdown", "id": "3e328848", "metadata": {}, "source": [ "### Using {class}`CustomSpinIce `\n", "\n", "{class}`CustomSpinIce ` can be used to quickly create a custom geometry.\n", "The {class}`CustomSpinIce ` class accepts a list of positions and angles for all the spins as the parameters `magnet_coords` and `magnet_angles`.\n", "\n", "Below we create a geometry on a square lattice in which the spin angles depend directly on their positions.\n", "The parameter `delta_angle` scales the amount of rotation per lattice spacing." ] }, { "cell_type": "code", "execution_count": null, "id": "c51b36e4", "metadata": {}, "outputs": [], "source": [ "from flatspin.model import CustomSpinIce\n", "\n", "# Size (cols, rows) of our geometry\n", "size = (10, 10)\n", "\n", "# Positions of spins\n", "lattice_spacing = 1\n", "x = lattice_spacing * np.arange(0, size[0])\n", "y = lattice_spacing * np.arange(0, size[1])\n", "xx, yy = np.meshgrid(x, y)\n", "xx = xx.ravel()\n", "yy = yy.ravel()\n", "pos = np.column_stack([xx, yy])\n", "\n", "# Angles of spins\n", "delta_angle = 10\n", "angle = (xx+yy) * delta_angle / lattice_spacing\n", "\n", "# Give the angles and positions to CustomSpinIce\n", "model = CustomSpinIce(magnet_coords=pos, magnet_angles=angle, radians=False)\n", "model.plot();" ] }, { "cell_type": "markdown", "id": "8587147f", "metadata": {}, "source": [ "While {class}`CustomSpinIce ` is one way of creating a custom geometry, it is not parametric.\n", "In other words, any modifications to the geometry must be made manually outside of the class.\n", "Consequently, it is cumbersome to explore variations of this geometry using, e.g., [`flatspin-run-sweep`](flatspin-run-sweep).\n", "In the next section, we will see how to extend flatspin with a new {class}`SpinIce ` class." ] }, { "cell_type": "markdown", "id": "ef4b9ade", "metadata": {}, "source": [ "### Extending SpinIce\n", "\n", "Fully parametric geometries can be created by creating a subclass of {class}`SpinIce `.\n", "Any new parameters should be introduced as keyword arguments to the `__init__` function of the subclass.\n", "The subclass should override {func}`_init_geometry() `, which should return a tuple `(pos, angle)` where `pos` is an array with the positions of the spins, and `angle` is an array with the rotations of the spins.\n", "\n", "Below we create a new subclass that provides a fully parametric version of the geometry we created earlier.\n", "We introduce a new parameter `delta_angle`, while `size` and `lattice_spacing` are already defined by the {class}`SpinIce ` base class." ] }, { "cell_type": "code", "execution_count": null, "id": "83487851", "metadata": {}, "outputs": [], "source": [ "from flatspin.model import SpinIce\n", "\n", "class MySpinIce(SpinIce):\n", " def __init__(self, *, delta_angle=10, **kwargs):\n", " self.delta_angle = delta_angle\n", "\n", " super().__init__(**kwargs)\n", "\n", " def _init_geometry(self):\n", " # size and lattice_spacing are SpinIce parameters\n", " size = self.size\n", " lattice_spacing = self.lattice_spacing\n", "\n", " # positions of spins\n", " x = lattice_spacing * np.arange(0, size[0])\n", " y = lattice_spacing * np.arange(0, size[1])\n", " xx, yy = np.meshgrid(x, y)\n", " xx = xx.ravel()\n", " yy = yy.ravel()\n", " pos = np.column_stack([xx, yy])\n", "\n", " # angles of spins\n", " delta_angle = np.deg2rad(self.delta_angle)\n", " angle = (xx+yy) * delta_angle / lattice_spacing\n", "\n", " # Generate labels for our geometry (optional)\n", " #self.labels = grid\n", "\n", " return pos, angle\n", "\n", " # The size of vertices in our geometry (optional)\n", " _vertex_size = (2, 2)" ] }, { "cell_type": "markdown", "id": "7db9847f", "metadata": {}, "source": [ "With our new `MySpinIce` class, we are ready to explore the parameter space:" ] }, { "cell_type": "code", "execution_count": null, "id": "aa443b3f", "metadata": {}, "outputs": [], "source": [ "for i, delta_angle in enumerate([0, 30, 60, 90]):\n", " model = MySpinIce(size=(10,10), delta_angle=delta_angle)\n", " plt.subplot(1, 4, i+1)\n", " plt.title(f\"{delta_angle}\")\n", " plt.axis('off')\n", " model.plot()" ] }, { "cell_type": "markdown", "id": "a72529a4", "metadata": {}, "source": [ "## Custom encoders\n", "\n", "An [encoder](encoders) translates logical input to an external field protocol.\n", "\n", "Input takes the form of arrays of shape `(n_inputs, input_dim)`.\n", "1D input arrays may be used as a shorthand for `(n_inputs, 1)`.\n", "\n", "The encoding process consists of one or more steps, where the output of one step is input to the next step:\n", "\n", "`input -> step1 -> step2 -> ... -> h_ext`\n", "\n", "In general, signals can take any shape as part of the encoding process.\n", "However, the last step must produce an output of either:\n", "\n", "1. `(time, 2)`: a global vector signal\n", "2. `(time, H, W, 2)`: a local vector signal on a grid\n", "\n", "Each step is a simple function taking a single `input` argument, and any number of parameters as keyword arguments:\n", "\n", "```python\n", "def step(input, param1=default1, param2=default2, ...):\n", " ...\n", "```\n", "\n", "```{note}\n", "The only non-keyword argument to a step function is `input`.\n", "Parameters are only allowed as keyword arguments, and **must** have default values.\n", "```\n", "\n", "The {class}`Encoder ` will inspect the signature of each step to discover the available parameters.\n", "The parameters can then be set during encoder initialization, or afterwards via {func}`set_params() `.\n", "Note that parameter names may overlap, in which case all matching parameters will be set to the same value.\n", "\n", "Custom encoders can be created by subclassing {class}`Encoder ` and provide a list of `steps`.\n", "\n", "Below we create a custom encoder where:\n", "1. Input is encoded as the amplitude of a global external field\n", "2. For each input, the angle of the field is incremented by a fixed amount `delta_angle`" ] }, { "cell_type": "code", "execution_count": null, "id": "1c8c1401", "metadata": {}, "outputs": [], "source": [ "from flatspin.encoder import Encoder\n", "\n", "def scale_step(input, H=1):\n", " return H * input\n", "\n", "def rotate_step(input, delta_angle=15):\n", " n_inputs = len(input)\n", " angles = np.arange(0, delta_angle * n_inputs, delta_angle)\n", " angles = np.deg2rad(angles)\n", " h_ext = input * np.column_stack([np.cos(angles), np.sin(angles)])\n", " return h_ext\n", "\n", "class MyEncoder(Encoder):\n", " steps = [scale_step, rotate_step]" ] }, { "cell_type": "markdown", "id": "7cc6c5d2", "metadata": {}, "source": [ "The two steps (1) and (2) are implemented by the functions `scale_step` and `rotate_step`, respectively.\n", "The steps are tied together in the new `MyEncoder` class." ] }, { "cell_type": "code", "execution_count": null, "id": "9db81820", "metadata": {}, "outputs": [], "source": [ "# Encoder automatically discovers the available parameters from the kwargs of the steps\n", "encoder = MyEncoder()\n", "print(encoder.get_params())" ] }, { "cell_type": "code", "execution_count": null, "id": "4797f2a1", "metadata": {}, "outputs": [], "source": [ "# Linear input from 0..1\n", "input = np.linspace(0, 1, 50, endpoint=False)\n", "h_ext = encoder(input)\n", "\n", "# Scatter plot of h_ext, where color indicates time\n", "plt.title('h_ext')\n", "plt.scatter(h_ext[:,0], h_ext[:,1], c=np.arange(len(h_ext)), marker='.', cmap='plasma')\n", "plt.axis('equal')\n", "plt.colorbar(label='time');" ] }, { "cell_type": "code", "execution_count": null, "id": "711feb2c", "metadata": {}, "outputs": [], "source": [ "# Four periods of a sine wave, scaled to the range 0.5..1\n", "input = np.linspace(0, 1, 360, endpoint=False)\n", "input = np.sin(-np.pi/2 + 8*np.pi*input)\n", "input = 1/2 + input/2\n", "\n", "plt.figure()\n", "plt.title('input')\n", "plt.plot(input)\n", "\n", "encoder.set_params(delta_angle=360/len(input))\n", "h_ext = encoder(input)\n", "\n", "plt.figure()\n", "plt.title('h_ext')\n", "plt.scatter(h_ext[:,0], h_ext[:,1], c=np.arange(len(h_ext)), marker='.', cmap='plasma')\n", "plt.axis('equal')\n", "plt.colorbar(label='time');" ] }, { "cell_type": "markdown", "id": "316f6792", "metadata": {}, "source": [ "The {mod}`flatspin.encoder` module contains a range of useful encoder steps.\n", "In fact, it already includes a step called {func}`scale ` which is functionally equivalent to our custom `scale_step` above, but with an additional parameter `H0` to specify an offset, so that the input is scaled from `H0..H`." ] }, { "cell_type": "code", "execution_count": null, "id": "ae8f1e56", "metadata": {}, "outputs": [], "source": [ "from flatspin.encoder import scale\n", "\n", "class MyEncoder2(Encoder):\n", " steps = [scale, rotate_step]\n", "\n", "encoder2 = MyEncoder2(delta_angle=360/len(input), H0=0.5, H=1.0)\n", "h_ext = encoder2(input)\n", "\n", "plt.title('h_ext')\n", "plt.scatter(h_ext[:,0], h_ext[:,1], c=np.arange(len(h_ext)), marker='.', cmap='plasma')\n", "plt.axis('equal')\n", "plt.colorbar(label='time');" ] }, { "cell_type": "markdown", "id": "daf567da", "metadata": {}, "source": [ "## Using custom models and encoders from the command-line\n", "\n", "The [command-line tools](cmdline) [`flatspin-run`](flatspin-run) and [`flatspin-run-sweep`](flatspin-run-sweep) support the use of custom models and encoders.\n", "\n", "To use your own model class, simply provide the full module path to `-m/--model`.\n", "Similarly, to use your own encoder class, provide the full module path to `-e/--encoder`.\n", "Any custom parameters can be set as usual with `-p/--param`.\n", "\n", "For example, placing the `MySpinIce` class in a file `mymodels.py`, and `MyEncoder` in a file `myencoders.py`, we can do:\n", "\n", "```bash\n", "flatspin-run -m mymodels.MySpinIce -p delta_angle=30 ... -e myencoders.MyEncoder -p [TODO]\n", "```" ] } ], "metadata": { "jupytext": { "formats": "md:myst", "text_representation": { "extension": ".md", "format_name": "myst", "format_version": 0.13, "jupytext_version": "1.11.5" } }, "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.10.6" }, "source_map": [ 15, 22, 31, 39, 49, 71, 78, 89, 122, 126, 133, 174, 189, 194, 200, 212, 230, 235, 248 ] }, "nbformat": 4, "nbformat_minor": 5 }