Implementing a disease module

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.

[NB. This describes a very basic module that does not have logging or interaction with the healthsystem or healthburden module. See the examples of 'mockitis' and 'chronicsyndrome' for those additions].

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-mymockitis

For example, I will use:

git checkout -b feature/tamuri-mymockitis

Copy the skeleton template for modules

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

Edit the src/tlo/methods/mymockitis.py file by changing the following lines:

class Skeleton(Module): → class MyMockitis(Module):

and

class Skeleton_Event(Module): → class MyMockitisEvent(Module):

There are also some other lines where Skeleton needs to be replaced by the name of your module (here, MyMockitis), i.e.:

assert isinstance(module, Skeleton) → assert isinstance(module, MyMockitis)

which occurs several times in the file. Ditto for any others which you find.

e.g. class Skeleton_LoggingEvent → class MyMockitisLoggingEvent

and

class HSI_Skeleton_Example_Interaction → class HSI_Mockitis_Example_Interaction

and

self.TREATMENT_ID = 'Skeleton_Example_Interaction' → self.TREATMENT_ID = 'HSI_MyMockitis_Example_Interaction'

Now check that the following code is underneath the SYMPTOMS block:

    def __init__(self, name=None, resourcefilepath=None):
        # NB. Parameters passed to the module can be inserted in the __init__ definition.

        super().__init__(name)
        self.resourcefilepath = resourcefilepath
        self.store = {'Proportion_infected': []}  # This line should be added if not present.

Create a file for testing the new module

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

touch tests/test_mymockitis.py

and add the following:

import os
from pathlib import Path

import pytest

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

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),
                 simplified_births.SimplifiedBirths(resourcefilepath=resourcefilepath))
    return sim


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


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

This provides the basic outline of the test.

Run the mymockitis test (demo)

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

This test can be run both using pytest or as a standalone script, suitable for development.

i.e. either:

pytest tests/test_mymockitis.py

or:

python3 tests/test_mymockitis.py

Start modifying the new disease module

Let's set the parameters of the module in the file mymockitis.py; replace the existing PARAMETERS block with:

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). Replace the existing PROPERTIES block with the following.

PROPERTIES = {
    'mi_is_infected': Property(Types.BOOL, 'Current status of mymockitis'),
    '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'),
}

NB Strictly, we should also explain the NT1 and NT2 categories in the Historical status line, the idea being that every category in your disease model has an informative string to explain what the relevant code means.

In the following, as we have just done above, replace any existing stub code in your mymockitis.py file with the fuller entries below.

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.01

Update the test with the new module

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

from tlo.methods import mymockitis

Then update the simulation() method.

@pytest.fixture
def simulation():
    sim = Simulation(start_date=start_date)
    sim.register(demography.Demography(resourcefilepath=resourcefilepath),
                 simplified_births.SimplifiedBirths(resourcefilepath=resourcefilepath))

    # Instantiate and add the mymockitis module to the simulation
    mymockitis_module = mymockitis.MyMockitis()
    sim.register(mymockitis_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 mymockitis.py 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 (NB values[:] is used so we don't lose categorical information and convert the entire column to a string 'N')
    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 = MyMockitisEvent(self)
    sim.schedule_event(event, sim.date + DateOffset(months=1))

    # add an event to log to screen
    sim.schedule_event(MyMockitisLoggingEvent(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 = df.iloc[index]
        # death_event = MyMockitisDeathEvent(self, individual)
        death_event = MyMockitisDeathEvent(self, index)
        self.sim.schedule_event(death_event, individual.mi_date_death)

We have not implemented some of these events yet, so add them after the MyMockitisEvent 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, logging
from tlo.events import PopulationScopeEventMixin, RegularEvent, Event, IndividualScopeEventMixin

And the empty event implementations are:

class MyMockitisDeathEvent(Event, IndividualScopeEventMixin):
    def __init__(self, module, individual):
        super().__init__(module, person_id=individual)
        assert isinstance(module, MyMockitis)

    def apply(self, individual):
        pass

class MyMockitisLoggingEvent(RegularEvent, PopulationScopeEventMixin):
    def __init__(self, module):
        super().__init__(module, frequency=DateOffset(months=1))
        assert isinstance(module, MyMockitis)

    def apply(self, population):
        pass

We can also change the MyMockitisEvent class apply method:

class MyMockitisEvent(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 inheritance of this disease:

def on_birth(self, mother, child):
    pass

At this point, your program should run without errors.

Implement the events for this module

We are going to implement three events for this example:

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

Add the blocks of code below into mymockitis.py, overwriting any existing blocks with the same name.

MyMockitisLoggingEvent

NB. This is not the logging event that is used for logging in the final version just a simple 'for development version'

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

    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['Proportion_infected'].append(proportion_infected)

        print('%s - MyMockitis: {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)

MyMockitisEvent

class MyMockitisEvent(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']
        assert isinstance(module, MyMockitis)

    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:
                death_event = MyMockitisDeathEvent(self.module, 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

MyMockitisDeathEvent

Here we are using the InstantaneousDeath class, so the imports at the top of the file should be as follows

import numpy as np
import pandas as pd

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

Now we create a Death Event

class MyMockitisDeathEvent(Event, IndividualScopeEventMixin):
    def __init__(self, module, individual):
        super().__init__(module, person_id=individual)
        assert isinstance(module, MyMockitis)

    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='MyMockitis')
            self.sim.schedule_event(death, self.sim.date)