Skip to content

Try some optimizations #117

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Kss2k wants to merge 7 commits into from
Closed

Try some optimizations #117

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
a4a6a24
Try some optimizations
Kss2k Jan 13, 2026
ee83aa4
Add benchmark workflow
Kss2k Jan 13, 2026
15b46ad
Fix multiplication ordering
Kss2k Jan 13, 2026
be5c115
Increase the number of benchmark iterations
Kss2k Jan 13, 2026
7b97242
Split benchmark workflow
Kss2k Jan 13, 2026
7751d65
Revert R/equations_lms.R
Kss2k Jan 13, 2026
67c0ba9
Make som small optimizations in `src/equations_qml.cpp`
Kss2k Jan 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
80 changes: 61 additions & 19 deletions R/equations_lms.R
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,10 +106,12 @@ estepLmsGroup <- function(submodel, lastQuad = NULL, recalcQuad = FALSE,
data.id <- data$data.split[[j]]
colidx <- data$colidx[[j]]

pj <- p[offset:end]
wm <- colSums(data.id * pj) / sum(pj)
X <- data.id - matrix(wm, nrow=nrow(data.id), ncol=ncol(data.id), byrow=TRUE)
wcov <- t(X) %*% (X * pj)
pj <- p[offset:end]
pj_sum <- sum(pj)
wm <- drop(crossprod(pj, data.id)) / pj_sum
centered <- sweep(data.id, 2, wm, "-")
weighted <- sweep(centered, 1, sqrt(pj), "*")
wcov <- crossprod(weighted)

wMeans[[i]][[j]] <- wm
wCovs[[i]][[j]] <- wcov
Expand Down Expand Up @@ -137,14 +139,32 @@ mstepLms <- function(theta, model, P,
optim.method = "L-BFGS-B",
epsilon = 1e-6,
...) {
getFilled <- local({
cache <- new.env(parent = emptyenv())
cache$theta <- NULL
cache$mod <- NULL

function(theta) {
if (!is.null(cache$theta) && identical(cache$theta, theta)) {
cache$mod
} else {
cache$theta <- theta
cache$mod <- fillModel(model = model, theta = theta, method = "lms")
cache$mod
}
}
})

gradient <- function(theta) {
modFilled <- getFilled(theta)
gradientCompLogLikLms(theta = theta, model = model, P = P, sign = -1,
epsilon = epsilon)
epsilon = epsilon, modFilled = modFilled)
}

objective <- function(theta) {
modFilled <- getFilled(theta)
compLogLikLms(theta = theta, model = model, P = P, sign = -1,
epsilon = epsilon)
epsilon = epsilon, modFilled = modFilled)
}

