Mixed-Type Trait Imputation • pigauto

Overview

Real comparative datasets contain many kinds of traits: body mass (continuous), clutch size (count), migratory status (yes/no), diet type (carnivore/herbivore/omnivore), threat status (LC < VU < EN). pigauto handles all of them in a single model — you do not need separate imputation runs for different column types.

Type	R class	Example	Notes
Continuous	`numeric`	Body mass, wing length	Auto-detected
Count	`integer`	Clutch size, litter size	Auto-detected
Binary	`factor` (2 levels)	Migratory yes/no	Auto-detected
Categorical	`factor` (>2 levels)	Diet type, lifestyle	Auto-detected
Ordinal	`ordered`	Threat status (LC < VU < EN)	Auto-detected
Proportion	`numeric` (0–1)	Habitat cover, diet fraction	Requires `trait_types = "proportion"` override
ZI count	`integer` (zero-inflated)	Parasite load, rare behaviour	Requires `trait_types = "zi_count"` override; experimental — accuracy more variable than other types
Multi-proportion	K `numeric` columns summing to 1	Diet composition, plumage-colour fractions, microbiome relative abundances	Requires `multi_proportion_groups = list(<name> = c("col1", ..., "colK"))`; encoded via centred log-ratio (CLR) + per-component z-score

The first five rows in this table are auto-detected from R column class; proportion, zi_count, and multi_proportion must be declared explicitly (trait_types or multi_proportion_groups). All eight share the same latent space — the phylogenetic baseline and GNN correction both operate in this space, and type-specific logic appears only at encoding, loss computation, and decoding.

Synthetic example

library(pigauto)
library(ape)

set.seed(42)
n <- 60
tree <- rtree(n)

traits <- data.frame(
  row.names = tree$tip.label,
  mass      = exp(rnorm(n, 3, 0.5)),
  clutch    = as.integer(rpois(n, 3) + 1L),
  migr      = factor(sample(c("no", "yes"), n, replace = TRUE)),
  diet      = factor(sample(c("herb", "carn", "omni"), n, replace = TRUE)),
  threat    = ordered(sample(c("LC", "VU", "EN"), n, replace = TRUE),
                      levels = c("LC", "VU", "EN"))
)

Preprocessing

preprocess_traits() auto-detects column types from R classes:

pd <- preprocess_traits(traits, tree)
print(pd)
#> pigauto_data
#>   Species: 60 
#>   Traits:  5 
#>   Types:   binary=1, categorical=1, continuous=1, count=1, ordinal=1 
#>   Latent columns: 7

The trait_map records each trait’s type, levels, latent column indices, and normalisation parameters:

str(pd$trait_map, max.level = 1)
#> List of 5
#>  $ mass  :List of 9
#>  $ clutch:List of 9
#>  $ migr  :List of 8
#>  $ diet  :List of 8
#>  $ threat:List of 8
pd$trait_map$diet
#> $name
#> [1] "diet"
#> 
#> $type
#> [1] "categorical"
#> 
#> $n_latent
#> [1] 3
#> 
#> $latent_cols
#> [1] 4 5 6
#> 
#> $levels
#> [1] "carn" "herb" "omni"
#> 
#> $log_transform
#> [1] FALSE
#> 
#> $mean
#> [1] NA
#> 
#> $sd
#> [1] NA

Creating splits

When trait_map is supplied, make_missing_splits() operates at the original-trait level. For categorical traits, all K one-hot columns are held out together:

spl <- make_missing_splits(pd$X_scaled, missing_frac = 0.20,
                           seed = 1, trait_map = pd$trait_map)
cat("Val cells (latent):", length(spl$val_idx), "\n")
#> Val cells (latent): 21
cat("Test cells (latent):", length(spl$test_idx), "\n")
#> Test cells (latent): 57

Baseline fitting

The baseline uses phylogenetic conditional MVN for continuous-family latent columns, and label-propagation or threshold/liability candidates for discrete-family columns:

bl <- fit_baseline(pd, tree, splits = spl)
dim(bl$mu)
#> [1] 60  7

Training

fit_pigauto() uses type-specific losses and trait-level corruption masking automatically when a trait_map is present:

fit <- fit_pigauto(
  pd, tree,
  splits = spl,
  epochs = 200L,
  eval_every = 50L,
  patience = 5L,
  verbose = FALSE,
  seed = 1
)
#> Warning: phylo_signal_gate requires the 'phytools' package; returning NA for
#> all traits.
#> Warning: Small validation set for 5 trait(s): mass (n=2), clutch (n=4), migr
#> (n=3), diet (n=3), threat (n=3). Calibrated gate and conformal scores will be
#> noisy; coverage may deviate from the 95%% target. See `?fit_pigauto` under
#> 'Calibration at small n' for smoothing options.
print(fit)
#> pigauto_fit
#>   Species : 60 
#>   Traits  : 5 -- mass, clutch, migr, diet, threat 
#>   Types   : binary=1, categorical=1, continuous=1, count=1, ordinal=1 
#>   Architecture: hidden_dim = 64 | k_eigen = 4 
#>   Best val loss : 1.3449 
#>   Test loss     : 1.2059 
#>   Gate calibration: yes
#>   Conformal scores: 3 traits

Prediction and decoding

predict() decodes latent predictions back to original types:

