Repository: starsimhub/starsim (language: python)

Issue:
Add ebola
Placeholder for Ebola, needs more testing.

Also fixes https://github.com/amath-idm/starsim/issues/232


Discussion / hints:

Fix patch:
diff --git a/starsim/disease.py b/starsim/disease.py
index 40176364..3cb42a48 100644
--- a/starsim/disease.py
+++ b/starsim/disease.py
@@ -415,6 +415,11 @@ def choose_source(df):
 
     def make_new_cases(self, sim):
         """ Add new cases of module, through transmission, incidence, etc. """
+        if not sim.networks:
+            warnmsg = f'Disease {self.name} does not transmit without a network.'
+            if sim.ti == 0: ss.warn(warnmsg, die=False)
+            return
+
         if not ss.options.multirng:
             # Determine new cases for singlerng
             new_cases, sources = self._make_new_cases_singlerng(sim)
diff --git a/starsim/diseases/__init__.py b/starsim/diseases/__init__.py
index 0ca7831a..8c0bf910 100644
--- a/starsim/diseases/__init__.py
+++ b/starsim/diseases/__init__.py
@@ -1,6 +1,10 @@
+from .sir           import *
+from .ncd           import *
+
+# Specific diseases listed alphabetically
 from .cholera       import *
+from .ebola         import *
 from .gonorrhea     import *
 from .hiv           import *
-from .ncd           import *
-from .sir           import *
+from .measles       import *
 from .syphilis      import *
