{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b1e96a86",
   "metadata": {},
   "source": [
    "# Custom Distributions"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c6552205",
   "metadata": {},
   "source": [
    "If the predefined distributions in QDFlow are insufficient, you can create your\n",
    "own distributions. To do this, you will need to extend the `Distribution` class\n",
    "and implement the `draw()` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4af3292f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tutorial_helper\n",
    "from qdflow.util import distribution\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "555a0b67",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a custom distribution\n",
    "class ChiSquare(distribution.Distribution):\n",
    "\n",
    "    # The init method should store any parameters that define the distribution\n",
    "    def __init__(self, degrees_freedom):\n",
    "        self.df = degrees_freedom\n",
    "\n",
    "    # The draw() method should return a single value if size is None,\n",
    "    # or a numpy array with shape size, if size is an int or tuple of ints\n",
    "    def draw(self, rng, size=None):\n",
    "        if size is None:\n",
    "            return rng.chisquare(self.df)\n",
    "        else:\n",
    "            return rng.chisquare(self.df, size=size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "075e4b46",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAARYAAAD/CAYAAAA0cfmxAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjcsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvTLEjVAAAAAlwSFlzAAAPYQAAD2EBqD+naQAAFqRJREFUeJzt3QlsVFUbxvG3UFoQLVCWlmIpgrKGRUERA/ppkbKIEEgEQUBFUAK4sCiNsqmxCoKsQowiMQEEEiEGDIsogsgiKCIKCEoDCAVlX7RAO1/ek9zJTClQy5mZzr3/X3LT3plL506n83DuOWfeE+Pz+XwCABaVsvnDAIBgARAStFgAWEewALCOYAFgHcECwDqCBYB1sfZ/pPvk5+fL4cOH5ZZbbpGYmJhInw4QETrl7ezZs5KSkiKlSl27TUKwFIGGSmpqqq3XB4hqBw8elFtvvfWaxxAsRaAtFecXmpCQYOfVAaLMmTNnzH+wzvvhWgiWInAufzRUCBZ4XUwRugPovAVgHcECwDqCBYB1BAsAdwXLunXrpHPnzmZcXDuEli5dGnT/k08+aW4P3Nq3bx90zIkTJ6R3796mU7VixYrSv39/OXfuXNAxO3bskDZt2kjZsmVNr/aECRPC8vwAr4posJw/f16aNm0qM2fOvOoxGiRHjhzxbwsWLAi6X0Pll19+kdWrV8uyZctMWA0cODBoiKxdu3aSlpYm27Ztk4kTJ8q4cePkgw8+COlzA7wsosPNHTp0MNu1xMfHS3JycqH37dq1S1asWCHff/+9tGjRwtw2ffp06dixo7z77rumJTRv3jy5ePGizJkzR+Li4qRRo0ayfft2mTx5clAABcrNzTVbYDiFU61Ry697TPbbncJyLoAr+1jWrl0r1apVk3r16smgQYPk+PHj/vs2btxoLn+cUFFt27Y10403b97sP+b+++83oeLIyMiQPXv2yMmTJwt9zKysLKlQoYJ/Y9Yt4KJg0cugTz75RNasWSPvvPOOfPPNN6aFk5eXZ+7PyckxoRMoNjZWEhMTzX3OMUlJSUHHOPvOMQVlZmbK6dOn/ZvOuAVQdCV65m3Pnj393zdu3FiaNGkiderUMa2Y9PT0kD2uXn7pBsCFLZaCateuLVWqVJF9+/aZfe17OXbsWNAxly9fNiNFTr+Mfj169GjQMc7+1fpuAHgoWA4dOmT6WKpXr272W7VqJadOnTKjPY6vvvrKlDlo2bKl/xgdKbp06ZL/GB1B0j6bSpUqReBZAO4X0WDR+SY6QqOb2r9/v/n+wIED5r6RI0fKpk2bJDs72/SzdOnSRW6//XbT+aoaNGhg+mEGDBggW7ZskQ0bNsiQIUPMJZSOCKlevXqZjlud36LD0gsXLpSpU6fKsGHDIvnUAVeLaLBs3bpV7rzzTrMpfbPr92PGjJHSpUubiW2PPvqo1K1b1wRD8+bNZf369UH9HzqcXL9+fdPnosPMrVu3DpqjoqM6q1atMqGl/3748OHm519tqBnAjYthJcTr03ksGlA6QhSOsgnMY0G0vw+iqo8FQHQgWABYR7AAsI5gAWAdwQLAOoIFgLc+K4TiD0lTVgGRRIsFgHUECwDrCBYA1hEsAKwjWABYx6hQmBXlA4ZAtKPFAsA6ggWAdQQLAOsIFgDWESwArCNYAFhHsACwjmABYB3BAsA6ggWAdQQLAOsIFgDWESwArCNYAFhHsACwjnosLsXC8ogkWiwArCNYAFhHsACwjmABYB3BAsA6ggWAdQQLAOsIFgDWESwArCNYAFhHsACwjmABYB3BAsA6ggWAdQQLAOsIFgDWESwArCNYAFhHsACwjmABYB3BAsBdwbJu3Trp3LmzpKSkSExMjCxdujTofp/PJ2PGjJHq1atLuXLlpG3btrJ3796gY06cOCG9e/eWhIQEqVixovTv31/OnTsXdMyOHTukTZs2UrZsWUlNTZUJEyaE5fkBXhXRYDl//rw0bdpUZs6cWej9GgDTpk2T2bNny+bNm6V8+fKSkZEh//77r/8YDZVffvlFVq9eLcuWLTNhNXDgQP/9Z86ckXbt2klaWpps27ZNJk6cKOPGjZMPPvggLM8R8KIYnzYLSgBtsSxZskS6du1q9vW0tCUzfPhwGTFihLnt9OnTkpSUJHPnzpWePXvKrl27pGHDhvL9999LixYtzDErVqyQjh07yqFDh8y/nzVrlrz66quSk5MjcXFx5phRo0aZ1tHu3bsLPZfc3FyzBYaTtnT08bVlFOr1fsIl++1OkT4FRBF9H1SoUKFI74MS28eyf/9+EwZ6+ePQJ9WyZUvZuHGj2devevnjhIrS40uVKmVaOM4x999/vz9UlLZ69uzZIydPniz0sbOyssxjOZuGCoCiK7HBoqGitIUSSPed+/RrtWrVgu6PjY2VxMTEoGMK+xmBj1FQZmamSWVnO3jwoMVnBrgfS6wWIj4+3mwAXNZiSU5ONl+PHj0adLvuO/fp12PHjgXdf/nyZTNSFHhMYT8j8DEAeCRYbrvtNvPGX7NmTVDnkfadtGrVyuzr11OnTpnRHsdXX30l+fn5pi/GOUZHii5duuQ/RkeQ6tWrJ5UqVQrrcwK8IqLBovNNtm/fbjanw1a/P3DggBklevHFF+XNN9+Uzz//XH7++Wfp27evGelxRo4aNGgg7du3lwEDBsiWLVtkw4YNMmTIEDNipMepXr16mY5bnd+iw9ILFy6UqVOnyrBhwyL51AFXi2gfy9atW+XBBx/07ztv9n79+pkh5ZdfftnMddF5Kdoyad26tRlO1olujnnz5pkwSU9PN6NB3bt3N3NfHDqqs2rVKhk8eLA0b95cqlSpYibdBc51AeDSeSxuGb+/HuaxwAvvg2K1WGrXrm0mpVWuXDnodm1V3HXXXfLHH38U58cizIoSckyiQ9j6WLKzsyUvL++K23W26p9//lmsEwHgHv+pxaKdqI6VK1eaZpFDg0ZHcGrVqmX3DAG4O1ic0RgdsdEO1kBlypQxoTJp0iS7ZwjA3cGi80OcOSbax6IjLABgpfNW55sAgPV5LNqfoptOqXdaMo45c+YU98cC8GqwjB8/Xl5//XVTrkCru2mfCwDcULBoRTedGdunT5/i/HMALleseSwXL16U++67z/7ZAPBusDzzzDMyf/58+2cDwLuXQlrMWotRf/nll9KkSRMzhyXQ5MmTbZ0fAK8Eiy6n0axZM/P9zp07g+6jIxdAsYLl66+/5jcHIPoqyAHwWItFizNd65JHy0MC8K5iBYvTv+LQerJaUlL7Wwp+OBGA9xQrWN57771Cb9elSwuumwzAe6z2sTzxxBN8TgiA3WDR5UwDC10D8KZiXQp169YtaF/rcR85csRU3R89erStcwPgpWAJLEmpdNkNXQBMP/Hcrl07W+cGwEvB8vHHH9s/EwCucUMLlunSprt27TLfN2rUSO68805b5wXAa8GiVeN0GdO1a9dKxYoV/WsK6cS5Tz/9VKpWrWr7PAG4PViGDh0qZ8+eNWsh6/rJ6tdffzWT455//nlZsGCB7fNEhLCoGcIWLLp+spZMcEJFNWzYUGbOnOnpztuStHwqEHXzWLR4dsEaLEpvK1hYG4D3FCtYHnroIXnhhRfk8OHD/tt0adWXXnpJ0tPTbZ4fAK8Ey4wZM8zK87ryYZ06dcymi5jpbdOnT7d/lgDc38eSmpoqP/zwg+ln2b17t7lN+1vatm1r+/wAuL3FonVWtJNWWyZaj+Xhhx82I0S63X333WYuy/r160N3tgDcFyxTpkyRAQMGSEJCQqHT/J999lkKaQP4b8Hy008/Sfv27a96vw4162xcAN72n4Ll6NGjhQ4zO2JjY+Wvv/6ycV4AvBIsNWrUuGK5j4LLguhazgC87T8FS8eOHU29FV2wrKB//vlHxo4dK4888ojN8wPg9uHm1157TT777DOpW7euDBkyxNRgUTrkrNP58/Ly5NVXXw3VuQJwY7AkJSXJd999J4MGDZLMzExTOU7p0HNGRoYJFz0GgLf95wlyaWlp8sUXX8jJkydl3759JlzuuOMOqVSpUmjOEIB3Cj1pkOikOAAoiCVWAVhHsACwjmABYB3BAsA6ggWAdQQLAOsIFgAla8EyQLFECKKqxTJu3DjzcYHArX79+v779cOQgwcPlsqVK8vNN98s3bt3N6UdAh04cEA6deokN910k1SrVk1Gjhwply9fjsCzAbyjxLdYtNyl1tYNrPni0FUBli9fLosXLzYV7PSDkd26dZMNGzaY+/VDkRoqycnJ5jNOR44ckb59+5qaMm+99VZEng/gBSU+WDRINBgKOn36tHz00Ucyf/58sxyJs1i9FvXetGmT3HvvvbJq1SqzQqMGk344slmzZvLGG2/IK6+8YlpDcXFxEXhGgPuV6EshtXfvXklJSZHatWtL7969zaWN0hKYly5dCloZQC+TatasKRs3bjT7+rVx48ZBn7jWT2FrMXBdHvZqcnNzzTGBGwCXBEvLli1l7ty5ZknXWbNmyf79+6VNmzZm3eicnBzT4nAWpXdoiOh9Sr8WLOPg7DvHFCYrK8tcWjmbLncCwCWXQh06dPB/36RJExM0WrZh0aJFUq5cuZA9rtaaGTZsmH9fWyyEC+CSFktB2jrR6nVaB0b7XS5evCinTp0KOkZHhZw+Gf1acJTI2S+s38YRHx9vljgJ3AC4NFjOnTsnv//+uynY3bx5czO6s2bNGv/9e/bsMX0wrVq1Mvv69eeff5Zjx475j1m9erUJCl14DYAHL4VGjBghnTt3Npc/ugC9FusuXbq0PP7446bvo3///uaSJTEx0YSFrsioYaIjQs46Rxogffr0kQkTJph+Fa3bq3NftFUCwIPBcujQIRMix48fl6pVq0rr1q3NULJ+r9577z0pVaqUmRinIzk64vP+++/7/72G0LJly0yNXg2c8uXLS79+/eT111+P4LMC3C/G51TExlVp5622kHTuzLX6W4oytR2Fy367E78al7wPoq6PBUB0IFgAWEewALCOYAFgHcECwDqCBYB1BAsA6wgWANYRLACsI1gAWEewALCOYAFgHcECwDqCBYB1BAsA6wgWAN6qIAfvYP1nd6HFAsA6ggWAdQQLAOsIFgDWESwArCNYAFhHsACwjmABYB3BAsA6Zt4iajA7N3rQYgFgHcECwDqCBYB1BAsA6+i8havQwVsy0GIBYB3BAsA6ggWAdQQLAOvovIXn0MEberRYAFhHsACwjmABYB3BAsA6ggWAdQQLAOsIFgDWMY8FKARzXW4MwQKEKHyy3+7k2d8tl0IArCNYAFhHsACwjj4WIERqXacPxs39MJ4KlpkzZ8rEiRMlJydHmjZtKtOnT5d77rkn0qcFD6tVhPApipIWUJ65FFq4cKEMGzZMxo4dKz/88IMJloyMDDl27FikTw1wnRifz+cTD2jZsqXcfffdMmPGDLOfn58vqampMnToUBk1atQ1/+2ZM2ekQoUKcvr0aUlISAj5/z5AKNxoq6ao7wPPXApdvHhRtm3bJpmZmf7bSpUqJW3btpWNGzdecXxubq7ZHPqLdH6x15Kfe8HqeQM21Xxp8TXv3zk+45r3O3//RWmLeCJY/v77b8nLy5OkpKSg23V/9+7dVxyflZUl48ePv+J2beEAblVhStGOO3v2rGm5iNeD5b/Slo32xzj0sunEiRNSuXJliYmJuWqaa/AcPHjwus3EaMFzih5nwvD3py0VDZWUlJTrHuuJYKlSpYqULl1ajh49GnS77icnJ19xfHx8vNkCVaxYsUiPpS+qW4LFwXOKHgkh/vu7XkvFU6NCcXFx0rx5c1mzZk1QK0T3W7VqFdFzA9zIEy0WpZc2/fr1kxYtWpi5K1OmTJHz58/LU089FelTA1zHM8HSo0cP+euvv2TMmDFmglyzZs1kxYoVV3ToFpdeOukcmYKXUNGM5xQ94kvY359n5rEACB9P9LEACC+CBYB1BAsA6wgWANYRLJbKMdSqVUvKli1rPuy4ZcsWiWbjxo0zM4wDt/r160s0WbdunXTu3NnMEtXzX7p0adD9OmahI4TVq1eXcuXKmc+N7d27V6L9eT355JNXvHbt27cP+3kSLDfIreUYGjVqJEeOHPFv3377rUQTnaOkr4WGfmEmTJgg06ZNk9mzZ8vmzZulfPny5nX7999/JZqfl9IgCXztFixYIGGnw80ovnvuucc3ePBg/35eXp4vJSXFl5WVFbW/1rFjx/qaNm3qcwv9M1+yZIl/Pz8/35ecnOybOHGi/7ZTp0754uPjfQsWLPBF6/NS/fr183Xp0sUXabRYLJRj0GZ0UcoxRBO9LNDmdu3ataV3795y4MABcYv9+/ebSZKBr5t+BkYvY6P9dVNr166VatWqSb169WTQoEFy/PhxCTeCJUTlGPQPN1rpG2zu3LlmZvKsWbPMG7FNmzbmk61u4Lw2bnvdnMugTz75xHwO7p133pFvvvlGOnToYP5Ow8kzU/pRdPqH6GjSpIkJmrS0NFm0aJH079+fX2UJ1rNnT//3jRs3Nq9fnTp1TCsmPT09bOdBiyWM5RiilZaMqFu3ruzbt0/cwHlt3P66Kb2U1b/TcL92BMsN8Eo5hnPnzsnvv/9uhmbd4LbbbjMBEvi6aaEkHR1y0+umDh06ZPpYwv3acSl0g9xYjmHEiBFmroRe/hw+fNgMpWvL7PHHH5doCsPA/6W1n2j79u2SmJgoNWvWlBdffFHefPNNueOOO0zQjB492nRWd+3aVaL1eSUmJpqSqt27dzfBqf8ZvPzyy3L77bebofSwivSwlBtMnz7dV7NmTV9cXJwZft60aZMvmvXo0cNXvXp183xq1Khh9vft2+eLJl9//bUZji246XCsM+Q8evRoX1JSkhlmTk9P9+3Zs8cXzc/rwoULvnbt2vmqVq3qK1OmjC8tLc03YMAAX05OTtjPk7IJAKyjjwWAdQQLAOsIFgDWESwArCNYAFhHsACwjmABYB3BAsA6ggUl3v/+9z8zBR/Rg2BBSOlnjq5Wc3X9+vWmJuuOHTt4FVyGYEFIaf2W1atXm0/ZFvTxxx+bD29qzRC4C8GCkHrkkUekatWqpiJdwU/pLl682HyaWD81XaNGDbnppptMcaLrFX8urDq91owJfIyDBw/KY489Zm7XT/126dJFsrOzLT87XA3BgpCKjY2Vvn37mjd94DLhGipaLvGJJ54wNW2WL18uO3fulIEDB0qfPn1uaAmVS5cumTIBt9xyi7nc2rBhg9x8883mkkzrFCP0CBaE3NNPP21qg2j91cDLIK0bojVftP5Ls2bNTLWzoUOHmgDQMpg3siSLFtz68MMPTQuoQYMG5vG0ILiWaEToESwIOV3s7L777pM5c+aYfS1UpC0J7X/RVssbb7xhAkAvWbRlsXLlyhtaFeCnn34yj6EtFv15uunP1jWDNOAQelSQQ1hoiGhrRBfa0taDFnh+4IEHTCX5qVOnmsp7Gi66cJgOLV/rkkX7WAIvq5zLn8D+G728mjdv3hX/Vvt7EHoEC8JCO1JfeOEFmT9/vlmeQte70YDQ/g/tWNW+FqWXML/99ps0bNjwqj9Lw0FX+AtcA+nChQv+/bvuustcDunaOgkJCSF+ZigMl0IIC70c6dGjh2RmZppQ0DWGldac1eHo7777Tnbt2iXPPvvsFdXzC3rooYdkxowZ8uOPP8rWrVvlueeekzJlyvjv1wXWtDK9BpZecmldWO1bef755wsd9oZ9BAvCejl08uRJM2KjhavVa6+9ZloYepvOsNUi0NcraD1p0iRJTU01i6j16tXLdP7qULVDv9fF07Vodrdu3UznrT629rHQggkPat4CsI4WCwDrCBYA1hEsAKwjWABYR7AAsI5gAWAdwQLAOoIFgHUECwDrCBYA1hEsAMS2/wMspfAQmqzL7QAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 250x250 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Create an instance of the custom distribution\n",
    "custom_dist = ChiSquare(4)\n",
    "\n",
    "num_samples = 10000 # How many times to sample from each distribution\n",
    "\n",
    "rng = np.random.default_rng(seed=5)\n",
    "\n",
    "samples = custom_dist.draw(rng, num_samples)\n",
    "\n",
    "# Plot results\n",
    "fig, ax = plt.subplots(figsize=(2.5,2.5))\n",
    "tutorial_helper.plot_dist_data(fig, ax, samples, bins=20)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "qdflow_venv",
   "language": "python",
   "name": "qdflow_venv"
  },
  "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}