jcm/sde_lib.py [257:277]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        )

    def sample_importance_weighted_time_for_likelihood(
        self, rng, shape, quantile=None, eps=1e-5, steps=100
    ):
        Z = self.likelihood_importance_cum_weight(self.T, eps=eps)
        if quantile is None:
            quantile = jax.random.uniform(rng, shape, minval=0, maxval=Z)
        lb = jnp.ones_like(quantile) * eps
        ub = jnp.ones_like(quantile) * self.T

        def bisection_func(carry, idx):
            lb, ub = carry
            mid = (lb + ub) / 2.0
            value = self.likelihood_importance_cum_weight(mid, eps=eps)
            lb = jnp.where(value <= quantile, mid, lb)
            ub = jnp.where(value <= quantile, ub, mid)
            return (lb, ub), idx

        (lb, ub), _ = jax.lax.scan(bisection_func, (lb, ub), jnp.arange(0, steps))
        return (lb + ub) / 2.0
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



jcm/sde_lib.py [360:380]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        )

    def sample_importance_weighted_time_for_likelihood(
        self, rng, shape, quantile=None, eps=1e-5, steps=100
    ):
        Z = self.likelihood_importance_cum_weight(self.T, eps=eps)
        if quantile is None:
            quantile = jax.random.uniform(rng, shape, minval=0, maxval=Z)
        lb = jnp.ones_like(quantile) * eps
        ub = jnp.ones_like(quantile) * self.T

        def bisection_func(carry, idx):
            lb, ub = carry
            mid = (lb + ub) / 2.0
            value = self.likelihood_importance_cum_weight(mid, eps=eps)
            lb = jnp.where(value <= quantile, mid, lb)
            ub = jnp.where(value <= quantile, ub, mid)
            return (lb, ub), idx

        (lb, ub), _ = jax.lax.scan(bisection_func, (lb, ub), jnp.arange(0, steps))
        return (lb + ub) / 2.0
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -