SearchLibrium 0.0.151__tar.gz → 0.0.155__tar.gz

{searchlibrium-0.0.151/src/SearchLibrium.egg-info → searchlibrium-0.0.155}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: SearchLibrium
-Version: 0.0.151
+Version: 0.0.155
 Summary: A Python package for econometric models driven by search
 Author: Mario Fernando Taco Morales, Alexander Paz, Prithvi Beeramole, Robert Burdett
 Author-email: Zeke Ahern <z.ahern@qut.edu.au>

{searchlibrium-0.0.151 → searchlibrium-0.0.155}/pyproject.toml

@@ -82,7 +82,7 @@ Documentation = "https://pypi.org/project/SearchLibrium/"
 realpython = "SearchLibrium.__main__:main"
 
 [tool.bumpver]
-current_version = "0.0.151"
+current_version = "0.0.155"
 version_pattern = "MAJOR.MINOR.PATCH"
 commit_message = "[skip ci] Bump version {old_version} -> {new_version}"
 commit = true

{searchlibrium-0.0.151 → searchlibrium-0.0.155}/src/SearchLibrium/MixedLogit.py

@@ -134,14 +134,13 @@ class MixedLogit(DiscreteChoiceModel):
     def setup(self, X, y, varnames=None, alts=None, isvars=None, transvars=None,
               transformation="boxcox", ids=None, weights=None, avail=None,
               randvars=None, panels=None, base_alt=None, fit_intercept=False,
-              init_coeff=None, maxiter=2000, correlated_vars=None,
+              init_coeff=None, maxiter=1000, correlated_vars=None,
               n_draws=1000, halton=True, minimise_func=None,
-              batch_size=None, halton_opts=None, ftol=1e-6,
+              batch_size=None, halton_opts=None, ftol=1e-12,
               gtol=1e-6, return_hess=True, return_grad=True, method="slsqp",
               save_fitted_params=True, mnl_init=True,
               de_init=False, de_popsize=4, de_maxiter=3, de_tol=0.5,
-              de_polish=False, l1_penalty=0.1, l2_penalty=0.5,
-              sd_penalty=0.001):
+              de_polish=False, l1_penalty=0.0):
         # {
         self.fit_intercept = fit_intercept
         # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -187,8 +186,6 @@ class MixedLogit(DiscreteChoiceModel):
         self.total_fun_eval = 0
         self.method = method.lower() if hasattr(method, 'lower') else method
         self.l1_penalty = float(l1_penalty)
-        self.l2_penalty = float(l2_penalty)
-        self.sd_penalty = float(sd_penalty)
         self.jac = self.return_grad  # scipy optimize parameter
         self.n_draws = n_draws
         self.batch_size = min(n_draws, batch_size) if batch_size is not None else n_draws
@@ -387,7 +384,9 @@ class MixedLogit(DiscreteChoiceModel):
         # 2x Kftrans - mean and lambda, 3x Krtrans - mean, s.d., lambda
         # Kchol, Kbw - relate to random variables, non-transformed
         # Kchol - cholesky matrix, Kbw the s.d. for random vars
-        n_coeff = self.Kf + self.Kr + self.Kchol + self.Kbw + 2 * self.Kftrans + 3 * self.Krtrans
+        n_coeff = (self.Kf + self.Kr + self.Kchol + self.Kbw
+                   + 2 * self.Kftrans + 3 * self.Krtrans
+                   + self._n_coeff_extra())
 
         # Initalise coefficients using a multinomial logit model
         if self.mnl_init and self.init_coeff is None:
@@ -411,30 +410,38 @@ class MixedLogit(DiscreteChoiceModel):
                       fit_intercept=False)
             mnl.fit()
 
