# Path: blackjax/mcmc/integrators.py
# class IntegratorState(NamedTuple):
#     """State of the trajectory integration.
# 
#     We keep the gradient of the potential energy to speedup computations.
#     """
# 
#     position: PyTree
#     momentum: PyTree
#     potential_energy: float
#     potential_energy_grad: PyTree
# 
# Path: blackjax/mcmc/proposal.py
# class Proposal(NamedTuple):
#     """Proposal for the next chain step.
#     state:
#         The trajectory state corresponding to this proposal.
#     energy:
#         The potential energy corresponding to the state.
#     weight:
#         Weight of the proposal. It is equal to the logarithm of the sum of the canonical
#         densities of each state :math:`e^{-H(z)}` along the trajectory.
#     sum_log_p_accept:
#         cumulated Metropolis-Hastings acceptance probabilty across entire trajectory.
#     """
# 
#     state: IntegratorState
#     energy: float
#     weight: float
#     sum_log_p_accept: float
# 
# def progressive_biased_sampling(rng_key, proposal, new_proposal):
#     """Baised proposal sampling.
# 
#     Unlike uniform sampling, biased sampling favors new proposals. It thus
#     biases the transition away from the trajectory's initial state.
# 
#     """
#     p_accept = jnp.clip(jnp.exp(new_proposal.weight - proposal.weight), a_max=1)
#     do_accept = jax.random.bernoulli(rng_key, p_accept)
#     new_weight = jnp.logaddexp(proposal.weight, new_proposal.weight)
#     new_sum_log_p_accept = jnp.logaddexp(
#         proposal.sum_log_p_accept, new_proposal.sum_log_p_accept
#     )
# 
#     return jax.lax.cond(
#         do_accept,
#         lambda _: Proposal(
#             new_proposal.state,
#             new_proposal.energy,
#             new_weight,
#             new_sum_log_p_accept,
#         ),
#         lambda _: Proposal(
#             proposal.state,
#             proposal.energy,
#             new_weight,
#             new_sum_log_p_accept,
#         ),
#         operand=None,
#     )
# 
# def progressive_uniform_sampling(rng_key, proposal, new_proposal):
#     # Using expit to compute exp(w1) / (exp(w0) + exp(w1))
#     p_accept = jax.scipy.special.expit(new_proposal.weight - proposal.weight)
#     do_accept = jax.random.bernoulli(rng_key, p_accept)
#     new_weight = jnp.logaddexp(proposal.weight, new_proposal.weight)
#     new_sum_log_p_accept = jnp.logaddexp(
#         proposal.sum_log_p_accept, new_proposal.sum_log_p_accept
#     )
# 
#     return jax.lax.cond(
#         do_accept,
#         lambda _: Proposal(
#             new_proposal.state,
#             new_proposal.energy,
#             new_weight,
#             new_sum_log_p_accept,
#         ),
#         lambda _: Proposal(
#             proposal.state,
#             proposal.energy,
#             new_weight,
#             new_sum_log_p_accept,
#         ),
#         operand=None,
#     )
# 
# def proposal_generator(
#     kinetic_energy: Callable, divergence_threshold: float
# ) -> Tuple[Callable, Callable]:
#     def new(state: IntegratorState) -> Proposal:
#         energy = state.potential_energy + kinetic_energy(state.momentum)
#         return Proposal(state, energy, 0.0, -np.inf)
# 
#     def update(initial_energy: float, state: IntegratorState) -> Tuple[Proposal, bool]:
#         """Generate a new proposal from a trajectory state.
# 
#         The trajectory state records information about the position in the state
#         space and corresponding potential energy. A proposal also carries a
#         weight that is equal to the difference between the current energy and
#         the previous one. It thus carries information about the previous state
#         as well as the current state.
# 
#         Parameters
#         ----------
#         initial_energy:
#             The initial energy.
#         state:
#             The new state.
# 
#         """
#         new_energy = state.potential_energy + kinetic_energy(state.momentum)
# 
#         delta_energy = initial_energy - new_energy
#         delta_energy = jnp.where(jnp.isnan(delta_energy), -jnp.inf, delta_energy)
#         is_transition_divergent = jnp.abs(delta_energy) > divergence_threshold
# 
#         # The weight of the new proposal is equal to H0 - H(z_new)
#         weight = delta_energy
#         # Acceptance statistic min(e^{H0 - H(z_new)}, 1)
#         sum_log_p_accept = jnp.minimum(delta_energy, 0.0)
# 
#         return (
#             Proposal(
#                 state,
#                 new_energy,
#                 weight,
#                 sum_log_p_accept,
#             ),
#             is_transition_divergent,
#         )
# 
#     return new, update
# 
# Path: blackjax/types.py



from typing import Callable, NamedTuple, Tuple
from blackjax.mcmc.integrators import IntegratorState
from blackjax.mcmc.proposal import (
    Proposal,
    progressive_biased_sampling,
    progressive_uniform_sampling,
    proposal_generator,
)
from blackjax.types import PRNGKey, PyTree
import jax
import jax.numpy as jnp
2. Sampling a new state from this sampled trajectory.

Step (1) ensures that the process is reversible and thus that detailed balance
is respected. The traditional implementation of HMC does not sample a
trajectory, but instead takes a fixed number of steps in the same direction and
flips the momentum of the last state.

We distinguish here between two different methods to sample trajectories: static
and dynamic sampling. In the static setting we sample trajectories with a fixed
number of steps, while in the dynamic setting the total number of steps is
determined by a dynamic termination criterion. Traditional HMC falls in the
former category, NUTS in the latter.

There are also two methods to sample proposals from these trajectories. In the
static setting we first build the trajectory and then sample a proposal from
this trajectory. In the progressive setting we update the proposal as the
trajectory is being sampled. While the former is faster, we risk saturating the
memory by keeping states that will subsequently be discarded.


References
----------
.. [1]: Betancourt, Michael. "A conceptual introduction to Hamiltonian Monte Carlo." arXiv preprint arXiv:1701.02434 (2017).

