Skip to content

Implementing a disease module

Jump to bottom
Asif Tamuri edited this page Oct 1, 2018 · 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

git checkout -b 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

touch tests/test_mockitis.py

and add the following:

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


def test_mockitis_simulation():
    sim = Simulation(start_date=Date(2010, 1, 1))
    mockitis_module = mockitis.Mockitis()
    sim.register(mockitis_module)
    sim.make_initial_population(n=500)
    sim.simulate(end_date=Date(2028, 1, 1))

Run the mockitis test (demo)

The test will fail with NotImplementedError!

Start modifying the new disease module

Let's set the parameters of the module:

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'),
    '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

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
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'] = '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_individuals = df[df.mi_is_infected].index

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

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

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

initialise_simulation(self, sim)

Add initial events to the event queue:

def initialise_simulation(self, sim):
    """some comments...
    """
    # 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=1))

    # 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 and we will implement below. You will need to import the Event and IndividualScopeEventMixin classes.

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

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 month
        super().__init__(module, frequency=DateOffset(months=1))

    def apply(self, population):
        # print 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=1))
        infected_in_last_month = mask.sum()
        mask = (df['mi_date_cure'] > self.sim.date - DateOffset(months=1))
        cured_in_last_month = mask.sum()

        print('Mockitis: {TotInf: %d; PropInf: %.3f; PrevMonth: {Inf: %d; Cured: %d} }' % (infected_total, proportion_infected, infected_in_last_month, cured_in_last_month))

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]
        currently_uninfected = df.index[~df.mi_is_infected]

        # 2. handle new infections
        new_infected = np.random.choice([True, False], size=len(currently_uninfected), p=[self.p_infection, 1 - self.p_infection])
        # if any are infected
        if new_infected.sum():
            new_infected = currently_uninfected[new_infected]
            df.loc[new_infected, 'mi_status'] = 'T2'
            df.loc[new_infected, 'mi_date_infected'] = self.sim.date
            df.loc[new_infected, 'mi_date_death'] = self.sim.date + pd.Timedelta(5, unit='Y')

            # schedule death events for newly infected individuals
            for index in new_infected:
                individual = self.sim.population[index]
                death_event = MockitisDeathEvent(self, individual)
                self.sim.schedule_event(death_event, individual.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, individual):
        pass

TLO Model Wiki

Clone this wiki locally