itlog 0.1.0__py3-none-any.whl

itlog/__init__.py

@@ -0,0 +1,56 @@
+"""itlog: JAX-backed discrete choice modeling."""
+
+from itlog.benchmarks.runner import (
+    BenchmarkResult,
+    available_engines,
+    benchmark_fit,
+    benchmark_objective_gradient,
+)
+from itlog.data import ChoiceDataset, Field, TensorData, wide_to_long_pylogit
+from itlog.expr import Parameter, Var
+from itlog.latex import expr_to_latex, model_to_latex, result_to_latex
+from itlog.models import (
+    BaseChoiceModel,
+    ClassMembership,
+    CrossNestedLogit,
+    LatentClass,
+    MixedLogit,
+    MultinomialLogit,
+    Nest,
+    NestedLogit,
+    OrderedLogit,
+    OrderedProbit,
+)
+from itlog.results import FitResult
+from itlog.suite import ModelReport, ModelSuite
+
+__all__ = [
+    "ChoiceDataset",
+    "Field",
+    "TensorData",
+    "Parameter",
+    "Var",
+    "MultinomialLogit",
+    "MixedLogit",
+    "NestedLogit",
+    "CrossNestedLogit",
+    "LatentClass",
+    "ClassMembership",
+    "Nest",
+    "OrderedLogit",
+    "OrderedProbit",
+    "BaseChoiceModel",
+    "FitResult",
+    "ModelSuite",
+    "ModelReport",
+    "expr_to_latex",
+    "model_to_latex",
+    "result_to_latex",
+    "wide_to_long_pylogit",
+    "BenchmarkResult",
+    "available_engines",
+    "benchmark_fit",
+    "benchmark_objective_gradient",
+]
+
+__version__ = "0.1.0"

itlog/backends/__init__.py

@@ -0,0 +1,5 @@
+"""NumPy and JAX estimation backends."""
+
+from itlog.backends.jax_backend import build_feature_arrays, make_jax_objective
+
+__all__ = ["make_jax_objective", "build_feature_arrays"]

itlog/backends/jax_backend.py

@@ -0,0 +1,67 @@
+"""JAX JIT objective/gradient factory."""
+
+from __future__ import annotations
+
+from typing import Callable, Dict, Tuple
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from itlog.core.objective import make_cached_val_grad
+from itlog.data import TensorData
+
+
+def make_jax_objective(
+    neg_ll_fn: Callable[[jnp.ndarray], jnp.ndarray],
+) -> Tuple[Callable[[np.ndarray], float], Callable[[np.ndarray], np.ndarray], Dict[str, int]]:
+    val_grad = jax.jit(jax.value_and_grad(neg_ll_fn))
+
+    def _eval(theta_np: np.ndarray):
+        val, grad = val_grad(jnp.asarray(theta_np))
+        return float(val), np.asarray(grad, dtype=np.float64)
+
+    return make_cached_val_grad(_eval)
+
+
+def build_feature_arrays(tensors: TensorData) -> Dict[str, jnp.ndarray]:
+    return {k: jnp.asarray(v) for k, v in tensors.features.items()}
+
+
+def make_mnl_objective(compiled, tensors: TensorData):
+    """JAX MNL objective/gradient builder (used for NumPy<->JAX parity checks)."""
+    from itlog.backends.likelihoods import mnl_neg_loglik
+
+    features = build_feature_arrays(tensors)
+    choice_idx = jnp.asarray(tensors.choice_idx)
+    availability = jnp.asarray(tensors.availability)
+
+    def neg_ll(theta):
+        return mnl_neg_loglik(theta, compiled, features, choice_idx, availability)
+
+    return make_jax_objective(neg_ll)
+
+
+def make_mxl_objective(compiled, tensors: TensorData, *, num_draws: int, seed: int):
+    """JAX MXL objective/gradient builder (used for NumPy<->JAX parity checks)."""
+    from itlog.models.mixed import _mxl_neg_loglik_jax
+
+    if tensors.panel_ids is None:
+        raise ValueError("Mixed logit requires panel_id")
+    features = build_feature_arrays(tensors)
+    choice_idx = jnp.asarray(tensors.choice_idx)
+    availability = jnp.asarray(tensors.availability)
+    unique_panels = np.unique(tensors.panel_ids)
+    n_panels = len(unique_panels)
+    n_random = len(compiled.random_param_names)
+    rng = np.random.default_rng(seed)
+    draws = jnp.asarray(rng.standard_normal((n_panels, num_draws, max(n_random, 1))))
+    panel_index_np = {pid: i for i, pid in enumerate(unique_panels)}
+    panel_index = jnp.asarray([panel_index_np[pid] for pid in tensors.panel_ids])
+
+    def neg_ll(theta):
+        return _mxl_neg_loglik_jax(
+            theta, compiled, features, choice_idx, availability, panel_index, draws, n_random
+        )
+
+    return make_jax_objective(neg_ll)