diff --git a/starsim/diseases/ebola.py b/starsim/diseases/ebola.py
new file mode 100644
index 00000000..ff970439
--- /dev/null
+++ b/starsim/diseases/ebola.py
@@ -0,0 +1,153 @@
+"""
+Define Ebola model.
+Adapted from https://github.com/optimamodel/gavi-outbreaks/blob/main/stisim/gavi/ebola.py
+Original version by@domdelport and @RomeshA
+"""
+
+import numpy as np
+import starsim as ss
+from starsim.diseases.sir import SIR
+
+__all__ = ['Ebola']
+
+
+class Ebola(SIR):
+
+    def __init__(self, pars=None, par_dists=None, *args, **kwargs):
+        """ Initialize with parameters """
+
+        pars = ss.omerge({
+            # Natural history parameters, all specified in days
+            'dur_exp2symp':     12.7,   # Add source
+            'dur_symp2sev':        6,   # Add source
+            'dur_sev2dead':      1.5,   # Add source
+            'dur_dead2buried':     2,   # Add source
+            'dur_symp2rec':       10,   # Add source
+            'dur_sev2rec':      10.4,   # Add source
+            'p_sev':             0.7,   # Add source
+            'p_death':          0.55,   # Add source
+            'p_safe_bury':      0.25,   # Probability of a safe burial - should be linked to diagnoses
+
+            # Initial conditions and beta
+            'init_prev': 0.005,
+            'beta': None,
+            'sev_factor': 2.2,
+            'unburied_factor': 2.1,
+        }, pars)
+
+        par_dists = ss.omerge({
+            'dur_exp2symp': ss.lognorm,
+            'dur_symp2sev': ss.lognorm,
+            'dur_sev2dead': ss.lognorm,
+            'dur_dead2buried': ss.lognorm,
+            'dur_symp2rec': ss.lognorm,
+            'dur_sev2rec': ss.lognorm,
+            'p_sev': ss.bernoulli,
+            'p_death': ss.bernoulli,
+            'p_safe_bury': ss.bernoulli,
+            'init_prev': ss.bernoulli,
+        }, par_dists)
+
+        super().__init__(pars=pars, par_dists=par_dists, *args, **kwargs)
+
+        # Boolean states
+        # SIR are added automatically, here we add E
+        self.exposed = ss.State('exposed', bool, False)
+        self.severe = ss.State('severe', bool, False)
+        self.recovered = ss.State('recovered', bool, False)
+        self.buried = ss.State('buried', bool, False)
+
+        # Timepoint states
+        self.ti_exposed = ss.State('ti_exposed', float, np.nan)
+        self.ti_severe = ss.State('ti_severe', float, np.nan)
+        self.ti_recovered = ss.State('ti_recovered', float, np.nan)
+        self.ti_dead = ss.State('ti_dead', float, np.nan)
+        self.ti_buried = ss.State('ti_buried', float, np.nan)
+
+        return
+
+    @property
+    def infectious(self):
+        return self.infected | self.exposed
+
+    def update_pre(self, sim):
+
+        # Progress exposed -> infected
+        infected = ss.true(self.exposed & (self.ti_infected <= sim.ti))
+        self.exposed[infected] = False
+        self.infected[infected] = True
+
+        # Progress infectious -> severe
+        severe = ss.true(self.infected & (self.ti_severe <= sim.ti))
+        self.severe[severe] = True
+
+        # Progress infected -> recovered
+        recovered = ss.true(self.infected & (self.ti_recovered <= sim.ti))
+        self.infected[recovered] = False
+        self.recovered[recovered] = True
+
+        # Progress severe -> recovered
+        recovered_sev = ss.true(self.severe & (self.ti_recovered <= sim.ti))
+        self.severe[recovered_sev] = False
+        self.recovered[recovered_sev] = True
+
+        # Trigger deaths
+        deaths = ss.true(self.ti_dead <= sim.year)
+        if len(deaths):
+            sim.people.request_death(deaths)
+
+        # Progress dead -> buried
+        buried = ss.true(self.ti_buried <= sim.ti)
+        self.buried[buried] = True
+
+
+    def set_prognoses(self, sim, uids, from_uids=None):
+        """ Set prognoses for those who get infected """
+        super().set_prognoses(sim, uids, from_uids)
+
+        self.susceptible[uids] = False
+        self.exposed[uids] = True
+        self.ti_exposed[uids] = sim.ti
+
+        p = self.pars
+
+        # Determine when exposed become infected
+        self.ti_infected[uids] = sim.ti + p.dur_exp2symp.rvs(uids)
+
+        # Determine who progresses to sever and when
+        sev_uids = p.p_sev.filter(uids)
+        self.ti_severe[sev_uids] = self.ti_infected[sev_uids] + p.dur_symp2sev.rvs(sev_uids)
+
+        # Determine who dies and who recovers and when
+        dead_uids = p.p_death.filter(sev_uids)
+        self.ti_dead[dead_uids] = self.ti_severe[dead_uids] + p.dur_sev2dead.rvs(dead_uids)
+        rec_sev_uids = np.setdiff1d(sev_uids, dead_uids)
+        self.ti_recovered[rec_sev_uids] = self.ti_severe[rec_sev_uids] + p.dur_sev2rec.rvs(rec_sev_uids)
+        rec_symp_uids = np.setdiff1d(uids, sev_uids)
+        self.ti_recovered[rec_symp_uids] = self.ti_infected[rec_symp_uids] + p.dur_symp2rec.rvs(rec_symp_uids)
+
+        # Determine time of burial - either immediate (safe burials) or after a delay (unsafe)
+        safe_buried = p.p_safe_bury.filter(dead_uids)
+        unsafe_buried = np.setdiff1d(dead_uids, safe_buried)
+        self.ti_buried[safe_buried] = self.ti_dead[safe_buried]
+        self.ti_buried[unsafe_buried] = self.ti_dead[unsafe_buried] + p.dur_dead2buried.rvs(unsafe_buried)
+
+        return
+
+    def update_death(self, sim, uids):
+        # Reset infected/recovered flags for dead agents
+        self.susceptible[uids] = False
+        self.exposed[uids] = False
+        self.infected[uids] = False
+        self.severe[uids] = False
+        self.recovered[uids] = False
+        return
+
+    def make_new_cases(self, sim):
+        self.rel_trans[self.infectious] = 1
+        self.rel_trans[self.severe] = self.pars['sev_factor']
+
+        # Unburied UIDs
+        unburied_uids = ss.true((self.ti_dead <= sim.ti) & (self.ti_buried > sim.ti))
+        self.rel_trans[unburied_uids] = self.pars['unburied_factor']
+        super().make_new_cases(sim)
diff --git a/starsim/diseases/measles.py b/starsim/diseases/measles.py
new file mode 100644
index 00000000..d7ba300c
--- /dev/null
+++ b/starsim/diseases/measles.py
@@ -0,0 +1,97 @@
+"""
+Define measles model.
+Adapted from https://github.com/optimamodel/gavi-outbreaks/blob/main/stisim/gavi/measles.py
+Original version by @alina-muellenmeister, @domdelport, and @RomeshA
+"""
+
+import numpy as np
+import starsim as ss
+from starsim.diseases.sir import SIR
+
+__all__ = ['Measles']
+
+
+class Measles(SIR):
+
+    def __init__(self, pars=None, par_dists=None, *args, **kwargs):
+        """ Initialize with parameters """
+
+        pars = ss.omerge({
+            # Natural history parameters, all specified in days
+            'dur_exp': 8,       # (days) - source: US CDC
+            'dur_inf': 11,      # (days) - source: US CDC
+            'p_death': 0.005,   # Probability of death
+
+            # Initial conditions and beta
+            'init_prev': 0.005,
+            'beta': None,
+        }, pars)
+
+        par_dists = ss.omerge({
+            'dur_exp': ss.norm,
+            'dur_inf': ss.norm,
+            'init_prev': ss.bernoulli,
+            'p_death': ss.bernoulli,
+        }, par_dists)
+
+        super().__init__(pars=pars, par_dists=par_dists, *args, **kwargs)
+
+        # Boolean states
+        # SIR are added automatically, here we add E
+        self.exposed = ss.State('exposed', bool, False)
+
+        # Timepoint states
+        self.ti_exposed = ss.State('ti_exposed', float, np.nan)
+
+        return
+
+    @property
+    def infectious(self):
+        return self.infected | self.exposed
+
+    def update_pre(self, sim):
+
+        # Progress exposed -> infected
+        infected = ss.true(self.exposed & (self.ti_infected <= sim.ti))
+        self.exposed[infected] = False
+        self.infected[infected] = True
+
+        # Progress infected -> recovered
+        recovered = ss.true(self.infected & (self.ti_recovered <= sim.ti))
+        self.infected[recovered] = False
+        self.recovered[recovered] = True
+
+        # Trigger deaths
+        deaths = ss.true(self.ti_dead <= sim.year)
+        if len(deaths):
+            sim.people.request_death(deaths)
+
+    def set_prognoses(self, sim, uids, from_uids=None):
+        """ Set prognoses for those who get infected """
+        super().set_prognoses(sim, uids, from_uids)
+
+        self.susceptible[uids] = False
+        self.exposed[uids] = True
+        self.ti_exposed[uids] = sim.ti
+
+        p = self.pars
+
+        # Determine when exposed become infected
+        self.ti_infected[uids] = sim.ti + p.dur_exp.rvs(uids)
+
+        # Determine who dies and who recovers and when
+        dead_uids = p.p_death.filter(uids)
+        self.ti_dead[dead_uids] = self.ti_infected[dead_uids] + p.dur_inf.rvs(dead_uids)
+        rec_uids = np.setdiff1d(uids, dead_uids)
+        self.ti_recovered[rec_uids] = self.ti_infected[rec_uids] + p.dur_inf.rvs(rec_uids)
+
+        return
+
+    def update_death(self, sim, uids):
+        # Reset infected/recovered flags for dead agents
+        self.susceptible[uids] = False
+        self.exposed[uids] = False
+        self.infected[uids] = False
+        self.recovered[uids] = False
+        return
+
diff --git a/starsim/diseases/sir.py b/starsim/diseases/sir.py
index 2310b2cb..8d24a664 100644
--- a/starsim/diseases/sir.py
+++ b/starsim/diseases/sir.py
@@ -7,6 +7,7 @@
 
 __all__ = ['SIR']
 
