diff --git a/README.md b/README.md index 9e5b75f..6dfb2a8 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,21 @@ A Python library for estimating discrete choice models where agents can select e This implements the model as described in Ophem, H.V., Stam, P. and Praag, B.V., 1999. [Multichoice Logit: Modeling Incomplete Preference Rankings of Classical Concerts](https://www.tandfonline.com/doi/abs/10.1080/07350015.1999.10524801). Journal of Business & Economic Statistics, 17(1), pp.117-128. +Built by [Thomas Monk](https://tdmonk.com), London School of Economics. + +## Citation + +If you use this package, please cite it as: + +``` +@misc{monk2025multe, + author = {Thomas Monk}, + title = {Multe: Multichoice Logit Estimation}, + howpublished = {\url{https://github.com/tmonk/multe}}, + year = {2025} +} +``` + ## Installation Install from PyPI: @@ -100,9 +115,9 @@ Model class with methods: - **`fit(X, y_single, y_dual, method='L-BFGS-B')`** - Fit model using MLE (recommended) - Returns `self` with fitted `coef_` attribute - Stores optimization details in `optimization_result_` -- `neg_log_likelihood(flat_beta, X, y_single, y_dual)` - Negative log-likelihood -- `gradient(flat_beta, X, y_single, y_dual)` - Analytical gradient -- `compute_standard_errors(flat_beta, X, y_single, y_dual)` - Standard errors +- `compute_standard_errors(X, y_single, y_dual, flat_beta=None)` - Standard errors (uses fitted params by default) +- `predict_proba(X, flat_beta=None)` - Single/dual choice probabilities +- `log_likelihood_contributions(X, y_single, y_dual, flat_beta=None)` - Per-observation log-likelihoods ### simulate_data(N, J, K, true_beta=None, mix_ratio=0.5, seed=42) Generate synthetic data following the RUM framework. diff --git a/examples/basic_example.py b/examples/basic_example.py index c9ff294..0739494 100644 --- a/examples/basic_example.py +++ b/examples/basic_example.py @@ -30,11 +30,12 @@ def main(): # Initial guess (zeros) init_beta = np.zeros((J - 1) * K) + single_idx, dual_idx = model._validate_data(X, y_single, y_dual) + res = minimize( - fun=model.neg_log_likelihood, - jac=model.gradient, + fun=lambda b: model._neg_log_likelihood(b, X, single_idx, dual_idx), + jac=lambda b: model._gradient(b, X, single_idx, dual_idx), x0=init_beta, - args=(X, y_single, y_dual), method="L-BFGS-B", options={"disp": True, "gtol": 1e-5}, ) @@ -43,7 +44,7 @@ def main(): # Compute Standard Errors print("Computing Standard Errors...") - std_errs = model.compute_standard_errors(res.x, X, y_single, y_dual) + std_errs = model.compute_standard_errors(X, y_single, y_dual, res.x) est_beta = res.x.reshape(J - 1, K) std_errs_reshaped = std_errs.reshape(J - 1, K) diff --git a/examples/benchmark.py b/examples/benchmark.py index b039fad..861b5aa 100644 --- a/examples/benchmark.py +++ b/examples/benchmark.py @@ -30,23 +30,25 @@ def benchmark_estimation(N, J, K, mix_ratio=0.5, method="BFGS", seed=42): model = MultichoiceLogit(J, K) init_beta = np.zeros((J - 1) * K) + # Prepare indices once for timing internals + single_idx, dual_idx = model._validate_data(X, y_single, y_dual) + # Time likelihood evaluation t0 = time.time() - _ = model.neg_log_likelihood(init_beta, X, y_single, y_dual) + _ = model._neg_log_likelihood(init_beta, X, single_idx, dual_idx) likelihood_time = time.time() - t0 # Time gradient evaluation t0 = time.time() - _ = model.gradient(init_beta, X, y_single, y_dual) + _ = model._gradient(init_beta, X, single_idx, dual_idx) gradient_time = time.time() - t0 # Optimize t0 = time.time() result = minimize( - fun=model.neg_log_likelihood, - jac=model.gradient, + fun=lambda b: model._neg_log_likelihood(b, X, single_idx, dual_idx), + jac=lambda b: model._gradient(b, X, single_idx, dual_idx), x0=init_beta, - args=(X, y_single, y_dual), method=method, options={"disp": False, "gtol": 1e-5, "maxiter": 1000}, ) @@ -60,7 +62,8 @@ def benchmark_estimation(N, J, K, mix_ratio=0.5, method="BFGS", seed=42): # Compute standard errors t0 = time.time() - model.compute_standard_errors(result.x, X, y_single, y_dual) + # Compute standard errors (runtime tracked, values unused in benchmark output) + _ = model.compute_standard_errors(result.x, X, y_single, y_dual) se_time = time.time() - t0 return { diff --git a/examples/csv_example.py b/examples/csv_example.py index 133e56b..e827d54 100644 --- a/examples/csv_example.py +++ b/examples/csv_example.py @@ -26,6 +26,7 @@ def save_data_to_csv(X, y_single, y_dual, filepath_prefix="data"): filepath_prefix: Prefix for output files """ N, K = X.shape + # Save covariates covariate_df = pd.DataFrame(X, columns=[f"x_{k}" for k in range(K)]) covariate_df.to_csv(f"{filepath_prefix}_covariates.csv", index=False) @@ -145,11 +146,14 @@ def main(): model = MultichoiceLogit(J, K) init_beta = np.zeros((J - 1) * K) + single_idx, dual_idx = model._validate_data( + X_loaded, y_single_loaded, y_dual_loaded + ) + result = minimize( - fun=model.neg_log_likelihood, - jac=model.gradient, + fun=lambda b: model._neg_log_likelihood(b, X_loaded, single_idx, dual_idx), + jac=lambda b: model._gradient(b, X_loaded, single_idx, dual_idx), x0=init_beta, - args=(X_loaded, y_single_loaded, y_dual_loaded), method="BFGS", options={"disp": False, "gtol": 1e-5}, ) @@ -179,7 +183,7 @@ def main(): # Step 6: Compute standard errors print("\n6. Computing standard errors...") std_errs = model.compute_standard_errors( - result.x, X_loaded, y_single_loaded, y_dual_loaded + X_loaded, y_single_loaded, y_dual_loaded, result.x ) std_errs_reshaped = std_errs.reshape(J - 1, K) @@ -206,7 +210,7 @@ def main(): os.remove("example_data_covariates.csv") os.remove("example_data_choices.csv") print("\nTemporary CSV files cleaned up.") - except OSError: + except FileNotFoundError: pass diff --git a/examples/simple_fit_example.py b/examples/simple_fit_example.py index d7317f6..762a1a9 100644 --- a/examples/simple_fit_example.py +++ b/examples/simple_fit_example.py @@ -43,7 +43,7 @@ def main(): # Compute standard errors print("\n5. Computing standard errors...") flat_coef = model.coef_.flatten() - std_errs = model.compute_standard_errors(flat_coef, X, y_single, y_dual) + std_errs = model.compute_standard_errors(X, y_single, y_dual, flat_coef) std_errs = std_errs.reshape(J - 1, K) print("\n Coefficients with Standard Errors:") diff --git a/multe/model.py b/multe/model.py index e1ecd08..bc0941e 100644 --- a/multe/model.py +++ b/multe/model.py @@ -2,13 +2,18 @@ Multichoice Logit Model Vectorized implementation for fast and accurate MLE estimation. +Supports single and dual (pairwise) discrete choices. """ +from __future__ import annotations + import typing -from typing import Any +import warnings +from typing import Any, Optional, Sequence import numpy as np import numpy.typing as npt +import scipy.sparse as sp from scipy.optimize import OptimizeResult, minimize from scipy.special import logsumexp @@ -16,16 +21,32 @@ CLIP_THRESHOLD = 1e-10 # Minimum probability value (avoid log(0)) HESSIAN_EPSILON = 1e-5 # Step size for Hessian finite differences +# Type alias for flexible dual-choice input formats +DualInput = ( + npt.NDArray[np.int8] + | npt.NDArray[np.int64] + | tuple[np.ndarray, np.ndarray, np.ndarray] + | sp.spmatrix +) + class MultichoiceLogit: """ Multichoice Logit discrete choice model with vectorized operations. + Supports both single choices (standard MNL) and dual/pairwise choices + using an inclusion-exclusion probability formulation. + Attributes: J (int): Total number of alternatives available. K (int): Number of covariates (features) for each alternative. coef_ (np.ndarray): Fitted coefficients of shape (J-1, K). Available after fit(). optimization_result_ (OptimizeResult): Full optimization result. Available after fit(). + + Example: + >>> model = MultichoiceLogit(num_alternatives=3, num_covariates=2) + >>> model.fit(X, y_single, y_dual) + >>> print(model.coef_) """ def __init__(self, num_alternatives: int, num_covariates: int) -> None: @@ -33,8 +54,8 @@ def __init__(self, num_alternatives: int, num_covariates: int) -> None: Initialize the Multichoice Logit model dimensions. Args: - num_alternatives (int): Total number of choices (J). - num_covariates (int): Number of independent variables/features (K). + num_alternatives: Total number of choices (J). Must be >= 2. + num_covariates: Number of independent variables/features (K). Must be >= 1. Raises: ValueError: If num_alternatives < 2 or num_covariates < 1. @@ -51,19 +72,21 @@ def __init__(self, num_alternatives: int, num_covariates: int) -> None: self.coef_: npt.NDArray[np.float64] | None = None self.optimization_result_: OptimizeResult | None = None - def transform_params( + def _transform_params( self, flat_beta: npt.NDArray[np.float64] ) -> npt.NDArray[np.float64]: """ - Reshapes a flat parameter vector into a (J, K) matrix, handling identification. + Reshape a flat parameter vector into a (J, K) matrix with identification constraint. + + The first alternative (index 0) is the reference category with coefficients + fixed to zero for identification. Args: - flat_beta (np.ndarray): 1D array of learnable parameters. - Size should be (J-1) * K. + flat_beta: 1D array of learnable parameters, size (J-1) * K. Returns: - np.ndarray: Matrix of shape (J, K) where row 0 is all zeros and - rows 1..J-1 contain the learned parameters. + Matrix of shape (J, K) where row 0 is zeros and rows 1..J-1 + contain the learned parameters. Raises: ValueError: If flat_beta has incorrect size. @@ -78,20 +101,21 @@ def transform_params( beta_fixed = np.zeros((1, self.K)) return np.vstack([beta_fixed, beta_free]) - def calculate_utilities( + def _calculate_utilities( self, X: npt.NDArray[np.float64], beta: npt.NDArray[np.float64] ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]: """ - Computes the deterministic utility V and the exponentiated utility a. + Compute deterministic utility V and exponentiated utility a. + + Uses the max-subtraction trick for numerical stability in exp(). Args: - X (np.ndarray): Covariate matrix of shape (N, K). - beta (np.ndarray): Parameter matrix of shape (J, K). + X: Covariate matrix of shape (N, K). + beta: Parameter matrix of shape (J, K). Returns: - tuple: - - V (np.ndarray): Deterministic utilities (N, J). - - a (np.ndarray): Exponentiated utilities exp(V) (N, J). + V: Deterministic utilities of shape (N, J). + a: Exponentiated utilities exp(V_stable) of shape (N, J). """ V = X @ beta.T # Numerical stability: subtract max V per row to avoid overflow @@ -99,153 +123,256 @@ def calculate_utilities( a = np.exp(V_stable) return V, a - def _prepare_data( + def _normalize_dual_indices( + self, y_dual: DualInput, *, N: int, J: int + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Convert dual choice inputs into index triplets (rows, s, t). + + Accepts: + - Dense tensors of shape (N, J, J) + - Scipy sparse matrices of shape (N, J*J) with row-major flattening + - Explicit index tuples (rows, s, t) + + Sparse flattening uses row-major order: column = s * J + t. + + Args: + y_dual: Dual choice data in any supported format. + N: Number of observations. + J: Number of alternatives. + + Returns: + Tuple of (row_indices, s_indices, t_indices) arrays. + + Raises: + ValueError: If tuple format is invalid. + TypeError: If y_dual type is unsupported. + """ + if isinstance(y_dual, tuple): + if len(y_dual) != 3: + raise ValueError("y_dual tuple must have length 3 (rows, s, t)") + rows, s_idx, t_idx = y_dual + if not (len(rows) == len(s_idx) == len(t_idx)): + raise ValueError("y_dual index arrays must have the same length") + return ( + np.asarray(rows, dtype=np.int64), + np.asarray(s_idx, dtype=np.int64), + np.asarray(t_idx, dtype=np.int64), + ) + + if sp.issparse(y_dual): + coo = y_dual.tocoo() + rows = coo.row + cols = coo.col + s_idx = cols // J + t_idx = cols % J + return rows.astype(np.int64), s_idx.astype(np.int64), t_idx.astype(np.int64) + + if isinstance(y_dual, np.ndarray): + dual_rows, dual_s, dual_t = np.nonzero(y_dual) + return dual_rows, dual_s, dual_t + + raise TypeError(f"Unsupported y_dual type: {type(y_dual)}") + + def _validate_data( self, + X: npt.NDArray[np.float64], y_single: npt.NDArray[np.int8], - y_dual: npt.NDArray[np.int8], + y_dual: DualInput, ) -> tuple[ tuple[np.ndarray, np.ndarray], tuple[np.ndarray, np.ndarray, np.ndarray] ]: """ - Pre-process data into sparse index format for faster iteration. + Validate input data and return prepared indices. + + Args: + X: Covariate matrix (N, K). + y_single: Single choice indicators (N, J). + y_dual: Dual choice indicators in any supported format. + + Returns: + single_indices: (row_indices, col_indices) for single choices. + dual_indices: (row_indices, s_indices, t_indices) for dual choices. + + Raises: + ValueError: If data dimensions or constraints are violated. """ + N = X.shape[0] + + if X.shape[1] != self.K: + raise ValueError(f"X must have {self.K} columns, got {X.shape[1]}") + + if y_single.shape != (N, self.J): + raise ValueError( + f"y_single must have shape ({N}, {self.J}), got {y_single.shape}" + ) + + # Validate binary values + if not np.isin(y_single, (0, 1)).all(): + raise ValueError("y_single must be binary (contain only 0 or 1).") + + # Normalize dual indices + dual_rows, dual_s, dual_t = self._normalize_dual_indices(y_dual, N=N, J=self.J) + + # Bounds checks for dual indices + if len(dual_rows) > 0: + if dual_rows.max() >= N or dual_rows.min() < 0: + raise ValueError("y_dual row indices out of bounds.") + if dual_s.max() >= self.J or dual_t.max() >= self.J: + raise ValueError("y_dual alternative indices out of bounds.") + if dual_s.min() < 0 or dual_t.min() < 0: + raise ValueError("y_dual alternative indices must be non-negative.") + + # Enforce upper triangle (s < t) and no diagonal + if np.any(dual_s == dual_t): + raise ValueError("y_dual diagonal must be zero (s != t).") + if np.any(dual_s > dual_t): + raise ValueError( + "y_dual must only have entries in upper triangle (s < t)." + ) + + # Binary checks for dense/sparse formats + if isinstance(y_dual, np.ndarray): + if y_dual.shape != (N, self.J, self.J): + raise ValueError( + f"y_dual must have shape ({N}, {self.J}, {self.J}), got {y_dual.shape}" + ) + if not np.isin(y_dual, (0, 1)).all(): + raise ValueError("y_dual must be binary (contain only 0 or 1).") + elif sp.issparse(y_dual): + if y_dual.shape != (N, self.J * self.J): + raise ValueError( + f"sparse y_dual must have shape ({N}, {self.J * self.J})" + ) + if y_dual.data.size and not np.isin(y_dual.data, (0, 1)).all(): + raise ValueError("Sparse y_dual must be binary.") + + # Check that each agent has exactly one choice single_rows, single_cols = np.nonzero(y_single) - dual_rows, dual_s, dual_t = np.nonzero(y_dual) + counts = np.bincount(single_rows, minlength=N) + if len(dual_rows) > 0: + counts += np.bincount(dual_rows, minlength=N) + + if not np.all(counts == 1): + invalid = np.where(counts != 1)[0] + raise ValueError( + f"Each agent must have exactly one choice. " + f"Agents with invalid choices: {invalid[:10]}" + + ("..." if len(invalid) > 10 else "") + ) + return (single_rows, single_cols), (dual_rows, dual_s, dual_t) def fit( self, X: npt.NDArray[np.float64], y_single: npt.NDArray[np.int8], - y_dual: npt.NDArray[np.int8], - init_beta: npt.NDArray[np.float64] | None = None, + y_dual: DualInput, + init_beta: Optional[npt.NDArray[np.float64]] = None, method: str = "L-BFGS-B", - options: dict[str, Any] | None = None, - ) -> "MultichoiceLogit": + options: Optional[dict[str, Any]] = None, + bounds: Optional[Sequence[tuple[float | None, float | None]]] = None, + constraints: Optional[Sequence[Any]] = None, + num_restarts: int = 0, + restart_scale: float = 0.5, + rng: Optional[np.random.Generator] = None, + ) -> MultichoiceLogit: """ Fit the multichoice logit model using maximum likelihood estimation. - This is a convenience method that wraps scipy.optimize.minimize with sensible - defaults. For more control over optimization, you can call neg_log_likelihood - and gradient directly with your own optimizer. - Args: - X (np.ndarray): Covariate matrix of shape (N, K). - y_single (np.ndarray): Binary matrix (N, J). y_single[i, j] = 1 if i chose j alone. - y_dual (np.ndarray): Binary tensor (N, J, J). y_dual[i, s, t] = 1 if i chose pair {s, t}. - init_beta (np.ndarray, optional): Initial parameter values (flat array of size (J-1)*K). - If None, initializes to zeros. - method (str): Optimization method for scipy.optimize.minimize. Default is 'L-BFGS-B'. - Other good options: 'BFGS', 'Newton-CG'. - options (dict, optional): Additional options to pass to the optimizer. - Default is {'gtol': 1e-5, 'maxiter': 1000}. + X: Covariate matrix of shape (N, K). + y_single: Binary matrix (N, J). y_single[i, j] = 1 if i chose j alone. + y_dual: Dual choice indicators. Accepts: + - Dense tensor (N, J, J) with y_dual[i, s, t] = 1 for pair {s, t} + - Sparse matrix (N, J*J) with row-major flattening + - Index triplet (rows, s, t) + init_beta: Initial parameter values, flat array of size (J-1)*K. + Defaults to zeros. + method: Optimization method for scipy.optimize.minimize. + Default 'L-BFGS-B'. Other options: 'BFGS', 'Newton-CG'. + options: Additional options for the optimizer. + Default: {'gtol': 1e-5, 'maxiter': 1000}. + bounds: Parameter bounds for scipy.optimize.minimize. + constraints: Constraints for scipy.optimize.minimize. + num_restarts: Number of random restarts beyond init_beta. + restart_scale: Standard deviation of normal noise for restart initialization. + rng: Random generator for restarts. Uses default_rng() if None. Returns: - self: Returns the instance itself for method chaining. + self: The fitted model instance (for method chaining). Raises: - ValueError: If data dimensions are incompatible or constraints are violated. + ValueError: If data dimensions are incompatible or constraints violated. RuntimeError: If optimization fails to converge. """ - # Set default options if options is None: options = {"gtol": 1e-5, "maxiter": 1000} - # Validate data - self._validate_data(X, y_single, y_dual) - - # Prepare data indices once - single_indices, dual_indices = self._prepare_data(y_single, y_dual) + # Validate data and prepare indices once + single_indices, dual_indices = self._validate_data(X, y_single, y_dual) # Initialize parameters + expected_size = (self.J - 1) * self.K if init_beta is None: - init_beta = np.zeros((self.J - 1) * self.K) - else: - # Validate initial parameters - expected_size = (self.J - 1) * self.K - if init_beta.size != expected_size: - raise ValueError( - f"init_beta must have size {expected_size}, got {init_beta.size}" - ) - - # Run optimization using optimized internal functions - result = minimize( - fun=self._neg_log_likelihood_fast, - jac=self._gradient_fast, - x0=init_beta, - args=(X, single_indices, dual_indices), - method=method, - options=options, - ) + init_beta = np.zeros(expected_size) + elif init_beta.size != expected_size: + raise ValueError( + f"init_beta must have size {expected_size}, got {init_beta.size}" + ) - # Check convergence - if not result.success: - raise RuntimeError( - f"Optimization failed to converge: {result.message}\n" - f"Try a different optimization method or adjust tolerance." + rng = rng or np.random.default_rng() + + def run_optimization(start_beta: npt.NDArray[np.float64]) -> OptimizeResult: + """Run optimization for a given starting point.""" + return minimize( + fun=self._neg_log_likelihood, + jac=self._gradient, + x0=start_beta, + args=(X, single_indices, dual_indices), + method=method, + bounds=bounds, + constraints=constraints, + options=options, ) - # Store results - self.coef_ = result.x.reshape(self.J - 1, self.K) - self.optimization_result_ = result + # Run optimizations and keep best result + best_result: OptimizeResult | None = None + start_points = [init_beta] + if num_restarts > 0: + noise = rng.normal(scale=restart_scale, size=(num_restarts, expected_size)) + start_points.extend(init_beta + noise_row for noise_row in noise) - return self + for start in start_points: + # Run optimization + result = run_optimization(start) + if best_result is None or result.fun < best_result.fun: + best_result = result - def _validate_data( - self, - X: npt.NDArray[np.float64], - y_single: npt.NDArray[np.int8], - y_dual: npt.NDArray[np.int8], - ) -> None: - """Validate input data dimensions and constraints.""" - N = X.shape[0] - - if X.shape[1] != self.K: - raise ValueError(f"X must have {self.K} columns, got {X.shape[1]}") - - if y_single.shape != (N, self.J): - raise ValueError( - f"y_single must have shape ({N}, {self.J}), got {y_single.shape}" - ) + assert best_result is not None - if y_dual.shape != (N, self.J, self.J): - raise ValueError( - f"y_dual must have shape ({N}, {self.J}, {self.J}), got {y_dual.shape}" + if not best_result.success: + raise RuntimeError( + f"Optimization failed to converge: {best_result.message}\n" + f"Try a different optimization method or adjust tolerance." ) - # Check that each agent has exactly one choice - single_choices = y_single.sum(axis=1) - dual_choices = y_dual.sum(axis=(1, 2)) - total_choices = single_choices + dual_choices - - if not np.all(total_choices == 1): - invalid = np.where(total_choices != 1)[0] - raise ValueError( - f"Each agent must have exactly one choice. " - f"Agents with invalid choices: {invalid[:10]}" - + ("..." if len(invalid) > 10 else "") - ) + # Store results + self.coef_ = best_result.x.reshape(self.J - 1, self.K) + self.optimization_result_ = best_result - # Check that dual choices are in upper triangle (s < t) - if np.any(y_dual): - lower_triangle = np.tril_indices(self.J, k=0) - if np.any(y_dual[:, lower_triangle[0], lower_triangle[1]]): - raise ValueError( - "y_dual must only have entries in upper triangle (s < t)" - ) + return self - def _neg_log_likelihood_fast( + def _neg_log_likelihood( self, flat_beta: npt.NDArray[np.float64], X: npt.NDArray[np.float64], single_indices: tuple[np.ndarray, np.ndarray], dual_indices: tuple[np.ndarray, np.ndarray, np.ndarray], ) -> float: - """ - Internal optimized NLL using pre-computed indices. - """ - beta = self.transform_params(flat_beta) - V, a = self.calculate_utilities(X, beta) + """Compute negative log-likelihood using pre-computed indices.""" + beta = self._transform_params(flat_beta) + V, a = self._calculate_utilities(X, beta) log_lik = 0.0 @@ -258,16 +385,16 @@ def _neg_log_likelihood_fast( V_chosen = V[row_idx, col_idx] # Extract full V for relevant rows - # We could use row_idx directly but logsumexp over all J is needed V_sub = V[row_idx] - # Use logsumexp for numerical stability + # Use logsumexp for numerical stability (avoids overflow in exp) log_sum = logsumexp(V_sub, axis=1) log_lik += np.sum(V_chosen - log_sum) # 2. Handle Dual Choices if len(dual_indices[0]) > 0: + # Extract all dual choice indices at once i_idx, s_idx, t_idx = dual_indices # Get utilities for all dual choices at once @@ -278,7 +405,8 @@ def _neg_log_likelihood_fast( a_sum = np.sum(a[i_idx], axis=1) # Shape: (n_dual,) R = a_sum - a_s - a_t # Shape: (n_dual,) - # Vectorized probability computation + # Vectorized probability computation using inclusion-exclusion: + # P = a_s/(a_s+R) + a_t/(a_t+R) - (a_s+a_t)/(a_s+a_t+R) D1 = a_s + R D2 = a_t + R D3 = a_s + a_t + R @@ -293,35 +421,18 @@ def _neg_log_likelihood_fast( return -log_lik - def neg_log_likelihood( - self, - flat_beta: npt.NDArray[np.float64], - X: npt.NDArray[np.float64], - y_single: npt.NDArray[np.int8], - y_dual: npt.NDArray[np.int8], - ) -> float: - """ - Computes the negative log-likelihood of the model. - Wrapper for public API compliance that computes indices on the fly. - """ - self._validate_data(X, y_single, y_dual) - single_indices, dual_indices = self._prepare_data(y_single, y_dual) - return self._neg_log_likelihood_fast(flat_beta, X, single_indices, dual_indices) - - def _gradient_fast( + def _gradient( self, flat_beta: npt.NDArray[np.float64], X: npt.NDArray[np.float64], single_indices: tuple[np.ndarray, np.ndarray], dual_indices: tuple[np.ndarray, np.ndarray, np.ndarray], ) -> npt.NDArray[np.float64]: - """ - Internal optimized gradient using pre-computed indices and matrix multiplication. - """ - beta = self.transform_params(flat_beta) - V, a = self.calculate_utilities(X, beta) + """Compute gradient using pre-computed indices and matrix multiplication.""" + beta = self._transform_params(flat_beta) + V, a = self._calculate_utilities(X, beta) - # Gradient buffer for all parameters (J, K) + # Gradient buffer for all parameters (J, K); flattened later to ((J-1)*K,) grad = np.zeros((self.J, self.K)) # 1. Single Choice Gradient (Standard MNL) @@ -336,20 +447,11 @@ def _gradient_fast( # Gradient = X * (y - p) # y is one-hot, so for chosen col y=1, else 0 - # We can do this by subtracting p from y, but y is sparse indices - # Easier: grad += X_sub.T @ (y_onehot - probs) - # But creating y_onehot is (N_sub, J). - # Memory efficient: grad += sum_i (delta_ij - p_ij) * x_i - # grad += X_sub[y=1] - X_sub.T @ probs - # Add positive term for chosen alternatives (y=1) - # X_sub corresponds to row_idx. - # We need to add X_i to grad[j] where j is chosen # Use np.add.at for sparse addition np.add.at(grad, col_idx, X_sub) # Subtract prob term for all alternatives - # grad -= probs.T @ X_sub grad -= probs.T @ X_sub # 2. Dual Choice Gradient @@ -357,7 +459,7 @@ def _gradient_fast( i_idx, s_idx, t_idx = dual_indices n_dual = len(i_idx) - # Get covariates and utilities + # Covariates and utilities for dual choices X_i = X[i_idx] # Shape: (n_dual, K) a_s = a[i_idx, s_idx] # Shape: (n_dual,) a_t = a[i_idx, t_idx] # Shape: (n_dual,) @@ -370,55 +472,54 @@ def _gradient_fast( D2 = a_t + R D3 = a_s + a_t + R - # Probability + # Probability (unclipped for gradient computation) P_raw = (a_s / D1) + (a_t / D2) - ((a_s + a_t) / D3) - # Clip mask + # Clip mask for safe division clipped_mask = P_raw >= CLIP_THRESHOLD - - # Safe division inv_P = np.zeros_like(P_raw) inv_P[clipped_mask] = 1.0 / P_raw[clipped_mask] - # Derivatives - dP_dVs = a_s * R * (1 / (D1**2) - 1 / (D3**2)) - dP_dVt = a_t * R * (1 / (D2**2) - 1 / (D3**2)) - common_r = (a_s + a_t) / (D3**2) - a_s / (D1**2) - a_t / (D2**2) + # Derivatives of P with respect to utilities + dP_dVs = a_s * R * (1 / (D1**2) - 1 / (D3**2)) # (n_dual,) + dP_dVt = a_t * R * (1 / (D2**2) - 1 / (D3**2)) # (n_dual,) + common_r = ( + (a_s + a_t) / (D3**2) - a_s / (D1**2) - a_t / (D2**2) + ) # (n_dual,) # Weights for s and t - w_s = inv_P * dP_dVs - w_t = inv_P * dP_dVt + w_s = inv_P * dP_dVs # (n_dual,) + w_t = inv_P * dP_dVt # (n_dual,) - # Weights for r (all alternatives) - w_r = inv_P * common_r + # Weights for r (all other alternatives) + w_r = inv_P * common_r # (n_dual,) + # Complexity: O(n_dual) weight computations; still cheaper than looping over J. - # For 'r' alternatives: grad += sum_i (w_r_i * a_ij) * X_i + # For 'r' alternatives (neither s nor t): grad += sum_i (w_r_i * a_ij) * X_i # But we must exclude j=s and j=t - # 1. Compute M = w_r[:, None] * a[i_idx] (Shape: n_dual, J) - M = w_r[:, np.newaxis] * a[i_idx] + a_i = a[i_idx] # (n_dual, J) + M = w_r[:, np.newaxis] * a_i # 2. Zero out columns s and t for each row - # Advanced indexing to set specific elements to 0 rows = np.arange(n_dual) M[rows, s_idx] = 0.0 M[rows, t_idx] = 0.0 # 3. Compute gradient contribution from r-terms using matrix multiplication - # grad += M.T @ X_i + # This is an O(n_dual * J * K) dense step; faster than looping for typical J. grad += M.T @ X_i - # 4. Add contributions from s and t + # 4. Add contributions from s and t (O(n_dual)) # grad[s] += sum(w_s * X_i) # grad[t] += sum(w_t * X_i) - - grad_contrib_s = w_s[:, np.newaxis] * X_i - grad_contrib_t = w_t[:, np.newaxis] * X_i + grad_contrib_s = w_s[:, np.newaxis] * X_i # (n_dual, K) + grad_contrib_t = w_t[:, np.newaxis] * X_i # (n_dual, K) np.add.at(grad, s_idx, grad_contrib_s) np.add.at(grad, t_idx, grad_contrib_t) - # Return negative gradient, remove fixed class 0 + # Return negative gradient for minimization, remove fixed class 0 return -grad[1:].flatten() def gradient( @@ -438,36 +539,60 @@ def gradient( def compute_standard_errors( self, - flat_beta: npt.NDArray[np.float64], X: npt.NDArray[np.float64], y_single: npt.NDArray[np.int8], - y_dual: npt.NDArray[np.int8], + y_dual: DualInput, + flat_beta: Optional[npt.NDArray[np.float64]] = None, + epsilon: Optional[float] = None, ) -> npt.NDArray[np.float64]: """ - Computes standard errors by approximating the Hessian of the negative log-likelihood - using central finite differences of the analytical gradient. + Compute standard errors via numerical Hessian approximation. + + Uses central finite differences of the analytical gradient to + approximate the Hessian, then inverts to get the covariance matrix. + + Args: + X: Covariate matrix (N, K). + y_single: Single-choice indicators (N, J). + y_dual: Dual-choice indicators (dense, sparse, or index triplet). + flat_beta: Parameter vector of size (J-1)*K. Uses fitted coef_ if None. + epsilon: Finite-difference step size. Defaults to HESSIAN_EPSILON. + + Returns: + 1D array of standard errors of size (J-1)*K. + + Raises: + ValueError: If model is unfitted and no parameters provided. + + Warns: + RuntimeWarning: If Hessian inverse has negative diagonal elements. """ - self._validate_data(X, y_single, y_dual) - single_indices, dual_indices = self._prepare_data(y_single, y_dual) + if flat_beta is None: + if self.coef_ is None: + raise ValueError( + "Model is not fitted. Provide flat_beta or call fit() first." + ) + flat_beta = self.coef_.flatten() + + single_indices, dual_indices = self._validate_data(X, y_single, y_dual) n_params = len(flat_beta) hessian = np.zeros((n_params, n_params)) + step = epsilon if epsilon is not None else HESSIAN_EPSILON # Central finite differences for Hessian for j in range(n_params): beta_plus = flat_beta.copy() - beta_plus[j] += HESSIAN_EPSILON + beta_plus[j] += step beta_minus = flat_beta.copy() - beta_minus[j] -= HESSIAN_EPSILON + beta_minus[j] -= step - grad_plus = self._gradient_fast(beta_plus, X, single_indices, dual_indices) - grad_minus = self._gradient_fast( - beta_minus, X, single_indices, dual_indices - ) + grad_plus = self._gradient(beta_plus, X, single_indices, dual_indices) + grad_minus = self._gradient(beta_minus, X, single_indices, dual_indices) # Central difference: O(ε²) accuracy - hessian[:, j] = (grad_plus - grad_minus) / (2 * HESSIAN_EPSILON) + hessian[:, j] = (grad_plus - grad_minus) / (2 * step) # Variance-Covariance Matrix is inverse of Hessian # Use pinv for stability with near-singular Hessians @@ -478,15 +603,141 @@ def compute_standard_errors( # If any variance is negative (numerical noise with singular hessian), warn if np.any(diag_cov < 0): - import warnings - warnings.warn( - "Hessian inverse has negative diagonal elements. Standard errors may be unreliable.", + "Hessian inverse has negative diagonal elements. " + "Standard errors may be unreliable.", RuntimeWarning, stacklevel=2, ) # Compute standard errors, setting invalid (negative variance) to NaN std_errs = np.sqrt(np.where(diag_cov >= 0, diag_cov, np.nan)) - return typing.cast(npt.NDArray[np.float64], std_errs) + + def predict_proba( + self, + X: npt.NDArray[np.float64], + flat_beta: Optional[npt.NDArray[np.float64]] = None, + ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]: + """ + Compute predicted probabilities for single and dual choices. + + Args: + X: Covariate matrix (N, K). + flat_beta: Parameter vector of size (J-1)*K. Uses fitted coef_ if None. + + Returns: + single_probs: (N, J) softmax probabilities for single choices. + dual_probs: (N, J, J) probabilities for unordered pairs. + Only upper triangle (s < t) is populated. + + Raises: + ValueError: If model is unfitted and no parameters provided. + """ + if flat_beta is None: + if self.coef_ is None: + raise ValueError( + "Model is not fitted. Provide flat_beta or call fit() first." + ) + flat_beta = self.coef_.flatten() + + beta = self._transform_params(flat_beta) + V, a = self._calculate_utilities(X, beta) + + # Single choice probabilities (softmax) + single_probs = a / np.sum(a, axis=1, keepdims=True) + + # Dual choice probabilities: compute for each pair s npt.NDArray[np.float64]: + """ + Compute per-observation log-likelihood contributions. + + Useful for diagnostics, cross-validation, or computing information criteria. + + Args: + X: Covariate matrix (N, K). + y_single: Single-choice indicators (N, J). + y_dual: Dual-choice indicators (dense, sparse, or index triplet). + flat_beta: Parameter vector of size (J-1)*K. Uses fitted coef_ if None. + + Returns: + Vector of per-observation log-likelihood contributions (N,). + """ + if flat_beta is None: + if self.coef_ is None: + raise ValueError( + "Model is not fitted. Provide flat_beta or call fit() first." + ) + flat_beta = self.coef_.flatten() + + single_indices, dual_indices = self._validate_data(X, y_single, y_dual) + + beta = self._transform_params(flat_beta) + V, a = self._calculate_utilities(X, beta) + contrib = np.zeros(X.shape[0]) + + if len(single_indices[0]) > 0: + row_idx, col_idx = single_indices + V_chosen = V[row_idx, col_idx] + log_sum = logsumexp(V[row_idx], axis=1) + contrib[row_idx] = V_chosen - log_sum + + if len(dual_indices[0]) > 0: + i_idx, s_idx, t_idx = dual_indices + a_s = a[i_idx, s_idx] + a_t = a[i_idx, t_idx] + a_sum = np.sum(a[i_idx], axis=1) + R = a_sum - a_s - a_t + probs = ( + (a_s / (a_s + R)) + (a_t / (a_t + R)) - ((a_s + a_t) / (a_s + a_t + R)) + ) + contrib[i_idx] = np.log(np.maximum(probs, CLIP_THRESHOLD)) + # contrib fills per-observation slots; untouched rows remain zero if no dual/single + + return contrib + + def log_likelihood( + self, + X: npt.NDArray[np.float64], + y_single: npt.NDArray[np.int8], + y_dual: DualInput, + flat_beta: Optional[npt.NDArray[np.float64]] = None, + ) -> float: + """ + Compute the total log-likelihood of the model. + + Args: + X: Covariate matrix (N, K). + y_single: Single-choice indicators (N, J). + y_dual: Dual-choice indicators (dense, sparse, or index triplet). + flat_beta: Parameter vector of size (J-1)*K. Uses fitted coef_ if None. + + Returns: + Scalar log-likelihood (not negated). + """ + return float( + np.sum(self.log_likelihood_contributions(X, y_single, y_dual, flat_beta)) + ) diff --git a/multe/simulate.py b/multe/simulate.py index f9b5873..da59c2c 100644 --- a/multe/simulate.py +++ b/multe/simulate.py @@ -6,6 +6,9 @@ Fully vectorized implementation for fast simulation. """ +from __future__ import annotations + +from typing import Optional, Tuple import numpy as np import numpy.typing as npt @@ -16,8 +19,10 @@ def simulate_data( K: int, true_beta: npt.NDArray[np.float64] | None = None, mix_ratio: float = 0.5, - seed: int = 42, -) -> tuple[ + seed: int | None = 42, + rng: np.random.Generator | None = None, + dtype: npt.DTypeLike = np.float64, +) -> Tuple[ npt.NDArray[np.float64], npt.NDArray[np.int8], npt.NDArray[np.int8], @@ -41,7 +46,9 @@ def simulate_data( If None, generated uniformly in [-1, 1]. mix_ratio (float): Fraction of population making single choices (0.0 to 1.0). Default is 0.5 (equal mix of single and dual choices). - seed (int): Random seed for reproducibility. Default is 42. + seed (int | None): Random seed for reproducibility. Ignored if rng is provided. + rng (np.random.Generator, optional): Use an existing RNG instead of seed. + dtype: dtype for generated covariates and parameters (default float64). Returns: tuple: (X, y_single, y_dual, true_beta_free) @@ -52,6 +59,7 @@ def simulate_data( Raises: ValueError: If N, J, or K are invalid, or if mix_ratio is not in [0, 1]. + ValueError: If true_beta has an incorrect shape. """ if N < 1: raise ValueError(f"N must be >= 1, got {N}") @@ -66,19 +74,19 @@ def simulate_data( f"true_beta must have shape ({J - 1}, {K}), got {true_beta.shape}" ) - rng = np.random.default_rng(seed) + rng = rng or np.random.default_rng(seed) # Generate covariates - X = rng.normal(size=(N, K)) + X = rng.normal(size=(N, K)).astype(dtype, copy=False) # Generate or Use Parameters if true_beta is None: - true_beta_free = rng.uniform(-1, 1, (J - 1, K)) + true_beta_free = rng.uniform(-1, 1, (J - 1, K)).astype(dtype, copy=False) else: - true_beta_free = true_beta + true_beta_free = true_beta.astype(dtype, copy=False) # Add fixed class 0 (identification constraint) - beta_full = np.vstack([np.zeros((1, K)), true_beta_free]) + beta_full = np.vstack([np.zeros((1, K), dtype=dtype), true_beta_free]) # Calculate Deterministic Utility V = X @ beta_full.T diff --git a/tests/test_model.py b/tests/test_model.py index 3c031f6..81ce1a7 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -33,7 +33,7 @@ def test_correct_shape(self): """Test that transform_params produces correct shape.""" model = MultichoiceLogit(num_alternatives=4, num_covariates=3) flat_beta = np.random.randn((4 - 1) * 3) - beta = model.transform_params(flat_beta) + beta = model._transform_params(flat_beta) assert beta.shape == (4, 3) assert np.allclose(beta[0], 0) # First row should be zeros @@ -44,7 +44,7 @@ def test_incorrect_size(self): flat_beta = np.random.randn(10) # Wrong size with pytest.raises(ValueError, match="flat_beta must have size"): - model.transform_params(flat_beta) + model._transform_params(flat_beta) class TestDataValidation: @@ -89,6 +89,21 @@ def test_wrong_y_dual_shape(self): with pytest.raises(ValueError, match="y_dual must have shape"): model._validate_data(X, y_single, y_dual_wrong) + def test_non_binary_inputs(self): + """Test that non-binary entries raise ValueError.""" + N, J, K = 20, 3, 2 + X, y_single, y_dual, _ = simulate_data(N, J, K, seed=42) + model = MultichoiceLogit(J, K) + + y_single[0, 0] = 2 + with pytest.raises(ValueError, match="binary"): + model._validate_data(X, y_single, y_dual) + + y_single[0, 0] = 1 # restore + y_dual[0, 1, 2] = 2 + with pytest.raises(ValueError, match="binary"): + model._validate_data(X, y_single, y_dual) + def test_multiple_choices_per_agent(self): """Test that multiple choices per agent raises ValueError.""" N, J, K = 10, 3, 2 @@ -135,6 +150,21 @@ def test_dual_choice_in_lower_triangle(self): with pytest.raises(ValueError, match="upper triangle"): model._validate_data(X, y_single, y_dual) + def test_dual_choice_on_diagonal(self): + """Test that diagonal dual entries raise ValueError.""" + N, J, K = 10, 3, 2 + model = MultichoiceLogit(J, K) + X = np.random.randn(N, K) + y_single = np.zeros((N, J), dtype=np.int8) + y_dual = np.zeros((N, J, J), dtype=np.int8) + y_single[:, 0] = 1 + + y_single[0, 0] = 0 + y_dual[0, 1, 1] = 1 + + with pytest.raises(ValueError, match="diagonal"): + model._validate_data(X, y_single, y_dual) + class TestNegLogLikelihood: """Test negative log-likelihood computation.""" @@ -146,7 +176,8 @@ def test_returns_scalar(self): model = MultichoiceLogit(J, K) flat_beta = true_beta.flatten() - nll = model.neg_log_likelihood(flat_beta, X, y_single, y_dual) + single_idx, dual_idx = model._validate_data(X, y_single, y_dual) + nll = model._neg_log_likelihood(flat_beta, X, single_idx, dual_idx) assert isinstance(nll, (float, np.floating)) @@ -157,7 +188,8 @@ def test_positive_value(self): model = MultichoiceLogit(J, K) flat_beta = true_beta.flatten() - nll = model.neg_log_likelihood(flat_beta, X, y_single, y_dual) + single_idx, dual_idx = model._validate_data(X, y_single, y_dual) + nll = model._neg_log_likelihood(flat_beta, X, single_idx, dual_idx) assert nll > 0 @@ -168,7 +200,8 @@ def test_all_single_choices(self): model = MultichoiceLogit(J, K) flat_beta = true_beta.flatten() - nll = model.neg_log_likelihood(flat_beta, X, y_single, y_dual) + single_idx, dual_idx = model._validate_data(X, y_single, y_dual) + nll = model._neg_log_likelihood(flat_beta, X, single_idx, dual_idx) assert nll > 0 assert np.sum(y_dual) == 0 # Verify no dual choices @@ -180,7 +213,8 @@ def test_all_dual_choices(self): model = MultichoiceLogit(J, K) flat_beta = true_beta.flatten() - nll = model.neg_log_likelihood(flat_beta, X, y_single, y_dual) + single_idx, dual_idx = model._validate_data(X, y_single, y_dual) + nll = model._neg_log_likelihood(flat_beta, X, single_idx, dual_idx) assert nll > 0 assert np.sum(y_single) == 0 # Verify no single choices @@ -203,15 +237,33 @@ def test_numerical_stability_large_utilities(self): flat_beta = large_beta.flatten() # Should not overflow or produce inf/nan - nll = model.neg_log_likelihood(flat_beta, X, y_single, y_dual) + single_idx, dual_idx = model._validate_data(X, y_single, y_dual) + nll = model._neg_log_likelihood(flat_beta, X, single_idx, dual_idx) assert np.isfinite(nll) assert nll > 0 # Gradient should also be stable - grad = model.gradient(flat_beta, X, y_single, y_dual) + grad = model._gradient(flat_beta, X, single_idx, dual_idx) assert np.all(np.isfinite(grad)) + def test_sparse_dual_input_equivalence(self): + """Test that tuple dual indices produce same NLL as dense tensor.""" + N, J, K = 150, 3, 2 + X, y_single, y_dual, true_beta = simulate_data(N, J, K, seed=42) + model = MultichoiceLogit(J, K) + + flat_beta = true_beta.flatten() + single_dense, dual_dense = model._validate_data(X, y_single, y_dual) + dense_nll = model._neg_log_likelihood(flat_beta, X, single_dense, dual_dense) + + rows, s_idx, t_idx = np.nonzero(y_dual) + tuple_input = (rows, s_idx, t_idx) + single_tuple, dual_tuple = model._validate_data(X, y_single, tuple_input) + tuple_nll = model._neg_log_likelihood(flat_beta, X, single_tuple, dual_tuple) + + assert np.isclose(dense_nll, tuple_nll) + class TestGradient: """Test gradient computation.""" @@ -223,7 +275,8 @@ def test_returns_correct_shape(self): model = MultichoiceLogit(J, K) flat_beta = true_beta.flatten() - grad = model.gradient(flat_beta, X, y_single, y_dual) + single_idx, dual_idx = model._validate_data(X, y_single, y_dual) + grad = model._gradient(flat_beta, X, single_idx, dual_idx) assert grad.shape == ((J - 1) * K,) @@ -234,7 +287,8 @@ def test_gradient_numerical_accuracy(self): model = MultichoiceLogit(J, K) flat_beta = true_beta.flatten() - analytical_grad = model.gradient(flat_beta, X, y_single, y_dual) + single_idx, dual_idx = model._validate_data(X, y_single, y_dual) + analytical_grad = model._gradient(flat_beta, X, single_idx, dual_idx) # Numerical gradient epsilon = 1e-5 @@ -245,8 +299,8 @@ def test_gradient_numerical_accuracy(self): beta_minus = flat_beta.copy() beta_minus[i] -= epsilon - nll_plus = model.neg_log_likelihood(beta_plus, X, y_single, y_dual) - nll_minus = model.neg_log_likelihood(beta_minus, X, y_single, y_dual) + nll_plus = model._neg_log_likelihood(beta_plus, X, single_idx, dual_idx) + nll_minus = model._neg_log_likelihood(beta_minus, X, single_idx, dual_idx) numerical_grad[i] = (nll_plus - nll_minus) / (2 * epsilon) @@ -266,7 +320,8 @@ def test_gradient_with_clipped_probabilities(self): flat_beta = extreme_beta.flatten() # Gradient should be finite even with clipped probabilities - grad = model.gradient(flat_beta, X, y_single, y_dual) + single_idx, dual_idx = model._validate_data(X, y_single, y_dual) + grad = model._gradient(flat_beta, X, single_idx, dual_idx) assert np.all(np.isfinite(grad)) @@ -279,14 +334,32 @@ def test_gradient_with_clipped_probabilities(self): beta_minus = flat_beta.copy() beta_minus[i] -= epsilon - nll_plus = model.neg_log_likelihood(beta_plus, X, y_single, y_dual) - nll_minus = model.neg_log_likelihood(beta_minus, X, y_single, y_dual) + nll_plus = model._neg_log_likelihood(beta_plus, X, single_idx, dual_idx) + nll_minus = model._neg_log_likelihood(beta_minus, X, single_idx, dual_idx) numerical_grad[i] = (nll_plus - nll_minus) / (2 * epsilon) # Should still be close even with clipped probabilities assert np.allclose(grad, numerical_grad, atol=1e-4) + def test_gradient_sparse_dual(self): + """Gradient with tuple dual input matches dense gradient.""" + N, J, K = 40, 3, 2 + X, y_single, y_dual, true_beta = simulate_data(N, J, K, mix_ratio=0.3, seed=42) + model = MultichoiceLogit(J, K) + + flat_beta = true_beta.flatten() + single_dense, dual_dense = model._validate_data(X, y_single, y_dual) + dense_grad = model._gradient(flat_beta, X, single_dense, dual_dense) + + dual_rows, dual_s, dual_t = np.nonzero(y_dual) + single_tuple, dual_tuple = model._validate_data( + X, y_single, (dual_rows, dual_s, dual_t) + ) + tuple_grad = model._gradient(flat_beta, X, single_tuple, dual_tuple) + + assert np.allclose(dense_grad, tuple_grad) + class TestComputeStandardErrors: """Test standard error computation.""" @@ -298,7 +371,7 @@ def test_returns_correct_shape(self): model = MultichoiceLogit(J, K) flat_beta = true_beta.flatten() - std_errs = model.compute_standard_errors(flat_beta, X, y_single, y_dual) + std_errs = model.compute_standard_errors(X, y_single, y_dual, flat_beta) assert std_errs.shape == ((J - 1) * K,) @@ -309,7 +382,7 @@ def test_positive_standard_errors(self): model = MultichoiceLogit(J, K) flat_beta = true_beta.flatten() - std_errs = model.compute_standard_errors(flat_beta, X, y_single, y_dual) + std_errs = model.compute_standard_errors(X, y_single, y_dual, flat_beta) # Standard errors should be positive (or NaN if singular) assert np.all((std_errs > 0) | np.isnan(std_errs)) @@ -324,7 +397,7 @@ def test_singular_hessian_warning(self): flat_beta = np.zeros((J - 1) * K) # Should produce RuntimeWarning in some cases (not always) - std_errs = model.compute_standard_errors(flat_beta, X, y_single, y_dual) + std_errs = model.compute_standard_errors(X, y_single, y_dual, flat_beta) assert std_errs.shape == ((J - 1) * K,) @@ -453,12 +526,18 @@ def test_estimation_recovers_parameters(self): model = MultichoiceLogit(J, K) init_beta = np.zeros((J - 1) * K) + single_idx, dual_idx = model._validate_data(X, y_single, y_dual) + + def fun(beta): + return model._neg_log_likelihood(beta, X, single_idx, dual_idx) + + def jac(beta): + return model._gradient(beta, X, single_idx, dual_idx) result = minimize( - fun=model.neg_log_likelihood, - jac=model.gradient, + fun=fun, + jac=jac, x0=init_beta, - args=(X, y_single, y_dual), method="BFGS", options={"gtol": 1e-5, "maxiter": 1000}, ) @@ -469,3 +548,54 @@ def test_estimation_recovers_parameters(self): # Mean absolute error should be reasonably small mae = np.mean(np.abs(est_beta - true_beta)) assert mae < 0.15 # Reasonable tolerance for N=1000 + + +class TestHelpers: + """Tests for helper APIs such as predict_proba and log_likelihood_contributions.""" + + def test_predict_proba_shapes_and_sums(self): + N, J, K = 50, 3, 2 + X, y_single, y_dual, true_beta = simulate_data(N, J, K, seed=42) + model = MultichoiceLogit(J, K) + + flat_beta = true_beta.flatten() + single_probs, dual_probs = model.predict_proba(X, flat_beta=flat_beta) + + assert single_probs.shape == (N, J) + assert dual_probs.shape == (N, J, J) + # Single probabilities should sum to 1 + np.testing.assert_allclose(np.sum(single_probs, axis=1), 1.0, atol=1e-6) + # Dual probabilities upper triangle should be non-negative + assert np.all( + dual_probs[:, np.triu_indices(J, k=1)[0], np.triu_indices(J, k=1)[1]] >= 0 + ) + + def test_log_likelihood_contributions_sum(self): + N, J, K = 80, 3, 2 + X, y_single, y_dual, true_beta = simulate_data(N, J, K, seed=42) + model = MultichoiceLogit(J, K) + flat_beta = true_beta.flatten() + + contributions = model.log_likelihood_contributions( + X, y_single, y_dual, flat_beta + ) + total = np.sum(contributions) + single_idx, dual_idx = model._validate_data(X, y_single, y_dual) + direct = -model._neg_log_likelihood(flat_beta, X, single_idx, dual_idx) + + assert contributions.shape == (N,) + np.testing.assert_allclose(total, direct, rtol=1e-6, atol=1e-6) + + def test_compute_standard_errors_custom_epsilon(self): + N, J, K = 60, 3, 2 + X, y_single, y_dual, true_beta = simulate_data(N, J, K, seed=42) + model = MultichoiceLogit(J, K) + + flat_beta = true_beta.flatten() + std_errs_default = model.compute_standard_errors(X, y_single, y_dual, flat_beta) + std_errs_smaller = model.compute_standard_errors( + X, y_single, y_dual, flat_beta, epsilon=1e-6 + ) + + assert std_errs_default.shape == std_errs_smaller.shape == ((J - 1) * K,) + assert np.all(np.isfinite(std_errs_default)) diff --git a/tests/test_simulate.py b/tests/test_simulate.py index 0a87114..3c85468 100644 --- a/tests/test_simulate.py +++ b/tests/test_simulate.py @@ -135,6 +135,20 @@ def test_seed_reproducibility(self): assert np.all(y_dual1 == y_dual2) assert np.allclose(beta1, beta2) + def test_rng_argument_reproducibility(self): + """Test that supplying rng yields reproducible draws independent of seed.""" + N, J, K = 50, 3, 2 + rng1 = np.random.default_rng(123) + rng2 = np.random.default_rng(123) + + X1, y_single1, y_dual1, beta1 = simulate_data(N, J, K, rng=rng1, seed=None) + X2, y_single2, y_dual2, beta2 = simulate_data(N, J, K, rng=rng2, seed=None) + + assert np.allclose(X1, X2) + assert np.all(y_single1 == y_single2) + assert np.all(y_dual1 == y_dual2) + assert np.allclose(beta1, beta2) + def test_different_seeds_different_results(self): """Test that different seeds produce different results.""" N, J, K = 100, 3, 2 @@ -149,6 +163,14 @@ def test_different_seeds_different_results(self): and np.all(y_dual1 == y_dual2) ) + def test_dtype_control(self): + """Test that dtype parameter controls output dtype.""" + N, J, K = 20, 3, 2 + X, y_single, y_dual, beta = simulate_data(N, J, K, seed=42, dtype=np.float32) + + assert X.dtype == np.float32 + assert beta.dtype == np.float32 + def test_binary_outputs(self): """Test that y_single and y_dual contain only 0 and 1.""" N, J, K = 100, 3, 2