Source code for biogeme.expressions.exp

"""Arithmetic expressions accepted by Biogeme:exp

Michel Bierlaire
10.04.2025 11:48
"""

from __future__ import annotations

import logging

import jax.numpy as jnp
import numpy as np
from biogeme.floating_point import MAX_EXP_ARG, MIN_EXP_ARG

from . import (
    ExpressionOrNumeric,
)
from .jax_utils import JaxFunctionType
from .unary_expressions import UnaryOperator

logger = logging.getLogger(__name__)




class exp(UnaryOperator):
    """
    exponential expression
    """

    def __init__(self, child: ExpressionOrNumeric) -> None:
        """Constructor

        :param child: first arithmetic expression
        :type child: biogeme.expressions.Expression
        """
        super().__init__(child)



    def deep_flat_copy(self) -> exp:
        """Provides a copy of the expression. It is deep in the sense that it generates copies of the children.
        It is flat in the sense that any `MultipleExpression` is transformed into the currently selected expression.
        The flat part is irrelevant for this expression.
        """
        copy_child = self.child.deep_flat_copy()
        return type(self)(child=copy_child)


    def __str__(self) -> str:
        return f'exp({self.child})'

    def __repr__(self) -> str:
        return f'exp({repr(self.child)})'



    def get_value(self) -> float:
        """Evaluates the value of the expression

        :return: value of the expression
        :rtype: float
        """
        return np.exp(self.child.get_value())




    def recursive_construct_jax_function(
        self, numerically_safe: bool
    ) -> JaxFunctionType:
        """
        Generates a function to be used by biogeme_jax. Must be overloaded by each
            expression
        :return: the function takes two parameters: the parameters, and one row
            of the database.
        """
        child_jax = self.child.recursive_construct_jax_function(
            numerically_safe=numerically_safe
        )

        if numerically_safe:

            def the_jax_function(
                parameters: jnp.ndarray,
                one_row: jnp.ndarray,
                the_draws: jnp.ndarray,
                the_random_variables: jnp.ndarray,
            ) -> jnp.ndarray:
                child_value = child_jax(
                    parameters, one_row, the_draws, the_random_variables
                )
                safe_value = jnp.clip(child_value, min=MIN_EXP_ARG, max=MAX_EXP_ARG)
                result = jnp.exp(safe_value)
                return result

            return the_jax_function

        def the_jax_function(
            parameters: jnp.ndarray,
            one_row: jnp.ndarray,
            the_draws: jnp.ndarray,
            the_random_variables: jnp.ndarray,
        ) -> jnp.ndarray:
            child_value = child_jax(
                parameters, one_row, the_draws, the_random_variables
            )
            result = jnp.exp(child_value)
            return result

        return the_jax_function