From 3908fe39e46793c8881c0075ea5f9085cd871599 Mon Sep 17 00:00:00 2001
From: Thomas Monk <246525+tmonk@users.noreply.github.com>
Date: Mon, 24 Nov 2025 23:13:45 +0000
Subject: [PATCH 1/5] =?UTF-8?q?=E2=80=A2=20Add=20sparse=20dual=20caching,?=
 =?UTF-8?q?=20new=20helpers,=20and=20validation/doc=20improvements?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

  - Reuse normalized dual indices across fit/likelihood/gradient/SEs to avoid double sparse conversion; support cached indices in _prepare_data.
  - Accept sparse/index dual inputs with stronger validation (binary, upper triangle) and document sparse flattening order (s*J+t); remove redundant diagonal check.
  - Add predict_proba and per-observation log-likelihood helpers; expose log_likelihood wrapper; extend simulate_data with rng/dtype controls.
  - Expand tests for validation edge cases, sparse equivalence, helpers, and reproducibility.
  - Keep dual gradient comments aligned with the vectorized implementation; all tests pass.
---
 .github/workflows/tests.yml    |   2 +-
 .pre-commit-config.yaml        |  22 ++
 examples/basic_example.py      |  30 +-
 examples/benchmark.py          | 105 ++---
 examples/csv_example.py        |  21 +-
 examples/simple_fit_example.py |  18 +-
 multe/model.py                 | 697 +++++++++++++++++++++++----------
 multe/simulate.py              |  36 +-
 tests/test_model.py            | 109 ++++++
 tests/test_simulate.py         |  22 ++
 10 files changed, 764 insertions(+), 298 deletions(-)
 create mode 100644 .pre-commit-config.yaml

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 42aa126..8136de2 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -30,4 +30,4 @@ jobs:
 
     - name: Run tests with pytest
       run: |
-        pytest tests/ -v --tb=short
\ No newline at end of file
+        pytest tests/ -v --tb=short
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 0000000..fcb3e78
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,22 @@
+# See https://pre-commit.com for more information
+# See https://pre-commit.com/hooks.html for more hooks
+repos:
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v6.0.0
+    hooks:
+    -   id: trailing-whitespace
+    -   id: end-of-file-fixer
+    -   id: check-yaml
+    -   id: check-added-large-files
+
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  # Ruff version.
+  rev: v0.14.6
+  hooks:
+    # Run the linter.
+    - id: ruff-check
+      types_or: [ python, pyi ]
+      args: [ --fix ]
+    # Run the formatter.
+    - id: ruff-format
+      types_or: [ python, pyi ]
diff --git a/examples/basic_example.py b/examples/basic_example.py
index 4564d55..f98e46a 100644
--- a/examples/basic_example.py
+++ b/examples/basic_example.py
@@ -18,27 +18,25 @@
 def main():
     # Settings
     N = 2000  # Number of observations
-    J = 4     # Number of alternatives
-    K = 3     # Number of covariates
+    J = 4  # Number of alternatives
+    K = 3  # Number of covariates
 
     print(f"Simulating Data (N={N}, J={J}, K={K})...")
-    X, y_single, y_dual, true_beta = simulate_data(
-        N, J, K, mix_ratio=0.5, seed=42
-    )
+    X, y_single, y_dual, true_beta = simulate_data(N, J, K, mix_ratio=0.5, seed=42)
 
     print("Starting Estimation...")
     model = MultichoiceLogit(J, K)
 
     # Initial guess (zeros)
-    init_beta = np.zeros((J-1) * K)
+    init_beta = np.zeros((J - 1) * K)
 
     res = minimize(
         fun=model.neg_log_likelihood,
         jac=model.gradient,
         x0=init_beta,
         args=(X, y_single, y_dual),
-        method='L-BFGS-B',
-        options={'disp': True, 'gtol': 1e-5}
+        method="L-BFGS-B",
+        options={"disp": True, "gtol": 1e-5},
     )
 
     print("\nOptimization Success:", res.success)
@@ -47,17 +45,19 @@ def main():
     print("Computing Standard Errors...")
     std_errs = model.compute_standard_errors(res.x, X, y_single, y_dual)
 
-    est_beta = res.x.reshape(J-1, K)
-    std_errs_reshaped = std_errs.reshape(J-1, K)
+    est_beta = res.x.reshape(J - 1, K)
+    std_errs_reshaped = std_errs.reshape(J - 1, K)
 
     print("\nComparison (Row 0 is fixed to 0, these are rows 1 to J-1):")
     # Header
     print("-" * 88)
-    print(f"{'True':<10} | {'Est':<10} | {'SE':<10} | {'t-stat':<10} | {'p-val':<10} | {'95% CI':<15}")
+    print(
+        f"{'True':<10} | {'Est':<10} | {'SE':<10} | {'t-stat':<10} | {'p-val':<10} | {'95% CI':<15}"
+    )
     print("-" * 88)
 
-    for i in range(J-1):
-        print(f"Alternative {i+1}:")
+    for i in range(J - 1):
+        print(f"Alternative {i + 1}:")
         for k in range(K):
             t_val = true_beta[i, k]
             e_val = est_beta[i, k]
@@ -72,7 +72,9 @@ def main():
             ci_upper = e_val + 1.96 * se
             ci_str = f"[{ci_lower:.2f}, {ci_upper:.2f}]"
 
-            print(f" {t_val:<9.4f} | {e_val:<9.4f}  | {se:<9.4f}  | {t_stat:<9.2f}  | {p_val:<9.4f}  | {ci_str:<15}")
+            print(
+                f" {t_val:<9.4f} | {e_val:<9.4f}  | {se:<9.4f}  | {t_stat:<9.2f}  | {p_val:<9.4f}  | {ci_str:<15}"
+            )
         print("-" * 88)
 
     mae = np.mean(np.abs(est_beta - true_beta))
diff --git a/examples/benchmark.py b/examples/benchmark.py
index 97bdd61..3578e2e 100644
--- a/examples/benchmark.py
+++ b/examples/benchmark.py
@@ -10,7 +10,7 @@
 from multe import MultichoiceLogit, simulate_data
 
 
-def benchmark_estimation(N, J, K, mix_ratio=0.5, method='BFGS', seed=42):
+def benchmark_estimation(N, J, K, mix_ratio=0.5, method="BFGS", seed=42):
     """
     Benchmark a single estimation run.
 
@@ -19,12 +19,14 @@ def benchmark_estimation(N, J, K, mix_ratio=0.5, method='BFGS', seed=42):
     """
     # Simulate data
     t0 = time.time()
-    X, y_single, y_dual, true_beta = simulate_data(N, J, K, mix_ratio=mix_ratio, seed=seed)
+    X, y_single, y_dual, true_beta = simulate_data(
+        N, J, K, mix_ratio=mix_ratio, seed=seed
+    )
     sim_time = time.time() - t0
 
     # Initialize model
     model = MultichoiceLogit(J, K)
-    init_beta = np.zeros((J-1) * K)
+    init_beta = np.zeros((J - 1) * K)
 
     # Time likelihood evaluation
     t0 = time.time()
@@ -44,59 +46,66 @@ def benchmark_estimation(N, J, K, mix_ratio=0.5, method='BFGS', seed=42):
         x0=init_beta,
         args=(X, y_single, y_dual),
         method=method,
-        options={'disp': False, 'gtol': 1e-5, 'maxiter': 1000}
+        options={"disp": False, "gtol": 1e-5, "maxiter": 1000},
     )
     opt_time = time.time() - t0
 
     # Compute accuracy
-    est_beta = result.x.reshape(J-1, K)
+    est_beta = result.x.reshape(J - 1, K)
     mae = np.mean(np.abs(est_beta - true_beta))
-    rmse = np.sqrt(np.mean((est_beta - true_beta)**2))
+    rmse = np.sqrt(np.mean((est_beta - true_beta) ** 2))
     max_error = np.max(np.abs(est_beta - true_beta))
 
     # Compute standard errors
     t0 = time.time()
-    std_errs = 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 {
-        'N': N, 'J': J, 'K': K,
-        'method': method,
-        'sim_time': sim_time,
-        'likelihood_time': likelihood_time,
-        'gradient_time': gradient_time,
-        'opt_time': opt_time,
-        'se_time': se_time,
-        'total_time': sim_time + opt_time + se_time,
-        'success': result.success,
-        'nit': result.nit,
-        'nfev': result.nfev,
-        'final_nll': result.fun,
-        'mae': mae,
-        'rmse': rmse,
-        'max_error': max_error
+        "N": N,
+        "J": J,
+        "K": K,
+        "method": method,
+        "sim_time": sim_time,
+        "likelihood_time": likelihood_time,
+        "gradient_time": gradient_time,
+        "opt_time": opt_time,
+        "se_time": se_time,
+        "total_time": sim_time + opt_time + se_time,
+        "success": result.success,
+        "nit": result.nit,
+        "nfev": result.nfev,
+        "final_nll": result.fun,
+        "mae": mae,
+        "rmse": rmse,
+        "max_error": max_error,
     }
 
 
 def print_results(results):
     """Pretty print benchmark results."""
-    print("\n" + "="*100)
-    print(f"{'N':<7} {'J':<4} {'K':<4} {'Method':<10} {'Opt(s)':<8} {'SE(s)':<8} {'Total(s)':<9} "
-          f"{'Iters':<6} {'FEval':<6} {'MAE':<8} {'RMSE':<8} {'MaxErr':<8} {'Success':<7}")
-    print("="*100)
+    print("\n" + "=" * 100)
+    print(
+        f"{'N':<7} {'J':<4} {'K':<4} {'Method':<10} {'Opt(s)':<8} {'SE(s)':<8} {'Total(s)':<9} "
+        f"{'Iters':<6} {'FEval':<6} {'MAE':<8} {'RMSE':<8} {'MaxErr':<8} {'Success':<7}"
+    )
+    print("=" * 100)
 
     for r in results:
