how to calculate manually propensity score weights for multinomial treatments where one of them is baseline I want to get intuition into the calculation of propensity scores (PS) and inverse probability of treatment weights (IPTW) for a multinomial treatment using multinomial regression. One of the treatments is the baseline treatment.
I am aware that powerful packages automate this but for my analyses, I need to be able to create my own function that allows me to:

*

*calculate manually the PS using a multinomial
regression model and

*calculate manually the IPTW from the estimated PS using the
relevant formulas.

Here I use the data from here where each of the three categories of the variable group represents a treatment, and the group == 1 represents the baseline treatment (control group) to which each of the other treatments are compared. The approach discussed there gives several PS per row, which is not what I want - all I need is only one PS and one IPTW per row as I get when I use twang and any other similar package.
data
library("survival")
require("survival")
library("nnet")
require("nnet")

set.seed(42)
days <- rpois(100, 3)
group <- sample(c(1,2,3), 100, replace=TRUE)
status <- rbinom(100,1,0.65)
demo1 <- rnorm(100,100,25)
demo2 <- rpois(100,10)
demo3 <- rbinom(100,1,0.67)

df <- data.frame(days, status, group, demo1, demo2,demo3)

Thanks in advance for any help.
 A: For the ATE, the IPW is the inverse of the model-predicted probability of being in the treatment actually received. Using multinom() in nnet, we can generate a multinomial regression model and then extract the predicted probabilities for each treatment group for each individual using predict(). Then we find the predicted probability for each individual that corresponds to the treatment actually received, and take its inverse. Finally, I use bal.tab() in cobalt to assess balance, demonstrating that balance has been achieved.
library("survival")
library("nnet")
library("cobalt")

set.seed(42)
days <- rpois(100, 3)
group <- factor(sample(c(1,2,3), 100, #needs to be a factor
                       replace=TRUE))
status <- rbinom(100,1,0.65)
demo1 <- rnorm(100,100,25)
demo2 <- rpois(100,10)
demo3 <- rbinom(100,1,0.67)

df <- data.frame(days, status, group, demo1, demo2,demo3)

fit <- multinom(group ~ demo1 + demo2 + demo3, data = df)
#> # weights:  15 (8 variable)
#> initial  value 109.861229 
#> iter  10 value 107.354113
#> final  value 107.353911 
#> converged

ps.mat <- predict(fit, type = 'probs')

w <- rep(0, nrow(df)) #inititalize weights

for (i in levels(group)) {
    w[group == i] <- 1/ps.mat[group == i, i]
}

bal.tab(group ~ demo1 + demo2 + demo3, data = df, 
        weights = w, un = TRUE)
#> Assuming "weighting". If not, specify with an argument to method.
#> Note: estimand and s.d.denom not specified; assuming ATE and pooled.
#> Balance summary across all treatment pairs
#>          Type Max.Diff.Un Max.Diff.Adj
#> demo1 Contin.      0.2034       0.0153
#> demo2 Contin.      0.3131       0.2557
#> demo3  Binary      0.0753       0.0306
#> 
#> Effective sample sizes
#>                 1      2      3
#> Unadjusted 31.000 30.000 39.000
#> Adjusted   30.736 25.553 37.823

Created on 2019-01-23 by the reprex package (v0.2.1)
For ATT weights, you simply multiply everyone's ATE weight by their probability of being in the focal (i.e., treated) group. Those in the focal group will now have a weight of 1, just like in the binary case. For example, if the focal group was group 2, you would simply do the following after running the above code:
w_att <- w*ps.mat[,"2"]

This multiplies each unit's weight by the probability of being in the focal group. For those in group 2, their original weight was equal to 1/ps.mat[,"2"] so multiplying by ps.mat[,"2"] will yield weights of 1 for those in group 2.
