generic_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "from dataclasses import dataclass, field\n",
    "from typing import Union, Tuple, Iterable, Optional, Any, Type, Callable, Dict\n",
    "from tqdm import tqdm\n",
    "from fancy_einsum import einsum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mpl.style.use('seaborn-v0_8')\n",
    "mpl.rcParams['figure.figsize'] = (15,10)\n",
    "fontsize = 20\n",
    "mpl.rcParams['font.size'] = fontsize\n",
    "mpl.rcParams['xtick.labelsize'] = fontsize\n",
    "mpl.rcParams['ytick.labelsize'] = fontsize\n",
    "mpl.rcParams['legend.fontsize'] = fontsize\n",
    "mpl.rcParams['axes.titlesize'] = fontsize\n",
    "mpl.rcParams['axes.labelsize'] = fontsize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Disease:\n",
    "    def __init__(self, beta: float, gamma: float, delta: float):\n",
    "        self.beta = beta\n",
    "        self.gamma = gamma\n",
    "        self.delta = delta\n",
    "        self.doubling_time, self.growth_rate, self.decay_rate = self.effective_growth_rate()\n",
    "\n",
    "    def effective_growth_rate(self):\n",
    "        '''\n",
    "        Returns the effective growth rate for the epidemic in the exponential phase\n",
    "        '''\n",
    "        b = self.beta\n",
    "        g = self.gamma\n",
    "        d = self.delta\n",
    "        growth_rate = 0.5 * (np.sqrt((d-g)**2 +4*d*b) - d - g)\n",
    "        doubling_time = np.log(2)/growth_rate\n",
    "        decay_rate = -0.5 * (d + g + np.sqrt((d-g)**2 +4*d*b))\n",
    "        return doubling_time, growth_rate, decay_rate\n",
    "    \n",
    "    def __str__(self):\n",
    "        return str(vars(self))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "covid = Disease(0.6,1/6,1/5.5)\n",
    "measles = Disease(2,1/10,1/8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def find_first_index(condition, axis = -1):\n",
    "    \"\"\"\n",
    "    Finds the index of the first time a condition is met for each row in a 2D numpy array.\n",
    "\n",
    "    Args:\n",
    "        array: A 2D numpy array with n rows and m columns.\n",
    "        condition: A boolean condition to be checked element-wise.\n",
    "\n",
    "    Returns:\n",
    "        A 1D numpy array containing the index of the first time the condition is met for each row in the input array.\n",
    "        Returns -1 if the condition is never met in a row.\n",
    "    \"\"\"\n",
    "    assert isinstance(condition, np.ndarray), \"array must be a numpy array\"\n",
    "    assert len(condition.shape) == 2, \"array must be a 2D numpy array\"\n",
    "    assert condition.shape[1] > 0, \"array must have at least one column\"\n",
    "\n",
    "    # Apply the condition element-wise to the input array and get the index of the first True value along axis 1\n",
    "    idx = np.argmax(condition, axis=axis)\n",
    "    \n",
    "    # If the condition is never met in a row, set the corresponding index to -1\n",
    "    idx[np.all(~condition, axis=axis)] = -1\n",
    "    \n",
    "    return idx\n",
    "\n",
    "def homogeneous_LR_matrix(n_groups: int, LR: float):\n",
    "    out = np.identity(n_groups,np.float64)\n",
    "    out[out == 0] = LR\n",
    "    return out\n",
    "\n",
    "def day_to_sim_step(day, delta_t):\n",
    "    #day should start at zero\n",
    "    steps_in_day = int(1/delta_t)\n",
    "    return np.round(steps_in_day * (day + 0.5)).astype(int)\n",
    "\n",
    "def sim_step_to_day(sim_step, delta_t):\n",
    "    return np.floor(sim_step * delta_t).astype(int)\n",
    "\n",
    "def quartiles(data: np.ndarray, axis: int = -1, squeeze = True):\n",
    "    'returns lower_quartile, median, upper_quartile'\n",
    "    return tuple(np.quantile(data, [0.25,0.5,0.75],axis=axis).squeeze())\n",
    "\n",
    "def list_to_str(l: Iterable, sep: str = ''):\n",
    "    out = ''\n",
    "    for i in l:\n",
    "        out+= str(i) + sep\n",
    "    return out[:-len(sep)]\n",
    "\n",
    "def dict_to_str(d: dict, sep: str = ', '):\n",
    "    out = ''\n",
    "    for key, val in d.items():\n",
    "        assert isinstance(val, list) or isinstance(val, str)\n",
    "        val_str = list_to_str(val,sep) if isinstance(val, list) else val\n",
    "        out += f'{key}: {val_str}{sep}'\n",
    "    return out[:-len(sep)]\n",
    "\n",
    "def force_iterable(input):\n",
    "    return input if isinstance(input, Iterable) else [input]\n",
    "\n",
    "def force_list(input):\n",
    "    return list(input) if isinstance(input, Iterable) and not isinstance(input,str) else [input]\n",
    "\n",
    "def show_fig(fig, figsavename):\n",
    "    if figsavename is not None:\n",
    "        fig.savefig(figsavename)\n",
    "\n",
    "def total_daily(array: np.ndarray, delta_t: float, axis: int = -1):\n",
    "    #the timeseries axis must be the last axis!\n",
    "    assert axis == -1, 'the timeseries axis must be the last axis!'\n",
    "    window = int(1/delta_t)\n",
    "    shape = array.shape\n",
    "    assert (len(shape) == 2) or (len(shape) == 1)\n",
    "    assert shape[-1] % window == 0\n",
    "    return array.reshape(shape[:-1] + (shape[-1]//window,window)).sum(axis=-1)\n",
    "\n",
    "def mean_daily(array: np.ndarray, delta_t: float, axis: int = -1):\n",
    "    #the timeseries axis must be the last axis!\n",
    "    assert axis == -1, 'the timeseries axis must be the last axis!'\n",
    "    window = int(1/delta_t)\n",
    "    shape = array.shape\n",
    "    assert (len(shape) == 2) or (len(shape) == 1)\n",
    "    assert shape[-1] % window == 0\n",
    "    return array.reshape(shape[:-1] + (shape[-1]//window,window)).mean(axis=-1)\n",
    "     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class City:\n",
    "    def __init__(self,\n",
    "                 N0s: np.ndarray,\n",
    "                 groups: list[str],\n",
    "                 compartments: str = 'SEIR',\n",
    "                 group_LR: Optional[np.ndarray] = None):\n",
    "        \"\"\"\n",
    "        Initializes the City object.\n",
    "\n",
    "        Args:\n",
    "        - N0s (np.ndarray): An array of initial population sizes for each group.\n",
    "        - groups (list): A list representing the groups in the model.\n",
    "        - compartments (str): A string representing the compartment model used in the model. Currently only 'SEIR' is supported.\n",
    "        - group_LR (np.ndarray): An optional array representing the matrix of inter-group interaction likelihood ratios for groups i and j. All diagonal elements should be one.\n",
    "\n",
    "        Returns:\n",
    "        - None\n",
    "        \"\"\"\n",
    "        self.n_groups = len(groups)\n",
    "        self.groups = groups\n",
    "\n",
    "        # Check that the length of N0s matches the number of groups\n",
    "        assert len(N0s) == self.n_groups, \"N0s should have length n_groups\"\n",
    "\n",
    "        self.N0s = N0s\n",
    "        self.N0 = self.N0s.sum()\n",
    "\n",
    "        assert compartments == 'SEIR', \"Only SEIR compartment model supported at the moment\"\n",
    "        self.model_type = compartments\n",
    "        self.compartments = [l for l in compartments]\n",
    "\n",
    "        # Check and set the group_LR matrix\n",
    "        if group_LR is None:\n",
    "            group_LR = np.identity(self.n_groups)\n",
    "        assert group_LR.shape == (self.n_groups, self.n_groups)\n",
    "        assert np.all(group_LR.diagonal() == 1)\n",
    "        assert np.all(group_LR == group_LR.T)\n",
    "        self.group_LR = group_LR\n",
    "\n",
    "        # Set the indices of the compartments\n",
    "        self.S_index = self.compartments.index('S')\n",
    "        self.E_index = self.compartments.index('E')\n",
    "        self.I_index = self.compartments.index('I')\n",
    "        self.R_index = self.compartments.index('R')\n",
    "\n",
    "        self.name: Optional[Union[int,str]] = None\n",
    "        self.disease: Optional[Disease] = None\n",
    "\n",
    "        self.data_cols = {'municipal': 'blue', 'arrivals': 'red', 'departures': 'green'}\n",
    "        self.error_bar_cols = {'municipal': 'cornflowerblue', 'arrivals': 'lightcoral', 'departures': 'mediumaquamarine'}\n",
    "\n",
    "    def reset_parameters(self, I0: int = 0, n_sims: int = 1, simulation_steps: int = 100):\n",
    "        \"\"\"\n",
    "        Sets self.municipal to an array of zeros of size (n_groups, n_compartments, n_sims, timesteps) for each compartment in the city\n",
    "        The only values which aren't zero are the initial value of I which is I0 and the initial value of S which is N-I0\n",
    "        The default value of I0 is the value set in the class initialization, but can be overridden\n",
    "        \"\"\"\n",
    "        # if I0 is None:\n",
    "        #     I0 = self.I0\n",
    "        self.I0 = I0\n",
    "        self.municipal = np.zeros((self.n_groups,len(self.compartments),self.n_sims, self.simulation_steps),dtype = np.int64)\n",
    "        self.arrivals = np.zeros((self.n_groups,len(self.compartments),self.n_sims, self.simulation_steps),dtype = np.int64)\n",
    "        self.departures = np.zeros((self.n_groups,len(self.compartments),self.n_sims, self.simulation_steps),dtype = np.int64)\n",
    "        self.community_infections = np.zeros((self.n_groups, self.n_sims, self.simulation_steps), dtype=np.int64)\n",
    "        self.initial_conditions()\n",
    "    \n",
    "    def initial_conditions(self):\n",
    "        raise ModuleNotFoundError\n",
    "    \n",
    "    def multiple_sims(self, delta_t: float, epidemic_time: Union[int,float], disease: Disease, I0: int = 0, n_sims: int = 100):\n",
    "        assert (1 / delta_t) % 1 == 0, \"1/delta_t must be an integer\"\n",
    "        self.n_sims = n_sims\n",
    "        self.delta_t = delta_t\n",
    "        self.epidemic_time = epidemic_time\n",
    "        p_recovery = 1 - np.exp( - delta_t * disease.gamma)\n",
    "        p_infectious = 1 - np.exp( - delta_t * disease.delta)\n",
    "        self.simulation_steps = int(epidemic_time // delta_t) + 1\n",
    "        self.times: np.ndarray = np.linspace(0, epidemic_time, self.simulation_steps)\n",
    "        self.scaled_times: np.ndarray = self.times /disease.doubling_time\n",
    "        self.reset_parameters(I0, n_sims, self.simulation_steps)\n",
    "        for sim_step in tqdm(range(1,self.simulation_steps)):\n",
    "            self.step_internal(disease.beta,\n",
    "                      delta_t,\n",
    "                      p_infectious,\n",
    "                      p_recovery,\n",
    "                      sim_step)\n",
    "        self.daily_flight_data()\n",
    "    \n",
    "    def __call__(self,\n",
    "                 delta_t: float,\n",
    "                 epidemic_time: Union[int,float],\n",
    "                 disease: Disease,\n",
    "                 I0: int = 0,\n",
    "                 n_sims: int = 100):\n",
    "        return self.multiple_sims(delta_t, epidemic_time, disease, I0, n_sims)\n",
    "\n",
    "    def step_internal(self,\n",
    "                      beta:float,\n",
    "                      delta_t: float,\n",
    "                      p_infectious: float,\n",
    "                      p_recovery: float,\n",
    "                      simulation_step: int):\n",
    "        N = self.municipal[...,simulation_step-1].sum(axis = 1)\n",
    "        \n",
    "        S = self.municipal[:,self.S_index,:,simulation_step-1]\n",
    "        E = self.municipal[:,self.E_index,:,simulation_step-1]\n",
    "        I = self.municipal[:,self.I_index,:,simulation_step-1]\n",
    "        R = self.municipal[:,self.R_index,:,simulation_step-1]\n",
    "\n",
    "\n",
    "        modified_I = einsum('group1 group2, group2 n_sims -> group1 n_sims', self.group_LR, I)\n",
    "        modified_N = einsum('group1 group2, group2 n_sims -> group1 n_sims', self.group_LR, N)\n",
    "        exposure_rate = beta * modified_I/modified_N\n",
    "        p_exposure = 1 - np.exp(- delta_t * exposure_rate)\n",
    "        n_exposed = np.random.binomial(S, p_exposure)\n",
    "        n_infectious = np.random.binomial(E,p_infectious)\n",
    "        n_recovered = np.random.binomial(I,p_recovery)\n",
    "\n",
    "        self.municipal[:,self.S_index,:,simulation_step] = S - n_exposed\n",
    "        self.municipal[:,self.E_index,:,simulation_step] = E + n_exposed - n_infectious\n",
    "        self.municipal[:,self.I_index,:,simulation_step] = I + n_infectious - n_recovered\n",
    "        self.municipal[:,self.R_index,:,simulation_step] = R + n_recovered\n",
    "        self.community_infections[..., simulation_step] = self.community_infections[..., simulation_step-1] + n_exposed\n",
    "\n",
    "    def daily_flight_data(self, moving_avg = False):\n",
    "        window = int(1/self.delta_t)\n",
    "        \n",
    "        if moving_avg:\n",
    "            kernel = np.zeros(2*self.simulation_steps-1)\n",
    "            kernel[self.simulation_steps-1:self.simulation_steps+window-1] = 1\n",
    "            strided_kernel = np.flip(np.lib.stride_tricks.sliding_window_view(kernel, self.simulation_steps).astype(int).copy(),axis=0)\n",
    "            outputs = [np.zeros_like(self.arrivals)] * 2\n",
    "            for i,data in tqdm(enumerate([self.arrivals, self.departures])):\n",
    "                output = einsum('groups compartments sims simsteps, kernelsteps simsteps -> groups compartments sims kernelsteps', data, strided_kernel)\n",
    "                # if i == 0:\n",
    "                #     print('data', data)\n",
    "                #     print('kernel', strided_kernel)\n",
    "                #     print('output', output)\n",
    "                outputs[i] = output\n",
    "            self.arrivals_moving_avg, self.departures_moving_avg = outputs[0], outputs[1]\n",
    "            \n",
    "        self.arrivals_daily_avg = self.arrivals.reshape((self.n_groups, len(self.compartments), self.n_sims, self.simulation_steps//window, window)).sum(axis = -1)\n",
    "        self.departures_daily_avg = self.departures.reshape((self.n_groups, len(self.compartments), self.n_sims, self.simulation_steps//window, window)).sum(axis = -1)\n",
    "\n",
    "    def select_travellers(self, daily_mixnumber: int, simulation_step: int):\n",
    "        raise ModuleNotFoundError\n",
    "\n",
    "    def plot_sims(self,\n",
    "                  times: Optional[np.ndarray] = None,\n",
    "                  cityname: Union[int,str] = 0, \n",
    "                  shift_index: Optional[np.ndarray] = None,\n",
    "                  separate_groups: bool = False,\n",
    "                  figsavename: Optional[str] = None,\n",
    "                  moving_avg = False,\n",
    "                  log = ''):\n",
    "        include_flight_data = 'arrivals' in dir(self)\n",
    "        subplots = self.n_groups if separate_groups else 1\n",
    "        if include_flight_data:\n",
    "            fig, axs = plt.subplots(3,subplots, figsize = (20,30))\n",
    "            axs = np.array(axs)\n",
    "        else:\n",
    "            fig, axs = plt.subplots(1,subplots, figsize = (20,10))\n",
    "            axs = np.expand_dims(np.array(axs),0)\n",
    "        \n",
    "        for ax in axs:\n",
    "            if 'x' in log:\n",
    "                ax.set_xscale('log')\n",
    "            if 'y' in log:\n",
    "                ax.set_yscale('log')\n",
    "        \n",
    "        if self.n_groups == 1:\n",
    "            axs = np.expand_dims(axs, -1)\n",
    "        if times is None:\n",
    "            times = self.times\n",
    "        if shift_index is None:\n",
    "            shift = np.zeros((self.n_sims,1))\n",
    "        else:\n",
    "            shift = times[shift_index].reshape((self.n_sims,1))\n",
    "        days = np.array(range(int(max(times))), dtype = np.float64)\n",
    "\n",
    "        times = times - shift + shift.mean()\n",
    "        days = days - shift + shift.mean()\n",
    "        \n",
    "        travel_times = times if moving_avg else days\n",
    "        arrivals = self.arrivals_moving_avg if moving_avg else self.arrivals_daily_avg\n",
    "        departures = self.departures_moving_avg if moving_avg else self.departures_daily_avg\n",
    "\n",
    "        cols = ['green', 'orange', 'red', 'blue']\n",
    "        labels = ['Municipal', 'Arrivals', 'Departures']\n",
    "        if separate_groups:\n",
    "            for i, group in enumerate(self.groups):\n",
    "                for j,compartment in tqdm(enumerate(self.compartments)):\n",
    "                    axs[0,i].plot(times[0],\n",
    "                                self.municipal[i,j,0],\n",
    "                                label = compartment,\n",
    "                                color = cols[j])\n",
    "                    if include_flight_data:\n",
    "                        axs[1,i].plot(travel_times[0],  \n",
    "                                    arrivals[i,j,0],\n",
    "                                    label = compartment,\n",
    "                                    color = cols[j])\n",
    "                        axs[2,i].plot(travel_times[0],\n",
    "                                    departures[i,j,0],\n",
    "                                    label = compartment,\n",
    "                                    color = cols[j])\n",
    "                    for k,datum in enumerate(self.municipal[i,j,1:]):\n",
    "                        axs[0,i].plot(times[k+1], datum, color = cols[j])\n",
    "                    if include_flight_data:\n",
    "                        for k, (arrival, departure) in enumerate(zip(arrivals[i,j,1:],departures[i,j,1:])):\n",
    "                            axs[1,i].plot(travel_times[k+1], arrival, color = cols[j])\n",
    "                            axs[2,i].plot(travel_times[k+1], departure, color = cols[j])\n",
    "                num_subplots = axs.shape[0]\n",
    "                for j in range(num_subplots):\n",
    "                    axs[j,i].legend()\n",
    "                    axs[j,i].set_title(f\"City {cityname}, {group}: {labels[j]}\")\n",
    "        else:\n",
    "            for j, compartment in tqdm(enumerate(self.compartments)):\n",
    "                axs[0].plot(times[0],\n",
    "                            self.municipal[:,j,0].sum(axis = 0),\n",
    "                            label = compartment,\n",
    "                            color = cols[j])\n",
    "                if include_flight_data:\n",
    "                    axs[1].plot(travel_times[0],\n",
    "                                arrivals[:,j,0].sum(axis = 0),\n",
    "                                label = compartment,\n",
    "                                color = cols[j])\n",
    "                    axs[2].plot(travel_times[0],\n",
    "                                departures[:,j,0].sum(axis = 0),\n",
    "                                label = compartment,\n",
    "                                color = cols[j])\n",
    "                for k,datum in enumerate(self.municipal[:,j,1:].sum(axis = 0)):\n",
    "                    axs[0].plot(times[k+1], datum, color = cols[j])\n",
    "                    if include_flight_data:\n",
    "                        for k, (arrival, departure) in enumerate(zip(arrivals[:,j,1:].sum(axis = 0),departures[:,j,1:].sum(axis = 0))):\n",
    "                            axs[1].plot(travel_times[k+1], arrival, color = cols[j])\n",
    "                            axs[2].plot(travel_times[k+1], departure, color = cols[j])\n",
    "            num_subplots = axs.shape[0]\n",
    "            for j in range(num_subplots):\n",
    "                axs[j].legend()\n",
    "                axs[j].set_title(f\"City {cityname}: {labels[j]}\")\n",
    "        show_fig(fig,figsavename)\n",
    "    \n",
    "    def peak_I_times(self):\n",
    "        return self.municipal[:,self.I_index].sum(axis = 0).argmax(axis = -1)\n",
    "\n",
    "    def __str__(self):\n",
    "        raise ModuleNotFoundError\n",
    "    \n",
    "    def sim_steps_from_times(self,times: Union[float,np.ndarray]):\n",
    "        return np.argmin(np.abs(self.times-np.expand_dims(times,-1)),axis=-1)\n",
    "    \n",
    "    def times_from_sim_steps(self,sim_steps: Union[int,np.ndarray]):\n",
    "        return sim_steps * self.delta_t\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FrequentFlyerCity(City):\n",
    "    def __init__(self,\n",
    "                 N0: int = 10**6,\n",
    "                 frequent_flyer_frac: float = 0.1,\n",
    "                 p_ff: Optional[float] = None,\n",
    "                 flying_LR: Optional[float] = None,\n",
    "                 group_LR: float = 5,\n",
    "                 compartments: str = 'SEIR'):\n",
    "        group_LR_matrix = homogeneous_LR_matrix(2,1/group_LR)\n",
    "        self.groups = ['normal', 'frequent_flyers']\n",
    "        self.N0s = np.array([N0 * (1 - frequent_flyer_frac), N0 * frequent_flyer_frac],dtype = np.int64)\n",
    "        super().__init__(self.N0s, self.groups, compartments, group_LR_matrix)\n",
    "        self.frequent_flyer_frac = frequent_flyer_frac\n",
    "        if (p_ff is None) and (flying_LR is None):\n",
    "            flying_LR = 10\n",
    "        assert (p_ff is None) != (flying_LR is None), \"Specify exactly one of p_ff OR flying_LR!\"\n",
    "        if flying_LR is not None:\n",
    "            self.flying_LR = flying_LR\n",
    "            self.p_ff = flying_LR * frequent_flyer_frac/ (flying_LR * frequent_flyer_frac + (1 - frequent_flyer_frac))\n",
    "        if p_ff is not None:\n",
    "            self.p_ff = p_ff\n",
    "            self.flying_LR = (p_ff / frequent_flyer_frac) / ((1-p_ff) / (1 - frequent_flyer_frac))\n",
    "\n",
    "    def initial_conditions(self):\n",
    "        I_n = np.random.binomial(self.I0, 1-self.frequent_flyer_frac, self.n_sims)\n",
    "        I_ff = self.I0 - I_n\n",
    "        self.municipal[0,2,:,0] = I_n\n",
    "        self.municipal[1,2,:,0] = I_ff\n",
    "        self.municipal[0,0,:,0] = self.N0s[0] - I_n\n",
    "        self.municipal[1,0,:,0] = self.N0s[1] - I_ff\n",
    "    \n",
    "    def select_travellers(self, daily_mixnumber: int, simulation_step: int):\n",
    "        \"Selects travellers based on a constant rate per person. This means number of travellers isn't preserved so the rate is modified to provide a force that pushes things back to the baseline population\"\n",
    "        # N = []\n",
    "        fractions = np.array([1 - self.p_ff, self.p_ff])\n",
    "        p_travel = 1 - np.exp( - self.delta_t * daily_mixnumber * fractions / self.N0s)\n",
    "        return np.random.binomial(self.municipal[...,simulation_step], p_travel.reshape((2,1,1)))\n",
    "\n",
    "    def __str__(self):\n",
    "        out = 'City Type:\\n FrequentFlyerCity'\n",
    "        out += f'\\nN:\\n {self.N0}'\n",
    "        out += f'\\nfrequent_flyer_frac:\\n {self.frequent_flyer_frac}'\n",
    "        out += f'\\nflying_LR:\\n {self.flying_LR}'\n",
    "        out += f'\\np_ff:\\n {self.p_ff}'\n",
    "        out += f'\\ngroup_LR:\\n {self.group_LR}'\n",
    "        return out\n",
    "\n",
    "    # def select_travellers(self, daily_mixnumber: int, simulation_step: int):\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BasicCity(City):\n",
    "    def __init__(self,\n",
    "                 N0: int = 10**6,\n",
    "                 compartments: str = 'SEIR'):\n",
    "        group_LR_matrix = np.array([[1.]])\n",
    "        self.groups = ['normal']\n",
    "        self.N0s = np.array([N0],dtype = np.int64)\n",
    "        super().__init__(self.N0s, self.groups, compartments, group_LR_matrix)\n",
    "    \n",
    "    def initial_conditions(self):\n",
    "        self.municipal[0,2,:,0] = self.I0\n",
    "        self.municipal[0,0,:,0] = self.N0s[0] - self.I0\n",
    "    \n",
    "    def __str__(self):\n",
    "        out = 'City Type:\\n BasicCity'\n",
    "        out += f'\\nN:\\n {self.N0}'\n",
    "        return out\n",
    "\n",
    "    def select_travellers(self, daily_mixnumber: int, simulation_step: int):\n",
    "        p_travel = 1 - np.exp(- self.delta_t * daily_mixnumber / self.N0)\n",
    "        return np.random.binomial(self.municipal[...,simulation_step],p_travel)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def default_axis_order():\n",
    "    return ['cities', 'datatypes', 'groups', 'compartments', 'sims', 'times']\n",
    "\n",
    "@dataclass\n",
    "class SimData:\n",
    "    array: np.ndarray\n",
    "    values_present: Dict\n",
    "    axis_order: list[str] = field(default_factory=default_axis_order)\n",
    "\n",
    "    def __getitem__(self,to_keep):\n",
    "        assert isinstance(self.values_present,dict)\n",
    "        values_to_keep = self.values_present.copy()\n",
    "        if isinstance(to_keep, list) or isinstance(to_keep, tuple):\n",
    "            for category, elemtype in zip(to_keep, values_to_keep.keys()):\n",
    "                values_to_keep[elemtype] = category\n",
    "        elif isinstance(to_keep, dict):\n",
    "            for key, value in to_keep.items():\n",
    "                values_to_keep[key] = value\n",
    "\n",
    "        new_all_values = {}\n",
    "        for key,value in self.values_present.items():\n",
    "            new_all_values[key] = []\n",
    "            for elem in value:\n",
    "                if elem in values_to_keep[key]:\n",
    "                    new_all_values[key].append(elem)\n",
    "\n",
    "        chosen_indices = {}\n",
    "        for key in self.values_present.keys():\n",
    "            chosen_indices[key] = np.zeros(len(self.values_present[key])).astype(bool)\n",
    "            for i,value in enumerate(self.values_present[key]):\n",
    "                if value in values_to_keep[key]:\n",
    "                    chosen_indices[key][i] = True\n",
    "\n",
    "            for value in values_to_keep[key]:\n",
    "                assert value in self.values_present[key], f'{value} is not a member of {key}, which currently only contains {self.values_present[key]}'\n",
    "        out_array = self.choose_array(chosen_indices)\n",
    "        return SimData(out_array, values_present = new_all_values, axis_order=self.axis_order)\n",
    "    \n",
    "    def __eq__(self,other):\n",
    "        array_bool = np.all(self.array - other.array == 0)\n",
    "        label_bool = self.values_present == other.all_labels\n",
    "        return array_bool and label_bool\n",
    "    \n",
    "    def filter(self, to_keep):\n",
    "        return self.__getitem__(to_keep)\n",
    "    \n",
    "    def choose_array(self, chosen_indices: dict):\n",
    "        out_array = self.array\n",
    "        for key, value in chosen_indices.items():\n",
    "            axis = self.axis_order.index(key)\n",
    "            out_array = np.compress(value, out_array, axis = axis)\n",
    "        return out_array\n",
    "    \n",
    "    def wrap_np_function(self,\n",
    "                         np_function: Callable,\n",
    "                         axis: Union[int, str, Iterable] = 0,\n",
    "                         keepdims: bool = False,\n",
    "                         SimData_out = True,\n",
    "                         **kwargs):\n",
    "        axislist = force_list(axis)\n",
    "        axisnums = [self.axis_order.index(ax) if isinstance(ax,str) else ax for ax in axislist]\n",
    "\n",
    "        if not SimData_out:\n",
    "            return np_function(self.array, axis = tuple(axisnums), keepdims = keepdims, **kwargs)\n",
    "        \n",
    "        axistypes = [ax if isinstance(ax,str) else self.axis_order[ax] for ax in axislist]\n",
    "        new_axes = self.axis_order if keepdims else [ax for ax in self.axis_order if ax not in axistypes]\n",
    "        new_labels = self.values_present.copy()\n",
    "        for ax in axistypes:\n",
    "            if not keepdims:\n",
    "                del new_labels[ax]\n",
    "            else:\n",
    "                new_labels[ax] = ['N/A']\n",
    "        new_arr = np_function(self.array, axis = tuple(axisnums), keepdims = keepdims, **kwargs)\n",
    "        return SimData(new_arr, new_labels, new_axes)\n",
    "    \n",
    "    def wrap_np_binary_operator(self, np_function: Callable, other, SimData_out: bool = True):\n",
    "        if isinstance(other, SimData):\n",
    "            assert self.axis_order == other.axis_order\n",
    "            assert self.values_present == other.values_present\n",
    "            other_arr = other.array\n",
    "        else:\n",
    "            other_arr = other\n",
    "        out_arr = np_function(self.array, other_arr)\n",
    "        return SimData(out_arr,self.values_present, self.axis_order) if SimData_out else out_arr\n",
    "\n",
    "    def sum(self, axis: Union[int, str, Iterable] = 0, keepdims: bool = False, SimData_out: bool = True):\n",
    "        return self.wrap_np_function(np.sum,axis,keepdims, SimData_out)\n",
    "    \n",
    "    def mean(self, axis: Union[int, str, Iterable] = 0, keepdims: bool = False, SimData_out: bool = True):\n",
    "        return self.wrap_np_function(np.mean,axis,keepdims, SimData_out)\n",
    "    \n",
    "    def std(self, axis: Union[int, str, Iterable] = 0, keepdims: bool = False, SimData_out: bool = True):\n",
    "        return self.wrap_np_function(np.std,axis,keepdims, SimData_out)\n",
    "    \n",
    "    def quartiles(self, axis: Union[int, str, Iterable] = 0, keepdims: bool = False, SimData_out: bool = True):\n",
    "        out = []\n",
    "        for i,quartile in enumerate([0.25,0.5,0.75]):\n",
    "            out.append(self.wrap_np_function(np.quantile,axis,keepdims, SimData_out, q=quartile))\n",
    "        return tuple(out)\n",
    "    \n",
    "    def __add__(self, other):\n",
    "        return self.wrap_np_binary_operator(np.add, other)\n",
    "    \n",
    "    def __truediv__(self, other):\n",
    "        return self.wrap_np_binary_operator(np.divide, other)\n",
    "    \n",
    "    def max(self, axis: Union[int, str, Iterable] = 0, keepdims: bool = False, SimData_out: bool = True):\n",
    "        return self.wrap_np_function(np.max,axis,keepdims, SimData_out)\n",
    "    \n",
    "    def argmax(self, axis: Union[int, str, Iterable] = 0, keepdims: bool = False, SimData_out: bool = True):\n",
    "        return self.wrap_np_function(np.argmax, axis, keepdims, SimData_out)\n",
    "    \n",
    "    def daily_avg(self):\n",
    "        old_times = self.values_present['times']\n",
    "        delta_t = old_times[1] - old_times[0]\n",
    "        window = round(1/delta_t)\n",
    "        \n",
    "        new_shape = list(self.array.shape)\n",
    "        new_shape[self.axis_order.index('times')] //= window\n",
    "        new_shape = tuple(new_shape+[window])\n",
    "\n",
    "        new_arr = self.array.reshape(new_shape).sum(axis = -1)\n",
    "        new_labels = self.values_present.copy()\n",
    "        new_labels['times'] = np.array(old_times).reshape((len(old_times)//window,window)).mean(axis=-1).tolist()\n",
    "        return SimData(new_arr, new_labels, self.axis_order)\n",
    "\n",
    "def plotprep(log = '',\n",
    "             n_figs: int = 1):\n",
    "    if n_figs > 1:\n",
    "        fig, ax = plt.subplots(n_figs,1, figsize = (15,10*n_figs))\n",
    "    else:\n",
    "        fig, ax = plt.subplots(1,1)\n",
    "    axs = force_iterable(ax)\n",
    "    assert fig is not None\n",
    "    for ax in axs:\n",
    "        if 'x' in log:\n",
    "            ax.set_xscale('log')\n",
    "        if 'y' in log:\n",
    "            ax.set_yscale('log')\n",
    "    return fig, axs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def hists_at_time(data: SimData,\n",
    "                  time: float,\n",
    "                  filters: list[dict],\n",
    "                  log: str = '',\n",
    "                  legend_labels: Optional[list[str]] = None,\n",
    "                  bins = 30,\n",
    "                  density = False,\n",
    "                  alpha: float = 0.5,\n",
    "                  figsavename: Optional[str] = None):\n",
    "\n",
    "    fig, axs = plotprep(log)\n",
    "    if legend_labels is None:\n",
    "        for i, filter in enumerate(filters):\n",
    "            print(f'Dataset {i+1}: {filter}')\n",
    "        legend_labels = [f'Dataset {i+1}' for i in range(len(filters))]\n",
    "    assert len(legend_labels) == len(filters)\n",
    "    summed_axes = tuple([i for i,axistype in enumerate(data.axis_order) if axistype != 'sims'])\n",
    "    time = min(data.values_present['times'], key=lambda x:abs(x-time))\n",
    "    for i,filter in enumerate(filters):\n",
    "        filter['times'] = [time]\n",
    "        values = data[filter].array.sum(axis=summed_axes)\n",
    "        \n",
    "        total_filter = filter.copy()\n",
    "        del total_filter['compartments']\n",
    "        total = data[total_filter].array.sum(axis=summed_axes)    \n",
    "        values = values/total\n",
    "        \n",
    "        axs[0].hist(values, bins = bins, label = legend_labels[i], density = density, alpha=alpha)\n",
    "    axs[0].set_title(f'Distributions at time {round(time)} days')\n",
    "    axs[0].set_xlabel('Number of People in Category')\n",
    "    ylabel = 'Probability Density' if density else 'Frequency'\n",
    "    axs[0].set_ylabel(ylabel)\n",
    "    axs[0].legend()\n",
    "    show_fig(fig, figsavename)\n",
    "\n",
    "def plot_avg_vals(datasets: Union[SimData, dict[str,SimData]],\n",
    "                  filters: list[dict],\n",
    "                  x_axis_type: str = 'times',\n",
    "                  log: str = '',\n",
    "                  error_bars: str = 'std',\n",
    "                  filter_labels: Optional[list[str]] = None,\n",
    "                  doubling_time: Optional[float] = None,\n",
    "                  figsavename: Optional[str] =None):\n",
    "    assert x_axis_type in ['times', 'scaled_times', 'total_infections']\n",
    "    assert error_bars in ['std', 'IQR', 'None']\n",
    "    fig, axs = plotprep(log)\n",
    "    ax = axs[0]\n",
    "    # if legend_labels is None:\n",
    "    #     legend_labels = [f'Dataset {i+1}' for i in range(len(filters))]\n",
    "    #     for i, filter in enumerate(filters):\n",
    "    #         print(f'Dataset {i+1}: {filter}')\n",
    "    caption = ''\n",
    "    if filter_labels is None:\n",
    "        filter_labels = [str(x) for x in range(1,len(filters)+1)]\n",
    "\n",
    "    is_dict = isinstance(datasets, dict)\n",
    "    if is_dict:\n",
    "        dataset_labels = list(datasets.keys())\n",
    "        datalist = list(datasets.values())\n",
    "    else:\n",
    "        dataset_labels = ['']\n",
    "        datalist = [datasets]\n",
    "\n",
    "    assert len(filter_labels) == len(filters)\n",
    "    assert len(dataset_labels) == len(datalist)\n",
    "    \n",
    "    for x, filter in enumerate(filters):\n",
    "        str_filter = dict_to_str(filter)\n",
    "        caption += f'Filter {x+1}: {str_filter}\\n'\n",
    "    if is_dict:\n",
    "        caption += '\\n'\n",
    "        for x, dataset_label in enumerate(dataset_labels):\n",
    "            caption += f'Dataset {x+1}: {dataset_label}\\n'\n",
    "        caption = caption[:-1]\n",
    "        \n",
    "\n",
    "\n",
    "    colors = ['blue', 'red', 'green', 'fuchsia', 'dimgrey', 'yellow','darkviolet', 'darkorange']\n",
    "    ecolors = ['cornflowerblue', 'lightcoral', 'palegreen', 'lightpink', 'lightgrey', 'lemonchiffon', 'thistle', 'navajowhite']\n",
    "    linecount, count = 1,1\n",
    "    total_lines = len(filters) * len(datalist)\n",
    "    for d, dataset in enumerate(datalist):\n",
    "        if x_axis_type == 'total_infections':\n",
    "            total_infections = dataset[create_filter(datatypes=['municipal'],compartments=['E','I','R'])]\n",
    "            total_infections = total_infections.sum(axis = ('cities','datatypes','groups','compartments'))\n",
    "            x_axis = total_infections.mean(axis='sims').array\n",
    "            x_std = total_infections.std(axis='sims').array\n",
    "        elif x_axis_type == 'times':\n",
    "            x_axis = np.array(dataset.values_present['times'])\n",
    "            x_std = 0\n",
    "        else:\n",
    "            assert isinstance(doubling_time, float), 'Specify Doubling time to use scaled_times!'\n",
    "            x_axis = np.array(dataset.values_present['times']) / doubling_time\n",
    "            x_std = 0\n",
    "\n",
    "        for f, filter in enumerate(filters):\n",
    "            values = dataset[filter].sum(axis = ('cities','datatypes','groups','compartments'))\n",
    "\n",
    "            total_filter = filter.copy()\n",
    "            del total_filter['compartments']\n",
    "            total = dataset[total_filter].sum(axis = ('cities','datatypes','groups','compartments'))\n",
    "            values = values/ (total + 1e-10)\n",
    "\n",
    "            assert isinstance(values, SimData)\n",
    "            means = values.mean(axis = 'sims').array\n",
    "            stds = values.std(axis = 'sims').array\n",
    "            lower_quartile, median, upper_quartile = values.quartiles(axis='sims',SimData_out=False)\n",
    "\n",
    "            label = f'Dataset {d+1}, Filter {f+1}'\n",
    "            if error_bars == 'std':\n",
    "                ax.errorbar(x_axis,means,yerr=stds, xerr = x_std, label = label, color = colors[linecount], ecolor = ecolors[linecount])\n",
    "            elif error_bars == 'IQR':\n",
    "                ax.errorbar(x_axis,median,yerr=np.array([median - lower_quartile, upper_quartile - median]), label = label + ' median', color = colors[linecount], ecolor = ecolors[linecount])\n",
    "                ax.plot(x_axis,means, '--', color = colors[linecount])\n",
    "                ax.plot()\n",
    "            else:\n",
    "                ax.plot(x_axis,means, color = colors[linecount], label = label)\n",
    "            print(f'{count}/{total_lines}')\n",
    "            linecount += 1\n",
    "            count += 1\n",
    "            linecount %= len(colors)\n",
    "\n",
    "    if x_axis_type == 'total_infections':\n",
    "        xlabel = f'Total Infections\\n\\n{caption}'\n",
    "    elif x_axis_type == 'times':\n",
    "        xlabel = f'Time (days)\\n\\n{caption}'\n",
    "    else:\n",
    "        xlabel = f'Doubling Times\\n\\n{caption}'\n",
    "    ax.set_xlabel(xlabel)\n",
    "    ax.set_ylabel(f'Fraction of population')\n",
    "    if error_bars == 'IQR':\n",
    "        plt.plot([],[],linestyle = '--',color='black',label = 'Means')\n",
    "    ax.legend()\n",
    "    show_fig(fig, figsavename)\n",
    "\n",
    "def plot_infection_ratio(dataset: SimData,\n",
    "                         filters: list[dict],\n",
    "                         x_axis_type: str = 'times',\n",
    "                         max_first: float = 0.1,\n",
    "                         min_first: Optional[float] = None,\n",
    "                         n_points: int = 100,\n",
    "                         log: str = '',\n",
    "                         error_bars: str = 'std',\n",
    "                         doubling_time: Optional[float] = None,\n",
    "                         figsavename: Optional[str] = None):\n",
    "    #Filters in format: denominator, numerator\n",
    "    assert x_axis_type in ['times', 'scaled_times', 'total_infections']\n",
    "    assert error_bars in ['std', 'IQR', 'None']\n",
    "    fig, axs = plotprep(log)\n",
    "    ax = axs[0]\n",
    "    if x_axis_type == 'total_infections':\n",
    "        total_infections = dataset[create_filter(datatypes=['municipal'],compartments=['E','I','R'])]\n",
    "        total_infections = total_infections.sum(axis = ('cities','datatypes','groups','compartments'))\n",
    "        x_axis = total_infections.mean(axis='sims').array\n",
    "        x_std = total_infections.std(axis='sims').array\n",
    "    elif x_axis_type == 'times':\n",
    "        x_axis = np.array(dataset.values_present['times'])\n",
    "        x_std = 0\n",
    "    else:\n",
    "        assert isinstance(doubling_time, float), 'Specify Doubling time to use scaled_times!'\n",
    "        x_axis = np.array(dataset.values_present['times']) / doubling_time\n",
    "        x_std = 0\n",
    "    data_1, data_2 = (dataset[f] for f in filters)\n",
    "    values_array = data_2.array/(data_1.array + 1e-10)\n",
    "    values = SimData(values_array, data_2.values_present)\n",
    "    assert isinstance(values, SimData)\n",
    "    means = values.mean(axis = 'sims').array.squeeze()\n",
    "    stds = values.std(axis = 'sims').array.squeeze()\n",
    "    lower_quartile, median, upper_quartile = values.quartiles(axis='sims',SimData_out=False)\n",
    "    if error_bars == 'std':\n",
    "        ax.errorbar(x_axis,means,yerr=stds, xerr = x_std)\n",
    "    elif error_bars == 'IQR':\n",
    "        ax.errorbar(x_axis,median,yerr=np.array([median - lower_quartile, upper_quartile - median]), label = 'median')\n",
    "        ax.plot(x_axis,means, '--', color = 'black', label = 'mean')\n",
    "        ax.plot()\n",
    "    else:\n",
    "        ax.plot(x_axis,means)\n",
    "    if x_axis_type == 'total_infections':\n",
    "        xlabel = 'Total Infections'\n",
    "    elif x_axis_type == 'times':\n",
    "        xlabel = 'Time (days)'\n",
    "    else:\n",
    "        xlabel = 'Doubling Times'\n",
    "    ax.set_xlabel(xlabel)\n",
    "    same_categories = {k:v for k,v in filters[0].items() if k in filters[1].keys() and v == filters[1][k]}\n",
    "    uniques = [{k:v for k,v in f.items() if k not in same_categories.keys()} for f in filters]\n",
    "    ylabel = f'{dict_to_str(same_categories)}. Ratio of {dict_to_str(uniques[1])} / {dict_to_str(uniques[0])}'\n",
    "    ax.set_ylabel(ylabel)\n",
    "    ax.legend()\n",
    "    show_fig(fig, figsavename)\n",
    "\n",
    "def joint_times_distribution(data: SimData,\n",
    "                       filters: list[dict],\n",
    "                       threshold: float,\n",
    "                       log: str = '',\n",
    "                       legend_labels: Optional[list[str]] = None,\n",
    "                       bins = 30,\n",
    "                       density = False,\n",
    "                       alpha: float = 0.5,\n",
    "                       figsavename: Optional[str] = None):\n",
    "    assert len(filters) == 2\n",
    "    if legend_labels is None:\n",
    "        for i, filter in enumerate(filters):\n",
    "            print(f'Dataset {i+1}: {filter}')\n",
    "        legend_labels = [f'Dataset {i+1}' for i in range(len(filters))]\n",
    "    assert len(legend_labels) == len(filters)\n",
    "    fig, axs = plotprep(log)\n",
    "    ax = axs[0]\n",
    "    \n",
    "    filtered = [data[filter].sum(axis=('cities','datatypes','groups','compartments')).array for filter in filters]\n",
    "\n",
    "    total_filters = [filter.copy() for filter in filters]\n",
    "    for f in total_filters:\n",
    "        del f['compartments']\n",
    "    totals = [data[filter].sum(axis=('cities','datatypes','groups','compartments')).array for filter in total_filters]\n",
    "    fractions = [a/b for a,b in zip(filtered,totals)]\n",
    "    steps = [find_first_index(arr>threshold) for arr in fractions]\n",
    "    x,y = tuple([np.array(data.values_present['times'])[step] for step in steps ])\n",
    "    ax.hist2d(x, y, bins=(bins,bins), cmap = plt.cm.jet)\n",
    "    ax.set_xlabel(f'Times in {legend_labels[0]}')\n",
    "    ax.set_ylabel(f'Times in {legend_labels[1]}')\n",
    "    ax.set_title(f'Time(days) till fraction in datasets reach {threshold}')\n",
    "    show_fig(fig, figsavename)\n",
    "\n",
    "def times_until_threshold(data: SimData,\n",
    "                          filter: dict,\n",
    "                          threshold: float):\n",
    "    filtered = data[filter].sum(axis=('cities','datatypes','groups','compartments')).array\n",
    "    total_filter = filter.copy()\n",
    "    del total_filter['compartments']\n",
    "    totals = data[total_filter].sum(axis=('cities','datatypes','groups','compartments')).array\n",
    "    fraction = filtered/np.maximum(totals, 1e-10)\n",
    "    return np.array(data.values_present['times'])[find_first_index(fraction>threshold)]\n",
    "\n",
    "def time_diff_at_threshold(data: SimData,\n",
    "                           filters: list[dict],\n",
    "                           threshold1: float,\n",
    "                           threshold2: Optional[float] = None):\n",
    "    assert len(filters) == 2\n",
    "    if threshold2 is None:\n",
    "        threshold2 = threshold1\n",
    "    thresholds = (threshold1,threshold2)\n",
    "    filtered = [data[filter].sum(axis=('cities','datatypes','groups','compartments')).array for filter in filters]\n",
    "\n",
    "    total_filters = [filter.copy() for filter in filters]\n",
    "    for f in total_filters:\n",
    "        del f['compartments']\n",
    "    totals = [data[filter].sum(axis=('cities','datatypes','groups','compartments')).array for filter in total_filters]\n",
    "    fractions = [a/np.maximum(b, 1e-10) for a,b in zip(filtered,totals)]\n",
    "    steps = [find_first_index(arr>threshold) for arr, threshold in zip(fractions,thresholds)]\n",
    "    x,y = tuple([np.array(data.values_present['times'])[step] for step in steps ])\n",
    "    return y - x\n",
    "\n",
    "def different_thresholds_diffs_data(data: SimData,\n",
    "                                    filters: list[dict],\n",
    "                                    thresholds1: np.ndarray = 10**np.linspace(-6,-1,24),\n",
    "                                    thresholds2: Optional[np.ndarray] = None):\n",
    "    if thresholds2 is None:\n",
    "        thresholds2 = thresholds1\n",
    "    times = np.zeros((len(thresholds1), len(thresholds2), 2, len(data.values_present['sims'])))\n",
    "    for i,threshold in tqdm(enumerate(thresholds1)):\n",
    "        times[i,:,0] = times_until_threshold(data, filters[0], threshold)\n",
    "    for i,threshold in tqdm(enumerate(thresholds2)):\n",
    "        times[:,i,1] = times_until_threshold(data, filters[1], threshold)\n",
    "    return times[:,:,1], times[:,:,0]\n",
    "\n",
    "def different_thresholds_diffs(data: SimData,\n",
    "                               filters: list[dict],\n",
    "                               x_thresholds: np.ndarray = 10**np.linspace(-6,-1,24),\n",
    "                               y_thresholds: Optional[np.ndarray] = None, \n",
    "                               log = 'xy',\n",
    "                               include_line = True,\n",
    "                               zero_centre = True,\n",
    "                               figsavename: Optional[str] = None):\n",
    "    fig, axs = plotprep(log)\n",
    "    if y_thresholds is None:\n",
    "        y_thresholds = x_thresholds\n",
    "    ax = axs[0]\n",
    "    a,b = different_thresholds_diffs_data(data, filters, x_thresholds, y_thresholds)\n",
    "    diffs = a - b\n",
    "    mean_diffs = diffs.mean(axis = -1)\n",
    "    positions = {'x': x_thresholds, 'y': y_thresholds}\n",
    "    edges = {}\n",
    "    for axis in ['x','y']:\n",
    "        if axis in log: \n",
    "            ps = np.log10(positions[axis])\n",
    "        else:\n",
    "            ps = positions[axis]\n",
    "        edge = [(ps[i]+ps[i+1])/2 for i in range(len(ps)-1)]\n",
    "        edge = np.array([2*ps[0]-edge[0]] + edge + [2*ps[-1]-edge[-1]])\n",
    "        if axis in log:\n",
    "            edge = 10 ** edge\n",
    "        edges[axis] = edge\n",
    "    y,x = np.meshgrid(edges['y'], edges['x'])\n",
    "    # a = np.ones((len(x)-1, len(y)-1))\n",
    "    # im = ax.pcolor(x, y, a, cmap = 'bwr')\n",
    "    # print(mean_diffs)\n",
    "    max_range = max(np.abs(np.min(mean_diffs)), np.abs(np.max(mean_diffs)))\n",
    "    if zero_centre:\n",
    "        im = ax.pcolor(x, y, mean_diffs, cmap = 'seismic', vmin = -max_range, vmax = max_range)\n",
    "    else:\n",
    "        im = ax.pcolor(x, y, mean_diffs, cmap = 'hot')\n",
    "    minval = max(x_thresholds.min(), y_thresholds.min())\n",
    "    maxval = min(x_thresholds.max(), y_thresholds.max())\n",
    "    xs = np.logspace(np.log10(minval), np.log10(maxval), 3)\n",
    "    if include_line:\n",
    "        ax.plot(xs,xs, color = 'black')\n",
    "    ax.set_xlabel(f'Threshold Fraction for {dict_to_str(filters[0])}')\n",
    "    ax.set_ylabel(f'Threshold Fraction for {dict_to_str(filters[1])}')\n",
    "    ax.set_title(f'Time to detection in ({dict_to_str(filters[0])}) - Time to detection ({dict_to_str(filters[1])})')\n",
    "    fig.colorbar(im, ax=ax)\n",
    "    show_fig(fig, figsavename)\n",
    "\n",
    "def threshold_ratio_diffs(data: SimData,\n",
    "                          filters: list[dict],\n",
    "                          filter_names: list[str] = ['Municipal', 'Airplane'],\n",
    "                          n_points: int = 30,\n",
    "                          xlims_log: list[Union[tuple,float]] = [(-1.5,-0.7)],\n",
    "                          ylims_log: Union[tuple,int] = (-0.7,-6),\n",
    "                          log = 'x',\n",
    "                          error_bars = 'std',\n",
    "                          figsavename: Optional[str] = None):\n",
    "    fig, axs = plotprep(log)\n",
    "    assert error_bars in ['std', 'IQR', 'None']\n",
    "    ax = axs[0]\n",
    "    for xlimlog in xlims_log:\n",
    "        x_points = 10**np.linspace(xlimlog[0],xlimlog[1],n_points) if isinstance(xlimlog,tuple) else np.ones(n_points) * 10 ** xlimlog\n",
    "        y_points = 10**np.linspace(ylims_log[0],ylims_log[1],n_points) if isinstance(ylims_log,tuple) else np.ones(n_points) * 10 ** ylims_log\n",
    "        ratios = y_points / x_points\n",
    "        diffs = np.zeros((n_points, len(data.values_present['sims'])))\n",
    "        for i in tqdm(range(n_points)):\n",
    "            diffs[i] = time_diff_at_threshold(data,filters,x_points[i], y_points[i])\n",
    "        means = diffs.mean(axis=1)\n",
    "        lower_quartiles, medians, upper_quartiles = quartiles(diffs, axis = 1)\n",
    "        stds = diffs.std(axis=1)\n",
    "        if error_bars == 'None':\n",
    "            if isinstance(xlimlog, float):\n",
    "                threshold = 10**xlimlog\n",
    "                label = f'{filter_names[1]} Threshold = {float(f\"{threshold:.1g}\"):g}'\n",
    "                ax.plot(ratios, means,label=label)\n",
    "            else:\n",
    "                ax.plot(ratios, means)\n",
    "        elif error_bars == 'std':\n",
    "            if isinstance(xlimlog, float):\n",
    "                threshold = 10**xlimlog\n",
    "                label = f'{filter_names[1]} Threshold = {float(f\"{threshold:.1g}\"):g}'\n",
    "                ax.errorbar(ratios, means,stds,label=label)\n",
    "            else:\n",
    "                ax.errorbar(ratios, means,stds)\n",
    "        else:\n",
    "            if isinstance(xlimlog, float):\n",
    "                threshold = 10**xlimlog\n",
    "                label = f'Median, {filter_names[1]} Threshold = {float(f\"{threshold:.1g}\"):g}'\n",
    "            else:\n",
    "                label = 'Median'\n",
    "            ax.errorbar(ratios, medians, [medians-lower_quartiles, upper_quartiles-medians], label = label)\n",
    "            ax.plot(ratios, means, '--', color='black')\n",
    "    ax.set_xlabel(f'Ratio of {filter_names[1]} threshold / {filter_names[0]} threshold')\n",
    "    ax.set_ylabel(f'Advantage of using {filter_names[1]} detection (days)')\n",
    "    if error_bars == 'IQR':\n",
    "        ax.plot('', '--', color = 'black', label = 'Means')\n",
    "    if len(xlims_log) > 1 or error_bars == 'IQR':\n",
    "        ax.legend()\n",
    "    ax.axvline(1, color='dimgray')\n",
    "    ax.axhline(0, color='dimgray')\n",
    "    show_fig(fig, figsavename)\n",
    "\n",
    "def time_diff_at_threshold_hist(data: SimData,\n",
    "                                filters: list[dict],\n",
    "                                threshold: float,\n",
    "                                log: str = '',\n",
    "                                legend_labels: Optional[list[str]] = None,\n",
    "                                bins = 30,\n",
    "                                density = False,\n",
    "                                daily = False,\n",
    "                                figsavename: Optional[str] = None):\n",
    "    assert len(filters) == 2\n",
    "    if legend_labels is None:\n",
    "        # for i, filter in enumerate(filters):\n",
    "        #     print(f'Dataset {i+1}: {filter}')\n",
    "        legend_labels = [f'Dataset {i+1}' for i in range(len(filters))]\n",
    "    assert len(legend_labels) == len(filters)\n",
    "    fig, axs = plotprep(log)\n",
    "    ax = axs[0]\n",
    "    \n",
    "    to_plot = time_diff_at_threshold(data, filters, threshold)\n",
    "    if not daily:\n",
    "        # to_plot = y-x\n",
    "        ax.hist(to_plot, bins=bins, density=density)\n",
    "        ylabel = 'Probability Density' if density else 'Frequency'\n",
    "        ax.set_ylabel(ylabel)\n",
    "    else:\n",
    "        to_plot = np.round(to_plot).astype(np.int64)\n",
    "        options = list(set(to_plot))\n",
    "        amounts = np.bincount(to_plot,minlength = len(options))\n",
    "        amounts = amounts[-len(options):]\n",
    "        pos = np.arange(len(options))\n",
    "        options = [str(option) for option in options]\n",
    "        ax.bar(pos,amounts)\n",
    "        ax.set_xticks(pos, options)\n",
    "        ax.set_ylabel('Frequency')\n",
    "    filters_str = [dict_to_str(filter) for filter in filters]\n",
    "    ax.set_xlabel(f'Time for {legend_labels[1]} - time for {legend_labels[0]} (days) \\n\\nDataset 1: {filters_str[0]}\\nDataset 2: {filters_str[1]}')\n",
    "    ax.set_title(f'Difference between time to reach {threshold} in the datasets')\n",
    "    show_fig(fig, figsavename)\n",
    "\n",
    "def differences_vs_threshold_data(datasets: dict[str,SimData],\n",
    "                     filters: list[dict],\n",
    "                     thresholds: Union[list[float], np.ndarray] = 10**np.linspace(-4,-0.7,20)):\n",
    "    assert len(filters) == 2\n",
    "    out = np.zeros((5,len(datasets.values()), len(thresholds)))\n",
    "    for i, dataset in enumerate(datasets.values()):\n",
    "        for j,threshold in tqdm(enumerate(thresholds)):\n",
    "            differences = time_diff_at_threshold(dataset, filters, threshold)\n",
    "            out[0,i,j] = differences.mean()\n",
    "            out[4,i,j], out[1,i,j], out[3,i,j] = quartiles(differences)#type: ignore\n",
    "            out[2,i,j] = differences.std() #type: ignore\n",
    "    return out\n",
    "\n",
    "def differences_vs_threshold(datasets: dict[str,SimData],\n",
    "                             calculated_data: np.ndarray,\n",
    "                             filters: list[dict],\n",
    "                             thresholds: Union[list[float], np.ndarray] = 10**np.linspace(-4,-0.7,21),\n",
    "                             log: str = 'x',\n",
    "                             error_bars: str = 'std',\n",
    "                             figsavename: Optional[str] = None):\n",
    "    assert error_bars in ['std', 'IQR', 'None']\n",
    "    assert len(filters) == 2\n",
    "    assert len(calculated_data) == 5\n",
    "    fig, axs = plotprep(log)\n",
    "    colors = ['blue', 'red', 'green', 'fuchsia', 'dimgrey', 'yellow','darkviolet', 'darkorange']\n",
    "    ecolors = ['cornflowerblue', 'lightcoral', 'palegreen', 'lightpink', 'lightgrey', 'lemonchiffon', 'thistle', 'navajowhite']\n",
    "    ax = axs[0]\n",
    "    legend_labels = list(datasets.keys())\n",
    "    means, medians, stds, upper_quartiles, lower_quartiles = tuple(calculated_data)\n",
    "\n",
    "    for i in tqdm(range(len(datasets))):\n",
    "        col_idx = i%len(colors)\n",
    "        if error_bars == 'None':\n",
    "            ax.plot(thresholds, means[i], label = legend_labels[i])\n",
    "        else:\n",
    "            errors = np.array([medians[i]-lower_quartiles[i], upper_quartiles[i]-medians[i]]) if error_bars == 'IQR' else stds[i] #type: ignore\n",
    "            y = means[i] if error_bars == 'std' else medians[i] #type: ignore\n",
    "            plt.errorbar(thresholds, y, errors, color = colors[col_idx], ecolor = ecolors[col_idx], label = legend_labels[i])\n",
    "    filters_str = [dict_to_str(filter) for filter in filters]\n",
    "    ax.set_xlabel(f'Threshold\\n\\nTime Difference between:\\nDataset 1: {filters_str[0]}\\nDataset 2: {filters_str[1]}')\n",
    "    ax.set_ylabel('Time Difference (days)')\n",
    "    ax.legend()\n",
    "    show_fig(fig, figsavename)\n",
    "\n",
    "def differences_vs_variable(datasets: dict[str,SimData],\n",
    "                             variables: list[int],\n",
    "                             calculated_data: np.ndarray,\n",
    "                             filters: list[dict],\n",
    "                             old_thresholds: np.ndarray,\n",
    "                             new_thresholds: np.ndarray,\n",
    "                             log: str = 'x',\n",
    "                             error_bars: str = 'std',\n",
    "                             legend_labels: Optional[list[str]] = None,\n",
    "                             x_name: str = 'Daily Mixing Rate',\n",
    "                             figsavename: Optional[str] = None):\n",
    "    assert error_bars in ['std', 'IQR', 'None']\n",
    "    assert len(filters) == 2\n",
    "    assert len(variables) == len(datasets)\n",
    "    assert len(calculated_data) == 5\n",
    "    fig, axs = plotprep(log)\n",
    "    colors = ['blue', 'red', 'green', 'fuchsia', 'dimgrey', 'yellow','darkviolet', 'darkorange']\n",
    "    ecolors = ['cornflowerblue', 'lightcoral', 'palegreen', 'lightpink', 'lightgrey', 'lemonchiffon', 'thistle', 'navajowhite']\n",
    "    ax = axs[0]\n",
    "    threshold_idxs = [np.argmin(np.abs(old_thresholds - t)) for t in new_thresholds]\n",
    "    thresholds = old_thresholds[threshold_idxs]\n",
    "    print(thresholds)\n",
    "    calculated_data = calculated_data[...,np.array(threshold_idxs)]\n",
    "    means, medians, stds, upper_quartiles, lower_quartiles = tuple(calculated_data)\n",
    "\n",
    "    if legend_labels is None:\n",
    "        legend_labels = [\"Threshold: {:.0e}\".format(threshold) for threshold in thresholds]\n",
    "    for i, threshold in enumerate(thresholds):\n",
    "        col_idx = i%len(colors)\n",
    "        if error_bars == 'None': \n",
    "            ax.plot(variables, means[:,i], label = legend_labels[i])\n",
    "        else:\n",
    "            errors = np.array([medians[:,i]-lower_quartiles[:,i], upper_quartiles[:,i]-medians[:,i]]) if error_bars == 'IQR' else stds[:,i] #type: ignore\n",
    "            y = means[:,i] if error_bars == 'std' else medians[:,i] #type: ignore\n",
    "            plt.errorbar(variables, y, errors, color = colors[col_idx], ecolor = ecolors[col_idx], label = legend_labels[i])\n",
    "    filters_str = [dict_to_str(filter) for filter in filters]\n",
    "    ax.set_xlabel(f'{x_name}\\n\\nTime Difference between:\\nDataset 1: {filters_str[0]}\\nDataset 2: {filters_str[1]}')\n",
    "    ax.set_ylabel('Time Difference (days)')\n",
    "    ax.legend()\n",
    "    show_fig(fig, figsavename)\n",
    "\n",
    "def differences_vs_variable_final_img(datasets: dict[str,SimData],\n",
    "                             variables: list[int],\n",
    "                             calculated_data: np.ndarray,\n",
    "                             filters: list[dict],\n",
    "                             old_thresholds: np.ndarray,\n",
    "                             new_thresholds: np.ndarray,\n",
    "                             log: str = 'x',\n",
    "                             error_bars: str = 'std',\n",
    "                             legend_labels: Optional[list[str]] = None,\n",
    "                             x_name: str = 'Daily Mixing Rate',\n",
    "                             figsavename: Optional[str] = None):\n",
    "    assert error_bars in ['std', 'IQR', 'None']\n",
    "    assert len(filters) == 2\n",
    "    assert len(variables) == len(datasets)\n",
    "    assert len(calculated_data) == 5\n",
    "    fig, axs = plotprep(log)\n",
    "    colors = ['blue', 'red', 'green', 'fuchsia', 'dimgrey', 'yellow','darkviolet', 'darkorange']\n",
    "    ecolors = ['cornflowerblue', 'lightcoral', 'palegreen', 'lightpink', 'lightgrey', 'lemonchiffon', 'thistle', 'navajowhite']\n",
    "    ax = axs[0]\n",
    "    threshold_idxs = [np.argmin(np.abs(old_thresholds - t)) for t in new_thresholds]\n",
    "    thresholds = old_thresholds[threshold_idxs]\n",
    "    print(thresholds)\n",
    "    calculated_data = calculated_data[...,np.array(threshold_idxs)]\n",
    "    means, medians, stds, upper_quartiles, lower_quartiles = tuple(calculated_data)\n",
    "\n",
    "    if legend_labels is None:\n",
    "        legend_labels = [\"Threshold: {:.0e}\".format(threshold) for threshold in thresholds]\n",
    "    for i, threshold in enumerate(thresholds):\n",
    "        col_idx = i%len(colors)\n",
    "        if error_bars == 'None': \n",
    "            ax.plot(variables, means[:,i], label = legend_labels[i])\n",
    "        else:\n",
    "            errors = np.array([medians[:,i]-lower_quartiles[:,i], upper_quartiles[:,i]-medians[:,i]]) if error_bars == 'IQR' else stds[:,i] #type: ignore\n",
    "            y = means[:,i] if error_bars == 'std' else medians[:,i] #type: ignore\n",
    "            plt.errorbar(variables, y, errors, color = colors[col_idx], ecolor = ecolors[col_idx], label = legend_labels[i])\n",
    "    filters_str = [dict_to_str(filter) for filter in filters]\n",
    "    ax.set_xlabel('Mixing Rate per Day \\nDetection threshold is chosen to be 0.2 for both detection approaches')\n",
    "    ax.set_ylabel('Advantage of using Airplane Detection (days)')\n",
    "    show_fig(fig, figsavename)\n",
    "\n",
    "def create_filter(cities: Optional[Iterable[Union[str,int]]]=None,\n",
    "                  datatypes: Optional[Iterable[Union[str,int]]]=None,\n",
    "                  groups: Optional[Iterable[Union[str,int]]]=None,\n",
    "                  compartments: Optional[Iterable[Union[str,int]]]=None,\n",
    "                  sims: Optional[Iterable[Union[str,int]]]=None,\n",
    "                  times: Optional[Iterable[Union[str,int]]]=None):\n",
    "    return {key:val for key, val in locals().items() if val is not None}\n",
    "\n",
    "def create_arrival_municipal_filter(city = 1, groups = None, compartments = 'I'):\n",
    "    datatypes = ['arrivals', 'municipal']\n",
    "    return [create_filter([city], [datatype], groups, list(compartments)) for datatype in datatypes]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class Travel(City):\n",
    "    def __init__(self,\n",
    "                 cities: list[City],\n",
    "                 mixmatrix: np.ndarray,\n",
    "                 city_names: Optional[list[str]] = None):\n",
    "        self.cities = cities\n",
    "        self.n_cities = len(cities)\n",
    "        self.city_names = list(range(self.n_cities)) if city_names is None else city_names\n",
    "        for i,city in enumerate(self.cities):\n",
    "            city.name = self.city_names[i]\n",
    "        self.mixmatrix = mixmatrix.astype(np.int64)\n",
    "        assert np.all(mixmatrix == mixmatrix.T), \"mixmatrix must be symmetrical!\"\n",
    "        assert np.trace(mixmatrix) == 0, \"mixmatrix must be traceless!\"\n",
    "        self.n_groups = self.cities[0].n_groups\n",
    "        self.compartments = list(self.cities[0].compartments)\n",
    "        self.I0s = []\n",
    "\n",
    "        self.colors = ['blue', 'red', 'green', 'fuchsia', 'dimgrey', 'yellow','darkviolet', 'darkorange']\n",
    "        self.ecolors = ['cornflowerblue', 'lightcoral', 'palegreen', 'lightpink', 'lightgrey', 'lemonchiffon', 'thistle', 'navajowhite']\n",
    "        self.ncolors = len(self.colors)\n",
    "  \n",
    "    def reset_parameters(self):\n",
    "        for i,city in enumerate(self.cities):\n",
    "            city.reset_parameters(I0 = self.I0s[i],\n",
    "                                  n_sims = self.n_sims,\n",
    "                                  simulation_steps = self.simulation_steps)\n",
    "    \n",
    "    def step(self,\n",
    "            beta:float,\n",
    "            delta_t: float,\n",
    "            p_infectious: float,\n",
    "            p_recovery: float,\n",
    "            simulation_step: int):\n",
    "        for i,city in enumerate(self.cities):\n",
    "            city.step_internal(beta,\n",
    "                               delta_t,\n",
    "                               p_infectious,\n",
    "                               p_recovery,\n",
    "                               simulation_step)\n",
    "        for i, city_i in enumerate(self.cities[:-1]):\n",
    "            for j, city_j in enumerate(self.cities[i+1:]):\n",
    "                j += i+1\n",
    "                mix = self.mixmatrix[i,j]\n",
    "                if mix == 0:\n",
    "                    continue\n",
    "                travel_ij = city_i.select_travellers(mix, simulation_step)\n",
    "                travel_ji = city_j.select_travellers(mix, simulation_step)\n",
    "                \n",
    "                city_i.municipal[...,simulation_step] += travel_ji - travel_ij\n",
    "                city_i.arrivals[...,simulation_step] += travel_ji\n",
    "                city_i.departures[...,simulation_step] += travel_ij\n",
    "                \n",
    "                city_j.municipal[...,simulation_step] += travel_ij - travel_ji\n",
    "                city_j.arrivals[...,simulation_step] += travel_ij\n",
    "                city_j.departures[...,simulation_step] += travel_ji\n",
    "    \n",
    "    def multiple_sims(self,\n",
    "                      delta_t: float,\n",
    "                      epidemic_time: Union[int,float],\n",
    "                      disease: Disease,\n",
    "                      I0s: Union[int,float,list[int],list[float],np.ndarray],\n",
    "                      n_sims: int = 100,\n",
    "                      moving_avg: bool = False):\n",
    "        if isinstance(I0s, np.ndarray):\n",
    "            assert I0s.dtype == np.int64\n",
    "        elif isinstance(I0s, int):\n",
    "            I0s = np.array([I0s] + [0] * (len(self.cities) - 1))\n",
    "        elif isinstance(I0s, float):\n",
    "            I0s = np.array([I0s * self.cities[0].N0] + [0] * (len(self.cities) - 1))\n",
    "        elif isinstance(I0s, list):\n",
    "            I0s = [I * self.cities[j].N0 if isinstance(I,float) else I for j,I in enumerate(I0s)]\n",
    "            I0s = np.array(I0s)\n",
    "        self.I0s = I0s\n",
    "        assert (1 / delta_t) % 1 == 0, \"1/delta_t must be an integer\"\n",
    "        self.start_city = self.cities[[I0 != 0 for I0 in I0s].index(True)]\n",
    "        p_recovery = 1 - np.exp( - delta_t * disease.gamma)\n",
    "        p_infectious = 1 - np.exp( - delta_t * disease.delta)\n",
    "        self.simulation_steps = int(epidemic_time / delta_t)\n",
    "        self.n_sims = n_sims\n",
    "        self.times = np.linspace(0, epidemic_time, self.simulation_steps)\n",
    "        self.scaled_times = self.times /disease.doubling_time\n",
    "        self.delta_t = delta_t\n",
    "        for city in self.cities:\n",
    "            city.delta_t = delta_t\n",
    "            city.n_sims = n_sims\n",
    "            city.epidemic_time = epidemic_time\n",
    "            city.times = self.times\n",
    "            city.scaled_times = self.scaled_times\n",
    "            city.simulation_steps = self.simulation_steps\n",
    "            city.disease = disease\n",
    "        self.reset_parameters()\n",
    "        for simulation_step in tqdm(range(1,self.simulation_steps)):\n",
    "            self.step(disease.beta,\n",
    "                      delta_t,\n",
    "                      p_infectious,\n",
    "                      p_recovery,\n",
    "                      simulation_step)\n",
    "        for city in tqdm(self.cities):\n",
    "            city.daily_flight_data(moving_avg = moving_avg)\n",
    "        \n",
    "        #make and return SimData\n",
    "        n_datatypes = 5 if moving_avg else 3\n",
    "        out_arr = np.zeros((self.n_cities, n_datatypes, self.n_groups, len(self.compartments), self.n_sims, self.simulation_steps))\n",
    "        if moving_avg:\n",
    "            for i,city in enumerate(self.cities):\n",
    "                out_arr[i] = np.array([city.municipal,\n",
    "                                    city.arrivals,\n",
    "                                    city.departures,\n",
    "                                    city.arrivals_moving_avg,\n",
    "                                    city.departures_moving_avg])\n",
    "            all_labels = {'cities': self.city_names,\n",
    "                        'datatypes': ['municipal', 'arrivals', 'departures', 'arrivals_moving_avg', 'departures_moving_avg'],\n",
    "                        'groups': self.cities[0].groups,\n",
    "                        'compartments': self.compartments,\n",
    "                        'sims': list(range(self.n_sims)),\n",
    "                        'times': self.times}\n",
    "        else:\n",
    "            for i,city in enumerate(self.cities):\n",
    "                out_arr[i] = np.array([city.municipal,\n",
    "                                    city.arrivals,\n",
    "                                    city.departures])\n",
    "            all_labels = {'cities': self.city_names,\n",
    "                        'datatypes': ['municipal', 'arrivals', 'departures'],\n",
    "                        'groups': self.cities[0].groups,\n",
    "                        'compartments': self.compartments,\n",
    "                        'sims': list(range(self.n_sims)),\n",
    "                        'times': self.times}\n",
    "\n",
    "        return SimData(out_arr, all_labels)\n",
    "\n",
    "    def __str__(self):\n",
    "        out = ''\n",
    "        for i, city in enumerate(self.cities):\n",
    "            out += f'City: {self.city_names[i]}\\n'\n",
    "            out += str(city) + '\\n'\n",
    "        if self.I0s is not None:\n",
    "            out += f'I0s: {self.I0s}\\n'\n",
    "        out += f'Mixmatrix: {self.mixmatrix}'\n",
    "        return out\n",
    "    \n",
    "    def plot_sims(self,\n",
    "                  to_shift: bool = True, #change to shift_by: Optional[Union[int,str]] = None, assert shift_by in city.groups\n",
    "                  included_cities: Optional[list[int]] = None,\n",
    "                  separate_groups: bool = True,\n",
    "                  figsavename: Optional[str] = None,\n",
    "                  moving_avg = False,\n",
    "                  log=''):\n",
    "\n",
    "        if included_cities is None:\n",
    "            included_cities = list(range(len(self.cities)))\n",
    "            included_names = included_cities\n",
    "        elif isinstance(included_cities[0],int):\n",
    "            included_names = included_cities\n",
    "        elif isinstance(included_cities[0],str):\n",
    "            included_names = included_cities\n",
    "            included_cities = [i for i, name in enumerate(self.city_names) if name in included_cities]\n",
    "        assert isinstance(included_names,list) #type: ignore\n",
    "        shift_index = None\n",
    "        if to_shift:\n",
    "            shift_index = self.start_city.municipal[:,self.start_city.I_index].sum(axis=0).argmax(axis = 1)\n",
    "        # for i, city in enumerate(self.cities):\n",
    "        for j,i in enumerate(included_cities):\n",
    "            city = self.cities[i]\n",
    "        # if i in included_cities:\n",
    "            if figsavename is None:\n",
    "                savename = None\n",
    "            else:\n",
    "                savename = f'{figsavename}_{i}'\n",
    "            print(f\"City {included_names[j]}:\")\n",
    "            city.plot_sims(times = self.times,\n",
    "                            cityname=i,\n",
    "                            shift_index = shift_index,\n",
    "                            separate_groups=separate_groups,\n",
    "                            moving_avg=moving_avg,\n",
    "                            log=log,\n",
    "                            figsavename = savename)\n",
    "            print('\\n')\n",
    "    \n",
    "    def __call__(self,\n",
    "                 delta_t: float,\n",
    "                 epidemic_time: Union[int,float],\n",
    "                 disease: Disease,\n",
    "                 I0s: Union[int,float,list[int],list[float], np.ndarray],\n",
    "                 n_sims: int = 100,\n",
    "                 moving_avg = True):\n",
    "        return self.multiple_sims(delta_t, epidemic_time, disease, I0s, n_sims, moving_avg = moving_avg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ReturnHomeCity(City):\n",
    "    def __init__(self,\n",
    "                 city_names: list[Union[int,str]],\n",
    "                 this_city_name: Union[int,str],\n",
    "                 mixmatrix: np.ndarray,\n",
    "                 N0: int = 10**6,\n",
    "                 group_LR: float = 0.1,\n",
    "                 trip_length: int = 10,\n",
    "                 p_go_home: float = 0.9,\n",
    "                 p_continue_travel: float = 0.05,\n",
    "                 compartments: str = 'SEIR'):\n",
    "        self.city_index = city_names.index(this_city_name)\n",
    "        self.n_cities = len(city_names)\n",
    "        self.NOs = np.ones(len(city_names),np.int64) * round(mixmatrix[self.city_index] * trip_length)\n",
    "        self.N0s[self.city_index] = N0 - self.NOs[:-1].sum()\n",
    "        assert self.N0s.sum() == N0\n",
    "        self.groups = city_names\n",
    "        self.n_groups = len(city_names)\n",
    "        group_LR_matrix = homogeneous_LR_matrix(self.n_groups, group_LR)\n",
    "        super().__init__(self.N0s, self.groups, compartments, group_LR_matrix)\n",
    "        self.trip_length = trip_length\n",
    "        self.p_go_home = p_go_home\n",
    "        self.p_continue_travel = p_continue_travel\n",
    "        \n",
    "    def initial_conditions(self):\n",
    "        #puts all initial infected people into the compartment of people who are in their home city\n",
    "        self.municipal[self.city_index,2,:,0] = self.I0\n",
    "    \n",
    "    def select_travellers(self,\n",
    "                          mix: int,\n",
    "                          destination_index: int,\n",
    "                          simulation_step: int):\n",
    "        p_travel = 1 - np.exp(- self.delta_t * mix / self.N0s[self.city_index])\n",
    "        p_return_home = 1 - np.exp(- self.delta_t * self.p_go_home / (self.trip_length * self.N0s[self.city_index]))\n",
    "        p_onward_travel = 1 - np.exp(- self.delta_t * self.p_continue_travel / (self.trip_length * self.N0s[self.city_index]))\n",
    "\n",
    "        to_travel = np.random.binomial(self.municipal[...,simulation_step],p_onward_travel)\n",
    "        to_travel[self.city_index] = np.random.binomial((self.municipal[self.city_index,...,simulation_step]), p_travel)\n",
    "        to_travel[destination_index] = np.random.binomial((self.municipal[destination_index,...,simulation_step]), p_return_home)\n",
    "        return to_travel \n",
    "\n",
    "    def select_settlers(self, simulation_step: int):\n",
    "        p_settle = 1 - np.exp(- self.delta_t * (1-self.p_go_home-self.p_continue_travel) / (self.trip_length * self.N0s[self.city_index] * (self.n_cities-1)))\n",
    "        to_settle = np.random.binomial(self.municipal[...,simulation_step], p_settle)\n",
    "        self.municipal[...,simulation_step] -= to_settle\n",
    "        self.municipal[self.city_index,...,simulation_step] += to_settle.sum(axis=0)\n",
    "\n",
    "class ReturnHomeTravel(Travel):\n",
    "    def __init__(self):\n",
    "        pass\n",
    "\n",
    "    def step(self,\n",
    "            beta:float,\n",
    "            delta_t: float,\n",
    "            p_infectious: float,\n",
    "            p_recovery: float,\n",
    "            simulation_step: int):\n",
    "        for i,city in enumerate(self.cities):\n",
    "            city.step_internal(beta,\n",
    "                               delta_t,\n",
    "                               p_infectious,\n",
    "                               p_recovery,\n",
    "                               simulation_step)\n",
    "        for i, city_i in enumerate(self.cities[:-1]):\n",
    "            assert isinstance(city_i, ReturnHomeCity)\n",
    "            for j, city_j in enumerate(self.cities[i+1:]):\n",
    "                j += i+1\n",
    "                mix = self.mixmatrix[i,j]\n",
    "                if mix == 0:\n",
    "                    continue\n",
    "                travel_ij = city_i.select_travellers(mix, j, simulation_step)\n",
    "                travel_ji = city_j.select_travellers(mix, i, simulation_step)\n",
    "                \n",
    "                city_i.municipal[...,simulation_step] += travel_ji - travel_ij\n",
    "                city_i.arrivals[...,simulation_step] += travel_ji\n",
    "                city_i.departures[...,simulation_step] += travel_ij\n",
    "                \n",
    "                city_j.municipal[...,simulation_step] += travel_ij - travel_ji\n",
    "                city_j.arrivals[...,simulation_step] += travel_ij\n",
    "                city_j.departures[...,simulation_step] += travel_ji\n",
    "            city_i.select_settlers(simulation_step)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class HomogeneousNetwork(Travel):\n",
    "    def __init__(self,\n",
    "                 citytype: Type,\n",
    "                 n_cities: int,\n",
    "                 mixnumber: int,\n",
    "                 *args, **kwargs):\n",
    "        cities = []\n",
    "        for i in range(n_cities):\n",
    "            cities.append(citytype(*args, **kwargs))\n",
    "        mixmatrix = (np.ones((n_cities,n_cities)) - np.identity(n_cities)).astype(np.int64) * mixnumber\n",
    "        super().__init__(cities, mixmatrix)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filters = create_arrival_municipal_filter()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "basic_datasets = {}\n",
    "mixing_rates = np.logspace(0,6,31,base=10)\n",
    "for mixing_rate in mixing_rates:\n",
    "    key = f'Mixing_Rate: {int(mixing_rate)}'\n",
    "    sim = HomogeneousNetwork(BasicCity,2,mixnumber=mixing_rate)\n",
    "    dataset = sim.multiple_sims(0.04,70,measles,100,1000,moving_avg = False)\n",
    "    daily = dataset.daily_avg()\n",
    "    basic_datasets[key] = daily\n",
    "    print(f'Mixing_Rate {int(mixing_rate)} complete')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ff_datasets2 = {}\n",
    "group_LRs = [1,2,3,6,10,20,30,60,100,200,300,600,1000]\n",
    "for group_LR in group_LRs:\n",
    "    key = f'group_LR: {group_LR}'\n",
    "    sim = HomogeneousNetwork(FrequentFlyerCity,2,mixnumber=3000,p_ff = 0.99, frequent_flyer_frac = 0.01, group_LR = group_LR)\n",
    "    dataset = sim.multiple_sims(0.04,70,measles,100,1000,moving_avg = False)\n",
    "    daily = dataset.daily_avg()\n",
    "    ff_datasets2[key] = daily\n",
    "    print(f'group_LR {group_LR} complete')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets_flyingLR = {}\n",
    "flying_LRs = np.array([1,3,10,30,100,300,1000,3000,10000,30000,100000])\n",
    "frequent_flyer_fracs = np.nan_to_num((flying_LRs**0.5 - 1)/(flying_LRs-1),nan=0.5)\n",
    "for frac, LR in zip(frequent_flyer_fracs,flying_LRs):\n",
    "    sim = HomogeneousNetwork(FrequentFlyerCity,2,mixnumber=3000,flying_LR = LR, frequent_flyer_frac = frac, group_LR = 1)\n",
    "    dataset = sim.multiple_sims(0.04,70,measles,100,1000,moving_avg = False)\n",
    "    daily = dataset.daily_avg()\n",
    "    datasets_flyingLR[LR] = daily\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "flying_LRs = np.array([300000,1000000])\n",
    "frequent_flyer_fracs = np.nan_to_num((flying_LRs**0.5 - 1)/(flying_LRs-1),nan=0.5)\n",
    "for frac, LR in zip(frequent_flyer_fracs,flying_LRs):\n",
    "    sim = HomogeneousNetwork(FrequentFlyerCity,2,mixnumber=3000,flying_LR = LR, frequent_flyer_frac = frac, group_LR = 1)\n",
    "    dataset = sim.multiple_sims(0.04,70,measles,100,1000,moving_avg = False)\n",
    "    daily = dataset.daily_avg()\n",
    "    datasets_flyingLR[LR] = daily"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "flying_LRs = np.array([1,3,10,30,100,300,1000,3000,10000,30000,100000])\n",
    "\n",
    "np.nan_to_num((flying_LRs**0.5 - 1)/(flying_LRs-1),nan=0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f1 = filters[1].copy()\n",
    "f1['groups'] = ['normal']\n",
    "f2 = filters[1].copy()\n",
    "f2['groups'] = ['frequent_flyers']\n",
    "f3 = filters[0].copy()\n",
    "f3['groups'] = ['normal']\n",
    "f4 = filters[0].copy()\n",
    "f4['groups'] = ['frequent_flyers']\n",
    "\n",
    "study = [f1,f2,f3,f4]\n",
    "study = [filters[0], filters[1],f1,f2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_avg_vals({k:v for k,v in datasets_flyingLR.items() if k in [1000000]},study, log = 'y')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ff_datasets = {}\n",
    "group_LRs = [1,3,10,30,100,300,1000,3000]\n",
    "for group_LR in group_LRs:\n",
    "    key = f'group_LR: {group_LR}'\n",
    "    sim = HomogeneousNetwork(FrequentFlyerCity,2,mixnumber=3000,p_ff = 0.99, frequent_flyer_frac = 0.01, group_LR = group_LR)\n",
    "    dataset = sim.multiple_sims(0.04,70,measles,100,1000,moving_avg = False)\n",
    "    daily = dataset.daily_avg()\n",
    "    ff_datasets[key] = daily\n",
    "    print(f'group_LR {group_LR} complete')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thresholds = np.array([0.2])\n",
    "calculation = differences_vs_threshold_data(basic_datasets, filters, thresholds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thresholds = 10**np.linspace(-3.5,-0.7,29)\n",
    "calculated_data3 = differences_vs_threshold_data(ff_datasets2, filters, thresholds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "differences_vs_variable(ff_datasets2,group_LRs,calculated_data3,filters,thresholds, 10**np.linspace(-3,-1,3),log= 'x', x_name='Mixing Likelihood Ratio')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rounded_fracs = ['%s' % float('%.1g' % num) for num in frequent_flyer_fracs]\n",
    "datasets_flyingLR_withkey = {f'Flying LR: {k}, ff frac: {rounded_fracs[i]}':v for i,(k,v) in enumerate(datasets_flyingLR.items())}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "differences_vs_threshold(datasets_flyingLR_withkey,calculated_data4,filters,thresholds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LRs = list(datasets_flyingLR.keys())[:-2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LRs=10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fracs = np.nan_to_num((np.array(LRs)**0.5 - 1)/(np.array(LRs)-1),nan=0.5)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "group_LRs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "differences_vs_variable(ff_datasets2,group_LRs,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "differences_vs_variable({k:v for k,v in datasets_flyingLR.items() if k in LRs},fracs.tolist(), calculated_data4, filters,thresholds, np.logspace(-5,-1,5),x_name='Frequent Flyer Fraction', log = 'x')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thresholds = 10**np.linspace(-7,-0.7,40)\n",
    "ff_calculated = differences_vs_threshold_data(ff_datasets2, filters, thresholds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LRs = [int(k.split(' ')[-1]) for k in ff_datasets2.keys()]\n",
    "differences_vs_variable(ff_datasets2,LRs, ff_calculated, filters,thresholds, np.logspace(-7,-1,7),x_name='group_LR')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reduced = {k:v for k,v in ff_datasets2.items() if any([i in k for i in ['1','10','100','1000']])}\n",
    "reduced_calculated = ff_calculated[:,0:8:2]\n",
    "differences_vs_threshold(ff_datasets2, ff_calculated,filters,thresholds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "group_filters = []\n",
    "for f in filters:\n",
    "    for group in ff_datasets['group_LR: 100'].all_labels['groups']:\n",
    "        g = f.copy()\n",
    "        g['groups'] = [group]\n",
    "        group_filters.append(g)\n",
    "group_filters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filters = create_arrival_municipal_filter(city = 1)\n",
    "filters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_avg_vals({k:v for k,v in ff_datasets2.items() if k in ['group_LR: 1', 'group_LR: 100']},group_filters[2:],log = 'y', error_bars='std')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "differences_vs_mixnumber(new_datasets, mixnumbers, calculated_data2, filters=filters,old_thresholds = thresholds, new_thresholds = 10**np.linspace(-5,-0.7,6))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "list(basic_datasets.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "differences_vs_variable_final_img(basic_datasets, list(mixing_rates)[1:], calculation,filters=filters,old_thresholds = thresholds, new_thresholds = thresholds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "differences_vs_variable(new_datasets, mixnumbers, calculated_data2,filters=filters,old_thresholds = thresholds, new_thresholds = 10**np.linspace(-5,-0.7,6))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = np.arange(10)\n",
    "np.argmin()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tuple(np.arange((24)).reshape(2,3,4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "differences_vs_threshold(new_datasets, *calculated_data, filters = filters, thresholds=thresholds, log = 'x', error_bars='std')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "differences_vs_threshold(list(new_datasets.values()),filters,legend_labels=list(new_datasets.keys()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = np.logspace(-3,-0.7,100,base=10)\n",
    "a[:20]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.all_labels['datatypes']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "quartiles(sim.cities[0].municipal[:,2,:,0],axis=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "investigate = data[create_filter(cities=[0],datatypes=['municipal'],compartments=['I'],times=[0.4802744425385934])]\n",
    "investigate.array.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test = HomogeneousNetwork(BasicCity, 2, mixnumber=1000)\n",
    "testdata = test.multiple_sims(0.04,70,measles,100,1000)\n",
    "testdata = testdata.daily_avg()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f\"{1.3456:1g}\":g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f'{float(f\"{1.3456:.3g}\"):g}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "threshold_ratio_diffs(testdata,filters,n_points = 200, error_bars='IQR')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "threshold_ratio_diffs(testdata,filters,n_points = 50, xlims_log = [-1,-2],error_bars='IQR')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "threshold_ratio_diffs(testdata,filters,n_points = 100, xlims_log = [-0.7,-1.8,-2.9],error_bars='std')#,figsavename='threshold_comparison.png')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "threshold_ratio_diffs(testdata,filters,n_points = 50, xlims_log = -2,error_bars='IQR')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f1 = filters[1].copy()\n",
    "f1['cities'] = [0]\n",
    "new_filters = [f1,filters[1].copy()]\n",
    "new_filters\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = 50\n",
    "different_thresholds_diffs(testdata, new_filters,np.logspace(-6.5,-0.7,bins,base=10), np.logspace(-6.5,-0.7,bins,base=10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filters\n",
    "new_filters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = 300\n",
    "different_thresholds_diffs(testdata2, new_filters,np.logspace(-6.5,-0.7,bins,base=10), np.logspace(-6.5,-0.7,bins,base=10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = ff_datasets2['group_LR: 1000']\n",
    "# thresholds = 10**np.linspace(-6,-0.7,8)\n",
    "# print(thresholds)\n",
    "bins = 300\n",
    "different_thresholds_diffs(data, filters, np.logspace(-3.8,-0.7,bins,base=10), np.logspace(-7,-0.7,bins,base=10), zero_centre=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.9.7"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "2faa98f88e05bac384706a76e854f856df6a3f5df786660c064036b10a55a4b5"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}