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,
  ...,
  ordinal_levels = NULL
)

# S3 method for class 'pigauto_result'
print(x, ...)

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 stochastic prediction-diagnostic draws (default 1). Must be a positive integer. With multiple draws, prediction adds independent Gaussian perturbations per latent cell using the selected baseline SD; GNN dropout also varies across runs when gnn = TRUE. pool_method combines decoded draws for the returned point prediction: completed is one data frame, not a list of completed datasets. Despite this legacy argument name, these are not joint phylogenetic posterior draws or validated imputations for downstream inference. See multi_impute and the multiple- imputation article for the separate workflows and their limits.

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. Covariate use depends on the selected route: a covariate-aware per-column BM path may use them. Joint MVN and threshold-joint baselines omit them and emit a separate baseline warning. One-versus-rest categorical baselines also omit them but do not trigger the separate baseline warning themselves. In a mixed fit, another selected joint route can still trigger it. With gnn = TRUE, covariates also enter the GNN and its additive linear component. predict_method = "per_column" selects predictions within a joint fit; it does not activate covariate regression. In multi-observation data the baseline aggregates observations and covariates to species level; with gnn = FALSE it does not predict observation-specific covariate effects within a species. 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. Predictions can vary between seeds; a fixed seed supports reproducibility, not a measure of predictive uncertainty.

multi_obs_aggregation

character. How to aggregate multiple observations per species before a joint 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 continuous-family baseline columns (continuous, count, proportion, and zi_count magnitude). "estimate" (default) estimates each eligible column's lambda within its selected joint or per-column fit. "fixed_1" uses lambda = 1. "cv" and "bayes" use alternative per-column estimators and cannot use the joint numeric route. Discrete-liability lambda is controlled separately by discrete_lambda. Exact joint prediction and optional cross-trait cell refinement also use a shared lambda_block for cross-trait computations, without replacing continuous-family column-specific lambdas. A covariate-aware per-column BM fit accepts "estimate" or fixed lambda; unsupported "cv"/"bayes" requests warn and use lambda = 1. See fit_baseline for route details and fixed-lambda controls.

discrete_lambda

character, "estimate" (default) or "fixed_1". Controls lambda for binary, ordinal, and zi_count gate liabilities in threshold-joint fits and for each categorical one-vs-rest fit. Exact joint predictions share lambda_block; categorical class-fit lambdas are not reported in $lambda_per_trait. See fit_baseline for route-specific behavior and the evaluated benefits and costs. Stored in the fitted model configuration.

joint_solver

character. Solver for the joint MVN, threshold-joint, and categorical one-vs-rest baselines. "inhouse" (default) uses pigauto's joint solver. "rphylopars" delegates to Rphylopars::phylopars(), using its lambda model under lambda_mode = "estimate", with automatic fallback to "inhouse" on failure or implausible output. This does not affect routes outside these joint fits. "cv"/"bayes" lambda modes instead use per-column numeric fits. See fit_baseline for dispatch and fallbacks.

predict_method

character. Prediction route for the in-house joint solver, passed to fit_baseline. "auto" (default) 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. With no validation split, exact prediction is requested but route-specific per-column fallbacks remain possible. "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. Maximum passes of optional cross-trait refinement of joint-baseline cell imputations using the estimated covariance matrix. 0L disables this refinement. The covariance step must shrink each iteration; otherwise the last accepted iterate is restored and $diverged is set. This has no effect on traits predicted by the "exact" route, which returns the analytic conditional directly. Under predict_method = "auto", it can therefore affect only traits routed to "per_column". Evaluate this opt-in setting on held-out cells from the intended data.

em_iterations

integer. Maximum passes through each liability-prior refinement loop for binary, ordinal, and one-vs-rest categorical baselines. Default 0L disables iterative refinement and uses a single-pass baseline. 1L also permits only the initial pass; 2L permits at most one feedback pass. After each pass, estimated liability variances inform the next prior. Refinement can stop early on convergence or a later fit failure. A categorical pass may require K one-vs-rest fits; predict_method = "auto" may fit two route candidates. Continuous-only joint MVN fits are unaffected. Passed to fit_baseline.

em_tol

numeric. Relative-Frobenius convergence tolerance for liability-prior refinement (default 1e-3). The binary/ordinal threshold-joint loop compares the variance diagonal by default, or the full covariance matrix when em_offdiag = TRUE. Categorical one-vs-rest refinement compares the vector of class variances. Passed to fit_baseline.

em_offdiag

logical. When TRUE and em_iterations >= 2L, binary/ordinal liability priors use conditional-MVN means and SDs based on the other traits and the full covariance matrix from the preceding pass. With the default FALSE, refinement updates per-trait prior SDs only. This option does not change categorical one-vs-rest refinement. Passed to fit_baseline.

pool_method

character. How to pool stochastic prediction-diagnostic 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" takes the arithmetic mean. "mode" 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. When TRUE, cap back-transformed predictions for log-transformed continuous, count, and zi_count magnitude traits at tm$obs_max * clamp_factor, where tm$obs_max is the observed maximum on the original scale, recorded during preprocessing. This optional cap limits extrapolation after exp() or expm1(); it is not an uncertainty bound. Default FALSE leaves these predictions uncapped.

clamp_factor

numeric scalar (>= 1). User-selected multiplier of the observed maximum for the upper cap when clamping is enabled. Default 5. Ignored when clamp_outliers = FALSE.

