SearchLibrium 0.0.261__tar.gz → 0.0.263__tar.gz

{searchlibrium-0.0.261/src/SearchLibrium.egg-info → searchlibrium-0.0.263}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: SearchLibrium
-Version: 0.0.261
+Version: 0.0.263
 Summary: Automated specification search for discrete choice models (logit, mixed logit, nested logit, RRM and more) using metaheuristic optimisation
 Author: Zeke Ahern, Mario Fernando Taco Morales, Alexander Paz, Prithvi Beeramole, Robert Burdett
 Maintainer-email: Zeke Ahern <z.ahern@qut.edu.au>

{searchlibrium-0.0.261 → searchlibrium-0.0.263}/pyproject.toml

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

{searchlibrium-0.0.261 → searchlibrium-0.0.263}/src/SearchLibrium/Halton.py

@@ -56,12 +56,14 @@ class Halton:
     Set use_sobol=True to use scrambled Sobol sequences instead.
     """
 
-    def __init__(self, primes=None, drop=100, shuffled=False, antithetic=False, use_sobol=False):
+    def __init__(self, primes=None, drop=100, shuffled=False,
+                 antithetic=False, use_sobol=False, seed=None):
         self.primes = primes
         self.drop = drop
         self.shuffled = shuffled
         self.antithetic = antithetic
         self.use_sobol = use_sobol
+        self.seed = seed
 
         # Default primes for Halton sequence
         if self.primes is None:
@@ -81,7 +83,9 @@ class Halton:
         if self.use_sobol:
             # Use Sobol sequence (newer method)
             base = n_draws // 2 if self.antithetic else n_draws
-            draws = _sobol_generate(sample_size, base, n_vars, shuffled=self.shuffled)
+            draws = _sobol_generate(
+                sample_size, base, n_vars, shuffled=self.shuffled,
+                seed=self.seed)
             if self.antithetic:
                 draws = np.concatenate([draws, 1.0 - draws], axis=2)
         else:
@@ -110,7 +114,7 @@ class Halton:
         return draws
 
 
-def _sobol_generate(sample_size, n_draws, n_vars, shuffled=False):
+def _sobol_generate(sample_size, n_draws, n_vars, shuffled=False, seed=None):
     """Generate scrambled Sobol (0,1) draws.
 
     Returns
@@ -123,7 +127,7 @@ def _sobol_generate(sample_size, n_draws, n_vars, shuffled=False):
         p2 = 1
         while p2 < total:
             p2 <<= 1
-        sobol = Sobol(d=n_vars, scramble=True)
+        sobol = Sobol(d=n_vars, scramble=True, seed=seed)
         flat = sobol.random(p2)[:total]  # (total, n_vars)
         draws = flat.reshape(sample_size, n_draws, n_vars).transpose(0, 2, 1)
     else:

{searchlibrium-0.0.261 → searchlibrium-0.0.263}/src/SearchLibrium/MixedLogit.py

@@ -122,10 +122,19 @@ def _panel_weighted_shares(prob_np, panel_info):
     return ind, np.mean(ind, axis=0)
 
 class MixedLogit(DiscreteChoiceModel):
-    def __init__(self, halton_opts=None, distributions=['n', 'ln', 'nln', 't', 'tn', 'u'], _jax=None):
+    def __init__(self, halton_opts=None, distributions=['n', 'ln', 'nln', 't', 'tn', 'u'], _jax=None,
+                 save_prob_full=None):
         super().__init__(_jax)
         self.descr = "MXL"
         self.halton_opts = halton_opts
+        # MEMORY: prob_full is the full (N, P, J, R) probability tensor built in
+        # post-process for the summary/plots. For a 13k-case/20-alt/1000-draw
+        # model that is ~2 GB, and it is never read during estimation, so it is
+        # now opt-in. Pass save_prob_full=True (or set
+        # SEARCHLIBRIUM_SAVE_PROB_FULL=1) to retain the previous behaviour.
+        if save_prob_full is None:
+            save_prob_full = bool(os.environ.get("SEARCHLIBRIUM_SAVE_PROB_FULL", ""))
+        self.save_prob_full = bool(save_prob_full)
         self.draws_generator = Draws(k=len(distributions), halton_opts=halton_opts, rvdist=distributions)
         self.random_parameters = RandomParameters(distributions or [])
         self.softmax_r = True
@@ -225,7 +234,8 @@ class MixedLogit(DiscreteChoiceModel):
               randvars=None, panels=None, base_alt=None, fit_intercept=False,
               init_coeff=None, maxiter=1000, correlated_vars=None,
               n_draws=1000, halton=True, minimise_func=None,
-              batch_size=None, halton_opts=None, ftol=1e-12,
+              batch_size=None, person_batch_size=None, halton_opts=None,
+              draw_seed=0, 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,
@@ -301,7 +311,15 @@ class MixedLogit(DiscreteChoiceModel):
         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
