test_gel.py

"""test_gel.py: framework to test gel implementations."""

import itertools
import os
import unittest

import cvxpy as cvx
import numpy as np
from scipy.spatial.distance import cosine
import torch

from gel.gelcd import (
    block_solve_agd,
    block_solve_newton,
    gel_solve as gel_solve_cd,
    make_A as make_A_cd,
)
from gel.gelfista import gel_solve as gel_solve_fista, make_A as make_A_fista


def gel_solve_cvx(As, y, l_1, l_2, ns):
    """Solve a group elastic net problem with cvx.

    Arguments:
        As: list of tensors.
        y: tensor.
        l_1, l_2: floats.
        ns: iterable.
    """
    # Convert everything to numpy
    dtype = As[0].dtype
    As = [A_j.cpu().numpy() for A_j in As]
    y = y.cpu().numpy()
    ns = np.array([int(n) for n in ns])

    # Create the b variables.
    b_0 = cvx.Variable()
    bs = []
    for _, n_j in zip(As, ns):
        bs.append(cvx.Variable(n_j))

    # Form g(b).
    Ab = sum(A_j * b_j for A_j, b_j in zip(As, bs))
    m = As[0].shape[0]
    g_b = cvx.square(cvx.norm(y - b_0 - Ab)) / (2 * m)

    # Form h(b).
    h_b = sum(
        np.sqrt(n_j) * (l_1 * cvx.norm(b_j) + l_2 * cvx.square(cvx.norm(b_j)))
        for n_j, b_j in zip(ns, bs)
    )

    # Build the optimization problem.
    obj = cvx.Minimize(g_b + h_b)
    problem = cvx.Problem(obj, constraints=None)

    problem.solve(solver="CVXOPT")

    b_0 = b_0.value.item()
    # Form B as returned by gel_solve.
    p = len(As)
    B = torch.zeros(p, int(max(ns)), dtype=dtype)
    for j in range(p):
        b_j = np.asarray(bs[j].value)
        B[j, : ns[j]] = torch.from_numpy(b_j)

    return b_0, B


def block_solve_cvx(r_j, A_j, a_1_j, a_2_j, m, b_j_init, verbose=False):
    # pylint: disable=unused-argument
    """Solve the gelcd optimization problem for a single block with cvx.

    b_j_init and verbose are ignored. b_j_init because cvx doesn't support it.
    verbose because it doesn't go together with tqdm.
    """
    # Convert everything to numpy.
    device = A_j.device
    dtype = A_j.dtype
    r_j = r_j.cpu().numpy()
    A_j = A_j.cpu().numpy()

    # Create the b_j variable.
    b_j = cvx.Variable(A_j.shape[1])

    # Form the objective.
    q_j = r_j - A_j * b_j
    obj_fun = cvx.square(cvx.norm(q_j)) / (2.0 * m)
    obj_fun += a_1_j * cvx.norm(b_j) + (a_2_j / 2.0) * cvx.square(cvx.norm(b_j))

    # Build the optimization problem.
    obj = cvx.Minimize(obj_fun)
    problem = cvx.Problem(obj, constraints=None)

    problem.solve(solver="CVXOPT", verbose=False)
    b_j = np.asarray(b_j.value)
    return torch.from_numpy(b_j).to(device, dtype)


def _b2vec(B, groups):
    """Convert B as returned by gel_solve functions to a single numpy vector."""
    d = sum(len(group_j) for group_j in groups)  # the total dimension
    b = np.zeros(d, dtype=B[0, 0].cpu().numpy().dtype)
    for j, group_j in enumerate(groups):
        b[group_j] = B[j, : len(group_j)].cpu().numpy()
    return b


class TestGelBirthwtBase:

    """Base class to test different gel_solve implementations with the birth
    weight data."""

    l_1_base = 4.0
    l_2_base = 0.5

    def __init__(self, device, dtype, *args, **kwargs):
        """Load data and solve with cvx to get ground truth solution."""
        super().__init__(*args, **kwargs)
        self.device = device
        self.dtype = dtype
        dtype_np = torch.rand(0, dtype=dtype).numpy().dtype

        data_dir = os.path.join(os.path.dirname(__file__), "data", "birthwt")
        self.X = np.loadtxt(
            os.path.join(data_dir, "X.csv"), skiprows=1, delimiter=",", dtype=dtype_np
        )
        self.y = np.loadtxt(os.path.join(data_dir, "y.csv"), skiprows=1, dtype=dtype_np)
        self.m = len(self.y)
        self.l_1 = self.l_1_base / (2 * self.m)
        self.l_2 = self.l_2_base / (2 * self.m)
        self.groups = [
            [0, 1, 2],
            [3, 4, 5],
            [6, 7],
            [8],
            [9, 10],
            [11],
            [12],
            [13, 14, 15],
        ]
        self.done_setup = False

    def setUp(self):
        if self.device.type == "cuda" and not torch.cuda.is_available():
            raise unittest.SkipTest("cuda unavailable")

        if self.done_setup:
            return

        self.ns = torch.tensor([len(g) for g in self.groups])
        self.p = len(self.groups)

        # Convert things to gel format.
        self.As = []
        for j in range(self.p):
            A_j = self.X[:, self.groups[j]]
            self.As.append(torch.from_numpy(A_j).to(self.device, self.dtype))
        self.yt = torch.from_numpy(self.y).to(self.device, self.dtype)

        # Solve with cvx.
        self.b_0_cvx, self.B_cvx = gel_solve_cvx(
            self.As, self.yt, self.l_1, self.l_2, self.ns
        )
        self.b_cvx = _b2vec(self.B_cvx, self.groups)
        self.obj_cvx = self._obj(self.b_0_cvx, self.b_cvx)

        self.done_setup = True

    def _obj(self, b_0, b):
        """Compute the objective function value for the given b_0, b."""
        r = self.y - b_0 - self.X @ b
        g_b = r @ r / (2.0 * self.m)
        b_j_norms = [np.linalg.norm(b[self.groups[j]], ord=2) for j in range(self.p)]
        h_b = self.l_1 * sum(
            np.sqrt(len(self.groups[j])) * b_j_norms[j] for j in range(self.p)
        )
        h_b += self.l_2 * sum(
            np.sqrt(len(self.groups[j])) * (b_j_norms[j] ** 2) for j in range(self.p)
        )
        return g_b + h_b

    def _compare_to_cvx(self, b_0, b, obj):
        """Compare the given solution to the cvx solution."""
        # pylint: disable=no-member
        self.assertAlmostEqual(obj, self.obj_cvx, places=2)
        self.assertAlmostEqual(b_0, self.b_0_cvx, places=2)
        if np.allclose(b, 0) or np.allclose(self.b_cvx, 0):
            for b_i, b_cvx_i in zip(b, self.b_cvx):
                self.assertAlmostEqual(b_i, b_cvx_i, places=2)
        else:
            self.assertAlmostEqual(cosine(b, self.b_cvx), 0, places=2)

    def _test_implementation(self, make_A, gel_solve, **gel_solve_kwargs):
        """Test the given implementation."""
        A = make_A(self.As, self.ns, self.device, self.dtype)
        b_0, B = gel_solve(A, self.yt, self.l_1, self.l_2, self.ns, **gel_solve_kwargs)
        b = _b2vec(B, self.groups)
        obj = self._obj(b_0, b)
        self._compare_to_cvx(b_0, b, obj)

    def test_fista(self):
        """Test the FISTA implementation of gel_solve."""
        self._test_implementation(
            make_A_fista,
            gel_solve_fista,
            t_init=0.1,
            ls_beta=0.9,
            max_iters=1000,
            rel_tol=1e-6,
        )

    def test_cd_cvx(self):
        """Test the CD implementation with cvx internal solver."""
        self._test_implementation(
            make_A_cd,
            gel_solve_cd,
            block_solve_fun=block_solve_cvx,
            block_solve_kwargs={},
            max_cd_iters=100,
            rel_tol=1e-6,
        )

    def test_cd_agd(self):
        """Test the CD implementation with AGD internal solver."""
        self._test_implementation(
            make_A_cd,
            gel_solve_cd,
            block_solve_fun=block_solve_agd,
            block_solve_kwargs={
                "t_init": 1,
                "ls_beta": 0.5,
                "max_iters": 100,
                "rel_tol": 1e-5,
            },
            max_cd_iters=100,
            rel_tol=1e-6,
        )

    def test_cd_newton(self):
        """Test the CD implementation with Newton internal solver."""
        # Compute the C_js and I_js.
        Cs = [(A_j.t() @ A_j) / self.m for A_j in self.As]
        Is = [torch.eye(n_j, device=self.device, dtype=self.dtype) for n_j in self.ns]
        self._test_implementation(
            make_A_cd,
            gel_solve_cd,
            block_solve_fun=block_solve_newton,
            block_solve_kwargs={
                "ls_alpha": 0.1,
                "ls_beta": 0.5,
                "max_iters": 10,
                "tol": 1e-10,
            },
            max_cd_iters=100,
            rel_tol=1e-6,
            Cs=Cs,
            Is=Is,
        )


def create_gel_birthwt_test(device_name, dtype, *mods):
    # I'm so sorry.
    device = torch.device(device_name)

    def __init__(self, *args, **kwargs):
        TestGelBirthwtBase.__init__(self, device, dtype, *args, **kwargs)
        for mod in mods:
            if mod == "l10":
                self.l_1 = 0
            elif mod == "l20":
                self.l_2 = 0
            elif mod == "nj1":
                self.groups = [[i] for i in range(self.X.shape[1])]
            else:
                raise RuntimeError("unrecognized mod: " + mod)

    _doc = "Test gel implementations on {} with {}".format(device_name, dtype)
    if mods:
        _doc += " (mods: " + ", ".join(mods) + ")"
    test_name = "TestGelBirthwt" + device_name.upper() + str(dtype)[-2:]
    if mods:
        test_name += "_" + "".join(str(m) for m in mods)

    globals()[test_name] = type(
        test_name,
        (TestGelBirthwtBase, unittest.TestCase),
        {"__init__": __init__, "__doc__": _doc},
    )


_mods = ["l10", "l20", "nj1"]
_mod_subsets = set(
    frozenset(s) for s in itertools.combinations_with_replacement(_mods, len(_mods))
)
_mod_subsets.add(frozenset())

for _device_name, _dtype, _mod_subset in itertools.product(
    ["cpu", "cuda"], [torch.float32, torch.float64], _mod_subsets
):
    create_gel_birthwt_test(_device_name, _dtype, *list(sorted(_mod_subset)))