Repository: starsimhub/starsim (language: python)

Issue:
Improving pregnancy configuration and fixing test_randomness
Fixing demographics so that users can override the fertility rate parameter even if using a subclass. Now `self.pars.p_fertility` in `init_pre` rather than doing everything in `__init__`.

Also changing `uint64` to signed `int64` in test_randomness.

### Description


### Checklist
- [ ] Code commented & docstrings added
- [ ] New tests were needed and have been added
- [ ] A new version number was needed & changelog has been updated
- [ ] A new PyPI version needs to be released


Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index c6980976..3d8035bb 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -7,6 +7,19 @@ What's new
 All notable changes to the codebase are documented in this file. Changes that may result in differences in model output, or are required in order to run an old parameter set with the current version, are flagged with the term "Regression information".
 
 
+Version 1.0.0 (2024-07-10)
+---------------------------
+- Official release of Starsim!
+- Adds a ``Calibration`` class, based on `Optuna <https://optuna.org>`_, to facilitate the calibration of Starsim models.
+- Adds ``mean()``, ``median()``, and ``plot()`` methods to ``MultiSim``.
+- Adds ``low`` and ``high`` attributes to ``Result`` objects.
+- Adds a ``flatten()`` method to ``Results``, allowing nested ``Results`` objects to be turned into flat dictionaries.
+- Removes duplicate UIDs among new infections, and adds a ``unique()`` method to ``ss.uids``.
+- Fixes a bug that prevented ``ss.lognorm_im()`` from using callable parameters.
+- Updates the default ``Sim`` string representation to be a single line; the more verbose version is available via ``sim.disp()``.
+- *GitHub info*: PR `581 <https://github.com/starsimhub/starsim/pull/581>`
+
+
 Version 0.5.10 (2024-07-03)
 ---------------------------
 - Adds two new common-random-number-safe networks. The first is an Erdős-Rényi network that is similar to ``RandomNet`` but parameterized differently. The second is a 2D spatial network with connectivity between agents within a given radius; these agents can also optionally move.
diff --git a/README.rst b/README.rst
index e8e1caf3..dbf26f01 100644
--- a/README.rst
+++ b/README.rst
@@ -1,13 +1,13 @@
 Starsim
 =======
 
-**Warning! Starsim is still in the early stages of development. It is being shared solely for transparency and to facilitate collaborative development. It is not ready to be used for real research or policy questions.**
+Starsim is an agent-based modeling framework for simulating the spread of diseases among agents via dynamic transmission networks. Starsim supports the co-transmission of multiple diseases at once, capturing how they interact biologically and behaviorally. Users can also include non-infectious diseases either on their own or as factors affecting infectious diseases. Starsim allows detailed modeling of mother-child relationships from conception to study birth-related diseases. Additionally, Starsim lets users compare different intervention strategies, like vaccines or treatments, to see their impact through various delivery methods such as mass campaigns or targeted outreach.
 
-Starsim is an agent-based modeling framework in which users can design and configure simulations of diseases (or other health states) that progress over time within each agent and pass from one agent to the next along dynamic transmission networks. The framework explicitly supports co-transmission of multiple pathogens, allowing users to concurrently simulate several diseases while capturing behavioral and biological interactions. Non-communicable diseases can be included as well, either as a co-factor for transmissible pathogens or as an independent exploration. Detailed modeling of mother-child relationships can be simulated from the timepoint of conception, enabling study of congenital diseases and associated birth outcomes. Finally, Starsim facilitates the comparison of one or more intervention scenarios to a baseline scenario in evaluating the impact of various products like vaccines, therapeutics, and novel diagnostics delivered via flexible routes including mass campaigns, screen and treat, and targeted outreach.
+Examples of diseases that have already been implemented in Starsim include sexually transmitted infections (HIV, HPV, and syphilis, including co-transmission), respiratory infections (tuberculosis and RSV), plus other diseases (Ebola and cholera) and underlying determinants of health (such as malnutrition).
 
-The framework is appropriate for simulating sexually transmitted infections (including syphilis, gonorrhea, chlamydia, HPV, and HIV, including co-transmission), respiratory infections (like RSV and tuberculosis), and other diseases and underlying determinants (such as Ebola, diabetes, and malnutrition).
+Note: Starsim is a general-purpose, multi-disease framework that builds on our previous suite of disease-specific models, which included `Covasim <https://covasim.org>`_, `HPVsim <https://hpvsim.org>`_, and `FPsim <https://fpsim.org>`_. In cases where a distinction needs to be made, Starsim is also known as the "Starsim framework", while this collection of other models is known as the "Starsim suite".
 
-Starsim is a general-purpose modeling framework that is part of the same suite of tools as `Covasim <https://covasim.org>`_, `HPVsim <https://hpvsim.org>`_, and `FPsim <https://fpsim.org>`_.
+For more information about Starsim, please see the `documentation <https://docs.starsim.org>`__.
 
 
 Requirements
@@ -29,32 +29,38 @@ Starsim can also be installed locally. To do this, clone first this repository,
 Usage and documentation
 -----------------------
 
-Documentation, including tutorials and an API reference, is available at https://docs.starsim.org. 
+Full documentation, including tutorials and an API reference, is available at https://docs.starsim.org. 
 
-If everything is working, the following Python commands will run a simulation with the simplest version of a Starsim model. We'll make a version of a classic SIR model::
+You can run a simple demo via::
 
   import starsim as ss
-  
+  ss.demo()
+
+Here is a slightly more realistic example of an SIR model with random connections between agents::
+
+  import starsim as ss
+
   # Define the parameters
   pars = dict(
-    n_agents = 5_000,   # Number of agents to simulate
-    networks = dict(    # *Networks* add detail on how agents interact w/ each other
-      type = 'random',  # Here, we use a 'random' network
-      n_contacts = 10   # Each person has an average of 10 contacts w/ other people  
-    ),
-    diseases = dict(    # *Diseases* add detail on what diseases to model
-      type = 'sir',     # Here, we're creating an SIR disease
-      init_prev = 0.1,  # Proportion of the population initially infected
-      beta = 0.5,       # Probability of transmission between contacts
-    )
+      n_agents = 5_000,     # Number of agents to simulate
+      networks = dict(      # Networks define how agents interact w/ each other
+          type = 'random',  # Here, we use a 'random' network
+          n_contacts = 10   # Each person has 10 contacts with other people  
+      ),
+      diseases = dict(      # *Diseases* add detail on what diseases to model
+          type = 'sir',     # Here, we're creating an SIR disease
+          init_prev = 0.01, # Proportion of the population initially infected
+          beta = 0.05,      # Probability of transmission between contacts
+      )
   )
-  
+
   # Make the sim, run and plot
   sim = ss.Sim(pars)
   sim.run()
-  sim.plot()
+  sim.plot() # Plot all the sim results
+  sim.diseases.sir.plot() # Plot the standard SIR curves
 
-More usage examples are available in the ``tests`` folder.
+More usage examples are available in the tutorials, as well as the ``tests`` folder.
 
 
 Model structure
@@ -66,6 +72,7 @@ The model consists of core classes including Sim, Run, People, State, Network, C
 
 The structure of the starsim folder is as follows, roughly in the order in which the modules are imported, building from most fundamental to most complex:
 
+• ``calibration.py``: Class to handle automated calibration of the model to data.
 •	``demographics.py``: Classes to transform initial condition input parameters for use in building and utilizing networks.
 •	``disease.py``: Classes to manage infection rate of spread, prevalence, waning effects, and other parameters for specific diseases.
 •	``distributions.py``: Classes that handle statistical distributions used throughout Starsim.
@@ -84,7 +91,7 @@ The structure of the starsim folder is as follows, roughly in the order in which
 •	``utils.py``: Helper functions.
 •	``version.py``: Version, date, and license information.
 
-The ``diseases`` folder within the Starsim package contains loading scripts for the epidemiological data specific to each respective disease.
+The ``diseases`` folder within the Starsim package contains definitions of different types of diseases, including STIs, Ebola, and cholera.
 
 
 Contributing
@@ -96,6 +103,4 @@ Questions or comments can be directed to `info@starsim.org <mailto:info@starsim.
 Disclaimer
 ----------
 
-The code in this repository was developed by IDM, the Burnet Institute, and other collaborators to support our joint research on flexible agent-based modeling. We've made it publicly available under the MIT License to provide others with a better understanding of our research and an opportunity to build upon it for their own work. We make no representations that the code works as intended or that we will provide support, address issues that are found, or accept pull requests. You are welcome to create your own fork and modify the code to suit your own modeling needs as permitted under the MIT License.
-
-
+The code in this repository was developed by `IDM <https://idmod.org>`_, the `Burnet Institute <https://burnet.edu.au>`_, and other collaborators to support our joint research on flexible agent-based modeling. We've made it publicly available under the MIT License to provide others with a better understanding of our research and an opportunity to build upon it for their own work. We make no representations that the code works as intended or that we will provide support, address issues that are found, or accept pull requests. You are welcome to create your own fork and modify the code to suit your own modeling needs as permitted under the MIT License.
diff --git a/docs/tutorials.rst b/docs/tutorials.rst
index 574924b2..88f9334e 100644
--- a/docs/tutorials.rst
+++ b/docs/tutorials.rst
@@ -20,4 +20,7 @@ Developer tutorials
     :maxdepth: 1
 
     tutorials/dev_tut_diseases.ipynb
+    tutorials/dev_tut_pars.ipynb
+    tutorials/dev_tut_analyzers.ipynb
+    tutorials/dev_tut_abm.ipynb
     
