From 976736c1505d07116e1a0d0ba03e75730295ecb6 Mon Sep 17 00:00:00 2001 From: Dimitra Maoutsa Date: Mon, 27 Jun 2022 16:00:20 +0200 Subject: [PATCH] updated to include periodic domain boundaries --- .../DeterministicParticleFlowControl.py | 3054 +++++++++-------- 1 file changed, 1547 insertions(+), 1507 deletions(-) diff --git a/DeterministicParticleFlowControl/DeterministicParticleFlowControl.py b/DeterministicParticleFlowControl/DeterministicParticleFlowControl.py index d33466e..13bae37 100644 --- a/DeterministicParticleFlowControl/DeterministicParticleFlowControl.py +++ b/DeterministicParticleFlowControl/DeterministicParticleFlowControl.py @@ -1,1507 +1,1547 @@ -# -*- coding: utf-8 -*- - -#Created on Sun Dec 12 00:02:39 2021 - -#@author: maout - - - - -# - -import time -import logging -import numpy as np -try: - import ot - POTTED = True -except ImportError: - POTTED = False - - -try: - import numba - NUMBED = True -except ImportError: - NUMBED = False - - - -from matplotlib import pyplot as plt -try: - import torch - TORCHED = True -except ImportError: - TORCHED = False - - -if __name__ == "DeterministicParticleFlowControl.DeterministicParticleFlowControl": - from .score_estimators.score_function_estimators import score_function_multid_seperate - from .reweighting.optimal_transport_reweighting import reweight_optimal_transport_multidim - if TORCHED: - from .score_estimators.score_function_estimators_pytorch import torched_score_function_multid_seperate_all_dims - #from .utils.utils_pytorch import set_device - -else: - ###this clause is for imports to work properly for the documendation - from score_estimators.score_function_estimators import score_function_multid_seperate - from reweighting.optimal_transport_reweighting import reweight_optimal_transport_multidim - if TORCHED: - from score_estimators.score_function_estimators_pytorch import torched_score_function_multid_seperate_all_dims - #from utils.utils_pytorch import set_device -from duecredit import due, BibTeX - - -if TORCHED: - __all__ = ["DPFC", "torched_DPFC"] -else: - __all__ = ["DPFC"] - -# Use duecredit (duecredit.org) to provide a citation to relevant work to -# be cited. This does nothing, unless the user has duecredit installed, -# And calls this with duecredit (as in `python -m duecredit script.py`): -due.cite(BibTeX(""" - @article{maoutsa2021deterministic, - title={Deterministic particle flows for constraining stochastic nonlinear systems}, - author={Maoutsa, Dimitra and Opper, Manfred}, - journal={arXiv preprint arXiv:2112.05735}, - year={2021} -} - """), - description="A deterministic barticle-based method for stochastic optimal control", - tags=["reference-implementation"], - path='DeterministicParticleFlowControl') - - - - -class DPFC(object): - """ - Deterministic particle flow control top-level class. - - Provides the necessary functions to sample the required probability - flows and estimate the controls. - - Attributes - ---------- - t1 : float - Initial time. - t2: float - end time point. - y1: array_like - initial position. - y2: array_like - terminal position. - f: function, callable - drift function handle. - g: float or array_like - diffusion coefficient or function handle. - N: int - number of particles/trajectories. - M: int - number of sparse points for grad log density estimation. - reweight: boolean - determines if reweighting will follow. - U: function, callable - reweighting function to be employed during reweighting, - dimensions :math:`dim_y1,t \\to 1`. - dens_est: str - - 'nonparametric' : non parametric density estimation (this was - used in the paper) - - TO BE ADDED: - - 'hermit1' : parametic density estimation empoying hermite - polynomials (physiscist's) - - 'hermit2' : parametic density estimation empoying hermite - polynomials (probabilists's) - - 'poly' : parametic density estimation empoying simple polynomials - - 'RBF' : parametric density estimation employing radial basis functions. - kern: str - type of kernel: 'RBF' or 'periodic' (only the 'RBF' was used and gives - robust results. Do not use 'periodic' yet!). - reject: boolean - parameter indicating whether non valid backward trajectories will be rejected. - plotting: boolean - parameter indicating whether bridge statistics will be plotted. - f_true: funtion, callable - in case of Brownian bridge reweighting this is the true forward drift - for simulating the forward dynamics. - brown_bridge: boolean, - determines if the reweighting concearns contstraint or reweighting with - respect to brownian bridge. - deterministic: boolean, - indicates the type of dynamics the particles will follow. - If False the flows are simulated with stochastic path sampling. - - Methods - ------- - forward_sampling_Otto: - Creates samples of the forward flow. - forward sampling(): - Samples the forward flow with stochatic particle trajectories. - f_seperate(x,t): - Drift for the deterministic propagation of partcles that are at time t - in position x. - backward_simulation(): - Sampling the backward density with stochastic particles. - reject_trajectories(): - Rejects backward trajectories that do not end up in the vicinity of the - initial point. - Run only if the instance is attribute "reject" is set to True. - Gives logging.warning messages. - forward_sampling_Otto_true(): - Relevant only when forward sampling happens with Brownian bridge. - """ - - - def __init__(self, t1, t2, y1, y2, f, g, N, M, reweight=False, U=None, dens_est='nonparametric', reject=True, kern='RBF', f_true=None, brown_bridge=False, deterministic=True): - - self.dim = y1.size # dimensionality of the system - self.t1 = t1 - self.t2 = t2 - self.y1 = y1 - self.y2 = y2 - - - ##density estimation stuff - self.kern = kern - if kern == 'periodic': - self.kern = 'RBF' - logging.warning('Please do not use periodic kernel yet!') - logging.warning('For all the numerical experiments RBF was used') - logging.warning('We changed your choice to RBF') - # DRIFT /DIFFUSION - self.f = f - self.g = g #scalar or array - - ### PARTICLE DISCRETISATION - self.N = N - - self.N_sparse = M - - self.dt = 0.001 #((t2-t1)/k) - ### reject unreasonable backward trajectories that do not return - ### to initial condition - self.reject = reject - ### indicator for what type of dynamics the particles follow - self.deterministic = deterministic - - - self.timegrid = np.arange(self.t1, self.t2+self.dt/2, self.dt) - self.k = self.timegrid.size - ### reweighting - self.brown_bridge = brown_bridge - self.reweight = reweight - if self.reweight: - self.U = U - if self.brown_bridge: - #storage for forward trajectories with true drift - self.Ztr = np.zeros((self.dim, self.N, self.k)) - self.f_true = f_true - - - - #storage for forward trajectories - self.Z = np.zeros((self.dim, self.N, self.k)) - #storage for backward trajectories - self.B = np.zeros((self.dim, self.N, self.k)) - self.ln_roD = [] ## storing the estimated forward logarithmic gradients - - - ##the stochastic sampling is provided for comparison - if self.deterministic: - self.forward_sampling_Otto() - ### if a Brownian bridge is used for forward sampling - if self.reweight and self.brown_bridge: - self.forward_sampling_Otto_true() - else: - self.forward_sampling() - ## the backward function selects internally for type of dynamics - self.backward_simulation() - # if self.reject: - # self.reject_trajectories() - - - def forward_sampling(self): - """ - Sampling forward probability flow with stochastic particle dynamics. - If reweighting is required at every time step the particles are - appropriatelly reweighted accordint to function :math:`U(x,t)` - - Returns - ------- - int - Returns 0 to make sure everything runs correctly. - The sampled density is stored in place in the array `self.Z`. - - """ - logging.info('Sampling forward...') - W = np.ones((self.N, 1))/self.N - for ti, tt in enumerate(self.timegrid): - - if ti == 0: - self.Z[0, :, 0] = self.y1[0] - self.Z[1, :, 0] = self.y1[1] - else: - for i in range(self.N): - #self.Z[:,i,:] = sdeint.itoint(self.f, self.g, self.Z[i,0], self.timegrid)[:,0] - self.Z[:, i, ti] = (self.Z[:, i, ti-1] + \ - self.dt* self.f(self.Z[:, i, ti-1]) + \ - (self.g)*np.random.normal(loc=0.0, scale=np.sqrt(self.dt), size=(self.dim,))) - - ###WEIGHT - if self.reweight == True: - if ti > 0: - W[:, 0] = np.exp(self.U(self.Z[:, :, ti])) - W = W/np.sum(W) - - ###REWEIGHT - Tstar = reweight_optimal_transport_multidim(self.Z[:, :, ti].T, W) - - self.Z[:, :, ti] = (self.Z[:, :, ti])@Tstar - - for di in range(self.dim): - self.Z[di, :, -1] = self.y2[di] - logging.info('Forward sampling done!') - return 0 - - - - - ### relevant only when forward trajectories follow brownian brifge - - ###this simulates forward trajectories with true f - def f_seperate_true(self, x, t): - """ - (Relevant only when forward sampling happens with Brownian bridge - reweighting) - Wrapper for the drift function of the deterministic particles with the - actual f (system drift) minus the logarithmic gradient term computed - on current particles positions. - Provided for easy integration, and can be passed to ode integrators. - - Parameters - ---------- - x : 2d-array, - Particle positions (dimension x number of particles). - t : float, - Time t within the [t1,t2] interval. - - Returns - ------- - 2d-array - Returns the deterministic forces required to ntegrate the particle - positions for one time step, - i.e. return :math:`f(x,t)-\\frac{1}{2}\\sigma^2\\nabla \\rho_t(x)`, - evaluated at the current positions x and t. - - """ - dimi, N = x.shape - bnds = np.zeros((dimi, 2)) - for ii in range(dimi): - bnds[ii] = [np.min(x[ii, :]), np.max(x[ii, :])] - #sum_bnds = np.sum(bnds) - - Sxx = np.array([np.random.uniform(low=bnd[0], high=bnd[1], size=(self.N_sparse)) for bnd in bnds]) - gpsi = np.zeros((dimi, N)) - lnthsc = 2*np.std(x, axis=1) - - for ii in range(dimi): - gpsi[ii, :] = score_function_multid_seperate(x.T, Sxx.T, False, C=0.001, which=1, l=lnthsc, which_dim=ii+1, kern=self.kern) - - return self.f_true(x, t)-0.5* self.g**2* gpsi - - - ### effective forward drift - estimated seperatelly for each dimension - #plain GP prior - def f_seperate(self, x, t): - """ - Computes the deterministic forces for the evolution of the deterministic - particles for the current particle positions, - ie. drift minus the logarithmic gradient term. - Is used as a wrapper for evolving the particles, - and can be provided to "any" ODE integrator. - - Parameters - ---------- - x : 2d-array, - Particle positions (dimension x number of particles). - t : float, - Time t within the [t1,t2] interval. - - Returns - ------- - 2d-array - Returns the deterministic forces required to ntegrate the particle - positions for one time step, - i.e. return :math:`f(x,t)-\\frac{1}{2}\\sigma^2\\nabla \\rho_t(x)`, - evaluated at the current positions x and t. - - """ - - - dimi, N = x.shape - ### detect min and max of forward flow for each dimension - ### we want to know the state space volume of the forward flow - bnds = np.zeros((dimi, 2)) - for ii in range(dimi): - bnds[ii] = [np.min(x[ii, :]), np.max(x[ii, :])] - sum_bnds = np.sum(bnds) ##this is for detecting if sth goes wrong i.e. trajectories explode - if np.isnan(sum_bnds) or np.isinf(sum_bnds): - ##if we get unreasoble bounds just plot the first 2 dimensions of the trajectories - plt.figure(figsize=(6, 4)), plt.plot(self.Z[0].T, self.Z[1].T, alpha=0.3) - plt.show() - - ##these are the inducing points - ## here we select them from a uniform distribution within the state space volume spanned from the forward flow - Sxx = np.array([np.random.uniform(low=bnd[0], high=bnd[1], size=(self.N_sparse)) for bnd in bnds]) - gpsi = np.zeros((dimi, N)) - lnthsc = 2*np.std(x, axis=1) - for ii in range(dimi): - gpsi[ii, :] = score_function_multid_seperate(x.T, Sxx.T, False, C=0.001, which=1, l=lnthsc, which_dim=ii+1, kern=self.kern) - - return self.f(x, t)-0.5* self.g**2* gpsi - - ###same as forward sampling but without reweighting - this is for bridge reweighting - ### not for constraint reweighting - def forward_sampling_Otto_true(self): - """ - (Relevant only when forward sampling happens with Brownian bridge - reweighting) - Same as forward sampling but without reweighting. - - Returns - ------- - int - Returns 0 to make sure everything runs correctly. - The sampled density is stored in place in the array `self.Ztr`. - - See also - --------- - DPFC.forward_sampling, DPFC.forward_sampling_Otto - - """ - logging.info('Sampling forward with deterministic particles and true drift...') - #W = np.ones((self.N,1))/self.N - for ti, tt in enumerate(self.timegrid): - - if ti == 0: - for di in range(self.dim): - self.Ztr[di, :, 0] = self.y1[di] - - elif ti == 1: #propagate one step with stochastic to avoid the delta function - #substract dt because I want the time at t-1 - self.Ztr[:, :, ti] = (self.Ztr[:, :, ti-1] + self.dt*self.f_true(self.Ztr[:, :, ti-1], tt-self.dt)+\ - (self.g)*np.random.normal(loc=0.0, scale=np.sqrt(self.dt), size=(self.dim, self.N))) - else: - self.Ztr[:, :, ti] = (self.Ztr[:, :, ti-1] + self.dt* self.f_seperate_true(self.Ztr[:, :, ti-1], tt-self.dt)) - - logging.info('Forward sampling with Otto true is ready!') - return 0 - - - - def forward_sampling_Otto(self): - """ - Samples the forward probability flow with deterministic particle - dynamics. - If required at every timestep a particle reweighting takes place - employing the weights obtained from the exponentiated path constraint - :math:`U(x,t)` - - Returns - ------- - int - Returns 0 to make sure everything runs correctly. - The sampled density is stored in place in the array `self.Z`. - - """ - logging.info('Sampling forward with deterministic particles...') - W = np.ones((self.N, 1))/self.N - for ti, tt in enumerate(self.timegrid): - if ti == 0: - for di in range(self.dim): - self.Z[di, :, 0] = self.y1[di] - if self.brown_bridge: - self.Z[di, :, -1] = self.y2[di] - ## we start forward trajectories for a delta function. - ##in principle we could start from an arbitrary distribution - ##if you want to start from a normal uncomen the following and comment the above initialisation for y1 - #self.Z[di,:,0] = np.random.normal(self.y1[di], 0.05, self.N) - elif ti == 1: #propagate one step with stochastic to avoid the delta function - #substract dt because I want the time at t-1 - self.Z[:, :, ti] = (self.Z[:, :, ti-1] + self.dt*self.f(self.Z[:, :, ti-1], tt-self.dt)+\ - (self.g)*np.random.normal(loc=0.0, scale=np.sqrt(self.dt), size=(self.dim, self.N))) - else: - self.Z[:, :, ti] = (self.Z[:, :, ti-1] + self.dt* self.f_seperate(self.Z[:, :, ti-1], tt-self.dt)) - ###REWEIGHT - if self.reweight == True: - if ti > 0: - - W[:, 0] = np.exp(self.U(self.Z[:, :, ti], tt)*self.dt) #-1 - W = W/np.sum(W) - - ###REWEIGHT - start = time.time() - Tstar = reweight_optimal_transport_multidim(self.Z[:, :, ti].T, W) - #print(Tstar) - if ti == 3: - stop = time.time() - logging.info('Timepoint: %d needed '%ti, stop-start) - self.Z[:, :, ti] = ((self.Z[:, :, ti])@Tstar) ##### - logging.info('Forward sampling with Otto is ready!') - return 0 - - def density_estimation(self, ti, rev_ti): - rev_t = rev_ti - grad_ln_ro = np.zeros((self.dim, self.N)) - lnthsc = 2*np.std(self.Z[:, :, rev_t], axis=1) - bnds = np.zeros((self.dim, 2)) - for ii in range(self.dim): - bnds[ii] = [max(np.min(self.Z[ii, :, rev_t]), np.min(self.B[ii, :, rev_ti])), min(np.max(self.Z[ii, :, rev_t]), np.max(self.B[ii, :, rev_ti]))] - sum_bnds = np.sum(bnds) - if np.isnan(sum_bnds) or np.isinf(sum_bnds): - plt.figure(figsize=(6, 4)), plt.plot(self.B[0].T, self.B[1].T, alpha=0.3) - plt.plot(self.y1[0], self.y1[1], 'go') - plt.show() - #sparse points - Sxx = np.array([np.random.uniform(low=bnd[0], high=bnd[1], size=(self.N_sparse)) for bnd in bnds]) - - for di in range(self.dim): - #estimate density from forward (Z) and evaluate at current postitions of backward particles (B) - grad_ln_ro[di, :] = score_function_multid_seperate(self.Z[:, :, rev_t].T, Sxx.T, func_out=True, C=0.001, which=1, l=lnthsc, which_dim=di+1, kern=self.kern)(self.B[:, :, rev_ti].T) - - - return grad_ln_ro - - - def bw_density_estimation(self, rev_ti): - """ - Estimates the logaritmic gradient of the backward flow evaluated at - particle positions of the backward flow. - - - Parameters - ---------- - - rev_ti : int, - indicates the time point in the timegrid where the estimation - will take place, i.e. for time t=self.timegrid[rev_ti. - - Returns - ------- - grad_ln_b: 2d-array, - with the logarithmic gradients of the time reversed - (backward) flow (dim x N) for the timestep `rev_ti`. - - """ - grad_ln_b = np.zeros((self.dim, self.N)) - lnthsc = 2*np.std(self.B[:, :, rev_ti], axis=1) - #print(ti, rev_ti, rev_ti-1) - bnds = np.zeros((self.dim, 2)) - for ii in range(self.dim): - bnds[ii] = [max(np.min(self.Z[ii, :, rev_ti]), np.min(self.B[ii, :, rev_ti])), min(np.max(self.Z[ii, :, rev_ti]), np.max(self.B[ii, :, rev_ti]))] - #sparse points - Sxx = np.array([np.random.uniform(low=bnd[0], high=bnd[1], size=(self.N_sparse)) for bnd in bnds]) - - for di in range(self.dim): - grad_ln_b[di, :] = score_function_multid_seperate(self.B[:, :, rev_ti].T, Sxx.T, func_out=False, C=0.001, which=1, l=lnthsc, which_dim=di+1, kern=self.kern) - # grad_ln_a = score_function_multid_seperate2(self.B[:, :, rev_ti].T, Sxx.T, func_out=False, C=0.001, which=1, l=lnthsc, which_dim=di+1, kern=self.kern) - # #np.testing.assert_array_equal(grad_ln_b[di, :], grad_ln_a) - # np.testing.assert_allclose(grad_ln_b[di, :], grad_ln_a) - return grad_ln_b - - - def backward_simulation(self): - """ - Sample time reversed flow with deterministic dynamics (or stochastic if - `self.deterministic == False`). - Trajectories are stored in place in `self.B` array of dimensionality - (dim x N x timegrid.size). - `self.B` does not require a timereversion at the end, everything - is stored in the correct order. - - Returns - ------- - int - Returns 0 to ensure everything was executed correctly. - - """ - - for ti, tt in enumerate(self.timegrid[:-1]): - if ti == 0: - for di in range(self.dim): - self.B[di, :, -1] = self.y2[di] - else: - - rev_ti = self.k -ti-1 - #density estimation of forward particles - grad_ln_ro = self.density_estimation(ti, rev_ti+1) - - if (ti == 1 and self.deterministic) or (not self.deterministic): - - self.B[:, :, rev_ti] = (self.B[:, :, rev_ti+1] -\ - self.f(self.B[:, :, rev_ti+1], self.timegrid[rev_ti+1])*self.dt + \ - self.dt*self.g**2*grad_ln_ro +\ - (self.g)*np.random.normal(loc=0.0, scale=np.sqrt(self.dt), size=(self.dim, self.N))) - else: - grad_ln_b = self.bw_density_estimation(rev_ti+1) - self.B[:, :, rev_ti] = (self.B[:, :, rev_ti+1] -\ - (self.f(self.B[:, :, rev_ti+1], self.timegrid[rev_ti+1])- self.g**2*grad_ln_ro +0.5*self.g**2*grad_ln_b)*self.dt) - - for di in range(self.dim): - self.B[di, :, 0] = self.y1[di] - return 0 - - - """ - def reject_trajectories(self): - - Reject backward trajectories that do not reach the vicinity of the - initial point. - Deletes in place relevant rows of the `self.B` array that contains - the time reversed trajectories. - - Returns - ------- - int - Returns 0. - - - fplus = self.y1+self.f(self.y1, self.t1)*self.dt+6*self.g**2 *np.sqrt(self.dt) - fminus = self.y1+self.f(self.y1, self.t1) *self.dt-6*self.g**2 *np.sqrt(self.dt) - reverse_order = np.zeros(self.dim) - #this is an indicator if along one of the dimensions fplus - #is smaller than fminus - for iii in range(self.dim): - if fplus[iii] < fminus[iii]: - reverse_order[iii] = 1 - to_delete = np.zeros(self.N) - ##these will be one if ith trajectory is out of bounds - - ## checking if out of bounds for each dim - for iii in range(self.dim): - if reverse_order[iii] == 0: - to_delete += np.logical_not(np.logical_and(self.B[iii, :, 1] < fplus[iii], self.B[iii, :, 1] > fminus[iii])) - elif reverse_order[iii] == 1: - to_delete += np.logical_not(np.logical_and(self.B[iii, :, 1] > fplus[iii], self.B[iii, :, 1] < fminus[iii])) - - sinx = np.where(to_delete >= 1)[0] - #sinx = np.where(np.logical_or(np.logical_not(np.logical_and(self.B[0, :, 1] < fplus[0], self.B[0, :, 1] > fminus[0])), np.logical_not(np.logical_and(self.B[0, :, 1] < fplus[0], self.B[0, :, 1] > fminus[0]))))[0] - #((self.B[1,:,-2]fminus[1]) ) ))[0] - - - #logging.warning("Identified %d invalid bridge trajectories "%len(sinx)) - # if self.reject: - # logging.warning("Deleting invalid trajectories...") - # sinx = sinx[::-1] - # for element in sinx: - # self.B = np.delete(self.B, element, axis=1) - return 0 - """ - - - def calculate_u(self, grid_x, ti): - """ - Computes the control at position(s) grid_x at timestep ti - (i.e. at time self.timegrid[ti]). - - Parameters - ---------- - grid_x : ndarray, - size d x number of points to be evaluated. - ti : int, - time index in timegrid for the computation of u. - - - Returns - ------- - u_t: ndarray, - same size as grid_x. These are the controls u(grid_x, t), - where t=self.timegrid[ti]. - - """ - #a = 0.001 - #grad_dirac = lambda x,di: - 2*(x[di] -self.y2[di])* - #np.exp(- (1/a**2)* (x[0]- self.y2[0])**2)/(a**3 *np.sqrt(np.pi)) - u_t = np.zeros(grid_x.T.shape) - - - lnthsc1 = 2*np.std(self.B[:, :, ti], axis=1) - lnthsc2 = 2*np.std(self.Z[:, :, ti], axis=1) - - - bnds = np.zeros((self.dim, 2)) - for ii in range(self.dim): - if self.reweight == False or self.brown_bridge == False: - bnds[ii] = [max(np.min(self.Z[ii, :, ti]), np.min(self.B[ii, :, ti])), min(np.max(self.Z[ii, :, ti]), np.max(self.B[ii, :, ti]))] - else: - bnds[ii] = [max(np.min(self.Ztr[ii, :, ti]), np.min(self.B[ii, :, ti])), min(np.max(self.Ztr[ii, :, ti]), np.max(self.B[ii, :, ti]))] - - if ti <= 5 or (ti >= self.k-5): - if self.reweight == False or self.brown_bridge == False: - ##for the first and last 5 timesteps, to avoid numerical singularities just assume gaussian densities - for di in range(self.dim): - mutb = np.mean(self.B[di, :, ti]) - stdtb = np.std(self.B[di, :, ti]) - mutz = np.mean(self.Z[di, :, ti]) - stdtz = np.std(self.Z[di, :, ti]) - u_t[di] = -(grid_x[:, di]- mutb)/stdtb**2 - (-(grid_x[:, di]- mutz)/stdtz**2) - elif self.reweight == True and self.brown_bridge == True: - for di in range(self.dim): - mutb = np.mean(self.B[di, :, ti]) - stdtb = np.std(self.B[di, :, ti]) - mutz = np.mean(self.Ztr[di, :, ti]) - stdtz = np.std(self.Ztr[di, :, ti]) - u_t[di] = -(grid_x[:, di]- mutb)/stdtb**2 - (-(grid_x[:, di]- mutz)/stdtz**2) - else: #if ti > 5: - ### clipping not used at the end but provided here for cases when - ### number of particles is small - ### and trajectories fall out of simulated flows - ### TO DO: add clipping as an option to be selected when - ### initialising the function - ### if point for evaluating control falls out of the region where we - ### have points, clip the points to - ### fall within the calculated region - we do not change the - ### position of the point, only the control value will be - ### calculated with clipped positions - bndsb = np.zeros((self.dim, 2)) - bndsz = np.zeros((self.dim, 2)) - for di in range(self.dim): - bndsb[di] = [np.min(self.B[di, :, ti]), np.max(self.B[di, :, ti])] - bndsz[di] = [np.min(self.Z[di, :, ti]), np.max(self.Z[di, :, ti])] - - ## clipping not used at the end! - ###cliping the values of points when evaluating the grad log p - grid_b = grid_x#np.clip(grid_x, bndsb[0], bndsb[1]) - grid_z = grid_x#np.clip(grid_x, bndsz[0], bndsz[1]) - - Sxx = np.array([np.random.uniform(low=bnd[0], high=bnd[1], size=(self.N_sparse)) for bnd in bnds]) - for di in range(self.dim): - score_Bw = score_function_multid_seperate(self.B[:, :, ti].T, Sxx.T, func_out=True, C=0.001, which=1, l=lnthsc1, which_dim=di+1, kern=self.kern)(grid_b) - if self.reweight == False or self.brown_bridge == False: - score_Fw = score_function_multid_seperate(self.Z[:, :, ti].T, Sxx.T, func_out=True, C=0.001, which=1, l=lnthsc2, which_dim=di+1, kern=self.kern)(grid_z) - else: - bndsztr = np.zeros((self.dim, 2)) - for ii in range(self.dim): - bndsztr[di] = [np.min(self.Ztr[di, :, ti]), np.max(self.Ztr[di, :, ti])] - grid_ztr = np.clip(grid_x, bndsztr[0], bndsztr[1]) - lnthsc3 = 2*np.std(self.Ztr[:, :, ti], axis=1) - score_Fw = score_function_multid_seperate(self.Ztr[:, :, ti].T, Sxx.T, func_out=True, C=0.001, which=1, l=lnthsc3, which_dim=di+1, kern=self.kern)(grid_ztr) - - u_t[di] = score_Bw - score_Fw - # for di in range(self.dim): - # u_t[di] = score_function_multid_seperate(self.B[:,:,ti].T,Sxx.T,func_out= True,C=0.001,which=1,l=lnthsc,which_dim=di+1, kern=self.kern)(grid_x.T) \ - # - score_function_multid_seperate(self.Z[:,:,ti].T,Sxx.T,func_out= True,C=0.001,which=1,l=lnthsc,which_dim=di+1, kern=self.kern)(grid_x.T) - - - return u_t - - - def check_if_covered(self, X, ti): - """ - Checks if test point X falls within forward and backward densities at - timepoint timegrid[ti]. - - - Parameters - ---------- - X : array 1x dim or Kxdim - Point in state space where control is evaluated. - ti : int - Index in timegrid array indicating the time within the - time interval [t1,t2]. - - Returns - ------- - Boolean variable - True if the text point X falls within the densities. - - """ - covered = True - bnds = np.zeros((self.dim, 2)) - for ii in range(self.dim): - bnds[ii] = [max(np.min(self.Z[ii, :, ti]), np.min(self.B[ii, :, ti])), min(np.max(self.Z[ii, :, ti]), np.max(self.B[ii, :, ti]))] - #bnds[ii] = [np.min(self.B[ii,:,ti]),np.max(self.B[ii,:,ti])] - - covered = covered * ((X[ii] >= bnds[ii][0]) and (X[ii] <= bnds[ii][1])) - - return covered - -#%% - -class torched_DPFC(object): - """ - Deterministic particle flow control top-level class implemented in pytorch. - - Provides the necessary functions to sample the required probability - flows and estimate the controls. - - Attributes - ---------- - t1 : float - Initial time. - t2: float - end time point. - y1: array_like - initial position. - y2: array_like - terminal position. - f: function, callable - drift function handle. - g: float or array_like - diffusion coefficient or function handle. - N: int - number of particles/trajectories. - M: int - number of sparse points for grad log density estimation. - reweight: boolean - determines if reweighting will follow. - U: function, callable - reweighting function to be employed during reweighting, - dimensions :math:`dim_y1,t \\to 1`. - dens_est: str - - 'nonparametric' : non parametric density estimation (this was - used in the paper) - - TO BE ADDED: - - 'hermit1' : parametic density estimation empoying hermite - polynomials (physiscist's) - - 'hermit2' : parametic density estimation empoying hermite - polynomials (probabilists's) - - 'poly' : parametic density estimation empoying simple polynomials - - 'RBF' : parametric density estimation employing radial basis functions. - kern: str - type of kernel: 'RBF' or 'periodic' (only the 'RBF' was used and gives - robust results. Do not use 'periodic' yet!). - reject: boolean - parameter indicating whether non valid backward trajectories will be rejected. - plotting: boolean - parameter indicating whether bridge statistics will be plotted. - f_true: funtion, callable - in case of Brownian bridge reweighting this is the true forward drift - for simulating the forward dynamics. - brown_bridge: boolean, - determines if the reweighting concearns contstraint or reweighting with - respect to brownian bridge. - deterministic: boolean, - indicates the type of dynamics the particles will follow. - If False the flows are simulated with stochastic path sampling. - device: string, - indicates the device where computations will be exacuted. - `cpu` or `gpu/cuda` or `tpu` if available. - - Methods - ------- - forward_sampling_Otto: - Creates samples of the forward flow. - forward sampling(): - Samples the forward flow with stochatic particle trajectories. - f_seperate(x,t): - Drift for the deterministic propagation of partcles that are at time t - in position x. - backward_simulation(): - Sampling the backward density with stochastic particles. - reject_trajectories(): - Rejects backward trajectories that do not end up in the vicinity of the - initial point. - Run only if the instance is attribute "reject" is set to True. - Gives logging.warning messages. - forward_sampling_Otto_true(): - Relevant only when forward sampling happens with Brownian bridge. - """ - - - def __init__(self, t1, t2, y1, y2, f, g, N, M, reweight=False, U=None, - dens_est='nonparametric', reject=True, kern='RBF', - f_true=None, brown_bridge=False, deterministic=True, - device=None): - - self.device = device - # dimensionality of the system - self.dim = torch.tensor(y1.size(dim=0), dtype=torch.int, device=self.device) # dimensionality of the system - self.t1 = t1 - self.t2 = t2 - if not torch.is_tensor(y1): - self.y1 = torch.tensor(y1, dtype=torch.float64, device=self.device) - self.y2 = torch.tensor(y2, dtype=torch.float64, device=self.device) - else: - self.y1 = y1 - self.y2 = y2 - - - ##density estimation stuff - self.kern = kern - if kern == 'periodic': - self.kern = 'RBF' - logging.warning('Please do not use periodic kernel yet!') - logging.warning('For all the numerical experiments RBF was used') - logging.warning('We changed your choice to RBF') - # DRIFT /DIFFUSION - self.f = f - self.g = torch.tensor(g, dtype=torch.float64, device=self.device) #scalar or array - - ### PARTICLE DISCRETISATION - self.N = torch.tensor(N, dtype=torch.int, device=self.device) - - self.N_sparse = torch.tensor(M, dtype=torch.int, device=self.device)#M - - self.dt = torch.tensor(0.001, dtype=torch.float64, device=self.device) - #((t2-t1)/k) - - ### reject unreasonable backward trajectories that do not return - ### to initial condition - self.reject = reject - ### indicator for what type of dynamics the particles follow - self.deterministic = deterministic - - - self.timegrid = torch.arange(self.t1, self.t2+self.dt/2, self.dt, - dtype=torch.float64, device=self.device) - self.k = torch.tensor(self.timegrid.size(dim=0), dtype=torch.int, device=self.device) - ### reweighting - self.brown_bridge = brown_bridge - self.reweight = reweight - if self.reweight: - self.U = U - if self.brown_bridge: - #storage for forward trajectories with true drift - self.Ztr = torch.zeros(self.dim, self.N, self.k, - dtype=torch.float64, device=self.device) - self.f_true = f_true - - - - #storage for forward trajectories - self.Z = torch.zeros(self.dim, self.N, self.k, dtype=torch.float64, - device=self.device) - #storage for backward trajectories - self.B = torch.zeros(self.dim, self.N, self.k, dtype=torch.float64, - device=self.device ) - self.ln_roD = [] ## storing the estimated forward logarithmic gradients - - - ##the stochastic sampling is provided for comparison - if self.deterministic: - self.forward_sampling_Otto() - - ### if a Brownian bridge is used for forward sampling - if self.reweight and self.brown_bridge: - self.forward_sampling_Otto_true() - else: - self.forward_sampling() - - ## the backward function selects internally for type of dynamics - self.backward_simulation() - # if self.reject: - # self.reject_trajectories() - - - def forward_sampling(self): - """ - Sampling forward probability flow with stochastic particle dynamics. - If reweighting is required at every time step the particles are - appropriatelly reweighted accordint to function :math:`U(x,t)` - - Returns - ------- - int - Returns 0 to make sure everything runs correctly. - The sampled density is stored in place in the array `self.Z`. - - """ - logging.info('Sampling forward...') - W = torch.ones(self.N, 1, dtype=torch.float64, device=self.device)/self.N - - for ti, tt in enumerate(self.timegrid): - if ti == 0: - self.Z[0, :, 0] = self.y1[0] - self.Z[1, :, 0] = self.y1[1] - else: - self.Z[:, :, ti] = self.Z[:, :, ti-1] + \ - self.dt*self.f(self.Z[:, :, ti-1], tt-self.dt)+\ - (self.g)*\ - torch.empty([self.dim, self.N]).normal_(mean=0, std=torch.sqrt(self.dt)) - ###WEIGHT - if self.reweight == True: - if ti > 0: - W[:, 0] = torch.exp(self.U(self.Z[:, :, ti])) - W = W/torch.sum(W) - - ###REWEIGHT with pot TO DO: - #Tstar = reweight_optimal_transport_multidim(self.Z[:, :, ti].T, W) - M = ot.dist(self.Z[:,:,ti].T, self.Z[:,:,ti].T) - M /= M.max() - a = W[:,0] - b = torch.ones_like(W[:,0], dtype=torch.float64, - device=self.device)/self.N - T2 = ot.emd(a, b, M) - self.Z[:, :, ti] = (self.N*self.Z[:, :, ti])@T2 - - #for di in range(self.dim): - #self.Z[di, :, -1] = self.y2[di] - logging.info('Forward sampling done!') - return 0 - - - - - ### relevant only when forward trajectories follow brownian brifge - - ###this simulates forward trajectories with true f - def f_seperate_true(self, x, t): - """ - (Relevant only when forward sampling happens with Brownian bridge - reweighting) - Wrapper for the drift function of the deterministic particles with the - actual f (system drift) minus the logarithmic gradient term computed - on current particles positions. - Provided for easy integration, and can be passed to ode integrators. - - Parameters - ---------- - x : 2d-array, - Particle positions (dimension x number of particles). - t : float, - Time t within the [t1,t2] interval. - - Returns - ------- - 2d-array - Returns the deterministic forces required to ntegrate the particle - positions for one time step, - i.e. return :math:`f(x,t)-\\frac{1}{2}\\sigma^2\\nabla \\rho_t(x)`, - evaluated at the current positions x and t. - - """ - dimi, N = x.shape - bnds = torch.zeros(dimi, 2, dtype=torch.float64, device=self.device) - for ii in range(dimi): - bnds[ii, 0] = torch.min(x[ii, :]) - bnds[ii, 1] = torch.max(x[ii, :]) - - Sxx = torch.distributions.Uniform(low=bnds[:, 0], high=bnds[:, 1]).sample([self.N_sparse]) - - - #gpsi = torch.zeros(dimi, N, dtype=torch.float64, device=self.device) - lnthsc = 2*torch.std(x, dim=1) - - - gpsi = torched_score_function_multid_seperate_all_dims(torch.t(x), - torch.t(Sxx), - func_out=False, - C=0.001, - l=lnthsc, - kern=self.kern, - device=self.device) - - return self.f_true(x, t)-0.5* self.g**2* gpsi - - - ### effective forward drift - estimated seperatelly for each dimension - #plain GP prior - def f_seperate(self, x, t): - """ - Computes the deterministic forces for the evolution of the deterministic - particles for the current particle positions, - ie. drift minus the logarithmic gradient term. - Is used as a wrapper for evolving the particles, - and can be provided to "any" ODE integrator. - - Parameters - ---------- - x : 2d-array, - Particle positions (dimension x number of particles). - t : float, - Time t within the [t1,t2] interval. - - Returns - ------- - 2d-array - Returns the deterministic forces required to ntegrate the particle - positions for one time step, - i.e. return :math:`f(x,t)-\\frac{1}{2}\\sigma^2\\nabla \\rho_t(x)`, - evaluated at the current positions x and t. - - - """ - - - dimi, N = x.shape - ### detect min and max of forward flow for each dimension - ### we want to know the state space volume of the forward flow - bnds = torch.zeros(dimi, 2, dtype=torch.float64, device=self.device) - - for ii in range(dimi): - bnds[ii, 0] = torch.min(x[ii, :]) - bnds[ii, 1] = torch.max(x[ii, :]) - # sum_bnds = np.sum(bnds) ##this is for detecting if sth goes wrong i.e. trajectories explode - # if np.isnan(sum_bnds) or np.isinf(sum_bnds): - # ##if we get unreasoble bounds just plot the first 2 dimensions of the trajectories - # plt.figure(figsize=(6, 4)), plt.plot(self.Z[0].T, self.Z[1].T, alpha=0.3) - # plt.show() - - ##these are the inducing points - ## here we select them from a uniform distribution within the state space volume spanned from the forward flow - Sxx = torch.distributions.Uniform(low=bnds[:, 0], high=bnds[:, 1]).sample([self.N_sparse]) - #print(Sxx.size()) - #print(torch.t(x).size()) - #gpsi = np.zeros((dimi, N)) - #print(Sxx.size()) - lnthsc = 2*torch.std(x, dim=1) - - gpsi = torched_score_function_multid_seperate_all_dims(torch.t(x), - Sxx, l=lnthsc, - func_out=False, - device=self.device) - - - - return self.f(x, t)-0.5* self.g**2* torch.t(gpsi) - - ###same as forward sampling but without reweighting - this is for bridge reweighting - ### not for constraint reweighting - def forward_sampling_Otto_true(self): - """ - (Relevant only when forward sampling happens with Brownian bridge - reweighting) - Same as forward sampling but without reweighting. - - Returns - ------- - int - Returns 0 to make sure everything runs correctly. - The sampled density is stored in place in the array `self.Ztr`. - - See also - --------- - DPFC.forward_sampling, DPFC.forward_sampling_Otto - - """ - logging.info('Sampling forward with deterministic particles and true drift...') - #W = np.ones((self.N,1))/self.N - for ti, tt in enumerate(self.timegrid): - - if ti == 0: - for di in range(self.dim): - self.Ztr[di, :, 0] = self.y1[di] - - elif ti == 1: #propagate one step with stochastic to avoid the delta function - #substract dt because I want the time at t-1 - self.Ztr[:, :, ti] = self.Ztr[:, :, ti-1] + \ - self.dt*self.f_true(self.Ztr[:, :, ti-1], tt-self.dt)+\ - (self.g)*torch.empty([self.dim, self.N]).normal_(mean=0, std=torch.sqrt(self.dt)) - else: - self.Ztr[:, :, ti] = self.Ztr[:, :, ti-1] + \ - self.dt* self.f_seperate_true(self.Ztr[:, :, ti-1], tt-self.dt) - - logging.info('Forward sampling with Otto true is ready!') - return 0 - - - - def forward_sampling_Otto(self): - """ - Samples the forward probability flow with deterministic particle - dynamics. - If required at every timestep a particle reweighting takes place - employing the weights obtained from the exponentiated path constraint - :math:`U(x,t)` - - Returns - ------- - int - Returns 0 to make sure everything runs correctly. - The sampled density is stored in place in the array `self.Z`. - - """ - logging.info('Sampling forward with deterministic particles...') - W = torch.ones(self.N, 1, dtype=torch.float64, device=self.device)/self.N - for ti, tt in enumerate(self.timegrid): - if ti == 0: - for di in range(self.dim): - self.Z[di, :, 0] = self.y1[di] - if self.brown_bridge: - self.Z[di, :, -1] = self.y2[di] - ## we start forward trajectories for a delta function. - ##in principle we could start from an arbitrary distribution - ##if you want to start from a normal uncomen the following and comment the above initialisation for y1 - #self.Z[di,:,0] = np.random.normal(self.y1[di], 0.05, self.N) - elif ti == 1: #propagate one step with stochastic to avoid the delta function - #substract dt because I want the time at t-1 - self.Z[:, :, ti] = self.Z[:, :, ti-1] + \ - self.dt*self.f(self.Z[:, :, ti-1], tt-self.dt)+\ - (self.g)*\ - torch.empty([self.dim, self.N], - dtype=torch.float64, - device=self.device).normal_(mean=torch.tensor(0, - dtype=torch.float64, - device=self.device), - std=torch.sqrt(self.dt)) - else: - self.Z[:, :, ti] = self.Z[:, :, ti-1] +\ - self.dt* self.f_seperate(self.Z[:, :, ti-1], tt-self.dt) - ###REWEIGHT - if self.reweight == True: - if ti > 0: - - W[:, 0] = torch.exp(self.U(self.Z[:, :, ti], tt)*self.dt) #-1 - W = W/torch.sum(W) - - ###REWEIGHT - start = time.time() - M = ot.dist(self.Z[:,:,ti].T, self.Z[:,:,ti].T) - M /= M.max() - a = W[:,0] - b = torch.ones_like(W[:,0], dtype=torch.float64, - device=self.device)/self.N - T2 = ot.emd(a, b, M) - self.Z[:, :, ti] = (self.N*self.Z[:, :, ti])@T2 - #Tstar = reweight_optimal_transport_multidim(self.Z[:, :, ti].T, W) - #print(Tstar) - if ti == 3: - stop = time.time() - logging.info('Timepoint: %d needed '%ti, stop-start) - - logging.info('Forward sampling with Otto is ready!') - return 0 - - def density_estimation(self, ti, rev_ti): - #print(ti) - rev_t = rev_ti - grad_ln_ro = torch.zeros(self.dim, self.N, dtype=torch.float64, - device=self.device) - lnthsc = 2*torch.std(self.Z[:, :, rev_t], dim=1) - bnds = torch.zeros(self.dim, 2, dtype=torch.float64, device=self.device) - for ii in range(self.dim): - if torch.min(self.B[ii, :, rev_ti]) == torch.max(self.B[ii, :, rev_ti]): - bnds[ii, 0] = torch.min(self.Z[ii, :, rev_t]) - bnds[ii, 1] = torch.max(self.Z[ii, :, rev_t]) - else: - bnds[ii, 0] = min(torch.min(self.Z[ii, :, rev_t]), torch.min(self.B[ii, :, rev_ti])) - bnds[ii, 1] = max(torch.max(self.Z[ii, :, rev_t]), torch.max(self.B[ii, :, rev_ti])) - - #sparse points - """ - print('low', bnds[:,0]) - print('high', bnds[:,1]) - plt.figure() - plt.subplot(1,2,1) - plt.plot(self.Z[0, :, rev_t].detach().numpy(), '.', c='grey') - plt.plot(self.B[0, :, rev_t].detach().numpy(), '.', c='maroon') - plt.subplot(1,2,2) - plt.plot(self.Z[1, :, rev_t].detach().numpy(), '.', c='grey') - plt.plot(self.B[1, :, rev_t].detach().numpy(), '.', c='maroon') - plt.show() - """ - Sxx = torch.distributions.Uniform(low=bnds[:, 0], high=bnds[:, 1]).sample([self.N_sparse]) - - #estimate density from forward (Z) and evaluate at current postitions of backward particles (B) - #grad_ln_ro = score_function_multid_seperate(self.Z[:, :, rev_t].T, Sxx.T, func_out=True, C=0.001, which=1, l=lnthsc, which_dim=di+1, kern=self.kern)(self.B[:, :, rev_ti].T) - grad_ln_ro = torch.t(torched_score_function_multid_seperate_all_dims(torch.t(self.Z[:, :, rev_t]), - Sxx, - func_out=True, - C=0.001, - l=lnthsc, - kern=self.kern, - device=self.device)(torch.t(self.B[:, :, rev_ti])) ) - - return grad_ln_ro - - - def bw_density_estimation(self, rev_ti): - """ - Estimates the logaritmic gradient of the backward flow evaluated at - particle positions of the backward flow. - - - Parameters - ---------- - - rev_ti : int, - indicates the time point in the timegrid where the estimation - will take place, i.e. for time t=self.timegrid[rev_ti. - - Returns - ------- - grad_ln_b: 2d-array, - with the logarithmic gradients of the time reversed - (backward) flow (dim x N) for the timestep `rev_ti`. - - """ - grad_ln_b = torch.zeros(self.dim, self.N, dtype=torch.float64, - device=self.device) - lnthsc = 2*torch.std(self.B[:, :, rev_ti], dim=1) - - bnds = torch.zeros(self.dim, 2, dtype=torch.float64, device=self.device) - for ii in range(self.dim): - bnds[ii, 0] = min(torch.min(self.Z[ii, :, rev_ti]), torch.min(self.B[ii, :, rev_ti])) - bnds[ii, 1] = max(torch.max(self.Z[ii, :, rev_ti]), torch.max(self.B[ii, :, rev_ti])) - - #sparse points - Sxx = torch.distributions.Uniform(low=bnds[:, 0], high=bnds[:, 1]).sample([self.N_sparse]) - - - - grad_ln_b = torched_score_function_multid_seperate_all_dims(torch.t(self.B[:, :, rev_ti]), - Sxx, - func_out=False, - C=0.001, - l=lnthsc, - kern=self.kern, - device=self.device) - return grad_ln_b - - - def backward_simulation(self): - """ - Sample time reversed flow with deterministic dynamics (or stochastic if - `self.deterministic == False`). - Trajectories are stored in place in `self.B` array of dimensionality - (dim x N x timegrid.size). - `self.B` does not require a timereversion at the end, everything - is stored in the correct order. - - Returns - ------- - int - Returns 0 to ensure everything was executed correctly. - - """ - - for ti, tt in enumerate(self.timegrid[:-1]): - if ti == 0: - for di in range(self.dim): - self.B[di, :, -1] = self.y2[di] - else: - - rev_ti = self.k -ti-1 - #density estimation of forward particles - - grad_ln_ro = self.density_estimation(ti, rev_ti+1) - - if (ti == 1 and self.deterministic) or (not self.deterministic): - - self.B[:, :, rev_ti] = self.B[:, :, rev_ti+1] -\ - self.f(self.B[:, :, rev_ti+1], self.timegrid[rev_ti+1])*self.dt + \ - self.dt*self.g**2*grad_ln_ro +\ - (self.g)*\ - torch.empty([self.dim, self.N], dtype=torch.float64, device=self.device).normal_(mean=0, std=torch.sqrt(self.dt)) - else: - - grad_ln_b = torch.t(self.bw_density_estimation(rev_ti+1)) - - self.B[:, :, rev_ti] = self.B[:, :, rev_ti+1] -\ - (self.f(self.B[:, :, rev_ti+1], self.timegrid[rev_ti+1])-\ - self.g**2*grad_ln_ro +0.5*self.g**2*grad_ln_b)*self.dt - - for di in range(self.dim): - self.B[di, :, 0] = self.y1[di] - return 0 - - - """ - def reject_trajectories(self): - - Reject backward trajectories that do not reach the vicinity of the - initial point. - Deletes in place relevant rows of the `self.B` array that contains - the time reversed trajectories. - - Returns - ------- - int - Returns 0. - - - fplus = self.y1+self.f(self.y1, self.t1)*self.dt+6*self.g**2 *np.sqrt(self.dt) - fminus = self.y1+self.f(self.y1, self.t1) *self.dt-6*self.g**2 *np.sqrt(self.dt) - reverse_order = np.zeros(self.dim) - #this is an indicator if along one of the dimensions fplus - #is smaller than fminus - for iii in range(self.dim): - if fplus[iii] < fminus[iii]: - reverse_order[iii] = 1 - to_delete = np.zeros(self.N) - ##these will be one if ith trajectory is out of bounds - - ## checking if out of bounds for each dim - for iii in range(self.dim): - if reverse_order[iii] == 0: - to_delete += np.logical_not(np.logical_and(self.B[iii, :, 1] < fplus[iii], self.B[iii, :, 1] > fminus[iii])) - elif reverse_order[iii] == 1: - to_delete += np.logical_not(np.logical_and(self.B[iii, :, 1] > fplus[iii], self.B[iii, :, 1] < fminus[iii])) - - sinx = np.where(to_delete >= 1)[0] - #sinx = np.where(np.logical_or(np.logical_not(np.logical_and(self.B[0, :, 1] < fplus[0], self.B[0, :, 1] > fminus[0])), np.logical_not(np.logical_and(self.B[0, :, 1] < fplus[0], self.B[0, :, 1] > fminus[0]))))[0] - #((self.B[1,:,-2]fminus[1]) ) ))[0] - - - #logging.warning("Identified %d invalid bridge trajectories "%len(sinx)) - # if self.reject: - # logging.warning("Deleting invalid trajectories...") - # sinx = sinx[::-1] - # for element in sinx: - # self.B = np.delete(self.B, element, axis=1) - return 0 - """ - - - def calculate_u(self, grid_x, ti): - """ - Computes the control at position(s) grid_x at timestep ti - (i.e. at time self.timegrid[ti]). - - Parameters - ---------- - grid_x : ndarray, - size d x number of points to be evaluated. - ti : int, - time index in timegrid for the computation of u. - - - Returns - ------- - u_t: ndarray, - same size as grid_x. These are the controls u(grid_x, t), - where t=self.timegrid[ti]. - - """ - #a = 0.001 - #grad_dirac = lambda x,di: - 2*(x[di] -self.y2[di])* - #np.exp(- (1/a**2)* (x[0]- self.y2[0])**2)/(a**3 *np.sqrt(np.pi)) - u_t = torch.zeros(grid_x.T.shape, dtype=torch.float64, - device=self.device) - - - lnthsc1 = 2*torch.std(self.B[:, :, ti], dim=1) - lnthsc2 = 2*torch.std(self.Z[:, :, ti], dim=1) - - - bnds = torch.zeros(self.dim, 2, dtype=torch.float64, - device=self.device) - - for ii in range(self.dim): - if self.reweight == False or self.brown_bridge == False: - bnds[ii, 0] = min(torch.min(self.Z[ii, :, ti]), torch.min(self.B[ii, :, ti])) - bnds[ii, 1] = max(torch.max(self.Z[ii, :, ti]), torch.max(self.B[ii, :, ti])) - - else: - bnds[ii, 0] = min(torch.min(self.Ztr[ii, :, ti]), torch.min(self.B[ii, :, ti])) - bnds[ii, 1] = max(torch.max(self.Ztr[ii, :, ti]), torch.max(self.B[ii, :, ti])) - - if ti <= 5 or (ti >= self.k-5): - if self.reweight == False or self.brown_bridge == False: - ##for the first and last 5 timesteps, to avoid numerical singularities just assume gaussian densities - for di in range(self.dim): - mutb = torch.mean(self.B[di, :, ti]) - stdtb = torch.std(self.B[di, :, ti]) - mutz = torch.mean(self.Z[di, :, ti]) - stdtz = torch.std(self.Z[di, :, ti]) - u_t[di] = -(grid_x[:, di]- mutb)/stdtb**2 - (-(grid_x[:, di]- mutz)/stdtz**2) - elif self.reweight == True and self.brown_bridge == True: - for di in range(self.dim): - mutb = torch.mean(self.B[di, :, ti]) - stdtb = torch.std(self.B[di, :, ti]) - mutz = torch.mean(self.Ztr[di, :, ti]) - stdtz = torch.std(self.Ztr[di, :, ti]) - u_t[di] = -(grid_x[:, di]- mutb)/stdtb**2 - (-(grid_x[:, di]- mutz)/stdtz**2) - else: #if ti > 5: - ### clipping not used at the end but provided here for cases when - ### number of particles is small - ### and trajectories fall out of simulated flows - ### TO DO: add clipping as an option to be selected when - ### initialising the function - ### if point for evaluating control falls out of the region where we - ### have points, clip the points to - ### fall within the calculated region - we do not change the - ### position of the point, only the control value will be - ### calculated with clipped positions - """ - bndsb =torch.zeros(self.dim, 2, dtype=torch.float64, - device=self.device) - bndsz = torch.zeros(self.dim, 2, dtype=torch.float64, - device=self.device) - for di in range(self.dim): - bndsb[ii, 0] = torch.min(self.B[ii, :, ti]) - bndsb[ii, 1] = torch.max(self.B[ii, :, ti]) - bnds[ii, 0] = torch.min(self.Z[ii, :, ti]) - bnds[ii, 1] = torch.max(self.Z[ii, :, ti]) - """ - - ## clipping not used at the end! - ###cliping the values of points when evaluating the grad log p - grid_b = grid_x#np.clip(grid_x, bndsb[0], bndsb[1]) - grid_z = grid_x#np.clip(grid_x, bndsz[0], bndsz[1]) - - Sxx = torch.distributions.Uniform(low=bnds[:, 0], high=bnds[:, 1]).sample([self.N_sparse]) - - #for di in range(self.dim): - score_Bw = torch.t(torched_score_function_multid_seperate_all_dims(torch.t(self.B[:, :, ti]), - Sxx, - func_out=True, - C=0.001, - l=lnthsc1, - kern=self.kern, - device=self.device)(grid_b) ) - - #score_Bw = score_function_multid_seperate(self.B[:, :, ti].T, Sxx.T, func_out=True, C=0.001, which=1, l=lnthsc1, which_dim=di+1, kern=self.kern)(grid_b) - if self.reweight == False or self.brown_bridge == False: - score_Fw = torch.t( torched_score_function_multid_seperate_all_dims(torch.t(self.Z[:, :, ti]), - Sxx, - func_out=True, - C=0.001, - l=lnthsc2, - kern=self.kern, - device=self.device)(grid_z) ) - #score_Fw = score_function_multid_seperate(self.Z[:, :, ti].T, Sxx.T, func_out=True, C=0.001, which=1, l=lnthsc2, which_dim=di+1, kern=self.kern)(grid_z) - else: - """ - bndsztr = torch.zeros(self.dim, 2, dtype=torch.float64, - device=self.device) - for ii in range(self.dim): - #bndsztr[di] = [torch.min(self.Ztr[di, :, ti]), torch.max(self.Ztr[di, :, ti])] - bndsztr[ii, 0] = torch.min(self.Ztr[ii, :, rev_t]) - bndsztr[ii, 1] = torch.max(self.Ztr[ii, :, rev_t]) - """ - grid_ztr = grid_x #np.clip(grid_x, bndsztr[0], bndsztr[1]) - lnthsc3 = 2*torch.std(self.Ztr[:, :, ti], dim=1) - #score_Fw = score_function_multid_seperate(self.Ztr[:, :, ti].T, Sxx.T, func_out=True, C=0.001, which=1, l=lnthsc3, which_dim=di+1, kern=self.kern)(grid_ztr) - score_Fw = torch.t( torched_score_function_multid_seperate_all_dims(torch.t(self.Ztr[:, :, ti]), - Sxx, - func_out=True, - C=0.001, - l=lnthsc3, - kern=self.kern, - device=self.device)(grid_ztr) ) - u_t = score_Bw - score_Fw - - return u_t - - - def check_if_covered(self, X, ti): - """ - Checks if test point X falls within forward and backward densities at - timepoint timegrid[ti]. - - - Parameters - ---------- - X : array 1x dim or Kxdim - Point in state space where control is evaluated. - ti : int - Index in timegrid array indicating the time within the - time interval [t1,t2]. - - Returns - ------- - Boolean variable - True if the text point X falls within the densities. - - """ - covered = True - bnds = torch.zeros(self.dim, 2, dtype=torch.float64, - device=self.device) - for ii in range(self.dim): - bnds[ii, 0] = min(torch.min(self.Z[ii, :, ti]), torch.min(self.B[ii, :, ti])) - bnds[ii, 1] = max(torch.max(self.Z[ii, :, ti]), torch.max(self.B[ii, :, ti])) - - - - covered = covered * ((X[ii] >= bnds[ii][0]) and (X[ii] <= bnds[ii][1])) - - return covered +# -*- coding: utf-8 -*- +""" +Created on Mon Jun 27 15:03:15 2022 + +@author: maout +""" + + +# -*- coding: utf-8 -*- + +#Created on Sun Dec 12 00:02:39 2021 + +#@author: maout + + + + +# + +import time +import logging +import numpy as np +try: + import ot + POTTED = True +except ImportError: + POTTED = False + + +try: + import numba + NUMBED = True +except ImportError: + NUMBED = False + + + +from matplotlib import pyplot as plt +try: + import torch + TORCHED = True +except ImportError: + TORCHED = False + + +if __name__ == "DeterministicParticleFlowControl.DeterministicParticleFlowControl": + from .score_estimators.score_function_estimators import score_function_multid_seperate + from .reweighting.optimal_transport_reweighting import reweight_optimal_transport_multidim + if TORCHED: + from .score_estimators.score_function_estimators_pytorch import torched_score_function_multid_seperate_all_dims + #from .utils.utils_pytorch import set_device + +else: + ###this clause is for imports to work properly for the documendation + from score_estimators.score_function_estimators import score_function_multid_seperate + from reweighting.optimal_transport_reweighting import reweight_optimal_transport_multidim + if TORCHED: + from score_estimators.score_function_estimators_pytorch import torched_score_function_multid_seperate_all_dims + #from utils.utils_pytorch import set_device +from duecredit import due, BibTeX + + +if TORCHED: + __all__ = ["DPFC", "torched_DPFC"] +else: + __all__ = ["DPFC"] + +# Use duecredit (duecredit.org) to provide a citation to relevant work to +# be cited. This does nothing, unless the user has duecredit installed, +# And calls this with duecredit (as in `python -m duecredit script.py`): +due.cite(BibTeX(""" + @article{maoutsa2021deterministic, + title={Deterministic particle flows for constraining stochastic nonlinear systems}, + author={Maoutsa, Dimitra and Opper, Manfred}, + journal={arXiv preprint arXiv:2112.05735}, + year={2021} +} + """), + description="A deterministic barticle-based method for stochastic optimal control", + tags=["reference-implementation"], + path='DeterministicParticleFlowControl') + + + + +class DPFC(object): + """ + Deterministic particle flow control top-level class. + + Provides the necessary functions to sample the required probability + flows and estimate the controls. + + Attributes + ---------- + t1 : float + Initial time. + t2: float + end time point. + y1: array_like + initial position. + y2: array_like + terminal position. + f: function, callable + drift function handle. + g: float or array_like + diffusion coefficient or function handle. + N: int + number of particles/trajectories. + M: int + number of sparse points for grad log density estimation. + reweight: boolean + determines if reweighting will follow. + U: function, callable + reweighting function to be employed during reweighting, + dimensions :math:`dim_y1,t \\to 1`. + dens_est: str + - 'nonparametric' : non parametric density estimation (this was + used in the paper) + - TO BE ADDED: + - 'hermit1' : parametic density estimation empoying hermite + polynomials (physiscist's) + - 'hermit2' : parametic density estimation empoying hermite + polynomials (probabilists's) + - 'poly' : parametic density estimation empoying simple polynomials + - 'RBF' : parametric density estimation employing radial basis functions. + kern: str + type of kernel: 'RBF' or 'periodic' (only the 'RBF' was used and gives + robust results. Do not use 'periodic' yet!). + reject: boolean + parameter indicating whether non valid backward trajectories will be rejected. + plotting: boolean + parameter indicating whether bridge statistics will be plotted. + f_true: funtion, callable + in case of Brownian bridge reweighting this is the true forward drift + for simulating the forward dynamics. + brown_bridge: boolean, + determines if the reweighting concearns contstraint or reweighting with + respect to brownian bridge. + deterministic: boolean, + indicates the type of dynamics the particles will follow. + If False the flows are simulated with stochastic path sampling. + b_type: string, + indicates the type of boundaries relevant for evolving the dynamics + Options: - normal: no boundaries + - kuramoto: periodic boundary [0, 2 pi] + + Methods + ------- + forward_sampling_Otto: + Creates samples of the forward flow. + forward sampling(): + Samples the forward flow with stochatic particle trajectories. + f_seperate(x,t): + Drift for the deterministic propagation of partcles that are at time t + in position x. + backward_simulation(): + Sampling the backward density with stochastic particles. + reject_trajectories(): + Rejects backward trajectories that do not end up in the vicinity of the + initial point. + Run only if the instance is attribute "reject" is set to True. + Gives logging.warning messages. + forward_sampling_Otto_true(): + Relevant only when forward sampling happens with Brownian bridge. + """ + + + def __init__(self, t1, t2, y1, y2, f, g, N, M, reweight=False, U=None, + dens_est='nonparametric', reject=True, kern='RBF', + f_true=None, brown_bridge=False, deterministic=True, + b_type='normal'): + + self.dim = y1.size # dimensionality of the system + self.t1 = t1 + self.t2 = t2 + self.y1 = y1 + self.y2 = y2 + + self.b_type = b_type ###determines type of domain boundaries + ##density estimation stuff + self.kern = kern + if kern == 'periodic': + self.kern = 'RBF' + logging.warning('Please do not use periodic kernel yet!') + logging.warning('For all the numerical experiments RBF was used') + logging.warning('We changed your choice to RBF') + # DRIFT /DIFFUSION + self.f = f + self.g = g #scalar or array + + ### PARTICLE DISCRETISATION + self.N = N + + self.N_sparse = M + + self.dt = 0.001 #((t2-t1)/k) + ### reject unreasonable backward trajectories that do not return + ### to initial condition + self.reject = reject + ### indicator for what type of dynamics the particles follow + self.deterministic = deterministic + + + self.timegrid = np.arange(self.t1, self.t2+self.dt/2, self.dt) + self.k = self.timegrid.size + ### reweighting + self.brown_bridge = brown_bridge + self.reweight = reweight + if self.reweight: + self.U = U + if self.brown_bridge: + #storage for forward trajectories with true drift + self.Ztr = np.zeros((self.dim, self.N, self.k)) + self.f_true = f_true + + + + #storage for forward trajectories + self.Z = np.zeros((self.dim, self.N, self.k)) + #storage for backward trajectories + self.B = np.zeros((self.dim, self.N, self.k)) + self.ln_roD = [] ## storing the estimated forward logarithmic gradients + + + ##the stochastic sampling is provided for comparison + if self.deterministic: + self.forward_sampling_Otto() + ### if a Brownian bridge is used for forward sampling + if self.reweight and self.brown_bridge: + self.forward_sampling_Otto_true() + else: + self.forward_sampling() + ## the backward function selects internally for type of dynamics + self.backward_simulation() + # if self.reject: + # self.reject_trajectories() + + + def forward_sampling(self): + """ + Sampling forward probability flow with stochastic particle dynamics. + If reweighting is required at every time step the particles are + appropriatelly reweighted accordint to function :math:`U(x,t)` + + Returns + ------- + int + Returns 0 to make sure everything runs correctly. + The sampled density is stored in place in the array `self.Z`. + + """ + logging.info('Sampling forward...') + W = np.ones((self.N, 1))/self.N + for ti, tt in enumerate(self.timegrid): + + if ti == 0: + self.Z[0, :, 0] = self.y1[0] + self.Z[1, :, 0] = self.y1[1] + else: + for i in range(self.N): + #self.Z[:,i,:] = sdeint.itoint(self.f, self.g, self.Z[i,0], self.timegrid)[:,0] + + self.Z[:, i, ti] = (self.Z[:, i, ti-1] + \ + self.dt* self.f(self.Z[:, i, ti-1]) + \ + (self.g)*np.random.normal(loc=0.0, scale=np.sqrt(self.dt), size=(self.dim,))) + + if self.b_type=='kuramoto': + self.Z[:, i, ti] = self.Z[:, i, ti] %(2*np.pi) + + ###WEIGHT + if self.reweight == True: + if ti > 0: + W[:, 0] = np.exp(self.U(self.Z[:, :, ti])) + W = W/np.sum(W) + + ###REWEIGHT + Tstar = reweight_optimal_transport_multidim(self.Z[:, :, ti].T, W) + + self.Z[:, :, ti] = (self.Z[:, :, ti])@Tstar + if self.b_type=='kuramoto': + self.Z[:, :, ti] = self.Z[:, :, ti] %(2*np.pi) + + logging.info('Forward sampling done!') + return 0 + + + + + ### relevant only when forward trajectories follow brownian brifge - + ###this simulates forward trajectories with true f + def f_seperate_true(self, x, t): + """ + (Relevant only when forward sampling happens with Brownian bridge + reweighting) + Wrapper for the drift function of the deterministic particles with the + actual f (system drift) minus the logarithmic gradient term computed + on current particles positions. + Provided for easy integration, and can be passed to ode integrators. + + Parameters + ---------- + x : 2d-array, + Particle positions (dimension x number of particles). + t : float, + Time t within the [t1,t2] interval. + + Returns + ------- + 2d-array + Returns the deterministic forces required to ntegrate the particle + positions for one time step, + i.e. return :math:`f(x,t)-\\frac{1}{2}\\sigma^2\\nabla \\rho_t(x)`, + evaluated at the current positions x and t. + + """ + dimi, N = x.shape + bnds = np.zeros((dimi, 2)) + for ii in range(dimi): + bnds[ii] = [np.min(x[ii, :]), np.max(x[ii, :])] + #sum_bnds = np.sum(bnds) + + Sxx = np.array([np.random.uniform(low=bnd[0], high=bnd[1], size=(self.N_sparse)) for bnd in bnds]) + gpsi = np.zeros((dimi, N)) + lnthsc = 2*np.std(x, axis=1) + + for ii in range(dimi): + gpsi[ii, :] = score_function_multid_seperate(x.T, Sxx.T, False, C=0.001, which=1, l=lnthsc, which_dim=ii+1, kern=self.kern) + + return self.f_true(x, t)-0.5* self.g**2* gpsi + + + ### effective forward drift - estimated seperatelly for each dimension + #plain GP prior + def f_seperate(self, x, t): + """ + Computes the deterministic forces for the evolution of the deterministic + particles for the current particle positions, + ie. drift minus the logarithmic gradient term. + Is used as a wrapper for evolving the particles, + and can be provided to "any" ODE integrator. + + Parameters + ---------- + x : 2d-array, + Particle positions (dimension x number of particles). + t : float, + Time t within the [t1,t2] interval. + + Returns + ------- + 2d-array + Returns the deterministic forces required to ntegrate the particle + positions for one time step, + i.e. return :math:`f(x,t)-\\frac{1}{2}\\sigma^2\\nabla \\rho_t(x)`, + evaluated at the current positions x and t. + + """ + + + dimi, N = x.shape + ### detect min and max of forward flow for each dimension + ### we want to know the state space volume of the forward flow + bnds = np.zeros((dimi, 2)) + for ii in range(dimi): + bnds[ii] = [np.min(x[ii, :]), np.max(x[ii, :])] + sum_bnds = np.sum(bnds) ##this is for detecting if sth goes wrong i.e. trajectories explode + if np.isnan(sum_bnds) or np.isinf(sum_bnds): + ##if we get unreasoble bounds just plot the first 2 dimensions of the trajectories + plt.figure(figsize=(6, 4)), plt.plot(self.Z[0].T, self.Z[1].T, alpha=0.3) + plt.show() + + ##these are the inducing points + ## here we select them from a uniform distribution within the state space volume spanned from the forward flow + Sxx = np.array([np.random.uniform(low=bnd[0], high=bnd[1], size=(self.N_sparse)) for bnd in bnds]) + gpsi = np.zeros((dimi, N)) + lnthsc = 2*np.std(x, axis=1) + for ii in range(dimi): + gpsi[ii, :] = score_function_multid_seperate(x.T, Sxx.T, False, C=0.001, which=1, l=lnthsc, which_dim=ii+1, kern=self.kern) + + return self.f(x, t)-0.5* self.g**2* gpsi + + ###same as forward sampling but without reweighting - this is for bridge reweighting + ### not for constraint reweighting + def forward_sampling_Otto_true(self): + """ + (Relevant only when forward sampling happens with Brownian bridge + reweighting) + Same as forward sampling but without reweighting. + + Returns + ------- + int + Returns 0 to make sure everything runs correctly. + The sampled density is stored in place in the array `self.Ztr`. + + See also + --------- + DPFC.forward_sampling, DPFC.forward_sampling_Otto + + """ + logging.info('Sampling forward with deterministic particles and true drift...') + #W = np.ones((self.N,1))/self.N + for ti, tt in enumerate(self.timegrid): + + if ti == 0: + for di in range(self.dim): + self.Ztr[di, :, 0] = self.y1[di] + + elif ti == 1: #propagate one step with stochastic to avoid the delta function + #substract dt because I want the time at t-1 + self.Ztr[:, :, ti] = (self.Ztr[:, :, ti-1] + self.dt*self.f_true(self.Ztr[:, :, ti-1], tt-self.dt)+\ + (self.g)*np.random.normal(loc=0.0, scale=np.sqrt(self.dt), size=(self.dim, self.N))) + else: + self.Ztr[:, :, ti] = (self.Ztr[:, :, ti-1] + self.dt* self.f_seperate_true(self.Ztr[:, :, ti-1], tt-self.dt)) + + if self.b_type=='kuramoto': + self.Ztr[:, :, ti] = self.Ztr[:, :, ti] %(2*np.pi) + logging.info('Forward sampling with Otto true is ready!') + return 0 + + + + def forward_sampling_Otto(self): + """ + Samples the forward probability flow with deterministic particle + dynamics. + If required at every timestep a particle reweighting takes place + employing the weights obtained from the exponentiated path constraint + :math:`U(x,t)` + + Returns + ------- + int + Returns 0 to make sure everything runs correctly. + The sampled density is stored in place in the array `self.Z`. + + """ + logging.info('Sampling forward with deterministic particles...') + W = np.ones((self.N, 1))/self.N + for ti, tt in enumerate(self.timegrid): + if ti == 0: + for di in range(self.dim): + self.Z[di, :, 0] = self.y1[di] + if self.brown_bridge: + self.Z[di, :, -1] = self.y2[di] + ## we start forward trajectories for a delta function. + ##in principle we could start from an arbitrary distribution + ##if you want to start from a normal uncomen the following and comment the above initialisation for y1 + #self.Z[di,:,0] = np.random.normal(self.y1[di], 0.05, self.N) + elif ti == 1: #propagate one step with stochastic to avoid the delta function + #substract dt because I want the time at t-1 + self.Z[:, :, ti] = (self.Z[:, :, ti-1] + self.dt*self.f(self.Z[:, :, ti-1], tt-self.dt)+\ + (self.g)*np.random.normal(loc=0.0, scale=np.sqrt(self.dt), size=(self.dim, self.N))) + else: + self.Z[:, :, ti] = (self.Z[:, :, ti-1] + self.dt* self.f_seperate(self.Z[:, :, ti-1], tt-self.dt)) + ##if kuramoto-periodic + if self.b_type=='kuramoto': + self.Z[:, :, ti] = self.Z[:, :, ti] %(2*np.pi) + ###REWEIGHT + if self.reweight == True: + if ti > 0: + + W[:, 0] = np.exp(self.U(self.Z[:, :, ti], tt)*self.dt) #-1 + W = W/np.sum(W) + + ###REWEIGHT + start = time.time() + Tstar = reweight_optimal_transport_multidim(self.Z[:, :, ti].T, W) + #print(Tstar) + if ti == 3: + stop = time.time() + logging.info('Timepoint: %d needed '%ti, stop-start) + self.Z[:, :, ti] = ((self.Z[:, :, ti])@Tstar) ##### + if self.b_type=='kuramoto': + self.Z[:, :, ti] = self.Z[:, :, ti] %(2*np.pi) + logging.info('Forward sampling with Otto is ready!') + return 0 + + def density_estimation(self, ti, rev_ti): + rev_t = rev_ti + grad_ln_ro = np.zeros((self.dim, self.N)) + lnthsc = 2*np.std(self.Z[:, :, rev_t], axis=1) + bnds = np.zeros((self.dim, 2)) + for ii in range(self.dim): + bnds[ii] = [max(np.min(self.Z[ii, :, rev_t]), np.min(self.B[ii, :, rev_ti])), min(np.max(self.Z[ii, :, rev_t]), np.max(self.B[ii, :, rev_ti]))] + sum_bnds = np.sum(bnds) + if np.isnan(sum_bnds) or np.isinf(sum_bnds): + plt.figure(figsize=(6, 4)), plt.plot(self.B[0].T, self.B[1].T, alpha=0.3) + plt.plot(self.y1[0], self.y1[1], 'go') + plt.show() + #sparse points + Sxx = np.array([np.random.uniform(low=bnd[0], high=bnd[1], size=(self.N_sparse)) for bnd in bnds]) + + for di in range(self.dim): + #estimate density from forward (Z) and evaluate at current postitions of backward particles (B) + grad_ln_ro[di, :] = score_function_multid_seperate(self.Z[:, :, rev_t].T, Sxx.T, func_out=True, C=0.001, which=1, l=lnthsc, which_dim=di+1, kern=self.kern)(self.B[:, :, rev_ti].T) + + + return grad_ln_ro + + + def bw_density_estimation(self, rev_ti): + """ + Estimates the logaritmic gradient of the backward flow evaluated at + particle positions of the backward flow. + + + Parameters + ---------- + + rev_ti : int, + indicates the time point in the timegrid where the estimation + will take place, i.e. for time t=self.timegrid[rev_ti. + + Returns + ------- + grad_ln_b: 2d-array, + with the logarithmic gradients of the time reversed + (backward) flow (dim x N) for the timestep `rev_ti`. + + """ + grad_ln_b = np.zeros((self.dim, self.N)) + lnthsc = 2*np.std(self.B[:, :, rev_ti], axis=1) + #print(ti, rev_ti, rev_ti-1) + bnds = np.zeros((self.dim, 2)) + for ii in range(self.dim): + bnds[ii] = [max(np.min(self.Z[ii, :, rev_ti]), np.min(self.B[ii, :, rev_ti])), min(np.max(self.Z[ii, :, rev_ti]), np.max(self.B[ii, :, rev_ti]))] + #sparse points + Sxx = np.array([np.random.uniform(low=bnd[0], high=bnd[1], size=(self.N_sparse)) for bnd in bnds]) + + for di in range(self.dim): + grad_ln_b[di, :] = score_function_multid_seperate(self.B[:, :, rev_ti].T, Sxx.T, func_out=False, C=0.001, which=1, l=lnthsc, which_dim=di+1, kern=self.kern) + # grad_ln_a = score_function_multid_seperate2(self.B[:, :, rev_ti].T, Sxx.T, func_out=False, C=0.001, which=1, l=lnthsc, which_dim=di+1, kern=self.kern) + # #np.testing.assert_array_equal(grad_ln_b[di, :], grad_ln_a) + # np.testing.assert_allclose(grad_ln_b[di, :], grad_ln_a) + return grad_ln_b + + + def backward_simulation(self): + """ + Sample time reversed flow with deterministic dynamics (or stochastic if + `self.deterministic == False`). + Trajectories are stored in place in `self.B` array of dimensionality + (dim x N x timegrid.size). + `self.B` does not require a timereversion at the end, everything + is stored in the correct order. + + Returns + ------- + int + Returns 0 to ensure everything was executed correctly. + + """ + + for ti, tt in enumerate(self.timegrid[:-1]): + if ti == 0: + for di in range(self.dim): + self.B[di, :, -1] = self.y2[di] + else: + + rev_ti = self.k -ti-1 + #density estimation of forward particles + grad_ln_ro = self.density_estimation(ti, rev_ti+1) + + if (ti == 1 and self.deterministic) or (not self.deterministic): + + self.B[:, :, rev_ti] = (self.B[:, :, rev_ti+1] -\ + self.f(self.B[:, :, rev_ti+1], self.timegrid[rev_ti+1])*self.dt + \ + self.dt*self.g**2*grad_ln_ro +\ + (self.g)*np.random.normal(loc=0.0, scale=np.sqrt(self.dt), size=(self.dim, self.N))) + else: + grad_ln_b = self.bw_density_estimation(rev_ti+1) + self.B[:, :, rev_ti] = (self.B[:, :, rev_ti+1] -\ + (self.f(self.B[:, :, rev_ti+1], self.timegrid[rev_ti+1])- self.g**2*grad_ln_ro +0.5*self.g**2*grad_ln_b)*self.dt) + + if self.b_type=='kuramoto': + self.B[:, :, rev_ti] = self.B[:, :, rev_ti] %(2*np.pi) + for di in range(self.dim): + self.B[di, :, 0] = self.y1[di] + return 0 + + + """ + def reject_trajectories(self): + + Reject backward trajectories that do not reach the vicinity of the + initial point. + Deletes in place relevant rows of the `self.B` array that contains + the time reversed trajectories. + + Returns + ------- + int + Returns 0. + + + fplus = self.y1+self.f(self.y1, self.t1)*self.dt+6*self.g**2 *np.sqrt(self.dt) + fminus = self.y1+self.f(self.y1, self.t1) *self.dt-6*self.g**2 *np.sqrt(self.dt) + reverse_order = np.zeros(self.dim) + #this is an indicator if along one of the dimensions fplus + #is smaller than fminus + for iii in range(self.dim): + if fplus[iii] < fminus[iii]: + reverse_order[iii] = 1 + to_delete = np.zeros(self.N) + ##these will be one if ith trajectory is out of bounds + + ## checking if out of bounds for each dim + for iii in range(self.dim): + if reverse_order[iii] == 0: + to_delete += np.logical_not(np.logical_and(self.B[iii, :, 1] < fplus[iii], self.B[iii, :, 1] > fminus[iii])) + elif reverse_order[iii] == 1: + to_delete += np.logical_not(np.logical_and(self.B[iii, :, 1] > fplus[iii], self.B[iii, :, 1] < fminus[iii])) + + sinx = np.where(to_delete >= 1)[0] + #sinx = np.where(np.logical_or(np.logical_not(np.logical_and(self.B[0, :, 1] < fplus[0], self.B[0, :, 1] > fminus[0])), np.logical_not(np.logical_and(self.B[0, :, 1] < fplus[0], self.B[0, :, 1] > fminus[0]))))[0] + #((self.B[1,:,-2]fminus[1]) ) ))[0] + + + #logging.warning("Identified %d invalid bridge trajectories "%len(sinx)) + # if self.reject: + # logging.warning("Deleting invalid trajectories...") + # sinx = sinx[::-1] + # for element in sinx: + # self.B = np.delete(self.B, element, axis=1) + return 0 + """ + + + def calculate_u(self, grid_x, ti): + """ + Computes the control at position(s) grid_x at timestep ti + (i.e. at time self.timegrid[ti]). + + Parameters + ---------- + grid_x : ndarray, + size d x number of points to be evaluated. + ti : int, + time index in timegrid for the computation of u. + + + Returns + ------- + u_t: ndarray, + same size as grid_x. These are the controls u(grid_x, t), + where t=self.timegrid[ti]. + + """ + #a = 0.001 + #grad_dirac = lambda x,di: - 2*(x[di] -self.y2[di])* + #np.exp(- (1/a**2)* (x[0]- self.y2[0])**2)/(a**3 *np.sqrt(np.pi)) + u_t = np.zeros(grid_x.T.shape) + + + lnthsc1 = 2*np.std(self.B[:, :, ti], axis=1) + lnthsc2 = 2*np.std(self.Z[:, :, ti], axis=1) + + + bnds = np.zeros((self.dim, 2)) + for ii in range(self.dim): + if self.reweight == False or self.brown_bridge == False: + bnds[ii] = [max(np.min(self.Z[ii, :, ti]), np.min(self.B[ii, :, ti])), min(np.max(self.Z[ii, :, ti]), np.max(self.B[ii, :, ti]))] + else: + bnds[ii] = [max(np.min(self.Ztr[ii, :, ti]), np.min(self.B[ii, :, ti])), min(np.max(self.Ztr[ii, :, ti]), np.max(self.B[ii, :, ti]))] + + if ti <= 5 or (ti >= self.k-5): + if self.reweight == False or self.brown_bridge == False: + ##for the first and last 5 timesteps, to avoid numerical singularities just assume gaussian densities + for di in range(self.dim): + mutb = np.mean(self.B[di, :, ti]) + stdtb = np.std(self.B[di, :, ti]) + mutz = np.mean(self.Z[di, :, ti]) + stdtz = np.std(self.Z[di, :, ti]) + u_t[di] = -(grid_x[:, di]- mutb)/stdtb**2 - (-(grid_x[:, di]- mutz)/stdtz**2) + elif self.reweight == True and self.brown_bridge == True: + for di in range(self.dim): + mutb = np.mean(self.B[di, :, ti]) + stdtb = np.std(self.B[di, :, ti]) + mutz = np.mean(self.Ztr[di, :, ti]) + stdtz = np.std(self.Ztr[di, :, ti]) + u_t[di] = -(grid_x[:, di]- mutb)/stdtb**2 - (-(grid_x[:, di]- mutz)/stdtz**2) + else: #if ti > 5: + ### clipping not used at the end but provided here for cases when + ### number of particles is small + ### and trajectories fall out of simulated flows + ### TO DO: add clipping as an option to be selected when + ### initialising the function + ### if point for evaluating control falls out of the region where we + ### have points, clip the points to + ### fall within the calculated region - we do not change the + ### position of the point, only the control value will be + ### calculated with clipped positions + bndsb = np.zeros((self.dim, 2)) + bndsz = np.zeros((self.dim, 2)) + for di in range(self.dim): + bndsb[di] = [np.min(self.B[di, :, ti]), np.max(self.B[di, :, ti])] + bndsz[di] = [np.min(self.Z[di, :, ti]), np.max(self.Z[di, :, ti])] + + ## clipping not used at the end! + ###cliping the values of points when evaluating the grad log p + grid_b = grid_x#np.clip(grid_x, bndsb[0], bndsb[1]) + grid_z = grid_x#np.clip(grid_x, bndsz[0], bndsz[1]) + + Sxx = np.array([np.random.uniform(low=bnd[0], high=bnd[1], size=(self.N_sparse)) for bnd in bnds]) + for di in range(self.dim): + score_Bw = score_function_multid_seperate(self.B[:, :, ti].T, Sxx.T, func_out=True, C=0.001, which=1, l=lnthsc1, which_dim=di+1, kern=self.kern)(grid_b) + if self.reweight == False or self.brown_bridge == False: + score_Fw = score_function_multid_seperate(self.Z[:, :, ti].T, Sxx.T, func_out=True, C=0.001, which=1, l=lnthsc2, which_dim=di+1, kern=self.kern)(grid_z) + else: + bndsztr = np.zeros((self.dim, 2)) + for ii in range(self.dim): + bndsztr[di] = [np.min(self.Ztr[di, :, ti]), np.max(self.Ztr[di, :, ti])] + grid_ztr = np.clip(grid_x, bndsztr[0], bndsztr[1]) + lnthsc3 = 2*np.std(self.Ztr[:, :, ti], axis=1) + score_Fw = score_function_multid_seperate(self.Ztr[:, :, ti].T, Sxx.T, func_out=True, C=0.001, which=1, l=lnthsc3, which_dim=di+1, kern=self.kern)(grid_ztr) + + u_t[di] = score_Bw - score_Fw + # for di in range(self.dim): + # u_t[di] = score_function_multid_seperate(self.B[:,:,ti].T,Sxx.T,func_out= True,C=0.001,which=1,l=lnthsc,which_dim=di+1, kern=self.kern)(grid_x.T) \ + # - score_function_multid_seperate(self.Z[:,:,ti].T,Sxx.T,func_out= True,C=0.001,which=1,l=lnthsc,which_dim=di+1, kern=self.kern)(grid_x.T) + + + return u_t + + + def check_if_covered(self, X, ti): + """ + Checks if test point X falls within forward and backward densities at + timepoint timegrid[ti]. + + + Parameters + ---------- + X : array 1x dim or Kxdim + Point in state space where control is evaluated. + ti : int + Index in timegrid array indicating the time within the + time interval [t1,t2]. + + Returns + ------- + Boolean variable - True if the text point X falls within the densities. + + """ + covered = True + bnds = np.zeros((self.dim, 2)) + for ii in range(self.dim): + bnds[ii] = [max(np.min(self.Z[ii, :, ti]), np.min(self.B[ii, :, ti])), min(np.max(self.Z[ii, :, ti]), np.max(self.B[ii, :, ti]))] + #bnds[ii] = [np.min(self.B[ii,:,ti]),np.max(self.B[ii,:,ti])] + + covered = covered * ((X[ii] >= bnds[ii][0]) and (X[ii] <= bnds[ii][1])) + + return covered + +#%% + +class torched_DPFC(object): + """ + Deterministic particle flow control top-level class implemented in pytorch. + + Provides the necessary functions to sample the required probability + flows and estimate the controls. + + Attributes + ---------- + t1 : float + Initial time. + t2: float + end time point. + y1: array_like + initial position. + y2: array_like + terminal position. + f: function, callable + drift function handle. + g: float or array_like + diffusion coefficient or function handle. + N: int + number of particles/trajectories. + M: int + number of sparse points for grad log density estimation. + reweight: boolean + determines if reweighting will follow. + U: function, callable + reweighting function to be employed during reweighting, + dimensions :math:`dim_y1,t \\to 1`. + dens_est: str + - 'nonparametric' : non parametric density estimation (this was + used in the paper) + - TO BE ADDED: + - 'hermit1' : parametic density estimation empoying hermite + polynomials (physiscist's) + - 'hermit2' : parametic density estimation empoying hermite + polynomials (probabilists's) + - 'poly' : parametic density estimation empoying simple polynomials + - 'RBF' : parametric density estimation employing radial basis functions. + kern: str + type of kernel: 'RBF' or 'periodic' (only the 'RBF' was used and gives + robust results. Do not use 'periodic' yet!). + reject: boolean + parameter indicating whether non valid backward trajectories will be rejected. + plotting: boolean + parameter indicating whether bridge statistics will be plotted. + f_true: funtion, callable + in case of Brownian bridge reweighting this is the true forward drift + for simulating the forward dynamics. + brown_bridge: boolean, + determines if the reweighting concearns contstraint or reweighting with + respect to brownian bridge. + deterministic: boolean, + indicates the type of dynamics the particles will follow. + If False the flows are simulated with stochastic path sampling. + device: string, + indicates the device where computations will be exacuted. + `cpu` or `gpu/cuda` or `tpu` if available. + b_type: string, + indicates the type of boundaries relevant for evolving the dynamics + Options: - normal: no boundaries + - kuramoto: periodic boundary [0, 2 pi] + + + Methods + ------- + forward_sampling_Otto: + Creates samples of the forward flow. + forward sampling(): + Samples the forward flow with stochatic particle trajectories. + f_seperate(x,t): + Drift for the deterministic propagation of partcles that are at time t + in position x. + backward_simulation(): + Sampling the backward density with stochastic particles. + reject_trajectories(): + Rejects backward trajectories that do not end up in the vicinity of the + initial point. + Run only if the instance is attribute "reject" is set to True. + Gives logging.warning messages. + forward_sampling_Otto_true(): + Relevant only when forward sampling happens with Brownian bridge. + """ + + + def __init__(self, t1, t2, y1, y2, f, g, N, M, reweight=False, U=None, + dens_est='nonparametric', reject=True, kern='RBF', + f_true=None, brown_bridge=False, deterministic=True, + device=None, b_type='normal'): + + self.device = device + # dimensionality of the system + self.dim = torch.tensor(y1.size(dim=0), dtype=torch.int, device=self.device) # dimensionality of the system + self.t1 = t1 + self.t2 = t2 + if not torch.is_tensor(y1): + self.y1 = torch.tensor(y1, dtype=torch.float64, device=self.device) + self.y2 = torch.tensor(y2, dtype=torch.float64, device=self.device) + else: + self.y1 = y1 + self.y2 = y2 + + self.b_type = b_type + ##density estimation stuff + self.kern = kern + if kern == 'periodic': + self.kern = 'RBF' + logging.warning('Please do not use periodic kernel yet!') + logging.warning('For all the numerical experiments RBF was used') + logging.warning('We changed your choice to RBF') + # DRIFT /DIFFUSION + self.f = f + self.g = torch.tensor(g, dtype=torch.float64, device=self.device) #scalar or array + + ### PARTICLE DISCRETISATION + self.N = torch.tensor(N, dtype=torch.int, device=self.device) + + self.N_sparse = torch.tensor(M, dtype=torch.int, device=self.device)#M + + self.dt = torch.tensor(0.001, dtype=torch.float64, device=self.device) + #((t2-t1)/k) + + ### reject unreasonable backward trajectories that do not return + ### to initial condition + self.reject = reject + ### indicator for what type of dynamics the particles follow + self.deterministic = deterministic + + + self.timegrid = torch.arange(self.t1, self.t2+self.dt/2, self.dt, + dtype=torch.float64, device=self.device) + self.k = torch.tensor(self.timegrid.size(dim=0), dtype=torch.int, device=self.device) + ### reweighting + self.brown_bridge = brown_bridge + self.reweight = reweight + if self.reweight: + self.U = U + if self.brown_bridge: + #storage for forward trajectories with true drift + self.Ztr = torch.zeros(self.dim, self.N, self.k, + dtype=torch.float64, device=self.device) + self.f_true = f_true + + + + #storage for forward trajectories + self.Z = torch.zeros(self.dim, self.N, self.k, dtype=torch.float64, + device=self.device) + #storage for backward trajectories + self.B = torch.zeros(self.dim, self.N, self.k, dtype=torch.float64, + device=self.device ) + self.ln_roD = [] ## storing the estimated forward logarithmic gradients + + + ##the stochastic sampling is provided for comparison + if self.deterministic: + self.forward_sampling_Otto() + + ### if a Brownian bridge is used for forward sampling + if self.reweight and self.brown_bridge: + self.forward_sampling_Otto_true() + else: + self.forward_sampling() + + ## the backward function selects internally for type of dynamics + self.backward_simulation() + # if self.reject: + # self.reject_trajectories() + + + def forward_sampling(self): + """ + Sampling forward probability flow with stochastic particle dynamics. + If reweighting is required at every time step the particles are + appropriatelly reweighted accordint to function :math:`U(x,t)` + + Returns + ------- + int + Returns 0 to make sure everything runs correctly. + The sampled density is stored in place in the array `self.Z`. + + """ + logging.info('Sampling forward...') + W = torch.ones(self.N, 1, dtype=torch.float64, device=self.device)/self.N + + for ti, tt in enumerate(self.timegrid): + if ti == 0: + self.Z[0, :, 0] = self.y1[0] + self.Z[1, :, 0] = self.y1[1] + else: + self.Z[:, :, ti] = self.Z[:, :, ti-1] + \ + self.dt*self.f(self.Z[:, :, ti-1], tt-self.dt)+\ + (self.g)*\ + torch.empty([self.dim, self.N]).normal_(mean=0, std=torch.sqrt(self.dt)) + ###WEIGHT + if self.reweight == True: + if ti > 0: + W[:, 0] = torch.exp(self.U(self.Z[:, :, ti])) + W = W/torch.sum(W) + + ###REWEIGHT with pot TO DO: + #Tstar = reweight_optimal_transport_multidim(self.Z[:, :, ti].T, W) + M = ot.dist(self.Z[:,:,ti].T, self.Z[:,:,ti].T) + M /= M.max() + a = W[:,0] + b = torch.ones_like(W[:,0], dtype=torch.float64, + device=self.device)/self.N + T2 = ot.emd(a, b, M) + self.Z[:, :, ti] = (self.N*self.Z[:, :, ti])@T2 + + #for di in range(self.dim): + #self.Z[di, :, -1] = self.y2[di] + logging.info('Forward sampling done!') + return 0 + + + + + ### relevant only when forward trajectories follow brownian brifge - + ###this simulates forward trajectories with true f + def f_seperate_true(self, x, t): + """ + (Relevant only when forward sampling happens with Brownian bridge + reweighting) + Wrapper for the drift function of the deterministic particles with the + actual f (system drift) minus the logarithmic gradient term computed + on current particles positions. + Provided for easy integration, and can be passed to ode integrators. + + Parameters + ---------- + x : 2d-array, + Particle positions (dimension x number of particles). + t : float, + Time t within the [t1,t2] interval. + + Returns + ------- + 2d-array + Returns the deterministic forces required to ntegrate the particle + positions for one time step, + i.e. return :math:`f(x,t)-\\frac{1}{2}\\sigma^2\\nabla \\rho_t(x)`, + evaluated at the current positions x and t. + + """ + dimi, N = x.shape + bnds = torch.zeros(dimi, 2, dtype=torch.float64, device=self.device) + for ii in range(dimi): + bnds[ii, 0] = torch.min(x[ii, :]) + bnds[ii, 1] = torch.max(x[ii, :]) + + Sxx = torch.distributions.Uniform(low=bnds[:, 0], high=bnds[:, 1]).sample([self.N_sparse]) + + + #gpsi = torch.zeros(dimi, N, dtype=torch.float64, device=self.device) + lnthsc = 2*torch.std(x, dim=1) + + + gpsi = torched_score_function_multid_seperate_all_dims(torch.t(x), + torch.t(Sxx), + func_out=False, + C=0.001, + l=lnthsc, + kern=self.kern, + device=self.device) + + return self.f_true(x, t)-0.5* self.g**2* gpsi + + + ### effective forward drift - estimated seperatelly for each dimension + #plain GP prior + def f_seperate(self, x, t): + """ + Computes the deterministic forces for the evolution of the deterministic + particles for the current particle positions, + ie. drift minus the logarithmic gradient term. + Is used as a wrapper for evolving the particles, + and can be provided to "any" ODE integrator. + + Parameters + ---------- + x : 2d-array, + Particle positions (dimension x number of particles). + t : float, + Time t within the [t1,t2] interval. + + Returns + ------- + 2d-array + Returns the deterministic forces required to ntegrate the particle + positions for one time step, + i.e. return :math:`f(x,t)-\\frac{1}{2}\\sigma^2\\nabla \\rho_t(x)`, + evaluated at the current positions x and t. + + + """ + + + dimi, N = x.shape + ### detect min and max of forward flow for each dimension + ### we want to know the state space volume of the forward flow + bnds = torch.zeros(dimi, 2, dtype=torch.float64, device=self.device) + + for ii in range(dimi): + bnds[ii, 0] = torch.min(x[ii, :]) + bnds[ii, 1] = torch.max(x[ii, :]) + # sum_bnds = np.sum(bnds) ##this is for detecting if sth goes wrong i.e. trajectories explode + # if np.isnan(sum_bnds) or np.isinf(sum_bnds): + # ##if we get unreasoble bounds just plot the first 2 dimensions of the trajectories + # plt.figure(figsize=(6, 4)), plt.plot(self.Z[0].T, self.Z[1].T, alpha=0.3) + # plt.show() + + ##these are the inducing points + ## here we select them from a uniform distribution within the state space volume spanned from the forward flow + Sxx = torch.distributions.Uniform(low=bnds[:, 0], high=bnds[:, 1]).sample([self.N_sparse]) + #print(Sxx.size()) + #print(torch.t(x).size()) + #gpsi = np.zeros((dimi, N)) + #print(Sxx.size()) + lnthsc = 2*torch.std(x, dim=1) + + gpsi = torched_score_function_multid_seperate_all_dims(torch.t(x), + Sxx, l=lnthsc, + func_out=False, + device=self.device) + + + + return self.f(x, t)-0.5* self.g**2* torch.t(gpsi) + + ###same as forward sampling but without reweighting - this is for bridge reweighting + ### not for constraint reweighting + def forward_sampling_Otto_true(self): + """ + (Relevant only when forward sampling happens with Brownian bridge + reweighting) + Same as forward sampling but without reweighting. + + Returns + ------- + int + Returns 0 to make sure everything runs correctly. + The sampled density is stored in place in the array `self.Ztr`. + + See also + --------- + DPFC.forward_sampling, DPFC.forward_sampling_Otto + + """ + logging.info('Sampling forward with deterministic particles and true drift...') + #W = np.ones((self.N,1))/self.N + for ti, tt in enumerate(self.timegrid): + + if ti == 0: + for di in range(self.dim): + self.Ztr[di, :, 0] = self.y1[di] + + elif ti == 1: #propagate one step with stochastic to avoid the delta function + #substract dt because I want the time at t-1 + self.Ztr[:, :, ti] = self.Ztr[:, :, ti-1] + \ + self.dt*self.f_true(self.Ztr[:, :, ti-1], tt-self.dt)+\ + (self.g)*torch.empty([self.dim, self.N]).normal_(mean=0, std=torch.sqrt(self.dt)) + else: + self.Ztr[:, :, ti] = self.Ztr[:, :, ti-1] + \ + self.dt* self.f_seperate_true(self.Ztr[:, :, ti-1], tt-self.dt) + if self.b_type=='kuramoto': + self.Ztr[:, :, ti] = self.Ztr[:, :, ti] %(2*np.pi) + logging.info('Forward sampling with Otto true is ready!') + return 0 + + + + def forward_sampling_Otto(self): + """ + Samples the forward probability flow with deterministic particle + dynamics. + If required at every timestep a particle reweighting takes place + employing the weights obtained from the exponentiated path constraint + :math:`U(x,t)` + + Returns + ------- + int + Returns 0 to make sure everything runs correctly. + The sampled density is stored in place in the array `self.Z`. + + """ + logging.info('Sampling forward with deterministic particles...') + W = torch.ones(self.N, 1, dtype=torch.float64, device=self.device)/self.N + for ti, tt in enumerate(self.timegrid): + if ti == 0: + for di in range(self.dim): + self.Z[di, :, 0] = self.y1[di] + if self.brown_bridge: + self.Z[di, :, -1] = self.y2[di] + ## we start forward trajectories for a delta function. + ##in principle we could start from an arbitrary distribution + ##if you want to start from a normal uncomen the following and comment the above initialisation for y1 + #self.Z[di,:,0] = np.random.normal(self.y1[di], 0.05, self.N) + elif ti == 1: #propagate one step with stochastic to avoid the delta function + #substract dt because I want the time at t-1 + self.Z[:, :, ti] = self.Z[:, :, ti-1] + \ + self.dt*self.f(self.Z[:, :, ti-1], tt-self.dt)+\ + (self.g)*\ + torch.empty([self.dim, self.N], + dtype=torch.float64, + device=self.device).normal_(mean=torch.tensor(0, + dtype=torch.float64, + device=self.device), + std=torch.sqrt(self.dt)) + else: + self.Z[:, :, ti] = self.Z[:, :, ti-1] +\ + self.dt* self.f_seperate(self.Z[:, :, ti-1], tt-self.dt) + if self.b_type=='kuramoto': + self.Z[:, :, ti] = self.Z[:, :, ti] %(2*np.pi) + + ###REWEIGHT + if self.reweight == True: + if ti > 0: + + W[:, 0] = torch.exp(self.U(self.Z[:, :, ti], tt)*self.dt) #-1 + W = W/torch.sum(W) + + ###REWEIGHT + start = time.time() + M = ot.dist(self.Z[:,:,ti].T, self.Z[:,:,ti].T) + M /= M.max() + a = W[:,0] + b = torch.ones_like(W[:,0], dtype=torch.float64, + device=self.device)/self.N + T2 = ot.emd(a, b, M) + self.Z[:, :, ti] = (self.N*self.Z[:, :, ti])@T2 + #Tstar = reweight_optimal_transport_multidim(self.Z[:, :, ti].T, W) + #print(Tstar) + if ti == 3: + stop = time.time() + logging.info('Timepoint: %d needed '%ti, stop-start) + if self.b_type=='kuramoto': + self.Z[:, :, ti] = self.Z[:, :, ti] %(2*np.pi) + logging.info('Forward sampling with Otto is ready!') + return 0 + + def density_estimation(self, ti, rev_ti): + #print(ti) + rev_t = rev_ti + grad_ln_ro = torch.zeros(self.dim, self.N, dtype=torch.float64, + device=self.device) + lnthsc = 2*torch.std(self.Z[:, :, rev_t], dim=1) + bnds = torch.zeros(self.dim, 2, dtype=torch.float64, device=self.device) + for ii in range(self.dim): + if torch.min(self.B[ii, :, rev_ti]) == torch.max(self.B[ii, :, rev_ti]): + bnds[ii, 0] = torch.min(self.Z[ii, :, rev_t]) + bnds[ii, 1] = torch.max(self.Z[ii, :, rev_t]) + else: + bnds[ii, 0] = min(torch.min(self.Z[ii, :, rev_t]), torch.min(self.B[ii, :, rev_ti])) + bnds[ii, 1] = max(torch.max(self.Z[ii, :, rev_t]), torch.max(self.B[ii, :, rev_ti])) + + #sparse points + """ + print('low', bnds[:,0]) + print('high', bnds[:,1]) + plt.figure() + plt.subplot(1,2,1) + plt.plot(self.Z[0, :, rev_t].detach().numpy(), '.', c='grey') + plt.plot(self.B[0, :, rev_t].detach().numpy(), '.', c='maroon') + plt.subplot(1,2,2) + plt.plot(self.Z[1, :, rev_t].detach().numpy(), '.', c='grey') + plt.plot(self.B[1, :, rev_t].detach().numpy(), '.', c='maroon') + plt.show() + """ + Sxx = torch.distributions.Uniform(low=bnds[:, 0], high=bnds[:, 1]).sample([self.N_sparse]) + + #estimate density from forward (Z) and evaluate at current postitions of backward particles (B) + #grad_ln_ro = score_function_multid_seperate(self.Z[:, :, rev_t].T, Sxx.T, func_out=True, C=0.001, which=1, l=lnthsc, which_dim=di+1, kern=self.kern)(self.B[:, :, rev_ti].T) + grad_ln_ro = torch.t(torched_score_function_multid_seperate_all_dims(torch.t(self.Z[:, :, rev_t]), + Sxx, + func_out=True, + C=0.001, + l=lnthsc, + kern=self.kern, + device=self.device)(torch.t(self.B[:, :, rev_ti])) ) + + return grad_ln_ro + + + def bw_density_estimation(self, rev_ti): + """ + Estimates the logaritmic gradient of the backward flow evaluated at + particle positions of the backward flow. + + + Parameters + ---------- + + rev_ti : int, + indicates the time point in the timegrid where the estimation + will take place, i.e. for time t=self.timegrid[rev_ti. + + Returns + ------- + grad_ln_b: 2d-array, + with the logarithmic gradients of the time reversed + (backward) flow (dim x N) for the timestep `rev_ti`. + + """ + grad_ln_b = torch.zeros(self.dim, self.N, dtype=torch.float64, + device=self.device) + lnthsc = 2*torch.std(self.B[:, :, rev_ti], dim=1) + + bnds = torch.zeros(self.dim, 2, dtype=torch.float64, device=self.device) + for ii in range(self.dim): + bnds[ii, 0] = min(torch.min(self.Z[ii, :, rev_ti]), torch.min(self.B[ii, :, rev_ti])) + bnds[ii, 1] = max(torch.max(self.Z[ii, :, rev_ti]), torch.max(self.B[ii, :, rev_ti])) + + #sparse points + Sxx = torch.distributions.Uniform(low=bnds[:, 0], high=bnds[:, 1]).sample([self.N_sparse]) + + + + grad_ln_b = torched_score_function_multid_seperate_all_dims(torch.t(self.B[:, :, rev_ti]), + Sxx, + func_out=False, + C=0.001, + l=lnthsc, + kern=self.kern, + device=self.device) + return grad_ln_b + + + def backward_simulation(self): + """ + Sample time reversed flow with deterministic dynamics (or stochastic if + `self.deterministic == False`). + Trajectories are stored in place in `self.B` array of dimensionality + (dim x N x timegrid.size). + `self.B` does not require a timereversion at the end, everything + is stored in the correct order. + + Returns + ------- + int + Returns 0 to ensure everything was executed correctly. + + """ + + for ti, tt in enumerate(self.timegrid[:-1]): + if ti == 0: + for di in range(self.dim): + self.B[di, :, -1] = self.y2[di] + else: + + rev_ti = self.k -ti-1 + #density estimation of forward particles + + grad_ln_ro = self.density_estimation(ti, rev_ti+1) + + if (ti == 1 and self.deterministic) or (not self.deterministic): + + self.B[:, :, rev_ti] = self.B[:, :, rev_ti+1] -\ + self.f(self.B[:, :, rev_ti+1], self.timegrid[rev_ti+1])*self.dt + \ + self.dt*self.g**2*grad_ln_ro +\ + (self.g)*\ + torch.empty([self.dim, self.N], dtype=torch.float64, device=self.device).normal_(mean=0, std=torch.sqrt(self.dt)) + else: + + grad_ln_b = torch.t(self.bw_density_estimation(rev_ti+1)) + + self.B[:, :, rev_ti] = self.B[:, :, rev_ti+1] -\ + (self.f(self.B[:, :, rev_ti+1], self.timegrid[rev_ti+1])-\ + self.g**2*grad_ln_ro +0.5*self.g**2*grad_ln_b)*self.dt + + if self.b_type=='kuramoto': + self.B[:, :, ti] = self.B[:, :, ti] %(2*np.pi) + for di in range(self.dim): + self.B[di, :, 0] = self.y1[di] + return 0 + + + """ + def reject_trajectories(self): + + Reject backward trajectories that do not reach the vicinity of the + initial point. + Deletes in place relevant rows of the `self.B` array that contains + the time reversed trajectories. + + Returns + ------- + int + Returns 0. + + + fplus = self.y1+self.f(self.y1, self.t1)*self.dt+6*self.g**2 *np.sqrt(self.dt) + fminus = self.y1+self.f(self.y1, self.t1) *self.dt-6*self.g**2 *np.sqrt(self.dt) + reverse_order = np.zeros(self.dim) + #this is an indicator if along one of the dimensions fplus + #is smaller than fminus + for iii in range(self.dim): + if fplus[iii] < fminus[iii]: + reverse_order[iii] = 1 + to_delete = np.zeros(self.N) + ##these will be one if ith trajectory is out of bounds + + ## checking if out of bounds for each dim + for iii in range(self.dim): + if reverse_order[iii] == 0: + to_delete += np.logical_not(np.logical_and(self.B[iii, :, 1] < fplus[iii], self.B[iii, :, 1] > fminus[iii])) + elif reverse_order[iii] == 1: + to_delete += np.logical_not(np.logical_and(self.B[iii, :, 1] > fplus[iii], self.B[iii, :, 1] < fminus[iii])) + + sinx = np.where(to_delete >= 1)[0] + #sinx = np.where(np.logical_or(np.logical_not(np.logical_and(self.B[0, :, 1] < fplus[0], self.B[0, :, 1] > fminus[0])), np.logical_not(np.logical_and(self.B[0, :, 1] < fplus[0], self.B[0, :, 1] > fminus[0]))))[0] + #((self.B[1,:,-2]fminus[1]) ) ))[0] + + + #logging.warning("Identified %d invalid bridge trajectories "%len(sinx)) + # if self.reject: + # logging.warning("Deleting invalid trajectories...") + # sinx = sinx[::-1] + # for element in sinx: + # self.B = np.delete(self.B, element, axis=1) + return 0 + """ + + + def calculate_u(self, grid_x, ti): + """ + Computes the control at position(s) grid_x at timestep ti + (i.e. at time self.timegrid[ti]). + + Parameters + ---------- + grid_x : ndarray, + size d x number of points to be evaluated. + ti : int, + time index in timegrid for the computation of u. + + + Returns + ------- + u_t: ndarray, + same size as grid_x. These are the controls u(grid_x, t), + where t=self.timegrid[ti]. + + """ + #a = 0.001 + #grad_dirac = lambda x,di: - 2*(x[di] -self.y2[di])* + #np.exp(- (1/a**2)* (x[0]- self.y2[0])**2)/(a**3 *np.sqrt(np.pi)) + u_t = torch.zeros(grid_x.T.shape, dtype=torch.float64, + device=self.device) + + + lnthsc1 = 2*torch.std(self.B[:, :, ti], dim=1) + lnthsc2 = 2*torch.std(self.Z[:, :, ti], dim=1) + + + bnds = torch.zeros(self.dim, 2, dtype=torch.float64, + device=self.device) + + for ii in range(self.dim): + if self.reweight == False or self.brown_bridge == False: + bnds[ii, 0] = min(torch.min(self.Z[ii, :, ti]), torch.min(self.B[ii, :, ti])) + bnds[ii, 1] = max(torch.max(self.Z[ii, :, ti]), torch.max(self.B[ii, :, ti])) + + else: + bnds[ii, 0] = min(torch.min(self.Ztr[ii, :, ti]), torch.min(self.B[ii, :, ti])) + bnds[ii, 1] = max(torch.max(self.Ztr[ii, :, ti]), torch.max(self.B[ii, :, ti])) + + if ti <= 5 or (ti >= self.k-5): + if self.reweight == False or self.brown_bridge == False: + ##for the first and last 5 timesteps, to avoid numerical singularities just assume gaussian densities + for di in range(self.dim): + mutb = torch.mean(self.B[di, :, ti]) + stdtb = torch.std(self.B[di, :, ti]) + mutz = torch.mean(self.Z[di, :, ti]) + stdtz = torch.std(self.Z[di, :, ti]) + u_t[di] = -(grid_x[:, di]- mutb)/stdtb**2 - (-(grid_x[:, di]- mutz)/stdtz**2) + elif self.reweight == True and self.brown_bridge == True: + for di in range(self.dim): + mutb = torch.mean(self.B[di, :, ti]) + stdtb = torch.std(self.B[di, :, ti]) + mutz = torch.mean(self.Ztr[di, :, ti]) + stdtz = torch.std(self.Ztr[di, :, ti]) + u_t[di] = -(grid_x[:, di]- mutb)/stdtb**2 - (-(grid_x[:, di]- mutz)/stdtz**2) + else: #if ti > 5: + ### clipping not used at the end but provided here for cases when + ### number of particles is small + ### and trajectories fall out of simulated flows + ### TO DO: add clipping as an option to be selected when + ### initialising the function + ### if point for evaluating control falls out of the region where we + ### have points, clip the points to + ### fall within the calculated region - we do not change the + ### position of the point, only the control value will be + ### calculated with clipped positions + """ + bndsb =torch.zeros(self.dim, 2, dtype=torch.float64, + device=self.device) + bndsz = torch.zeros(self.dim, 2, dtype=torch.float64, + device=self.device) + for di in range(self.dim): + bndsb[ii, 0] = torch.min(self.B[ii, :, ti]) + bndsb[ii, 1] = torch.max(self.B[ii, :, ti]) + bnds[ii, 0] = torch.min(self.Z[ii, :, ti]) + bnds[ii, 1] = torch.max(self.Z[ii, :, ti]) + """ + + ## clipping not used at the end! + ###cliping the values of points when evaluating the grad log p + grid_b = grid_x#np.clip(grid_x, bndsb[0], bndsb[1]) + grid_z = grid_x#np.clip(grid_x, bndsz[0], bndsz[1]) + + Sxx = torch.distributions.Uniform(low=bnds[:, 0], high=bnds[:, 1]).sample([self.N_sparse]) + + #for di in range(self.dim): + score_Bw = torch.t(torched_score_function_multid_seperate_all_dims(torch.t(self.B[:, :, ti]), + Sxx, + func_out=True, + C=0.001, + l=lnthsc1, + kern=self.kern, + device=self.device)(grid_b) ) + + #score_Bw = score_function_multid_seperate(self.B[:, :, ti].T, Sxx.T, func_out=True, C=0.001, which=1, l=lnthsc1, which_dim=di+1, kern=self.kern)(grid_b) + if self.reweight == False or self.brown_bridge == False: + score_Fw = torch.t( torched_score_function_multid_seperate_all_dims(torch.t(self.Z[:, :, ti]), + Sxx, + func_out=True, + C=0.001, + l=lnthsc2, + kern=self.kern, + device=self.device)(grid_z) ) + #score_Fw = score_function_multid_seperate(self.Z[:, :, ti].T, Sxx.T, func_out=True, C=0.001, which=1, l=lnthsc2, which_dim=di+1, kern=self.kern)(grid_z) + else: + """ + bndsztr = torch.zeros(self.dim, 2, dtype=torch.float64, + device=self.device) + for ii in range(self.dim): + #bndsztr[di] = [torch.min(self.Ztr[di, :, ti]), torch.max(self.Ztr[di, :, ti])] + bndsztr[ii, 0] = torch.min(self.Ztr[ii, :, rev_t]) + bndsztr[ii, 1] = torch.max(self.Ztr[ii, :, rev_t]) + """ + grid_ztr = grid_x #np.clip(grid_x, bndsztr[0], bndsztr[1]) + lnthsc3 = 2*torch.std(self.Ztr[:, :, ti], dim=1) + #score_Fw = score_function_multid_seperate(self.Ztr[:, :, ti].T, Sxx.T, func_out=True, C=0.001, which=1, l=lnthsc3, which_dim=di+1, kern=self.kern)(grid_ztr) + score_Fw = torch.t( torched_score_function_multid_seperate_all_dims(torch.t(self.Ztr[:, :, ti]), + Sxx, + func_out=True, + C=0.001, + l=lnthsc3, + kern=self.kern, + device=self.device)(grid_ztr) ) + u_t = score_Bw - score_Fw + + return u_t + + + def check_if_covered(self, X, ti): + """ + Checks if test point X falls within forward and backward densities at + timepoint timegrid[ti]. + + + Parameters + ---------- + X : array 1x dim or Kxdim + Point in state space where control is evaluated. + ti : int + Index in timegrid array indicating the time within the + time interval [t1,t2]. + + Returns + ------- + Boolean variable - True if the text point X falls within the densities. + + """ + covered = True + bnds = torch.zeros(self.dim, 2, dtype=torch.float64, + device=self.device) + for ii in range(self.dim): + bnds[ii, 0] = min(torch.min(self.Z[ii, :, ti]), torch.min(self.B[ii, :, ti])) + bnds[ii, 1] = max(torch.max(self.Z[ii, :, ti]), torch.max(self.B[ii, :, ti])) + + + + covered = covered * ((X[ii] >= bnds[ii][0]) and (X[ii] <= bnds[ii][1])) + + return covered