feat: add agent parameters, more info labels

app.py

@@ -7,14 +7,18 @@ from ui.tabs.with_agent_tab import WithAgentTab
 with gr.Blocks(theme=gr.themes.Glass()) as interface:
     gr.Markdown("**Social distancing sim**")
     gr.Markdown(
-        "Graph based disease spread simulator. Uses https://github.com/garethjns/social-distancing-sim - "
-        "see [README](https://github.com/garethjns/social-distancing-sim/blob/master/README.MD) "
-        "for details."
+        """Graph based disease spread simulator. Uses https://github.com/garethjns/social-distancing-sim -
+        see [README](https://github.com/garethjns/social-distancing-sim/blob/master/README.MD)
+        for details.
+
+        Individuals in the population are represented as graph nodes, with edges modelling connections 
+        allowing the possibility of disease transmission. between individuals. 
+        """
     )
     gr.Markdown(f"App version: {__version__} using `social-distancing-sim==0.11.2`")
-    _ = SimpleExampleTab().build()
-    _ = SinglePopTab().build()
-    _ = WithAgentTab().build()
+    SimpleExampleTab().build()
+    SinglePopTab().build()
+    WithAgentTab().build()
 
 if __name__ == "__main__":
     interface.launch()

ui/__init__.py

@@ -1 +1 @@
-__version__ = "0.2.0"
+__version__ = "0.3.0"

ui/functions/run_gym_env.py

@@ -12,18 +12,28 @@ from social_distancing_sim.templates.template_base import TemplateBase
 
 import gradio as gr
 
+AGENT_TYPE_NAMES = Literal[
+    "dummy",
+    "distancing",
+    "vaccination",
+    "treatment",
+    "masking",
+]
 
