Skip to content

Implementing a disease module

Jump to bottom
Stefan Piatek edited this page Aug 5, 2019 · 55 revisions

Here we will describe how to implement a fictional disease module on the TLO framework. We assume that you've manage to successfully clone the TLOmodel from the Github repository and set up the Python environment as described here.

Ensure you have the latest version of the code

There are multiple ways to interact with the git code repository. You can use the command-line tool git or the built-in tool available in PyCharm (free academic license for professional version). There are also several GUI tools (e.g. SourceTree) available.

Pull the latest version of the master branch:

# go to the where the code respository is located on your machine
cd TLOmodel
# make sure you're on the master branch
git checkout master
# update the master branch
git pull

Create a new branch to implement our disease module

Replace YOUR_NAME with your own name:

git checkout -b feature/YOUR_NAME-mockitis

For example, I will use:

git checkout -b feature/tamuri-mockitis

Copy the skeleton template for modules

cp src/tlo/methods/skeleton.py src/tlo/methods/mockitis.py

Edit the src/tlo/methods/mockitis.py file and edit the following lines:

class Skeleton(Module): → class Mockitis(Module):

and

class SkeletonEvent(Module): → class MockitisEvent(Module):

Create a file for testing the new module

You can use the command line, or create the file in your editor:

touch tests/test_mockitis.py

and add the following:

import os
from pathlib import Path

import pytest

from tlo import Simulation, Date
from tlo.methods import demography

resourcefilepath = Path(os.path.dirname(__file__)) / '../resources'

start_date = Date(2010, 1, 1)
end_date = Date(2015, 1, 1)
popsize = 1000


@pytest.fixture(scope='module')
def simulation():
    sim = Simulation(start_date=start_date)
    sim.register(demography.Demography(resourcefilepath=resourcefilepath))
    return sim


def test_mockitis_simulation(simulation):
    simulation.make_initial_population(n=popsize)
    simulation.simulate(end_date=end_date)


if __name__ == '__main__':
    simulation = simulation()
    test_mockitis_simulation(simulation)

This provides the basic outline of the test. This test can be run both using pytest or as a standalone script, suitable for development.

Run the mockitis test (demo)

The file should run without error. All it will do is load the [current] demography module.

Start modifying the new disease module

Let's set the parameters of the module in the file mockitis.py:

PARAMETERS = {
    'p_infection': Parameter(Types.REAL, 'Probability that an uninfected individual becomes infected'),
    'p_cure': Parameter(Types.REAL, 'Probability that an infected individual is cured'),
    'initial_prevalence': Parameter(Types.REAL, 'Prevalence of the disease in the population'),
}

