from andi_datasets.models_theory import models_theory
Firsts steps in anomalous diffusion
Characterizing anomalous diffusion and ergodicity breaking from single trajectories
In this notebook you will find the basic tools to learn how to deal in practice with the topic of diffusion. First, few simple comments about the andi_datasets
package. We will use this python package to create trajectories of different anomalous diffusion models. You can find all details in the documentation webpage.
Let’s start by importing the models_theory
module and checking the available anomalous diffusion models.
= 1
dimensions = models_theory()._oneD()
models
= inspect.getmembers(models, inspect.ismethod)
available_models
print('The availailabe models for generating trajectories in '+str(dimensions)+'D are:\n')
for x in available_models:
print('- ' + x[0] + '\n')
The availailabe models for generating trajectories in 1D are:
- attm
- ctrw
- fbm
- lw
- sbm
The andi-datasets
package allows us then to generate trajectories according to these five models. We will learn more about them along the notebook. To generate a single trajectory of, for instance, Fractional Brownian motion (fbm) with anomalous diffusion exponent \(\alpha=0.5\) we just need to do:
= 200 # Length of the trajectories
T = 0.5 # Anomalous diffusion exponent
alpha = 2 # This corresponds to the index of fbm in the previous list
model
= models.fbm(T = T, alpha = alpha)
traj plt.plot(traj)
Moreover, the andi_datasets
library allows us to create bigger datasets of trajectories with the datasets_theory
module. Below we show how to do this. There, dataset
is a numpy
matrix whose first two columns indicate the exponent and model of each trajectory (now labels
) and the rest of the matrix is the trajectory itself (now trajs
). Let’s check how the trajectories look like:
from andi_datasets.datasets_theory import datasets_theory
= datasets_theory()
DT
= 5 # Number of trajectories
N
= DT.create_dataset(T = T, N_models = N, exponents = alpha, models = [2])
dataset
= dataset[:,:2]
labels = dataset[:,2:]
trajs
# dispay a few
for traj in trajs:
plt.plot(traj)'time')
plt.xlabel('position') plt.ylabel(
Text(0, 0.5, 'position')
Anomalous diffusion exponent
In this section we will show the basic tools to asses the anomalous diffusion exponent of a set of trajectories.
Time averaged mean squared displacement (TA-MSD)
We will now create a function that calculates the TA-MSD for a given trajectory:
(Note that this function calculates the TA-MSD for a single trajectory. Using the theoretical definition and the function below, can you rewrite the funcion such that it calculates the TA-MSD for a set of trajectories in an efficient way?)
def TAMSD(traj, t_lags):
'''
Calculates the time average mean squared displacement of a set of trajectories
Inputs: - traj: trajectory to calculate TA-MSD
- t_lags: time lags used for the TA-MSD
Outputs: - TA-MSD
'''
= np.zeros_like(t_lags, dtype= float)
tamsd for idx, t in enumerate(t_lags):
for p in range(len(traj)-t):
+= (traj[p]-traj[p+t])**2
tamsd[idx] /= len(traj)-t
tamsd[idx] return tamsd
Let’s use the dataset we created above and calculate the TAMSD of each trajectory and plot it in log-scale. As the trajectories have 200 points, let’s choose the time lages \(\Delta \in [1,20]\). In addition, we will also plot two lines: one linearly proportional to the \(\Delta\) and one proportional to \(\Delta^\alpha\). This will give you information about the ergodicity! Can you tell from here if fbm is ergodic or not?
= np.arange(1, 101)
t_lags
for traj in trajs[:5]:
= TAMSD(traj, t_lags)
tamsd
plt.loglog(t_lags,tamsd)
0.1*t_lags, c = 'k', label = r'$\sim \Delta$')
plt.loglog(t_lags, 0.1*t_lags**alpha, c = 'k', ls = '--', label = r'$\sim \Delta^\alpha$')
plt.loglog(t_lags, r'Time lag ($\Delta$)')
plt.xlabel('TA-MSD')
plt.ylabel( plt.legend()
<matplotlib.legend.Legend>
We can also calculate the ensemble averaged (EA-MSD) and time-ensemble averaged mean squared displacement (TEA-MSD):
## EA-MSD ##
= np.mean(trajs**2, axis = 0)
eamsd
plt.loglog(eamsd)1,T), 0.1*np.arange(1,T), c = 'k', label = r'$\sim$ time')
plt.loglog(np.arange(1,T), 0.1*np.arange(1,T)**alpha, c = 'k', ls = '--', label = r'$\sim$ time$^\alpha$')
plt.loglog(np.arange(r'time')
plt.xlabel('EA-MSD')
plt.ylabel( plt.legend()
<matplotlib.legend.Legend>
## TEA-MSD ##
= np.arange(1, 21)
t_lags = np.zeros_like(t_lags, dtype = 'float64')
teamsd for traj in trajs:
+= TAMSD(traj, t_lags)
teamsd /= N
teamsd
= 3)
plt.loglog(t_lags,teamsd, lw 0.07*t_lags, c = 'k', label = r'$\sim \Delta$')
plt.loglog(t_lags, 0.07*t_lags**alpha, c = 'k', ls = '--', label = r'$\sim \Delta^\alpha$')
plt.loglog(t_lags, r'time lag ($\Delta$)')
plt.xlabel('TEA-MSD')
plt.ylabel( plt.legend()
<matplotlib.legend.Legend>
A visual inspection seems to tell us that FBM is ergodic! To confirm it, let’s fit the TA-MSD and check if we recover the anomalous diffusion exponent we set above. For that, we will use a fitting function from numpy
, forcing a linear fit:
= []
fitted_alpha for traj in trajs:
= TAMSD(traj, t_lags)
tamsd 1)[0])
fitted_alpha.append(np.polyfit(np.log(t_lags), np.log(tamsd),
range(N), fitted_alpha, label = r'Fitted $\alpha$')
plt.scatter(= r'True $\alpha$')
plt.axhline(alpha, label
plt.legend()'Trajectory number')
plt.xlabel(r'$\alpha$')
plt.ylabel(0,1) plt.ylim(
(0.0, 1.0)
The TA-MSD fit seemed to work, but rather poorly… What can be improved? Solve the exercises to learn about it!
Trajectories with localization error
The trajectories we used aboved were examples from a purely theoretical fbm. However, in real life things are not so clean and in many experiments we have the appearance of localization noise. A simple way of simulating this is to add Gaussian noise to each of the positions of a trajectory. Here is an example with noise with standard deviation $ _{} = 0.5$:
= 0.3
sigma_n = trajs + np.random.randn(N,T)*sigma_n
noisy_trajs
0,:], label = 'Original trajectory')
plt.plot(trajs[0,:], label = 'Noisy trajectory', alpha = 0.8)
plt.plot(noisy_trajs['Time')
plt.xlabel('Position')
plt.ylabel( plt.legend()
<matplotlib.legend.Legend>
Let’s see what happens when we try to predict the anomalous diffusion exponent for this noisy trajectories:
= []
fitted_alpha_noisy for traj in noisy_trajs:
= TAMSD(traj, t_lags)
tamsd 1)[0])
fitted_alpha_noisy.append(np.polyfit(np.log(t_lags), np.log(tamsd),
range(N), fitted_alpha, label = r'Fitted $\alpha$ (normal trajs.)')
plt.scatter(range(N), fitted_alpha_noisy, label = r'Fitted $\alpha$ (noisy trajs.)')
plt.scatter(
= r'True $\alpha$')
plt.axhline(alpha, label
plt.legend()'Trajectory number')
plt.xlabel(r'$\alpha$')
plt.ylabel(0,1) plt.ylim(
(0.0, 1.0)
You can see that the noise has largely erased the information of the anomalous diffusion exponent!! This shows the need for better techniques if one wants to correctly characterize a trajectory. In the last part of this notebook we will show a novel approach based on ML, with much better results! Having an independent estimation of the noise allows for a correction and provides a better quantification of \(\alpha\):
= []
fitted_alpha_noisy_corr for traj in noisy_trajs:
= TAMSD(traj, t_lags) - 2*(sigma_n**2)
tamsd = np.where(tamsd>0)
i_ok 1)[0])
fitted_alpha_noisy_corr.append(np.polyfit(np.log(t_lags[i_ok]), np.log(tamsd[i_ok]),
range(N), fitted_alpha, label = r'Fitted $\alpha$ (normal trajs.)')
plt.scatter(range(N), fitted_alpha_noisy, label = r'Fitted $\alpha$ (noisy trajs.)')
plt.scatter(range(N), fitted_alpha_noisy_corr, label = r'Fitted $\alpha$ (noisy trajs. with corr.)')
plt.scatter(
= r'True $\alpha$')
plt.axhline(alpha, label
plt.legend()'Trajectory number')
plt.xlabel(r'$\alpha$')
plt.ylabel(0,1) plt.ylim(
(0.0, 1.0)
Ergodicity breaking
In the cells above we have focused on an ergodic model, FBM. However, we now know that there exist diffusion models that break ergodicity. This means that the time average mean squared displacement (TA-MSD) and the ensemble average mean squared displacement (EA-MSD) are not longer equal. Let’s explore this intriguing feature with a very well known non-ergodic model: continuous time random walk (CTRW). We will start by creating a dataset of CTRW trajectories. As we will be perfoming ensemble averages, we will need much more trajectories.
= 200 # Length of the trajectories
T = 1000 # Number of trajectories
N = 0.5 # Anomalous diffusion exponent
alpha
# Now we want CTRW trajectories (see list of available models above):
= 4
model
= DT.create_dataset(T = T, N_models=N, exponents = alpha, models = model)
dataset = dataset[:,:2]
labels = dataset[:,2:] trajs
This is how CTRW trajectories look like. You will recognize the characteristic waiting times.
for traj in trajs[:5,:]:
plt.plot(traj)'Time')
plt.xlabel('Position') plt.ylabel(
Text(0, 0.5, 'Position')
Ensemble averaged mean squared displacement (EA-MSD)
To showcase the appearance of ergodicity breaking, let’s calculate the EA-MSD and the ensemble average of the TA-MSD (TEA-MSD). For the former, if all your trajectories start at zero, it is as easy as to do the mean of the position squared! For the latter, we will use the function we define previously. We will use the same time lags we used for the FBM case.
## EA-MSD ##
= np.mean(trajs**2, axis = 0)
eamsd
## TEA-MSD ##
= np.arange(1, 21)
t_lags = np.zeros_like(t_lags, dtype = 'float64')
teamsd for traj in trajs:
+= TAMSD(traj, t_lags)
teamsd /= N
teamsd
= plt.subplots(1,2, figsize = (10, 5), constrained_layout = True)
fig, ax
0].loglog(t_lags, teamsd)
ax[0].loglog(t_lags, 1.1*teamsd[0]*t_lags, c = 'k', label = r'$\sim \Delta$')
ax[0].loglog(t_lags, 1.1*teamsd[0]*t_lags**alpha, c = 'k', ls = '--', label = r'$\sim \Delta^\alpha$')
ax[0], xlabel =r'Time lag ($\Delta$)', ylabel = 'TEA-MSD')
plt.setp(ax[0].legend()
ax[
1].loglog(eamsd)
ax[1].loglog(np.arange(1,T), eamsd[1]*np.arange(1,T), c = 'k', label = r'$\sim t$')
ax[1].loglog(np.arange(1,T), eamsd[1]*np.arange(1,T)**alpha, c = 'k', ls = '--', label = r'$\sim t^\alpha$')
ax[1], xlabel =r'Time', ylabel = 'EA-MSD')
plt.setp(ax[1].legend() ax[
<matplotlib.legend.Legend>
As we can clearly see, the TEA-MSD is linearly proportional to the time lags \(\Delta\) while the EA-MSD is propotional to \(t^\alpha\), showcasing the break of ergodicity.
Ergodicity breaking parameter
A quantitative approach to ergodicity breaking is the calculation of the ergodicity breaking parameter (EB). Here is how to calculate it:
def EB(trajs, t_lag):
= []
eb # We loop over T, which is the point at which we cut the trajectory
for T in range(t_lag+1, trajs.shape[1]):
= trajs[:,:T]
cut_trajs # Now we calculate the TAMSD for every trajectory at the
# given t_lag
= np.zeros(trajs.shape[0])
tamsd for idx, traj in enumerate(cut_trajs):
for p in range(len(traj)-t_lag):
+= (traj[p]-traj[p+t_lag])**2
tamsd[idx] /= len(traj)-t_lag
tamsd # Compute EB and append
**2)/np.mean(tamsd)**2-1)
eb.append(np.mean(tamsdreturn eb
Let’s compute the EB for similar trajectories to the ones we used in the previous section. The expected value at \(t\rightarrow\infty\) is given by Eq.(8) of this paper and for \(\alpha = 0.5\) we should have EB\(\approx 0.570\). As a bonus, let’s create some FBM trajectories to check that for ergodic processes, the EB goes to zero, even when the anomalous diffusion exponent is different than one. We will use as time lag \(\Delta = 2\). You can check the effect of varying \(\Delta\) on the EB in the previous paper (see e.g. Fig. 2(b)).
Note: the EB is much slower because you need a couple of extra loops. If you managed to create the efficient function for the TA-MSD, now is the time to use it!
= 200 # Length of the trajectories
T = 500 # Number of trajectories
N = 0.5 # Anomalous diffusion exponent
alpha
= DT.create_dataset(T = T, N_models=N, exponents = alpha, models = 2)[:,2:]
fbm = DT.create_dataset(T = T, N_models=N, exponents = alpha, models = 4)[:,2:] ctrw
= EB(ctrw, t_lag = 2)
EB_ctrw = EB(fbm, t_lag = 2) EB_fbm
= 'CTRW')
plt.plot(EB_ctrw, label = 'FBM')
plt.plot(EB_fbm, label 0.570, c = 'k', alpha = 0.4, label = r'EB($t\rightarrow\infty$) CTRW')
plt.axhline('EB(t)')
plt.ylabel('Time')
plt.xlabel( plt.legend()
<matplotlib.legend.Legend>
As you can see, the EB for the FBM trajectories quickly goes to 0. However, for the CTRW we see a value different that zero, that will converge to the theoretical value if we consider much longer trajectories. Feel free to explore how long you have to wait until reaching such value!
Machine learning approach to single trajectory characterization
As we just saw, analyzing diffusion processes at the single trajectory level has many problems related to, e.g. noise, trajectory length, ergodicity breaking,… In the last years we have seen how machine learning can indeed solve many of this. Here we will show an example on how to use a convolutional neural network to decide wether the anomalous diffusion exponent of a dataset of trajectories.
Training and test dataset
Let’s start by creating the datasets that we will use for training and evaluating the trained machine. We will consider for this example a dataset of FBM trajectories, with exponents \(\in [0.2,1.4]\) and lenght 200. For training we need a big dataset, so we will generate 10000 trajectories:
= 100
num_alphas = np.linspace(0.2,1.81, num_alphas)
exponents = 100 # this is the number of trajectories per exponent, which makes a total of 10000
N_models = 200 # Length of the trajectories
T
= DT.create_dataset(T = T, N_models=N_models, exponents = exponents, models = 2) dataset
We will now transform the previous dataset to a machine learning dataset which we can then input to the neural network. The shape constraints are set by the machines we will later use. Moreover, we will split the dataset into a training and test datasets with a 80%/20% ratio. Moreover, we will use the function normalize
from the andi-datasets
package to ensure that all the trajectories are in the same scale. This function divides the displacements of a trajectory by their standard deviation. This changes the diffusion coefficient of any trajectory to \(D=1\), but the anomalous diffusion exponent remains untouched. Moreover, it is also important to shuffle our dataset.
from andi_datasets.utils_trajectories import normalize
= int(0.8*dataset.shape[0])
ratio
np.random.shuffle(dataset)# NN inputs: trajectories
= normalize(dataset[:ratio, 2:]).reshape(ratio, T, 1)
X_train = normalize(dataset[ratio:, 2:]).reshape(N-ratio, T, 1)
X_test
# NN outputs: their anomalous exponent
= dataset[:ratio, 1]
Y_train = dataset[ratio:, 1] Y_test
Neural network architecture
from keras.models import Sequential, Model
from keras.layers import Dense, Dropout, Conv1D, Flatten, BatchNormalization
from keras.optimizers import Adam
from keras.regularizers import l2 as regularizer_l2
= Sequential()
model
=3, kernel_size=3 ,strides=1,
model.add(Conv1D(filters=(T, 1),
input_shape= 'uniform',
kernel_initializer= 'relu', kernel_regularizer = regularizer_l2(l = 0.001)))
activation=-1, momentum=0.99, epsilon=0.001))
model.add(BatchNormalization(axis
=8, kernel_size=5 ,strides=1,
model.add(Conv1D(filters= 'uniform',
kernel_initializer= 'relu', kernel_regularizer = regularizer_l2(l = 0.001)))
activation=-1, momentum=0.99, epsilon=0.001))
model.add(BatchNormalization(axis
=3, kernel_size=2 ,strides=1,
model.add(Conv1D(filters= 'uniform',
kernel_initializer= 'relu', kernel_regularizer = regularizer_l2(l = 0.001)))
activation=-1, momentum=0.99, epsilon=0.001))
model.add(BatchNormalization(axis
model.add(Flatten())0.5))
model.add(Dropout(64*2, activation='sigmoid', kernel_regularizer = regularizer_l2(l = 0.001)))
model.add(Dense(=-1, momentum=0.99, epsilon=0.001))
model.add(BatchNormalization(axis
0.5))
model.add(Dropout(64, activation='sigmoid'))
model.add(Dense(=-1, momentum=0.99, epsilon=0.001)) model.add(BatchNormalization(axis
As we want to extract the anomalous diffusion exponent, the output layer needs to have size one. It is also important that its activation function is linear.
1, activation='linear')) model.add(Dense(
Now that we have all the layers, we can take a look at a visual summary of it. The plot below also gives us the input/output shape for every layer!
from keras.utils.vis_utils import plot_model
=True, show_layer_names = False) plot_model(model, show_shapes
Now we compile the model. For the loss function, we will use the mean absolute error (MAE). Another possible choice is the mean squared error (MSE). We will keep track of it as metric. We will use Adam as optimizer.
compile(loss='mean_absolute_error',
model.=Adam()) optimizer
At last, we train!
= 200
batch_size = 150
epochs
= model.fit(X_train, Y_train,
training =batch_size,
batch_size=epochs,
epochs=2,
verbose=(X_test, Y_test)) validation_data
In order to see if the training worked, let’s take a look at the training and test dataset loss:
= training.history['loss']
acc = training.history['val_loss']
val_acc
len(training.history['loss'])), acc, label='Training loss')
plt.plot(np.arange(len(training.history['loss'])), val_acc,label='Validation loss')
plt.plot(np.arange('Epochs')
plt.xlabel('Accuracy')
plt.ylabel( plt.legend()
<matplotlib.legend.Legend>
Let’s create a validation dataset with few exponents and check the predictions:
= np.linspace(0.5, 1.5, 4)
validation_exponents
= DT.create_dataset(T = T, N_models=1000, exponents = validation_exponents, models = 2)
val_dataset = normalize(val_dataset[:, 2:]).reshape(val_dataset.shape[0], T, 1)
X_val = val_dataset[:, 1]
Y_val = model.predict(X_val) preds_val
125/125 [==============================] - 0s 2ms/step
Now we can plot de distribution of predicted exponents for each of the groundtruth ones:
= plt.subplots(figsize = (15,3))
fig, ax for idx, u in enumerate(np.unique(Y_val)):
== u], facecolor = f'C{idx}', bins = 100, alpha = 0.4, label = 'Predicted value' if idx == 0 else '')
ax.hist(preds_val[Y_val = '--', c = f'C{idx}', label = 'True value' if idx == 0 else '')
plt.axvline(u, ls
plt.legend()'Anomalous exponent') plt.xlabel(
Text(0.5, 0, 'Anomalous exponent')
As we can see the methods is able to correctly predict the exponent! The bigger the error, the more spread this distributions will be. Can you improve the model? Go for it!
Extracting other parameters with ML
Now that we defined our network, it is really easy to transform it for any other task. For instance, let’s see how to create a machine that distinguishes between ergodic and non-ergodic trajectories. First we need to create a new dataset with all the andi-datasets
diffusion models. Due to the constraints of the package, we will need to create the subdiffusive and superdiffusive datasets separately and then merge them:
= 200 # this is the number of trajectories per exponent, which makes a total of 10000
N = 200 # Length of the trajectories
T
= np.arange(0.2, 1, 0.2) # We only consider subdiffusion
subdiffusive_exponents = [0, 1, 2, 4] # All models but LW
subdiffusive_models
= DT.create_dataset(T = T, N_models = N, exponents = subdiffusive_exponents, models = subdiffusive_models)
sub_dataset
= np.arange(1, 2, 0.2) # We only consider superdiffusion
superdiffusive_exponents = [2, 3, 4] # All models but CTRW and ATTM
superdiffusive_models
= DT.create_dataset(T = T, N_models = N, exponents = superdiffusive_exponents, models = superdiffusive_models)
super_dataset
= np.vstack((sub_dataset, super_dataset)) dataset
The labels now need to be transform such that for ATTM, CTRW and SBM we have 0 (non-ergodic) and for FBM and LW we will have 1 (ergodic):
= int(0.8*dataset.shape[0])
ratio
np.random.shuffle(dataset)# NN inputs: trajectories
= normalize(dataset[:ratio, 2:]).reshape(ratio, T, 1)
X_train = normalize(dataset[ratio:, 2:]).reshape(N-ratio, T, 1)
X_test
# NN outputs
# First we take the models
= dataset[:ratio, 0]
Y_train = dataset[ratio:, 0]
Y_test
= np.argwhere((Y_train == 2) | (Y_train == 3))
where_ergodic = np.argwhere((Y_train == 0) | (Y_train == 1) | (Y_train == 4))
where_nonergodic = 0
Y_train[where_ergodic] = 0
Y_train[where_nonergodic]
= np.argwhere((Y_test == 2) | (Y_test == 3))
where_ergodic = np.argwhere((Y_test == 0) | (Y_test == 1) | (Y_test == 4))
where_nonergodic = 0
Y_test[where_ergodic] = 0 Y_test[where_nonergodic]
Finally, we can use the same architecture as before, but will need to change the last layer to two neurons with softmax activation, as e.g.:
2, activation='softmax')) model.add(Dense(
Now you are on your own, can you train this network?