-            # mnl estimates -> mxl needs to add stdev to random variables
-            self.init_coeff = mnl.coeff_est
-
-            lower = self.Kf + 2 * self.Kftrans + self.Kr
-            upper = lower + self.Krtrans
-
-            # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-            # ERROR HANDLING (FIX) - WHEN LESS COEFFICIENTS THAN "kf + kr"
-            if lower > len(self.init_coeff):
-                # {
-                additional_elements_needed = lower - len(self.init_coeff)
-                extra = np.full(additional_elements_needed, 0.1)
-                self.init_coeff = np.concatenate((self.init_coeff, extra))
-            # }
-            # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-            arr = self.init_coeff[:lower]
-            # Use simple initialization matching searchlogit for better convergence
+            # Build init_coeff from MNL estimates, mapping each variable
+            # to its correct MXL slot (Bf for fixed, Br_b for random).
+            bf_init = np.zeros(self.Kf)
+            br_b_init = np.zeros(self.Kr)
+            bf_pos = br_pos = 0
+            for ii in range(len(self.varnames)):
+                if not mnl.fxidx[ii]:
+                    continue
+                mnl_idx = int(np.sum(mnl.fxidx[:ii]))
+                mnl_val = mnl.coeff_est[mnl_idx] if mnl_idx < len(mnl.coeff_est) else 0.1
+                if self.fxidx[ii]:
+                    bf_init[bf_pos] = mnl_val
+                    bf_pos += 1
+                elif self.rvidx[ii]:
+                    br_b_init[br_pos] = mnl_val
+                    br_pos += 1
+
+            mnl_Kf = int(np.sum(mnl.fxidx))
+            mnl_Kftrans = int(np.sum(mnl.fxtransidx))
+            bftrans_b = mnl.coeff_est[mnl_Kf:mnl_Kf + mnl_Kftrans] if mnl_Kftrans > 0 else np.array([])
+            bftrans_l = mnl.coeff_est[mnl_Kf + mnl_Kftrans:mnl_Kf + 2 * mnl_Kftrans] if mnl_Kftrans > 0 else np.array([])
+
+            arr = np.concatenate([bf_init, br_b_init] + list(self._init_pad_arrays()) + [bftrans_b, bftrans_l])
             rep = np.repeat(0.1, self.Kchol + self.Kbw)
-            self.init_coeff = np.concatenate((arr, rep, self.init_coeff[lower:upper],))
+            self.init_coeff = np.concatenate([arr, rep])
 
             if self.Krtrans:
-                rep = np.repeat(0.1, self.Krtrans)
-                self.init_coeff = np.concatenate((self.init_coeff, rep, self.init_coeff[-self.Krtrans:]))
+                self.init_coeff = np.concatenate([
+                    self.init_coeff,
+                    np.repeat(0.1, self.Krtrans),
+                    np.repeat(0.1, self.Krtrans),
+                ])
         # }
 
         betas = np.repeat(0.1, n_coeff) if self.init_coeff is None else self.init_coeff
@@ -454,7 +461,8 @@ class MixedLogit(DiscreteChoiceModel):
             "flmbda": (lmda_bound, self.Kftrans),
             "br_trans_b": (any_bound, self.Krtrans),
             "br_trans_w": (any_bound, self.Krtrans),
-            "rlmbda": (lmda_bound, self.Krtrans)
+            "rlmbda": (lmda_bound, self.Krtrans),
+            **self._bound_extra(),
         }
         bnds = [[bound[1][0]] * bound[1][1] for bound in bound_dict.items() if bound[1][1] > 0]
         bnds = list(itertools.chain.from_iterable(bnds))
@@ -553,7 +561,8 @@ class MixedLogit(DiscreteChoiceModel):
             "flmbda": (lmda_bound, self.Kftrans),
             "br_trans_b": (any_bound, self.Krtrans),
             "br_trans_w": (any_bound, self.Krtrans),
-            "rlmbda": (lmda_bound, self.Krtrans)
+            "rlmbda": (lmda_bound, self.Krtrans),
+            **self._bound_extra(),
         }
 
         # This code makes a specific number of copies of each range
@@ -570,10 +579,13 @@ class MixedLogit(DiscreteChoiceModel):
         if getattr(self, '_jax', False):
             jax_result = self.optimize_jax(betas, draws, drawstrans)
             if jax_result is not None:
-                beta_segment_names = ["Bf", "Br_b", "chol", "Br_w", "Bftrans",
-                                    "flmbda", "Brtrans_b", "Brtrans_w", "rlmda"]
-                iterations = [self.Kf, self.Kr, self.Kchol, self.Kbw, self.Kftrans,
-                            self.Kftrans, self.Krtrans, self.Krtrans, self.Krtrans]
+                extra_names, extra_counts = self._beta_segment_extra()
+                beta_segment_names = (["Bf", "Br_b"] + list(extra_names)
+                                     + ["chol", "Br_w", "Bftrans",
+                                        "flmbda", "Brtrans_b", "Brtrans_w", "rlmda"])
+                iterations = ([self.Kf, self.Kr] + list(extra_counts)
+                              + [self.Kchol, self.Kbw, self.Kftrans,
+                                 self.Kftrans, self.Krtrans, self.Krtrans, self.Krtrans])
                 self.var_list = self.split_betas(jax_result['x'], iterations, beta_segment_names)
                 self.chol_mat = self.construct_chol_mat(
                     self.var_list['chol'], self.var_list['Br_w'], self.var_list['Brtrans_w'])
