"""Arithmetic expressions accepted by Biogeme: sampled cross-nested logit."""
from __future__ import annotations
import logging
from itertools import chain
from typing import TYPE_CHECKING
import jax
import jax.numpy as jnp
import numpy as np
from biogeme.exceptions import BiogemeError
from biogeme.floating_point import JAX_FLOAT
from .base_expressions import Expression, LogitTuple
from .convert import validate_and_convert
from .jax_utils import JaxFunctionType
if TYPE_CHECKING:
from biogeme.nests import NestsForCrossNestedLogit, OldNestsForCrossNestedLogit
from . import ExpressionOrNumeric
logger = logging.getLogger(__name__)
def _index_of(key: jnp.ndarray, keys: jnp.ndarray) -> jnp.ndarray:
"""Return the index of a key in a vector of identifiers."""
return jnp.argmax(keys == key)
[docs]
class LogSampleCrossNested(Expression):
"""Log probability for a sampled cross-nested logit model.
This expression is designed for the sampling-of-alternatives framework.
It represents the same expression currently generated by
``GenerateModel.get_cross_nested_logit``:
.. math::
\\log P(0 \\mid S)
=
K_0 - \\log \\sum_{i \\in S} \\exp(K_i),
where
.. math::
K_i =
V_i - \\omega_i + \\log G_i,
and
.. math::
G_i =
\\sum_m
\\alpha_{im}^{\\mu_m}
\\exp((\\mu_m - 1) V_i)
B_m^{1/\\mu_m - 1},
with
.. math::
B_m =
\\sum_{j \\in S^{MEV}}
w_j \\alpha_{jm}^{\\mu_m}
\\exp(\\mu_m V_j).
The expression avoids constructing a large tree of ``ConditionalSum``,
``ConditionalTermTuple``, ``exp``, ``logzero``, and ``loglogit``
expressions.
"""
def __init__(
self,
utilities: dict[int, ExpressionOrNumeric],
log_probabilities: dict[int, ExpressionOrNumeric],
alphas: dict[str, dict[int, ExpressionOrNumeric]],
mev_utilities: dict[int, ExpressionOrNumeric],
mev_weights: dict[int, ExpressionOrNumeric],
mev_alphas: dict[str, dict[int, ExpressionOrNumeric]],
nests: NestsForCrossNestedLogit | OldNestsForCrossNestedLogit,
choice: ExpressionOrNumeric = 0,
):
"""Constructor.
:param utilities: utility expressions for the main sampled alternatives,
indexed by their position in the sample. The chosen alternative is
normally index 0.
:param log_probabilities: log sampling probability corrections for the
main sampled alternatives.
:param alphas: allocation parameters for the main sampled alternatives.
The first key is the nest name, and the second key is the sample
position.
:param mev_utilities: utility expressions for the MEV sample.
:param mev_weights: MEV sampling weights.
:param mev_alphas: allocation parameters for the MEV sampled
alternatives. The first key is the nest name, and the second key is
the MEV sample position.
:param nests: cross-nested-logit nest structure.
:param choice: expression identifying the chosen alternative in the
main sampled set. Defaults to 0.
"""
from biogeme.nests import NestsForCrossNestedLogit
Expression.__init__(self)
if not isinstance(nests, NestsForCrossNestedLogit):
logger.warning(
'It is recommended to define the nests of the cross-nested '
'logit model using OneNestForCrossNestedLogit and '
'NestsForCrossNestedLogit.'
)
nests = NestsForCrossNestedLogit(
choice_set=list(utilities),
tuple_of_nests=nests,
)
ok, message = nests.check_validity()
if not ok:
raise BiogemeError(message)
self._is_complex = True
self.nests = nests
self._check_same_keys(
reference=utilities,
candidate=log_probabilities,
candidate_name='log_probabilities',
)
if not utilities:
raise BiogemeError('The dictionary of sampled utilities cannot be empty.')
if not mev_utilities:
raise BiogemeError('The dictionary of MEV utilities cannot be empty.')
self.nest_names = tuple(nest.name for nest in nests)
self._check_alpha_structure(
alphas=alphas,
expected_outer_keys=set(self.nest_names),
expected_inner_keys=set(utilities),
name='alphas',
)
self._check_alpha_structure(
alphas=mev_alphas,
expected_outer_keys=set(self.nest_names),
expected_inner_keys=set(mev_utilities),
name='mev_alphas',
)
self._check_same_keys(
reference=mev_utilities,
candidate=mev_weights,
candidate_name='mev_weights',
)
self.utilities: dict[int, Expression] = {
sample_id: validate_and_convert(expression)
for sample_id, expression in utilities.items()
}
self.log_probabilities: dict[int, Expression] = {
sample_id: validate_and_convert(expression)
for sample_id, expression in log_probabilities.items()
}
self.alphas: dict[str, dict[int, Expression]] = {
nest_name: {
sample_id: validate_and_convert(expression)
for sample_id, expression in alpha_dict.items()
}
for nest_name, alpha_dict in alphas.items()
}
self.mev_utilities: dict[int, Expression] = {
sample_id: validate_and_convert(expression)
for sample_id, expression in mev_utilities.items()
}
self.mev_weights: dict[int, Expression] = {
sample_id: validate_and_convert(expression)
for sample_id, expression in mev_weights.items()
}
self.mev_alphas: dict[str, dict[int, Expression]] = {
nest_name: {
sample_id: validate_and_convert(expression)
for sample_id, expression in alpha_dict.items()
}
for nest_name, alpha_dict in mev_alphas.items()
}
self.choice: Expression = validate_and_convert(choice)
self.sample_ids = list(self.utilities.keys())
self.sample_keys = jnp.array(self.sample_ids, dtype=JAX_FLOAT)
self.mev_sample_ids = list(self.mev_utilities.keys())
self.utility_values = tuple(self.utilities[i] for i in self.sample_ids)
self.log_probability_values = tuple(
self.log_probabilities[i] for i in self.sample_ids
)
self.alpha_values = tuple(
tuple(self.alphas[nest_name][i] for i in self.sample_ids)
for nest_name in self.nest_names
)
self.mev_utility_values = tuple(
self.mev_utilities[i] for i in self.mev_sample_ids
)
self.mev_weight_values = tuple(self.mev_weights[i] for i in self.mev_sample_ids)
self.mev_alpha_values = tuple(
tuple(self.mev_alphas[nest_name][i] for i in self.mev_sample_ids)
for nest_name in self.nest_names
)
self.nest_parameters = tuple(
validate_and_convert(nest.nest_param) for nest in nests
)
self.children.append(self.choice)
for expression in self.utility_values:
self.children.append(expression)
for expression in self.log_probability_values:
self.children.append(expression)
for alpha_row in self.alpha_values:
for expression in alpha_row:
self.children.append(expression)
for expression in self.mev_utility_values:
self.children.append(expression)
for expression in self.mev_weight_values:
self.children.append(expression)
for alpha_row in self.mev_alpha_values:
for expression in alpha_row:
self.children.append(expression)
for expression in self.nest_parameters:
self.children.append(expression)
@staticmethod
def _check_same_keys(
reference: dict[int, ExpressionOrNumeric],
candidate: dict[int, ExpressionOrNumeric],
candidate_name: str,
) -> None:
"""Check that two dictionaries use the same keys."""
if set(reference) != set(candidate):
missing = set(reference) - set(candidate)
unknown = set(candidate) - set(reference)
raise BiogemeError(
f'The dictionary {candidate_name} must contain exactly the same '
f'sample identifiers as the utility dictionary. '
f'Missing entries: {missing}. Unknown entries: {unknown}.'
)
@staticmethod
def _check_alpha_structure(
alphas: dict[str, dict[int, ExpressionOrNumeric]],
expected_outer_keys: set[str],
expected_inner_keys: set[int],
name: str,
) -> None:
"""Check the two-level alpha dictionary structure."""
if set(alphas) != expected_outer_keys:
missing = expected_outer_keys - set(alphas)
unknown = set(alphas) - expected_outer_keys
raise BiogemeError(
f'The dictionary {name} must contain exactly the nest names. '
f'Missing entries: {missing}. Unknown entries: {unknown}.'
)
for nest_name, alpha_dict in alphas.items():
if set(alpha_dict) != expected_inner_keys:
missing = expected_inner_keys - set(alpha_dict)
unknown = set(alpha_dict) - expected_inner_keys
raise BiogemeError(
f'The dictionary {name}[{nest_name}] must contain exactly '
f'the sample identifiers. Missing entries: {missing}. '
f'Unknown entries: {unknown}.'
)
[docs]
def deep_flat_copy(self) -> LogSampleCrossNested:
"""Deep flat copy."""
return type(self)(
utilities={
key: value.deep_flat_copy() for key, value in self.utilities.items()
},
log_probabilities={
key: value.deep_flat_copy()
for key, value in self.log_probabilities.items()
},
alphas={
nest_name: {
key: value.deep_flat_copy() for key, value in alpha_dict.items()
}
for nest_name, alpha_dict in self.alphas.items()
},
mev_utilities={
key: value.deep_flat_copy() for key, value in self.mev_utilities.items()
},
mev_weights={
key: value.deep_flat_copy() for key, value in self.mev_weights.items()
},
mev_alphas={
nest_name: {
key: value.deep_flat_copy() for key, value in alpha_dict.items()
}
for nest_name, alpha_dict in self.mev_alphas.items()
},
nests=self.nests,
choice=self.choice.deep_flat_copy(),
)
[docs]
def logit_choice_avail(self) -> list[LogitTuple]:
"""Return availability structures appearing in this expression."""
return list(
chain.from_iterable(child.logit_choice_avail() for child in self.children)
)
[docs]
def get_value(self) -> float:
"""Evaluate the sampled cross-nested logit log probability using NumPy."""
choice = int(self.choice.get_value())
if choice not in self.utilities:
raise BiogemeError(
f'Alternative {choice} does not appear in the sampled utilities: '
f'{self.utilities.keys()}'
)
utilities = np.asarray(
[expression.get_value() for expression in self.utility_values],
dtype=float,
)
log_probabilities = np.asarray(
[expression.get_value() for expression in self.log_probability_values],
dtype=float,
)
alphas = np.asarray(
[
[expression.get_value() for expression in alpha_row]
for alpha_row in self.alpha_values
],
dtype=float,
)
mev_utilities = np.asarray(
[expression.get_value() for expression in self.mev_utility_values],
dtype=float,
)
mev_weights = np.asarray(
[expression.get_value() for expression in self.mev_weight_values],
dtype=float,
)
mev_alphas = np.asarray(
[
[expression.get_value() for expression in alpha_row]
for alpha_row in self.mev_alpha_values
],
dtype=float,
)
nest_parameters = np.asarray(
[expression.get_value() for expression in self.nest_parameters],
dtype=float,
)
# --- Begin numerically robust computation ---
log_mev_sums = []
for nest_index, mu_m in enumerate(nest_parameters):
mev_weights_for_nest = mev_weights * mev_alphas[nest_index, :] ** mu_m
positive_terms = mev_weights_for_nest > 0.0
if not np.any(positive_terms):
log_mev_sums.append(-np.inf)
continue
mev_log_terms = (
np.log(mev_weights_for_nest[positive_terms])
+ mu_m * mev_utilities[positive_terms]
)
max_term = np.max(mev_log_terms)
log_mev_sum = max_term + np.log(np.sum(np.exp(mev_log_terms - max_term)))
log_mev_sums.append(log_mev_sum)
log_mev_sums = np.asarray(log_mev_sums, dtype=float)
log_gi_terms = np.full(
(len(nest_parameters), len(utilities)),
-np.inf,
dtype=float,
)
for nest_index, mu_m in enumerate(nest_parameters):
positive_alpha = alphas[nest_index, :] > 0.0
finite_mev_sum = np.isfinite(log_mev_sums[nest_index])
if not finite_mev_sum:
continue
log_gi_terms[nest_index, positive_alpha] = (
mu_m * np.log(alphas[nest_index, positive_alpha])
+ (mu_m - 1.0) * utilities[positive_alpha]
+ ((1.0 / mu_m) - 1.0) * log_mev_sums[nest_index]
)
max_log_gi_terms = np.max(log_gi_terms, axis=0)
finite_log_gi = np.isfinite(max_log_gi_terms)
log_gi = np.zeros_like(utilities)
log_gi[finite_log_gi] = max_log_gi_terms[finite_log_gi] + np.log(
np.sum(
np.exp(
log_gi_terms[:, finite_log_gi]
- max_log_gi_terms[finite_log_gi][None, :]
),
axis=0,
)
)
# If all log_gi_terms for an alternative are -inf, log_gi remains 0.0 (log(1))
# --- End numerically robust computation ---
kernels = utilities - log_probabilities + log_gi
choice_index = self.sample_ids.index(choice)
chosen_kernel = kernels[choice_index]
max_kernel = np.max(kernels)
log_denominator = max_kernel + np.log(np.sum(np.exp(kernels - max_kernel)))
return chosen_kernel - log_denominator
def __str__(self) -> str:
entries = ', '.join(
f'{sample_id}:{self.utilities[sample_id]}'
f'-{self.log_probabilities[sample_id]}'
for sample_id in self.sample_ids
)
return f'{self.get_class_name()}[choice={self.choice}; kernels=({entries})]'
[docs]
def recursive_construct_jax_function(
self,
numerically_safe: bool,
) -> JaxFunctionType:
"""Generate a compact JAX function for sampled cross-nested logit."""
utility_functions = tuple(
expression.recursive_construct_jax_function(
numerically_safe=numerically_safe
)
for expression in self.utility_values
)
log_probability_functions = tuple(
expression.recursive_construct_jax_function(
numerically_safe=numerically_safe
)
for expression in self.log_probability_values
)
alpha_functions = tuple(
tuple(
expression.recursive_construct_jax_function(
numerically_safe=numerically_safe
)
for expression in alpha_row
)
for alpha_row in self.alpha_values
)
mev_utility_functions = tuple(
expression.recursive_construct_jax_function(
numerically_safe=numerically_safe
)
for expression in self.mev_utility_values
)
mev_weight_functions = tuple(
expression.recursive_construct_jax_function(
numerically_safe=numerically_safe
)
for expression in self.mev_weight_values
)
mev_alpha_functions = tuple(
tuple(
expression.recursive_construct_jax_function(
numerically_safe=numerically_safe
)
for expression in alpha_row
)
for alpha_row in self.mev_alpha_values
)
nest_parameter_functions = tuple(
expression.recursive_construct_jax_function(
numerically_safe=numerically_safe
)
for expression in self.nest_parameters
)
choice_function = self.choice.recursive_construct_jax_function(
numerically_safe=numerically_safe
)
sample_keys = self.sample_keys
def evaluate_all(
functions,
parameters: jnp.ndarray,
one_row: jnp.ndarray,
the_draws: jnp.ndarray,
the_random_variables: jnp.ndarray,
) -> jnp.ndarray:
return jnp.stack(
[
function(parameters, one_row, the_draws, the_random_variables)
for function in functions
],
axis=0,
)
def evaluate_matrix(
matrix_functions,
parameters: jnp.ndarray,
one_row: jnp.ndarray,
the_draws: jnp.ndarray,
the_random_variables: jnp.ndarray,
) -> jnp.ndarray:
return jnp.stack(
[
evaluate_all(
row_functions,
parameters,
one_row,
the_draws,
the_random_variables,
)
for row_functions in matrix_functions
],
axis=0,
)
def the_jax_function(
parameters: jnp.ndarray,
one_row: jnp.ndarray,
the_draws: jnp.ndarray,
the_random_variables: jnp.ndarray,
) -> jnp.ndarray:
utilities = evaluate_all(
utility_functions,
parameters,
one_row,
the_draws,
the_random_variables,
)
log_probabilities = evaluate_all(
log_probability_functions,
parameters,
one_row,
the_draws,
the_random_variables,
)
alphas = evaluate_matrix(
alpha_functions,
parameters,
one_row,
the_draws,
the_random_variables,
)
mev_utilities = evaluate_all(
mev_utility_functions,
parameters,
one_row,
the_draws,
the_random_variables,
)
mev_weights = evaluate_all(
mev_weight_functions,
parameters,
one_row,
the_draws,
the_random_variables,
)
mev_alphas = evaluate_matrix(
mev_alpha_functions,
parameters,
one_row,
the_draws,
the_random_variables,
)
nest_parameters = evaluate_all(
nest_parameter_functions,
parameters,
one_row,
the_draws,
the_random_variables,
)
# Shapes:
# utilities: (S,)
# log_probabilities: (S,)
# alphas: (M, S)
# mev_utilities: (R,)
# mev_weights: (R,)
# mev_alphas: (M, R)
# nest_parameters: (M,)
safe_mev_alphas = jnp.where(mev_alphas > 0.0, mev_alphas, 1.0)
mev_alpha_power = jnp.where(
mev_alphas > 0.0,
safe_mev_alphas ** nest_parameters[:, None],
0.0,
)
mev_log_sums = jax.nn.logsumexp(
nest_parameters[:, None] * mev_utilities[None, :],
axis=1,
b=mev_weights[None, :] * mev_alpha_power,
)
safe_alphas = jnp.where(alphas > 0.0, alphas, 1.0)
main_alpha_power = jnp.where(
alphas > 0.0,
safe_alphas ** nest_parameters[:, None],
0.0,
)
gi_log_terms = (nest_parameters[:, None] - 1.0) * utilities[None, :] + (
(1.0 / nest_parameters) - 1.0
)[:, None] * mev_log_sums[:, None]
log_gi = jax.nn.logsumexp(
gi_log_terms,
axis=0,
b=main_alpha_power,
)
log_gi = jnp.where(
jnp.isneginf(log_gi),
jnp.array(0.0, dtype=JAX_FLOAT),
log_gi,
)
kernels = utilities - log_probabilities + log_gi
choice_id = choice_function(
parameters,
one_row,
the_draws,
the_random_variables,
)
choice_index = _index_of(choice_id, sample_keys)
chosen_kernel = kernels[choice_index]
return chosen_kernel - jax.nn.logsumexp(kernels)
return the_jax_function