---
title: Simulation and estimation of partition models - Example
  script with random data
author: "Marion Hoffman"
date: "28/09/2020"
output:
  pdf_document: default
  html_document: default
---

```{r setup, include=FALSE}
knitr::opts_chunk$set(echo = TRUE)
```

In this example script, we use the Exponential Random Partition Model to simulate cross-sectional partitions and estimate the model for one observation (here, synthetic data). More information can be found in the manual of this package present on the github repository stocnet/erpm or by looking up the documentation of specific functions. 

This script can be used as a starting point for anyone wishing to use the model on cross-sectional observations of partitions, like teams, interaction groups, political parties, animal herds, ... 

## 0. Load dependencies

The following packages and scripts are required to fully run this example script. 

```{r loads, message=FALSE, warning=FALSE, results='hide'}
library(ERPM)

library(numbers) # for calculations of Bell numbers
library(gmp) # for calculations of Stirling numbers
library(mclust) # for calculations of Rand indexes

library(ggplot2) # for plots
theme_set(theme_minimal())
```

## 1. Make some general calculations

### 1.1 Calculate the size of the space of partitions

For reasonable values of n, we can calculate exactly the total number of possible partitions (this is the Bell number), for example with n = 6,

```{r}
n <- 6
bell(n)
```

or the number of partitions with k = 2 groups for example,

```{r}
k <- 2
Stirling2(n,k)
```

or the number of partitions with groups of size comprised bewteen two values size_min and size_max,

```{r}
size_min <- 2
size_max <- 4
Bell_constraints(n,size_min,size_max)
```

or the number of partitions with k = 2 groups and groups of size comprised bewteen two values size_min and size_max.

```{r}
Stirling2_constraints(n,k,size_min,size_max)
```


### 1.2 Calculate the expected size of a random partition

We can also compute the average size (i.e., the number of groups) of a partition (under a null model),

```{r}
compute_averagesize(n)
```

### 1.3 Enumerate all partitions

For a low number of nodes (below 10 for example), one could enumerate all possible partitions,

```{r echo=TRUE}
all_partitions <- find_all_partitions(n)
```

and look at the number of partitions with a certain size structure (for example, how many partitions have 2 groups of 3, or 6 groups of 1?).

```{r}
counts_partition_classes <- count_classes(all_partitions)
```

### 1.4 Calculate a distance measure between two partitions

One can use the Rand distance to evaluate the distance between two partitions, for example:

```{r include=TRUE, echo=FALSE}
p1 <- all_partitions[1,]
p2 <- all_partitions[2,]
rand_12 <- adjustedRandIndex(p1,p2)
```



## 2. Simulate partitions

### 2.1 Define nodesets and attributes

Here we define an arbitrary set of n = 6 nodes with attributes, and an arbitrary covariate matrix.

```{r}
n <- 6 
nodes <- data.frame(label = c("A","B","C","D","E","F"),
                    gender = c(1,1,2,1,2,2),
                    age = c(20,22,25,30,30,31)) 
friendship <- matrix(c(0, 1, 1, 1, 0, 0,
                       1, 0, 0, 0, 1, 0,
                       1, 0, 0, 0, 1, 0,
                       1, 0, 0, 0, 0, 0,
                       0, 1, 1, 0, 0, 1,
                       0, 0, 0, 0, 1, 0), 6, 6, TRUE) 
```

### 2.2 Define a model and simulate

First, we need to choose the effects we want to include (see manual for all effect names). For example we set four (which is of course not reasonable for 6 nodes):

```{r}
effects <- list(names = c("num_groups","same","diff","tie"),
                   objects = c("partition","gender","age","friendship"))
objects <- list()
objects[[1]] <- list(name = "friendship", object = friendship)
```

and we can set parameter values for each of these effects.

```{r}
parameters <- c(-0.2, 0.2, -0.1, 0.5)
```


Now we can generate our simulated sample, by setting the desired additional parameters for the Metropolis sampler and choosing a starting point for the chain (first.partition):

```{r}
nsteps <- 100
sample <- draw_Metropolis_single(theta = parameters, 
                          first.partition = c(1,1,2,2,3,3), 
                          nodes = nodes, 
                          effects = effects, 
                          objects = objects, 
                          burnin = 100, 
                          thining = 10, 
                          num.steps = nsteps, 
                          neighborhood = c(0,1,0), 
                          numgroups.allowed = 1:n,
                          numgroups.simulated = 1:n,
                          sizes.allowed = 1:n,
                          sizes.simulated = 1:n,
                          return.all.partitions = TRUE)
```

### 2.3 Trace plots

We can check the mixing of the chain with autocorrelations (they should be below 0.4). For example, for the first statistic:

```{r}
s <- 1
cor(sample$draws[1:(nsteps-1),s], sample$draws[2:nsteps,s])
```


and with the trace plots;

```{r}
s <- 1
ggplot(data = data.frame(step = 1:nsteps,
                         stat = sample$draws[,s])) +
  geom_line(aes(x=step,y=stat)) 
```


