Repository: starsimhub/starsim (language: python)

Issue:
Starsim tries to attempt an invalid modulo operation
`  File "/Users/devengokhale/Documents/GitHub/starsim/starsim/sim.py", line 292, in start_step
    if not (self.ti % int(1.0 / self.verbose)):
            ~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~
ZeroDivisionError: integer modulo by zero`




Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 3edf66af..c1d5110b 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -7,6 +7,12 @@ 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 2.1.0 (2024-11-01)
+---------------------------
+- Adds a new approach to disease transmission called mixing pools. A mixing pool is a "mean field" coupling wherein susceptible agents are exposed to the average infectious agent. The user can create a single mixing pool using the ``MixingPool`` class, or create many pools using ``MixingPools``. Such mixing pools could be used to simulate contact matrices, for example as published by Prem et al.
+- *GitHub info*: PR `724 <https://github.com/starsimhub/starsim/pull/724>`_
+
+
 Version 2.0.0 (2024-10-01)
 --------------------------
 
diff --git a/docs/tutorials.rst b/docs/tutorials.rst
index a02506c2..b121e51d 100644
--- a/docs/tutorials.rst
+++ b/docs/tutorials.rst
@@ -9,8 +9,8 @@ User tutorials
     tutorials/tut_intro.ipynb
     tutorials/tut_buildsim.ipynb
     tutorials/tut_demographics.ipynb
-    tutorials/tut_networks.ipynb
     tutorials/tut_diseases.ipynb
+    tutorials/tut_transmission.ipynb
     tutorials/tut_interventions.ipynb
 
 Developer tutorials
diff --git a/docs/tutorials/tut_diseases.ipynb b/docs/tutorials/tut_diseases.ipynb
index 6b9ec07b..d85cbbf2 100644
--- a/docs/tutorials/tut_diseases.ipynb
+++ b/docs/tutorials/tut_diseases.ipynb
@@ -4,7 +4,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# T5 - Diseases\n",
+    "# T4 - Diseases\n",
     "\n",
     "Diseases are the cornerstone of almost any Starsim analysis. This is where you add details about what disease you are modeling, including transmissibility, natural history, and mortality outcomes.      \n",
     "\n",