-        print(f"{r['N']:<7} {r['J']:<4} {r['K']:<4} {r['method']:<10} "
-              f"{r['opt_time']:<8.3f} {r['se_time']:<8.3f} {r['total_time']:<9.3f} "
-              f"{r['nit']:<6} {r['nfev']:<6} "
-              f"{r['mae']:<8.4f} {r['rmse']:<8.4f} {r['max_error']:<8.4f} "
-              f"{'✓' if r['success'] else '✗':<7}")
-    print("="*100)
+        print(
+            f"{r['N']:<7} {r['J']:<4} {r['K']:<4} {r['method']:<10} "
+            f"{r['opt_time']:<8.3f} {r['se_time']:<8.3f} {r['total_time']:<9.3f} "
+            f"{r['nit']:<6} {r['nfev']:<6} "
+            f"{r['mae']:<8.4f} {r['rmse']:<8.4f} {r['max_error']:<8.4f} "
+            f"{'✓' if r['success'] else '✗':<7}"
+        )
+    print("=" * 100)
 
 
 def main():
     print("Multichoice Logit Estimation Benchmark")
-    print("="*100)
+    print("=" * 100)
 
     # Test different problem sizes
     problem_sizes = [
@@ -108,7 +117,7 @@ def main():
     ]
 
     # Test different optimization methods
-    methods = ['BFGS', 'L-BFGS-B', 'Newton-CG']
+    methods = ["BFGS", "L-BFGS-B", "Newton-CG"]
 
     results = []
 
@@ -117,7 +126,7 @@ def main():
     print("-" * 100)
     for N, J, K in problem_sizes:
         print(f"Running: N={N}, J={J}, K={K}...", end=" ", flush=True)
-        r = benchmark_estimation(N, J, K, method='BFGS')
+        r = benchmark_estimation(N, J, K, method="BFGS")
         results.append(r)
         print(f"✓ ({r['opt_time']:.2f}s)")
 
@@ -139,8 +148,8 @@ def main():
     print("-" * 100)
     for mix_ratio in [0.2, 0.5, 0.8]:
         print(f"Running: mix_ratio={mix_ratio}...", end=" ", flush=True)
-        r = benchmark_estimation(2000, 4, 3, mix_ratio=mix_ratio, method='BFGS')
-        r['mix_ratio'] = mix_ratio
+        r = benchmark_estimation(2000, 4, 3, mix_ratio=mix_ratio, method="BFGS")
+        r["mix_ratio"] = mix_ratio
         results.append(r)
         print(f"✓ ({r['opt_time']:.2f}s)")
 
@@ -150,16 +159,24 @@ def main():
     # Summary statistics
     print("\nSummary:")
     print("-" * 100)
-    bfgs_results = [r for r in results if r['method'] == 'BFGS' and 'mix_ratio' not in r]
+    bfgs_results = [
+        r for r in results if r["method"] == "BFGS" and "mix_ratio" not in r
+    ]
     if bfgs_results:
-        avg_time_per_1k = np.mean([r['opt_time'] / (r['N']/1000) for r in bfgs_results])
-        avg_mae = np.mean([r['mae'] for r in bfgs_results])
-        avg_rmse = np.mean([r['rmse'] for r in bfgs_results])
-
-        print(f"Average optimization time per 1000 observations: {avg_time_per_1k:.3f}s")
+        avg_time_per_1k = np.mean(
+            [r["opt_time"] / (r["N"] / 1000) for r in bfgs_results]
+        )
+        avg_mae = np.mean([r["mae"] for r in bfgs_results])
+        avg_rmse = np.mean([r["rmse"] for r in bfgs_results])
+
+        print(
+            f"Average optimization time per 1000 observations: {avg_time_per_1k:.3f}s"
+        )
         print(f"Average MAE: {avg_mae:.4f}")
         print(f"Average RMSE: {avg_rmse:.4f}")
-        print(f"Success rate: {sum(r['success'] for r in bfgs_results) / len(bfgs_results) * 100:.1f}%")
+        print(
+            f"Success rate: {sum(r['success'] for r in bfgs_results) / len(bfgs_results) * 100:.1f}%"
+        )
 
 
 if __name__ == "__main__":
diff --git a/examples/csv_example.py b/examples/csv_example.py
index 1450e5c..a1a0c6c 100644
--- a/examples/csv_example.py
+++ b/examples/csv_example.py
@@ -25,7 +25,6 @@ def save_data_to_csv(X, y_single, y_dual, filepath_prefix="data"):
         filepath_prefix: Prefix for output files
     """
     N, K = X.shape
-    J = y_single.shape[1]
 
     # Save covariates
     covariate_df = pd.DataFrame(X, columns=[f"x_{k}" for k in range(K)])
@@ -62,7 +61,9 @@ def save_data_to_csv(X, y_single, y_dual, filepath_prefix="data"):
     choices_df = pd.DataFrame(choices_list)
     choices_df.to_csv(f"{filepath_prefix}_choices.csv", index=False)
 
-    print(f"Data saved to {filepath_prefix}_covariates.csv and {filepath_prefix}_choices.csv")
+    print(
+        f"Data saved to {filepath_prefix}_covariates.csv and {filepath_prefix}_choices.csv"
+    )
 
 
 def load_data_from_csv(filepath_prefix="data", J=None):
@@ -165,12 +166,12 @@ def main():
     print("\nTrue vs Estimated Parameters:")
     print("-" * 60)
     for j in range(J - 1):
-        print(f"Alternative {j+1}:")
+        print(f"Alternative {j + 1}:")
         for k in range(K):
             print(
-                f"  Covariate {k}: True={true_beta[j,k]:7.4f}, "
-                f"Est={est_beta[j,k]:7.4f}, "
-                f"Error={abs(true_beta[j,k] - est_beta[j,k]):7.4f}"
+                f"  Covariate {k}: True={true_beta[j, k]:7.4f}, "
+                f"Est={est_beta[j, k]:7.4f}, "
+                f"Error={abs(true_beta[j, k] - est_beta[j, k]):7.4f}"
             )
     print("-" * 60)
     print(f"Mean Absolute Error: {mae:.4f}")
@@ -185,12 +186,12 @@ def main():
     print("\nParameter Estimates with Standard Errors:")
     print("-" * 60)
     for j in range(J - 1):
-        print(f"Alternative {j+1}:")
+        print(f"Alternative {j + 1}:")
         for k in range(K):
             t_stat = est_beta[j, k] / std_errs_reshaped[j, k]
             print(
-                f"  Covariate {k}: {est_beta[j,k]:7.4f} "
-                f"(SE: {std_errs_reshaped[j,k]:6.4f}, "
+                f"  Covariate {k}: {est_beta[j, k]:7.4f} "
+                f"(SE: {std_errs_reshaped[j, k]:6.4f}, "
                 f"t: {t_stat:6.2f})"
             )
 
@@ -205,7 +206,7 @@ def main():
         os.remove("example_data_covariates.csv")
         os.remove("example_data_choices.csv")
         print("\nTemporary CSV files cleaned up.")
-    except:
+    except FileNotFoundError:
         pass
 
 
diff --git a/examples/simple_fit_example.py b/examples/simple_fit_example.py
index 838e8e5..feb4504 100644
--- a/examples/simple_fit_example.py
+++ b/examples/simple_fit_example.py
@@ -29,9 +29,9 @@ def main():
     print("\n3. Fitted coefficients (model.coef_):")
     print("-" * 70)
     print(f"   Shape: {model.coef_.shape}")
-    print(f"\n   Values:")
+    print("\n   Values:")
     for j in range(J - 1):
-        print(f"   Alternative {j+1}: {model.coef_[j]}")
+        print(f"   Alternative {j + 1}: {model.coef_[j]}")
 
     # Compare to true parameters
     print("\n4. Comparison to true parameters:")
@@ -48,13 +48,21 @@ def main():
     print("\n   Coefficients with Standard Errors:")
     print("-" * 70)
     for j in range(J - 1):
-        print(f"   Alternative {j+1}:")
+        print(f"   Alternative {j + 1}:")
         for k in range(K):
             coef = model.coef_[j, k]
             se = std_errs[j, k]
             t_stat = coef / se
-            sig = "***" if abs(t_stat) > 2.576 else ("**" if abs(t_stat) > 1.96 else ("*" if abs(t_stat) > 1.645 else ""))
-            print(f"      Covariate {k}: {coef:7.4f} (SE: {se:6.4f}, t: {t_stat:6.2f}) {sig}")
+            sig = (
+                "***"
+                if abs(t_stat) > 2.576
+                else (
+                    "**" if abs(t_stat) > 1.96 else ("*" if abs(t_stat) > 1.645 else "")
+                )
+            )
+            print(
+                f"      Covariate {k}: {coef:7.4f} (SE: {se:6.4f}, t: {t_stat:6.2f}) {sig}"
+            )
 
     print("\n" + "=" * 70)
     print("Example completed! ✓")
diff --git a/multe/model.py b/multe/model.py
index 0889800..86f1e76 100644
--- a/multe/model.py
+++ b/multe/model.py
@@ -4,16 +4,28 @@
 Vectorized implementation for fast and accurate MLE estimation.
 """
 
-from typing import Tuple, Optional, Dict, Any
+from __future__ import annotations
+
+import typing
+from typing import Any, Optional, Sequence
+
 import numpy as np
 import numpy.typing as npt
-from scipy.optimize import minimize, OptimizeResult
+import scipy.sparse as sp
+from scipy.optimize import OptimizeResult, minimize
 from scipy.special import logsumexp
 
 # Numerical constants for stability and accuracy
 CLIP_THRESHOLD = 1e-10  # Minimum probability value (avoid log(0))
 HESSIAN_EPSILON = 1e-5  # Step size for Hessian finite differences
 