match_observed

character, one of c("none", "pmm"). 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" leaves predictions unmatched.

pmm_K

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

safety_floor

logical. Default FALSE. When TRUE, calibration searches a three-way convex combination of phylogenetic baseline, GNN, and grand-mean candidates. The selected candidate is no worse than the mean-corner candidate on the validation cells under the calibration metric. An additive fixed-effect term remains when present, so that corner is not always a pure grand mean; multi_proportion groups are excluded. This is a validation-set safeguard, not a guarantee for future data. With FALSE, calibration uses the baseline/GNN line with r_cal_mean = 0.

phylo_signal_gate, phylo_signal_threshold, phylo_signal_method

Pass-through to fit_pigauto(). phylo_signal_gate = FALSE by default; enable it to apply the documented training-data signal override after calibration. Assess the selected setting on held-out data.

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 full-data baseline fitted using all observed cells (baseline_full, with no validation/test hold-out); result$evaluation and every other scorer keep using the held-out baseline, so evaluation never leaks test cells. User covariates enter the GNN and its additive fixed-effect term, and may also enter a selected per-column numeric BM fallback, including when gnn = FALSE. Joint MVN and threshold-joint routes omit covariates and emit a separate baseline warning. One-versus-rest categorical routes also omit them but do not trigger that warning themselves. In a mixed fit, another selected joint route can still trigger it. Covariate use does not imply improved predictions. Without a validation split (missing_frac = 0) the fit is pure baseline and carries no conformal scores. In the reported simulated and real-data comparison at the evaluated baseline settings, GNN-on did not lower imputation error and took 20 to 90 times longer to fit (see the changelog for that study). These results do not establish performance on new data.

...

Additional arguments passed by impute() to fit_pigauto(); ignored by print.pigauto_result().

ordinal_levels

named list giving the scientific order of ordinal levels, e.g. list(status = c("LC", "NT", "VU", "EN", "CR")). Each named column is coerced to an ordered factor with exactly those levels before type detection (so no trait_types entry is needed), and $completed returns it as an ordered factor. Observed values absent from the declared levels are an error. Same semantics as the ordinal_levels argument of read_traits. Default NULL.

x

A pigauto_result object returned by impute().

Value

An object of class "pigauto_result" with components:

completed: The input traits data.frame. Observed values are preserved. Missing cells are filled when the input row is included in the fit and a prediction is available; other missing cells remain NA.
imputed_mask: Logical matrix in original input row order, with one column per input trait. The species_col column, when supplied, is excluded. TRUE marks an originally missing target in a row matched to the fitted data; a marked cell can remain NA if no prediction is available. Observed cells and rows outside the fitted tree are FALSE.
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. For an untransformed continuous trait, the usual global per-trait score yields the same interval width across rows; it does not adapt to each species' number of observations. Other scoring methods or transformations can change this pattern.
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.

Invisibly returns x.

Details

A single completed data set is useful for point prediction, but analysing it as if its imputed values were observed ignores imputation uncertainty and can also change estimates. See multi_impute and the multiple-imputation article for downstream workflows and their limitations.

Trait type auto-detection

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

numeric: continuous (automatically log-transformed when all observed values are positive)
integer: count, except that columns with any negative observed value are detected as continuous
factor with two levels: binary
unordered factor with more than two levels: categorical
ordered factor with two levels: binary; ordered factors with other numbers of levels: ordinal
character: converted to a factor, then treated as binary or categorical according to its number of levels
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 a named trait_types override when the inferred type does not match the measurement. This 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.

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

With safety_floor = TRUE, post-training calibration grid searches a 3-way convex combination of the Brownian-motion baseline, the GNN delta, and the mean baseline. The simplex is sampled at step 0.05 (231 candidates per latent column). Because the corner (0, 0, 1) is always in the grid, the selected candidate cannot be worse than the mean-corner candidate on the validation cells under the calibration metric. An additive fixed-effect term is retained when present, so this is not always a pure grand mean; multi_proportion groups are excluded. This safeguard applies only to the validation cells used for calibration. The fit object stores: r_cal_bm, r_cal_gnn, r_cal_mean (each a named numeric of length p_latent), and mean_baseline_per_col.

With the default safety_floor = FALSE, calibration uses no mean term: r_cal_mean = 0 and r_cal_bm = 1 - r_cal_gnn.

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 mask marks originally missing targets in fitted rows; some selected cells can remain NA when no prediction was available. To extract target values by trait without coercing mixed columns:


  mask <- result$imputed_mask
  imputed_values <- lapply(colnames(mask), function(nm) {
    result$completed[[nm]][mask[, nm]]
  })
  names(imputed_values) <- colnames(mask)

For successful fills only, also exclude NA values from each trait vector.

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 about 78 percent Resident, 14 percent Partial, and 8 percent Full at n=300) can have its calibrated gate 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.

Printing

print.pigauto_result() reports species, traits and types, the total and per-trait targeted-cell counts, and a route-specific inference note. Held-out test metrics are shown for continuous and count traits (RMSE and, when available, Pearson correlation), ordinal traits (Spearman correlation), and binary and categorical traits (accuracy). The printout does not show metrics for proportion, zero-inflated count, or multi-proportion traits; see result$evaluation or evaluate_imputation() for the available evaluation metrics. Completed values are stored in result$completed. Some targeted cells can remain NA when a prediction is unavailable. The method ignores ... and returns x invisibly.

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
# }