# -*- coding: utf-8 -*-
"""
Created on Fri Feb 18 17:24:21 2022
@author: voumardt
"""
import numba
import numpy as np
import pyfftw
from scipy.constants import c
import time as timing
import matplotlib.pyplot as plt
[docs]class Model():
"""
Parent class for the different nonlinear Schrödinger equations. Provide
general utility, optimization and setup functions for the physics
happening in the child classes.
"""
def __init__(self, waveguide, light):
"""
Construct model class. Some attributes are initialized in other methods.
Parameters
----------
waveguide : Waveguide
Object containing the material parameters.
light : Light
Object containing the light parameters.
"""
self.waveguide = waveguide
self.light = light
self.half_pts = self.waveguide.t_pts//2
def _setup_aaspm(self):
"""
Setup storage, ffts and selection array for anti-aliasing spm
"""
# padded fft
self.spm_pad_factor = 2 # On the Elimination of Aliasing in a Finite-Difference Schemes by Filtering High Wavenumber
self._fft_a_aaspm = pyfftw.empty_aligned(self.spm_pad_factor*self.waveguide.t_pts, dtype='complex128')
self._fft_b_aaspm = pyfftw.empty_aligned(self.spm_pad_factor*self.waveguide.t_pts, dtype='complex128')
self._fft_c_aaspm = pyfftw.empty_aligned(self.spm_pad_factor*self.waveguide.t_pts, dtype='complex128')
self._fft_d_aaspm = pyfftw.empty_aligned(self.spm_pad_factor*self.waveguide.t_pts, dtype='complex128')
self.fft_spm = pyfftw.FFTW(self._fft_a_aaspm, self._fft_b_aaspm)
self.ifft_spm = pyfftw.FFTW(self._fft_c_aaspm, self._fft_d_aaspm, direction='FFTW_BACKWARD')
# selection arrays
self.selection_array_spm = np.zeros(self.spm_pad_factor*self.waveguide.t_pts, dtype=bool)
self.selection_array_spm[:self.half_pts] = True
self.selection_array_spm[-self.half_pts:] = True
self.storage_array_spm = np.zeros(self.spm_pad_factor*self.waveguide.t_pts, dtype='complex128')
# Raman term
self.raman_effect = np.zeros(self.spm_pad_factor*self.waveguide.t_pts, dtype='complex128')
self.raman_effect[:self.half_pts] = self.waveguide.raman_effect[:self.half_pts]
self.raman_effect[-self.half_pts:] = self.waveguide.raman_effect[self.half_pts:]
self.raman_effect += (1 - self.waveguide.raman_fraction)
self.raman_effect *= 3*self.spm_pad_factor**2
def _setup_aashg(self):
"""
Setup storage, selection array and exponents for anti-aliasing chi 2
frequency conversion
"""
# padded fft
self.shg_pad_exp = np.ceil(np.log2(3/2 + self.waveguide.center_omega/2/np.amax(self.waveguide.rel_omega)))
self.shg_pts_pad = int(self.waveguide.t_pts*2**self.shg_pad_exp)
self._fft_a_shg = pyfftw.empty_aligned(self.shg_pts_pad, dtype='complex128')
self._fft_b_shg = pyfftw.empty_aligned(self.shg_pts_pad, dtype='complex128')
self._fft_c_shg = pyfftw.empty_aligned(self.shg_pts_pad, dtype='complex128')
self._fft_d_shg = pyfftw.empty_aligned(self.shg_pts_pad, dtype='complex128')
self.fft_shg = pyfftw.FFTW(self._fft_a_shg, self._fft_b_shg)
self.ifft_shg = pyfftw.FFTW(self._fft_c_shg, self._fft_d_shg, direction='FFTW_BACKWARD')
# selection arrays
self.selection_array_shg = np.zeros(self.shg_pts_pad, dtype=bool)
self.selection_array_shg[:self.half_pts] = True
self.selection_array_shg[-self.half_pts:] = True
self.storage_array_shg = np.zeros(self.shg_pts_pad, dtype='complex128')
# frequency shifting term
self.shg_time = np.linspace(np.amin(self.waveguide.time), np.amax(self.waveguide.time), self.shg_pts_pad)
self.exponent_sfg = numbaexp(1j*self.waveguide.center_omega*self.shg_time)*2**self.shg_pad_exp
self.exponent_dfg = numbaconj(self.exponent_sfg)
def _setup_aathg(self):
"""
Setup storage, selection array and exponents for anti-aliasing chi 3
frequency conversion
"""
# padded fft
self.thg_pad_exp = np.ceil(np.log2(2 + self.waveguide.center_omega/np.amax(self.waveguide.rel_omega)))
self.thg_pts_pad = int(self.waveguide.t_pts*2**self.thg_pad_exp)
self._fft_a_thg = pyfftw.empty_aligned(self.thg_pts_pad, dtype='complex128')
self._fft_b_thg = pyfftw.empty_aligned(self.thg_pts_pad, dtype='complex128')
self._fft_c_thg = pyfftw.empty_aligned(self.thg_pts_pad, dtype='complex128')
self._fft_d_thg = pyfftw.empty_aligned(self.thg_pts_pad, dtype='complex128')
self.fft_thg = pyfftw.FFTW(self._fft_a_thg, self._fft_b_thg)
self.ifft_thg = pyfftw.FFTW(self._fft_c_thg, self._fft_d_thg, direction='FFTW_BACKWARD')
# selection arrays
self.selection_array_thg = np.zeros(self.thg_pts_pad, dtype=bool)
self.selection_array_thg[:self.half_pts] = True
self.selection_array_thg[-self.half_pts:] = True
self.storage_array_thg = np.zeros(self.thg_pts_pad, dtype='complex128')
# frequency shifting terms
self.thg_time = np.linspace(np.amin(self.waveguide.time), np.amax(self.waveguide.time), self.thg_pts_pad)
self.exponent_thg = numbaexp(2j*self.waveguide.center_omega*self.thg_time)*(2**self.thg_pad_exp)**2
"""
Nonlinear operations
"""
def _aaspm(self, field_f):
"""
Compute anti-aliasing self-phase modulation term with Raman effect.
Parameters
----------
field_f : array
Field in frequency domain.
Returns
-------
array
Anti-aliased SPM term in frequency domain.
"""
self.storage_array_spm[:self.half_pts] = field_f[:self.half_pts]
self.storage_array_spm[-self.half_pts:] = field_f[self.half_pts:]
field_t_pad = numbacopy(self.ifft_spm(self.storage_array_spm))
int_f = self.fft_spm(numbaabs2(field_t_pad))
spm = self.ifft_spm(int_f*self.raman_effect)*field_t_pad
return self.fft_spm(spm)[self.selection_array_spm]
def _aash(self, field_f):
"""
Compute combined anti-aliasing sum frequency generation and difference
frequency generation
Parameters
----------
field_f : array
Field in frequency domain.
Returns
-------
array
Anti-aliased SFG and DFG in frequency domain.
"""
self.storage_array_shg[:self.half_pts] = field_f[:self.half_pts]
self.storage_array_shg[-self.half_pts:] = field_f[self.half_pts:]
sht = numbasht(self.ifft_shg(self.storage_array_shg), self.exponent_sfg)
return self.fft_shg(sht)[self.selection_array_shg]
def _aathg(self, field_f):
"""
Compute combined anti-aliasing triple and inverse triple harmonic.
Parameters
----------
field_f : array
Field in time domain.
Returns
-------
array
Anti-aliased triple and inverse triple harmonic in frequency domain.
"""
self.storage_array_thg[:self.half_pts] = field_f[:self.half_pts]
self.storage_array_thg[-self.half_pts:] = field_f[self.half_pts:]
field_t_mult = numbathg(self.ifft_thg(self.storage_array_thg), self.exponent_thg)
return self.fft_thg(field_t_mult)[self.selection_array_thg]
"""
Propagation functions
"""
[docs] def nonlinear_term(self, field):
"""
Nonlinear term, implemented in child classes for different interaction
types.
"""
raise NotImplementedError
[docs] def dispersion_step(self, field, dz):
"""
Compute dispersion step with field in frequency domain.
Parameters
----------
field : array
Field in frequency domain.
dz : float
Step size.
Returns
-------
array
Dispersion affected field in frequency domain.
"""
return numbaexp(-1j*self.waveguide.k*dz)*field
"""
Model classes
"""
[docs]class Spm(Model):
"""
Self phase modulation only physical model.
"""
def __init__(self, waveguide, light):
"""
Class constructor.
Parameters
-------
waveguide : Waveguide
Object with the material properties of the propagation medium.
light : Light
Object containing the optical field.
"""
Model.__init__(self, waveguide, light)
self._setup_aaspm()
self._recommended_solver = 'ERK4IP'
[docs] def nonlinear_term(self, field_f):
"""
Nonlinear function with spm
Parameters
-------
field_f : array
Analytical envelope of the electric field in frequency domain.
Returns
-------
array
Nonlinear evolution term.
"""
tot_spm = self._aaspm(field_f)
return self.waveguide.rhs_prefactor*tot_spm*self.waveguide.chi_3/4
[docs]class Chi3(Model):
"""
Triple sum-frequency physical model.
"""
def __init__(self, waveguide, light):
"""
Class constructor.
Parameters
-------
waveguide : Waveguide
Object with the material properties of the propagation medium.
light : Light
Object containing the optical field.
"""
Model.__init__(self, waveguide, light)
self._setup_aathg()
[docs] def nonlinear_term(self, field_f):
"""
Nonlinear function with spm and chi 3
Parameters
-------
field_f : array
Analytical envelope of the electric field in frequency domain.
Returns
-------
array
Nonlinear evolution term.
"""
thg_term = self._aathg(field_f)
return self.waveguide.rhs_prefactor*thg_term*self.waveguide.chi_3/4
[docs]class SpmChi3(Model):
"""
Self phase modulation and triple harmonic physical model.
"""
def __init__(self, waveguide, light):
"""
Class constructor.
Parameters
-------
waveguide : Waveguide
Object with the material properties of the propagation medium.
light : Light
Object containing the optical field.
"""
Model.__init__(self, waveguide, light)
self._setup_aaspm()
self._setup_aathg()
[docs] def nonlinear_term(self, field_f):
"""
Nonlinear function with spm and chi 3
Parameters
-------
field_f : array
Analytical envelope of the electric field in frequency domain.
Returns
-------
array
Nonlinear evolution term.
"""
tot_spm = self._aaspm(field_f)
thg_term = self._aathg(field_f)
return self.waveguide.rhs_prefactor*(thg_term + tot_spm)*self.waveguide.chi_3/4
[docs]class Chi2(Model):
"""
Sum frequency generation and difference frequency generation only physical
model.
"""
def __init__(self, waveguide, light):
"""
Class constructor.
Parameters
-------
waveguide : Waveguide
Object with the material properties of the propagation medium.
light : Light
Object containing the optical field.
"""
Model.__init__(self, waveguide, light)
self._setup_aashg()
[docs] def nonlinear_term(self, field_f):
"""
Nonlinear function with chi 2
Parameters
-------
field_f : array
Analytical envelope of the electric field in frequency domain.
Returns
-------
array
Nonlinear evolution term.
"""
shg_term = self._aash(field_f)
return self.waveguide.rhs_prefactor*shg_term*self.waveguide.chi_2/2
[docs]class SpmChi2(Model):
"""
Self phase modulation, sum and difference frequency
generation physical model.
"""
def __init__(self, waveguide, light):
"""
Class constructor.
Parameters
-------
waveguide : Waveguide
Object with the material properties of the propagation medium.
light : Light
Object containing the optical field.
"""
Model.__init__(self, waveguide, light)
self._setup_aaspm()
self._setup_aashg()
[docs] def nonlinear_term(self, field_f):
"""
Nonlinear function with spm and chi 2
Parameters
-------
field_f : array
Analytical envelope of the electric field in frequency domain.
Returns
-------
array
Nonlinear evolution term.
"""
tot_spm = self._aaspm(field_f)
shg_term = self._aash(field_f)
return self.waveguide.rhs_prefactor*(tot_spm*self.waveguide.chi_3/4 + shg_term*self.waveguide.chi_2/2)
[docs]class Chi2Chi3(Model):
"""
Triple-sum, sum and difference frequency
generation physical model.
"""
def __init__(self, waveguide, light):
"""
Class constructor.
Parameters
-------
waveguide : Waveguide
Object with the material properties of the propagation medium.
light : Light
Object containing the optical field.
"""
Model.__init__(self, waveguide, light)
self._setup_aashg()
self._setup_aathg()
[docs] def nonlinear_term(self, field_f):
"""
Nonlinear function with chi 2 and chi 3
Parameters
-------
field_f : array
Analytical envelope of the electric field in frequency domain.
Returns
-------
array
Nonlinear evolution term.
"""
shg_term = self._aash(field_f)
thg_term = self._aathg(field_f)
return self.waveguide.rhs_prefactor*(thg_term*self.waveguide.chi_3/4 + shg_term*self.waveguide.chi_2/2)
[docs]class SpmChi2Chi3(Model):
"""
Self phase modulation, triple harmonic, sum and difference frequency
generation physical model.
"""
def __init__(self, waveguide, light):
"""
Class constructor.
Parameters
-------
waveguide : Waveguide
Object with the material properties of the propagation medium.
light : Light
Object containing the optical field.
"""
Model.__init__(self, waveguide, light)
self._setup_aaspm()
self._setup_aashg()
self._setup_aathg()
[docs] def nonlinear_term(self, field_f):
"""
Nonlinear function with spm, chi 2 and chi 3
Parameters
-------
field_f : array
Analytical envelope of the electric field in frequency domain.
Returns
-------
array
Nonlinear evolution term.
"""
tot_spm = self._aaspm(field_f)
thg_term = self._aathg(field_f)
shg_term = self._aash(field_f)
return self.waveguide.rhs_prefactor*((thg_term + tot_spm)*self.waveguide.chi_3/4 + shg_term*self.waveguide.chi_2/2)
"""
Optimized functions
"""
### Numba exponential
[docs]@numba.njit([numba.complex128[:](numba.complex128[:]),
numba.complex64[:](numba.complex64[:])])
def numbaexp(x):
"""Compute numba optimized exponential of input vector."""
return np.exp(x)
### Numba absolute value squared
@numba.vectorize([numba.float64(numba.complex128),
numba.float32(numba.complex64)])
def numbaabs2(x):
"""Compute numba optimized absolute norm square of input vector."""
return x.real**2 + x.imag**2
### Numba copy
[docs]@numba.njit([numba.complex128[:](numba.complex128[:]),
numba.complex64[:](numba.complex64[:])])
def numbacopy(x):
"""Create a numba optimized copy of a vector."""
return np.copy(x)
### Numba conjugation
[docs]@numba.njit([numba.complex128[:](numba.complex128[:]),
numba.complex64[:](numba.complex64[:])])
def numbaconj(x):
"""Compute numba optimized complex conjugate of a vector."""
return np.conj(x)
### Numba shg
@numba.vectorize([numba.complex128(numba.complex128, numba.complex128),
numba.complex64(numba.complex64, numba.complex64)])
def numbasht(x, y):
"""Compute numba optimized sh term."""
return x*x*y + 2*(x.real**2 + x.imag**2)*(y.real - 1j*y.imag)
### Numba thg
@numba.vectorize([numba.complex128(numba.complex128, numba.complex128),
numba.complex64(numba.complex64, numba.complex64)])
def numbathg(x, y):
"""Compute numba optimized th term."""
return x*x*x*y + 3*(x.real**2 + x.imag**2)*(x.real - 1j*x.imag)*(y.real - 1j*y.imag)