## 3. Estimate for an observed partition

### 3.1 Define the observation

```{r}
partition <- c(1,1,2,2,2,3)
```

### 3.2 Estimate

The average number of groups expected at random is lower than 3 (see section 1), so let's set an initial estimate for the number of  groups parameter negative, and leave the others to zero.
The burnin and thining are chosen from trace plots to have a good mixing.
For phase 1, we don't need a large sample, because it is just for the scaling matrix which is not very sensitive to the number of samples.
For phase 2, we have to define a number of subphases and the number of iterations in the subphases. There will be (2.52)^k iterations in each subphase k, multiplied by a constant multiplication factor, that we need to choose big enough so that statistics are crossed. When approaching the correct estimates, we can increasingly reduce the number of steps of phase 2.
For phase 3, we need a lot of samples to have correct estimates of the final distribution. When we approach the right estimates, it is the most important phase, and its results should be re-used to restart the estimation from a good starting point.
The multiplication factor and gain factor were chosen as in the ERGM book, but they can be adjusted by looking at the evolution of parameters in phase 2.

```{r}
startingestimates <- c(-2,0,0,0)
estimation <- estimate_ERPM(partition, 
                          nodes, 
                          objects, 
                          effects, 
                          startingestimates = startingestimates, 
                          burnin = 100, 
                          thining = 20,
                          length.p1 = 500, # number of samples in phase 1
                          multiplication.iter.p2 = 20,  # multiplication factor for the number of iteration in phase 2 subphases 
                          num.steps.p2 = 4, # number of phase 2 subphases
                          length.p3 = 1000) # number of samples in phase 3

# get results table
estimation
```

We get some intermediary objects calculated in the different phases

```{r}
objects.phase1 <- estimation$objects.phase1
objects.phase2 <- estimation$objects.phase2
objects.phase3 <- estimation$objects.phase3
```


The convergence should be as small as possible, below 0.1 for example. We can retry to estimate the model starting from the result of phase 3 and using the covariance calculated in phase 3 (then phase 1 is skipped). It can be good to reduce the gain factor, and make the subphases longer, or lengthen phase 3:


```{r}
startingestimates <- estimation$results$est
startingcovariance <- objects.phase3$inv.zcov
startingscaling <- objects.phase3$inv.scaling
estimation <- estimate_ERPM(partition, 
                          nodes, 
                          objects, 
                          effects, 
                          startingestimates = startingestimates, 
                          gainfactor = 0.05, 
                          burnin = 500, 
                          thining = 20,
                          length.p1 = 100, 
                          multiplication.iter.p2 = 50, 
                          num.steps.p2 = 3, 
                          length.p3 = 2000,
                          inv.zcov = startingcovariance,
                          inv.scaling = startingscaling)

# get results table
estimation
```

Convergence improved!

### 3.3 Estimate plots 

To assess convergence problems, we can have a look at how estimates evolve during phase 2.


```{r}
ggplot(data = data.frame(step = 1:nrow(objects.phase2$estimates),
                         stat = array(objects.phase2$estimates),
                         parameter = as.character(rep(seq(1,ncol(objects.phase2$estimates)), each=nrow(objects.phase2$estimates))))) +
  geom_line(aes(x=step,y=stat,colour=parameter)) 
```



### 3.4 Goodness of fit

We can check how the model reproduces statistics of the observed data. First we simulate the estimated model (with the option of returning all partitions!):

```{r}
nsimulations <- 1000
simulations <- draw_Metropolis_single(theta = estimation$results$est, 
                          first.partition = partition, 
                          nodes = nodes, 
                          effects = effects, 
                          objects = objects, 
                          burnin = 100, 
                          thining = 20, 
                          num.steps = nsimulations, 
                          neighborhood = c(0,1,0), 
                          sizes.allowed = 1:n,
                          sizes.simulated = 1:n,
                          return.all.partitions = TRUE)
```

We can first check the group size distribution:

```{r}
observedsizes <- rep(0,n)
for(size in 1:n){
  observedsizes[size] <- length(which(table(partition)==size))
}

allsizes <- matrix(0,nrow=nsimulations,ncol=n)
meansizes <- matrix(0,n)
for(size in 1:n){
  for(simu in 1:nsimulations){
    allsizes[simu,size] <- length(which(table(simulations$all.partitions[simu,])==size))
  }
  meansizes[size] <- mean(allsizes[,size])
}

df <- data.frame(simulation = 1:nsimulations,
                 simulated_stat = array(allsizes),
                 size = as.character(rep(seq(1,n),each=nsimulations)),
                 observed_stat = rep(observedsizes,each=nsimulations),
                 mean_stat = rep(meansizes,each=nsimulations))
ggplot(df, aes(factor(size), simulated_stat)) + 
    geom_violin() +
    geom_point(aes(x=factor(size),y=mean_stat, colour = "simulated", shape= "simulated")) +
    geom_point(aes(x=factor(size),y=observed_stat, colour = "observed", shape= "observed"), stroke= 1.5) +
    labs(x = "Group size",
         y = "Group count",
         color="Legend",
         shape="Legend") +
    scale_color_manual(values = c(simulated="black",observed="red"), labels=c("observed value","simulated average"))+
    scale_shape_manual(values = c(simulated=19,observed=5), labels=c("observed value","simulated average")) +
    theme(legend.position = "right")

```