diff --git a/docs/tutorials/dev_tut_abm.ipynb b/docs/tutorials/dev_tut_abm.ipynb
new file mode 100644
index 00000000..17067680
--- /dev/null
+++ b/docs/tutorials/dev_tut_abm.ipynb
@@ -0,0 +1,158 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Developer tutorial: ABM examples\n",
+    "\n",
+    "This tutorial describes how Starsim can be used to reproduce examples from other Python ABMs.\n",
+    "## Mesa: Wealth model\n",
+    "\n",
+    "Although Starsim is intended primarily for modeling disease, it can also be used as a general-purpose agent-based model.  This example illustrates a simple \"weath model\", in which each agent starts with a single unit of wealth, and on each timestep, every agent with more than zero wealth gives one unit of wealth to another agent.\n",
+    "\n",
+    "This tutorial is adapted from the following example:\n",
+    "\n",
+    "https://mesa.readthedocs.io/en/stable/tutorials/intro_tutorial.html\n",
+    "\n",
+    "### Setting up the model\n",
+    "\n",
+    "We could define the wealth model as any type of module, since they all can store states and update them. Here we will define wealth as a subclass of `ss.Intervention` (though it could equally well be a subclass of `ss.Demographics` or even `ss.Disease`, if you are so inclined). All we need to do is update the `wealth` state (which we can store inside the \"intervention\"), and we can also use this class to track the wealth distribution over time and plot it. The full model looks like this:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Imports\n",
+    "import numpy as np\n",
+    "import starsim as ss\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "# Define the model\n",
+    "class WealthModel(ss.Intervention):\n",
+    "    \"\"\" A simple wealth transfer model\"\"\"\n",
+    "    \n",
+    "    def init_post(self, bins=10):\n",
+    "        \"\"\" Define custom model attributes \"\"\"\n",
+    "        super().init_post()\n",
+    "        self.npts = self.sim.npts # Number of timepoints\n",
+    "        self.n_agents = len(sim.people) # Number of agents\n",
+    "        self.wealth = np.ones(self.n_agents) # Initial wealth of each agent\n",
+    "        self.bins = np.arange(bins+1) # Bins used for plotting\n",
+    "        self.wealth_dist = np.zeros((self.npts, len(self.bins)-1)) # Wealth distribution over time\n",
+    "        return\n",
+    "        \n",
+    "    def apply(self, sim):\n",
+    "        \"\"\" Transfer wealth between agents -- core model logic \"\"\"\n",
+    "        self.wealth_hist() # Store the wealth at this time point\n",
+    "        givers = self.wealth > 0 # People need wealth to be givers\n",
+    "        receivers = np.random.choice(sim.people.uid, size=givers.sum()) # Anyone can be a receiver\n",
+    "        self.wealth[givers] -= 1 # Givers are unique, so can use vectorized version\n",
+    "        for receive in receivers: # Vectorized version is: np.add.at(sim.people.wealth.raw, receivers, 1)\n",
+    "            self.wealth[receive] += 1\n",
+    "        return\n",
+    "    \n",
+    "    def wealth_hist(self):\n",
+    "        \"\"\" Calculate the wealth histogram \"\"\"\n",
+    "        ti = self.sim.ti # Current timestep\n",
+    "        self.wealth_dist[ti,:], _ = np.histogram(self.wealth, bins=self.bins)\n",
+    "        return\n",
+    "    \n",
+    "    def plot(self):\n",
+    "        \"\"\" Plot a 2D histogram of the final wealth distribution \"\"\"\n",
+    "        plt.figure()\n",
+    "        plt.bar(self.bins[:-1], self.wealth_dist[-1,:])\n",
+    "        plt.title('Wealth distribution at final time point')\n",
+    "        plt.xlabel('Wealth')\n",
+    "        plt.ylabel('Number of agents')\n",
+    "        plt.show()\n",
+    "        return\n",
+    "    \n",
+    "    def plot3d(self):\n",
+    "        \"\"\" Plot a 3D heatmap of the wealth distribution over time \"\"\"\n",
+    "        plt.figure()\n",
+    "        plt.pcolor(self.wealth_dist.T, cmap='turbo')\n",
+    "        plt.title('Wealth distribution over time')\n",
+    "        plt.xlabel('Time')\n",
+    "        plt.ylabel('Wealth')\n",
+    "        plt.colorbar().set_label('Number of agents', rotation=270)\n",
+    "        plt.show()\n",
+    "        return\n",
+    "\n",
+    "# Create sim inputs, including the wealth model\n",
+    "wealth = WealthModel()\n",
+    "pars = dict(\n",
+    "    n_agents = 100, # Number of agents\n",
+    "    start = 0,\n",
+    "    end = 100,\n",
+    "    interventions = wealth,\n",
+    ")\n",
+    "\n",
+    "# Run the model\n",
+    "sim = ss.Sim(pars, copy_inputs=False) # copy_inputs=False lets us reuse the \"wealth\" object from above\n",
+    "sim.run()\n",
+    "\n",
+    "# Plot the results\n",
+    "wealth.plot()\n",
+    "wealth.plot3d()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Comparison with Mesa\n",
+    "\n",
+    "While the implementation in Starsim is similar to Mesa, there are a couple key differences:\n",
+    "- Because Starsim's people object is vectorized, the wealth definition and update is vectorized as well.\n",
+    "- Both Mesa and Starsim versions of the model are quite simple, but there is a little less boilerplate in the Starsim version."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "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.11.7"
+  },
+  "pycharm": {
+   "stem_cell": {
+    "cell_type": "raw",
+    "metadata": {
+     "collapsed": false
+    },
+    "source": []
+   }
+  },
+  "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": 4
+}
diff --git a/docs/tutorials/dev_tut_analyzers.ipynb b/docs/tutorials/dev_tut_analyzers.ipynb
new file mode 100644
index 00000000..facfa12c
--- /dev/null
+++ b/docs/tutorials/dev_tut_analyzers.ipynb
@@ -0,0 +1,129 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Developer tutorial: Analyzers"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Reporting results\n",
+    "\n",
+    "Each Starsim module can have its own results, which get added to the full list of results in the Sim object. For example, the `ss.Pregnancy` module adds results like `sim.results.pregnancy.pregnant`, and the `ss.HIV` module adds results like `sim.results.hiv.new_infections`. If you are writing your own module, you can add whatever custom results you want. However, another option is to create an `Analyzer` to store results that you might need for one particular analysis but won't need all the time. An `Analyzer` is very similar to other Starsim modules in its structure, but the general idea of an analyzer is that it gets called at the end of a timestep, and reports of the state of things after everything else has been updated without changing any of the module states itself.\n",
+    "\n",
+    "Suppose we wanted to create an analyzer that would report on the number of new HIV infections in pregnant women:\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import starsim as ss\n",
+    "import pandas as pd\n",
+    "\n",
+    "class HIV_in_pregnancy(ss.Analyzer):\n",
+    "    def __init__(self, *args, **kwargs):\n",
+    "        super().__init__(*args, **kwargs)\n",
+    "        self.requires = [ss.HIV, ss.Pregnancy]\n",
+    "        self.name = 'hiv_in_pregnancy'\n",
+    "        return\n",
+    "    \n",
+    "    def init_pre(self, sim):\n",
+    "        super().init_pre(sim)\n",
+    "        npts = self.sim.npts\n",
+    "        self.results += [\n",
+    "            ss.Result(self.name, 'new_infections_pregnancy', npts, dtype=float, scale=True),\n",
+    "        ]\n",
+    "        return\n",
+    "\n",
+    "    def apply(self, sim):\n",
+    "        ti = sim.ti\n",
+    "        hiv = sim.diseases.hiv\n",
+    "        pregnant = sim.demographics.pregnancy.pregnant\n",
+    "        newly_infected = hiv.ti_infected == ti\n",
+    "        self.results['new_infections_pregnancy'][ti] = len((newly_infected & pregnant).uids)\n",
+    "        return\n",
+    "\n",
+    "pregnancy = ss.Pregnancy(pars=dict(fertility_rate=pd.read_csv('test_data/nigeria_asfr.csv')))\n",
+    "hiv = ss.HIV(beta={'mfnet':[0.5,0.25]})\n",
+    "sim = ss.Sim(diseases=hiv, networks='mfnet', demographics=pregnancy, analyzers=HIV_in_pregnancy())\n",
+    "sim.run()\n",
+    "print(f'Total infections among pregnant women: {sim.results.hiv_in_pregnancy.new_infections_pregnancy.sum()}')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Analyzers are ideal for adding custom results, and because they get added to the sim in the same way as any other result, they also get automatically exported in the same format, e.g. using `sim.export_df()`.\n",
+    "\n",
+    "Here's a plot of the results from our HIV in pregnancy analyzer:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "res = sim.results\n",
+    "\n",
+    "plt.figure()\n",
+    "plt.plot(res.yearvec, res.hiv_in_pregnancy.new_infections_pregnancy)\n",
+    "plt.title('HIV infections acquired during pregnancy')\n",
+    "plt.show();"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.11"
+  },
+  "pycharm": {
+   "stem_cell": {
+    "cell_type": "raw",
+    "metadata": {
+     "collapsed": false
+    },
+    "source": []
+   }
+  },
+  "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": 4
+}
diff --git a/docs/tutorials/dev_tut_diseases.ipynb b/docs/tutorials/dev_tut_diseases.ipynb
index 69e93b6f..919e04e0 100644
--- a/docs/tutorials/dev_tut_diseases.ipynb
+++ b/docs/tutorials/dev_tut_diseases.ipynb
@@ -4,7 +4,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Developer tutorial - Diseases"
+    "# Developer tutorial: Diseases"
    ]
   },
   {
diff --git a/docs/tutorials/dev_tut_pars.ipynb b/docs/tutorials/dev_tut_pars.ipynb
new file mode 100644
index 00000000..e83147bb
--- /dev/null
+++ b/docs/tutorials/dev_tut_pars.ipynb
@@ -0,0 +1,139 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Developer tutorial: Flexible parameters"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Defining default parameters\n",
+    "\n",
+    "When you create a module using Starsim, you have the opportunity to define the default format for parameters. Let's look at an example from the SIR disease model:\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import starsim as ss\n",
+    "\n",
+    "class SIR(ss.SIR):\n",
+    "    def __init__(self, pars=None, **kwargs):\n",
+    "        super().__init__()\n",
+    "        self.label = 'SIRdemo'\n",
+    "        self.default_pars(\n",
+    "            beta = 0.1,\n",
+    "            init_prev = ss.bernoulli(p=0.01),\n",
+    "            dur_inf = ss.lognorm_ex(mean=6),\n",
+    "            p_death = ss.bernoulli(p=0.01),\n",
+    "        )\n",
+    "        self.update_pars(pars, **kwargs)\n",
+    "        \n",
+    "        self.add_states(\n",
+    "            ss.BoolArr('recovered', label='Recovered'),\n",
+    "            ss.FloatArr('ti_recovered', label='Time of recovery'),\n",
+    "            ss.FloatArr('ti_dead', label='Time of death'),\n",
+    "        )\n",
+    "        return"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The point of `self.default_pars()` is to set the ground truth for the format that the parameters should take. When users enter their own parameters for defining an instance of this module, the parameter values they enter will be processed within `self.update_pars()` and will be checked for consistency with the format provided in the `default_pars`. For example, the parameter `p_death` in the SIR example above is specified within the default_pars as a Bernoulli distribution. It would be perfectly legitimate to create an instance of the SIR model using any of the following formats:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sir1 = ss.SIR(p_death=0.02)\n",
+    "sir2 = ss.SIR(p_death=ss.bernoulli(p=0.2))\n",
+    "sir3 = ss.SIR(pars=dict(p_death=0.03))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "However, it would NOT be ok to create an SIR model with e.g. `ss.SIR(p_death=ss.lognorm_ex(4))`, because if a distribution is defined as a Bernoulli in the default_pars, it can't be changed. This is only the case for Bernoulli distributions; other distributions can be changed, e.g. `ss.SIR(dur_inf=ss.normal(4))` would be okay.\n",
+    "\n",
+    "## Using callable parameters\n",
+    " \n",
+    "One of the most flexible aspects of how Starsim's distributions are defined is that they can take callable functions as parameter values. For example, in reality the duration of infection of a disease might vary by age. We could model this as follows:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sir = ss.SIR(dur_inf=ss.normal(loc=10))  # Define an SIR model with a default duration of 10 days\n",
+    "sir.pars.dur_inf.set(loc = lambda self, sim, uids: sim.people.age[uids] / 10)  # Change the mean duration so it scales with age\n",
+    "sim = ss.Sim(n_agents=1e3, diseases=sir, networks='random')\n",
+    "sim.run()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Using similar logic, any other parameter could be set to depend on anything that the sim is aware of, including time or agent properties like age, sex, or health attributes.\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.11"
+  },
+  "pycharm": {
+   "stem_cell": {
+    "cell_type": "raw",
+    "metadata": {
+     "collapsed": false
+    },
+    "source": []
+   }
+  },
+  "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": 4
+}
diff --git a/starsim/__init__.py b/starsim/__init__.py
index 1e873df3..19c32818 100644
--- a/starsim/__init__.py
+++ b/starsim/__init__.py
@@ -16,6 +16,7 @@
 from .diseases      import *
 from .sim           import *
 from .run           import *