-        self.halton_opts = halton_opts  # Store halton options for draw generation
+        if person_batch_size is None:
+            self.person_batch_size = None
+        else:
+            self.person_batch_size = max(1, int(person_batch_size))
+        self._person_batch_active = False
+        self.draw_seed = int(draw_seed)
+        self.halton_opts = dict(halton_opts or {})
+        if self.person_batch_size is not None:
+            self.halton_opts.setdefault("seed", self.draw_seed)
 
         # CRITICAL FIX: Recreate draws_generator with proper halton_opts AND rvdist
         # The __init__ creates it with halton_opts=None, but setup() gets the actual options
@@ -309,7 +327,8 @@ class MixedLogit(DiscreteChoiceModel):
         # AND to pass the correct random variable distributions (e.g., 'ln' for lognormal)
         # IMPORTANT: Order rvdist by varnames order, not dictionary order
         rvdist = [randvars[var] for var in varnames if var in randvars] if randvars else []
-        self.draws_generator = Draws(k=len(randvars or {}), halton_opts=halton_opts, rvdist=rvdist)
+        self.draws_generator = Draws(
+            k=len(randvars or {}), halton_opts=self.halton_opts, rvdist=rvdist)
 
         self.randvarsdict = randvars  # random variables not transformed
 
@@ -812,7 +831,16 @@ class MixedLogit(DiscreteChoiceModel):
     def fit(self):
         # {
         # Generate draws:
-        draws, drawstrans = self.generate_draws(self.N, self.n_draws, self.halton)
+        streaming_people = (
+            self.person_batch_size is not None
+            and self.person_batch_size < self.N
+        )
+        if streaming_people:
+            # Person chunks regenerate deterministic draws during each
+            # objective evaluation, avoiding an (N, Kr, R) draw bank.
+            draws, drawstrans = None, None
+        else:
+            draws, drawstrans = self.generate_draws(self.N, self.n_draws, self.halton)
         self.draws, self.drawstrans = draws, drawstrans  # Record generated values
         # Explicit caller-supplied starts are respected verbatim (no ASC
         # share seeding); the MNL block below may still overwrite
@@ -829,9 +857,15 @@ class MixedLogit(DiscreteChoiceModel):
                 from numba_engine import effective_engine as _mx_eff_eng
             except ImportError:
                 from .numba_engine import effective_engine as _mx_eff_eng
-            use_numba = (_mx_eff_eng(self) == 'numba')
+            # Numba is the default CPU engine. JAX must be selected
+            # explicitly; unsupported specifications fall back to SciPy.
+            use_numba = (_mx_eff_eng(self) != 'jax')
         except Exception:
-            use_numba = (getattr(self, 'engine', None) == 'numba')
+            use_numba = (getattr(self, 'engine', None) != 'jax')
+        if streaming_people:
+            # The streaming path regenerates person-chunk draws and is not
+            # compatible with the current full-array Numba adapter.
+            use_numba = False
 
         # 2x Kftrans - mean and lambda, 3x Krtrans - mean, s.d., lambda
         # Kchol, Kbw - relate to random variables, non-transformed
@@ -1052,11 +1086,12 @@ class MixedLogit(DiscreteChoiceModel):
                 if _v: print(f"[MXL] Differential evolution initialization failed: {e}")
 
         # '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
-        if dev.using_gpu:  # {
+        if dev.using_gpu and draws is not None:  # {
             self.X, self.y = dev.convert_array_gpu(self.X), dev.convert_array_gpu(self.y)
             self.panel_info = dev.convert_array_gpu(self.panel_info)
             draws = dev.convert_array_gpu(draws)
-            drawstrans = dev.convert_array_gpu(drawstrans)
+            if drawstrans is not None:
+                drawstrans = dev.convert_array_gpu(drawstrans)
             self.weights = dev.convert_array_gpu(self.weights) if self.weights is not None else None
             self.avail = dev.convert_array_gpu(self.avail) if self.avail is not None else None
         # }
@@ -1104,7 +1139,8 @@ class MixedLogit(DiscreteChoiceModel):
         # SOLVE OPTIMISATION PROBLEM - COMPUTATIONALLY TIME CONSUMING!
         # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
             # ---- JAX fast path ----
-        if (getattr(self, '_jax', False) and not use_numba
+        if (getattr(self, '_jax', False) and not streaming_people
+                and not use_numba
                 and self.minimise_func is None):
             jax_result = self.optimize_jax(betas, draws, drawstrans)
             if jax_result is not None:
@@ -1173,7 +1209,8 @@ class MixedLogit(DiscreteChoiceModel):
                 self.choice_pred_prob = np.mean(p, axis=3)
                 self.ind_pred_prob, self.pred_prob = _panel_weighted_shares(
                     self.choice_pred_prob, self.panel_info)
-                self.prob_full = p
+                if self.save_prob_full:
+                    self.prob_full = p
                 self._refresh_corr_from_varlist()
                 self.post_process(jax_result, self.Xnames, self.N)
                 return
@@ -1234,6 +1271,16 @@ class MixedLogit(DiscreteChoiceModel):
                 except Exception:
                     pass
 
+        if draws is None:
+            # Estimation stayed streamed; materialize full draws only for the
+            # final fitted-parameter and prediction summaries.
+            draws, drawstrans = self.generate_draws(
+                self.N, self.n_draws, self.halton)
+            self.draws, self.drawstrans = draws, drawstrans
+            self.p = self.compute_probabilities(
+                result['x'], self.X, self.panel_info, draws, drawstrans,
+                self.avail, self.var_list, self.chol_mat)
+
         if self.save_fitted_params:
             self.compute_fitted_params(self.y, self.p, self.panel_info, self.Br)
 
@@ -1251,7 +1298,8 @@ class MixedLogit(DiscreteChoiceModel):
             self.ind_pred_prob, self.pred_prob = _panel_weighted_shares(
                 self.choice_pred_prob, self.panel_info)
 
-            self.prob_full = p
+            if self.save_prob_full:
+                self.prob_full = p
         # }
         self._refresh_corr_from_varlist()
         self.post_process(result, self.Xnames, self.N)
@@ -1413,7 +1461,8 @@ class MixedLogit(DiscreteChoiceModel):
 
         sim_p = jnp.mean(pch, axis=1)                        # (N,)
         sim_p = jnp.clip(sim_p, 1e-300, None)
-        ll = -jnp.sum(jnp.log(sim_p))
+        valid_person = jnp.any(panel_info_jax > 0, axis=1)
+        ll = -jnp.sum(jnp.where(valid_person, jnp.log(sim_p), 0.0))
         # Tiny L2 penalty — makes the Hessian diagonally dominant so it
         # never becomes singular, even with near-constant variables or
         # simulation noise in the mixed-logit probability integration.
@@ -1499,7 +1548,11 @@ class MixedLogit(DiscreteChoiceModel):
     # process; during a heterogeneous GA search that is one retained model
     # per distinct parameter-count combination.
     # ------------------------------------------------------------------
-    _JIT_CACHE_MAXSIZE = 16
+    # MEMORY: cache size is in *live models*, each of which can pin its X/y
+    # arrays and a compiled XLA executable. The old default of 16 was tuned for
+    # speed; 4 keeps most of the reuse benefit while capping the worst case.
+    # Override with SEARCHLIBRIUM_JIT_CACHE_MAXSIZE (0 disables the cache).
+    _JIT_CACHE_MAXSIZE = max(0, int(os.environ.get("SEARCHLIBRIUM_JIT_CACHE_MAXSIZE", "4")))
     _mxl_jit_cache: "OrderedDict[tuple, tuple]" = OrderedDict()
 
     @staticmethod
@@ -1562,6 +1615,17 @@ class MixedLogit(DiscreteChoiceModel):
             pass
         gc.collect()
 
+    def release_memory(self):
+        """Free per-fit caches and large post-process tensors this model holds.
+
+        Safe to call between metaheuristic candidates: it clears the process-
+        wide JIT/XLA cache and drops ``prob_full`` unless it was explicitly
+        saved. Fitted coefficients, standard errors and covariance are kept.
+        """
+        if not getattr(self, "save_prob_full", False):
+            self.prob_full = None
+        type(self)._clear_jit_cache()
+
     def _run_jaxopt(self, betas, X_jax, y_jax, pi_jax, draws_jax,
                     _fn, _cache_key):
         """JAX likelihood + ``jaxopt.LBFGS`` solver (experiment).
@@ -1640,6 +1704,125 @@ class MixedLogit(DiscreteChoiceModel):
                 print(f"[JAX MXL optimizer] jaxopt run failed ({_e}).")
             return None
 
+    def _optimize_jax_streamed(self, betas):
+        """Fit plain MXL with bounded person chunks and deterministic draws."""
+        try:
+            import os
+            import jax
+            import jax.numpy as jnp
+            import numpy as _np
+            from scipy.optimize import minimize as sp_min
+
+            chunk_size = int(self.person_batch_size)
+            X_np = _np.asarray(self.X)
+            y_np = _np.asarray(self.y)
+            if y_np.ndim > 3:
+                y_np = y_np[..., 0]
+            pi_np = _np.asarray(self.panel_info)
+            N, P, J, _ = X_np.shape
+            Kf, Kr = int(self.Kf), int(self.Kr)
+            Kchol, Kbw = int(self.Kchol), int(self.Kbw)
+            fxidx = jnp.asarray(self.fxidx, dtype=bool)
+            rvidx = jnp.asarray(self.rvidx, dtype=bool)
+            rvdist_names = [d for d in self.rvdist if d is not False]
+            extra = self._jax_negloglik_extra_kwargs()
+
+            def _pad(value, shape):
+                result = _np.zeros(shape, dtype=_np.float64)
+                result[tuple(slice(0, n) for n in value.shape)] = value
+                return result
+
+            def _chunk_inputs(chunk_index):
+                start = chunk_index * chunk_size
+                end = min(start + chunk_size, N)
+                n = end - start
+                X_chunk = _pad(X_np[start:end], (chunk_size, P, J, X_np.shape[3]))
+                y_chunk = _pad(y_np[start:end], (chunk_size, P, J))
+                pi_chunk = _pad(pi_np[start:end], (chunk_size, P))
+
+                state = _np.random.get_state()
+                try:
+                    _np.random.seed(self.draw_seed + chunk_index)
+                    if self.halton and hasattr(self.draws_generator, "halton"):
+                        self.draws_generator.halton.seed = (
+                            self.draw_seed + chunk_index)
+                    draws_chunk, _ = self.generate_draws(
+                        n, self.n_draws, self.halton)
+                finally:
+                    _np.random.set_state(state)
+                draws_chunk = _pad(
+                    _np.asarray(draws_chunk),
+                    (chunk_size, int(self.Kr), self.n_draws))
+                return (
+                    jnp.asarray(X_chunk),
+                    jnp.asarray(y_chunk),
+                    jnp.asarray(pi_chunk),
+                    jnp.asarray(draws_chunk),
+                )
+
+            def _chunk_objective(b, X_j, y_j, pi_j, draws_j):
+                return self._jax_mxl_negloglik(
+                    b, X_j, y_j, pi_j, draws_j,
+                    fxidx, rvidx, Kf, Kr, Kchol, Kbw,
+                    rvdist_names, int(self.correlationLength),
+                    **extra)
+
+            compiled = jax.jit(jax.value_and_grad(_chunk_objective))
+            chunk_args = [
+                _chunk_inputs(i)
+                for i in range((N + chunk_size - 1) // chunk_size)
+            ]
+            valid_counts = [
+                min(chunk_size, N - i * chunk_size)
+                for i in range(len(chunk_args))
+            ]
+
+            def _objective(betas_np):
+                b = jnp.asarray(betas_np, dtype=jnp.float64)
+                value = 0.0
+                gradient = jnp.zeros_like(b)
+                for args, valid in zip(chunk_args, valid_counts):
+                    chunk_value, chunk_gradient = compiled(b, *args)
+                    # Padded people contribute zero because their panel mask
+                    # is zero; valid is retained for clarity and future
+                    # weighted-panel support.
+                    del valid
+                    value = value + chunk_value
+                    gradient = gradient + chunk_gradient
+                value, gradient = jax.block_until_ready((value, gradient))
+                return float(value), _np.asarray(gradient, dtype=_np.float64)
+
+            options = {'maxiter': self.maxiter, 'disp': False}
+            if self.method in ('bfgs', 'l-bfgs-b'):
+                options['gtol'] = self.gtol
+                options['ftol'] = self.ftol
+            else:
+                options['ftol'] = self.ftol
+            result = sp_min(
+                _objective, betas, jac=True, method=self.method,
+                options=options)
+
+            if self.return_hess:
+                def _hessian_objective(b):
+                    value = 0.0
+                    for args in chunk_args:
+                        value = value + _chunk_objective(b, *args)
+                    return value
+
+                hessian_fn = jax.jit(jax.hessian(_hessian_objective))
+                hessian = _np.asarray(
+                    hessian_fn(jnp.asarray(result['x'], dtype=jnp.float64)),
+                    dtype=float)
+                try:
+                    result['hess_inv'] = _np.linalg.inv(hessian)
+                except _np.linalg.LinAlgError:
+                    result['hess_inv'] = _np.linalg.pinv(hessian)
+            return result
+        except Exception as exc:
+            if not _SL_QUIET:
+                print(f"[JAX streamed MXL] falling back to scipy: {exc}")
+            return None
+
     def optimize_jax(self, betas, draws, drawstrans):
         """JAX-accelerated optimisation for the Mixed Logit model.
 
@@ -1655,6 +1838,10 @@ class MixedLogit(DiscreteChoiceModel):
                 print("[JAX MXL optimizer] heterogeneity detected; falling back to scipy.")
             return None
 
+        if (self.person_batch_size is not None
+                and self.person_batch_size < self.N):
+            return self._optimize_jax_streamed(betas)
+
         try:
             import os
             os.environ.setdefault("JAX_ENABLE_X64", "True")
@@ -1914,12 +2101,88 @@ class MixedLogit(DiscreteChoiceModel):
                 print(f"[JAX MXL masked] falling back: {e}")
             return None
 
+    def _person_batch_slice(self, value, start, end):
+        """Slice an observation-level array while preserving scalar settings."""
+        if value is None:
+            return None
+        try:
+            if np.ndim(value) > 0 and np.shape(value)[0] == self.N:
+                return value[start:end]
+        except (TypeError, IndexError):
+            pass
+        return value
+
+    def _stream_person_batches(self, betas, X, y, panel_info, draws,
+                               drawstrans, weights, avail, batch_size):
+        """Evaluate the objective in bounded person chunks.
+
+        When ``draws`` is absent, each chunk receives deterministic random
+        draws derived from ``draw_seed`` and its chunk index. This keeps every
+        optimizer evaluation reproducible without retaining all draws.
+        """
+        original_n = self.N
+        total_value = 0.0
+        total_gradient = None
+        has_gradient = bool(self.return_grad)
+        try:
+            self._person_batch_active = True
+            for chunk_index, start in enumerate(
+                    range(0, original_n, self.person_batch_size)):
+                end = min(start + self.person_batch_size, original_n)
+                self.N = end - start
+                chunk_X = X[start:end]
+                chunk_y = y[start:end]
+                chunk_panel = panel_info[start:end]
+                chunk_weights = self._person_batch_slice(weights, start, end)
+                chunk_avail = self._person_batch_slice(avail, start, end)
+
+                if draws is None:
+                    rng_state = np.random.get_state()
+                    try:
+                        np.random.seed(self.draw_seed + chunk_index)
+                        if self.halton and hasattr(self.draws_generator, "halton"):
+                            self.draws_generator.halton.seed = (
+                                self.draw_seed + chunk_index)
+                        chunk_draws, chunk_drawstrans = self.generate_draws(
+                            self.N, self.n_draws, self.halton)
+                    finally:
+                        np.random.set_state(rng_state)
+                else:
+                    chunk_draws = draws[start:end]
+                    chunk_drawstrans = (
+                        drawstrans[start:end] if drawstrans is not None
+                        else drawstrans
+                    )
+
+                result = self.get_loglik_gradient(
+                    betas, chunk_X, chunk_y, chunk_panel,
+                    chunk_draws, chunk_drawstrans, chunk_weights,
+                    chunk_avail, batch_size)
+                total_value += float(result[0])
+                if has_gradient and len(result) > 1:
+                    if total_gradient is None:
+                        total_gradient = np.zeros_like(result[1], dtype=float)
+                    total_gradient += np.asarray(result[1], dtype=float)
+        finally:
+            self.N = original_n
+            self._person_batch_active = False
+
+        if has_gradient:
+            return total_value, total_gradient
+        return (total_value,)
+
     def get_loglik_gradient(self, betas, X, y, panel_info, draws,
                             drawstrans, weights, avail, batch_size):
         # {
         """ Fixed and random parameters are handled separately to
         speed up the estimation and the results are concatenated.
         """
+        if (not self._person_batch_active
+                and self.person_batch_size is not None
+                and self.person_batch_size < self.N):
+            return self._stream_person_batches(
+                betas, X, y, panel_info, draws, drawstrans,
+                weights, avail, batch_size)
         # Segregating initial values to fixed betas (Bf),
         # random beta means (Br_b)
         # for both non-transformed and transformed variables
@@ -1999,11 +2262,10 @@ class MixedLogit(DiscreteChoiceModel):
 
         # _, gr_b, gr_w, pch = np.zeros((N, self.Kf)), np.zeros((N, self.Kr)), np.zeros((N, self.Kr)), []  # Batch data
 
-        _ = np.zeros((N, self.Kf))  # CONCEPTUAL ERROR. WHY CREATE IF NO NAME?
         gr_b = np.zeros((N, self.Kr))
         gr_w = np.zeros((N, self.Kr))
         g_all = np.zeros((N, len(betas)))
-        pch = []  # Batch data
+        lik_sum = np.zeros(N, dtype=float)
         for batch in range(n_batches):
             # {
             a = batch * batch_size
@@ -2025,6 +2287,8 @@ class MixedLogit(DiscreteChoiceModel):
 
             # Thresholds to avoid divide by zero warnings
             pch_batch = truncate_lower(pch_batch, min_comp_val)
+            draw_count = b - a
+            lik_sum += np.sum(pch_batch, axis=1)
 
             # Observed probability minus predicted probability
             ymp = y - p  # (N, P, J, R)
@@ -2263,12 +2527,10 @@ class MixedLogit(DiscreteChoiceModel):
                 g, pch_batch = dev.to_cpu(g), dev.to_cpu(pch_batch)
             # '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
 
-            pch.append(pch_batch)
-            g_all += g
+            g_all += g * draw_count
         # }
 
-        pch = np.concatenate(pch, axis=-1)
-        lik = pch.mean(axis=1)  # (N,)
+        lik = lik_sum / float(self.n_draws)  # (N,)
         loglik = np.log(lik)
         if weights is not None: loglik = loglik * weights
         loglik = loglik.sum()
@@ -2296,7 +2558,7 @@ class MixedLogit(DiscreteChoiceModel):
                 g *= weights[:, 0][:, None]  # Use only the first column and convert to column vector
         # }
 
-        g = np.sum(g, axis=0) / n_batches  # (K, )
+        g = np.sum(g_all, axis=0) / float(self.n_draws)  # (K, )
         g = g - self.regularize_l1_grad(betas)
         # L2 ridge penalty gradient: the objective subtracts
         # reg_penalty * sum(betas^2) from the log-likelihood, so its
@@ -2645,14 +2907,10 @@ class MixedLogit(DiscreteChoiceModel):
 
         # INITIALISE
         XBf = np.zeros((self.N, self.P, self.J))
-        XBr = np.zeros((self.N, self.P, self.J, self.batch_size))  # NOT USED?
-        V = np.zeros((self.N, self.P, self.J, self.batch_size))  # NOT USED?
 
         # '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
         if dev.using_gpu:  # {
             XBf = dev.convert_array_gpu(XBf)
-            XBr = dev.convert_array_gpu(XBr)
-            V = dev.convert_array_gpu(V)
         # }
         # '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
 
@@ -2692,10 +2950,16 @@ class MixedLogit(DiscreteChoiceModel):
             Br = self.draws_generator.apply_distribution(Br, self.rvdist)
             self.Br = Br  # save Br to use later
             Xr = X[:, :, :, self.rvidx]
-            XBr = dev.cust_einsum('npjk,nkr -> npjr', Xr, Br).astype(float)  # (N, P, J, R)
-            V = XBf[:, :, :, None] + XBr  # Add an extra dimension to XBf and then add XBr
+            # Build the utility directly. Keeping a separate XBr and a
+            # preallocated V needlessly retains two (N, P, J, R) work arrays.
+            V = dev.cust_einsum('npjk,nkr -> npjr', Xr, Br).astype(float)
+            V += XBf[:, :, :, None]
         else:
             self.Br = Br_b[None, :, None]
+            V = np.broadcast_to(
+                XBf[:, :, :, None],
+                (self.N, self.P, self.J, draws.shape[-1]),
+            ).copy()
 
         #  transformations for variables with fixed coeffs
         if self.Kftrans != 0:

{searchlibrium-0.0.261 → searchlibrium-0.0.263}/src/SearchLibrium/destination_distances.py

@@ -167,6 +167,7 @@ def build_case_alt_distances(
     dest_lat_col: str = "dest_lat",
     dest_lon_col: str = "dest_lon",
     offset_col: str | None = "distance_offset_km",
+    member_zones_col: str | None = "member_suburbs",
 ) -> pd.DataFrame:
     """Build larch-style per-(case, alternative) distances.
 
@@ -204,7 +205,12 @@ def build_case_alt_distances(
         raise KeyError(f"build_case_alt_distances: alternatives missing columns {missing}")
 
     c = cases[req_cases].copy()
-    a = alternatives[req_alts + ([offset_col] if offset_col and offset_col in alternatives.columns else [])].copy()
+    extra_alt_cols = []
+    if offset_col and offset_col in alternatives.columns:
+        extra_alt_cols.append(offset_col)
+    if member_zones_col and member_zones_col in alternatives.columns:
+        extra_alt_cols.append(member_zones_col)
+    a = alternatives[req_alts + extra_alt_cols].copy()
     c["_oz"] = c[orig_zone_col].astype(str).str.strip().str.upper()
     a["_dz"] = a[dest_zone_col].astype(str).str.strip().str.upper()
     if offset_col and offset_col in a.columns:
@@ -225,24 +231,41 @@ def build_case_alt_distances(
     # Skim index lookup (zone label -> matrix row).
     skim_lookup: dict[str, int] = {}
     if skim_zone_ids is not None and skim_dist_mat is not None:
-        skim_lookup = {z: i for i, z in enumerate(skim_zone_ids)}
+        skim_lookup = {str(z).strip().upper(): i
+                       for i, z in enumerate(skim_zone_ids)}
     o_idx = np.array([skim_lookup.get(z, -1) for z in c["_oz"].to_numpy()], dtype=int)
     d_idx = np.array([skim_lookup.get(z, -1) for z in a["_dz"].to_numpy()], dtype=int)
 
     skim_d = np.full((n_c, n_a), np.nan, dtype=float)
     skim_t = np.full((n_c, n_a), np.nan, dtype=float)
-    if skim_dist_mat is not None and (o_idx >= 0).any() and (d_idx >= 0).any():
+    if skim_dist_mat is not None and (o_idx >= 0).any():
         sd = np.asarray(skim_dist_mat, dtype=float)
         st = np.asarray(skim_time_mat, dtype=float) if skim_time_mat is not None else None
         ok_c = np.where(o_idx >= 0)[0]
-        ok_a = np.where(d_idx >= 0)[0]
-        sub = sd[np.ix_(o_idx[ok_c], d_idx[ok_a])]
-        usable = (sub > 0.0) & (sub < float(max_km))
-        block = np.where(usable, sub, np.nan)
-        skim_d[np.ix_(ok_c, ok_a)] = block
-        if st is not None and st.shape == sd.shape:
-            sub_t = st[np.ix_(o_idx[ok_c], d_idx[ok_a])]
-            skim_t[np.ix_(ok_c, ok_a)] = np.where(usable, sub_t, np.nan)
+        ok_a = np.arange(n_a) if member_zones_col and member_zones_col in a.columns \
+            else np.where(d_idx >= 0)[0]
+        if member_zones_col and member_zones_col in a.columns:
+            for ci in ok_c:
+                for ai in ok_a:
+                    members = [z.strip().upper() for z in str(
+                        a.iloc[ai][member_zones_col] or "").split(",") if z.strip()]
+                    member_idx = [skim_lookup[z] for z in members if z in skim_lookup]
+                    if not member_idx:
+                        member_idx = [d_idx[ai]]
+                    values = sd[o_idx[ci], member_idx]
+                    valid = np.isfinite(values) & (values > 0.0) & (values < float(max_km))
+                    if valid.any():
+                        pick = int(np.flatnonzero(valid)[np.argmin(values[valid])])
+                        skim_d[ci, ai] = values[valid][pick]
+                        if st is not None and st.shape == sd.shape:
+                            skim_t[ci, ai] = st[o_idx[ci], member_idx[np.flatnonzero(valid)[pick]]]
+        else:
+            sub = sd[np.ix_(o_idx[ok_c], d_idx[ok_a])]
+            usable = (sub > 0.0) & (sub < float(max_km))
+            skim_d[np.ix_(ok_c, ok_a)] = np.where(usable, sub, np.nan)
+            if st is not None and st.shape == sd.shape:
+                sub_t = st[np.ix_(o_idx[ok_c], d_idx[ok_a])]
+                skim_t[np.ix_(ok_c, ok_a)] = np.where(usable, sub_t, np.nan)
 
     # Haversine fallback from case-specific origin coords (broadcast).
     with np.errstate(invalid="ignore"):

{searchlibrium-0.0.261 → searchlibrium-0.0.263}/src/SearchLibrium/latent_class.py

@@ -1780,6 +1780,7 @@ class LatentClassMixedLogit(DiscreteChoiceModel):
         member_params_spec=None,
         class_params_spec=None,
         panels=None,
+        retain_all=False,
         ident_penalty_kwargs=None,
         **kwargs,
     ):
@@ -1797,8 +1798,17 @@ class LatentClassMixedLogit(DiscreteChoiceModel):
             Variables used in the class membership equation.
         member_params_spec : list or array, optional
             Per-class membership variable specification.
+        retain_all : bool
+            MEMORY: when False (default) only the best model is kept alive and
+            the returned list contains just that model; every intermediate
+            class-count model is released as the search advances. Each model
+            otherwise pins its X/y/draws design, JAX device copies and cached
+            JIT objectives, so retaining all of them across a GA search leaks
+            hundreds of MB per candidate. Set True to recover the previous
+            behaviour (all models returned and stored on
+            ``best_model.search_results``).
         """
-        fitted_models = []
+        fitted_models = [] if retain_all else None
         best_model = None
         _best_nsig = float('inf')
         prev_model = None
@@ -1833,7 +1843,8 @@ class LatentClassMixedLogit(DiscreteChoiceModel):
                 de_seed=de_seed,
             )
             model.get_loglik_null()