pred <- predict(fit, return_se = TRUE)
head(pred$imputed)
#>       mass clutch migr diet threat
#> t58 17.909      4  yes carn     VU
#> t8  17.909      4  yes carn     VU
#> t36 17.909      4  yes carn     VU
#> t4  17.909      4  yes carn     VU
#> t22 17.909      4   no carn     VU
#> t18 17.909      4   no carn     VU

For binary and categorical traits, class probabilities are available:

# Binary: probability of "yes"
head(pred$probabilities$migr)
#> [1] 0.5179941 0.7875313 0.7875313 0.7875313 0.2124687 0.2124687

# Categorical: probability of each diet class
head(pred$probabilities$diet)
#>          carn      herb      omni
#> t58 0.4558713 0.2464169 0.2977118
#> t8  0.4558713 0.2464169 0.2977118
#> t36 0.4558713 0.2464169 0.2977118
#> t4  0.4558713 0.2464169 0.2977118
#> t22 0.4558713 0.2464169 0.2977118
#> t18 0.4558713 0.2464169 0.2977118

The SE matrix provides type-appropriate uncertainty:

head(pred$se)
#>     mass clutch      migr      diet    threat
#> t58    0      0 0.4820059 0.5441287 0.4998422
#> t8     0      0 0.2124687 0.5441287 0.4735485
#> t36    0      0 0.2124687 0.5441287 0.4902633
#> t4     0      0 0.2124687 0.5441287 0.4746017
#> t22    0      0 0.2124687 0.5441287 0.4892949
#> t18    0      0 0.2124687 0.5441287 0.4497299

Evaluation

evaluate_imputation() dispatches type-specific metrics:

ev <- evaluate_imputation(pred, pd$X_scaled, spl)
#> Warning in stats::cor(t_j[ok], p_j[ok]): the standard deviation is zero
#> Warning in stats::cor(t_j[ok], p_j[ok]): the standard deviation is zero
#> Warning in stats::cor(t_j[ok], p_j[ok]): the standard deviation is zero
#> Warning in stats::cor(t_j[ok], p_j[ok]): the standard deviation is zero
print(ev)
#>    split  trait        type  n      rmse pearson_r coverage_95       mae
#> 1    val   mass  continuous  2 0.6996585        NA           1        NA
#> 2    val clutch       count  4 1.0938335        NA          NA 0.9784612
#> 3    val   migr      binary  3        NA        NA          NA        NA
#> 4    val   diet categorical  3        NA        NA          NA        NA
#> 5    val threat     ordinal  3 0.6551186        NA          NA        NA
#> 6   test   mass  continuous 13 0.9332707        NA           1        NA
#> 7   test clutch       count 10 0.7248662        NA          NA 0.6104951
#> 8   test   migr      binary 10        NA        NA          NA        NA
#> 9   test   diet categorical  6        NA        NA          NA        NA
#> 10  test threat     ordinal  6 1.2144578        NA          NA        NA
#>    spearman_rho  accuracy     brier zero_accuracy aitchison rmse_clr
#> 1            NA        NA        NA            NA        NA       NA
#> 2            NA        NA        NA            NA        NA       NA
#> 3            NA 0.6666667 0.2169137            NA        NA       NA
#> 4            NA 0.3333333        NA            NA        NA       NA
#> 5    -0.8660254        NA        NA            NA        NA       NA
#> 6            NA        NA        NA            NA        NA       NA
#> 7            NA        NA        NA            NA        NA       NA
#> 8            NA 0.6000000 0.2472688            NA        NA       NA
#> 9            NA 0.1666667        NA            NA        NA       NA
#> 10   -0.4011887        NA        NA            NA        NA       NA
#>    simplex_mae
#> 1           NA
#> 2           NA
#> 3           NA
#> 4           NA
#> 5           NA
#> 6           NA
#> 7           NA
#> 8           NA
#> 9           NA
#> 10          NA

Experimental analysis-aware multiple imputation

multi_impute_analysis() requires the downstream formula before imputation. The initial backend supports exactly one incomplete continuous covariate under MAR. It does not yet impute the mixed discrete and compositional outcomes shown earlier in this vignette for downstream inference.

# Precompute nonlinear auxiliaries explicitly when required.
analysis_data$z_sq <- analysis_data$z^2
mi <- multi_impute_analysis(
  data = analysis_data, formula = y ~ x + z,
  missing = "x", model = "glm", m = 50L,
  auxiliary = "z_sq", seed = 1L
)

# Fit a downstream model to each
fits <- with_imputations(mi, function(d) {
  glm(y ~ x + z, data = d, family = binomial)
})

# Pool with Rubin's rules
pool_mi(fits)

pool_mi() returns a tidy data.frame with estimate, std.error, p.value, df (classical Rubin degrees of freedom by default, or Barnard-Rubin when df_fun supplies complete-data residual degrees of freedom), fmi (fraction of missing information), and riv (relative increase in variance) per coefficient.

The "lm" branch uses proper Bayesian normal-regression MI, "glm" uses smcfcs for a binomial-logit model, and "lmer" uses jomo.smc for one Gaussian random intercept. Fixed effects only: multiple incomplete columns, missing discrete covariates, MNAR, interactions/nonlinear terms involving the incomplete covariate, random slopes, variance components, correlations, BLUPs/conditional modes, and latent loadings are unsupported.

For prediction diagnostics, the lower-level predict(fit, n_imputations = 5L) interface returns stochastic datasets. These are not supported for downstream inference.