itlog/backends/likelihoods.py

@@ -0,0 +1,220 @@
+"""JAX log-likelihood kernels (backend-neutral array API)."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+import jax
+import jax.numpy as jnp
+
+from itlog.expr.compile import CompiledModel, CompiledTerm
+
+NEG_INF = -1e20
+
+
+def compute_utilities(
+    theta: Any,
+    compiled: CompiledModel,
+    features: Dict[str, Any],
+    *,
+    beta_override: Optional[Any] = None,
+) -> Any:
+    n_obs = next(iter(features.values())).shape[0]
+    n_alt = len(compiled.alt_labels)
+    V = jnp.zeros((n_obs, n_alt))
+    alt_to_idx = {label: i for i, label in enumerate(compiled.alt_labels)}
+
+    for alt_label, util in compiled.utilities.items():
+        ai = alt_to_idx[alt_label]
+        v_alt = jnp.zeros(n_obs)
+        for term in util.terms:
+            v_alt = v_alt + _eval_term(theta, term, features, ai, beta_override)
+        V = V.at[:, ai].set(v_alt)
+    return V
+
+
+def _eval_term(
+    theta: Any,
+    term: CompiledTerm,
+    features: Dict[str, Any],
+    alt_idx: int,
+    beta_override: Optional[Any],
+) -> Any:
+    n_obs = features[next(iter(features))].shape[0]
+
+    if term.feature_name is None:
+        if term.param_idx is not None:
+            if beta_override is not None:
+                return beta_override[:, term.param_idx]
+            return jnp.full(n_obs, theta[term.param_idx])
+        return jnp.full(n_obs, term.const_value)
+
+    x = features[term.feature_name][:, alt_idx]
+    if term.param_idx is not None:
+        if beta_override is not None:
+            beta = beta_override[:, term.param_idx]
+        else:
+            beta = jnp.full(n_obs, theta[term.param_idx])
+        return beta * x
+    return term.const_value * x
+
+
+def mnl_neg_loglik(
+    theta: Any,
+    compiled: CompiledModel,
+    features: Dict[str, Any],
+    choice_idx: Any,
+    availability: Any,
+) -> Any:
+    V = compute_utilities(theta, compiled, features)
+    V = jnp.where(availability, V, NEG_INF)
+    chosen = V[jnp.arange(V.shape[0]), choice_idx]
+    log_denom = jax.scipy.special.logsumexp(V, axis=1)
+    return -jnp.sum(chosen - log_denom)
+
+
+def nested_neg_loglik(
+    theta: Any,
+    compiled: CompiledModel,
+    features: Dict[str, Any],
+    choice_idx: Any,
+    availability: Any,
+    nest_for_alt: Any,
+    nest_scales: Any,
+    scale_idx: Any,
+) -> Any:
+    """Nested logit with nest scales in (0, 1] via sigmoid reparam in theta."""
+    V = compute_utilities(theta, compiled, features)
+    V = jnp.where(availability, V, NEG_INF)
+    n_obs, n_alt = V.shape
+    n_nests = nest_scales.shape[0]
+
+    mu = jnp.where(
+        scale_idx >= 0,
+        jax.nn.sigmoid(nest_scales) * 0.999 + 0.001,
+        1.0,
+    )
+    log_iv = jnp.full((n_obs, n_nests), -jnp.inf)
+    for m in range(n_nests):
+        in_nest = nest_for_alt == m
+        Vm = jnp.where(in_nest[None, :], V / mu[m], NEG_INF)
+        log_iv = log_iv.at[:, m].set(jax.scipy.special.logsumexp(Vm, axis=1))
+    chosen_nest = nest_for_alt[choice_idx]
+    mu_c = mu[chosen_nest]
+    Vc = V[jnp.arange(n_obs), choice_idx]
+    log_iv_c = log_iv[jnp.arange(n_obs), chosen_nest]
+    log_p = (Vc / mu_c) + (mu_c - 1.0) * log_iv_c
+    log_denom = jax.scipy.special.logsumexp(mu[None, :] * log_iv, axis=1)
+    log_p = log_p - log_denom
+    return -jnp.sum(log_p)
+
+
+def cnl_neg_loglik(
+    theta: Any,
+    compiled: CompiledModel,
+    features: Dict[str, Any],
+    choice_idx: Any,
+    availability: Any,
+    alpha: Any,
+    nest_scales: Any,
+    scale_idx: Any,
+) -> Any:
+    """Cross-nested logit: alpha[n_alt, n_nests], nest_scales raw -> sigmoid."""
+    V = compute_utilities(theta, compiled, features)
+    V = jnp.where(availability, V, NEG_INF)
+    n_obs, n_alt = V.shape
+    n_nests = nest_scales.shape[0]
+
+    alpha_norm = alpha / jnp.maximum(jnp.sum(alpha, axis=1, keepdims=True), 1e-300)
+    mu = jnp.where(
+        scale_idx >= 0,
+        jax.nn.sigmoid(nest_scales) * 0.999 + 0.001,
+        1.0,
+    )
+
+    log_iv = jnp.full((n_obs, n_nests), -jnp.inf)
+    for m in range(n_nests):
+        Vm = V / mu[m] + jnp.log(jnp.maximum(alpha_norm[:, m], 1e-300))[None, :]
+        log_iv = log_iv.at[:, m].set(jax.scipy.special.logsumexp(Vm, axis=1))
+
+    c = choice_idx
+    Vc = V[jnp.arange(n_obs), c]
+    log_terms = jnp.full((n_obs, n_nests), -jnp.inf)
+    for m in range(n_nests):
+        log_terms = log_terms.at[:, m].set(
+            jnp.log(jnp.maximum(alpha_norm[c, m], 1e-300))
+            + Vc / mu[m]
+            + (mu[m] - 1.0) * log_iv[:, m]
+        )
+    log_p = jax.scipy.special.logsumexp(log_terms, axis=1) - jax.scipy.special.logsumexp(
+        mu[None, :] * log_iv, axis=1
+    )
+    return -jnp.sum(log_p)
+
+
+def latent_class_neg_loglik(
+    theta: Any,
+    compiled_list: list[CompiledModel],
+    features: Dict[str, Any],
+    choice_idx: Any,
+    availability: Any,
+    membership_X: Any,
+    membership_param_idx: Any,
+) -> Any:
+    """Latent class MNL with softmax membership."""
+    n_classes = len(compiled_list)
+    n_obs = choice_idx.shape[0]
+
+    class_ll = []
+    for compiled in compiled_list:
+        V = compute_utilities(theta, compiled, features)
+        V = jnp.where(availability, V, NEG_INF)
+        chosen = V[jnp.arange(n_obs), choice_idx]
+        log_denom = jax.scipy.special.logsumexp(V, axis=1)
+        class_ll.append(chosen - log_denom)
+    class_ll = jnp.stack(class_ll, axis=1)
+
+    # membership logits: reference class 0
+    if membership_param_idx.shape[0] == 0:
+        log_pi = jnp.zeros((n_obs, n_classes))
+        log_pi = log_pi.at[:, 0].set(0.0)
+    else:
+        logits = jnp.zeros((n_obs, n_classes - 1))
+        for k, idxs in enumerate(membership_param_idx):
+            for j, pidx in enumerate(idxs):
+                logits = logits.at[:, k].add(theta[pidx] * membership_X[:, j])
+        log_pi = jnp.concatenate([jnp.zeros((n_obs, 1)), logits], axis=1)
+        log_pi = log_pi - jax.scipy.special.logsumexp(log_pi, axis=1, keepdims=True)
+
+    log_mix = jax.scipy.special.logsumexp(log_pi + class_ll, axis=1)
+    return -jnp.sum(log_mix)
+
+
+def ordered_neg_loglik(
+    theta: Any,
+    X: Any,
+    y_idx: Any,
+    n_thresholds: int,
+    link: str = "logit",
+) -> Any:
+    """Ordered logit/probit with monotone thresholds via cumulative softplus."""
+    n_beta = X.shape[1]
+    beta = theta[:n_beta]
+    raw_tau = theta[n_beta:]
+    tau = jnp.cumsum(jax.nn.softplus(raw_tau))
+    eta = X @ beta
+    k = y_idx.astype(jnp.int32)
+
+    if link == "logit":
+        cdf = jax.scipy.special.expit
+    else:
+        cdf = jax.scipy.stats.norm.cdf
+
+    big = 1e20
+    bounds_low = jnp.concatenate([jnp.array([-big]), tau])
+    bounds_high = jnp.concatenate([tau, jnp.array([big])])
+    cdf_high = cdf(bounds_high[None, :] - eta[:, None])
+    cdf_low = cdf(bounds_low[None, :] - eta[:, None])
+    cat_probs = jnp.maximum(cdf_high - cdf_low, 1e-300)
+    log_p = jnp.log(jnp.take_along_axis(cat_probs, y_idx[:, None], axis=1)).squeeze(1)
+    return -jnp.sum(log_p)