-            fitted_models.append(model)
+            if fitted_models is not None:
+                fitted_models.append(model)
 
             # Pareto-style significance prioritisation:
             #  1) prefer fewer insignificant variable groups,
@@ -1858,8 +1869,14 @@ class LatentClassMixedLogit(DiscreteChoiceModel):
                 _best_nsig = nsig
             prev_model = model
 
-        best_model.search_results = fitted_models
-        return best_model, fitted_models
+        if fitted_models is not None:
+            best_model.search_results = fitted_models
+            return best_model, fitted_models
+        # Default path: keep only the best model alive. Drop the warm-start
+        # chain's trailing reference so intermediate models can be collected.
+        best_model.search_results = None
+        prev_model = None
+        return best_model, [best_model]
 
     # ------------------------------------------------------------------
     # Internal helpers for numerical standard errors

{searchlibrium-0.0.261 → searchlibrium-0.0.263}/src/SearchLibrium/mixed_logit.py

@@ -54,6 +54,7 @@ in a panel or longitudinal dataset.
 ''' LIBRARIES                                                  '''
 ''' ---------------------------------------------------------- '''
 import itertools
+import os
 import numpy as np
 import scipy.stats as ss
 from scipy.optimize import minimize
@@ -273,6 +274,8 @@ class MixedLogit(DiscreteChoiceModel):
         self.rvdist = None          # Distributions for randvar
         self.correlated_vars, self.halton_opts = None, None
         self.batch_size, self.minimise_func = None, None
