Impute missing phylogenetic traits (convenience wrapper)

One-call interface to the full pigauto pipeline: preprocessing, baseline fitting, GNN training, and prediction. For fine-grained control, use the individual functions (preprocess_traits, fit_baseline, fit_pigauto, etc.) directly.

Usage

impute(
  traits,
  tree,
  species_col = NULL,
  trait_types = NULL,
  multi_proportion_groups = NULL,
  log_transform = TRUE,
  missing_frac = 0.25,
  n_imputations = 1L,
  covariates = NULL,
  epochs = 2000L,
  verbose = TRUE,
  seed = NULL,
  multi_obs_aggregation = c("hard", "soft"),
  lambda_mode = c("estimate", "fixed_1", "cv", "bayes"),
  joint_solver = c("inhouse", "rphylopars"),
  predict_method = c("auto", "exact", "per_column"),
  joint_refine_iter = 0L,
  em_iterations = 0L,
  em_tol = 0.001,
  em_offdiag = FALSE,
  pool_method = c("median", "mean", "mode"),
  clamp_outliers = FALSE,
  clamp_factor = 5,
  match_observed = c("none", "pmm"),
  pmm_K = 5L,
  safety_floor = TRUE,
  phylo_signal_gate = TRUE,
  phylo_signal_threshold = 0.2,
  phylo_signal_method = "lambda",
  conformal_split_val = FALSE,
  gnn = FALSE,
  ...
)

Arguments

traits

data.frame with species as rownames and trait columns, or (when species_col is supplied) a data.frame with a species column that may have multiple rows per species. Supported column types: numeric (continuous), integer (count), factor (binary/categorical), ordered (ordinal), character (factor → binary/categorical), logical (binary). See the Trait type auto-detection section below.

tree

object of class "phylo".

species_col

character. Name of the column in traits that identifies species. When supplied, multiple observations per species are supported. Default NULL uses row names (one row per species).

trait_types

named character vector overriding the auto-detected type for specific trait columns, e.g. c(Survival = "proportion", Parasites = "zi_count"). Required for the two types that cannot be inferred from R class (see Trait type auto-detection below). Default NULL (auto).

multi_proportion_groups

named list declaring compositional (multi_proportion) traits, e.g. list(colour = c("black", "blue", "red", "yellow")). Each list element names a group and gives the K trait columns that form a simplex (rows summing to 1). Encoded via CLR + per-component z-score. Multi_proportion traits cannot be declared through trait_types — use this argument instead. Default NULL (no multi_proportion groups).

log_transform

logical. Auto-log positive continuous columns (default TRUE).

missing_frac

numeric. Fraction of observed cells held out for validation/test evaluation (default 0.25). Set to 0 to skip splitting (all cells used for training, no evaluation).

n_imputations

integer. Number of MC-dropout imputation sets (default 1). Values > 1 enable between-imputation uncertainty.

covariates

data.frame or matrix of environmental covariates (fully observed — no NAs). Covariates are conditioners: they inform imputation but are not themselves imputed. Numeric/integer columns are z-scored; factor/ordered columns are one-hot encoded automatically. If a variable has missing values, include it in traits instead. Rows are aligned to the same species / observation order as traits after tree matching (single-obs: rownames vs tip labels; multi-obs: species_col if present, otherwise input-row order). Matching nrow alone does not pair species to environment. Default NULL (no covariates).

epochs

integer. Maximum GNN training epochs (default 2000).

verbose

logical. Print progress (default TRUE).

seed

optional integer. When supplied, makes the stochastic fitting and prediction steps reproducible; the default NULL uses the current RNG stream.

multi_obs_aggregation

character. How to aggregate multiple observations per species before the Level-C baseline. "hard" (default) thresholds binary proportions at 0.5 and uses argmax for categorical. "soft" preserves species-level proportions and dispatches a soft E-step so that intermediate class frequencies contribute fractional liability evidence. Passed to fit_baseline.

lambda_mode