+from .calibration   import *
 from .samples       import *
 
 # Assign the root folder
diff --git a/starsim/calibration.py b/starsim/calibration.py
new file mode 100644
index 00000000..677ce47b
--- /dev/null
+++ b/starsim/calibration.py
@@ -0,0 +1,496 @@
+"""
+Define the calibration class
+"""
+
+import os
+import numpy as np
+import pandas as pd
+import sciris as sc
+import starsim as ss
+
+
+__all__ = ['Calibration']
+
+
+def import_optuna():
+    """ A helper function to import Optuna, which is an optional dependency """
+    try:
+        import optuna as op # Import here since it's slow
+    except ModuleNotFoundError as E: # pragma: no cover
+        errormsg = f'Optuna import failed ({str(E)}), please install first (pip install optuna)'
+        raise ModuleNotFoundError(errormsg)
+    return op
+
+
+def compute_gof(actual, predicted, normalize=True, use_frac=False, use_squared=False, as_scalar='none', eps=1e-9, skestimator=None, estimator=None, **kwargs):
+    """
+    Calculate the goodness of fit. By default use normalized absolute error, but
+    highly customizable. For example, mean squared error is equivalent to
+    setting normalize=False, use_squared=True, as_scalar='mean'.
+
+    Args:
+        actual      (arr):   array of actual (data) points
+        predicted   (arr):   corresponding array of predicted (model) points
+        normalize   (bool):  whether to divide the values by the largest value in either series
+        use_frac    (bool):  convert to fractional mismatches rather than absolute
+        use_squared (bool):  square the mismatches
+        as_scalar   (str):   return as a scalar instead of a time series: choices are sum, mean, median
+        eps         (float): to avoid divide-by-zero
+        skestimator (str):   if provided, use this scikit-learn estimator instead
+        estimator   (func):  if provided, use this custom estimator instead
+        kwargs      (dict):  passed to the scikit-learn or custom estimator
+
+    Returns:
+        gofs (arr): array of goodness-of-fit values, or a single value if as_scalar is True
+
+    **Examples**::
+
+        x1 = np.cumsum(np.random.random(100))
+        x2 = np.cumsum(np.random.random(100))
+
+        e1 = compute_gof(x1, x2) # Default, normalized absolute error
+        e2 = compute_gof(x1, x2, normalize=False, use_frac=False) # Fractional error
+        e3 = compute_gof(x1, x2, normalize=False, use_squared=True, as_scalar='mean') # Mean squared error
+        e4 = compute_gof(x1, x2, skestimator='mean_squared_error') # Scikit-learn's MSE method
+        e5 = compute_gof(x1, x2, as_scalar='median') # Normalized median absolute error -- highly robust
+    """
+
+    # Handle inputs
+    actual    = np.array(sc.dcp(actual), dtype=float)
+    predicted = np.array(sc.dcp(predicted), dtype=float)
+
+    # Scikit-learn estimator is supplied: use that
+    if skestimator is not None: # pragma: no cover
+        try:
+            import sklearn.metrics as sm
+            sklearn_gof = getattr(sm, skestimator) # Shortcut to e.g. sklearn.metrics.max_error
+        except ImportError as E:
+            raise ImportError(f'You must have scikit-learn >=0.22.2 installed: {str(E)}')
+        except AttributeError:
+            raise AttributeError(f'Estimator {skestimator} is not available; see https://scikit-learn.org/stable/modules/model_evaluation.html#scoring-parameter for options')
+        gof = sklearn_gof(actual, predicted, **kwargs)
+        return gof
+
+    # Custom estimator is supplied: use that
+    if estimator is not None:
+        try:
+            gof = estimator(actual, predicted, **kwargs)
+        except Exception as E:
+            errormsg = f'Custom estimator "{estimator}" must be a callable function that accepts actual and predicted arrays, plus optional kwargs'
+            raise RuntimeError(errormsg) from E
+        return gof
+
+    # Default case: calculate it manually
+    else:
+        # Key step -- calculate the mismatch!
+        gofs = abs(np.array(actual) - np.array(predicted))
+
+        if normalize and not use_frac:
+            actual_max = abs(actual).max()
+            if actual_max>0:
+                gofs /= actual_max
+
+        if use_frac:
+            if (actual<0).any() or (predicted<0).any():
+                print('Warning: Calculating fractional errors for non-positive quantities is ill-advised!')
+            else:
+                maxvals = np.maximum(actual, predicted) + eps
+                gofs /= maxvals
+
+        if use_squared:
+            gofs = gofs**2
+
+        if as_scalar == 'sum':
+            gofs = np.sum(gofs)
+        elif as_scalar == 'mean':
+            gofs = np.mean(gofs)
+        elif as_scalar == 'median':
+            gofs = np.median(gofs)
+
+        return gofs
+
+
+class Calibration(sc.prettyobj):
+    """
+    A class to handle calibration of STIsim simulations. Uses the Optuna hyperparameter
+    optimization library (optuna.org), which must be installed separately (via
+    pip install optuna).
+    Args:
+        sim          (Sim)  : the simulation to calibrate
+        data         (df)   : pandas dataframe
+        calib_pars   (dict) : a dictionary of the parameters to calibrate of the format dict(key1=[best, low, high])
+        fit_args     (dict) : a dictionary of options that are passed to sim.compute_fit() to calculate the goodness-of-fit
+        par_samplers (dict) : an optional mapping from parameters to the Optuna sampler to use for choosing new points for each; by default, suggest_float
+        n_trials     (int)  : the number of trials per worker
+        n_workers    (int)  : the number of parallel workers (default: maximum
+        total_trials (int)  : if n_trials is not supplied, calculate by dividing this number by n_workers)
+        name         (str)  : the name of the database (default: 'hpvsim_calibration')
+        db_name      (str)  : the name of the database file (default: 'hpvsim_calibration.db')
+        keep_db      (bool) : whether to keep the database after calibration (default: false)
+        storage      (str)  : the location of the database (default: sqlite)
+        rand_seed    (int)  : if provided, use this random seed to initialize Optuna runs (for reproducibility)
+        label        (str)  : a label for this calibration object
+        die          (bool) : whether to stop if an exception is encountered (default: false)
+        verbose      (bool) : whether to print details of the calibration
+        kwargs       (dict) : passed to hpv.Calibration()
+
+    Returns:
+        A Calibration object
+    """
+    def __init__(self, sim, data, calib_pars=None, weights=None, fit_args=None, par_samplers=None, n_trials=None, n_workers=None,
+                total_trials=None, name=None, db_name=None, estimator=None, keep_db=None, storage=None, rand_seed=None,
+                 sampler=None, label=None, die=False, verbose=True):
+
+        import multiprocessing as mp
+
+        # Handle run arguments
+        if n_trials  is None: n_trials  = 20
+        if n_workers is None: n_workers = mp.cpu_count()
+        if name      is None: name      = 'starsim_calibration'
+        if db_name   is None: db_name   = f'{name}.db'
+        if keep_db   is None: keep_db   = False
+        if storage   is None: storage   = f'sqlite:///{db_name}'
+        if total_trials is not None: n_trials = int(np.ceil(total_trials/n_workers))
+        self.run_args   = sc.objdict(n_trials=int(n_trials), n_workers=int(n_workers), name=name, db_name=db_name,
+                                     keep_db=keep_db, storage=storage, rand_seed=rand_seed, sampler=sampler)
+
+        # Handle other inputs
+        self.label          = label
+        self.sim            = sim
+        self.calib_pars     = calib_pars
+        self.weights        = weights
+        self.fit_args       = sc.mergedicts(fit_args)
+        self.par_samplers   = sc.mergedicts(par_samplers)
+        self.die            = die
+        self.verbose        = verbose
+        self.calibrated     = False
+
+        # Load data -- this is expecting a dataframe with a column for 'year' and other columns for to sim results
+        if not isinstance(data, pd.DataFrame):
+            errormsg = 'Please pass data as a pandas dataframe'
+            raise ValueError(errormsg)
+        self.target_data = data
+        self.target_data.set_index('year', inplace=True)
+
+        # Temporarily store a filename
+        self.tmp_filename = 'tmp_calibration_%05i.obj'
+
+        # Initialize sim
+        if not self.sim.initialized:
+            self.sim.initialize()
+
+        # Figure out which sim results to get
+        self.sim_result_list = self.target_data.columns.values.tolist()
+
+        return
+
+    def run_sim(self, calib_pars=None, label=None):
+        """ Create and run a simulation """
+        sim = sc.dcp(self.sim)
+        if label: sim.label = label
+
+        sim = self.translate_pars(sim, calib_pars=calib_pars)
+
+        # Run the sim
+        try:
+            sim.run()
+            return sim
+
+        except Exception as E:
+            if self.die:
+                raise E
+            else:
+                print(f'Encountered error running sim!\nParameters:\n{calib_pars}\nTraceback:\n{sc.traceback()}')
+                output = None
+                return output
+
+    @staticmethod
+    def translate_pars(sim=None, calib_pars=None):
+        """ Take the nested dict of calibration pars and modify the sim """
+        for modtype in calib_pars.keys():
+            for dkey, dpars in calib_pars[modtype].items():
+                for dparkey, dparval in dpars.items():
+                    targetpar = sim[modtype][dkey].pars[dparkey]
+                    if sc.isnumber(targetpar):
+                        sim[modtype][dkey].pars[dparkey] = dparval
+                    elif isinstance(targetpar, ss.Dist):
+                        sim[modtype][dkey].pars[dparkey].set(dparval)
+                    else:
+                        errormsg = 'Type not implemented'
+                        raise ValueError(errormsg)
+
+        return sim
+
+    def trial_to_sim_pars(self, pardict=None, trial=None):
+        """
+        Take in an optuna trial and sample from pars, after extracting them from the structure they're provided in
+        Different use cases:
+            - pardict is self.calib_pars, i.e. {'diseases':{'hiv':{'art_efficacy':[0.96, 0.9, 0.99]}}}, need to sample
+            - pardict is self.initial_pars, i.e. {'diseases':{'hiv':{'art_efficacy':[0.96, 0.9, 0.99]}}}, pull 1st vals
+            - pardict is self.best_pars, i.e. {'diseases':{'hiv':{'art_efficacy':0.96786}}}, pull single vals
+        """
+        pars = sc.dcp(pardict)
+        flattened_inputs = sc.flattendict(pars)
+
+        structured_outputs = sc.dcp(self.calib_pars)
+        flattened_outputs = sc.flattendict(structured_outputs)
+
+        for key in flattened_outputs.keys():
+            sampler_key = '_'.join(key)
+
+            if key in flattened_inputs.keys():
+                val = flattened_inputs[key]
+            else:
+                val = pars[sampler_key]
+
+            if sc.isnumber(val) and trial is None:
+                sc.setnested(structured_outputs, list(key), val)
+
+            else:
+                low, high = val[1], val[2]
+                step = val[3] if len(val) > 3 else None
+
+                if trial is not None:
+                    if key in self.par_samplers:  # If a custom sampler is used, get it now (Not working properly for now)
+                        try:
+                            sampler_fn = getattr(trial, self.par_samplers[key])
+                        except Exception as E:
+                            errormsg = 'The requested sampler function is not found: ensure it is a valid attribute of an Optuna Trial object'
+                            raise AttributeError(errormsg) from E
+                    else:
+                        sampler_fn = trial.suggest_float
+                    sc.setnested(structured_outputs, list(key), sampler_fn(sampler_key, low, high, step=step))
+                else:
+                    sc.setnested(structured_outputs, list(key), val[0])
+
+        return structured_outputs
+
+    def run_trial(self, trial, save=True):
+        """ Define the objective for Optuna """
+        if self.calib_pars is not None:
+            calib_pars = self.trial_to_sim_pars(self.calib_pars, trial)
+        else:
+            calib_pars = None
+        sim = self.run_sim(calib_pars)
+
+        # Export results
+        df_res = sim.export_df()
+        df_res['year'] = np.floor(np.round(df_res.index, 1)).astype(int)
+        sim_results = sc.objdict()
+
+        for skey in self.sim_result_list:
+            if 'prevalence' in skey:
+                model_output = df_res.groupby(by='year')[skey].mean()
+            else:
+                model_output = df_res.groupby(by='year')[skey].sum()
+            sim_results[skey] = model_output.values
+
+        sim_results['year'] = model_output.index.values
+        # Store results in temporary files
+        if save:
+            filename = self.tmp_filename % trial.number
+            sc.save(filename, sim_results)
+
+        # Compute fit
+        fit = self.compute_fit(sim)
+        return fit
+
+    def compute_fit(self, sim):
+        """ Compute goodness-of-fit """
+        fit = 0
+        df_res = sim.export_df()
+        df_res['year'] = np.floor(np.round(df_res.index, 1)).astype(int)
+        for skey in self.sim_result_list:
+            if 'prevalence' in skey:
+                model_output = df_res.groupby(by='year')[skey].mean()
+            else:
+                model_output = df_res.groupby(by='year')[skey].sum()
+
+            data = self.target_data[skey]
+            combined = pd.merge(data, model_output, how='left', on='year')
+            combined['diffs'] = combined[skey+'_x'] - combined[skey+'_y']
+            gofs = compute_gof(combined.dropna()[skey+'_x'], combined.dropna()[skey+'_y'])
+
+            losses = gofs  #* self.weights[skey]
+            mismatch = losses.sum()
+            fit += mismatch
+
+        return fit
+
+    def worker(self):
+        """ Run a single worker """
+        op = import_optuna()
+        if self.verbose:
+            op.logging.set_verbosity(op.logging.DEBUG)
+        else:
+            op.logging.set_verbosity(op.logging.ERROR)
+        study = op.load_study(storage=self.run_args.storage, study_name=self.run_args.name, sampler = self.run_args.sampler)
+        output = study.optimize(self.run_trial, n_trials=self.run_args.n_trials, callbacks=None)
+        return output
+
+    def run_workers(self):
+        """ Run multiple workers in parallel """
+        if self.run_args.n_workers > 1: # Normal use case: run in parallel
+            output = sc.parallelize(self.worker, iterarg=self.run_args.n_workers)
+        else: # Special case: just run one
+            output = [self.worker()]
+        return output
+
+    def remove_db(self):
+        """
+        Remove the database file if keep_db is false and the path exists.
+        """
+        try:
+            if 'sqlite' in self.run_args.storage:
+                # Delete the file from disk
+                if os.path.exists(self.run_args.db_name):
+                    os.remove(self.run_args.db_name)
+                if self.verbose:
+                    print(f'Removed existing calibration file {self.run_args.db_name}')
+            else:
+                # Delete the study from the database e.g., mysql
+                op = import_optuna()
+                op.delete_study(study_name=self.run_args.name, storage=self.run_args.storage)
+                if self.verbose:
+                    print(f'Deleted study {self.run_args.name} in {self.run_args.storage}')
+        except Exception as E:
+            print('Could not delete study, skipping...')
+            print(str(E))
+        return
+
+    def make_study(self):
+        """ Make a study, deleting one if it already exists """
+        op = import_optuna()
+        if not self.run_args.keep_db:
+            self.remove_db()
+        if self.run_args.rand_seed is not None:
+            sampler = op.samplers.RandomSampler(self.run_args.rand_seed)
+            sampler.reseed_rng()
+            raise NotImplementedError('Implemented but does not work')
+        else:
+            sampler = None
+        print(self.run_args.storage)
+        output = op.create_study(storage=self.run_args.storage, study_name=self.run_args.name, sampler=sampler)
+        return output
+
+    def calibrate(self, calib_pars=None, confirm_fit=False, load=True, tidyup=True, **kwargs):
+        """
+        Perform calibration.
+        Args:
+            calib_pars (dict): if supplied, overwrite stored calib_pars
+            confirm_fit (bool): if True, run simulations with parameters from before and after calibration
+            load (bool): whether to load existing trials from the database (if rerunning the same calibration)
+            tidyup (bool): whether to delete temporary files from trial runs
+            verbose (bool): whether to print output from each trial
+            kwargs (dict): if supplied, overwrite stored run_args (n_trials, n_workers, etc.)
+        """
+        op = import_optuna()
+
+        # Load and validate calibration parameters
+        if calib_pars is not None:
+            self.calib_pars = calib_pars
+        self.run_args.update(kwargs) # Update optuna settings
+
+        # Run the optimization
+        t0 = sc.tic()
+        self.make_study()
+        self.run_workers()
+        study = op.load_study(storage=self.run_args.storage, study_name=self.run_args.name, sampler = self.run_args.sampler)
+        self.best_pars = sc.objdict(study.best_params)
+        self.elapsed = sc.toc(t0, output=True)
+
+        self.sim_results = []
+        if load:
+            print('Loading saved results...')
+            for trial in study.trials:
+                n = trial.number
+                try:
+                    filename = self.tmp_filename % trial.number
+                    results = sc.load(filename)
+                    self.sim_results.append(results)
+                    if tidyup:
+                        try:
+                            os.remove(filename)
+                            print(f'    Removed temporary file {filename}')
+                        except Exception as E:
+                            errormsg = f'Could not remove {filename}: {str(E)}'
+                            print(errormsg)
+                    print(f'  Loaded trial {n}')
+                except Exception as E:
+                    errormsg = f'Warning, could not load trial {n}: {str(E)}'
+                    print(errormsg)
+
+        # Compare the results
+        self.initial_pars = self.trial_to_sim_pars(pardict=self.calib_pars)
+        self.parse_study(study)
+
+        # Tidy up
+        self.calibrated = True
+        if not self.run_args.keep_db:
+            self.remove_db()
+        
+        # Optionally compute the sims before and after the fit
+        if confirm_fit:
+            self.confirm_fit()
+
+        return self
+    
+    def confirm_fit(self):
+        """ Run before and after simulations to validate the fit """
+        before_pars = self.initial_pars
+        after_pars  = self.trial_to_sim_pars(pardict=self.best_pars)
+        self.before_sim = self.run_sim(calib_pars=before_pars, label='Before calibration')
+        self.after_sim  = self.run_sim(calib_pars=after_pars, label='After calibration')
+        self.before_fit = self.compute_fit(self.before_sim)
+        self.after_fit  = self.compute_fit(self.after_sim)
+        return self.before_fit, self.after_fit
+
+    def parse_study(self, study):
+        """Parse the study into a data frame -- called automatically """
+        best = study.best_params
+        self.best_pars = best
+
+        print('Making results structure...')
+        results = []
+        n_trials = len(study.trials)
+        failed_trials = []
+        for trial in study.trials:
+            data = {'index':trial.number, 'mismatch': trial.value}
+            for key,val in trial.params.items():
+                data[key] = val
+            if data['mismatch'] is None:
+                failed_trials.append(data['index'])
+            else:
+                results.append(data)
+        print(f'Processed {n_trials} trials; {len(failed_trials)} failed')
+
+        keys = ['index', 'mismatch'] + list(best.keys())
+        data = sc.objdict().make(keys=keys, vals=[])
+        for i,r in enumerate(results):
+            for key in keys:
+                if key not in r:
+                    warnmsg = f'Key {key} is missing from trial {i}, replacing with default'
+                    print(warnmsg)
+                    r[key] = best[key]
+                data[key].append(r[key])
+        self.data = data
+        self.df = pd.DataFrame.from_dict(data)
+        self.df = self.df.sort_values(by=['mismatch']) # Sort
+        return
+
+    def to_json(self, filename=None, indent=2, **kwargs):
+        """ Convert the results to JSON """
+        order = np.argsort(self.df['mismatch'])
+        json = []
+        for o in order:
+            row = self.df.iloc[o,:].to_dict()
+            rowdict = dict(index=row.pop('index'), mismatch=row.pop('mismatch'), pars={})
+            for key,val in row.items():
+                rowdict['pars'][key] = val
+            json.append(rowdict)
+        self.json = json
+        if filename:
+            return sc.savejson(filename, json, indent=indent, **kwargs)
+        else:
+            return json
+
diff --git a/starsim/demographics.py b/starsim/demographics.py
index 8a9506c1..a768c3c7 100644
--- a/starsim/demographics.py
+++ b/starsim/demographics.py
@@ -70,7 +70,9 @@ def init_pre(self, sim):
     def standardize_birth_data(self):
         """ Standardize/validate birth rates - handled in an external file due to shared functionality """
         birth_rate = ss.standardize_data(data=self.pars.birth_rate, metadata=self.metadata)