@@ -603,64 +615,11 @@ class MixedLogit(DiscreteChoiceModel):
         args = (self.X, self.y, self.panel_info, draws, drawstrans, self.weights, self.avail, self.batch_size)
         bounds = bnds if self.method == "L-BFGS-B" else None
         options = {'gtol': self.gtol, 'maxiter': self.maxiter, 'disp': False}
-
-        # ── SD penalty wrapper: regularise SD params away from exactly zero ──
-        # When an SD parameter hits zero, the gradient vanishes and BFGS gets stuck.
-        # A small L2 penalty that activates below a threshold keeps SDs bounded
-        # away from zero without affecting the well-identified estimates.
-        if self.sd_penalty > 0 and self.Kbw > 0:
-            # SD param indices: end of beta vector, before Kftrans and Krtrans
-            sd_start = self.Kf + self.Kr + self.Kchol
-            sd_end   = sd_start + self.Kbw
-            _orig_obj = self.get_loglik_gradient
-
-            def _penalised_obj(betas, *obj_args):
-                f, g = _orig_obj(betas, *obj_args)
-                for i in range(sd_start, sd_end):
-                    if i >= len(betas):
-                        break
-                    sd_val = betas[i]
-                    if sd_val < 0.05:
-                        p = self.sd_penalty * (sd_val - 0.05) ** 2
-                        f += p
-                        if g is not None and len(g) > i:
-                            g[i] += 2.0 * self.sd_penalty * (sd_val - 0.05)
-                return f, g
-
-            obj_fn = _penalised_obj
-        else:
-            obj_fn = self.get_loglik_gradient
-
-        result = minimise_func(obj_fn, betas, jac=self.jac, method=self.method,
+        result = minimise_func(self.get_loglik_gradient, betas, jac=self.jac, method=self.method,
                                args=args, tol=self.ftol, bounds=bounds, options=options)
-        print(f"[MXL] Primary minimization ({self.method}): success={result.get('success', None)}, fun={result.get('fun', float('nan')):.6g}, nit={result.get('nit', '?')}")
+        print(f"[MXL] Minimization completed: success={result.get('success', None)}, fun={result.get('fun', float('nan')):.6g}, nit={result.get('nit', '?')}")
         if 'x' in result:
-            print(f"[MXL] Primary betas first_values={np.asarray(result['x'])[:min(8, len(result['x']))]!r}")
-
-        # ── BFGS polish: after non-BFGS methods (SLSQP, Powell, etc.),
-        # run a quick BFGS to get the inverse Hessian for standard errors.
-        if (self.method not in ('bfgs', 'l-bfgs-b')
-                and 'x' in result
-                and result.get('success', False)):
-            try:
-                polish_opts = {'gtol': self.gtol * 10, 'maxiter': min(self.maxiter // 5, 200), 'disp': False}
-                polish_bnds = bnds  # re-use same bounds
-                # Use original objective (no SD penalty) for clean Hessian
-                polish_result = minimise_func(
-                    self.get_loglik_gradient,
-                    np.asarray(result['x']).copy(),
-                    jac=True, method='L-BFGS-B',
-                    args=args, tol=self.ftol,
-                    bounds=polish_bnds, options=polish_opts,
-                )
-                if polish_result.get('fun', float('inf')) < result.get('fun', float('inf')) + 0.1:
-                    print(f"[MXL] BFGS polish: success={polish_result.get('success')}, "
-                          f"fun={polish_result.get('fun', float('nan')):.6g}")
-                    result = polish_result
-                else:
-                    print(f"[MXL] BFGS polish skipped (worse objective)")
-            except Exception as e:
-                print(f"[MXL] BFGS polish failed: {e}")
+            print(f"[MXL] Minimization final betas first_values={np.asarray(result['x'])[:min(8, len(result['x']))]!r}")
         # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
         if hasattr(self, 'method') and self.method == "L-BFGS-B":  # {
@@ -720,7 +679,7 @@ class MixedLogit(DiscreteChoiceModel):
     @staticmethod
     def _jax_mxl_negloglik(betas, X_jax, y_jax, panel_info_jax, draws_jax,
                             fxidx, rvidx, Kf, Kr, Kchol, Kbw, rvdist_names,
-                            correlationLength):
+                            correlationLength, **kwargs):
         """JAX simulation-based log-likelihood for Mixed Logit (standard case).
 
         Handles fixed and random (normally/lognormally distributed) parameters.
@@ -849,12 +808,14 @@ class MixedLogit(DiscreteChoiceModel):
             _N = X_jax.shape[0]
             _P = X_jax.shape[1] if X_jax.ndim > 1 else 1
             _J = X_jax.shape[2] if X_jax.ndim > 2 else 1
-            _cache_key = (_N, _P, _J, Kf, Kr, Kchol, Kbw)
+            _cache_key = (_N, _P, _J, Kf, Kr, Kchol, Kbw) + self._jax_cache_key_extra()
             _compiled = self._mxl_jit_cache.get(_cache_key)
             if _compiled is None:
+                _extra = self._jax_negloglik_extra_kwargs()
                 _fn = lambda b, _X, _y, _pi, _dr: self._jax_mxl_negloglik(
                     b, _X, _y, _pi, _dr,
-                    fxidx, rvidx, Kf, Kr, Kchol, Kbw, rvdist_names, correlationLength)
+                    fxidx, rvidx, Kf, Kr, Kchol, Kbw, rvdist_names, correlationLength,
+                    **_extra)
                 _compiled = jax.jit(jax.value_and_grad(_fn))
                 self._mxl_jit_cache[_cache_key] = _compiled
             # ────────────────────────────────────────────────────────
@@ -1668,4 +1629,33 @@ class MixedLogit(DiscreteChoiceModel):
             pch2_sd_test[i, :] = np.std(pch2[batch_slice, :], axis=0)
         # }
         self.pch2_res, self.pch2_sd_test = pch2_res, pch2_sd_test
-    # }
+    # }
+
+    # ── Extension hooks for subclasses (e.g. MixedLogitGSE) ────────
+    def _n_coeff_extra(self) -> int:
+        """Extra parameter count beyond base MXL params."""
+        return 0
+
+    def _bound_extra(self) -> dict:
+        """Extra bound dictionary entries for subclasses."""
+        return {}
+
+    def _beta_segment_extra(self):
+        """Returns (extra_names, extra_counts) for split_betas."""
+        return [], []
+
+    def _augment_br(self, Br, var_list):
+        """Modify random coefficients. Subclasses may add gradient terms."""
+        return Br
+
+    def _jax_negloglik_extra_kwargs(self):
+        """Extra keyword arguments forwarded to JAX negloglik fn."""
+        return {}
+
+    def _jax_cache_key_extra(self):
+        """Extra tuple elements for the JAX compilation cache key."""
+        return ()
+
+    def _init_pad_arrays(self):
+        """Extra arrays to insert between Br_b and Bftrans in init_coeff."""
+        return []

searchlibrium-0.0.155/src/SearchLibrium/MixedLogitGSE.py

@@ -0,0 +1,249 @@
+"""MixedLogitGSE — Gradient Score Enhanced Mixed Logit.
+
+Extends MixedLogit with gradient-latent loadings.
+
+  beta_nk = mu_k + (gamma_k + tau_k * z_nk) * g_nk + sigma_k * eta_nk
+
+Parameter layout (no transforms, no correlation):
+  [Br_b(Kr) | gamma(Kr) | gamma_w(Kr) | Br_w(Kr)]   (12 or 16 params)
+
+Usage:
+  mxl = MixedLogitGSE()
+  mxl.setup(..., gradient_scores=g_avg, random_gamma=False)
+  mxl.fit()
+
+Larch equivalence (kept in sync with test_runs_tours stage-5 mode model):
+  The model β_nk = μ_k + γ_k·g_nk + σ_k·η_nk decomposes term-by-term into a
+  plain Larch mixed logit, since β_nk multiplies x_k in the utility:
+      μ_k·x_k          → ordinary fixed utility term         P('c_k')  * X('x_k')
+      γ_k·(g_nk·x_k)   → systematic gradient interaction     P('g_k')  * X('x_k * g')
+      σ_k·(η_nk·x_k)   → Normal random coefficient           mixtures.append(Normal('c_k','s_k'))
+  random_gamma=True adds τ_k·z_nk·g_nk, i.e. promoting the γ_k loadings to
+  random coefficients too (Normal mixtures on the g_k params in the Larch build).
+  See stage5_destination_choice._estimate_mixed_logit for the mirrored build.
+"""
+
+import numpy as np
+import jax.numpy as jnp
+import jax
+
+from SearchLibrium.MixedLogit import MixedLogit
+
+try:
+    from ._device import device as dev
+except ImportError:
+    from _device import device as dev
+
+
+class MixedLogitGSE(MixedLogit):
+    """Mixed Logit with Gradient Score Enhanced latent loadings."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.Kgrad = 0
+        self.Kgrad_w = 0
+        self.gamma_draws = None
+
+    def setup(self, *args, gradient_scores=None, random_gamma=False, **kwargs):
+        self.random_gamma = bool(random_gamma)
+        if gradient_scores is not None:
+            # Standardise each gradient column to mean 0 / std 1 (O(1)) so the
+            # γ loadings are on a common scale and can't numerically overflow.
+            # Centering matters: an uncentred score folds a spurious γ·mean/std
+            # shift into μ; centring keeps μ interpretable as the population mean.
+            # A constant (zero-variance) column would divide by ~0 and blow up to
+            # ~1e8 — guard it by leaving such columns at 0 (they carry no gradient
+            # information anyway).
+            g = np.asarray(gradient_scores, dtype=float)
+            gmu = g.mean(axis=0, keepdims=True)
+            gsd = g.std(axis=0, keepdims=True)
+            gsd_safe = np.where(gsd < 1e-8, 1.0, gsd)
+            self.gradient_scores = (g - gmu) / gsd_safe
+            n_const = int(np.sum(gsd < 1e-8))
+            print(f"[MixedLogitGSE] Gradient loadings enabled (Kgrad={len(kwargs.get('randvars', {}))}, "
+                  f"random_gamma={self.random_gamma}, "
+                  f"raw std range=[{float(gsd.min()):.2f}, {float(gsd.max()):.2f}] -> "
+                  f"centred+standardised"
+                  + (f", {n_const} constant column(s) zeroed" if n_const else "") + ")")
+        else:
+            self.gradient_scores = None
+        super().setup(*args, **kwargs)
+
+    # ── Hook overrides ─────────────────────────────────────────────
+
+    def _n_coeff_extra(self) -> int:
+        return self.Kgrad + self.Kgrad_w
+
+    def _bound_extra(self) -> dict:
+        inf = float("inf")
+        return {"grad": ((-inf, inf), self.Kgrad),
+                "grad_w": ((0, inf), self.Kgrad_w)}
+
+    def _beta_segment_extra(self):
+        return (["grad", "grad_w"], [self.Kgrad, self.Kgrad_w])
+
+    def _jax_cache_key_extra(self):
+        return (self.Kgrad, self.Kgrad_w)
+
+    def _init_pad_arrays(self):
+        pads = []
+        if self.Kgrad > 0:
+            pads.append(np.repeat(0.1, self.Kgrad))
+        if self.Kgrad_w > 0:
+            pads.append(np.repeat(0.1, self.Kgrad_w))
+        return pads
+
+    def _jax_negloglik_extra_kwargs(self):
+        gd = (jnp.array(self.gamma_draws) if self.Kgrad_w > 0
+              and self.gamma_draws is not None else None)
+        gs = (jnp.array(self.gradient_scores) if self.gradient_scores is not None
+              else None)
+        return dict(Kgrad=self.Kgrad, Kgrad_w=self.Kgrad_w,
+                    gradient_scores=gs, gamma_draws=gd)
+
+    # ── Fit override ───────────────────────────────────────────────
+
+    def fit(self):
+        if self.gradient_scores is not None:
+            self.Kgrad = self.Kr
+            self.Kgrad_w = self.Kr if self.random_gamma else 0
+            if self.Kgrad_w > 0:
+                # Seed the random-γ (τ·z) draws so a fit is reproducible across
+                # runs — unseeded np.random.randn made every fit non-deterministic
+                # (and non-comparable in a search). Reuse the model's own seed
+                # when set, mixing in Kgrad so it differs from the η draws' seed.
+                _seed = getattr(self, "seed", None)
+                rng = np.random.default_rng(None if _seed is None
+                                            else int(_seed) + 10007 * self.Kgrad)
+                self.gamma_draws = rng.standard_normal((self.N, self.Kgrad))
+        # Force JAX-only (scipy fallback doesn't handle gradient loadings)
+        self._jax = True
+        super().fit()
+
+    # ── JAX loglik with GSE terms ──────────────────────────────────
+
+    @staticmethod
+    def _jax_mxl_negloglik(betas, X_jax, y_jax, panel_info_jax, draws_jax,
+                            fxidx, rvidx, Kf, Kr, Kchol, Kbw, rvdist_names,
+                            correlationLength,
+                            Kgrad=0, Kgrad_w=0,
+                            gradient_scores=None, gamma_draws=None,
+                            **kwargs):
+        """GSE loglik: β_nk = μ_k + γ_k·g_nk + σ_k·η_nk."""
+        # Split
+        Bf       = betas[:Kf]
+        Br_b     = betas[Kf:Kf + Kr]
+        gamma_mu = betas[Kf + Kr:Kf + Kr + Kgrad] if Kgrad > 0 else jnp.array([])
+        gamma_w  = (betas[Kf + Kr + Kgrad:Kf + Kr + Kgrad + Kgrad_w]
+                    if Kgrad_w > 0 else jnp.array([]))
+        offset   = Kf + Kr + Kgrad + Kgrad_w
+        chol_v   = betas[offset:offset + Kchol]
+        Br_w     = betas[offset + Kchol:offset + Kchol + Kbw]
+
+        # Cholesky
+        chol_mat = jnp.zeros((Kr, Kr))
+        idx = 0
+        for r in range(correlationLength):
+            for c in range(r + 1):
+                chol_mat = chol_mat.at[r, c].set(chol_v[idx]); idx += 1
+        for k in range(Kbw):
+            chol_mat = chol_mat.at[correlationLength + k,
+                                   correlationLength + k].set(jnp.abs(Br_w[k]))
+
+        # Random coefficients
+        N = X_jax.shape[0]
+        Br = Br_b[:, None] + jnp.einsum("kl,nlr->nkr", chol_mat, draws_jax[:, :Kr, :])
+
+        # GSE gradient-latent term (standardised scores, O(1))
+        if Kgrad > 0 and gradient_scores is not None:
+            if Kgrad_w > 0 and gamma_draws is not None:
+                gamma_rnd = gamma_mu[None, :] + gamma_w[None, :] * gamma_draws
+            else:
+                gamma_rnd = gamma_mu[None, :]
+            Br = Br + gamma_rnd[:, :, None] * gradient_scores[:, :, None]
+
+        # Distribution transforms
+        for k, dist in enumerate(rvdist_names):
+            if dist == "ln":
+                Br = Br.at[:, k, :].set(jnp.exp(Br[:, k, :]))
+            elif dist == "tn":
+                Br = Br.at[:, k, :].set(jnp.abs(Br[:, k, :]))
+            elif dist == "u":
+                Br = Br.at[:, k, :].set(Br_b[k] + Br_w[k] * (draws_jax[:, k, :] - 0.5))
+
+        # Utility
+        Xr = X_jax[:, :, :, rvidx]
+        if Kf > 0:
+            Xf = X_jax[:, :, :, fxidx]
+            V = (jnp.einsum("npjk,k->npj", Xf, Bf)[:, :, :, None]
+                 + jnp.einsum("npjk,nkr->npjr", Xr, Br))
+        else:
+            V = jnp.einsum("npjk,nkr->npjr", Xr, Br)
+
+        V   = V - jnp.max(V, axis=2, keepdims=True)
+        eV  = jnp.exp(V)
+        p   = eV / jnp.sum(eV, axis=2, keepdims=True)
+
+        pch   = jnp.sum(y_jax[:, :, :, None] * p, axis=2)
+        pch   = jnp.prod(pch, axis=1)
+        pch   = jnp.clip(pch, 1e-300, None)
+        sim_p = jnp.mean(pch, axis=1)
+        sim_p = jnp.clip(sim_p, 1e-300, None)
+        return -jnp.sum(jnp.log(sim_p))
+
+    # ── Scipy compute_probabilities override ────────────────────────
+
+    def compute_probabilities(self, betas, X, panel_info, draws, drawstrans,
+                              avail, var_list, chol_mat):
+        """GSE override: add gradient-latent term to Br."""
+        vals = list(var_list.values())
+        Bf   = vals[0]
+        Br_b = vals[1]
+        # GSE segments sit between Br_b and chol/Br_w
+        has_gse = len(vals) > 9
+        gamma   = vals[2] if has_gse else np.array([])
+        gamma_w = vals[3] if has_gse else np.array([])
+        chol_v  = vals[4] if has_gse else vals[2]
+        Br_w    = vals[5] if has_gse else vals[3]
+        Bftrans = vals[6] if has_gse else vals[4] if len(vals) > 4 else np.array([])
+        flmbda  = vals[7] if has_gse else vals[5] if len(vals) > 5 else np.array([])
+
+        if dev.using_gpu:
+            Bf = dev.convert_array_gpu(Bf)
+            Br_b = dev.convert_array_gpu(Br_b)
+            Br_w = dev.convert_array_gpu(Br_w)
+            if len(Bftrans) > 0: Bftrans = dev.convert_array_gpu(Bftrans)
+            if len(flmbda) > 0: flmbda = dev.convert_array_gpu(flmbda)
+            if len(gamma) > 0:
+                gamma = dev.convert_array_gpu(gamma)
+                gamma_w = dev.convert_array_gpu(gamma_w)
+
+        XBf = np.zeros((self.N, self.P, self.J))
+        if dev.using_gpu: XBf = dev.convert_array_gpu(XBf)
+
+        if self.Kf != 0:
+            Xf = X[:, :, :, self.fxidx]
+            XBf = dev.cust_einsum("npjk,k -> npj", Xf, Bf).astype(float)
+
+        XBr = np.zeros((self.N, self.P, self.J, draws.shape[2]))
+        if self.Kr != 0:
+            tmp = dev.np.matmul(chol_mat[:self.Kr, :self.Kr], draws)
+            Br = Br_b[None, :, None] + tmp
+
+            if self.Kgrad > 0 and self.gradient_scores is not None:
+                if (self.Kgrad_w > 0 and self.gamma_draws is not None
+                        and len(gamma_w) > 0):
+                    g_rnd = gamma[None, :] + gamma_w[None, :] * self.gamma_draws
+                else:
+                    g_rnd = gamma[None, :] if len(gamma) > 0 else np.zeros((1, self.Kr))
+                Br = Br + g_rnd[:, :, None] * self.gradient_scores[:, :, None]
+
+            Br = self.draws_generator.apply_distribution(Br, self.rvdist)
+            self.Br = Br
+            Xr = X[:, :, :, self.rvidx].astype(float)
+            XBr = dev.cust_einsum("npjk,nkr -> npjr", Xr, Br)
+
+        V = XBf[:, :, :, None] + XBr if self.Kf != 0 else XBr
+        eV = np.exp(np.clip(V, -700, 700))
+        probs = eV / np.clip(eV.sum(axis=2, keepdims=True), 1e-300, None)
+        return probs

{searchlibrium-0.0.151 → searchlibrium-0.0.155}/src/SearchLibrium/__init__.py

@@ -106,6 +106,7 @@ try:
     from .latent_class import LatentClassMixedLogit
     from .mdcev import MDCEVFitResult, MDCEVModel
     from .multinomial_probit import MultinomialProbit
+    from .MixedLogitGSE import MixedLogitGSE
     from .RandomP import RandomParameters
     from .constraints_builder import ConstraintBuilder, create_constraints
     from .search import Parameters
@@ -131,6 +132,7 @@ except ImportError as e:
     from latent_class import LatentClassMixedLogit
     from mdcev import MDCEVFitResult, MDCEVModel
     from multinomial_probit import MultinomialProbit
+    from MixedLogitGSE import MixedLogitGSE
     from RandomP import RandomParameters
     from constraints_builder import ConstraintBuilder, create_constraints
     from search import Parameters

{searchlibrium-0.0.151 → searchlibrium-0.0.155}/src/SearchLibrium/constraints_builder.py

@@ -135,6 +135,79 @@ class ConstraintBuilder:
         self._print(f"Mutually exclusive group: {variables}")
         return self
 
+    # ========================
+    # Alternative-Specific Variable Constraints
+    # ========================
+
+    def asvar_alt_spec(self, variable: str, alts: List[str]) -> 'ConstraintBuilder':
+        """
+        Break alternative-specific variables into subsets so they don't all
+        apply uniformly to the model. For each variable, creates
+        alternative-specific dummy columns only for the given *alts*, while
+        keeping the original variable for the remaining alternatives as a
+        generic coefficient. The search will then select among these
+        alt-specific columns as part of the variable search.
+
+        Parameters
+        ----------
+        variable : str
+            Base variable name (e.g. 'TIME', 'COST').
+        alts : list of str
+            Alternatives for which alt-specific coefficients should be
+            created (e.g. ['Car', 'Train']). Other alternatives share a
+            single generic coefficient.
+
+        Returns
+        -------
+        ConstraintBuilder
+            Self for method chaining
+
+        Example
+        -------
+        >>> constraints.asvar_alt_spec('TIME', ['Car', 'Train'])
+        >>> constraints.asvar_alt_spec('COST', ['Bus', 'Car'])
+
+        This creates ``TIME_Car``, ``TIME_Train``, ``COST_Bus``, ``COST_Car``
+        in the active variable pool while keeping ``TIME`` and ``COST`` as
+        generic columns for the remaining alternatives.
+        """
+        if 'asvar_alt_spec' not in self._constraints:
+            self._constraints['asvar_alt_spec'] = {}
+        self._constraints['asvar_alt_spec'][variable] = list(alts)
+        self._print(f"Alt-specific subset: '{variable}' → {alts}")
+        return self
+
+    def incompatible_specs(self, *alt_var_names: str) -> 'ConstraintBuilder':
+        """
+        Declare that certain alternative-specific variable names (e.g.,
+        ``TIME_Car`` and ``COST_Car``) are incompatible and must not appear
+        together in the same model. Each call creates one exclusion group.
+
+        This is applied **after** the alternative-specific dummies are
+        created, so you can declare constraints on the resulting column
+        names.
+
+        Parameters
+        ----------
+        *alt_var_names : str
+            Alt-specific column names forming a mutually exclusive group.
+
+        Returns
+        -------
+        ConstraintBuilder
+            Self for method chaining
+
+        Example
+        -------
+        >>> constraints.incompatible_specs('TIME_Car', 'COST_Car')
+        >>> constraints.incompatible_specs('FARE_Train', 'TIME_Train')
+        """
+        if 'incompatible_specs' not in self._constraints:
+            self._constraints['incompatible_specs'] = []
+        self._constraints['incompatible_specs'].append(list(alt_var_names))
+        self._print(f"Incompatible alt-specific group: {alt_var_names}")
+        return self
+
     def min_behavioral(self, min_count: int, *variables: str) -> 'ConstraintBuilder':
         """
         Require at least ``min_count`` variables from a behavioural pool

{searchlibrium-0.0.151 → searchlibrium-0.0.155}/src/SearchLibrium/harmony.py

@@ -949,10 +949,11 @@ class HarmonySearch(Search):
         """Identify groups of variables that measure similar constructs
         and can be substituted for one another."""
         groups = []
-        # Group 1: destination size measures
+        # Group 1: destination size measures (now includes per-type ACDC)
         size_measures = [v for v in asvars
                          if v in {'ln_dest_trips', 'osm_destination_pull',
-                                  'acdc_aggl', 'osm_poi_intensity'}]
+                                  'acdc_aggl', 'osm_poi_intensity'}
+                         or v.startswith('acdc_aggl_')]
         for i in range(len(size_measures)):
             for j in range(i + 1, len(size_measures)):
                 groups.append((size_measures[i], size_measures[j]))
@@ -972,6 +973,25 @@ class HarmonySearch(Search):
             for j in range(i + 1, len(osm_measures)):
                 groups.append((osm_measures[i], osm_measures[j]))
 
+        # Group 4: Per-type ACDC agglomeration terms (substitutable by OSM type)
+        per_type_aggl = [v for v in asvars if v.startswith('acdc_aggl_')]
+        for i in range(len(per_type_aggl)):
+            for j in range(i + 1, len(per_type_aggl)):
+                groups.append((per_type_aggl[i], per_type_aggl[j]))
+
+        # Group 5: Per-type ACDC competition terms (substitutable by OSM type)
+        per_type_comp = [v for v in asvars if v.startswith('acdc_comp_')]
+        for i in range(len(per_type_comp)):
+            for j in range(i + 1, len(per_type_comp)):
+                groups.append((per_type_comp[i], per_type_comp[j]))
+
+        # Group 6: Cross-substitution — agglomeration <-> competition for same type
+        for t in ['retail', 'food', 'leisure', 'parks', 'transit', 'cycleway', 'culture']:
+            aggl_var = f'acdc_aggl_{t}'
+            comp_var = f'acdc_comp_{t}'
+            if aggl_var in asvars and comp_var in asvars:
+                groups.append((aggl_var, comp_var))
+
         return groups
 
     def _get_forced_vars(self):