-def _run_and_replay(sim):
-    sim.run()
-    if sim.save:
-        sim.env.replay()
+AGENT_TYPE_MAP = {
+    "dummy": agent.DummyAgent,
+    "distancing": agent.DistancingPolicyAgent,
+    "vaccination": agent.VaccinationAgent,
+    "treatment": agent.TreatmentAgent,
+    "masking": agent.MaskingAgent,
+}
 
 
 class _EnvTemplate(TemplateBase):
     def build(self):
         seed = np.random.randint(int(1e6))
         return env.Environment(
-            name="visual_compare_basic_agents_custom_env",
+            name="custom_env",
             disease=env.Disease(
                 name="COVID-19",
                 virulence=0.01,
@@ -50,38 +60,57 @@ class CustomEnv(GymEnv):
     template = _EnvTemplate()
 
 
-AGENT_TYPE_NAMES = Literal[
-    "dummy",
-    "distancing",
-    "vaccination",
-    "treatment",
-    "masking",
-]
-
-AGENT_TYPE_MAP = {
-    "dummy": agent.DummyAgent,
-    "distancing": agent.DistancingPolicyAgent,
-    "vaccination": agent.VaccinationAgent,
-    "treatment": agent.TreatmentAgent,
-    "masking": agent.MaskingAgent,
-}
-
-
 def run_gym_env(
     name: str,
     steps: int,
     test_rate: float,
     agents: list[AGENT_TYPE_NAMES],
-    n_actions: int = 6,
-    progress=gr.Progress(track_tqdm=True),
+    n_actions: int,
+    isolate_start: int,
+    reconnect_start: int,
+    isolate_end: int,
+    reconnect_end: int,
+    vaccination_start: int,
+    vaccination_end: int,
+    treatment_start: int,
+    treatment_end: int,
+    masking_start: int,
+    remove_mask_start: int,
+    masking_end: int,
+    remove_mask_end: int,
+    progress=gr.Progress(track_tqdm=True),  # noqa
 ):
-
     if agents is None:
         # Set no agent
         agents = ["dummy"]
 
+    start_step = "start_step"
+    end_step = "end_step"
+    agent_params = {
+        "dummy": {start_step: 0, end_step: steps},
+        "distancing": {
+            start_step: {"isolate": isolate_start, "reconnect": reconnect_start},
+            end_step: {"isolate": isolate_end, "reconnect": reconnect_end},
+        },
+        "vaccination": {
+            start_step: {"vaccinate": vaccination_start},
+            end_step: {"vaccinate": vaccination_end},
+        },
+        "treatment": {
+            start_step: {"treat": treatment_start},
+            end_step: {"treat": treatment_end},
+        },
+        "masking": {
+            start_step: {
+                "provide_mask": masking_start,
+                "remove_mask": remove_mask_start,
+            },
+            end_step: {"provide_mask": masking_end, "remove_mask": remove_mask_end},
+        },
+    }
+
     # Prepare agents
-    prepared_agents = [AGENT_TYPE_MAP[a] for a in agents]
+    prepared_agents = [(AGENT_TYPE_MAP[a], agent_params[a]) for a in agents]
 
     # Prepare a custom environment
     env_name = f"SDS-CustomEnv{np.random.randint(int(2e6))}-v0"
@@ -99,8 +128,9 @@ def run_gym_env(
         agent_class(
             actions_per_turn=adjusted_actions_per_turn,
             name=f"{agent_class.__name__} - {n_actions} actions",
+            **agent_params,
         )
-        for agent_class in prepared_agents
+        for agent_class, agent_params in prepared_agents
     ]
 
     if n_agents == 1:
@@ -120,15 +150,7 @@ def run_gym_env(
     )
 
     # Run all the prepared Sims
-    _run_and_replay(sim)
+    sim.run()
+    sim.env.replay()
 
     return pathlib.Path(f"{sim.save_path}/replay.gif")
-
-
-if __name__ == "__main__":
-    run_gym_env(
-        name="test",
-        steps=5,
-        test_rate=1,
-        agents=["dummy"],
-    )

ui/functions/run_single_population.py

@@ -8,6 +8,8 @@ import gradio as gr
 def run_single_population(
     name: str,
     steps: int,
+    n_initial_infections: int,
+    random_infection_chance: float,
     test_rate: float,
     community_n: int,
     community_size_mean: int,
@@ -65,6 +67,8 @@ def run_single_population(
         healthcare=healthcare,
         environment_plotting=environment_plotting,
         observation_space=observation_space,
+        initial_infections=n_initial_infections,
+        random_infection_chance=random_infection_chance,
     )
     # Turn logging on
     pop.log_to_file = True

ui/labels/disease_labels.py

@@ -1,5 +1,16 @@
 TITLE = "**Disease parameters**"
 VIRULENCE = "Virulence"
+VIRULENCE_INFO = (
+    "Increased virulence makes the disease more likely to "
+    "be transferred from an infected node to a connected uninfected node each turn."
+)
 RECOVERY_RATE = "Recovery rate"
+RECOVERY_RATE_INFO = "The likelihood an individual will recover at the end of the infection period, rather than dying."
 IMMUNITY_MEAN = "Immunity after recovery"
+IMMUNITY_MEAN_INFO = (
+    "The average disease immunity gained by an individual that survives infection."
+)
 IMMUNITY_DECAY_MEAN = "Immunity decay rate"
+IMMUNITY_DECAY_MEAN_INFO = (
+    "The average rate at which a individual's immunity decays after recovery."
+)

ui/labels/healthcare_labels.py

@@ -1,3 +1,11 @@
 TITLE = "**Healthcare parameters**"
 CAPACITY = "Capacity"
+CAPACITY_INFO = (
+    "The capacity of the healthcare system. When there are are more "
+    "infected nodes requiring treatment than there is capacity for, chance of death is increased"
+)
 MAX_PENALTY = "Max penalty when over capacity"
+MAX_PENALTY_INFO = (
+    "The maximum amount by which the chance of an infected node dying is "
+    "increased when the healthcare system is over capacity."
+)

ui/labels/population_labels.py

@@ -1,6 +1,12 @@
 TITLE = "**Population parameters**"
 COMMUNITY_N = "Number of communities"
+COMMUNITY_N_INFO = (
+    "Communities are cluster of graph nodes, with generally more connections "
+    "between members than with non-members."
+)
 MEAN_COMMUNITY_SIZE = "Mean community size"
-CONNECTION_RATE = "Connection rate within communities"
+MEAN_COMMUNITY_SIZE_INFO = (
+    "Mean number of members (nodes) in each community (cluster of nodes)."
+)
 COMMUNITY_P_IN = "Connection rate within communities"
 COMMUNITY_P_OUT = "Connection rate between communities"

ui/labels/simulation_labels.py

@@ -1,4 +1,22 @@
 TITLE = "**Simulation parameters**"
 OUTPUT_PATH = "Output path"
 STEPS = "Number of steps"
+STEPS_INFO = "The number of steps to run the simulation for"
 TEST_RATE = "Test rate"
+TEST_RATE_INFO = (
+    "The rate of testing applied at each step - node disease state is not fully known "
+    "without testing, whereas some state is always know (e.g. death)."
+    "Testing is partly random, with infected nodes being tested more often than uninfected, "
+    "based on the assumption that having symptoms causes an individual to actively seek testing."
+)
+TEST_RATE_EXTRA_INFO = (
+    "When test rate == 1, the ground truth is fully observed and plotted. "
+    "When it's <1, the ground truth graph is shown on the left, and the observed graph shown on the right."
+)
+INITIAL_INFECTIONS = "Number of initial infections"
+INITIAL_INFECTIONS_INFO = "The number of individuals randomly infected on step 0"
+RANDOM_INFECTION_CHANCE = "Random infection chance"
+RANDOM_INFECTION_CHANCE_INFO = (
+    "Chance each step that an individual is randomly infected, even without being "
+    "connected to another infected individual"
+)

ui/tabs/compare_pop_tab.py

@@ -1,8 +1,6 @@
 from dataclasses import dataclass
 from typing import Callable
 
-from gradio.components import FormComponent
-
 import gradio as gr
 from ui.labels import simulation_labels
 from ui.tabs.tab_builder import ITabBuilder
@@ -16,7 +14,7 @@ class ComparePopTab(ITabBuilder):
     output_format = "gif"
     button_label: str = "run"
 
-    def build(self) -> tuple[list[FormComponent], list[gr.Image], gr.Button]:
+    def build(self) -> None:
         with gr.Tab(self.tab_name):
             inputs = []
             with gr.Row():
@@ -29,5 +27,3 @@ class ComparePopTab(ITabBuilder):
 
             button = gr.Button(self.button_label)
             button.click(self.f, inputs=inputs, outputs=outputs)
-
-        return inputs, outputs, button

ui/tabs/simple_example_tab.py

@@ -1,10 +1,7 @@
 from dataclasses import dataclass
 from typing import Callable
 
-from gradio.components import FormComponent
-
 import gradio as gr
-from gradio import Image
 from ui.functions.run_single_population import run_single_population
 from ui.labels import population_labels, simulation_labels
 from ui.tabs.tab_builder import ITabBuilder
@@ -16,9 +13,9 @@ class SimpleExampleTab(ITabBuilder):
     tab_name = "Simple example"
     default_output_path: str = "./gradio/simple-example/"
     output_format = "gif"
-    button_label: str = "run"
+    button_label: str = "Run"
 
-    def build(self) -> tuple[list[FormComponent], list[Image], gr.Button]:
+    def build(self) -> None:
 
         with gr.Tab(self.tab_name):
             inputs = []
@@ -42,6 +39,23 @@ class SimpleExampleTab(ITabBuilder):
                                 step=1,
                                 value=20,
                                 label=simulation_labels.STEPS,
+                                info=simulation_labels.STEPS_INFO,
+                            ),
+                            gr.Slider(
+                                minimum=1,
+                                maximum=100,
+                                step=1,
+                                value=10,
+                                label=simulation_labels.INITIAL_INFECTIONS,
+                                info=simulation_labels.INITIAL_INFECTIONS_INFO,
+                            ),
+                            gr.Slider(
+                                minimum=0,
+                                maximum=1,
+                                step=0.01,
+                                value=0.02,
+                                label=simulation_labels.RANDOM_INFECTION_CHANCE,
+                                info=simulation_labels.RANDOM_INFECTION_CHANCE_INFO,
                             ),
                             gr.Slider(
                                 minimum=0.0,
@@ -49,6 +63,7 @@ class SimpleExampleTab(ITabBuilder):
                                 step=0.01,
                                 value=1.0,
                                 label=simulation_labels.TEST_RATE,
+                                info=f"{simulation_labels.TEST_RATE_INFO} {simulation_labels.TEST_RATE_EXTRA_INFO}",
                                 interactive=False,
                             ),
                         ]
@@ -61,8 +76,9 @@ class SimpleExampleTab(ITabBuilder):
                                 minimum=1,
                                 maximum=100,
                                 step=1,
-                                value=3,
+                                value=6,
                                 label=population_labels.COMMUNITY_N,
+                                info=population_labels.COMMUNITY_N_INFO,
                             ),
                             gr.Slider(
                                 minimum=1,
@@ -70,13 +86,13 @@ class SimpleExampleTab(ITabBuilder):
                                 step=1,
                                 value=10,
                                 label=population_labels.MEAN_COMMUNITY_SIZE,
+                                info=population_labels.MEAN_COMMUNITY_SIZE_INFO,
                             ),
                         ]
                     )
 