-        return birth_rate.xs(0,level='age')
+        if isinstance(birth_rate, (pd.Series, pd.DataFrame)):
+            return birth_rate.xs(0,level='age')
+        return birth_rate
 
     def init_results(self):
         npts = self.sim.npts
@@ -93,10 +95,14 @@ def get_births(self):
         sim = self.sim
         p = self.pars
 
-        available_years = p.birth_rate.index
-        year_ind = sc.findnearest(available_years, sim.year)
-        nearest_year = available_years[year_ind]
-        this_birth_rate = p.birth_rate.loc[nearest_year]
+        if isinstance(p.birth_rate, (pd.Series, pd.DataFrame)):
+            available_years = p.birth_rate.index
+            year_ind = sc.findnearest(available_years, sim.year)
+            nearest_year = available_years[year_ind]
+            this_birth_rate = p.birth_rate.loc[nearest_year]
+        else:
+            this_birth_rate = p.birth_rate
+
         scaled_birth_prob = this_birth_rate * p.units * p.rel_birth * sim.pars.dt
         scaled_birth_prob = np.clip(scaled_birth_prob, a_min=0, a_max=1)
         n_new = int(np.floor(sim.people.alive.count() * scaled_birth_prob))
@@ -178,8 +184,9 @@ def __init__(self, pars=None, metadata=None, **kwargs):
     def standardize_death_data(self):
         """ Standardize/validate death rates - handled in an external file due to shared functionality """
         death_rate = ss.standardize_data(data=self.pars.death_rate, metadata=self.metadata)
