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

# In[1]:


get_ipython().run_line_magic('load_ext', 'memory_profiler')


# In[2]:


import numpy as np
import matplotlib
import matplotlib.pyplot as plt
get_ipython().run_line_magic('matplotlib', 'inline')


# In[3]:


from dipy.data import read_stanford_labels
import nibabel as nib
hardi_img, gtab, labels_img = read_stanford_labels()
data = hardi_img.get_data()
candidate_sl = [s[0] for s in nib.trackvis.read('./probabilistic_small_sphere.trk', points_space='voxel')[0]]


# In[4]:


import dipy.tracking.life as life
fiber_model = life.FiberModel(gtab)


# In[5]:


len(candidate_sl)


# In[6]:


for i in np.arange(0, 19):
    this = 2 ** i
    print("Number of streamlines: %s"%this)
    get_ipython().run_line_magic('memit', 'fiber_fit = fiber_model.fit(data, candidate_sl[:this], affine=np.eye(4))')


# In[9]:


m = [3261.71, 3261.73, 3261.75, 3261.76, 3261.76, 3261.76,
     3261.76, 3261.76, 3261.76, 3261.76, 3261.76, 3276.14,
     3321.05, 3414.68, 3617.24, 4721.73, 8826.69, 20247.30, 55127.31]     


# In[10]:


fig, ax = plt.subplots()
ax.plot(2 ** np.arange(19), m, 'o-')
ax.set_ylabel('Memory usage (MiB)')
ax.set_xlabel('# streamlines')


# In[ ]: