SearchLibrium 0.0.152__tar.gz → 0.0.153__tar.gz

{searchlibrium-0.0.152/src/SearchLibrium.egg-info → searchlibrium-0.0.153}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: SearchLibrium
-Version: 0.0.152
+Version: 0.0.153
 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.152 → searchlibrium-0.0.153}/pyproject.toml

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

{searchlibrium-0.0.152 → searchlibrium-0.0.153}/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.153/src/SearchLibrium/MixedLogitGSE.py

@@ -0,0 +1,221 @@
+"""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()
+"""
+
+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 to O(1) to prevent numerical overflow
+            g = np.asarray(gradient_scores, dtype=float)
+            self.gradient_scores = g / (g.std(axis=0, keepdims=True) + 1e-8)
+            gstd = g.std(axis=0)
+            print(f"[MixedLogitGSE] Gradient loadings enabled (Kgrad={len(kwargs.get('randvars', {}))}, "
+                  f"random_gamma={self.random_gamma}, "
+                  f"gradient std range=[{gstd.min():.2f}, {gstd.max():.2f}] -> standardised)")
+        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:
+                self.gamma_draws = np.random.randn(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.153/src/SearchLibrium/version.txt

@@ -0,0 +1 @@
+0.0.153

{searchlibrium-0.0.152 → searchlibrium-0.0.153/src/SearchLibrium.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: SearchLibrium
-Version: 0.0.152
+Version: 0.0.153
 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.152 → searchlibrium-0.0.153}/src/SearchLibrium.egg-info/SOURCES.txt

@@ -4,6 +4,7 @@ pyproject.toml
 src/SearchLibrium/BEHier.py
 src/SearchLibrium/Halton.py
 src/SearchLibrium/MixedLogit.py
+src/SearchLibrium/MixedLogitGSE.py
 src/SearchLibrium/Mode_Activity_Nested.py
 src/SearchLibrium/RandomP.py
 src/SearchLibrium/SEARCH_SM_MARIO.py
@@ -54,4 +55,7 @@ src/SearchLibrium/bhhh/minimize.py
 src/SearchLibrium/data/Swissmetro_final.csv
 src/SearchLibrium/data/TravelMode.csv
 src/SearchLibrium/data/electricity.csv
-src/data_generator/latent_class_gen.py
+src/data_generator/demonstrate_generation.py
+src/data_generator/fit_recovery.py
+src/data_generator/latent_class_gen.py
+src/data_generator/print_synthetic_data.py

searchlibrium-0.0.153/src/data_generator/demonstrate_generation.py

@@ -0,0 +1,85 @@
+"""
+demonstrate_generation.py
+=========================
+Demonstrates how to generate latent class choice data with membership variables.
+Run this to inspect the generated dataset and parameter structure.
+"""
+
+import sys
+import os
+import numpy as np
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from data_generator.latent_class_gen import AdvancedLatentClassGenerator
+
+
+def main(seed=42):
+    print("=" * 70)
+    print("  Data Generation Demo (with Membership Variables)")
+    print("=" * 70)
+
+    gen = AdvancedLatentClassGenerator(
+        n_classes=3,
+        n_alternatives=3,
+        n_individuals=500,
+        n_choice_tasks=2,
+        scale_separation=2.5,
+        n_noise_vars=2,
+        n_weak_vars=1,
+        n_collinear_vars=1,
+        n_membership_vars=3,
+        membership_scale=2.0,
+        random_state=seed,
+    )
+
+    print(f"\n  Classes:         {gen.K}")
+    print(f"  Alternatives:    {gen.J}")
+    print(f"  Individuals:     {gen.N}")
+    print(f"  Choice tasks:    {gen.T}")
+    print(f"  Shared vars:     {gen.shared_vars}")
+    print(f"  Membership vars: {gen.membership_var_names}")
+    print(f"  Noise vars:      {gen.noise_vars}")
+    print(f"  Weak vars:       {gen.weak_vars}")
+    print(f"  Collinear vars:  {gen.collinear_vars}")
+
+    print(f"\n  Class-specific utility vars:")
+    for k in range(gen.K):
+        print(f"    Class {k}: {gen.class_specific_vars.get(k, [])}")
+
+    print(f"\n  True Utility Coefficients (betas):")
+    for k in range(gen.K):
+        print(f"    Class {k}:")
+        for v, b in gen.parameters[k].items():
+            print(f"      {v:>25s} = {b:+.4f}")
+
+    print(f"\n  True Membership Coefficients (gammas):")
+    if gen.gammas is not None:
+        for c in range(gen.K - 1):
+            for m in range(gen.n_membership_vars):
+                print(f"    Class_{c + 1}_{gen.membership_var_names[m]} = {gen.gammas[c, m]:+.4f}")
+        print(f"    Class_{gen.K} (reference) all = 0")
+
+    df, true_classes = gen.generate()
+
+    print(f"\n  Generated data: {len(df)} rows")
+    print(f"  Columns: {list(df.columns)}")
+    print(f"  Class distribution:")
+    for k in range(gen.K):
+        n_k = (true_classes == k).sum()
+        print(f"    Class {k}: {n_k} ({n_k / gen.N:.1%})")
+
+    print(f"\n  First 5 rows:")
+    print(df.head(10).to_string())
+
+    print(f"\n  Individual-level membership vars (first 5 individuals):")
+    mem_cols = gen.membership_var_names
+    for n in range(min(5, gen.N)):
+        vals = {v: df.loc[df["individual"] == n, v].iloc[0] for v in mem_cols}
+        print(f"    Individual {n} (true_class={df.loc[df['individual'] == n, 'true_class'].iloc[0]}): {vals}")
+
+    print(f"\n  Done. Use the generated data with SearchLibrium's LatentClassMixedLogit.")
+    return gen, df, true_classes
+
+
+if __name__ == "__main__":
+    main()