This is a draft PR to close #1551. I'm opening the PR as a draft to get feedback on the overall design and implementation before I work on polishing.

This PR adds the argument z_mix = FALSE to pp_mixture with the following effects:

• the default behavior (z_mix = FALSE) should be unchanged.
• For z_mix = TRUE:
  • If the model family is a zero-inflation or hurdle family, return the zero-inflation weights in the same format as the weights for mixture families. The first mixture component is the zero-inflation state; the final component is the non-inflated state, and (in the case of zero_one_inflated_beta) the middle component is the one-inflated state.
  • If the model is a mixture family that contains a zero-inflated or hurdle model as one of the mixture components, then the output is like the default output, but the array slice for the zero-inflated component is split into two array slices (three in the case of a zero-one-inflated beta, but this should be irrelevant unless incorrect mixture probabilities returned when mixtures include hurdle families #1552 is resolved in a way that admits the zero-one-inflated-beta in mixture families) following the format described in the previous bullet.
  • If the model family is not itself zero-inflated, and doesn't have any zero-inflated components, then this should error informatively.

TODO (I'll take care of these, but could use feedback from @paul-buerkner where tagged):

• Verify correctness of output and write tests.
• Feedback from @paul-buerkner on a good naming convention for the array slices corresponding to different zi/hu mixture components, and implement. Suggestion: P(K = 2 & z = 1 | Y), where z = 1 means a structural zero, z = 2 means no inflation, and z = 3 means a structural 1.
• Confirm that this works (or fails gracefully) when confronted with multivariate response models.
• Feedback from @paul-buerkner on whether to retain minor helper functions and where to put them.

A final design note: if preferred, it is possible to construct an implementation that is easier (for me) to grasp, but that supports only garden-variety zi/hu families and not mixture families that include zi/hu families as components of the mixture. It would be possible to insert this simpler implementation inside of pp_mixture() via the same z_mix argument. A functioning example of how this branch of the control flow inside of pp_mixture might look internally is available at #1551.

Here's a snippet to play with:

remotes::install_github("jsocolar/brms", ref = "zi-mixture-probs")
library(brms)
set.seed(2)
dat <- data.frame(
  y = c(rep(0, 40), rpois(30, 3), rnbinom(30, 4, .4)),
  x = rnorm(100)
)
mix <- mixture(poisson, zero_inflated_negbinomial)
fit1 <- brm(y ~ x, dat, family = mix, 
           backend = "cmdstanr", 
           cores = 4)

fit2 <- brm(y ~ 1, dat, family = zero_inflated_poisson, 
            backend = "cmdstanr", 
            cores = 4)

pp_mixture(fit1)
pp_mixture(fit1, z_mix = TRUE)
pp_mixture(fit2, z_mix = FALSE)
pp_mixture(fit2, z_mix = TRUE)

out1 <- pp_mixture(fit1, summary = FALSE)
out2 <- pp_mixture(fit1, summary = FALSE, z_mix = TRUE)
out3 <- pp_mixture(fit2, summary = FALSE, z_mix = TRUE)

all(round(apply(out1, c(1,2), sum), 10) == 1)
all(round(apply(out2, c(1,2), sum), 10) == 1)
all(round(apply(out3, c(1,2), sum), 10) == 1)