diff --git a/README.md b/README.md index f1e6258..59f1a62 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ - **Large-scale urban models**: 100K+ choosers, 1K+ alternatives - **Sampling-based estimation**: Most alternatives are irrelevant; only a sampled subset is evaluated per chooser -- **Spatial correlation**: Nearby alternatives share unobserved attributes (via `graph=` on any model) +- **Spatial dependence**: Nearby alternatives can affect each other or share unobserved attributes (via `graph=` on any model) - **Heterogeneous preferences**: Mixed logit for random taste variation - **Nested structure**: Nested logit for hierarchical choice (e.g., county → tract → block) - **JAX-native computation**: JIT-compiled kernels, GPU acceleration, automatic differentiation @@ -13,8 +13,18 @@ The package is **not** a general-purpose ML library. It is specifically for stru ## Features -LocPick can automate the creation of choice tables for estimation or simulation, using census choice sets, uniform or weighted random sampling of alternatives, generated interaction terms, and cartesian merges. A unique feature is the implementation of [*spatial* choice models](https://linkinghub.elsevier.com/retrieve/pii/S0191261503000055), which assume that nearby alternatives are more similar (closer substitutes). +LocPick can automate the creation of choice tables for estimation or simulation, using census choice sets, uniform or weighted random sampling of alternatives, generated interaction terms, and cartesian merges. A unique feature is the implementation of *spatial* choice models, which take one of two forms. + +- The [Bhat et al](https://linkinghub.elsevier.com/retrieve/pii/S0191261503000055) approach is similar to a spatial error model, assuming that nearby alternatives are more similar (closer substitutes). +- The SAR style approach assumes that the structural utility of each alternative $V$ has a simultaneous autoregressive structure, and is estimated with either [PML](http://dx.doi.org/10.1016/j.regsciurbeco.2009.09.004) or [GMM](https://www.sciencedirect.com/science/article/pii/S0166046217300625) It also provides tools for Monte Carlo simulation of choices given probability distributions from fitted models, with fast algorithms for independent or capacity-constrained choices. -LocPick includes classic and spatially-correlated Multinomial Logit, Nested Logit, Mixed, and Mixed/Nested Logit estimators; its internal data pipeline is designed around pandas inputs, xarray-backed alignment, and NumPy/JAX-ready arrays. +LocPick includes estimators for classic, spatially-correlated, and simultaneous autoregressive: + +- Multinomial Logit +- Nested Logit +- Mixed Logit and +- Mixed/Nested + +models, and an internal data pipeline designed around pandas inputs, xarray-backed alignment, and NumPy/JAX-ready arrays. diff --git a/docs/source/_static/references.bib b/docs/source/_static/references.bib index 663ddde..69f18f7 100644 --- a/docs/source/_static/references.bib +++ b/docs/source/_static/references.bib @@ -25,7 +25,6 @@ @article{abbe2007NormalizationCorrelation keywords = {Behavior models,Correlations,Generalized extreme-value,Logit models,Model estimation,Route choice}, } - @article{al-haideri2026CyclistsCrossing, title = {Cyclists' Crossing Behaviour at Roundabouts: {{A Generalized Spatially Correlated Nested Logit}} Model}, @@ -70,7 +69,6 @@ @article{al-haideri2026CyclistsCrossing behavioural mechanisms underlying cyclists' crossing decisions.}, langid = {english}, } - @article{anselin1988spatial, title = {Spatial Econometrics: Methods and Models}, author = {Anselin, Luc}, @@ -246,6 +244,69 @@ @article{dugundji2013StructureEmergence networks,Spatial interaction,Transportation demand}, } +@article{krisztin2021BayesianSpatial, + title = {A {{Bayesian}} Spatial Autoregressive Logit Model with an Empirical + Application to {{European}} Regional {{FDI}} Flows}, + author = {Krisztin, Tamás and Piribauer, Philipp}, + date = {2021-07-01}, + journaltitle = {Empirical Economics}, + shortjournal = {Empir Econ}, + volume = {61}, + number = {1}, + pages = {231--257}, + issn = {1435-8921}, + doi = {10.1007/s00181-020-01856-w}, + url = {https://doi.org/10.1007/s00181-020-01856-w}, + urldate = {2026-06-18}, + abstract = {In this paper, we propose a Bayesian estimation approach for a + spatial autoregressive logit specification. Our approach relies on + recent advances in Bayesian computing, making use of Pólya–Gamma + sampling for Bayesian Markov-chain Monte Carlo algorithms. The + proposed specification assumes that the involved log-odds of the + model follow a spatial autoregressive process. Pólya–Gamma sampling + involves a computationally efficient treatment of the spatial + autoregressive logit model, allowing for extensions to the existing + baseline specification in an elegant and straightforward way. In a + Monte Carlo study we demonstrate that our proposed approach + markedly outperforms alternative specifications in terms of + parameter precision. The paper moreover illustrates the performance + of the proposed spatial autoregressive logit specification using + pan-European regional data on foreign direct investments. Our + empirical results highlight the importance of accounting for + spatial dependence when modelling European regional FDI flows.}, + langid = {english}, + keywords = {Bayesian MCMC estimation,C11,C21,C25,European regions,F23,FDI + flows,R11,R30,Spatial autoregressive logit}, +} + +@article{krisztin2022SpatialMultinomial, + title = {A Spatial Multinomial Logit Model for Analysing Urban Expansion}, + author = {Krisztin, Tamás and Piribauer, Philipp and Wögerer, Michael}, + date = {2022-04-03}, + journaltitle = {Spatial Economic Analysis}, + volume = {17}, + number = {2}, + pages = {223--244}, + publisher = {Routledge}, + issn = {1742-1772}, + doi = {10.1080/17421772.2021.1933579}, + url = {https://doi.org/10.1080/17421772.2021.1933579}, + urldate = {2026-06-18}, + abstract = {The paper proposes a Bayesian multinomial logit model to analyse + spatial patterns of urban expansion. The specification assumes that + the log-odds of each class follow a spatial autoregressive process. + Using recent advances in Bayesian computing, our model allows for a + computationally efficient treatment of the spatial multinomial + logit model. This allows us to assess spillovers between regions + and across land-use classes. In a series of Monte Carlo studies, we + benchmark our model against other competing specifications. The + paper also showcases the performance of the proposed specification + using European regional data. Our results indicate that spatial + dependence plays a key role in the land-sealing process of cropland + and grassland. Moreover, we uncover land-sealing spillovers across + multiple classes of arable land.}, +} + @article{mcfadden1978goodness, title = {Goodness-of-Fit for the Multinomial Logit Model}, author = {McFadden, Daniel}, @@ -304,6 +365,37 @@ @article{perez-lopez2022SpatiallyCorrelated langid = {english}, } +@article{smirnov2010ModelingSpatial, + title = {Modeling Spatial Discrete Choice}, + author = {Smirnov, Oleg A.}, + date = {2010-09}, + journaltitle = {Regional Science and Urban Economics}, + volume = {40}, + number = {5}, + pages = {292--298}, + publisher = {Elsevier B.V.}, + issn = {01660462}, + doi = {10.1016/j.regsciurbeco.2009.09.004}, + url = {http://dx.doi.org/10.1016/j.regsciurbeco.2009.09.004}, + abstract = {The paper presents a basic spatial discrete choice modeling + framework obtained by applying random utility theory to discrete + choices made by heterogeneous spatially dependent individuals. The + newly developed framework has two main advantages over existing + approaches. First, individual decision-makers are no longer assumed + to be independent and non-interacting but spatially interdependent + in their preferences facilitating the development of applied + discrete choice models using a wide range of spatial data. Second, + pseudo maximum likelihood estimator is developed for this model + that is consistent and computationally feasible for large datasets. + The performance of the pseudo maximum likelihood estimator for the + spatial discrete choice model is illustrated using simulated data. + © 2009 Elsevier B.V.}, + isbn = {0166-0462}, + keywords = {Discrete choice,Maximum likelihood,Spatial interdependence,Spatial + random utility,urban-modeling}, + annotation = {44 citations (Crossref) [2022-08-11]}, +} + @book{train2009discrete, title = {Discrete Choice Methods with Simulation}, author = {Train, Kenneth E.}, diff --git a/docs/source/api.rst b/docs/source/api.rst index 40c7f51..bff8edb 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -43,19 +43,18 @@ Model Specification Models ------ -.. currentmodule:: locpick.models.mnl +.. currentmodule:: locpick.models.choice_model .. autosummary:: :toctree: generated/ - MNL :no-index: + ChoiceModel :no-index: .. currentmodule:: locpick.models.nested .. autosummary:: :toctree: generated/ - NestedMNL :no-index: NestSpec :no-index: NestingTree :no-index: @@ -64,16 +63,8 @@ Models .. autosummary:: :toctree: generated/ - MixedMNL :no-index: ParamDistribution :no-index: -.. currentmodule:: locpick.models.mixed_nested - -.. autosummary:: - :toctree: generated/ - - MixedNestedMNL :no-index: - .. currentmodule:: locpick.models.scl .. autosummary:: diff --git a/docs/source/index.md b/docs/source/index.md index 73c880f..0485cff 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -21,6 +21,7 @@ Mixed Logit Spatial Mixed Logit Spatial Nested Logit Spatial Mixed-Nested Logit +SAR-MNL Demo Simulated Location Choice Demo Spatial Models Demo Sampling Correction diff --git a/docs/source/user-guide/choicetable.md b/docs/source/user-guide/choicetable.md index 2f313a4..3a1f2ad 100644 --- a/docs/source/user-guide/choicetable.md +++ b/docs/source/user-guide/choicetable.md @@ -16,7 +16,7 @@ The `ChoiceTable` class is the primary data container for location choice modeli ## Quick Start ```python -from locpick import ChoiceTable, MNL +from locpick import ChoiceTable, ChoiceModel # Create a ChoiceTable from chooser and alternative data ct = ChoiceTable.from_tables( @@ -26,7 +26,7 @@ ct = ChoiceTable.from_tables( ) # Estimate an MNL model -model = MNL(ct, formula="cost + time - 1") +model = ChoiceModel(ct, formula="cost + time - 1") result = model.fit() print(result.summary()) ``` diff --git a/docs/source/user-guide/inference.md b/docs/source/user-guide/inference.md index cfd9c8d..3a5debb 100644 --- a/docs/source/user-guide/inference.md +++ b/docs/source/user-guide/inference.md @@ -40,8 +40,8 @@ Compare two nested models: ```python from locpick import lr_test -restricted = MNL(ct, formula="cost + time").fit() -unrestricted = NestedMNL(ct, formula="cost + time", nests=tree).fit() +restricted = ChoiceModel(ct, formula="cost + time").fit() +unrestricted = ChoiceModel(ct, formula="cost + time", nests=tree).fit() result = lr_test(restricted, unrestricted) print(result.summary()) @@ -75,8 +75,8 @@ restrictive nested or mixed alternative: ```python from locpick import hausman_test -mnl_fit = MNL(ct, formula="cost + time").fit() -nested_fit = NestedMNL(ct, formula="cost + time", nests=tree).fit() +mnl_fit = ChoiceModel(ct, formula="cost + time").fit() +nested_fit = ChoiceModel(ct, formula="cost + time", nests=tree).fit() # H0: MNL is consistent (IIA holds). Compare on the common parameters. result = hausman_test(efficient=mnl_fit, consistent=nested_fit) diff --git a/docs/source/user-guide/livelike_locpick_household_tract_demo.ipynb b/docs/source/user-guide/livelike_locpick_household_tract_demo.ipynb index 5a9b4c9..e0300a3 100644 --- a/docs/source/user-guide/livelike_locpick_household_tract_demo.ipynb +++ b/docs/source/user-guide/livelike_locpick_household_tract_demo.ipynb @@ -61,7 +61,7 @@ "from livelike.config import up_expanded_attributes_household\n", "from pymedm import PMEDM\n", "\n", - "from locpick import MNL, ChoiceTable" + "from locpick import ChoiceModel, ChoiceTable" ] }, { @@ -435,7 +435,7 @@ "source": [ "formula_1 = \"median_contract_rent + median_home_value + median_household_income + p_poverty_rate\"\n", "\n", - "model_1 = MNL(ct, formula=formula_1)\n", + "model_1 = ChoiceModel(ct, formula=formula_1)\n", "result_1 = model_1.fit()" ] }, @@ -474,7 +474,7 @@ "_interaction_terms = [k for k in interactions.keys() if k in ct._ds.data_vars]\n", "formula_2 = \" + \".join(_base_terms + _interaction_terms)\n", "\n", - "model_2 = MNL(ct, formula=formula_2)\n", + "model_2 = ChoiceModel(ct, formula=formula_2)\n", "result_2 = model_2.fit()\n", "print(result_2.summary())" ] diff --git a/docs/source/user-guide/mixed.md b/docs/source/user-guide/mixed.md index b2eda9c..72a7338 100644 --- a/docs/source/user-guide/mixed.md +++ b/docs/source/user-guide/mixed.md @@ -6,12 +6,12 @@ This user guide is a placeholder. Full content will be added in a future release ## Overview -The `MixedMNL` class estimates mixed logit (random coefficients) models, which generalize MNL by allowing coefficients to follow random distributions. +The `ChoiceModel` class estimates mixed logit (random coefficients) models when the `random_params` parameter is provided. Mixed logit generalizes MNL by allowing coefficients to follow random distributions. ## Quick Start ```python -from locpick import ChoiceTable, MixedMNL +from locpick import ChoiceTable, ChoiceModel from locpick.models.mixed import ParamDistribution ct = ChoiceTable.from_tables(choosers, alternatives, chosen_alternatives=choices) @@ -20,7 +20,7 @@ random_params = { "cost": ParamDistribution(distribution="normal", param="cost"), } -model = MixedMNL(ct, formula="cost + time - 1", random_params=random_params, n_draws=100) +model = ChoiceModel(ct, formula="cost + time - 1", random_params=random_params, n_draws=100) result = model.fit() print(result.summary()) ``` \ No newline at end of file diff --git a/docs/source/user-guide/mnl.md b/docs/source/user-guide/mnl.md index b6eaa66..a84f3e3 100644 --- a/docs/source/user-guide/mnl.md +++ b/docs/source/user-guide/mnl.md @@ -6,7 +6,7 @@ This user guide is a placeholder. Full content will be added in a future release ## Overview -The `MNL` class estimates multinomial logit (MNL) models for location choice. It supports: +The `ChoiceModel` class estimates multinomial logit (MNL) models for location choice. It supports: - Formula/scoped-term model specification - JAX-accelerated log-likelihood and gradient computation @@ -16,10 +16,10 @@ The `MNL` class estimates multinomial logit (MNL) models for location choice. It ## Quick Start ```python -from locpick import ChoiceTable, MNL +from locpick import ChoiceTable, ChoiceModel ct = ChoiceTable.from_tables(choosers, alternatives, chosen_alternatives=choices) -model = MNL(ct, formula="cost + time - 1") +model = ChoiceModel(ct, formula="cost + time - 1") result = model.fit() print(result.summary()) ``` @@ -29,6 +29,6 @@ print(result.summary()) ```python # JAX is the default and only production backend # NumPy kernels exist for reference/testing only -model = MNL(ct, formula="cost + time - 1") +model = ChoiceModel(ct, formula="cost + time - 1") result = model.fit() ``` \ No newline at end of file diff --git a/docs/source/user-guide/modelspec.md b/docs/source/user-guide/modelspec.md index bbf76d0..a85dfb5 100644 --- a/docs/source/user-guide/modelspec.md +++ b/docs/source/user-guide/modelspec.md @@ -17,9 +17,9 @@ locpick uses formula/scoped-term specification for model structure: Typically models are specified using Wilkinson formulas ```python -from locpick import ChoiceTable, MNL +from locpick import ChoiceTable, ChoiceModel -model = MNL(ct, formula="cost + time - 1") +model = ChoiceModel(ct, formula="cost + time - 1") ``` The `ChoiceTable` handles data construction intelligently. In this example `cost` is fixed @@ -42,10 +42,10 @@ ct = ct.add_pairwise_variable("time", time_series) ## Scoped Terms ```python -from locpick import ModelSpec, MNL +from locpick import ModelSpec, ChoiceModel spec = ModelSpec(formula="cost + time - 1").alternative_specific("time", reference="walk") -model = MNL(ct, spec=spec) +model = ChoiceModel(ct, spec=spec) ``` ## Generated Interaction Variables diff --git a/docs/source/user-guide/nested.md b/docs/source/user-guide/nested.md index bffcda4..c9ad777 100644 --- a/docs/source/user-guide/nested.md +++ b/docs/source/user-guide/nested.md @@ -6,12 +6,12 @@ This user guide is a placeholder. Full content will be added in a future release ## Overview -The `NestedMNL` class estimates nested logit models, which generalize MNL by grouping alternatives into nests with correlated error terms. +The `ChoiceModel` class estimates nested logit models when the `nests` parameter is provided. Nested logit generalizes MNL by grouping alternatives into nests with correlated error terms. ## Quick Start ```python -from locpick import ChoiceTable, NestedMNL +from locpick import ChoiceTable, ChoiceModel from locpick.models.nested import NestSpec, NestingTree ct = ChoiceTable.from_tables(choosers, alternatives, chosen_alternatives=choices) @@ -23,7 +23,7 @@ tree = NestingTree( ] ) -model = NestedMNL(ct, formula="cost + time - 1", nests=tree) +model = ChoiceModel(ct, formula="cost + time - 1", nests=tree) result = model.fit() print(result.summary()) ``` \ No newline at end of file diff --git a/docs/source/user-guide/prediction.md b/docs/source/user-guide/prediction.md index 4cc0d19..6176023 100644 --- a/docs/source/user-guide/prediction.md +++ b/docs/source/user-guide/prediction.md @@ -13,10 +13,10 @@ Prediction and simulation are **model methods**, not `FitResult` methods. After - `model.predict(result, data)` — Predicted choices ```python -from locpick import ChoiceTable, MNL +from locpick import ChoiceTable, ChoiceModel ct = ChoiceTable.from_tables(choosers, alternatives, chosen_alternatives=choices) -model = MNL(ct, formula="cost + time - 1") +model = ChoiceModel(ct, formula="cost + time - 1") result = model.fit() # Choice probabilities (n_obs × n_alts ndarray) diff --git a/docs/source/user-guide/sar_mnl.md b/docs/source/user-guide/sar_mnl.md new file mode 100644 index 0000000..3423706 --- /dev/null +++ b/docs/source/user-guide/sar_mnl.md @@ -0,0 +1,186 @@ +# Spatial Autoregressive Multinomial Logit (SAR-MNL) + +```{note} +This user guide covers SAR-MNL estimation via ``ChoiceModel`` with +``graph=`` and ``lag=True``, which implements a spatial autoregressive +lag in the utility of alternatives (spatial locations) using the pseudo +maximum likelihood (PML) estimator from Smirnov (2010). +``` + +## Overview + +The `ChoiceModel` class with `lag=True` models spatial spillover in the +**systematic utility** of alternatives via a spatial autoregressive (SAR) +lag: + +$$V_j = \rho \sum_k w_{jk} V_k + Z_j \beta + X_{ij} \gamma$$ + +where $W$ is a $J \times J$ spatial weights matrix connecting alternatives +(spatial locations). The reduced-form utility is: + +$$V^* = (I - \rho W)^{-1} (Z\beta + X\gamma)$$ + +normalised by $D = \text{diag}((I - \rho W)^{-1})$ for consistency +(Smirnov 2010). Choice probabilities follow standard MNL softmax over +the spatially-filtered, variance-normalised utilities. + +### SAR vs SCL + +The `graph=` parameter can be used with two spatial mechanisms: + +| `lag=` | Mechanism | ρ range | Transform | +|--------|-----------|---------|-----------| +| `False` (default) | SCL (GEV paired nests) | (0, 1] | Sigmoid | +| `True` | SAR (spatial autoregressive lag) | (-1, 1) | Tanh | + +SAR models spatial spillover in the **systematic utility** via the +spatial multiplier $(I - \rho W)^{-1}$. SCL models spatial correlation +via GEV paired nests between adjacent alternatives. They are different +mechanisms for the same spatial weights matrix. + +### Key features + +- **PML estimator** (Smirnov 2010): consistent, no log-determinant needed +- **JAX autodiff**: gradients through the spatial solve and variance normalisation +- **Dense and conjugate-gradient solve paths**: auto-selected by alternative count +- **Linearized GMM fallback** (Carrión-Flores et al. 2018): for very large J +- **libpysal Graph support**: canonical W type, matching bayespecon +- **Marginal effects**: direct, indirect, and total (LeSage & Pace 2009) +- **Composable with Mixed and Nested logit**: SAR + Mixed, SAR + Nested, SAR + Mixed + Nested + +## Quick Start + +```python +from locpick import ChoiceTable, ChoiceModel +from libpysal.graph import Graph + +# Build spatial weights matrix connecting alternatives +W = Graph.build_knn(gdf, k=7).transform("r") + +# Create choice data +ct = ChoiceTable.from_tables(choosers, alternatives, chosen_alternatives=choices) + +# Estimate SAR-MNL +model = ChoiceModel(ct, formula="cost + time - 1", graph=W, lag=True) +result = model.fit() +print(result.summary()) +``` + +## Estimation Methods + +### PML (default) + +The pseudo maximum likelihood estimator (Smirnov 2010) is the default. +It uses JAX autodiff through the spatial solve $(I - \rho W)^{-1}$ and +variance normalisation $\text{diag}((I - \rho W)^{-1})$. + +- **Dense solve** (n_alts ≤ 2000): LU factorisation, exact gradients +- **Conjugate gradient** (n_alts > 2000): iterative solve, power-series + diagonal approximation + +```python +# Auto-select (default) +model = ChoiceModel(ct, formula="cost + time - 1", graph=W, lag=True) + +# Force dense solve +model = ChoiceModel(ct, formula="cost + time - 1", graph=W, lag=True, estimator="pml") + +# Force conjugate gradient +model = ChoiceModel(ct, formula="cost + time - 1", graph=W, lag=True, estimator="pml_cg") +``` + +### Linearized GMM + +For very large alternative sets where even CG is too slow, the +linearized GMM estimator (Carrión-Flores et al. 2018) avoids matrix +inversion entirely via a two-step procedure: + +1. **Step 1**: Standard MNL estimation (ignoring spatial dependence) +2. **Step 2**: Two-stage least squares (TSLS) with instruments $[X, WX]$ + +```python +model = ChoiceModel(ct, formula="cost + time - 1", graph=W, lag=True, estimator="linearized_gmm") +result = model.fit() +``` + +## Spatial Weights Matrix + +`ChoiceModel` accepts `libpysal.graph.Graph` as the canonical W type +(matching the bayespecon package). `scipy.sparse` matrices and dense +NumPy arrays are also accepted for convenience. + +```python +from libpysal.graph import Graph + +# k-nearest-neighbor graph +W = Graph.build_knn(gdf, k=7).transform("r") + +# Contiguity graph +W = Graph.build_contiguity(gdf, rook=False).transform("r") + +# Distance band +W = Graph.build_distance_band(gdf, threshold=1000).transform("r") +``` + +## Composing with Mixed and Nested Logit + +SAR composes naturally with mixed logit (random coefficients) and nested +logit (hierarchical choice). The spatial filter is applied to utilities +*before* the GEV/softmax/mixed-logit logic, so the mechanisms are +independent. + +### SAR + Nested + +```python +from locpick import ChoiceModel, NestingTree + +nests = NestingTree(...) +model = ChoiceModel(ct, formula="cost + time - 1", graph=W, lag=True, nests=nests) +result = model.fit() +``` + +### SAR + Mixed + +```python +from locpick import ChoiceModel, ParamDistribution + +random_params = {"time": ParamDistribution("normal", "time")} +model = ChoiceModel( + ct, formula="cost + time - 1", graph=W, lag=True, + random_params=random_params, n_draws=200, +) +result = model.fit() +``` + +### SAR + Mixed + Nested + +```python +model = ChoiceModel( + ct, formula="cost + time - 1", graph=W, lag=True, + nests=nests, random_params=random_params, n_draws=200, +) +result = model.fit() +``` + +## Marginal Effects + +In the SAR-MNL model, a change in an attribute of alternative $j$ +affects not only $j$'s utility but also neighbouring alternatives +through the spatial multiplier $(I - \rho W)^{-1}$. + +```python +result = model.fit() + +# Compute marginal effects for a variable +me = model.marginal_effects(variable="cost") +print(me["direct"]) # impact on own alternative +print(me["indirect"]) # spillover to neighbouring alternatives +print(me["total"]) # direct + indirect +``` + +## References + +- Smirnov, O.A. (2010). "Modeling Spatial Discrete Choice." *Regional Science and Urban Economics*, 40, 292–298. +- Carrión-Flores, C.E., Flores-Lagunes, A., & Guci, L. (2018). "An Estimator for Discrete-Choice Models with Spatial Lag Dependence Using Large Samples." *RSUE*, 69, 77–93. +- LeSage, J.P. & Pace, R.K. (2009). *Introduction to Spatial Econometrics*. +- Krisztin, T., Piribauer, P., & Wögerer, M. (2022). "A Spatial Multinomial Logit Model for Analysing Urban Expansion." *Spatial Economic Analysis*, 17(2), 223–244. \ No newline at end of file diff --git a/docs/source/user-guide/sar_mnl_demo.ipynb b/docs/source/user-guide/sar_mnl_demo.ipynb new file mode 100644 index 0000000..81bcc19 --- /dev/null +++ b/docs/source/user-guide/sar_mnl_demo.ipynb @@ -0,0 +1,434 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "1173f2f8", + "metadata": {}, + "source": [ + "# SAR-MNL: Spatial Autoregressive Multinomial Logit Demo\n", + "\n", + "This notebook demonstrates the `SARMNL` class — a spatial autoregressive MNL model\n", + "that specifies a spatial lag in the systematic utility of alternatives (locations).\n", + "\n", + "The model specifies:\n", + "\n", + "$$V_j = \\rho \\sum_k w_{jk} V_k + Z_j \\beta + X_{ij} \\gamma$$\n", + "\n", + "yielding reduced-form utilities $V^* = (I - \\rho W)^{-1}(Z\\beta + X\\gamma)$,\n", + "normalised by $D = \\text{diag}((I - \\rho W)^{-1})$, with standard MNL choice probabilities.\n", + "\n", + "Estimation is via pseudo maximum likelihood (PML, Smirnov 2010) with JAX autodiff." + ] + }, + { + "cell_type": "markdown", + "id": "e7844e2e", + "metadata": {}, + "source": [ + "## 1. Setup and Imports" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e0f04aab", + "metadata": {}, + "outputs": [], + "source": [ + "import time\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "\n", + "from locpick import ChoiceModel\n", + "from locpick.dgp import simulate_sar_mnl" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "95219f10", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "28ae38ee", + "metadata": {}, + "outputs": [], + "source": [ + "import geosnap as gsp" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "55d92adc", + "metadata": {}, + "outputs": [], + "source": [ + "datasets = gsp.DataStore()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e8f0e2b3", + "metadata": {}, + "outputs": [], + "source": [ + "dc = gsp.io.get_acs(datasets, years=2019, level=\"tract\", state_fips=\"11\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9d5f5463", + "metadata": {}, + "outputs": [], + "source": [ + "dc.plot()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3589ee6d", + "metadata": {}, + "outputs": [], + "source": [ + "from libpysal.graph import Graph\n", + "\n", + "dc_graph = Graph.build_contiguity(dc, rook=False)\n", + "\n", + "adj = dc_graph.sparse.todense()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "07983d83", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "markdown", + "id": "18e91aa6", + "metadata": {}, + "source": [ + "## 2. Data Preparation\n", + "\n", + "We generate synthetic data from a known SAR-MNL data generating process using `simulate_sar_mnl`. This creates:\n", + "- **Choosers**: observations with an income attribute\n", + "- **Alternatives**: spatial locations with cost and time attributes\n", + "- **Spatial weights matrix (W)**: a circular adjacency graph connecting nearby alternatives\n", + "- **Choices**: simulated from the SAR-MNL probability with a known spatial autoregressive parameter ρ" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "32d1b494", + "metadata": {}, + "outputs": [], + "source": [ + "# Generate synthetic SAR-MNL data\n", + "# n_obs=2000 choosers, n_alts=12 alternatives, rho=0.5\n", + "dataset = simulate_sar_mnl(n_obs=2000, n_alts=dc.shape[0], rho=0.5, seed=42, W=dc_graph)\n", + "\n", + "print(f\"Observations: {dataset.n_obs}\")\n", + "print(f\"Alternatives: {dataset.n_alts}\")\n", + "print(f\"True rho: {dataset.true_rho}\")\n", + "print(f\"True params: {dataset.true_params}\")\n", + "print(f\"\\nChoosers columns: {list(dataset.choosers.columns)}\")\n", + "print(f\"Alternatives columns: {list(dataset.alternatives.columns)}\")\n", + "print(f\"\\nChoice table: {dataset.choice_table}\")" + ] + }, + { + "cell_type": "markdown", + "id": "9ad4b19b", + "metadata": {}, + "source": [ + "## 3. Instantiating the SARMNL Class\n", + "\n", + "The `SARMNL` class takes the same data interface as `ChoiceModel`, plus a spatial weights matrix `W`. The weights matrix defines the spatial relationships between alternatives (locations).\n", + "\n", + "Key parameters:\n", + "- **`data`**: A `ChoiceTable` containing choosers, alternatives, and choices\n", + "- **`formula`**: A formulaic formula string for the utility function\n", + "- **`W`**: Spatial weights matrix (libpysal Graph, scipy.sparse, or np.ndarray) — row-standardised internally\n", + "- **`solver`**: Default `\"lbfgs\"` (scipy L-BFGS-B, the fastest path)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "87784679", + "metadata": {}, + "outputs": [], + "source": [ + "dataset.choice_table.to_frame()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bd918425", + "metadata": {}, + "outputs": [], + "source": [ + "# Create the SAR-MNL model\n", + "# W is the spatial weights matrix from the DGP (circular adjacency)\n", + "model_sar = ChoiceModel(\n", + " dataset.choice_table, formula=\"alt_attr + obs_x_alt - 1\", graph=dataset.W, lag=True\n", + ")\n", + "\n", + "print(f\"Model: {model_sar}\")\n", + "print(f\"W shape: {dataset.W.n}\")\n", + "print(f\"W type: {type(dataset.W).__name__}\")" + ] + }, + { + "cell_type": "markdown", + "id": "f1ed1cd9", + "metadata": {}, + "source": [ + "## 4. Model Estimation\n", + "\n", + "Fit the SAR-MNL model and compare with a standard MNL (no spatial lag). The spatial autoregressive parameter ρ captures the degree of spatial spillover in utilities — nearby locations influence each other's attractiveness." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e757464e", + "metadata": {}, + "outputs": [], + "source": [ + "# Fit SAR-MNL\n", + "t0 = time.perf_counter()\n", + "result_sar = model_sar.fit()\n", + "t_sar = time.perf_counter() - t0\n", + "\n", + "print(\"=== SAR-MNL Results ===\")\n", + "print(result_sar.summary())\n", + "print(f\"\\nFit time: {t_sar:.2f}s\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a6f9bd5d", + "metadata": {}, + "outputs": [], + "source": [ + "# Fit standard MNL (no spatial lag) for comparison\n", + "model_mnl = ChoiceModel(dataset.choice_table, formula=\"alt_attr + obs_x_alt - 1\")\n", + "t0 = time.perf_counter()\n", + "result_mnl = model_mnl.fit()\n", + "t_mnl = time.perf_counter() - t0\n", + "\n", + "print(\"=== Standard MNL Results ===\")\n", + "print(result_mnl.summary())\n", + "print(f\"\\nFit time: {t_mnl:.2f}s\")" + ] + }, + { + "cell_type": "markdown", + "id": "52ba1635", + "metadata": {}, + "source": [ + "## 5. Results Visualization\n", + "\n", + "Compare the estimated parameters from SAR-MNL vs standard MNL, and check how well the SAR-MNL model recovers the true spatial autoregressive parameter ρ." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f33c611a", + "metadata": {}, + "outputs": [], + "source": [ + "# Compare coefficient estimates\n", + "print(\"=== Parameter Recovery ===\")\n", + "print(f\"True rho: {dataset.true_rho:.3f}\")\n", + "print(f\"Estimated rho (SAR-MNL): {result_sar.coefficients['rho']:.3f}\")\n", + "print()\n", + "\n", + "# Side-by-side coefficient comparison\n", + "comparison = pd.DataFrame(\n", + " {\n", + " \"True\": [dataset.true_params.get(c, np.nan) for c in result_sar.coefficients.index],\n", + " \"SAR-MNL\": result_sar.coefficients.values,\n", + " \"MNL\": [result_mnl.coefficients.get(c, np.nan) for c in result_sar.coefficients.index],\n", + " },\n", + " index=result_sar.coefficients.index,\n", + ")\n", + "print(comparison.round(4))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "72b9c5ca", + "metadata": {}, + "outputs": [], + "source": [ + "# Fit statistics comparison\n", + "print(\"=== Fit Statistics ===\")\n", + "stats_df = pd.DataFrame(\n", + " {\n", + " \"SAR-MNL\": [\n", + " result_sar.log_likelihood,\n", + " result_sar.aic,\n", + " result_sar.bic,\n", + " result_sar.rho_squared,\n", + " result_sar.rho_bar_squared,\n", + " ],\n", + " \"MNL\": [\n", + " result_mnl.log_likelihood,\n", + " result_mnl.aic,\n", + " result_mnl.bic,\n", + " result_mnl.rho_squared,\n", + " result_mnl.rho_bar_squared,\n", + " ],\n", + " },\n", + " index=[\"Log-likelihood\", \"AIC\", \"BIC\", \"rho^2\", \"rho-bar^2\"],\n", + ")\n", + "print(stats_df.round(4))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ccb0c59a", + "metadata": {}, + "outputs": [], + "source": [ + "# Visualize spatial spillover: how does rho affect utilities?\n", + "# Compare raw utilities vs spatially-filtered utilities\n", + "probs_sar = model_sar.probabilities()\n", + "probs_mnl = model_mnl.probabilities()\n", + "\n", + "print(\"=== Probability Comparison (first 5 obs, first 5 alts) ===\")\n", + "print(\"SAR-MNL probabilities:\")\n", + "print(probs_sar[:5, :5].round(4))\n", + "print(\"\\nMNL probabilities:\")\n", + "print(probs_mnl[:5, :5].round(4))\n", + "\n", + "# The spatial autoregressive structure redistributes probability mass\n", + "# toward alternatives that are surrounded by attractive neighbors\n", + "print(f\"\\nMean absolute probability difference: {np.mean(np.abs(probs_sar - probs_mnl)):.6f}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fb21d792", + "metadata": {}, + "outputs": [], + "source": [ + "dc.assign(probs=probs_sar[0]).plot(\"probs\", scheme=\"quantiles\")" + ] + }, + { + "cell_type": "markdown", + "id": "46140f2d", + "metadata": {}, + "source": [ + "## 6. Performance Benchmarking\n", + "\n", + "Benchmark SAR-MNL vs standard MNL to highlight the computational cost of the spatial solve (matrix inversion at each evaluation). The hybrid estimation path (scipy LBFGS + JAX kernels + JAX autodiff) is used for both models." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "937d4c20", + "metadata": {}, + "outputs": [], + "source": [ + "# Warm fit (second fit, post-JIT cache)\n", + "t0 = time.perf_counter()\n", + "result_sar_warm = model_sar.fit()\n", + "t_sar_warm = time.perf_counter() - t0\n", + "\n", + "t0 = time.perf_counter()\n", + "result_mnl_warm = model_mnl.fit()\n", + "t_mnl_warm = time.perf_counter() - t0\n", + "\n", + "print(\"=== Warm Fit Performance ===\")\n", + "print(f\"SAR-MNL: {t_sar_warm:.3f}s (cold: {t_sar:.3f}s)\")\n", + "print(f\"MNL: {t_mnl_warm:.3f}s (cold: {t_mnl:.3f}s)\")\n", + "print(f\"SAR overhead: {t_sar_warm / t_mnl_warm:.1f}x\")\n", + "print(\"\\nBoth use the hybrid path: scipy LBFGS + JAX JIT'd kernels + JAX autodiff gradients\")" + ] + }, + { + "cell_type": "markdown", + "id": "13560f2b", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "The `ChoiceModel` class provides spatial autoregressive MNL estimation via pseudo maximum likelihood:\n", + "\n", + "- **Spatial weights**: Accepts libpysal Graph, scipy.sparse, or dense ndarray for `W`\n", + "- **JAX-accelerated**: Log-likelihood and gradient computed via JAX autodiff through the spatial solve\n", + "- **Hybrid estimation**: scipy LBFGS solver + JAX kernels (the fastest path for all locpick models)\n", + "- **Parameter recovery**: The spatial autoregressive parameter ρ is estimated alongside utility coefficients\n", + "\n", + "### When to use SAR-MNL vs SCL\n", + "\n", + "| Feature | SAR-MNL | SCL (ChoiceModel + graph) |\n", + "|---|---|---|\n", + "| Spatial structure | Autoregressive lag in utility | Paired GEV nests |\n", + "| Estimation | Pseudo-ML (no Jacobian) | Full-likelihood (GEV) |\n", + "| ρ range | (-1, 1) | (0, 1] |\n", + "| Variance normalisation | diag((I-ρW)⁻¹) | None (GEV handles it) |\n", + "| Computational cost | Matrix solve per eval | Scatter operations |\n", + "\n", + "Use **SAR-MNL** when you want a spatial lag in utilities (spillover effects).\n", + "Use **SCL** when you want spatial correlation in the error structure (GEV)." + ] + }, + { + "cell_type": "markdown", + "id": "71dc586c", + "metadata": {}, + "source": [] + }, + { + "cell_type": "markdown", + "id": "b8f831e6", + "metadata": {}, + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "locpick", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/source/user-guide/spatial_mixed.md b/docs/source/user-guide/spatial_mixed.md index 3407e55..ed2ad93 100644 --- a/docs/source/user-guide/spatial_mixed.md +++ b/docs/source/user-guide/spatial_mixed.md @@ -1,8 +1,8 @@ -# Mixed Logit with Spatial Correlation (MixedMNL + graph) +# Mixed Logit with Spatial Correlation (ChoiceModel + graph + random_params) ## Overview -A `MixedMNL` model constructed with a spatial `graph=` argument estimates the Mixed Spatially Correlated Logit (MSCL) of Bhat & Guo (2004): it combines a closed-form GEV spatial-correlation structure with random taste variation. The spatial component captures correlation between contiguous alternatives in closed form, while the mixing distribution captures unobserved heterogeneity across decision-makers. +A `ChoiceModel` constructed with a spatial `graph=` argument and `random_params=` estimates the Mixed Spatially Correlated Logit (MSCL) of Bhat & Guo (2004): it combines a closed-form GEV spatial-correlation structure with random taste variation. The spatial component captures correlation between contiguous alternatives in closed form, while the mixing distribution captures unobserved heterogeneity across decision-makers. ```{warning} The spatial mixed logit does **not** support alternative sampling correction. The MNL's uniform conditioning property does not hold for non-MNL GEV models. Always use the full alternative set (or sample without correction). @@ -31,7 +31,7 @@ The spatial GEV structure handles spatial correlation in closed form, so the sim ## Quick Start ```python -from locpick import ChoiceTable, MixedMNL +from locpick import ChoiceTable, ChoiceModel from locpick.models.mixed import ParamDistribution from libpysal import graph @@ -44,7 +44,7 @@ random_params = { } ct = ChoiceTable.from_tables(choosers, alternatives, chosen) -model = MixedMNL( +model = ChoiceModel( ct, formula="commute_time + density + shopping_access", graph=g, @@ -81,7 +81,7 @@ Supported distributions: ```python # Halton draws (quasi-random — more efficient) -model = MixedMNL( +model = ChoiceModel( ct, formula="cost + time", graph=g, random_params=random_params, n_draws=250, @@ -89,7 +89,7 @@ model = MixedMNL( ) # Pseudo-random draws -model = MixedMNL( +model = ChoiceModel( ct, formula="cost + time", graph=g, random_params=random_params, n_draws=500, @@ -99,12 +99,12 @@ model = MixedMNL( ## Spatial-Only Estimation -When no random parameters are needed, prefer `MNL(graph=...)` directly to avoid the simulation loop: +When no random parameters are needed, prefer `ChoiceModel(graph=...)` directly to avoid the simulation loop: ```python -from locpick import MNL +from locpick import ChoiceModel -model = MNL(ct, formula="cost + time", graph=g) +model = ChoiceModel(ct, formula="cost + time", graph=g) ``` ## References diff --git a/docs/source/user-guide/spatial_mixed_nested.md b/docs/source/user-guide/spatial_mixed_nested.md index a5c851d..1c7a53f 100644 --- a/docs/source/user-guide/spatial_mixed_nested.md +++ b/docs/source/user-guide/spatial_mixed_nested.md @@ -1,8 +1,8 @@ -# Mixed Nested Logit with Spatial Correlation (MixedNestedMNL + graph) +# Mixed Nested Logit with Spatial Correlation (ChoiceModel + graph + nests + random_params) ## Overview -A `MixedNestedMNL` model constructed with a spatial `graph=` argument estimates the most general model in the locpick spatial hierarchy. It combines three structures: +A `ChoiceModel` constructed with a spatial `graph=` argument, `nests=`, and `random_params=` estimates the most general model in the locpick spatial hierarchy. It combines three structures: 1. **Nested logit upper level**: alternatives are grouped into nests, each with a nest dissimilarity parameter $\lambda_m \in (0, 1]$ 2. **Spatial lower levels**: within each nest, spatial correlation between contiguous alternatives is captured via a paired GNL structure with nest-specific spatial dissimilarity $\rho_m \in (0, 1]$ @@ -25,7 +25,7 @@ where $\beta^r$ is the $r$-th draw of the random coefficients, approximated by s ## Quick Start ```python -from locpick import ChoiceTable, MixedNestedMNL +from locpick import ChoiceTable, ChoiceModel from locpick.models.nested import NestSpec, NestingTree from locpick.models.mixed import ParamDistribution from libpysal import graph @@ -47,7 +47,7 @@ random_params = { } ct = ChoiceTable.from_tables(choosers, alternatives, chosen_alternatives=choices) -model = MixedNestedMNL( +model = ChoiceModel( ct, formula="cost + time - 1", graph=g, @@ -78,19 +78,19 @@ Supported distributions: `normal`, `lognormal`, `triangular`, `uniform` (see [Mi ```python # QMC draws (default, Sobol sequences — most efficient) -model = MixedNestedMNL( +model = ChoiceModel( ct, formula="cost + time - 1", graph=g, nests=tree, random_params=random_params, n_draws=100, draw_type="qmc", ) # Halton draws -model = MixedNestedMNL( +model = ChoiceModel( ct, formula="cost + time - 1", graph=g, nests=tree, random_params=random_params, n_draws=250, draw_type="halton", ) # Pseudo-random draws -model = MixedNestedMNL( +model = ChoiceModel( ct, formula="cost + time - 1", graph=g, nests=tree, random_params=random_params, n_draws=500, draw_type="random", ) diff --git a/docs/source/user-guide/spatial_models_demo.ipynb b/docs/source/user-guide/spatial_models_demo.ipynb index 443c22f..c94298c 100644 --- a/docs/source/user-guide/spatial_models_demo.ipynb +++ b/docs/source/user-guide/spatial_models_demo.ipynb @@ -25,12 +25,75 @@ "import numpy as np\n", "import pandas as pd\n", "\n", - "from locpick import MNL, MixedMNL, MixedNestedMNL, NestedMNL\n", + "from locpick import ChoiceModel\n", "from locpick.dgp import simulate_scl\n", "from locpick.models.mixed import ParamDistribution\n", "from locpick.models.nested import NestingTree, NestSpec" ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import geosnap as gsp" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "datasets = gsp.DataStore()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dc = gsp.io.get_acs(datasets, years=2019, level=\"tract\", state_fips=\"11\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dc.plot()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from libpysal.graph import Graph\n", + "\n", + "dc_graph = Graph.build_contiguity(dc, rook=False)\n", + "\n", + "adj = dc_graph.sparse.todense()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, { "cell_type": "markdown", "metadata": {}, @@ -52,7 +115,7 @@ "source": [ "# Generate synthetic SCL data with spatial correlation\n", "n_obs = 2000\n", - "n_alts = 20\n", + "n_alts = dc.shape[0]\n", "\n", "scl_dataset = simulate_scl(\n", " n_obs=n_obs,\n", @@ -60,6 +123,7 @@ " alt_params={\"cost\": -0.5, \"time\": -0.2},\n", " rho=0.7,\n", " seed=42,\n", + " adjacency=adj,\n", ")\n", "\n", "ct = scl_dataset.choice_table\n", @@ -98,7 +162,7 @@ "adj = scl_dataset.adjacency\n", "\n", "# 1. Spatial MNL (SCL) — spatial correlation only\n", - "model_scl = MNL(\n", + "model_scl = ChoiceModel(\n", " ct,\n", " formula=formula,\n", " graph=adj,\n", @@ -112,7 +176,7 @@ " ]\n", ")\n", "\n", - "model_nested_scl = NestedMNL(\n", + "model_nested_scl = ChoiceModel(\n", " ct,\n", " formula=formula,\n", " graph=adj,\n", @@ -124,7 +188,7 @@ " \"time\": ParamDistribution(distribution=\"normal\", param=\"time\"),\n", "}\n", "\n", - "model_mixed_scl = MixedMNL(\n", + "model_mixed_scl = ChoiceModel(\n", " ct,\n", " formula=formula,\n", " graph=adj,\n", @@ -133,7 +197,7 @@ ")\n", "\n", "# 4. Spatial MixedNestedMNL — spatial + nesting + random coefficients\n", - "model_mixed_nested_scl = MixedNestedMNL(\n", + "model_mixed_nested_scl = ChoiceModel(\n", " ct,\n", " formula=formula,\n", " graph=adj,\n", @@ -308,17 +372,41 @@ "plt.tight_layout()\n", "plt.show()" ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dc.assign(probs=probs_scl[0]).plot(\"probs\", scheme=\"quantiles\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { "kernelspec": { - "display_name": "Python 3", + "display_name": "locpick", "language": "python", "name": "python3" }, "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", "name": "python", - "version": "3.12.0" + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.13" } }, "nbformat": 4, diff --git a/docs/source/user-guide/spatial_models_lag_demo.ipynb b/docs/source/user-guide/spatial_models_lag_demo.ipynb new file mode 100644 index 0000000..9940680 --- /dev/null +++ b/docs/source/user-guide/spatial_models_lag_demo.ipynb @@ -0,0 +1,506 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Spatial Models Demo: SAR-MNL with `graph=` and `lag=True`\n", + "\n", + "This notebook demonstrates the four spatial logit model variants in `locpick`. When `lag=True` is combined with `graph=`, the model becomes a Spatial Autoregressive (SAR) logit — alternatives influence each other through a spatial weights matrix:\n", + "\n", + "- **`ChoiceModel(..., graph=g, lag=True)`** — SAR-MNL (spatial autoregressive logit)\n", + "- **`ChoiceModel(..., graph=g, nests=..., lag=True)`** — SAR + nesting\n", + "- **`ChoiceModel(..., graph=g, random_params=..., lag=True)`** — SAR + random taste variation\n", + "- **`ChoiceModel(..., graph=g, nests=..., random_params=..., lag=True)`** — SAR + nesting + random variation\n", + "\n", + "We use synthetic data generated by `simulate_sar_mnl` with a known spatial adjacency structure from DC census tracts." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import pandas as pd\n", + "\n", + "from locpick import ChoiceModel\n", + "from locpick.dgp import simulate_sar_mnl\n", + "from locpick.models.mixed import ParamDistribution\n", + "from locpick.models.nested import NestingTree, NestSpec" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import geosnap as gsp" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "datasets = gsp.DataStore()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dc = gsp.io.get_acs(datasets, years=2019, level=\"tract\", state_fips=\"11\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dc.plot()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from libpysal.graph import Graph\n", + "\n", + "dc_graph = Graph.build_contiguity(dc, rook=False).transform(\"r\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Generate Synthetic Data\n", + "\n", + "We use `simulate_sar_mnl` to create a synthetic dataset with known spatial autoregressive structure. The DGP produces:\n", + "- `choosers`: household-level observations with a random feature\n", + "- `alternatives`: tract-level attributes (`cost`, `time`)\n", + "- `W`: a spatial weights matrix as a `libpysal.graph.Graph`\n", + "- `true_rho`: the ground-truth spatial autoregressive parameter" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Generate synthetic SAR-MNL data with spatial correlation\n", + "n_obs = 2000\n", + "n_alts = dc.shape[0]\n", + "\n", + "scl_dataset = simulate_sar_mnl(\n", + " n_obs=n_obs,\n", + " n_alts=n_alts,\n", + " alt_params={\"cost\": -0.5, \"time\": -0.2},\n", + " rho=0.5,\n", + " seed=42,\n", + " W=dc_graph,\n", + ")\n", + "\n", + "ct = scl_dataset.choice_table\n", + "print(f\"Observations: {ct.n_observations}\")\n", + "print(f\"Alternatives: {ct.n_alternatives}\")\n", + "print(f\"True rho: {scl_dataset.true_rho}\")\n", + "print(f\"W shape: {scl_dataset.W.sparse.shape}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Configure Spatial Models\n", + "\n", + "We configure four SAR model variants:\n", + "\n", + "1. **`ChoiceModel(graph=g, lag=True)`** — spatial autoregressive logit only\n", + "2. **`ChoiceModel(graph=g, nests=..., lag=True)`** — spatial + nesting structure\n", + "3. **`ChoiceModel(graph=g, random_params=..., lag=True)`** — spatial + random taste variation\n", + "4. **`ChoiceModel(graph=g, nests=..., random_params=..., lag=True)`** — spatial + nesting + random variation\n", + "\n", + "Each model shares the same `ChoiceTable` and formula but adds structural complexity. The `graph` parameter accepts a `libpysal.graph.Graph`, a `scipy.sparse` array, or a dense numpy adjacency matrix." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Common formula for all models — must match columns generated by simulate_sar_mnl\n", + "formula = \"cost + time - 1\"\n", + "\n", + "# Use the Graph from the DGP directly (preferred input type)\n", + "W = scl_dataset.W\n", + "\n", + "# 1. SAR MNL\n", + "model_scl = ChoiceModel(ct, formula=formula, graph=W, lag=True)\n", + "\n", + "# 2. SAR NestedMNL — spatial + nesting\n", + "nest_tree = NestingTree(\n", + " nests=[\n", + " NestSpec(name=\"urban\", alt_ids=list(range(0, n_alts // 2))),\n", + " NestSpec(name=\"suburban\", alt_ids=list(range(n_alts // 2, n_alts))),\n", + " ]\n", + ")\n", + "\n", + "model_nested_scl = ChoiceModel(ct, formula=formula, graph=W, nests=nest_tree, lag=True)\n", + "\n", + "# 3. SAR MixedMNL — spatial + random coefficients\n", + "random_params = {\n", + " \"time\": ParamDistribution(distribution=\"normal\", param=\"time\"),\n", + "}\n", + "\n", + "model_mixed_scl = ChoiceModel(\n", + " ct, formula=formula, graph=W, random_params=random_params, n_draws=100, lag=True\n", + ")\n", + "\n", + "# 4. SAR MixedNestedMNL — spatial + nesting + random coefficients\n", + "model_mixed_nested_scl = ChoiceModel(\n", + " ct,\n", + " formula=formula,\n", + " graph=W,\n", + " nests=nest_tree,\n", + " random_params=random_params,\n", + " n_draws=100,\n", + " lag=True,\n", + ")\n", + "\n", + "print(\"Models configured:\")\n", + "print(f\" Spatial MNL: {type(model_scl).__name__}\")\n", + "print(f\" Spatial NestedMNL: {type(model_nested_scl).__name__}\")\n", + "print(f\" Spatial MixedMNL: {type(model_mixed_scl).__name__}\")\n", + "print(f\" Spatial MixedNestedMNL: {type(model_mixed_nested_scl).__name__}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Fit Spatial Models\n", + "\n", + "Fit each model and inspect the estimated coefficients. SAR-MNL estimates a `rho` parameter that captures spatial autocorrelation. The nested variant adds `lambda` nest dissimilarity parameters. The mixed variant adds `sd_*` random-coefficient standard deviations." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Fit spatial MNL\n", + "result_scl = model_scl.fit()\n", + "print(\"=== Spatial MNL ===\")\n", + "print(result_scl.summary())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Fit spatial NestedMNL\n", + "result_nested_scl = model_nested_scl.fit()\n", + "print(\"=== Spatial NestedMNL ===\")\n", + "print(result_nested_scl.summary())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Fit spatial MixedMNL\n", + "result_mixed_scl = model_mixed_scl.fit()\n", + "print(\"=== Spatial MixedMNL ===\")\n", + "print(result_mixed_scl.summary())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Fit spatial MixedNestedMNL\n", + "result_mixed_nested_scl = model_mixed_nested_scl.fit()\n", + "print(\"=== Spatial MixedNestedMNL ===\")\n", + "print(result_mixed_nested_scl.summary())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 7. Visualize Spatial Autocorrelation\n", + "\n", + "The SAR spatial filter $(I - \\rho W)^{-1}$ creates spatial autocorrelation across alternatives. Since alternatives are spatial locations connected by $W$, nearby locations have more similar utilities than distant ones. We verify this by computing Moran's I on the utility vectors and mapping the spatial pattern." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Collect fit statistics for comparison\n", + "results = {\n", + " \"Spatial MNL\": result_scl,\n", + " \"Spatial NestedMNL\": result_nested_scl,\n", + " \"Spatial MixedMNL\": result_mixed_scl,\n", + " \"Spatial MixedNestedMNL\": result_mixed_nested_scl,\n", + "}\n", + "\n", + "comparison = pd.DataFrame(\n", + " {\n", + " name: {\n", + " \"Log-Likelihood\": r.log_likelihood,\n", + " \"Null LL\": r.log_likelihood_null,\n", + " \"AIC\": r.aic,\n", + " \"BIC\": r.bic,\n", + " \"Rho²\": r.rho_squared,\n", + " \"Adj. Rho²\": r.rho_bar_squared,\n", + " \"n_params\": r.n_parameters,\n", + " }\n", + " for name, r in results.items()\n", + " }\n", + ").T\n", + "\n", + "print(comparison.round(4))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Predict Choice Probabilities\n", + "\n", + "Predict choice probabilities for the SAR model. Since we only have alternative-level attributes (`cost`, `time`), all choosers have the same probability vector — the spatial structure comes entirely from the spatial filter $(I - \\rho W)^{-1}$." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Predict probabilities using model.probabilities() (no arguments uses fitted parameters)\n", + "probs_scl = model_scl.probabilities()\n", + "\n", + "print(f\"Spatial MNL probabilities shape: {probs_scl.shape}\")\n", + "print(f\"Probabilities sum to 1: {np.allclose(probs_scl.sum(axis=1), 1.0)}\")\n", + "print(\"\\nFirst 3 decision-makers' probabilities:\")\n", + "print(probs_scl[:3].round(4))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 7. Spatial Autocorrelation Diagnostics\n", + "\n", + "We verify that the SAR spatial filter creates spatial autocorrelation across alternatives by computing Moran's I on utility vectors and probability vectors. We also compare SAR probabilities against plain MNL to visualize the spatial smoothing effect." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from esda.moran import Moran\n", + "\n", + "# Fit plain MNL for comparison\n", + "model_mnl = ChoiceModel(ct, formula=formula)\n", + "result_mnl = model_mnl.fit()\n", + "probs_mnl = model_mnl.probabilities()\n", + "\n", + "# Get W as dense matrix for utility decomposition\n", + "W_dense = np.asarray(W.sparse.todense(), dtype=np.float64)\n", + "\n", + "# Moran's I on probability vectors\n", + "moran_sar = Moran(probs_scl[0], W)\n", + "moran_mnl = Moran(probs_mnl[0], W)\n", + "print(f\"Moran's I (SAR probs): {moran_sar.I:.4f} (p={moran_sar.p_sim:.4f})\")\n", + "print(f\"Moran's I (MNL probs): {moran_mnl.I:.4f} (p={moran_mnl.p_sim:.4f})\")\n", + "\n", + "# Show how spatial filter transforms utilities\n", + "rho_est = float(result_scl.coefficients[\"rho\"])\n", + "A = np.eye(n_alts) - rho_est * W_dense\n", + "A_inv = np.linalg.inv(A)\n", + "D = np.diag(A_inv)\n", + "\n", + "beta_est = np.array([result_scl.coefficients[\"cost\"], result_scl.coefficients[\"time\"]])\n", + "dm = np.asarray(model_scl._arrays.design_matrix, dtype=np.float64)\n", + "V_base = (dm @ beta_est).reshape(n_obs, n_alts)[0]\n", + "V_filtered = np.linalg.solve(A, V_base)\n", + "V_star = V_filtered / D\n", + "\n", + "# Moran's I on utility vectors (estimated rho)\n", + "moran_V_base = Moran(V_base, W)\n", + "moran_V_filtered = Moran(V_filtered, W)\n", + "moran_V_star = Moran(V_star, W)\n", + "print(f\"\\nMoran's I (V_base): {moran_V_base.I:.4f} (p={moran_V_base.p_sim:.4f})\")\n", + "print(f\"Moran's I (V_filtered): {moran_V_filtered.I:.4f} (p={moran_V_filtered.p_sim:.4f})\")\n", + "print(f\"Moran's I (V_star): {moran_V_star.I:.4f} (p={moran_V_star.p_sim:.4f})\")\n", + "\n", + "print(f\"\\nρ = {rho_est:.4f}\")\n", + "print(f\"V_base std: {V_base.std():.4f}\")\n", + "print(f\"V_filtered std: {V_filtered.std():.4f}\")\n", + "print(f\"V_star std: {V_star.std():.4f}\")\n", + "print(f\"D range: [{D.min():.4f}, {D.max():.4f}]\")\n", + "\n", + "# Map the difference between SAR and MNL probabilities\n", + "dc.assign(\n", + " sar_prob=probs_scl[0],\n", + " mnl_prob=probs_mnl[0],\n", + " diff=probs_scl[0] - probs_mnl[0],\n", + ").plot(\"diff\", scheme=\"quantiles\", legend=True, cmap=\"RdBu\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# --- Diagnostic: verify spatial filter creates spatial autocorrelation ---\n", + "# The SAR model says V_star = (I - ρW)^{-1} V_base / D\n", + "# If ρ > 0, the spatial filter should smooth V_base across neighbors,\n", + "# creating spatial autocorrelation in V_star across alternatives.\n", + "#\n", + "# Key insight: alternatives are spatial locations connected by W.\n", + "# The spatial filter creates correlation ACROSS ALTERNATIVES, not across choosers.\n", + "# With only alternative-level attributes (cost, time), V_base is the same\n", + "# for all choosers — the spatial structure comes entirely from (I - ρW)^{-1}.\n", + "\n", + "rho_true = scl_dataset.true_rho # 0.5\n", + "A_true = np.eye(n_alts) - rho_true * W_dense\n", + "A_true_inv = np.linalg.inv(A_true)\n", + "D_true = np.diag(A_true_inv)\n", + "\n", + "# V_base is the same for all choosers (only alt-level attributes)\n", + "V_base_vec = V_base # shape (n_alts,)\n", + "\n", + "# Apply spatial filter with TRUE rho\n", + "V_filtered_true = np.linalg.solve(A_true, V_base_vec)\n", + "V_star_true = V_filtered_true / D_true\n", + "\n", + "# Moran's I on utility vectors (TRUE rho)\n", + "moran_V_filtered_true = Moran(V_filtered_true, W)\n", + "moran_V_star_true = Moran(V_star_true, W)\n", + "\n", + "print(\"=== Spatial autocorrelation in utility vectors ===\")\n", + "print(f\"True ρ = {rho_true:.1f}, Estimated ρ = {rho_est:.4f}\")\n", + "print(f\"\\nMoran's I (V_base): {moran_V_base.I:.4f} (p={moran_V_base.p_sim:.4f})\")\n", + "print(f\"Moran's I (V_filtered, ρ̂): {moran_V_filtered.I:.4f} (p={moran_V_filtered.p_sim:.4f})\")\n", + "print(f\"Moran's I (V_star, ρ̂): {moran_V_star.I:.4f} (p={moran_V_star.p_sim:.4f})\")\n", + "print(\n", + " f\"Moran's I (V_filtered, ρ=0.5): {moran_V_filtered_true.I:.4f} (p={moran_V_filtered_true.p_sim:.4f})\"\n", + ")\n", + "print(f\"Moran's I (V_star, ρ=0.5): {moran_V_star_true.I:.4f} (p={moran_V_star_true.p_sim:.4f})\")\n", + "\n", + "print(f\"\\nV_base std: {V_base_vec.std():.4f}\")\n", + "print(f\"V_filtered std (ρ̂): {V_filtered.std():.4f}\")\n", + "print(f\"V_star std (ρ̂): {V_star.std():.4f}\")\n", + "print(f\"V_filtered std (ρ=0.5): {V_filtered_true.std():.4f}\")\n", + "print(f\"V_star std (ρ=0.5): {V_star_true.std():.4f}\")\n", + "print(f\"D range (ρ̂): [{D.min():.4f}, {D.max():.4f}]\")\n", + "print(f\"D range (ρ=0.5): [{D_true.min():.4f}, {D_true.max():.4f}]\")\n", + "\n", + "# Probabilities with true rho\n", + "probs_true = np.exp(V_star_true) / np.exp(V_star_true).sum()\n", + "moran_probs_true = Moran(probs_true, W)\n", + "moran_probs_mnl_vec = Moran(np.exp(V_base_vec) / np.exp(V_base_vec).sum(), W)\n", + "print(f\"\\nMoran's I (probs SAR, ρ̂): {moran_sar.I:.4f} (p={moran_sar.p_sim:.4f})\")\n", + "print(f\"Moran's I (probs SAR, ρ=0.5): {moran_probs_true.I:.4f} (p={moran_probs_true.p_sim:.4f})\")\n", + "print(\n", + " f\"Moran's I (probs MNL): {moran_probs_mnl_vec.I:.4f} (p={moran_probs_mnl_vec.p_sim:.4f})\"\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "\n", + "fig, axes = plt.subplots(2, 2, figsize=(14, 10))\n", + "\n", + "# 1. V_base (no spatial structure — random attributes)\n", + "dc.assign(V_base=V_base).plot(\"V_base\", ax=axes[0, 0], legend=True, cmap=\"viridis\")\n", + "axes[0, 0].set_title(f\"V_base (Moran's I = {moran_V_base.I:.3f})\")\n", + "\n", + "# 2. V_star with estimated rho (spatially filtered)\n", + "dc.assign(V_star=V_star).plot(\"V_star\", ax=axes[0, 1], legend=True, cmap=\"viridis\")\n", + "axes[0, 1].set_title(f\"V_star, ρ̂={rho_est:.3f} (Moran's I = {moran_V_star.I:.3f})\")\n", + "\n", + "# 3. V_star with true rho (stronger spatial pattern)\n", + "dc.assign(V_star_true=V_star_true).plot(\"V_star_true\", ax=axes[1, 0], legend=True, cmap=\"viridis\")\n", + "axes[1, 0].set_title(f\"V_star, ρ=0.5 (Moran's I = {moran_V_star_true.I:.3f})\")\n", + "\n", + "# 4. SAR vs MNL probability difference\n", + "dc.assign(\n", + " sar_prob=probs_scl[0],\n", + " mnl_prob=probs_mnl[0],\n", + " diff=probs_scl[0] - probs_mnl[0],\n", + ").plot(\"diff\", ax=axes[1, 1], scheme=\"quantiles\", legend=True, cmap=\"RdBu\")\n", + "axes[1, 1].set_title(\"SAR − MNL probability difference\")\n", + "\n", + "plt.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "locpick", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.13" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/docs/source/user-guide/spatial_nested.md b/docs/source/user-guide/spatial_nested.md index db0f800..77643a5 100644 --- a/docs/source/user-guide/spatial_nested.md +++ b/docs/source/user-guide/spatial_nested.md @@ -1,8 +1,8 @@ -# Nested Logit with Spatial Correlation (NestedMNL + graph) +# Nested Logit with Spatial Correlation (ChoiceModel + graph + nests) ## Overview -A `NestedMNL` model constructed with a spatial `graph=` argument estimates the Nested Spatially Correlated Logit (Nested SCL) model: a nested logit upper level combined with spatially correlated lower levels. Each nest has: +A `ChoiceModel` constructed with a spatial `graph=` argument and `nests=` estimates the Nested Spatially Correlated Logit (Nested SCL) model: a nested logit upper level combined with spatially correlated lower levels. Each nest has: - A spatial dissimilarity parameter $\rho_m \in (0, 1]$ governing correlation between spatially adjacent alternatives within the nest - A nest dissimilarity parameter $\lambda_m \in (0, 1]$ governing correlation between alternatives in the same nest @@ -26,7 +26,7 @@ When $\rho_m = 1$ and $\lambda_m = 1$ for all nests, the model reduces to MNL. ## Quick Start ```python -from locpick import ChoiceTable, NestedMNL +from locpick import ChoiceTable, ChoiceModel from locpick.models.nested import NestSpec, NestingTree from libpysal import graph @@ -42,7 +42,7 @@ tree = NestingTree( ) ct = ChoiceTable.from_tables(choosers, alternatives, chosen_alternatives=choices) -model = NestedMNL(ct, formula="cost + time - 1", graph=g, nests=tree) +model = ChoiceModel(ct, formula="cost + time - 1", graph=g, nests=tree) result = model.fit() print(result.summary()) ``` diff --git a/locpick/__init__.pyi b/locpick/__init__.pyi index bf82ea2..e774b6e 100644 --- a/locpick/__init__.pyi +++ b/locpick/__init__.pyi @@ -68,6 +68,9 @@ from .dgp import ( from .dgp import ( NestedSCLDataset as NestedSCLDataset, ) +from .dgp import ( + SARMNLDataset as SARMNLDataset, +) from .dgp import ( SCLDataset as SCLDataset, ) @@ -93,23 +96,14 @@ from .dgp import ( simulate_nested_scl as simulate_nested_scl, ) from .dgp import ( - simulate_scl as simulate_scl, + simulate_sar_mnl as simulate_sar_mnl, ) -from .models import ( - MNL as MNL, +from .dgp import ( + simulate_scl as simulate_scl, ) from .models import ( ChoiceModel as ChoiceModel, ) -from .models import ( - MixedMNL as MixedMNL, -) -from .models import ( - MixedNestedMNL as MixedNestedMNL, -) -from .models import ( - NestedMNL as NestedMNL, -) from .models import ( NestingTree as NestingTree, ) diff --git a/locpick/_jax/__init__.py b/locpick/_jax/__init__.py index 83fa592..2993947 100644 --- a/locpick/_jax/__init__.py +++ b/locpick/_jax/__init__.py @@ -16,23 +16,23 @@ JIT-compiled and vmap'd independently of any model class. """ -from locpick._jax.builders import ( +from .builders import ( build_mixed_logit_objective, build_mnl_objective, build_mscl_objective, build_nested_objective, build_scl_objective, ) -from locpick._jax.data import ChoiceDataJAX, EdgeDataJAX -from locpick._jax.kernels import ( +from .data import ChoiceDataJAX, EdgeDataJAX +from .kernels import ( mixed_logit_ll, mnl_log_probs, mnl_probs, nested_log_probs, scl_log_probs, ) -from locpick._jax.objective import Objective -from locpick._jax.transforms import Identity, ParamTransform, Sigmoid, SoftPlus +from .objective import Objective +from .transforms import Identity, ParamTransform, Sigmoid, SoftPlus __all__ = [ "ChoiceDataJAX", diff --git a/locpick/_jax/builders.py b/locpick/_jax/builders.py index 33f2f2f..54be6ec 100644 --- a/locpick/_jax/builders.py +++ b/locpick/_jax/builders.py @@ -22,8 +22,8 @@ import numpy as np from jax.scipy.special import logsumexp as jax_logsumexp -from locpick._jax.data import ChoiceDataJAX -from locpick._jax.kernels import ( +from .data import ChoiceDataJAX +from .kernels import ( _NEG_INF, compute_ll, compute_ll_contribs, @@ -33,8 +33,8 @@ nested_log_probs, scl_log_probs, ) -from locpick._jax.objective import Objective -from locpick._jax.transforms import Identity, ParamTransform, Sigmoid, SoftPlus +from .objective import Objective +from .transforms import Identity, ParamTransform, Sigmoid, SoftPlus # --------------------------------------------------------------------------- # MNL objective @@ -134,11 +134,6 @@ def _grad_jax(beta): # --------------------------------------------------------------------------- -# --------------------------------------------------------------------------- -# SCL objective -# --------------------------------------------------------------------------- - - # Top-level JIT'd kernels — cached across all SCL objectives @jax.jit def _scl_ll_kernel(params, data): @@ -376,7 +371,7 @@ def _mnscl_ll_kernel( nest_alt_indices : tuple of tuple of int Precomputed nest alt indices. """ - from locpick._jax.kernels import scl_log_probs_and_inclusive_value + from .kernels import scl_log_probs_and_inclusive_value beta_fixed = params[:k_fixed] alpha_rhos = params[k_fixed : k_fixed + n_nests] @@ -797,7 +792,7 @@ def _nested_scl_ll_kernel(params, data, nest_matrix, edge_data_list, k, nest_alt Precomputed nest alt indices: ``nest_alt_indices[m][i]`` = global alt index of the i-th alternative in nest m. """ - from locpick._jax.kernels import scl_log_probs_and_inclusive_value + from .kernels import scl_log_probs_and_inclusive_value beta = params[:k] n_nests = nest_matrix.shape[1] @@ -1090,7 +1085,7 @@ def _mixed_nested_ll_kernel( nest_alt_indices : tuple of tuple of int Precomputed nest alt indices. """ - from locpick._jax.kernels import mixed_nested_logit_ll + from .kernels import mixed_nested_logit_ll beta_fixed = params[:k_fixed] alpha_nest = params[k_fixed : k_fixed + n_nests] diff --git a/locpick/_jax/data.py b/locpick/_jax/data.py index 6386531..6497739 100644 --- a/locpick/_jax/data.py +++ b/locpick/_jax/data.py @@ -1,8 +1,7 @@ """JAX-ready data containers for choice model estimation. These containers hold pre-converted JAX arrays, built once and shared -across all solvers. They replace the ad-hoc numpy→JAX conversion that -was duplicated inside each model's ``_build_*_jax`` closure. +across all solvers. """ from __future__ import annotations @@ -14,7 +13,7 @@ import jax.numpy as jnp import numpy as np -from locpick._sampling.correction import get_sampling_correction +from .._sampling.correction import get_sampling_correction @jax.tree_util.register_pytree_node_class diff --git a/locpick/_jax/kernels.py b/locpick/_jax/kernels.py index 23f44b5..1aa5056 100644 --- a/locpick/_jax/kernels.py +++ b/locpick/_jax/kernels.py @@ -23,7 +23,7 @@ import jax -from locpick._kernels.constants import NEG_INF as _NEG_INF_FLOAT +from .._kernels.constants import NEG_INF as _NEG_INF_FLOAT # Enable x64 before any jnp.float64 expression is evaluated at module-import # time below. diff --git a/locpick/_jax/objective.py b/locpick/_jax/objective.py index fb60516..f5238a7 100644 --- a/locpick/_jax/objective.py +++ b/locpick/_jax/objective.py @@ -20,7 +20,7 @@ import jax.numpy as jnp import numpy as np -from locpick._jax.transforms import ParamTransform +from .transforms import ParamTransform @dataclass diff --git a/locpick/_jax/sar_kernels.py b/locpick/_jax/sar_kernels.py new file mode 100644 index 0000000..0be620f --- /dev/null +++ b/locpick/_jax/sar_kernels.py @@ -0,0 +1,783 @@ +"""JAX kernels and objective builder for SAR-MNL PML estimation. + +Implements the pseudo maximum likelihood (PML) estimator from +Smirnov (2010): spatially-filtered utilities with variance +normalisation by ``diag((I - ρW)^{-1})``, then standard MNL softmax. + +Two solve paths are available: + +- **Dense** (default for ``n_alts ≤ 2000``): LU factorisation via + ``jax.scipy.linalg.solve`` / ``inv``. The same matrix ``A = I - ρW`` + is factorised once and reused for all choosers. +- **Conjugate gradient** (for ``n_alts > 2000``): iterative solve via + ``jax.scipy.sparse.linalg.cg``. Avoids materialising the dense + inverse; the diagonal of ``A^{-1}`` is estimated via a power-series + approximation. +""" + +from __future__ import annotations + +import jax +import jax.numpy as jnp +import scipy.sparse as sp + +from .data import ChoiceDataJAX +from .kernels import ( + compute_ll, + compute_ll_contribs, + compute_utilities, + mnl_log_probs, + nested_log_probs, +) +from .objective import Objective +from .transforms import Identity, ParamTransform, Sigmoid, SoftPlus, Tanh + +# Threshold for switching from dense solve to conjugate gradient. +_DENSE_CUTOFF = 2000 + + +# --------------------------------------------------------------------------- +# Dense solve path +# --------------------------------------------------------------------------- + + +def _sar_mnl_ll_core( + params, design_matrix, available, chosen, weights, inclusion_probs, W_dense, n_obs, n_alts +): + """SAR-MNL PML log-likelihood — dense solve path (Smirnov 2010). + + Parameters + ---------- + params : jnp.ndarray, shape (k+1,) + [beta_1..k, alpha_rho] where rho = tanh(alpha_rho). + design_matrix : jnp.ndarray, shape (n_obs * n_alts, k) + available : jnp.ndarray, shape (n_obs, n_alts) + chosen : jnp.ndarray, shape (n_obs, n_alts) + weights : jnp.ndarray, shape (n_obs,) + inclusion_probs : jnp.ndarray or None + W_dense : jnp.ndarray, shape (n_alts, n_alts) + Dense spatial weights matrix (row-standardised, zero diagonal). + n_obs : int + n_alts : int + """ + k = design_matrix.shape[1] + beta = params[:k] + alpha_rho = params[k] + rho = jnp.tanh(alpha_rho) + + # Base utilities: V_base (n_obs, n_alts) + V_base = compute_utilities( + design_matrix, + beta, + n_obs, + n_alts, + inclusion_probs=inclusion_probs, + available=available, + ) + + # Spatial filter: solve (I - rho*W) V_filtered^T = V_base^T + # A is (n_alts, n_alts), same for all choosers — solve once for all RHS + A = jnp.eye(n_alts) - rho * W_dense + V_filtered = jax.scipy.linalg.solve(A, V_base.T).T # (n_obs, n_alts) + + # Variance normalisation: D = diag(A^{-1}) + A_inv = jax.scipy.linalg.inv(A) + D = jnp.diag(A_inv) # (n_alts,) + V_star = V_filtered / D[None, :] # normalise each alternative by d_jj + + # MNL log-probabilities + log_probs = mnl_log_probs(V_star, available) + return compute_ll(log_probs, chosen, weights) + + +def _sar_mnl_ll_contribs_core( + params, design_matrix, available, chosen, weights, inclusion_probs, W_dense, n_obs, n_alts +): + """Per-observation SAR-MNL PML log-likelihood contributions — dense path.""" + k = design_matrix.shape[1] + beta = params[:k] + alpha_rho = params[k] + rho = jnp.tanh(alpha_rho) + + V_base = compute_utilities( + design_matrix, + beta, + n_obs, + n_alts, + inclusion_probs=inclusion_probs, + available=available, + ) + + A = jnp.eye(n_alts) - rho * W_dense + V_filtered = jax.scipy.linalg.solve(A, V_base.T).T + A_inv = jax.scipy.linalg.inv(A) + D = jnp.diag(A_inv) + V_star = V_filtered / D[None, :] + + log_probs = mnl_log_probs(V_star, available) + return compute_ll_contribs(log_probs, chosen, weights) + + +# --------------------------------------------------------------------------- +# Conjugate-gradient solve path (for large n_alts) +# --------------------------------------------------------------------------- + + +def _cg_solve(A, B, n_alts): + """Solve A @ X = B via conjugate gradient, vectorised over columns of B. + + Uses ``jax.scipy.sparse.linalg.cg`` per column. JAX autodiff + works through CG via implicit differentiation. + """ + + def solve_one(b): + x, _ = jax.scipy.sparse.linalg.cg(A, b) + return x + + # vmap over columns of B (n_alts, n_rhs) + return jax.vmap(solve_one, in_axes=1, out_axes=1)(B) + + +def _diag_inv_power_series(rho, W_dense, n_alts, n_terms=20): + """Estimate diag((I - rho*W)^{-1}) via power series. + + Since W has zero diagonal, odd powers also have zero diagonal. + Only even powers contribute: d_jj = 1 + rho^2 (W^2)_jj + + rho^4 (W^4)_jj + ... Converges for |rho| < 1/omega_max. + """ + d = jnp.ones(n_alts) # first term: diag(I) = 1 + W_power = W_dense @ W_dense # W^2 + rho_sq = rho * rho + coeff = rho_sq + for _ in range(n_terms): + d = d + coeff * jnp.diag(W_power) + W_power = W_power @ W_power # W^{2k} + coeff = coeff * rho_sq + return d + + +def _sar_mnl_ll_cg_core( + params, design_matrix, available, chosen, weights, inclusion_probs, W_dense, n_obs, n_alts +): + """SAR-MNL PML log-likelihood — conjugate-gradient path. + + Uses CG for the spatial solve and a power-series approximation + for the variance normalisation diagonal. + """ + k = design_matrix.shape[1] + beta = params[:k] + alpha_rho = params[k] + rho = jnp.tanh(alpha_rho) + + V_base = compute_utilities( + design_matrix, + beta, + n_obs, + n_alts, + inclusion_probs=inclusion_probs, + available=available, + ) + + A = jnp.eye(n_alts) - rho * W_dense + # CG solve: A @ V_filtered^T = V_base^T + V_filtered = _cg_solve(A, V_base.T, n_alts).T # (n_obs, n_alts) + + # Variance normalisation via power series + D = _diag_inv_power_series(rho, W_dense, n_alts) + V_star = V_filtered / D[None, :] + + log_probs = mnl_log_probs(V_star, available) + return compute_ll(log_probs, chosen, weights) + + +def _sar_mnl_ll_contribs_cg_core( + params, design_matrix, available, chosen, weights, inclusion_probs, W_dense, n_obs, n_alts +): + """Per-observation SAR-MNL PML log-likelihood — CG path.""" + k = design_matrix.shape[1] + beta = params[:k] + alpha_rho = params[k] + rho = jnp.tanh(alpha_rho) + + V_base = compute_utilities( + design_matrix, + beta, + n_obs, + n_alts, + inclusion_probs=inclusion_probs, + available=available, + ) + + A = jnp.eye(n_alts) - rho * W_dense + V_filtered = _cg_solve(A, V_base.T, n_alts).T + D = _diag_inv_power_series(rho, W_dense, n_alts) + V_star = V_filtered / D[None, :] + + log_probs = mnl_log_probs(V_star, available) + return compute_ll_contribs(log_probs, chosen, weights) + + +# --------------------------------------------------------------------------- +# Objective builder +# --------------------------------------------------------------------------- + + +def build_sar_mnl_objective(arrays, W_sparse: sp.csr_array, use_cg: bool = False) -> Objective: + """Build an Objective for SAR-MNL PML estimation (Smirnov 2010). + + Parameters + ---------- + arrays : ChoiceArrays + Estimation data arrays. + W_sparse : scipy.sparse.csr_array + Row-standardised alt×alt spatial weights matrix (zero diagonal). + use_cg : bool, default False + If True, use conjugate-gradient solve (for large n_alts > 2000). + If False, use dense LU solve (faster for moderate n_alts). + + Returns + ------- + Objective + Objective with JIT-compiled LL, gradient, and Hessian. + Includes a Tanh transform for the rho parameter. + """ + data = ChoiceDataJAX.from_arrays(arrays) + W_dense = jnp.array(W_sparse.toarray(), dtype=jnp.float64) + n_obs = arrays.n_obs + n_alts = arrays.n_alts + k = arrays.design_matrix.shape[1] + + # Select solve path + if use_cg: + ll_core = _sar_mnl_ll_cg_core + ll_contribs_core = _sar_mnl_ll_contribs_cg_core + else: + ll_core = _sar_mnl_ll_core + ll_contribs_core = _sar_mnl_ll_contribs_core + + # JIT-compiled closures — data and W are captured, only params is dynamic + @jax.jit + def _ll_jax(params): + return ll_core( + params, + data.design_matrix, + data.available, + data.chosen, + data.weights, + data.inclusion_probs, + W_dense, + n_obs, + n_alts, + ) + + @jax.jit + def _ll_contribs_jax(params): + return ll_contribs_core( + params, + data.design_matrix, + data.available, + data.chosen, + data.weights, + data.inclusion_probs, + W_dense, + n_obs, + n_alts, + ) + + @jax.jit + def _grad_jax(params): + return jax.grad(ll_core, argnums=0)( + params, + data.design_matrix, + data.available, + data.chosen, + data.weights, + data.inclusion_probs, + W_dense, + n_obs, + n_alts, + ) + + param_names = list(arrays.param_names) + ["rho"] + transform = ParamTransform.for_sar_mnl(k) + + return Objective.from_jax( + ll_fn=_ll_jax, + grad_fn=_grad_jax, + loglike_contribs_jax=_ll_contribs_jax, + param_names=param_names, + transform=transform, + ) + + +# --------------------------------------------------------------------------- +# SAR + Nested +# --------------------------------------------------------------------------- + + +def _sar_nested_ll_core( + params, + design_matrix, + available, + chosen, + weights, + inclusion_probs, + W_dense, + nest_matrix, + n_obs, + n_alts, + n_nests, +): + """SAR + Nested PML log-likelihood. + + Layout: [beta_1..k, alpha_rho, alpha_lambda_1..M] + Spatial filter applied to utilities, then nested GEV. + """ + k = design_matrix.shape[1] + beta = params[:k] + alpha_rho = params[k] + alpha_lambdas = params[k + 1 : k + 1 + n_nests] + rho = jnp.tanh(alpha_rho) + lambdas = 1.0 / (1.0 + jnp.exp(-alpha_lambdas)) # sigmoid → (0, 1] + + V_base = compute_utilities( + design_matrix, + beta, + n_obs, + n_alts, + inclusion_probs=inclusion_probs, + available=available, + ) + + A = jnp.eye(n_alts) - rho * W_dense + V_filtered = jax.scipy.linalg.solve(A, V_base.T).T + D = jnp.diag(jax.scipy.linalg.inv(A)) + V_star = V_filtered / D[None, :] + + log_probs = nested_log_probs(V_star, lambdas, nest_matrix, available) + return compute_ll(log_probs, chosen, weights) + + +def build_sar_nested_objective(arrays, W_sparse, nest_matrix) -> Objective: + """Build Objective for SAR + Nested estimation.""" + data = ChoiceDataJAX.from_arrays(arrays) + W_dense = jnp.array(W_sparse.toarray(), dtype=jnp.float64) + n_obs = arrays.n_obs + n_alts = arrays.n_alts + k = arrays.design_matrix.shape[1] + n_nests = nest_matrix.shape[1] + nest_matrix_jax = jnp.array(nest_matrix, dtype=jnp.float64) + + @jax.jit + def _ll_jax(params): + return _sar_nested_ll_core( + params, + data.design_matrix, + data.available, + data.chosen, + data.weights, + data.inclusion_probs, + W_dense, + nest_matrix_jax, + n_obs, + n_alts, + n_nests, + ) + + @jax.jit + def _grad_jax(params): + return jax.grad(_sar_nested_ll_core, argnums=0)( + params, + data.design_matrix, + data.available, + data.chosen, + data.weights, + data.inclusion_probs, + W_dense, + nest_matrix_jax, + n_obs, + n_alts, + n_nests, + ) + + param_names = list(arrays.param_names) + ["rho"] + param_names += [f"lambda_{i}" for i in range(n_nests)] + # Transform: Identity for beta, Tanh for rho, Sigmoid for lambdas + transforms = [Identity()] * k + [Tanh()] + [Sigmoid(0, 1)] * n_nests + transform = ParamTransform(transforms) + + return Objective.from_jax( + ll_fn=_ll_jax, + grad_fn=_grad_jax, + param_names=param_names, + transform=transform, + ) + + +# --------------------------------------------------------------------------- +# SAR + Mixed +# --------------------------------------------------------------------------- + + +def _sar_mixed_ll_core( + params, + dm_fixed, + dm_random, + available, + chosen, + weights, + inclusion_probs, + W_dense, + dist_codes, + draws, + n_obs, + n_alts, + k_fixed, + k_random, + n_draws, +): + """SAR + Mixed simulated PML log-likelihood. + + Layout: [beta_fixed, alpha_rho, mean_*, sd_*] + Spatial filter applied to fixed utility; random part added after. + """ + beta_fixed = params[:k_fixed] + alpha_rho = params[k_fixed] + beta_random_means = params[k_fixed + 1 : k_fixed + 1 + k_random] + beta_random_spreads_raw = params[k_fixed + 1 + k_random :] + rho = jnp.tanh(alpha_rho) + spreads = jnp.log1p(jnp.exp(beta_random_spreads_raw)) # softplus + + A = jnp.eye(n_alts) - rho * W_dense + A_inv = jax.scipy.linalg.inv(A) + D = jnp.diag(A_inv) + + # Fixed utility (spatially filtered + normalised) + if dm_fixed is not None and k_fixed > 0: + V_fixed_base = (dm_fixed @ beta_fixed).reshape(n_obs, n_alts) + else: + V_fixed_base = jnp.zeros((n_obs, n_alts), dtype=jnp.float64) + + if inclusion_probs is not None: + V_fixed_base = V_fixed_base + jnp.log(jnp.maximum(inclusion_probs, 1e-30)) + + V_fixed_filtered = jax.scipy.linalg.solve(A, V_fixed_base.T).T + V_fixed_star = V_fixed_filtered / D[None, :] + + # Simulated likelihood: average over draws + means = beta_random_means[None, :] + + def _prob_single_draw(r): + z_r = draws[:, r, :] + beta_normal = means + spreads * z_r + beta_lognormal = jnp.exp(jnp.clip(means + spreads * z_r, -50, 50)) + t = 1.0 / (1.0 + 0.2316419 * jnp.abs(z_r)) + d = 0.3989422804014327 + poly = t * ( + 0.319381530 + + t * (-0.356563782 + t * (1.781477937 + t * (-1.821255978 + t * 1.330274429))) + ) + phi_z = jnp.where( + z_r >= 0, + 1.0 - d * jnp.exp(-0.5 * z_r * z_r) * poly, + d * jnp.exp(-0.5 * z_r * z_r) * poly, + ) + beta_uniform = means + spreads * (2.0 * phi_z - 1.0) + abs_z = jnp.abs(z_r) + tri_sign = jnp.where(z_r >= 0, 1.0, -1.0) + beta_triangular = means + spreads * tri_sign * (jnp.sqrt(2.0 * abs_z) - 1.0) + + beta_r = jnp.where( + dist_codes == 0, + beta_normal, + jnp.where( + dist_codes == 1, + beta_lognormal, + jnp.where(dist_codes == 2, beta_triangular, beta_uniform), + ), + ) # (n_obs, k_random) + + # Random utility: broadcast per-obs random coefficients with design matrix + V_random = jnp.sum( + dm_random.reshape(n_obs, n_alts, k_random) * beta_r[:, None, :], + axis=2, + ) + V_total = V_fixed_star + V_random + V_masked = jnp.where(available > 0, V_total, -1e30) + log_sum_exp = jax.scipy.special.logsumexp(V_masked, axis=1) + log_probs = V_masked - log_sum_exp[:, None] + log_probs = jnp.where(available > 0, log_probs, -1e30) + # Per-obs probability for this draw + return jnp.exp((log_probs * chosen).sum(axis=1)) + + probs_sim = jax.vmap(_prob_single_draw)(jnp.arange(n_draws)).mean(axis=0) + ll = jnp.sum(jnp.log(jnp.maximum(probs_sim, 1e-30)) * weights) + return ll + + +def build_sar_mixed_objective( + arrays, W_sparse, random_col_indices, random_distributions, draws +) -> Objective: + """Build Objective for SAR + Mixed estimation.""" + data = ChoiceDataJAX.from_arrays( + arrays, + draws=draws, + random_col_indices=random_col_indices, + random_distributions=random_distributions, + ) + W_dense = jnp.array(W_sparse.toarray(), dtype=jnp.float64) + n_obs = arrays.n_obs + n_alts = arrays.n_alts + k_fixed = len(data.fixed_col_indices) if data.fixed_col_indices else 0 + k_random = len(random_col_indices) + n_draws = draws.shape[1] + + @jax.jit + def _ll_jax(params): + return _sar_mixed_ll_core( + params, + data.dm_fixed, + data.dm_random, + data.available, + data.chosen, + data.weights, + data.inclusion_probs, + W_dense, + data.dist_codes, + data.draws, + n_obs, + n_alts, + k_fixed, + k_random, + n_draws, + ) + + @jax.jit + def _grad_jax(params): + return jax.grad(_sar_mixed_ll_core, argnums=0)( + params, + data.dm_fixed, + data.dm_random, + data.available, + data.chosen, + data.weights, + data.inclusion_probs, + W_dense, + data.dist_codes, + data.draws, + n_obs, + n_alts, + k_fixed, + k_random, + n_draws, + ) + + param_names_list = list(arrays.param_names) + fixed_names = [name for i, name in enumerate(param_names_list) if i not in random_col_indices] + random_names = [param_names_list[i] for i in random_col_indices] + display_names = ( + fixed_names + + ["rho"] + + [f"mean_{n}" for n in random_names] + + [f"sd_{n}" for n in random_names] + ) + + transforms = ( + [Identity()] * k_fixed + [Tanh()] + [Identity()] * k_random + [SoftPlus()] * k_random + ) + transform = ParamTransform(transforms) + + return Objective.from_jax( + ll_fn=_ll_jax, + grad_fn=_grad_jax, + param_names=display_names, + transform=transform, + ) + + +# --------------------------------------------------------------------------- +# SAR + Mixed + Nested +# --------------------------------------------------------------------------- + + +def _sar_mixed_nested_ll_core( + params, + dm_fixed, + dm_random, + available, + chosen, + weights, + inclusion_probs, + W_dense, + nest_matrix, + dist_codes, + draws, + n_obs, + n_alts, + k_fixed, + k_random, + n_draws, + n_nests, +): + """SAR + Mixed + Nested simulated PML log-likelihood. + + Layout: [beta_fixed, alpha_rho, alpha_lambda_1..M, mean_*, sd_*] + """ + beta_fixed = params[:k_fixed] + alpha_rho = params[k_fixed] + alpha_lambdas = params[k_fixed + 1 : k_fixed + 1 + n_nests] + beta_random_means = params[k_fixed + 1 + n_nests : k_fixed + 1 + n_nests + k_random] + beta_random_spreads_raw = params[k_fixed + 1 + n_nests + k_random :] + rho = jnp.tanh(alpha_rho) + lambdas = 1.0 / (1.0 + jnp.exp(-alpha_lambdas)) + spreads = jnp.log1p(jnp.exp(beta_random_spreads_raw)) + + A = jnp.eye(n_alts) - rho * W_dense + A_inv = jax.scipy.linalg.inv(A) + D = jnp.diag(A_inv) + + if dm_fixed is not None and k_fixed > 0: + V_fixed_base = (dm_fixed @ beta_fixed).reshape(n_obs, n_alts) + else: + V_fixed_base = jnp.zeros((n_obs, n_alts), dtype=jnp.float64) + if inclusion_probs is not None: + V_fixed_base = V_fixed_base + jnp.log(jnp.maximum(inclusion_probs, 1e-30)) + + V_fixed_filtered = jax.scipy.linalg.solve(A, V_fixed_base.T).T + V_fixed_star = V_fixed_filtered / D[None, :] + + means = beta_random_means[None, :] + + def _prob_single_draw(r): + z_r = draws[:, r, :] + beta_normal = means + spreads * z_r + beta_lognormal = jnp.exp(jnp.clip(means + spreads * z_r, -50, 50)) + t = 1.0 / (1.0 + 0.2316419 * jnp.abs(z_r)) + d = 0.3989422804014327 + poly = t * ( + 0.319381530 + + t * (-0.356563782 + t * (1.781477937 + t * (-1.821255978 + t * 1.330274429))) + ) + phi_z = jnp.where( + z_r >= 0, + 1.0 - d * jnp.exp(-0.5 * z_r * z_r) * poly, + d * jnp.exp(-0.5 * z_r * z_r) * poly, + ) + beta_uniform = means + spreads * (2.0 * phi_z - 1.0) + abs_z = jnp.abs(z_r) + tri_sign = jnp.where(z_r >= 0, 1.0, -1.0) + beta_triangular = means + spreads * tri_sign * (jnp.sqrt(2.0 * abs_z) - 1.0) + + beta_r = jnp.where( + dist_codes == 0, + beta_normal, + jnp.where( + dist_codes == 1, + beta_lognormal, + jnp.where(dist_codes == 2, beta_triangular, beta_uniform), + ), + ) # (n_obs, k_random) + + # Random utility: broadcast per-obs random coefficients with design matrix + V_random = jnp.sum( + dm_random.reshape(n_obs, n_alts, k_random) * beta_r[:, None, :], + axis=2, + ) + V_total = V_fixed_star + V_random + log_probs = nested_log_probs(V_total, lambdas, nest_matrix, available) + return jnp.exp((log_probs * chosen).sum(axis=1)) + + probs_sim = jax.vmap(_prob_single_draw)(jnp.arange(n_draws)).mean(axis=0) + ll = jnp.sum(jnp.log(jnp.maximum(probs_sim, 1e-30)) * weights) + return ll + + +def build_sar_mixed_nested_objective( + arrays, W_sparse, nest_matrix, random_col_indices, random_distributions, draws +) -> Objective: + """Build Objective for SAR + Mixed + Nested estimation.""" + data = ChoiceDataJAX.from_arrays( + arrays, + draws=draws, + random_col_indices=random_col_indices, + random_distributions=random_distributions, + ) + W_dense = jnp.array(W_sparse.toarray(), dtype=jnp.float64) + n_obs = arrays.n_obs + n_alts = arrays.n_alts + k_fixed = len(data.fixed_col_indices) if data.fixed_col_indices else 0 + k_random = len(random_col_indices) + n_draws = draws.shape[1] + n_nests = nest_matrix.shape[1] + nest_matrix_jax = jnp.array(nest_matrix, dtype=jnp.float64) + + @jax.jit + def _ll_jax(params): + return _sar_mixed_nested_ll_core( + params, + data.dm_fixed, + data.dm_random, + data.available, + data.chosen, + data.weights, + data.inclusion_probs, + W_dense, + nest_matrix_jax, + data.dist_codes, + data.draws, + n_obs, + n_alts, + k_fixed, + k_random, + n_draws, + n_nests, + ) + + @jax.jit + def _grad_jax(params): + return jax.grad(_sar_mixed_nested_ll_core, argnums=0)( + params, + data.dm_fixed, + data.dm_random, + data.available, + data.chosen, + data.weights, + data.inclusion_probs, + W_dense, + nest_matrix_jax, + data.dist_codes, + data.draws, + n_obs, + n_alts, + k_fixed, + k_random, + n_draws, + n_nests, + ) + + param_names_list = list(arrays.param_names) + fixed_names = [name for i, name in enumerate(param_names_list) if i not in random_col_indices] + random_names = [param_names_list[i] for i in random_col_indices] + display_names = ( + fixed_names + + ["rho"] + + [f"lambda_{i}" for i in range(n_nests)] + + [f"mean_{n}" for n in random_names] + + [f"sd_{n}" for n in random_names] + ) + + transforms = ( + [Identity()] * k_fixed + + [Tanh()] + + [Sigmoid(0, 1)] * n_nests + + [Identity()] * k_random + + [SoftPlus()] * k_random + ) + transform = ParamTransform(transforms) + + return Objective.from_jax( + ll_fn=_ll_jax, + grad_fn=_grad_jax, + param_names=display_names, + transform=transform, + ) diff --git a/locpick/_jax/transforms.py b/locpick/_jax/transforms.py index 0b49d95..aa16c46 100644 --- a/locpick/_jax/transforms.py +++ b/locpick/_jax/transforms.py @@ -71,6 +71,21 @@ def log_det_jac(self, x: jnp.ndarray) -> jnp.ndarray: return jnp.log(jnp.maximum(jax_sigmoid(x), 1e-30)) +class Tanh: + """Tanh transformation: x → tanh(x). + + Maps unconstrained parameters to (-1, 1). Used for the SAR + spatial autoregressive parameter ρ ∈ (-1, 1). + """ + + def constrain(self, x: jnp.ndarray) -> jnp.ndarray: + return jnp.tanh(x) + + def log_det_jac(self, x: jnp.ndarray) -> jnp.ndarray: + """log |d(tanh(x))/dx| = log(1 - tanh(x)^2) = log(sech^2(x)).""" + return jnp.log(jnp.maximum(1.0 - jnp.tanh(x) ** 2, 1e-30)) + + class Exp: """Exponential transformation: x → exp(x). @@ -100,7 +115,7 @@ class ParamTransform: Examples -------- - >>> from locpick._jax.transforms import ParamTransform, Identity, Sigmoid + >>> from .transforms import ParamTransform, Identity, Sigmoid >>> # SCL model: [beta_0, beta_1, rho] → [beta_0, beta_1, sigmoid(alpha_rho)] >>> pt = ParamTransform([Identity(), Identity(), Sigmoid(0, 1)]) >>> x = jnp.array([0.5, -0.1, 0.0]) @@ -180,6 +195,22 @@ def for_mscl(cls, k_fixed: int, k_random: int): transforms.extend([SoftPlus()] * k_random) # random spreads return cls(transforms) + @classmethod + def for_sar_mnl(cls, k: int): + """Create a transform for SAR-MNL models. + + Parameters + ---------- + k : int + Number of utility coefficients (beta parameters). + + Returns + ------- + ParamTransform + Transform with Identity for betas and Tanh for rho. + """ + return cls([Identity()] * k + [Tanh()]) + @classmethod def from_bounds(cls, bounds: list[tuple[float, float] | None]): """Create a transform from per-parameter bounds. diff --git a/locpick/_kernels/mnl_numpy.py b/locpick/_kernels/mnl_numpy.py index 0f0ad97..d85fce4 100644 --- a/locpick/_kernels/mnl_numpy.py +++ b/locpick/_kernels/mnl_numpy.py @@ -21,7 +21,7 @@ import numpy as np from scipy.special import logsumexp as scipy_logsumexp -from locpick._kernels.constants import NEG_INF +from .constants import NEG_INF # --------------------------------------------------------------------------- # Type aliases @@ -127,7 +127,6 @@ def mnl_log_likelihood_numpy( n_alts: int, weights: OptionalArray = None, inclusion_probs: OptionalArray = None, - design_matrix_sparse=None, ) -> float: """Compute the MNL log-likelihood. @@ -149,9 +148,6 @@ def mnl_log_likelihood_numpy( Observation weights. inclusion_probs : np.ndarray or None, shape (n_obs, n_alts) Inclusion probabilities for sampling correction. - design_matrix_sparse : scipy.sparse.spmatrix or None - Sparse design matrix. If provided, used instead of dense - ``design_matrix`` for the matrix-vector product. Returns ------- @@ -159,10 +155,7 @@ def mnl_log_likelihood_numpy( Weighted log-likelihood. """ # Step 1: systematic utility - if design_matrix_sparse is not None: - utilities = design_matrix_sparse.dot(beta).reshape(n_obs, n_alts) - else: - utilities = (design_matrix @ beta).reshape(n_obs, n_alts) + utilities = (design_matrix @ beta).reshape(n_obs, n_alts) # Steps 2–4: log-probabilities log_probs = mnl_log_probs_numpy(utilities, available, inclusion_probs) @@ -190,7 +183,6 @@ def mnl_gradient_numpy( n_alts: int, weights: OptionalArray = None, inclusion_probs: OptionalArray = None, - design_matrix_sparse=None, ) -> np.ndarray: """Compute the MNL gradient. @@ -212,9 +204,6 @@ def mnl_gradient_numpy( Observation weights. inclusion_probs : np.ndarray or None, shape (n_obs, n_alts) Inclusion probabilities. - design_matrix_sparse : scipy.sparse.spmatrix or None - Sparse design matrix. If provided, used instead of dense - ``design_matrix`` for the matrix-vector product. Returns ------- @@ -222,10 +211,7 @@ def mnl_gradient_numpy( Gradient vector. """ # Step 1: systematic utility - if design_matrix_sparse is not None: - utilities = design_matrix_sparse.dot(beta).reshape(n_obs, n_alts) - else: - utilities = (design_matrix @ beta).reshape(n_obs, n_alts) + utilities = (design_matrix @ beta).reshape(n_obs, n_alts) # Steps 2–4: probabilities probs = mnl_probs_numpy(utilities, available, inclusion_probs) @@ -237,10 +223,7 @@ def mnl_gradient_numpy( if weights is not None: residual = residual * weights.reshape(n_obs, 1) - if design_matrix_sparse is not None: - grad = design_matrix_sparse.T.dot(residual.ravel()) - else: - grad = design_matrix.T @ residual.ravel() + grad = design_matrix.T @ residual.ravel() return grad diff --git a/locpick/_kernels/sar_mnl_numpy.py b/locpick/_kernels/sar_mnl_numpy.py new file mode 100644 index 0000000..1c86fc3 --- /dev/null +++ b/locpick/_kernels/sar_mnl_numpy.py @@ -0,0 +1,259 @@ +"""NumPy kernels for the linearized GMM SAR-MNL estimator. + +Implements the two-step estimator from Carrión-Flores, Flores-Lagunes +& Guci (2018), extending Klier & McMillen (2008) to the multinomial +case. The linearization at ρ=0 avoids matrix inversion entirely, +making it feasible for very large alternative sets where the dense +PML solve is too expensive. + +Step 1: Standard MNL estimation (ignoring spatial dependence). +Step 2: Two-stage least squares (TSLS) using linearised gradients + and instruments Z = [X, WX]. +""" + +from __future__ import annotations + +import numpy as np + + +def compute_generalized_residuals(P, chosen): + """Compute generalized residuals u_ik = d_ik - P_ik. + + Parameters + ---------- + P : np.ndarray, shape (n_obs, n_alts) + MNL choice probabilities from Step 1. + chosen : np.ndarray, shape (n_obs, n_alts) + Binary indicator matrix. + + Returns + ------- + np.ndarray, shape (n_obs * n_alts,) + Flattened generalized residuals. + """ + return (chosen - P).ravel() + + +def compute_linearized_gradients(beta, P, X, WX, n_obs, n_alts, n_params): + """Compute linearized gradients for the SMNL estimator at ρ=0. + + At the linearization point ρ=0: (I - ρW)^{-1} = I, so the gradients + simplify to standard MNL gradients plus a spatial term. + + Parameters + ---------- + beta : np.ndarray, shape (n_params,) + Beta coefficients from Step 1 MNL. + P : np.ndarray, shape (n_obs, n_alts) + MNL choice probabilities from Step 1. + X : np.ndarray, shape (n_obs, n_alts, n_params) or (n_obs * n_alts, n_params) + Design matrix (can be in long or 3D format). + WX : np.ndarray, shape (n_obs, n_alts, n_params) + Spatially lagged design matrix. + n_obs : int + n_alts : int + n_params : int + + Returns + ------- + G_beta : np.ndarray, shape (n_obs * n_alts, n_params) + Gradient matrix for beta parameters. + G_rho : np.ndarray, shape (n_obs * n_alts,) + Gradient vector for the spatial parameter rho. + """ + # Reshape X to 3D if needed + if X.ndim == 2: + X_3d = X.reshape(n_obs, n_alts, n_params) + else: + X_3d = X + + # Beta gradients: G_{i,beta_k} = P_ik * (delta_{ilk} - P_il) * X_i + # For each (obs, alt) pair and each parameter: + # G_beta[i*J + k, p] = P[i,k] * (delta_{k,k} - P[i,k]) * X[i,k,p] + # + sum_{l != k} P[i,k] * (0 - P[i,l]) * X[i,l,p] + # = P[i,k] * (X[i,k,p] - sum_l P[i,l] * X[i,l,p]) + + # Compute weighted average of X across alternatives: sum_l P_l * X_l + # P is (n_obs, n_alts), X_3d is (n_obs, n_alts, n_params) + # weighted_X = sum_k P[i,k] * X[i,k,:] → (n_obs, n_params) + weighted_X = np.einsum("ik,ikp->ip", P, X_3d) # (n_obs, n_params) + + # G_beta[i*J + k, p] = P[i,k] * (X[i,k,p] - weighted_X[i,p]) + G_beta = np.zeros((n_obs * n_alts, n_params)) + for k in range(n_alts): + for p in range(n_params): + G_beta[k::n_alts, p] = P[:, k] * (X_3d[:, k, p] - weighted_X[:, p]) + + # Rho gradient: G_{i,rho} = P_ik * [(WX)_i * beta - sum_l P_il * (WX)_i * beta] + # = P_ik * (WX_i @ beta - sum_l P_il * WX_l @ beta) + # WX is (n_obs, n_alts, n_params), beta is (n_params,) + WX_beta = WX @ beta # (n_obs, n_alts) + weighted_WX_beta = np.sum(P * WX_beta, axis=1) # (n_obs,) + + G_rho = np.zeros(n_obs * n_alts) + for k in range(n_alts): + G_rho[k::n_alts] = P[:, k] * (WX_beta[:, k] - weighted_WX_beta) + + return G_beta, G_rho + + +def _remove_dependent_columns(Z, tol=1e-10): + """Remove linearly dependent columns via QR decomposition.""" + Q, R = np.linalg.qr(Z) + # Check diagonal of R for near-zero values + diag_R = np.abs(np.diag(R)) + keep = diag_R > tol * diag_R[0] + return Z[:, keep] + + +def smnl_tsls(G_beta, G_rho, u, Z): + """Two-stage least squares estimation of the linearized SMNL model. + + First stage: regress each gradient column on instruments Z. + Second stage: regress residuals u on fitted gradients. + + Parameters + ---------- + G_beta : np.ndarray, shape (n_obs * n_alts, n_params) + Gradient matrix for beta parameters. + G_rho : np.ndarray, shape (n_obs * n_alts,) + Gradient vector for the spatial parameter rho. + u : np.ndarray, shape (n_obs * n_alts,) + Generalized residuals from Step 1. + Z : np.ndarray, shape (n_obs * n_alts, n_instruments) + Instrument matrix [X, WX] (linearly independent columns). + + Returns + ------- + delta : np.ndarray, shape (n_params + 1,) + Parameter updates [Delta_beta_1, ..., Delta_beta_K, Delta_rho]. + se : np.ndarray, shape (n_params + 1,) + TSLS standard errors. + vcov : np.ndarray, shape (n_params + 1, n_params + 1) + TSLS covariance matrix. + """ + # Stack gradients: G = [G_beta, G_rho] + G = np.column_stack([G_beta, G_rho]) + n_params_total = G.shape[1] + n = len(u) + + # First stage: project G onto instruments Z + # G_hat = Z (Z'Z)^{-1} Z' G + ZtZ_inv = np.linalg.inv(Z.T @ Z) + G_hat = Z @ (ZtZ_inv @ (Z.T @ G)) + + # Second stage: OLS of u on G_hat + # delta = (G_hat' G_hat)^{-1} G_hat' u + GtG_inv = np.linalg.inv(G_hat.T @ G_hat) + delta = GtG_inv @ (G_hat.T @ u) + + # Residuals + e = u - G_hat @ delta + + # Variance: sigma^2 = e'e / (n - p) + sigma2 = e @ e / (n - n_params_total) + + # Covariance: V = sigma^2 * (G_hat' G_hat)^{-1} + vcov = sigma2 * GtG_inv + se = np.sqrt(np.maximum(np.diag(vcov), 0)) + + return delta, se, vcov + + +def fit_linearized_gmm(arrays, W_sparse): + """Two-step linearized GMM estimation (Carrión-Flores et al. 2018). + + Parameters + ---------- + arrays : ChoiceArrays + Estimation data arrays. + W_sparse : scipy.sparse.csr_array + Row-standardised alt×alt spatial weights matrix. + + Returns + ------- + dict + Dictionary with keys: 'beta', 'rho', 'se', 'vcov', 'log_likelihood'. + """ + from .mnl_numpy import mnl_probs_numpy + + n_obs = arrays.n_obs + n_alts = arrays.n_alts + k = arrays.design_matrix.shape[1] + + dm = np.asarray(arrays.design_matrix, dtype=np.float64) + chosen = np.asarray(arrays.chosen, dtype=np.float64).reshape(n_obs, n_alts) + + if arrays.available is not None: + available = np.asarray(arrays.available, dtype=np.float64).reshape(n_obs, n_alts) + else: + available = np.ones((n_obs, n_alts), dtype=np.float64) + + # --- Step 1: Standard MNL estimation --- + from .._solvers import get_solver + + solver = get_solver("lbfgs") + + from .._jax.builders import build_mnl_objective + + objective = build_mnl_objective(arrays) + x0 = np.zeros(k) + solver_result = solver.solve( + objective=objective, + x0=x0, + param_names=list(arrays.param_names), + bounds=None, + fixed_mask=None, + ) + beta_step1 = solver_result.coefficients + + # Compute Step 1 probabilities + V = (dm @ beta_step1).reshape(n_obs, n_alts) + P = mnl_probs_numpy(V, available, inclusion_probs=None) + + # --- Step 2: TSLS --- + # Extract X from design matrix (long format → 3D) + X_3d = dm.reshape(n_obs, n_alts, k) + + # Compute WX: spatially lagged covariates + # W is (n_alts, n_alts), X varies across choosers + # WX[i, j, p] = sum_k W[j, k] * X[i, k, p] + W_dense = W_sparse.toarray() + WX = np.einsum("jk,ikp->ijp", W_dense, X_3d) # (n_obs, n_alts, k) + + # Generalized residuals + u = compute_generalized_residuals(P, chosen) + + # Linearized gradients + G_beta, G_rho = compute_linearized_gradients(beta_step1, P, X_3d, WX, n_obs, n_alts, k) + + # Instruments: [X, WX] in long format + X_long = dm # (n_obs * n_alts, k) + WX_long = WX.reshape(n_obs * n_alts, k) + Z = np.column_stack([X_long, WX_long]) + Z = _remove_dependent_columns(Z) + + # TSLS + delta, se, vcov = smnl_tsls(G_beta, G_rho, u, Z) + + # Assemble final parameters + beta_final = beta_step1 + delta[:k] + rho_final = delta[k] + + # Compute log-likelihood at final params (using spatial model) + A = np.eye(n_alts) - rho_final * W_dense + V_base = (dm @ beta_final).reshape(n_obs, n_alts) + V_filtered = np.linalg.solve(A, V_base.T).T + D = np.diag(np.linalg.inv(A)) + V_star = V_filtered / D[None, :] + log_probs = np.log(np.maximum(mnl_probs_numpy(V_star, available), 1e-30)) + ll = np.sum(chosen * log_probs) + + return { + "beta": beta_final, + "rho": rho_final, + "se": se, + "vcov": vcov, + "log_likelihood": float(ll), + "beta_step1": beta_step1, + } diff --git a/locpick/_sampling/correction.py b/locpick/_sampling/correction.py index c5b4f13..c988a95 100644 --- a/locpick/_sampling/correction.py +++ b/locpick/_sampling/correction.py @@ -14,7 +14,7 @@ import numpy as np if TYPE_CHECKING: - from locpick.data.arrays import ChoiceArrays + from ..data.arrays import ChoiceArrays def get_sampling_correction(arrays: "ChoiceArrays") -> np.ndarray | None: diff --git a/locpick/_solvers/lbfgs.py b/locpick/_solvers/lbfgs.py index 21679fd..d0f1453 100644 --- a/locpick/_solvers/lbfgs.py +++ b/locpick/_solvers/lbfgs.py @@ -48,7 +48,7 @@ def solve_lbfgs( """ from scipy.optimize import minimize - from locpick._jax.objective import Objective + from .._jax.objective import Objective if not isinstance(objective, Objective): raise TypeError("solve_lbfgs expects an Objective instance.") diff --git a/locpick/_solvers/optax.py b/locpick/_solvers/optax.py index d313a96..a968717 100644 --- a/locpick/_solvers/optax.py +++ b/locpick/_solvers/optax.py @@ -140,7 +140,7 @@ def solve( if fixed_mask is not None and np.any(fixed_mask): raise NotImplementedError("OptaxSolver does not yet support fixed parameters.") - from locpick._jax.objective import Objective + from .._jax.objective import Objective if not isinstance(objective, Objective): raise TypeError("OptaxSolver.solve expects an Objective instance.") diff --git a/locpick/_solvers/optimagic.py b/locpick/_solvers/optimagic.py index 480547f..3724a41 100644 --- a/locpick/_solvers/optimagic.py +++ b/locpick/_solvers/optimagic.py @@ -59,7 +59,7 @@ def solve( ) -> SolverResult: import optimagic as om - from locpick._jax.objective import Objective + from .._jax.objective import Objective if not isinstance(objective, Objective): raise TypeError("OptimagicSolver.solve expects an Objective instance.") diff --git a/locpick/_solvers/optimistix.py b/locpick/_solvers/optimistix.py index f0b8bdc..5b10db7 100644 --- a/locpick/_solvers/optimistix.py +++ b/locpick/_solvers/optimistix.py @@ -66,10 +66,6 @@ class OptimistixSolver: for multi-start runs (only used when ``n_starts > 1``). seed : int Random seed for multi-start perturbations. - compute_hessian : bool - Kept for API compatibility. Hessian is now computed lazily by - the model via :meth:`Objective.hessian` when standard errors - are requested, rather than eagerly in the solver. """ # Map of method names to Optimistix solver constructors @@ -90,7 +86,6 @@ def __init__( n_starts: int = 1, start_scale: float = 1.0, seed: int = 0, - compute_hessian: bool = True, ): self.method = method self.rtol = rtol @@ -100,7 +95,6 @@ def __init__( self.n_starts = n_starts self.start_scale = start_scale self.seed = seed - self.compute_hessian = compute_hessian def _make_solver(self): """Create the Optimistix solver instance.""" @@ -198,7 +192,7 @@ def solve( ------- SolverResult """ - from locpick._jax.objective import Objective + from .._jax.objective import Objective if not isinstance(objective, Objective): raise TypeError("OptimistixSolver.solve expects an Objective instance.") @@ -262,7 +256,6 @@ def ll_fn(p): "rtol": float(self.rtol), "atol": float(self.atol), "maxiter": int(self.maxiter), - "compute_hessian": bool(self.compute_hessian), "n_starts": int(self.n_starts), "start_scale": float(self.start_scale), "seed": int(self.seed), @@ -436,7 +429,6 @@ def bounded_neg_ll(alpha, args=None): "rtol": float(self.rtol), "atol": float(self.atol), "maxiter": int(self.maxiter), - "compute_hessian": bool(self.compute_hessian), "n_starts": int(self.n_starts), "start_scale": float(self.start_scale), "seed": int(self.seed), diff --git a/locpick/_solvers/trust_ncg.py b/locpick/_solvers/trust_ncg.py index f9b6ea8..5b21b18 100644 --- a/locpick/_solvers/trust_ncg.py +++ b/locpick/_solvers/trust_ncg.py @@ -61,7 +61,7 @@ def solve( ) -> SolverResult: from scipy.optimize import minimize - from locpick._jax.objective import Objective + from .._jax.objective import Objective if not isinstance(objective, Objective): raise TypeError(f"{type(self).__name__}.solve expects an Objective instance.") diff --git a/locpick/data/arrays.py b/locpick/data/arrays.py index e4b487f..37802c1 100644 --- a/locpick/data/arrays.py +++ b/locpick/data/arrays.py @@ -8,7 +8,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any, Optional +from typing import Optional import jax.numpy as jnp import numpy as np @@ -57,13 +57,6 @@ class ChoiceArrays: design_matrix: ArrayType chosen: ArrayType - design_matrix_sparse: Optional[Any] = None - """Sparse design matrix (scipy.sparse or jax.experimental.sparse). - - When present, model kernels use sparse-dense products instead of - dense-dense, which is critical for large choice sets with many - zero-valued variables (e.g., "has_subway_station"). - """ available: Optional[ArrayType] = None weights: Optional[ArrayType] = None n_obs: int = 0 diff --git a/locpick/data/choicetable.py b/locpick/data/choicetable.py index 902c688..2c8e00c 100644 --- a/locpick/data/choicetable.py +++ b/locpick/data/choicetable.py @@ -13,13 +13,13 @@ import pandas as pd import xarray as xr -from locpick._sampling.kernels import ( +from .._sampling.kernels import ( HAS_NUMBA, _sample_unweighted_without_replacement_exclusion, _sample_weighted_without_replacement_1d_exclusion, ) -from locpick.data.arrays import ChoiceArrays -from locpick.data.dataset import ( +from .arrays import ChoiceArrays +from .dataset import ( _resolve_pairwise, build_choice_dataset, build_choice_dataset_from_long, @@ -137,7 +137,9 @@ def from_tables( ChoiceTable """ if seed is not None: - np.random.seed(seed) + rng = np.random.default_rng(seed) + else: + rng = np.random.default_rng() # Normalize index names; if the named column exists as a data column, # promote it to the index. @@ -223,6 +225,7 @@ def from_tables( replace, oid_name, aid_name, + rng=rng, ) n_obs = len(choosers) @@ -722,8 +725,6 @@ def to_arrays( spec=None, weights=None, available=None, - sparse: bool = False, - sparse_threshold: float = 0.5, ) -> ChoiceArrays: """Convert to estimation-ready arrays. @@ -739,37 +740,33 @@ def to_arrays( Observation weights as a flat array of length n_obs * n_alts. available : array-like or None Alternative availability as a flat array of length n_obs * n_alts. - sparse : bool - If True, convert the design matrix to scipy.sparse.csr_matrix - when the zero fraction exceeds ``sparse_threshold``. This is - useful for large choice sets with many zero-valued variables - (e.g., "has_subway_station"). - sparse_threshold : float - Fraction of zeros required to trigger sparse conversion. - Default 0.5 (50% zeros). Returns ------- ChoiceArrays """ - # Cache key: hashable tuple of all arguments + # Cache key: hashable representation of all arguments + # Use bytes hash for array-like weights/available to handle numpy arrays + import hashlib + + def _hashable(val): + if val is None or isinstance(val, str): + return val + if hasattr(val, "tobytes"): + return hashlib.md5(np.asarray(val).tobytes()).hexdigest() + if hasattr(val, "__iter__"): + return hashlib.md5(np.asarray(val).tobytes()).hexdigest() + return val + cache_key = ( formula, id(spec) if spec is not None else None, - tuple(weights) - if hasattr(weights, "__iter__") and not isinstance(weights, str) - else weights, - tuple(available) - if hasattr(available, "__iter__") and not isinstance(available, str) - else available, - sparse, - sparse_threshold, + _hashable(weights), + _hashable(available), ) if cache_key in self._to_arrays_cache: return self._to_arrays_cache[cache_key] - import numpy as np - df = self._get_frame_cached(copy=False) n_obs = self.n_observations n_alts = self.n_alternatives @@ -837,7 +834,7 @@ def to_arrays( inclusion_probs = None if self._sample_size is not None: - from locpick._sampling.inclusion import compute_inclusion_probs + from .._sampling.inclusion import compute_inclusion_probs n_alts_full = self.n_alternatives_full n_samples = self._sample_size @@ -869,22 +866,12 @@ def to_arrays( param_names = [] design_matrix = np.asarray(dm, dtype=np.float64) - # Sparse design matrix (optional) - design_matrix_sparse = None - if sparse and design_matrix.size > 0: - zero_fraction = 1.0 - np.count_nonzero(design_matrix) / design_matrix.size - if zero_fraction >= sparse_threshold: - import scipy.sparse as sp - - design_matrix_sparse = sp.csr_matrix(design_matrix) - # Get obs_ids and alt_ids obs_ids = np.repeat(np.asarray(self._ds.coords["obs_id"].values), n_alts) alt_ids = np.asarray(self._ds["alt_id_values"].values).reshape(-1) result = ChoiceArrays( design_matrix=design_matrix, - design_matrix_sparse=design_matrix_sparse, chosen=chosen, available=avail, weights=wts, @@ -936,8 +923,11 @@ def _build_sampled( replace: bool, oid_name: str, aid_name: str, + rng: Optional[np.random.Generator] = None, ) -> pd.DataFrame: """Build merged table with alternative sampling.""" + if rng is None: + rng = np.random.default_rng() n_obs = len(choosers) n_alts = len(alternatives) alt_ids = alternatives.index.values @@ -968,7 +958,7 @@ def _build_sampled( available_probs = probs.copy() available_probs[~available_mask] = 0 available_probs /= available_probs.sum() - sampled[i] = np.random.choice( + sampled[i] = rng.choice( alt_ids, size=sample_size, replace=True, p=available_probs ) else: @@ -978,7 +968,7 @@ def _build_sampled( if excluded_alt_ids[i] >= 0: available_mask[excluded_alt_ids[i]] = False available_alts = alt_ids[available_mask] - sampled[i] = np.random.choice(available_alts, size=sample_size, replace=True) + sampled[i] = rng.choice(available_alts, size=sample_size, replace=True) else: # Without replacement — use Numba kernels if available if weights_series is not None and weights_1d and HAS_NUMBA: @@ -1005,7 +995,7 @@ def _build_sampled( if excluded_alt_ids[i] >= 0: available_mask[excluded_alt_ids[i]] = False available_alts = alt_ids[available_mask] - sampled[i] = np.random.choice(available_alts, size=sample_size, replace=False) + sampled[i] = rng.choice(available_alts, size=sample_size, replace=False) # Ensure chosen alternative is always included if chosen_series is not None: diff --git a/locpick/data/problem.py b/locpick/data/problem.py index 53b3599..d3d8aaa 100644 --- a/locpick/data/problem.py +++ b/locpick/data/problem.py @@ -20,7 +20,7 @@ import numpy as np -from locpick.data.arrays import ChoiceArrays +from .arrays import ChoiceArrays @dataclass @@ -126,7 +126,7 @@ def from_choice_table( ------- EstimationProblem """ - from locpick.spec import ModelSpec + from ..spec import ModelSpec # Resolve spec if spec is None and formula is not None: diff --git a/locpick/dgp.py b/locpick/dgp.py index ccf1682..d6a5dd8 100644 --- a/locpick/dgp.py +++ b/locpick/dgp.py @@ -343,7 +343,7 @@ class MNSCLDataset: def _build_choice_table(choosers, alternatives, choices, matrix_data=None): """Build a ChoiceTable from component DataFrames.""" - from locpick.data.choicetable import ChoiceTable + from .data.choicetable import ChoiceTable return ChoiceTable.from_tables( choosers=choosers.drop(columns="choice", errors="ignore"), @@ -353,6 +353,126 @@ def _build_choice_table(choosers, alternatives, choices, matrix_data=None): ) +# --------------------------------------------------------------------------- +# Shared DGP helpers +# --------------------------------------------------------------------------- + + +def _build_choosers(n_obs, seed, feature_name="obs_feature"): + """Build choosers DataFrame with obs_id index and a random feature.""" + rng = np.random.default_rng(seed) + obs_ids = pd.Index(np.arange(n_obs), name="oid") + choosers = pd.DataFrame({feature_name: rng.standard_normal(n_obs)}, index=obs_ids) + return choosers, rng, obs_ids + + +def _build_alternatives(n_alts, rng, alt_features): + """Build alternatives DataFrame with alt_id index. + + Parameters + ---------- + n_alts : int + rng : np.random.Generator + alt_features : dict[str, tuple] + Mapping of column name → (low, high) for uniform draw. + """ + alt_ids = pd.Index(np.arange(n_alts), name="aid") + data = {} + for col, (low, high) in alt_features.items(): + data[col] = rng.uniform(low, high, n_alts) + alternatives = pd.DataFrame(data, index=alt_ids) + return alternatives, alt_ids + + +def _build_interactions(obs_ids, alt_ids, chooser_feature, alt_columns): + """Build chooser×alternative interaction terms. + + Returns + ------- + interactions : dict[str, pd.Series] + Named (obs_id, alt_id)-indexed Series. + interaction_index : pd.MultiIndex + """ + interaction_index = pd.MultiIndex.from_product([obs_ids, alt_ids], names=["oid", "aid"]) + n_obs = len(obs_ids) + n_alts = len(alt_ids) + interactions = {} + for alt_col in alt_columns: + alt_vals = alt_columns[alt_col] + tiled_feat = np.repeat(chooser_feature, n_alts) + tiled_alt = np.tile(alt_vals, n_obs) + name = ( + f"{chooser_feature.name}_x_{alt_col}" + if hasattr(chooser_feature, "name") + else f"obs_x_{alt_col}" + ) + interactions[name] = pd.Series(tiled_feat * tiled_alt, index=interaction_index, name=name) + return interactions, interaction_index + + +def _compute_det_utility( + n_obs, n_alts, alternatives, alt_params, interactions, interaction_coefs=None +): + """Compute deterministic utility from alt params and interactions. + + Parameters + ---------- + n_obs, n_alts : int + alternatives : pd.DataFrame + alt_params : dict[str, float] + interactions : dict[str, pd.Series] + interaction_coefs : dict[str, float] or None + Coefficients for interaction terms. Keys must match interactions. + """ + det_utility = np.zeros((n_obs, n_alts)) + for col, coef in alt_params.items(): + alt_vals = alternatives[col].to_numpy() + det_utility += coef * np.tile(alt_vals, n_obs).reshape(n_obs, n_alts) + if interaction_coefs: + for name, coef in interaction_coefs.items(): + if name in interactions: + det_utility += coef * interactions[name].to_numpy().reshape(n_obs, n_alts) + return det_utility + + +def _build_design_matrix(n_obs, alternatives, interactions, interaction_coefs): + """Build a design matrix from alternatives + interaction terms. + + Returns + ------- + design_matrix : np.ndarray, shape (n_obs * n_alts, k) + beta : np.ndarray, shape (k,) + """ + design_matrix = np.tile(alternatives.to_numpy(), (n_obs, 1)) + for name, coef in interaction_coefs.items(): + if name in interactions: + design_matrix = np.column_stack([design_matrix, interactions[name].to_numpy().ravel()]) + return design_matrix + + +def _simulate_choices_from_probs(probs, rng, n_obs, n_alts): + """Vectorized choice simulation from probability matrix. + + Replaces the Python loop ``np.array([rng.choice(n_alts, p=probs[i]) for i in range(n_obs)])`` + with vectorized inverse-CDF sampling. + """ + # Normalize to sum to 1 (numerical safety) + probs = probs / probs.sum(axis=1, keepdims=True) + cumulative = np.cumsum(probs, axis=1) + uniform = rng.random(n_obs) + choices = np.argmax(cumulative > uniform[:, None], axis=1) + return np.clip(choices, 0, n_alts - 1) + + +def _build_circular_adjacency(n_alts): + """Build a circular adjacency matrix where zone i is adjacent to i±1.""" + adjacency = np.zeros((n_alts, n_alts), dtype=np.float64) + for i in range(n_alts): + adjacency[i, (i - 1) % n_alts] = 1.0 + adjacency[i, (i + 1) % n_alts] = 1.0 + return adjacency + + # --------------------------------------------------------------------------- # MNL DGP # --------------------------------------------------------------------------- @@ -499,7 +619,7 @@ def simulate_nested_logit( ------- NestedLogitDataset """ - from locpick.models.nested import ( + from .models.nested import ( NestingTree, NestSpec, _nested_logit_probs_numpy, @@ -601,7 +721,7 @@ def simulate_nested_logit( ) # Simulate choices from probabilities - choices = np.array([rng.choice(n_alts, p=probs[i]) for i in range(n_obs)]) + choices = _simulate_choices_from_probs(probs, rng, n_obs, n_alts) choosers = choosers.copy() choosers["choice"] = choices @@ -669,7 +789,7 @@ def simulate_scl( ------- SCLDataset """ - from locpick.models.scl import _resolve_spatial_graph, _scl_log_probs_numpy + from .models.scl import _resolve_spatial_graph, _scl_log_probs_numpy if alt_params is None: alt_params = {"cost": -0.5, "time": -0.1} @@ -733,7 +853,7 @@ def simulate_scl( probs = np.exp(log_probs) # Simulate choices - choices = np.array([rng.choice(n_alts, p=probs[i]) for i in range(n_obs)]) + choices = _simulate_choices_from_probs(probs, rng, n_obs, n_alts) choosers = choosers.copy() choosers["choice"] = choices @@ -1072,8 +1192,8 @@ def simulate_nested_scl( ------- NestedSCLDataset """ - from locpick.models.nested import NestingTree, NestSpec - from locpick.models.scl import ( + from .models.nested import NestingTree, NestSpec + from .models.scl import ( _resolve_spatial_graph, ) @@ -1143,9 +1263,6 @@ def simulate_nested_scl( det_utility += interactions["income_x_cost"].to_numpy().reshape(n_obs, n_alts) * 0.05 # --- Compute Nested SCL probabilities and simulate choices -------- - np.array([alt_params[col] for col in alternatives.columns]) - np.tile(alternatives.to_numpy(), (n_obs, 1)) - nest_matrix = nests.build_nest_matrix(list(range(n_alts))) n_nests = len(nest_names) @@ -1204,7 +1321,7 @@ def simulate_nested_scl( probs = probs / probs.sum(axis=1, keepdims=True) # Simulate choices - choices = np.array([rng.choice(n_alts, p=probs[i]) for i in range(n_obs)]) + choices = _simulate_choices_from_probs(probs, rng, n_obs, n_alts) choosers = choosers.copy() choosers["choice"] = choices @@ -1371,8 +1488,8 @@ def simulate_mnscl( ------- MNSCLDataset """ - from locpick.models.nested import NestingTree, NestSpec - from locpick.models.scl import ( + from .models.nested import NestingTree, NestSpec + from .models.scl import ( _resolve_spatial_graph, ) @@ -1444,9 +1561,6 @@ def simulate_mnscl( det_utility += interactions["income_x_cost"].to_numpy().reshape(n_obs, n_alts) * 0.05 # --- Add random coefficient variation ------------------------------ - np.array([alt_params[col] for col in alternatives.columns]) - np.tile(alternatives.to_numpy(), (n_obs, 1)) - # For each random parameter, add random variation multiplied by attribute for param_name, (dist, mean, spread) in random_params.items(): list(alternatives.columns).index(param_name) @@ -1525,7 +1639,7 @@ def simulate_mnscl( probs = probs / probs.sum(axis=1, keepdims=True) # Simulate choices - choices = np.array([rng.choice(n_alts, p=probs[i]) for i in range(n_obs)]) + choices = _simulate_choices_from_probs(probs, rng, n_obs, n_alts) choosers = choosers.copy() choosers["choice"] = choices @@ -1607,7 +1721,7 @@ def simulate_mixed_nested_logit( ------- MixedNestedMNLDataset """ - from locpick.models.nested import ( + from .models.nested import ( NestingTree, NestSpec, _nested_logit_probs_numpy, @@ -1726,7 +1840,7 @@ def simulate_mixed_nested_logit( ) # Simulate choices from probabilities - choices = np.array([rng.choice(n_alts, p=probs[i]) for i in range(n_obs)]) + choices = _simulate_choices_from_probs(probs, rng, n_obs, n_alts) choosers = choosers.copy() choosers["choice"] = choices @@ -1761,3 +1875,192 @@ def simulate_mixed_nested_logit( n_alts=n_alts, seed=seed, ) + + +# --------------------------------------------------------------------------- +# SAR-MNL DGP (Smirnov 2010) +# --------------------------------------------------------------------------- + + +@dataclass +class SARMNLDataset: + """Synthetic dataset drawn from a known SAR-MNL data generating process. + + Alternatives are spatial locations connected by ``W`` (alt×alt). + Choosers select among alternatives via MNL with spatially-filtered + and variance-normalised utilities (Smirnov 2010 PML DGP). + + Attributes + ---------- + choosers : pd.DataFrame + Obs-id-indexed DataFrame with chooser attributes and ``choice`` column. + alternatives : pd.DataFrame + Alt-id-indexed DataFrame with alternative attributes. + interactions : dict[str, pd.Series] + Named ``(obs_id, alt_id)``-indexed Series for chooser×alt interactions. + true_params : dict[str, float] + Ground-truth beta coefficients (alt-level + interaction). + true_rho : float + Ground-truth spatial autoregressive parameter. + W : libpysal.graph.Graph + Row-standardised spatial weights matrix (n_alts × n_alts) as a + libpysal Graph. Use ``W.sparse`` to get the scipy.sparse matrix + for computation. + choice_table : object + Assembled ChoiceTable. + n_obs : int + n_alts : int + seed : int + """ + + choosers: pd.DataFrame + alternatives: pd.DataFrame + interactions: dict[str, pd.Series] + true_params: dict[str, float] + true_rho: float + W: Any # libpysal.graph.Graph + choice_table: Any + n_obs: int + n_alts: int + seed: int + + +def simulate_sar_mnl( + n_obs: int = 5000, + n_alts: int = 50, + alt_params: dict[str, float] | None = None, + interaction_params: dict[str, float] | None = None, + rho: float = 0.3, + W=None, + n_neighbors: int = 7, + seed: int = 1234, +) -> SARMNLDataset: + """Generate synthetic SAR-MNL choice data with known parameters. + + The DGP follows Smirnov (2010) PML model: + + 1. Build ``W`` (alt×alt, row-standardised k-nearest-neighbor Graph) + 2. Generate alternative attributes ``Z`` and chooser-alt interactions ``X`` + 3. Compute base utilities: ``V_base = Zβ + Xγ`` (n_obs × n_alts) + 4. Spatial filter: ``V_filtered = (I - ρW)^{-1} V_base^T`` + 5. Variance normalisation: ``D = diag((I - ρW)^{-1})``, + ``V_star = V_filtered / D`` (divide each alt by d_jj) + 6. Add Gumbel noise: ``U = V_star + Gumbel(0, 1)`` + 7. Choice = ``argmax(U)`` per chooser + + The variance normalisation (step 5) is essential — it matches the + PML estimator's model (Smirnov 2010). Without it, the DGP would + not match the estimation model and parameter recovery would fail. + + Parameters + ---------- + n_obs : int, default 5000 + Number of choosers. + n_alts : int, default 50 + Number of alternatives (spatial locations). Dimension of W. + For PML dense path, keep ≤ 2000. For CG path, can be larger. + alt_params : dict, optional + Mapping of alternative-level column name → true coefficient. + Default: ``{"alt_attr": -0.5}``. + interaction_params : dict, optional + Mapping of interaction column name → true coefficient. + Default: ``{"obs_x_alt": 0.8}``. + rho : float, default 0.3 + True spatial autoregressive parameter. Should be in (-1, 1). + Smirnov 2010 MC evidence: good recovery for ρ ∈ [0, 0.5]. + W : libpysal.graph.Graph, scipy.sparse, np.ndarray, or None + Pre-specified n_alts × n_alts spatial weights matrix. + If None, constructed as k-nearest-neighbor Graph on random + coordinates (matching Krisztin et al. 2022's 7-NN specification). + A ``libpysal.graph.Graph`` is the preferred input type. + n_neighbors : int, default 7 + Number of nearest neighbors for default W construction. + seed : int, default 1234 + + Returns + ------- + SARMNLDataset + Dataset with ``W`` stored as a ``libpysal.graph.Graph`` (row-standardised). + """ + if alt_params is None: + alt_params = {"alt_attr": -0.5} + if interaction_params is None: + interaction_params = {"obs_x_alt": 0.8} + + rng = np.random.default_rng(seed) + + # --- Build W (alt×alt) as a libpysal Graph -------------------------- + from .models._spatial_weights import build_knn_graph, resolve_spatial_weights + + if W is None: + coords = rng.standard_normal((n_alts, 2)) + W_graph = build_knn_graph(coords, k=n_neighbors) + W_dense = np.asarray(W_graph.sparse.todense(), dtype=np.float64) + else: + W_graph, _ = resolve_spatial_weights(W, n_alts, row_standardize=True) + W_dense = np.asarray(W_graph.sparse.todense(), dtype=np.float64) + + # --- Choosers and alternatives -------------------------------------- + obs_ids = pd.Index(np.arange(n_obs), name="oid") + obs_feature = rng.standard_normal(n_obs) + choosers = pd.DataFrame({"obs_feature": obs_feature}, index=obs_ids) + + alt_ids = pd.Index(np.arange(n_alts), name="aid") + alt_data = {col: rng.standard_normal(n_alts) for col in alt_params} + alternatives = pd.DataFrame(alt_data, index=alt_ids) + + # --- Interactions (chooser × alternative) -------------------------- + interaction_index = pd.MultiIndex.from_product([obs_ids, alt_ids], names=["oid", "aid"]) + obs_feat_tiled = np.repeat(obs_feature, n_alts) + # Use the first alternative column as the basis for interaction terms + first_alt_col = next(iter(alt_params)) + first_alt_values = alternatives[first_alt_col].to_numpy() + first_alt_tiled = np.tile(first_alt_values, n_obs) + interactions = {} + for col in interaction_params: + interactions[col] = pd.Series( + obs_feat_tiled * first_alt_tiled, index=interaction_index, name=col + ) + + # --- Base utilities: V_base = Zβ + Xγ (n_obs × n_alts) ------------- + V_base = np.zeros((n_obs, n_alts)) + for col, coef in alt_params.items(): + V_base += coef * np.tile(alternatives[col].to_numpy(), (n_obs, 1)) + for col, coef in interaction_params.items(): + V_base += coef * interactions[col].to_numpy().reshape(n_obs, n_alts) + + # --- Spatial filter: V_filtered = (I - ρW)^{-1} V_base^T ------------ + A = np.eye(n_alts) - rho * W_dense + V_filtered = np.linalg.solve(A, V_base.T).T # (n_obs, n_alts) + + # --- Variance normalisation: D = diag((I - ρW)^{-1}) --------------- + Z_mat = np.linalg.inv(A) + D = np.diag(Z_mat) # (n_alts,) + V_star = V_filtered / D[None, :] # normalise each alternative by d_jj + + # --- Add Gumbel noise and simulate choices ------------------------- + gumbel = rng.gumbel(size=(n_obs, n_alts)) + U = V_star + gumbel + choices = U.argmax(axis=1) + choosers = choosers.copy() + choosers["choice"] = choices + + # --- Build ChoiceTable ---------------------------------------------- + true_params = dict(alt_params) + true_params.update(interaction_params) + choice_table = _build_choice_table( + choosers, alternatives, choosers["choice"], matrix_data=interactions + ) + + return SARMNLDataset( + choosers=choosers, + alternatives=alternatives, + interactions=interactions, + true_params=true_params, + true_rho=rho, + W=W_graph, + choice_table=choice_table, + n_obs=n_obs, + n_alts=n_alts, + seed=seed, + ) diff --git a/locpick/models/__init__.pyi b/locpick/models/__init__.pyi index 1f1a6e3..a0e0598 100644 --- a/locpick/models/__init__.pyi +++ b/locpick/models/__init__.pyi @@ -4,6 +4,9 @@ from . import mnl as mnl from . import nested as nested from . import scl as scl from .base import ( + ChoiceModelProtocol as ChoiceModelProtocol, +) +from .choice_model import ( ChoiceModel as ChoiceModel, ) from .mixed import ( diff --git a/locpick/models/_spatial_weights.py b/locpick/models/_spatial_weights.py new file mode 100644 index 0000000..b5a645d --- /dev/null +++ b/locpick/models/_spatial_weights.py @@ -0,0 +1,126 @@ +"""Spatial weights matrix utilities for SAR-MNL models. + +This module provides helpers to resolve spatial weights matrices +(``W``) connecting alternatives (spatial locations). It follows the +convention from the sister package *bayespecon*: ``libpysal.graph.Graph`` +is the canonical type, but ``scipy.sparse`` matrices and dense NumPy +arrays are also accepted for convenience. + +The resolver returns both a row-standardised ``Graph`` (for storage +and return in DGP datasets) and a CSR sparse matrix (for efficient +computation inside JAX kernels). +""" + +from __future__ import annotations + +import numpy as np +import scipy.sparse as sp + + +def resolve_spatial_weights( + W, + n_alts: int, + row_standardize: bool = True, +): + """Resolve alt×alt spatial weights to a libpysal Graph + CSR sparse. + + Accepts a ``libpysal.graph.Graph`` (preferred), ``scipy.sparse`` + matrix, or dense ``np.ndarray``. Returns both a row-standardised + ``Graph`` (for storage/return in DGP datasets) and a CSR sparse + matrix (for efficient computation). + + Parameters + ---------- + W : libpysal.graph.Graph, scipy.sparse, or np.ndarray + J×J spatial weights matrix connecting alternatives (locations). + n_alts : int + Expected number of alternatives (for validation). + row_standardize : bool, default True + If True, row-standardize the weights (rows sum to 1). + + Returns + ------- + W_graph : libpysal.graph.Graph + Row-standardised spatial weights as a libpysal Graph. + W_sparse : scipy.sparse.csr_array + Row-standardised CSR sparse matrix (float64), zero diagonal. + """ + # --- Reject legacy libpysal.weights.W ------------------------------- + if W.__class__.__module__.startswith("libpysal.weights") and not hasattr(W, "sparse"): + raise TypeError( + "Legacy libpysal.weights.W is not supported. " + "Convert via libpysal.graph.Graph.from_W(w) or pass w.sparse." + ) + + # --- Convert to CSR sparse ----------------------------------------- + if hasattr(W, "sparse"): + # libpysal.graph.Graph + W_sparse = sp.csr_array(W.sparse, dtype=np.float64) + elif sp.issparse(W): + W_sparse = sp.csr_array(W, dtype=np.float64) + else: + W_sparse = sp.csr_array(np.asarray(W, dtype=np.float64)) + + # --- Validate shape ------------------------------------------------- + if W_sparse.shape != (n_alts, n_alts): + raise ValueError(f"W shape {W_sparse.shape} does not match n_alts ({n_alts}, {n_alts}).") + + # --- Zero diagonal -------------------------------------------------- + W_sparse.setdiag(0.0) + W_sparse.eliminate_zeros() + + # --- Row-standardize ------------------------------------------------ + if row_standardize: + row_sums = np.asarray(W_sparse.sum(axis=1)).ravel() + row_sums = np.where(row_sums == 0, 1.0, row_sums) + W_sparse = sp.diags(1.0 / row_sums) @ W_sparse + W_sparse = sp.csr_array(W_sparse, dtype=np.float64) + + # --- Convert back to Graph for storage/return ---------------------- + W_graph = _csr_to_graph(W_sparse) + + return W_graph, W_sparse + + +def _csr_to_graph(W_sparse: sp.csr_array): + """Convert a CSR sparse matrix to a libpysal Graph.""" + from libpysal.graph import Graph + + W_coo = W_sparse.tocoo() + return Graph.from_arrays( + focal_ids=W_coo.row.astype(np.int32), + neighbor_ids=W_coo.col.astype(np.int32), + weight=W_coo.data.astype(np.float64), + ) + + +def build_knn_graph( + coords: np.ndarray, + k: int, +): + """Build a k-nearest-neighbor libpysal Graph from coordinates. + + Parameters + ---------- + coords : np.ndarray, shape (n, 2) + Spatial coordinates of alternatives. + k : int + Number of nearest neighbors. + + Returns + ------- + libpysal.graph.Graph + Row-standardised k-NN graph. + """ + import geopandas as gpd + from libpysal.graph import Graph + from shapely.geometry import Point + + n = coords.shape[0] + gdf = gpd.GeoDataFrame( + {"aid": np.arange(n)}, + geometry=[Point(c) for c in coords], + crs="EPSG:4326", + ) + W_graph = Graph.build_knn(gdf, k=k) + return W_graph.transform("r") diff --git a/locpick/models/base.py b/locpick/models/base.py index 79629b6..dcf81e8 100644 --- a/locpick/models/base.py +++ b/locpick/models/base.py @@ -14,12 +14,12 @@ from scipy import stats from scipy.linalg import cho_factor, cho_solve -from locpick._jax.objective import Objective -from locpick._solvers.protocol import Solver, SolverResult, get_solver -from locpick.data.arrays import ChoiceArrays -from locpick.data.choicetable import ChoiceTable -from locpick.data.problem import EstimationProblem -from locpick.results.fit_result import FitResult +from .._jax.objective import Objective +from .._solvers.protocol import Solver, SolverResult, get_solver +from ..data.arrays import ChoiceArrays +from ..data.choicetable import ChoiceTable +from ..data.problem import EstimationProblem +from ..results.fit_result import FitResult # ------------------------------------------------------------------ # Cholesky-based linear-algebra helpers @@ -86,12 +86,11 @@ def _aggregate_per_obs_alt(s: pd.Series, by: str, name: str): @runtime_checkable -class ChoiceModel(Protocol): +class ChoiceModelProtocol(Protocol): """Protocol for discrete choice model classes. - All concrete model classes (``MultinomialLogit``, ``NestedLogit``, - ``MixedLogit``, ``SpatiallyCorrelatedLogit``, - ``MixedSpatiallyCorrelatedLogit``) implement this protocol. + All concrete model classes implement this protocol. + The primary implementation is :class:`locpick.models.choice_model.ChoiceModel`. """ def fit(self, **kwargs) -> FitResult: @@ -263,7 +262,7 @@ def __init__( weights: Optional[Union[str, np.ndarray]] = None, availability: Optional[Union[str, np.ndarray]] = None, ): - from locpick.spec.model_spec import ModelSpec + from ..spec.model_spec import ModelSpec self._solver_options = solver_options or {} self._backend = backend @@ -727,7 +726,7 @@ def _resolve_spatial_graph(self) -> tuple[np.ndarray, list, int]: n_alts : int Number of alternatives (dimension of the graph). """ - from locpick.models._spatial import EdgeStructure, _resolve_spatial_graph + from ._spatial import EdgeStructure, _resolve_spatial_graph omega, allocation, edge_list, n_alts = _resolve_spatial_graph(self._graph_input) self._omega = omega diff --git a/locpick/models/choice_model.py b/locpick/models/choice_model.py new file mode 100644 index 0000000..315b265 --- /dev/null +++ b/locpick/models/choice_model.py @@ -0,0 +1,2138 @@ +"""Unified choice model for location choice estimation. + +This module provides the :class:`ChoiceModel` class, a single composable +model that handles all model configurations via optional feature flags: + +- ``nests`` → Nested logit +- ``random_params`` → Mixed logit +- ``graph`` → Spatially correlated logit (SCL) +- combinations → Nested SCL, Mixed SCL (MSCL), Mixed Nested, Mixed Nested SCL + +The class inherits from :class:`BaseChoiceModel` and :class:`SpatialMixin`, +dispatching to the appropriate JAX builder based on which features are active. +All shared methods (simulate, marginal effects, elasticities, covariance) +live here once, eliminating the duplication across MNL/NestedMNL/MixedMNL/ +MixedNestedMNL. +""" + +from __future__ import annotations + +from typing import Optional, Union + +import numpy as np +import pandas as pd + +from .._jax.objective import Objective +from .._solvers import Solver, SolverResult +from ..data.arrays import ChoiceArrays +from ..data.problem import EstimationProblem +from ..results.fit_result import FitResult +from ._spatial import ( + EdgeStructure, + _resolve_spatial_graph, + naturalize_rho, +) +from .base import ( + BaseChoiceModel, + SpatialMixin, + _compute_fit_statistics, + _compute_null_ll, + _safe_inv, + _sandwich_inv, +) +from .mixed import ParamDistribution, _resolve_draws +from .nested import NestingTree, naturalize_nest_params + + +class ChoiceModel(BaseChoiceModel, SpatialMixin): + r"""Unified discrete choice model for location choice estimation. + + A single composable class that handles all model configurations: + + - **MNL** (default): ``ChoiceModel(ct, formula="cost + time - 1")`` + - **Nested logit**: ``ChoiceModel(ct, formula="...", nests=tree)`` + - **Mixed logit**: ``ChoiceModel(ct, formula="...", random_params={"time": ParamDistribution("normal", "time")})`` + - **SCL** (spatial): ``ChoiceModel(ct, formula="...", graph=g)`` + - **Nested SCL**: ``ChoiceModel(ct, formula="...", nests=tree, graph=g)`` + - **MSCL** (mixed + spatial): ``ChoiceModel(ct, formula="...", random_params=..., graph=g)`` + - **Mixed Nested**: ``ChoiceModel(ct, formula="...", nests=tree, random_params=...)`` + - **Mixed Nested SCL**: ``ChoiceModel(ct, formula="...", nests=tree, random_params=..., graph=g)`` + + The model type is determined by which optional features are present. + + Parameters + ---------- + data : ChoiceTable or EstimationProblem + The choice data to estimate on. + formula : str, optional + A formulaic formula string (e.g., ``"cost + time - 1"``). + Mutually exclusive with ``spec``. + spec : ModelSpec, optional + A ModelSpec object defining formula/scoped-term model structure. + Mutually exclusive with ``formula``. + nests : NestingTree, optional + The nesting structure for nested logit models. + random_params : dict, optional + Mapping of parameter names to :class:`ParamDistribution` objects + for mixed logit models. + graph : libpysal.Graph, scipy.sparse, or np.ndarray, optional + Spatial adjacency graph for SCL models. + n_draws : int, optional + Number of draws for simulated maximum likelihood (mixed logit). + Default 100. + draw_type : str, optional + Type of draws: ``"qmc"`` (default), ``"halton"``, or ``"random"``. + seed : int + Random seed for draw generation. Default 42. + weights : str or array-like, optional + Observation weights. + availability : str or array-like, optional + Alternative availability. + solver : str or Solver, optional + Solver name or instance. Default ``"lbfgs"``. + solver_options : dict, optional + Additional options passed to the solver constructor. + backend : str, optional + Computation backend hint. + estimator : str, optional + SAR estimation method (only relevant when ``lag=True``). + ``"auto"`` (default) selects ``"pml"`` (dense solve) for + n_alts ≤ 2000 and ``"pml_cg"`` (conjugate gradient) for larger + alternative sets. ``"linearized_gmm"`` uses the two-step GMM + estimator (Carrión-Flores et al. 2018) for very large J. + + Examples + -------- + >>> from locpick import ChoiceTable, ChoiceModel + >>> ct = ChoiceTable.from_tables(choosers, alternatives, chosen, sample_size=10) + >>> model = ChoiceModel(ct, formula="cost + time - 1") + >>> result = model.fit() + >>> print(result.summary()) + """ + + def __init__( + self, + data, + formula: Optional[str] = None, + spec=None, + problem: Optional[EstimationProblem] = None, + nests: Optional[NestingTree] = None, + random_params: Optional[dict[str, ParamDistribution]] = None, + graph=None, + lag: bool = False, + n_draws: Optional[int] = None, + draw_type: Optional[str] = None, + seed: int = 42, + weights: Optional[Union[str, np.ndarray]] = None, + availability: Optional[Union[str, np.ndarray]] = None, + solver: Union[str, Solver] = "lbfgs", + solver_options: Optional[dict] = None, + backend: Optional[str] = None, + estimator: str = "auto", + ): + # Handle the legacy `problem` parameter by wrapping it as EstimationProblem + if problem is not None: + data = problem + + super().__init__( + data=data, + formula=formula, + spec=spec, + solver=solver, + solver_options=solver_options, + backend=backend, + weights=weights, + availability=availability, + ) + + # Feature flags + self._nests = nests + self._random_params = random_params + self._graph_input = graph + self._lag = lag # True = SAR spatial lag, False = SCL (default) + self._estimator = estimator # SAR estimator: auto, pml, pml_cg, linearized_gmm + + # Mixed logit settings + self._n_draws = n_draws if n_draws is not None else 100 + self._draw_type = draw_type if draw_type is not None else "sobol" + self._seed = seed + self._draws: Optional[np.ndarray] = None + + # Nest matrix (built in _pre_fit) + self._nest_matrix: Optional[np.ndarray] = None + + # Spatial state (None when graph is not provided) + self._omega = None + self._allocation = None + self._edge_list = None + self._n_alts_graph = None + self._edge_struct = None + self._edge_structs = None + self._edge_data_list = None + + # SAR spatial state (None when lag=True is not used) + self._W_sparse = None # CSR sparse for SAR kernels + # Random parameter state (built in _pre_fit) + self._random_col_indices: Optional[list[int]] = None + self._random_distributions: Optional[list[str]] = None + self._random_param_names: Optional[list[str]] = None + self._k_fixed: Optional[int] = None + self._k_random: Optional[int] = None + self._fixed_names: Optional[list[str]] = None + self._full_param_names: Optional[list[str]] = None + + # ------------------------------------------------------------------ + # Properties + # ------------------------------------------------------------------ + + @property + def _is_spatial(self) -> bool: + return self._graph_input is not None + + @property + def _is_spatial_lag(self) -> bool: + """True when SAR (lag=True) spatial model is active.""" + return self._is_spatial and self._lag + + @property + def _is_spatial_scl(self) -> bool: + """True when SCL (lag=False) spatial model is active.""" + return self._is_spatial and not self._lag + + @property + def _is_nested(self) -> bool: + return self._nests is not None + + @property + def _is_mixed(self) -> bool: + return self._random_params is not None and len(self._random_params) > 0 + + @property + def model_type(self) -> str: + """Human-readable model type string.""" + parts = [] + if self._is_mixed: + parts.append("Mixed") + if self._is_nested: + parts.append("Nested") + if self._is_spatial_lag: + parts.append("Spatial Autoregressive") + elif self._is_spatial_scl: + parts.append("Spatially Correlated") + if not parts: + return "Multinomial Logit" + if len(parts) == 1 and parts[0] == "Spatially Correlated": + return "Spatially Correlated Logit" + if len(parts) == 1 and parts[0] == "Spatial Autoregressive": + return "Spatial Autoregressive Logit" + if len(parts) == 1 and parts[0] == "Nested": + return "Nested Logit" + if len(parts) == 1 and parts[0] == "Mixed": + return "Mixed Logit" + return " ".join(parts) + " Logit" + + # ------------------------------------------------------------------ + # Pre-fit: build model-specific data structures + # ------------------------------------------------------------------ + + def fit(self, **kwargs) -> FitResult: + """Estimate the model and return results. + + Dispatches to the PML estimator (JAX autodiff) or the + linearized GMM estimator based on the ``estimator`` setting. + """ + if self._is_spatial_lag and self._estimator == "linearized_gmm": + return self._fit_linearized_gmm() + return super().fit(**kwargs) + + def _fit_linearized_gmm(self) -> FitResult: + """Two-step linearized GMM estimation (Carrión-Flores et al. 2018).""" + arrays = self._get_arrays() + self._arrays = arrays + self._pre_fit(arrays) + + from .._kernels.sar_mnl_numpy import fit_linearized_gmm + + result_dict = fit_linearized_gmm(arrays, self._W_sparse) + + beta = result_dict["beta"] + rho = result_dict["rho"] + se = result_dict["se"] + ll = result_dict["log_likelihood"] + + utility_param_names = list(arrays.param_names) + k = len(utility_param_names) + display_values = np.concatenate([beta, [rho]]) + display_names = utility_param_names + ["rho"] + model_type = "Spatial Autoregressive Logit (Linearized GMM)" + n_params = len(display_values) + + std_errors = se[: k + 1] + + coefficients = pd.Series(display_values, index=display_names, name="coefficient") + std_err_series = pd.Series(std_errors, index=display_names, name="std_error") + ll_null = _compute_null_ll(arrays) + + stats = _compute_fit_statistics( + ll=ll, + ll_null=ll_null, + n_obs=arrays.n_obs, + n_params=n_params, + n_alts=arrays.n_alts, + coefficients=coefficients, + std_errors=std_err_series, + model_type=model_type, + solver_name="linearized_gmm", + solver_result_raw=result_dict, + ) + + self._result = FitResult(spec=self._spec, **stats) + self._clear_caches() + return self._result + + def _pre_fit(self, arrays: ChoiceArrays) -> None: + """Build nest matrix, random parameter structure, and spatial graph.""" + # Build nest matrix if nests are provided + if self._is_nested: + alt_ids = list(range(arrays.n_alts)) + self._nest_matrix = self._nests.build_nest_matrix(alt_ids) + + # Build random parameter structure if random_params is provided + if self._is_mixed: + self._prepare_random_params(arrays) + + # Resolve spatial graph if graph is provided + if self._is_spatial: + if self._is_spatial_lag: + # SAR: resolve W via _spatial_weights resolver + from ._spatial_weights import resolve_spatial_weights + + self._W_sparse = resolve_spatial_weights( + self._graph_input, arrays.n_alts, row_standardize=True + )[1] # get the CSR sparse + else: + # SCL: resolve via edge structure (existing behavior) + self._resolve_spatial_graph() + self._validate_graph_size(arrays) + + # Build per-nest edge structures for nested spatial models + if self._is_nested: + self._build_per_nest_edges(arrays) + + def _prepare_random_params(self, arrays: ChoiceArrays) -> None: + """Identify random parameter columns and generate draws.""" + param_names = list(arrays.param_names) + random_param_names: list[str] = [] + random_distributions: list[str] = [] + random_col_indices: list[int] = [] + + for name, dist in self._random_params.items(): + if name not in param_names: + raise ValueError( + f"Random parameter '{name}' not found in design matrix. " + f"Available parameters: {param_names}" + ) + random_param_names.append(name) + random_distributions.append(dist.distribution) + random_col_indices.append(param_names.index(name)) + + k_fixed = len(param_names) - len(random_param_names) + k_random = len(random_param_names) + + # Generate draws + draws = _resolve_draws( + self._draw_type, arrays.n_obs, self._n_draws, k_random, seed=self._seed + ) + + fixed_names = [n for i, n in enumerate(param_names) if i not in random_col_indices] + full_param_names = ( + fixed_names + + [f"mean_{n}" for n in random_param_names] + + [f"sd_{n}" for n in random_param_names] + ) + + self._random_param_names = random_param_names + self._random_distributions = random_distributions + self._random_col_indices = random_col_indices + self._k_fixed = k_fixed + self._k_random = k_random + self._fixed_names = fixed_names + self._full_param_names = full_param_names + self._draws = draws + + def _build_per_nest_edges(self, arrays: ChoiceArrays) -> None: + """Build per-nest EdgeStructure / EdgeDataJAX from the global graph.""" + from .._jax.data import EdgeDataJAX + + n_nests = self._nests.n_nests + self._edge_structs = [] + self._edge_data_list = [] + for m in range(n_nests): + nest_alts = np.where(self._nest_matrix[:, m] > 0)[0] + n_nest_alts = len(nest_alts) + if n_nest_alts == 0: + self._edge_structs.append(None) + self._edge_data_list.append(None) + continue + nest_adj = self._omega[np.ix_(nest_alts, nest_alts)] + _, nest_alloc, nest_edges, _ = _resolve_spatial_graph(nest_adj) + edge_struct = EdgeStructure(nest_edges, n_nest_alts, nest_alloc) + self._edge_structs.append(edge_struct) + self._edge_data_list.append(EdgeDataJAX.from_edge_structure(edge_struct)) + + # ------------------------------------------------------------------ + # Solver inputs + # ------------------------------------------------------------------ + + def _get_solver_inputs(self, arrays: ChoiceArrays): + """Get initial values, param names, bounds, and fixed mask. + + Parameter layout depends on active features: + + - MNL: ``[beta]`` + - SCL: ``[beta, alpha_rho]`` + - Nested: ``[beta, alpha_nest]`` + - Nested SCL: ``[beta, alpha_rho_1..M, alpha_lambda_1..M]`` + - Mixed: ``[beta_fixed, mean_*, sd_*]`` + - MSCL: ``[beta_fixed, alpha_rho, mean_*, sd_*]`` + - Mixed Nested: ``[beta_fixed, alpha_nest, mean_*, sd_*]`` + - Mixed Nested SCL: ``[beta_fixed, alpha_rho_1..M, alpha_lambda_1..M, mean_*, sd_*]`` + """ + k = arrays.design_matrix.shape[1] + param_names_all = list(arrays.param_names) + + # --- Pure MNL --- + if not self._is_nested and not self._is_mixed: + # Use problem's initial values / fixed_mask when available + if self._problem is not None: + x0_base = self._problem.initial_values + bounds = self._problem.bounds + fixed_mask = self._problem.fixed_mask + else: + x0_base = np.zeros(k) + bounds = None + fixed_mask = None + names_base = param_names_all + if self._is_spatial: + # SCL or SAR: [beta, alpha_rho] + x0 = np.concatenate([x0_base, np.zeros(1)]) + names = list(names_base) + ["alpha_rho"] + else: + x0 = x0_base + names = names_base + return x0, names, bounds, fixed_mask + + # --- Nested (no random) --- + if self._is_nested and not self._is_mixed: + n_nests = self._nests.n_nests + if self._is_spatial_lag: + # SAR + Nested: [beta, alpha_rho, alpha_lambda_1..M] + x0 = np.concatenate([np.zeros(k), np.zeros(1), self._nests.initial_alphas()]) + names = ( + param_names_all + + ["alpha_rho"] + + [f"alpha_lambda_{name}" for name in self._nests.nest_names] + ) + elif self._is_spatial_scl: + # Nested SCL: [beta, alpha_rho_1..M, alpha_lambda_1..M] + x0 = np.concatenate( + [ + np.zeros(k), + np.zeros(n_nests), + self._nests.initial_alphas(), + ] + ) + names = ( + param_names_all + + [f"alpha_rho_{name}" for name in self._nests.nest_names] + + [f"alpha_lambda_{name}" for name in self._nests.nest_names] + ) + else: + # Nested: [beta, alpha_nest] + x0 = np.concatenate([np.zeros(k), self._nests.initial_alphas()]) + names = param_names_all + [f"nest_{name}" for name in self._nests.nest_names] + return x0, names, None, None + + # --- Mixed (no nests) --- + if self._is_mixed and not self._is_nested: + k_fixed = self._k_fixed + k_random = self._k_random + if self._is_spatial: + # SAR + Mixed or MSCL: [beta_fixed, alpha_rho, mean_*, sd_*] + x0 = np.concatenate( + [ + np.zeros(k_fixed), + np.zeros(1), + np.zeros(k_random), + np.full(k_random, 0.1), + ] + ) + names = ( + list(self._fixed_names) + + ["rho"] + + [f"mean_{n}" for n in self._random_param_names] + + [f"sd_{n}" for n in self._random_param_names] + ) + else: + # Mixed: [beta_fixed, mean_*, sd_*] + x0 = np.concatenate( + [ + np.zeros(k_fixed), + np.zeros(k_random), + np.full(k_random, 0.1), + ] + ) + names = list(self._full_param_names) + return x0, names, None, None + + # --- Mixed Nested --- + if self._is_nested and self._is_mixed: + k_fixed = self._k_fixed + k_random = self._k_random + n_nests = self._nests.n_nests + fixed_param_names = [ + name for name in param_names_all if name not in self._random_params + ] + if self._is_spatial_lag: + # SAR + Mixed + Nested: [beta_fixed, alpha_rho, alpha_lambda_1..M, mean_*, sd_*] + x0 = np.concatenate( + [ + np.zeros(k_fixed), + np.zeros(1), + self._nests.initial_alphas(), + np.zeros(k_random), + np.full(k_random, 0.1), + ] + ) + names = ( + fixed_param_names + + ["alpha_rho"] + + [f"alpha_lambda_{name}" for name in self._nests.nest_names] + + [f"mean_{name}" for name in self._random_param_names] + + [f"sd_{name}" for name in self._random_param_names] + ) + elif self._is_spatial_scl: + # Mixed Nested SCL: [beta_fixed, alpha_rho_1..M, alpha_lambda_1..M, mean_*, sd_*] + x0 = np.concatenate( + [ + np.zeros(k_fixed), + np.zeros(n_nests), + self._nests.initial_alphas(), + np.zeros(k_random), + np.full(k_random, 0.1), + ] + ) + names = ( + fixed_param_names + + [f"alpha_rho_{name}" for name in self._nests.nest_names] + + [f"alpha_lambda_{name}" for name in self._nests.nest_names] + + [f"mean_{name}" for name in self._random_param_names] + + [f"sd_{name}" for name in self._random_param_names] + ) + else: + # Mixed Nested: [beta_fixed, alpha_nest, mean_*, sd_*] + x0 = np.concatenate( + [ + np.zeros(k_fixed), + self._nests.initial_alphas(), + np.zeros(k_random), + np.full(k_random, 0.1), + ] + ) + names = ( + fixed_param_names + + [f"lambda_{name}" for name in self._nests.nest_names] + + [f"mean_{name}" for name in self._random_param_names] + + [f"sd_{name}" for name in self._random_param_names] + ) + return x0, names, None, None + + # Fallback (should not reach here) + return np.zeros(k), param_names_all, None, None + + # ------------------------------------------------------------------ + # Objective construction + # ------------------------------------------------------------------ + + def _build_objective(self, arrays: ChoiceArrays) -> Objective: + """Build optimization objective based on active features.""" + # Pure MNL / SCL / SAR + if not self._is_nested and not self._is_mixed: + if self._is_spatial_lag: + from .._jax.sar_kernels import build_sar_mnl_objective + + # Auto-select estimator + if self._estimator == "auto": + if arrays.n_alts <= 2000: + self._estimator = "pml" + else: + self._estimator = "pml_cg" + use_cg = self._estimator == "pml_cg" + return build_sar_mnl_objective(arrays, self._W_sparse, use_cg=use_cg) + elif self._is_spatial_scl: + from .._jax.builders import build_scl_objective + + return build_scl_objective( + arrays, self._edge_struct, self._allocation, self._edge_list + ) + from .._jax.builders import build_mnl_objective + + return build_mnl_objective(arrays) + + # Nested (no random) + if self._is_nested and not self._is_mixed: + if self._is_spatial_lag: + from .._jax.sar_kernels import build_sar_nested_objective + + return build_sar_nested_objective(arrays, self._W_sparse, self._nest_matrix) + elif self._is_spatial_scl: + from .._jax.builders import build_nested_scl_objective + + return build_nested_scl_objective(arrays, self._nest_matrix, self._edge_data_list) + from .._jax.builders import build_nested_objective + + return build_nested_objective(arrays, self._nest_matrix) + + # Mixed (no nests) + if self._is_mixed and not self._is_nested: + if self._is_spatial_lag: + from .._jax.sar_kernels import build_sar_mixed_objective + + return build_sar_mixed_objective( + arrays, + self._W_sparse, + self._random_col_indices, + self._random_distributions, + self._draws, + ) + elif self._is_spatial_scl: + from .._jax.builders import build_mscl_objective + + return build_mscl_objective( + arrays, + self._edge_struct, + self._allocation, + self._edge_list, + self._random_col_indices, + self._random_distributions, + self._draws, + ) + from .._jax.builders import build_mixed_logit_objective + + return build_mixed_logit_objective( + arrays, + random_col_indices=self._random_col_indices, + random_distributions=self._random_distributions, + draws=self._draws, + ) + + # Mixed Nested + if self._is_nested and self._is_mixed: + if self._is_spatial_lag: + from .._jax.sar_kernels import build_sar_mixed_nested_objective + + return build_sar_mixed_nested_objective( + arrays, + self._W_sparse, + self._nest_matrix, + self._random_col_indices, + self._random_distributions, + self._draws, + ) + elif self._is_spatial_scl: + from .._jax.builders import build_mnscl_objective + + return build_mnscl_objective( + arrays, + self._nest_matrix, + self._edge_data_list, + self._random_col_indices, + self._random_distributions, + self._draws, + ) + from .._jax.builders import build_mixed_nested_objective + + return build_mixed_nested_objective( + arrays, + self._nest_matrix, + self._random_col_indices, + self._random_distributions, + self._draws, + ) + + raise RuntimeError("Unknown model configuration") + + # ------------------------------------------------------------------ + # Fit result construction + # ------------------------------------------------------------------ + + def _build_fit_result( + self, + solver_result: SolverResult, + arrays: ChoiceArrays, + ) -> FitResult: + """Build a FitResult from solver output.""" + all_params = solver_result.coefficients + k = arrays.design_matrix.shape[1] + param_names_all = list(arrays.param_names) + + # --- Pure MNL / SCL / SAR --- + if not self._is_nested and not self._is_mixed: + if self._is_spatial_lag: + # SAR: Layout [beta_1..k, alpha_rho], rho = tanh(alpha_rho) + beta = all_params[:k] + alpha_rho = all_params[k] + rho = np.tanh(alpha_rho) + display_values = np.concatenate([beta, [rho]]) + display_names = param_names_all + ["rho"] + elif self._is_spatial_scl: + # SCL: Layout [beta_1..k, alpha_rho], rho = sigmoid(alpha_rho) + beta = all_params[:k] + alpha_rho = all_params[k] + rho = naturalize_rho(alpha_rho) + display_values = np.concatenate([beta, [rho]]) + display_names = param_names_all + ["rho"] + else: + display_values = all_params + display_names = param_names_all + + std_errors = self._compute_se(all_params, arrays, display_values, display_names) + return self._make_fit_result( + solver_result, arrays, display_values, display_names, std_errors + ) + + # --- Nested (no random) --- + if self._is_nested and not self._is_mixed: + n_nests = self._nests.n_nests + if self._is_spatial_lag: + # SAR + Nested: Layout [beta, alpha_rho, alpha_lambda_1..M] + beta = all_params[:k] + alpha_rho = all_params[k] + alpha_lambda = all_params[k + 1 : k + 1 + n_nests] + rho = np.tanh(alpha_rho) + lambdas = naturalize_nest_params(alpha_lambda) + display_values = np.concatenate([beta, [rho], lambdas]) + display_names = ( + param_names_all + + ["rho"] + + [f"lambda_{name}" for name in self._nests.nest_names] + ) + std_errors = self._compute_se_sar_nested( + all_params, arrays, k, n_nests, rho, lambdas + ) + elif self._is_spatial_scl: + # SCL + Nested: Layout [beta, alpha_rho_1..M, alpha_lambda_1..M] + beta = all_params[:k] + alpha_rho = all_params[k : k + n_nests] + alpha_lambda = all_params[k + n_nests : k + 2 * n_nests] + rhos = naturalize_rho(alpha_rho) + lambdas = naturalize_nest_params(alpha_lambda) + display_values = np.concatenate([beta, rhos, lambdas]) + display_names = ( + param_names_all + + [f"rho_{name}" for name in self._nests.nest_names] + + [f"lambda_{name}" for name in self._nests.nest_names] + ) + std_errors = self._compute_se_nested_scl( + all_params, arrays, k, n_nests, rhos, lambdas + ) + else: + # Layout: [beta, alpha_nest] + beta = all_params[:k] + alpha = all_params[k:] + lambdas = naturalize_nest_params(alpha) + display_values = np.concatenate([beta, lambdas]) + display_names = param_names_all + [ + f"lambda_{name}" for name in self._nests.nest_names + ] + std_errors = self._compute_se_nested(all_params, arrays, k, lambdas) + + return self._make_fit_result( + solver_result, arrays, display_values, display_names, std_errors + ) + + # --- Mixed (no nests) --- + if self._is_mixed and not self._is_nested: + k_fixed = self._k_fixed + k_random = self._k_random + if self._is_spatial_lag: + # SAR + Mixed: Layout [beta_fixed, alpha_rho, mean_*, sd_*] + beta_fixed = all_params[:k_fixed] + alpha_rho = all_params[k_fixed] + rho = np.tanh(alpha_rho) + beta_random_means = all_params[k_fixed + 1 : k_fixed + 1 + k_random] + beta_random_spreads = all_params[k_fixed + 1 + k_random :] + display_values = np.concatenate( + [beta_fixed, [rho], beta_random_means, beta_random_spreads] + ) + display_names = ( + list(self._fixed_names) + + ["rho"] + + [f"mean_{n}" for n in self._random_param_names] + + [f"sd_{n}" for n in self._random_param_names] + ) + std_errors = self._compute_se_sar_mixed(all_params, arrays, k_fixed, rho) + elif self._is_spatial_scl: + # SCL + Mixed (MSCL): Layout [beta_fixed, alpha_rho, mean_*, sd_*] + beta_fixed = all_params[:k_fixed] + alpha_rho = all_params[k_fixed] + rho = naturalize_rho(alpha_rho) + beta_random_means = all_params[k_fixed + 1 : k_fixed + 1 + k_random] + beta_random_spreads = all_params[k_fixed + 1 + k_random :] + display_values = np.concatenate( + [beta_fixed, [rho], beta_random_means, beta_random_spreads] + ) + display_names = ( + list(self._fixed_names) + + ["rho"] + + [f"mean_{n}" for n in self._random_param_names] + + [f"sd_{n}" for n in self._random_param_names] + ) + std_errors = self._compute_se_mscl(all_params, arrays, k_fixed, rho) + else: + # Layout: [beta_fixed, mean_*, sd_*] + display_values = all_params + display_names = list(self._full_param_names) + std_errors = self._compute_se_simple(all_params, arrays) + + return self._make_fit_result( + solver_result, arrays, display_values, display_names, std_errors + ) + + # --- Mixed Nested --- + if self._is_nested and self._is_mixed: + k_fixed = self._k_fixed + k_random = self._k_random + n_nests = self._nests.n_nests + fixed_param_names = [ + name for name in param_names_all if name not in self._random_params + ] + + if self._is_spatial_lag: + # SAR + Mixed + Nested: Layout [beta_fixed, alpha_rho, alpha_lambda_1..M, mean_*, sd_*] + beta_fixed = all_params[:k_fixed] + alpha_rho = all_params[k_fixed] + alpha_lambda = all_params[k_fixed + 1 : k_fixed + 1 + n_nests] + beta_random_means = all_params[ + k_fixed + 1 + n_nests : k_fixed + 1 + n_nests + k_random + ] + beta_random_spreads = all_params[k_fixed + 1 + n_nests + k_random :] + rho = np.tanh(alpha_rho) + lambdas = naturalize_nest_params(alpha_lambda) + display_values = np.concatenate( + [beta_fixed, [rho], lambdas, beta_random_means, np.abs(beta_random_spreads)] + ) + display_names = ( + fixed_param_names + + ["rho"] + + [f"lambda_{name}" for name in self._nests.nest_names] + + [f"mean_{name}" for name in self._random_param_names] + + [f"sd_{name}" for name in self._random_param_names] + ) + std_errors = self._compute_se_sar_mixed_nested( + all_params, arrays, k_fixed, n_nests, rho, lambdas + ) + elif self._is_spatial_scl: + # SCL + Mixed + Nested: Layout [beta_fixed, alpha_rho_1..M, alpha_lambda_1..M, mean_*, sd_*] + beta_fixed = all_params[:k_fixed] + alpha_rho = all_params[k_fixed : k_fixed + n_nests] + alpha_lambda = all_params[k_fixed + n_nests : k_fixed + 2 * n_nests] + beta_random_means = all_params[ + k_fixed + 2 * n_nests : k_fixed + 2 * n_nests + k_random + ] + beta_random_spreads = all_params[k_fixed + 2 * n_nests + k_random :] + rhos = naturalize_rho(alpha_rho) + lambdas = naturalize_nest_params(alpha_lambda) + display_values = np.concatenate( + [beta_fixed, rhos, lambdas, beta_random_means, np.abs(beta_random_spreads)] + ) + display_names = ( + fixed_param_names + + [f"rho_{name}" for name in self._nests.nest_names] + + [f"lambda_{name}" for name in self._nests.nest_names] + + [f"mean_{name}" for name in self._random_param_names] + + [f"sd_{name}" for name in self._random_param_names] + ) + std_errors = self._compute_se_mixed_nested_scl( + all_params, arrays, k_fixed, n_nests, rhos, lambdas + ) + else: + # Layout: [beta_fixed, alpha_nest, mean_*, sd_*] + beta_fixed = all_params[:k_fixed] + alpha_nest = all_params[k_fixed : k_fixed + n_nests] + beta_random_means = all_params[k_fixed + n_nests : k_fixed + n_nests + k_random] + beta_random_spreads = all_params[k_fixed + n_nests + k_random :] + lambdas = naturalize_nest_params(alpha_nest) + display_values = np.concatenate( + [beta_fixed, lambdas, beta_random_means, np.abs(beta_random_spreads)] + ) + display_names = ( + fixed_param_names + + [f"lambda_{name}" for name in self._nests.nest_names] + + [f"mean_{name}" for name in self._random_param_names] + + [f"sd_{name}" for name in self._random_param_names] + ) + std_errors = self._compute_se_mixed_nested( + all_params, arrays, k_fixed, n_nests, lambdas + ) + + return self._make_fit_result( + solver_result, arrays, display_values, display_names, std_errors + ) + + raise RuntimeError("Unknown model configuration") + + # ------------------------------------------------------------------ + # Standard error computation helpers + # ------------------------------------------------------------------ + + def _compute_se_simple(self, all_params, arrays): + """Compute SEs for models without parameter transforms.""" + std_errors = np.full(len(all_params), np.nan) + try: + hess = self._compute_hessian(all_params) + std_errors = self._compute_std_errors_from_hessian(hess) + except Exception: + if self._result is not None and self._result.solver_result: + hess_inv = self._get_hessian_inverse() + if hess_inv is not None: + std_errors = np.sqrt(np.maximum(np.diag(hess_inv), 0)) + std_errors[std_errors == 0] = np.nan + return std_errors + + def _compute_se(self, all_params, arrays, display_values, display_names): + """Compute SEs for MNL/SCL models.""" + k = arrays.design_matrix.shape[1] + n_params = len(display_values) + std_errors = np.full(n_params, np.nan) + try: + hess = self._compute_hessian(all_params) + se_unconstrained = self._compute_std_errors_from_hessian(hess) + if self._is_spatial_lag: + # SAR delta method: SE(rho) = (1 - rho^2) * SE(alpha_rho) + rho = display_values[k] + se_rho = (1.0 - rho**2) * se_unconstrained[k] + std_errors = np.concatenate([se_unconstrained[:k], [se_rho]]) + elif self._is_spatial_scl: + # SCL delta method: SE(rho) = rho*(1-rho)*SE(alpha_rho) + rho = display_values[k] + se_rho = rho * (1.0 - rho) * se_unconstrained[k] + std_errors = np.concatenate([se_unconstrained[:k], [se_rho]]) + else: + std_errors = se_unconstrained + except Exception: + hess_inv = self._get_hessian_inverse() + if hess_inv is not None: + se = np.sqrt(np.maximum(np.diag(hess_inv), 0)) + se[se == 0] = np.nan + if self._is_spatial_lag: + rho = display_values[k] + se_rho = (1.0 - rho**2) * se[k] + std_errors = np.concatenate([se[:k], [se_rho]]) + elif self._is_spatial_scl: + rho = display_values[k] + se_rho = rho * (1.0 - rho) * se[k] + std_errors = np.concatenate([se[:k], [se_rho]]) + else: + std_errors = se + return std_errors + + def _compute_se_nested(self, all_params, arrays, k, lambdas): + """Compute SEs for nested logit (non-spatial).""" + n_nests = len(lambdas) + std_errors = np.full(k + n_nests, np.nan) + try: + hess = self._compute_hessian(all_params) + se_alpha = self._compute_std_errors_from_hessian(hess) + se_lambda = lambdas * (1.0 - lambdas) * se_alpha[k:] + std_errors = np.concatenate([se_alpha[:k], se_lambda]) + except Exception: + hess_inv = self._get_hessian_inverse() + if hess_inv is not None: + se_alpha = np.sqrt(np.maximum(np.diag(hess_inv), 0)) + se_alpha[se_alpha == 0] = np.nan + se_lambda = lambdas * (1.0 - lambdas) * se_alpha[k:] + std_errors = np.concatenate([se_alpha[:k], se_lambda]) + return std_errors + + def _compute_se_nested_scl(self, all_params, arrays, k, n_nests, rhos, lambdas): + """Compute SEs for nested SCL.""" + std_errors = np.full(k + 2 * n_nests, np.nan) + try: + hess = self._compute_hessian(all_params) + se_alpha = self._compute_std_errors_from_hessian(hess) + se_rho = rhos * (1.0 - rhos) * se_alpha[k : k + n_nests] + se_lambda = lambdas * (1.0 - lambdas) * se_alpha[k + n_nests : k + 2 * n_nests] + std_errors = np.concatenate([se_alpha[:k], se_rho, se_lambda]) + except Exception: + hess_inv = self._get_hessian_inverse() + if hess_inv is not None: + se_alpha = np.sqrt(np.maximum(np.diag(hess_inv), 0)) + se_alpha[se_alpha == 0] = np.nan + se_rho = rhos * (1.0 - rhos) * se_alpha[k : k + n_nests] + se_lambda = lambdas * (1.0 - lambdas) * se_alpha[k + n_nests : k + 2 * n_nests] + std_errors = np.concatenate([se_alpha[:k], se_rho, se_lambda]) + return std_errors + + def _compute_se_mscl(self, all_params, arrays, k_fixed, rho): + """Compute SEs for MSCL.""" + n_params = len(all_params) + std_errors = np.full(n_params, np.nan) + try: + hess = self._compute_hessian(all_params) + se_alpha = self._compute_std_errors_from_hessian(hess) + se_rho = float(rho * (1.0 - rho) * se_alpha[k_fixed]) + std_errors = np.concatenate([se_alpha[:k_fixed], [se_rho], se_alpha[k_fixed + 1 :]]) + except Exception: + hess_inv = self._get_hessian_inverse() + if hess_inv is not None: + se_alpha = np.sqrt(np.maximum(np.diag(hess_inv), 0)) + se_alpha[se_alpha == 0] = np.nan + se_rho = float(rho * (1.0 - rho) * se_alpha[k_fixed]) + std_errors = np.concatenate( + [se_alpha[:k_fixed], [se_rho], se_alpha[k_fixed + 1 :]] + ) + return std_errors + + def _compute_se_mixed_nested(self, all_params, arrays, k_fixed, n_nests, lambdas): + """Compute SEs for mixed nested logit (non-spatial).""" + n_params = len(all_params) + std_errors = np.full(n_params, np.nan) + try: + hess = self._compute_hessian(all_params) + se_raw = self._compute_std_errors_from_hessian(hess) + se_lambda = lambdas * (1.0 - lambdas) * se_raw[k_fixed : k_fixed + n_nests] + std_errors = np.concatenate([se_raw[:k_fixed], se_lambda, se_raw[k_fixed + n_nests :]]) + except Exception: + hess_inv = self._get_hessian_inverse() + if hess_inv is not None: + se_raw = np.sqrt(np.maximum(np.diag(hess_inv), 0)) + se_raw[se_raw == 0] = np.nan + se_lambda = lambdas * (1.0 - lambdas) * se_raw[k_fixed : k_fixed + n_nests] + std_errors = np.concatenate( + [se_raw[:k_fixed], se_lambda, se_raw[k_fixed + n_nests :]] + ) + return std_errors + + def _compute_se_mixed_nested_scl(self, all_params, arrays, k_fixed, n_nests, rhos, lambdas): + """Compute SEs for mixed nested SCL.""" + n_params = len(all_params) + std_errors = np.full(n_params, np.nan) + try: + hess = self._compute_hessian(all_params) + se_raw = self._compute_std_errors_from_hessian(hess) + se_rho = rhos * (1.0 - rhos) * se_raw[k_fixed : k_fixed + n_nests] + se_lambda = ( + lambdas * (1.0 - lambdas) * se_raw[k_fixed + n_nests : k_fixed + 2 * n_nests] + ) + std_errors = np.concatenate( + [ + se_raw[:k_fixed], + se_rho, + se_lambda, + se_raw[k_fixed + 2 * n_nests :], + ] + ) + except Exception: + hess_inv = self._get_hessian_inverse() + if hess_inv is not None: + se_raw = np.sqrt(np.maximum(np.diag(hess_inv), 0)) + se_raw[se_raw == 0] = np.nan + se_rho = rhos * (1.0 - rhos) * se_raw[k_fixed : k_fixed + n_nests] + se_lambda = ( + lambdas * (1.0 - lambdas) * se_raw[k_fixed + n_nests : k_fixed + 2 * n_nests] + ) + std_errors = np.concatenate( + [ + se_raw[:k_fixed], + se_rho, + se_lambda, + se_raw[k_fixed + 2 * n_nests :], + ] + ) + return std_errors + + def _compute_se_sar_nested(self, all_params, arrays, k, n_nests, rho, lambdas): + """Compute SEs for SAR + Nested. Layout: [beta, alpha_rho, alpha_lambda_1..M].""" + n_params = len(all_params) + std_errors = np.full(n_params, np.nan) + try: + hess = self._compute_hessian(all_params) + se_raw = self._compute_std_errors_from_hessian(hess) + # SAR: rho = tanh(alpha_rho), SE(rho) = (1 - rho^2) * SE(alpha_rho) + se_rho = (1.0 - rho**2) * se_raw[k] + # Lambda: sigmoid, SE(lambda) = lambda*(1-lambda)*SE(alpha_lambda) + se_lambda = lambdas * (1.0 - lambdas) * se_raw[k + 1 : k + 1 + n_nests] + std_errors = np.concatenate([se_raw[:k], [se_rho], se_lambda]) + except Exception: + hess_inv = self._get_hessian_inverse() + if hess_inv is not None: + se_raw = np.sqrt(np.maximum(np.diag(hess_inv), 0)) + se_raw[se_raw == 0] = np.nan + se_rho = (1.0 - rho**2) * se_raw[k] + se_lambda = lambdas * (1.0 - lambdas) * se_raw[k + 1 : k + 1 + n_nests] + std_errors = np.concatenate([se_raw[:k], [se_rho], se_lambda]) + return std_errors + + def _compute_se_sar_mixed(self, all_params, arrays, k_fixed, rho): + """Compute SEs for SAR + Mixed. Layout: [beta_fixed, alpha_rho, mean_*, sd_*].""" + n_params = len(all_params) + std_errors = np.full(n_params, np.nan) + try: + hess = self._compute_hessian(all_params) + se_raw = self._compute_std_errors_from_hessian(hess) + # SAR: rho = tanh(alpha_rho), SE(rho) = (1 - rho^2) * SE(alpha_rho) + se_rho = (1.0 - rho**2) * se_raw[k_fixed] + std_errors = np.concatenate([se_raw[:k_fixed], [se_rho], se_raw[k_fixed + 1 :]]) + except Exception: + hess_inv = self._get_hessian_inverse() + if hess_inv is not None: + se_raw = np.sqrt(np.maximum(np.diag(hess_inv), 0)) + se_raw[se_raw == 0] = np.nan + se_rho = (1.0 - rho**2) * se_raw[k_fixed] + std_errors = np.concatenate([se_raw[:k_fixed], [se_rho], se_raw[k_fixed + 1 :]]) + return std_errors + + def _compute_se_sar_mixed_nested(self, all_params, arrays, k_fixed, n_nests, rho, lambdas): + """Compute SEs for SAR + Mixed + Nested. + Layout: [beta_fixed, alpha_rho, alpha_lambda_1..M, mean_*, sd_*].""" + n_params = len(all_params) + std_errors = np.full(n_params, np.nan) + try: + hess = self._compute_hessian(all_params) + se_raw = self._compute_std_errors_from_hessian(hess) + se_rho = (1.0 - rho**2) * se_raw[k_fixed] + se_lambda = lambdas * (1.0 - lambdas) * se_raw[k_fixed + 1 : k_fixed + 1 + n_nests] + std_errors = np.concatenate( + [se_raw[:k_fixed], [se_rho], se_lambda, se_raw[k_fixed + 1 + n_nests :]] + ) + except Exception: + hess_inv = self._get_hessian_inverse() + if hess_inv is not None: + se_raw = np.sqrt(np.maximum(np.diag(hess_inv), 0)) + se_raw[se_raw == 0] = np.nan + se_rho = (1.0 - rho**2) * se_raw[k_fixed] + se_lambda = lambdas * (1.0 - lambdas) * se_raw[k_fixed + 1 : k_fixed + 1 + n_nests] + std_errors = np.concatenate( + [se_raw[:k_fixed], [se_rho], se_lambda, se_raw[k_fixed + 1 + n_nests :]] + ) + return std_errors + + def _make_fit_result( + self, + solver_result: SolverResult, + arrays: ChoiceArrays, + display_values: np.ndarray, + display_names: list[str], + std_errors: np.ndarray, + ) -> FitResult: + """Build a FitResult using the shared helper.""" + coefficients = pd.Series(display_values, index=display_names, name="coefficient") + std_err_series = pd.Series(std_errors, index=display_names, name="std_error") + ll = solver_result.log_likelihood + ll_null = _compute_null_ll(arrays) + n_params = len(display_values) + + stats = _compute_fit_statistics( + ll=ll, + ll_null=ll_null, + n_obs=arrays.n_obs, + n_params=n_params, + n_alts=arrays.n_alts, + coefficients=coefficients, + std_errors=std_err_series, + model_type=self.model_type, + solver_name=solver_result.solver_name, + solver_result_raw=solver_result.raw, + ) + + return FitResult(spec=self._spec, **stats) + + # ------------------------------------------------------------------ + # Prediction + # ------------------------------------------------------------------ + + def probabilities(self, data=None, beta=None, alpha=None) -> np.ndarray: + """Compute choice probabilities. + + Parameters + ---------- + data : ChoiceTable or None + Data to predict on. If None, uses estimation data. + beta : np.ndarray or None + Parameter vector. If None, uses estimated values. + alpha : np.ndarray or None + Nest parameters (for nested models). If None, uses estimated values. + + Returns + ------- + np.ndarray, shape (n_obs, n_alts) + Choice probabilities. + """ + if self._arrays is None: + raise RuntimeError("Model must be estimated before prediction.") + + arrays = self._arrays + if data is not None: + from ..data.choicetable import ChoiceTable + + if not isinstance(data, ChoiceTable): + raise TypeError("data must be a ChoiceTable") + arrays = data.to_arrays( + formula=self._spec.formula, + spec=self._spec if self._spec.formula is None else None, + ) + + # Dispatch based on model type + if self._is_nested or self._is_mixed: + return self._probabilities_complex(arrays, data, beta, alpha) + return self._probabilities_mnl(arrays, data, beta) + + def _probabilities_mnl(self, arrays, data, beta) -> np.ndarray: + """Compute MNL, SCL, or SAR probabilities.""" + from .._kernels.mnl_numpy import mnl_probs_numpy + + if self._is_spatial_lag: + # SAR: spatial filter + variance normalisation + k = arrays.design_matrix.shape[1] + if beta is None: + coef_vals = np.asarray(self._result.coefficients.values, dtype=np.float64) + beta_use = coef_vals[:k] + rho = float(coef_vals[k]) + else: + beta = np.asarray(beta, dtype=np.float64) + beta_use = beta[:k] + rho = ( + float(beta[k]) if beta.size > k else float(self._result.coefficients.values[k]) + ) + + dm = np.asarray(arrays.design_matrix, dtype=np.float64) + n_obs = arrays.n_obs + n_alts = arrays.n_alts + W_dense = np.asarray(self._W_sparse.toarray(), dtype=np.float64) + + V_base = (dm @ beta_use).reshape(n_obs, n_alts) + from .._sampling.correction import apply_sampling_correction + + V_base = apply_sampling_correction(V_base, arrays) + + A = np.eye(n_alts) - rho * W_dense + V_filtered = np.linalg.solve(A, V_base.T).T + D = np.diag(np.linalg.inv(A)) + V_star = V_filtered / D[None, :] + + if arrays.available is not None: + available = np.asarray(arrays.available, dtype=np.float64).reshape(n_obs, n_alts) + else: + available = np.ones((n_obs, n_alts), dtype=np.float64) + + return mnl_probs_numpy(V_star, available, inclusion_probs=None) + + if self._is_spatial_scl: + from .scl import _scl_log_probs_numpy + + k = arrays.design_matrix.shape[1] + if beta is None: + coef_vals = np.asarray(self._result.coefficients.values, dtype=np.float64) + beta_use = coef_vals[:k] + rho = float(coef_vals[k]) + else: + beta = np.asarray(beta, dtype=np.float64) + beta_use = beta[:k] + rho = ( + float(beta[k]) if beta.size > k else float(self._result.coefficients.values[k]) + ) + + from .._sampling.correction import get_sampling_correction + + log_probs = _scl_log_probs_numpy( + beta_use, + rho, + np.asarray(arrays.design_matrix, dtype=np.float64), + self._allocation, + self._edge_list, + arrays.n_obs, + arrays.n_alts, + available=arrays.available, + inclusion_probs=get_sampling_correction(arrays), + ) + return np.exp(log_probs) + + if beta is None: + beta = np.asarray(self._result.coefficients.values, dtype=np.float64) + + dm = np.asarray(arrays.design_matrix, dtype=np.float64) + n_obs = arrays.n_obs + n_alts = arrays.n_alts + + utilities = (dm @ beta).reshape(n_obs, n_alts) + from .._sampling.correction import apply_sampling_correction + + utilities = apply_sampling_correction(utilities, arrays) + + if arrays.available is not None: + available = np.asarray(arrays.available, dtype=np.float64).reshape(n_obs, n_alts) + else: + available = np.ones((n_obs, n_alts), dtype=np.float64) + + return mnl_probs_numpy(utilities, available, inclusion_probs=None) + + def _probabilities_complex(self, arrays, data, beta, alpha) -> np.ndarray: + """Compute probabilities for nested/mixed models (NumPy fallback).""" + # For nested models, use the NumPy kernel + if self._is_nested and not self._is_mixed: + from .nested import _nested_logit_probs_numpy + + k = arrays.design_matrix.shape[1] + n_nests = self._nests.n_nests + if beta is None: + beta = np.asarray(self._result.coefficients.values[:k], dtype=np.float64) + if alpha is None: + lambda_vals = self._result.coefficients.values[k : k + n_nests] + lambda_vals = np.clip(lambda_vals, 1e-10, 1.0 - 1e-10) + alpha = np.log(lambda_vals / (1.0 - lambda_vals)) + elif beta.size > k: + # Full parameter vector passed — split into beta and alpha + alpha = beta[k : k + n_nests] + beta = beta[:k] + if alpha is None: + alpha = np.zeros(n_nests) + + from .._sampling.correction import get_sampling_correction + + return _nested_logit_probs_numpy( + np.asarray(beta, dtype=np.float64), + np.asarray(alpha, dtype=np.float64), + np.asarray(arrays.design_matrix, dtype=np.float64), + self._nest_matrix, + arrays.n_obs, + arrays.n_alts, + available=arrays.available, + inclusion_probs=get_sampling_correction(arrays), + ) + + # For mixed models, use the NumPy kernel + if self._is_mixed and not self._is_nested: + from .mixed import _mixed_logit_probs_numpy + + k_fixed = self._k_fixed + k_random = self._k_random + + if beta is None: + beta_fixed = self._result.coefficients.values[:k_fixed] + beta_random_means = self._result.coefficients.values[k_fixed : k_fixed + k_random] + beta_random_spreads = self._result.coefficients.values[k_fixed + k_random :] + else: + beta_fixed = beta[:k_fixed] + beta_random_means = beta[k_fixed : k_fixed + k_random] + beta_random_spreads = beta[k_fixed + k_random :] + + from .._sampling.correction import get_sampling_correction + + return _mixed_logit_probs_numpy( + beta_fixed, + beta_random_means, + beta_random_spreads, + self._random_distributions, + self._draws, + np.asarray(arrays.design_matrix, dtype=np.float64), + self._random_col_indices, + arrays.n_obs, + arrays.n_alts, + available=arrays.available, + inclusion_probs=get_sampling_correction(arrays), + ) + + # Mixed nested — use NumPy fallback + if self._is_nested and self._is_mixed: + return self._probabilities_mixed_nested_numpy(arrays, beta, alpha) + + raise RuntimeError("Unknown model configuration for prediction") + + def _probabilities_mixed_nested_numpy(self, arrays, beta, alpha) -> np.ndarray: + """Compute mixed nested logit probabilities (NumPy fallback).""" + from .._sampling.correction import get_sampling_correction + from .nested import _nested_logit_probs_numpy + + k_total = arrays.design_matrix.shape[1] + k_fixed = self._k_fixed + k_random = self._k_random + n_nests = self._nests.n_nests + + if beta is None: + beta_fixed = self._result.coefficients.values[:k_fixed] + else: + beta_fixed = beta[:k_fixed] + + if alpha is None: + lambda_vals = self._result.coefficients.values[k_fixed : k_fixed + n_nests] + lambda_vals = np.clip(lambda_vals, 1e-10, 1.0 - 1e-10) + alpha = np.log(lambda_vals / (1.0 - lambda_vals)) + + beta_random_means = self._result.coefficients.values[ + k_fixed + n_nests : k_fixed + n_nests + k_random + ] + beta_random_spreads = self._result.coefficients.values[k_fixed + n_nests + k_random :] + + dm = np.asarray(arrays.design_matrix, dtype=np.float64) + n_obs = arrays.n_obs + n_alts = arrays.n_alts + available = arrays.available + inclusion_probs = get_sampling_correction(arrays) + + dm_fixed = dm[:, [i for i in range(k_total) if i not in self._random_col_indices]] + dm_random = dm[:, self._random_col_indices] + + if dm_fixed.shape[1] > 0 and len(beta_fixed) > 0: + v_fixed = (dm_fixed @ beta_fixed).reshape(n_obs, n_alts) + else: + v_fixed = np.zeros((n_obs, n_alts)) + + if inclusion_probs is not None: + sr = np.asarray(inclusion_probs, dtype=np.float64).reshape(n_obs, n_alts) + v_fixed = v_fixed + np.log(np.maximum(sr, 1e-30)) + + if available is not None: + avail = np.asarray(available, dtype=np.float64).reshape(n_obs, n_alts) + else: + avail = np.ones((n_obs, n_alts), dtype=np.float64) + + n_draws = self._draws.shape[1] + probs_sum = np.zeros((n_obs, n_alts), dtype=np.float64) + + for r in range(n_draws): + beta_random_r = np.zeros((n_obs, k_random)) + for p in range(k_random): + z_p = self._draws[:, r, p] + mean_p = beta_random_means[p] + spread_p = abs(beta_random_spreads[p]) + dist = self._random_distributions[p] + + if dist == "normal": + beta_random_r[:, p] = mean_p + spread_p * z_p + elif dist == "lognormal": + exponent = mean_p + spread_p * z_p + beta_random_r[:, p] = np.exp(np.clip(exponent, -50, 50)) + elif dist == "triangular": + from scipy.stats import norm as norm_dist + + u = norm_dist.cdf(z_p) + mask = u <= 0.5 + beta_random_r[:, p] = np.where( + mask, + mean_p + spread_p * (np.sqrt(2 * u) - 1), + mean_p + spread_p * (1 - np.sqrt(2 * (1 - u))), + ) + elif dist == "uniform": + from scipy.stats import norm as norm_dist + + u = norm_dist.cdf(z_p) + beta_random_r[:, p] = mean_p + spread_p * (2 * u - 1) + + v_random = np.sum( + dm_random.reshape(n_obs, n_alts, k_random) * beta_random_r[:, None, :], + axis=2, + ) + V = v_fixed + v_random + + probs_r = _nested_logit_probs_numpy( + np.zeros(k_total), # beta not used — V is precomputed + np.asarray(alpha, dtype=np.float64), + V.reshape(-1, 1) if V.size > 0 else np.zeros((n_obs * n_alts, 1)), + self._nest_matrix, + n_obs, + n_alts, + available=avail, + inclusion_probs=None, + ) + probs_sum += probs_r + + return probs_sum / n_draws + + def utilities(self, data=None, beta=None) -> np.ndarray: + """Compute deterministic utilities. + + Parameters + ---------- + data : ChoiceTable or None + Data to predict on. If None, uses estimation data. + beta : np.ndarray or None + Utility coefficients. If None, uses estimated values. + + Returns + ------- + np.ndarray, shape (n_obs, n_alts) + Deterministic utilities. + """ + if self._arrays is None: + raise RuntimeError("Model must be estimated before prediction.") + + arrays = self._arrays + if data is not None: + arrays = data.to_arrays( + formula=self._spec.formula, + spec=self._spec if self._spec.formula is None else None, + ) + + if beta is None: + k = arrays.design_matrix.shape[1] + beta = np.asarray(self._result.coefficients.values[:k], dtype=np.float64) + + dm = np.asarray(arrays.design_matrix, dtype=np.float64) + n_obs = arrays.n_obs + n_alts = arrays.n_alts + + V = (dm @ beta).reshape(n_obs, n_alts) + from .._sampling.correction import apply_sampling_correction + + V = apply_sampling_correction(V, arrays) + return V + + # ------------------------------------------------------------------ + # Simulation (vectorized) + # ------------------------------------------------------------------ + + def simulate(self, data=None, n_draws: int = 1, seed: Optional[int] = None) -> pd.DataFrame: + """Simulate choices from the estimated model. + + Uses vectorized inverse-CDF sampling — no Python loops over + observations. + + Parameters + ---------- + data : ChoiceTable or None + Data to simulate on. If None, uses estimation data. + n_draws : int, optional + Number of simulation draws per observation. Default 1. + seed : int or None, optional + Random seed for reproducibility. + + Returns + ------- + pd.DataFrame + Simulated choices with columns ``draw``, ``obs_id``, + ``alt_id``, and ``probability``. + """ + from ..data.choicetable import ChoiceTable + + if self._arrays is None: + raise RuntimeError("Model must be estimated before simulation.") + + arrays = self._arrays + ct = self._data + if data is not None: + if not isinstance(data, ChoiceTable): + raise TypeError("data must be a ChoiceTable") + arrays = data.to_arrays( + formula=self._spec.formula, + spec=self._spec if self._spec.formula is None else None, + ) + ct = data + + rng = np.random.default_rng(seed) + probs = self.probabilities(data=data) + probs = probs / probs.sum(axis=1, keepdims=True) + n_obs = arrays.n_obs + n_alts = arrays.n_alts + + df = ct.to_frame() + alt_ids = df[ct.alt_id_col].values.reshape(n_obs, n_alts) + obs_ids = df[ct.obs_id_col].values.reshape(n_obs, n_alts)[:, 0] + + # Vectorized simulation: draw all choices at once + cumulative_probs = np.cumsum(probs, axis=1) + uniform_draws = rng.random((n_draws, n_obs)) + chosen_indices = np.argmax( + cumulative_probs[None, :, :] > uniform_draws[:, :, None], axis=2 + ) + chosen_indices = np.clip(chosen_indices, 0, n_alts - 1) + + chosen_alts = alt_ids[np.arange(n_obs), chosen_indices] + chosen_probs = probs[np.arange(n_obs), chosen_indices] + + # Build results DataFrame (vectorized) + results = pd.DataFrame( + { + "draw": np.repeat(np.arange(n_draws), n_obs), + ct.obs_id_col: np.tile(obs_ids, n_draws), + ct.alt_id_col: chosen_alts.T.ravel(), + "probability": chosen_probs.T.ravel(), + } + ) + return results + + # ------------------------------------------------------------------ + # Marginal Effects + # ------------------------------------------------------------------ + + def _resolve_me_data(self, data=None): + """Resolve data for marginal effects computation.""" + from ..data.choicetable import ChoiceTable + + if self._arrays is None: + raise RuntimeError("Model must be estimated before computing marginal effects.") + + ct = self._data + if data is not None: + if not isinstance(data, ChoiceTable): + raise TypeError("data must be a ChoiceTable") + ct = data + + probs = self.probabilities(data=data) + df = ct.to_frame() + index = pd.MultiIndex.from_arrays( + [df[ct.obs_id_col].values, df[ct.alt_id_col].values], + names=[ct.obs_id_col, ct.alt_id_col], + ) + return ct, probs, df, index + + def marginal_effects(self, data=None, variable: Optional[str] = None): + """Compute average direct, indirect, and total marginal effects. + + In the SAR-MNL model (``lag=True``), a change in an attribute of + alternative *j* affects not only *j*'s utility but also neighbouring + alternatives through the spatial multiplier + :math:`(I - \\rho W)^{-1}`. + + Following LeSage & Pace (2009): + + - **Direct effect**: impact on own alternative. + - **Indirect effect**: spillover to neighbouring alternatives. + - **Total effect**: direct + indirect. + + Parameters + ---------- + data : ChoiceTable or None + Data to compute marginal effects on. If None, uses + estimation data. + variable : str + Name of the variable to compute marginal effects for. + + Returns + ------- + dict + Dictionary with keys ``"direct"``, ``"indirect"``, + ``"total"``, each mapping to a ``pd.Series`` indexed by + alternative ID. + """ + if not self._is_spatial_lag: + raise ValueError( + "marginal_effects() is only available for SAR models (lag=True). " + "Use marginal_effect() for non-spatial models." + ) + if self._arrays is None: + raise RuntimeError("Model must be estimated before computing marginal effects.") + + from ..data.choicetable import ChoiceTable + + ct = self._data + arrays = self._arrays + if data is not None: + if not isinstance(data, ChoiceTable): + raise TypeError("data must be a ChoiceTable") + arrays = data.to_arrays( + formula=self._spec.formula, + spec=self._spec if self._spec.formula is None else None, + ) + ct = data + + n_obs = arrays.n_obs + n_alts = arrays.n_alts + k = arrays.design_matrix.shape[1] + + coef_vals = np.asarray(self._result.coefficients.values, dtype=np.float64) + beta = coef_vals[:k] + rho = float(coef_vals[k]) + + param_names = list(arrays.param_names) + if variable not in param_names: + raise ValueError(f"Variable '{variable}' not found in parameters: {param_names}") + beta_r = beta[param_names.index(variable)] + + probs = self.probabilities(data=data) + W_dense = np.asarray(self._W_sparse.toarray(), dtype=np.float64) + A = np.eye(n_alts) - rho * W_dense + Z_mat = np.linalg.inv(A) + + direct = np.zeros(n_alts) + indirect = np.zeros(n_alts) + for k_alt in range(n_alts): + direct[k_alt] = ( + beta_r * Z_mat[k_alt, k_alt] * np.mean(probs[:, k_alt] * (1 - probs[:, k_alt])) + ) + for j_alt in range(n_alts): + if j_alt != k_alt: + indirect[k_alt] += ( + -beta_r * Z_mat[k_alt, j_alt] * np.mean(probs[:, k_alt] * probs[:, j_alt]) + ) + + total = direct + indirect + + df = ct.to_frame() + alt_ids = df[ct.alt_id_col].values.reshape(n_obs, n_alts)[0] + + return { + "direct": pd.Series(direct, index=alt_ids, name=f"direct_{variable}"), + "indirect": pd.Series(indirect, index=alt_ids, name=f"indirect_{variable}"), + "total": pd.Series(total, index=alt_ids, name=f"total_{variable}"), + } + + def marginal_effect(self, data=None, variable: Optional[str] = None) -> pd.Series: + """Compute direct marginal effects for a variable. + + For MNL: :math:`(1 - P_{qi}) \\beta_x`. + + For nested logit: :math:`P_i (1 - \\lambda_m P_{i|m}) \\beta_x` + where :math:`P_{i|m}` is the conditional probability within nest m. + + For mixed logit: :math:`E_z[(1 - P_i(z)) \\beta_x]` via simulation + over draws. + + For SCL: raises ``NotImplementedError`` (derivation pending). + + Parameters + ---------- + data : ChoiceTable or None + Data to compute marginal effects on. + variable : str + Name of the variable. + + Returns + ------- + pd.Series + Direct marginal effects, indexed by (obs_id, alt_id). + """ + ct, probs, df, index = self._resolve_me_data(data) + beta = self._result.coefficients.get(variable, 0.0) + + if self._is_nested: + # Nested logit: P_i * (1 - lambda_m * P_{i|m}) * beta + n_obs = probs.shape[0] + n_alts = probs.shape[1] + n_nests = self._nests.n_nests + + # Get lambda values from estimated coefficients + k = self._arrays.design_matrix.shape[1] if self._arrays is not None else None + if k is None and data is not None: + k = data.to_arrays( + formula=self._spec.formula, + spec=self._spec if self._spec.formula is None else None, + ).design_matrix.shape[1] + + # Extract lambda values (naturalized) + if self._is_spatial: + # Nested SCL: [beta, rho_1..M, lambda_1..M] + lambda_vals = self._result.coefficients.values[k + n_nests : k + 2 * n_nests] + else: + # Nested: [beta, lambda_1..M] + lambda_vals = self._result.coefficients.values[k : k + n_nests] + + # Compute conditional probabilities P_{i|m} = P_i / P_m + # P_m = sum of P_i for alts in nest m + nest_matrix = self._nest_matrix # (n_alts, n_nests) + P_nest = probs @ nest_matrix # (n_obs, n_nests) + # P_{i|m} = P_i / P_m (avoid division by zero) + # For each alt, find which nest it belongs to + alt_in_nest = nest_matrix.sum(axis=1) > 0 # (n_alts,) bool + + # lambda for each alternative (from its nest) + long_lambda = np.ones(n_alts) + for m in range(n_nests): + mask = nest_matrix[:, m] > 0 + long_lambda[mask] = lambda_vals[m] + + # P_{i|m} for each (obs, alt) + P_i_given_m = np.zeros_like(probs) + for m in range(n_nests): + mask = nest_matrix[:, m] > 0 + if not mask.any(): + continue + P_m = P_nest[:, m : m + 1] # (n_obs, 1) + P_i_given_m[:, mask] = probs[:, mask] / np.maximum(P_m, 1e-30) + + # Marginal effect: P_i * (1 - lambda_m * P_{i|m}) * beta + me = probs * (1 - long_lambda[None, :] * P_i_given_m) * beta + me = me.ravel() + + # For root nest alternatives (not in any nest), use MNL formula + if not alt_in_nest.all(): + root_mask = ~alt_in_nest + me_2d = me.reshape(n_obs, n_alts) + me_2d[:, root_mask] = (1 - probs[:, root_mask]) * beta + me = me_2d.ravel() + + elif self._is_mixed: + # Mixed logit: E_z[(1 - P_i(z)) * beta] via simulation + # For now, use the MNL approximation with mean coefficients + # (proper implementation requires per-draw probability computation) + me = (1 - probs.ravel()) * beta + + elif self._is_spatial: + # SCL: derivation pending + raise NotImplementedError( + "Marginal effects for SCL models are not yet implemented. " + "The MNL approximation is incorrect for spatially correlated logit." + ) + + else: + # MNL: (1 - P_i) * beta + me = (1 - probs.ravel()) * beta + + return pd.Series(me, index=index, name=f"marginal_effect_{variable}") + + def cross_marginal_effect(self, data=None, variable: Optional[str] = None) -> pd.Series: + """Compute cross-marginal effects for a variable. + + For MNL: :math:`-P_i \\beta_x`. + + For nested logit: :math:`-P_i \\lambda_m P_{i|m} \\beta_x`. + + For SCL: raises ``NotImplementedError``. + + Parameters + ---------- + data : ChoiceTable or None + variable : str + + Returns + ------- + pd.Series + Cross-marginal effects, indexed by (obs_id, alt_id). + """ + ct, probs, df, index = self._resolve_me_data(data) + beta = self._result.coefficients.get(variable, 0.0) + + if self._is_nested: + # Nested logit cross-ME: -P_i * lambda_m * P_{i|m} * beta + n_obs = probs.shape[0] + n_alts = probs.shape[1] + n_nests = self._nests.n_nests + + k = self._arrays.design_matrix.shape[1] if self._arrays is not None else None + if k is None and data is not None: + k = data.to_arrays( + formula=self._spec.formula, + spec=self._spec if self._spec.formula is None else None, + ).design_matrix.shape[1] + + if self._is_spatial: + lambda_vals = self._result.coefficients.values[k + n_nests : k + 2 * n_nests] + else: + lambda_vals = self._result.coefficients.values[k : k + n_nests] + + nest_matrix = self._nest_matrix + P_nest = probs @ nest_matrix + alt_in_nest = nest_matrix.sum(axis=1) > 0 + + long_lambda = np.ones(n_alts) + for m in range(n_nests): + mask = nest_matrix[:, m] > 0 + long_lambda[mask] = lambda_vals[m] + + P_i_given_m = np.zeros_like(probs) + for m in range(n_nests): + mask = nest_matrix[:, m] > 0 + if not mask.any(): + continue + P_m = P_nest[:, m : m + 1] + P_i_given_m[:, mask] = probs[:, mask] / np.maximum(P_m, 1e-30) + + cross_me = -probs * long_lambda[None, :] * P_i_given_m * beta + cross_me = cross_me.ravel() + + if not alt_in_nest.all(): + root_mask = ~alt_in_nest + cross_me_2d = cross_me.reshape(n_obs, n_alts) + cross_me_2d[:, root_mask] = -probs[:, root_mask] * beta + cross_me = cross_me_2d.ravel() + + elif self._is_spatial: + raise NotImplementedError( + "Cross-marginal effects for SCL models are not yet implemented." + ) + + else: + # MNL and mixed (approximation): -P_i * beta + cross_me = -probs.ravel() * beta + + return pd.Series(cross_me, index=index, name=f"cross_marginal_effect_{variable}") + + def elasticity(self, data=None, variable: Optional[str] = None) -> pd.Series: + """Compute direct elasticities for a variable. + + For MNL: :math:`(1 - P_{qi}) \\beta_x x_{qi}`. + + For nested logit: :math:`P_i (1 - \\lambda_m P_{i|m}) \\beta_x x_{qi}`. + + For SCL: raises ``NotImplementedError``. + + Parameters + ---------- + data : ChoiceTable or None + variable : str + + Returns + ------- + pd.Series + Direct elasticities, indexed by (obs_id, alt_id). + """ + ct, probs, df, index = self._resolve_me_data(data) + x = df[variable].values + + if self._is_spatial and not self._is_nested: + raise NotImplementedError("Elasticities for SCL models are not yet implemented.") + + # For MNL, nested, and mixed: elasticity = marginal_effect * x + me = self.marginal_effect(data=data, variable=variable) + elasticities = me.values * x + + return pd.Series(elasticities, index=index, name=f"elasticity_{variable}") + + def cross_elasticity(self, data=None, variable: Optional[str] = None) -> pd.Series: + """Compute cross-elasticities for a variable. + + For MNL: :math:`-P_i \\beta_x x_{ij}`. + + For SCL: raises ``NotImplementedError``. + + Parameters + ---------- + data : ChoiceTable or None + variable : str + + Returns + ------- + pd.Series + Cross-elasticities, indexed by (obs_id, alt_id). + """ + ct, probs, df, index = self._resolve_me_data(data) + x = df[variable].values + + if self._is_spatial and not self._is_nested: + raise NotImplementedError("Cross-elasticities for SCL models are not yet implemented.") + + # cross_elasticity = cross_marginal_effect * x + cme = self.cross_marginal_effect(data=data, variable=variable) + cross_elast = cme.values * x + + return pd.Series(cross_elast, index=index, name=f"cross_elasticity_{variable}") + + # ------------------------------------------------------------------ + # Covariance estimation + # ------------------------------------------------------------------ + + def covariance_robust(self, data=None) -> np.ndarray: + """Compute the sandwich (Huber-White) robust covariance matrix. + + Parameters + ---------- + data : ChoiceTable or None + + Returns + ------- + np.ndarray, shape (n_parameters, n_parameters) + Sandwich (robust) covariance matrix. + """ + from ..data.choicetable import ChoiceTable + + if self._arrays is None: + raise RuntimeError("Model must be estimated first.") + + arrays = self._arrays + if data is not None: + if not isinstance(data, ChoiceTable): + raise TypeError("data must be a ChoiceTable") + arrays = data.to_arrays( + formula=self._spec.formula, + spec=self._spec if self._spec.formula is None else None, + ) + + scores = self._observation_scores(arrays) + B = scores.T @ scores + H_inv = self._get_hessian_inverse() + + if H_inv is None: + return _safe_inv(B) + + return _sandwich_inv(H_inv, B) + + def covariance_clustered(self, data=None, groups=None) -> np.ndarray: + """Compute cluster-robust (Rogers) covariance matrix. + + Parameters + ---------- + data : ChoiceTable or None + groups : array-like, shape (n_obs,) + Cluster/group identifiers. + + Returns + ------- + np.ndarray, shape (n_parameters, n_parameters) + Cluster-robust covariance matrix. + """ + from ..data.choicetable import ChoiceTable + + if self._arrays is None: + raise RuntimeError("Model must be estimated first.") + + arrays = self._arrays + if data is not None: + if not isinstance(data, ChoiceTable): + raise TypeError("data must be a ChoiceTable") + arrays = data.to_arrays( + formula=self._spec.formula, + spec=self._spec if self._spec.formula is None else None, + ) + + if groups is None: + raise ValueError("groups must be provided for cluster-robust covariance.") + + scores = self._observation_scores(arrays) + groups = np.asarray(groups) + unique_groups = np.unique(groups) + n_params = scores.shape[1] + + B_clustered = np.zeros((n_params, n_params)) + for g in unique_groups: + mask = groups == g + g_c = scores[mask].sum(axis=0) + B_clustered += np.outer(g_c, g_c) + + H_inv = self._get_hessian_inverse() + + if H_inv is None: + return _safe_inv(B_clustered) + + return _sandwich_inv(H_inv, B_clustered) + + def std_errors_robust(self, data=None) -> pd.Series: + """Compute sandwich (Huber-White) robust standard errors.""" + cov = self.covariance_robust(data=data) + se = np.sqrt(np.maximum(np.diag(cov), 0)) + se[se == 0] = np.nan + return pd.Series(se, index=self._result.coefficients.index, name="std_error_robust") + + def std_errors_clustered(self, data=None, groups=None) -> pd.Series: + """Compute cluster-robust standard errors.""" + cov = self.covariance_clustered(data=data, groups=groups) + se = np.sqrt(np.maximum(np.diag(cov), 0)) + se[se == 0] = np.nan + return pd.Series(se, index=self._result.coefficients.index, name="std_error_clustered") + + # ------------------------------------------------------------------ + # Observation scores + # ------------------------------------------------------------------ + + def _observation_scores(self, arrays) -> np.ndarray: + """Compute observation-level score (gradient) vectors. + + Uses analytical MNL scores for MNL models, JAX jacrev for models + with a JAX objective (spatial/nested/mixed), and finite differences + as a last-resort fallback. + """ + cache_key = id(arrays) + if cache_key in self._observation_scores_cache: + return self._observation_scores_cache[cache_key] + + if not self._is_spatial and not self._is_nested and not self._is_mixed: + scores = self._mnl_observation_scores(arrays) + elif self._objective is not None and self._objective.jax_fn is not None: + scores = self._jax_observation_scores(arrays) + else: + scores = self._finite_diff_observation_scores(arrays) + + self._observation_scores_cache[cache_key] = scores + return scores + + def _mnl_observation_scores(self, arrays) -> np.ndarray: + """Compute MNL observation scores analytically.""" + from .._kernels.mnl_numpy import mnl_observation_scores_numpy + + dm = np.asarray(arrays.design_matrix, dtype=np.float64) + chosen = np.asarray(arrays.chosen, dtype=np.float64).reshape(arrays.n_obs, arrays.n_alts) + beta = np.asarray(self._result.coefficients.values, dtype=np.float64) + n_obs = arrays.n_obs + n_alts = arrays.n_alts + + if arrays.available is not None: + available = np.asarray(arrays.available, dtype=np.float64).reshape(n_obs, n_alts) + else: + available = np.ones((n_obs, n_alts), dtype=np.float64) + + from .._sampling.correction import get_sampling_correction + + inclusion_probs = get_sampling_correction(arrays) + + weights = None + if arrays.weights is not None: + weights = np.asarray(arrays.weights, dtype=np.float64) + + return mnl_observation_scores_numpy( + beta=beta, + design_matrix=dm, + chosen=chosen, + available=available, + n_obs=n_obs, + n_alts=n_alts, + weights=weights, + inclusion_probs=inclusion_probs, + ) + + def _finite_diff_observation_scores(self, arrays) -> np.ndarray: + """Compute observation scores via finite differences (fallback).""" + eps = 1e-5 + n_params = len(self._result.coefficients) + n_obs = arrays.n_obs + + full_params = np.asarray(self._result.coefficients.values, dtype=np.float64).copy() + chosen = np.asarray(arrays.chosen, dtype=np.float64).reshape(n_obs, arrays.n_alts) + scores = np.zeros((n_obs, n_params)) + + for j in range(n_params): + p_plus = full_params.copy() + p_plus[j] += eps + p_minus = full_params.copy() + p_minus[j] -= eps + + probs_plus = self.probabilities(data=None, beta=p_plus) + probs_minus = self.probabilities(data=None, beta=p_minus) + + ll_plus = np.log(np.maximum(np.sum(probs_plus * chosen, axis=1), 1e-30)) + ll_minus = np.log(np.maximum(np.sum(probs_minus * chosen, axis=1), 1e-30)) + scores[:, j] = (ll_plus - ll_minus) / (2 * eps) + + return scores + + def _jax_observation_scores(self, arrays) -> np.ndarray: + """Compute observation scores via JAX jacrev on per-obs LL contributions. + + This uses ``jax.jacrev`` on the objective's per-observation log-likelihood + contributions, giving exact gradients in a single backward pass — O(1) + instead of O(n_params) probability evaluations. + """ + import jax.numpy as jnp + + # Get per-observation LL contribution function from the objective + if self._objective is None or self._objective.jax_fn is None: + return self._finite_diff_observation_scores(arrays) + + # Try score_contribs (requires loglike_contribs_jax to be set) + try: + contribs_fn = self._objective.score_contribs + beta = jnp.asarray(self._result.coefficients.values, dtype=jnp.float64) + scores = np.asarray(contribs_fn(beta)) + return scores + except (ValueError, AttributeError, TypeError): + pass + + # Fall back to finite differences + return self._finite_diff_observation_scores(arrays) + + # ------------------------------------------------------------------ + # Convenience + # ------------------------------------------------------------------ + + def __repr__(self) -> str: + status = "estimated" if self._result is not None else "not estimated" + formula_str = self._formula or "custom spec" + return f"ChoiceModel(formula='{formula_str}', {status})" diff --git a/locpick/models/mixed.py b/locpick/models/mixed.py index 36d0e12..3051c71 100644 --- a/locpick/models/mixed.py +++ b/locpick/models/mixed.py @@ -52,15 +52,15 @@ import numpy as np import pandas as pd -from locpick._jax.objective import Objective -from locpick._kernels.constants import NEG_INF -from locpick._solvers import Solver, SolverResult -from locpick.data.arrays import ChoiceArrays -from locpick.models._spatial import ( +from .._jax.objective import Objective +from .._kernels.constants import NEG_INF +from .._solvers import Solver, SolverResult +from ..data.arrays import ChoiceArrays +from ..results.fit_result import FitResult +from ._spatial import ( naturalize_rho, ) -from locpick.models.base import BaseChoiceModel, SpatialMixin, _safe_inv, _sandwich_inv -from locpick.results.fit_result import FitResult +from .base import BaseChoiceModel, SpatialMixin, _safe_inv, _sandwich_inv # --------------------------------------------------------------------------- # Distribution specifications @@ -439,21 +439,7 @@ def _mixed_logit_probs_numpy( log_probs_draws = np.zeros((n_obs, n_draws, n_alts), dtype=np.float64) for r in range(n_draws): - # Realise random coefficients for this draw - beta_r = np.zeros(k_random) - for p in range(k_random): - z_p = draws[:, r, p] # (n_obs,) - mean_p = np.full(n_obs, beta_random_means[p]) - spread_p = np.full(n_obs, beta_random_spreads[p]) - beta_r[p] = _apply_distribution( - z_p[:, None], - mean_p[:, None], - spread_p[:, None], - random_distributions[p], - ).ravel()[0] # scalar for this draw - - # Actually, we need per-observation random coefficients - # beta_r[n, p] = mean_p + spread_p * z[n, r, p] + # Realise per-observation random coefficients for this draw beta_random_r = np.zeros((n_obs, k_random)) for p in range(k_random): z_p = draws[:, r, p] # (n_obs,) @@ -630,90 +616,6 @@ def _mixed_logit_ll_numpy( return float(np.sum(log_sim_probs)) -def _mixed_logit_gradient_numpy( - params: np.ndarray, - random_col_indices: list[int], - k_fixed: int, - k_random: int, - random_distributions: list[str], - draws: np.ndarray, - design_matrix: np.ndarray, - chosen: np.ndarray, - n_obs: int, - n_alts: int, - available: Optional[np.ndarray] = None, - inclusion_probs: Optional[np.ndarray] = None, - weights: Optional[np.ndarray] = None, -) -> np.ndarray: - """Compute mixed logit gradient via finite differences (NumPy backend). - - Parameters - ---------- - params : np.ndarray, shape (k_fixed + 2 * k_random,) - Full parameter vector: [beta_fixed, beta_random_means, beta_random_spreads]. - random_col_indices, k_fixed, k_random, random_distributions, draws, design_matrix, chosen, n_obs, n_alts, available, inclusion_probs, weights - See :func:`_mixed_logit_ll_numpy`. - - Returns - ------- - np.ndarray, shape (k_fixed + 2 * k_random,) - Gradient of the simulated log-likelihood. - """ - eps = 1e-5 - n_params = len(params) - grad = np.zeros(n_params) - - for i in range(n_params): - params_plus = params.copy() - params_plus[i] += eps - params_minus = params.copy() - params_minus[i] -= eps - - def _unpack(p): - bf = p[:k_fixed] - rm = p[k_fixed : k_fixed + k_random] - rs = p[k_fixed + k_random :] - return bf, rm, rs - - bf_plus, rm_plus, rs_plus = _unpack(params_plus) - bf_minus, rm_minus, rs_minus = _unpack(params_minus) - - ll_plus = _mixed_logit_ll_numpy( - bf_plus, - rm_plus, - rs_plus, - random_distributions, - draws, - design_matrix, - chosen, - random_col_indices, - n_obs, - n_alts, - available=available, - inclusion_probs=inclusion_probs, - weights=weights, - ) - ll_minus = _mixed_logit_ll_numpy( - bf_minus, - rm_minus, - rs_minus, - random_distributions, - draws, - design_matrix, - chosen, - random_col_indices, - n_obs, - n_alts, - available=available, - inclusion_probs=inclusion_probs, - weights=weights, - ) - - grad[i] = (ll_plus - ll_minus) / (2 * eps) - - return grad - - # --------------------------------------------------------------------------- # MixedLogit model class # --------------------------------------------------------------------------- @@ -755,7 +657,7 @@ class MixedMNL(BaseChoiceModel, SpatialMixin): Examples -------- >>> from locpick import ChoiceTable - >>> from locpick.models.mixed import MixedLogit, ParamDistribution + >>> from .mixed import MixedLogit, ParamDistribution >>> ct = ChoiceTable.from_tables(choosers, alternatives, chosen) >>> model = MixedLogit( ... ct, formula="cost + time - 1", @@ -910,11 +812,9 @@ def _build_objective(self, arrays: ChoiceArrays) -> Objective: """Build optimization objective for mixed logit estimation.""" random_col_indices = self._random_col_indices random_distributions = self._random_distributions - k_fixed = self._k_fixed - k_random = self._k_random if self._is_spatial: - from locpick._jax.builders import build_mscl_objective + from .._jax.builders import build_mscl_objective return build_mscl_objective( arrays, @@ -928,7 +828,7 @@ def _build_objective(self, arrays: ChoiceArrays) -> Objective: backend = (self._backend or os.environ.get("LOCPICK_MIXED_BACKEND", "")).lower() if backend != "numpy": - from locpick._jax.builders import build_mixed_logit_objective + from .._jax.builders import build_mixed_logit_objective return build_mixed_logit_objective( arrays, @@ -937,58 +837,8 @@ def _build_objective(self, arrays: ChoiceArrays) -> Objective: draws=self._draws, ) - dm = np.asarray(arrays.design_matrix, dtype=np.float64) - chosen = np.asarray(arrays.chosen, dtype=np.float64) - n_obs = arrays.n_obs - n_alts = arrays.n_alts - available = arrays.available - weights = arrays.weights - - from locpick._sampling.correction import get_sampling_correction - - inclusion_probs = get_sampling_correction(arrays) - - def ll_fn(params): - beta_fixed = params[:k_fixed] - beta_random_means = params[k_fixed : k_fixed + k_random] - beta_random_spreads = params[k_fixed + k_random :] - return _mixed_logit_ll_numpy( - beta_fixed, - beta_random_means, - beta_random_spreads, - random_distributions, - self._draws, - dm, - chosen, - random_col_indices, - n_obs, - n_alts, - available=available, - inclusion_probs=inclusion_probs, - weights=weights, - ) - - def grad_fn(params): - return _mixed_logit_gradient_numpy( - params, - random_col_indices, - k_fixed, - k_random, - random_distributions, - self._draws, - dm, - chosen, - n_obs, - n_alts, - available=available, - inclusion_probs=inclusion_probs, - weights=weights, - ) - - return Objective.from_numpy( - ll_fn=ll_fn, - grad_fn=grad_fn, - param_names=list(self._full_param_names), + raise NotImplementedError( + "MixedMNL NumPy backend has been removed. Use ChoiceModel (JAX backend)." ) def _build_fit_result( @@ -1289,7 +1139,7 @@ def utilities(self, data=None, beta_fixed=None, beta_random_means=None): V = (dm @ beta_full).reshape(n_obs, n_alts) # Add sampling correction if present - from locpick._sampling.correction import apply_sampling_correction + from .._sampling.correction import apply_sampling_correction V = apply_sampling_correction(V, arrays) @@ -1321,7 +1171,7 @@ def simulate(self, data=None, n_draws: int = 1, seed: Optional[int] = None) -> p Simulated choices with columns ``draw``, ``obs_id``, ``alt_id``, and ``probability``. """ - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if self._arrays is None: raise RuntimeError("Model must be estimated before simulation.") @@ -1390,7 +1240,7 @@ def marginal_effect(self, data=None, variable: Optional[str] = None) -> pd.Serie pd.Series Direct marginal effects, indexed by (obs_id, alt_id). """ - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if self._arrays is None: raise RuntimeError("Model must be estimated before computing marginal effects.") @@ -1436,7 +1286,7 @@ def cross_marginal_effect(self, data=None, variable: Optional[str] = None) -> pd pd.Series Cross-marginal effects, indexed by (obs_id, alt_id). """ - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if self._arrays is None: raise RuntimeError("Model must be estimated before computing marginal effects.") @@ -1490,7 +1340,7 @@ def elasticity(self, data=None, variable: Optional[str] = None) -> pd.Series: pd.Series Direct elasticities, indexed by (obs_id, alt_id). """ - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if self._arrays is None: raise RuntimeError("Model must be estimated before computing elasticities.") @@ -1544,7 +1394,7 @@ def cross_elasticity(self, data=None, variable: Optional[str] = None) -> pd.Seri pd.Series Cross-elasticities, indexed by (obs_id, alt_id). """ - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if self._arrays is None: raise RuntimeError("Model must be estimated before computing elasticities.") @@ -1592,7 +1442,7 @@ def covariance_robust(self, data=None) -> np.ndarray: np.ndarray, shape (n_parameters, n_parameters) Sandwich (robust) covariance matrix. """ - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if self._arrays is None: raise RuntimeError("Model must be estimated first.") @@ -1631,7 +1481,7 @@ def covariance_clustered(self, data=None, groups=None) -> np.ndarray: np.ndarray, shape (n_parameters, n_parameters) Cluster-robust covariance matrix. """ - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if self._arrays is None: raise RuntimeError("Model must be estimated first.") @@ -1773,7 +1623,7 @@ def probabilities( draws = self._draws # Resolve canonical sampling correction tensor. - from locpick._sampling.correction import get_sampling_correction + from .._sampling.correction import get_sampling_correction sampling_correction = get_sampling_correction(arrays) diff --git a/locpick/models/mixed_nested.py b/locpick/models/mixed_nested.py index 7ea9c67..dcae25a 100644 --- a/locpick/models/mixed_nested.py +++ b/locpick/models/mixed_nested.py @@ -46,23 +46,23 @@ import numpy as np import pandas as pd -from locpick._jax.objective import Objective -from locpick._solvers import Solver, SolverResult -from locpick.data.arrays import ChoiceArrays -from locpick.models._spatial import ( +from .._jax.objective import Objective +from .._solvers import Solver, SolverResult +from ..data.arrays import ChoiceArrays +from ..results.fit_result import FitResult +from ._spatial import ( EdgeStructure, _resolve_spatial_graph, naturalize_rho, ) -from locpick.models.base import ( +from .base import ( BaseChoiceModel, SpatialMixin, _compute_fit_statistics, _compute_null_ll, ) -from locpick.models.mixed import ParamDistribution, _resolve_draws -from locpick.models.nested import NestingTree, naturalize_nest_params -from locpick.results.fit_result import FitResult +from .mixed import ParamDistribution, _resolve_draws +from .nested import NestingTree, naturalize_nest_params class MixedNestedMNL(BaseChoiceModel, SpatialMixin): @@ -106,9 +106,9 @@ class MixedNestedMNL(BaseChoiceModel, SpatialMixin): Examples -------- >>> from locpick import ChoiceTable - >>> from locpick.models.mixed_nested import MixedNestedMNL - >>> from locpick.models.nested import NestingTree, NestSpec - >>> from locpick.models.mixed import ParamDistribution + >>> from .mixed_nested import MixedNestedMNL + >>> from .nested import NestingTree, NestSpec + >>> from .mixed import ParamDistribution >>> ct = ChoiceTable.from_tables(choosers, alternatives, chosen) >>> nests = NestingTree([ ... NestSpec("transit", alt_ids=[0, 1, 2]), @@ -220,7 +220,7 @@ def _pre_fit(self, arrays: ChoiceArrays) -> None: self._resolve_spatial_graph() self._validate_graph_size(arrays) - from locpick._jax.data import EdgeDataJAX + from .._jax.data import EdgeDataJAX n_nests = self._nests.n_nests self._edge_structs = [] @@ -301,7 +301,7 @@ def _build_objective(self, arrays: ChoiceArrays) -> Objective: ) if self._is_spatial: - from locpick._jax.builders import build_mnscl_objective + from .._jax.builders import build_mnscl_objective return build_mnscl_objective( arrays, @@ -315,7 +315,7 @@ def _build_objective(self, arrays: ChoiceArrays) -> Objective: # Try JAX backend first backend = (self._backend or os.environ.get("LOCPICK_MIXED_NESTED_BACKEND", "")).lower() if backend != "numpy": - from locpick._jax.builders import build_mixed_nested_objective + from .._jax.builders import build_mixed_nested_objective return build_mixed_nested_objective( arrays, @@ -546,7 +546,7 @@ def probabilities(self, data=None, beta=None, alpha=None): arrays = self._arrays if data is not None: - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if not isinstance(data, ChoiceTable): raise TypeError("data must be a ChoiceTable") @@ -562,7 +562,7 @@ def _probabilities_jax(self, arrays, beta=None, alpha=None): """Compute probabilities using JAX backend.""" import jax.numpy as jnp - from locpick._jax.data import ChoiceDataJAX + from .._jax.data import ChoiceDataJAX k_total = arrays.design_matrix.shape[1] k_fixed = k_total - len(self._random_col_indices) @@ -662,7 +662,7 @@ def _probabilities_jax(self, arrays, beta=None, alpha=None): V = np.asarray(v_fixed) + v_random # Nested logit probabilities for this draw - from locpick.models.nested import _nested_logit_probs_numpy + from .nested import _nested_logit_probs_numpy _nested_logit_probs_numpy( np.concatenate([beta_fixed, np.zeros(0)]), # beta only, no nest params in utility @@ -688,7 +688,7 @@ def _probabilities_jax(self, arrays, beta=None, alpha=None): def _probabilities_numpy(self, arrays, beta=None, alpha=None): """Compute probabilities using NumPy backend (fallback).""" - from locpick._sampling.correction import get_sampling_correction + from .._sampling.correction import get_sampling_correction k_total = arrays.design_matrix.shape[1] k_fixed = k_total - len(self._random_col_indices) @@ -869,7 +869,7 @@ def utilities(self, data=None, beta=None): arrays = self._arrays if data is not None: - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if not isinstance(data, ChoiceTable): raise TypeError("data must be a ChoiceTable") @@ -898,7 +898,7 @@ def utilities(self, data=None, beta=None): V = np.zeros((n_obs, n_alts)) # Add sampling correction if present - from locpick._sampling.correction import apply_sampling_correction + from .._sampling.correction import apply_sampling_correction V = apply_sampling_correction(V, arrays) diff --git a/locpick/models/mnl.py b/locpick/models/mnl.py index b69b0cb..9a87912 100644 --- a/locpick/models/mnl.py +++ b/locpick/models/mnl.py @@ -12,13 +12,14 @@ import numpy as np import pandas as pd -from locpick._solvers import Solver, SolverResult -from locpick.data.arrays import ChoiceArrays -from locpick.data.problem import EstimationProblem -from locpick.models._spatial import ( +from .._solvers import Solver, SolverResult +from ..data.arrays import ChoiceArrays +from ..data.problem import EstimationProblem +from ..results.fit_result import FitResult +from ._spatial import ( naturalize_rho, ) -from locpick.models.base import ( +from .base import ( BaseChoiceModel, SpatialMixin, _compute_fit_statistics, @@ -26,7 +27,6 @@ _safe_inv, _sandwich_inv, ) -from locpick.results.fit_result import FitResult class MNL(BaseChoiceModel, SpatialMixin): @@ -158,7 +158,7 @@ def probabilities(self, data=None, beta=None): np.ndarray, shape (n_obs, n_alts) Choice probabilities for each observation and alternative. """ - from locpick._kernels.mnl_numpy import mnl_probs_numpy + from .._kernels.mnl_numpy import mnl_probs_numpy if self._arrays is None: raise RuntimeError("Model must be estimated before prediction.") @@ -171,7 +171,7 @@ def probabilities(self, data=None, beta=None): ) if self._is_spatial: - from locpick.models.scl import _scl_log_probs_dispatch + from .scl import _scl_log_probs_dispatch k = arrays.design_matrix.shape[1] if beta is None: @@ -186,7 +186,7 @@ def probabilities(self, data=None, beta=None): else: rho = float(self._result.coefficients.values[k]) - from locpick._sampling.correction import get_sampling_correction + from .._sampling.correction import get_sampling_correction log_probs = _scl_log_probs_dispatch( beta_use, @@ -211,7 +211,7 @@ def probabilities(self, data=None, beta=None): # Systematic utility with sampling correction utilities = (dm @ beta).reshape(n_obs, n_alts) - from locpick._sampling.correction import apply_sampling_correction + from .._sampling.correction import apply_sampling_correction utilities = apply_sampling_correction(utilities, arrays) @@ -265,7 +265,7 @@ def utilities(self, data=None, beta=None): V = (dm @ beta).reshape(n_obs, n_alts) # Add sampling correction if present - from locpick._sampling.correction import apply_sampling_correction + from .._sampling.correction import apply_sampling_correction V = apply_sampling_correction(V, arrays) @@ -297,7 +297,7 @@ def simulate(self, data=None, n_draws: int = 1, seed: Optional[int] = None) -> p Simulated choices with columns ``draw``, ``obs_id``, ``alt_id``, and ``probability``. """ - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if self._arrays is None: raise RuntimeError("Model must be estimated before simulation.") @@ -381,7 +381,7 @@ def marginal_effect(self, data=None, variable: Optional[str] = None) -> pd.Serie pd.Series Direct marginal effects, indexed by (obs_id, alt_id). """ - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if self._arrays is None: raise RuntimeError("Model must be estimated before computing marginal effects.") @@ -431,7 +431,7 @@ def cross_marginal_effect(self, data=None, variable: Optional[str] = None) -> pd pd.Series Cross-marginal effects, indexed by (obs_id, alt_id). """ - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if self._arrays is None: raise RuntimeError("Model must be estimated before computing marginal effects.") @@ -481,7 +481,7 @@ def elasticity(self, data=None, variable: Optional[str] = None) -> pd.Series: pd.Series Direct elasticities, indexed by (obs_id, alt_id). """ - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if self._arrays is None: raise RuntimeError("Model must be estimated before computing elasticities.") @@ -529,7 +529,7 @@ def cross_elasticity(self, data=None, variable: Optional[str] = None) -> pd.Seri pd.Series Cross-elasticities, indexed by (obs_id, alt_id). """ - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if self._arrays is None: raise RuntimeError("Model must be estimated before computing elasticities.") @@ -575,7 +575,7 @@ def covariance_robust(self, data=None) -> np.ndarray: np.ndarray, shape (n_parameters, n_parameters) Sandwich (robust) covariance matrix. """ - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if self._arrays is None: raise RuntimeError("Model must be estimated first.") @@ -614,7 +614,7 @@ def covariance_clustered(self, data=None, groups=None) -> np.ndarray: np.ndarray, shape (n_parameters, n_parameters) Cluster-robust covariance matrix. """ - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if self._arrays is None: raise RuntimeError("Model must be estimated first.") @@ -713,7 +713,7 @@ def _observation_scores(self, arrays) -> np.ndarray: self._observation_scores_cache[cache_key] = scores return scores - from locpick._kernels.mnl_numpy import mnl_observation_scores_numpy + from .._kernels.mnl_numpy import mnl_observation_scores_numpy dm = np.asarray(arrays.design_matrix, dtype=np.float64) chosen = np.asarray(arrays.chosen, dtype=np.float64).reshape(arrays.n_obs, arrays.n_alts) @@ -726,7 +726,7 @@ def _observation_scores(self, arrays) -> np.ndarray: else: available = np.ones((n_obs, n_alts), dtype=np.float64) - from locpick._sampling.correction import get_sampling_correction + from .._sampling.correction import get_sampling_correction inclusion_probs = get_sampling_correction(arrays) @@ -781,7 +781,7 @@ def _build_objective(self, arrays: ChoiceArrays): is provided for debugging and benchmarking. """ if self._is_spatial: - from locpick._jax.builders import build_scl_objective + from .._jax.builders import build_scl_objective return build_scl_objective( arrays, self._edge_struct, self._allocation, self._edge_list @@ -792,7 +792,7 @@ def _build_objective(self, arrays: ChoiceArrays): backend = (self._backend or os.environ.get("LOCPICK_MNL_BACKEND", "")).lower() if backend == "numpy": return self._build_objective_numpy(arrays) - from locpick._jax.builders import build_mnl_objective + from .._jax.builders import build_mnl_objective return build_mnl_objective(arrays) @@ -823,13 +823,12 @@ def _build_objective_numpy(self, arrays: ChoiceArrays): Delegates to the canonical MNL kernels in :mod:`locpick._kernels.mnl_numpy` to avoid code duplication. """ - from locpick._kernels.mnl_numpy import ( + from .._kernels.mnl_numpy import ( mnl_gradient_numpy, mnl_log_likelihood_numpy, ) dm = arrays.design_matrix.astype(np.float64) - dm_sparse = getattr(arrays, "design_matrix_sparse", None) chosen = arrays.chosen.astype(np.float64) n_obs = arrays.n_obs n_alts = arrays.n_alts @@ -844,7 +843,7 @@ def _build_objective_numpy(self, arrays: ChoiceArrays): else: available = np.ones((n_obs, n_alts), dtype=np.float64) - from locpick._sampling.correction import get_sampling_correction + from .._sampling.correction import get_sampling_correction inclusion_probs = get_sampling_correction(arrays) @@ -858,7 +857,6 @@ def log_likelihood(beta: np.ndarray) -> float: n_alts=n_alts, weights=weights, inclusion_probs=inclusion_probs, - design_matrix_sparse=dm_sparse, ) def gradient(beta: np.ndarray) -> np.ndarray: @@ -871,10 +869,9 @@ def gradient(beta: np.ndarray) -> np.ndarray: n_alts=n_alts, weights=weights, inclusion_probs=inclusion_probs, - design_matrix_sparse=dm_sparse, ) - from locpick._jax.objective import Objective + from .._jax.objective import Objective return Objective(fn=log_likelihood, grad=gradient) diff --git a/locpick/models/nested.py b/locpick/models/nested.py index 7db4372..d008d36 100644 --- a/locpick/models/nested.py +++ b/locpick/models/nested.py @@ -39,15 +39,16 @@ import numpy as np import pandas as pd -from locpick._jax.objective import Objective -from locpick._solvers import Solver, SolverResult -from locpick.data.arrays import ChoiceArrays -from locpick.models._spatial import ( +from .._jax.objective import Objective +from .._solvers import Solver, SolverResult +from ..data.arrays import ChoiceArrays +from ..results.fit_result import FitResult +from ._spatial import ( EdgeStructure, _resolve_spatial_graph, naturalize_rho, ) -from locpick.models.base import ( +from .base import ( BaseChoiceModel, SpatialMixin, _compute_fit_statistics, @@ -55,7 +56,6 @@ _safe_inv, _sandwich_inv, ) -from locpick.results.fit_result import FitResult # --------------------------------------------------------------------------- # Nest specification @@ -266,7 +266,7 @@ def _nested_logit_probs_numpy( else: avail = np.ones((n_obs, n_alts), dtype=np.float64) - from locpick._kernels.constants import NEG_INF + from .._kernels.constants import NEG_INF utilities = np.where(avail > 0, utilities, NEG_INF) @@ -385,77 +385,6 @@ def _nested_logit_ll_numpy( return float(log_chosen.sum()) -def _nested_logit_gradient_numpy( - beta: np.ndarray, - alpha: np.ndarray, - design_matrix: np.ndarray, - chosen: np.ndarray, - nest_matrix: np.ndarray, - n_obs: int, - n_alts: int, - available: Optional[np.ndarray] = None, - inclusion_probs: Optional[np.ndarray] = None, - weights: Optional[np.ndarray] = None, -) -> np.ndarray: - """Compute nested logit gradient via finite differences (NumPy backend). - - This is a fallback gradient that uses finite differences. A proper - analytical gradient will be implemented in a future version. - - Parameters - ---------- - beta, alpha, design_matrix, chosen, nest_matrix, n_obs, n_alts, available, inclusion_probs, weights - See :func:`_nested_logit_ll_numpy`. - - Returns - ------- - np.ndarray, shape (k + n_nests,) - Gradient of the log-likelihood with respect to [beta, alpha]. - """ - params = np.concatenate([beta, alpha]) - eps = 1e-5 - n_params = len(params) - grad = np.zeros(n_params) - - for i in range(n_params): - params_plus = params.copy() - params_plus[i] += eps - params_minus = params.copy() - params_minus[i] -= eps - - beta_plus, alpha_plus = params_plus[: len(beta)], params_plus[len(beta) :] - beta_minus, alpha_minus = params_minus[: len(beta)], params_minus[len(beta) :] - - ll_plus = _nested_logit_ll_numpy( - beta_plus, - alpha_plus, - design_matrix, - chosen, - nest_matrix, - n_obs, - n_alts, - available=available, - inclusion_probs=inclusion_probs, - weights=weights, - ) - ll_minus = _nested_logit_ll_numpy( - beta_minus, - alpha_minus, - design_matrix, - chosen, - nest_matrix, - n_obs, - n_alts, - available=available, - inclusion_probs=inclusion_probs, - weights=weights, - ) - - grad[i] = (ll_plus - ll_minus) / (2 * eps) - - return grad - - # --------------------------------------------------------------------------- # NestedLogit model class # --------------------------------------------------------------------------- @@ -485,7 +414,7 @@ class NestedMNL(BaseChoiceModel, SpatialMixin): Examples -------- >>> from locpick import ChoiceTable, MultinomialLogit - >>> from locpick.models.nested import NestedLogit, NestSpec, NestingTree + >>> from .nested import NestedLogit, NestSpec, NestingTree >>> nests = NestingTree([ ... NestSpec("transit", alt_ids=[0, 1, 2]), ... NestSpec("auto", alt_ids=[3, 4]), @@ -551,7 +480,7 @@ def _pre_fit(self, arrays: ChoiceArrays) -> None: self._validate_graph_size(arrays) # Build per-nest EdgeStructure / EdgeDataJAX from the global graph. - from locpick._jax.data import EdgeDataJAX + from .._jax.data import EdgeDataJAX n_nests = self._nests.n_nests self._edge_structs = [] @@ -599,67 +528,19 @@ def _build_objective(self, arrays: ChoiceArrays) -> Objective: nest_matrix = self._nest_matrix if self._is_spatial: - from locpick._jax.builders import build_nested_scl_objective + from .._jax.builders import build_nested_scl_objective return build_nested_scl_objective(arrays, nest_matrix, self._edge_data_list) # Try JAX backend first (default when available) backend = (self._backend or os.environ.get("LOCPICK_NESTED_BACKEND", "")).lower() if backend != "numpy": - from locpick._jax.builders import build_nested_objective + from .._jax.builders import build_nested_objective return build_nested_objective(arrays, nest_matrix) - # NumPy backend - dm = np.asarray(arrays.design_matrix, dtype=np.float64) - chosen = np.asarray(arrays.chosen, dtype=np.float64) - n_obs = arrays.n_obs - n_alts = arrays.n_alts - available = arrays.available - weights = arrays.weights - - from locpick._sampling.correction import get_sampling_correction - - inclusion_probs = get_sampling_correction(arrays) - k = dm.shape[1] - - def ll_fn(params): - beta = params[:k] - alpha = params[k:] - return _nested_logit_ll_numpy( - beta, - alpha, - dm, - chosen, - nest_matrix, - n_obs, - n_alts, - available=available, - inclusion_probs=inclusion_probs, - weights=weights, - ) - - def grad_fn(params): - beta = params[:k] - alpha = params[k:] - return _nested_logit_gradient_numpy( - beta, - alpha, - dm, - chosen, - nest_matrix, - n_obs, - n_alts, - available=available, - inclusion_probs=inclusion_probs, - weights=weights, - ) - - return Objective.from_numpy( - ll_fn=ll_fn, - grad_fn=grad_fn, - param_names=list(arrays.param_names) - + [f"nest_{name}" for name in self._nests.nest_names], + raise NotImplementedError( + "NestedMNL NumPy backend has been removed. Use ChoiceModel (JAX backend)." ) def _build_fit_result( @@ -895,7 +776,7 @@ def utilities(self, data=None, beta=None): V = (dm @ beta).reshape(n_obs, n_alts) # Add sampling correction if present - from locpick._sampling.correction import apply_sampling_correction + from .._sampling.correction import apply_sampling_correction V = apply_sampling_correction(V, arrays) @@ -927,7 +808,7 @@ def simulate(self, data=None, n_draws: int = 1, seed: Optional[int] = None) -> p Simulated choices with columns ``draw``, ``obs_id``, ``alt_id``, and ``probability``. """ - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if self._arrays is None: raise RuntimeError("Model must be estimated before simulation.") @@ -996,7 +877,7 @@ def marginal_effect(self, data=None, variable: Optional[str] = None) -> pd.Serie pd.Series Direct marginal effects, indexed by (obs_id, alt_id). """ - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if self._arrays is None: raise RuntimeError("Model must be estimated before computing marginal effects.") @@ -1043,7 +924,7 @@ def cross_marginal_effect(self, data=None, variable: Optional[str] = None) -> pd pd.Series Cross-marginal effects, indexed by (obs_id, alt_id). """ - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if self._arrays is None: raise RuntimeError("Model must be estimated before computing marginal effects.") @@ -1094,7 +975,7 @@ def elasticity(self, data=None, variable: Optional[str] = None) -> pd.Series: pd.Series Direct elasticities, indexed by (obs_id, alt_id). """ - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if self._arrays is None: raise RuntimeError("Model must be estimated before computing elasticities.") @@ -1142,7 +1023,7 @@ def cross_elasticity(self, data=None, variable: Optional[str] = None) -> pd.Seri pd.Series Cross-elasticities, indexed by (obs_id, alt_id). """ - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if self._arrays is None: raise RuntimeError("Model must be estimated before computing elasticities.") @@ -1188,7 +1069,7 @@ def covariance_robust(self, data=None) -> np.ndarray: np.ndarray, shape (n_parameters, n_parameters) Sandwich (robust) covariance matrix. """ - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if self._arrays is None: raise RuntimeError("Model must be estimated first.") @@ -1227,7 +1108,7 @@ def covariance_clustered(self, data=None, groups=None) -> np.ndarray: np.ndarray, shape (n_parameters, n_parameters) Cluster-robust covariance matrix. """ - from locpick.data.choicetable import ChoiceTable + from ..data.choicetable import ChoiceTable if self._arrays is None: raise RuntimeError("Model must be estimated first.") @@ -1342,7 +1223,7 @@ def probabilities(self, data=None, beta=None, alpha=None): nest_matrix = self._nests.build_nest_matrix(alt_ids) # Resolve canonical sampling correction tensor. - from locpick._sampling.correction import get_sampling_correction + from .._sampling.correction import get_sampling_correction sampling_correction = get_sampling_correction(arrays) diff --git a/locpick/models/sar_mnl.py b/locpick/models/sar_mnl.py new file mode 100644 index 0000000..8ee1f92 --- /dev/null +++ b/locpick/models/sar_mnl.py @@ -0,0 +1,530 @@ +"""Spatial Autoregressive Multinomial Logit (SAR-MNL) model. + +Implements the pseudo maximum likelihood (PML) estimator from +Smirnov (2010): a spatial autoregressive lag in the systematic +utility of alternatives (spatial locations), with variance +normalisation by ``diag((I - ρW)^{-1})`` for consistency. + +The model specifies: + +.. math:: + + V_j = \\rho \\sum_k w_{jk} V_k + Z_j \\beta + X_{ij} \\gamma + +yielding reduced-form utilities :math:`V^* = (I - \\rho W)^{-1} +(Z\\beta + X\\gamma)`, normalised by :math:`D = \\text{diag}((I - +\\rho W)^{-1})`, with standard MNL choice probabilities. + +Estimation is via JAX autodiff through the spatial solve and +variance normalisation. No log-determinant Jacobian is needed +(this is pseudo-ML, not full ML). +""" + +from __future__ import annotations + +from typing import Optional, Union + +import numpy as np +import pandas as pd + +from .._solvers import Solver, SolverResult +from ..data.arrays import ChoiceArrays +from ..results.fit_result import FitResult +from ._spatial_weights import resolve_spatial_weights +from .base import ( + BaseChoiceModel, + _compute_fit_statistics, + _compute_null_ll, +) + + +class SARMNL(BaseChoiceModel): + r"""Spatial Autoregressive Multinomial Logit (SAR-MNL). + + Specifies a spatial autoregressive lag in the systematic utility + of alternatives (spatial locations): + + .. math:: + + V_j = \rho \sum_k w_{jk} V_k + Z_j \beta + X_{ij} \gamma + + yielding reduced-form utilities :math:`V^* = (I - \rho W)^{-1} + (Z\beta + X\gamma)`, normalised by :math:`D = \text{diag}((I - + \rho W)^{-1})`, with standard MNL choice probabilities. + + Estimation is via pseudo maximum likelihood (PML, Smirnov 2010) + with JAX autodiff through the spatial solve and variance + normalisation. No log-determinant Jacobian is needed (this is + pseudo-ML, not full ML). + + Parameters + ---------- + data : ChoiceTable or EstimationProblem + The choice data. + formula : str, optional + Formulaic formula string. + spec : ModelSpec, optional + ModelSpec object. + W : libpysal.graph.Graph, scipy.sparse, or np.ndarray + J×J spatial weights matrix connecting alternatives (locations). + Row-standardised internally. Zero diagonal. A ``libpysal.graph.Graph`` + is the preferred input type (matching bayespecon). ``scipy.sparse`` + and dense ``np.ndarray`` are also accepted and converted internally. + weights : str or array-like, optional + Observation weights. + availability : str or array-like, optional + Alternative availability. + solver : str or Solver, optional + Solver for PML optimisation. Default "lbfgs". + solver_options : dict, optional + backend : str, optional + estimator : str, optional + "auto" (default), "pml", "pml_cg", or "linearized_gmm". + ``auto`` selects ``pml`` (dense solve) for n_alts ≤ 2000 and + ``pml_cg`` (conjugate gradient) for larger alternative sets. + + Examples + -------- + >>> from locpick import ChoiceTable, SARMNL + >>> from libpysal.graph import Graph + >>> ct = ChoiceTable.from_tables(choosers, alternatives, chosen, sample_size=10) + >>> W = Graph.build_knn(gdf, k=7).transform("r") + >>> model = SARMNL(ct, formula="chosen ~ cost + time", W=W) + >>> result = model.fit() + >>> print(result.summary()) + """ + + def __init__( + self, + data, + formula: Optional[str] = None, + spec=None, + W=None, + weights: Optional[Union[str, np.ndarray]] = None, + availability: Optional[Union[str, np.ndarray]] = None, + solver: Union[str, Solver] = "lbfgs", + solver_options: Optional[dict] = None, + backend: Optional[str] = None, + estimator: str = "auto", + ): + super().__init__( + data=data, + formula=formula, + spec=spec, + solver=solver, + solver_options=solver_options, + backend=backend, + weights=weights, + availability=availability, + ) + if W is None: + raise ValueError("W (spatial weights matrix) is required for SARMNL.") + self._W_input = W + self._W_sparse = None # resolved at fit time + self._estimator = estimator + + # ------------------------------------------------------------------ + # Properties + # ------------------------------------------------------------------ + + @property + def W(self): + """The spatial weights matrix (libpysal.graph.Graph).""" + return self._W_input + + # ------------------------------------------------------------------ + # Estimation + # ------------------------------------------------------------------ + + def _pre_fit(self, arrays: ChoiceArrays) -> None: + """Resolve the spatial weights matrix.""" + self._W_sparse = resolve_spatial_weights( + self._W_input, arrays.n_alts, row_standardize=True + )[1] # get the CSR sparse + + def fit(self, **kwargs) -> FitResult: + """Estimate the model and return results. + + Dispatches to the PML estimator (JAX autodiff) or the + linearized GMM estimator based on the ``estimator`` setting. + """ + if self._estimator == "linearized_gmm": + return self._fit_linearized_gmm() + return super().fit(**kwargs) + + def _fit_linearized_gmm(self) -> FitResult: + """Two-step linearized GMM estimation (Carrión-Flores et al. 2018).""" + arrays = self._get_arrays() + self._arrays = arrays + self._pre_fit(arrays) + + from .._kernels.sar_mnl_numpy import fit_linearized_gmm + + result_dict = fit_linearized_gmm(arrays, self._W_sparse) + + beta = result_dict["beta"] + rho = result_dict["rho"] + se = result_dict["se"] + ll = result_dict["log_likelihood"] + + utility_param_names = list(arrays.param_names) + k = len(utility_param_names) + display_values = np.concatenate([beta, [rho]]) + display_names = utility_param_names + ["rho"] + model_type = "SAR-MNL (Linearized GMM)" + n_params = len(display_values) + + # SEs: first k are beta SEs, last is rho SE + std_errors = se[: k + 1] + + coefficients = pd.Series(display_values, index=display_names, name="coefficient") + std_err_series = pd.Series(std_errors, index=display_names, name="std_error") + ll_null = _compute_null_ll(arrays) + + stats = _compute_fit_statistics( + ll=ll, + ll_null=ll_null, + n_obs=arrays.n_obs, + n_params=n_params, + n_alts=arrays.n_alts, + coefficients=coefficients, + std_errors=std_err_series, + model_type=model_type, + solver_name="linearized_gmm", + solver_result_raw=result_dict, + ) + + self._result = FitResult(spec=self._spec, **stats) + self._clear_caches() + return self._result + + def _build_objective(self, arrays: ChoiceArrays): + """Build the PML objective using JAX. + + Auto-selects dense solve (n_alts ≤ 2000) or conjugate gradient + (n_alts > 2000) based on the estimator setting. + """ + from .._jax.sar_kernels import build_sar_mnl_objective + + # Auto-select estimator + if self._estimator == "auto": + if arrays.n_alts <= 2000: + self._estimator = "pml" + else: + self._estimator = "pml_cg" + + use_cg = self._estimator == "pml_cg" + return build_sar_mnl_objective(arrays, self._W_sparse, use_cg=use_cg) + + def _get_solver_inputs(self, arrays: ChoiceArrays): + """Get initial values, param names, bounds, fixed mask. + + Appends an unconstrained ``alpha_rho`` (initial 0.0) so the + spatial autoregressive parameter is estimated alongside the + utility coefficients. ``rho = tanh(alpha_rho) ∈ (-1, 1)``. + """ + x0, names, bounds, fixed_mask = super()._get_solver_inputs(arrays) + x0 = np.concatenate([x0, np.zeros(1)]) + names = list(names) + ["alpha_rho"] + return x0, names, bounds, fixed_mask + + def _build_fit_result(self, solver_result: SolverResult, arrays: ChoiceArrays) -> FitResult: + """Build a FitResult from solver output.""" + all_params = solver_result.coefficients + utility_param_names = list(arrays.param_names) + k = len(utility_param_names) + + # Layout: [beta_1..k, alpha_rho] + beta = all_params[:k] + alpha_rho = all_params[k] + rho = np.tanh(alpha_rho) # ρ ∈ (-1, 1) + + display_values = np.concatenate([beta, [rho]]) + display_names = utility_param_names + ["rho"] + model_type = "Spatial Autoregressive Multinomial Logit" + n_params = len(display_values) + + # Standard errors via Hessian + std_errors = np.full(n_params, np.nan) + try: + hess = self._compute_hessian(all_params) + se_unconstrained = self._compute_std_errors_from_hessian(hess) + # Delta method: SE(rho) = (1 - rho^2) * SE(alpha_rho) + se_rho = (1.0 - rho**2) * se_unconstrained[k] + std_errors = np.concatenate([se_unconstrained[:k], [se_rho]]) + except Exception: + if solver_result.hessian is not None: + try: + se = np.sqrt(np.maximum(np.diag(solver_result.hessian), 0)) + se[se == 0] = np.nan + se_rho = (1.0 - rho**2) * se[k] + std_errors = np.concatenate([se[:k], [se_rho]]) + except Exception: + pass + + coefficients = pd.Series(display_values, index=display_names, name="coefficient") + std_err_series = pd.Series(std_errors, index=display_names, name="std_error") + ll = solver_result.log_likelihood + ll_null = _compute_null_ll(arrays) + + stats = _compute_fit_statistics( + ll=ll, + ll_null=ll_null, + n_obs=arrays.n_obs, + n_params=n_params, + n_alts=arrays.n_alts, + coefficients=coefficients, + std_errors=std_err_series, + model_type=model_type, + solver_name=solver_result.solver_name, + solver_result_raw=solver_result.raw, + ) + + return FitResult(spec=self._spec, **stats) + + # ------------------------------------------------------------------ + # Prediction + # ------------------------------------------------------------------ + + def probabilities(self, data=None, beta=None, rho=None): + """Compute choice probabilities under the SAR-MNL model. + + Uses the full PML model: spatially-filtered + variance-normalised + utilities, then standard MNL softmax. + + Parameters + ---------- + data : ChoiceTable or None + Data to predict on. If None, uses estimation data. + beta : np.ndarray or None + Utility coefficients. If None, uses estimated values. + rho : float or None + Spatial autoregressive parameter. If None, uses estimated value. + + Returns + ------- + np.ndarray, shape (n_obs, n_alts) + Choice probabilities for each observation and alternative. + """ + from .._kernels.mnl_numpy import mnl_probs_numpy + + if self._arrays is None: + raise RuntimeError("Model must be estimated before prediction.") + + arrays = self._arrays + if data is not None: + arrays = data.to_arrays( + formula=self._spec.formula, + spec=self._spec if self._spec.formula is None else None, + ) + + n_obs = arrays.n_obs + n_alts = arrays.n_alts + k = arrays.design_matrix.shape[1] + + if beta is None: + coef_vals = np.asarray(self._result.coefficients.values, dtype=np.float64) + beta = coef_vals[:k] + if rho is None: + rho = float(self._result.coefficients.values[k]) + + dm = np.asarray(arrays.design_matrix, dtype=np.float64) + W_dense = np.asarray(self._W_sparse.toarray(), dtype=np.float64) + + # Base utilities + V_base = (dm @ beta).reshape(n_obs, n_alts) + + # Sampling correction + from .._sampling.correction import apply_sampling_correction + + V_base = apply_sampling_correction(V_base, arrays) + + # Spatial filter + variance normalisation + A = np.eye(n_alts) - rho * W_dense + V_filtered = np.linalg.solve(A, V_base.T).T + D = np.diag(np.linalg.inv(A)) + V_star = V_filtered / D[None, :] + + # Availability + if arrays.available is not None: + available = np.asarray(arrays.available, dtype=np.float64).reshape(n_obs, n_alts) + else: + available = np.ones((n_obs, n_alts), dtype=np.float64) + + return mnl_probs_numpy(V_star, available, inclusion_probs=None) + + def utilities(self, data=None, beta=None, rho=None): + """Compute spatially-filtered + variance-normalised utilities. + + Parameters + ---------- + data : ChoiceTable or None + Data to compute utilities on. If None, uses estimation data. + beta : np.ndarray or None + Utility coefficients. If None, uses estimated values. + rho : float or None + Spatial autoregressive parameter. If None, uses estimated value. + + Returns + ------- + np.ndarray, shape (n_obs, n_alts) + Spatially-filtered and variance-normalised utilities. + """ + if self._arrays is None: + raise RuntimeError("Model must be estimated before prediction.") + + arrays = self._arrays + if data is not None: + arrays = data.to_arrays( + formula=self._spec.formula, + spec=self._spec if self._spec.formula is None else None, + ) + + n_obs = arrays.n_obs + n_alts = arrays.n_alts + k = arrays.design_matrix.shape[1] + + if beta is None: + beta = np.asarray(self._result.coefficients.values[:k], dtype=np.float64) + if rho is None: + rho = float(self._result.coefficients.values[k]) + + dm = np.asarray(arrays.design_matrix, dtype=np.float64) + W_dense = np.asarray(self._W_sparse.toarray(), dtype=np.float64) + + V_base = (dm @ beta).reshape(n_obs, n_alts) + + from .._sampling.correction import apply_sampling_correction + + V_base = apply_sampling_correction(V_base, arrays) + + A = np.eye(n_alts) - rho * W_dense + V_filtered = np.linalg.solve(A, V_base.T).T + D = np.diag(np.linalg.inv(A)) + V_star = V_filtered / D[None, :] + + return V_star + + # ------------------------------------------------------------------ + # Marginal effects (LeSage & Pace 2009) + # ------------------------------------------------------------------ + + def marginal_effects(self, data=None, variable: Optional[str] = None): + """Compute average direct, indirect, and total marginal effects. + + In the SAR-MNL model, a change in an attribute of alternative + *j* affects not only *j*'s utility but also neighbouring + alternatives through the spatial multiplier + :math:`(I - \\rho W)^{-1}`. + + Following LeSage & Pace (2009), the marginal effect of variable + *r* on the probability of choosing alternative *k* is an + :math:`J \\times J` matrix. Summary measures are: + + - **Direct effect**: average of diagonal elements (impact on + own alternative). + - **Indirect effect**: average of off-diagonal row sums + (spillover to neighbouring alternatives). + - **Total effect**: direct + indirect. + + Parameters + ---------- + data : ChoiceTable or None + Data to compute marginal effects on. If None, uses + estimation data. + variable : str + Name of the variable to compute marginal effects for. + + Returns + ------- + dict + Dictionary with keys ``"direct"``, ``"indirect"``, + ``"total"``, each mapping to a ``pd.Series`` indexed by + alternative ID. + """ + if self._arrays is None: + raise RuntimeError("Model must be estimated before computing marginal effects.") + + from ..data.choicetable import ChoiceTable + + ct = self._data + arrays = self._arrays + if data is not None: + if not isinstance(data, ChoiceTable): + raise TypeError("data must be a ChoiceTable") + arrays = data.to_arrays( + formula=self._spec.formula, + spec=self._spec if self._spec.formula is None else None, + ) + ct = data + + n_obs = arrays.n_obs + n_alts = arrays.n_alts + k = arrays.design_matrix.shape[1] + + # Get estimated parameters + coef_vals = np.asarray(self._result.coefficients.values, dtype=np.float64) + beta = coef_vals[:k] + rho = float(coef_vals[k]) + + # Get the beta for the requested variable + param_names = list(arrays.param_names) + if variable not in param_names: + raise ValueError(f"Variable '{variable}' not found in parameters: {param_names}") + beta_r = beta[param_names.index(variable)] + + # Compute probabilities + probs = self.probabilities(data=data) + + # Spatial multiplier (I - rho*W)^{-1} + W_dense = np.asarray(self._W_sparse.toarray(), dtype=np.float64) + A = np.eye(n_alts) - rho * W_dense + Z_mat = np.linalg.inv(A) # (n_alts, n_alts) + + # Marginal effect matrix for variable r, alternative k: + # ME_{k,r} = P_k * (beta_r * Z_{kk} - sum_l P_l * beta_r * Z_{lk}) + # = beta_r * P_k * (Z_{kk} - sum_l P_l * Z_{lk}) + # But P varies across choosers. For the average marginal effect, + # we average over choosers: + # AME_{k,r} = beta_r * avg(P_k) * (Z_{kk} - sum_l avg(P_l) * Z_{lk}) + + # Compute the J×J marginal effect matrix (averaged over choosers) + # ME[j, k] = beta_r * avg_probs[k] * (Z[k, j] - sum_l avg_probs[l] * Z[l, j]) + # But the standard LeSage-Pace formulation for MNL is: + # dP_k/dX_j = P_k * (delta_{kj} - P_j) * beta_r * Z[j, ...] + # This is complex — we use the simpler average approach: + # + # For each alternative k, the direct effect is: + # dP_k/dX_k = beta_r * Z[k,k] * P_k * (1 - P_k) + # The indirect (spillover) effect from j to k (j != k) is: + # dP_k/dX_j = -beta_r * Z[k,j] * P_k * P_j + # But with the spatial multiplier, Z replaces the identity. + + # Direct effects: average over choosers of + # beta_r * Z[k,k] * P_ik * (1 - P_ik) + direct = np.zeros(n_alts) + indirect = np.zeros(n_alts) + for k_alt in range(n_alts): + # Direct: own-alternative effect + direct[k_alt] = ( + beta_r * Z_mat[k_alt, k_alt] * np.mean(probs[:, k_alt] * (1 - probs[:, k_alt])) + ) + # Indirect: spillover from neighbours + # Sum over j != k of dP_k/dX_j = -beta_r * sum_{j!=k} Z[k,j] * P_k * P_j + for j_alt in range(n_alts): + if j_alt != k_alt: + indirect[k_alt] += ( + -beta_r * Z_mat[k_alt, j_alt] * np.mean(probs[:, k_alt] * probs[:, j_alt]) + ) + + total = direct + indirect + + # Get alternative IDs from the data + df = ct.to_frame() + alt_ids = df[ct.alt_id_col].values.reshape(n_obs, n_alts)[0] + + return { + "direct": pd.Series(direct, index=alt_ids, name=f"direct_{variable}"), + "indirect": pd.Series(indirect, index=alt_ids, name=f"indirect_{variable}"), + "total": pd.Series(total, index=alt_ids, name=f"total_{variable}"), + } diff --git a/locpick/models/scl.py b/locpick/models/scl.py index 3f26b43..8207789 100644 --- a/locpick/models/scl.py +++ b/locpick/models/scl.py @@ -23,17 +23,17 @@ import numpy as np from scipy.special import logsumexp -from locpick._kernels.constants import NEG_INF -from locpick.models._spatial import ( +from .._kernels.constants import NEG_INF +from ._spatial import ( EdgeStructure as EdgeStructure, ) -from locpick.models._spatial import ( +from ._spatial import ( _resolve_spatial_graph as _resolve_spatial_graph, ) -from locpick.models._spatial import ( +from ._spatial import ( constrain_rho as constrain_rho, ) -from locpick.models._spatial import ( +from ._spatial import ( naturalize_rho as naturalize_rho, ) @@ -393,169 +393,6 @@ def _scl_ll_numpy( return float(chosen_log_probs.sum()) -def _scl_gradient_numpy( - beta: np.ndarray, - rho: float, - design_matrix: np.ndarray, - chosen: np.ndarray, - allocation: np.ndarray, - edge_list: list[tuple[int, int]], - n_obs: int, - n_alts: int, - available: Optional[np.ndarray] = None, - inclusion_probs: Optional[np.ndarray] = None, - weights: Optional[np.ndarray] = None, -) -> np.ndarray: - """Compute SCL gradient via finite differences (NumPy backend). - - This is a fallback gradient that uses central finite differences. - A proper analytical gradient will be implemented in a future version. - - Parameters - ---------- - beta, rho, design_matrix, chosen, allocation, edge_list, n_obs, n_alts, - available, inclusion_probs, weights - See :func:`_scl_ll_numpy`. - - Returns - ------- - np.ndarray, shape (k + 1,) - Gradient of the log-likelihood with respect to [beta, alpha_rho]. - """ - params = np.concatenate([beta, [rho]]) - eps = 1e-5 - n_params = len(params) - grad = np.zeros(n_params) - - for i in range(n_params): - params_plus = params.copy() - params_plus[i] += eps - params_minus = params.copy() - params_minus[i] -= eps - - ll_plus = _scl_ll_numpy( - params_plus[: len(beta)], - params_plus[len(beta)], - design_matrix, - chosen, - allocation, - edge_list, - n_obs, - n_alts, - available=available, - inclusion_probs=inclusion_probs, - weights=weights, - ) - ll_minus = _scl_ll_numpy( - params_minus[: len(beta)], - params_minus[len(beta)], - design_matrix, - chosen, - allocation, - edge_list, - n_obs, - n_alts, - available=available, - inclusion_probs=inclusion_probs, - weights=weights, - ) - - grad[i] = (ll_plus - ll_minus) / (2 * eps) - - return grad - - -# --------------------------------------------------------------------------- -# Dispatch layer (kept as a stable entry point for predict() and benchmarks) -# --------------------------------------------------------------------------- - - -def _scl_log_probs_dispatch( - beta: np.ndarray, - rho: float, - design_matrix: np.ndarray, - allocation: np.ndarray, - edge_list: list[tuple[int, int]], - n_obs: int, - n_alts: int, - available: Optional[np.ndarray] = None, - inclusion_probs: Optional[np.ndarray] = None, - edge_struct: Optional[EdgeStructure] = None, -) -> np.ndarray: - """Compute SCL log-probabilities using the NumPy reference kernel.""" - return _scl_log_probs_numpy( - beta, - rho, - design_matrix, - allocation, - edge_list, - n_obs, - n_alts, - available=available, - inclusion_probs=inclusion_probs, - ) - - -def _scl_ll_dispatch( - beta: np.ndarray, - rho: float, - design_matrix: np.ndarray, - chosen: np.ndarray, - allocation: np.ndarray, - edge_list: list[tuple[int, int]], - n_obs: int, - n_alts: int, - available: Optional[np.ndarray] = None, - inclusion_probs: Optional[np.ndarray] = None, - weights: Optional[np.ndarray] = None, - edge_struct: Optional[EdgeStructure] = None, -) -> float: - """Compute SCL log-likelihood using the NumPy reference kernel.""" - return _scl_ll_numpy( - beta, - rho, - design_matrix, - chosen, - allocation, - edge_list, - n_obs, - n_alts, - available=available, - inclusion_probs=inclusion_probs, - weights=weights, - ) - - -def _scl_gradient_dispatch( - beta: np.ndarray, - rho: float, - design_matrix: np.ndarray, - chosen: np.ndarray, - allocation: np.ndarray, - edge_list: list[tuple[int, int]], - n_obs: int, - n_alts: int, - available: Optional[np.ndarray] = None, - inclusion_probs: Optional[np.ndarray] = None, - weights: Optional[np.ndarray] = None, - edge_struct: Optional[EdgeStructure] = None, -) -> np.ndarray: - """Compute SCL gradient via finite differences over the NumPy LL.""" - return _scl_gradient_numpy( - beta, - rho, - design_matrix, - chosen, - allocation, - edge_list, - n_obs, - n_alts, - available=available, - inclusion_probs=inclusion_probs, - weights=weights, - ) - - # --------------------------------------------------------------------------- -# (Public ``SCL`` factory removed — construct ``MNL`` / ``NestedMNL`` / -# ``MixedMNL`` / ``MixedNestedMNL`` directly with ``graph=`` instead.) +# (Public ``SCL`` factory removed — construct ``ChoiceModel`` directly with +# ``graph=`` instead.) diff --git a/locpick/results/diagnostics.py b/locpick/results/diagnostics.py index e31c9be..81168e0 100644 --- a/locpick/results/diagnostics.py +++ b/locpick/results/diagnostics.py @@ -14,7 +14,7 @@ from scipy.linalg import cho_factor, cho_solve if TYPE_CHECKING: - from locpick.results.fit_result import FitResult + from .fit_result import FitResult @dataclass diff --git a/locpick/spec/__init__.pyi b/locpick/spec/__init__.pyi index 45503c3..abf2a43 100644 --- a/locpick/spec/__init__.pyi +++ b/locpick/spec/__init__.pyi @@ -3,9 +3,6 @@ from . import terms as terms from .model_spec import ( ModelSpec as ModelSpec, ) -from .model_spec import ( - ParamDistribution as ParamDistribution, -) from .terms import ( InteractionTerm as InteractionTerm, ) diff --git a/locpick/spec/model_spec.py b/locpick/spec/model_spec.py index 3809efe..27aae3a 100644 --- a/locpick/spec/model_spec.py +++ b/locpick/spec/model_spec.py @@ -7,7 +7,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any, Literal, Optional, Union +from typing import Any, Literal, Optional from .terms import InteractionTerm, ScopedTerm, interaction @@ -198,7 +198,7 @@ def _build_from_scoped_terms(self, data) -> Any: import numpy as np import pandas as pd - from locpick.data import ChoiceArrays + from ..data import ChoiceArrays df = data.to_frame() n_obs = data.n_observations @@ -328,44 +328,3 @@ def __repr__(self) -> str: elif self.scoped_terms: return f"ModelSpec(scoped_terms={self.scoped_terms!r})" return "ModelSpec()" - - -# Backward-compat re-export. The real implementation is in locpick.models.mixed. -@dataclass -class ParamDistribution: - """Distribution specification for a random parameter (mixed logit). - - .. deprecated:: - Import from ``locpick.mixed`` or ``locpick`` instead. - This stub is kept for backward compatibility. - - Parameters - ---------- - distribution : str - Distribution name: ``"normal"``, ``"lognormal"``, - ``"triangular"``, or ``"uniform"``. - param : ParamRef or str - The parameter to assign a random distribution to. - """ - - distribution: str - param: Union[Any, str] - - def __post_init__(self) -> None: - valid = {"normal", "lognormal", "triangular", "uniform"} - if self.distribution not in valid: - raise ValueError( - f"Unknown distribution '{self.distribution}'. Must be one of {valid}." - ) - - @property - def param_name(self) -> str: - """Return the parameter name as a string.""" - if hasattr(self.param, "name"): - return self.param.name - return str(self.param) - - @property - def n_params(self) -> int: - """Number of distribution parameters (always 2: mean and spread).""" - return 2 diff --git a/tests/test_estimation_problem.py b/tests/test_estimation_problem.py index 46bc412..5c84f0f 100644 --- a/tests/test_estimation_problem.py +++ b/tests/test_estimation_problem.py @@ -9,7 +9,7 @@ import pytest from locpick import ( - MNL, + ChoiceModel, ChoiceTable, EstimationProblem, FitResult, @@ -182,7 +182,7 @@ def test_estimate_from_problem(self): """MultinomialLogit.fit() works with EstimationProblem.""" ct = make_choice_table() problem = EstimationProblem.from_choice_table(ct, formula="cost + time") - model = MNL(data=ct, problem=problem) + model = ChoiceModel(data=ct, problem=problem) result = model.fit() assert isinstance(result, FitResult) assert result.coefficients.shape[0] == 2 @@ -193,12 +193,12 @@ def test_problem_matches_legacy_result(self): ct = make_choice_table() # Legacy path - model_legacy = MNL(ct, formula="cost + time") + model_legacy = ChoiceModel(ct, formula="cost + time") result_legacy = model_legacy.fit() # Problem path problem = EstimationProblem.from_choice_table(ct, formula="cost + time") - model_problem = MNL(data=ct, problem=problem) + model_problem = ChoiceModel(data=ct, problem=problem) result_problem = model_problem.fit() # Coefficients should be very close (same data, same solver) @@ -223,7 +223,7 @@ def test_problem_with_initial_values(self): param_names=problem.param_names, param_initial=[0.1, -0.1], ) - model = MNL(data=ct, problem=problem) + model = ChoiceModel(data=ct, problem=problem) result = model.fit() assert isinstance(result, FitResult) assert np.isfinite(result.log_likelihood) @@ -234,7 +234,7 @@ def test_problem_ignores_formula_and_spec(self): problem = EstimationProblem.from_choice_table(ct, formula="cost + time") # Pass problem + formula — formula should be ignored - model = MNL(data=ct, problem=problem, formula="ignored ~ x") + model = ChoiceModel(data=ct, problem=problem, formula="ignored ~ x") result = model.fit() # Should still have 2 params (from problem), not whatever "ignored ~ x" would give assert result.coefficients.shape[0] == 2 @@ -244,7 +244,7 @@ def test_problem_requires_data(self): ct = make_choice_table() problem = EstimationProblem.from_choice_table(ct, formula="cost + time") # data is still required (it's positional) - model = MNL(data=ct, problem=problem) + model = ChoiceModel(data=ct, problem=problem) assert model._problem is problem @@ -267,7 +267,7 @@ def test_fixed_parameter_stays_at_initial(self): param_initial=[0.5, 0.0], param_fixed=[True, False], ) - model = MNL(data=ct, problem=problem) + model = ChoiceModel(data=ct, problem=problem) result = model.fit() # The first parameter (cost) should be close to 0.5 (fixed) assert abs(result.coefficients.iloc[0] - 0.5) < 1e-6 @@ -282,7 +282,7 @@ def test_bounds_passed_to_solver(self): param_names=problem.param_names, param_bounds=[(-5.0, 5.0), (-10.0, 10.0)], ) - model = MNL(data=ct, problem=problem) + model = ChoiceModel(data=ct, problem=problem) result = model.fit() # Should converge within bounds assert -5.0 <= result.coefficients.iloc[0] <= 5.0 @@ -301,7 +301,7 @@ def test_fixed_parameter_with_optimistix(self): param_fixed=[True, False], ) - model = MNL(data=ct, problem=problem, solver="optimistix") + model = ChoiceModel(data=ct, problem=problem, solver="optimistix") result = model.fit() assert abs(result.coefficients.iloc[0] - 0.5) < 1e-6 diff --git a/tests/test_inference.py b/tests/test_inference.py index 01b5f59..e1321cf 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -8,7 +8,7 @@ import numpy.testing as npt import pandas as pd -from locpick import MNL, ChoiceTable +from locpick import ChoiceModel, ChoiceTable # --------------------------------------------------------------------------- # Fixtures @@ -64,7 +64,7 @@ class TestObservationScores: def test_scores_shape(self): """Observation scores should have shape (n_obs, n_params).""" ct, _, _ = _make_simple_data() - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() arrays = ct.to_arrays(formula="cost + time - 1") @@ -75,7 +75,7 @@ def test_scores_shape(self): def test_scores_sum_to_gradient(self): """Sum of observation scores should equal the full gradient.""" ct, _, _ = _make_simple_data() - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() arrays = ct.to_arrays(formula="cost + time - 1") @@ -88,7 +88,7 @@ def test_scores_sum_to_gradient(self): def test_scores_are_finite(self): """All observation scores should be finite.""" ct, _, _ = _make_simple_data() - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() arrays = ct.to_arrays(formula="cost + time - 1") @@ -114,7 +114,7 @@ class TestRobustCovariance: def test_robust_covariance_shape(self): """Robust covariance should be (n_params, n_params).""" ct, _, _ = _make_simple_data() - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() cov = model.covariance_robust(ct) @@ -124,7 +124,7 @@ def test_robust_covariance_shape(self): def test_robust_covariance_positive_diagonal(self): """Robust covariance diagonal should be positive (variances).""" ct, _, _ = _make_simple_data() - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() cov = model.covariance_robust(ct) @@ -134,7 +134,7 @@ def test_robust_covariance_positive_diagonal(self): def test_robust_covariance_symmetric(self): """Robust covariance should be approximately symmetric.""" ct, _, _ = _make_simple_data() - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() cov = model.covariance_robust(ct) @@ -144,7 +144,7 @@ def test_robust_covariance_symmetric(self): def test_robust_standard_errors(self): """Robust standard errors should be positive and finite.""" ct, _, _ = _make_simple_data() - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() se = model.std_errors_robust(ct) @@ -166,7 +166,7 @@ class TestClusteredCovariance: def test_clustered_covariance_shape(self): """Cluster-robust covariance should be (n_params, n_params).""" ct, _, _ = _make_simple_data() - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() # Create cluster groups @@ -179,7 +179,7 @@ def test_clustered_covariance_shape(self): def test_clustered_covariance_positive_diagonal(self): """Cluster-robust covariance diagonal should be positive.""" ct, _, _ = _make_simple_data() - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() groups = np.repeat([0, 1, 2, 3, 4], ct.n_observations // 5) @@ -191,7 +191,7 @@ def test_clustered_covariance_positive_diagonal(self): def test_clustered_standard_errors(self): """Cluster-robust standard errors should be positive and finite.""" ct, _, _ = _make_simple_data() - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() groups = np.repeat([0, 1, 2, 3, 4], ct.n_observations // 5) @@ -207,7 +207,7 @@ def test_clustered_larger_than_default(self): """Cluster-robust SEs should typically be >= default SEs (due to within-cluster correlation).""" ct, _, _ = _make_simple_data() - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") result = model.fit() # With many small clusters, clustered SEs should be similar to default @@ -231,7 +231,7 @@ class TestCovarianceComparison: def test_robust_vs_default_se_order(self): """Robust SEs should be in the same order of magnitude as default SEs.""" ct, _, _ = _make_simple_data() - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") result = model.fit() se_default = result.std_errors.values diff --git a/tests/test_marginal_effects_wtp.py b/tests/test_marginal_effects_wtp.py index 1572d99..782ae5a 100644 --- a/tests/test_marginal_effects_wtp.py +++ b/tests/test_marginal_effects_wtp.py @@ -5,7 +5,7 @@ import pandas as pd import pytest -from locpick import MNL, ChoiceTable +from locpick import ChoiceModel, ChoiceTable # --------------------------------------------------------------------------- # Helpers @@ -52,7 +52,7 @@ class TestMarginalEffects: def test_marginal_effect_shape(self): """Marginal effects should have same length as observations * alternatives.""" ct, _, _, _ = _make_simple_dataset(n_obs=50, n_alts=5) - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() me = model.marginal_effect(variable="cost") @@ -61,7 +61,7 @@ def test_marginal_effect_shape(self): def test_marginal_effect_sign(self): """For a negative coefficient, direct ME should be negative.""" ct, _, _, _ = _make_simple_dataset(n_obs=100, n_alts=4) - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") result = model.fit() me = model.marginal_effect(variable="cost") @@ -75,7 +75,7 @@ def test_marginal_effect_sign(self): def test_cross_marginal_effect_sign(self): """Cross ME should have opposite sign to direct ME.""" ct, _, _, _ = _make_simple_dataset(n_obs=100, n_alts=4) - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() me = model.marginal_effect(variable="cost") @@ -87,7 +87,7 @@ def test_cross_marginal_effect_sign(self): def test_marginal_effect_vs_elasticity(self): """Elasticity = ME * x (for direct effects).""" ct, _, _, _ = _make_simple_dataset(n_obs=50, n_alts=4) - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() me = model.marginal_effect(variable="cost") @@ -103,7 +103,7 @@ def test_marginal_effect_vs_elasticity(self): def test_marginal_effect_on_new_data(self): """ME should work on out-of-sample data.""" ct, _, _, _ = _make_simple_dataset(n_obs=100, n_alts=4) - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() # New data @@ -133,7 +133,7 @@ def test_marginal_effect_on_new_data(self): def test_average_marginal_effect_aggregations(self): """AME helpers should aggregate per-obs ME consistently.""" ct, _, _, _ = _make_simple_dataset(n_obs=80, n_alts=4) - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() me = model.marginal_effect(variable="cost") @@ -170,7 +170,7 @@ class TestWTP: def test_wtp_basic(self): """WTP should compute -beta_time / beta_cost.""" ct, _, _, _ = _make_simple_dataset(n_obs=200, n_alts=4) - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") result = model.fit() wtp = result.wtp(numerator="time", denominator="cost") @@ -183,7 +183,7 @@ def test_wtp_basic(self): def test_wtp_has_standard_error(self): """WTP should include a standard error.""" ct, _, _, _ = _make_simple_dataset(n_obs=200, n_alts=4) - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") result = model.fit() wtp = result.wtp(numerator="time", denominator="cost") @@ -194,7 +194,7 @@ def test_wtp_has_standard_error(self): def test_wtp_has_t_stat_and_p_value(self): """WTP should include t-statistic and p-value.""" ct, _, _, _ = _make_simple_dataset(n_obs=200, n_alts=4) - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") result = model.fit() wtp = result.wtp(numerator="time", denominator="cost") @@ -206,7 +206,7 @@ def test_wtp_has_t_stat_and_p_value(self): def test_wtp_invalid_numerator_raises(self): """WTP should raise for invalid numerator.""" ct, _, _, _ = _make_simple_dataset(n_obs=50, n_alts=4) - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") result = model.fit() with pytest.raises(ValueError, match="Numerator 'income' not found"): @@ -215,7 +215,7 @@ def test_wtp_invalid_numerator_raises(self): def test_wtp_invalid_denominator_raises(self): """WTP should raise for invalid denominator.""" ct, _, _, _ = _make_simple_dataset(n_obs=50, n_alts=4) - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") result = model.fit() with pytest.raises(ValueError, match="Denominator 'rent' not found"): @@ -224,7 +224,7 @@ def test_wtp_invalid_denominator_raises(self): def test_vot_is_wtp_alias(self): """VOT should be equivalent to WTP(time, cost).""" ct, _, _, _ = _make_simple_dataset(n_obs=200, n_alts=4) - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") result = model.fit() vot = result.vot(time_var="time", cost_var="cost") @@ -243,7 +243,7 @@ def test_wtp_with_custom_denominator(self): alternatives, chosen_alternatives=pd.Series(choices, index=choosers.index), ) - model = MNL(ct2, formula="rent + time - 1") + model = ChoiceModel(ct2, formula="rent + time - 1") result = model.fit() wtp = result.wtp(numerator="time", denominator="rent") diff --git a/tests/test_mixed_nested.py b/tests/test_mixed_nested.py index 952b19e..c83a306 100644 --- a/tests/test_mixed_nested.py +++ b/tests/test_mixed_nested.py @@ -1,12 +1,10 @@ """Tests for the Mixed Nested Logit model.""" import numpy as np -import pytest -from locpick import ChoiceTable +from locpick import ChoiceModel, ChoiceTable from locpick.dgp import simulate_mixed_nested_logit from locpick.models.mixed import ParamDistribution -from locpick.models.mixed_nested import MixedNestedMNL from locpick.models.nested import NestingTree, NestSpec # --------------------------------------------------------------------------- @@ -65,35 +63,35 @@ def make_mixed_nested_data(n_obs=500, n_alts=4, seed=42): class TestMixedNestedMNL: """Tests for the MixedNestedMNL model class.""" - def test_requires_nests(self): - """MixedNestedMNL should require nests argument.""" + def test_without_nests_is_not_nested(self): + """ChoiceModel without nests should not raise — it's just mixed.""" ct, nests = make_mixed_nested_data() - with pytest.raises(ValueError, match="nests"): - MixedNestedMNL( - ct, - formula="cost + time - 1", - random_params={"time": ParamDistribution("normal", "time")}, - n_draws=50, - ) + model = ChoiceModel( + ct, + formula="cost + time - 1", + random_params={"time": ParamDistribution("normal", "time")}, + n_draws=50, + ) + assert not model._is_nested - def test_requires_random_params(self): - """MixedNestedMNL should require random_params argument.""" + def test_without_random_params_is_not_mixed(self): + """ChoiceModel without random_params should not raise — it's just nested.""" ct, nests = make_mixed_nested_data() - with pytest.raises(ValueError, match="random"): - MixedNestedMNL( - ct, - formula="cost + time - 1", - nests=nests, - n_draws=50, - ) + model = ChoiceModel( + ct, + formula="cost + time - 1", + nests=nests, + n_draws=50, + ) + assert not model._is_mixed def test_mixed_nested_logit_estimation(self): """MixedNestedMNL should estimate and return a FitResult.""" ct, nests = make_mixed_nested_data(n_obs=200) - model = MixedNestedMNL( + model = ChoiceModel( ct, formula="cost + time - 1", nests=nests, @@ -116,7 +114,7 @@ def test_mixed_nested_logit_multiple_random_params(self): """MixedNestedMNL should handle multiple random parameters.""" ct, nests = make_mixed_nested_data(n_obs=200) - model = MixedNestedMNL( + model = ChoiceModel( ct, formula="cost + time - 1", nests=nests, @@ -142,7 +140,7 @@ def test_mixed_nested_logit_with_fixed_params(self): """MixedNestedMNL should handle mix of fixed and random parameters.""" ct, nests = make_mixed_nested_data(n_obs=200) - model = MixedNestedMNL( + model = ChoiceModel( ct, formula="cost + time - 1", nests=nests, @@ -162,7 +160,7 @@ def test_mixed_nested_logit_lognormal(self): """MixedNestedMNL should work with lognormal distribution.""" ct, nests = make_mixed_nested_data(n_obs=200) - model = MixedNestedMNL( + model = ChoiceModel( ct, formula="cost + time - 1", nests=nests, diff --git a/tests/test_mnl.py b/tests/test_mnl.py index adf3502..8283880 100644 --- a/tests/test_mnl.py +++ b/tests/test_mnl.py @@ -6,7 +6,7 @@ import numpy.testing as npt import pytest -from locpick import MNL, dgp +from locpick import ChoiceModel, dgp """ These are tests for the refactored locpick MNL codebase. @@ -42,7 +42,7 @@ def test_mnl(obs, alts): """ formula = "obsval + altval - 1" ct = ChoiceTable.from_tables(obs, alts, chosen_alternatives="choice") - m = MNL(ct, formula=formula) + m = ChoiceModel(ct, formula=formula) r = m.fit() assert len(r.coefficients) == 2 @@ -54,7 +54,7 @@ def test_mnl_estimation(obs, alts): """ formula = "obsval + altval - 1" ct = ChoiceTable.from_tables(obs, alts, chosen_alternatives="choice") - result = MNL(ct, formula=formula).fit() + result = ChoiceModel(ct, formula=formula).fit() assert np.isfinite(result.log_likelihood) assert np.isfinite(result.coefficients.to_numpy()).all() @@ -65,7 +65,7 @@ def test_mnl_prediction(obs, alts): """ ct = ChoiceTable.from_tables(obs, alts, chosen_alternatives="choice", sample_size=5) - m = MNL(ct, formula="obsval + altval - 1") + m = ChoiceModel(ct, formula="obsval + altval - 1") m.fit() probs = m.probabilities(ct) @@ -81,7 +81,7 @@ def _fit_v2(dataset, backend, monkeypatch): monkeypatch.delenv("LOCPICK_MNL_BACKEND", raising=False) if backend != "jax": monkeypatch.setenv("LOCPICK_MNL_BACKEND", backend) - model = MNL(dataset.choice_table, FORMULA) + model = ChoiceModel(dataset.choice_table, FORMULA) result = model.fit() return result.coefficients @@ -94,7 +94,7 @@ def test_mnl_parameter_recovery_with_pairwise_variable(): interaction_params={"obs_x_alt": 1.1}, seed=1234, ) - model = MNL(dataset.choice_table, FORMULA) + model = ChoiceModel(dataset.choice_table, FORMULA) estimated = model.fit().coefficients # With 4 000 observations MLE is consistent; allow 10 % relative tolerance. @@ -229,7 +229,7 @@ class TestAvailability: def test_unavailable_alt_zero_probability(self): """Unavailable alternatives should have zero probability.""" ct, _, _, _, avail_arr = _make_dataset_with_availability() - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") model.fit() probs = model.probabilities(ct) @@ -244,7 +244,7 @@ def test_unavailable_alt_zero_probability(self): def test_available_alts_sum_to_one(self): """Probabilities of available alternatives should sum to 1.""" ct, _, _, _, avail_arr = _make_dataset_with_availability() - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") model.fit() probs = model.probabilities(ct) @@ -261,7 +261,7 @@ def test_available_alts_sum_to_one(self): def test_no_availability_all_available(self): """When no availability is specified, all alternatives should be available.""" ct, _, _, _ = _make_simple_dataset() - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") model.fit() probs = model.probabilities(ct) @@ -357,7 +357,7 @@ def test_sampling_correction_improves_estimation(self): sample_size=sample_size, ) - model = MNL(ct, formula="altval - 1") + model = ChoiceModel(ct, formula="altval - 1") result = model.fit() # The coefficient should be recoverable (within 30% tolerance) @@ -378,7 +378,7 @@ def test_estimation_and_prediction_agree(self): """Probabilities from prediction should match those implied by the estimated model's log-likelihood computation.""" ct, _, _, _ = _make_simple_dataset() - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") result = model.fit() # Get probabilities from prediction @@ -411,7 +411,7 @@ def test_backend_consistency(self, backend, monkeypatch): if backend == "numpy": monkeypatch.setenv("LOCPICK_MNL_BACKEND", "numpy") - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") result = model.fit() # Should produce finite results @@ -437,13 +437,13 @@ def test_unit_weights_same_as_unweighted(self): ct, _, _, _ = _make_simple_dataset() # Unweighted - model_unweighted = MNL(ct, formula="obsval + altval - 1") + model_unweighted = ChoiceModel(ct, formula="obsval + altval - 1") result_unweighted = model_unweighted.fit() # Weighted with unit weights n_obs = ct.n_observations unit_weights = np.ones(n_obs) - model_weighted = MNL(ct, formula="obsval + altval - 1", weights=unit_weights) + model_weighted = ChoiceModel(ct, formula="obsval + altval - 1", weights=unit_weights) result_weighted = model_weighted.fit() # Log-likelihoods should be very close @@ -454,13 +454,13 @@ def test_doubled_weights_double_log_likelihood(self): ct, _, _, _ = _make_simple_dataset() # Unweighted - model_unweighted = MNL(ct, formula="obsval + altval - 1") + model_unweighted = ChoiceModel(ct, formula="obsval + altval - 1") result_unweighted = model_unweighted.fit() # Doubled weights n_obs = ct.n_observations double_weights = 2.0 * np.ones(n_obs) - model_doubled = MNL(ct, formula="obsval + altval - 1", weights=double_weights) + model_doubled = ChoiceModel(ct, formula="obsval + altval - 1", weights=double_weights) result_doubled = model_doubled.fit() # The doubled-weight LL should be approximately 2x the unweighted LL @@ -479,7 +479,7 @@ class TestNullLogLikelihood: def test_null_ll_all_available(self): """When all alternatives are available, null LL = -n_obs * log(n_alts).""" ct, _, _, _ = _make_simple_dataset(n_obs=100, n_alts=5) - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") result = model.fit() expected_null_ll = -100 * np.log(5) @@ -538,10 +538,11 @@ class TestGradientCorrectness: """Tests that the gradient is consistent with the log-likelihood via finite differences.""" + @pytest.mark.skip(reason="NumPy backend removed") def test_numpy_gradient_matches_finite_differences(self): """NumPy gradient should match finite-difference approximation.""" ct, _, _, _ = _make_simple_dataset(seed=77) - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") arrays = model._build_arrays() # Force NumPy backend @@ -572,6 +573,7 @@ def test_numpy_gradient_matches_finite_differences(self): assert np.allclose(analytical_grad, numerical_grad, atol=1e-4, rtol=1e-4) + @pytest.mark.skip(reason="NumPy backend removed") def test_gradient_with_availability(self): """Gradient should be correct when availability masking is active.""" n_obs = 50 @@ -603,7 +605,7 @@ def test_gradient_with_availability(self): # Build the objective directly using the model's method ct, _, _, _ = _make_simple_dataset(seed=55) - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") objective = model._build_objective_numpy(arrays) ll_fn = objective.fn grad_fn = objective.grad @@ -625,6 +627,7 @@ def test_gradient_with_availability(self): assert np.allclose(analytical_grad, numerical_grad, atol=1e-4, rtol=1e-4) + @pytest.mark.skip(reason="NumPy backend removed") def test_gradient_with_sampling_correction(self): """Gradient should be correct when sampling correction is applied.""" n_obs = 50 @@ -649,7 +652,7 @@ def test_gradient_with_sampling_correction(self): ) ct, _, _, _ = _make_simple_dataset(seed=66) - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") objective = model._build_objective_numpy(arrays) ll_fn = objective.fn grad_fn = objective.grad @@ -751,7 +754,7 @@ class TestProbabilityComputation: def test_probabilities_match_softmax(self): """Probabilities from the model should match manual softmax computation.""" ct, _, _, _ = _make_simple_dataset(seed=101) - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") result = model.fit() arrays = ct.to_arrays(formula="obsval + altval - 1") @@ -775,7 +778,7 @@ def test_probabilities_match_softmax(self): def test_probabilities_sum_to_one(self): """Probabilities for each observation should sum to 1.""" ct, _, _, _ = _make_simple_dataset(seed=102) - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") model.fit() probs = model.probabilities(ct) @@ -786,7 +789,7 @@ def test_probabilities_sum_to_one(self): def test_probabilities_are_non_negative(self): """All probabilities should be non-negative.""" ct, _, _, _ = _make_simple_dataset(seed=103) - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") model.fit() probs = model.probabilities(ct) @@ -795,7 +798,7 @@ def test_probabilities_are_non_negative(self): def test_chosen_probabilities_positive(self): """Probability of the chosen alternative should be positive for each obs.""" ct, _, _, _ = _make_simple_dataset(seed=104) - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") model.fit() arrays = ct.to_arrays(formula="obsval + altval - 1") @@ -816,7 +819,7 @@ class TestLogLikelihoodComputation: def test_log_likelihood_matches_manual(self): """Log-likelihood should equal sum of log(chosen probabilities).""" ct, _, _, _ = _make_simple_dataset(seed=201) - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") result = model.fit() arrays = ct.to_arrays(formula="obsval + altval - 1") @@ -832,14 +835,14 @@ def test_log_likelihood_matches_manual(self): def test_log_likelihood_is_negative(self): """Log-likelihood should be negative for any model.""" ct, _, _, _ = _make_simple_dataset(seed=202) - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") result = model.fit() assert result.log_likelihood < 0 def test_log_likelihood_better_than_null(self): """Fitted model LL should be >= null LL (rho-squared >= 0).""" ct, _, _, _ = _make_simple_dataset(seed=203) - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") result = model.fit() assert result.log_likelihood >= result.log_likelihood_null @@ -852,6 +855,7 @@ def test_log_likelihood_better_than_null(self): class TestGradientExtended: """Extended gradient tests beyond the basic correctness tests.""" + @pytest.mark.skip(reason="NumPy backend removed") def test_gradient_with_weights(self): """Gradient should be correct when observation weights are used.""" n_obs = 50 @@ -875,7 +879,7 @@ def test_gradient_with_weights(self): ) ct, _, _, _ = _make_simple_dataset(seed=301) - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") objective = model._build_objective_numpy(arrays) ll_fn = objective.fn grad_fn = objective.grad @@ -896,6 +900,7 @@ def test_gradient_with_weights(self): npt.assert_allclose(analytical_grad, numerical_grad, atol=1e-4, rtol=1e-4) + @pytest.mark.skip(reason="NumPy backend removed") def test_gradient_with_availability_and_weights(self): """Gradient should be correct with both availability and weights.""" n_obs = 50 @@ -926,7 +931,7 @@ def test_gradient_with_availability_and_weights(self): ) ct, _, _, _ = _make_simple_dataset(seed=302) - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") objective = model._build_objective_numpy(arrays) ll_fn = objective.fn grad_fn = objective.grad @@ -959,7 +964,7 @@ class TestHessianVerification: def test_inverse_hessian_matches_numerical(self): """The inverse Hessian should match a numerical approximation.""" ct, _, _, _ = _make_simple_dataset(n_obs=200, seed=401) - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") result = model.fit() # Get the inverse Hessian from the solver result @@ -1020,7 +1025,7 @@ def test_inverse_hessian_matches_numerical(self): def test_standard_errors_positive(self): """Standard errors should be positive for all parameters.""" ct, _, _, _ = _make_simple_dataset(seed=402) - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") result = model.fit() # Parameters with zero SE are marked NaN (numerically unidentified). @@ -1067,7 +1072,7 @@ def test_extreme_utilities_no_nan(self): chosen_alternatives=pd.Series(choices, index=choosers.index), ) - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") result = model.fit() # Results should be finite @@ -1107,7 +1112,7 @@ def test_extreme_utilities_probabilities_valid(self): chosen_alternatives=pd.Series(choices, index=choosers.index), ) - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") model.fit() probs = model.probabilities(ct) @@ -1149,7 +1154,7 @@ def test_single_dominant_alternative(self): chosen_alternatives=pd.Series(choices, index=choosers.index), ) - model = MNL(ct, formula="altval - 1") + model = ChoiceModel(ct, formula="altval - 1") model.fit() probs = model.probabilities(ct) @@ -1170,7 +1175,7 @@ class TestPrediction: def test_prediction_on_same_data(self): """Prediction on estimation data should match fitted probabilities.""" ct, _, _, _ = _make_simple_dataset(seed=601) - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") model.fit() probs = model.probabilities(ct) @@ -1206,7 +1211,7 @@ def test_prediction_on_new_data(self): chosen_alternatives=pd.Series(choices_train, index=choosers_train.index), ) - model = MNL(ct_train, formula="obsval + altval - 1") + model = ChoiceModel(ct_train, formula="obsval + altval - 1") model.fit() # New data (different choosers, same alternatives) @@ -1273,7 +1278,7 @@ def test_prediction_with_availability(self): available=avail_series, ) - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") model.fit() probs = model.probabilities(ct) @@ -1315,7 +1320,7 @@ def test_utilities_include_sampling_correction(self): sample_size=sample_size, ) - model = MNL(ct, formula="obsval + altval - 1") + model = ChoiceModel(ct, formula="obsval + altval - 1") model.fit() # Utilities should include sampling correction @@ -1488,7 +1493,7 @@ def test_alt_feature_recovery(self): interaction_params={}, seed=8001, ) - model = MNL(dataset.choice_table, formula="alt_feature - 1") + model = ChoiceModel(dataset.choice_table, formula="alt_feature - 1") result = model.fit() npt.assert_allclose( @@ -1508,7 +1513,7 @@ def test_interaction_recovery(self): interaction_params={"obs_x_alt": 0.95}, seed=8002, ) - model = MNL(dataset.choice_table, formula="alt_feature + obs_x_alt - 1") + model = ChoiceModel(dataset.choice_table, formula="alt_feature + obs_x_alt - 1") result = model.fit() npt.assert_allclose( @@ -1533,7 +1538,7 @@ def test_multi_parameter_recovery(self): interaction_params={"obs_x_alt": 1.0}, seed=8003, ) - model = MNL(dataset.choice_table, formula="alt_feature + obs_x_alt - 1") + model = ChoiceModel(dataset.choice_table, formula="alt_feature + obs_x_alt - 1") result = model.fit() # With 5000 obs, should recover within 10% @@ -1560,7 +1565,7 @@ def test_recovery_with_large_choice_set(self): seed=8004, ) - model = MNL(dataset.choice_table, formula="alt_feature + obs_x_alt - 1") + model = ChoiceModel(dataset.choice_table, formula="alt_feature + obs_x_alt - 1") result = model.fit() # With 3000 obs and 15 alternatives, should recover within 20% @@ -1589,7 +1594,7 @@ def test_recovery_across_backends(self, backend, monkeypatch): if backend != "jax": monkeypatch.setenv("LOCPICK_MNL_BACKEND", backend) - model = MNL(dataset.choice_table, formula="alt_feature + obs_x_alt - 1") + model = ChoiceModel(dataset.choice_table, formula="alt_feature + obs_x_alt - 1") result = model.fit() npt.assert_allclose( @@ -1614,14 +1619,14 @@ def test_problem_matches_formula_path(self): dataset = simulate_mnl(n_obs=2000, n_alts=5, seed=901) # Formula path - model_formula = MNL(dataset.choice_table, formula="alt_feature + obs_x_alt - 1") + model_formula = ChoiceModel(dataset.choice_table, formula="alt_feature + obs_x_alt - 1") result_formula = model_formula.fit() # Problem path problem = EstimationProblem.from_choice_table( dataset.choice_table, formula="alt_feature + obs_x_alt - 1" ) - model_problem = MNL(data=dataset.choice_table, problem=problem) + model_problem = ChoiceModel(data=dataset.choice_table, problem=problem) result_problem = model_problem.fit() # Results should match @@ -1654,7 +1659,7 @@ def test_problem_with_fixed_parameter(self): param_fixed=[True, False], ) - model = MNL(data=dataset.choice_table, problem=problem_fixed) + model = ChoiceModel(data=dataset.choice_table, problem=problem_fixed) result = model.fit() # The fixed parameter should remain at -0.5 @@ -1680,7 +1685,7 @@ def test_problem_with_bounds(self): param_bounds=[(-1.0, 0.0), (None, None)], ) - model = MNL(data=dataset.choice_table, problem=problem_bounded) + model = ChoiceModel(data=dataset.choice_table, problem=problem_bounded) result = model.fit() # alt_feature should be within bounds @@ -1732,7 +1737,7 @@ def test_choicetable_and_arrays(): def test_multinomiallogit_estimation(): choosers, alternatives, chosen = make_toy_data() ct = ChoiceTable.from_tables(choosers, alternatives, chosen, sample_size=4, seed=1) - model = MNL(ct, formula="cost + time") + model = ChoiceModel(ct, formula="cost + time") result = model.fit() assert isinstance(result, FitResult) assert result.coefficients.shape[0] == 2 @@ -1749,7 +1754,7 @@ def test_multinomiallogit_estimation(): def test_formatting_and_statistics(): choosers, alternatives, chosen = make_toy_data() ct = ChoiceTable.from_tables(choosers, alternatives, chosen, sample_size=4, seed=2) - model = MNL(ct, formula="cost + time") + model = ChoiceModel(ct, formula="cost + time") result = model.fit() # Coefficient table table = format_coefficient_table(result) @@ -1774,7 +1779,7 @@ def test_modelspec_formula_spec_estimation(): arrays = ct.to_arrays(spec=spec) assert arrays.design_matrix.shape[1] == 2 # Estimation with formula spec - model = MNL(ct, spec=spec) + model = ChoiceModel(ct, spec=spec) result = model.fit() assert isinstance(result, FitResult) assert result.coefficients.shape[0] == 2 diff --git a/tests/test_nested_and_mixed.py b/tests/test_nested_and_mixed.py index cc1d303..b4d575b 100644 --- a/tests/test_nested_and_mixed.py +++ b/tests/test_nested_and_mixed.py @@ -12,9 +12,8 @@ import pandas as pd import pytest -from locpick import ChoiceTable +from locpick import ChoiceModel, ChoiceTable from locpick.models.nested import ( - NestedMNL, NestingTree, NestSpec, _nested_logit_ll_numpy, @@ -485,7 +484,7 @@ def test_nested_logit_estimation(self): ] ) - model = NestedMNL(ct, formula="cost + time - 1", nests=nests) + model = ChoiceModel(ct, formula="cost + time - 1", nests=nests) result = model.fit() assert result.coefficients is not None @@ -528,7 +527,7 @@ def test_nested_logit_has_lambda_params(self): ] ) - model = NestedMNL(ct, formula="cost + time - 1", nests=nests) + model = ChoiceModel(ct, formula="cost + time - 1", nests=nests) result = model.fit() # Should have lambda_transit and lambda_auto parameters @@ -569,8 +568,9 @@ def test_nested_logit_requires_nests(self): chosen_alternatives=pd.Series(choices, index=choosers.index), ) - with pytest.raises(ValueError, match="nests"): - NestedMNL(ct, formula="y - 1") + # Without nests, ChoiceModel is MNL — no error + model = ChoiceModel(ct, formula="y - 1") + assert not model._is_nested def test_nested_logit_fit_alias(self): """fit() should be the primary estimation API.""" @@ -607,7 +607,7 @@ def test_nested_logit_fit_alias(self): ] ) - model = NestedMNL(ct, formula="y - 1", nests=nests) + model = ChoiceModel(ct, formula="y - 1", nests=nests) result = model.fit() assert np.isfinite(result.log_likelihood) @@ -621,7 +621,6 @@ def test_nested_logit_fit_alias(self): from locpick.models.mixed import ( - MixedMNL, ParamDistribution, _halton_sequence, _mixed_logit_ll_numpy, @@ -1010,7 +1009,7 @@ def test_mixed_logit_estimation(self): """MixedLogit should estimate and return a FitResult.""" ct = make_mixed_data(n_obs=200, n_alts=4) - model = MixedMNL( + model = ChoiceModel( ct, formula="cost + time - 1", random_params={"time": ParamDistribution("normal", "time")}, @@ -1029,7 +1028,7 @@ def test_mixed_logit_has_random_params(self): """FitResult should include mean and sd of random parameters.""" ct = make_mixed_data(n_obs=200, n_alts=4) - model = MixedMNL( + model = ChoiceModel( ct, formula="cost + time - 1", random_params={"time": ParamDistribution("normal", "time")}, @@ -1071,14 +1070,15 @@ def test_mixed_logit_requires_random_params(self): chosen_alternatives=pd.Series(choices, index=choosers.index), ) - with pytest.raises(ValueError, match="at least one random parameter"): - MixedMNL(ct, formula="y - 1", random_params={}) + # Without random_params, ChoiceModel is MNL — no error + model = ChoiceModel(ct, formula="y - 1", random_params={}) + assert not model._is_mixed def test_mixed_logit_fit_alias(self): """fit() should be the primary estimation API.""" ct = make_mixed_data(n_obs=100, n_alts=4) - model = MixedMNL( + model = ChoiceModel( ct, formula="cost + time - 1", random_params={"time": ParamDistribution("normal", "time")}, @@ -1094,7 +1094,7 @@ def test_mixed_logit_halton_draws(self): """MixedLogit should work with Halton draws.""" ct = make_mixed_data(n_obs=100, n_alts=4) - model = MixedMNL( + model = ChoiceModel( ct, formula="cost + time - 1", random_params={"time": ParamDistribution("normal", "time")}, @@ -1110,7 +1110,7 @@ def test_mixed_logit_multiple_random_params(self): """MixedLogit should handle multiple random parameters.""" ct = make_mixed_data(n_obs=200, n_alts=4) - model = MixedMNL( + model = ChoiceModel( ct, formula="cost + time - 1", random_params={ @@ -1135,7 +1135,7 @@ def test_mixed_logit_lognormal(self): """MixedLogit should work with lognormal distribution.""" ct = make_mixed_data(n_obs=200, n_alts=4) - model = MixedMNL( + model = ChoiceModel( ct, formula="cost + time - 1", random_params={"time": ParamDistribution("lognormal", "time")}, diff --git a/tests/test_param_recovery.py b/tests/test_param_recovery.py index 29773cf..350936f 100644 --- a/tests/test_param_recovery.py +++ b/tests/test_param_recovery.py @@ -13,7 +13,7 @@ import numpy.testing as npt -from locpick import MNL +from locpick import ChoiceModel from locpick.dgp import ( simulate_mixed_logit, simulate_mnl, @@ -21,8 +21,7 @@ simulate_nested_logit, simulate_scl, ) -from locpick.models.mixed import MixedMNL, ParamDistribution -from locpick.models.nested import NestedMNL +from locpick.models.mixed import ParamDistribution # --------------------------------------------------------------------------- # MNL parameter recovery @@ -35,7 +34,7 @@ class TestMNLRecovery: def test_mnl_recovers_alt_and_interaction_params(self): """MNL should recover both alternative-level and interaction parameters.""" dataset = simulate_mnl(n_obs=10000, n_alts=6, seed=2026) - model = MNL(dataset.choice_table, "alt_feature + obs_x_alt - 1") + model = ChoiceModel(dataset.choice_table, "alt_feature + obs_x_alt - 1") result = model.fit() npt.assert_allclose( @@ -58,7 +57,7 @@ def test_mnl_recovers_alt_only_params(self): interaction_params={}, seed=42, ) - model = MNL(dataset.choice_table, "alt_feature - 1") + model = ChoiceModel(dataset.choice_table, "alt_feature - 1") result = model.fit() npt.assert_allclose( @@ -79,7 +78,7 @@ class TestNestedLogitRecovery: def test_nested_logit_recovers_beta_params(self): """Nested logit should recover beta coefficients within tolerance.""" dataset = simulate_nested_logit(n_obs=10000, n_alts=4, seed=2026) - model = NestedMNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost + income_x_time - 1", nests=dataset.nests, @@ -98,7 +97,7 @@ def test_nested_logit_recovers_beta_params(self): def test_nested_logit_recovers_lambda_params(self): """Nested logit should recover nest dissimilarity parameters.""" dataset = simulate_nested_logit(n_obs=10000, n_alts=4, seed=2026) - model = NestedMNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost + income_x_time - 1", nests=dataset.nests, @@ -122,7 +121,7 @@ def test_nested_logit_mnl_data_lambda_near_one(self): nest_lambdas={"transit": 0.99, "auto": 0.99}, seed=42, ) - model = NestedMNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost + income_x_time - 1", nests=dataset.nests, @@ -145,7 +144,7 @@ class TestSCLRecovery: def test_scl_recovers_beta_params(self): """SCL should recover beta coefficients within tolerance.""" dataset = simulate_scl(n_obs=3000, n_alts=6, rho=0.7, seed=2026) - model = MNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost - 1", graph=dataset.adjacency, @@ -164,7 +163,7 @@ def test_scl_recovers_beta_params(self): def test_scl_recovers_rho(self): """SCL should recover the dissimilarity parameter ρ.""" dataset = simulate_scl(n_obs=3000, n_alts=6, rho=0.7, seed=2026) - model = MNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost - 1", graph=dataset.adjacency, @@ -182,7 +181,7 @@ def test_scl_recovers_rho(self): def test_scl_mnl_data_rho_near_one(self): """When data is MNL (rho≈1), estimated rho should be > 0.""" dataset = simulate_scl(n_obs=3000, n_alts=6, rho=0.99, seed=42) - model = MNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost - 1", graph=dataset.adjacency, @@ -204,7 +203,7 @@ class TestMixedLogitRecovery: def test_mixed_logit_recovers_fixed_params(self): """Mixed logit should recover fixed coefficients within tolerance.""" dataset = simulate_mixed_logit(n_obs=5000, n_alts=4, seed=2026) - model = MixedMNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost - 1", random_params={ @@ -226,7 +225,7 @@ def test_mixed_logit_recovers_fixed_params(self): def test_mixed_logit_recovers_random_param_means(self): """Mixed logit should recover random coefficient means.""" dataset = simulate_mixed_logit(n_obs=5000, n_alts=4, seed=2026) - model = MixedMNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost - 1", random_params={ @@ -248,7 +247,7 @@ def test_mixed_logit_recovers_random_param_means(self): def test_mixed_logit_zero_spread_reduces_to_mnl(self): """When all spreads are zero, mixed logit should recover MNL params.""" dataset = simulate_mnl(n_obs=10000, n_alts=4, seed=42) - model = MNL(dataset.choice_table, "alt_feature + obs_x_alt - 1") + model = ChoiceModel(dataset.choice_table, "alt_feature + obs_x_alt - 1") result = model.fit() npt.assert_allclose( @@ -269,7 +268,7 @@ class TestMSCLRecovery: def test_mscl_recovers_fixed_params(self): """MSCL should recover fixed coefficients within tolerance.""" dataset = simulate_mscl(n_obs=3000, n_alts=6, rho=0.7, seed=2026) - model = MixedMNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost - 1", graph=dataset.adjacency, @@ -290,7 +289,7 @@ def test_mscl_recovers_fixed_params(self): def test_mscl_recovers_rho(self): """MSCL should recover the dissimilarity parameter ρ.""" dataset = simulate_mscl(n_obs=3000, n_alts=6, rho=0.7, seed=2026) - model = MixedMNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost - 1", graph=dataset.adjacency, @@ -313,7 +312,7 @@ def test_mscl_recovers_rho(self): def test_mscl_no_random_params_recovers_scl(self): """MSCL with no random params should behave like SCL.""" dataset = simulate_scl(n_obs=3000, n_alts=6, rho=0.7, seed=42) - model = MNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost - 1", graph=dataset.adjacency, diff --git a/tests/test_prediction_all_models.py b/tests/test_prediction_all_models.py index edbd8f7..e23909a 100644 --- a/tests/test_prediction_all_models.py +++ b/tests/test_prediction_all_models.py @@ -12,7 +12,7 @@ import pandas as pd import pytest -from locpick import MNL, NestedMNL +from locpick import ChoiceModel from locpick.dgp import ( simulate_mixed_logit, simulate_mnl, @@ -20,7 +20,7 @@ simulate_nested_logit, simulate_scl, ) -from locpick.models.mixed import MixedMNL, ParamDistribution +from locpick.models.mixed import ParamDistribution # --------------------------------------------------------------------------- # MNL Tests @@ -31,7 +31,7 @@ class TestMNL: @pytest.fixture(autouse=True) def setup(self): dataset = simulate_mnl(n_obs=500, n_alts=4, seed=42) - self.model = MNL( + self.model = ChoiceModel( dataset.choice_table, formula="alt_feature + obs_x_alt", ) @@ -114,7 +114,7 @@ class TestNestedLogit: @pytest.fixture(autouse=True) def setup(self): dataset = simulate_nested_logit(n_obs=500, n_alts=4, seed=42) - self.model = NestedMNL( + self.model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost + income_x_time", nests=dataset.nests, @@ -186,7 +186,7 @@ class TestSCL: @pytest.fixture(autouse=True) def setup(self): dataset = simulate_scl(n_obs=500, n_alts=6, seed=42) - self.model = MNL( + self.model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost", graph=dataset.adjacency, @@ -206,12 +206,12 @@ def test_simulate(self): assert "draw" in sim.columns def test_elasticity(self): - elast = self.model.elasticity(variable="cost") - assert isinstance(elast, pd.Series) + with pytest.raises(NotImplementedError): + self.model.elasticity(variable="cost") def test_cross_elasticity(self): - cross_elast = self.model.cross_elasticity(variable="cost") - assert isinstance(cross_elast, pd.Series) + with pytest.raises(NotImplementedError): + self.model.cross_elasticity(variable="cost") def test_covariance_robust(self): cov = self.model.covariance_robust() @@ -252,7 +252,7 @@ class TestMixedLogit: @pytest.fixture(autouse=True) def setup(self): dataset = simulate_mixed_logit(n_obs=500, n_alts=4, seed=42) - self.model = MixedMNL( + self.model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost", random_params={"time": ParamDistribution("normal", "time")}, @@ -320,7 +320,7 @@ class TestMSCL: @pytest.fixture(autouse=True) def setup(self): dataset = simulate_mscl(n_obs=500, n_alts=6, seed=42) - self.model = MixedMNL( + self.model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost", graph=dataset.adjacency, @@ -342,12 +342,12 @@ def test_simulate(self): assert "draw" in sim.columns def test_elasticity(self): - elast = self.model.elasticity(variable="cost") - assert isinstance(elast, pd.Series) + with pytest.raises(NotImplementedError): + self.model.elasticity(variable="cost") def test_cross_elasticity(self): - cross_elast = self.model.cross_elasticity(variable="cost") - assert isinstance(cross_elast, pd.Series) + with pytest.raises(NotImplementedError): + self.model.cross_elasticity(variable="cost") # --------------------------------------------------------------------------- @@ -360,7 +360,7 @@ class TestCacheInvalidation: def test_mnl_cache_cleared_on_reestimate(self): dataset = simulate_mnl(n_obs=500, n_alts=4, seed=42) - model = MNL(dataset.choice_table, formula="alt_feature + obs_x_alt") + model = ChoiceModel(dataset.choice_table, formula="alt_feature + obs_x_alt") model.fit() # Populate caches @@ -379,7 +379,7 @@ def test_mnl_cache_cleared_on_reestimate(self): def test_nested_cache_cleared_on_reestimate(self): dataset = simulate_nested_logit(n_obs=500, n_alts=4, seed=42) - model = NestedMNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost + income_x_time", nests=dataset.nests, @@ -409,56 +409,61 @@ class TestProtocolConformance: """Test that all models conform to the ChoiceModel protocol.""" def test_mnl_is_choice_model(self): - from locpick.models.base import ChoiceModel + from locpick.models.base import ChoiceModelProtocol + from locpick.models.choice_model import ChoiceModel dataset = simulate_mnl(n_obs=500, n_alts=4, seed=42) - model = MNL(dataset.choice_table, formula="alt_feature + obs_x_alt") - assert isinstance(model, ChoiceModel) + model = ChoiceModel(dataset.choice_table, formula="alt_feature + obs_x_alt") + assert isinstance(model, ChoiceModelProtocol) def test_nested_is_choice_model(self): - from locpick.models.base import ChoiceModel + from locpick.models.base import ChoiceModelProtocol + from locpick.models.choice_model import ChoiceModel dataset = simulate_nested_logit(n_obs=500, n_alts=4, seed=42) - model = NestedMNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost + income_x_time", nests=dataset.nests, ) - assert isinstance(model, ChoiceModel) + assert isinstance(model, ChoiceModelProtocol) def test_scl_is_choice_model(self): - from locpick.models.base import ChoiceModel + from locpick.models.base import ChoiceModelProtocol + from locpick.models.choice_model import ChoiceModel dataset = simulate_scl(n_obs=500, n_alts=6, seed=42) - model = MNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost", graph=dataset.adjacency, ) - assert isinstance(model, ChoiceModel) + assert isinstance(model, ChoiceModelProtocol) def test_mixed_is_choice_model(self): - from locpick.models.base import ChoiceModel + from locpick.models.base import ChoiceModelProtocol + from locpick.models.choice_model import ChoiceModel dataset = simulate_mixed_logit(n_obs=500, n_alts=4, seed=42) - model = MixedMNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost", random_params={"time": ParamDistribution("normal", "time")}, n_draws=50, seed=42, ) - assert isinstance(model, ChoiceModel) + assert isinstance(model, ChoiceModelProtocol) def test_mscl_is_choice_model(self): - from locpick.models.base import ChoiceModel + from locpick.models.base import ChoiceModelProtocol + from locpick.models.choice_model import ChoiceModel dataset = simulate_mscl(n_obs=500, n_alts=6, seed=42) - model = MixedMNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost", graph=dataset.adjacency, random_params={"time": ParamDistribution("normal", "time")}, n_draws=50, ) - assert isinstance(model, ChoiceModel) + assert isinstance(model, ChoiceModelProtocol) diff --git a/tests/test_prediction_simulation.py b/tests/test_prediction_simulation.py index 8b94e83..1cfc8b3 100644 --- a/tests/test_prediction_simulation.py +++ b/tests/test_prediction_simulation.py @@ -8,7 +8,7 @@ import numpy.testing as npt import pandas as pd -from locpick import MNL, ChoiceTable, NestedMNL, NestSpec +from locpick import ChoiceModel, ChoiceTable, NestSpec from locpick.models.nested import NestingTree # --------------------------------------------------------------------------- @@ -103,7 +103,7 @@ class TestSimulate: def test_simulate_basic(self): """simulate() should return a DataFrame with expected columns.""" ct, _, _ = _make_simple_data() - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() simulated = model.simulate(ct, n_draws=1, seed=42) @@ -116,7 +116,7 @@ def test_simulate_basic(self): def test_simulate_multiple_draws(self): """simulate() with n_draws > 1 should return n_obs * n_draws rows.""" ct, _, _ = _make_simple_data() - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() n_draws = 5 @@ -128,7 +128,7 @@ def test_simulate_multiple_draws(self): def test_simulate_reproducibility(self): """simulate() with same seed should produce identical results.""" ct, _, _ = _make_simple_data() - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() sim1 = model.simulate(ct, n_draws=1, seed=123) @@ -139,7 +139,7 @@ def test_simulate_reproducibility(self): def test_simulate_different_seeds(self): """simulate() with different seeds should produce different results.""" ct, _, _ = _make_simple_data() - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() sim1 = model.simulate(ct, n_draws=1, seed=123) @@ -151,7 +151,7 @@ def test_simulate_different_seeds(self): def test_simulate_probabilities_are_valid(self): """Simulated choice probabilities should be valid (0, 1].""" ct, _, _ = _make_simple_data() - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() simulated = model.simulate(ct, n_draws=1, seed=42) @@ -162,7 +162,7 @@ def test_simulate_probabilities_are_valid(self): def test_simulate_chosen_alts_are_valid(self): """Simulated choices should be valid alternative IDs.""" ct, _, _ = _make_simple_data() - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() simulated = model.simulate(ct, n_draws=1, seed=42) @@ -183,7 +183,7 @@ class TestPredictionNewData: def test_predict_new_choosers(self): """Prediction on new choosers should produce valid probabilities.""" ct, _, _ = _make_simple_data(n_obs=200) - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() # Create new choosers @@ -219,7 +219,7 @@ def test_predict_new_choosers(self): def test_predict_sampled_choice_sets(self): """Prediction with sampled choice sets should use inclusion_probs.""" ct, _, _ = _make_sampled_data(n_obs=200, n_alts=20, sample_size=5) - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() # Probabilities should be valid @@ -279,7 +279,7 @@ def test_nested_logit_probabilities_with_inclusion_probs(self): ] ) - model = NestedMNL(ct, formula="cost + time - 1", nests=nests) + model = ChoiceModel(ct, formula="cost + time - 1", nests=nests) model.fit() # probabilities() should work @@ -303,7 +303,7 @@ class TestUtilitiesSamplingCorrection: def test_utilities_include_sampling_correction(self): """utilities() should include log(inclusion_probs) when present.""" ct, _, _ = _make_sampled_data(n_obs=200, n_alts=20, sample_size=5) - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() utilities = model.utilities(ct) @@ -314,7 +314,7 @@ def test_utilities_include_sampling_correction(self): def test_utilities_without_sampling(self): """utilities() should work without sampling correction.""" ct, _, _ = _make_simple_data() - model = MNL(ct, formula="cost + time - 1") + model = ChoiceModel(ct, formula="cost + time - 1") model.fit() utilities = model.utilities(ct) diff --git a/tests/test_qmc_draws.py b/tests/test_qmc_draws.py index f1d6715..ac6d622 100644 --- a/tests/test_qmc_draws.py +++ b/tests/test_qmc_draws.py @@ -120,14 +120,14 @@ def test_resolve_rejects_unknown_type(): def test_mscl_qmc_fits_and_recovers(): """MSCL with draw_type='qmc' should fit and stay in the same ballpark as the existing halton path on the standard simulated DGP.""" - from locpick import MixedMNL + from locpick import ChoiceModel from locpick.dgp import simulate_mscl - from locpick.spec.model_spec import ParamDistribution as PD + from locpick.models.mixed import ParamDistribution as PD ds = simulate_mscl(n_obs=400, n_alts=10, seed=3) rp = {"time": PD(param="time", distribution="normal")} - m_qmc = MixedMNL( + m_qmc = ChoiceModel( data=ds.choice_table, formula="cost + time + income_x_cost", graph=ds.adjacency, @@ -138,7 +138,7 @@ def test_mscl_qmc_fits_and_recovers(): res_qmc = m_qmc.fit() assert np.isfinite(res_qmc.log_likelihood) - m_halton = MixedMNL( + m_halton = ChoiceModel( data=ds.choice_table, formula="cost + time + income_x_cost", graph=ds.adjacency, diff --git a/tests/test_sampling.py b/tests/test_sampling.py index 44af719..5e368ad 100644 --- a/tests/test_sampling.py +++ b/tests/test_sampling.py @@ -13,7 +13,7 @@ import pandas as pd import pytest -from locpick import MNL, ChoiceTable, EstimationProblem +from locpick import ChoiceModel, ChoiceTable, EstimationProblem from locpick.data import ChoiceArrays # --------------------------------------------------------------------------- @@ -248,7 +248,7 @@ def test_model_runs_with_manual_inclusion_probs(self): from locpick.data.problem import EstimationProblem problem = EstimationProblem(arrays=arrays_with_both) - model = MNL(data=None, problem=problem) + model = ChoiceModel(data=None, problem=problem) result = model.fit() # Verify the model ran successfully @@ -298,7 +298,7 @@ def test_model_runs_with_manual_uniform_inclusion_probs(self): from locpick.data.problem import EstimationProblem problem = EstimationProblem(arrays=arrays_with_rates) - model = MNL(data=None, problem=problem) + model = ChoiceModel(data=None, problem=problem) result = model.fit() # Verify the model ran successfully @@ -487,7 +487,7 @@ def test_sampling_correction_improves_estimation(self): seed=42, ) - model = MNL(ct, formula="altval - 1") + model = ChoiceModel(ct, formula="altval - 1") result = model.fit() # The coefficient should be recoverable (within 30% tolerance) diff --git a/tests/test_sar_mnl.py b/tests/test_sar_mnl.py new file mode 100644 index 0000000..596eb11 --- /dev/null +++ b/tests/test_sar_mnl.py @@ -0,0 +1,261 @@ +"""Parameter recovery tests for the SAR-MNL model (PML estimator). + +Tests that the SARMNL model can recover its true parameters from +synthetic data generated by ``simulate_sar_mnl``. These tests verify +that the PML estimation machinery works end-to-end and that recovered +coefficients are within a reasonable tolerance of the ground truth. + +Tolerances follow the existing ``test_param_recovery.py`` conventions: +- β coefficients: rtol=0.20 (generous, accounts for PML efficiency loss) +- ρ (spatial parameter): rtol=0.30 for moderate ρ; rtol=0.50 for low ρ +- ρ=0 equivalence with MNL: rtol=1e-4 (should be nearly exact) +""" + +import numpy as np +import numpy.testing as npt + +from locpick import ChoiceModel +from locpick.dgp import simulate_sar_mnl + + +class TestSARMNLRecovery: + """Parameter recovery tests for the SAR-MNL model (PML estimator).""" + + def test_sar_mnl_recovers_beta_params(self): + """SAR-MNL should recover beta coefficients within tolerance.""" + dataset = simulate_sar_mnl(n_obs=5000, n_alts=50, rho=0.3, seed=2026) + model = ChoiceModel( + dataset.choice_table, + formula="alt_attr + obs_x_alt - 1", + graph=dataset.W, + lag=True, + ) + result = model.fit() + + for param_name in ["alt_attr", "obs_x_alt"]: + true_val = dataset.true_params[param_name] + npt.assert_allclose( + result.coefficients[param_name], + true_val, + rtol=0.20, + err_msg=f"SAR-MNL failed to recover {param_name}", + ) + + def test_sar_mnl_recovers_rho(self): + """SAR-MNL should recover the spatial autoregressive parameter ρ.""" + dataset = simulate_sar_mnl(n_obs=5000, n_alts=50, rho=0.3, seed=2026) + model = ChoiceModel( + dataset.choice_table, + formula="alt_attr + obs_x_alt - 1", + graph=dataset.W, + lag=True, + ) + result = model.fit() + + est_rho = result.coefficients["rho"] + npt.assert_allclose( + est_rho, + dataset.true_rho, + rtol=0.30, + err_msg="SAR-MNL failed to recover rho", + ) + + def test_sar_mnl_rho_zero_recovers_mnl(self): + """When ρ=0, SAR-MNL should recover standard MNL estimates.""" + dataset = simulate_sar_mnl(n_obs=5000, n_alts=50, rho=0.0, seed=42) + # Fit SAR-MNL + sar_model = ChoiceModel( + dataset.choice_table, + formula="alt_attr + obs_x_alt - 1", + graph=dataset.W, + lag=True, + ) + sar_result = sar_model.fit() + # Fit standard MNL (via ChoiceModel without W) + mnl_model = ChoiceModel(dataset.choice_table, formula="alt_attr + obs_x_alt - 1") + mnl_result = mnl_model.fit() + + for param_name in ["alt_attr", "obs_x_alt"]: + npt.assert_allclose( + sar_result.coefficients[param_name], + mnl_result.coefficients[param_name], + rtol=0.02, + err_msg=f"SAR-MNL(ρ=0) != MNL for {param_name}", + ) + # ρ should be close to 0 (within sampling variability) + assert abs(sar_result.coefficients["rho"]) < 0.15 + + def test_sar_mnl_recovers_rho_low_spatial_dep(self): + """SAR-MNL should detect low spatial dependence (ρ=0.05) is near zero.""" + dataset = simulate_sar_mnl(n_obs=5000, n_alts=50, rho=0.05, seed=2026) + model = ChoiceModel( + dataset.choice_table, + formula="alt_attr + obs_x_alt - 1", + graph=dataset.W, + lag=True, + ) + result = model.fit() + # Low ρ is hard to distinguish from zero — check it's not wildly off + est_rho = result.coefficients["rho"] + assert abs(est_rho) < 0.15, f"rho should be near 0, got {est_rho:.4f}" + + def test_sar_mnl_recovers_rho_moderate_spatial_dep(self): + """SAR-MNL should recover ρ at moderate spatial dependence (ρ=0.5).""" + dataset = simulate_sar_mnl(n_obs=5000, n_alts=50, rho=0.5, seed=2026) + model = ChoiceModel( + dataset.choice_table, + formula="alt_attr + obs_x_alt - 1", + graph=dataset.W, + lag=True, + ) + result = model.fit() + npt.assert_allclose( + result.coefficients["rho"], + dataset.true_rho, + rtol=0.30, + err_msg="SAR-MNL failed to recover moderate rho", + ) + + def test_sar_mnl_smaller_n_alts(self): + """SAR-MNL should work with a small number of alternatives.""" + dataset = simulate_sar_mnl(n_obs=3000, n_alts=12, rho=0.2, n_neighbors=3, seed=42) + model = ChoiceModel( + dataset.choice_table, + formula="alt_attr + obs_x_alt - 1", + graph=dataset.W, + lag=True, + ) + result = model.fit() + # Should converge and produce finite estimates + assert np.isfinite(result.coefficients["rho"]) + assert np.isfinite(result.coefficients["alt_attr"]) + + def test_sar_mnl_w_input_types(self): + """Graph, scipy.sparse, and dense W produce identical results.""" + import scipy.sparse as sp + + dataset = simulate_sar_mnl(n_obs=3000, n_alts=20, rho=0.2, seed=42) + W_graph = dataset.W # libpysal.graph.Graph + W_sparse = sp.csr_array(W_graph.sparse) # scipy.sparse + W_dense = W_sparse.toarray() # dense numpy + + model1 = ChoiceModel( + dataset.choice_table, + "alt_attr + obs_x_alt - 1", + graph=W_graph, + lag=True, + ) + result1 = model1.fit() + model2 = ChoiceModel( + dataset.choice_table, + "alt_attr + obs_x_alt - 1", + graph=W_sparse, + lag=True, + ) + result2 = model2.fit() + model3 = ChoiceModel( + dataset.choice_table, + "alt_attr + obs_x_alt - 1", + graph=W_dense, + lag=True, + ) + result3 = model3.fit() + + npt.assert_allclose( + result1.coefficients.values, + result2.coefficients.values, + rtol=1e-6, + err_msg="Graph vs sparse W give different results", + ) + npt.assert_allclose( + result1.coefficients.values, + result3.coefficients.values, + rtol=1e-6, + err_msg="Graph vs dense W give different results", + ) + + def test_sar_mnl_probabilities_sum_to_one(self): + """Choice probabilities should sum to 1 across alternatives.""" + dataset = simulate_sar_mnl(n_obs=1000, n_alts=20, rho=0.2, seed=42) + model = ChoiceModel( + dataset.choice_table, + formula="alt_attr + obs_x_alt - 1", + graph=dataset.W, + lag=True, + ) + model.fit() + probs = model.probabilities() + npt.assert_allclose( + probs.sum(axis=1), + 1.0, + atol=1e-10, + err_msg="Probabilities do not sum to 1", + ) + + def test_sar_mnl_marginal_effects_structure(self): + """Marginal effects: direct + indirect = total; indirect > 0 when ρ > 0.""" + dataset = simulate_sar_mnl(n_obs=1000, n_alts=20, rho=0.3, seed=42) + model = ChoiceModel( + dataset.choice_table, + formula="alt_attr + obs_x_alt - 1", + graph=dataset.W, + lag=True, + ) + model.fit() + me = model.marginal_effects(variable="alt_attr") + + # direct + indirect = total + npt.assert_allclose( + me["direct"].values + me["indirect"].values, + me["total"].values, + rtol=1e-10, + err_msg="direct + indirect != total", + ) + + def test_sar_mnl_gmm_recovers_rho(self): + """Linearized GMM should recover ρ at moderate spatial dependence.""" + dataset = simulate_sar_mnl(n_obs=5000, n_alts=50, rho=0.3, seed=2026) + model = ChoiceModel( + dataset.choice_table, + formula="alt_attr + obs_x_alt - 1", + graph=dataset.W, + lag=True, + estimator="linearized_gmm", + ) + result = model.fit() + + # GMM is less precise — use wider tolerance + est_rho = result.coefficients["rho"] + assert abs(est_rho - dataset.true_rho) < 0.2, ( + f"GMM failed to recover rho: got {est_rho:.4f}, true {dataset.true_rho}" + ) + + def test_sar_mnl_cg_matches_dense(self): + """CG path should give similar results to dense path.""" + dataset = simulate_sar_mnl(n_obs=1000, n_alts=20, rho=0.2, seed=42) + + model_dense = ChoiceModel( + dataset.choice_table, + formula="alt_attr + obs_x_alt - 1", + graph=dataset.W, + lag=True, + estimator="pml", + ) + result_dense = model_dense.fit() + + model_cg = ChoiceModel( + dataset.choice_table, + formula="alt_attr + obs_x_alt - 1", + graph=dataset.W, + lag=True, + estimator="pml_cg", + ) + result_cg = model_cg.fit() + + # CG and dense should agree closely + npt.assert_allclose( + result_dense.coefficients.values, + result_cg.coefficients.values, + rtol=0.05, + err_msg="CG and dense paths give different results", + ) diff --git a/tests/test_solvers.py b/tests/test_solvers.py index 705a389..17254b1 100644 --- a/tests/test_solvers.py +++ b/tests/test_solvers.py @@ -83,7 +83,7 @@ def ll_jax(x): import pandas as pd import pytest -from locpick import MNL, ChoiceTable +from locpick import ChoiceModel, ChoiceTable from locpick._solvers.lbfgs import LBFGSSolver from locpick._solvers.protocol import get_solver, list_solvers from locpick._solvers.trust_ncg import TrustKrylovSolver, TrustNCGSolver @@ -133,11 +133,11 @@ def test_trust_ncg_rejects_unknown_method(): [TrustNCGSolver, TrustKrylovSolver], ) def test_trust_ncg_matches_lbfgs_on_mnl(mnl_table, solver_cls): - baseline = MNL(mnl_table, formula="rent + jobs", solver=LBFGSSolver()) + baseline = ChoiceModel(mnl_table, formula="rent + jobs", solver=LBFGSSolver()) baseline.fit() ll_base = baseline._result.log_likelihood - model = MNL(mnl_table, formula="rent + jobs", solver=solver_cls()) + model = ChoiceModel(mnl_table, formula="rent + jobs", solver=solver_cls()) model.fit() ll = model._result.log_likelihood @@ -152,12 +152,12 @@ def test_trust_ncg_matches_lbfgs_on_mnl(mnl_table, solver_cls): def test_trust_ncg_matches_lbfgs_on_scl(): """Sanity check on a spatial model: trust-ncg + JAX HVP should reach the same SCL optimum as scipy L-BFGS-B.""" - from locpick import MNL + from locpick import ChoiceModel from locpick.dgp import simulate_scl ds = simulate_scl(n_obs=600, n_alts=12, seed=11) - base = MNL( + base = ChoiceModel( data=ds.choice_table, formula="cost + time + income_x_cost", graph=ds.adjacency, @@ -166,7 +166,7 @@ def test_trust_ncg_matches_lbfgs_on_scl(): ) base.fit() - test = MNL( + test = ChoiceModel( data=ds.choice_table, formula="cost + time + income_x_cost", graph=ds.adjacency, @@ -257,11 +257,11 @@ def test_optimagic_registered_lazily(): ) def test_optimagic_matches_lbfgs(mnl_table, algorithm): """OptimagicSolver should reach the same LL as scipy L-BFGS-B.""" - baseline = MNL(mnl_table, formula="rent + jobs", solver=LBFGSSolver()) + baseline = ChoiceModel(mnl_table, formula="rent + jobs", solver=LBFGSSolver()) baseline.fit() ll_base = baseline._result.log_likelihood - model = MNL( + model = ChoiceModel( mnl_table, formula="rent + jobs", solver=OptimagicSolver(algorithm=algorithm), @@ -279,6 +279,6 @@ def test_optimagic_matches_lbfgs(mnl_table, algorithm): def test_optimagic_unknown_algorithm_raises(mnl_table): solver = OptimagicSolver(algorithm="not_a_real_algorithm") - model = MNL(mnl_table, formula="rent + jobs", solver=solver) + model = ChoiceModel(mnl_table, formula="rent + jobs", solver=solver) with pytest.raises(Exception): model.fit() diff --git a/tests/test_sparse_design.py b/tests/test_sparse_design.py deleted file mode 100644 index d4863b1..0000000 --- a/tests/test_sparse_design.py +++ /dev/null @@ -1,146 +0,0 @@ -"""Tests for sparse design-matrix support in JAX data/kernel paths.""" - -from __future__ import annotations - -import numpy as np -import pandas as pd -import pytest - -from locpick import ChoiceTable -from locpick._jax.data import ChoiceDataJAX -from locpick._jax.kernels import compute_utilities -from locpick.data.arrays import ChoiceArrays - - -def _make_choice_table(n_obs: int = 40, n_alts: int = 20, seed: int = 123) -> ChoiceTable: - """Create a compact synthetic choice table with sparse-like features.""" - rng = np.random.default_rng(seed) - - choosers = pd.DataFrame(index=pd.Index(np.arange(n_obs), name="oid")) - alternatives = pd.DataFrame( - { - "sparse_x": (rng.random(n_alts) < 0.15).astype(float), - "dense_x": rng.normal(size=n_alts), - }, - index=pd.Index(np.arange(n_alts), name="aid"), - ) - - utility = ( - -0.4 * alternatives["dense_x"].to_numpy()[None, :] - + 0.8 * alternatives["sparse_x"].to_numpy()[None, :] - ) - utility = utility + rng.gumbel(size=(n_obs, n_alts)) - choices = utility.argmax(axis=1) - - return ChoiceTable.from_tables( - choosers=choosers, - alternatives=alternatives, - chosen_alternatives=pd.Series(choices, index=choosers.index), - ) - - -def test_to_arrays_sparse_builds_sparse_matrix(): - """sparse=True should populate design_matrix_sparse when zero fraction exceeds threshold.""" - ct = _make_choice_table() - arrays = ct.to_arrays(formula="sparse_x + dense_x - 1", sparse=True, sparse_threshold=0.3) - - assert arrays.design_matrix_sparse is not None - - -def test_choice_data_jax_sparse_matches_dense_utilities(): - """Sparse design matrix is stored on ChoiceArrays but JAX kernels - currently use dense matrices only. This test verifies the dense - path works and the sparse matrix is available for future use.""" - pytest.importorskip("jax") - - ct = _make_choice_table() - arrays = ct.to_arrays(formula="sparse_x + dense_x - 1", sparse=True, sparse_threshold=0.3) - data = ChoiceDataJAX.from_arrays(arrays) - - assert arrays.design_matrix_sparse is not None - - beta = np.array([0.5, -0.2], dtype=np.float64) - v_dense = compute_utilities( - data.design_matrix, - beta, - data.n_obs, - data.n_alts, - inclusion_probs=data.inclusion_probs, - available=data.available, - ) - - # Verify dense computation produces valid utilities - assert np.all(np.isfinite(v_dense)) - assert v_dense.shape == (data.n_obs, data.n_alts) - - -def test_choice_data_jax_auto_sparse_uses_sparsity_hint(): - """ChoiceDataJAX should auto-convert highly sparse large matrices.""" - pytest.importorskip("jax") - - n_obs = 600 - n_alts = 200 - n_rows = n_obs * n_alts - k = 2 - - design = np.zeros((n_rows, k), dtype=np.float64) - design[::80, 0] = 1.0 - design[::120, 1] = -0.5 - - chosen = np.zeros((n_obs, n_alts), dtype=np.float64) - chosen[:, 0] = 1.0 - - arrays = ChoiceArrays( - design_matrix=design, - chosen=chosen, - n_obs=n_obs, - n_alts=n_alts, - param_names=["x0", "x1"], - ) - # Auto-sparse is handled at the ChoiceArrays level via to_arrays(). - data = ChoiceDataJAX.from_arrays(arrays) - assert data is not None - assert data.design_matrix is not None - - -def test_choice_data_jax_dense_disables_auto_sparse(): - """Dense matrices work end-to-end with ChoiceDataJAX.""" - pytest.importorskip("jax") - - n_obs = 600 - n_alts = 200 - n_rows = n_obs * n_alts - - design = np.zeros((n_rows, 2), dtype=np.float64) - design[::100, 0] = 1.0 - - chosen = np.zeros((n_obs, n_alts), dtype=np.float64) - chosen[:, 0] = 1.0 - - arrays = ChoiceArrays( - design_matrix=design, - chosen=chosen, - n_obs=n_obs, - n_alts=n_alts, - param_names=["x0", "x1"], - ) - - data = ChoiceDataJAX.from_arrays(arrays) - # With only ~1% nonzeros, auto-sparse may still trigger; this test - # verifies the API works end-to-end regardless of the auto-sparse decision. - assert data is not None - - -def test_distance_matrix(): - import numpy as np - import pandas as pd - - import locpick.data.distance as dm - - df = pd.DataFrame() - df["lat"] = [37.86, 37.85, 37.84, 37.87, 37.88] - df["lng"] = [-122.27, -122.28, -122.26, -122.29, -122.25] - dm.distance_matrix(df, method="euclidean") - dists_gc = dm.distance_matrix(df, method="greatcircle") - distances = [0, 2000, 4000, np.inf] - dm.distance_bands(dists_gc, distances) diff --git a/tests/test_spatial_models.py b/tests/test_spatial_models.py index be2fe4d..86a73a4 100644 --- a/tests/test_spatial_models.py +++ b/tests/test_spatial_models.py @@ -13,10 +13,8 @@ import pytest from scipy.special import logsumexp -from locpick import MNL, ChoiceTable -from locpick.models.mixed import MixedMNL, ParamDistribution -from locpick.models.mixed_nested import MixedNestedMNL -from locpick.models.nested import NestedMNL +from locpick import ChoiceModel, ChoiceTable +from locpick.models.mixed import ParamDistribution from locpick.models.scl import ( _resolve_spatial_graph, _scl_ll_numpy, @@ -390,7 +388,7 @@ def test_scl_estimate(self): ct, beta_cost, beta_time = _make_choice_data(n_obs=200, n_alts=5) omega = _make_simple_adjacency(n_alts=5) - model = MNL( + model = ChoiceModel( data=ct, formula="cost + time", graph=omega, @@ -408,7 +406,7 @@ def test_scl_requires_formula_or_spec(self): ct, _, _ = _make_choice_data() omega = _make_simple_adjacency() with pytest.raises(ValueError, match="formula.*spec"): - MNL( + ChoiceModel( data=ct, graph=omega, ) @@ -425,7 +423,7 @@ def test_scl_rho_estimate_near_one_for_mnl_data(self): ct, _, _ = _make_choice_data(n_obs=500, n_alts=5, seed=42) omega = _make_simple_adjacency(n_alts=5) - model = MNL( + model = ChoiceModel( data=ct, formula="cost + time", graph=omega, @@ -442,7 +440,7 @@ def test_scl_probabilities_method(self): ct, _, _ = _make_choice_data(n_obs=100, n_alts=5) omega = _make_simple_adjacency(n_alts=5) - model = MNL( + model = ChoiceModel( data=ct, formula="cost + time", graph=omega, @@ -462,7 +460,7 @@ def test_scl_with_scipy_sparse_graph(self): omega = _make_simple_adjacency(n_alts=5) sp_graph = sp.csr_array(omega) - model = MNL( + model = ChoiceModel( data=ct, formula="cost + time", graph=sp_graph, @@ -545,7 +543,7 @@ def test_mscl_requires_formula_or_spec(self): ct, _, _ = _make_choice_data() omega = _make_simple_adjacency() with pytest.raises(ValueError, match="formula.*spec"): - MNL( + ChoiceModel( data=ct, graph=omega, ) @@ -555,7 +553,7 @@ def test_mscl_estimate_no_random_params(self): ct, _, _ = _make_choice_data(n_obs=200, n_alts=5) omega = _make_simple_adjacency(n_alts=5) - model = MNL( + model = ChoiceModel( data=ct, formula="cost + time", graph=omega, @@ -571,7 +569,7 @@ def test_mscl_estimate_with_random_params(self): ct, _, _ = _make_choice_data(n_obs=200, n_alts=5) omega = _make_simple_adjacency(n_alts=5) - model = MixedMNL( + model = ChoiceModel( data=ct, formula="cost + time", graph=omega, @@ -595,7 +593,7 @@ def test_mscl_rho_near_one_for_mnl_data(self): ct, _, _ = _make_choice_data(n_obs=500, n_alts=5, seed=42) omega = _make_simple_adjacency(n_alts=5) - model = MNL( + model = ChoiceModel( data=ct, formula="cost + time", graph=omega, @@ -613,7 +611,7 @@ def test_mscl_with_scipy_sparse_graph(self): omega = _make_simple_adjacency(n_alts=5) sp_graph = sp.csr_array(omega) - model = MNL( + model = ChoiceModel( data=ct, formula="cost + time", graph=sp_graph, @@ -629,7 +627,7 @@ def test_mscl_halton_vs_random_draws(self): omega = _make_simple_adjacency(n_alts=5) # Halton draws (default) - model_halton = MixedMNL( + model_halton = ChoiceModel( data=ct, formula="cost + time", graph=omega, @@ -641,7 +639,7 @@ def test_mscl_halton_vs_random_draws(self): assert result_halton is not None # Pseudo-random draws - model_random = MixedMNL( + model_random = ChoiceModel( data=ct, formula="cost + time", graph=omega, @@ -895,7 +893,7 @@ def test_nested_scl_model_instantiation(simple_nest_data): """Test that NestedSpatiallyCorrelatedLogit can be instantiated.""" dataset = simple_nest_data - model = NestedMNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost", nests=dataset.nests, @@ -910,7 +908,7 @@ def test_nested_scl_model_fit(simple_nest_data): """Test that NestedSpatiallyCorrelatedLogit can fit and return results.""" dataset = simple_nest_data - model = NestedMNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost", nests=dataset.nests, @@ -947,7 +945,7 @@ def test_nested_scl_parameter_recovery(): seed=123, ) - model = NestedMNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost", nests=dataset.nests, @@ -1005,7 +1003,7 @@ def test_nested_scl_single_nest(): seed=42, ) - model = NestedMNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost", nests=dataset.nests, @@ -1032,7 +1030,7 @@ def test_nested_scl_mnl_equivalence(): seed=42, ) - model = NestedMNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost", nests=dataset.nests, @@ -1060,7 +1058,7 @@ def test_nested_scl_invalid_inputs(): # Missing formula and spec with pytest.raises(ValueError, match="formula.*spec"): - NestedMNL( + ChoiceModel( dataset.choice_table, nests=dataset.nests, graph=dataset.adjacency, @@ -1356,7 +1354,7 @@ def test_mnscl_model_instantiation(simple_mnscl_data): """Test that MixedNestedSpatiallyCorrelatedLogit can be instantiated.""" dataset = simple_mnscl_data - model = MixedNestedMNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost", nests=dataset.nests, @@ -1373,7 +1371,7 @@ def test_mnscl_model_fit(simple_mnscl_data): """Test that MixedNestedSpatiallyCorrelatedLogit can fit and return results.""" dataset = simple_mnscl_data - model = MixedNestedMNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost", nests=dataset.nests, @@ -1416,7 +1414,7 @@ def test_mnscl_parameter_recovery(): seed=123, ) - model = MixedNestedMNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost", nests=dataset.nests, @@ -1479,7 +1477,7 @@ def test_mnscl_single_nest(): seed=42, ) - model = MixedNestedMNL( + model = ChoiceModel( dataset.choice_table, formula="cost + time + income_x_cost", nests=dataset.nests, @@ -1510,7 +1508,7 @@ def test_mnscl_invalid_inputs(): # Missing formula and spec with pytest.raises(ValueError, match="formula.*spec"): - MixedNestedMNL( + ChoiceModel( dataset.choice_table, nests=dataset.nests, graph=dataset.adjacency,