Worked out how to do additional seeds in a demographic way. Also now working on the multi-demographic population object and output

Author: chryswoods

.env

@@ -0,0 +1,2 @@
+PYTHONPATH=build/lib.macosx-10.9-x86_64-3.7:$PYTHONPATH
+

src/metawards/_demographics.py

@@ -85,13 +85,46 @@ def add(self, demographic: Demographic):
         """Add a demographic to the set to be modelled"""
         if demographic.name in self._names:
             raise ValueError(
-                    f"There is already a demographic called "
-                    f"{demographic.name} in this set. Please rename "
-                    f"and try again.")
+                f"There is already a demographic called "
+                f"{demographic.name} in this set. Please rename "
+                f"and try again.")
 
         self.demographics.append(demographic)
         self._names[demographic.name] = len(self.demographics) - 1
 
+    def get_name(self, item):
+        """Return the name of the demographic at 'item'"""
+        return self.demographics[self.get_index(item)].name
+
+    def get_index(self, item):
+        """Return the index of the passed item"""
+        try:
+            item = int(item)
+        except Exception:
+            pass
+
+        if isinstance(item, str):
+            try:
+                return self._names[item]
+            except Exception:
+                pass
+
+        elif isinstance(item, int):
+            try:
+                if self.demographics[item] is not None:
+                    return item
+            except Exception:
+                pass
+
+        elif isinstance(item, Demographic):
+            for i, d in enumerate(self.demographics):
+                if item == d:
+                    return i
+
+        # haven't found the item
+        raise KeyError(f"There is no demographic is this set that "
+                       f"matches {item}.")
+
     @staticmethod
     def load(name: str = None,
              repository: str = None,
@@ -177,7 +210,7 @@ def load(name: str = None,
         play_ratios = data.get("play_ratios", [])
 
         if (len(demographics) != len(work_ratios) or
-           len(demographics) != len(play_ratios)):
+                len(demographics) != len(play_ratios)):
             raise ValueError(
                 f"The number of work_ratios ({len(work_ratios)}) must "
                 f"equal to number of play_ratios "

src/metawards/_population.py

@@ -75,6 +75,32 @@ def __str__(self):
         else:
             return s
 
+    def summary(self, demographics=None):
+        """Return a short summary string that is suitable to be printed
+           out during a model run
+
+           Returns
+           -------
+           summary: str
+             The short summary string
+        """
+        summary = f"S: {self.susceptibles}  E: {self.latent}  " \
+                  f"I: {self.total}  R: {self.recovereds}  " \
+                  f"IW: {self.n_inf_wards}  T_POP: {self.population}"
+
+        if self.subpops is None or len(self.subpops) == 0:
+            return summary
+
+        subs = []
+        for i, subpop in enumerate(self.subpops):
+            if demographics is not None:
+                name = demographics.get_name(i)
+                subs.append(f"name: {name}  {subpop.summary()}")
+            else:
+                subs.append(f"index: {i}  {subpop.summary()}")
+
+        return f"{summary}\n  " + "\n  ".join(subs)
+
 
 @_dataclass
 class Populations:

src/metawards/extractors/_output_core.pyx

@@ -713,13 +713,6 @@ def output_core_serial(network: Network, population: Population,
             f"{S} vs {susceptibles}, {E} vs {latent}, {I} vs {total}, "
             f"{R} vs {recovereds}")
 
-    print(f"S: {susceptibles}    ", end="")
-    print(f"E: {latent}    ", end="")
-    print(f"I: {total}    ", end="")
-    print(f"R: {recovereds}    ", end="")
-    print(f"IW: {n_inf_wards[0]}   ", end="")
-    print(f"TOTAL POPULATION {susceptibles+latent+total+recovereds}")
-
     if population is not None:
         population.susceptibles = susceptibles
         population.total = total
@@ -772,6 +765,7 @@ def output_core(network: _Union[Network, Networks],
     if isinstance(network, Network):
         output_func(network=network, population=population,
                     workspace=workspace, infections=infections)
+        print(population.summary())
 
     elif isinstance(network, Networks):
         if population.subpops is None:
@@ -784,6 +778,8 @@ def output_core(network: _Union[Network, Networks],
                         infections=infections.subinfs[i],
                         **kwargs)
 
+        print(population.summary(demographics=network.demographics))
+
         # aggregate the infection information from across
         # the different demographics
         network.aggregate(infections=infections)

src/metawards/iterators/_advance_additional.py

@@ -1,5 +1,8 @@
 
+from typing import Union as _Union
+
 from .._network import Network
+from .._networks import Networks
 from .._population import Population
 from ..utils._profiler import Profiler
 from .._infections import Infections
@@ -25,7 +28,13 @@ def _load_additional_seeds(filename: str):
 
             # yes, this is really the order of the seeds - "t num loc"
             # is in the file as "t loc num"
-            seeds.append((int(words[0]), int(words[2]), int(words[1])))
+            if len(words) == 4:
+                seeds.append((int(words[0]), int(words[2]),
+                              int(words[1]), words[3]))
+            else:
+                seeds.append((int(words[0]), int(words[2]),
+                              int(words[1]), None))
+
             print(seeds[-1])
             line = FILE.readline()
 
@@ -40,7 +49,7 @@ def _load_additional_seeds(filename: str):
 _additional_seeds = None
 
 
-def setup_additional_seeds(network: Network,
+def setup_additional_seeds(network: _Union[Network, Networks],
                            profiler: Profiler,
                            **kwargs):
     """Setup function that reads in the additional seeds held
@@ -51,7 +60,7 @@ def setup_additional_seeds(network: Network,
 
        Parameters
        ----------
-       network: Network
+       network: Network or Networks
          The network to be seeded
        profiler: Profiler
          Profiler used to profile this function
@@ -70,7 +79,7 @@ def setup_additional_seeds(network: Network,
     p = p.stop()
 
 
-def advance_additional_serial(network: Network,
+def advance_additional_serial(network: _Union[Network, Networks],
                               population: Population,
                               infections: Infections,
                               profiler: Profiler,
@@ -80,7 +89,7 @@ def advance_additional_serial(network: Network,
 
        Parameters
        ----------
-       network: Network
+       network: Network or Networks
          The network being modelled
        population: Population
          The population experiencing the outbreak - also contains the day
@@ -93,24 +102,44 @@ def advance_additional_serial(network: Network,
          Arguments that aren't used by this advancer
     """
 
-    wards = network.nodes
-
-    play_infections = infections.play
-    infections = infections.work
-
     # The 'setup_additional_seeds' function should have loaded
     # all additional seeds into this global '_additional_seeds' variable
     global _additional_seeds
 
     p = profiler.start("additional_seeds")
     for seed in _additional_seeds:
         if seed[0] == population.day:
-            if wards.play_suscept[seed[1]] < seed[2]:
+            ward = seed[1]
+            num = seed[2]
+
+            if isinstance(network, Networks):
+                demographic = seed[3]
+
+                if demographic is None:
+                    # have to choose a demographic randomly - choose based
+                    # on the relative populations of the demographics
+                    demographic = 0
+                else:
+                    demographic = network.demographics.get_index(demographic)
+
+                wards = network.subnets[demographic]
+                play_infections = infections.subinfs[demographic].play
+            else:
+                demographic = None
+                wards = network.wards
+                play_infections = infections.play
+
+            if wards.play_suscept[ward] < num:
                 print(f"Not enough susceptibles in ward for seeding")
             else:
-                wards.play_suscept[seed[1]] -= seed[2]
-                print(f"seeding play_infections[0][{seed[1]}] += {seed[2]}")
-                play_infections[0][seed[1]] += seed[2]
+                wards.play_suscept[ward] -= num
+                if demographic is not None:
+                    print(f"seeding demographic {demographic} "
+                          f"play_infections[0][{ward}] += {num}")
+                else:
+                    print(f"seeding play_infections[0][{ward}] += {num}")
+
+                play_infections[0][ward] += num
     p.stop()
 
 

src/metawards/iterators/_setup_imports.py

@@ -29,7 +29,11 @@ def setup_seed_specified_ward(network: _Union[Network, Networks],
     """
     if isinstance(network, Networks):
         # only seed the overall ward
-        network = network.overall
+        raise NotImplementedError(
+            "Seeding a specific ward for a multi-demographic Network is "
+            "not yet supported. If you need this please raise an issue "
+            "at https://github.com/metawards/MetaWards/issues and we will "
+            "work as quickly as we can to implement it.")
 
     wards = network.nodes
     links = network.links
@@ -66,12 +70,21 @@ def setup_seed_all_wards(network: _Union[Network, Networks],
        in params.daily_imports
     """
     if isinstance(network, Networks):
-        # only seed the overall network
-        network = network.overall
+        if network.overall.params.daily_imports == 0:
+            return
+
+        raise NotImplementedError(
+            "Daily seeding of a multi-demographic network is not yet "
+            "supported. If you need this, please raise an issue on "
+            "https://github.com/metawards/MetaWards/issues and we will "
+            "work as quickly as we can to implement it.")
 
     wards = network.nodes
     params = network.params
 
+    if params.daily_imports == 0:
+        return
+
     play_infections = infections.play
     infections = infections.work