+DualInput = (
+    npt.NDArray[np.int8]
+    | npt.NDArray[np.int64]
+    | tuple[np.ndarray, np.ndarray, np.ndarray]
+    | sp.spmatrix
+)
+
 
 class MultichoiceLogit:
     """
@@ -46,10 +58,12 @@ def __init__(self, num_alternatives: int, num_covariates: int) -> None:
         self.K = num_covariates
 
         # Fitted attributes (set by fit method)
-        self.coef_: Optional[npt.NDArray[np.float64]] = None
-        self.optimization_result_: Optional[OptimizeResult] = None
+        self.coef_: npt.NDArray[np.float64] | None = None
+        self.optimization_result_: OptimizeResult | None = None
 
-    def transform_params(self, flat_beta: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+    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.
 
@@ -76,7 +90,7 @@ def transform_params(self, flat_beta: npt.NDArray[np.float64]) -> npt.NDArray[np
 
     def calculate_utilities(
         self, X: npt.NDArray[np.float64], beta: npt.NDArray[np.float64]
-    ) -> Tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+    ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
         """
         Computes the deterministic utility V and the exponentiated utility a.
 
@@ -95,14 +109,160 @@ def calculate_utilities(
         a = np.exp(V_stable)
         return V, a
 
+    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, scipy sparse matrices (shape N x J*J),
+        or explicit index tuples.
+
+        Sparse flattening uses row-major order: column = s * J + t.
+        """
+        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("Unsupported y_dual type")
+
+    def _validate_binary(self, arr: np.ndarray, name: str) -> None:
+        """Ensure array contains only 0/1 values."""
+        if not np.isin(arr, (0, 1)).all():
+            raise ValueError(f"{name} must be binary (contain only 0 or 1).")
+
+    def _validate_data(
+        self,
+        X: npt.NDArray[np.float64],
+        y_single: npt.NDArray[np.int8],
+        y_dual: DualInput,
+        *,
+        dual_indices: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None,
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """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}"
+            )
+
+        self._validate_binary(y_single, "y_single")
+
+        # Normalize dual indices for validation
+        if dual_indices is None:
+            dual_rows, dual_s, dual_t = self._normalize_dual_indices(
+                y_dual, N=N, J=self.J
+            )
+        else:
+            dual_rows, dual_s, dual_t = dual_indices
+
+        # Bounds checks
+        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 or dual_s.min() < 0:
+                raise ValueError("y_dual alternative indices out of bounds.")
+
+        # Enforce upper triangle only and no diagonal
+        if np.any(dual_s == dual_t):
+            raise ValueError("y_dual diagonal must be zero.")
+        if np.any(dual_s > dual_t):
+            # Check that dual choices are in upper triangle (s < t)
+            raise ValueError("y_dual must only have entries in upper triangle (s < t).")
+        # Check that dual choices are in upper triangle (s < t)
+
+        # Binary checks for dense/sparse
+        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}"
+                )
+            self._validate_binary(y_dual, "y_dual")
+        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 = np.nonzero(y_single)[0]
+        counts = np.bincount(single_rows, minlength=N)
+        if len(dual_rows) > 0:
+            counts += np.bincount(dual_rows, minlength=N)
+
+        # Check that each agent has exactly one choice
+        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 dual_rows, dual_s, dual_t
+
+    def _prepare_data(
+        self,
+        y_single: npt.NDArray[np.int8],
+        y_dual: DualInput,
+        *,
+        N: Optional[int] = None,
+        J: Optional[int] = None,
+        dual_indices: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None,
+    ) -> tuple[
+        tuple[np.ndarray, np.ndarray], tuple[np.ndarray, np.ndarray, np.ndarray]
+    ]:
+        """
+        Pre-process data into sparse index format for faster iteration.
+        Supports dense (N,J,J) tensors, scipy sparse matrices of shape (N, J*J),
+        or explicit index tuples (rows, s, t).
+        """
+        single_rows, single_cols = np.nonzero(y_single)
+        if dual_indices is None:
+            dual_rows, dual_s, dual_t = self._normalize_dual_indices(
+                y_dual, N=N if N is not None else y_single.shape[0], J=J or self.J
+            )
+        else:
+            dual_rows, dual_s, dual_t = dual_indices
+        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],
+        y_dual: DualInput,
         init_beta: Optional[npt.NDArray[np.float64]] = None,
         method: str = "L-BFGS-B",
-        options: Optional[Dict[str, Any]] = None,
+        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.
@@ -114,13 +274,19 @@ def fit(
         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}.
+            y_dual (DualInput): Binary tensor (N, J, J), sparse matrix (N, J*J),
+                                or index triplet (rows, s, t) for dual choices.
             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}.
+            bounds (Sequence, optional): Bounds to pass to scipy.optimize.minimize.
+            constraints (Sequence, optional): Constraints to pass to minimize.
+            num_restarts (int): Number of random restarts to perform beyond init_beta.
+            restart_scale (float): Scale of normal noise for restart initialization.
+            rng (np.random.Generator, optional): Random generator for restarts.
 
         Returns:
             self: Returns the instance itself for method chaining.
@@ -131,7 +297,7 @@ def fit(
 
         Example:
             >>> from multe import MultichoiceLogit, simulate_data
-            >>> X, y_single, y_dual, true_beta = simulate_data(N=1000, J=3, K=2)
+            >>> X, y_single, y_dual, _ = simulate_data(N=1000, J=3, K=2)
             >>> model = MultichoiceLogit(num_alternatives=3, num_covariates=2)
             >>> model.fit(X, y_single, y_dual)
             >>> print(model.coef_)  # Estimated coefficients
@@ -140,130 +306,100 @@ def fit(
         if options is None:
             options = {"gtol": 1e-5, "maxiter": 1000}
 
+        # Validate data and prepare indices once
+        dual_indices = self._validate_data(X, y_single, y_dual)
+        single_indices, dual_indices = self._prepare_data(
+            y_single, y_dual, N=X.shape[0], J=self.J, dual_indices=dual_indices
+        )
+
         # Initialize parameters
+        expected_size = (self.J - 1) * self.K
         if init_beta is None:
-            init_beta = np.zeros((self.J - 1) * self.K)
+            init_beta = np.zeros(expected_size)
         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
-        result = minimize(
-            fun=self.neg_log_likelihood,
-            jac=self.gradient,
-            x0=init_beta,
-            args=(X, y_single, y_dual),
-            method=method,
-            options=options,
-        )
+        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_fast,
+                jac=self._gradient_fast,
+                x0=start_beta,
+                args=(X, single_indices, dual_indices),
+                method=method,
+                bounds=bounds,
+                constraints=constraints,
+                options=options,
+            )
+
+        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_i for noise_i in noise)
+
+        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
+
+        assert best_result is not None
 
         # Check convergence
-        if not result.success:
+        if not best_result.success:
             raise RuntimeError(
-                f"Optimization failed to converge: {result.message}\n"
+                f"Optimization failed to converge: {best_result.message}\n"
                 f"Try a different optimization method or adjust tolerance."
             )
 
         # Store results
-        self.coef_ = result.x.reshape(self.J - 1, self.K)
-        self.optimization_result_ = result
+        self.coef_ = best_result.x.reshape(self.J - 1, self.K)
+        self.optimization_result_ = best_result
 
         return self
 
-    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}"
-            )
-
-        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}"
-            )
-
-        # 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 "")
-            )
-
-        # 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)"
-                )
-
-    def neg_log_likelihood(
+    def _neg_log_likelihood_fast(
         self,
         flat_beta: npt.NDArray[np.float64],
         X: npt.NDArray[np.float64],
-        y_single: npt.NDArray[np.int8],
-        y_dual: npt.NDArray[np.int8],
+        single_indices: tuple[np.ndarray, np.ndarray],
+        dual_indices: tuple[np.ndarray, np.ndarray, np.ndarray],
     ) -> float:
         """
-        Computes the negative log-likelihood of the model for minimization.
-        Vectorized version for faster computation.
-
-        Args:
-            flat_beta (np.ndarray): 1D array of parameters optimization is performed on.
-            X (np.ndarray): Covariates 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}.
-
-        Returns:
-            float: The negative log-likelihood value (scalar).
-
-        Raises:
-            ValueError: If data dimensions are incompatible or constraints are violated.
+        Internal optimized NLL using pre-computed indices.
         """
-        self._validate_data(X, y_single, y_dual)
         beta = self.transform_params(flat_beta)
         V, a = self.calculate_utilities(X, beta)
 
         log_lik = 0.0
 
         # 1. Handle Single Choices (MNL)
-        single_choice_mask = y_single.sum(axis=1) > 0
-        if np.any(single_choice_mask):
-            V_sub = V[single_choice_mask]
-            y_sub = y_single[single_choice_mask]
-            term1 = np.sum(y_sub * V_sub)
+        # single_indices = (row_idx, col_idx)
+        if len(single_indices[0]) > 0:
+            row_idx, col_idx = single_indices
+
+            # Extract V for chosen alternatives
+            V_chosen = V[row_idx, col_idx]
+
+            # Extract full V for relevant rows
+            V_sub = V[row_idx]
+
             # Use logsumexp for numerical stability (avoids overflow in exp)
-            term2 = np.sum(logsumexp(V_sub, axis=1))
-            log_lik += (term1 - term2)
+            log_sum = logsumexp(V_sub, axis=1)
 
-        # 2. Handle Dual Choices
-        dual_indices = np.argwhere(y_dual > 0)
+            log_lik += np.sum(V_chosen - log_sum)
 
-        if len(dual_indices) > 0:
+        # 2. Handle Dual Choices
+        if len(dual_indices[0]) > 0:
             # Extract all dual choice indices at once
-            i_idx = dual_indices[:, 0]
-            s_idx = dual_indices[:, 1]
-            t_idx = dual_indices[:, 2]
+            i_idx, s_idx, t_idx = dual_indices
 
             # Get utilities for all dual choices at once
             a_s = a[i_idx, s_idx]  # Shape: (n_dual,)
@@ -273,79 +409,92 @@ def neg_log_likelihood(
             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)
-            p1 = a_s / (a_s + R)
-            p2 = a_t / (a_t + R)
-            p3 = (a_s + a_t) / (a_s + a_t + R)
+            # Vectorized probability computation
+            D1 = a_s + R
+            D2 = a_t + R
+            D3 = a_s + a_t + R
 
-            probs = p1 + p2 - p3
+            probs = (a_s / D1) + (a_t / D2) - ((a_s + a_t) / D3)
 
             # Safety clipping for numerical stability
-            probs = np.maximum(probs, CLIP_THRESHOLD)
+            probs = np.maximum(probs, CLIP_THRESHOLD)  # Clipped for likelihood
 
             # Sum log probabilities
             log_lik += np.sum(np.log(probs))
 
         return -log_lik
 
-    def gradient(
+    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],
-    ) -> npt.NDArray[np.float64]:
+        y_dual: DualInput,
+    ) -> float:
         """
-        Computes the analytical gradient (Jacobian) of the negative log-likelihood.
+        Computes the negative log-likelihood of the model.
+        Wrapper for public API compliance that computes indices on the fly.
 
         Args:
-            flat_beta (np.ndarray): 1D array of parameters.
-            X (np.ndarray): Covariates of shape (N, K).
-            y_single (np.ndarray): Single choice indicators (N, J).
-            y_dual (np.ndarray): Dual choice indicators (N, J, J).
+            flat_beta: Parameter vector ((J-1)*K).
+            X: Covariate matrix (N, K).
+            y_single: Single-choice indicators (N, J).
+            y_dual: Dual-choice indicators (dense, sparse, or index triplet).
 
         Returns:
-            np.ndarray: Flattened gradient vector of shape ((J-1)*K, ).
+            Scalar negative log-likelihood.
+        """
+        dual_indices = self._validate_data(X, y_single, y_dual)
+        single_indices, dual_indices = self._prepare_data(
+            y_single, y_dual, N=X.shape[0], J=self.J, dual_indices=dual_indices
+        )
+        return self._neg_log_likelihood_fast(flat_beta, X, single_indices, dual_indices)
 
