.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "auto_examples/core/chirp_slowing.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_auto_examples_core_chirp_slowing.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_auto_examples_core_chirp_slowing.py:


Chirped slowing
===============

Slowing of a molecular beam using radiation pressure.

- Rate equations
- 100 particles
- ...

.. GENERATED FROM PYTHON SOURCE LINES 12-23

.. code-block:: Python

    from MoleCool import System, np, open_object, save_object, plt, c
    import os.path
    from MoleCool.tools import gaussian
    from copy import deepcopy

    def return_fun(system):
        return dict(
            vx_end  = system.v[0,-1], 
            photons  = system.photons[0],
            )


.. GENERATED FROM PYTHON SOURCE LINES 24-31

Initialize levels from BaF
--------------------------

We load the predefined constants of :math:`^{138}\text{BaF}` which are stored
in the repository's folder `constants` in the file ``138BaF.json``.
For the electronic ground and excited states, the states with vibrational
quantum number ``v=0`` are loaded from the constants file.

.. GENERATED FROM PYTHON SOURCE LINES 31-39

.. code-block:: Python


    if __name__ == '__main__':    
        if not os.path.isfile("chirp_slowing.pkl"):
    
            system      = System(load_constants='138BaF')
            system.levels.add_all_levels(v_max=3)
            del system.levels['B']
        

.. GENERATED FROM PYTHON SOURCE LINES 40-42

Velocity distribution
---------------------

.. GENERATED FROM PYTHON SOURCE LINES 42-47

.. code-block:: Python
   :dedent: 1

        
            system.set_v0(
                v0          = np.linspace(70, 310, 3000),
                direction   = 'x',
                )

.. GENERATED FROM PYTHON SOURCE LINES 48-50

Laser wavelengths and chirp rates
---------------------------------

.. GENERATED FROM PYTHON SOURCE LINES 50-69

.. code-block:: Python
   :dedent: 1


            beta_arr    = np.arange(0,3*4+1,3) * 1e6/1e-3 #MHz per ms
            lamb_arr    = np.array([system.levels.wavelengths.iloc[i,j]
                                    for i,j in zip([0,1,2,3],[0,0,1,2])])*1e-9
            freq_arr    = c/lamb_arr
            freq_chirp  = -190/(c+190) * freq_arr
            lamb_start  = c/ ( freq_arr + freq_chirp)
        
            sidebands_kwargs = dict(
                sinusoidal  = dict(sidebands    = [-2,-1,1,2],
                                   ratios       = [0.8,1,1,0.8],
                                   mod_freq     = 38.4e6,
                                   offset_freq  = 19e6,
                                   ),
                perfect_fit = dict(sidebands    = [94.9, 66.9, -39.5, -56.5],
                                   ratios       = [28,   15,    17,    40],
                                   mod_freq     = 1e6,
                                   ),
                )

.. GENERATED FROM PYTHON SOURCE LINES 70-72

Calculation for two different types of sidebands
------------------------------------------------

.. GENERATED FROM PYTHON SOURCE LINES 72-102

.. code-block:: Python
   :dedent: 1

            # system.multiprocessing['processes'] = 30  
            labels  = ['sinusoidal', 'perfect_fit']
            data    = dict()
            for label in labels:
            
                del system.lasers[:]
            
                for lamb in lamb_start:
                    system.lasers.add_sidebands(
                        lamb    = lamb,
                        I       = 7000,
                        pol     = 'lin',
                        beta    = beta_arr*lamb_arr[0]/lamb,
                        k       = [-1,0,0],
                        r_k     = [0,0,0], 
                        **sidebands_kwargs[label],
                        )
            
                system.calc_rateeqs(
                    t_int           = 9e-3,
                    method          = 'LSODA',
                    magn_remixing   = True,
                    magn_strength   = 8,
                    trajectory      = True,
                    verbose         = False,
                    return_fun      = return_fun,
                    )
            
                data[label] = deepcopy(system)
            save_object(data, "chirp_slowing")

.. GENERATED FROM PYTHON SOURCE LINES 103-105

Loading results
---------------

.. GENERATED FROM PYTHON SOURCE LINES 105-147

.. code-block:: Python
   :dedent: 1


        data = open_object("chirp_slowing")
    
        plt.figure()
    
        for ls, label in zip([':', '-'], ['sinusoidal', 'perfect_fit']):
        
            system      = data[label]
            v0          = system.v0
            beta_arr    = system.lasers.getarr('beta')[0]
            vx_end      = system.results.vals['vx_end'].T
            photons     = system.results.vals['photons'].T
    
            mu          = np.array([190,0])
            sigma       = np.array([32.5448,  3.2545])
            fac         = np.exp(-(v0[:,0]-mu[0])**2/(2*sigma[0]**2))
            # fac         = gaussian(v0[:,0], x0=mu[0], std=sigma[0])
            bins        = 200
            y_offset    = 20
    
            for i,beta in enumerate(beta_arr):
                out = np.histogram(v0[:,0], bins=bins, weights=fac)
                # plt.plot(out[1][1:], out[0]+i*y_offset,
                #         color='grey',ls='-', alpha=0.4)
                plt.plot(out[1][1:], out[0]*0+i*y_offset,
                        color='k',ls='-', alpha=0.4)
            
                plt.fill_between(out[1][1:], i*y_offset, out[0]+i*y_offset,
                                color='grey',ls='-', alpha=0.1)
            
                out = np.histogram(vx_end[i,:], bins=bins, weights=fac)
                plt.plot(out[1][1:],out[0]+i*y_offset,
                        label='$\\beta$ = {:.0f}'.format(beta/(1e6/1e-3)) if ls=='-' else None,
                        color=f"C{i}", ls=ls)
            
        plt.xlim(81, 259)
        plt.ylim(bottom=-y_offset*0.1)
        plt.xlabel('Velocity (m/s)')
        plt.ylabel('Number of molecules per bin')

        plt.legend(title='Chirping rate\n(MHz/ms)')


.. GENERATED FROM PYTHON SOURCE LINES 148-151

.. image:: /_figures/core/chirp_slowing_fig1.svg
   :alt: Demo plot
   :align: center


.. _sphx_glr_download_auto_examples_core_chirp_slowing.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: chirp_slowing.ipynb <chirp_slowing.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: chirp_slowing.py <chirp_slowing.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: chirp_slowing.zip <chirp_slowing.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_