For continuous variables, we can check the intra-class correlation of the variable (linked to the statistic used in the model, but not equal):

```{r}
allicc <- rep(0,nsimulations)
for(simu in 1:nsimulations){
  allicc[simu] <- icc(simulations$all.partitions[simu,],nodes$age)
}

df <- data.frame(simulation = 1:nsimulations,
                 simulated_stat = allicc,
                 observed_stat = rep(icc(partition,nodes$age),nsimulations),
                 mean_stat = rep(mean(allicc,na.rm=TRUE),nsimulations))
ggplot(df) + 
    geom_violin(aes(x=1,y=simulated_stat), fill = NA) +
    geom_point(aes(x=1, y=mean_stat, colour = "simulated", shape= "simulated"), stroke= 1.5) +
    geom_point(aes(x=1, y=observed_stat, colour = "observed", shape= "observed"), stroke= 1.5) +
    labs(x = "",
         y = "Intraclass correlation coefficient (intragroup)",
         color="",
         shape="",
         linetype="") +
    scale_color_manual(values = c(observed="red", simulated="black"), labels=c("Observed value","Simulated average"))+
    scale_shape_manual(values = c(observed=5, simulated=16), labels=c("Observed value","Simulated average")) +
    theme(legend.position = "right",
          legend.key.height = unit(0.4,"in"),
          text = element_text(size=10))

```

We can also check the density of pairs with same gender in the groups (linked to the statistic used in the model, but not equal):

```{r}
alldensities <- rep(0,nsimulations)
for(simu in 1:nsimulations){
  alldensities[simu] <- same_pairs(simulations$all.partitions[simu,], nodes$gender, 'avg_pergroup')
}

df <- data.frame(simulation = 1:nsimulations,
                 simulated_stat = alldensities,
                 observed_stat = rep(same_pairs(partition, nodes$gender, 'avg_pergroup'),nsimulations),
                 mean_stat = rep(mean(alldensities,na.rm=TRUE),nsimulations))
ggplot(df) + 
    geom_violin(aes(x=1,y=simulated_stat), fill = NA) +
    geom_point(aes(x=1, y=mean_stat, colour = "simulated", shape= "simulated"), stroke= 1.5) +
    geom_point(aes(x=1, y=observed_stat, colour = "observed", shape= "observed"), stroke= 1.5) +
    labs(x = "",
         y = "Intraclass correlation coefficient (intragroup)",
         color="",
         shape="",
         linetype="") +
    scale_color_manual(values = c(observed="red", simulated="black"), labels=c("Observed value","Simulated average"))+
    scale_shape_manual(values = c(observed=5, simulated=16), labels=c("Observed value","Simulated average")) +
    theme(legend.position = "right",
          legend.key.height = unit(0.4,"in"),
          text = element_text(size=10))

```

And it is also good to check other statistics even less linked to the model, for example on attributes that are not related to any effect in the model.

### 3.5 Estimated log-likelihood and AIC

Finally, we can estimate the log-likelihood and AIC of the model (useful to compare two models for example). First we need to estimate the ML estimates of a simple model with only one parameter for number of groups (this parameter should be in the model!).

```{r}
likelihood_function <- function(x){ exp(x*max(partition)) / compute_numgroups_denominator(n,x)}
curve(likelihood_function, from=-2, to=0)
parameter_base <- optimize(likelihood_function, interval=c(-2, 0), maximum=TRUE)
parameters_basemodel <- c(parameter_base$maximum,0,0,0)
```


Then we can get our estimated logL and AIC.

```{r}
logL_AIC <- estimate_logL(partition,
                         nodes,
                         effects, 
                         objects,
                         theta = estimation$results$est,
                         theta_0 = parameters_basemodel,
                         M = 3,
                         num.steps = 200,
                         burnin = 100,
                         thining = 20)
logL_AIC$logL
logL_AIC$AIC
```

To check that there were enough steps in the algorithm (M), we can plot the statistics distributions of the intermediate models sampled and see whether they overlap (if not, we need more steps). For example for the stat about age:

```{r}
stat <- 3
ggplot(data =  data.frame(stat = c(logL_AIC$draws[[1]]$draws[,stat], logL_AIC$draws[[2]]$draws[,stat], logL_AIC$draws[[3]]$draws[,stat]),
             draw = c(rep("1",100), rep("2",100), rep("3",100)))) +
    geom_density(aes(x=stat, fill=draw),alpha=0.2) +
    scale_fill_brewer(palette = "Spectral") +
    theme_minimal()
```



NOTE: with larger data, things can be parallelized. The only available "automatic" way to do it now is to specify the options "cpus" in the estimation function. In that case, phase 1 and 3 samples are parallelized.