-        death_rate = death_rate.unstack(level='age')
-        assert not death_rate.isna().any(axis=None) # For efficiency, we assume that the age bins are the same for all years in the input dataset
+        if isinstance(death_rate, (pd.Series, pd.DataFrame)):
+            death_rate = death_rate.unstack(level='age')
+            assert not death_rate.isna().any(axis=None) # For efficiency, we assume that the age bins are the same for all years in the input dataset
         return death_rate
 
     @staticmethod # Needs to be static since called externally, although it sure looks like a class method!
@@ -257,9 +264,9 @@ class Pregnancy(Demographics):
     def __init__(self, pars=None, metadata=None, **kwargs):
         super().__init__()
         self.default_pars(
-            dur_pregnancy = 0.75,   # Duration for pre-natal transmission
-            dur_postpartum = ss.lognorm_ex(0.5, 0.5),   # Duration for post-natal transmission (e.g. via breastfeeding)
-            fertility_rate = 0,     # See make_fertility_prob_function
+            dur_pregnancy = 0.75, # Duration for pre-natal transmission
+            dur_postpartum = ss.lognorm_ex(0.5, 0.5), # Duration for post-natal transmission (e.g. via breastfeeding)
+            fertility_rate = 0, # Can be a number of Pandas DataFrame
             rel_fertility = 1,
             maternal_death_prob = ss.bernoulli(0),
             sex_ratio = ss.bernoulli(0.5), # Ratio of babies born female
@@ -268,6 +275,8 @@ def __init__(self, pars=None, metadata=None, **kwargs):
             units = 1e-3, # Assumes fertility rates are per 1000. If using percentages, switch this to 1
         )
         self.update_pars(pars, **kwargs)
+
+        self.pars.p_fertility = ss.bernoulli(p=0) # Placeholder, see make_fertility_prob_fn
         
         # Other, e.g. postpartum, on contraception...
         self.add_states(
@@ -289,11 +298,6 @@ def __init__(self, pars=None, metadata=None, **kwargs):
         )
         self.choose_slots = None # Distribution for choosing slots; set in self.initialize()
 
-        # Process data, which may be provided as a number, dict, dataframe, or series
-        # If it's a number it's left as-is; otherwise it's converted to a dataframe
-        self.fertility_rate_data = self.standardize_fertility_data()
-        self.pars.fertility_rate = ss.bernoulli(self.make_fertility_prob_fn)
-
         # For results tracking
         self.n_pregnancies = 0
         self.n_births = 0
@@ -349,14 +353,19 @@ def standardize_fertility_data(self):
         Standardize/validate fertility rates
         """
         fertility_rate = ss.standardize_data(data=self.pars.fertility_rate, metadata=self.metadata)
-        fertility_rate = fertility_rate.unstack()
-        max_age = fertility_rate.columns.max()
-        fertility_rate[max_age+1] = 0
-        assert not fertility_rate.isna().any(axis=None) # For efficiency, we assume that the age bins are the same for all years in the input dataset
+        if isinstance(fertility_rate, (pd.Series, pd.DataFrame)):
+            fertility_rate = fertility_rate.unstack()
+            assert not fertility_rate.isna().any(axis=None) # For efficiency, we assume that the age bins are the same for all years in the input dataset
         return fertility_rate
 
     def init_pre(self, sim):
         super().init_pre(sim)
+
+        # Process data, which may be provided as a number, dict, dataframe, or series
+        # If it's a number it's left as-is; otherwise it's converted to a dataframe
+        self.fertility_rate_data = self.standardize_fertility_data()
+        self.pars.p_fertility.set(p=self.make_fertility_prob_fn)
+
         low = sim.pars.n_agents + 1
         high = int(sim.pars.slot_scale*sim.pars.n_agents)
         self.choose_slots = ss.randint(low=low, high=high, sim=sim, module=self)
@@ -409,7 +418,7 @@ def update_states(self):
 
                 # Validation
                 if not np.array_equal(new_mother_uids, deliveries.uids):
-                    errormsg = f'IDs of new mothers do not match IDs of new deliveries.'
+                    errormsg = 'IDs of new mothers do not match IDs of new deliveries'
                     raise ValueError(errormsg)
 
                 # Create durations and start dates, and add connections
@@ -435,7 +444,7 @@ def make_pregnancies(self):
         # People eligible to become pregnant. We don't remove pregnant people here, these
         # are instead handled in the fertility_dist logic as the rates need to be adjusted
         eligible_uids = self.sim.people.female.uids
-        conceive_uids = self.pars.fertility_rate.filter(eligible_uids)
+        conceive_uids = self.pars.p_fertility.filter(eligible_uids)
 
         # Validation
         if np.any(self.pregnant[conceive_uids]):
diff --git a/starsim/disease.py b/starsim/disease.py
index 150e2270..13f69262 100644
--- a/starsim/disease.py
+++ b/starsim/disease.py
@@ -295,8 +295,9 @@ def make_new_cases(self):
         # Tidy up
         if len(new_cases) and len(sources):
             new_cases = ss.uids.cat(new_cases)
-            sources = ss.uids.cat(sources)
-            networks = np.concatenate(networks)
+            new_cases, inds = new_cases.unique(return_index=True)
+            sources = ss.uids.cat(sources)[inds]
+            networks = np.concatenate(networks)[inds]
         else:
             new_cases = np.empty(0, dtype=int)
             sources = np.empty(0, dtype=int)
diff --git a/starsim/diseases/sir.py b/starsim/diseases/sir.py
index 25a6cc87..939909cc 100644
--- a/starsim/diseases/sir.py
+++ b/starsim/diseases/sir.py
@@ -3,6 +3,7 @@
 """
 
 import numpy as np
+import sciris as sc
 import matplotlib.pyplot as pl
 import starsim as ss
 
@@ -77,12 +78,17 @@ def update_death(self, uids):
         self.recovered[uids] = False
         return
 
-    def plot(self):
+    def plot(self, plot_kw=None):
         """ Default plot for SIR model """
         fig = pl.figure()
+        plot_kw = sc.mergedicts(dict(lw=2, alpha=0.8), plot_kw)
         for rkey in ['n_susceptible', 'n_infected', 'n_recovered']:
-            pl.plot(self.results[rkey], label=self.results[rkey].label)
-        pl.legend()
+            pl.plot(self.sim.results.yearvec, self.results[rkey], label=self.results[rkey].label, **plot_kw)
+        pl.legend(frameon=False)
+        pl.xlabel('Year')
+        pl.ylabel('Number of people')
+        sc.boxoff()
+        sc.commaticks()
         return fig
     
 
diff --git a/starsim/distributions.py b/starsim/distributions.py
index 6cf11b68..12624086 100644
--- a/starsim/distributions.py
+++ b/starsim/distributions.py
@@ -661,8 +661,10 @@ def __init__(self, mean=0.0, sigma=1.0, **kwargs):
         super().__init__(distname='lognormal', dist=sps.lognorm, mean=mean, sigma=sigma, **kwargs)
         return
     
-    def sync_pars(self):
+    def sync_pars(self, call=True):
         """ Translate between NumPy and SciPy parameters """
+        if call:
+            self.call_pars()
         p = self._pars
         spars = sc.dictobj()
         spars.s = p.sigma
@@ -717,7 +719,7 @@ def convert_ex_to_im(self):
     def sync_pars(self):
         """ Convert from overlying to underlying parameters, then translate to SciPy """
         self.convert_ex_to_im()
-        spars = lognorm_im.sync_pars(self) # Borrow sync_pars from lognorm_im
+        spars = lognorm_im.sync_pars(self, call=False) # Borrow sync_pars from lognorm_im
         return spars
     
 
diff --git a/starsim/modules.py b/starsim/modules.py
index 51bec676..e3b24234 100644
--- a/starsim/modules.py
+++ b/starsim/modules.py
@@ -39,7 +39,7 @@ def find_modules(key=None):
             except:
                 pass
     return modules if key is None else modules[key]
-    
+
 
 class Module(sc.quickobj):
 
@@ -82,7 +82,7 @@ def update_pars(self, pars, **kwargs):
         self.pars.update(matches)
                 
         # Update module attributes
-        metadata = {key:pars.pop(key, None) for key in ['name', 'label', 'requires']}
+        metadata = {key:pars.pop(key, getattr(self, key, None)) for key in ['name', 'label', 'requires']}
         self.set_metadata(**metadata)
         
         # Should be no remaining pars
