Source code for pychi.solvers

# -*- coding: utf-8 -*-
"""
Created on Fri Feb 18 17:54:20 2022

@author: voumardt
"""
import numba
import numpy as np
from scipy.constants import c
import time as timing


[docs]class Solver(): def __init__(self, model, z_pts=500, local_error=0.00001, adaptive_factor=1.1, max_dz=None, method=None): """ Solver class. Provides a stepper which performs integration of the PDE as well as error control and results saving. Parameters ---------- model : Model object Physical model of the problem to be integrated. z_pts : int, optional Number of results to keep along the propagation axis. The default is 500. local_error : float, optional Tolerated local error. The default is 0.00001. adaptive_factor : float, optional Adaptive factor for the step size control. Bigger values lead to more aggressive step size control. The default is 1.1. max_dz : float, optional Maximum step size allowed. Particularly useful when simulating quasi-phase matched structures. The default is None. method: string Method used for integration. Can be RK4IP, ERK4IP or DP5IP. Default is DP5IP. """ self.model = model self.z_pts = z_pts self.local_error = local_error self.adaptive_factor = adaptive_factor self.max_dz = max_dz self.z_save = np.linspace(0, self.model.waveguide.length, self.z_pts) self.method_name = self._set_method_name(method) def _set_method_name(self, method): """ Helper function to choose an appropriate solver. If a method is not provided by the user, it will look into the model used to see if there is a recommended solver. Otherwise, defaults to the model that should be mostly used, i.e. DP5IP. Parameters ---------- method : string Name of the solver. Returns ------- string Name of the solver. """ if method is None: if hasattr(self.model, '_recommended_solver'): return self.model._recommended_solver else: return 'DP5IP' else: return method def _compute_error(self, field_1, field_2): """ Error estimation for adaptive step size Parameters ---------- field_1 : array Reference field. field_2 : array Test field. Returns ------- float Error estimate between the two fields. """ return numbanorm(field_1 - field_2)/numbanorm(field_1) def _dormand_prince_step(self, field, dz): """ Dormand Prince order 5(4) adaptive scheme in the interaction picture. This new stepper uses the interaction picture in an embedded scheme of order 5. Tests have shown it to be up to four times faster than RK4IP. Parameters ---------- field : array Current field. dz : float Current step size. Returns ------- array New field estimate. float Error estimate. """ if self._k_7 is None: k_1 = dz*self.model.dispersion_step(self.model.nonlinear_term(field), dz) else: k_1 = self.model.dispersion_step(self._k_7, dz) a_int = self.model.dispersion_step(field, dz) k_2 = dz*self.model.dispersion_step(self.model.nonlinear_term(self.model.dispersion_step(a_int + k_1/5, -4*dz/5)), 4*dz/5) k_3 = dz*self.model.dispersion_step(self.model.nonlinear_term(self.model.dispersion_step(a_int + 3*k_1/40 + 9*k_2/40, -7*dz/10)), 7*dz/10) k_4 = dz*self.model.dispersion_step(self.model.nonlinear_term(self.model.dispersion_step(a_int + 44*k_1/45 - 56*k_2/15 + 32*k_3/9, -dz/5)), dz/5) k_5 = dz*self.model.dispersion_step(self.model.nonlinear_term(self.model.dispersion_step(a_int + 19372*k_1/6561 - 25360*k_2/2187 + 64448*k_3/6561 - 212*k_4/729, -dz/9)), dz/9) k_6 = dz*self.model.nonlinear_term(a_int + 9017*k_1/3168 - 355*k_2/33 + 46732*k_3/5247 + 49*k_4/176 - 5103*k_5/18656) sol_5 = a_int + 35*k_1/384 + 500*k_3/1113 + 125*k_4/192 - 2187*k_5/6784 + 11*k_6/84 self._k_7 = dz*self.model.nonlinear_term(sol_5) sol_4 = a_int + 5179*k_1/57600 + 7571*k_3/16695 + 393*k_4/640 - 92097*k_5/339200 + 187*k_6/2100 + self._k_7/40 error = self._compute_error(sol_5, sol_4) return sol_5, error def _rk546m(self, field, dz): """ Some other Dormand-Prince method that has one less step but does not leverage the interaction picture and does not have the FSAL property. Usually worse than the other dormand-prince method. References: DOI: 10.1016/0771-050X(80)90013-3 Parameters ---------- field : array Current field. dz : float Current step size. Returns ------- array New field estimate. float Error estimate. """ k_1 = dz*self.model.nonlinear_term(field) k_2 = dz*self.model.dispersion_step(self.model.nonlinear_term(self.model.dispersion_step(field + k_1/5, dz/5)), -dz/5) k_3 = dz*self.model.dispersion_step(self.model.nonlinear_term(self.model.dispersion_step(field + 3*k_1/40 + 9*k_2/40, 3*dz/10)), -3*dz/10) k_4 = dz*self.model.dispersion_step(self.model.nonlinear_term(self.model.dispersion_step(field + 3*k_1/10 - 9*k_2/10 + 6*k_3/5, 3*dz/5)), -3*dz/5) k_5 = dz*self.model.dispersion_step(self.model.nonlinear_term(self.model.dispersion_step(field + 226*k_1/729 - 25*k_2/27 + 880*k_3/729 + 55*k_4/729, 2*dz/3)), -2*dz/3) k_6 = dz*self.model.dispersion_step(self.model.nonlinear_term(self.model.dispersion_step(field - 181*k_1/270 + 5*k_2/2 - 266*k_3/297 - 91*k_4/27 + 189*k_5/55, dz)), -dz) sol_5 = self.model.dispersion_step(field + 19*k_1/216 + 1000*k_3/2079 - 125*k_4/216 + 81*k_5/88 + 5*k_6/56, dz) sol_4 = self.model.dispersion_step(field + 31*k_1/540 + 190*k_3/297 - 145*k_4/108 + 351*k_5/220 + k_6/20, dz) error = self._compute_error(sol_5, sol_4) return sol_5, error def _adams_bashforth54_step(self, field, dz): """To be implemented""" pass def _rk4ip_single_step(self, field, dz): """ Runge-Kutta order 4 step in the interaction picture (RK4IP) Parameters ---------- field : array Current field. dz : float Current step size. Returns ------- array New field estimate. """ dispersion_op = numbaexp(-1j*self.model.waveguide.k*dz/2) a_int = dispersion_op*field k_1 = dz*dispersion_op*self.model.nonlinear_term(field) k_2 = dz*self.model.nonlinear_term(a_int + k_1/2) k_3 = dz*self.model.nonlinear_term(a_int + k_2/2) k_4 = dz*self.model.nonlinear_term(dispersion_op*(a_int + k_3)) return dispersion_op*(a_int + k_1/6 + k_2/3 + k_3/3) + k_4/6 def _rk4ip_step(self, field, dz): """ Runge-Kutta order 4 adaptive scheme in the interaction picture (RK4IP) with step doubling and combining. Provided for comparison but slower than the dormand-prince method. References: DOI: 10.1109/JLT.2007.909373 DOI: 10.1109/JLT.2009.2021538 Parameters ---------- field : array Current field. dz : float Current step size. Returns ------- array New field estimate. float Error estimate. """ sol_single = self._rk4ip_single_step(field, dz) sol_double = self._rk4ip_single_step(field, dz/2) sol_double = self._rk4ip_single_step(sol_double, dz/2) error = self._compute_error(sol_double, sol_single) sol_5 = (sol_double*16 - sol_single)/15 return sol_5, error def _erk4ip_step(self, field, dz): """ Embedded Runge-Kutta order 4(3) adaptive scheme in the interaction picture (ERK4IP). Can be faster than DP5IP if accuracy is irrelevant, for example in SPM only cases. Otherwise, DP5IP performs better. References: DOI: 10.1016/j.cpc.2012.12.020 Parameters ---------- field : array Current field. dz : float Current step size. Returns ------- array New field estimate. float Error estimate. """ dispersion_op = numbaexp(-1j*self.model.waveguide.k*dz/2) a_int = dispersion_op*field if self._k_5 is None: k_1 = dz*dispersion_op*self.model.nonlinear_term(field) else: k_1 = dispersion_op*self._k_5 k_2 = dz*self.model.nonlinear_term(a_int + k_1/2) k_3 = dz*self.model.nonlinear_term(a_int + k_2/2) k_4 = dz*self.model.nonlinear_term(dispersion_op*(a_int + k_3)) beta = dispersion_op*(a_int + k_1/6 + k_2/3 + k_3/3) sol_4 = beta + k_4/6 self._k_5 = dz*self.model.nonlinear_term(sol_4) sol_3 = beta + k_4/15 + self._k_5/10 error = self._compute_error(sol_4, sol_3) return sol_4, error def _integrate_to_z(self, z): """ Integrate to z with adaptive step size. Parameters ------- z : float Target propagation length. """ while self.stock_z < z: # Actuate position inside the waveguide self.model.waveguide.z = self.stock_z # Step of dz vec_eval, error = self.step(self.stock_vec, self.dz) # Error control if error >= 2*self.local_error: self.dz /= 2 print('Discarded') if self.method_name == 'DP5IP': self._k_7 = None if self.method_name == 'ERK4IP': self._k_5 = None continue elif self.local_error <= error < 2*self.local_error: self.stock_z += self.dz self.dz /= self.adaptive_factor elif 0.5*self.local_error <= error < self.local_error: self.stock_z += self.dz else: self.stock_z += self.dz if self.max_dz is None or self.dz < self.max_dz: self.dz *= self.adaptive_factor self.stock_vec = vec_eval
[docs] def solve(self): """ Integrate over waveguide length with adaptive step size """ # Initialization self.spectrum = np.zeros((self.z_pts, self.model.waveguide.t_pts), dtype='complex128') self.dz = np.mean(np.diff(self.z_save))*np.sqrt(np.sqrt(self.local_error)) self.stock_z = 0 self.stock_vec = np.fft.fft(self.model.light.field_t_in) if self.method_name == 'DP5IP': self._k_7 = None self.step = self._dormand_prince_step elif self.method_name == 'ERK4IP': self._k_5 = None self.step = self._erk4ip_step elif self.method_name == 'RK4IP': self.step = self._rk4ip_step elif self.method_name == 'RK546M': self.step = self._rk546m timer = timing.perf_counter() # Total integration for i, z in enumerate(self.z_save): print('z: {}/{}'.format(i+1, self.z_pts)) # Integrate over dz self._integrate_to_z(z) # Save spectrum self.spectrum[i, :] = np.fft.fftshift(self.stock_vec)\ *numbaexp(1j*(self.model.waveguide.betas[0] + self.model.waveguide.betas[1]*self.model.waveguide.rel_omega)*self.stock_z) # Save waveform print(f"{timing.perf_counter() - timer} seconds for integration.") self.waveform = np.fft.ifft(np.fft.ifftshift(self.spectrum, axes=1), axis=1) self.model.light._set_propagation(self.z_save, self.model.waveguide.freq, self.spectrum, self.waveform)
""" Optimized functions """ ### Numba exponential
[docs]@numba.njit([numba.complex128[:](numba.complex128[:]),numba.complex64[:](numba.complex64[:])]) def numbaexp(x): return np.exp(x)
### Numba norm
[docs]@numba.njit([numba.float64(numba.complex128[:])]) def numbanorm(field): return np.linalg.norm(field)