SearchLibrium 0.0.85__tar.gz → 0.0.86__tar.gz

{searchlibrium-0.0.85 → searchlibrium-0.0.86}/PKG-INFO

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

{searchlibrium-0.0.85 → searchlibrium-0.0.86}/pyproject.toml

@@ -59,7 +59,7 @@ Homepage = "https://github.com/zahern/HypothesisX"
 realpython = "SearchLibrium.__main__:main"
 
 [tool.bumpver]
-current_version = "0.0.85"
+current_version = "0.0.86"
 version_pattern = "MAJOR.MINOR.PATCH"
 commit_message = "[skip ci] Bump version {old_version} -> {new_version}"
 commit = true

{searchlibrium-0.0.85 → searchlibrium-0.0.86}/src/SearchLibrium/mdcev.py

@@ -9,11 +9,18 @@ conditions.
 
 The class is intended as a practical bridge between the current scalar budget
 models and a fuller MDCEV pipeline. It includes both a stable heuristic fit
-and a likelihood-based quasi-MLE refinement.
+and a JAX-accelerated quasi-MLE refinement with exact automatic differentiation.
+
+JAX is used for:
+- Bisection solver: JIT-compiled + vmapped over observations (fast predict/simulate)
+- fit_mle gradients: exact autodiff via jax.grad instead of scipy FD approximations
+
+A pure-numpy fallback is provided when JAX is not installed.
 """
 
 from __future__ import annotations
 
+import os as _os
 from dataclasses import dataclass
 from typing import Iterable, Optional
 
@@ -21,6 +28,17 @@ import numpy as np
 import pandas as pd
 from scipy.optimize import minimize
 
+# ── JAX optional import ──────────────────────────────────────────────────────
+_os.environ.setdefault("JAX_PLATFORMS", "cpu")
+try:
+    import jax
+    import jax.numpy as jnp
+    from jax import jit as _jit, vmap as _vmap
+    jax.config.update("jax_enable_x64", True)
+    _JAX_AVAILABLE = True
+except Exception:
+    _JAX_AVAILABLE = False
+
 
 def _as_2d_float(array_like) -> np.ndarray:
     arr = np.asarray(array_like, dtype=float)
@@ -42,12 +60,60 @@ class MDCEVFitResult:
     mean_budget: float
 
 
+# ── JAX bisection kernel (module-level, JIT-compiled once) ───────────────────
+if _JAX_AVAILABLE:
+    def _jax_bisect_single(budget, util_idx, alpha, gamma):
+        """Translated-utility KKT bisection for one observation.
+
+        Fully JAX-jittable and differentiable: uses ``jax.lax.scan`` for both
+        the initial bracket search and the bisection, so gradients flow through
+        the solver and ``fit_mle`` can use exact autodiff.
+        """
+        _tol = 1e-9
+        weights = jnp.exp(jnp.clip(util_idx, -40.0, 40.0))
+
+        def _alloc_sum(lam):
+            power = 1.0 / jnp.clip(1.0 - alpha, _tol, None)
+            raw = jnp.power(weights / jnp.maximum(lam, _tol), power) - gamma
+            return jnp.maximum(raw, 0.0).sum()
+
+        # Find upper bracket: keep doubling hi until alloc_sum(hi) <= budget.
+        # scan (max 60 steps) is differentiable; once condition is false hi stays fixed.
+        init_hi = jnp.maximum(jnp.max(weights), 1.0)
+
+        def _double_step(hi, _):
+            return jnp.where(_alloc_sum(hi) > budget, hi * 2.0, hi), None
+
+        hi, _ = jax.lax.scan(_double_step, init_hi, None, length=60)
+
+        # 80-step bisection
+        def _bisect_step(carry, _):
+            lo, hi = carry
+            mid = 0.5 * (lo + hi)
+            go_lo = _alloc_sum(mid) > budget
+            return (jnp.where(go_lo, mid, lo), jnp.where(go_lo, hi, mid)), None
+
+        (_, hi_f), _ = jax.lax.scan(_bisect_step, (_tol, hi), None, length=80)
+
+        power = 1.0 / jnp.clip(1.0 - alpha, _tol, None)
+        raw = jnp.power(weights / jnp.maximum(hi_f, _tol), power) - gamma
+        allocation = jnp.maximum(raw, 0.0)
+        total = allocation.sum()
+        allocation = jnp.where(total > _tol, allocation * (budget / total), allocation)
+        # zero-budget guard
+        return jnp.where(budget <= _tol, jnp.zeros_like(allocation), allocation)
+
+    # Batch version: vmap over (budget, util_idx); alpha and gamma are shared
+    _jax_bisect_batch = _jit(_vmap(_jax_bisect_single, in_axes=(0, 0, None, None)))
+
+
 class MDCEVModel:
     """Translated-utility MDCEV-style allocator.
 
     Parameters are learned from observed budget shares using stable moment-based
