mb_solve_json = """
{
  "atom": {
    "decays": [
      {
        "channels": [[0, 1]],
        "rate": 1.0
      }
    ],
    "fields": [
      {
        "coupled_levels": [[0, 1]],
        "rabi_freq": 1.0e-3,
        "rabi_freq_t_args": {
          "ampl": 1.0,
          "centre": 0.0,
          "fwhm": 1.0
        },
        "rabi_freq_t_func": "gaussian"
      }
    ],
    "num_states": 2
  },
  "t_min": -2.0,
  "t_max": 10.0,
  "t_steps": 1000,
  "z_min": -0.2,
  "z_max": 1.2,
  "z_steps": 50,
  "interaction_strengths": [
    1.0
  ],
  "savefile": "spectral-analysis"
}
"""

from maxwellbloch import mb_solve
mbs = mb_solve.MBSolve().from_json_str(mb_solve_json)

Omegas_zt, states_zt = mbs.mbsolve(recalc=False)

"""
Explanation: Spectral Analysis
We can study the response of a quantised system to an input field using spectral analysis. For linear systems, analytic response functions are known. We can perform spectral analysis by transforming into the frequency domain.
Here we'll look at how to use the spectral module to obtain the absorption and dispersion profiles by transforming into the frequency domain using a discrete Fourier transform.
First we'll define a typical two-level problem. We'll need to use a high resolution in the time domain if we want to consider high frequencies above and below resonance. So here, note that we use a lot of time steps. We wouldn't usually want this! It will take a long time to save and result in a large savefile.
End of explanation
"""


fig = plt.figure(1, figsize=(16, 6))
ax = fig.add_subplot(111)
cmap_range = np.linspace(0.0, 1.0e-3, 11)
cf = ax.contourf(mbs.tlist, mbs.zlist, 
                 np.abs(mbs.Omegas_zt[0]/(2*np.pi)), 
                 cmap_range, cmap=plt.cm.Blues)
ax.set_title('Rabi Frequency ($\Gamma / 2\pi $)')
ax.set_xlabel('Time ($1/\Gamma$)')
ax.set_ylabel('Distance ($L$)')
for y in [0.0, 1.0]:
    ax.axhline(y, c='grey', lw=1.0, ls='dotted')
plt.colorbar(cf);

"""
Explanation: Results in the Time Domain
First we'll look at the results in the time domain as usual.
End of explanation
"""


from maxwellbloch import spectral, utility

pal = sns.color_palette('deep')

fig = plt.figure(2, figsize=(16, 6))
ax = fig.add_subplot(111)
freq_list = spectral.freq_list(mbs)
ax.plot(freq_list, spectral.absorption(mbs, field_idx=0), 
        label='Absorption', c=pal[0], lw=3.0)
ax.plot(freq_list, spectral.dispersion(mbs, field_idx=0), 
        label='Dispersion', c=pal[1], lw=3.0)
rabi_freq_abs_0 = np.abs(spectral.rabi_freq(mbs, 0))[0]
ax.plot(freq_list, 
        rabi_freq_abs_0/np.max(rabi_freq_abs_0),
        label='Pulse, Normalised', ls='dashed', c=pal[7])

ax.set_xlim(-3.0, 3.0)
ax.set_ylim(-1.0, 1.0)
ax.set_xlabel('Frequency ($\Gamma$)')
ax.legend();

"""
Explanation: As we can see, it's the typical absorption picture we've seen before.
Results in the Frequency Domain
The envelope of a short Gaussian pulse on resonance is equivalent to a superposition of waves across a wide range of frequencies. We can analyse the spectral response of the medium to the field by transforming the simulated Rabi frequency in the time domain, $\Omega(z, t)$, to the frequency domain, $\Omega(z, \omega)$.
We determine the absoprtion profile by measuring the intensity of the pulse that arrives at the back of the medium relative to the intensity put in at the front. This we can relate to the imaginary part of the optical susceptibility $\chi_I$ via
$$
\frac{k}{2} \chi_I(\omega) z = -\log \frac{| \Omega (z, \omega) |}{| \Omega(z = 0, \omega) |}
$$
where $k$ is the wavenumber.
We determine the dispersion profile by measuring the phase shift of the pulse that arrives at the back of the medium relative to the pulse that arrived at the front. This we can relate to the real part of optical susceptibility  $\chi_R$ via
$$
\frac{k}{2} \chi_R(\omega) z = \phi(z, \omega) - \phi(z=0, \omega)
$$
where $\phi$ is the complex phase defined via
$$
\Omega(z, \omega) = |\Omega|\mathrm{e}^{\mathrm{i}\phi}.
$$
See section 2.5 of my thesis for an explanation of how the susceptibility is derived.
Methods to provide the dispersion and absorption of a given field are provided in the spectral module.
End of explanation
"""


interaction_strength = mbs.interaction_strengths[0]
decay_rate = mbs.atom.decays[0]['rate']
absorption_linear_known = spectral.absorption_two_linear_known(freq_list, 
    interaction_strength, decay_rate)
dispersion_linear_known = spectral.dispersion_two_linear_known(freq_list, 
    interaction_strength, decay_rate)

fig = plt.figure(4, figsize=(16, 6))
ax = fig.add_subplot(111)

ax.plot(freq_list, spectral.absorption(mbs, 0, -1), 
        label='Absorption', lw=5.0, c=pal[0])
ax.plot(freq_list, spectral.dispersion(mbs, 0, -1), 
        label='Dispersion', lw=5.0, c=pal[1])

ax.plot(freq_list, absorption_linear_known, ls='dashed', c='white', lw=2.0)
ax.plot(freq_list, dispersion_linear_known, ls='dashed', c='white', lw=2.0)

# Widths
hm, r1, r2 = utility.half_max_roots(freq_list, spectral.absorption(mbs, field_idx=0))
plt.hlines(y=hm, xmin=r1, xmax=r2, linestyle='dotted', color=pal[0])
plt.annotate('FWHM: ' + '%0.2f'%(r2 - r1), xy=(r2, hm), color=pal[0],
              xycoords='data', xytext=(5, 5), textcoords='offset points');

ax.set_xlim(-3.0, 3.0)
ax.set_ylim(-1.0, 1.0)
ax.set_xlabel('Frequency ($\Gamma$)')

ax.legend();

# Plot residuals
fig = plt.figure(figsize=(16, 2))
ax = fig.add_subplot(111)
ax.plot(freq_list, spectral.absorption(mbs, 0, -1) - absorption_linear_known, 
        label='Absorption', lw=2.0, c=pal[0])
ax.set_xlim(-3.0, 3.0)
ax.set_ylim(-3e-2, 3e-2)
ax.set_xlabel('Frequency ($\Gamma$)');

"""
Explanation: The input pulse profile (normalised) is shown in grey. In blue we see the familiar Lorentzian absorption profile of homogeneous broadening due to the spontaneous decay. The full-width at half maximum (FWHM) of the peak is the decay rate, $1 \Gamma$. The associated dispersion profile is shown in orange, describing the phase of the input field relative to the frequency of the oscillator as it passes over resonance from lagging to leading.
Comparing with Analytic Results
For a weak input field like this one, the field is in the regime of linear optics and it is possible to use the weak probe approximation to derive an analytic lineshape
$$
\frac{k}{2} \chi({\omega}) = -Ng \frac{1}{\mathrm{i}\frac{\Gamma}{2} + \omega}.
$$
For details on how this is derived, see equation (2.61) of my thesis.
These analytic absorption and dispersion profiles for weak fields are also included in the spectral module, so you can check if you are really in the linear regime.
End of explanation
"""