-        Raises:
-            ValueError: If data dimensions are incompatible or constraints are violated.
+    def _gradient_fast(
+        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.
         """
-        self._validate_data(X, y_single, y_dual)
         beta = self.transform_params(flat_beta)
         V, a = self.calculate_utilities(X, beta)
 
+        # Gradient buffer for all parameters (J, K); flattened later to ((J-1)*K,)
         # Gradient buffer for all parameters (J, K)
         grad = np.zeros((self.J, self.K))
 
         # 1. Single Choice Gradient (Standard MNL)
-        single_mask = y_single.sum(axis=1) > 0
-        if np.any(single_mask):
-            X_sub = X[single_mask]
-            a_sub = a[single_mask]
-            y_sub = y_single[single_mask]
-            probs = a_sub / np.sum(a_sub, axis=1, keepdims=True)
-            error = y_sub - probs
-            grad += error.T @ X_sub
+        if len(single_indices[0]) > 0:
+            row_idx, col_idx = single_indices
 
-        # 2. Dual Choice Gradient
-        dual_indices = np.argwhere(y_dual > 0)
+            X_sub = X[row_idx]
+            a_sub = a[row_idx]
 
-        if len(dual_indices) > 0:
-            # Extract indices
-            i_idx = dual_indices[:, 0]
-            s_idx = dual_indices[:, 1]
-            t_idx = dual_indices[:, 2]
+            # Probabilities P(y=j) = exp(V_j) / sum(exp(V_k))
+            probs = a_sub / np.sum(a_sub, axis=1, keepdims=True)
 
+            # Add positive term for chosen alternatives (y=1)
+            np.add.at(grad, col_idx, X_sub)
+
+            # Subtract prob term for all alternatives
+            grad -= probs.T @ X_sub
+        # 2. Dual Choice Gradient
+        if len(dual_indices[0]) > 0:
+            i_idx, s_idx, t_idx = dual_indices
             n_dual = len(i_idx)
 
-            # Get covariates and utilities for all 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,)
+            # Covariates/utilities for dual choices
+            X_i = X[i_idx]  # (n_dual, K)
+            a_s = a[i_idx, s_idx]  # (n_dual,)
+            a_t = a[i_idx, t_idx]  # (n_dual,)
 
-            # Compute R and denominators
-            a_sum = np.sum(a[i_idx], axis=1)  # Shape: (n_dual,)
+            a_sum = np.sum(a[i_idx], axis=1)  # (n_dual,)
             R = a_sum - a_s - a_t
 
             D1 = a_s + R
@@ -353,122 +502,244 @@ def gradient(
             D3 = a_s + a_t + R
 
             # Probability (unclipped for gradient computation)
-            P_raw = (a_s/D1) + (a_t/D2) - ((a_s+a_t)/D3)
-
-            # Clip for numerical stability (use module constant)
-            P = np.maximum(P_raw, CLIP_THRESHOLD)  # Clipped for likelihood
+            P_raw = (a_s / D1) + (a_t / D2) - ((a_s + a_t) / D3)
 
-            # Mask for where clipping occurred (gradient should be 0)
             clipped_mask = P_raw >= CLIP_THRESHOLD
+            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))  # (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,)
+
+            w_s = inv_P * dP_dVs  # (n_dual,)
+            w_t = inv_P * dP_dVt  # (n_dual,)
+            w_r = inv_P * common_r  # (n_dual,)
+            # Complexity: O(n_dual) weight computations; still cheaper than looping over J.
+
+            # 'r' alternatives (neither s nor t): grad += sum_i (w_r_i * a_ij) * X_i
+            a_i = a[i_idx]  # (n_dual, J)
+            M = w_r[:, np.newaxis] * a_i
+            rows = np.arange(n_dual)
+            M[rows, s_idx] = 0.0
+            M[rows, t_idx] = 0.0
+
+            # This is an O(n_dual * J * K) dense step; faster than looping for typical J.
+            grad += M.T @ X_i
+
+            # Add contributions from s and t (O(n_dual))
+            grad_contrib_s = w_s[:, np.newaxis] * X_i  # (n_dual, K)
+            grad_contrib_t = w_t[:, np.newaxis] * X_i  # (n_dual, K)
 
-            # Precompute derivatives
-            dP_dVs = a_s * R * (1/(D1**2) - 1/(D3**2))  # Shape: (n_dual,)
-            dP_dVt = a_t * R * (1/(D2**2) - 1/(D3**2))  # Shape: (n_dual,)
-            common_r = (a_s+a_t)/(D3**2) - a_s/(D1**2) - a_t/(D2**2)  # Shape: (n_dual,)
-
-            # Compute gradient contributions for s and t
-            # For alternative s (gradient = 0 if probability was clipped)
-            grad_weight_s = np.where(clipped_mask, (1/P_raw) * dP_dVs, 0.0)  # Shape: (n_dual,)
-            grad_contrib_s = grad_weight_s[:, np.newaxis] * X_i  # Shape: (n_dual, K)
-
-            # For alternative t (gradient = 0 if probability was clipped)
-            grad_weight_t = np.where(clipped_mask, (1/P_raw) * dP_dVt, 0.0)  # Shape: (n_dual,)
-            grad_contrib_t = grad_weight_t[:, np.newaxis] * X_i  # Shape: (n_dual, K)
-
-            # Complexity: O(n_dual) - faster than O(J * n_dual) with loops
             np.add.at(grad, s_idx, grad_contrib_s)
             np.add.at(grad, t_idx, grad_contrib_t)
-
-            # For alternatives that are neither s nor t (the 'r' alternatives)
-            # Memory: O(n_dual * J * K) tensor, but faster than looping
-            a_i = a[i_idx]  # Shape: (n_dual, J)
-            grad_weight_r = np.where(clipped_mask, (1/P_raw) * common_r, 0.0)  # Shape: (n_dual,)
-
-            # Create mask: True where alternative is neither s nor t for each observation
-            # Shape: (n_dual, J) - True if alternative j is 'r' for observation i
-            is_r = np.ones((n_dual, self.J), dtype=bool)
-            is_r[np.arange(n_dual), s_idx] = False
-            is_r[np.arange(n_dual), t_idx] = False
-
-            # Gradient contributions for 'r' alternatives (neither s nor t)
-            # 
-            # For each dual choice observation i and each 'r' alternative j:
-            #   grad_contrib[j] += (1/P) * common_r * a[i,j] * X[i]
-            #
-            # We vectorize this as a 3D tensor multiplication:
-            #   grad_weight_r: (n_dual,)     -> (n_dual, 1, 1)  scalar weight per obs
-            #   a_i:           (n_dual, J)   -> (n_dual, J, 1)  utility weight per alt
-            #   X_i:           (n_dual, K)   -> (n_dual, 1, K)  covariates per obs
-            #   result:        (n_dual, J, K)                   gradient contrib per obs/alt
-            #
-            # Memory: O(n_dual * J * K) — for large datasets, consider chunked processing
-            grad_contrib_r_all = (grad_weight_r[:, np.newaxis, np.newaxis] * # Shape: (n_dual, 1, 1)
-                                  a_i[:, :, np.newaxis] * # Shape: (n_dual, J, 1)
-                                  X_i[:, np.newaxis, :])  # Shape: (n_dual, 1, K)
-
-            # Zero out contributions from s and t alternatives
-            grad_contrib_r_all[~is_r] = 0
-
-            # Sum over dual choices and add to gradient for each alternative
-            grad += grad_contrib_r_all.sum(axis=0)  # Sum over n_dual, result: (J, K)
-
         # Return negative gradient for minimization, remove fixed class 0
         return -grad[1:].flatten()
 
+    def gradient(
+        self,
+        flat_beta: npt.NDArray[np.float64],
+        X: npt.NDArray[np.float64],
+        y_single: npt.NDArray[np.int8],
+        y_dual: DualInput,
+    ) -> npt.NDArray[np.float64]:
+        """
+        Computes the analytical gradient (Jacobian).
+        Wrapper for public API compliance that computes indices on the fly.
+
+        Args:
+            flat_beta: Parameter vector ((J-1)*K).
+            X: Covariate matrix (N, K).
+            y_single: Single-choice indicators (N, J).
+            y_dual: Dual-choice indicators (dense, sparse, or index triplet).
+
+        Returns:
+            Gradient vector ((J-1)*K,).
+        """
+        dual_indices = self._validate_data(X, y_single, y_dual)
+        single_indices, dual_indices = self._prepare_data(
+            y_single, y_dual, N=X.shape[0], J=self.J, dual_indices=dual_indices
+        )
+        return self._gradient_fast(flat_beta, X, single_indices, dual_indices)
+
     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,
+        *,
+        epsilon: float | None = 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.
 
-        Uses central differences for O(ε²) accuracy vs O(ε) for forward differences.
-
         Args:
-            flat_beta (np.ndarray): Optimal parameters (flattened).
-            X, y_single, y_dual: Data required for gradient calculation.
+            flat_beta: Parameter vector ((J-1)*K).
+            X: Covariate matrix (N, K).
+            y_single: Single-choice indicators (N, J).
+            y_dual: Dual-choice indicators (dense, sparse, or index triplet).
+            epsilon: Optional finite-difference step; defaults to HESSIAN_EPSILON.
 
         Returns:
-            np.ndarray: Standard errors for the parameters.
+            1D array of standard errors ((J-1)*K,).
 
         Raises:
-            ValueError: If data dimensions are incompatible or constraints are violated.
-            RuntimeWarning: If Hessian is singular or ill-conditioned.
+            ValueError: If data validation fails.
         """
-        self._validate_data(X, y_single, y_dual)
+        dual_indices = self._validate_data(X, y_single, y_dual)
+        single_indices, dual_indices = self._prepare_data(
+            y_single, y_dual, N=X.shape[0], J=self.J, dual_indices=dual_indices
+        )
+
         n_params = len(flat_beta)
         hessian = np.zeros((n_params, n_params))
+        step = HESSIAN_EPSILON if epsilon is None else float(epsilon)
 
         # Central finite differences for Hessian (more accurate than forward)
         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(beta_plus, X, y_single, y_dual)
-            grad_minus = self.gradient(beta_minus, X, y_single, y_dual)
+            grad_plus = self._gradient_fast(beta_plus, X, single_indices, dual_indices)
+            grad_minus = self._gradient_fast(
+                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
-        try:
-            cov_matrix = np.linalg.inv(hessian)
-            std_errs = np.sqrt(np.diag(cov_matrix))
-        except np.linalg.LinAlgError as e:
+        # Use pinv for stability with near-singular Hessians
+        cov_matrix = np.linalg.pinv(hessian)
+
+        # Check if diagonal elements are positive (valid variance)
+        diag_cov = np.diag(cov_matrix)
+
+        # If any variance is negative (numerical noise with singular hessian), warn
+        if np.any(diag_cov < 0):
             import warnings
 
             warnings.warn(
-                "Hessian is singular or ill-conditioned. Standard errors may be unreliable.",
+                "Hessian inverse has negative diagonal elements. Standard errors may be unreliable.",
                 RuntimeWarning,
                 stacklevel=2,
             )
-            std_errs = np.full(n_params, np.nan)
 
-        return std_errs
+        # 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 for new data.
+
+        Args:
+            X: Covariate matrix (N, K)
+            flat_beta: Optional parameter vector ((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 (upper triangle populated)
+
+        Raises:
+            ValueError: If the model is unfitted and no parameters are provided.
+        """
+        if flat_beta is None:
+            if self.coef_ is None:
+                raise ValueError("Model is not fitted. Provide flat_beta or call fit.")
+            flat_beta = self.coef_.flatten()
+
+        beta = self.transform_params(flat_beta)
+        V, a = self.calculate_utilities(X, beta)
+
+        single_probs = a / np.sum(a, axis=1, keepdims=True)
+
+        # Dual probabilities: compute for each pair s<t
+        N, J = X.shape[0], self.J
+        dual_probs = np.zeros((N, J, J))
+        a_sum = np.sum(a, axis=1)
+
+        for s in range(J - 1):
+            a_s = a[:, s]
+            for t in range(s + 1, J):
+                a_t = a[:, t]
+                R = a_sum - a_s - a_t
+                D1 = a_s + R
+                D2 = a_t + R
+                D3 = a_s + a_t + R
+                probs = (a_s / D1) + (a_t / D2) - ((a_s + a_t) / D3)
+                dual_probs[:, s, t] = np.maximum(probs, CLIP_THRESHOLD)
+
+        return single_probs, dual_probs
+
+    def log_likelihood_contributions(
+        self,
+        flat_beta: npt.NDArray[np.float64],
+        X: npt.NDArray[np.float64],
+        y_single: npt.NDArray[np.int8],
+        y_dual: DualInput,
+    ) -> npt.NDArray[np.float64]:
+        """
+        Return per-observation log-likelihood contributions.
+
+        Args:
+            flat_beta: Parameter vector ((J-1)*K).
+            X: Covariate matrix (N, K).
+            y_single: Single-choice indicators (N, J).
+            y_dual: Dual-choice indicators (dense, sparse, or index triplet).
+
+        Returns:
+            Vector of per-observation log-likelihood contributions (N,).
+        """
+        dual_indices = self._validate_data(X, y_single, y_dual)
+        single_indices, dual_indices = self._prepare_data(
+            y_single, y_dual, N=X.shape[0], J=self.J, dual_indices=dual_indices
+        )
+        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,
+        flat_beta: npt.NDArray[np.float64],
+        X: npt.NDArray[np.float64],
+        y_single: npt.NDArray[np.int8],
+        y_dual: DualInput,
+    ) -> float:
+        """Convenience wrapper returning the sum of log-likelihood contributions."""
+        return float(
+            np.sum(self.log_likelihood_contributions(flat_beta, X, y_single, y_dual))
+        )
diff --git a/multe/simulate.py b/multe/simulate.py
index 19211b0..9f1ee79 100644
--- a/multe/simulate.py
+++ b/multe/simulate.py
@@ -6,6 +6,8 @@
 Fully vectorized implementation for fast simulation.
 """
 
+from __future__ import annotations
+
 from typing import Optional, Tuple
 import numpy as np
 import numpy.typing as npt
@@ -17,7 +19,9 @@ def simulate_data(
     K: int,
     true_beta: Optional[npt.NDArray[np.float64]] = None,
     mix_ratio: float = 0.5,
-    seed: int = 42,
+    seed: int | None = 42,
+    rng: np.random.Generator | None = None,
+    dtype: npt.DTypeLike = np.float64,
 ) -> Tuple[
     npt.NDArray[np.float64],
     npt.NDArray[np.int8],
@@ -42,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)
@@ -53,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}")
@@ -64,22 +71,22 @@ def simulate_data(
         raise ValueError(f"mix_ratio must be in [0, 1], got {mix_ratio}")
     if true_beta is not None and true_beta.shape != (J - 1, K):
         raise ValueError(
-            f"true_beta must have shape ({J-1}, {K}), got {true_beta.shape}"
+            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
@@ -97,8 +104,12 @@ def simulate_data(
     top_k_values = np.take_along_axis(U, top_k_indices, axis=1)
     sorted_within_top = np.argsort(top_k_values, axis=1)
 
-    best_idx = np.take_along_axis(top_k_indices, sorted_within_top[:, -1:], axis=1).flatten()
-    second_best_idx = np.take_along_axis(top_k_indices, sorted_within_top[:, -2:-1], axis=1).flatten()
+    best_idx = np.take_along_axis(
+        top_k_indices, sorted_within_top[:, -1:], axis=1
+    ).flatten()
+    second_best_idx = np.take_along_axis(
+        top_k_indices, sorted_within_top[:, -2:-1], axis=1
+    ).flatten()
 
     # Randomly assign choice mode (single vs dual) based on mix_ratio
     mode_choice = rng.binomial(1, mix_ratio, N).astype(bool)
@@ -121,7 +132,10 @@ def simulate_data(
 
         # Ensure s < t (swap if needed)
         swap_mask = s_indices > t_indices
-        s_indices[swap_mask], t_indices[swap_mask] = t_indices[swap_mask], s_indices[swap_mask]
+        s_indices[swap_mask], t_indices[swap_mask] = (
+            t_indices[swap_mask],
+            s_indices[swap_mask],
+        )
 
         y_dual[row_indices_dual, s_indices, t_indices] = 1
 
diff --git a/tests/test_model.py b/tests/test_model.py
index c00b1dc..e1376a8 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -88,6 +88,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
@@ -134,6 +149,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."""
@@ -211,6 +241,21 @@ def test_numerical_stability_large_utilities(self):
         grad = model.gradient(flat_beta, X, y_single, y_dual)
         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()
+        dense_nll = model.neg_log_likelihood(flat_beta, X, y_single, y_dual)
+
+        rows, s_idx, t_idx = np.nonzero(y_dual)
+        tuple_input = (rows, s_idx, t_idx)
+        tuple_nll = model.neg_log_likelihood(flat_beta, X, y_single, tuple_input)
+
+        assert np.isclose(dense_nll, tuple_nll)
+
 
 class TestGradient:
     """Test gradient computation."""
@@ -286,6 +331,20 @@ def test_gradient_with_clipped_probabilities(self):
         # 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()
+        dense_grad = model.gradient(flat_beta, X, y_single, y_dual)
+
+        dual_rows, dual_s, dual_t = np.nonzero(y_dual)
+        tuple_grad = model.gradient(flat_beta, X, y_single, (dual_rows, dual_s, dual_t))
+
+        assert np.allclose(dense_grad, tuple_grad)
+
 
 class TestComputeStandardErrors:
     """Test standard error computation."""
@@ -468,3 +527,53 @@ 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(
+            flat_beta, X, y_single, y_dual
+        )
+        total = np.sum(contributions)
+        direct = -model.neg_log_likelihood(flat_beta, X, y_single, y_dual)
+
+        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(flat_beta, X, y_single, y_dual)
+        std_errs_smaller = model.compute_standard_errors(
+            flat_beta, X, y_single, y_dual, 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 b0a8d3e..50ec735 100644
--- a/tests/test_simulate.py
+++ b/tests/test_simulate.py
@@ -134,6 +134,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
@@ -148,6 +162,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

From 9c0e11c7ec9b93b4956025a99b427d7a4f14b1c7 Mon Sep 17 00:00:00 2001
From: Thomas Monk <246525+tmonk@users.noreply.github.com>
Date: Mon, 24 Nov 2025 23:13:55 +0000
Subject: [PATCH 2/5] docs: add author information and citation details to
 README.

---
 README.md | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 94f6753..03eab26 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.1198/07350019919290156). 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:
@@ -111,4 +126,4 @@ Returns: `X, y_single, y_dual, true_beta`
 
 ## Acknowledgements
 
-Many thanks to [Alan Manning](https://www.alan-manning.com/) for his guidance and support with this project.
\ No newline at end of file
+Many thanks to [Alan Manning](https://www.alan-manning.com/) for his guidance and support with this project.

From c78c6d17fcdaf06569f9a62b36c3ff106cc3ed9f Mon Sep 17 00:00:00 2001
From: Thomas Monk <246525+tmonk@users.noreply.github.com>
Date: Mon, 24 Nov 2025 23:39:32 +0000
Subject: [PATCH 3/5] feat: Refactor model internals to use private methods for
 utility and parameter transformation, introduce a dedicated data validation
 and index preparation method, and update examples to reflect these changes
 and the `compute_standard_errors` signature.

---
 README.md                      |   5 +-
 examples/basic_example.py      |   2 +-
 examples/benchmark.py          |   7 +-
 examples/csv_example.py        |   2 +-
 examples/simple_fit_example.py |   2 +-
 multe/model.py                 | 460 +++++++++++++++++----------------
 tests/test_model.py            |  52 ++--
 7 files changed, 277 insertions(+), 253 deletions(-)

diff --git a/README.md b/README.md
index 03eab26..3ef0d50 100644
--- a/README.md
+++ b/README.md
@@ -116,8 +116,9 @@ Model class with methods:
   - 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 f98e46a..f9ab7d4 100644
--- a/examples/basic_example.py
+++ b/examples/basic_example.py
@@ -43,7 +43,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 3578e2e..83519e2 100644
--- a/examples/benchmark.py
+++ b/examples/benchmark.py
@@ -28,14 +28,17 @@ 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
diff --git a/examples/csv_example.py b/examples/csv_example.py
index a1a0c6c..feeb63b 100644
--- a/examples/csv_example.py
+++ b/examples/csv_example.py
@@ -179,7 +179,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)
 
diff --git a/examples/simple_fit_example.py b/examples/simple_fit_example.py
index feb4504..166305a 100644
--- a/examples/simple_fit_example.py
+++ b/examples/simple_fit_example.py
@@ -42,7 +42,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 86f1e76..4619493 100644
--- a/multe/model.py
+++ b/multe/model.py
@@ -2,11 +2,13 @@
 Multichoice Logit Model
 
 Vectorized implementation for fast and accurate MLE estimation.
+Supports single and dual (pairwise) discrete choices.
 """
 
 from __future__ import annotations
 
 import typing
+import warnings
 from typing import Any, Optional, Sequence
 
 import numpy as np
@@ -19,6 +21,7 @@
 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]
@@ -31,11 +34,19 @@ 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:
@@ -43,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.
@@ -61,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.
@@ -88,20 +101,28 @@ def transform_params(
         beta_fixed = np.zeros((1, self.K))
         return np.vstack([beta_fixed, beta_free])
 
-    def calculate_utilities(
+    # Public alias for compatibility
+    def transform_params(
+        self, flat_beta: npt.NDArray[np.float64]
+    ) -> npt.NDArray[np.float64]:
+        """Public wrapper for parameter reshaping (kept for compatibility)."""
+        return self._transform_params(flat_beta)
+
+    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
@@ -109,15 +130,37 @@ def calculate_utilities(
         a = np.exp(V_stable)
         return V, a
 
+    # Public alias for compatibility
+    def calculate_utilities(
+        self, X: npt.NDArray[np.float64], beta: npt.NDArray[np.float64]
+    ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+        """Public wrapper for utility calculation (kept for compatibility)."""
+        return self._calculate_utilities(X, beta)
+
     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, scipy sparse matrices (shape N x J*J),
-        or explicit index tuples.
+
+        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:
@@ -143,22 +186,31 @@ def _normalize_dual_indices(
             dual_rows, dual_s, dual_t = np.nonzero(y_dual)
             return dual_rows, dual_s, dual_t
 
-        raise TypeError("Unsupported y_dual type")
-
-    def _validate_binary(self, arr: np.ndarray, name: str) -> None:
-        """Ensure array contains only 0/1 values."""
-        if not np.isin(arr, (0, 1)).all():
-            raise ValueError(f"{name} must be binary (contain only 0 or 1).")
+        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: DualInput,
-        *,
-        dual_indices: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None,
-    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-        """Validate input data dimensions and constraints."""
+    ) -> tuple[
+        tuple[np.ndarray, np.ndarray], tuple[np.ndarray, np.ndarray, np.ndarray]
+    ]:
+        """
+        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:
@@ -169,38 +221,38 @@ def _validate_data(
                 f"y_single must have shape ({N}, {self.J}), got {y_single.shape}"
             )
 
-        self._validate_binary(y_single, "y_single")
+        # 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 for validation
-        if dual_indices is None:
-            dual_rows, dual_s, dual_t = self._normalize_dual_indices(
-                y_dual, N=N, J=self.J
-            )
-        else:
-            dual_rows, dual_s, dual_t = dual_indices
+        # Normalize dual indices
+        dual_rows, dual_s, dual_t = self._normalize_dual_indices(y_dual, N=N, J=self.J)
 
-        # Bounds checks
+        # 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 or dual_s.min() < 0:
+            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 only and no diagonal
-        if np.any(dual_s == dual_t):
-            raise ValueError("y_dual diagonal must be zero.")
-        if np.any(dual_s > dual_t):
-            # Check that dual choices are in upper triangle (s < t)
-            raise ValueError("y_dual must only have entries in upper triangle (s < t).")
-        # Check that dual choices are in upper triangle (s < t)
+            # 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
+        # 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}"
                 )
-            self._validate_binary(y_dual, "y_dual")
+            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(
@@ -210,12 +262,11 @@ def _validate_data(
                 raise ValueError("Sparse y_dual must be binary.")
 
         # Check that each agent has exactly one choice
-        single_rows = np.nonzero(y_single)[0]
+        single_rows, single_cols = np.nonzero(y_single)
         counts = np.bincount(single_rows, minlength=N)
         if len(dual_rows) > 0:
             counts += np.bincount(dual_rows, minlength=N)
 
-        # Check that each agent has exactly one choice
         if not np.all(counts == 1):
             invalid = np.where(counts != 1)[0]
             raise ValueError(
@@ -223,31 +274,7 @@ def _validate_data(
                 f"Agents with invalid choices: {invalid[:10]}"
                 + ("..." if len(invalid) > 10 else "")
             )
-        return dual_rows, dual_s, dual_t
 
-    def _prepare_data(
-        self,
-        y_single: npt.NDArray[np.int8],
-        y_dual: DualInput,
-        *,
-        N: Optional[int] = None,
-        J: Optional[int] = None,
-        dual_indices: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None,
-    ) -> tuple[
-        tuple[np.ndarray, np.ndarray], tuple[np.ndarray, np.ndarray, np.ndarray]
-    ]:
-        """
-        Pre-process data into sparse index format for faster iteration.
-        Supports dense (N,J,J) tensors, scipy sparse matrices of shape (N, J*J),
-        or explicit index tuples (rows, s, t).
-        """
-        single_rows, single_cols = np.nonzero(y_single)
-        if dual_indices is None:
-            dual_rows, dual_s, dual_t = self._normalize_dual_indices(
-                y_dual, N=N if N is not None else y_single.shape[0], J=J or self.J
-            )
-        else:
-            dual_rows, dual_s, dual_t = dual_indices
         return (single_rows, single_cols), (dual_rows, dual_s, dual_t)
 
     def fit(
@@ -263,73 +290,58 @@ def fit(
         num_restarts: int = 0,
         restart_scale: float = 0.5,
         rng: Optional[np.random.Generator] = None,
-    ) -> "MultichoiceLogit":
+    ) -> 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 (DualInput): Binary tensor (N, J, J), sparse matrix (N, J*J),
-                                or index triplet (rows, s, t) for dual choices.
-            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}.
-            bounds (Sequence, optional): Bounds to pass to scipy.optimize.minimize.
-            constraints (Sequence, optional): Constraints to pass to minimize.
-            num_restarts (int): Number of random restarts to perform beyond init_beta.
-            restart_scale (float): Scale of normal noise for restart initialization.
-            rng (np.random.Generator, optional): Random generator for restarts.
+            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.
-
-        Example:
-            >>> from multe import MultichoiceLogit, simulate_data
-            >>> X, y_single, y_dual, _ = simulate_data(N=1000, J=3, K=2)
-            >>> model = MultichoiceLogit(num_alternatives=3, num_covariates=2)
-            >>> model.fit(X, y_single, y_dual)
-            >>> print(model.coef_)  # Estimated coefficients
         """
-        # Set default options
         if options is None:
             options = {"gtol": 1e-5, "maxiter": 1000}
 
         # Validate data and prepare indices once
-        dual_indices = self._validate_data(X, y_single, y_dual)
-        single_indices, dual_indices = self._prepare_data(
-            y_single, y_dual, N=X.shape[0], J=self.J, dual_indices=dual_indices
-        )
+        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(expected_size)
-        else:
-            # Validate initial parameters
-            if init_beta.size != expected_size:
-                raise ValueError(
-                    f"init_beta must have size {expected_size}, got {init_beta.size}"
-                )
+        elif init_beta.size != expected_size:
+            raise ValueError(
+                f"init_beta must have size {expected_size}, got {init_beta.size}"
+            )
 
         rng = rng or np.random.default_rng()
 
         def run_optimization(start_beta: npt.NDArray[np.float64]) -> OptimizeResult:
-            # Run optimization for a given starting point
+            """Run optimization for a given starting point."""
             return minimize(
-                fun=self._neg_log_likelihood_fast,
-                jac=self._gradient_fast,
+                fun=self._neg_log_likelihood,
+                jac=self._gradient,
                 x0=start_beta,
                 args=(X, single_indices, dual_indices),
                 method=method,
@@ -338,11 +350,12 @@ def run_optimization(start_beta: npt.NDArray[np.float64]) -> OptimizeResult:
                 options=options,
             )
 
+        # 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_i for noise_i in noise)
+            start_points.extend(init_beta + noise_row for noise_row in noise)
 
         for start in start_points:
             # Run optimization
@@ -352,7 +365,6 @@ def run_optimization(start_beta: npt.NDArray[np.float64]) -> OptimizeResult:
 
         assert best_result is not None
 
-        # Check convergence
         if not best_result.success:
             raise RuntimeError(
                 f"Optimization failed to converge: {best_result.message}\n"
@@ -365,18 +377,16 @@ def run_optimization(start_beta: npt.NDArray[np.float64]) -> OptimizeResult:
 
         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
 
@@ -411,7 +421,6 @@ def _neg_log_likelihood_fast(
 
             # 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)
-            # Vectorized probability computation
             D1 = a_s + R
             D2 = a_t + R
             D3 = a_s + a_t + R
@@ -419,7 +428,7 @@ def _neg_log_likelihood_fast(
             probs = (a_s / D1) + (a_t / D2) - ((a_s + a_t) / D3)
 
             # Safety clipping for numerical stability
-            probs = np.maximum(probs, CLIP_THRESHOLD)  # Clipped for likelihood
+            probs = np.maximum(probs, CLIP_THRESHOLD)
 
             # Sum log probabilities
             log_lik += np.sum(np.log(probs))
@@ -434,39 +443,23 @@ def neg_log_likelihood(
         y_dual: DualInput,
     ) -> float:
         """
-        Computes the negative log-likelihood of the model.
-        Wrapper for public API compliance that computes indices on the fly.
-
-        Args:
-            flat_beta: Parameter vector ((J-1)*K).
-            X: Covariate matrix (N, K).
-            y_single: Single-choice indicators (N, J).
-            y_dual: Dual-choice indicators (dense, sparse, or index triplet).
-
-        Returns:
-            Scalar negative log-likelihood.
+        Public wrapper: compute negative log-likelihood with on-the-fly indices.
         """
-        dual_indices = self._validate_data(X, y_single, y_dual)
-        single_indices, dual_indices = self._prepare_data(
-            y_single, y_dual, N=X.shape[0], J=self.J, dual_indices=dual_indices
-        )
-        return self._neg_log_likelihood_fast(flat_beta, X, single_indices, dual_indices)
+        single_indices, dual_indices = self._validate_data(X, y_single, y_dual)
+        return self._neg_log_likelihood(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); flattened later to ((J-1)*K,)
-        # Gradient buffer for all parameters (J, K)
         grad = np.zeros((self.J, self.K))
 
         # 1. Single Choice Gradient (Standard MNL)
@@ -479,22 +472,26 @@ def _gradient_fast(
             # Probabilities P(y=j) = exp(V_j) / sum(exp(V_k))
             probs = a_sub / np.sum(a_sub, axis=1, keepdims=True)
 
+            # Gradient = X * (y - p)
+            # y is one-hot, so for chosen col y=1, else 0
             # Add positive term for chosen alternatives (y=1)
             np.add.at(grad, col_idx, X_sub)
 
             # Subtract prob term for all alternatives
             grad -= probs.T @ X_sub
+
         # 2. Dual Choice Gradient
         if len(dual_indices[0]) > 0:
             i_idx, s_idx, t_idx = dual_indices
             n_dual = len(i_idx)
 
-            # Covariates/utilities for dual choices
-            X_i = X[i_idx]  # (n_dual, K)
-            a_s = a[i_idx, s_idx]  # (n_dual,)
-            a_t = a[i_idx, t_idx]  # (n_dual,)
+            # 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,)
 
-            a_sum = np.sum(a[i_idx], axis=1)  # (n_dual,)
+            # Compute R and denominators
+            a_sum = np.sum(a[i_idx], axis=1)  # Shape: (n_dual,)
             R = a_sum - a_s - a_t
 
             D1 = a_s + R
@@ -504,103 +501,94 @@ def _gradient_fast(
             # Probability (unclipped for gradient computation)
             P_raw = (a_s / D1) + (a_t / D2) - ((a_s + a_t) / D3)
 
+            # Clip mask for safe division
             clipped_mask = P_raw >= CLIP_THRESHOLD
             inv_P = np.zeros_like(P_raw)
             inv_P[clipped_mask] = 1.0 / P_raw[clipped_mask]
 
-            # Derivatives
+            # 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  # (n_dual,)
             w_t = inv_P * dP_dVt  # (n_dual,)
+
+            # 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.
 
-            # 'r' alternatives (neither s nor t): 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
+            # 1. Compute M = w_r[:, None] * a[i_idx]   (Shape: n_dual, J)
             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
             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
             # This is an O(n_dual * J * K) dense step; faster than looping for typical J.
             grad += M.T @ X_i
 
-            # Add contributions from s and t (O(n_dual))
+            # 4. Add contributions from s and t (O(n_dual))
             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 for minimization, remove fixed class 0
         return -grad[1:].flatten()
 
-    def gradient(
-        self,
-        flat_beta: npt.NDArray[np.float64],
-        X: npt.NDArray[np.float64],
-        y_single: npt.NDArray[np.int8],
-        y_dual: DualInput,
-    ) -> npt.NDArray[np.float64]:
-        """
-        Computes the analytical gradient (Jacobian).
-        Wrapper for public API compliance that computes indices on the fly.
-
-        Args:
-            flat_beta: Parameter vector ((J-1)*K).
-            X: Covariate matrix (N, K).
-            y_single: Single-choice indicators (N, J).
-            y_dual: Dual-choice indicators (dense, sparse, or index triplet).
-
-        Returns:
-            Gradient vector ((J-1)*K,).
-        """
-        dual_indices = self._validate_data(X, y_single, y_dual)
-        single_indices, dual_indices = self._prepare_data(
-            y_single, y_dual, N=X.shape[0], J=self.J, dual_indices=dual_indices
-        )
-        return self._gradient_fast(flat_beta, X, single_indices, dual_indices)
-
     def compute_standard_errors(
         self,
-        flat_beta: npt.NDArray[np.float64],
         X: npt.NDArray[np.float64],
         y_single: npt.NDArray[np.int8],
         y_dual: DualInput,
-        *,
-        epsilon: float | None = None,
+        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:
-            flat_beta: Parameter vector ((J-1)*K).
             X: Covariate matrix (N, K).
             y_single: Single-choice indicators (N, J).
             y_dual: Dual-choice indicators (dense, sparse, or index triplet).
-            epsilon: Optional finite-difference step; defaults to HESSIAN_EPSILON.
+            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 ((J-1)*K,).
+            1D array of standard errors of size (J-1)*K.
 
         Raises:
-            ValueError: If data validation fails.
+            ValueError: If model is unfitted and no parameters provided.
+
+        Warns:
+            RuntimeWarning: If Hessian inverse has negative diagonal elements.
         """
-        dual_indices = self._validate_data(X, y_single, y_dual)
-        single_indices, dual_indices = self._prepare_data(
-            y_single, y_dual, N=X.shape[0], J=self.J, dual_indices=dual_indices
-        )
+        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 = HESSIAN_EPSILON if epsilon is None else float(epsilon)
+        step = epsilon if epsilon is not None else HESSIAN_EPSILON
 
-        # Central finite differences for Hessian (more accurate than forward)
+        # Central finite differences for Hessian
         for j in range(n_params):
             beta_plus = flat_beta.copy()
             beta_plus[j] += step
@@ -608,10 +596,8 @@ def compute_standard_errors(
             beta_minus = flat_beta.copy()
             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 * step)
@@ -625,17 +611,15 @@ 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(
@@ -644,30 +628,34 @@ def predict_proba(
         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 for new data.
+        Compute predicted probabilities for single and dual choices.
 
         Args:
-            X: Covariate matrix (N, K)
-            flat_beta: Optional parameter vector ((J-1)*K). Uses fitted coef_ if None.
+            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 (upper triangle populated)
+            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 the model is unfitted and no parameters are provided.
+            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.")
+                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)
+        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 probabilities: compute for each pair s<t
+        # Dual choice probabilities: compute for each pair s<t
         N, J = X.shape[0], self.J
         dual_probs = np.zeros((N, J, J))
         a_sum = np.sum(a, axis=1)
@@ -687,29 +675,36 @@ def predict_proba(
 
     def log_likelihood_contributions(
         self,
-        flat_beta: npt.NDArray[np.float64],
         X: npt.NDArray[np.float64],
         y_single: npt.NDArray[np.int8],
         y_dual: DualInput,
+        flat_beta: Optional[npt.NDArray[np.float64]] = None,
     ) -> npt.NDArray[np.float64]:
         """
-        Return per-observation log-likelihood contributions.
+        Compute per-observation log-likelihood contributions.
+
+        Useful for diagnostics, cross-validation, or computing information criteria.
 
         Args:
-            flat_beta: Parameter vector ((J-1)*K).
             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,).
         """
-        dual_indices = self._validate_data(X, y_single, y_dual)
-        single_indices, dual_indices = self._prepare_data(
-            y_single, y_dual, N=X.shape[0], J=self.J, dual_indices=dual_indices
-        )
-        beta = self.transform_params(flat_beta)
-        V, a = self.calculate_utilities(X, beta)
+        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:
@@ -734,12 +729,23 @@ def log_likelihood_contributions(
 
     def log_likelihood(
         self,
-        flat_beta: npt.NDArray[np.float64],
         X: npt.NDArray[np.float64],
         y_single: npt.NDArray[np.int8],
         y_dual: DualInput,
+        flat_beta: Optional[npt.NDArray[np.float64]] = None,
     ) -> float:
-        """Convenience wrapper returning the sum of log-likelihood contributions."""
+        """
+        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(flat_beta, X, y_single, y_dual))
+            np.sum(self.log_likelihood_contributions(X, y_single, y_dual, flat_beta))
         )
diff --git a/tests/test_model.py b/tests/test_model.py
index e1376a8..62dab57 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -238,7 +238,8 @@ def test_numerical_stability_large_utilities(self):
         assert nll > 0
 
         # Gradient should also be stable
-        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))
 
     def test_sparse_dual_input_equivalence(self):
@@ -267,7 +268,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,)
 
@@ -278,7 +280,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
@@ -289,8 +292,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)
 
@@ -310,7 +313,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))
 
@@ -323,8 +327,8 @@ 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)
 
@@ -338,10 +342,14 @@ def test_gradient_sparse_dual(self):
         model = MultichoiceLogit(J, K)
 
         flat_beta = true_beta.flatten()
-        dense_grad = model.gradient(flat_beta, X, y_single, y_dual)
+        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)
-        tuple_grad = model.gradient(flat_beta, X, y_single, (dual_rows, dual_s, dual_t))
+        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)
 
@@ -356,7 +364,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,)
 
@@ -367,7 +375,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))
@@ -382,7 +390,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,)
 
 
@@ -511,12 +519,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},
         )
