{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "be58fc76",
   "metadata": {},
   "source": [
    "# Goniometer geometry\n",
    "We can investigate how to convert to and from the lab and translated/rotated reference frames here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9488be37",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "\n",
    "import anri"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c68dfbfa",
   "metadata": {},
   "source": [
    "Let's get some useful basis vectors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce03a8d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "e1, e2, e3 = jnp.eye(3)\n",
    "e1, e2, e3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b25ee8d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "vecs = jnp.stack((e1, e2, e3))\n",
    "vecs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcdc56aa",
   "metadata": {},
   "source": [
    "And we'll whip up a quick function to plot some 3D vectors in the lab frame.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3eff7856",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_vectors_3d(vec, origin=None, legend=''):\n",
    "    if origin is None:\n",
    "        origin = jnp.zeros_like(vec)\n",
    "    ax = plt.figure(figsize=(8,8)).add_subplot(projection='3d',)\n",
    "    ax.set_proj_type('ortho')\n",
    "    ax.quiver([0, 0, 0], [0, 0, 0], [0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1], color='k', label='Lab basis vectors')\n",
    "    uv = vec - origin\n",
    "    ax.quiver(origin[:, 0], origin[:, 1], origin[:, 2],\n",
    "    uv[:, 0], uv[:, 1], uv[:, 2], color=list('rgbcmyk'[:len(vec)]), label=legend)\n",
    "    ax.set_xlim(-1, 1)\n",
    "    ax.set_ylim(-1, 1)\n",
    "    ax.set_zlim(-1, 1)\n",
    "    ax.set_aspect('equal')\n",
    "    ax.set_box_aspect((1,1,1))\n",
    "    ax.set(xlabel='X', ylabel='Y', zlabel='Z', title='Static lab frame')\n",
    "\n",
    "    ax.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50e5a2a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_vectors_3d(vecs, None, '')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdb2fda7",
   "metadata": {},
   "source": [
    "Now let's investigate how it looks when we rotate the goniometer. \n",
    "We can take the basis vectors for the sample (rotating) frame, and map them into the lab frame:\n",
    "\n",
    "$R_z(\\phi) \\cdot v_{\\text{sample}} = v_{\\text{lab}}$\n",
    "\n",
    "To apply this transform to many vectors at once, we need to vmap the core function over an extra axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11863d77",
   "metadata": {},
   "outputs": [],
   "source": [
    "# broadcast over (npks, 3)\n",
    "sample_to_lab_vec = jax.vmap(anri.geom.sample_to_lab, in_axes=(0, None, None, None, None, None))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0143f8fa",
   "metadata": {},
   "source": [
    "We can try a small omega rotation (10 degrees) to see how it'll look:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ded7b81",
   "metadata": {},
   "outputs": [],
   "source": [
    "omega = 20\n",
    "origin_lab = sample_to_lab_vec(jnp.zeros_like(vecs), omega, 0, 0, 0, 0)\n",
    "vec_lab = sample_to_lab_vec(vecs, omega, 0, 0, 0, 0)\n",
    "plot_vectors_3d(vec_lab,origin_lab, fr'$\\omega={omega}$')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "776dbe07",
   "metadata": {},
   "source": [
    "How about the wedge and omega angles?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e71cddf4",
   "metadata": {},
   "outputs": [],
   "source": [
    "omega = 0\n",
    "wedge = 20\n",
    "origin_lab = sample_to_lab_vec(jnp.zeros_like(vecs), omega, wedge, 0, 0, 0)\n",
    "vec_lab = sample_to_lab_vec(vecs, omega, wedge, 0, 0, 0)\n",
    "plot_vectors_3d(vec_lab,origin_lab, f'Wedge={wedge}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b16f4edd",
   "metadata": {},
   "outputs": [],
   "source": [
    "omega = 0\n",
    "wedge = 0\n",
    "chi = 20\n",
    "origin_lab = sample_to_lab_vec(jnp.zeros_like(vecs), omega, wedge, chi, 0, 0)\n",
    "vec_lab = sample_to_lab_vec(vecs, omega, wedge, chi, 0, 0)\n",
    "plot_vectors_3d(vec_lab,origin_lab, fr'$\\chi={chi}$')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee738597",
   "metadata": {},
   "source": [
    "We can also investigate the effect of changing dty, which translates the whole diffractometer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d6bdfee",
   "metadata": {},
   "outputs": [],
   "source": [
    "omega = 0\n",
    "wedge = 0\n",
    "chi = 0\n",
    "dty = -0.5\n",
    "origin_lab = sample_to_lab_vec(jnp.zeros_like(vecs), omega, wedge, chi, dty, 0)\n",
    "vec_lab = sample_to_lab_vec(vecs, omega, wedge, chi, dty, 0)\n",
    "plot_vectors_3d(vec_lab,origin_lab, f'dty={dty}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c8a08a1",
   "metadata": {},
   "source": [
    "As $y_0$ is the motor value of `dty` when the rotation axis hits the beam, we can investigate what happens when it's non-zero:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a5f3396",
   "metadata": {},
   "outputs": [],
   "source": [
    "omega = 0\n",
    "wedge = 0\n",
    "chi = 0\n",
    "dty = 0\n",
    "y0 = 0.5\n",
    "origin_lab = sample_to_lab_vec(jnp.zeros_like(vecs), omega, wedge, chi, dty, y0)\n",
    "vec_lab = sample_to_lab_vec(vecs, omega, wedge, chi, dty, y0)\n",
    "plot_vectors_3d(vec_lab,origin_lab, f'$y_0={y0}$, dty={dty}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0de6ada2",
   "metadata": {},
   "outputs": [],
   "source": [
    "omega = 0\n",
    "wedge = 0\n",
    "chi = 0\n",
    "dty = 0.5\n",
    "y0 = 0.5\n",
    "origin_lab = sample_to_lab_vec(jnp.zeros_like(vecs), omega, wedge, chi, dty, y0)\n",
    "vec_lab = sample_to_lab_vec(vecs, omega, wedge, chi, dty, y0)\n",
    "plot_vectors_3d(vec_lab,origin_lab, f'$y_0={y0}$, dty={dty}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edc3573a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "anri",
   "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.14.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}