"""




class Trajectory(NamedTuple):

# Path: blackjax/mcmc/integrators.py
# class IntegratorState(NamedTuple):
#     """State of the trajectory integration.
# 
#     We keep the gradient of the potential energy to speedup computations.
#     """
# 
#     position: PyTree
#     momentum: PyTree
#     potential_energy: float
#     potential_energy_grad: PyTree
# 
# Path: blackjax/mcmc/proposal.py
# class Proposal(NamedTuple):
#     """Proposal for the next chain step.
#     state:
#         The trajectory state corresponding to this proposal.
#     energy:
#         The potential energy corresponding to the state.
#     weight:
#         Weight of the proposal. It is equal to the logarithm of the sum of the canonical
#         densities of each state :math:`e^{-H(z)}` along the trajectory.
#     sum_log_p_accept:
#         cumulated Metropolis-Hastings acceptance probabilty across entire trajectory.
#     """
# 
#     state: IntegratorState
#     energy: float
#     weight: float
#     sum_log_p_accept: float
# 
# def progressive_biased_sampling(rng_key, proposal, new_proposal):
#     """Baised proposal sampling.
# 
#     Unlike uniform sampling, biased sampling favors new proposals. It thus
#     biases the transition away from the trajectory's initial state.
# 
#     """
#     p_accept = jnp.clip(jnp.exp(new_proposal.weight - proposal.weight), a_max=1)
#     do_accept = jax.random.bernoulli(rng_key, p_accept)
#     new_weight = jnp.logaddexp(proposal.weight, new_proposal.weight)
#     new_sum_log_p_accept = jnp.logaddexp(
#         proposal.sum_log_p_accept, new_proposal.sum_log_p_accept
#     )
# 
#     return jax.lax.cond(
#         do_accept,
#         lambda _: Proposal(
#             new_proposal.state,
#             new_proposal.energy,
#             new_weight,
#             new_sum_log_p_accept,
#         ),
#         lambda _: Proposal(
#             proposal.state,
#             proposal.energy,
#             new_weight,
#             new_sum_log_p_accept,
#         ),
#         operand=None,
#     )
# 
# def progressive_uniform_sampling(rng_key, proposal, new_proposal):
#     # Using expit to compute exp(w1) / (exp(w0) + exp(w1))
#     p_accept = jax.scipy.special.expit(new_proposal.weight - proposal.weight)
#     do_accept = jax.random.bernoulli(rng_key, p_accept)
#     new_weight = jnp.logaddexp(proposal.weight, new_proposal.weight)
#     new_sum_log_p_accept = jnp.logaddexp(
#         proposal.sum_log_p_accept, new_proposal.sum_log_p_accept
#     )
# 
#     return jax.lax.cond(
#         do_accept,
#         lambda _: Proposal(
#             new_proposal.state,
#             new_proposal.energy,
#             new_weight,
#             new_sum_log_p_accept,
#         ),
#         lambda _: Proposal(
#             proposal.state,
#             proposal.energy,
#             new_weight,
#             new_sum_log_p_accept,
#         ),
#         operand=None,
#     )
# 
# def proposal_generator(
#     kinetic_energy: Callable, divergence_threshold: float
# ) -> Tuple[Callable, Callable]:
#     def new(state: IntegratorState) -> Proposal:
#         energy = state.potential_energy + kinetic_energy(state.momentum)
#         return Proposal(state, energy, 0.0, -np.inf)
# 
#     def update(initial_energy: float, state: IntegratorState) -> Tuple[Proposal, bool]:
#         """Generate a new proposal from a trajectory state.
# 
#         The trajectory state records information about the position in the state
#         space and corresponding potential energy. A proposal also carries a
#         weight that is equal to the difference between the current energy and
#         the previous one. It thus carries information about the previous state
#         as well as the current state.
# 
#         Parameters
#         ----------
#         initial_energy:
#             The initial energy.
#         state:
#             The new state.
# 
#         """
#         new_energy = state.potential_energy + kinetic_energy(state.momentum)
# 
#         delta_energy = initial_energy - new_energy
#         delta_energy = jnp.where(jnp.isnan(delta_energy), -jnp.inf, delta_energy)
#         is_transition_divergent = jnp.abs(delta_energy) > divergence_threshold
# 
#         # The weight of the new proposal is equal to H0 - H(z_new)
#         weight = delta_energy
#         # Acceptance statistic min(e^{H0 - H(z_new)}, 1)
#         sum_log_p_accept = jnp.minimum(delta_energy, 0.0)
# 
#         return (
#             Proposal(
#                 state,
#                 new_energy,
#                 weight,
#                 sum_log_p_accept,
#             ),
#             is_transition_divergent,
#         )
# 
#     return new, update
# 
# Path: blackjax/types.py



# Path: blackjax/mcmc/trajectory.py
from typing import Callable, NamedTuple, Tuple
from blackjax.mcmc.integrators import IntegratorState
from blackjax.mcmc.proposal import (
    Proposal,
    progressive_biased_sampling,
    progressive_uniform_sampling,
    proposal_generator,
)
from blackjax.types import PRNGKey, PyTree
import jax
import jax.numpy as jnp

2. Sampling a new state from this sampled trajectory.

Step (1) ensures that the process is reversible and thus that detailed balance
is respected. The traditional implementation of HMC does not sample a
trajectory, but instead takes a fixed number of steps in the same direction and
flips the momentum of the last state.

We distinguish here between two different methods to sample trajectories: static
and dynamic sampling. In the static setting we sample trajectories with a fixed
number of steps, while in the dynamic setting the total number of steps is
determined by a dynamic termination criterion. Traditional HMC falls in the
former category, NUTS in the latter.

There are also two methods to sample proposals from these trajectories. In the
static setting we first build the trajectory and then sample a proposal from
this trajectory. In the progressive setting we update the proposal as the
trajectory is being sampled. While the former is faster, we risk saturating the
memory by keeping states that will subsequently be discarded.


References
----------
.. [1]: Betancourt, Michael. "A conceptual introduction to Hamiltonian Monte Carlo." arXiv preprint arXiv:1701.02434 (2017).

"""




class Trajectory(NamedTuple):

    leftmost_state: IntegratorState