-                outputs = [gr.Image(format=self.output_format)]
-
-            button = gr.Button(self.button_label)
-            button.click(self.f, inputs=inputs, outputs=outputs)
+                with gr.Column():
+                    outputs = [gr.Image(format=self.output_format)]
 
-        return inputs, outputs, button
+                    button = gr.Button(self.button_label)
+                    button.click(self.f, inputs=inputs, outputs=outputs)

ui/tabs/single_pop_tab.py

@@ -1,9 +1,6 @@
 from dataclasses import dataclass
 from typing import Callable
 
-from gradio.components.base import FormComponent
-from gradio.components.image import Image
-
 import gradio as gr
 from ui.functions.run_single_population import run_single_population
 from ui.labels import (
@@ -21,16 +18,15 @@ class SinglePopTab(ITabBuilder):
     tab_name = "Passive population"
     default_output_path: str = "./gradio/single-population-example/"
     output_format = "gif"
-    button_label: str = "run"
+    button_label: str = "Run"
+
+    def build(self) -> None:
 
-    def build(
-        self,
-    ) -> tuple[list[FormComponent], list[Image], gr.Button]:
         with gr.Tab(self.tab_name):
             inputs = []
             gr.Markdown(
                 "Run a customized simulation. "
-                "With Test rate < 1.0, results will show both ground truth and observed space."
+                "With Test rate <= 1.0, results will show both ground truth and observed space."
             )
 
             with gr.Row():
@@ -50,6 +46,23 @@ class SinglePopTab(ITabBuilder):
                                         step=1,
                                         value=50,
                                         label=simulation_labels.STEPS,
+                                        info=simulation_labels.STEPS_INFO,
+                                    ),
+                                    gr.Slider(
+                                        minimum=1,
+                                        maximum=100,
+                                        step=1,
+                                        value=2,
+                                        label=simulation_labels.INITIAL_INFECTIONS,
+                                        info=simulation_labels.INITIAL_INFECTIONS_INFO,
+                                    ),
+                                    gr.Slider(
+                                        minimum=0,
+                                        maximum=1,
+                                        step=0.01,
+                                        value=0.02,
+                                        label=simulation_labels.RANDOM_INFECTION_CHANCE,
+                                        info=simulation_labels.RANDOM_INFECTION_CHANCE_INFO,
                                     ),
                                     gr.Slider(
                                         minimum=0.0,
@@ -57,6 +70,7 @@ class SinglePopTab(ITabBuilder):
                                         step=0.01,
                                         value=0.25,
                                         label=simulation_labels.TEST_RATE,
+                                        info=simulation_labels.TEST_RATE_INFO,
                                     ),
                                 ]
                             )