diff --git a/docs/tutorials/tut_networks.ipynb b/docs/tutorials/tut_networks.ipynb
deleted file mode 100644
index 6a75c8eb..00000000
--- a/docs/tutorials/tut_networks.ipynb
+++ /dev/null
@@ -1,211 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# T4 - Networks\n",
-    "\n",
-    "Many infectious diseases transmit when people come into contact with one another. In the basic SIR example that we covered in the first two tutorials, we used a 'random' contact network, which creates random pairings of people. Generally speaking, the structure of network that you use depends on what disease you're modeling. In addition to the 'random' network, Starim contains some basic sexual networks for modeling sexually transmitted infections (STIs) and a 'maternal' network for modeling mother-to-child transmission.  \n",
-    "\n",
-    "In this tutorial, we'll demonstrate a couple of the built-in networks, and give an overview of how you might create your own for your own disease modeling work."
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Overview of Starsim's network structure\n",
-    "\n",
-    "Starsim's networks are stored in the `networks.py` file. At the top of this file you'll see the `Network` class, which provides the template for all the other networks. At a high level, all you need to know is that network stores various *parameters* that determine how pairs of people form and dissolve, and a collection of all the *contacts* that exist at a given point in time. Let's look at a simple example of the `random` network:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "metadata": {
-    "tags": []
-   },
-   "outputs": [],
-   "source": [
-    "import starsim as ss\n",
-    "\n",
-    "# Make random network\n",
-    "random = ss.RandomNet()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "collapsed": false
-   },
-   "source": [
-    "The network we just made stores two very important things: the list of contacts (found in `random.contacts`) and the parameters (`random.pars`).  \n",
-    "\n",
-    "If we just make a network on its own like this, it's not very interesting. Like all components of Starsim, it needs to be initialized within a sim before the database of contacts gets populated:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "network = ss.RandomNet()\n",
-    "sir = ss.SIR()\n",
-    "sim = ss.Sim(n_agents=2_000, diseases=sir, networks=network)\n",
-    "sim.init()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "collapsed": false
-   },
-   "source": [
-    "Now `sim.networks.random.contacts` has all the information about which agents (referenced by their UID) are paired with which other agents.\n",
-    "\n",
-    "## Sexual networks\n",
-    "Starsim contains a few basic sexual networks, including an `mf` network for modeling heterosexual partnerships and an `msm` network for men having sex with men. Modeling STIs is a straightforward extension of the SIR example we looked at previously. Here's an example modeling syphilis:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "metadata": {
-    "collapsed": false
-   },
-   "outputs": [],
-   "source": [
-    "# Create the disease\n",
-    "syph = ss.Syphilis(\n",
-    "    beta = {'mf': [0.25, 0.15]},    # See notes below \n",
-    ")\n",
-    "\n",
-    "# Create a heterosexual network\n",
-    "mf = ss.MFNet(\n",
-    "    duration=1/24,  # Using very short relationships so that we get more transmission \n",
-    "    acts=80,        # Distribution of the number of coital acts/year\n",
-    ")\n",
-    "\n",
-    "pars = dict(start=2000, dur=20, dt=1/12)\n",
-    "sim = ss.Sim(pars=pars, diseases=syph, networks=mf)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "collapsed": false
-   },
-   "source": [
-    "The main difference here is that we've added the line `syph.pars['beta'] = {'mf': [0.25, 0.15]}`. Because STIs are generally more transmissible from males to females than vice versa, we can specify `beta` as a list, and we can also allow it to vary between networks. For instance, if we also had an MSM network here, we could set \n",
-    "```\n",
-    "syph.pars.update(beta = {'mf': [0.25, 0.15], 'msm': [0.4, 0.4]})\n",
-    "```\n",
-    "The sexual network itself is structurally similar to the `random` network that we looked at previously, and can be customized by changing the `pars` dictionary that's used to construct it.\n",
-    " \n",
-    "## Maternal networks\n",
-    "Maternal networks are useful for capturing mother-to-child transmission. If you want to use them, you also need to model pregnancy. Here's an example showing how you would do this for syphilis:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "metadata": {
-    "collapsed": false
-   },
-   "outputs": [],
-   "source": [
-    "syph = ss.Syphilis(\n",
-    "    beta = {'mf': [0.25, 0.15], 'maternal': [0.99, 0]}\n",
-    ")\n",
-    "\n",
-    "# Make demographic modules\n",
-    "pregnancy = ss.Pregnancy(fertility_rate=20)\n",
-    "death = ss.Deaths(death_rate=15)\n",
-    "\n",
-    "# Make maternal network\n",
-    "maternal = ss.MaternalNet()\n",
-    "\n",
-    "# Make sim\n",
-    "sim = ss.Sim(pars=pars, diseases=syph, networks=[mf, maternal], demographics=[pregnancy, death])"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "collapsed": false
-   },
-   "source": [
-    "## Making your own network\n",
-    "If you want to make your own network, a good place to start is by inheriting from one of the templates in networks.py. For example, if you wanted to make an `mf` network that has more age structure, you could do something like this:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 11,
-   "metadata": {
-    "collapsed": false
-   },
-   "outputs": [],
-   "source": [
-    "class age_mf(ss.MFNet):\n",
-    "    def add_pairs(self, people, ti=None):\n",
-    "        # Custom code to select pairs based on age\n",
-    "        return "
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "collapsed": false
-   },
-   "source": [
-    "This network would operate exactly like the existing `mf` network, but you could control how new pairs were added.\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.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/tut_transmission.ipynb b/docs/tutorials/tut_transmission.ipynb
new file mode 100644
index 00000000..1da4546f
--- /dev/null
+++ b/docs/tutorials/tut_transmission.ipynb
@@ -0,0 +1,482 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# T5 - Transmission\n",
+    "\n",
+    "Many infectious diseases transmit when people come into contact with one another. In the basic SIR example that we covered in the first two tutorials, we used a 'random' contact network, which creates random pairings of people. Generally speaking, the structure of network that you use depends on what disease you're modeling.\n",
+    "\n",
+    "Starsim contains two different approaches to disease transmission.\n",
+    "\n",
+    "**Contact network:** Disease transmits on individual person-to-person relationships. These are networks in the sense that each agent (person) can be viewed as a node and each contact is an edge. Networks are undirected in the sense that an edge will allow transmission in both directions, however you can make transmission stronger one way (p1-->p2) than the other (p2-->p1). The framework includes several pre-built contact network algorithms, each of which manages a list of edges that can change on each simulation step. In addition to the 'random' network, the framework include algorithms representing basic sexual networks for modeling sexually transmitted infections (STIs) and a 'maternal' network for modeling mother-to-child transmission.\n",
+    "\n",
+    "**Mixing pools:** Simulate well-mixed transmission from one group of agents to another. Each mixing pool simulates directed transmission from a group of source agents to a group of destination agents. Instead of simulating individual edges like in a contact network, mixing pools first compute the total \"contagion\" that is shed by the source group. This total contagion value is normalized by the number of agents in the source group to arrive at a value that represents a typical infectious source agent. Then, each destination agent is exposed to this typical (average) infectious agent, likely resulting in new infections in the destination group. Multiple mixing pools can be run simultaneously to simulate heterogeneous transmission between various source and destination groups. Mixing pools could be used with age contact matrices like those developed by Prem et al.\n",
+    "\n",
+    "Note that contact networks and mixing pools can be used together in a single simulation.\n",
+    "\n",
+    "In this tutorial, we'll demonstrate a couple of the built-in contact networks, and give an overview of how you might create your own for your own disease modeling work. We'll also demonstrate how to use mixing pools."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Overview of Starsim's contact network structure\n",
+    "\n",
+    "Starsim's networks are stored in the `networks.py` file. At the top of this file you'll see the `Network` class, which provides the template for all the other networks. At a high level, all you need to know is that network stores various *parameters* that determine how pairs of people form and dissolve, and a collection of all the *edges* that exist at a given point in time. Let's look at a simple example of the `random` network:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "import starsim as ss\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "# Make random network\n",
+    "random = ss.RandomNet()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "collapsed": false
+   },
+   "source": [
+    "The network we just made stores two very important things: the list of contacts (found in `random.edges`) and the parameters (`random.pars`).  \n",
+    "\n",
+    "If we just make a network on its own like this, it's not very interesting. Like all components of Starsim, it needs to be initialized within a `sim` before the database of contacts gets populated:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "net = ss.RandomNet()\n",
+    "sir = ss.SIR()\n",
+    "sim = ss.Sim(n_agents=100, diseases=sir, networks=net) # Decreasing n_agents so we can plot the network below\n",
+    "sim.init()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "collapsed": false
+   },
+   "source": [
+    "We can easily take a look at the list of edges between agents. All networks are stored in `sim.networks`, and individual networks can be found at their class name after conversion to lower case. So for this example, `sim.networks.randomnet.edges` has all the information about which agents (referenced by their UID) are paired with which other agents as well as a per-edge transmission multiplier called `beta` and a per-edge duration, `dur`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sim.networks.randomnet"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Networks provide some basic functionality, for example to identify contacts. Additionally, the edge list can be quickly converted to a Pandas data frame for further analysis and manipulation."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "net = sim.networks.randomnet\n",
+    "\n",
+    "# Use the \"find_contacts\" function to find all contacts of agent with uid 0\n",
+    "contacts_of_agent_0 = net.find_contacts([0])\n",
+    "print('The following agents were identified as contacts of agent 0:', contacts_of_agent_0)\n",
+    "\n",
+    "# Convert to a pandas DataFrame and manually find all edges involving agent with uid 0\n",
+    "df = net.to_df()\n",
+    "uid0_edges = df.loc[(df['p1']==0) | (df['p2']==0)]\n",
+    "print('\\nEdges involving UID 0 extracted from the edge list data frame:')\n",
+    "display(uid0_edges)\n",
+    "\n",
+    "import networkx as nx\n",
+    "G = net.to_graph()\n",
+    "nx.draw(G, with_labels=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "## Sexual networks\n",
+    "Starsim contains a few basic sexual networks, including an `mf` network for modeling heterosexual partnerships and an `msm` network for men having sex with men. Modeling STIs is a straightforward extension of the SIR example we looked at previously. Here's an example modeling syphilis:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [],
+   "source": [
+    "# Create the disease\n",
+    "syph = ss.Syphilis(\n",
+    "    beta = {'mf': [0.25, 0.15]},    # See notes below \n",
+    ")\n",
+    "\n",
+    "# Create a heterosexual network\n",
+    "mf = ss.MFNet(\n",
+    "    duration=1/24,  # Using very short relationships so that we get more transmission \n",
+    "    acts=80,        # Distribution of the number of coital acts/year\n",
+    ")\n",
+    "\n",
+    "pars = dict(start=2000, dur=20, dt=1/12)\n",
+    "sim = ss.Sim(pars=pars, diseases=syph, networks=mf)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "collapsed": false
+   },
+   "source": [
+    "The main difference here is that we've added the line `syph.pars['beta'] = {'mf': [0.25, 0.15]}`. Because STIs are generally more transmissible from males to females than vice versa, we can specify `beta` as a list, and we can also allow it to vary between networks. For instance, if we also had an MSM network here, we could set \n",
+    "```\n",
+    "syph.pars.update(beta = {'mf': [0.25, 0.15], 'msm': [0.4, 0.4]})\n",
+    "```\n",
+    "The sexual network itself is structurally similar to the `random` network that we looked at previously, and can be customized by changing the `pars` dictionary that's used to construct it.\n",
+    " \n",
+    "## Maternal networks\n",
+    "Maternal networks are useful for capturing mother-to-child transmission. If you want to use them, you also need to model pregnancy. As new agents are born via pregnancy, edges are automatically added between mother and prenatal child in the maternal network. Here's an example showing how you would do this for syphilis:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [],
+   "source": [
+    "syph = ss.Syphilis(\n",
+    "    beta = {'mf': [0.25, 0.15], 'maternal': [0.99, 0]}\n",
+    ")\n",
+    "\n",
+    "# Make demographic modules\n",
+    "pregnancy = ss.Pregnancy(fertility_rate=20)\n",
+    "death = ss.Deaths(death_rate=15)\n",
+    "\n",
+    "# Make maternal network\n",
+    "maternal = ss.MaternalNet()\n",
+    "\n",
+    "# Make sim\n",
+    "sim = ss.Sim(pars=pars, diseases=syph, networks=[mf, maternal], demographics=[pregnancy, death])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "collapsed": false
+   },
+   "source": [
+    "## Making your own contact network\n",
+    "If you want to make your own network, a good place to start is by inheriting from one of the templates in networks.py. For example, if you wanted to make an `mf` network that has more age structure, you could do something like this:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [],
+   "source": [
+    "class age_mf(ss.MFNet):\n",
+    "    def add_pairs(self, people, ti=None):\n",
+    "        # Custom code to select pairs based on age\n",
+    "        return "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "collapsed": false
+   },
+   "source": [
+    "This network would operate exactly like the existing `mf` network, but you could control how new pairs were added.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Mixing Pools"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let's compare a simple SIR model configured with transmission using a single Mixing Pool to a similar model using a contact network. Mixing pools are not the same as contact networks, so the results are not expected to be identical."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Transmission using a mixing pool\n",
+    "mp = ss.MixingPool(beta=ss.beta(0.1), contacts=ss.poisson(lam=3)) # A single mixing pool\n",
+    "sir = ss.SIR() # Beta doesn't matter, that's just for networks\n",
+    "sim1 = ss.Sim(diseases=sir, networks=mp, verbose=0, label='Mixing Pool')\n",
+    "\n",
+    "# Transmission using a network\n",
+    "net = ss.RandomNet(n_contacts=ss.poisson(lam=3))\n",
+    "sir = ss.SIR(beta=ss.beta(0.1))\n",
+    "sim2 = ss.Sim(diseases=sir, networks=net, verbose=0, label='Contact Network')\n",
+    "\n",
+    "msim = ss.MultiSim([sim1, sim2]).run()\n",
+    "msim.plot();"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Each MixingPool has several configuration parameters, including `src` and `dst`. The value of these parameters can either be an array of uids, like `ss.uids([0,1,2])` or a callable function that takes `sim` as an argument and returns an array of uids (or a BoolArr) like `lambda sim: sim.people.age<10`. One helper class we have created is called `AgeGroup`, which takes `low` and `high` ages as parameters. This class has some performance advantages when using multiple pools via `MixingPools` through caching.\n",
+    "\n",
+    "While mixing pools do not represent individual agent-to-agent connections, each agent does have an \"effective\" number of contacts, typically drawn from a Poisson distribution. Individuals retain their number of contacts within each pool for the duration of the simulation. Effective contacts are used on the acquisition sides. On the transmission side, the \"average\" infection level, called `trans`, against which susceptible agents are challenged is computed as:\n",
+    "```python\n",
+    "trans = np.mean(disease.infectious[self.src_uids] * disease.rel_trans[self.src_uids])\n",
+    "```\n",
+    "And on the acquisition side, each susceptible agent can be thought of as connected to `eff_contacts` average infectious agents, computed as:\n",
+    "```python\n",
+    "acq = self.eff_contacts[self.dst_uids] * disease.susceptible[self.dst_uids] * disease.rel_sus[self.dst_uids]\n",
+    "```\n",
+    "Ultimately, the acquisition probability for each susceptible agent in `dst_uids` is computed as\n",
+    "```python\n",
+    "p = self.pars.beta * trans * acq\n",
+    "```\n",
+    "\n",
+    "You can see the `beta` parameter here, which is another pool parameter and should use `ss.beta()` to automatically adjust for time units. Mixing pools do not currently use the `beta` parameters configured in individual disease modules."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mp = ss.MixingPool(\n",
+    "    diseases = 'sir', # Use this pool only for SIR, not other diseases (HIV in this example)\n",
+    "    src = lambda sim: sim.people.age < 15, # Infection sources will be agents under age 15\n",
+    "    dst = ss.AgeGroup(low=15, high=None), # Here demonstrating the AgeGroup class, identical to lambda sim: sim.people.age >= 15\n",
+    "    contacts = ss.poisson(lam=2), # Contacts will be poisson distributed with a mean of 2\n",
+    "    beta = ss.beta(0.2), # And we set beta to 0.2, automatically adjusting for time units via the ss.beta wrapper\n",
+    ")\n",
+    "sim = ss.Sim(diseases=['sir', 'hiv'], networks=mp) # Two diseases, only \"sir\" will have transmission\n",
+    "sim.run()\n",
+    "sim.plot();"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "`MixingPools` (plural) allows you to configure multiple individual `MixingPool` pools. This first example shows how to do age mixing. Like `MixingPool`, `MixingPools` also takes `src` and `dst`, however each of these parameters must be a dictionary with key specifying the name of each group and value as in the singlular mixing pool proving the uids by a callable function or list of uids."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "lows = np.arange(0, 80, 5)\n",
+    "groups = ss.ndict([ss.AgeGroup(low=low, high=low+5) for low in lows[:-1]])\n",
+    "groups.append(ss.AgeGroup(low=lows[-1], high=None)) # 80+\n",
+    "n_groups = len(groups)\n",
+    "\n",
+    "# Let's create a random contact matrix. If using contact matrices from Prem et al., put data here.\n",
+    "cm = np.random.random((n_groups, n_groups)) + 3*np.diag(np.random.rand(n_groups))\n",
+    "\n",
+    "print('The first axis (rows) of the contact matrix is the SOURCE group')\n",
+    "print('The second axis (columns) of the contact matrix is the DESTINATION group')\n",
+    "plt.imshow(cm)\n",
+    "\n",
+    "mps = ss.MixingPools(\n",
+    "    contact_matrix = cm,\n",
+    "    beta = ss.beta(0.2),\n",
+    "    src = groups,\n",
+    "    dst = groups,\n",
+    ")\n",
+    "\n",
+    "# Track new infections by age using an Analyzer, you'll learn about these in a subsequent tutorial\n",
+    "class New_Infections_By_Age(ss.Analyzer):\n",
+    "    def __init__(self, bins, **kwargs):\n",
+    "        super().__init__(**kwargs)\n",
+    "        self.bins = bins\n",
+    "\n",
+    "    def init_post(self):\n",
+    "        super().init_post()\n",
+    "        self.new_cases = np.zeros((self.npts, len(self.bins)-1))\n",
+    "\n",
+    "    def step(self):\n",
+    "        new_inf = self.sim.diseases.sir.ti_infected == self.ti\n",
+    "        if not new_inf.any(): return\n",
+    "        self.new_cases[self.ti, :] = np.histogram(self.sim.people.age[new_inf], bins=self.bins)[0]\n",
+    "\n",
+    "    def plot(self):\n",
+    "        fig, ax = plt.subplots()\n",
+    "        ax.set_prop_cycle('color', plt.cm.nipy_spectral(np.linspace(0, 1, len(self.bins))))\n",
+    "\n",
+    "        for i, (b1, b2) in enumerate(zip(self.bins[:-1], self.bins[1:])):\n",
+    "            ax.plot(self.sim.timevec, self.new_cases[:,i], label=f'Age {b1}-{b2}')\n",
+    "        ax.legend(); ax.set_xlabel('Year'); ax.set_ylabel('New Infections')\n",
+    "\n",
+    "az = New_Infections_By_Age(np.concatenate([lows, [1000]]))\n",
+    "\n",
+    "sir = ss.SIR()\n",
+    "sim = ss.Sim(diseases=sir, networks=mps, analyzers=az, dur=5, dt=1/4, n_agents=2000, verbose=0)\n",
+    "sim.run()\n",
+    "sim.analyzers[0].plot()\n",
+    "print('^^^ The default age distribution is uniform, so the 75+ group contains many more people than the other age groups, explaining the larger number of new infections.')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Mixing pools work well with properties other than age. Here's an example that mixes by socio-economic status (SES)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from enum import IntEnum\n",
+    "class SES(IntEnum):\n",
+    "    LOW = 0\n",
+    "    MID = 1\n",
+    "    HIGH = 2\n",
+    "\n",
+    "# Create 5,000 people adding 50% LOW, 30% MID, and 20% HIGH SES\n",
+    "ses = ss.FloatArr('ses', default=ss.choice(a=[SES.LOW, SES.MID, SES.HIGH], p=[0.5, 0.3, 0.2]))\n",
+    "ppl = ss.People(n_agents=5_000, extra_states=ses)\n",
+    "\n",
+    "# Create the mixing pools, note that the contact matrix need not be symmetric\n",
+    "mps = ss.MixingPools(\n",
+    "    src = {s.name: lambda sim, s=s: ss.uids(sim.people.ses == s) for s in [SES.LOW, SES.MID, SES.HIGH]},\n",
+    "    dst = {s.name: lambda sim, s=s: ss.uids(sim.people.ses == s) for s in [SES.LOW, SES.MID]}, # People in the HIGH group have not route to acquire new infections, for demonstration purposes here\n",
+    "\n",
+    "    # src on rows (1st dimension), dst on cols (2nd dimension)\n",
+    "    contact_matrix = np.array([\n",
+    "        [2.50, 0.00], # LOW->LOW,  LOW->MID\n",
+    "        [0.05, 1.75], # MID->LOW,  MID->MID\n",
+    "        [0.00, 0.15], # HIGH->LOW, HIGH->MID\n",
+    "    ]),\n",
+    "\n",
+    "    beta = ss.beta(0.2),\n",
+    ")\n",
+    "\n",
+    "# Create an analyzer to track results by SES group\n",
+    "class New_Infections_By_SES(ss.Analyzer):\n",
+    "    def init_results(self):\n",
+    "        self.new_cases = np.zeros((self.npts, len(SES)))\n",
+    "\n",
+    "    def step(self):\n",
+    "        new_inf = self.sim.diseases.sir.ti_infected == self.ti\n",
+    "        if not new_inf.any():\n",
+    "            return\n",
+    "\n",
+    "        for ses in SES:\n",
+    "            self.new_cases[self.ti, ses] = np.count_nonzero(new_inf & (self.sim.people.ses==ses))\n",
+    "\n",
+    "az = New_Infections_By_SES()\n",
+    "\n",
+    "def seeding(self, sim, uids):\n",
+    "    p = np.zeros(len(uids))\n",
+    "    high_ses = ss.uids(sim.people.ses == SES.HIGH)\n",
+    "    p[high_ses] = 0.1 # Seed infections into 10% of SES HIGH\n",
+    "    return p\n",
+    "sir = ss.SIR(init_prev = ss.bernoulli(p=seeding)) # The probability of being selected as a seed is determined on a per-agent basis by the above 'seeding' function\n",
+    "\n",
+    "sim = ss.Sim(people=ppl, diseases=sir, networks=mps, analyzers=az, dt=1/12, dur=35, verbose=0)\n",
+    "sim.run()\n",
+    "\n",
+    "\n",
+    "fig, ax = plt.subplots()\n",
+    "new_cases = sim.analyzers[0].new_cases\n",
+    "for ses in [SES.LOW, SES.MID, SES.HIGH]:\n",
+    "    ax.plot(sim.results.timevec, new_cases[:,ses], label=ses.name)\n",
+    "ax.legend(); ax.set_xlabel('Year'); ax.set_ylabel('New Infections');"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "In the above figure, infections are seeded into the HIGH group at the beginning of the simulation, explaining the spike in new cases at that time. HIGH SES can only infect MID SES, so the MID (orange) wave follows. LOW SES can only become infected by MID SES in this example, so the blue wave follows. This example intentionally demonstrates the directional nature of mixing pools and includes strong mixing within groups and bottlenecks between groups to make the point."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "IDM",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  },
+  "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/pyproject.toml b/pyproject.toml
index 2bf7067d..c5529abc 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -34,12 +34,12 @@ classifiers = [
 ]
 
 dependencies = [
-    "numpy",
-    "scipy",
+    "numpy>=1.24.0",
     "pandas>=2.0.0",
     "sciris>=3.2.0",
-    "matplotlib",
     "numba",
+    "scipy",
+    "matplotlib",
     "networkx",
     "optuna"
 ]
diff --git a/starsim/demographics.py b/starsim/demographics.py
index db0357d2..1663d041 100644
--- a/starsim/demographics.py
+++ b/starsim/demographics.py
@@ -110,7 +110,7 @@ def update_results(self):
         # Calculate crude birth rate (CBR)
         inv_units = 1.0/self.pars.units
         births_per_year = self.n_births/self.sim.dt_year
-        denom = self.sim.results.n_alive[self.sim.ti]
+        denom = self.sim.results.n_alive[self.ti]
         self.results.cbr[self.ti] = inv_units*births_per_year/denom
         return
 
@@ -186,8 +186,11 @@ def make_death_prob_fn(self, sim, uids):
         """ Take in the module, sim, and uids, and return the probability of death for each UID on this timestep """
 
         drd = self.death_rate_data
+        time_factor = ss.time_ratio(unit1=self.unit, unit2='year')
         if sc.isnumber(drd):
             death_rate = drd
+            if isinstance(death_rate, ss.TimePar):
+                time_factor = 1 # Time conversion performed automatically by TimePar
         else:
 
             ppl = sim.people
@@ -197,7 +200,7 @@ def make_death_prob_fn(self, sim, uids):
             assert len(uids) == len(ppl.auids)
 
             available_years = drd.index.get_level_values('year')
-            year_ind = sc.findnearest(available_years, sim.now) # TODO: make work with different timesteps
+            year_ind = sc.findnearest(available_years, sim.now_year) # TODO: make work with different timesteps
             nearest_year = available_years[year_ind]
 
             death_rate = np.empty(uids.shape, dtype=ss_float_)
@@ -215,7 +218,7 @@ def make_death_prob_fn(self, sim, uids):
                 death_rate[:] = s.values[binned_ages]
 
         # Scale from rate to probability. Consider an exponential here.
-        death_prob = death_rate * (self.pars.units * self.pars.rel_death * sim.pars.dt)
+        death_prob = death_rate * (self.pars.units * self.pars.rel_death * sim.pars.dt) * time_factor
         death_prob = np.clip(death_prob, a_min=0, a_max=1)
 
         return death_prob
@@ -303,10 +306,13 @@ def make_fertility_prob_fn(self, sim, uids):
         frd = self.fertility_rate_data
         fertility_rate = np.zeros(len(sim.people.uid.raw), dtype=ss_float_)
 
+        time_factor = ss.time_ratio(unit1=self.unit, unit2='year')
         if sc.isnumber(frd):
             fertility_rate[uids] = self.fertility_rate_data
+            if isinstance(frd, ss.TimePar):
+                time_factor = 1 # Time conversion performed automatically by TimePar
         else:
-            year_ind = sc.findnearest(frd.index, sim.now-self.pars.dur_pregnancy) # TODO: make time-unit-aware
+            year_ind = sc.findnearest(frd.index, self.now_year-self.pars.dur_pregnancy.to('year')) # TODO: make time-unit-aware
             nearest_year = frd.index[year_ind]
 
             # Assign agents to age bins
@@ -333,7 +339,7 @@ def make_fertility_prob_fn(self, sim, uids):
 
         # Scale from rate to probability
         invalid_age = (age < self.pars.min_age) | (age > self.pars.max_age)
-        fertility_prob = fertility_rate * (self.pars.units * self.pars.rel_fertility * sim.pars.dt)
+        fertility_prob = fertility_rate * (self.pars.units * self.pars.rel_fertility * self.dt) * time_factor
         fertility_prob[(~self.fecund).uids] = 0 # Currently infecund women cannot become pregnant
         fertility_prob[uids[invalid_age]] = 0 # Women too young or old cannot become pregnant
         fertility_prob = np.clip(fertility_prob[uids], a_min=0, a_max=1)
@@ -382,12 +388,12 @@ def init_results(self):
         return
 
     def step(self):
-        if self.sim.ti == 0 and self.pars.burnin:
+        if self.ti == 0 and self.pars.burnin:
             dtis = np.arange(np.ceil(-1 * self.pars.dur_pregnancy), 0, 1).astype(int)
             for dti in dtis:
-                self.sim.ti = dti
+                self.ti = dti
                 self.do_step()
-            self.sim.ti = 0
+            self.ti = 0
         new_uids = self.do_step()
 
         return new_uids
@@ -396,7 +402,7 @@ def do_step(self):
         """ Perform all updates """
         self.update_states()
         conceive_uids = self.make_pregnancies()
-        self.n_pregnancies = len(conceive_uids)
+        self.n_pregnancies += len(conceive_uids) # += to handle burn-in
         new_uids = self.make_embryos(conceive_uids)
         return new_uids
 
@@ -454,6 +460,9 @@ def make_pregnancies(self):
         eligible_uids = self.sim.people.female.uids
         conceive_uids = self.pars.p_fertility.filter(eligible_uids)
 
+        if len(conceive_uids) == 0:
+            return ss.uids()
+
         # Validation
         if np.any(self.pregnant[conceive_uids]):
             which_uids = conceive_uids[self.pregnant[conceive_uids]]
@@ -461,8 +470,7 @@ def make_pregnancies(self):
             raise ValueError(errormsg)
 
         # Set prognoses for the pregnancies
-        if len(conceive_uids) > 0:
-            self.set_prognoses(conceive_uids)
+        self.set_prognoses(conceive_uids)
         return conceive_uids
 
     def make_embryos(self, conceive_uids):
@@ -491,8 +499,8 @@ def make_embryos(self, conceive_uids):
                     start = np.full(n_unborn, fill_value=self.ti)
                     layer.add_pairs(conceive_uids, new_uids, dur=durs, start=start)
 
-        if self.sim.ti < 0:
-            people.age[new_uids] += -self.sim.ti * self.dt # Age to ti=0, TODO: Check!
+        if self.ti < 0:
+            people.age[new_uids] += -self.ti * self.sim.dt_year # Age to ti=0
 
         return new_uids
 
@@ -511,7 +519,7 @@ def set_prognoses(self, uids):
         self.ti_pregnant[uids] = ti
 
         # Outcomes for pregnancies
-        dur_preg = np.ones(len(uids))*self.pars.dur_pregnancy  # Duration in years
+        dur_preg = np.ones(len(uids))*self.pars.dur_pregnancy
         dur_postpartum = self.pars.dur_postpartum.rvs(uids)
         dead = self.pars.p_maternal_death.rvs(uids)
         self.ti_delivery[uids] = ti + dur_preg # Currently assumes maternal deaths still result in a live baby
@@ -524,7 +532,7 @@ def set_prognoses(self, uids):
 
     def finish_step(self):
         super().finish_step()
-        death_uids = ss.uids(self.sim.people.ti_dead <= self.sim.ti)
+        death_uids = ss.uids(self.sim.people.ti_dead <= self.ti)
         if len(death_uids) == 0:
             return
 
@@ -553,10 +561,15 @@ def update_results(self):
         ti = self.ti
         self.results['pregnancies'][ti] = self.n_pregnancies
         self.results['births'][ti] = self.n_births
+
+        # Reset for the next step
+        self.n_pregnancies = 0
+        self.n_births = 0
+
         return
 
     def finalize(self):
         super().finalize()
         n_alive = self.sim.results.n_alive
-        self.results['cbr'] = 1/self.pars.units * np.divide(self.results['births'] / self.sim.dt_year, n_alive, where=n_alive>0)
+        self.results['cbr'][:] = 1/self.pars.units * np.divide(self.results['births'] / self.sim.dt_year, n_alive, where=n_alive>0)
         return
diff --git a/starsim/disease.py b/starsim/disease.py
index 378aaeba..e340b34f 100644
--- a/starsim/disease.py
+++ b/starsim/disease.py
@@ -152,7 +152,7 @@ class Infection(Disease):
     operate on to capture co-infection
     """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, pars=None, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.define_states(
             ss.State('susceptible', default=True, label='Susceptible'),
@@ -162,6 +162,11 @@ def __init__(self, *args, **kwargs):
             ss.FloatArr('ti_infected', label='Time of infection' ),
         )
 
+        self.define_pars(
+            init_prev = None, # Replace None with a ss.bernoulli to seed infections
+        )
+        self.update_pars(pars, **kwargs)
+
         # Define random number generator for determining transmission
         self.trans_rng = ss.multi_random('source', 'target')
         return
@@ -296,7 +301,7 @@ def infect(self):
 
         for i, (nkey,net) in enumerate(self.sim.networks.items()):
             nk = ss.standardize_netkey(nkey)
-            if len(net): # Skip networks with no edges
+            if isinstance(net, ss.Network) and len(net): # Skip networks with no edges
                 edges = net.edges
                 p1p2b0 = [edges.p1, edges.p2, betamap[nk][0]] # Person 1, person 2, beta 0
                 p2p1b1 = [edges.p2, edges.p1, betamap[nk][1]] # Person 2, person 1, beta 1
diff --git a/starsim/diseases/ncd.py b/starsim/diseases/ncd.py
index d55fe67e..841eda99 100644
--- a/starsim/diseases/ncd.py
+++ b/starsim/diseases/ncd.py
@@ -55,7 +55,8 @@ def step_state(self):
         ti = self.ti
         deaths = (self.ti_dead == ti).uids
         self.sim.people.request_death(deaths)
-        self.log.add_data(deaths, died=True)
+        if self.pars.log:
+            self.log.add_data(deaths, died=True)
         self.results.new_deaths[ti] = len(deaths) # Log deaths attributable to this module
         return
 
diff --git a/starsim/distributions.py b/starsim/distributions.py
index 42bb243f..3162c969 100644
--- a/starsim/distributions.py
+++ b/starsim/distributions.py
@@ -11,18 +11,17 @@
 __all__ = ['link_dists', 'make_dist', 'dist_list', 'Dists', 'Dist']
 
 
-def str2int(string, modulo=1_000_000):
+def str2int(string, modulo=1_000_000_000):
     """
-    Convert a string to an int
+    Convert a string to an int to use as a random seed; not for the user
 
     Cannot use Python's built-in hash() since it's randomized for strings. Hashlib
-    is 5x slower than int.from_bytes(string.encode(), byteorder='big'), but should
+    (sc.sha) is 5x slower than int.from_bytes(string.encode(), byteorder='big'), but should
     only add a couple milliseconds to a typical sim.
     """
-    integer = int.from_bytes(string.encode(), byteorder='big')
-    # integer = sc.sha(string, asint=True) # Hash the string to an integer
-    out = integer % modulo # Don't need all of it, this is more user-friendly
-    return out
+    integer = sc.sha(string, asint=True) # Hash the string to an integer
+    seed = integer % modulo # Ensure a user-friendly representation user-friendly
+    return seed
 
 
 def link_dists(obj, sim, module=None, overwrite=False, init=False, **kwargs): # TODO: actually link the distributions to the modules! Currently this only does the opposite, but should have mod.dists as well
@@ -85,7 +84,10 @@ def init(self, obj=None, base_seed=None, sim=None, force=False):
             raise ValueError(errormsg)
 
         # Do not look for distributions in the people states, since they shadow the "real" states
-        skip = id(sim.people._states) if sim is not None else None
+        skip = dict(
+            ids=id(sim.people._states) if sim is not None else None,
+            keys='module',
+        )
 
         # Find and initialize the distributions
         self.dists = sc.search(obj, type=Dist, skip=skip, flatten=True)
@@ -1017,7 +1019,7 @@ def filter(self, uids=None, both=False):
         """ Return UIDs that correspond to True, or optionally return both True and False """
         if uids is None:
             uids = self.sim.people.auids # All active UIDs
-        elif isinstance(uids, ss.BoolArr):
+        elif isinstance(uids, (ss.BoolArr, ss.IndexArr)):
             uids = uids.uids
 
         bools = self.rvs(uids)
diff --git a/starsim/interventions.py b/starsim/interventions.py
index 0fae1f8b..baa4528d 100644
--- a/starsim/interventions.py
+++ b/starsim/interventions.py
@@ -464,4 +464,4 @@ class campaign_vx(BaseVaccination, CampaignDelivery):
     Campaign vaccination - an instance of base vaccination combined with campaign delivery.
     See base classes for a description of input arguments.
     """
-    pass
\ No newline at end of file
+    pass
diff --git a/starsim/modules.py b/starsim/modules.py
index e8039aae..cbf5e598 100644
--- a/starsim/modules.py
+++ b/starsim/modules.py
@@ -5,6 +5,7 @@
 import sciris as sc
 import starsim as ss
 from functools import partial
+import datetime as dt
 
 __all__ = ['module_map', 'find_modules', 'Module', 'Analyzer', 'Connector']
 
@@ -55,7 +56,7 @@ class Module(sc.quickobj):
         dt (float): the timestep (e.g. 1.0, 0.1); inherits from sim if not supplied
     """
 
-    def __init__(self, name=None, label=None, unit=None, dt=None):
+    def __init__(self, name=None, label=None, unit=None, dt=None, **kwargs):
         # Handle parameters
         self.pars = ss.Pars() # Usually populated via self.define_pars()
         self.set_metadata(name, label) # Usually reset as part of self.update_pars()
@@ -206,11 +207,30 @@ def init_time_pars(self, force=False):
     def now(self):
         """ Return the current time, i.e. the time vector at the current timestep """
         try:
-            return self.timevec[self.ti]
+            if self.ti >= 0:
+                return self.timevec[self.ti]
+            else:
+                if sc.isnumber(self.sim.pars.start):
+                    return self.sim.pars.start + self.ti * self.dt * ss.time_ratio(unit1=self.unit, unit2='year')
+                else:
+                    assert isinstance(self.sim.pars.start, dt.date), 'Expected a datetime'
+                    if self.unit == 'day':
+                        return self.sim.pars.start + dt.timedelta(days=int(self.ti * self.dt))
+                    else:
+                        assert self.unit == 'year', 'Expected unit of "year"'
+                        return self.sim.pars.start + dt.timedelta(days=int(365.25 * self.ti * self.dt))
         except Exception as E:
             ss.warn(f'Encountered exception when getting the current time in {self.name}: {E}')
             return None
 
+    @property
+    def now_year(self):
+        """ Like now, but convert datetime to floating point year """
+        now = self.now
+        if isinstance(now, dt.date):
+            return sc.datetoyear(now)
+        return now
+
     def start_step(self):
         """ Tasks to perform at the beginning of the step """
         if self.dists is not None: # Will be None if no distributions are defined
diff --git a/starsim/networks.py b/starsim/networks.py
index 80f19b5f..088a6fad 100644
--- a/starsim/networks.py
+++ b/starsim/networks.py
@@ -15,11 +15,18 @@
 
 
 # Specify all externally visible functions this file defines; see also more definitions below
-__all__ = ['Network', 'DynamicNetwork', 'SexualNetwork']
+__all__ = ['Route', 'Network', 'DynamicNetwork', 'SexualNetwork']
 
 # %% General network classes
 
-class Network(ss.Module):
+class Route(ss.Module):
+    """
+    A transmission route -- e.g., a network, mixing pool, environmental transmission, etc.
+    """
+    pass
+
+
+class Network(Route):
     """
     A class holding a single network of contact edges (connections) between people
     as well as methods for updating these.
@@ -987,3 +994,265 @@ class PostnatalNet(MaternalNet):
     def __init__(self, key_dict=None, prenatal=False, postnatal=True, **kwargs):
         super().__init__(key_dict=key_dict, prenatal=prenatal, postnatal=postnatal, **kwargs)
         return
+
+
+__all__ += ['AgeGroup', 'MixingPools', 'MixingPool']
+
+class AgeGroup(sc.prettyobj):
+    """ A simple age-based filter that returns uids of agents that match the criteria """
+    def __init__(self, low, high, do_cache=True):
+        self.low = low
+        self.high = high
+
+        self.do_cache = do_cache
+        self.uids = None # Cached
+        self.ti_cache = -1
+
+        self.name = repr(self)
+        return
+
+    def __call__(self, sim):
+        if (not self.do_cache) or (self.ti_cache != sim.ti):
+            in_group = sim.people.age >= self.low
+            if self.high is not None:
+                in_group = in_group & (sim.people.age < self.high)
+            self.uids = ss.uids(in_group)
+            self.ti_cache = sim.ti
+        return self.uids
+
+    def __repr__(self):
+        return f'age({self.low}-{self.high})'
+
+
+class MixingPools(Route):
+    """
+    A container for creating a rectangular array of MixingPool instances
+
+    By default, separates the population into <15 and >15 age groups.
+
+    Args:
+        diseases (str): the diseases that transmit via these mixing pools
+        src (inds): source agents; can be AgeGroup(), ss.uids(), or lambda(sim); None indicates all alive agents
+        dst (inds): destination agents; as above
+        beta (float): overall transmission via these mixing pools
+        contacts (array): the relative connectivity between different mixing pools (can be float or Dist)
+
+    **Example**::
+
+        import starsim as ss
+        mps = ss.MixingPools(
+            diseases = 'sis',
+            beta = 0.1,
+            src = {'0-15': ss.AgeGroup(0, 15), '15+': ss.AgeGroup(15, None)},
+            dst = {'0-15': ss.AgeGroup(0, 15), '15+': ss.AgeGroup(15, None)},
+            contacts = [[2.4, 0.49], [0.91, 0.16]],
+        )
+        sim = ss.Sim(diseases='sis', networks=mps).run()
+        sim.plot()
+    """
+    def __init__(self, pars=None, **kwargs):
+        super().__init__(**kwargs)
+        self.define_pars(
+            diseases = None,
+            src = None,
+            dst = None,
+            beta = ss.beta(0.1),
+            contacts = None,
+        )
+        self.update_pars(pars, **kwargs)
+        self.validate_pars()
+        self.pools = []
+        return
+
+    def __len__(self):
+        try:
+            return len(self.pools)
+        except:
+            return 0
+
+    def validate_pars(self):
+        """ Check that src and dst have correct types, and contacts is the correct shape """
+        p = self.pars
+
+        # Validate src and dst
+        if not isinstance(p.src, dict):
+            errormsg = f'src must be a provided as a dictionary, not {type(self.pars.src)}'
+            raise TypeError(errormsg)
+        if not isinstance(p.dst, dict):
+            raise TypeError(f'dst must be a provided as a dictionary, not {type(self.pars.src)}')
+        p.src = sc.objdict(p.src)
+        p.dst = sc.objdict(p.dst)
+
+        # Validate the contacts
+        p.contacts = np.array(p.contacts)
+        actual = p.contacts.shape
+        expected = (len(p.src), len(p.dst))
+        if actual != expected:
+            errormsg = f'The number of source and destination groups must match the number of rows and columns in the mixing matrix, but {actual} != {expected}.'
+            raise ValueError(errormsg)
+
+        return
+
+    def init_pre(self, sim):
+        super().init_pre(sim)
+        p = self.pars
+
+        self.pools = []
+        for i,sk,src in p.src.enumitems():
+            for j,dk,dst in p.dst.enumitems():
+                contacts = p.contacts[i,j]
+                if sc.isnumber(contacts): # If it's a number, convert to a distribution
+                    contacts = ss.poisson(lam=contacts)
+                name = f'pool:{sk}->{dk}'
+                mp = MixingPool(name=name, diseases=p.diseases, beta=p.beta, contacts=contacts, src=src, dst=dst)
+                mp.init_pre(sim) # Initialize the pool
+                self.pools.append(mp)
+        return
+
+    def init_post(self):
+        """ Initialize each mixing pool """
+        for mp in self.pools:
+            mp.init_post()
+        return
+
+    def step(self):
+        """ Step each mixing pool """
+        for mp in self.pools:
+            mp.step()
+        return
+
+    def remove_uids(self, uids):
+        """ Remove UIDs from each mixing pool """
+        for mp in self.pools:
+            mp.remove_uids(uids)
+        return
+
+
+class MixingPool(Route):
+    """
+    Define a single mixing pool; can be used as a drop-in replacement for a network.
+
+    Args:
+        diseases (str): the diseases that transmit via this mixing pool
+        src (inds): source agents; can be AgeGroup(), ss.uids(), or lambda(sim); None indicates all alive agents
+        dst (inds): destination agents; as above
+        beta (float): overall transmission
+        contacts (Dist): the number of effective contacts of the destination agents
+
+    **Example**::
+
+        import starsim as ss
+
+        # Set the parameters
+        mp_pars = dict(
+            src = lambda sim: sim.people.male, # only males are infectious
+            dst = None, # all agents are susceptible
+            beta = ss.beta(0.2),
+            contacts = ss.poisson(lam=4),
+        )
+
+        # Seed 5% of the male population
+        def p_init(self, sim, uids):
+            return 0.05*sim.people.male
+
+        # Create and run the sim
+        sis = ss.SIS(init_prev=p_init)
+        mp = ss.MixingPool(mp_pars)
+        sim = ss.Sim(diseases=sis, networks=mp)
+        sim.run()
+        sim.plot()
+    """
+    def __init__(self, pars=None, **kwargs):
+        super().__init__(**kwargs)
+
+        self.define_pars(
+            diseases = None,
+            src = None,
+            dst = None, # Same as src
+            beta = ss.beta(0.2),
+            contacts = ss.poisson(lam=1),
+        )
+        self.update_pars(pars, **kwargs)
+
+        self.define_states(
+            ss.FloatArr('eff_contacts', default=self.pars.contacts, label='Effective number of contacts')
+        )
+
+        self.pars.diseases = sc.promotetolist(self.pars.diseases)
+        self.diseases = None
+        self.src_uids = None
+        self.dst_uids = None
+
+        self.p_acquire = ss.bernoulli(p=0) # Placeholder value
+
+        return
+
+    def __len__(self):
+        try:
+            return len(self.pars.dst)
+        except:
+            return 0
+
+    def init_post(self):
+        super().init_post()
+
+        if len(self.pars.diseases) == 0:
+            self.diseases = [d for d in self.sim.diseases.values() if isinstance(d, ss.Infection)] # Assume the user wants all communicable diseases
+        else:
+            self.diseases = []
+            for d in self.pars.diseases:
+                if not isinstance(d, str):
+                    raise TypeError(f'Diseases can be specified as ss.Disease objects or strings, not {type(d)}')
+                if d not in self.sim.diseases:
+                    raise KeyError(f'Could not find disease with name {d} in the list of diseases.')
+
+                dis = self.sim.diseases[d]
+                if not isinstance(dis, ss.Infection):
+                    raise TypeError(f'Cannot create a mixing pool for disease {d}. Mixing pools only work for communicable diseases.')
+                self.diseases.append(dis)
+
+            if len(self.diseases) == 0:
+                raise ValueError('You must specify at least one transmissible disease to use mixing pools')
+        return
+
+    def get_uids(self, func_or_array):
+        if func_or_array is None:
+            return self.sim.people.auids
+        elif callable(func_or_array):
+            return func_or_array(self.sim)
+        elif isinstance(func_or_array, ss.uids):
+            return func_or_array
+        raise Exception('src must be either a callable function, e.g. lambda sim: ss.uids(sim.people.age<5), or an array of uids.')
+
+    def remove_uids(self, uids):
+        """ If UIDs are supplied explicitly, remove them if people die """
+        for key in ['src', 'dst']:
+            inds = self.pars[key]
+            if isinstance(inds, ss.uids):
+                self.pars[key] = inds.remove(uids)
+        return
+
+    def step(self):
+        super().step()
+        self.src_uids = self.get_uids(self.pars.src)
+        self.dst_uids = self.get_uids(self.pars.dst)
+
+        if self.pars.beta == 0:
+            return 0
+
+        if len(self.src_uids) == 0 or len(self.dst_uids) == 0:
+            return 0
+
+        n_new_cases = 0
+        for disease in self.diseases:
+            trans = np.mean(disease.infectious[self.src_uids] * disease.rel_trans[self.src_uids])
+            acq = self.eff_contacts[self.dst_uids] * disease.susceptible[self.dst_uids] * disease.rel_sus[self.dst_uids]
+            p = self.pars.beta * trans * acq #1 - np.exp(-self.pars.beta * trans * acq)
+
+            self.p_acquire.set(p=p)
+            new_cases = self.p_acquire.filter(self.dst_uids)
+            n_new_cases += len(new_cases)
+
+            disease.set_prognoses(new_cases)
+
+        return n_new_cases
diff --git a/starsim/parameters.py b/starsim/parameters.py
index db378ca4..af1ea19b 100644
--- a/starsim/parameters.py
+++ b/starsim/parameters.py
@@ -68,6 +68,8 @@ def update(self, pars=None, create=False, **kwargs):
                     self._update_dist(key, old, new)
                 elif callable(old): # It's a function: update directly
                     self[key] = new
+                elif isinstance(old, dict):
+                    self[key] = new # Take dictionaries directly, without warning the user
                 else: # Everything else; not used currently but could be
                     warnmsg = f'No known mechanism for handling {type(old)} → {type(new)}; using default'
                     ss.warn(warnmsg)
@@ -101,13 +103,7 @@ def _update_timepar(self, key, old, new):
 
         # It's a TimePar, e.g. dur_inf = ss.dur(6); use directly
         if isinstance(new, ss.TimePar):
-            new_cls = new.__class__
-            old_cls = old.__class__
-            if new_cls == old_cls:
-                self[key] = new
-            else:
-                errormsg = f'When updating a time parameter, the new class ({new_cls}) should match the original class ({old_cls})'
-                raise TypeError(errormsg)
+            self[key] = new
 
         # It's a single number, e.g. dur_inf = 6; set parameters
         elif isinstance(new, Number):
@@ -369,8 +365,12 @@ def validate_demographics(self):
     def validate_networks(self):
         """ Validate networks """
         # Don't allow more than one prenatal or postnatal network
-        prenatal_nets  = {k:nw for k,nw in self.networks.items() if nw.prenatal}
-        postnatal_nets = {k:nw for k,nw in self.networks.items() if nw.postnatal}
+        prenatal_nets = []
+        postnatal_nets = []
+        for k,nw in self.networks.items():
+            if isinstance(nw, ss.Network):
+                if nw.prenatal: prenatal_nets.append(k)
+                if nw.postnatal: postnatal_nets.append(k)
         if len(prenatal_nets) > 1:
             errormsg = f'Starsim currently only supports one prenatal network; prenatal networks are: {prenatal_nets.keys()}'
             raise ValueError(errormsg)
diff --git a/starsim/run.py b/starsim/run.py
index 064591c1..af561da7 100644
--- a/starsim/run.py
+++ b/starsim/run.py
@@ -429,10 +429,6 @@ def single_run(sim, ind=0, reseed=True, keep_people=False, run_args=None, sim_ar
         sim.pars['rand_seed'] += ind  # Reset the seed, otherwise no point of parallel runs
         ss.set_seed() # Note: may not be needed
 
-    if verbose >= 1:
-        verb = 'Running' if do_run else 'Creating'
-        print(f'{verb} a simulation using seed={sim["rand_seed"]}')
-
     # Handle additional arguments
     for key, val in sim_args.items():
         print(f'Processing {key}:{val}')
diff --git a/starsim/sim.py b/starsim/sim.py
index 155df378..3499d841 100644
--- a/starsim/sim.py
+++ b/starsim/sim.py
@@ -5,6 +5,7 @@
 import numpy as np
 import sciris as sc
 import starsim as ss
+import datetime as dt
 import matplotlib.pyplot as plt
 
 __all__ = ['Sim', 'AlreadyRunError', 'demo', 'diff_sims', 'check_sims_match']
@@ -115,11 +116,30 @@ def __repr__(self):
     def now(self):
         """ Return the current time, i.e. the time vector at the current timestep """
         try:
-            ti = min(self.ti, len(self.timevec)-1) # During integration, ti can go one past the end of the time vector
-            return self.timevec[ti]
+            if self.ti >= 0:
+                ti = min(self.ti, len(self.timevec)-1) # During integration, ti can go one past the end of the time vector
+                return self.timevec[ti]
+            else:
+                if sc.isnumber(self.pars.start):
+                    return self.pars.start + self.ti * self.dt * ss.time_ratio(unit1=self.unit, unit2='year')
+                else:
+                    assert isinstance(self.pars.start, dt.date), 'Expected a datetime'
+                    if self.unit == 'day':
+                        return self.pars.start + dt.timedelta(days=int(self.ti * self.dt))
+                    else:
+                        assert self.unit == 'year', 'Expected unit of "year"'
+                        return self.pars.start + dt.timedelta(days=int(365.25 * self.ti * self.dt))
         except Exception as E:
-            ss.warn(f'Encountered exception in sim when getting the current time: {E}')
-            return None
+            ss.warn(f'Encountered exception when getting the current time in {self.name}: {E}')
+            return None 
+
+    @property
+    def now_year(self):
+        """ Like now, but convert datetime to floating point year """
+        now = self.now
+        if isinstance(now, dt.date):
+            return sc.datetoyear(now)
+        return now
 
     @property
     def modules(self):
@@ -266,7 +286,7 @@ def start_step(self):
             timepoint = self.timevec[self.ti]
             timelabel = f'{timepoint:0.1f}' if isinstance(timepoint, float) else str(timepoint) # TODO: fix
             string = f'  Running {simlabel}{timelabel} ({self.ti:2.0f}/{self.npts}) ({self.elapsed:0.2f} s) '
-            if self.verbose >= 2:
+            if self.verbose >= 1:
                 sc.heading(string)
             elif self.verbose > 0:
                 if not (self.ti % int(1.0 / self.verbose)):
@@ -598,6 +618,8 @@ def plot(self, key=None, fig=None, style='fancy', show_data=True, fig_kw=None, p
                 ax.set_title(res.full_label)
                 ax.set_xlabel('Time')
 
+        if self.label is not None:
+            fig.suptitle(self.label)
         sc.figlayout(fig=fig)
 
         return fig
diff --git a/starsim/timepars.py b/starsim/timepars.py
index bb730c45..813a253b 100644
--- a/starsim/timepars.py
+++ b/starsim/timepars.py
@@ -7,7 +7,7 @@
 
 # Classes that are externally visible
 __all__ = ['time_units', 'time_ratio', 'date_add', 'date_diff', 'make_timevec', 'make_abs_tvec',
-           'TimePar', 'dur', 'days', 'years', 'rate', 'perday', 'peryear', 'time_prob', 'beta']
+           'TimePar', 'dur', 'days', 'years', 'rate', 'perday', 'peryear', 'time_prob', 'beta', 'rate_prob']
 
 
 #%% Helper functions
@@ -97,7 +97,7 @@ def make_timevec(start, stop, dt, unit):
     else:
         day_delta = time_ratio(unit1=unit, dt1=dt, unit2='day', dt2=1.0, as_int=True)
         if day_delta > 0:
-            timevec = sc.daterange(start, stop, interval={'days':day_delta})
+            timevec = np.array(sc.daterange(start, stop, interval={'days':day_delta}))
         else:
             errormsg = f'Timestep {dt} is too small; must be at least 1 day'
             raise ValueError(errormsg)
@@ -108,15 +108,16 @@ def make_abs_tvec(tv, unit, sim_unit):
     """ Convert a module time vector into a numerical time array with the same units as the sim """
 
     # It's an array of days or years: easy
-    if sc.isarray(tv):
+
+    if sc.isnumber(tv[0]):
         ratio = time_ratio(unit1=unit, unit2=sim_unit)
         abstv = tv*ratio # Get the units right
         abstv -= abstv[0] # Start at 0
 
     # It's a date: convert to fractional years and then subtract the
     else:
-        yearvec = [sc.datetoyear(d) for d in tv]
-        absyearvec = np.array(yearvec) - yearvec[0] # Subtract start date
+        yearvec = np.vectorize(sc.datetoyear)(tv)
+        absyearvec = yearvec - yearvec[0] # Subtract start date
         abstv = absyearvec*time_ratio(unit1='year', unit2=sim_unit)
 
     # Round to the value of epsilon; alternative to np.round(abstv/eps)*eps, which has floating point error
@@ -201,7 +202,32 @@ def __repr__(self):
         else:
             parentstr = ''
 
-        return f'ss.{name}({self.v}, unit={self.unit}{parentstr}{xstr})'
+        default_dt = sc.ifelse(self.self_dt, 1.0) == 1.0
+        if not default_dt:
+            dtstr = f', self_dt={self.self_dt}'
+        else:
+            dtstr = ''
+
+        # Rather than ss.dur(3, unit='day'), dispaly as ss.days(3)
+        prefixstr = 'ss.'
+        key = (name, self.unit)
+        mapping = {
+            ('dur',  'day'):  'days',
+            ('dur',  'year'): 'years',
+            ('rate', 'day'):  'perday',
+            ('rate', 'year'): 'peryear',
+        }
+
+        if key in mapping and default_dt:
+            prefixstr += mapping[key]
+            unitstr = ''
+        else:
+            prefixstr += name
+            unitstr = f', unit={self.unit}'
+
+        suffixstr = unitstr + parentstr + dtstr + xstr
+
+        return f'{prefixstr}({self.v}{suffixstr})'
 
     @property
     def isarray(self):
@@ -252,10 +278,12 @@ def to(self, unit=None, dt=None):
         new = self.asnew()
         unit = sc.ifelse(unit, self.parent_unit, self.unit)
         parent_dt = sc.ifelse(dt, 1.0)
-        new.factor = time_ratio(unit1=self.unit, dt1=self.self_dt, unit2=unit, dt2=parent_dt) # Calculaet the new factor
+        new.factor = time_ratio(unit1=self.unit, dt1=self.self_dt, unit2=unit, dt2=parent_dt) # Calculate the new factor
         new.update_values() # Update values
         new.v = new.values # Reset the base value
         new.factor = 1.0 # Reset everything else to be 1
+        new.unit = unit
+        new.self_dt = parent_dt
         new.parent_unit = unit
         new.parent_dt = parent_dt
         return new
@@ -305,43 +333,41 @@ def update_values(self):
         return self.values
 
 
-class days(dur):
+def days(v, parent_unit=None, parent_dt=None):
     """ Shortcut to ss.dur(value, units='day') """
-    def __init__(self, v, parent_unit=None, parent_dt=None):
-        super().__init__(v=v, unit='day', parent_unit=parent_unit, parent_dt=parent_dt)
-        return
+    return dur(v=v, unit='day', parent_unit=parent_unit, parent_dt=parent_dt)
 
 
-class years(dur):
+def years(v, parent_unit=None, parent_dt=None):
     """ Shortcut to ss.dur(value, units='year') """
-    def __init__(self, v, parent_unit=None, parent_dt=None):
-        super().__init__(v=v, unit='year', parent_unit=parent_unit, parent_dt=parent_dt)
-        return
+    return dur(v=v, unit='year', parent_unit=parent_unit, parent_dt=parent_dt)
 
 
-class rate(TimePar): # TODO: should all rates just be time_prob?
+class rate(TimePar):
     """ Any number that acts like a rate; can be greater than 1 """
     def update_values(self):
         self.values = self.v/self.factor
         return self.values
 
 
-class perday(rate):
+def perday(v, parent_unit=None, parent_dt=None):
     """ Shortcut to ss.rate(value, units='day') """
-    def __init__(self, v, parent_unit=None, parent_dt=None):
-        super().__init__(v=v, unit='day', parent_unit=parent_unit, parent_dt=parent_dt)
-        return
+    return rate(v=v, unit='day', parent_unit=parent_unit, parent_dt=parent_dt)
 
 
-class peryear(rate):
+def peryear(v, parent_unit=None, parent_dt=None):
     """ Shortcut to ss.rate(value, units='year') """
-    def __init__(self, v, parent_unit=None, parent_dt=None):
-        super().__init__(v=v, unit='year', parent_unit=parent_unit, parent_dt=parent_dt)
-        return
+    return rate(v=v, unit='year', parent_unit=parent_unit, parent_dt=parent_dt)
 
 
 class time_prob(TimePar):
-    """ A probability over time (a.k.a. a cumulative hazard rate); must be >0 and <1 """
+    """
+    A probability over time (a.k.a. a cumulative hazard rate); must be >=0 and <=1.
+
+    Note: ``ss.time_prob()`` converts one cumulative hazard rate to another with a
+    different time unit. ``ss.rate_prob()`` converts an exponential rate to a cumulative
+    hazard rate.
+    """
     def update_values(self):
         v = self.v
         if self.isarray:
@@ -364,6 +390,32 @@ def update_values(self):
         return self.values
 
 
+class rate_prob(TimePar):
+    """
+    An instantaneous rate converted to a probability; must be >=0.
+
+    Note: ``ss.time_prob()`` converts one cumulative hazard rate to another with a
+    different time unit. ``ss.rate_prob()`` converts an exponential rate to a cumulative
+    hazard rate.
+    """
+    def update_values(self):
+        v = self.v
+        if self.isarray:
+            self.values = v.copy()
+            inds = v > 0.0
+            if inds.sum():
+                self.values[inds] = 1 - np.exp(-v[inds]/self.factor)
+        else:
+            if v == 0:
+                self.values = 0
+            elif v > 0:
+                self.values = 1 - np.exp(-v/self.factor)
+            else:
+                errormsg = f'Invalid value {self.value} for {self}: must be >=0'
+                raise ValueError(errormsg)
+        return self.values
+
+
 class beta(time_prob):
     """ A container for beta (i.e. the disease transmission rate) """
     pass
diff --git a/starsim/version.py b/starsim/version.py
index 12426412..e741b03b 100644
--- a/starsim/version.py
+++ b/starsim/version.py
@@ -4,6 +4,6 @@
 
 __all__ = ['__version__', '__versiondate__', '__license__']
 
-__version__ = '2.0.0'
-__versiondate__ = '2024-10-01'
+__version__ = '2.1.0'
+__versiondate__ = '2024-11-01'
 __license__ = f'Starsim {__version__} ({__versiondate__}) — © 2023-2024 by IDM'


Test patch:
diff --git a/tests/baseline.json b/tests/baseline.json
index 8c6fbc9b..9eb981f2 100644
--- a/tests/baseline.json
+++ b/tests/baseline.json
@@ -1,26 +1,26 @@
 {
   "summary": {
     "timevec": 2010.0,
-    "births_new": 65.02970297029702,
-    "births_cumulative": 3224.019801980198,
-    "births_cbr": 29.723743148428454,
-    "deaths_new": 44.97029702970297,
-    "deaths_cumulative": 2229.108910891089,
-    "deaths_cmr": 20.555949238807358,
-    "sir_n_susceptible": 2520.4356435643563,
-    "sir_n_infected": 3626.366336633663,
-    "sir_n_recovered": 4792.089108910891,
-    "sir_prevalence": 0.33888892397890347,
-    "sir_new_infections": 130.36633663366337,
-    "sir_cum_infections": 13167.0,
-    "sis_n_susceptible": 4740.722772277228,
-    "sis_n_infected": 6244.287128712871,
-    "sis_prevalence": 0.5587172284995694,
-    "sis_new_infections": 183.87128712871288,
-    "sis_cum_infections": 18571.0,
-    "sis_rel_sus": 0.5182920328461298,
-    "n_alive": 10938.89108910891,
-    "new_deaths": 46.118811881188115,
-    "cum_deaths": 4608.0
+    "births_new": 65.03960396039604,
+    "births_cumulative": 3223.4554455445545,
+    "births_cbr": 29.720841844309653,
+    "deaths_new": 45.18811881188119,
+    "deaths_cumulative": 2235.079207920792,
+    "deaths_cmr": 20.646834320732697,
+    "sir_n_susceptible": 2597.3960396039606,
+    "sir_n_infected": 3631.5940594059407,
+    "sir_n_recovered": 4712.693069306931,
+    "sir_prevalence": 0.33872902205702377,
+    "sir_new_infections": 129.88118811881188,
+    "sir_cum_infections": 13118.0,
+    "sis_n_susceptible": 4705.277227722772,
+    "sis_n_infected": 6282.574257425743,
+    "sis_prevalence": 0.5618315778700211,
+    "sis_new_infections": 183.97029702970298,
+    "sis_cum_infections": 18581.0,
+    "sis_rel_sus": 0.516166645111424,
+    "n_alive": 10941.683168316831,
+    "new_deaths": 46.16831683168317,
+    "cum_deaths": 4623.0
   }
 }
\ No newline at end of file
diff --git a/tests/benchmark.json b/tests/benchmark.json
index 992ff85b..36d4d8fd 100644
--- a/tests/benchmark.json
+++ b/tests/benchmark.json
@@ -1,12 +1,12 @@
 {
   "time": {
-    "initialize": 0.034,
-    "run": 0.913
+    "initialize": 0.036,
+    "run": 0.932
   },
   "parameters": {
     "n_agents": 10000,
     "dur": 20,
     "dt": 0.2
   },
-  "cpu_performance": 0.9555946088102025
+  "cpu_performance": 0.9253830638711784
 }
\ No newline at end of file
diff --git a/tests/test_mixingpools.py b/tests/test_mixingpools.py
new file mode 100644
index 00000000..c63fd949
--- /dev/null
+++ b/tests/test_mixingpools.py
@@ -0,0 +1,203 @@
+"""
+Test Sim API
+"""
+
+# %% Imports and settings
+import pytest
+import starsim as ss
+import sciris as sc
+import matplotlib.pyplot as plt
+import numpy as np
+import sys
+
+n_agents = 1_000
+do_plot = False
+sc.options(interactive=False) # Assume not running interactively
+
+
+def test_single_defaults(do_plot=do_plot):
+    """ Test a single MixingPool using defaults """
+    test_name = sys._getframe().f_code.co_name
+    sc.heading(f'Testing {test_name}...')
+    mp = ss.MixingPool()
+    sir = ss.SIR()
+    sim = ss.Sim(diseases=sir, networks=mp, label=test_name)
+    sim.run()
+
+    if do_plot: sim.plot()
+
+    assert(sim.results.sir['cum_infections'][-1] > sim.results.sir['cum_infections'][0]) # There were infections
+    return sim
+
+
+def test_single_uids(do_plot=do_plot):
+    """ Test a single MixingPool by UIDS, pre-initialization configuration """
+    test_name = sys._getframe().f_code.co_name
+    sc.heading(f'Testing {test_name}...')
+
+    n_agents = 10_000
+    k = n_agents // 2
+    mp = ss.MixingPool(
+        src = ss.uids(np.arange(k)),
+        dst = ss.uids(np.arange(k,n_agents))
+    )
+    sir = ss.SIR()
+    sim = ss.Sim(n_agents=n_agents, diseases=sir, networks=mp, label=test_name)
+    sim.run()
+
+    if do_plot: sim.plot()
+
+    assert sim.results.sir['cum_infections'][-1] > sim.results.sir['cum_infections'][0], 'There were no new infections in the simulation'
+    return sim
+
+
+def test_single_ncd():
+    """ Test a single MixingPool with a ncd """
+    test_name = sys._getframe().f_code.co_name
+    sc.heading(f'Testing {test_name}...')
+    mp_pars = {
+        'src': ss.AgeGroup(0, 15),
+        'dst': ss.AgeGroup(15, None),
+        'beta': ss.beta(0.15),
+        'contacts': ss.poisson(lam=5),
+        'diseases': 'ncd'
+    }
+    mp = ss.MixingPool(mp_pars)
+
+    ncd = ss.NCD()
+    sim = ss.Sim(diseases=ncd, networks=mp, label=test_name)
+    with pytest.raises(Exception):
+        sim.run()
+    return sim
+
+
+def test_single_missing_disease():
+    """ Test a single MixingPool with a missing disease """
+    test_name = sys._getframe().f_code.co_name
+    sc.heading(f'Testing {test_name}...')
+    mp_pars = {
+        'src': ss.AgeGroup(0, 15),
+        'dst': ss.AgeGroup(15, None),
+        'beta': ss.beta(0.15),
+        'contacts': ss.poisson(lam=5),
+        'diseases': 'hiv'
+    }
+    mp = ss.MixingPool(mp_pars)
+
+    sir = ss.SIR()
+    sim = ss.Sim(diseases=sir, networks=mp, label=test_name)
+    with pytest.raises(Exception):
+        sim.run()
+
+    return sim
+
+
+def test_single_age(do_plot=do_plot):
+    """ Test a single MixingPool by age """
+    # Incidence must decline because 0-15 --> 15+ transmission only
+    test_name = sys._getframe().f_code.co_name
+    sc.heading(f'Testing {test_name}...')
+    mp_pars = {
+        'src': ss.AgeGroup(0, 15),
+        'dst': ss.AgeGroup(15, None),
+        'beta': ss.beta(0.15),
+        'contacts': ss.poisson(lam=5),
+    }
+    mp = ss.MixingPool(mp_pars)
+
+    sir = ss.SIR()
+    sim = ss.Sim(diseases=sir, networks=mp, label=test_name)
+    sim.run()
+
+    if do_plot: sim.plot()
+
+    assert(sim.results.sir['cum_infections'][-1] > sim.results.sir['cum_infections'][0]) # There were infections
+    return sim
+
+
+def test_single_sex(do_plot=do_plot):
+    """ Test a single MixingPool by sex """
+    # Incidence must decline because M --> F transmission only
+    test_name = sys._getframe().f_code.co_name
+    sc.heading(f'Testing {test_name}...')
+    mp_pars = {
+        'src': lambda sim: sim.people.female, # female to male (only) transmission
+        'dst': lambda sim: sim.people.male,
+        'beta': ss.beta(0.2),
+        'contacts': ss.poisson(lam=4),
+    }
+    mp = ss.MixingPool(mp_pars)
+
+    sir = ss.SIR(init_prev=ss.bernoulli(p=lambda self, sim, uids: 0.05*sim.people.female)) # Seed 5% of the female population
+    sim = ss.Sim(diseases=sir, networks=mp, label=test_name)
+    sim.run()
+
+    if do_plot: sim.plot()
+
+    assert(sim.results.sir['cum_infections'][-1] > sim.results.sir['cum_infections'][0]) # There were infections
+    assert(sim.people.male[sim.diseases.sir.ti_infected>0].all()) # All new infections should be in men
+    return sim
+
+
+def test_multi_defaults(do_plot=do_plot):
+    """ Test MixingPools using defaults """
+    test_name = sys._getframe().f_code.co_name
+    sc.heading(f'Testing {test_name}...')
+    mps = ss.MixingPools(src={'all':None}, dst={'all':None}, contacts=[[1]])
+    sir = ss.SIR()
+    sim = ss.Sim(diseases=sir, networks=mps, label=test_name)
+    sim.run()
+
+    if do_plot: sim.plot()
+
+    assert(sim.results.sir['cum_infections'][-1] > sim.results.sir['cum_infections'][0]) # There were infections
+    return sim
+
+
+def test_multi(do_plot=do_plot):
+    """ Test MixingPools """
+    test_name = sys._getframe().f_code.co_name
+    sc.heading(f'Testing {test_name}...')
+
+    groups = {
+        'Female': lambda sim: sim.people.female,
+        'Male': lambda sim: sim.people.male,
+    }
+
+    mps_pars = dict(
+        contacts = np.array([[1.4, 0.5], [1.2, 0.7]]),
+        beta = ss.beta(0.2),
+        src = groups,
+        dst = groups,
+    )
+    mps = ss.MixingPools(mps_pars)
+
+    sir = ss.SIR()
+    sim = ss.Sim(diseases=sir, networks=mps, label=test_name)
+    sim.run()
+
+    if do_plot: sim.plot()
+
+    assert(sim.results.sir['cum_infections'][-1] > sim.results.sir['cum_infections'][0]) # There were infections
+    return sim
+
+
+if __name__ == '__main__':
+    do_plot = True
+    sc.options(interactive=do_plot)
+    T = sc.timer()
+
+    sim0 = test_single_defaults(do_plot)
+    sim1 = test_single_uids(do_plot)
+    sim2 = test_single_ncd()
+    sim3 = test_single_missing_disease()
+    sim4 = test_single_age(do_plot)
+    sim5 = test_single_sex(do_plot)
+
+    sim6 = test_multi_defaults(do_plot)
+    sim7 = test_multi(do_plot)
+
+    T.toc()
+
+    if do_plot:
+        plt.show()