character. Pagel-lambda mode for the BM baseline. "estimate" (default) fits a per-trait Pagel's lambda on each continuous-family (BM-eligible) latent column – continuous, count, proportion, and zi_count magnitude, via the joint MVN / threshold-joint baseline's own lambda_cols machinery when that joint path fires, or a per-column re-fit otherwise. Discrete traits (binary, categorical, zi gate) AND ordinal have no discrete-trait analogue of Pagel's lambda of their own – there is no lambda_k to estimate for them. Under predict_method = "per_column" (see below) they stay fixed at lambda = 1 in every path, matching pre-S3 behaviour. Under predict_method = "exact" (and for traits the default "auto" routes to "exact") they instead share the joint fit's lambda_block – the exact conditional's covariance model uses ONE shared phylogenetic correlation matrix R(lambda_block) for every column, so the Sigma estimate feeding it must itself come from an internally consistent init, not a mix of R(1) for discrete columns and R(lambda_k) for continuous ones (Shinichi's decision, docs/dev-log/exact-default/S3-default-report.md). "fixed_1" preserves the pre-lambda Brownian correlation matrix everywhere; "cv" and "bayes" are alternative per-column estimators that force continuous-family columns onto the per-column BM path (no joint analogue). When predict_method = "exact" or joint_refine_iter > 0, the joint (multi-trait) prediction path additionally uses the single shared lambda_block for cross-trait computations that need one common phylogenetic correlation matrix, never overriding a continuous-family column's own estimated lambda_k. Passed to fit_baseline and stored in the fitted model config. When covariates are supplied, the covariate-aware BM path (bm_impute_col_with_cov()) accepts a numeric lambda or "estimate", so lambda_mode %in% c("estimate", "fixed_1") reaches it and each covariate-aware BM-eligible column gets its own estimated lambda; it does NOT accept "cv" / "bayes", which fall back to lambda = 1 for BM-eligible columns with a warning from fit_baseline.

joint_solver

character. Which solver estimates the joint MVN / threshold-joint / OVR categorical baselines. "inhouse" (default) is byte-identical to prior releases; "rphylopars" delegates to Rphylopars::phylopars()'s converged REML fit, with automatic fallback to "inhouse" on failure. Passed to fit_baseline and stored in the fitted model config. See docs/dev-log/2026-08-16-continuous-gap-diagnosis.md.

predict_method

character. Prediction route for the in-house joint solver, passed to fit_baseline. "auto" (default, S5b) fits the baseline with both the "exact" and "per_column" routes and picks, per trait, whichever has the lower loss on that trait's validation cells (mean squared error on the z-scored latent scale for continuous/count/proportion/ordinal/zi_count magnitude; mean log-loss for binary/zi_count gate; mean multinomial log-loss for categorical), falling back to "exact" on a tie or when fewer than 5 validation cells belong to that trait (including when there is no validation split at all). "exact" uses the full cross-trait conditional mean and variance of vec(L) ~ MVN(0, Sigma %x% R(lambda_block)) (Hadfield & Nakagawa, 2010 sparse precision form), each column GLS-mean-centred at lambda_block before the solve; discrete liability columns share lambda_block under this route (see lambda_mode above). Falls back to "per_column" above roughly 20000 unknown cells (roughly 4000 species at 5 traits), on a singular/unusable Sigma, or when fewer than 2 joint columns or no Henderson sparse precision are available; prints a one-time message() per R session when predict_method was left at its default, or a warning() every time when "exact" was requested explicitly. "per_column" retains the original per-column conditional prediction route (no cross-trait borrowing in the prediction step). None of the three options change covariance estimation or the "rphylopars" solver. The route actually used is recorded in fit$model_config$predict_method_used ("exact", "per_column", or "auto") and, per trait, in fit$model_config$predict_method_by_trait. Under "auto" with real validation cells, fit$model_config $route_val_n / $score_val_n record, per trait, how many of its validation cells chose the route versus were reserved for gate calibration and conformal scoring (see fit_baseline's predict_method docs).

joint_refine_iter

