---
title: Post-process data in a trajectory
author:
  Stefan Widgren
  <a href="https://orcid.org/0000-0001-5745-2284">
    <img src="https://info.orcid.org/wp-content/uploads/2019/11/orcid_16x16.png"
         alt="ORCID logo"
         width="16"
         height="16"
         style="border-style:none;" />
  </a>
output:
  html_vignette:
    toc: true
    toc_depth: 3
vignette: >
  %\VignetteIndexEntry{Post-process data in a trajectory}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

After a model is created, a simulation is started with a call to the
`run()` function. The function returns a modified model object
containing a single stochastic solution trajectory. This trajectory
includes the state of each compartment recorded at every time-point
specified in `tspan`.

This vignette introduces the functionality in `SimInf` to post-process
and explore this trajectory data.

Let us first load the SimInf package.

```{r}
library(SimInf)
```

## Extract trajectory data with `trajectory()`

Most modeling studies require custom data analysis beyond simple
plotting.  To support this, `SimInf` provides the `trajectory()`
method to extract the raw data as a `data.frame`. This is useful if
you need to:

- Perform custom statistical calculations (e.g., time to peak).
- Export data to CSV for use in other software.
- Combine results from multiple simulation runs.

Let's simulate 10 days of data from an SIR model with 6 nodes. For
reproducibility, we set the seed and specify the number of threads.

```{r}
set.seed(123)
set_num_threads(1)

u0 <- data.frame(
  S = c(100, 101, 102, 103, 104, 105),
  I = c(1, 2, 3, 4, 5, 6),
  R = c(0, 0, 0, 0, 0, 0)
)

model  <- SIR(
  u0 = u0,
  tspan = 1:10,
  beta = 0.16,
  gamma = 0.077
)

result <- run(model)
```

Extract the full trajectory data (all compartments, all nodes).

```{r}
trajectory(result)
```

Extract the number of recovered individuals (R) in the first node
only.

```{r}
trajectory(result, compartments = "R", index = 1)
```

Extract the number of recovered individuals in the first and third
nodes.

```{r}
trajectory(result, compartments = "R", index = c(1, 3))
```

## Calculate prevalence from a trajectory using `prevalence()`

The `prevalence()` function calculates the proportion of individuals
with the disease.  It takes a model object and a formula:

- **Left-hand side (LHS):** Compartments representing "cases" (e.g.,
    I).
- **Right-hand side (RHS):** Compartments representing the "at-risk"
  population (e.g., S + I + R).

The function also supports a `level` argument to change the
aggregation level:

- `level = 1` (default): Prevalence aggregated over all nodes (global).
- `level = 2`: Proportion of nodes that have at least one case.
- `level = 3`: Prevalence calculated within each node (returns a
  matrix).

Let's determine the proportion of infected individuals in the total
population.

```{r}
prevalence(result, I ~ S + I + R)
```

Identical result is obtained with the shorthand `I ~ .` (where `.`
means "all compartments").

```{r}
prevalence(result, I ~ .)
```

Calculate the proportion of nodes that are infected (at least one I
individual).

```{r}
prevalence(result, I ~ S + I + R, level = 2)
```

Calculate the prevalence **within each node** individually.

```{r}
prevalence(result, I ~ S + I + R, level = 3)
```

## Visualize a trajectory with `plot()`

The `plot()` function provides a quick way to inspect the outcome. It
can display:

- The median and quantile range across all nodes.
- Individual trajectories for specific nodes.
- Prevalence curves.

*Note: Since the simulation is stochastic, the exact lines shown below
 will vary unless set.seed() is used.*

### Aggregated View (Median and Range)

Plot the median and interquartile range (IQR) of all compartments.

```{r, fig.width=7, fig.height=4, fig.align="left"}
plot(result)
```

Plot the median and the middle 95\% quantile range.

```{r, fig.width=7, fig.height=4, fig.align="left"}
plot(result, range = 0.95)
```

Plot only the infected individuals (I).

```{r, fig.width=7, fig.height=4, fig.align="left"}
plot(result, "I")
```

Use formula notation to plot the infected individuals.

```{r, fig.width=7, fig.height=4, fig.align="left"}
plot(result, ~I)
```

### Individual Node View

Plot the trajectories for the first three nodes. We use `range =
FALSE` to suppress the shaded median/range bands and show the
individual lines.

```{r, fig.width=7, fig.height=4, fig.align="left"}
plot(result, index = 1:3, range = FALSE)
```

Use `type = "l"` to draw a line.

```{r, fig.width=7, fig.height=4, fig.align="left"}
plot(result, index = 1:3, range = FALSE, type = "l")
```

Plot the infected individuals in the first node only.

```{r, fig.width=7, fig.height=4, fig.align="left"}
plot(result, "I", index = 1, range = FALSE)
```

### Prevalence Plots

Plot the proportion of infected individuals in the population.

```{r, fig.width=7, fig.height=4, fig.align="left"}
plot(result, I ~ S + I + R)
```

Plot the proportion of nodes with infected individuals (`level = 2`).

```{r, fig.width=7, fig.height=4, fig.align="left"}
plot(result, I ~ S + I + R, level = 2)
```

Plot the median and IQR of the prevalence **within in each node**
(`level = 3`).

```{r, fig.width=7, fig.height=4, fig.align="left"}
plot(result, I ~ S + I + R, level = 3)
```

Plot the prevalence in the first three nodes.

```{r, fig.width=7, fig.height=4, fig.align="left"}
plot(result, I ~ S + I + R, level = 3, index = 1:3, range = FALSE)
```

## Summary

- Use `trajectory()` to extract raw data for custom analysis.
- Use `prevalence()` to calculate disease proportions at different
  aggregation levels.
- Use `plot()` for quick visual inspection of medians, ranges, or
  individual trajectories.

To find more details on the plot method for `SimInf_model` objects,
run:

```{r, eval=FALSE}
help("plot,SimInf_model-method", package = "SimInf")
```