%load_ext memory_profiler

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
%matplotlib inline

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(
    '/Users/arokem/source/dipy/doc/examples/lr-superiorfrontal.trk',
                                                  points_space='voxel')[0]]

candidate_sl = candidate_sl * 6

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

len(candidate_sl)

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

m = [353.93, 355.88, 356.99, 347.34, 297.30, 316.82, 348.77, 
     439.76, 547.12, 805.96, 1322.99, 2254.77, 3012.15, 4300.14, 4553.35, 4592.92]     

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