@@ -70,6 +84,7 @@ class SinglePopTab(ITabBuilder):
                                         step=1,
                                         value=50,
                                         label=population_labels.COMMUNITY_N,
+                                        info=population_labels.COMMUNITY_N_INFO,
                                     ),
                                     gr.Slider(
                                         minimum=1,
@@ -77,6 +92,7 @@ class SinglePopTab(ITabBuilder):
                                         step=1,
                                         value=15,
                                         label=population_labels.MEAN_COMMUNITY_SIZE,
+                                        info=population_labels.MEAN_COMMUNITY_SIZE_INFO,
                                     ),
                                     gr.Slider(
                                         minimum=-0.0,
@@ -106,6 +122,7 @@ class SinglePopTab(ITabBuilder):
                                     step=0.001,
                                     value=0.005,
                                     label=disease_labels.VIRULENCE,
+                                    info=disease_labels.VIRULENCE_INFO,
                                 ),
                                 gr.Slider(
                                     minimum=-0.0,
@@ -113,6 +130,7 @@ class SinglePopTab(ITabBuilder):
                                     step=0.1,
                                     value=0.99,
                                     label=disease_labels.RECOVERY_RATE,
+                                    info=disease_labels.RECOVERY_RATE_INFO,
                                 ),
                                 gr.Slider(
                                     minimum=-0.0,
@@ -120,6 +138,7 @@ class SinglePopTab(ITabBuilder):
                                     step=0.1,
                                     value=0.66,
                                     label=disease_labels.IMMUNITY_MEAN,
+                                    info=disease_labels.IMMUNITY_MEAN_INFO,
                                 ),
                                 gr.Slider(
                                     minimum=-0.0,
@@ -127,6 +146,7 @@ class SinglePopTab(ITabBuilder):
                                     step=0.1,
                                     value=0.1,
                                     label=disease_labels.IMMUNITY_DECAY_MEAN,
+                                    info=disease_labels.IMMUNITY_DECAY_MEAN_INFO,
                                 ),
                             ]
                         )
