#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
solve Bloch-Lorenz equations for single-mode laser,
coupled to two-level medium.
"""

from pylab import *
from scipy.integrate import odeint

global pumpRate, gammaI, gammaP, kappa, deltaC, delta

Q_plot = True
xsize = 9.5
ysize = xsize*0.5*(sqrt(5.) - 1.)
mysize = 20; tsize = mysize - 4

def BlochLorenz(y, t):
"""
Bloch-Lorenz three-variable Bloch-Lorenz equations
Variables are [F, P, inv] with
F = field (complex amplitude), damping kappa, detuning deltaC
P = polarization (complex), damping gammaP, detuning delta
inv = inversion (real), damping 2 gammaI, pumping pumpRate
"""
global pumpRate, gammaI, gammaP, kappa, deltaC, delta

# restore components of variable vector
F = y;
P = y;
inv = y;

# simplest semiclassical laser model with medium and cavity
# dynamics, must be re-interpreted as real function for
# Python ODE integrator (see r_BlochLorentz below)
#
Fdot = -(kappa + 1j*deltaC)*F + 1j*kappa*P;
Pdot = -(gammaP + 1j*delta)*P - 1j*gammaI*(F*inv);
invDot = -2*gammaI*inv + pumpRate + gammaI*imag(conj(F)*P);

# F = 1j*kappa*P/(kappa + 1j*deltaC);
# P = -i*gammaP*(F*inv)/(gammaP + 1j*delta);
# inv = (pumpRate + gammaI*imag(conj(F)*P))/(2*gammaI);

# collect output variables in vector
return [ Fdot, Pdot, invDot ]

def r_BlochLorenz(x, t, *args):
# map defined on real parameters
y = x.view(complex128)
dydt = BlochLorenz(y, t, *args)
return asarray(dydt, dtype=complex128).view(float64)

def steadyState(t_end, dt_end = 0.1, Q_error = False):
"""
"""
global pumpRate, gammaI, gammaP, kappa, deltaC, delta
# initial data
#
F_0 = 0.3*1j # field amplitude, need some small "seeding value"
P_0 = 0.0*1j # polarisation
inv_0 = 0. # inversion, a negative value could also be used, but
# its magnitude would be somewhat arbitrary
#
# parameters
#
gammaI, gammaP = 1., 10.
kappa = 1.0*gammaP
deltaC = -1.1*kappa # laser (cavity, field) detuning
delta = 1.1*kappa # medium polarisation detuning
f_shift = 1. # common frequency shift does not affect intensities
delta += f_shift
deltaC += f_shift

# time interval
nb_t = 100
t = linspace(0., t_end, nb_t)
# convert complex initial data to real
y_0 = array([F_0, P_0, inv_0], dtype = complex)
outode = odeint(r_BlochLorenz, y_0.view(float64), t, args = (), full_output = 1)
outode = outode.view(complex128)
#
F_t = outode[:,0]
P_t = outode[:,1]
inv_t = real(outode[:,2])
# imInv_t = imag(outode[:,2]) # checked that imaginary part is zero as it should

# 'measure' stationary values from last 10% of time
t_meas = (1. - dt_end)*t_end #
i_meas = find(t >= t_meas)
I_stat = mean(abs(F_t[i_meas])**2)
dI_stat = std(abs(F_t[i_meas])**2)
P2_stat = mean(abs(P_t[i_meas])**2)
dP2_stat = std(abs(P_t[i_meas])**2)
inv_stat = mean(inv_t[i_meas])
dInv_stat = std(inv_t[i_meas])

if Q_error:
print('relative errors in stationary states')
print('dI / I = %4.4g' %(dI_stat/I_stat))
print('dP^2 / P^2 = %4.4g' %(dP2_stat/P2_stat))
print('d inv / inv = %4.4g' %(dInv_stat/inv_stat))

return I_stat, P2_stat, inv_stat

if True:
# initial data
F_0 = 0.3*1j # field amplitude
P_0 = 0.01*1j # polarisation
P_0 = 0.0*1j # polarisation
inv_0 = 0. # inversion

# parameters
#
gammaI, gammaP = 1., 2.
pumpRate = 5.2*gammaI
kappa = 2.0*gammaP
deltaC = -0.1*kappa # laser (cavity, field) detuning
delta = 0.1*kappa # medium polarisation detuning
f_shift = 0.1*kappa # common frequency shift does not affect intensities
delta += f_shift
deltaC += f_shift

# time interval
nb_t = 300
t_end = 30./gammaI
t = linspace(0., t_end, nb_t)

# --- differential equation solver ---
# odeint evaluates the solution immediately (no object syntax required)
# convert complex initial data to real
y_0 = array([F_0, P_0, inv_0], dtype = complex)
outode = odeint(r_BlochLorenz, y_0.view(float64), t, args = (), full_output = 1)
outode = outode.view(complex128)
#
F_t = outode[:,0]
P_t = outode[:,1]
inv_t = real(outode[:,2])
# imInv_t = imag(outode[:,2]) # checked that imaginary part is zero as it should

# 'measure' stationary values from last 10% of time
t_meas = 0.9*t_end #
i_meas = find(t >= t_meas)
I_stat = mean(abs(F_t[i_meas])**2)
dI_stat = std(abs(F_t[i_meas])**2)
P2_stat = mean(abs(P_t[i_meas])**2)
dP2_stat = std(abs(P_t[i_meas])**2)
inv_stat = mean(inv_t[i_meas])
dInv_stat = std(inv_t[i_meas])
print('Bloch-Lorenz model solved to time gamma_I t = %3.3g ' %(gammaI*t_end))
print('stationary data')
print('   field intensity %4.4g +- %4.4g' %(I_stat, dI_stat))
print("   polar'n intensity %4.4g +- %4.4g" %(P2_stat, dP2_stat))
print('   inversion %4.4g +- %4.4g' %(inv_stat, dInv_stat))

if Q_plot:
figure(1, figsize = (xsize, ysize), tight_layout = True)
plot( gammaI*t, inv_t, label = 'inversion' )
plot( gammaI*t, abs(F_t)**2, lw = 2, label = 'intensity' )
plot( gammaI*t, abs(P_t)**2, '--', label = '|pol|^2' )
xlabel( r'$\mathsf{time\ } \gamma_I t$', size = mysize)
legend(loc = 'upper right', fontsize = mysize)
xticks(size = tsize); yticks(size = tsize)

Q_pumpScan = False

if Q_pumpScan:
# scan pump rate and make a plot of laser threshold
nb_p = 15
pump_s = linspace(0.2*pumpRate, 1.5*pumpRate, nb_p)
I_s = zeros(nb_p)
inv_s = zeros(nb_p)
for ii, pumpRate in enumerate(pump_s):
print ii,
I_stat, P2_stat, inv_stat = steadyState(t_end, Q_error = True)
I_s[ii] = I_stat
inv_s[ii] = inv_stat
#
figure(2, figsize = (xsize, ysize), tight_layout = True)
plot( pump_s / gammaI, I_s, label = 'intensity' )
plot( pump_s / gammaI, inv_s, label = 'inversion' )
xlabel( r'$\mathsf{pump\ rate\ } R / \gamma_I$', size = mysize)
legend(loc = 'upper left', fontsize = mysize)
xticks(size = tsize); yticks(size = tsize)

show(block=False)  Solutions of the Lorentz model: intensity of the cavity mode, inversion of the medium, modulus squared of the medium polarisation. Arbitrary units, time in units of the decay time for the inversion. (left) large pumping rate (right) smaller pumping rate, just above threshold. For even smaller rates, the intensity stays zero at all times. Initial values: medium in ground state, cavity field with small amplitude. For large pumping, one observes relaxation oscillations before the system settles into a stationary state.