@@ -237,4 +237,5 @@ class Connector(Module):
     Because connectors can do anything, they have no specified structure: it is
     up to the user to define how they behave.    
     """
-    pass
+    def update(self, sim):
+        pass
diff --git a/starsim/network.py b/starsim/network.py
index 1eaa9209..0708bdc8 100644
--- a/starsim/network.py
+++ b/starsim/network.py
@@ -10,7 +10,7 @@
 import scipy.optimize as spo
 import scipy.spatial as spsp
 
-# CK: need to remove this, but slows down the code otherwise
+# This has a significant impact on runtime, surprisingly
 ss_float_ = ss.dtypes.float
 ss_int_ = ss.dtypes.int
 
@@ -905,12 +905,11 @@ def add_pairs(self):
 
         dist_mat = spsp.distance_matrix(loc_m[:, np.newaxis], loc_f[:, np.newaxis])
         ind_m, ind_f = spo.linear_sum_assignment(dist_mat)
-        n_pairs = len(ind_f)
 
         # Finalize pairs
         p1 = available_m[ind_m]
         p2 = available_f[ind_f]
-        beta = np.ones(n_pairs) # TODO: Allow custom beta
+        beta = np.ones(len(p1)) # TODO: Allow custom beta
         dur_vals = self.pars.duration.rvs(p1)
         act_vals = self.pars.acts.rvs(p1)
 
diff --git a/starsim/parameters.py b/starsim/parameters.py
index 301ccb17..87ee237f 100644
--- a/starsim/parameters.py
+++ b/starsim/parameters.py
@@ -177,7 +177,6 @@ def __init__(self, **kwargs):
         self.slot_scale = 5      # Random slots will be assigned to newborn agents between min=n_agents and max=slot_scale*n_agents
 
         # Demographic parameters
-        self.location   = None  #  NOT CURRENTLY FUNCTIONAL - what demographics to use
         self.birth_rate = None
         self.death_rate = None
         self.use_aging  = None # True if demographics, false otherwise
@@ -244,7 +243,7 @@ def validate_agents(self):
     
     def validate_total_pop(self):
         """ Ensure one but not both of total_pop and pop_scale are defined """
-        if self.total_pop is not None:  # If no pop_scale has been provided, try to get it from the location
+        if self.total_pop is not None:  # In future, if no pop_scale has been provided, try to get it from the location
             if self.pop_scale is not None:
                 errormsg = f'You can define total_pop ({self.total_pop}) or pop_scale ({self.pop_scale}), but not both, since one is calculated from the other'
                 raise ValueError(errormsg)
diff --git a/starsim/results.py b/starsim/results.py
index b5890c3b..789de967 100644
--- a/starsim/results.py
+++ b/starsim/results.py
@@ -18,14 +18,18 @@ class Result(np.ndarray):
         'module', # The name of the module (e.g. hiv)
         'scale', # Whether or not the result scales with population size (e.g. True)
         'label', # The human-readable label for the result (e.g. Number of infections)
+        'low', # The lower bound for the values (used with MultiSim)
+        'high', # The upper bound for the values
     ]
     
-    def __new__(cls, module=None, name=None, shape=None, dtype=None, scale=None, label=None):
+    def __new__(cls, module=None, name=None, shape=None, dtype=None, scale=None, label=None, low=None, high=None):
         arr = np.zeros(shape=shape, dtype=dtype).view(cls)
         arr.name = name
         arr.module = module
         arr.scale = scale
         arr.label = label
+        arr.low = low
+        arr.high = high
         return arr
     
     
@@ -51,7 +55,13 @@ def __array_wrap__(self, obj, **kwargs):
             return super().__array_wrap__(obj, **kwargs)
     
     def to_df(self):
-        return sc.dataframe({self.name:self})
+        data = {self.name:self}
+        if self.low is not None:
+            data['low']  = self.low
+        if self.high is not None:
+            data['high'] = self.high
+        df = sc.dataframe(data)
+        return df
     
 
 class Results(ss.ndict):
@@ -76,6 +86,9 @@ def append(self, arg, key=None):
         super().append(result, key=key)
         return
     
+    def flatten(self, sep='_'):
+        return sc.objdict(sc.flattendict(self, sep=sep))
+    
     def to_df(self):
         pass
     
diff --git a/starsim/run.py b/starsim/run.py
index 9d9b6c22..a2a5a284 100644
--- a/starsim/run.py
+++ b/starsim/run.py
@@ -4,6 +4,7 @@
 
 import numpy as np
 import sciris as sc
+import pylab as pl
 import starsim as ss
 
 __all__ = ['MultiSim', 'single_run', 'multi_run', 'parallel']
@@ -149,7 +150,6 @@ def reduce(self, quantiles=None, use_mean=False, bounds=None, output=False):
             msim.reduce()
             msim.summarize()
         """
-
         if use_mean:
             if bounds is None:
                 bounds = 2
@@ -170,35 +170,37 @@ def reduce(self, quantiles=None, use_mean=False, bounds=None, output=False):
         reduced_sim.metadata = dict(parallelized=True, combined=False, n_runs=n_runs, quantiles=quantiles,
                                     use_mean=use_mean, bounds=bounds)  # Store how this was parallelized
 
-        # Perform the statistics
+        # Calculate the statistics
         raw = {}
+        exclude = ['yearvec'] # Exclude things that we don't want to include
 
-        rkeys = reduced_sim.results.keys()
+        rflat = reduced_sim.results.flatten()
+        rkeys = list(rflat.keys())
+        rkeys = [k for k in rkeys if k not in exclude]
 
         for rkey in rkeys:
-            raw[rkey] = np.zeros((reduced_sim.res_npts, len(self.sims)))
+            raw[rkey] = np.zeros((reduced_sim.npts, len(self.sims)))
             for s, sim in enumerate(self.sims):
-                vals = sim.results[rkey].values
-                raw[rkey][:, s] = vals
+                flat = sim.results.flatten()
+                raw[rkey][:, s] = flat[rkey]
 
         for rkey in rkeys:
-            results = reduced_sim.results
             if use_mean:
                 r_mean = np.mean(raw[rkey], axis=1)
                 r_std = np.std(raw[rkey], axis=1)
-                results[rkey].values[:] = r_mean
-                results[rkey].low = r_mean - bounds * r_std
-                results[rkey].high = r_mean + bounds * r_std
+                rflat[rkey][:] = r_mean
+                rflat[rkey].low = r_mean - bounds * r_std
+                rflat[rkey].high = r_mean + bounds * r_std
             else:
-                results[rkey].values[:] = np.quantile(raw[rkey], q=0.5, axis=1)
-                results[rkey].low = np.quantile(raw[rkey], q=quantiles['low'], axis=1)
-                results[rkey].high = np.quantile(raw[rkey], q=quantiles['high'], axis=1)
+                rflat[rkey][:] = np.quantile(raw[rkey], q=0.5, axis=1)
+                rflat[rkey].low = np.quantile(raw[rkey], q=quantiles['low'], axis=1)
+                rflat[rkey].high = np.quantile(raw[rkey], q=quantiles['high'], axis=1)
 
         # Compute and store final results
-        reduced_sim.compute_summary()
+        reduced_sim.summarize()
         self.orig_base_sim = self.base_sim
         self.base_sim = reduced_sim
-        self.results = reduced_sim.results
+        self.results = rflat
         self.summary = reduced_sim.summary
         self.which = 'reduced'
 