Set the properties of the individual this module provides. (I've added a prefix to avoid name clashes with other modules)

PROPERTIES = {
    'mi_is_infected': Property(Types.BOOL, 'Current status of mockitis'),
    'mi_status': Property(Types.CATEGORICAL,
                          'Historical status: N=never; T1=type 1; T2=type 2; P=previously',
                          categories=['N', 'T1', 'T2', 'P', 'NT1', 'NT2']),
    'mi_date_infected': Property(Types.DATE, 'Date of latest infection'),
    'mi_date_death': Property(Types.DATE, 'Date of death of infected individual'),
    'mi_date_cure': Property(Types.DATE, 'Date an infected individual was cured'),
}

read_parameters(self, data_folder)

Here is where you would set any parameter values and read data files required for running the module and initialising the population.

For now, let's set the parameter values directly:

def read_parameters(self, data_folder):
    """some commments....
    """
    self.parameters['p_infection'] = 0.01
    self.parameters['p_cure'] = 0.01
    self.parameters['initial_prevalence'] = 0.0

Update the test with the new module

Update test_mockitis.py to load and work with the new module. Add the import statement for loading mockitis:

from tlo.methods import mockitis

Then update the simulation() method.

@pytest.fixture
def simulation():
    sim = Simulation(start_date=start_date)
    core_module = demography.Demography(workbook_path=path)
    sim.register(core_module)

    # Instantiate and add the mockitis module to the simulation
    mockitis_module = mockitis.Mockitis()
    sim.register(mockitis_module)

    return sim

Run the test. It should fail with a 'NotImplementedError'!

initialise_population(self, population)

The simulator calls this method to initialise the properties of the individual in the population for this disease module.

This method uses the Pandas and Numpy libraries, so you will need to add import statements at the top of the file:

import numpy as np
import pandas as pd

Then we implement the initialise_population method:

def initialise_population(self, population):
    """Some comments...
    """
    df = population.props   # a shortcut to the dataframe storing data for individiuals

    df['mi_is_infected'] = False   # default: no individuals infected
    df['mi_status'].values[:] = 'N'   # default: never infected
    df['mi_date_infected'] = pd.NaT   # default: not a time
    df['mi_date_death'] = pd.NaT   # default: not a time
    df['mi_date_cure'] = pd.NaT   # default: not a time

    # randomly selected some individuals as infected
    initial_infected = self.parameters['initial_prevalence']
    initial_uninfected = 1 - initial_infected
    df['mi_is_infected'] = np.random.choice([True, False], size=len(df), p=[initial_infected, initial_uninfected])

    # get all the infected individuals
    infected_count = df.mi_is_infected.sum()

    # date of infection of infected individuals
    infected_years_ago = np.random.exponential(scale=5, size=infected_count)  # sample years in the past
    # pandas requires 'timedelta' type for date calculations
    infected_td_ago = pd.to_timedelta(infected_years_ago, unit='y')

    # date of death of the infected individuals (in the future)
    death_years_ahead = np.random.exponential(scale=2, size=infected_count)
    death_td_ahead = pd.to_timedelta(death_years_ahead, unit='y')

    # set the properties of infected individuals
    df.loc[df.mi_is_infected, 'mi_date_infected'] = self.sim.date - infected_td_ago
    df.loc[df.mi_is_infected, 'mi_date_death'] = self.sim.date + death_td_ahead

    df.loc[df.mi_is_infected & (df.age_years > 15), 'mi_status'] = 'T1'
    df.loc[df.mi_is_infected & (df.age_years <= 15), 'mi_status'] = 'T2'

Run the test again. What happens?

initialise_simulation(self, sim)

Add initial events to the event queue:

def initialise_simulation(self, sim):
    """Get ready for simulation start.

    This method is called just before the main simulation loop begins, and after all
    modules have read their parameters and the initial population has been created.
    It is a good place to add initial events to the event queue.
    """
    # add the basic event (we will implement below)
    event = MockitisEvent(self)
    sim.schedule_event(event, sim.date + DateOffset(months=1))

    # add an event to log to screen
    sim.schedule_event(MockitisLoggingEvent(self), sim.date + DateOffset(months=6))

    # add the death event of infected individuals
    df = sim.population.props
    infected_individuals = df[df.mi_is_infected].index
    for index in infected_individuals:
        individual = self.sim.population[index]
        death_event = MockitisDeathEvent(self, individual)
        self.sim.schedule_event(death_event, individual.mi_date_death)

We have not implemented some of these events yet, so add them after the MockitisEvent class using the implementation below. You will need to import the Event and IndividualScopeEventMixin classes at the top of the file, so your import lines should now look as follows:

import pandas as pd
import numpy as np

from tlo import DateOffset, Module, Parameter, Property, Types
from tlo.events import PopulationScopeEventMixin, RegularEvent, Event, IndividualScopeEventMixin

And the empty event implementations are:

class MockitisDeathEvent(Event, IndividualScopeEventMixin):
    def __init__(self, module, individual):
        super().__init__(module, person=individual)

    def apply(self, individual):
        pass

class MockitisLoggingEvent(RegularEvent, PopulationScopeEventMixin):
    def __init__(self, module):
        super().__init__(module, frequency=DateOffset(months=1))

    def apply(self, population):
        pass

We can also change the MockitisEvent class apply method:

class MockitisEvent(RegularEvent, PopulationScopeEventMixin):
    ....
    def apply(self, population):
        pass

Run the test again.

on_birth(self, mother, child)

Called by the simulation any time an individual is born. This is where we can pass properties from mother to child. Let's imagine that there is no inheritence of this disease:

def on_birth(self, mother, child):
    pass

Implement the events for this module

We are going to implement three events for this example:

  1. MockitisEvent: the standard event that updated the infected/cured individuals periodically throughout the simulation run.
  2. MockitisDeathEvent: event scheduled for an infected individual, triggered on their death by infection
  3. MockitisLoggingEvent: an event to print a line to screen for debugging (useful during dev.)

MockitisLoggingEvent

class MockitisLoggingEvent(RegularEvent, PopulationScopeEventMixin):
    def __init__(self, module):
        """comments...
        """
        # run this event every year
        self.repeat = 12
        super().__init__(module, frequency=DateOffset(months=self.repeat))

    def apply(self, population):
        # get some summary statistics
        df = population.props

        infected_total = df.mi_is_infected.sum()
        proportion_infected = infected_total / len(df)

        mask = (df['mi_date_infected'] > self.sim.date - DateOffset(months=self.repeat))
        infected_in_last_month = mask.sum()
        mask = (df['mi_date_cure'] > self.sim.date - DateOffset(months=self.repeat))
        cured_in_last_month = mask.sum()

        counts = {'N': 0, 'T1': 0, 'T2': 0, 'P': 0}
        counts.update(df['mi_status'].value_counts().to_dict())
        status = 'Status: { N: %(N)d; T1: %(T1)d; T2: %(T2)d; P: %(P)d }' % counts

        self.module.store.append(proportion_infected)

        print('%s - Mockitis: {TotInf: %d; PropInf: %.3f; PrevMonth: {Inf: %d; Cured: %d}; %s }' %
              (self.sim.date,
               infected_total,
               proportion_infected,
               infected_in_last_month,
               cured_in_last_month,
               status), flush=True)

MockitisEvent

class MockitisEvent(RegularEvent, PopulationScopeEventMixin):

    def __init__(self, module):
        super().__init__(module, frequency=DateOffset(months=1))
        self.p_infection = module.parameters['p_infection']
        self.p_cure = module.parameters['p_cure']

    def apply(self, population):
        df = population.props

        # 1. get (and hold) index of currently infected and uninfected individuals
        currently_infected = df.index[df.mi_is_infected & df.is_alive]
        currently_uninfected = df.index[~df.mi_is_infected & df.is_alive]

        # 2. handle new infections
        now_infected = np.random.choice([True, False], size=len(currently_uninfected),
                                        p=[self.p_infection, 1 - self.p_infection])
        # if any are infected
        if now_infected.sum():
            infected_idx = currently_uninfected[now_infected]

            df.loc[infected_idx, 'mi_date_infected'] = self.sim.date
            df.loc[infected_idx, 'mi_date_death'] = self.sim.date + pd.Timedelta(25, unit='Y')
            df.loc[infected_idx, 'mi_date_cure'] = pd.NaT
            df.loc[infected_idx, 'mi_is_infected'] = True

            infected_lte_15 = df.index[df.age_years <= 15].intersection(infected_idx)
            infected_gt_15 = df.index[df.age_years > 15].intersection(infected_idx)
            df.loc[infected_gt_15, 'mi_status'] = 'NT1'
            df.loc[infected_lte_15, 'mi_status'] = 'NT2'

            # schedule death events for newly infected individuals
            for person_index in infected_idx:
                individual = self.sim.population[index]
                death_event = MockitisDeathEvent(self, person_index)
                self.sim.schedule_event(death_event, df.at[person_index, 'mi_date_death'])

        # 3. handle cures
        cured = np.random.choice([True, False], size=len(currently_infected), p=[self.p_cure, 1 - self.p_cure])

        if cured.sum():
            cured = currently_infected[cured]
            df.loc[cured, 'mi_is_infected'] = False
            df.loc[cured, 'mi_status'] = 'P'
            df.loc[cured, 'mi_date_death'] = pd.NaT
            df.loc[cured, 'mi_date_cure'] = self.sim.date

MockitisDeathEvent (to do!)

class MockitisDeathEvent(Event, IndividualScopeEventMixin):
    def __init__(self, module, individual):
        super().__init__(module, person=individual)

    def apply(self, person_id):
        df = self.sim.population.props  # shortcut to the dataframe

        # Apply checks to ensure that this death should occur
        if df.at[person_id, 'mi_status'] == 'C':

            # Fire the centralised death event:
            death = InstantaneousDeath(self.module, person_id, cause='Mockitis')
            self.sim.schedule_event(death, self.sim.date)

TLO Model Wiki

Clone this wiki locally