if (optimizer == "nlminb") {
Expand Down Expand Up @@ -172,9 +192,10 @@ mstepLms <- function(theta, model, P,
}


compLogLikLms <- function(theta, model, P, sign = -1, ...) {
compLogLikLms <- function(theta, model, P, sign = -1, modFilled = NULL, ...) {
tryCatch({
modFilled <- fillModel(model = model, theta = theta, method = "lms")
if (is.null(modFilled))
modFilled <- fillModel(model = model, theta = theta, method = "lms")

ll <- 0
for (g in seq_len(model$info$n.groups)) {
Expand Down Expand Up @@ -211,24 +232,34 @@ compLogLikLmsGroup <- function(submodel, P, sign = -1, ...) {



gradientCompLogLikLms <- function(theta, model, P, sign = -1, epsilon = 1e-6) {
gradientCompLogLikLms <- function(theta, model, P, sign = -1, epsilon = 1e-6,
modFilled = NULL) {
gradientAllLogLikLms(theta = theta, model = model, P = P, sign = sign,
epsilon = epsilon, FGRAD = gradLogLikLmsCpp, FOBJECTIVE = compLogLikLmsGroup)
epsilon = epsilon, FGRAD = gradLogLikLmsCpp,
FOBJECTIVE = compLogLikLmsGroup, modFilled = modFilled)
}


gradientAllLogLikLms <- function(theta, model, P, sign = -1, epsilon = 1e-6,
FGRAD, FOBJECTIVE) {
FGRAD, FOBJECTIVE, modFilled = NULL) {
hasCovModel <- model$params$gradientStruct$hasCovModel

if (hasCovModel) gradient <- \(...) complicatedGradientAllLogLikLms(..., FOBJECTIVE = FOBJECTIVE)
else gradient <- \(...) simpleGradientAllLogLikLms(..., FGRAD = FGRAD)
if (hasCovModel) {
gradient <- \(...) complicatedGradientAllLogLikLms(...,
FOBJECTIVE = FOBJECTIVE, modFilled = modFilled, theta.base = theta)
} else {
gradient <- \(...) simpleGradientAllLogLikLms(...,
FGRAD = FGRAD, modFilled = modFilled)
}

c(gradient(theta = theta, model = model, P = P, sign = sign, epsilon = epsilon))
}


complicatedGradientAllLogLikLms <- function(theta, model, P, sign = -1, epsilon = 1e-4, FOBJECTIVE, sum = TRUE, ...) {
complicatedGradientAllLogLikLms <- function(theta, model, P, sign = -1,
epsilon = 1e-4, FOBJECTIVE,
sum = TRUE, modFilled = NULL,
theta.base = NULL, ...) {
params <- model$params

SELECT_THETA_LAB <- params$SELECT_THETA_LAB
Expand All @@ -246,9 +277,17 @@ complicatedGradientAllLogLikLms <- function(theta, model, P, sign = -1, epsilon

grad <- matrix(0, nrow = n, ncol = k, dimnames = list(NULL, names(theta)))

FOBJECTIVE_GROUP <- function(theta, g) {
modFilled <- fillModel(theta = theta, model = model, method = "lms")
FOBJECTIVE(submodel = modFilled$models[[g]], P = P$P_GROUPS[[g]], sign = sign, ...)
theta.base <- if (is.null(theta.base)) theta else theta.base
use_cache <- !is.null(modFilled)

FOBJECTIVE_GROUP <- function(theta.eval, g) {
if (use_cache && identical(theta.eval, theta.base)) {
submodel <- modFilled$models[[g]]
} else {
modFilledEval <- fillModel(theta = theta.eval, model = model, method = "lms")
submodel <- modFilledEval$models[[g]]
}
FOBJECTIVE(submodel = submodel, P = P$P_GROUPS[[g]], sign = sign, ...)
}

for (g in seq_len(model$info$n.groups)) {
Expand Down Expand Up @@ -276,10 +315,13 @@ complicatedGradientAllLogLikLms <- function(theta, model, P, sign = -1, epsilon
}


simpleGradientAllLogLikLms <- function(theta, model, P, sign = -1, epsilon = 1e-6, FGRAD, N.FGRAD = 1L) {
simpleGradientAllLogLikLms <- function(theta, model, P, sign = -1, epsilon = 1e-6,
FGRAD, N.FGRAD = 1L, modFilled = NULL) {
# simple gradient which should work if constraints are well-behaved functions
# which can be derivated by Deriv::Deriv, and there is no covModel
modelR <- fillModel(model=model, theta=theta, method="lms")
modelR <- if (is.null(modFilled))
fillModel(model = model, theta = theta, method = "lms")
else modFilled
locations <- model$params$gradientStruct$locations
Jacobian <- model$params$gradientStruct$Jacobian
nlinDerivs <- model$params$gradientStruct$nlinDerivs
Expand Down
35 changes: 19 additions & 16 deletions src/equations_lms.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ inline arma::vec make_zvec(unsigned k, unsigned numXis, const arma::vec& z) {

struct LMSModel {
arma::mat A, Oxx, Oex, Ie, lY, lX, tY, tX, Gx, Ge,
a, beta0, Psi, d, e;
a, beta0, Psi, d, e, Oi;
unsigned k = 0;
unsigned numXis = 0;

Expand Down Expand Up @@ -133,31 +133,33 @@ struct LMSModel {
Psi = Rcpp::as<arma::mat>(matrices["psi"]);
d = Rcpp::as<arma::mat>(matrices["thetaDelta"]);
e = Rcpp::as<arma::mat>(matrices["thetaEpsilon"]);
Oi = make_Oi(k, numXis);
}

arma::vec mu(const arma::vec& z) const {
const arma::vec zVec = make_zvec(k, numXis, z);
const arma::mat kronZ = arma::kron(Ie, beta0 + A * zVec);
const arma::mat Binv = arma::inv(Ie - Ge - kronZ.t() * Oex);

const arma::vec muXi = beta0 + A * zVec;
const arma::vec muEta = Binv * (a +
Gx * (beta0 + A * zVec) +
kronZ.t() * Oxx * (beta0 + A * zVec));
const arma::vec muXi = beta0 + A * zVec;
const arma::mat kronZ = arma::kron(Ie, muXi);
const arma::mat B = Ie - Ge - kronZ.t() * Oex;
const arma::vec rhs = a + Gx * muXi + kronZ.t() * Oxx * muXi;
const arma::vec muEta = arma::solve(B, rhs, arma::solve_opts::fast);

return tX + lX * arma::join_cols(muXi, muEta);
}

arma::mat Sigma(const arma::vec& z) const {
const arma::vec zVec = make_zvec(k, numXis, z);
const arma::mat kronZ = arma::kron(Ie, beta0 + A * zVec);
const arma::mat Binv = arma::inv(Ie - Ge - kronZ.t() * Oex);

const arma::mat Oi = make_Oi(k, numXis);
const arma::mat Eta = Binv * (Gx * A + kronZ.t() * Oxx * A);
const arma::mat varXi = A * Oi * A.t();
const arma::mat varEta = Eta * Oi * Eta.t() + Binv * Psi * Binv.t();
const arma::mat covXiEta = A * Oi * Eta.t();
const arma::vec muXi = beta0 + A * zVec;
const arma::mat kronZ = arma::kron(Ie, muXi);
const arma::mat B = Ie - Ge - kronZ.t() * Oex;
const arma::mat rhs = Gx * A + kronZ.t() * Oxx * A;
const arma::mat Eta = arma::solve(B, rhs, arma::solve_opts::fast);
const arma::mat AOi = A * Oi;
const arma::mat varXi = AOi * A.t();
const arma::mat tmp = arma::solve(arma::trans(B), Psi, arma::solve_opts::fast);
const arma::mat noise = arma::solve(B, tmp, arma::solve_opts::fast);
Copy link

@chatgpt-codex-connector chatgpt-codex-connector bot Jan 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Restore correct order for B^{-1} Psi B^{-T}

The new computation of noise changes the multiplication order from B^{-1} * Psi * B^{-T} to B^{-1} * B^{-T} * Psi because tmp = solve(trans(B), Psi) yields B^{-T} * Psi and then solve(B, tmp) left-multiplies by B^{-1}. These are only equivalent when Psi commutes with B^{-T} (e.g., Psi is diagonal in the same basis), which is not generally true. For non-diagonal Psi, this produces an incorrect covariance term and will skew Sigma(z)/likelihood results. Consider solving in a way that preserves B^{-1} * Psi * B^{-T} (e.g., solve left for B^{-1} * Psi, then right-solve for B^{-T}).

Useful? React with 👍 / 👎.

const arma::mat varEta = Eta * Oi * Eta.t() + noise;
const arma::mat covXiEta = AOi * Eta.t();

const arma::mat vcovXiEta = arma::join_cols(
arma::join_rows(varXi, covXiEta),
Expand Down Expand Up @@ -185,6 +187,7 @@ struct LMSModel {
c.Psi = arma::mat(Psi);
c.d = arma::mat(d);
c.e = arma::mat(e);
c.Oi = arma::mat(Oi);

return c;
}
Expand Down
Loading