+        if not hasattr(self, "save_prob_full"):
+            self.save_prob_full = bool(os.environ.get("SEARCHLIBRIUM_SAVE_PROB_FULL", ""))
     # }
 
     ''' -------------------------------------------------------------------- '''
@@ -781,11 +784,17 @@ class MixedLogit(DiscreteChoiceModel):
             self.ind_pred_prob, self.pred_prob = _panel_weighted_shares(
                 self.choice_pred_prob, self.panel_info)
 
-            self.prob_full = p
+            if self.save_prob_full:
+                self.prob_full = p
         # }
         self.post_process(result, self.Xnames, self.N)
     # }
 
+    def release_memory(self):
+        """Drop prob_full if not explicitly saved. Safe between GA candidates."""
+        if not getattr(self, "save_prob_full", False):
+            self.prob_full = None
+
     ''' ---------------------------------------------------------- '''
     ''' Function. This code segments the betas array into different 
         parts and stores these parts in var_list under keys from 

{searchlibrium-0.0.261 → searchlibrium-0.0.263}/src/SearchLibrium/multinomial_logit.py

@@ -3,6 +3,7 @@ IMPLEMENTATION: MULTINOMIAL AND CONDITIONAL LOGIT MODEL
 """""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
 from sklearn.preprocessing import RobustScaler
 from addicty import  Dict
+import os
 
 """
 BACKGROUND - MULTINOMIAL LOGIT
@@ -1019,7 +1020,9 @@ class MultinomialLogit(DiscreteChoiceModel):
     # owning model instead of capturing ``self``, so an entry never pins a
     # dead model (and its X/y arrays) for the life of the process.
     # ------------------------------------------------------------------
-    _JIT_CACHE_MAXSIZE = 16
+    _JIT_CACHE_MAXSIZE = max(
+        0, int(os.environ.get("SEARCHLIBRIUM_MNL_JIT_CACHE_MAXSIZE", "4"))
+    )
     _jit_cache: "OrderedDict[tuple, tuple]" = OrderedDict()
 
     @classmethod

{searchlibrium-0.0.261 → searchlibrium-0.0.263}/src/SearchLibrium/ordered_logit.py

@@ -69,6 +69,8 @@ from scipy import stats
 from scipy.optimize import minimize
 
 import inspect
+import os
+from collections import OrderedDict
 
 minval = 1E-30
 
@@ -2481,7 +2483,15 @@ class SA_ORDLOG():
         self.mod = OrderedLogit(X=X, y=y, J=J, distr='logit', start=None, fit_intercept=True)
 
         # SA specific parameters
-        self.archive = {}  # Define a dictionary
+        # MEMORY: the archive memoises spec -> loglik so repeated specs skip
+        # the expensive refit. It used to be an unbounded dict, so a long run
+        # retained one entry per distinct predictor mask (2**K in the limit).
+        # It is now an LRU-bounded OrderedDict; evicting an entry only costs a
+        # future refit, never correctness. Set
+        # SEARCHLIBRIUM_ORDLOG_ARCHIVE_MAX=0 to restore unbounded behaviour.
+        self.archive = OrderedDict()
+        self.archive_maxsize = int(
+            os.environ.get("SEARCHLIBRIUM_ORDLOG_ARCHIVE_MAX", "50000"))
         self.tI = 100
         self.tF = 0.001
         self.t = self.tI
@@ -2544,6 +2554,12 @@ class SA_ORDLOG():
             print("Solution already seen")
             candidate_score = self.archive[key]
 
+        # LRU touch + bounded eviction (see archive_maxsize above).
+        self.archive.move_to_end(key)
+        _amax = getattr(self, "archive_maxsize", 0)
+        while _amax and len(self.archive) > _amax:
+            self.archive.popitem(last=False)
+
         # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
         if accept_change(self.current_score, candidate_score, self.t):

{searchlibrium-0.0.261 → searchlibrium-0.0.263}/src/SearchLibrium/test_lc_de.py

@@ -231,6 +231,7 @@ def run_class_search(
         max_classes=max_classes,
         criterion="bic",
         warm_start=True,
+        retain_all=True,
         de_init=de_init,
         de_popsize=de_popsize,
         de_maxiter=de_maxiter,

searchlibrium-0.0.263/src/SearchLibrium/version.txt

@@ -0,0 +1 @@
+0.0.263

{searchlibrium-0.0.261 → searchlibrium-0.0.263/src/SearchLibrium.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: SearchLibrium
-Version: 0.0.261
+Version: 0.0.263
 Summary: Automated specification search for discrete choice models (logit, mixed logit, nested logit, RRM and more) using metaheuristic optimisation
 Author: Zeke Ahern, Mario Fernando Taco Morales, Alexander Paz, Prithvi Beeramole, Robert Burdett
 Maintainer-email: Zeke Ahern <z.ahern@qut.edu.au>

searchlibrium-0.0.261/src/SearchLibrium/version.txt

@@ -1 +0,0 @@
-0.0.261