@@ -226,6 +228,55 @@ def median(self, quantiles=None, **kwargs):
             kwargs (dict): passed to reduce()
         """
         return self.reduce(use_mean=False, quantiles=quantiles, **kwargs)
+    
+    def plot(self, key=None, fig=None, fig_kw=None, plot_kw=None, fill_kw=None):
+        """ 
+        Plot all results in the MultiSim object.
+        
+        If the MultiSim object has been reduced (i.e. mean or median), then plot
+        the best value and uncertainty bound. Otherwise, plot individual sims.
+        
+        Args:
+            key (str): the results key to plot (by default, all)
+            fig (Figure): if provided, plot results into an existing figure
+            fig_kw (dict): passed to ``plt.subplots()``
+            plot_kw (dict): passed to ``plt.plot()``
+            fill_kw (dict): passed to ``plt.fill_between()``
+        """
+        # Has not been reduced yet, plot individual sim
+        if self.which is None:
+            fig = None
+            alpha = 0.7 if len(self.sims) < 5 else 0.5
+            plot_kw = sc.mergedicts({'alpha':alpha}, plot_kw)
+            for sim in self.sims:
+                fig = sim.plot(key=key, fig=fig, fig_kw=fig_kw, plot_kw=plot_kw)
+            pl.legend()
+        
+        # Has been reduced, plot with uncertainty bounds
+        else:
+            flat = self.results
+            yearvec = flat.pop('yearvec')
+            n_cols = np.ceil(np.sqrt(len(flat))) # TODO: remove duplication with sim.plot()
+            default_figsize = np.array([8, 6])
+            figsize_factor = np.clip((n_cols-3)/6+1, 1, 1.5) # Scale the default figure size based on the number of rows and columns
+            figsize = default_figsize*figsize_factor
+            fig_kw = sc.mergedicts({'figsize':figsize}, fig_kw)
+            fig_kw = sc.mergedicts(fig_kw)
+            fill_kw = sc.mergedicts({'alpha':0.2}, fill_kw)
+            plot_kw = sc.mergedicts({'lw':2, 'alpha':0.8}, plot_kw)
+            with sc.options.with_style('simple'):
+                if key is not None:
+                    flat = {k:v for k,v in flat.items() if k.startswith(key)}
+                fig, axs = sc.getrowscols(len(flat), make=True, **fig_kw)
+                    
+                # Do the plotting
+                for ax, (key, res) in zip(axs.flatten(), flat.items()):
+                    ax.fill_between(yearvec, res.low, res.high, **fill_kw)
+                    ax.plot(yearvec, res, **plot_kw)
+                    ax.set_title(getattr(res, 'label', key)) 
+                    ax.set_xlabel('Year')
+                
+        return fig
 
 
 def single_run(sim, ind=0, reseed=True, keep_people=False, run_args=None, sim_args=None,
diff --git a/starsim/sim.py b/starsim/sim.py
index 9c995a1d..c3f628db 100644
--- a/starsim/sim.py
+++ b/starsim/sim.py
@@ -12,7 +12,7 @@
 __all__ = ['Sim', 'AlreadyRunError', 'demo', 'diff_sims', 'check_sims_match']
 
 
-class Sim(sc.prettyobj):
+class Sim:
 
     def __init__(self, pars=None, label=None, people=None, demographics=None, diseases=None, networks=None,
                  interventions=None, analyzers=None, connectors=None, copy_inputs=True, **kwargs):
@@ -44,6 +44,38 @@ def __setitem__(self, key, value):
         """ Allow dict-like access, e.g. sim['created'] = sc.now() """
         return setattr(self, key, value)
     
+    def __repr__(self):
+        """ Show a quick version of the sim """
+        # Try a more custom repr first
+        try:
+            n = int(self.pars.n_agents)
+            moddict = {}
+            for modkey in ss.module_map().keys():
+                if hasattr(self, modkey):
+                    thismodtype = self[modkey]
+                elif modkey in self.pars:
+                    thismodtype = self.pars[modkey]
+                else:
+                    thismodtype = {}
+                if sc.isiterable(thismodtype) and len(thismodtype):
+                    moddict[modkey] = sc.strjoin(thismodtype.keys())
+            if len(moddict):
+                modulestr = ''
+                for k,mstr in moddict.items():
+                    modulestr += f'; {k}={mstr}'
+            else:
+                modulestr = ''
+            if not self.initialized:
+                modulestr += '; not initialized'
+            string = f'Sim(n={n:n}{modulestr})'
+        
+        # Or just use default
+        except Exception as E:
+            ss.warn(f'Error displaying custom sim repr, falling back to default: {E}')
+            string = sc.prepr(self, vals=False)
+            
+        return string
+    
     def initialize(self, **kwargs):
         """ Perform all initializations for the sim; most heavy lifting is done by the parameters """
         # Validation and initialization
@@ -147,7 +179,7 @@ def init_results(self):
 
     @property
     def modules(self):
-        # Return iterator over all Module instances (stored in standard places) in the Sim
+        """ Return iterator over all Module instances (stored in standard places) in the Sim """
         products = [intv.product for intv in self.interventions.values() if
                     hasattr(intv, 'product') and isinstance(intv.product, ss.Product)]
         return itertools.chain(
@@ -354,7 +386,8 @@ def get_result(res, func):
         return summary
     
     def disp(self):
-        print(self.summary)
+        """ Print a full version of the sim """
+        sc.pr(self)
         return
     
     def shrink(self, skip_attrs=None, in_place=True):
@@ -677,19 +710,64 @@ def to_json(self, filename=None, keys=None, tostring=False, indent=2, verbose=Fa
             output = sc.savejson(filename=filename, obj=d, indent=indent, *args, **kwargs)
 
         return output
-
-    def plot(self, key=None):
-        with sc.options.with_style('fancy'):
-            flat = sc.flattendict(self.results, sep=': ')
+    
+    def plot(self, key=None, fig=None, style='fancy', fig_kw=None, plot_kw=None):
+        """ 
+        Plot all results in the Sim object
+        
+        Args:
+            key (str): the results key to plot (by default, all)
+            fig (Figure): if provided, plot results into an existing figure
+            style (str): the plotting style to use (default "fancy"; other options are "simple", None, or any Matplotlib style)
+            fig_kw (dict): passed to ``plt.subplots()``
+            plot_kw (dict): passed to ``plt.plot()``
+        
+        """
+        
+        # Configuration
+        flat = self.results.flatten()
+        n_cols = np.ceil(np.sqrt(len(flat))) # Number of columns of axes
+        default_figsize = np.array([8, 6])
+        figsize_factor = np.clip((n_cols-3)/6+1, 1, 1.5) # Scale the default figure size based on the number of rows and columns
+        figsize = default_figsize*figsize_factor
+        fig_kw = sc.mergedicts({'figsize':figsize}, fig_kw)
+        plot_kw = sc.mergedicts({'lw':2}, plot_kw)
+        modmap = {m.name:m for m in self.modules} # Find modules
+        
+        # Do the plotting
+        with sc.options.with_style(style):
+            
             yearvec = flat.pop('yearvec')
             if key is not None:
                 flat = {k:v for k,v in flat.items() if k.startswith(key)}
-            fig, axs = sc.getrowscols(len(flat), make=True)
-            for ax, (k, v) in zip(axs.flatten(), flat.items()):
-                ax.plot(yearvec, v)
-                # get label or use the key
-                ax.set_title(getattr(v, 'label', k))
+            
+            # Get the figure
+            if fig is None:
+                fig, axs = sc.getrowscols(len(flat), make=True, **fig_kw)
+                if isinstance(axs, np.ndarray):
+                    axs = axs.flatten()
+            else:
+                axs = fig.axes
+            if not sc.isiterable(axs):
+                axs = [axs]
+                
+            # Do the plotting
+            for ax, (key, res) in zip(axs, flat.items()):
+                ax.plot(yearvec, res, **plot_kw, label=self.label)
+                title = getattr(res, 'label', key)
+                if res.module != 'sim':
+                    try:
+                        mod = modmap[res.module]
+                        modtitle = mod.__class__.__name__
+                        assert res.module == modtitle.lower() # Only use the class name if the module name is the default
+                    except:
+                        modtitle = res.module
+                    title = f'{modtitle}: {title}'
+                ax.set_title(title) 
                 ax.set_xlabel('Year')
+            
+        sc.figlayout(fig=fig)
+                
         return fig
 
 
diff --git a/starsim/states.py b/starsim/states.py
index 84870e5a..c9c45d17 100644
--- a/starsim/states.py
+++ b/starsim/states.py
@@ -187,7 +187,7 @@ def auids(self):
             return self.people.auids
         except:
             if not self.initialized:
-                ss.warn('Trying to access non-initialized States object')
+                ss.warn('Trying to access non-initialized Arr object; in most cases, Arr objects need to be initialized with a Sim object, but set skip_init=True if this is intentional.')
             return uids(np.arange(len(self.raw)))
     
     def count(self):
@@ -292,7 +292,13 @@ def false(self):
 
 
 class FloatArr(Arr):
-    """ Subclass of Arr with defaults for floats """
+    """
+    Subclass of Arr with defaults for floats and ints.
+    
+    Note: Starsim does not support integer arrays by default since they introduce
+    ambiguity in dealing with NaNs, and float arrays are suitable for most purposes.
+    If you really want an integer array, you can use the default Arr class instead.    
+    """
     def __init__(self, name=None, nan=np.nan, **kwargs):
         super().__init__(name=name, dtype=ss_float, nan=nan, coerce=False, **kwargs)
         return
@@ -348,7 +354,7 @@ def split(self):
 
     
 class IndexArr(Arr):
-    """ A special class of IndexArr used for UIDs and RNG IDs """
+    """ A special class of Arr used for UIDs and RNG IDs; not to be used as an integer array (for that, use FloatArr) """
     def __init__(self, name=None, label=None):
         super().__init__(name=name, dtype=ss_int, default=None, nan=-1, label=label, coerce=False, skip_init=True)
         self.raw = uids(self.raw)
@@ -392,12 +398,12 @@ def __new__(cls, arr=None):
         return np.asarray(arr, dtype=ss_int).view(cls) # Handle everything else
 
     def concat(self, other, **kw): # TODO: why can't they both be called cat()?
-        """ Equivalent to np.concatenate(), but return correct type """
+        """ Equivalent to np.concatenate(), but return correct type; see ss.uids.cat() for the class method """
         return np.concatenate([self, other], **kw).view(self.__class__)
 
     @classmethod
     def cat(cls, *args, **kw):
-        """ Equivalent to np.concatenate(), but return correct type """
+        """ Equivalent to np.concatenate(), but return correct type; see ss.uids.concat() for the instance method """
         arrs = args[0] if len(args) == 1 else args
         return np.concatenate(arrs, **kw).view(cls)
 
@@ -428,6 +434,15 @@ def xor(self, other, **kw):
     def to_numpy(self):
         """ Return a view as a standard NumPy array """
         return self.view(np.ndarray)
+    
+    def unique(self, return_index=False):
+        """ Return unique UIDs; equivalent to np.unique() """
+        if return_index:
+            arr, index = np.unique(self, return_index=True)
+            return arr.view(self.__class__), index
+        else:
+            arr = np.unique(self).view(self.__class__)
+            return arr
 
     # Implement collection of operators
     def __and__(self, other): return self.intersect(other)
diff --git a/starsim/utils.py b/starsim/utils.py
index 71b5e75f..2045f050 100644
--- a/starsim/utils.py
+++ b/starsim/utils.py
@@ -267,8 +267,7 @@ def standardize_data(data=None, metadata=None, min_year=1800, out_of_range=0, de
     if isinstance(data, ss.Dist):
         return data
     elif sc.isnumber(data):
-        index = pd.MultiIndex.from_arrays([[default_year, default_year], [-np.inf, 0]], names=['year','age'])
-        return pd.Series(index=index, data=[out_of_range, data])
+        return data
 
     # Convert series and dataframe inputs into dicts
     if isinstance(data, pd.Series) or isinstance(data, pd.DataFrame):
diff --git a/starsim/version.py b/starsim/version.py
index 54f75802..24910079 100644
--- a/starsim/version.py
+++ b/starsim/version.py
@@ -1,9 +1,9 @@
-'''
+"""
 Version and license information.