+
 class SIR(ss.Infection):
     """
     Example SIR model
@@ -17,7 +18,6 @@ class SIR(ss.Infection):
     """
 
     def __init__(self, pars=None, par_dists=None, *args, **kwargs):
-
         pars = ss.omerge({
             'dur_inf': 1,
             'init_prev': 0.1,
diff --git a/starsim/sim.py b/starsim/sim.py
index a4a668f5..896fb268 100644
--- a/starsim/sim.py
+++ b/starsim/sim.py
@@ -253,20 +253,31 @@ def convert_plugins(self, plugin_class, plugin_name=None):
 
         if plugin_name is None: plugin_name = plugin_class.__name__.lower()
 
+        # Get lower-case names of all subclasses
+        known_plugins = [n.__name__.lower() for n in ss.all_subclasses(plugin_class)]
+
         # Figure out if it's in the sim pars or provided directly
         attr_plugins = getattr(self, plugin_name)  # Get any plugins that have been provided directly
         if attr_plugins is None or len(attr_plugins) == 0:  # None have been provided directly
-            if self.pars.get(plugin_name) and len(
-                    self.pars[plugin_name]):  # See if they've been provided in the pars dict
-                plugins = ss.ndict(self.pars[plugin_name])
+
+            # See if they've been provided in the pars dict
+            if self.pars.get(plugin_name):
+
+                par_plug = self.pars[plugin_name]
+
+                # String: convert to ndict
+                if isinstance(par_plug, str):
+                    plugins = ss.ndict(dict(name=par_plug))
+
+                # List or dict: convert to ndict
+                elif sc.isiterable(par_plug) and len(par_plug):
+                    plugins = ss.ndict(par_plug)
+
             else:  # Not provided directly or in pars
                 plugins = {}
         else:
             plugins = attr_plugins
 
-        # Convert
-        known_plugins = [n.__name__.lower() for n in ss.all_subclasses(plugin_class)]
-
         processed_plugins = sc.autolist()
         for plugin in plugins.values():
 


Test patch:
diff --git a/tests/test_diseases.py b/tests/test_diseases.py
index 0c293d6a..ff72488b 100644
--- a/tests/test_diseases.py
+++ b/tests/test_diseases.py
@@ -6,10 +6,10 @@
 import starsim as ss
 import matplotlib.pyplot as plt
 import scipy.stats as sps
-import numpy as np
+import sciris as sc
 
 test_run = True
-n_agents = [2_000, 10_000][test_run]
+n_agents = [10_000, 2_000][test_run]
 
 
 def test_sir():
@@ -76,16 +76,24 @@ def test_ncd():
     return sim
 
 
-def test_cholera():
+def test_gavi():
     ss.options(multirng=False)
-    sim = ss.Sim(n_agents=n_agents, networks=ss.random(), diseases=ss.Cholera())
-    sim.run()
-    return sim
+    sims = sc.autolist()
+    for disease in ['cholera', 'measles', 'ebola']:
+        pars = dict(
+            diseases=disease,
+            n_agents=n_agents,
+            networks='random',
+        )
+        sim = ss.Sim(pars)
+        sim.run()
+        sims += sim
+    return sims
 
 
 if __name__ == '__main__':
     ss.options(multirng=False)
     sim1 = test_sir()
     sim2 = test_ncd()
-    sim3 = test_cholera()
+    sims = test_gavi()
     plt.show()