Impute missing phylogenetic traits (convenience wrapper)

One-call interface for preprocessing, phylogenetic baseline fitting, and prediction. The default gnn = FALSE uses the baseline alone. Set gnn = TRUE to also train a GNN and calibrate its contribution. 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,
  k_eigen = "auto",
  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"),
  discrete_lambda = c("estimate", "fixed_1"),
  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 = FALSE,
  phylo_signal_gate = FALSE,
  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).

k_eigen

integer or "auto". Number of spectral node features used to build the phylogenetic graph. The default "auto" scales the dimension to tree size; an integer overrides that choice.

log_transform

logical. Log-transform every continuous trait whose observed values are all positive (default TRUE), then back-transform the imputations. Set it to FALSE when traits are already on a log scale (for example log body mass): otherwise they are logged a second time and imputed on a log(log) scale, which can bias analyses done on the log scale. With a mix of raw and already-logged traits, log the raw ones yourself and set FALSE. In a real-data check (PanTHERIA, 4,027 mammal species, four already-logged traits, posterior draws), the default shifted Rubin-pooled regression slopes by up to 3.9 standard errors of the complete-data slope, in a direction that depended on which values were missing; with FALSE the largest shift was 1.8.

missing_frac

numeric. Fraction of observed cells held out for validation/test evaluation (default 0.25). Must be in [0, 1). 0 adds no extra missingness (no val/test hold-out).

n_imputations

integer. Number of MC-dropout imputation sets (default 1). Must be a positive integer. 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). Must be a positive integer.

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. Binary, ordinal and categorical traits get their own lambda by default (discrete_lambda = "estimate"); the rest of this paragraph describes discrete_lambda = "fixed_1", where 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.

discrete_lambda

character, "estimate" (default) or "fixed_1". With "estimate", the binary, ordinal and zero-inflated-count gate liability columns of the threshold-joint baseline, and the one-vs-rest column of each categorical class, get their own Pagel's lambda, so a weak-signal discrete trait shrinks toward its prevalence instead of copying close relatives. "fixed_1" restores the previous behaviour (binary and ordinal at lambda = 1, or at the shared joint lambda on the "exact" route; categorical at 1). Evidence (simulated Brownian-motion and OU data, n = 100 to 1000, 30\ missing; docs/dev-log/discrete-lambda/README.md): estimating discrete lambda alone added about 0.01 to 0.04 discrete accuracy at Pagel's lambda of 0.3 or below; together with safety_floor = FALSE and phylo_signal_gate = FALSE (also the defaults) the gain was 0.04 to 0.12, and binary and categorical accuracy matched or beat BACE's. Costs: with lambda near 1 and n >= 1000, accuracy -0.002 to -0.005 and Brier score +0.010 to +0.016; on AVONET300's categorical traits, Brier +0.014 to +0.020 with accuracy unchanged or better. Ordinal accuracy still trails BACE's at lambda = 0.3. The categorical one-vs-rest lambdas are not reported in $lambda_per_trait. See fit_pigauto().

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. Default FALSE (it was TRUE from v0.9.1.9002 until 0.11.0.9002; see fit_pigauto() for why it changed; set TRUE to restore). When TRUE, 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. phylo_signal_gate now defaults to FALSE (it was TRUE from v0.9.1.9003 until 0.11.0.9002): with Pagel's $\lambda$ estimated for every trait type the gate no longer improved accuracy in any tested setting, and at $\lambda$ of 0 to 0.1 it was worse than predicting the mean. Set phylo_signal_gate = TRUE to restore the old behaviour.

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 when safety_floor = TRUE). With the defaults (safety_floor = FALSE, phylo_signal_gate = FALSE) this is the pure traditional-stats arm. 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 (opt-in since 0.11.0.9002; it was the default from v0.9.1.9002), 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: 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
# }