"""Utility recovery via linear programming (Afriat's inequalities)."""
from __future__ import annotations
import time
from typing import Callable, cast
import numpy as np
from numpy.typing import NDArray
from scipy.optimize import linprog
from prefgraph.core.session import ConsumerSession
from prefgraph.core.result import UtilityRecoveryResult
from prefgraph.core.exceptions import OptimizationError
def recover_utility(
session: ConsumerSession,
tolerance: float = 1e-8,
solver: str = "highs",
) -> UtilityRecoveryResult:
"""
Recover utility values satisfying Afriat's inequalities using LP.
If the data satisfies GARP, we can recover utility values U_k and
Lagrange multipliers (marginal utility of money) lambda_k such that
Afriat's inequalities hold:
U_k <= U_l + lambda_l * p_l @ (x_k - x_l) for all k, l
This is solved as a linear programming feasibility problem:
- Variables: U_1, ..., U_T, lambda_1, ..., lambda_T (2T variables)
- Constraints: Afriat inequalities for all k, l pairs
- Objective: Minimize sum of variables (for a centered solution)
The recovered utility function is piecewise linear and concave.
Args:
session: ConsumerSession with prices and quantities
tolerance: Numerical tolerance for LP solver
solver: scipy.optimize.linprog method name
Returns:
UtilityRecoveryResult with utility values and multipliers
Example:
>>> import numpy as np
>>> from prefgraph import ConsumerSession, recover_utility
>>> prices = np.array([[1.0, 2.0], [2.0, 1.0]])
>>> quantities = np.array([[4.0, 1.0], [1.0, 4.0]])
>>> session = ConsumerSession(prices=prices, quantities=quantities)
>>> result = recover_utility(session)
>>> if result.success:
... print(f"Utility values: {result.utility_values}")
"""
start_time = time.perf_counter()
T = session.num_observations
# prices/quantities are Optional on the dataclass but always populated by
# ConsumerSession.__post_init__, so this is a cast (not a runtime guard).
P = cast("NDArray[np.float64]", session.prices)
Q = cast("NDArray[np.float64]", session.quantities)
# Number of variables: U_1, ..., U_T, lambda_1, ..., lambda_T
n_vars = 2 * T
# Build inequality constraints: A_ub @ x <= b_ub
# Afriat inequality: U_k - U_l - lambda_l * p_l @ (x_k - x_l) <= 0
# Rearranged for LP format
constraints_A = []
constraints_b = []
for k in range(T):
for obs_l in range(T):
if k == obs_l:
continue
# Constraint: U_k - U_l - lambda_l * [p_l @ (x_k - x_l)] <= 0
# Variables are [U_0, ..., U_{T-1}, lambda_0, ..., lambda_{T-1}]
row = np.zeros(n_vars)
row[k] = 1.0 # Coefficient for U_k
row[obs_l] = -1.0 # Coefficient for U_l
# Coefficient for lambda_l: -p_l @ (x_k - x_l) = p_l @ (x_l - x_k)
diff = Q[obs_l] - Q[k] # x_l - x_k
lambda_coef = P[obs_l] @ diff # p_l @ (x_l - x_k)
row[T + obs_l] = lambda_coef
constraints_A.append(row)
constraints_b.append(0.0)
A_ub = np.array(constraints_A)
b_ub = np.array(constraints_b)
# Bounds:
# - U_k >= 0 (utility can be normalized to be non-negative)
# - lambda_k > 0 (strictly positive marginal utility of money)
epsilon = 1e-6
bounds = [(0, None)] * T + [(epsilon, None)] * T # U >= 0, lambda > 0
# Objective: minimize sum of lambdas (utility values don't need minimizing)
# This finds a solution with minimal marginal utilities of money
c = np.zeros(n_vars)
c[T:] = 1.0 # Only minimize lambdas
# Solve LP using HiGHS solver
try:
# scipy-stubs types `method` as a Literal of solver names, but linprog
# accepts an arbitrary runtime string; passing our `solver: str` param
# is valid scipy usage the stub cannot express.
result = linprog( # type: ignore[call-overload]
c,
A_ub=A_ub,
b_ub=b_ub,
bounds=bounds,
method=solver,
options={"presolve": True},
)
except Exception as e:
computation_time = (time.perf_counter() - start_time) * 1000
return UtilityRecoveryResult(
success=False,
utility_values=None,
lagrange_multipliers=None,
lp_status=f"LP solver error: {str(e)}",
residuals=None,
computation_time_ms=computation_time,
)
computation_time = (time.perf_counter() - start_time) * 1000
if result.success:
U = result.x[:T]
lambdas = result.x[T:]
# Compute residuals for verification
residuals = _compute_afriat_residuals(U, lambdas, P, Q)
return UtilityRecoveryResult(
success=True,
utility_values=U,
lagrange_multipliers=lambdas,
lp_status=result.message,
residuals=residuals,
computation_time_ms=computation_time,
)
else:
return UtilityRecoveryResult(
success=False,
utility_values=None,
lagrange_multipliers=None,
lp_status=result.message,
residuals=None,
computation_time_ms=computation_time,
)
def _compute_afriat_residuals(
U: NDArray[np.float64],
lambdas: NDArray[np.float64],
P: NDArray[np.float64],
Q: NDArray[np.float64],
) -> NDArray[np.float64]:
"""
Compute residuals of Afriat inequalities for verification.
Residual[k,l] = U_l + lambda_l * p_l @ (x_k - x_l) - U_k
If constraints are satisfied, all residuals should be >= 0.
Args:
U: Utility values
lambdas: Lagrange multipliers
P: Price matrix
Q: Quantity matrix
Returns:
T x T matrix of residuals
"""
T = len(U)
residuals = np.zeros((T, T))
for k in range(T):
for obs_l in range(T):
# Afriat inequality: U_k <= U_l + lambda_l * p_l @ (x_k - x_l)
# Residual = RHS - LHS = U_l + lambda_l * p_l @ (x_k - x_l) - U_k
diff = Q[k] - Q[obs_l]
residuals[k, obs_l] = U[obs_l] + lambdas[obs_l] * (P[obs_l] @ diff) - U[k]
return residuals
def construct_afriat_utility(
session: ConsumerSession,
utility_result: UtilityRecoveryResult,
) -> Callable[[NDArray[np.float64]], float]:
"""
Construct the Afriat utility function from recovered parameters.
The Afriat utility function is piecewise linear and concave:
u(x) = min_k { U_k + lambda_k * p_k @ (x - x_k) }
This function rationalizes the observed data: for each observation k,
the bundle x_k maximizes u(x) subject to the budget constraint p_k @ x <= p_k @ x_k.
Args:
session: ConsumerSession with the choice data
utility_result: Result from recover_utility
Returns:
Callable function u(x) that takes a bundle and returns utility
Raises:
ValueError: If utility recovery was not successful
Example:
>>> result = recover_utility(session)
>>> if result.success:
... u = construct_afriat_utility(session, result)
... # Evaluate utility at a new bundle
... new_bundle = np.array([2.0, 3.0])
... utility = u(new_bundle)
"""
if not utility_result.success:
raise OptimizationError(
"Cannot construct utility from failed recovery. "
f"LP status: {utility_result.lp_status}. "
"Hint: Check data consistency with compute_integrity_score() first. "
"If integrity is low, the behavior may be too inconsistent for utility recovery."
)
if (
utility_result.utility_values is None
or utility_result.lagrange_multipliers is None
):
raise OptimizationError(
"Utility values or multipliers are None despite successful LP. "
"This may indicate a numerical issue. Try adjusting tolerance."
)
# Non-None is guaranteed by the success/None checks above; the explicit
# annotations let the nested closure see concrete arrays (mypy does not
# carry flow narrowing into nested function scopes).
U: NDArray[np.float64] = utility_result.utility_values
lambdas: NDArray[np.float64] = utility_result.lagrange_multipliers
P = cast("NDArray[np.float64]", session.prices)
Q = cast("NDArray[np.float64]", session.quantities)
T = len(U)
def afriat_utility(x: NDArray[np.float64]) -> float:
"""
Evaluate Afriat utility at bundle x.
u(x) = min_k { U_k + lambda_k * p_k @ (x - x_k) }
Args:
x: Bundle (array of quantities for each good)
Returns:
Utility value
"""
x = np.asarray(x, dtype=np.float64)
values = []
for k in range(T):
val = U[k] + lambdas[k] * (P[k] @ (x - Q[k]))
values.append(val)
# min over numpy scalars yields a numpy float; the declared float
# return type covers it (np.float64 subclasses float).
return cast(float, min(values))
return afriat_utility
def predict_demand(
session: ConsumerSession,
utility_result: UtilityRecoveryResult,
new_prices: NDArray[np.float64],
budget: float,
n_goods: int | None = None,
) -> NDArray[np.float64] | None:
"""
Predict demand at new prices using the recovered utility function.
Uses the Afriat utility function to find the utility-maximizing bundle
at the given prices and budget.
This is a simple grid search implementation suitable for small problems.
For production use with many goods, consider using scipy.optimize.
Args:
session: ConsumerSession with choice data
utility_result: Result from recover_utility
new_prices: Price vector for the new budget set
budget: Total budget available
n_goods: Number of goods (inferred from prices if not provided)
Returns:
Predicted demand bundle, or None if optimization fails
"""
if not utility_result.success:
return None
u = construct_afriat_utility(session, utility_result)
new_prices = np.asarray(new_prices, dtype=np.float64)
if n_goods is None:
n_goods = len(new_prices)
if n_goods > 3:
# For more than 3 goods, use scipy optimization
from scipy.optimize import minimize
def neg_utility(x: NDArray[np.float64]) -> float:
return -u(x)
# Budget constraint: p @ x <= budget
constraints = {"type": "ineq", "fun": lambda x: budget - new_prices @ x}
bounds = [(0, budget / p) for p in new_prices]
# Start from equal allocation
x0 = np.full(n_goods, budget / (n_goods * np.mean(new_prices)))
# scipy-stubs types `constraints` as the _ConstraintDict TypedDict, but
# a plain {"type": ..., "fun": ...} dict is valid scipy input the stub
# overloads cannot match.
result = minimize( # type: ignore[call-overload]
neg_utility,
x0,
method="SLSQP",
bounds=bounds,
constraints=constraints,
)
if result.success:
return cast("NDArray[np.float64]", result.x)
return None
else:
# For 2-3 goods, use grid search for simplicity
best_x = None
best_utility = float("-inf")
# Create grid
n_points = 50
max_quantities = budget / new_prices
if n_goods == 2:
for q0 in np.linspace(0, max_quantities[0], n_points):
remaining = budget - new_prices[0] * q0
if remaining >= 0:
q1 = remaining / new_prices[1]
x = np.array([q0, q1])
util = u(x)
if util > best_utility:
best_utility = util
best_x = x
return best_x
# =============================================================================
# TECH-FRIENDLY ALIASES
# =============================================================================
# fit_latent_values: Tech-friendly name for recover_utility
fit_latent_values = recover_utility
"""
Fit latent preference values from user behavior.
This is the tech-friendly alias for recover_utility.
Extracts latent preference values that can be used as:
- Features for ML models
- Inputs to counterfactual simulations
- User embeddings for personalization
Example:
>>> from prefgraph import BehaviorLog, fit_latent_values
>>> result = fit_latent_values(user_log)
>>> if result.converged:
... features = result.latent_values # Use as ML features
Returns:
LatentValueResult with latent_values array
"""
# build_value_function: Tech-friendly name for construct_afriat_utility
build_value_function = construct_afriat_utility
"""
Build a callable value function from fitted latent values.
This is the tech-friendly alias for construct_afriat_utility.
Returns a function that estimates the value/utility of any action vector.
Example:
>>> from prefgraph import BehaviorLog, fit_latent_values, build_value_function
>>> result = fit_latent_values(user_log)
>>> value_fn = build_value_function(user_log, result)
>>> estimated_value = value_fn(new_action_vector)
"""
# predict_choice: Tech-friendly name for predict_demand
predict_choice = predict_demand
"""
Predict what action a user will take given new costs and a resource limit.
This is the tech-friendly alias for predict_demand.
Uses the fitted latent value model to predict user behavior under
new conditions (counterfactual analysis).
Example:
>>> from prefgraph import BehaviorLog, fit_latent_values, predict_choice
>>> result = fit_latent_values(user_log)
>>> predicted_action = predict_choice(user_log, result, new_costs, budget)
"""