Plan for mixed binary-continuous multivariate outcomes

[Gattocrucco]

• Work out the math (@miaoqingyu2)
  • Adapt Talhouk et al. (2012) "Efficient Bayesian Inference..." to the mixed continuous-binary case
    • There are other methods for multivariate probit but this seems the best one
    • Hopefully it can be adapted to the mixed case
    • It's used in Goh et al. (2024) "Joint Model for Handling..." that adapt it to BART already, their code is https://github.com/yongchengoh/missBART.
      • => I was wrong, they don't adapt it to BART, they just use it for linear regression. So now we don't know if this algorithm can be adapted at all to BART, even before considering the mixed binary-continuous variant.
  • How to support missing outcome components
    • New: In the easy case with no correlations first (@Gattocrucco)
    • I guess this requires messing with the leaves, while the above just touches the outcome layer
    • Right now we need only missing continuous outcomes, if supporting missing binary outcomes turns out to be a problem, drop it
    • I haven't checked the literature
• Implementation
  • Bridge PR to change the interfaces (@Gattocrucco, Adapt interface for planned mixed-outcome implementation #98)
    • This PR does not add functionality; it just changes the interfaces to make them forward-compatible with the additions, such that other work can continue in parallel after this PR
    • This PR changes at once the low-level and high-level interfaces
    • Low-level interface in mcmcstep
      • Replace State.y with State.binary_y, because y is actually used only in the binary case
      • In the continuous case, State.binary_y is set to None
      • init() now always wants as input a float y, no boolean allowed. For binary, the convention is that y != 0 means True, as usual in Python and C
      • init() stores binary_y as booleans though, binary_y is always booleans because that's what's needed in step_z
      • To specify whether regression is binary or continuous, init() has a new parameter outcome_type: OutcomeType = 'continuous'
      • OutcomeType is an Enum with two string values 'continuous' and 'binary'
    • High-level interface Bart
      • type is replaced by outcome_type, the same format used in the low-level interface
      • y_train is always floats, like in the low-level interface
      • Drop the x_test parameter from __init__
      • Drop all prediction attributes: yhat_train, prob_test, etc.
      • The x_test argument of Bart.predict can be also be the string 'train' to use x_train
      • Bart.predict() as an additional parameter format: PredictFormat = 'mean'
      • PredictFormat is a string-valued Enum with values 'mean', 'mean_samples', 'outcome_samples', 'latent_samples'
      • Example to understand what these settings mean: in the binary case, 'latent_samples' will return ndpost x m samples of the sum-of-trees, 'mean_samples' squashes the sum-of-trees with probit, 'outcome_samples' samples binary outcomes based on the probability in 'mean_samples', and 'mean' is the (m,) vector of averages of mean_samples that gives the posterior mean conditional on a given X.
    • Compatibility interface mc_gbart
      • Adapt this internally to the changes in Bart
  • PR: multivariate binary in mcmcstep, without correlations (@Gattocrucco, Multivariate binary support in mcmcstep #99 )
    • Without correlations first because it's easier
    • For step_z, simply vmap over outcome component
  • PR: mixed binary-continuous in mcmcstep, without correlations (@Gattocrucco, Mixed binary-continuous outcomes w/o correlations, internals only #100)
    • Again without correlations first because it's easier
    • outcome_type in init() can also have type Sequence[OutcomeType], matching the number of outcomes in y
    • State.z and State.binary_y have only as many rows as the number of 'binary' outcome components, or are None if there are none
    • State has a new attribute binary_indices, as long as the number of binary components, with the indices of the binary components in the list of all components; this attribute is filled in by init() and used by step_z to update only those rows of resid that correspond to binary outcomes
    • init() enforces that error_cov_scale is diagonal with error_if
    • The rows/columns in error_cov_scale corresponding to binary outcomes are ignored
    • error_cov_inv is still stored as a matrix but it's always diagonal
    • step_error_cov_inv_uv is repeated for all continuous outcomes
  • PR: expose new functionality in the Bart interface (@Gattocrucco, Expose multivariate binary features in the high-level interface #103)
  • New: PR: missing outcome components, w/o correlations (@Gattocrucco, Missing outcome components #127)
  • PR: multivariate binary in mcmcstep, with correlations (@miaoqingyu2)
    • drop the support for diagonal-only altogether, I see no reason to keep it around
  • PR: multivariate binary-continuous in mcmcstep, with correlations (@miaoqingyu2)
    • merge this with the previous PR if it's convenient
  • PR: expose new functionality in the Bart interface (@Gattocrucco)
  • PR: partially missing outcomes in mcmcstep, with correlations (@miaoqingyu2)
  • PR: expose new functionality in the Bart interface (@Gattocrucco)

[Gattocrucco]

How to implement missing outcome components:

• for probit components, the latent z is already an imputation; on missing subunits, draw z from a full normal instead of a truncated normal just always use heteroskedasticity
• for normal components, use component-wise heteroskedasticity, setting the precision scale to 0; the MCMC always multiplies by or sums the precision scale, so this works out of the box