pychi.solvers — pychi 0.0.1 documentation

# 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

    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

        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
    @numba.njit([numba.complex128[:](numba.complex128[:]),numba.complex64[:](numba.complex64[:])])
    def numbaexp(x):
        return np.exp(x)

    ### Numba norm
    @numba.njit([numba.float64(numba.complex128[:])])
    def numbanorm(field):
        return np.linalg.norm(field)