-'''
+"""
 
 __all__ = ['__version__', '__versiondate__', '__license__']
 
-__version__ = '0.5.10'
-__versiondate__ = '2024-07-03'
+__version__ = '1.0.0'
+__versiondate__ = '2024-07-10'
 __license__ = f'Starsim {__version__} ({__versiondate__}) — © 2023-2024 by IDM'


Test patch:
diff --git a/tests/requirements.txt b/tests/requirements.txt
index fce27433..a2bea3cf 100644
--- a/tests/requirements.txt
+++ b/tests/requirements.txt
@@ -1,4 +1,5 @@
 pytest>=7.3
 pytest-cov
 pytest-env
-pytest-xdist
\ No newline at end of file
+pytest-xdist
+optuna
\ No newline at end of file
diff --git a/tests/test_calibration.py b/tests/test_calibration.py
new file mode 100644
index 00000000..31999001
--- /dev/null
+++ b/tests/test_calibration.py
@@ -0,0 +1,148 @@
+"""
+Test calibration
+"""
+
+#%% Imports and settings
+import sciris as sc
+import starsim as ss
+
+do_plot = 1
+do_save = 0
+n_agents = 2e3
+
+
+#%% Helper functions
+
+def make_sim():
+    hiv = ss.HIV(
+        beta = {'random': [0.01]*2, 'maternal': [1, 0]},
+        init_prev = 0.15,
+    )
+    pregnancy = ss.Pregnancy(fertility_rate=20)
+    death = ss.Deaths(death_rate=10)
+    random = ss.RandomNet(n_contacts=4)
+    maternal = ss.MaternalNet()
+
+    sim = ss.Sim(
+        dt = 1,
+        n_agents = n_agents,
+        total_pop = 9980999,
+        start = 1990,
+        n_years = 40,
+        diseases = [hiv],
+        networks = [random, maternal],
+        demographics = [pregnancy, death],
+    )
+
+    return sim
+
+
+def make_data():
+    """ Define the calibration target data """
+    target_data = [
+    ['year', 'n_alive', 'hiv.prevalence', 'hiv.n_infected', 'hiv.new_infections', 'hiv.new_deaths'],
+    [  1990,  10432409,        0.0699742,          730000 ,               210000,           25000],
+    [  1991,  10681008,        0.0851979,          910000 ,               220000,           33000],
+    [  1992,  10900511,        0.1009127,          1100000,               220000,           43000],
+    [  1993,  11092775,        0.1081785,          1200000,               210000,           53000],
+    [  1994,  11261752,        0.1154349,          1300000,               200000,           63000],
+    [  1995,  11410721,        0.1226916,          1400000,               180000,           74000],
+    [  1996,  11541215,        0.1299689,          1500000,               160000,           84000],
+    [  1997,  11653254,        0.1287194,          1500000,               150000,           94000],
+    [  1998,  11747079,        0.1362040,          1600000,               140000,           100000],
+    [  1999,  11822722,        0.1353326,          1600000,               130000,           110000],
+    [  2000,  11881482,        0.1346633,          1600000,               120000,           120000],
+    [  2001,  11923906,        0.1341842,          1600000,               110000,           130000],
+    [  2002,  11954293,        0.1254779,          1500000,               100000,           130000],
+    [  2003,  11982219,        0.1251854,          1500000,               94000 ,           130000],
+    [  2004,  12019911,        0.1164734,          1400000,               89000 ,           120000],
+    [  2005,  12076697,        0.1159257,          1400000,               83000 ,           120000],
+    [  2006,  12155496,        0.1069475,          1300000,               78000 ,           110000],
+    [  2007,  12255920,        0.1060711,          1300000,               74000 ,           93000],
+    [  2008,  12379553,        0.1050118,          1300000,               69000 ,           80000],
+    [  2009,  12526964,        0.0957933,          1200000,               65000 ,           68000],
+    [  2010,  12697728,        0.0945050,          1200000,               62000 ,           54000],
+    [  2011,  12894323,        0.0930642,          1200000,               56000 ,           42000],
+    [  2012,  13115149,        0.0914972,          1200000,               49000 ,           34000],
+    [  2013,  13350378,        0.0973755,          1300000,               47000 ,           28000],
+    [  2014,  13586710,        0.0956817,          1300000,               45000 ,           25000],
+    [  2015,  13814642,        0.0941030,          1300000,               44000 ,           24000],
+    [  2016,  14030338,        0.0926563,          1300000,               43000 ,           23000],
+    [  2017,  14236599,        0.0913139,          1300000,               34000 ,           23000],
+    [  2018,  14438812,        0.0900351,          1300000,               27000 ,           22000],
+    [  2019,  14645473,        0.0920401,          1347971,               23000 ,           None],
+    [  2020,  14862927,        0.0874659,          1300000,               20000 ,           None],
+    [  2021,  15085870,        0.0861733,          1300000,               19000 ,           None],
+    [  2022,  15312158,        0.0848998,          1300000,               17000 ,           None],
+    ]
+    df = sc.dataframe(target_data[1:], columns=target_data[0])
+    return df
+
+#%% Define the tests
+
+def test_calibration(do_plot=True):
+    sc.heading('Testing calibration')
+
+    # Define the calibration parameters
+    calib_pars = dict(
+        diseases = dict(
+            hiv = dict(
+                init_prev = [0.15, 0.01, 0.30],
+            ),
+        ),
+        networks = dict(
+            randomnet = dict(
+                n_contacts = [4, 2, 10],
+            ),
+        ),
+    )
+
+    # Make the sim and data
+    sim = make_sim()
+    data = make_data()
+
+    # Define weights for the data
+    weights = {
+        'n_alive':            1.0,
+        'hiv.prevalence':     1.0,
+        'hiv.n_infected':     1.0,
+        'hiv.new_infections': 1.0,
+        'hiv.new_deaths':     1.0,
+    }
+
+    # Make the calibration
+    calib = ss.Calibration(
+        calib_pars = calib_pars,
+        sim = sim,
+        data = data,
+        weights = weights,
+        total_trials = 4,
+        n_workers = 2,
+        die = True
+    )
+
+    # Perform the calibration
+    sc.printcyan('\nPeforming calibration...')
+    calib.calibrate(confirm_fit=False)
+
+    # Confirm
+    sc.printcyan('\nConfirming fit...')
+    calib.confirm_fit()
+    print(f'Fit with original pars: {calib.before_fit:n}')
+    print(f'Fit with best-fit pars: {calib.after_fit:n}')
+    if calib.after_fit <= calib.before_fit:
+        print('✓ Calibration improved fit')
+    else:
+        print('✗ Calibration did not improve fit, but this sometimes happens stochastically and is not necessarily an error')
+    
+    return sim, calib
+
+
+#%% Run as a script
+if __name__ == '__main__':
+
+    T = sc.timer()
+
+    sim, calib = test_calibration()
+
+    T.toc()
diff --git a/tests/test_networks.py b/tests/test_networks.py
index ce4bf690..1ceb0eae 100644
--- a/tests/test_networks.py
+++ b/tests/test_networks.py
@@ -62,6 +62,7 @@ def test_random():
     o = sc.objdict(nw1=nw1, nw2=nw2, nw3=nw3)
     return o
 
+
 def test_erdosrenyi():
     sc.heading('Testing Erdos-Renyi network')
 
@@ -98,7 +99,7 @@ def test_ER(n, p, nw, alpha=0.01):
     # Manual creation
     p = 0.1
     nw1 = ss.ErdosRenyiNet(p=p)
-    s1 = ss.Sim(n_agents=small, networks=nw1, copy_inputs=False).initialize() # This initializes the network
+    ss.Sim(n_agents=small, networks=nw1, copy_inputs=False).initialize() # This initializes the network
     test_ER(small, p, nw1)
 
     # Automatic creation as part of sim
@@ -119,6 +120,7 @@ def test_ER(n, p, nw, alpha=0.01):
     o = sc.objdict(nw1=nw1, nw2=nw2, nw3=nw3)
     return o
 
+
 def test_disk():
     sc.heading('Testing Disk network')
 
@@ -179,20 +181,29 @@ def test_static():
     return o
 
 
+def test_null():
+    sc.heading('Testing NullNet...')
+    people = ss.People(n_agents=small)
+    network = ss.NullNet()
+    sir = ss.SIR(pars=dict(dur_inf=10, beta=0.1))
+    sim = ss.Sim(diseases=sir, people=people, networks=network)
+    sim.run()
+    return sim
+
+
+
 # %% Run as a script
 if __name__ == '__main__':
     do_plot = True
     sc.options(interactive=do_plot)
-
-    # Start timing
-    T = sc.tic()
+    T = sc.timer()
 
     # Run tests
-    man = test_manual()
-    rnd = test_random()
-    sta = test_static()
-    er = test_erdosrenyi()
-    d = test_disk()
-
-    sc.toc(T)
-    print('Done.')
\ No newline at end of file
+    man  = test_manual()
+    rand = test_random()
+    stat = test_static()
+    erdo = test_erdosrenyi()
+    disk = test_disk()
+    null = test_null()
+
+    T.toc()
\ No newline at end of file
diff --git a/tests/test_randomness.py b/tests/test_randomness.py
index 0f543112..120bb9e9 100644
--- a/tests/test_randomness.py
+++ b/tests/test_randomness.py
@@ -291,8 +291,8 @@ def test_combine_rands(do_plot=False):
     atol = 1e-3
     target = 0.5
     np.random.seed(2)
-    a = np.random.randint(low=0, high=np.iinfo(np.uint64).max, dtype=np.uint64, size=n)
-    b = np.random.randint(low=0, high=np.iinfo(np.uint64).max, dtype=np.uint64, size=n)
+    a = np.random.randint(np.iinfo(np.int64).max, size=n)
+    b = np.random.randint(np.iinfo(np.int64).max, size=n)
     c = ss.combine_rands(a, b)
     if do_plot:
         pl.figure()
diff --git a/tests/test_run.py b/tests/test_run.py
new file mode 100644
index 00000000..f9f47dbe
--- /dev/null
+++ b/tests/test_run.py
@@ -0,0 +1,80 @@
+"""
+Test run
+"""
+
+# %% Imports and settings
+import starsim as ss
+import sciris as sc
+import numpy as np
+import matplotlib.pyplot as plt
+
+n_agents = 1_000
+do_plot = False
+sc.options(interactive=False) # Assume not running interactively
+
+
+def make_sim_pars():
+    pars = sc.objdict(
+        n_agents = n_agents,
+        demographics = True,
+        networks = 'random',
+        diseases = sc.objdict(type='sir', beta=0.1), # To allow for modification later
+    )
+    return pars
+
+
+def test_parallel():
+    """ Test running two identical sims in parallel """
+    sc.heading('Testing parallel...')
+    pars = make_sim_pars()
+
+    # Check that two identical sims match
+    msim = ss.MultiSim([ss.Sim(pars, label='Sim1'), ss.Sim(pars, label='Sim2')])
+    msim.run(keep_people=True)
+    s1, s2 = msim.sims
+    assert np.allclose(s1.summary[:], s2.summary[:], rtol=0, atol=0, equal_nan=True)
+
+    # Check that two non-identical sims don't match
+    pars2 = sc.dcp(pars)
+    pars2.diseases.beta *= 2
+    sims = ss.MultiSim([ss.Sim(pars, label='Sim1'), ss.Sim(pars2, label='Sim2')])
+    sims.run(keep_people=True)
+    s1, s2 = sims.sims
+    assert not np.allclose(s1.summary[:], s2.summary[:], rtol=0, atol=0, equal_nan=True)
+
+    return s1, s2
+
+
+def test_multisim():
+    """ Check MultiSim methods """
+    sc.heading('Testing MultiSim')
+    pars = make_sim_pars()
+    msim = ss.MultiSim(ss.Sim(pars))
+    msim.run(n_runs=4)
+    
+    # Plot individual sims
+    msim.plot()
+    
+    # Reduce and plot mean
+    msim.mean()
+    msim.plot()
+    
+    # Reduce and plot median
+    msim.median()
+    msim.plot()
+
+    return msim
+
+
+if __name__ == '__main__':
+    do_plot = True
+    sc.options(interactive=do_plot)
+    T = sc.timer()
+
+    # s1, s2 = test_parallel()
+    msim = test_multisim()
+
+    T.toc()
+    
+    if do_plot:
+        plt.show()
diff --git a/tests/test_sim.py b/tests/test_sim.py
index ae5d04cf..b45e127c 100644
--- a/tests/test_sim.py
+++ b/tests/test_sim.py
@@ -5,7 +5,6 @@
 # %% Imports and settings
 import starsim as ss
 import sciris as sc
-import numpy as np
 import matplotlib.pyplot as plt
 import pytest
 
@@ -240,28 +239,6 @@ def test_components(do_plot=do_plot):
     return sim
 
 
-def test_parallel():
-    """ Test running two identical sims in parallel """
-    sc.heading('Testing parallel...')
-    pars = make_sim_pars()
-
-    # Check that two identical sims match
-    sims = ss.MultiSim([ss.Sim(pars, label='Sim1'), ss.Sim(pars, label='Sim2')])
-    sims.run(keep_people=True)
-    s1, s2 = sims.sims
-    assert np.allclose(s1.summary[:], s2.summary[:], rtol=0, atol=0, equal_nan=True)
-
-    # Check that two non-identical sims don't match
-    pars2 = sc.dcp(pars)
-    pars2.diseases.beta *= 2
-    sims = ss.MultiSim([ss.Sim(pars, label='Sim1'), ss.Sim(pars2, label='Sim2')])
-    sims.run(keep_people=True)
-    s1, s2 = sims.sims
-    assert not np.allclose(s1.summary[:], s2.summary[:], rtol=0, atol=0, equal_nan=True)
-
-    return s1, s2
-
-
 if __name__ == '__main__':
     do_plot = True
     sc.options(interactive=do_plot)
@@ -275,8 +252,7 @@ def test_parallel():
     sim5b, sim5i = test_simple_vax(do_plot=do_plot)
     sim6 = test_shared_product(do_plot=do_plot)
     sim7 = test_components(do_plot=do_plot)
-    sim8a, sim8b = test_parallel()
-    
+
     T.toc()
     
     if do_plot: