%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]]

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

for i in np.arange(0, 14):
    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.95, 355.89, 355.79, 346.62, 297.62, 316.05, 347.95, 
     442.94, 546.53, 805.21, 1322.86, 2107.50, 3229.17, 4551.69]

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

fig, ax = plt.subplots()
ax.set_ylabel('Memory usage (MiB)')
ax.set_xlabel('# streamlines')
xx = np.linspace(0, 1000000, 2)
ax.plot(xx, np.polyval(p, xx), '--')