{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    " \n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "from cil.optimisation.algorithms import PDHG\n",
    "from cil.optimisation.utilities import Sampler, SamplerRandom\n",
    "import time "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0  5 10 15 20  1  6 11 16  2  7 12 17  3  8 13 18  4  9 14 19  0]\n"
     ]
    }
   ],
   "source": [
    "sampler=Sampler.staggered(21,5)\n",
    "print(sampler.get_samples(22))\n",
    "for _ in range(21):\n",
    "    sampler.next()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "44\n",
      "37\n",
      "40\n",
      "42\n",
      "46\n",
      "35\n",
      "10\n",
      "47\n",
      "3\n",
      "28\n",
      "9\n",
      "[44, 37, 40, 42, 46, 35, 10, 47, 3, 28, 9, 25, 11, 18, 43, 8, 41, 47, 42, 29, 35, 9, 4, 19, 34]\n",
      "Sampler that selects from a list of indices {0, 1, …, S-1}, where S is the number of indices. \n",
      "Type : from_function \n",
      "Current iteration number : 11 \n",
      "Number of indices : 50 \n",
      "Probability weights : [0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "def test_function(iteration_number):\n",
    "    if iteration_number<500:\n",
    "        #print(iteration_number)\n",
    "        np.random.seed(iteration_number)\n",
    "        return(np.random.choice(49,1)[0])\n",
    "    else:\n",
    "        np.random.seed(iteration_number)\n",
    "        return(np.random.choice(50,1)[0])\n",
    "        \n",
    "\n",
    "sampler=Sampler.from_function(50, test_function)\n",
    "for _ in range(11):\n",
    "    print(sampler.next())\n",
    "print(list(sampler.get_samples(25)))\n",
    "print(sampler)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n",
      "1\n",
      "1\n",
      "0\n",
      "0\n",
      "11\n",
      "5\n",
      "9\n",
      "8\n",
      "3\n",
      "1\n",
      "[1, 1, 1, 0, 0, 11, 5, 9, 8, 3, 1, 1, 1, 0, 0, 11, 5, 9, 8, 3, 1, 1, 1, 0, 0]\n",
      "Sampler that selects from a list of indices {0, 1, …, S-1}, where S is the number of indices. \n",
      "Type : from_function \n",
      "Current iteration number : 11 \n",
      "Number of indices : 13 \n",
      "Probability weights : [0.2, 0.3, 0.0, 0.1, 0.0, 0.1, 0.0, 0.0, 0.1, 0.1, 0.0, 0.1, 0.0] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "custom_list=[1,1,1,0,0,11,5,9,8,3]\n",
    "num_indices=13\n",
    "\n",
    "def test_function(iteration_number, custom_list=custom_list):\n",
    "    return(custom_list[iteration_number%len(custom_list)])\n",
    "\n",
    "#calculate prob weights \n",
    "temp_list = []\n",
    "for i in range(num_indices):\n",
    "    temp_list.append(custom_list.count(i))\n",
    "total = sum(temp_list)\n",
    "prob_weights = [x/total for x in temp_list]\n",
    "\n",
    "sampler=Sampler.from_function(num_indices=num_indices, function=test_function, prob_weights=prob_weights)\n",
    "for _ in range(11):\n",
    "    print(sampler.next())\n",
    "print(list(sampler.get_samples(25)))\n",
    "print(sampler)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9]\n",
      "0\n",
      "1\n",
      "2\n",
      "3\n",
      "4\n",
      "5\n",
      "6\n",
      "7\n",
      "8\n",
      "9\n",
      "0\n",
      "Sampler that selects from a list of indices {0, 1, …, S-1}, where S is the number of indices. \n",
      "Type : sequential \n",
      "Current iteration number : 12 \n",
      "Number of indices : 10 \n",
      "Probability weights : [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "sampler=Sampler.sequential(10)\n",
    "print(sampler.get_samples(20))\n",
    "for _ in range(11):\n",
    "    print(sampler.next())\n",
    "\n",
    "next(sampler)\n",
    "np.array([0 ,1 ,2 ,3, 4, 5, 6 ,7, 8, 9, 0 ,1,2 ,3 ,4, 5, 6 ,7, 8, 9])\n",
    "print(sampler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2, 5, 0, 1, 4, 3, 6, 1, 6, 0, 4, 2, 3, 5, 5, 6, 2, 4, 0, 1, 3, 0, 2, 6, 3]\n"
     ]
    }
   ],
   "source": [
    "sampler=Sampler.random_without_replacement(7, seed=1)\n",
    "print(list(sampler.get_samples(25)))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0  8  4 12  2 10  6 14  1  9  5 13  3 11  7 15]\n",
      "0\n",
      "8\n",
      "4\n",
      "12\n",
      "2\n",
      "10\n",
      "6\n",
      "14\n",
      "1\n",
      "9\n",
      "5\n",
      "13\n",
      "3\n",
      "11\n",
      "7\n",
      "Sampler that selects from a list of indices {0, 1, …, S-1}, where S is the number of indices. \n",
      "Type : herman_meyer \n",
      "Current iteration number : 15 \n",
      "Number of indices : 16 \n",
      "Probability weights : [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "sampler=Sampler.herman_meyer(16)\n",
    "print(sampler.get_samples(16))\n",
    "for _ in range(15):\n",
    "    print(next(sampler))\n",
    "\n",
    "\n",
    "print(sampler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n",
      "4\n",
      "1\n",
      "4\n",
      "2\n",
      "3\n",
      "3\n",
      "2\n",
      "1\n",
      "0\n",
      "0\n",
      "3\n",
      "2\n",
      "0\n",
      "4\n",
      "1\n",
      "2\n",
      "1\n",
      "3\n",
      "2\n",
      "2\n",
      "1\n",
      "1\n",
      "1\n",
      "1\n",
      "0\n",
      "0\n",
      "4\n",
      "4\n",
      "0\n",
      "[1, 4, 1, 4, 2, 3, 3, 2, 1, 0, 0, 3, 2, 0, 4, 1, 2, 1, 3, 2, 2, 1, 1, 1, 1]\n",
      "Sampler that selects from a list of indices {0, 1, …, S-1}, where S is the number of indices. \n",
      "Type : random_with_replacement \n",
      "Current iteration number : 30 \n",
      "Number of indices : 5 \n",
      "Probability weights : [0.2, 0.2, 0.2, 0.2, 0.2] \n",
      "Seed : 5 \n",
      "\n"
     ]
    }
   ],
   "source": [
    "sampler=Sampler.random_with_replacement(5, seed=5)\n",
    "for _ in range(30):\n",
    "    print(sampler.next())\n",
    "print(list(sampler.get_samples(25)))\n",
    "print(sampler)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0, 2, 0, 3, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "Sampler that selects from a list of indices {0, 1, …, S-1}, where S is the number of indices. \n",
      "Type : random_with_replacement \n",
      "Current iteration number : 0 \n",
      "Number of indices : 4 \n",
      "Probability weights : [0.7, 0.1, 0.1, 0.1] \n",
      "Seed : 5 \n",
      "\n"
     ]
    }
   ],
   "source": [
    "sampler=Sampler.random_with_replacement(4, [0.7,0.1,0.1,0.1], seed=5)\n",
    "print(list(sampler.get_samples(25)))\n",
    "print(sampler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0  4  8 12 16]\n",
      "0\n",
      "4\n",
      "8\n",
      "12\n",
      "16\n",
      "20\n",
      "1\n",
      "5\n",
      "9\n",
      "13\n",
      "17\n",
      "2\n",
      "6\n",
      "10\n",
      "14\n",
      "[ 0  4  8 12 16]\n",
      "Sampler that selects from a list of indices {0, 1, …, S-1}, where S is the number of indices. \n",
      "Type : staggered \n",
      "Current iteration number : 15 \n",
      "Number of indices : 21 \n",
      "Probability weights : [0.047619047619047616, 0.047619047619047616, 0.047619047619047616, 0.047619047619047616, 0.047619047619047616, 0.047619047619047616, 0.047619047619047616, 0.047619047619047616, 0.047619047619047616, 0.047619047619047616, 0.047619047619047616, 0.047619047619047616, 0.047619047619047616, 0.047619047619047616, 0.047619047619047616, 0.047619047619047616, 0.047619047619047616, 0.047619047619047616, 0.047619047619047616, 0.047619047619047616, 0.047619047619047616] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "sampler=Sampler.staggered(21,4)\n",
    "print(sampler.get_samples(5))\n",
    "for _ in range(15):\n",
    "    print(sampler.next())\n",
    "print(sampler.get_samples(5))\n",
    "\n",
    "print(sampler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0 1 3 2 1 3 2 0 2 3 0 1 0 3 1 2 1 0 2 3 0 1 3 2 2 3 0 1 0 1 3 2 3 1 2 0 2\n",
      " 3 1 0 0 3 1 2 1 3 2 0 2 3 0 1 2 0 3 1 2 0 3 1 1 2 0 3 2 1 0 3 0 3 2 1 1 3\n",
      " 2 0 3 0 1 2 4 3 2 1 0 3 4 2 1 0 4 1 0 3 2 4 2 0 1 3 3 0 4 2 1 4 0 2 1 3 2\n",
      " 4 0 3 1 3 1 0 4 2]\n"
     ]
    }
   ],
   "source": [
    "#Custom sampler example\n",
    "\n",
    "class DelayedRegularisationSampler(SamplerRandom):\n",
    "    def function(self, iteration_number):\n",
    "        \"\"\" For each iteration number this function samples from a randomly generated list in order. Every num_indices the list is re-created. For the first aproximately 20*(num_indices -1) iterations the last index is never called.  \"\"\"\n",
    "        if iteration_number < 20*(self.num_indices - 1):\n",
    "            location = iteration_number % (self.num_indices - 1)\n",
    "            if location == 0:\n",
    "                self._sampling_list = self._generator.choice(self.num_indices-1, self.num_indices - 1, p=[\n",
    "                                                             1/(self.num_indices-1)]*(self.num_indices-1), replace=self.replace)\n",
    "        else:\n",
    "            location = iteration_number % self.num_indices\n",
    "            if location == 0:\n",
    "                self._sampling_list = self._generator.choice(\n",
    "                    self.num_indices, self.num_indices, p=self.prob_weights, replace=self.replace)\n",
    "        out = self._sampling_list[location]\n",
    "        return out\n",
    "\n",
    "    def __init__(self, num_indices,  seed=None, replace=False, prob=None,  sampling_type='Mantid Sampler'):\n",
    "        super(DelayedRegularisationSampler, self).__init__(\n",
    "            num_indices,  seed, replace, prob,  sampling_type)\n",
    "        \n",
    "sampler=DelayedRegularisationSampler(5)\n",
    "print(sampler.get_samples(120))\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cil",
   "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"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}