"""Arithmetic expressions accepted by Biogeme: sampled 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 NestsForNestedLogit, OldNestsForNestedLogit
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 LogSampledNested(Expression):
"""Log probability for a sampled nested logit model.
This expression is designed for the sampling-of-alternatives framework.
It represents the same expression currently generated by
``GenerateModel.get_nested_logit``:
.. math::
\\log P(0 \\mid S)
=
K_0 - \\log \\sum_{i \\in S} \\exp(K_i),
where the corrected kernel is
.. math::
K_i =
V_i
- \\omega_i
+
M_i,
and, if alternative :math:`i` belongs to nest :math:`m`,
.. math::
M_i =
(\\mu_m - 1) V_i
+
\\left(\\frac{1}{\\mu_m} - 1\\right)
\\log
\\sum_{j \\in S^{MEV}_m}
w_j \\exp(\\mu_m V_j).
Here :math:`\\omega_i` is the log sampling probability correction and
:math:`w_j` is the MEV sampling weight.
The expression avoids constructing a large tree of ``ConditionalSum``,
``ConditionalTermTuple``, ``BelongsTo``, ``exp``, and ``log`` expressions.
"""
def __init__(
self,
utilities: dict[int, ExpressionOrNumeric],
log_probabilities: dict[int, ExpressionOrNumeric],
alternative_ids: dict[int, ExpressionOrNumeric],
mev_utilities: dict[int, ExpressionOrNumeric],
mev_weights: dict[int, ExpressionOrNumeric],
mev_alternative_ids: dict[int, ExpressionOrNumeric],
nests: NestsForNestedLogit | OldNestsForNestedLogit,
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 alternative_ids: expressions giving the true alternative ids of
the main sampled alternatives.
:param mev_utilities: utility expressions for the MEV sample.
:param mev_weights: MEV sampling weights.
:param mev_alternative_ids: expressions giving the true alternative ids
of the MEV sampled alternatives.
:param nests: nested-logit nest structure.
:param choice: expression identifying the chosen alternative in the
main sampled set. Defaults to 0.
"""
from biogeme.nests import NestsForNestedLogit
Expression.__init__(self)
if not isinstance(nests, NestsForNestedLogit):
logger.warning(
'It is recommended to define the nests of the nested logit '
'model using OneNestForNestedLogit and NestsForNestedLogit.'
)
nests = NestsForNestedLogit(
choice_set=list(utilities),
tuple_of_nests=nests,
)
ok, message = nests.check_partition()
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',
)
self._check_same_keys(
reference=utilities,
candidate=alternative_ids,
candidate_name='alternative_ids',
)
self._check_same_keys(
reference=mev_utilities,
candidate=mev_weights,
candidate_name='mev_weights',
)
self._check_same_keys(
reference=mev_utilities,
candidate=mev_alternative_ids,
candidate_name='mev_alternative_ids',
)
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.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.alternative_ids: dict[int, Expression] = {
sample_id: validate_and_convert(expression)
for sample_id, expression in alternative_ids.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_alternative_ids: dict[int, Expression] = {
sample_id: validate_and_convert(expression)
for sample_id, expression in mev_alternative_ids.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.alternative_id_values = tuple(
self.alternative_ids[i] for i in self.sample_ids
)
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_alternative_id_values = tuple(
self.mev_alternative_ids[i] for i in self.mev_sample_ids
)
self.nest_parameters = tuple(
validate_and_convert(nest.nest_param) for nest in nests
)
self.nest_alternative_ids = tuple(
jnp.array(list(nest.list_of_alternatives), dtype=JAX_FLOAT)
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 expression in self.alternative_id_values:
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 expression in self.mev_alternative_id_values:
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}.'
)
[docs]
def deep_flat_copy(self) -> LogSampledNested:
"""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()
},
alternative_ids={
key: value.deep_flat_copy()
for key, value in self.alternative_ids.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_alternative_ids={
key: value.deep_flat_copy()
for key, value in self.mev_alternative_ids.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 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,
)
alternative_ids = np.asarray(
[expression.get_value() for expression in self.alternative_id_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_alternative_ids = np.asarray(
[expression.get_value() for expression in self.mev_alternative_id_values],
dtype=float,
)
mev_terms = np.zeros_like(utilities)
for nest, nest_parameter in zip(self.nests, self.nest_parameters, strict=True):
mu_m = nest_parameter.get_value()
nest_ids = set(nest.list_of_alternatives)
mev_membership = np.asarray(
[alternative_id in nest_ids for alternative_id in mev_alternative_ids],
dtype=float,
)
mev_sum = np.sum(
mev_membership * mev_weights * np.exp(mu_m * mev_utilities)
)
if mev_sum <= 0.0:
log_mev_sum = -np.inf
else:
log_mev_sum = np.log(mev_sum)
main_membership = np.asarray(
[alternative_id in nest_ids for alternative_id in alternative_ids],
dtype=float,
)
nest_term = (mu_m - 1.0) * utilities + ((1.0 / mu_m) - 1.0) * log_mev_sum
mev_terms += main_membership * nest_term
kernels = utilities - log_probabilities + mev_terms
choice_index = self.sample_ids.index(choice)
chosen_kernel = kernels[choice_index]
return chosen_kernel - float(jax.nn.logsumexp(jnp.asarray(kernels)))
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 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
)
alternative_id_functions = tuple(
expression.recursive_construct_jax_function(
numerically_safe=numerically_safe
)
for expression in self.alternative_id_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_alternative_id_functions = tuple(
expression.recursive_construct_jax_function(
numerically_safe=numerically_safe
)
for expression in self.mev_alternative_id_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
nest_alternative_ids = self.nest_alternative_ids
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 belongs_to(
ids: jnp.ndarray,
alternatives: jnp.ndarray,
) -> jnp.ndarray:
"""Check if each id belongs to the static set of alternatives."""
return jnp.any(ids[:, None] == alternatives[None, :], axis=1)
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,
)
alternative_ids = evaluate_all(
alternative_id_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_alternative_ids = evaluate_all(
mev_alternative_id_functions,
parameters,
one_row,
the_draws,
the_random_variables,
)
nest_parameters = evaluate_all(
nest_parameter_functions,
parameters,
one_row,
the_draws,
the_random_variables,
)
mev_terms = jnp.zeros_like(utilities)
for nest_index, alternatives in enumerate(nest_alternative_ids):
mu_m = nest_parameters[nest_index]
mev_membership = belongs_to(
ids=mev_alternative_ids,
alternatives=alternatives,
).astype(JAX_FLOAT)
log_mev_sum = jax.nn.logsumexp(
mu_m * mev_utilities,
b=mev_membership * mev_weights,
)
main_membership = belongs_to(
ids=alternative_ids,
alternatives=alternatives,
).astype(JAX_FLOAT)
nest_term = (mu_m - 1.0) * utilities + (
(1.0 / mu_m) - 1.0
) * log_mev_sum
mev_terms = mev_terms + main_membership * nest_term
kernels = utilities - log_probabilities + mev_terms
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