{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Field visualisation using `k3d`\n",
    "\n",
    "**Note:** If you experience any problems in plotting with `k3d`, please make sure you run the Jupyter notebook in Google Chrome.\n",
    "\n",
    "There are several ways how a field can be visualised, using:\n",
    "\n",
    "- `matplotlib`\n",
    "-`k3d`\n",
    "- `holoviews`\n",
    "- vtk-based libraries, e.g. `pyvista`\n",
    "\n",
    "\n",
    "`k3d` provides three-dimensional interactive plots of fields inside Jupyter notebook.\n",
    "\n",
    "Let us say we have a sample, which is an ellipsoid\n",
    "\n",
    "$$\\frac{x^2}{a^2} + \\frac{y^2}{b^2} + \\frac{z^2}{c^2} <= 1$$\n",
    "\n",
    "with $a=5\\,\\text{nm}$, $b=3\\,\\text{nm}$, and $c=2\\,\\text{nm}$. The space is discretised into cells with dimensions $(0.5\\,\\text{nm}, 0.5\\,\\text{nm}, 0.5\\,\\text{nm})$. The value of the field at $(x, y, z)$ point is $(-cy, cx, cz)$, with $c=10^{9}$. The norm of the field inside the cylinder is $10^{6}$.\n",
    "\n",
    "Let us first build that field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import discretisedfield as df\n",
    "\n",
    "a, b, c = 5e-9, 3e-9, 2e-9\n",
    "cell = (0.5e-9, 0.5e-9, 0.5e-9)\n",
    "\n",
    "mesh = df.Mesh(p1=(-a, -b, -c), p2=(a, b, c), cell=cell)\n",
    "\n",
    "\n",
    "def norm_fun(pos):\n",
    "    x, y, z = pos\n",
    "    if (x / a) ** 2 + (y / b) ** 2 + (z / c) ** 2 <= 1:\n",
    "        return 1e6\n",
    "    else:\n",
    "        return 0\n",
    "\n",
    "\n",
    "def value_fun(pos):\n",
    "    x, y, z = pos\n",
    "    c = 1e9\n",
    "    return (-c * y, c * x, c * z)\n",
    "\n",
    "\n",
    "field = df.Field(mesh, nvdim=3, value=value_fun, norm=norm_fun)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The most basic plot we can show is the plot of all the cells where the value is non-zero. This can be useful, to inspect the domain created, by plotting the norm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": [
     "nbval-ignore-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "40b18c1d5d9b4a708d5f2e655c42d90b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "field.norm.k3d.nonzero()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The plot is interactive, so it can be manipulated using a mouse. To change the color of voxels, we can pass the new color via `color` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": [
     "nbval-ignore-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ae2126d6ab034cff9d239df80daf1c42",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "field.norm.k3d.nonzero(color=0x27AE60)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we can plot a scalar field. For plotting a scalar field, we are using `discretisedfield.Field.k3d.scalar()` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "Cannot plot nvdim=3 field.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[4], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mfield\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mk3d\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mscalar\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/repos/ubermag-devtools/repos/discretisedfield/discretisedfield/plotting/k3d_field.py:248\u001b[0m, in \u001b[0;36mK3dField.scalar\u001b[0;34m(self, plot, filter_field, cmap, multiplier, interactive_field, **kwargs)\u001b[0m\n\u001b[1;32m    246\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdata\u001b[38;5;241m.\u001b[39mnvdim \u001b[38;5;241m!=\u001b[39m \u001b[38;5;241m1\u001b[39m:\n\u001b[1;32m    247\u001b[0m     msg \u001b[38;5;241m=\u001b[39m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCannot plot nvdim=\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdata\u001b[38;5;241m.\u001b[39mnvdim\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m field.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m--> 248\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(msg)\n\u001b[1;32m    250\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m plot \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    251\u001b[0m     plot \u001b[38;5;241m=\u001b[39m k3d\u001b[38;5;241m.\u001b[39mplot()\n",
      "\u001b[0;31mValueError\u001b[0m: Cannot plot nvdim=3 field."
     ]
    }
   ],
   "source": [
    "field.k3d.scalar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An exception was raised because we attempted to plot a vector field using voxels. Therefore, we first need to extract a component of the field. Let us plot the $x$ component."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": [
     "nbval-ignore-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "472a4cb6c87947d2b86833f64d8878a7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "field.x.k3d.scalar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, we can see that the points which we consider to be outside the sample are also plotted. This is because, `discretisedfield.Field.k3d.scalar()` method cannot determine the points where norm is zero from the passed scalar field. Therefore, we need to pass `field.norm` as the  `filter_field.`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": [
     "nbval-ignore-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "71ee39c7358943d4954350b37dab6465",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "field.x.k3d.scalar(filter_field=field.norm, multiplier=1e-6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By cascading operations, we can similarly plot the slice of the ellipsoid at $z=0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": [
     "nbval-ignore-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "00226357ee9842c4a34e970c16ad7adc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "field.sel(x=(0, 0)).orientation.z.k3d.scalar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To further modify the plot, keyword arguments for `k3d.factory.voxels()` function are accepted. Please refer to its [documentation](https://k3d-jupyter.org/k3d.html#k3d.factory.voxels)\n",
    "\n",
    "Next, we can plot the vector field itself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": [
     "nbval-ignore-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "99bd02cd14dc432ea603b1660bfaf1fb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "field.k3d.vector()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, points at the discretisation cell centres are plotted together with vectors to help understand the structure of the field. However, they can be deactivated by passing `points=False`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "tags": [
     "nbval-ignore-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4a536ddbbece46f297f1d17d53abc3bd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "field.k3d.vector(points=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is difficult to understand the vector field from this plot. By cascading, we can plot its slice at $x=0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "tags": [
     "nbval-ignore-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b9b6ccb57ac64a6f99843c986cb5de4d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "field.orientation.sel(x=(0, 0)).k3d.vector()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To improve the understanding of the plot, we can now colour the vectors plotted. For that, we need to pass a scalar field, according to which the vectors will be coloured."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "tags": [
     "nbval-ignore-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "78356f0bb6c34046a9fa8c7c8cce37c2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "field.orientation.sel(x=(0, 0)).k3d.vector(color_field=field.x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To further modify the plot, keyword arguments for `k3d.factory.vectors()` function are accepted. Plese refer to its [documentation](https://k3d-jupyter.org/k3d.html#k3d.factory.vectors).\n",
    "\n",
    "### Multiple visualisation on the same plot\n",
    "\n",
    "Sometimes, it is necessary to show, for example, multiple planes of the sample on the same plot. This can be done by exposing the `k3d.plot` and passing it to different plotting methods. For instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "tags": [
     "nbval-ignore-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "96ab93a13ae546e2a1bfc529859778ec",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import k3d\n",
    "\n",
    "plot = k3d.plot()\n",
    "field.sel(x=(-3e-9, -3e-9)).k3d.vector(plot=plot, color_field=field.z)\n",
    "field.sel(x=(0, 0)).k3d.vector(plot=plot, color_field=field.z, cmap=\"hsv\")\n",
    "field.sel(x=(3e-9, 3e-9)).k3d.vector(plot=plot, color_field=field.z)\n",
    "plot.display()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `k3d` interactive plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "tags": [
     "nbval-ignore-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "515d04747f6a461c9f347f7f2bc87b20",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "field.sel(z=(0, 0)).z.k3d.scalar(filter_field=field.norm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "tags": [
     "nbval-ignore-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "25968ff867ef4144b1676c94b732fb95",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(SelectionSlider(description='z (nm)', index=4, options=((-1.75, -1.7500000000000002e-09)…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "af9bf33d3e1345f380638909aaa994a3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import k3d\n",
    "\n",
    "plot1 = k3d.plot()\n",
    "\n",
    "\n",
    "@df.interact(z=field.mesh.slider(\"z\", continuous_update=True))\n",
    "def myplot(z):\n",
    "    field.z.sel(z=(z, z)).k3d.scalar(\n",
    "        filter_field=field.norm, interactive_field=field, plot=plot1\n",
    "    )\n",
    "\n",
    "\n",
    "plot1.display()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "tags": [
     "nbval-ignore-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "12c5e53809b44343a440fd137cd5d937",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(SelectionSlider(description='z (nm)', index=4, options=((-1.75, -1.7500000000000002e-09)…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b95f3d2b0dcf4a8b96952fde0a80efc2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot = k3d.plot()\n",
    "\n",
    "\n",
    "@df.interact(\n",
    "    z=field.mesh.slider(\"z\", continuous_update=True), axis=field.mesh.axis_selector()\n",
    ")\n",
    "def myplot(z, axis):\n",
    "    getattr(field, axis).sel(z=(z, z)).k3d.scalar(\n",
    "        filter_field=field.norm, interactive_field=field, plot=plot\n",
    "    )\n",
    "\n",
    "\n",
    "plot.display()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "tags": [
     "nbval-ignore-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a9cb230674b94ecc9373da7c5390ae19",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(SelectionSlider(description='y (nm)', index=6, options=((-2.75, -2.75e-09), (-2.25, -2.2…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8003f2309fa74ce8a5ec7606f126e13f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot = k3d.plot()\n",
    "\n",
    "\n",
    "@df.interact(\n",
    "    y=field.mesh.slider(\"y\", continuous_update=True), axis=field.mesh.axis_selector()\n",
    ")\n",
    "def myplot(y, axis):\n",
    "    field.sel(y=(y, y)).k3d.vector(\n",
    "        color_field=getattr(field, axis), interactive_field=field, plot=plot\n",
    "    )\n",
    "\n",
    "\n",
    "plot.display()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}