update readme and other files

covasim/README.md

@@ -6,7 +6,8 @@ This file describes the expected behavior of each parameter in the model. Note:
 * `pop_scale`    = Multiplicative scale for results. Test: run 2 sims, set to 10, `sim.results['cum_exposed']` in 2nd sim should be 10x higher than first.
 * `pop_size`     = Nmber of people in the simulation. Test: `len(sim.people)` should equal this number.
 * `pop_infected` = Initial number of people infected. Test: if 0, there should be no infections; if equals `n`, should be no _new_ infections.
-* `pop_type`   = Whether or not to use the `synthpops` library for contact matrices. Consult that library's documentation for tests.
+* `pop_type`   = The type of population structure to use with the model; default `'random'`, using Seattle demographics; other options are `'hybrid'` (structured households and random school, work, and community contacts), `'clustered'` (all clustered contacts), and `'synthpops'` (requires the `synthpops` library, but uses data-based population structure)
+* `location`   = Which country or state to load demographic data from (optional)
 
 ## Simulation parameters
 * `start_day`    = The calendar date of the start day of the simulation.

covasim/data/loaders.py

@@ -26,7 +26,7 @@ def get_age_distribution(location=None):
     countries = [entry["country"] for entry in json] # Pull out available countries
 
     # Set parameters
-    max_age = 120
+    max_age = 99
     if location is None:
         location = countries
     else:

covasim/parameters.py

@@ -27,6 +27,7 @@ def make_pars(set_prognoses=False, prog_by_age=True, use_layers=False, **kwargs)
     pars['pop_size']     = 20e3 # Number ultimately susceptible to CoV
     pars['pop_infected'] = 10 # Number of initial infections
     pars['pop_type']     = 'random' # What type of population data to use -- random (fastest), synthpops (best), realistic (compromise), or clustered (not recommended)
+    pars['location']     = None # What location to load data from -- default Seattle
 
     # Simulation parameters
     pars['start_day']  = '2020-03-01' # Start day of the simulation

covasim/population.py

@@ -160,17 +160,20 @@ def make_people(sim, verbose=None, die=True, reset=False):
     return
 
 
-def make_randpop(sim, location=None, age_data=None, sex_ratio=0.5, microstructure=False):
+def make_randpop(sim, age_data=None, sex_ratio=0.5, microstructure=False):
     ''' Make a random population, without contacts '''
 
     pop_size = int(sim['pop_size']) # Number of people
 
     # Load age data based on 2018 Seattle demographics by default
-    if age_data is None:
-        if location is not None:
+    location = sim['location']
+    if location is not None:
+        try:
             age_data = cvdata.loaders.get_age_distribution(location)
-        else:
-            age_data = cvd.default_age_data
+        except ValueError as E:
+            print(f'Could not load data for requested location "{location}" ({str(E)}), using default')
+    if age_data is None:
+        age_data = cvd.default_age_data
 
     # Handle sexes and ages
     uids = np.arange(pop_size, dtype=int)

tests/devtests/dev_test_age_distributions.py

@@ -0,0 +1,27 @@
+import covasim as cv
+import pylab as pl
+
+loc1 = 'Sweden'
+loc2 = 'Somalia'
+
+sim1 = cv.Sim(location=loc1)
+sim2 = cv.Sim(location=loc2)
+
+sim1.initialize()
+sim2.initialize()
+
+ages1 = sim1.people.extract('age')
+ages2 = sim2.people.extract('age')
+
+n = 100
+fig = pl.figure()
+
+pl.subplot(2,1,1)
+pl.hist(ages1, n)
+pl.title(loc1)
+
+pl.subplot(2,1,2)
+pl.hist(ages2, n)
+pl.title(loc2)
+pl.xlabel('Age')
+pl.ylabel('Count')

tests/test_age_structure.py

@@ -1,69 +1,35 @@
 '''
-Simple example usage for the Covid-19 agent-based model
+Test different age structures
 '''
 
 #%% Imports and settings
-import pylab as pl
 import sciris as sc
-import covasim as cova
-
-doplot = 1
-figsize = (20,16)
+import covasim as cv
+import pytest
 
 
 #%% Define the tests
 
-def test_age_structure(doplot=False): # If being run via pytest, turn off
-    pass # Not implemented yet
-
-    # # Create and run the simulation without age structure
-    # sims = sc.objdict()
-
-    # sims.without = cova.Sim()
-    # sims.without['usepopdata'] = 0
-    # sims.without.run(verbose=1)
+def test_age_structure(): # If being run via pytest, turn off
 
-    # # ...and with
-    # sims.withdata = cova.Sim()
-    # sims.withdata['usepopdata'] = 1
-    # sims.withdata.run(verbose=1)
+    available = 'Lithuania'
+    not_available = 'Ruritania'
 
-    # # Calculate ages
-    # ages = sc.objdict()
-    # for key in sims.keys():
-    #     ages[key] = []
-    #     for person in sims[key].people.values():
-    #         ages[key].append(person.age)
+    age_data = cv.data.loaders.get_age_distribution(available)
 
-    # # Optionally plot
-    # if doplot:
-    #     nbins = 100
+    with pytest.raises(ValueError):
+        cv.data.loaders.get_age_distribution(not_available)
 
-    #     # Plot epi results
-    #     for sim in sims.values():
-    #         sim.plot()
+    return age_data
 
-    #     # Plot ages
-    #     pl.figure(figsize=figsize)
-    #     for a,key,age in ages.enumitems():
-    #         if key == 'without':
-    #             title = f'Normally-distributed age distribution'
-    #         elif key == 'withdata':
-    #             title = f'Age distribution based on Seattle census data'
-    #         pl.subplot(2,1,a+1)
-    #         pl.hist(ages[a], nbins)
-    #         pl.title(title)
-    #         pl.xlabel('Age')
-    #         pl.ylabel('Number of people')
-
-    # return sims, ages
 
 
 #%% Run as a script
 if __name__ == '__main__':
     sc.tic()
-    print('Not implemented yet')
-    # sims, ages = test_age_structure(doplot=doplot)
+
+    age_data = test_age_structure()
+
     sc.toc()