complete sar mix/nested objectives

locpick/_jax/kernels.py

@@ -562,6 +562,81 @@ def _ll_single_draw(r):
     return jnp.sum(log_L_sim * weights)
 
 
+def mixed_logit_ll_contribs(
+    V_fixed: jnp.ndarray,
+    dm_random: jnp.ndarray,
+    beta_random_means: jnp.ndarray,
+    beta_random_spreads: jnp.ndarray,
+    dist_codes: jnp.ndarray,
+    draws: jnp.ndarray,
+    chosen: jnp.ndarray,
+    weights: jnp.ndarray,
+    available: jnp.ndarray,
+    n_obs: int,
+    n_alts: int,
+    k_random: int,
+    n_draws: int,
+) -> jnp.ndarray:
+    """Compute mixed logit per-observation simulated log-likelihood contributions.
+
+    Same as :func:`mixed_logit_ll` but returns per-observation ``(n_obs,)``
+    array instead of a scalar sum.  Used for BHHH/robust standard errors
+    via ``jax.jacrev``.
+
+    Returns
+    -------
+    jnp.ndarray, shape (n_obs,)
+        Per-observation simulated log-likelihood contributions.
+    """
+    means = beta_random_means[None, :]
+    spreads = beta_random_spreads[None, :]
+
+    def _ll_single_draw(r):
+        z_r = draws[:, r, :]
+        beta_normal = means + spreads * z_r
+        beta_lognormal = jnp.exp(jnp.clip(means + spreads * z_r, -50.0, 50.0))
+        t = 1.0 / (1.0 + 0.2316419 * jnp.abs(z_r))
+        d = 0.3989422804014327
+        poly = t * (
+            0.319381530
+            + t * (-0.356563782 + t * (1.781477937 + t * (-1.821255978 + t * 1.330274429)))
+        )
+        phi_z = jnp.where(
+            z_r >= 0,
+            1.0 - d * jnp.exp(-0.5 * z_r * z_r) * poly,
+            d * jnp.exp(-0.5 * z_r * z_r) * poly,
+        )
+        beta_uniform = means + spreads * (2.0 * phi_z - 1.0)
+        mask = phi_z <= 0.5
+        beta_triangular = jnp.where(
+            mask,
+            means + spreads * (jnp.sqrt(2.0 * phi_z) - 1.0),
+            means + spreads * (1.0 - jnp.sqrt(2.0 * (1.0 - phi_z))),
+        )
+        dist = dist_codes[None, :]
+        beta_random_r = jnp.where(
+            dist == 0,
+            beta_normal,
+            jnp.where(
+                dist == 1,
+                beta_lognormal,
+                jnp.where(dist == 2, beta_triangular, beta_uniform),
+            ),
+        )
+        v_random = jnp.sum(
+            dm_random.reshape(n_obs, n_alts, k_random) * beta_random_r[:, None, :],
+            axis=2,
+        )
+        V = V_fixed + v_random
+        log_probs = mnl_log_probs(V, available)
+        log_L_n = (log_probs * chosen).sum(axis=1)
+        return log_L_n
+
+    log_L_all = jax.vmap(_ll_single_draw, in_axes=0)(jnp.arange(n_draws))
+    log_L_sim = jax_logsumexp(log_L_all, axis=0) - jnp.log(float(n_draws))
+    return log_L_sim * weights
+
+
 # ---------------------------------------------------------------------------
 # Mixed nested logit kernel
 # ---------------------------------------------------------------------------
@@ -694,6 +769,82 @@ def _ll_single_draw(r):
     return jnp.sum(log_L_sim * weights)
 
 
+def mixed_nested_logit_ll_contribs(
+    V_fixed: jnp.ndarray,
+    dm_random: jnp.ndarray,
+    beta_random_means: jnp.ndarray,
+    beta_random_spreads: jnp.ndarray,
+    dist_codes: jnp.ndarray,
+    draws: jnp.ndarray,
+    lambdas: jnp.ndarray,
+    nest_matrix: jnp.ndarray,
+    chosen: jnp.ndarray,
+    weights: jnp.ndarray,
+    available: jnp.ndarray,
+    n_obs: int,
+    n_alts: int,
+    k_random: int,
+    n_draws: int,
+    n_nests: int,
+) -> jnp.ndarray:
+    """Compute mixed nested logit per-observation LL contributions.
+
+    Same as :func:`mixed_nested_logit_ll` but returns per-observation
+    ``(n_obs,)`` array.  Used for BHHH/robust standard errors.
+    """
+    means = beta_random_means[None, :]
+    spreads = beta_random_spreads[None, :]
+    long_lambda = jnp.ones(n_alts, dtype=jnp.float64)
+    for m in range(n_nests):
+        mask_m = nest_matrix[:, m] > 0
+        long_lambda = jnp.where(mask_m, lambdas[m], long_lambda)
+
+    def _ll_single_draw(r):
+        z_r = draws[:, r, :]
+        beta_normal = means + spreads * z_r
+        beta_lognormal = jnp.exp(jnp.clip(means + spreads * z_r, -50.0, 50.0))
+        t = 1.0 / (1.0 + 0.2316419 * jnp.abs(z_r))
+        d = 0.3989422804014327
+        poly = t * (
+            0.319381530
+            + t * (-0.356563782 + t * (1.781477937 + t * (-1.821255978 + t * 1.330274429)))
+        )
+        phi_z = jnp.where(
+            z_r >= 0,
+            1.0 - d * jnp.exp(-0.5 * z_r * z_r) * poly,
+            d * jnp.exp(-0.5 * z_r * z_r) * poly,
+        )
+        beta_uniform = means + spreads * (2.0 * phi_z - 1.0)
+        mask = phi_z <= 0.5
+        beta_triangular = jnp.where(
+            mask,
+            means + spreads * (jnp.sqrt(2.0 * phi_z) - 1.0),
+            means + spreads * (1.0 - jnp.sqrt(2.0 * (1.0 - phi_z))),
+        )
+        dist = dist_codes[None, :]
+        beta_random_r = jnp.where(
+            dist == 0,
+            beta_normal,
+            jnp.where(
+                dist == 1,
+                beta_lognormal,
+                jnp.where(dist == 2, beta_triangular, beta_uniform),
+            ),
+        )
+        v_random = jnp.sum(
+            dm_random.reshape(n_obs, n_alts, k_random) * beta_random_r[:, None, :],
+            axis=2,
+        )
+        V = V_fixed + v_random
+        log_probs = nested_log_probs(V, lambdas, nest_matrix, available)
+        log_L_n = (log_probs * chosen).sum(axis=1)
+        return log_L_n
+
+    log_L_all = jax.vmap(_ll_single_draw, in_axes=0)(jnp.arange(n_draws))
+    log_L_sim = jax_logsumexp(log_L_all, axis=0) - jnp.log(float(n_draws))
+    return log_L_sim * weights
+
+
 # ---------------------------------------------------------------------------
 # Log-likelihood computation (shared across models)
 # ---------------------------------------------------------------------------