-    heuristics, then predictions are produced by solving the translated-utility
-    KKT system with a bisection search on the shadow price.
+    heuristics (``fit``), then optionally refined via quasi-MLE with JAX
+    automatic differentiation (``fit_mle``). Predictions use a JAX-jitted and
+    vmapped bisection solver when JAX is available, with a pure-numpy fallback.
     """
 
     def __init__(
@@ -152,87 +218,139 @@ class MDCEVModel:
     ):
         """Likelihood-based parameter refinement.
 
-        The objective is a Gaussian log-likelihood on log allocations around
-        translated-utility MDCEV deterministic predictions. This is a practical
-        quasi-MLE refinement that preserves the MDCEV budget constraint while
-        improving fit over pure moments.
+        Uses JAX automatic differentiation for exact analytic gradients when
+        JAX is available, replacing the slow scipy finite-difference
+        approximation. Falls back to scipy FD when JAX is not installed.
         """
         self.fit(allocations, labels=labels)
 
         y = _as_2d_float(allocations)
-        budgets = y.sum(axis=1)
+        budgets_np = y.sum(axis=1)
         n_alt = y.shape[1]
-
-        free_base_idx = [i for i in range(n_alt) if i != self.outside_good]
-
-        def _pack(base, alpha, gamma, sigma):
-            b = np.asarray(base, dtype=float)
-            a = np.asarray(alpha, dtype=float)
-            g = np.asarray(gamma, dtype=float)
-
+        alpha_floor = self.alpha_floor
+        alpha_cap = self.alpha_cap
+        gamma_floor = self.gamma_floor
+        tol = self.tol
+        og = self.outside_good
+        has_og = og is not None and 0 <= og < n_alt
+        free_base_idx = [i for i in range(n_alt) if not has_og or i != og]
+
+        def _pack_np(base, alpha, gamma, sigma):
             p = []
-            p.extend(b[free_base_idx].tolist())
-            p.extend(np.log(np.clip((a - self.alpha_floor) / np.clip(self.alpha_cap - a, self.tol, None), self.tol, None)).tolist())
-            p.extend(np.log(np.clip(g, self.gamma_floor, None)).tolist())
+            p.extend(np.asarray(base)[free_base_idx].tolist())
+            p.extend(np.log(np.clip(
+                (np.asarray(alpha) - alpha_floor) / np.clip(alpha_cap - np.asarray(alpha), tol, None),
+                tol, None)).tolist())
+            p.extend(np.log(np.clip(np.asarray(gamma), gamma_floor, None)).tolist())
             p.append(np.log(max(float(sigma), 1e-3)))
             return np.asarray(p, dtype=float)
 
-        def _unpack(theta):
-            theta = np.asarray(theta, dtype=float)
-            o = 0
-
-            base = self.baseline_utility_.copy()
-            for idx in free_base_idx:
-                base[idx] = theta[o]
-                o += 1
-            if self.outside_good is not None and 0 <= self.outside_good < n_alt:
-                base[self.outside_good] = 0.0
-
-            alpha_raw = theta[o:o + n_alt]
-            o += n_alt
-            alpha_sig = 1.0 / (1.0 + np.exp(-alpha_raw))
-            alpha = self.alpha_floor + (self.alpha_cap - self.alpha_floor) * alpha_sig
-
-            gamma_raw = theta[o:o + n_alt]
-            o += n_alt
-            gamma = np.maximum(np.exp(gamma_raw), self.gamma_floor)
-
-            sigma = max(np.exp(theta[o]), 1e-3)
-            return base, alpha, gamma, sigma
-
-        def _neg_loglike(theta):
-            base, alpha, gamma, sigma = _unpack(theta)
-
-            old_b, old_a, old_g = self.baseline_utility_, self.alpha_, self.gamma_
-            self.baseline_utility_, self.alpha_, self.gamma_ = base, alpha, gamma
-            try:
-                mu = np.zeros_like(y)
-                for i, b in enumerate(budgets):
-                    mu[i] = self._solve_budget(float(b), base)
-            finally:
-                self.baseline_utility_, self.alpha_, self.gamma_ = old_b, old_a, old_g
-
-            log_y = np.log(np.clip(y, self.tol, None))
-            log_mu = np.log(np.clip(mu, self.tol, None))
-            resid = log_y - log_mu
-            ll = -0.5 * resid.size * np.log(2.0 * np.pi * sigma * sigma)
-            ll -= 0.5 * np.sum((resid / sigma) ** 2)
-            ll -= l2_penalty * np.sum(theta * theta)
-            return -float(ll)
-
-        theta0 = _pack(self.baseline_utility_, self.alpha_, self.gamma_, sigma=0.5)
-        res = minimize(
-            _neg_loglike,
-            theta0,
-            method="L-BFGS-B",
-            options={"maxiter": int(maxiter), "ftol": 1e-9},
-        )
+        theta0 = _pack_np(self.baseline_utility_, self.alpha_, self.gamma_, sigma=0.5)
+
+        if _JAX_AVAILABLE:
+            # ── JAX path: exact gradients via autodiff ────────────────────
+            y_jax = jnp.array(y)
+            B_jax = jnp.array(budgets_np)
+            free_idx_jax = jnp.array(free_base_idx)
+
+            def _unpack_jax(theta):
+                o = 0
+                n_free = len(free_base_idx)
+                free_base = theta[o:o + n_free]
+                o += n_free
+                base = jnp.zeros(n_alt).at[free_idx_jax].set(free_base)
+
+                alpha_raw = theta[o:o + n_alt]
+                o += n_alt
+                alpha = alpha_floor + (alpha_cap - alpha_floor) * jax.nn.sigmoid(alpha_raw)
+
+                gamma_raw = theta[o:o + n_alt]
+                o += n_alt
+                gamma = jnp.maximum(jnp.exp(gamma_raw), gamma_floor)
+
+                sigma = jnp.maximum(jnp.exp(theta[o]), 1e-3)
+                return base, alpha, gamma, sigma
+
+            @_jit
+            def _neg_loglike_jax(theta):
+                base, alpha, gamma, sigma = _unpack_jax(theta)
+                util_matrix = jnp.broadcast_to(base[None, :], (len(B_jax), n_alt))
+                preds = _jax_bisect_batch(B_jax, util_matrix, alpha, gamma)
+                log_y = jnp.log(jnp.clip(y_jax, tol, None))
+                log_p = jnp.log(jnp.clip(preds, tol, None))
+                resid = log_y - log_p
+                ll = -0.5 * resid.size * jnp.log(2.0 * jnp.pi * sigma * sigma)
+                ll -= 0.5 * jnp.sum((resid / sigma) ** 2)
+                ll -= l2_penalty * jnp.sum(theta * theta)
+                return -ll
+
+            val_and_grad = _jit(jax.value_and_grad(_neg_loglike_jax))
+
+            def _scipy_obj(theta_np):
+                val, grad = val_and_grad(jnp.array(theta_np, dtype=jnp.float64))
+                return float(val), np.asarray(grad, dtype=np.float64)
+
+            res = minimize(
+                _scipy_obj,
+                theta0,
+                method="L-BFGS-B",
+                jac=True,
+                options={"maxiter": int(maxiter), "ftol": 1e-9},
+            )
+            base_j, alpha_j, gamma_j, sigma_j = _unpack_jax(jnp.array(res.x))
+            self.baseline_utility_ = np.asarray(base_j)
+            self.alpha_ = np.asarray(alpha_j)
+            self.gamma_ = np.asarray(gamma_j)
+            self.noise_sigma_ = float(sigma_j)
+
+        else:
+            # ── numpy / scipy FD fallback (original behaviour) ────────────
+            def _unpack_np(theta):
+                theta = np.asarray(theta, dtype=float)
+                o = 0
+                base = self.baseline_utility_.copy()
+                for idx in free_base_idx:
+                    base[idx] = theta[o]
+                    o += 1
+                if has_og:
+                    base[og] = 0.0
+                alpha_raw = theta[o:o + n_alt]; o += n_alt
+                alpha = alpha_floor + (alpha_cap - alpha_floor) / (1.0 + np.exp(-alpha_raw))
+                gamma_raw = theta[o:o + n_alt]; o += n_alt
+                gamma = np.maximum(np.exp(gamma_raw), gamma_floor)
+                sigma = max(np.exp(theta[o]), 1e-3)
+                return base, alpha, gamma, sigma
+
+            def _neg_loglike_np(theta):
+                base, alpha, gamma, sigma = _unpack_np(theta)
+                old_b, old_a, old_g = self.baseline_utility_, self.alpha_, self.gamma_
+                self.baseline_utility_, self.alpha_, self.gamma_ = base, alpha, gamma
+                try:
+                    mu = np.zeros_like(y)
+                    for i, b in enumerate(budgets_np):
+                        mu[i] = self._solve_budget(float(b), base)
+                finally:
+                    self.baseline_utility_, self.alpha_, self.gamma_ = old_b, old_a, old_g
+                log_y = np.log(np.clip(y, tol, None))
+                log_mu = np.log(np.clip(mu, tol, None))
+                resid = log_y - log_mu
+                ll = -0.5 * resid.size * np.log(2.0 * np.pi * sigma * sigma)
+                ll -= 0.5 * np.sum((resid / sigma) ** 2)
+                ll -= l2_penalty * np.sum(theta * theta)
+                return -float(ll)
+
+            res = minimize(
+                _neg_loglike_np,
+                theta0,
+                method="L-BFGS-B",
+                options={"maxiter": int(maxiter), "ftol": 1e-9},
+            )
+            base, alpha, gamma, sigma = _unpack_np(res.x)
+            self.baseline_utility_ = base
+            self.alpha_ = alpha
+            self.gamma_ = gamma
+            self.noise_sigma_ = float(sigma)
 
-        base, alpha, gamma, sigma = _unpack(res.x)
-        self.baseline_utility_ = base
-        self.alpha_ = alpha
-        self.gamma_ = gamma
-        self.noise_sigma_ = float(sigma)
         self.mle_success_ = bool(res.success)
         self.mle_message_ = str(res.message)
         return self
@@ -267,6 +385,16 @@ class MDCEVModel:
         budgets_arr = np.asarray(budgets, dtype=float).reshape(-1)
         shifts = self._prepare_utility_shift(utility_shift, len(budgets_arr))
 
+        if _JAX_AVAILABLE:
+            util_matrix = jnp.array(self.baseline_utility_[None, :] + shifts)
+            preds = _jax_bisect_batch(
+                jnp.array(budgets_arr),
+                util_matrix,
+                jnp.array(self.alpha_),
+                jnp.array(self.gamma_),
+            )
+            return np.asarray(preds)
+
         predictions = np.zeros((len(budgets_arr), len(self.baseline_utility_)), dtype=float)
         for row_idx, budget in enumerate(budgets_arr):
             predictions[row_idx] = self._solve_budget(budget, self.baseline_utility_ + shifts[row_idx])
@@ -280,6 +408,17 @@ class MDCEVModel:
         rng = np.random.default_rng(random_state)
 
         sims = np.zeros((n_draws, len(budgets_arr), len(self.baseline_utility_)), dtype=float)
+
+        if _JAX_AVAILABLE:
+            B_jax = jnp.array(budgets_arr)
+            alpha_jax = jnp.array(self.alpha_)
+            gamma_jax = jnp.array(self.gamma_)
+            for draw_idx in range(n_draws):
+                shocks = rng.gumbel(loc=0.0, scale=1.0, size=shifts.shape)
+                util_matrix = jnp.array(self.baseline_utility_[None, :] + shifts + shocks)
+                sims[draw_idx] = np.asarray(_jax_bisect_batch(B_jax, util_matrix, alpha_jax, gamma_jax))
+            return sims
+
         for draw_idx in range(n_draws):
             shocks = rng.gumbel(loc=0.0, scale=1.0, size=shifts.shape)
             for row_idx, budget in enumerate(budgets_arr):
@@ -304,6 +443,7 @@ class MDCEVModel:
         return shift_arr
 
     def _solve_budget(self, budget: float, utility_index: np.ndarray) -> np.ndarray:
+        """Pure-numpy bisection fallback — used when JAX is not available."""
         if budget <= self.tol:
             return np.zeros(len(self.baseline_utility_), dtype=float)
 

{searchlibrium-0.0.85 → searchlibrium-0.0.86}/src/SearchLibrium/siman.py

@@ -616,7 +616,7 @@ class SA(Search):
     def copy_solution(self, sol):
     # {
         logging.info('normal copy')
-        copy_sol = sol.copy()
+        copy_sol = copy.deepcopy(sol)  # Deep copy to prevent shared mutable state (lists, dicts inside solution)
         return copy_sol
     # }
 
@@ -998,15 +998,12 @@ class SA(Search):
             print("Initialized best_sol with the first solution")
             return
 
-        self.no_impr = 0
         if self.nb_crit == 1:
         # {
             if is_better(sol.obj(0), self.best_sol.obj(0), self.param.sign_crit(0)):
                 self.best_sol = self.copy_solution(sol)
-
-                print('new best')
-
                 self.no_impr = 0
+                print('new best')
         # }
         else:
             self.archive = self.update_archive(sol)
@@ -1087,7 +1084,7 @@ class SA(Search):
         # {
             print("ARCHIVE STATIC. RESTORE A NON DOMINATED SOLUTION.")
             choice = np.random.randint(len(self.archive))
-            self.current_sol = self.archive[choice]
+            self.current_sol = self.copy_solution(self.archive[choice])  # Deep copy to avoid aliasing archive entry
             self.no_impr = 0
         # }
     # }

searchlibrium-0.0.86/src/SearchLibrium/version.txt

@@ -0,0 +1 @@
+0.0.86

{searchlibrium-0.0.85 → searchlibrium-0.0.86}/src/SearchLibrium.egg-info/PKG-INFO

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

searchlibrium-0.0.85/src/SearchLibrium/version.txt

@@ -1 +0,0 @@
-0.0.85