Restricted regression and model-agnostic edit restrictions with vimpute

Coefficient-constrained regression

For one target, let \(y_O\) contain its observed responses, \(X_O\) the corresponding predictor design matrix, and \(X_M\) the design matrix for rows whose target is missing. With the default per_observation = FALSE, restricted least squares estimates

\[ \widehat\beta =\operatorname*{arg\,min}_{\beta}\frac12\|y_O-X_O\beta\|_2^2 \quad\text{subject to}\quad C\beta\le d,\qquad H\beta=g. \]

The matrices \(C,H\) and right-hand sides \(d,g\) come from model_rules after substituting the other columns. For example, row-wise target bounds \(\ell_i\le x_i^T\beta\le u_i\) become linear inequalities on the same coefficient vector. An equality on a prediction becomes an equality on that vector. The resulting missing-target predictions are \(X_M\widehat\beta\). All constrained rows share coefficients, so individually feasible bounds can still be incompatible with a single fitted regression model.

With per_observation = TRUE, each missing row \(i\) gets its own coefficient vector \(\widehat\beta_i\). Each fit uses the same observed training responses and design matrix, but applies only the constraints for row \(i\). The prediction is \(x_i^T\widehat\beta_i\). Two businesses with identical model predictors can therefore receive different turnover estimates when their observed accounting components require different totals. This entails one constrained fit per missing row for each target.

With robust = TRUE, the same coefficient constraints are retained and the least-squares objective is replaced by a standardized Huber loss:

\[ \min_{\beta}\sum_{i\in O}\rho_k\!\left( \frac{y_i-x_i^T\beta}{s}\right),\qquad \rho_k(r)= \begin{cases} \tfrac12 r^2,& |r|\le k,\\ k|r|-\tfrac12 k^2,& |r|>k. \end{cases} \]

Here \(s\) is the residual scale used by the restricted learner and \(k\) is huber_k. Robust fitting changes the influence of observed outliers; it does not change which coefficient restrictions are imposed.

Initializing linked missing values in one row

Suppose turnover \(T\) and intermediate consumption \(C\) are both missing for a business whose observed gross value added is 30. The edit requires \(T-C=30\). If the observed column medians are 100 for turnover and 80 for consumption, filling each cell with its median would give a difference of only 20.

When both targets use per_observation = TRUE, vimpute() uses these medians as starting points for a joint calculation. With just this edit, it chooses the pair minimizing \((T-100)^2 + (C-80)^2\) subject to \(T-C=30\). The result is 105 for turnover and 75 for consumption. If other edits also constrain these amounts, such as a turnover total fixed by observed domestic sales and exports, the initialization enforces those edits at the same time.

These values supply missing predictors during prediction. If turnover is visited first, its coefficient fit can use 75 for the still-missing consumption value. When consumption is visited, its fit uses the turnover just imputed. The training responses come from observed data. Joint initialization is triggered only when an edit references multiple missing per-observation targets in that row.

Fixed observations and the feasible region

For record \(i\), let \(M_i\) denote the originally eligible numeric imputation cells, including cells selected by makeNA, and \(O_i\) the fixed cells. Columns excluded from imputation are fixed too. Let \(o_i\) be their values and \(\widehat x_{i,M_i}\) the imputations after prediction and any configured stochastic operations. Different entries may come from different learners.

Consider full-record linear edits

\[ \widetilde A_{i,M_i}x_{i,M_i} +\widetilde A_{i,O_i}o_i\le\widetilde b_i, \qquad \widetilde E_{i,M_i}x_{i,M_i} +\widetilde E_{i,O_i}o_i=\widetilde f_i. \]

Substituting fixed values gives

\[ A_i=\widetilde A_{i,M_i},\quad b_i=\widetilde b_i-\widetilde A_{i,O_i}o_i, \qquad E_i=\widetilde E_{i,M_i},\quad f_i=\widetilde f_i-\widetilde E_{i,O_i}o_i. \]

The feasible region is \(\mathcal C_i=\{x:A_ix\le b_i,\ E_ix=f_i\}\). Conditions evaluated from fixed categorical or numeric values select the active linear edits for that record. Fixed observations are substituted exactly; they are not assigned large penalty weights. A violation involving only fixed values is therefore infeasible and raises an error.

Weighted value projection

With positive diagonal weights \(W_i=\operatorname{diag}(w_{ij}:j\in M_i)\), the generic stage solves independently for each record

\[ \boxed{ x^*_{i,M_i} =\operatorname*{arg\,min}_{x\in\mathcal C_i} \frac12(x-\widehat x_{i,M_i})^T W_i(x-\widehat x_{i,M_i}) }. \]

For a nonempty supported closed convex feasible region, positive weights give a unique projection. Feasible predictions are preserved exactly without a solver call. Cells absent from the active constraints also remain unchanged. An empty feasible region produces an error; solver status, finite solutions and absolute constraint residuals are checked before accepting any adjustment. The implementation accepts residuals within tolerance and represents strict inequalities using an interior margin.

Equal weights measure changes in original units. Scaled weights use \(w_{ij}=1/s_j^2\), where \(s_j\) is an observed variable scale, so the objective measures squared standardized adjustments. Larger weights make changes more costly. An alternative is to use weights \(1/\widehat\sigma_{ij}^2\) based on predictive uncertainty. Automatic predictive-uncertainty weighting is not implemented: observed scales are not predictive standard deviations, and the current named weight vector is common across records.

ECOS formulation

Let \(D_i=W_i^{1/2}\). The identity-design least-squares formulation from the setup is

\[ X_i=D_i,\qquad y_i=D_i\widehat x_{i,M_i},\qquad \|X_ix-y_i\|_2^2 =(x-\widehat x_{i,M_i})^TW_i(x-\widehat x_{i,M_i}). \]

