Skip to content

Commit

Permalink
Add tests for diagnostic, make the functions more robust to types, fi…
Browse files Browse the repository at this point in the history
…x a small bug in constraint
  • Loading branch information
JamesYang007 committed Oct 21, 2024
1 parent ca811b2 commit 99803f1
Show file tree
Hide file tree
Showing 4 changed files with 170 additions and 18 deletions.
6 changes: 5 additions & 1 deletion adelie/constraint.py
Original file line number Diff line number Diff line change
Expand Up @@ -249,7 +249,11 @@ def linear(
else:
assert not (vars is None)

A_dtype = np.float32 if "32" in A.__class__.__name__ else np.float64
A_dtype = (
np.float32
if isinstance(A, MatrixConstraintBase32) else
np.float64
)

lower, l_dtype = _coerce_dtype(lower, dtype)
upper, u_dtype = _coerce_dtype(upper, dtype)
Expand Down
56 changes: 45 additions & 11 deletions adelie/diagnostic.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,7 @@ def predict(
linear_preds : (L, n) or (L, n, K) ndarray
Linear predictions.
"""
intercepts = np.atleast_1d(intercepts)
is_multi = len(intercepts.shape) == 2
if is_multi:
K = intercepts.shape[1]
Expand All @@ -93,12 +94,21 @@ def predict(
n = X.rows()
y_shape = (n,)

dtype = (
np.float32
if isinstance(X, MatrixNaiveBase32) else
np.float64
)

if offsets is None:
offsets = np.zeros(y_shape)
offsets = np.zeros(y_shape, dtype=dtype)

if isinstance(betas, np.ndarray):
betas = np.atleast_2d(betas)

L = betas.shape[0]

etas = np.zeros((L,) + y_shape, order="C")
etas = np.zeros((L,) + y_shape, order="C", dtype=dtype)
if isinstance(betas, np.ndarray):
for i in range(etas.shape[0]):
X.btmul(0, X.cols(), betas[i], etas[i].ravel())
Expand Down Expand Up @@ -213,8 +223,15 @@ def objective(
group_sizes = np.concatenate([groups, [p]], dtype=int)
group_sizes = group_sizes[1:] - group_sizes[:-1]


dtype = (
np.float32
if isinstance(X, MatrixNaiveBase32) else
np.float64
)

if penalty is None:
penalty = np.sqrt(group_sizes)
penalty = np.sqrt(group_sizes).astype(dtype)

etas = predict(
X=X_raw,
Expand All @@ -228,7 +245,7 @@ def objective(
objs = np.array([
glm.loss(etas[i])
for i in range(etas.shape[0])
])
], dtype=dtype)

# relative to saturated model
if relative:
Expand All @@ -238,9 +255,15 @@ def objective(
if add_penalty:
penalty_f = None
if isinstance(betas, np.ndarray):
penalty_f = core.solver.compute_penalty_dense
penalty_f = {
np.float32: core.solver.compute_penalty_dense_32,
np.float64: core.solver.compute_penalty_dense_64,
}[dtype]
elif isinstance(betas, csr_matrix):
penalty_f = core.solver.compute_penalty_sparse
penalty_f = {
np.float32: core.solver.compute_penalty_sparse_32,
np.float64: core.solver.compute_penalty_sparse_64,
}[dtype]
objs += lmdas * penalty_f(
groups,
group_sizes,
Expand Down Expand Up @@ -283,7 +306,12 @@ def residuals(
--------
adelie.diagnostic.predict
"""
resids = np.empty(etas.shape)
dtype = (
np.float32
if isinstance(glm, (GlmBase32, GlmMultiBase32)) else
np.float64
)
resids = np.empty(etas.shape, dtype=dtype)
for eta, resid in zip(etas, resids):
glm.gradient(eta, resid)
return resids
Expand Down Expand Up @@ -345,7 +373,13 @@ def gradients(
X = matrix.dense(X, method="naive", n_threads=n_threads)
grad_shape = (X.cols(),)

grads = np.empty((resids.shape[0],) + grad_shape)
dtype = (
np.float32
if isinstance(X, MatrixNaiveBase32) else
np.float64
)

grads = np.empty((resids.shape[0],) + grad_shape, dtype=dtype)
ones = np.ones(np.prod(resids.shape[1:]))
for i in range(grads.shape[0]):
X.mul(resids[i].ravel(), ones, grads[i].ravel())
Expand Down Expand Up @@ -463,9 +497,9 @@ def gradient_norms(
dual_groups = render_dual_groups(constraints)
mu_grads = np.zeros(grads.shape, dtype=dtype)
for k in range(L):
beta_curr = betas[k].toarray()[0]
mu_curr = duals[k].toarray()[0]
mu_grads_curr = mu_grads[k]
beta_curr = betas[k].toarray()[0].astype(dtype)
mu_curr = duals[k].toarray()[0].astype(dtype)
mu_grads_curr = mu_grads[k].astype(dtype)
for constraint, g, gs, dg in zip(
constraints,
groups,
Expand Down
14 changes: 8 additions & 6 deletions adelie/src/py_solver.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,8 @@ namespace ad = adelie_core;
// =================================================================
template <class T>
ad::util::rowvec_type<T> compute_penalty_sparse(
const Eigen::Ref<const ad::util::rowvec_type<int>>& groups,
const Eigen::Ref<const ad::util::rowvec_type<int>>& group_sizes,
const Eigen::Ref<const ad::util::rowvec_type<Eigen::Index>>& groups,
const Eigen::Ref<const ad::util::rowvec_type<Eigen::Index>>& group_sizes,
const Eigen::Ref<const ad::util::rowvec_type<T>>& penalty,
T alpha,
const Eigen::SparseMatrix<T, Eigen::RowMajor>& betas,
Expand Down Expand Up @@ -52,8 +52,8 @@ ad::util::rowvec_type<T> compute_penalty_sparse(

template <class T>
ad::util::rowvec_type<T> compute_penalty_dense(
const Eigen::Ref<const ad::util::rowvec_type<int>>& groups,
const Eigen::Ref<const ad::util::rowvec_type<int>>& group_sizes,
const Eigen::Ref<const ad::util::rowvec_type<Eigen::Index>>& groups,
const Eigen::Ref<const ad::util::rowvec_type<Eigen::Index>>& group_sizes,
const Eigen::Ref<const ad::util::rowvec_type<T>>& penalty,
T alpha,
const Eigen::Ref<const ad::util::rowmat_type<T>>& betas,
Expand Down Expand Up @@ -92,6 +92,8 @@ ad::util::rowvec_type<T> compute_penalty_dense(
void register_solver(py::module_& m)
{
/* helper functions */
m.def("compute_penalty_sparse", &compute_penalty_sparse<double>);
m.def("compute_penalty_dense", &compute_penalty_dense<double>);
m.def("compute_penalty_sparse_32", &compute_penalty_sparse<float>);
m.def("compute_penalty_sparse_64", &compute_penalty_sparse<double>);
m.def("compute_penalty_dense_32", &compute_penalty_dense<float>);
m.def("compute_penalty_dense_64", &compute_penalty_dense<double>);
}
112 changes: 112 additions & 0 deletions tests/test_diagnostic.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,112 @@
import adelie as ad
import numpy as np
import pytest
import scipy


def generate_data(n, p, L, K, beta_type, dtype, seed):
np.random.seed(seed)
X = np.random.normal(0, 1, (n, p)).astype(dtype)
y = X @ np.random.normal(0, 1, (p, K)) + np.random.normal(0, 1, (n, K))
y = y.astype(dtype)
betas = np.random.uniform(-1, 1, (L, p*K)).astype(dtype)
if beta_type == "sparse":
betas = scipy.sparse.csr_matrix(betas, dtype=dtype)
if K == 1:
y = y.squeeze(axis=1)
intercepts = np.random.normal(0, 1, (L,)).astype(dtype)
else:
intercepts = np.random.normal(0, 1, (L, K)).astype(dtype)
lmdas = np.sort(np.random.uniform(0, 1, L))[::-1]

return {
"X": X,
"y": y,
"betas": betas,
"intercepts": intercepts,
"lmdas": lmdas,
}


def predict(X, betas, intercepts):
p = X.shape[1]
L = betas.shape[0]
K = 1 if len(intercepts.shape) == 1 else intercepts.shape[1]
if K != 1:
if isinstance(betas, scipy.sparse.csr_matrix):
betas = betas.toarray()
betas = betas.reshape((L, p, K))
Xbetas = np.einsum("ij,ljk->lik", X, betas)
else:
Xbetas = betas @ X.T
return intercepts[:, None] + Xbetas


@pytest.mark.filterwarnings("ignore: Detected matrix to be C-contiguous.")
@pytest.mark.parametrize("dtype", [np.float32, np.float64])
@pytest.mark.parametrize("beta_type", ["dense", "sparse"])
@pytest.mark.parametrize("K", [1, 3])
@pytest.mark.parametrize("L", [1, 2])
@pytest.mark.parametrize("p", [1, 10])
@pytest.mark.parametrize("n", [2, 20])
def test_predict(n, p, L, K, beta_type, dtype, seed=0):
data = generate_data(n, p, L, K, beta_type, dtype, seed)
X = data["X"]
betas = data["betas"]
intercepts = data["intercepts"]
actual = ad.diagnostic.predict(X, betas, intercepts)
expected = predict(X, betas, intercepts)

assert np.allclose(actual, expected, atol=1e-6)


@pytest.mark.filterwarnings("ignore: Detected matrix to be C-contiguous.")
@pytest.mark.parametrize("dtype", [np.float32, np.float64])
@pytest.mark.parametrize("beta_type", ["dense", "sparse"])
@pytest.mark.parametrize("K", [1, 3])
@pytest.mark.parametrize("L", [1, 2])
@pytest.mark.parametrize("p", [1, 10])
@pytest.mark.parametrize("n", [2, 20])
def test_objective(n, p, L, K, beta_type, dtype, seed=0):
def _objective(X, glm, betas, intercepts, lmdas):
etas = predict(X, betas, intercepts)
losses = np.array([glm.loss(eta) - glm.loss_full() for eta in etas])
if isinstance(betas, scipy.sparse.csr_matrix):
betas = betas.toarray()
if K == 1:
penalty = np.sum(np.abs(betas), axis=-1)
else:
betas = betas.reshape((L, p, K))
penalty = np.sum(np.linalg.norm(betas, axis=-1), axis=-1)
return losses + lmdas * np.sqrt(K) * penalty

data = generate_data(n, p, L, K, beta_type, dtype, seed)
X = data["X"]
y = data["y"]
betas = data["betas"]
intercepts = data["intercepts"]
lmdas = data["lmdas"]
if K == 1:
glm = ad.glm.gaussian(y)
else:
glm = ad.glm.multigaussian(y)

actual = ad.diagnostic.objective(X, glm, betas, intercepts, lmdas)
expected = _objective(X, glm, betas, intercepts, lmdas)

assert np.allclose(actual, expected, atol=1e-6)


@pytest.mark.filterwarnings("ignore: Detected matrix to be C-contiguous.")
@pytest.mark.parametrize("dtype", [np.float32, np.float64])
def test_diagnostic(dtype, seed=0):
n = 10
p = 10
np.random.seed(seed)
X = np.random.normal(0, 1, (n, p)).astype(dtype)
y = X @ np.random.normal(0, 1, p) + np.random.normal(0, 1, n)
y = y.astype(dtype)
constraints = [None] * p
constraints[0] = ad.constraint.lower(np.full(1, -1, dtype=dtype))
state = ad.grpnet(X, ad.glm.gaussian(y), progress_bar=False)
ad.diagnostic.diagnostic(state)

0 comments on commit 99803f1

Please sign in to comment.