jkirkby3 / pymle Goto Github PK

class EulerDensity(TransitionDensity):
    def __init__(self, model: Model1D):
        """
        Class which represents the Euler approximation transition density for a model
        :param model: the SDE model, referenced during calls to the transition density
        """
        super().__init__(model=model)

    def __call__(self,
                 x0: Union[float, np.ndarray],
                 xt: Union[float, np.ndarray],
                 t: float) -> Union[float, np.ndarray]:
        """
        The transition density obtained via Euler expansion
        :param x0: float or array, the current value
        :param xt: float or array, the value to transition to  (must be same dimension as x0)
        :param t: float, the time of observing Xt
        :return: probability (same dimension as x0 and xt)
        """
        sig2t = (self._model.diffusion(x0, t) ** 2) * 2 * t
        mut = x0 + self._model.drift(x0, t) * t
        return np.exp(-(xt - mut) ** 2 / sig2t) / np.sqrt(np.pi * sig2t)

    def eulerdensity(self,
                 x0: Union[float, np.ndarray],
                 xt: Union[float, np.ndarray],
                 t: float,
                 dt: float) -> Union[float, np.ndarray]:
        """
        The transition density obtained via Euler expansion
        :param x0: float or array, the current value
        :param xt: float or array, the value to transition to  (must be same dimension as x0)
        :param t: float, the time of observing Xt
        :param dt: float, the time setps
        :return: probability (same dimension as x0 and xt)
        """
        sig2t = (self.diffusion(x0, t) ** 2) * 2 * dt
        
        mut = x0 + self.drift(x0, t) * dt
        
        # print(np.exp(-(xt - mut) ** 2 / sig2t) / np.sqrt(np.pi * sig2t))
        
        return np.exp(-(xt - mut) ** 2 / sig2t) / np.sqrt(np.pi * sig2t)

Kt = (1 / t) * np.log(1 + self._model.drift(x0, t) * (np.exp(self._model.drift_x(x0, t) * t) - 1) / (
                    x0 * self._model.drift_x(x0, t)))
Vt = sig ** 2 * (np.exp(2 * Kt * t) - 1) / (2 * Kt)
Vt = np.sqrt(Vt)