#!/usr/bin/env python
# coding: utf-8

# In[1]:


import wot


# # Notebook 3: Computing long-range couplings
# 
# In Notebook 2 we saw how to compute transport matrices $\pi_{t_i,t_{i+1}}$ connecting adjacent time-points $(t_i,t_{i+1})$. Over short time-scales, this OT matrix is a good approximation to the true temporal coupling:
# 
# $$ \pi_{t_i,t_{i+1}} \approx \gamma_{t_i,t_{i+1}}. $$
# 
# In this notebook we show how to infer transitions over a longer time interval $(t_i,t_j)$.
# To do this, we assume the developmental stochastic process ${\mathbb{P}}_t$ is Markov. This means that the couplings $\gamma_{s,t}$ satisfy the following algebraic relationship:
# 
# $$ \gamma_{t_1,t_3}(x,z) = \iint \gamma_{t_1,t_2}(x,y) \gamma_{t_2,t_3}(y,z) dy$$
# 
# for any $t_1 < t_2 < t_3$. The integral on the right hand side becomes a finite sum when we have a finite sample of cells. In fact, it's a matrix multiplication! 
# 
# Therefore we can infer long-range transitions by composing transport maps as follows:
# 
# $$\gamma_{t_i,t_j} = \gamma_{t_i,t_{i+1}} \circ \gamma_{t_{i+1},t_{i+2}} \circ \cdots \circ \gamma_{t_{j-1},t_{j}} 
# \approx \pi_{t_i,t_{i+1}} \circ \pi_{t_{i+1},t_{i+2}} \circ \cdots \circ \pi_{t_{j-1},t_{j}}.$$
# 
# Here $\circ$ denotes matrix multiplication. This is of course different than directly computing OT over longer time-intervals:
# 
# $$ \pi_{t_1,t_2} \circ \pi_{t_2,t_3} \ne \pi_{t_1,t_3}.$$
# 
# In the following code block we construct a TransportMapModel from a directory of pre-computed transport matrices. 

# In[2]:


tmap_model = wot.tmap.TransportMapModel.from_directory('tmaps/serum')


# # Composing couplings
# 
# We can now easily compute the coupling between any pair of time-points as follows:

# In[3]:


gamma_8_10 = tmap_model.get_coupling(8, 10)
gamma_8_10


# The resulting temporal coupling $\gamma_{t_i,t_j}$ has a row for each cell at time $t_i$ and a column for each cell at time $t_j$. 
# Just like for a short-range coupling, the units are "transported mass". So a value of $\gamma_{t_i,t_j}(x,y) = 0.2$ means that cell $x$ will have on average $0.2$ descendants with expression profile similar to $y$ at time $t_j$.
# Note that the sum of a row shows the total number of descendants that a cell will have at time $t_j$.