Hail_ordinal_reg_dummy_0.py

gistfile1.txt

# https://discuss.hail.is/t/ordinal-regression/1387/13

import numpy as np
rng = np.random.default_rng(seed=42)
import hail as hl

hl.init()

# dummy data generation

def sigmoid(arr):
    result = np.empty(arr.shape, dtype='float')
    positive = arr > 0
    negative = ~positive
    result[positive] = 1 / (1 + np.exp(-arr[positive]))
    exp_arr = np.exp(arr[negative])
    result[negative] = exp_arr / (exp_arr + 1)
    return result

k = 3
m = 4
n = 123
_beta = np.array([0.1, 0.2, 1, 2])
_theta = np.array([10, 15])
X = rng.uniform(0, 8, (n, m))

_Xbeta = X @ _beta
_p = sigmoid(_theta[None, :] - _Xbeta[:, None] + 2 * rng.standard_normal(n)[:, None])
y = 1 + (_p < 0.5).sum(axis=1)


# some definitions for Hail

hl_y = hl.nd.array(y)
hl_X = hl.nd.array(X)

def hl_gte(nd_left, nd_right):
    # >=
    return hl.map(
        lambda _: hl.float(_>=0),
        (nd_left - nd_right)
    )

def hl_scalar_sigmoid(x):
    return hl.if_else(
        x > 0,
        1 / (1 + hl.exp(-x)),
        hl.rbind(
            hl.exp(x),
            lambda _: _ / (1 + _)
        )
    )

def hl_sigmoid(hl_nd):
    return hl_nd.map(hl_scalar_sigmoid)

def hl_log(nd):
    return nd.map(lambda _: hl.log(_))

def hl_logit(nd):
    return hl_log(nd / (1 - nd))

def hl_mean(nd, dim):
    return nd.sum(axis=dim) / nd.shape[dim]

def hl_diff(nd):
    return nd[1:] - nd[:-1]

def hl_get_t(k):
    # lower triangle, np.tril
    i = hl.nd.arange(1, k)
    return hl_gte(
        i.reshape(-1, 1),
        i.reshape(1, -1)
    )

def hl_get_s(y, k):
    return 2 * hl_gte(
        hl.nd.arange(1, k).reshape(1, -1),
        y.reshape(-1, 1)
    ) - 1

def hl_maximum(nd_left, nd_right):
    return hl.nd.array(hl.zip(
        nd_left._data_array(),
        nd_right._data_array()
    ).map(
        lambda _: hl.max(*_)
    ))

def hl_get_gamma_0(y, k, m):
    t = hl_gte(
        hl.nd.arange(1, k)[:, None],
        y[None, :]
    )
    theta = hl_logit(
        hl_mean(t, 1)
    )
    eta = hl_diff(theta)
    return hl.nd.hstack([
        theta[:1],
        eta,
        hl.nd.zeros(m)
    ])

hl_get_eta_beta_0 = hl_get_gamma_0

def hl_get_gamma_lower_bound(k, m):
    return hl.nd.hstack([
        hl.nd.array([-np.inf]),
        hl.nd.array([0] * (k-2), dtype=hl.dtype('float')),
        hl.nd.array([-np.inf] * m)
    ])

hl_get_eta_beta_lower_bound = hl_get_gamma_lower_bound


# ordinal reg in hail

def hl_get_hess_negjac(X, t, s, p):
    
    _1 = ((1 - p) * s)[:, :, None]
    negdLdeta = (_1 * t[None, :, :]).sum(axis=(0, 1))
    negdLdbeta = -(_1 * X[:, None, :]).sum(axis=(0, 1))
    negjac = hl.nd.hstack([negdLdeta, negdLdbeta])
    
    _2 = p * (1 - p)
    d2Ldeta2 = (_2.sum(axis=0) * t[:, :, None] * t[:, None, :]).sum(axis=0)
    _3 = (_2[:, :, None] * t[None, :, :]).sum(axis=1)[:, :, None]
    d2Ldetadbeta = (_3 * X[:, None, :]).sum(axis=0)
    _4 = _2.sum(axis=1)[:, None, None] * X[:, :, None] * X[:, None, :]
    d2Ldbeta2 = _4.sum(axis=0)
    hess = hl.nd.vstack([
        hl.nd.hstack([
            d2Ldeta2, d2Ldetadbeta
        ]),
        hl.nd.hstack([
            d2Ldetadbeta.T, d2Ldbeta2
        ])
    ])
    return hess, negjac


def _hl_get_eta_beta(_, eta_beta, eta_beta_lower_bound, X, t, s, tol, k):
    eta = eta_beta[:k-1]
    beta = eta_beta[k-1:]
    p = hl_sigmoid(s * (
        (t @ eta)[None, :] - (X @ beta)[:, None]
    ))
    delta_eta_beta = hl.nd.solve(
        *hl_get_hess_negjac(X, t, s, p)
    )
    eta_beta = eta_beta + delta_eta_beta
    eta_beta = hl_maximum(eta_beta, eta_beta_lower_bound)
    converged = hl.max(
        hl.abs(delta_eta_beta)._data_array()
    ) < tol
    return hl.if_else(
        converged,
        eta_beta,
        _(eta_beta, eta_beta_lower_bound, X, t, s, tol, k)
    )

def hl_get_eta_beta(X, y, tol, k, m):
    eta_beta_0 = hl_get_eta_beta_0(y, k, m)
    eta_beta_lower_bound = hl_get_eta_beta_lower_bound(k, m)
    t = hl_get_t(k)
    s = hl_get_s(y, k)
    return hl.experimental.loop(
        _hl_get_eta_beta,
        hl.dtype('ndarray<float64, 1>'),
        eta_beta_0, eta_beta_lower_bound, X, t, s, tol, k
    )

hl_eta_beta = hl_get_eta_beta(hl_X, hl_y, 1e-6, k, m)

# Error
print(hl_eta_beta.collect())