@@ -140,6 +160,7 @@ class SinglePopTab(ITabBuilder):
                                     step=1,
                                     value=200,
                                     label=healthcare_labels.CAPACITY,
+                                    info=healthcare_labels.CAPACITY_INFO,
                                 ),
                                 gr.Slider(
                                     minimum=0.0,
@@ -147,6 +168,7 @@ class SinglePopTab(ITabBuilder):
                                     step=0.01,
                                     value=0.5,
                                     label=healthcare_labels.MAX_PENALTY,
+                                    info=healthcare_labels.MAX_PENALTY_INFO,
                                 ),
                             ]
                         )
@@ -154,5 +176,3 @@ class SinglePopTab(ITabBuilder):
             outputs = [gr.Image(format=self.output_format)]
             button = gr.Button(self.button_label)
             button.click(self.f, inputs=inputs, outputs=outputs)
-
-        return inputs, outputs, button

ui/tabs/tab_builder.py

@@ -1,15 +1,11 @@
 import abc
 from typing import Callable
 
-from gradio.components import FormComponent
-
-import gradio as gr
-
 
 class ITabBuilder(abc.ABC):
     f: Callable
     name: str
 
     @abc.abstractmethod
-    def build(self) -> tuple[list[FormComponent], list[gr.Image], gr.Button]:
+    def build(self) -> None:
         """Compose the tab components."""

ui/tabs/with_agent_tab.py

@@ -1,8 +1,6 @@
 from dataclasses import dataclass
 from typing import Callable
 
-from gradio.components import FormComponent
-
 import gradio as gr
 from ui.functions.run_gym_env import run_gym_env
 from ui.labels import simulation_labels
