{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Basic workflow\n",
    "\n",
    "This page demonstrates how to perform basic operations in Atomica. First, we will set up the notebook environment - the commands below are typically not required in user scripts:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "import sys\n",
    "sys.path.append('..')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To start with, import Atomica itself. It is often also useful to import `numpy` and `matplotlib`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import atomica as at\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Starting an application\n",
    "\n",
    "The first step in starting a new application is to write a Framework file. This can be done by copying one of the templates in the `atomica/library` folder (either `framework_template.xlsx` or `framework_template_advanced.xlsx`) and implementing your model. Further guidance on this is provided separately in the framework documentation. \n",
    "\n",
    "After writing the Framework, the next step is to generate a databook. This is performed in three steps\n",
    "\n",
    "1. Load the framework into a `ProjectFramework` Python instance\n",
    "2. Use the framework to make a new `ProjectData` instance\n",
    "3. Save the `ProjectData` instance to a spreadsheet\n",
    "\n",
    "In this example, we will load an existing framework from the library. You can use `at.LIBRARY_PATH` to refer to the folder containing the library Excel files shipped with Atomica:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "F = at.ProjectFramework(at.LIBRARY_PATH / 'tb_framework.xlsx') # Load the Framework\n",
    "D = at.ProjectData.new(F,pops=2, tvec=np.arange(2000,2018), transfers=0) \n",
    "D.save('new_databook.xlsx')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `ProjectData` class in Python can be thought of as an equivalent representation of the databook - you can edit the databook in Excel, which will result in changes to the `ProjectData` variable when the spreadsheet is loaded, and you can modify the `ProjectData` in Python and then write a modified spreadsheet. `ProjectData` has a number of methods that you can use to modify the databook, to do things like\n",
    "\n",
    "- Add or remove populations\n",
    "- Change the time span of the databook\n",
    "\n",
    "To perform these operations, you can load in a databook using `ProjectData.from_spreadsheet()`. This lets you load in a databook given a particular framework. It is not required that the databook be completed prior to loading - you only need to complete the databook in its entirity if you want to use the databook in a project. So for example, to add an additional population and a transfer to this newly created databook, we could use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "D = at.ProjectData.from_spreadsheet('new_databook.xlsx', framework=F)\n",
    "D.add_pop('pris','Prisoners')\n",
    "D.add_transfer('aging','Aging')\n",
    "D.save('new_databook_2.xlsx')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a project\n",
    "\n",
    "Once you have completed the framework file and databook, you can create a project that can be used to run simulations and analyses. To do this, simply create a `Project` instance, passing in the file names for the framework and databook. Here we will use a pre-filled databook from the library:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "P = at.Project(framework=at.LIBRARY_PATH / 'tb_framework.xlsx', databook=at.LIBRARY_PATH / 'tb_databook.xlsx')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When you create a project, a default simulation is automatically run. You can subsequently run simulations using `P.run_sim()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = P.run_sim(parset='default', result_name='Default parset')\n",
    "P.results.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When you run a simulation, by default it is automatically copied into the project, as well as being returned. Specifying the result name is optional, but recommended because it helps to keep track of the simulations when comparing and plotting them. We can now plot the result to show the compartment sizes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d = at.PlotData(res,pops='0-4',project=P)\n",
    "at.plot_series(d,plot_type='stacked', data=P.data, legend_mode='separate');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For full details on plotting, please refer to the full plotting documentation [here](https://atomica.tools/docs/master/examples/Plotting.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calibrating the model\n",
    "\n",
    "Model calibration can be performed in one of two ways - either manually, or automatically\n",
    "\n",
    "### Manual calibration\n",
    "\n",
    "Manual calibration of the model proceeds in three steps\n",
    "\n",
    "1. Make a new ParameterSet (e.g., by copying an existing one)\n",
    "2. Modify the calibration scale factors in that ParameterSet\n",
    "3. Run a simulation using the new parameter set\n",
    "\n",
    "The commands to do this are shown below, for an example where the force of infection has been decreased:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_parset = P.parsets.copy('default','manually_calibrated')\n",
    "new_parset.pars['foi_out'].meta_y_factor = 0.8 # Decrease infectiousness of all populations\n",
    "new_parset.pars['foi_in'].y_factor['0-4'] = 2.0 # Increase susceptibility of young children\n",
    "res_manually_calibrated = P.run_sim(parset='manually_calibrated', result_name='Manually calibrated')\n",
    "d = at.PlotData([res,res_manually_calibrated],outputs='ac_inf',project=P)\n",
    "at.plot_series(d, axis='results');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Automatic calibration\n",
    "\n",
    "To perform an automatic calibration, simply use `P.calibrate()` specifying the amount of time to run the calibration for, and the name of the new calibrated parset to create."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "P.calibrate(max_time=10, parset='default', new_name=\"auto_calibrated\", save_to_project=True);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can then run a simulation with the calibrated parset by passing the name of the new parset to `run_sim`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res_auto_calibrated = P.run_sim(parset='auto_calibrated',result_name='Automatically calibrated')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding programs\n",
    "\n",
    "The programs system allows parameter values to be overwritten based on spending on a set of programs. To get started, you will first need a program spreadsheet (progbook). The progbook is specific to a framework and a databook, because it refers to both the compartments and parameters of the model (from the framework) as well as the populations (from the databook). \n",
    "\n",
    "You can make a new progbook using the `.make_progbook()` method of the project:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "P.make_progbook(progbook_path='example_progbook.xlsx', progs=4, data_start=2014, data_end=2018)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After filling out the progbook, you can load it into the project using the `.load_progbook()` method. Here, we will load in a pre-filled progbook from the library:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "P.load_progbook(at.LIBRARY_PATH / 'tb_progbook.xlsx')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This progbook has been added to the list of available progsets:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "P.progsets.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Running a simulation with programs requires one additional piece of information - a `ProgramInstructions` instance that specifies\n",
    "\n",
    "- What years the programs are active\n",
    "- Any overwrites to spending or coverage\n",
    "\n",
    "In our case, we might just want to run a simulation with programs starting in 2018, so we can create a `ProgramInstructions` instance accordingly, and then use it to run the simulation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "instructions = at.ProgramInstructions(start_year=2018)\n",
    "res_progs = P.run_sim(parset='default',progset='default',progset_instructions=instructions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reconciliation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reconciliation is an operation that aims to change the properties of programs (such as their unit costs) such that the program-calculated parameter values optimally match the databook parameter values in the year the programs become active (or some other specified year). The reconciliation operation can therefore be treated as a mapping from one progset to another. To perform reconciliation, use the `reconcile` function directly, passing in:\n",
    "\n",
    "- the parameter set you want to match\n",
    "- the program set to modify\n",
    "- the reconciliation year\n",
    "- a specification of which aspects of the program set to modify (e.g. unit cost, program outcomes)\n",
    "\n",
    "The reconcile function returns a new progset, which you can store in the project if desired, or otherwise work with independently:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "P.progsets['reconciled'] = at.reconcile(project=P, parset='default', progset='default', reconciliation_year=2018, unit_cost_bounds=0.05)[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can then run simulations with the modified program set. You can also save the new programset to a progbook if you wish to edit it further in Excel:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "P.progsets['reconciled'].save('reconciled_progset.xlsx');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scenarios\n",
    "\n",
    "A scenario involves overriding some aspect of the simulation that would otherwise be specified in the databook or progbook. There are three kinds of scenarios\n",
    "\n",
    "- Parameter scenarios, when you want to test the effect of a specific parameter value\n",
    "- Budget scenarios, when you want to examine the outcomes of specific spending values\n",
    "- Coverage scenarios, when you want to examine the effect of specific program coverages irrespective of spending\n",
    "\n",
    "Examples of these scenarios are shown below:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameter scenarios"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scvalues = dict()\n",
    "scvalues['b_rate'] = dict()\n",
    "scvalues['b_rate']['0-4'] = dict()\n",
    "scvalues['b_rate']['0-4'][\"t\"] = [2015, 2020, 2035]\n",
    "scvalues['b_rate']['0-4'][\"y\"] = [270000, 220000, 220000]\n",
    "scen = P.make_scenario(which='parameter', name=\"Reduced births\", scenario_values=scvalues)\n",
    "res_par_scen = scen.run(P, P.parsets[\"default\"]);\n",
    "\n",
    "# Plot the parameter and compare to scenario input values\n",
    "d = at.PlotData(res_par_scen,outputs='b_rate',pops='0-4')\n",
    "at.plot_series(d)\n",
    "plt.scatter(scvalues['b_rate']['0-4'][\"t\"],scvalues['b_rate']['0-4'][\"y\"],color='r')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Budget scenarios\n",
    "\n",
    "To run a program-related scenario, such as a budget or coverage scenario, it is not necessary to construct a `Scenario` object. Instead, you can directly create and use the program instructions that define the scenario:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alloc = P.progsets[0].get_alloc(2018)\n",
    "doubled_budget = {x:v*2 for x,v in alloc.items()}\n",
    "instructions = at.ProgramInstructions(start_year=2018,alloc=doubled_budget)\n",
    "res_baseline = P.run_sim(parset='default',progset='default',progset_instructions=at.ProgramInstructions(start_year=2018),result_name='Baseline')\n",
    "res_budget_scen = P.run_sim(parset='default',progset='default',progset_instructions=instructions,result_name='Doubled');\n",
    "\n",
    "d = at.PlotData.programs([res_baseline,res_budget_scen]).interpolate(2018)\n",
    "at.plot_bars(d,stack_outputs='all');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, you can create a full scenario object by storing the `instructions` in a `CombinedScenario`. The `CombinedScenario` optionally allows you to mix parameter and program scenarios."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scen = P.make_scenario(which='combined', name=\"Doubled (scen)\", instructions=instructions)\n",
    "res_combined_scen = scen.run(P, parset='default',progset='default')\n",
    "\n",
    "d = at.PlotData.programs([res_baseline,res_budget_scen, res_combined_scen]).interpolate(2018)\n",
    "at.plot_bars(d,stack_outputs='all');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Coverage scenarios\n",
    "\n",
    "With coverage scenarios, the program instructions override a program's coverage. Therefore, the spending values and coverage values may not match up with what is entered in the program book. If running coverage scenarios, take care not to use the spending values for such results - typically this is not a problem, because if you did have a particular spending amount in mind, then it would be better to use a budget scenario. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "half_coverage = {x:0.5 for x in P.progsets[0].programs.keys()}\n",
    "instructions = at.ProgramInstructions(start_year=2018,coverage=half_coverage)\n",
    "scen = at.CombinedScenario(name='Reduced coverage',instructions=instructions)\n",
    "res_cov_scen = scen.run(P,parset='default',progset='default');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimization\n",
    "\n",
    "The role of optimization is to produce a set of program spending overwrites that improves the model output in some way. It is thus an operation that maps one set of program instructions to another, where the optimized program instructions contain the optimized allocation. An optimization consists of three parts\n",
    "\n",
    "- `adjustments` that specify what parts of the program instructions to change, and how to change them\n",
    "- `measurables` that define optimality (e.g. reducing new infections, maximizing people alive)\n",
    "- `constraints` that must be satisfied, such as fixed total spending\n",
    "\n",
    "An `Optimization` object contains these three items, as well any additional parameters specific to the optimization algorithm (e.g. the optimization method, the maximum run time).\n",
    "\n",
    "The `optimize` function uses the `Optimization` to modify a particular set of program instructions. It therefore takes in\n",
    "\n",
    "- A parset and progset to use\n",
    "- A program instructions instance to optimize\n",
    "- An optimization object, that specifies how to perform the optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "instructions = at.ProgramInstructions(alloc=P.progsets[0],start_year=2020) # Instructions for default spending\n",
    "adjustments = [at.SpendingAdjustment(x,2020,'rel',0.,2.) for x in instructions.alloc.keys()]\n",
    "measurables = at.MaximizeCascadeStage(None,2020)\n",
    "constraints = at.TotalSpendConstraint() # Cap total spending in all years\n",
    "optimization = at.Optimization(name='default', adjustments=adjustments, measurables=measurables,constraints=constraints,maxtime=10) # Evaluate from 2020 to end of simulation\n",
    "optimized_instructions = at.optimize(P, optimization, parset=P.parsets[\"default\"], progset=P.progsets['default'], instructions=instructions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function returns a set of optimized instructions, that can then be used to run a simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res_optimized = P.run_sim(parset='default',progset='default',progset_instructions=optimized_instructions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For more details on the optimization system, see the [general documentation](https://atomica.tools/docs/master/general/Optimization.html) on optimization."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:atomica311]",
   "language": "python",
   "name": "conda-env-atomica311-py"
  },
  "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.11.5"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}