integer, default 0L. Enables cross-trait refinement of the joint baseline's cell imputations using the estimated Sigma (the in-house solver's max_iter EM cell-refinement; R/joint_mvn_solver.R). 0L preserves current behaviour byte-for-byte. The refinement is guarded: the Sigma step must shrink each iteration, or the loop rolls back to the last good iterate and sets $diverged. Assess this opt-in control with held-out evaluation on the intended data. Has no effect on any trait predicted by the "exact" route (see fit_baseline's predict_method docs); under the default predict_method = "auto" it therefore only applies to traits "auto" routes to "per_column". Passed to fit_baseline and stored in the fitted model config.

em_iterations

integer. Phase 6 EM iterations for the threshold-joint baseline (binary + ordinal + OVR categorical). Default 0L preserves v0.9.1 behaviour byte-for-byte. When >= 2L, the BM rate $\Sigma$ learned by the joint solver (joint_solver; the in-house solver by default) at iteration $k$ is fed back as the per-trait prior SD at iteration $k+1$, up to em_iterations times or until em_tol convergence. Passed to fit_baseline.

em_tol

numeric. Relative-Frobenius convergence tolerance for the Phase 6 / 7 EM loop. Default 1e-3.

em_offdiag

logical. Phase 7 opt-in: when TRUE AND em_iterations >= 2L, each liability cell's prior uses the full conditional-MVN from $\Sigma$'s off-diagonal entries, so that observing one discrete trait shifts (not just tightens) the prior on correlated other traits. Binary + ordinal only; OVR categorical stays on Phase 6 diagonal. Default FALSE. Passed to fit_baseline.

pool_method

character. How to pool multiple imputation draws (n_imputations > 1) for count, proportion, and zi_count magnitude traits: "median" (default) takes the per-cell median of the M decoded draws — robust to dropout-noisy latents amplified by expm1() / plogis() decoders. "mean" restores the pre-v0.9.2 arithmetic-mean pooling. "mode" (Phase H, v0.9.1.9010+) is intended for ordinal traits: per-cell majority vote across the M draws, avoiding the integer-mean-round bias toward middle classes. For continuous-family traits, "mode" falls back to "median". Binary / categorical / multi_proportion traits always pool by probability average; unaffected by this argument. See issue #40.

clamp_outliers

logical. Phase G (v0.9.1.9011+). When TRUE, post-back-transform predictions for log-transformed continuous, count, and zi_count magnitude traits are capped at tm$obs_max * clamp_factor (and tm$obs_max is the observed maximum on the original scale, recorded at preprocess time). Targets the AVONET Mass tail-extrapolation mode documented in useful/MEMO_2026-05-01_avonet_mass_diag.md where a $+3$-$4 \sigma$ latent overshoot becomes a 50x-100x value error after expm1(). Default FALSE preserves v0.9.1 behaviour exactly.

clamp_factor

numeric scalar (>= 1). Multiplicative factor on the observed maximum used by clamp_outliers. Default 5 (Tukey-style outlier definition: anything >= 5x the observed max is implausible). Ignored when clamp_outliers = FALSE.

match_observed

character, one of c("none", "pmm"). Phase G' (v0.9.1.9012+). Pass-through to predict.pigauto_fit. When "pmm", uses Predictive Mean Matching for log-transformed continuous, count, zi_count magnitude, and proportion traits: imputed values are drawn from the observed value pool, never extrapolated.

When to use: PMM is a niche feature. pigauto already provides conformal prediction intervals calibrated against held-out residuals and stochastic prediction draws for diagnostics. PMM is only worth enabling for: (a) methodological comparison against mice, or (b) workflows that specifically require imputed values to come from the observed data pool. For tail safety, prefer clamp_outliers = TRUE. PMM failed the downstream fixed-effect redesign pilot and is unsupported for inference; see multi_impute_analysis() for the analysis-aware backend.

Default "none" preserves pre-G' behaviour.

pmm_K

integer (>= 1). Donor pool size for PMM. Default 5L (mice convention). Ignored when match_observed = "none".

safety_floor

logical. When TRUE (default since v0.9.1.9002), calibration searches the 3-way simplex r_BM * BM + r_GNN * GNN + r_MEAN * MEAN so the grand mean is always in the candidate set. Under the validation metric used for calibration, the selected candidate cannot be worse than that grand-mean corner on the validation cells. This is a validation safeguard, not a guarantee about future held-out data. When FALSE, the v0.9.1 1-D calibration is used exactly. See the Safety floor section below.

phylo_signal_gate, phylo_signal_threshold, phylo_signal_method

Pass-through to fit_pigauto(). See that help page for details.

conformal_split_val

logical. Default FALSE. When TRUE, the validation cells are split 50/50 per trait into a gate-calibration half and a conformal-scoring half, so the conformal quantile is not post-selected on the same cells the gate was tuned on. Costs interval width at small validation sizes; see fit_pigauto for details. Previously reachable only via ...; exposed here for discoverability.

gnn

logical. When TRUE, trains the attention-based GNN correction (see fit_pigauto). When FALSE (default), no GNN is constructed or trained – impute() makes no torch/GPU calls, and the fit is the phylogenetic baseline alone (optionally re-weighted against a grand-mean floor). safety_floor and phylo_signal_gate keep their usual semantics; the pure traditional-stats arm is gnn = FALSE, safety_floor = FALSE, phylo_signal_gate = FALSE. Production predictions (result$completed, result$prediction) use a tax-free baseline_full fit on ALL observed cells (no val/test hold-out); result$evaluation and every other scorer keep using the held-out baseline, so evaluation never leaks test cells. User covariates enter pigauto only through the GNN, so under gnn = FALSE they are ignored (with a warning); without a validation split (missing_frac = 0) the fit is pure baseline and carries no conformal scores. Default FALSE since pigauto 0.11.0.9001: on simulated and real data at the current baseline defaults the GNN did not lower imputation error and took 20 to 90 times longer to fit (docs/dev-log/arc/2026-10-02-campaign-gnn-rerun.md). Set gnn = TRUE to train it, for example to use covariates, which only the GNN uses, or to restore the behaviour of earlier versions.

...

additional arguments passed to fit_pigauto.

Value

An object of class "pigauto_result" with components:

completed: The input traits data.frame with observed values preserved and only missing cells filled in. This is the primary output – typically what users want.
imputed_mask: Logical matrix (same shape as completed) that is TRUE for cells that were imputed (originally NA) and FALSE for observed cells.
prediction: A pigauto_pred object from predict.pigauto_fit containing raw model predictions for every cell (observed + missing), per-cell uncertainty ($se – see predict.pigauto_fit's "What $se means per trait type" section; it is not uniformly a standard error), class probabilities, and conformal intervals.
fit: The trained pigauto_fit object.
baseline: The phylogenetic baseline.
data: The preprocessed pigauto_data object.
splits: The val/test splits (or NULL if missing_frac = 0).
evaluation: Evaluation metrics on test set (or NULL).
check: A compact pigauto_check preflight summary.

Trait type auto-detection

pigauto infers each trait's type from its R class — no trait_types argument is needed for most data:

R class	pigauto type
`numeric`	continuous (auto-log if all-positive)
`integer`	count
`factor` with 2 levels	binary
`factor` (unordered) with >2 levels	categorical
`ordered` / `factor(..., ordered = TRUE)`	ordinal
`character`	→ factor → binary or categorical
`logical`	binary

Two types cannot be inferred from class alone and must be declared via trait_types:

"proportion": A numeric bounded 0–1, e.g. survival rate: trait_types = c(Survival = "proportion").
"zi_count": An integer with excess zeros, e.g. parasite count: trait_types = c(Parasites = "zi_count").

Use the trait_types argument directly (it is an explicit parameter, not a ... pass-through).

Traits vs covariates

The distinction is functional, not ontological: a trait is something you want to impute (NA values allowed in traits); a covariate is something you use to sharpen imputation accuracy (must be fully observed, passed via covariates). The same variable can be either depending on the scientific question.

Examples:

IUCN status with Data Deficient species → put it in traits as ordered(c("LC","NT","VU","EN","CR")) so pigauto predicts the unknown categories.
IUCN status fully known for all species → pass as a covariate to inform imputation of other traits (e.g. body mass, range size).
Realm / biome (factor) → pass as a covariate; pigauto one-hot encodes factor columns automatically (v0.6.1+).

Variables that belong in traits: anything with missing values you care about predicting. Variables that belong in covariates: fully observed, exogenous to the trait space (geography, climate, habitat, experimental treatment).

Safety floor (v0.9.1.9002+)

With safety_floor = TRUE (the new default), the post-training calibration grid searches a 3-way convex combination of the Brownian-motion baseline, the GNN delta, and the per-trait grand mean. The simplex is sampled at step 0.05 (231 candidates per latent column). Because the corner (0, 0, 1) — pure grand mean — is always in the grid, the selected candidate cannot be worse than the grand-mean corner on the validation cells under the calibration metric. The fit object gains four new slots: r_cal_bm, r_cal_gnn, r_cal_mean (each a named numeric of length p_latent), and mean_baseline_per_col.

Set safety_floor = FALSE to reproduce the pre-v0.9.1.9002 1-D calibration bit-identically (no mean term; r_cal_mean = 0; r_cal_bm = 1 - r_cal_gnn). See specs/2026-04-23-safety-floor-mean-gate-design.md for the design rationale and plans/2026-04-23-safety-floor-mean-gate.md for the implementation plan.

What gets imputed (read this first)

pigauto only imputes cells that are NA in the input. Observed cells are preserved as-is in result$completed. The slot result$prediction$imputed contains the model's prediction for every cell – observed and missing alike – and is intended for diagnostics (e.g. checking calibration on training cells), not as the imputed-values output. The imputed values themselves are result$completed[result$imputed_mask].

Common pitfall. If you call impute() on a fully observed trait matrix (no NAs among matched target rows), there is nothing to impute. impute() stops before graph construction or training and directs you to cross_validate() for predictive assessment. To exercise imputation, supply a real missing target cell.

Imbalanced K-class traits. At default settings (n_imputations = 1L, pool_method = "median"), a small ordinal / categorical trait whose marginal distribution is heavily skewed (e.g. AVONET Migration is ~78\ ~14\ collapse onto a corner that predicts the majority class everywhere. Compare candidate settings, including pool_method = "mode", with cross_validate() on the intended trait table rather than applying a generic setting.

Examples

# \donttest{
data(avonet300, tree300)
tree <- ape::keep.tip(tree300, tree300$tip.label[seq_len(30L)])
df <- avonet300[match(tree$tip.label, avonet300$Species_Key),
                c("Mass", "Wing.Length"), drop = FALSE]
rownames(df) <- tree$tip.label
df$Mass[seq_len(3L)] <- NA_real_
result <- impute(df, tree, epochs = 5L, verbose = FALSE)
#> Warning: phylo_signal_gate requires the 'phytools' package; returning NA for all traits.
#> Warning: Small validation set for 2 trait(s): Mass (n=4), Wing.Length (n=3). Calibrated gate and conformal scores will be noisy for these trait(s). 95% split-conformal coverage is NOT achievable for 2 trait(s) with fewer than 19 validation cells (Mass (n=4), Wing.Length (n=3)): the achievable ceiling is n_val / (n_val + 1), which only reaches 0.95 at n_val >= 19. See `?fit_pigauto` under 'Calibration at small n' for smoothing options.
head(result$completed)
#>                              Mass Wing.Length
#> Nothoprocta_pentlandii   103.3614       148.2
#> Eudromia_formosa         257.1174       223.9
#> Rhea_americana          2305.3834       604.5
#> Ptilopachus_petrosus     193.0000       117.2
#> Bambusicola_fytchii      313.7000       142.8
#> Francolinus_psilolaemus  478.3000       167.3
# }