@@ -556,7 +570,7 @@ def test_log_likelihood_contributions_sum(self):
         flat_beta = true_beta.flatten()
 
         contributions = model.log_likelihood_contributions(
-            flat_beta, X, y_single, y_dual
+            X, y_single, y_dual, flat_beta
         )
         total = np.sum(contributions)
         direct = -model.neg_log_likelihood(flat_beta, X, y_single, y_dual)
@@ -570,9 +584,9 @@ def test_compute_standard_errors_custom_epsilon(self):
         model = MultichoiceLogit(J, K)
 
         flat_beta = true_beta.flatten()
-        std_errs_default = model.compute_standard_errors(flat_beta, X, y_single, y_dual)
+        std_errs_default = model.compute_standard_errors(X, y_single, y_dual, flat_beta)
         std_errs_smaller = model.compute_standard_errors(
-            flat_beta, X, y_single, y_dual, epsilon=1e-6
+            X, y_single, y_dual, flat_beta, epsilon=1e-6
         )
 
         assert std_errs_default.shape == std_errs_smaller.shape == ((J - 1) * K,)

From 4d2f9325f0f822f175d79aea282c7bd859a001bd Mon Sep 17 00:00:00 2001
From: Thomas Monk <246525+tmonk@users.noreply.github.com>
Date: Mon, 24 Nov 2025 23:43:46 +0000
Subject: [PATCH 4/5] refactor: Remove redundant public compatibility wrappers
 for `transform_params`, `calculate_utilities`, and `neg_log_likelihood`.

