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

# # Tutorial on JIT versus Scipy execution within Parcels

# This very brief tutorial is meant to highlight the potentially very big difference between the computational time required to run Parcels in JIT (Just-In-Time compliation) versus in Scipy mode.
# 
# In the code snippet below, we use `AdvectionRK4` to advect 100 particles in the peninsula `FieldSet`. We first do it in JIT mode (by setting `ptype=JITParticle` in the declaration of `pset`) and then we also do it in Scipy mode (by setting `ptype=ScipyParticle` in the declaration of `pset`).
# 
# In both cases, we advect the particles for 1 hour, with a timestep of 30 seconds.
# 
# To measure the computational time, we use the `timer` module.

# In[2]:


from parcels import FieldSet, ParticleSet, JITParticle, ScipyParticle
from parcels import AdvectionRK4
from parcels import timer
from datetime import timedelta as delta

timer.root = timer.Timer('root')

timer.fieldset = timer.Timer('fieldset creation', parent=timer.root)
fieldset = FieldSet.from_parcels('Peninsula_data/peninsula', allow_time_extrapolation=True)
timer.fieldset.stop()

ptype = {'scipy': ScipyParticle, 'jit': JITParticle}
ptimer = {'scipy': timer.Timer('scipy', parent=timer.root, start=False),
          'jit': timer.Timer('jit', parent=timer.root, start=False)}

for p in ['scipy', 'jit']:
    pset = ParticleSet.from_line(fieldset=fieldset, pclass=ptype[p], size=100, start=(0.1, 0.1), finish=(0.1, 0.4))

    ptimer[p].start()
    pset.execute(AdvectionRK4, runtime=delta(hours=1), dt=delta(seconds=30))
    ptimer[p].stop()

timer.root.stop()
timer.root.print_tree()


# As you can see above, **Scipy mode took more than 20 times as long** (17.2 seconds versus 0.79 seconds) as the JIT mode!!
# 
# This is just an illustrative example, depending on the number of calls to `AdvectionRK4`, the size of the `FieldSet`, the size of the `pset`, the ratio between `dt` and `outputdt` in the `.execute` etc, the difference between JIT and Scipy can vary significantly. However, JIT will almost always be faster!
# 
# So why does Parcels support both JIT and Scipy mode then? Because Scipy is easier to debug when writing custom kernels, so can provide faster development of new features.

# *As an aside, you may wonder why we use the `time.time` module, and not the timeit module, to time the runs above. That's because it affects the AST of the kernels, causing errors in JIT mode.*