The implementation uses the equivalent centered coordinates \(z=D_i(x-\widehat x_{i,M_i})\) and the second-order cone problem

\[ \min_{z,t}t\quad\text{subject to}\quad \|z\|_2\le t, \] \[ A_iD_i^{-1}z\le b_i-A_i\widehat x_{i,M_i},\qquad E_iD_i^{-1}z=f_i-E_i\widehat x_{i,M_i}. \]

Minimizing the norm gives the same adjusted values as minimizing half its square. Each problem involves only one record’s relevant imputed cells; no dense optimization matrix spans all missing cells in the dataset.

An accounting example with an analytical solution

For equality-only edits with independent rows of \(E\), the solution is

\[ x^*=\widehat x-W^{-1}E^T (EW^{-1}E^T)^{-1}(E\widehat x-f). \]

Consider turnover \(T\), costs \(C\) and profit \(P\), with \(T=C+P\). If all three are imputed as \((100,70,40)\) and \(W=I\), then \(E=(1,-1,-1)\) and the residual is \(-10\). The projection is

\[ (T^*,C^*,P^*)= \left(\frac{310}{3},\frac{200}{3},\frac{110}{3}\right). \]

If turnover is observed at 100, it stays exactly 100 and only \((C,P)\) can move. Equal weights give \((C^*,P^*)=(65,35)\). For arbitrary positive weights,

\[ C^*=70-\frac{10w_P}{w_C+w_P},\qquad P^*=40-\frac{10w_C}{w_C+w_P}. \]

For example, user-supplied \(w_C=1/15^2\) and \(w_P=1/10^2\) give \(C^*\approx63.077\) and \(P^*\approx36.923\). If nonnegativity constraints are also imposed, these examples already satisfy them; otherwise active bounds must be included in the joint optimization.

Sequential and stochastic imputation

For a sequential sweep \(s\), write the learner-and-uncertainty update as \(\widehat x^{(s)}=F_s(x^{(s-1)},\varepsilon_s)\). The completed sweep is then projected record by record:

\[ x^{(s)}_i=P_{\mathcal C_i}^{W_i} \bigl(\widehat x^{(s)}_i\bigr). \]

Convergence and the next sweep use \(x^{(s)}\), so they reflect the adjusted values. With sequential = FALSE, this sequence runs once. With multiple imputation, it runs separately for every imputed dataset. The order is

\[ \text{prediction}\longrightarrow \text{stochastic imputation, if configured}\longrightarrow \text{joint projection}\longrightarrow \text{convergence or return}. \]

The restricted learner retains its existing uncert = "none" override. For other learners, projection can move a PMM draw away from its donor or alter the distribution of a residual or normal-error draw. Sample-and-project generally does not sample from the constrained conditional law,

\[ p(x_M\mid x_O,\mathcal C) \propto p(x_M\mid x_O)\,\mathbf 1\{x\in\mathcal C\}. \]

This expression applies when the admissible set has positive probability. For exact equalities under continuous predictions, the equality surface has probability zero; a constrained sampling model needs a suitable parameterization or conditional distribution on that surface. Constrained probabilistic sampling and edit-aware model training remain research extensions, rather than guarantees of the present projection method.

Linked missing values in 10,000 records

In the next example the same 200 businesses are missing turnover, intermediate consumption and personnel costs. Gross value added is observed, so the edit gross_value_added == turnover_total - intermediate_consumption links two missing cells in every affected row. The observed turnover components, cost components and payroll components determine the three totals. This lets us check recovery against the held-out values as well as compliance with the edits.

targets <- c("turnover_total", "intermediate_consumption", "personnel_costs")
original <- lse_synthetic[seq_len(10000L), edit_cols]
set.seed(20260630)
rows <- sample.int(nrow(original), 600L)[seq_len(200L)]
masked <- original
masked[rows, targets] <- NA_real_

Run both restricted fits with the formulas defined above. On this mask the shared fit raises “Restricted regression equality constraints are inconsistent.” The per-observation fit fills all 600 cells.

fit_restricted <- function(per_observation) {
  vimpute(
    masked, method = "restricted", formula = restricted_formulas[targets],
    pmm = FALSE, sequential = FALSE,
    learner_params = list(restricted = list(
      rules = lse_rules, per_observation = per_observation
    ))
  )
}
tryCatch(fit_restricted(FALSE), error = conditionMessage)
row_fit <- fit_restricted(TRUE)

For ranger, pass the same validator through edit_rules. After the forest predicts all three amounts, the projector adjusts the missing cells jointly within each row. The observed industry code also participates in the full validator through the membership rule using setdiff(LETTERS[2:19], "O").

forest_fit <- vimpute(
  masked, method = "ranger", edit_rules = lse_rules,
  learner_params = list(ranger = list(num.trees = 15L, num.threads = 1L)),
  pmm = FALSE, sequential = FALSE, seed = 8
)

Both completed datasets pass all 920,000 record-by-rule checks. The following checks also verify that each imputation flag identifies the intended row and that each recovered amount matches its held-out value within numerical tolerance.

for (fit in list(row_fit, forest_fit)) {
  completed <- as.data.frame(fit)
  for (v in targets) {
    stopifnot(
      identical(which(completed[[paste0(v, "_imp")]]), sort(rows)),
      isTRUE(all.equal(completed[[v]][rows], original[[v]][rows], tolerance = 1e-7))
    )
  }
  checks <- validate::values(validate::confront(completed[edit_cols], lse_rules))
  stopifnot(identical(dim(checks), c(10000L, 92L)), all(checks))
}