---
 multe/model.py | 31 ++++---------------------------
 1 file changed, 4 insertions(+), 27 deletions(-)

diff --git a/multe/model.py b/multe/model.py
index 4619493..14c0119 100644
--- a/multe/model.py
+++ b/multe/model.py
@@ -101,13 +101,6 @@ def _transform_params(
         beta_fixed = np.zeros((1, self.K))
         return np.vstack([beta_fixed, beta_free])
 
-    # Public alias for compatibility
-    def transform_params(
-        self, flat_beta: npt.NDArray[np.float64]
-    ) -> npt.NDArray[np.float64]:
-        """Public wrapper for parameter reshaping (kept for compatibility)."""
-        return self._transform_params(flat_beta)
-
     def _calculate_utilities(
         self, X: npt.NDArray[np.float64], beta: npt.NDArray[np.float64]
     ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
@@ -130,13 +123,6 @@ def _calculate_utilities(
         a = np.exp(V_stable)
         return V, a
 
-    # Public alias for compatibility
-    def calculate_utilities(
-        self, X: npt.NDArray[np.float64], beta: npt.NDArray[np.float64]
-    ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
-        """Public wrapper for utility calculation (kept for compatibility)."""
-        return self._calculate_utilities(X, beta)
-
     def _normalize_dual_indices(
         self, y_dual: DualInput, *, N: int, J: int
     ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
@@ -435,19 +421,6 @@ def _neg_log_likelihood(
 
         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: DualInput,
-    ) -> float:
-        """
-        Public wrapper: compute negative log-likelihood with on-the-fly indices.
-        """
-        single_indices, dual_indices = self._validate_data(X, y_single, y_dual)
-        return self._neg_log_likelihood(flat_beta, X, single_indices, dual_indices)
-
     def _gradient(
         self,
         flat_beta: npt.NDArray[np.float64],
@@ -475,6 +448,7 @@ def _gradient(
             # Gradient = X * (y - p)
             # y is one-hot, so for chosen col y=1, else 0
             # Add positive term for chosen alternatives (y=1)
+            # Use np.add.at for sparse addition
             np.add.at(grad, col_idx, X_sub)
 
             # Subtract prob term for all alternatives
@@ -522,6 +496,7 @@ def _gradient(
             # Complexity: O(n_dual) weight computations; still cheaper than looping over J.
 
             # 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)
             a_i = a[i_idx]  # (n_dual, J)
             M = w_r[:, np.newaxis] * a_i
@@ -536,6 +511,8 @@ def _gradient(
             grad += M.T @ X_i
 
             # 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  # (n_dual, K)
             grad_contrib_t = w_t[:, np.newaxis] * X_i  # (n_dual, K)
 

From 84ef95a94b661c322fe1de21850fe0b9fcc18ee1 Mon Sep 17 00:00:00 2001
From: Thomas Monk <246525+tmonk@users.noreply.github.com>
Date: Mon, 24 Nov 2025 23:48:25 +0000
Subject: [PATCH 5/5] emove legacy public wrappers and align callers

  - Drop compatibility wrappers for transform_params/calculate_utilities/neg_log_likelihood and update examples/tests to use internal APIs with cached indices.
  - Document sparse dual flattening order (s*J+t), keep validation tight, remove redundant diagonal check.
  - Fix lint nits in examples (unused vars, bare except) and keep benchmark timing using internal nll/grad.
  - Refresh README API list to match current public surface; tests now pass with new signatures.
---
 README.md                 |  1 -
 examples/basic_example.py |  7 ++++---
 examples/benchmark.py     |  5 ++---
 examples/csv_example.py   |  9 ++++++---
 tests/test_model.py       | 29 ++++++++++++++++++-----------
 5 files changed, 30 insertions(+), 21 deletions(-)

diff --git a/README.md b/README.md
index 3ef0d50..f330f63 100644
--- a/README.md
+++ b/README.md
@@ -115,7 +115,6 @@ 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
 - `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
diff --git a/examples/basic_example.py b/examples/basic_example.py
index f9ab7d4..2ab68a8 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},
     )
diff --git a/examples/benchmark.py b/examples/benchmark.py
index 83519e2..eb86371 100644
--- a/examples/benchmark.py
+++ b/examples/benchmark.py
@@ -44,10 +44,9 @@ def benchmark_estimation(N, J, K, mix_ratio=0.5, method="BFGS", seed=42):
     # 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},
     )
diff --git a/examples/csv_example.py b/examples/csv_example.py
index feeb63b..cfeaed0 100644
--- a/examples/csv_example.py
+++ b/examples/csv_example.py
@@ -145,11 +145,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},
     )
diff --git a/tests/test_model.py b/tests/test_model.py
index 62dab57..cdb75ed 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -32,7 +32,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
@@ -43,7 +43,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:
@@ -175,7 +175,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))
 
@@ -186,7 +187,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
 
@@ -197,7 +199,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
@@ -209,7 +212,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
@@ -232,13 +236,13 @@ 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
-        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))
 
@@ -249,11 +253,13 @@ def test_sparse_dual_input_equivalence(self):
         model = MultichoiceLogit(J, K)
 
         flat_beta = true_beta.flatten()
-        dense_nll = model.neg_log_likelihood(flat_beta, X, y_single, y_dual)
+        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)
-        tuple_nll = model.neg_log_likelihood(flat_beta, X, y_single, tuple_input)
+        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)
 
@@ -573,7 +579,8 @@ def test_log_likelihood_contributions_sum(self):
             X, y_single, y_dual, flat_beta
         )
         total = np.sum(contributions)
-        direct = -model.neg_log_likelihood(flat_beta, X, y_single, y_dual)
+        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)