@@ -15,36 +13,39 @@ class WithAgentTab(ITabBuilder):
     tab_name = "With agent(s)"
     default_output_path: str = "./gradio/agent-example/"
     output_format = "gif"
-    button_label: str = "run"
+    button_label: str = "Run"
 
     def build(
         self,
-    ) -> tuple[list[FormComponent], list[gr.Image], gr.Button]:
-
+    ) -> None:
+        inputs = []
         with gr.Tab(self.tab_name):
-            inputs = []
             gr.Markdown(
-                "Run a simulation with fixed parameters and optional agents. "
-                "Agents and number of actions they perform each turn can be configured."
-            )
+                """
+                Run a simulation with fixed parameters and optional agents.
 
+                Agents have fixed efficiency, and number of actions they perform each turn can be configured.
+                """
+            )
             with gr.Row():
                 with gr.Column():
                     gr.Markdown(simulation_labels.TITLE)
-
+                    inputs.append(
+                        gr.Text(
+                            value=self.default_output_path,
+                            label=simulation_labels.OUTPUT_PATH,
+                        )
+                    )
+                    steps_slider_component = gr.Slider(
+                        minimum=5,
+                        maximum=200,
+                        step=1,
+                        value=100,
+                        label=simulation_labels.STEPS,
+                    )
                     inputs.extend(
                         [
-                            gr.Text(
-                                value=self.default_output_path,
-                                label=simulation_labels.OUTPUT_PATH,
-                            ),
-                            gr.Slider(
-                                minimum=5,
-                                maximum=200,
-                                step=1,
-                                value=100,
-                                label=simulation_labels.STEPS,
-                            ),
+                            steps_slider_component,
                             gr.Slider(
                                 minimum=0.0,
                                 maximum=1.0,
@@ -56,34 +57,168 @@ class WithAgentTab(ITabBuilder):
                         ]
                     )
 
-            with gr.Row():
-                with gr.Accordion("Agents"):
-                    gr.Markdown("In here")
-                    inputs.extend(
-                        [
-                            gr.Checkboxgroup(
-                                choices=[
-                                    ("Dummy", "dummy"),
-                                    ("Distancing", "distancing"),
-                                    ("Vaccination", "vaccination"),
-                                    ("Treatment", "treatment"),
-                                    ("Masking", "masking"),
-                                ],
-                                value=["distancing"],
-                            ),
-                            gr.Slider(
-                                minimum=0,
-                                maximum=20,
-                                step=1,
-                                value=6,
-                                label="Total number of actions allowed per turn",
-                                interactive=True,
-                            ),
-                        ]
-                    )
+                    with gr.Accordion("Agents"):
+                        agents_selection_component = gr.Checkboxgroup(
+                            choices=[
+                                ("Distancing", "distancing"),
+                                ("Vaccination", "vaccination"),
+                                ("Treatment", "treatment"),
+                                ("Masking", "masking"),
+                            ],
+                            value=["distancing"],
+                            label="Enabled agents",
+                        )
+                        inputs.extend(
+                            [
+                                agents_selection_component,
+                                gr.Slider(
+                                    minimum=0,
+                                    maximum=20,
+                                    step=1,
+                                    value=6,
+                                    label="Total number of actions allowed per turn",
+                                    info="Number of actions per turn, split equally between active agents (rounded down where required.)",
+                                    interactive=True,
+                                ),
+                            ]
+                        )
+
+                        with gr.Group():
+                            gr.Markdown(
+                                """
+                                **Distancing agent params**
+
+                                The distancing agent disconnects nodes wth 95% efficiency."""
+                            )
+                            inputs.extend(
+                                [
+                                    gr.Slider(
+                                        label="Start isolating on step",
+                                        minimum=steps_slider_component.minimum,
+                                        value=5,
+                                        maximum=steps_slider_component.maximum,
+                                        interactive=True,
+                                    ),
+                                    gr.Slider(
+                                        label="Stop isolating on step",
+                                        minimum=steps_slider_component.minimum,
+                                        value=50,
+                                        maximum=steps_slider_component.maximum,
+                                        interactive=True,
+                                    ),
+                                    gr.Slider(
+                                        label="Start reconnecting on step",
+                                        minimum=steps_slider_component.minimum,
+                                        value=40,
+                                        maximum=steps_slider_component.maximum,
+                                        interactive=True,
+                                    ),
+                                    gr.Slider(
+                                        label="Stop reconnecting on step",
+                                        minimum=steps_slider_component.minimum,
+                                        value=steps_slider_component.maximum,
+                                        maximum=steps_slider_component.maximum,
+                                        interactive=True,
+                                    ),
+                                ]
+                            )
+
+                        with gr.Group():
+                            gr.Markdown(
+                                """
+                                **Vaccination agent params**
 
-            outputs = [gr.Image(format=self.output_format)]
-            button = gr.Button(self.button_label)
-            button.click(self.f, inputs=inputs, outputs=outputs)
+                                The vaccination agent adds 95% immunity to targeted uninfected nodes."""
+                            )
+                            inputs.extend(
+                                [
+                                    gr.Slider(
+                                        label="Start vaccinating on step",
+                                        minimum=steps_slider_component.minimum,
+                                        value=5,
+                                        maximum=steps_slider_component.maximum,
+                                        interactive=True,
+                                    ),
+                                    gr.Slider(
+                                        label="Stop vaccinating on step",
+                                        minimum=steps_slider_component.minimum,
+                                        value=50,
+                                        maximum=steps_slider_component.maximum,
+                                        interactive=True,
+                                    ),
+                                ]
+                            )
+
+                        with gr.Group():
+                            gr.Markdown(
+                                """
+                                **Vaccination agent params**
+
+                                 "The treatment agent treats infected nodes with 60% immediate recovery,
+                                 and a boost to recovery rate otherwise."""
+                            )
+                            inputs.extend(
+                                [
+                                    gr.Slider(
+                                        label="Start treating on step",
+                                        minimum=steps_slider_component.minimum,
+                                        value=5,
+                                        maximum=steps_slider_component.maximum,
+                                        interactive=True,
+                                    ),
+                                    gr.Slider(
+                                        label="Stop treating on step",
+                                        minimum=steps_slider_component.minimum,
+                                        value=50,
+                                        maximum=steps_slider_component.maximum,
+                                        interactive=True,
+                                    ),
+                                ]
+                            )
+
+                        with gr.Group():
+                            gr.Markdown(
+                                """
+                                **Masking agent params**
+
+                                The masking agent gives nodes masks (indicated by circles in plot),
+                                which reduces disease spread to/from the node by 25%."""
+                            )
+                            inputs.extend(
+                                [
+                                    gr.Slider(
+                                        label="Start masking on step",
+                                        minimum=steps_slider_component.minimum,
+                                        value=5,
+                                        maximum=steps_slider_component.maximum,
+                                        interactive=True,
+                                    ),
+                                    gr.Slider(
+                                        label="Stop masking on step",
+                                        minimum=steps_slider_component.minimum,
+                                        value=50,
+                                        maximum=steps_slider_component.maximum,
+                                        interactive=True,
+                                    ),
+                                    gr.Slider(
+                                        label="Start removing masks on step",
+                                        minimum=steps_slider_component.minimum,
+                                        value=40,
+                                        maximum=steps_slider_component.maximum,
+                                        interactive=True,
+                                    ),
+                                    gr.Slider(
+                                        label="Stop removing masks on step",
+                                        minimum=steps_slider_component.minimum,
+                                        value=steps_slider_component.maximum,
+                                        maximum=steps_slider_component.maximum,
+                                        interactive=True,
+                                    ),
+                                ]
+                            )
+
+                with gr.Column():
+                    outputs = [gr.Image(format=self.output_format)]
 
-        return inputs, outputs, button
+                    button = gr.Button(self.button_label)
+                    button.click(self.f, inputs=inputs, outputs=outputs)