pychi.materials — pychi 0.0.1 documentation

# Source code for pychi.materials

    # -*- coding: utf-8 -*-
    """
    Created on Fri Aug 13 12:07:43 2021

    @author: Thibault
    """
    from inspect import signature
    from math import factorial
    import matplotlib.pyplot as plt
    import numpy as np
    from scipy.interpolate import interp1d
    from scipy.constants import c
    eps_0 = 8.8541878128e-12

    class Waveguide():
        """
        Waveguide class. Contains all material related information, i.e. the
        refractive index curve versus frequency, the Raman parameters, the
        nonlinear coefficients, the effective area and length. This object
        should be the first object instantiated, as the time/frequency axes of
        the simulation are derived here from the available refractive
        index data.
        """

        def __init__(self, frequency, n_eff, chi_2, chi_3,
                     effective_area, length, raman_fraction=0.18,
                     raman_tau_1=0.0122e-12, raman_tau_2=0.032e-12,
                     t_pts=2**14):
            """
            Construct the waveguide instance.

            Parameters
            ----------
            frequency : array
                Frequency axis on which the effective refractive index is given.
            n_eff : array
                Effective refractive index of the material.
            chi_2 : float
                Order 2 material nonlinearity.
            chi_3 : float
                Order 3 material nonlinearity.
            effective_area : float
                Effective area at the pump wavelength.
            length : float
                Length of the material.
            raman_fraction : float
                Fractional contribution of the Raman effect to the Kerr effect.
                The default is 0.18.
            raman_tau_1 : float
                Raman effect primary time scale. The default is 0.0122e-12.
            raman_tau_2 : float
                Raman effect secondary time scale. The default is 0.032e-12.
            t_pts : float, optional
                Number of points in the time and frequency axes, will be rounded
                to nearest higher power of 2. The default is 2**14.
            """
            self.z = 0
            self.t_pts = int(2**np.ceil(np.log2(t_pts)))
            self.length = length
            self.data_freq = frequency
            self.center_freq, self.freq = self._compute_frequency_axis()
            self.center_omega, self.omega, self.rel_omega = self._compute_omega_axis()
            self.time, self.delta_t = self._compute_time_axis()

            self.n_eff = n_eff

            self.chi_2 = chi_2
            self.chi_3 = chi_3
            self.eff_area = effective_area

            # Standard Raman response for fibers
            self.raman_fraction = raman_fraction
            self.raman_tau_1 = raman_tau_1
            self.raman_tau_2 = raman_tau_2
            self.raman_effect = self._compute_raman_effect()

        @property
        def n_eff(self):
            return self._n_eff(self.z)

        @n_eff.setter
        def n_eff(self, n_eff):
            self._n_eff = self._make_callable(n_eff)
            self.n_eff_inter = self._interpolate_n_eff()
            self.betas = self.compute_betas()

        @property
        def chi_2(self):
            return self._chi_2(self.z)

        @chi_2.setter
        def chi_2(self, chi_2):
            self._chi_2 = self._make_callable(chi_2)

        @property
        def chi_3(self):
            return self._chi_3(self.z)

        @chi_3.setter
        def chi_3(self, chi_3):
            self._chi_3 = self._make_callable(chi_3)

        @property
        def k(self):
            return np.fft.fftshift(self.n_eff*self.omega/c)

        @property
        def rhs_prefactor(self):
            return -1j*np.fft.fftshift(self.omega/(2*self.n_eff*c))

        def _make_callable(self, var):
            """
            Transform an input vector, array, float or callable into a z-callable.
            The user can then provide z-dependant quantities.

            Note that a provided callable should have as arguments exactly at most
            'z' and 'freq', e.g. var(z, freq), var(freq, z), var(z) or var(freq).
            A bit wonky, but couldn't figure out a better way to do that. Another
            alternative would be to implement derived classes with overwritten
            properties.

            Note that a provided array should have first dimension coinciding with
            the frequency axis given along the refractive index data. The second
            dimension is assumed to be evenly distributed along the whole length
            of the waveguide.

            Note that if a vector is provided it will be assumed that it describes
            a z dependent quantity, unless it has the same dimension as the
            frequency axis. Wavelength dependent only quantities can also be
            provided through callables.

            Parameters
            ----------
            var : float, vector, array or callable
                Some simulation parameter.

            Returns
            -------
            function
                A z-callable version of the provided variable.
            """
            # Case when the argument is callable
            if callable(var):
                return self._parse_callable(var)
            # Case when the argument is an array or a vector
            elif hasattr(var, '__len__'):
                return self._parse_array(var)
            # Otherwise (float)
            else:
                return lambda _: var

        def _parse_array(self, var):
            """
            Parse the correct arguments for a callable parameter.

            Parameters
            ----------
            var : array
                A one or two dimensional array representing z and/or frequency
                dependence of a parameter.

            Returns
            -------
            function
                A z-callable function.
            """
            # Case of an array
            if len(np.shape(var)) == 2:
                arr_inter = interp1d(self.data_freq, var.T, kind='cubic',
                                     bounds_error=False, fill_value='extrapolate')
                arr = arr_inter(self.freq).T
                z = np.linspace(0, self.length, np.shape(arr)[-1])
                arr_inter = interp1d(z, arr, kind='cubic', bounds_error=False,
                                     fill_value='extrapolate')
                return arr_inter
            # Case of a vector
            elif len(np.shape(var)) == 1:
                if len(var) == len(self.data_freq):
                    arr_inter = interp1d(self.data_freq, var, kind='cubic',
                                         bounds_error=False, fill_value='extrapolate')
                    arr = arr_inter(self.freq)
                    return lambda _: arr
                else:
                    z = np.linspace(0, self.length, len(var))
                    var_inter = interp1d(z, var, kind='cubic', bounds_error=False,
                                         fill_value='extrapolate')
                    return var_inter
            else:
                raise ValueError

        def _parse_callable(self, var):
            """
            Parse the correct arguments for a callable parameter.

            Parameters
            ----------
            var : callable
                A z or frequency (or both) dependent callable.

            Returns
            -------
            function
                A z-callable function.
            """
            # Get funtion signature
            z_flag = False
            freq_flag = False
            order_flag = False
            args = signature(var).parameters
            for arg in args:
                if arg == 'z':
                    z_flag = True
                if arg == 'freq':
                    freq_flag = True
                    if z_flag:
                        order_flag = True

            # Reorder arguments and evaluate the frequency axis
            if z_flag:
                if freq_flag:
                    if order_flag:
                        return lambda _: var(_, self.freq)
                    else:
                        return lambda _: var(self.freq, _)
                else:
                    return var
            else:
                if freq_flag:
                    return lambda _: var(self.freq)
                else:
                    return lambda _: var

        def _compute_frequency_axis(self):
            """
            Define a frequency axis covering exactly the data range available.
            """
            min_freq = np.amin(self.data_freq)
            max_freq = np.amax(self.data_freq)
            center_freq = min_freq/2 + max_freq/2
            freq = np.linspace(min_freq, max_freq, self.t_pts)
            return center_freq, freq

        def _compute_omega_axis(self):
            """
            Compute the angular frequency axis associated to the frequency axis
            """
            center_omega = 2*np.pi*self.center_freq
            omega = 2*np.pi*self.freq
            rel_omega = omega - center_omega
            return center_omega, omega, rel_omega

        def _compute_time_axis(self):
            """
            Compute the time axis associated to the frequency axis
            """
            time = np.fft.fftshift(np.fft.fftfreq(self.t_pts, np.mean(np.diff(self.freq))))
            delta_t = np.mean(np.diff(time))
            return time, delta_t

        def _interpolate_n_eff(self):
            """
            Interpolate the refractive index data on the simulation frequency axis
            """
            n_eff_inter = interp1d(self.freq, self.n_eff, kind='cubic',
                                   bounds_error=False, fill_value='extrapolate')
            return n_eff_inter

        def _compute_raman_effect(self):
            """
            Compute the Raman term for the simulation frequency axis
            """
            raman_time = (self.raman_tau_1**2 + self.raman_tau_2**2)
            raman_time = raman_time/self.raman_tau_1/self.raman_tau_2**2
            raman_time *= np.exp(-np.abs(self.time)/self.raman_tau_2)*np.sin(self.time/self.raman_tau_1)
            raman_time[self.time < 0] = 0
            raman_effect = np.conjugate(self.t_pts*np.fft.ifft(np.fft.fftshift(raman_time)))
            raman_effect *= self.delta_t*self.raman_fraction
            return raman_effect

        def compute_betas(self, wavelength=None, order=6):
            """
            Compute beta coefficients at a given wavelength and order.

            Parameters
            ----------
            wavelength : float, optional
                Wavelength at which the beta coefficients are computed. The default is
                the center of the simulation axis.
            order : int, optional
                Highest order of the beta coefficients to be computed. The default is 6.

            Returns
            -------
            beta : array
                Beta coefficients of the material at the given wavelength.
            """
            if wavelength is None:
                center_omega = self.center_omega
            else:
                center_omega = 2*np.pi*c/wavelength

            om = 2*np.pi*self.freq
            dk = np.real(self.n_eff)*om/c
            b_inter = interp1d(om, dk, kind='cubic')
            beta = np.zeros(order+1)
            beta[0] = b_inter(center_omega)
            for i in np.arange(order):
                dk = np.diff(dk)
                dom = np.diff(om)
                om = om[:-1] + dom/2
                der = dk/dom
                der_inter = interp1d(om, der, kind='cubic')
                beta[i+1] = der_inter(center_omega)
                dk = der
            return beta

        def set_betas(self, betas, wavelength):
            """
            Convenience only. Allow the use to provide the beta coefficients
            instead of n_eff. Should be avoided.
            """
            rel_om = self.omega - 2*np.pi*c/wavelength
            n = np.ones(self.t_pts)*betas[0]
            factorial = 1
            for i, b in enumerate(betas[1:]):
                factorial *= (i+1)
                n += b*rel_om**(i+1)/factorial
            n *= c/self.omega
            self.n_eff = n

        def set_n2(self, n2, wavelength=None):
            """
            Convenience only. Allow the user to give the nonlinear index instead
            of chi3. Overwrite the initialized chi3 value accordingly.

            Parameters
            ----------
            n2 : float
                Nonlinear index.
            """
            if wavelength is None:
                freq = self.center_freq
            else:
                freq = c/wavelength
            self.chi_3 = 4*n2*eps_0*c*self.n_eff_inter(freq)**2/3

        def set_gamma(self, gamma, wavelength=None):
            """
            Convenience only. Allow the user to give the nonlinear coefficient instead
            of chi3. Overwrite the initialized chi3 value accordingly.

            Parameters
            ----------
            gamma : float
                Nonlinear coefficient.
            """
            if wavelength is None:
                freq = self.center_freq
            else:
                freq = c/wavelength
            self.chi_3 = 2*self.n_eff_inter(freq)**2*c**2*eps_0*gamma*self.eff_area/3/np.pi/freq

        def plot_refractive_index(self, savename=None):
            """
            Plot refractive index
            """
            k = self.n_eff*2*np.pi*self.freq/c
            inverse_group_velocity = np.gradient(k, 2*np.pi*self.freq)
            group_velocity_dispersion = np.gradient(inverse_group_velocity, 2*np.pi*self.freq)

            fig, axs = plt.subplots(3)
            axs[0].plot(self.freq, self.n_eff)
            axs[0].set_ylabel('Refractive index')
            axs[1].plot(self.freq, inverse_group_velocity)
            axs[1].set_ylabel('IGV [s/m]')
            axs[2].plot(self.freq, group_velocity_dispersion)
            axs[2].set_ylabel('GVD [s^2/m]')
            axs[2].set_xlabel('Frequency [Hz]')
            fig.tight_layout()