Comprehensive Analyzer Plots Example

This notebook demonstrates all available plotting methods in the DwtAnalyzer class from the tbsim package. It is adapted from the original Python script for interactive exploration.

import tbsim as mtb
import starsim as ss
import sciris as sc
import matplotlib.pyplot as plt
import numpy as np
import scipy.stats as stats
TBS = mtb.TBS

Build a TB simulation with dwell time analyzer

def build_tbsim(sim_pars=None):
    sim_params = dict(
        start = ss.date('2013-01-01'),
        stop = ss.date('2016-12-31'),
        rand_seed=123,
        dt=ss.days(7),
    )
    if sim_pars is not None:
        sim_params.update(sim_pars)
    pop = ss.People(n_agents=1000)
    tb_params = dict(
        beta=ss.peryear(0.025),
        init_prev=ss.bernoulli(p=0.25),
        rel_sus_latentslow=0.2,
    )
    tb = mtb.TB(tb_params)
    net = ss.RandomNet(dict(n_contacts=ss.poisson(lam=5), dur=0))
    dwell_analyzer = mtb.DwtAnalyzer(adjust_to_unit=True, unit=1.0, scenario_name='comprehensive_plots_example')
    sim = ss.Sim(
        people=pop,
        networks=net,
        diseases=tb,
        pars=sim_params,
        analyzers=dwell_analyzer,
    )
    sim.pars.verbose = 0
    return sim

Run the simulation and extract the analyzer

print("Building and running TB simulation...")
sim_tb = build_tbsim()
sim_tb.run()
analyzer = sim_tb.analyzers[0]

Building and running TB simulation...
No dwell time data available to plot.
No data provided, or data is corrupted
===> Dwell time logs saved to:
 /Users/mine/newgit/newtbsim/docs/tutorials/results/comprehensiveplotsexample-0908161007.csv

1. Sankey Diagrams

Demonstrate Sankey diagrams for state transitions.

# Basic Sankey diagram for all agents
analyzer.sankey_agents()

Data type cannot be displayed: application/vnd.plotly.v1+json

# Sankey diagram with dwell times
analyzer.sankey_dwelltimes(subtitle="State Transitions with Dwell Times")

# Sankey diagrams by age groups
analyzer.sankey_agents_by_age_subplots(bins=[0, 5, 15, 30, 50, 200], scenario="Age-stratified Analysis")

# Sankey diagrams with even age ranges
analyzer.sankey_agents_even_age_ranges(number_of_plots=3, scenario="Even Age Distribution")

2. Network Graphs

Visualize state transition networks.

analyzer.graph_state_transitions(subtitle="State Transition Network", colormap='tab20')

analyzer.graph_state_transitions_curved(subtitle="Curved State Transitions", colormap='plasma')

3. Histograms and Distributions

Explore dwell time distributions.

help(analyzer.histogram_with_kde)

4. Interactive Bar Charts

Interactive bar charts for state transitions and reinfections.

analyzer.barchar_all_state_transitions_interactive(
    dwell_time_bins=[0, 30, 90, 180, 365, float('inf')],
    filter_states=['-1.None', '0.Latent Slow', '1.Latent Fast', '2.Active Presymp']
)

analyzer.reinfections_age_bins_bars_interactive(
    target_states=[0.0, 1.0],
    barmode='group',
    scenario="Age-stratified Reinfection Analysis"
)

analyzer.reinfections_percents_bars_interactive(
    target_states=[0.0, 1.0],
    scenario="Population Reinfection Analysis"
)

analyzer.reinfections_bystates_bars_interactive(
    target_states=[0.0, 1.0],
    scenario="State Transition Reinfection Analysis",
    barmode='group'
)

5. Stacked Bar Charts

Stacked bar charts for cumulative time and dwell time analysis.

analyzer.stacked_bars_states_per_agent_static()
analyzer.stackedbars_dwelltime_state_interactive(bin_size=5, num_bins=15)
analyzer.stackedbars_subplots_state_transitions(bin_size=2, num_bins=25)

6. Custom Transition Analysis

Custom transition subplots.

custom_transitions = {
    '-1.0.None': ['0.0.Latent Slow', '1.0.Latent Fast'],
    '0.0.Latent Slow': ['2.0.Active Presymp', '-1.0.None'],
    '1.0.Latent Fast': ['2.0.Active Presymp', '-1.0.None']
}
analyzer.subplot_custom_transitions(transitions_dict=custom_transitions)

7. Survival Analysis

Kaplan-Meier survival curve for dwell times.

analyzer.plot_kaplan_meier(dwell_time_col='dwell_time')

8. Using DwtPlotter Directly

Demonstrate additional plots using the DwtPlotter class.

file_path = analyzer.file_path
print(f'Generated data file: {file_path}')
plotter = mtb.DwtPlotter(file_path=file_path)
plotter.histogram_with_kde(subtitle="From Generated File")
plotter.sankey_agents(subtitle="From Generated File")

9. Post Processor Demonstration

Example usage of the DwtPostProcessor for multiple simulation results. (This is informational; actual usage requires multiple result files.)

# Example usage (commented out):
# postproc = mtb.DwtPostProcessor(directory='results', prefix='Baseline')
# postproc.sankey_agents(subtitle="Aggregated Results")
# postproc.histogram_with_kde(subtitle="Aggregated Distributions")
# postproc.reinfections_percents_bars_interactive(
#     target_states=[0.0, 1.0],
#     scenario="Aggregated Reinfection Analysis"
# )