Repository navigation
Added standalone AI REML Package - #77
Draft
adityasyam wants to merge 8 commits into
Draft
adityasyam wants to merge 8 commits into
adityasyam wants to merge 8 commits into
Conversation
dcdehaas
marked this pull request as draft
October 8, 2026 12:04
dcdehaas
reviewed
Oct 12, 2026
| return _hutchpp(operator, num_vectors, rng) | ||
| if method == "xtrace": | ||
| if num_vectors < 2: | ||
| return _hutchinson(operator, num_vectors, rng) |
Member
There was a problem hiding this comment.
if this is a requirement of _xtrace, that it needs >= 2 vectors, then it seems like this should be in the _xtrace() function.
Member
There was a problem hiding this comment.
same with the other checks above, like num_vectors < 3
| ``"exact"``. | ||
| :param rng: Random generator; a fresh default one is used if omitted. | ||
| """ | ||
| operator = aslinearoperator(operator) |
Member
There was a problem hiding this comment.
this doesn't seem necessary - either your type annotation is wrong (operator is not always a LinearOperator) or this is redundant?
| from scipy.linalg import solve_triangular | ||
| from scipy.sparse.linalg import LinearOperator, aslinearoperator | ||
|
|
||
| __all__ = ["TraceEstimate", "estimate_trace", "rademacher"] |
Member
There was a problem hiding this comment.
you don't necessarily need this - its just another thing to keep consistent later. The default Python behavior is to export all symbols that do not begin with "_". The downside of this default behavior is that your imports above will be exported - e.g., aireml.trace.aslinearoperator will be a thing. However, this really only affects the use case of from aireml.trace import *, which is best avoided!
I suggest removing this, since in __init__.py you already list the ones you want exported to the aireml top-level module:
from .trace import estimate_trace, TraceEstimate
If you want a nice "clean" symbol list, you can just change your imports above to be:
from scipy.linalg import _solve_triangular
from scipy.sparse.linalg import _LinearOperator, _aslinearoperator
This is totally an optional change. Its akin to "do you mow your grass to 1.5 or 2 inches high?"
| The Nystrom preconditioner for ``V = A + sigma^2 I`` needs a sketch | ||
| ``A @ Omega``. During AI-REML, ``A = sum_i theta_i V_i`` changes every | ||
| iteration but the operators ``V_i`` do not, so each ``V_i @ Omega`` is | ||
| computed once here and recombined for free as the coefficients move. |
Member
There was a problem hiding this comment.
notation is a little confusing here -- what is the relationship between Omega and theta_i?
|
|
||
| def __init__( | ||
| self, | ||
| operators: Sequence[LinearOperator], |
Member
There was a problem hiding this comment.
is operators a good name, or is there something more specific that can be used? What do the operators represent?
| :param rng: Random generator; a fresh default one is used if omitted. | ||
| """ | ||
| operator = aslinearoperator(operator) | ||
| rows, cols = operator.shape |
Member
There was a problem hiding this comment.
use operator_size() to replace the next 3 lines?
|
|
||
| The caller is responsible for the operator actually being symmetric | ||
| positive semi-definite; that is assumed throughout and is not checked, | ||
| because checking would require materializing the matrix. |
Member
There was a problem hiding this comment.
I'm not sure I understand this comment. If matrix is a LinearOperator, it has a shape associated with it - and everything assumes that shape is correct (if its not, the matmul will fail). If matrix is an array, then the shape is correct by def.
|
|
||
| ArrayLike = Union[numpy.ndarray, LinearOperator] | ||
| #: Variance-component values: anything array-like of floats. | ||
| Coefficients = Union[Sequence[float], numpy.ndarray] |
Member
There was a problem hiding this comment.
what's the advantage of supporting more than just ndarray?
also, I think that numpy.typing.NDArray is the correct annotation to use here.
| def _matmat(self, other: numpy.ndarray) -> numpy.ndarray: | ||
| result = numpy.zeros((self.shape[0], other.shape[1]), dtype=self.dtype) | ||
| for coefficient, operator in zip(self.coefficients, self.operators): | ||
| if coefficient == 0.0: |
Member
There was a problem hiding this comment.
why does 0.0 get special behavior? if this is set elsewhere as an explicit "skip this coefficient", then maybe nan would be better to use. if it is not, then it seems unlikely that an exact 0.0 would occur?
| ) | ||
|
|
||
|
|
||
| class LinearCombinationOperator(LinearOperator): |
Member
There was a problem hiding this comment.
this comment is totally optional / for-your-information. was useful for me to look into this, for my own learning...
this kind of neat, and I think pylops may be worth exploring more in general. This operator can be simplified by using pylops.FunctionOperator, which lets you just define a matvec function f and wrap it in a LinearOperator. Here is my simple proof of concept:
from aireml.operators import LinearCombinationOperator
from pylops import FunctionOperator
from scipy.sparse.linalg import aslinearoperator
import numpy
# Test inputs
x1 = numpy.random.standard_normal((100, 100))
x2 = numpy.random.standard_normal((100, 100))
op_list = list(map(aslinearoperator, [x1, x2]))
coeffs = [10, 19]
y = numpy.random.standard_normal((100, 1))
# Get result using existing LinearCombinationOperator
op1 = LinearCombinationOperator(op_list, coeffs)
z1 = op1 @ y
# Re-implementation using a simple function and pylops.FunctionOperator
def lin_combo(op_list, coeffs):
assert len(op_list) == len(coeffs)
def f(other):
result = numpy.zeros_like(other)
for lop, coeff in zip(op_list, coeffs):
result += (lop @ other) * coeff
return result
return f
op2 = FunctionOperator(lin_combo(op_list, coeffs), 100)
z2 = op2 @ y
numpy.testing.assert_allclose(z1, z2)
print(z1.shape, z2.shape)
Up to you whether you want to try to reduce the size of some of the linear operators using this approach. My personal philosophy for code is "the less, the better" -- pylops only depends on scipy and numpy, so it is a good dependency to add.
Member
There was a problem hiding this comment.
Relatedly, I don't think there is any reason to re-use this operator with the set_coefficients() method? nothing in the setup of this operator is expensive, computationally.
Implements the average-information REML algorithm of Lee et al. (2026),
"Genetic prediction with ARG-powered linear algebra" (GENETICS 233(1),
iyag074), as a self-contained package that depends only on numpy and
scipy and has no notion of GRG, ARGs, or genotype file formats.
The relatedness structure enters only through matrix-vector products, so
it is supplied as a scipy.sparse.linalg.LinearOperator; nothing in the
fit is O(N^2). With V = sum_i theta_i V_i and P the REML projection,
gradient_i = trace(P V_i) - y^T P V_i P y
AI_ij = y^T P V_i P V_j P y
theta <- theta - AI^-1 gradient
The trace terms present in both the Hessian and the Fisher information
cancel in their average, so the curvature is exact and only the gradient
is stochastic.
Contents:
* aireml/trace.py - Hutchinson and XTrace (Epperly et al. 2024)
estimators for trace(P V_i)
* aireml/solvers.py - preconditioned CG with a randomized Nystrom
preconditioner (Frangella et al. 2021), sketched
once per component and reused across iterations;
a dense reference solver for small problems
* aireml/model.py - products with P, the gradient, and the average
information at a given theta
* aireml/reml.py - randomized Haseman-Elston initialization, the
AI-REML loop with step halving and the paper's
averaging stopping rule, standard errors from the
average information, and in- and out-of-sample
BLUPs
* aireml/operators.py - LinearOperator helpers, including building a GRM
operator from any N x M genotype operator
Verification in test/aireml/ (60 tests): the analytic gradient matches
finite differences of the REML objective, the average information matches
the explicit triple product, AI-REML with exact traces reaches the same
optimum as a brute-force Nelder-Mead maximization of the restricted
likelihood, the matrix-free stochastic path lands within a few percent of
that optimum, estimates are unbiased across replicates with standard
errors that track their empirical spread, and BLUPs match the closed form.
examples/aireml_benchmark.py measures runtime scaling and compares
AI-REML with Haseman-Elston.
CI and prep_for_commit.sh now run black, flake8 and mypy over aireml/.
Implements Hutch++ (Meyer et al. 2021, Algorithm 1) as a third option for trace(P V_i), the one quantity in the AI-REML gradient that cannot be reduced to a handful of matrix-vector products. With a budget of k products and c = k // 3: sketch Q from A S, take trace(Q^T A Q) exactly on that subspace, and add a Hutchinson correction on its orthogonal complement using independent probes G. Unbiased because G is independent of the sketch. Written inline so the package keeps its numpy + scipy only dependency footprint, but cross-checked against pylops.utils.estimators.trace_hutchpp, which implements the same algorithm; that test skips when pylops is absent. Also makes the matrix-vector cost explicit, since it differs per estimator and a comparison at equal num_vectors would be unfair: num_vectors is the product budget for hutchinson and hutchpp (matching pylops' neval) but the number of test vectors for xtrace, which spends two products each. TraceEstimate.num_matvecs reports what was actually consumed. Measured, at equal product budget: * As a trace estimator, the sketch-and-correct methods are an order of magnitude better than plain Hutchinson -- relative SD 1.000 (hutchinson), 0.069 (hutchpp), 0.053 (xtrace) on a 400x400 matrix with eigenvalues k^-1.5 over 300 repeats. * Inside a whole AI-REML fit the difference washes out. At N=1500 over 12 seeds all three converge every time and agree on h2 to within their mutual uncertainty (SD 0.0018, 0.0022, 0.0017), because the average information is exact and the stopping rule averages the last 15 iterates. * The remaining difference is wall clock, tracking dense overhead rather than solves: 14.1s (xtrace), 12.9s (hutchpp), 12.2s (hutchinson). XTrace's leave-one-out downdate costs a triangular solve and several m-by-m products per estimate. xtrace stays the default: lowest variance per product, the choice in Lee et al. (2026), and a ~9% wall-clock gap on one benchmark is too thin to flip a default. Adds examples/aireml_benchmark.py --trace-compare, which reproduces both tables, and six tests covering unbiasedness, the three-way budget split, the minimum budget, non-symmetric operators, the equal-budget variance ranking, and the pylops cross-check. Separately fixes six mypy errors on numpy.finfo(float).tiny that predate this change and only surface with current mypy, which CI installs unpinned.
Total wall clock is a confounded way to compare trace estimators inside a full AI-REML fit: a run that happens to need one extra iteration spends about 6% more products, which is larger than the per-product difference between the estimators. Observed directly -- one run drew 16.7 / 17.7 / 17.0 iterations for xtrace / hutchpp / hutchinson, so hutchpp finished slower in total despite being cheaper per product. Adds a ms/1k-matvecs column to the --trace-compare fit table to divide the iteration count out, and corrects the README, which claimed a ~9% wall-clock advantage for hutchpp on the strength of a single run whose iteration counts happened to coincide. What actually holds, timed in isolation against a cheap operator so the dense bookkeeping is visible rather than buried under the solves: XTrace does strictly more dense work per estimate -- both QR the sketch, but XTrace also needs the triangular solve for R^-T, a column normalization, and several m-by-m products for the leave-one-out terms. That overhead is large when a product is trivial and fades as the product gets costlier (+222% at a rank-5 operator, +36% at rank-50, within noise by rank-200). Since every product in a real fit is a CG solve, the end-to-end difference is not reliably measurable, and the README now says so instead of quoting a number.
Code review on PR aprilweilab#77. No behavioral change: variance components, BLUPs and fixed effects are bit-for-bit identical before and after, verified across the dense/exact path and all three trace estimators. Substantive fixes: * Each trace estimator now enforces the minimum budget its own algorithm implies, instead of the dispatcher doing it: Hutchinson needs one probe, XTrace two (leaving one out of a single-vector sketch leaves nothing to build Q_{-i} from), Hutch++ three (its budget splits three ways). Previously XTrace silently fell back to Hutchinson below two vectors, which hid a caller mistake. * estimate_trace reuses operator_size() rather than repeating the square-shape check, and drops the redundant num_vectors < 1 guard. * LinearCombinationOperator no longer skips coefficients that are exactly 0.0. There was no sentinel meaning attached to zero, and an exact 0.0 is unlikely enough that the branch bought nothing. * Dropped LinearCombinationOperator.set_coefficients(). Construction is pure bookkeeping -- it rewraps the components and stores the coefficients, touching no data -- so CovarianceSolver.update() builds a fresh instance rather than mutating one. The caching that does matter, the per-component Nystrom sketch, is untouched. Clarity: * NystromSketch documents how Omega relates to theta_i: Omega is a fixed matrix of random probe directions with no connection to the variance components, and because the sketch is linear in the coefficients, A @ Omega = sum_i theta_i (V_i @ Omega) -- which is why each V_i @ Omega is computed once and reweighted thereafter. * Renamed the generic `operators` parameter and attribute to `components` throughout the solvers, matching the V_i of V = sum_i theta_i V_i. All call sites were positional, so no caller changes. * Removed the Sphinx "#:" attribute-comment syntax, which was used inconsistently and read as a typo; those attributes are now documented in the enclosing docstrings. * Rewrote the as_symmetric_operator note that conflated shape with symmetry. Shape is carried by the operator and is checked; symmetry and positive semi-definiteness are assumed and unverifiable without materializing the matrix, and violating them yields wrong estimates rather than an error. * ArrayLike says why it is wider than an array -- LinearOperator is the case that matters -- and the annotations now use numpy.typing.NDArray. * Dropped the module-level __all__ lists. Python already hides underscore-prefixed names, and aireml/__init__.py curates the public surface; keeping both meant keeping two lists in sync. Also folded the repeated five-way matvec/matmat/rmatvec/rmatmat/dtype LinearOperator boilerplate into one _from_matmat() helper, which is the simplification pylops.FunctionOperator offers while keeping the package on numpy and scipy only, and deleted check_operator_shapes(), which was never referenced.
Reviewer suggestion on aprilweilab#77 was to replace the operator definitions with pylops.FunctionOperator, on the grounds that pylops only needs numpy and scipy so it is a cheap dependency. Measured it rather than assumed, and it is not cheap for this package. FunctionOperator supplies only _matvec, so the inherited matmat falls back to looping over columns one at a time -- f() invoked 50 times at width 1 rather than once at width 50. At N = 4000 on a 50-column block that is 8.70 ms against 2.00 ms, a 4.4x penalty on the operation this package leans on hardest: trace estimation and the Nystrom sketch both apply operators to 30-100 columns at once. The penalty would be worse for a real ARG or GRG operator, where a single block product amortizes one graph traversal across every column. pylops' operator algebra (a * Op1 + b * Op2) degrades identically; scipy's own algebra does not. A pylops.LinearOperator subclass with an explicit _matmat does preserve block products, including through the algebra, so pylops is usable -- but that subclass is what this package already has, so adopting pylops would buy no code reduction while still costing the dependency. And the dependency is not free: importing pylops adds 76 MB of resident memory (42.6 -> 118.5 MB) and 821 ms of startup (350 -> 1171 ms) over numpy and scipy alone, for a package whose selling point is being light enough to vendor into another project. No code change. The _from_matmat docstring and the package README now record this so the next reader does not have to re-measure it. pylops stays an optional test-only cross-check for the Hutch++ implementation, which skips when it is absent.
aireml_benchmark.py runs many fits inside one process, which is fine for
wall clock but useless for memory: the high-water mark it would report is
whatever the largest fit needed plus everything the earlier fits left
behind. examples/aireml_timing.py does exactly one solve per process, so
an external profiler measures that solve alone.
Three modes:
* make - simulate and store a dataset (sparse genotypes plus a
phenotype drawn from the LMM), once per size, so the timed run
does not pay for simulation
* solve - load one dataset and fit it, reporting load time, fit time,
iterations, products, the estimates, and its own peak RSS;
--json for machine-readable output
* trials - drive a sweep: make any missing datasets, re-run each solve in
a fresh process under a profiler, parse the peak RSS back out,
and print one table with per-individual time and memory plus
fitted scaling exponents
Profiler handling is the fiddly part and is done once here. GNU time
reports "Maximum resident set size (kbytes)" under -v; BSD and macOS time
report "maximum resident set size" in *bytes* under -l, which is an easy
way to be wrong by a factor of 1024. Both dialects are detected and
parsed, and macOS /usr/bin/time has no -v at all, so gtime is preferred
when present. ru_maxrss has the same bytes-vs-kilobytes split and is
handled alongside.
The solve mode also accepts --grg, taking relatedness straight off a GRG
through grapp's existing SciPyStdXOperator, and trials --grg adds a
"grapp pca -d 10" row as a memory baseline on the same data. That is the
apples-to-apples comparison against an established GRG method. It needs
grapp's own dependencies (pygrgl, pandas), which aireml does not; the
sweep runs without them and reports the rows as unavailable rather than
failing. aireml itself still imports nothing from grapp.
examples/aireml_timing.py defaults --data-dir to timing-data/, so running a sweep from the repo root leaves a few hundred MB of .npz genotype matrices as untracked files.
Generating the N = 64000 dataset took over 150 seconds and was still running when killed; at N = 128000 it would have been far worse. Two causes, both in make_dataset: scipy.sparse.random samples positions without replacement from the full size-by-variants grid, so asking for 3.2M nonzeros out of 320M cells cost minutes. Sampling positions directly with replacement is O(nnz) instead. About 0.5% of positions collide and simply sum, which is harmless for synthetic genotypes and leaves nnz within 0.5% of target. The group-structure block was built with nested Python loops extending lists, 25 groups x 50 columns x thousands of members at a time. Replaced with numpy.tile and numpy.repeat per group, which produces a bit-identical matrix (verified at N = 2000 and 8000) and is 11x faster on that step. Generation now takes 3.2s at N = 64000 and 7.3s at N = 128000, against 150s+ before for the smaller of the two. Because the random stream changed, datasets cached by the old code would no longer match freshly generated ones. meta.json now records a generator version, and trials regenerates a dataset whose cached version, variant count, density or group count does not match what was asked for, rather than silently mixing old and new data in one table. Previously reported timings stay comparable: nnz is unchanged by construction, so only the particular random values differ.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Added standalone AI REML Package