jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def theory_xi_2D(self, params, rmax=85, **kwargs):
        r"""
        Compute the theory model prediction as a function of distances perpendicular and parallel to the line of sight

        Parameters
        ----------
        params : dict
            Dictionary containing parameter values at which to evaluate the model

        rmax : float, default=85
            Maximum distance (in each direction) out to which model should be calculated

        kwargs : dict
            Optional settings to override those in :attr:`model`

        Returns
        -------
        xi_model : instance of :class:`scipy.interpolate.bisplev`
            Representing the redshift-space model prediction :math:`\xi^s(s_\perp, s_{||})`
        """

        sperp = np.linspace(0.01, rmax)
        spar = np.linspace(-rmax, rmax)  # allow for potential odd functions
        def smu(sigma, pi):
            s = np.sqrt(sigma**2 + pi**2)
            return s, pi/s
        s, mu = distance_modulus(*np.meshgrid(sperp, spar))
        xi_smu = self.theory_xi(s, mu, params, **kwargs)
        tck = si.bisplrep(sperp, spar, xi_smu)
        xi_model = lambda x, y: si.bisplev(x, y, tck)
        return xi_model
//  [python]
def interp2d_wrapper(data_path, band_index_to_read, east_out_axis, north_out_axis, fill_value):

    data_driver = gdal.Open(data_path, GA_ReadOnly)
    data_in = data_driver.GetRasterBand(band_index_to_read).ReadAsArray()
    geotransform = data_driver.GetGeoTransform()
    east_axis_in = geotransform[0] + geotransform[1] * np.arange(data_driver.RasterXSize)
    north_axis_in = geotransform[3] + geotransform[5] * np.arange(data_driver.RasterYSize)
    data_driver = None

    # interpolation function
    min_input = np.nanmin(data_in)
    nan_mask = np.isnan(data_in)
    data_in[nan_mask] = -9999

    interp2d_fun = interp2d(east_axis_in, north_axis_in[::-1], data_in[::-1, :], fill_value=fill_value)
    interp2d_fun_mask = interp2d(east_axis_in, north_axis_in[::-1], nan_mask[::-1, :], fill_value=float(0))

    # interpolate
    data_out = interp2d_fun(east_out_axis, north_out_axis[::-1])[::-1, :]
    mask_out = np.ceil(interp2d_fun_mask(east_out_axis, north_out_axis[::-1]))[::-1, :]

    mask_out = np.logical_and(mask_out, np.isnan(data_out))
    data_out[mask_out.astype(int)] = -9999

    data_out[data_out < min_input] = np.nan
    data_out[mask_out.astype(int)] = np.nan

    return data_out
//  [python]
def interp2d_wrapper(data_path, band_index_to_read, east_out_axis, north_out_axis, fill_value):

    data_driver = gdal.Open(data_path, GA_ReadOnly)
    data_in = data_driver.GetRasterBand(band_index_to_read).ReadAsArray()
    geotransform = data_driver.GetGeoTransform()
    east_axis_in = geotransform[0] + geotransform[1] * np.arange(data_driver.RasterXSize)
    north_axis_in = geotransform[3] + geotransform[5] * np.arange(data_driver.RasterYSize)
    data_driver = None

    # interpolation function
    min_input = np.nanmin(data_in)
    nan_mask = np.isnan(data_in)
    data_in[nan_mask] = -9999

    tck = splprep([east_axis_in, north_axis_in[::-1]], data_in[::-1, :], s=0)
    interp2d_fun = lambda x, y: bisplev(x, y, tck)
    tck_mask = splprep([east_axis_in, north_axis_in[::-1]], nan_mask[::-1, :], s=0)
    interp2d_fun_mask = lambda x, y: bisplev(x, y, tck_mask)


    # interpolate
    data_out = interp2d_fun(east_out_axis, north_out_axis[::-1])[::-1, :]
    mask_out = np.ceil(interp2d_fun_mask(east_out_axis, north_out_axis[::-1]))[::-1, :]

    mask_out = np.logical_and(mask_out, np.isnan(data_out))
    data_out[mask_out.astype(int)] = -9999

    data_out[data_out < min_input] = np.nan
    data_out[mask_out.astype(int)] = np.nan

    return data_out
//  [python]
def get_obs_background_field(self, Ub, Vb, grid=None):
        """
        Contains ability to handle 1-D varying wind field. However, recommend
        only supplying scalars as arguments.
        Ub, Vb = Background U & V winds
        grid = 1D array of ranges to consider (from VAD)
        """
        iflag = False
        if hasattr(Ub, '__len__') or hasattr(Vb, '__len__'):
            rngf = (self.obs_xf**2 + self.obs_yf**2)**0.5
        # Background U
        if Ub is None:
            self.obs_Ub = 0.0
        elif hasattr(Ub, '__len__'):
            if np.ndim(Ub) == 2:
                # Assumes that Ub grid matches analysis grid
                f = scipy.interpolate.interp2d(
                    self.analysis_x, self.analysis_y, Ub, kind='linear')
                x1, y1, xii, yii = self._get_sorted_coords()
                tmpUb = f(x1, y1)
                self.obs_Ub = tmpUb[yii, xii]
                iflag = True
            elif grid is None:
                self.obs_Ub = 0.0
            else:
                self.obs_Ub = np.interp(rngf, grid, Ub)
        else:
            self.obs_Ub = Ub
        # Background V
        if Vb is None:
            self.obs_Vb = 0.0
        elif hasattr(Vb, '__len__'):
            if np.ndim(Vb) == 2:
                # Assumes that Vb grid matches analysis grid
                f = scipy.interpolate.interp2d(
                    self.analysis_x, self.analysis_y, Vb, kind='linear')
                # Check to see if we already have sorted coords
                if not iflag:
                    x1, y1, xii, yii = self._get_sorted_coords()
                tmpVb = f(x1, y1)
                self.obs_Vb = tmpVb[yii, xii]
            elif grid is None:
                self.obs_Vb = 0.0
            else:
                self.obs_Vb = np.interp(rngf, grid, Vb)
        else:
            self.obs_Vb = Vb
        self.compute_beta_and_m()
        # Following is for Beta as non-radar convention (0 = due East)
        self.obs_vrbf = self.obs_Ub * np.cos(self.obs_Beta) + \
            self.obs_Vb * np.sin(self.obs_Beta)
//  [python]
def get_obs_background_field(self, Ub, Vb, grid=None):
        """
        Contains ability to handle 1-D varying wind field. However, recommend
        only supplying scalars as arguments.
        Ub, Vb = Background U & V winds
        grid = 1D array of ranges to consider (from VAD)
        """
        iflag = False
        if hasattr(Ub, '__len__') or hasattr(Vb, '__len__'):
            rngf = (self.obs_xf**2 + self.obs_yf**2)**0.5
        # Background U
        if Ub is None:
            self.obs_Ub = 0.0
        elif hasattr(Ub, '__len__'):
            if np.ndim(Ub) == 2:
                # Assumes that Ub grid matches analysis grid
                x, y = np.mgrid[0:self.analysis_x.size, 0:self.analysis_y.size]
                tck = scipy.interpolate.bisplev(self.analysis_x, self.analysis_y, Ub)
                x1, y1, xii, yii = self._get_sorted_coords()
                tmpUb = scipy.interpolate.bisplev(x1, y1, tck)
                self.obs_Ub = tmpUb[yii, xii]
                iflag = True
            elif grid is None:
                self.obs_Ub = 0.0
            else:
                self.obs_Ub = np.interp(rngf, grid, Ub)
        else:
            self.obs_Ub = Ub
        # Background V
        if Vb is None:
            self.obs_Vb = 0.0
        elif hasattr(Vb, '__len__'):
            if np.ndim(Vb) == 2:
                # Assumes that Vb grid matches analysis grid
                x, y = np.mgrid[0:self.analysis_x.size, 0:self.analysis_y.size]
                tck = scipy.interpolate.bisplev(self.analysis_x, self.analysis_y, Vb)
                # Check to see if we already have sorted coords
                if not iflag:
                    x1, y1, xii, yii = self._get_sorted_coords()
                tmpVb = scipy.interpolate.bisplev(x1, y1, tck)
                self.obs_Vb = tmpVb[yii, xii]
            elif grid is None:
                self.obs_Vb = 0.0
            else:
                self.obs_Vb = np.interp(rngf, grid, Vb)
        else:
            self.obs_Vb = Vb
        self.compute_beta_and_m()
        # Following is for Beta as non-radar convention (0 = due East)
        self.obs_vrbf = self.obs_Ub * np.cos(self.obs_Beta) + \
            self.obs_Vb * np.sin(self.obs_Beta)
//  [python]
def image_rec_function(m, wavel, flux, inc_all, angle_all, V_model_TOT, q_interp, bounds, PATH_OUTPUT_IMAGE_REC, resolution):

    res_F = flux[m]

    inc =  inc_all[m]
    ang  =  angle_all[m]


    V_model= V_model_TOT
    
    CF = np.sum([V_model[i]*res_F[i] for i in range(len(res_F))],axis=0)     
    F_tot = np.sum([res_F[k] for k in range(len(res_F))],axis=0)
    
    Vis_tmp = CF/F_tot

    # A convoluer par une gaussienne  

    q, image  = imaging_model(q_interp, Vis_tmp, max(q_interp)/np.sqrt(2)) # max(q_interp)/np.sqrt(2) weird artifacts test with Gaussian convolution
    
    # np.sqrt(2) is here to cut the image to not see the circle shape
    
    window = np.outer(general_gaussian(len(image),1.5,len(image)/14),general_gaussian(len(image),1.5,len(image)/14))

    # len(image)/14 is the standard deviation of the Gaussian with a number of point equal to len(image)

    image = window * image
    
    # Above, I modulate the fourier transform image by a gaussian ==> IMPROVEMENT == convolved by a Door function ending with smoothed wings. 
    
    qu = q
    qv = q
    
    qu_inc, qv_inc, q_inc = coord_inclination(qu, qv, inc)
    f_inc = scipy.interpolate.interp2d(qu, qv, image)
    image_inc = f_inc(qu_inc,qv_inc)
    
    qu_rot, qv_rot, q_rot = inv_coord_rotation(qu_inc, qv_inc, ang)
    image_rot = rotate(image_inc,ang, reshape= False)
    
    wr, hr = rotatedRectWithMaxArea(np.shape(image_rot)[0], np.shape(image_rot)[1], radians(inc))
    wr_lower  = int(np.floor(np.shape(image_rot)[0]/2-wr/2))
    wr_higher = int(np.floor(np.shape(image_rot)[0]/2+wr/2))
    x, image = FFT_to_image(q[wr_lower:wr_higher], np.real(image_rot[wr_lower:wr_higher,wr_lower:wr_higher]))


    # plt.figure()
    # plt.imshow(np.abs(image_rot), extent=[min(qu),max(qu),min(qv),max(qv)])
    
    
    # x, image = FFT_to_image(q, np.real(image_rot))
    
    # print(x)
            
    image = np.abs(np.fft.ifftshift(image))/np.amax(np.real(image))
    
    plot_intensity_model_image(x, image, wavel[m], bounds, SAVE_OUTPUT=PATH_OUTPUT_IMAGE_REC)
    
    func = interp2d(x, x, image)
    
    x_new = np.linspace(-bounds,bounds, resolution)
    
    cropped = func(x_new, x_new)
    

    return  cropped, x_new
//  [python]
def image_rec_function(m, wavel, flux, inc_all, angle_all, V_model_TOT, q_interp, bounds, PATH_OUTPUT_IMAGE_REC, resolution):

    res_F = flux[m]

    inc =  inc_all[m]
    ang  =  angle_all[m]


    V_model= V_model_TOT
    
    CF = np.sum([V_model[i]*res_F[i] for i in range(len(res_F))],axis=0)     
    F_tot = np.sum([res_F[k] for k in range(len(res_F))],axis=0)
    
    Vis_tmp = CF/F_tot

    # A convoluer par une gaussienne  

    q, image  = imaging_model(q_interp, Vis_tmp, max(q_interp)/np.sqrt(2)) # max(q_interp)/np.sqrt(2) weird artifacts test with Gaussian convolution
    
    # np.sqrt(2) is here to cut the image to not see the circle shape
    
    window = np.outer(general_gaussian(len(image),1.5,len(image)/14),general_gaussian(len(image),1.5,len(image)/14))

    # len(image)/14 is the standard deviation of the Gaussian with a number of point equal to len(image)

    image = window * image
    
    # Above, I modulate the fourier transform image by a gaussian ==> IMPROVEMENT == convolved by a Door function ending with smoothed wings. 
    
    qu = q
    qv = q
    
    qu_inc, qv_inc, q_inc = coord_inclination(qu, qv, inc)
    f_inc = scipy.interpolate.bisplev(qu, qv, image)
    image_inc = f_inc(qu_inc,qv_inc)
    
    qu_rot, qv_rot, q_rot = inv_coord_rotation(qu_inc, qv_inc, ang)
    image_rot = rotate(image_inc,ang, reshape= False)
    
    wr, hr = rotatedRectWithMaxArea(np.shape(image_rot)[0], np.shape(image_rot)[1], radians(inc))
    wr_lower  = int(np.floor(np.shape(image_rot)[0]/2-wr/2))
    wr_higher = int(np.floor(np.shape(image_rot)[0]/2+wr/2))
    x, image = FFT_to_image(q[wr_lower:wr_higher], np.real(image_rot[wr_lower:wr_higher,wr_lower:wr_higher]))


    # plt.figure()
    # plt.imshow(np.abs(image_rot), extent=[min(qu),max(qu),min(qv),max(qv)])
    
    
    # x, image = FFT_to_image(q, np.real(image_rot))
    
    # print(x)
            
    image = np.abs(np.fft.ifftshift(image))/np.amax(np.real(image))
    
    plot_intensity_model_image(x, image, wavel[m], bounds, SAVE_OUTPUT=PATH_OUTPUT_IMAGE_REC)
    
    func = interp2d(x, x, image)
    
    x_new = np.linspace(-bounds,bounds, resolution)
    
    cropped = func(x_new, x_new)
    

    return  cropped, x_new
//  [python]
def xi_2D_from_multipoles(self, multipoles, poles, rmax=85):
        r"""
        Use input Legendre multipoles to generate a 2D ccf as a function of distances perpendicular and parallel to the
        line of sight

        Parameters
        ----------
        multipoles : dict
            Dictionary of callable multipole functions accepting distance modulus argument, with keys being the Legendre order

        poles : list
            The Legendre multipole orders, should match the keys in `multipoles`

        rmax : float, default=85
            Maximum distance (in each direction) out to which to interpolate

        Returns
        -------
        xi_2D : instance of :class:`scipy.interpolate.interp2d`
            Representing the 2D ccf :math:`\xi(r_\perp, r_{||})`
        """

        sperp = np.linspace(0.01, rmax)
        spar = np.linspace(-rmax, rmax)  # allow for potential odd functions
        def smu(sigma, pi):
            s = np.sqrt(sigma**2 + pi**2)
            return s, pi/s
        s, mu = distance_modulus(*np.meshgrid(sperp, spar))

        xi_2D_grid = np.zeros_like(s)
        for ell in poles:
            xi_2D_grid = xi_2D_grid + multipoles[f'{ell}'](s) * legendre(ell)(mu)
        xi_2D = si.interp2d(sperp, spar, xi_2D_grid)
        return xi_2D
//  [python]
def xi_2D_from_multipoles(self, multipoles, poles, rmax=85):
        r"""
        Use input Legendre multipoles to generate a 2D ccf as a function of distances perpendicular and parallel to the
        line of sight

        Parameters
        ----------
        multipoles : dict
            Dictionary of callable multipole functions accepting distance modulus argument, with keys being the Legendre order

        poles : list
            The Legendre multipole orders, should match the keys in `multipoles`

        rmax : float, default=85
            Maximum distance (in each direction) out to which to interpolate

        Returns
        -------
        xi_2D : instance of :class:`scipy.interpolate.bisplev`
            Representing the 2D ccf :math:`\xi(r_\perp, r_{||})`
        """

        sperp = np.linspace(0.01, rmax)
        spar = np.linspace(-rmax, rmax)  # allow for potential odd functions
        def smu(sigma, pi):
            s = np.sqrt(sigma**2 + pi**2)
            return s, pi/s
        s, mu = distance_modulus(*np.meshgrid(sperp, spar))

        xi_2D_grid = np.zeros_like(s)
        for ell in poles:
            xi_2D_grid = xi_2D_grid + multipoles[f'{ell}'](s) * legendre(ell)(mu)
        tck = si.bisplrep(sperp, spar, xi_2D_grid)
        xi_2D = lambda x, y: si.bisplev(x, y, tck)
        return xi_2D
//  [python]
def bilinear_interp_2d(target_point, pts, values):
    quad_xs = [p[0] for p in pts]
    quad_ys = [p[1] for p in pts]
    interp_f = interpolate.interp2d(quad_xs, quad_ys, values, kind='linear')
    result = interp_f(target_point[0], target_point[1])[0]
    return result
//  [python]
def bilinear_interp_2d(target_point, pts, values):
    quad_xs = [p[0] for p in pts]
    quad_ys = [p[1] for p in pts]
    interp_f = interpolate.bisplev(quad_xs, quad_ys, values)
    result = interp_f(target_point[0], target_point[1])[0]
    return result
//  [python]
def three_d_surface_plot(Xdata,Ydata,Zdata):
    Ydata=[float(i) for i in Ydata]        #since 'POSITION_SENSORS' radial stations are stored as strings
    Xdata=np.asarray(Xdata)
    Ydata=np.asarray(Ydata) 
    f = scipy.interpolate.interp2d(Ydata, Xdata, Zdata, kind='cubic')    
    Xdata_smooth = np.linspace(min(Xdata), max(Xdata), 1*len(Xdata))               #smoothing the xdata distribution by increasing the resolution to 5 times
    Ydata_smooth = np.linspace(min(Ydata), max(Ydata), 1*len(Ydata))
    r_new,theta_new=np.meshgrid(Ydata_smooth, Xdata_smooth) 
    Zdata_smooth=f(Ydata_smooth,Xdata_smooth)         
    # Make the plot
    fig = plt.figure()
    ax = fig.gca(projection='3d')
    surf2=ax.plot_surface(r_new,theta_new, Zdata_smooth, cmap=plt.cm.jet, rcount=len(Zdata_smooth), ccount=len(Zdata_smooth[0]))
    ax.set_ylabel('r')
    ax.set_xlabel(r'$\Psi$')
    fig.colorbar( surf2, shrink=0.5, aspect=5)
    # Rotate it
    ax.view_init(30, 220)
//  [python]
def three_d_surface_plot(Xdata,Ydata,Zdata):
    Ydata=[float(i) for i in Ydata]        #since 'POSITION_SENSORS' radial stations are stored as strings
    Xdata=np.asarray(Xdata)
    Ydata=np.asarray(Ydata) 
    tck = scipy.interpolate.bisplev(Ydata, Xdata, s=0)
    Xdata_smooth = np.linspace(min(Xdata), max(Xdata), 1*len(Xdata))               #smoothing the xdata distribution by increasing the resolution to 5 times
    Ydata_smooth = np.linspace(min(Ydata), max(Ydata), 1*len(Ydata))
    r_new,theta_new=np.meshgrid(Ydata_smooth, Xdata_smooth) 
    Zdata_smooth=scipy.interpolate.bisplev(Ydata_smooth, Xdata_smooth,tck)         
    # Make the plot
    fig = plt.figure()
    ax = fig.gca(projection='3d')
    surf2=ax.plot_surface(r_new,theta_new, Zdata_smooth, cmap=plt.cm.jet, rcount=len(Zdata_smooth), ccount=len(Zdata_smooth[0]))
    ax.set_ylabel('r')
    ax.set_xlabel(r'$\Psi$')
    fig.colorbar( surf2, shrink=0.5, aspect=5)
    # Rotate it
    ax.view_init(30, 220)
//  [python]
def plot_2d_tri_func(X, T, V, fig=None, kind="linear", N=200, colorbar=True):
    import scipy.interpolate

    x, y = X[:, 0], X[:, 1]
    fig = fig or plt.figure()
    # plt.triplot(x, y, triangles=T)
    f = sp.interpolate.interp2d(x, y, V, kind=kind, fill_value=0)
    xs, ys = np.linspace(min(x), max(x), num=N), np.linspace(min(y), max(y), num=N)
    zs = f(xs, ys)
    plt.pcolor(xs, ys, zs, cmap="Greys_r")
    if colorbar:
        plt.colorbar()
    plt.axis("off")
//  [python]
def plot_2d_tri_func(X, T, V, fig=None, kind="linear", N=200, colorbar=True):
    import scipy.interpolate
    import numpy as np
    import matplotlib.pyplot as plt

    x, y = X[:, 0], X[:, 1]
    fig = fig or plt.figure()
    # plt.triplot(x, y, triangles=T)
    tck = scipy.interpolate.bisplrep(x, y, V, kx=1, ky=1, s=0)
    xs, ys = np.linspace(min(x), max(x), num=N), np.linspace(min(y), max(y), num=N)
    zs = scipy.interpolate.bisplev(xs, ys, tck)
    plt.pcolor(xs, ys, zs, cmap="Greys_r")
    if colorbar:
        plt.colorbar()
    plt.axis("off")
//  [python]
def interpolate_by_points(self, tth, phi):
        """Interpolate by points."""
        l_phi_b = self.nu
        l_tth_b = self.gamma
        ll_int_b = self.intensity
        ll_int_b = [[float(ll_int_b[_2][_1]) for _2 in range(len(ll_int_b))]
                    for _1 in range(len(ll_int_b[0]))]
        if len(l_tth_b) == 0:
            int_2d = numpy.zeros((tth.size, phi.size), dtype=float)
        else:
            phi_b = numpy.array(l_phi_b, dtype=float)
            tth_b = numpy.array(l_tth_b, dtype=float)
            int_b = numpy.array(ll_int_b, dtype=float)
            func = scipy.interpolate.interp2d(tth_b, phi_b, int_b)
            # tth_2d, phi_2d = numpy.meshgrid(tth, phi, indexing="ij")
            int_2d = func(tth, phi)
            int_2d = int_2d.transpose()
        return int_2d
//  [python]
def interpolate_by_points(self, tth, phi):
        """Interpolate by points."""
        l_phi_b = self.nu
        l_tth_b = self.gamma
        ll_int_b = self.intensity
        ll_int_b = [[float(ll_int_b[_2][_1]) for _2 in range(len(ll_int_b))]
                    for _1 in range(len(ll_int_b[0]))]
        if len(l_tth_b) == 0:
            int_2d = numpy.zeros((tth.size, phi.size), dtype=float)
        else:
            phi_b = numpy.array(l_phi_b, dtype=float)
            tth_b = numpy.array(l_tth_b, dtype=float)
            int_b = numpy.array(ll_int_b, dtype=float)
            tck = scipy.interpolate.bisplrep(tth_b, phi_b, int_b)
            int_2d = scipy.interpolate.bisplev(tth, phi, tck)
            int_2d = int_2d.transpose()
        return int_2d
//  [python]
def set_input_rates_low_resolution(self, syn_type, positions):
        """
        Interpolate input_rates towards smaller resolution for each input cell

        Interpolation needs to be done on each input individually so we
        have to loop over the total number of input neurons.

        Note: Only needed for Gaussian process input.
        For normal input we make use of get_rates_function.
        NB: For a boxside switch experimetn with GRF input, n_total would not
        be the correct number.

        Parameters
        ----------
        syn_type : string
            Type of the synapse
        positions : ndarray
            Positions on which the high resultion input rates are defined
        """
        if self.dimensions  == 1:
            self.input_rates_low_resolution[syn_type] = np.empty(
                                (self.spacing, self.synapses[syn_type].n_total))
            for i in np.arange(self.synapses[syn_type].n_total):
                self.input_rates_low_resolution[syn_type][:,i] = np.interp(
                    np.squeeze(self.positions_grid),
                    np.squeeze(positions),
                    self.input_rates[syn_type][:,i])

        elif self.dimensions == 2:
            linspace_low_resolution = np.linspace(-self.radius, self.radius, self.spacing)
            linspace = np.squeeze(positions)[0,:,0]
            self.input_rates_low_resolution[syn_type] = np.empty(
                                (self.spacing, self.spacing,
                                 self.synapses[syn_type].n_total))
            for i in np.arange(self.synapses[syn_type].n_total):
                self.input_rates_low_resolution[syn_type][...,i] = scipy.interpolate.interp2d(
                    linspace,
                    linspace,
                    self.input_rates[syn_type][...,i])(linspace_low_resolution,
                                                     linspace_low_resolution)
//  [python]
def set_input_rates_low_resolution(self, syn_type, positions):
        """
        Interpolate input_rates towards smaller resolution for each input cell

        Interpolation needs to be done on each input individually so we
        have to loop over the total number of input neurons.

        Note: Only needed for Gaussian process input.
        For normal input we make use of get_rates_function.
        NB: For a boxside switch experimetn with GRF input, n_total would not
        be the correct number.

        Parameters
        ----------
        syn_type : string
            Type of the synapse
        positions : ndarray
            Positions on which the high resultion input rates are defined
        """
        if self.dimensions  == 1:
            self.input_rates_low_resolution[syn_type] = np.empty(
                                (self.spacing, self.synapses[syn_type].n_total))
            for i in np.arange(self.synapses[syn_type].n_total):
                self.input_rates_low_resolution[syn_type][:,i] = np.interp(
                    np.squeeze(self.positions_grid),
                    np.squeeze(positions),
                    self.input_rates[syn_type][:,i])

        elif self.dimensions == 2:
            linspace_low_resolution = np.linspace(-self.radius, self.radius, self.spacing)
            linspace = np.squeeze(positions)[0,:,0]
            self.input_rates_low_resolution[syn_type] = np.empty(
                                (self.spacing, self.spacing,
                                 self.synapses[syn_type].n_total))
            for i in np.arange(self.synapses[syn_type].n_total):
                x, y = np.mgrid[-self.radius:self.radius:self.spacing*1j, -self.radius:self.radius:self.spacing*1j]
                tck = scipy.interpolate.bisplev(linspace,linspace,scipy.interpolate.bisplrep(linspace,linspace,self.input_rates[syn_type][...,i]))
                self.input_rates_low_resolution[syn_type][...,i] = scipy.interpolate.bisplev(x,y,tck)
//  [python]
def project_2(source: CoordMask, target: CoordMask):

    s_mask = source.index_mask  # .astype(np.float64,casting='safe')
    #print(s_mask.mean())
    s_n_lat, s_n_lon = s_mask.shape

    t_mask = target.index_mask
    t_n_lat, t_n_lon = t_mask.shape

    t_tr = target.tr
    s_tr = source.tr
    # we first create the (unified) lat lon coordinates of the source array
    # and interpolate the source mask on them
    lats = np.array(list(map(s_tr.i2lat, range(s_n_lat))))
    lons = np.array(list(map(s_tr.i2lon, range(s_n_lon))))
    f = interpolate.interp2d(x=lons, y=lats, z=s_mask)
    # now we apply this interpolating function to the target grid
    # points
    target_lats = np.array(list(map(t_tr.i2lat, range(t_n_lat))))
    target_lons = np.array(list(map(t_tr.i2lon, range(t_n_lon))))
    # order lats and lons since f returns an array that assumes this anyway
    otlats, p_lat, p_lat_inv = target.ordered_lats()
    otlons, p_lon, p_lon_inv = target.ordered_lons()
    float_grid = p_lat_inv @ f(otlons, otlats) @ p_lon_inv
    #print(float_grid.mean())
    projected_mask = float_grid > 0.5
    # from IPython import embed; embed()
    return CoordMask(index_mask=np.logical_or(projected_mask, t_mask), tr=t_tr)
//  [python]
def project_2(source: CoordMask, target: CoordMask):

    s_mask = source.index_mask  # .astype(np.float64,casting='safe')
    #print(s_mask.mean())
    s_n_lat, s_n_lon = s_mask.shape

    t_mask = target.index_mask
    t_n_lat, t_n_lon = t_mask.shape

    t_tr = target.tr
    s_tr = source.tr
    # we first create the (unified) lat lon coordinates of the source array
    # and interpolate the source mask on them
    lats = np.array(list(map(s_tr.i2lat, range(s_n_lat))))
    lons = np.array(list(map(s_tr.i2lon, range(s_n_lon))))
    tck = interpolate.bisplrep(lons, lats, s_mask)
    # now we apply this interpolating function to the target grid
    # points
    target_lats = np.array(list(map(t_tr.i2lat, range(t_n_lat))))
    target_lons = np.array(list(map(t_tr.i2lon, range(t_n_lon))))
    # order lats and lons since f returns an array that assumes this anyway
    otlats, p_lat, p_lat_inv = target.ordered_lats()
    otlons, p_lon, p_lon_inv = target.ordered_lons()
    float_grid = p_lat_inv @ interpolate.bisplev(otlons, otlats, tck) @ p_lon_inv
    #print(float_grid.mean())
    projected_mask = float_grid > 0.5
    # from IPython import embed; embed()
    return CoordMask(index_mask=np.logical_or(projected_mask, t_mask), tr=t_tr)
//  [python]
def _second_order_quant_plot(
    fig, ax, quantiles_list, ale, mark_empty=True, n_interp=50, **kwargs
):
    """Second order ALE plot.

    Parameters
    ----------
    ax : matplotlib.axes.Axes
        Axis to plot onto.
    quantiles_list : array-like
        ALE quantiles for the first (`quantiles_list[0]`) and second
        (`quantiles_list[1]`) features.
    ale : masked array
        ALE to plot. Where `ale.mask` is True, this denotes bins where no samples were
        available. See `mark_empty`.
    mark_empty : bool, optional
        If True, plot rectangles over bins that did not contain any samples.
    n_interp : [2-iterable of] int, optional
        The number of interpolated samples generated from `ale` prior to contour
        plotting. Two integers may be given to specify different interpolation steps
        for the two features.
    **kwargs : contourf properties, optional
        Additional keyword parameters are passed to `ax.contourf`.

    Raises
    ------
    ValueError
        If `n_interp` values were not integers.
    ValueError
        If more than 2 values were given for `n_interp`.

    """
    centres_list = [_get_centres(quantiles) for quantiles in quantiles_list]
    n_x, n_y = _check_two_ints(n_interp)
    x = np.linspace(centres_list[0][0], centres_list[0][-1], n_x)
    y = np.linspace(centres_list[1][0], centres_list[1][-1], n_y)

    X, Y = np.meshgrid(x, y, indexing="xy")
    ale_interp = scipy.interpolate.interp2d(centres_list[0], centres_list[1], ale.T)
    CF = ax.contourf(X, Y, ale_interp(x, y), cmap="bwr", levels=30, alpha=0.7, **kwargs)

    if mark_empty and np.any(ale.mask):
        # Do not autoscale, so that boxes at the edges (contourf only plots the bin
        # centres, not their edges) don't enlarge the plot.
        plt.autoscale(False)
        # Add rectangles to indicate cells without samples.
        for i, j in zip(*np.where(ale.mask)):
            ax.add_patch(
                Rectangle(
                    [quantiles_list[0][i], quantiles_list[1][j]],
                    quantiles_list[0][i + 1] - quantiles_list[0][i],
                    quantiles_list[1][j + 1] - quantiles_list[1][j],
                    linewidth=1,
                    edgecolor="k",
                    facecolor="none",
                    alpha=0.4,
                )
            )
    fig.colorbar(CF)
//  [python]
def _second_order_quant_plot(
    fig, ax, quantiles_list, ale, mark_empty=True, n_interp=50, **kwargs
):
    """Second order ALE plot.

    Parameters
    ----------
    ax : matplotlib.axes.Axes
        Axis to plot onto.
    quantiles_list : array-like
        ALE quantiles for the first (`quantiles_list[0]`) and second
        (`quantiles_list[1]`) features.
    ale : masked array
        ALE to plot. Where `ale.mask` is True, this denotes bins where no samples were
        available. See `mark_empty`.
    mark_empty : bool, optional
        If True, plot rectangles over bins that did not contain any samples.
    n_interp : [2-iterable of] int, optional
        The number of interpolated samples generated from `ale` prior to contour
        plotting. Two integers may be given to specify different interpolation steps
        for the two features.
    **kwargs : contourf properties, optional
        Additional keyword parameters are passed to `ax.contourf`.

    Raises
    ------
    ValueError
        If `n_interp` values were not integers.
    ValueError
        If more than 2 values were given for `n_interp`.

    """
    centres_list = [_get_centres(quantiles) for quantiles in quantiles_list]
    n_x, n_y = _check_two_ints(n_interp)
    x = np.linspace(centres_list[0][0], centres_list[0][-1], n_x)
    y = np.linspace(centres_list[1][0], centres_list[1][-1], n_y)

    X, Y = np.meshgrid(x, y, indexing="xy")
    tck = scipy.interpolate.bisplrep(centres_list[0], centres_list[1], ale.T)
    ale_interp = scipy.interpolate.bisplev(x, y, tck)
    CF = ax.contourf(X, Y, ale_interp, cmap="bwr", levels=30, alpha=0.7, **kwargs)

    if mark_empty and np.any(ale.mask):
        # Do not autoscale, so that boxes at the edges (contourf only plots the bin
        # centres, not their edges) don't enlarge the plot.
        plt.autoscale(False)
        # Add rectangles to indicate cells without samples.
        for i, j in zip(*np.where(ale.mask)):
            ax.add_patch(
                Rectangle(
                    [quantiles_list[0][i], quantiles_list[1][j]],
                    quantiles_list[0][i + 1] - quantiles_list[0][i],
                    quantiles_list[1][j + 1] - quantiles_list[1][j],
                    linewidth=1,
                    edgecolor="k",
                    facecolor="none",
                    alpha=0.4,
                )
            )
    fig.colorbar(CF)
//  [python]
def get_dispersion(coords, dir_='C:\\Users\\olivar.DESKTOP-LS58OMC\\Work\\0_tesis\\phase_velocity\\tomo_code\\output', wave='rayleigh'):
        lon, lat = coords
        obs_ph_vel = []
        obs_lin_freq = []
        std = []

        for top_dir, sub_dir, files in os.walk(dir_):
            for file in files:
                full_path = os.path.join(top_dir, file)
                if full_path[-6:] == '.ttomo' and wave in file:
                    result_dict = pickle.load(open(full_path, 'rb'))
                    try:
                        freq = 1/result_dict['period']
                    except KeyError:
                        freq = 1/float(file.split('s')[0])
                    m_opt = result_dict['m_opt']
                    grid = result_dict['grid']
                    cov_map = result_dict['cov_map']
                    lons = grid[0,:][:,0]
                    lats = grid[:,0][:,1]
                    m_opt_interp = scint.interp2d(lons, lats, m_opt)
                    cov_map_interp = scint.interp2d(lons, lats, cov_map)

                    obs_ph_vel.append(m_opt_interp(lon, lat)[0])
                    obs_lin_freq.append(freq)
                    std.append(np.sqrt(cov_map_interp(lon, lat)[0]))

        return np.array(obs_ph_vel)[np.argsort(obs_lin_freq)], np.array(obs_lin_freq)[np.argsort(obs_lin_freq)], std
//  [python]
def get_dispersion(coords, dir_='C:\\Users\\olivar.DESKTOP-LS58OMC\\Work\\0_tesis\\phase_velocity\\tomo_code\\output', wave='rayleigh'):
        lon, lat = coords
        obs_ph_vel = []
        obs_lin_freq = []
        std = []

        for top_dir, sub_dir, files in os.walk(dir_):
            for file in files:
                full_path = os.path.join(top_dir, file)
                if full_path[-6:] == '.ttomo' and wave in file:
                    result_dict = pickle.load(open(full_path, 'rb'))
                    try:
                        freq = 1/result_dict['period']
                    except KeyError:
                        freq = 1/float(file.split('s')[0])
                    m_opt = result_dict['m_opt']
                    grid = result_dict['grid']
                    cov_map = result_dict['cov_map']
                    lons = grid[0,:][:,0]
                    lats = grid[:,0][:,1]
                    m_opt_interp = scint.bisplev(lons,lats,m_opt)
                    cov_map_interp = scint.bisplev(lons,lats,cov_map)

                    obs_ph_vel.append(m_opt_interp(lon, lat)[0])
                    obs_lin_freq.append(freq)
                    std.append(np.sqrt(cov_map_interp(lon, lat)[0]))

        return np.array(obs_ph_vel)[np.argsort(obs_lin_freq)], np.array(obs_lin_freq)[np.argsort(obs_lin_freq)], std
//  [python]
def __init__(
        self, spec_arr, eng=None, in_eng=None, rs=None,
        dlnz=-1, spec_type='dNdE', rebin_eng=None, with_interp_func = False
    ):

        super().__init__(
            spec_arr, eng=eng, in_eng=in_eng, rs=rs,spec_type=spec_type, rebin_eng=rebin_eng
        )
        self.dlnz = dlnz

        if with_interp_func:
            non_zero_grid = self.grid_vals
            # set zero values to some small value for log interp.
            non_zero_grid[np.abs(non_zero_grid) < 1e-100] = 1e-200

            # interp_grid  = np.log(non_zero_grid)

            interp_grid  = non_zero_grid

            self.interp_func = interpolate.interp2d(
                np.log(self.in_eng), np.log(self.eng),
                np.transpose(interp_grid), bounds_error = False,
                fill_value = 1e-200
            )
//  [python]
def __init__(
        self, spec_arr, eng=None, in_eng=None, rs=None,
        dlnz=-1, spec_type='dNdE', rebin_eng=None, with_interp_func = False
    ):

        super().__init__(
            spec_arr, eng=eng, in_eng=in_eng, rs=rs,spec_type=spec_type, rebin_eng=rebin_eng
        )
        self.dlnz = dlnz

        if with_interp_func:
            non_zero_grid = self.grid_vals
            # set zero values to some small value for log interp.
            non_zero_grid[np.abs(non_zero_grid) < 1e-100] = 1e-200

            # interp_grid  = np.log(non_zero_grid)

            interp_grid  = non_zero_grid

            self.interp_func = interpolate.bisplev(
                np.log(self.in_eng), np.log(self.eng),
                interpolate.bisplrep(np.log(self.in_eng), np.log(self.eng),
                                      np.transpose(interp_grid), kx=1, ky=1)
            )
//  [python]
def interpolate_forcing_data(input_variables,sample_points_lat_lon,znew_tmp,cubes,k):
    single_input_variable = input_variables[k]
    cube_year = cubes[k]
    znew_tmp[:] = np.nan
    print('processing '+single_input_variable)
    for j in range(np.shape(cube_year)[0]):
        # met_cubes[single_input_variable]['day']={}
        # progress_bar(j,np.shape(cube_year)[0])
        # different interpolation options commented out
        f = interpolate.interp2d(cube_year.coord('longitude').points, cube_year.coord('latitude').points, cube_year[j].data, kind='linear')
        znew_tmp[j,:] = [f(sample_points_lat_lon['lon'].values[i], sample_points_lat_lon['lat'].values[i])[0] for i in range(len(sample_points_lat_lon['lat'].values))]
        #https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.RectBivariateSpline.html
        # note, moving away from RectBivariateSpline because it can generate negative values
        #f = RectBivariateSpline(cube_year.coord('latitude').points,cube_year.coord('longitude').points, cube_year[j].data)
        #znew_tmp[j,:] = [f(sample_points_lat_lon['lat'].values[i], sample_points_lat_lon['lon'].values[i])[0] for i in range(len(sample_points_lat_lon['lat'].values))]
        #https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.Rbf.html
        # f = interp.Rbf(X,Y, cube_year[j].data, function='linear', smooth=0)  # default smooth=0 for interpolation
        # znew[j,k,:] = [f(sample_points_lat_lon['lat'].values[i], sample_points_lat_lon['lon'].values[i]) for i in range(len(sample_points_lat_lon['lat'].values))]
    # return [znew_tmp.reshape(j*len(sample_points_lat_lon['lat'].values))]
    return [[znew_tmp]]
//  [python]
def interpolate_forcing_data(input_variables,sample_points_lat_lon,znew_tmp,cubes,k):
    single_input_variable = input_variables[k]
    cube_year = cubes[k]
    znew_tmp[:] = np.nan
    print('processing '+single_input_variable)
    for j in range(np.shape(cube_year)[0]):
        # met_cubes[single_input_variable]['day']={}
        # progress_bar(j,np.shape(cube_year)[0])
        # different interpolation options commented out
        tck = interpolate.bisplrep(cube_year.coord('longitude').points, cube_year.coord('latitude').points, cube_year[j].data)
        znew_tmp[j,:] = [interpolate.bisplev(sample_points_lat_lon['lon'].values[i], sample_points_lat_lon['lat'].values[i], tck) for i in range(len(sample_points_lat_lon['lat'].values))]
        #https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.RectBivariateSpline.html
        # note, moving away from RectBivariateSpline because it can generate negative values
        #f = RectBivariateSpline(cube_year.coord('latitude').points,cube_year.coord('longitude').points, cube_year[j].data)
        #znew_tmp[j,:] = [f(sample_points_lat_lon['lat'].values[i], sample_points_lat_lon['lon'].values[i])[0] for i in range(len(sample_points_lat_lon['lat'].values))]
        #https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.Rbf.html
        # f = interp.Rbf(X,Y, cube_year[j].data, function='linear', smooth=0)  # default smooth=0 for interpolation
        # znew[j,k,:] = [f(sample_points_lat_lon['lat'].values[i], sample_points_lat_lon['lon'].values[i]) for i in range(len(sample_points_lat_lon['lat'].values))]
    # return [znew_tmp.reshape(j*len(sample_points_lat_lon['lat'].values))]
    return [[znew_tmp]]
//  [python]
def resample2d(array:np.ndarray, new_size:int, kind='linear'):
    """resample2d

    Resamples a 2D (or 3D) array with :func:`scipy.interpolate.interp2d` to `new_size`. If input is 2D, we'll loop over the final axis.

    Parameters
    ----------
    array: np.ndarray
        Array to be interpolated. Ideally axis have the same size.
    new_size: int
        New size of array
    kind: str, optional
        Interpolation method, by default 'linear'

    Returns
    ----------
    np.ndarray
        If 2D: resampled array of shape `(new_size,new_size)`
        If 3D: resampled array of shape `(new_size,new_size, array.shape[-1])`
    """
    # set up interpolater
    x = np.array(range(array.shape[0]))
    y = np.array(range(array.shape[1]))

    # define new grid
    xnew = np.linspace(0, x.shape[0], new_size)
    ynew = np.linspace(0, y.shape[0], new_size)

    if array.ndim > 2:
        new = np.zeros((new_size,new_size,array.shape[-1]))

        for dd in range(array.shape[-1]):
            f = interpolate.interp2d(x, y, array[...,dd], kind=kind)
            new[...,dd] = f(xnew,ynew)

        return new    
    else:
        f = interpolate.interp2d(x, y, array, kind=kind)
        return f(xnew,ynew)
//  [python]
def resample2d(array:np.ndarray, new_size:int, kind='linear'):
    """resample2d

    Resamples a 2D (or 3D) array with :func:`scipy.interpolate.interp2d` to `new_size`. If input is 2D, we'll loop over the final axis.

    Parameters
    ----------
    array: np.ndarray
        Array to be interpolated. Ideally axis have the same size.
    new_size: int
        New size of array
    kind: str, optional
        Interpolation method, by default 'linear'

    Returns
    ----------
    np.ndarray
        If 2D: resampled array of shape `(new_size,new_size)`
        If 3D: resampled array of shape `(new_size,new_size, array.shape[-1])`
    """
    # set up interpolater
    x = np.array(range(array.shape[0]))
    y = np.array(range(array.shape[1]))

    # define new grid
    xnew = np.linspace(0, x.shape[0]-1, new_size)
    ynew = np.linspace(0, y.shape[0]-1, new_size)
    X,Y = np.meshgrid(x,y)
    tck = interpolate.bisplrep(X, Y, array, s=0)


    if array.ndim > 2:
        new = np.zeros((new_size,new_size,array.shape[-1]))

        for dd in range(array.shape[-1]):
            new[...,dd] = interpolate.bisplev(xnew,ynew,tck)

        return new    
    else:
        return interpolate.bisplev(xnew,ynew,tck)
//  [python]
def _create_interpolator(u_intensities):

    """This function creates the interpolator object to use when computing double-differential underground intensities."""

    global _interpolator

    _interpolator = sciint.interp2d(
        constants._SLANT_DEPTHS,
        constants.ANGLES_FOR_S_FLUXES,
        np.log(u_intensities.T),
        kind="linear",
        fill_value=-np.inf,
    )

    return _interpolator
//  [python]
def _create_interpolator(u_intensities):

    """This function creates the interpolator object to use when computing double-differential underground intensities."""

    global _interpolator

    x, y = np.meshgrid(constants._SLANT_DEPTHS, constants.ANGLES_FOR_S_FLUXES)
    tck = splprep([x.ravel(), y.ravel(), np.log(u_intensities.T).ravel()], s=0)
    _interpolator = lambda x, y: bisplev(x, y, tck)

    return _interpolator
//  [python]
def _second_order_quant_plot(ax, quantiles, ALE, **kwargs):
    print(ALE)
    x = np.linspace(quantiles[0][0], quantiles[0][-1], 50)
    y = np.linspace(quantiles[1][0], quantiles[1][-1], 50)
    X, Y = np.meshgrid(x, y)
    # ax.imshow(ALE, interpolation="bilinear", cmap='viridis')
    ALE_interp = scipy.interpolate.interp2d(quantiles[0], quantiles[1], ALE)
    CF = ax.contourf(X, Y, ALE_interp(x, y), cmap="bwr", levels=30, alpha=0.7)
    plt.colorbar(CF)
//  [python]
def _second_order_quant_plot(ax, quantiles, ALE, **kwargs):
    print(ALE)
    x = np.linspace(quantiles[0][0], quantiles[0][-1], 50)
    y = np.linspace(quantiles[1][0], quantiles[1][-1], 50)
    X, Y = np.meshgrid(x, y)
    # ax.imshow(ALE, interpolation="bilinear", cmap='viridis')
    ALE_interp = scipy.interpolate.bisplev(quantiles[0], quantiles[1], ALE)
    CF = ax.contourf(X, Y, ALE_interp, cmap="bwr", levels=30, alpha=0.7)
    plt.colorbar(CF)
//  [python]
def imresize(D, dims, kind='cubic'):
    """
    Resize a floating point image
    Parameters
    ----------
    D : ndarray(M1, N1)
        Original image
    dims : tuple(M2, N2)
        The dimensions to which to resize
    kind : string
        The kind of interpolation to use
    Returns
    -------
    D2 : ndarray(M2, N2)
        A resized array
    """
    M, N = dims
    x1 = np.array(0.5 + np.arange(D.shape[1]), dtype=np.float32)/D.shape[1]
    y1 = np.array(0.5 + np.arange(D.shape[0]), dtype=np.float32)/D.shape[0]
    x2 = np.array(0.5 + np.arange(N), dtype=np.float32)/N
    y2 = np.array(0.5 + np.arange(M), dtype=np.float32)/M
    f = scipy.interpolate.interp2d(x1, y1, D, kind=kind)
    return f(x2, y2)
//  [python]
def imresize(D, dims, kind='cubic'):
    """
    Resize a floating point image
    Parameters
    ----------
    D : ndarray(M1, N1)
        Original image
    dims : tuple(M2, N2)
        The dimensions to which to resize
    kind : string
        The kind of interpolation to use
    Returns
    -------
    D2 : ndarray(M2, N2)
        A resized array
    """
    M, N = dims
    x1 = np.array(0.5 + np.arange(D.shape[1]), dtype=np.float32)/D.shape[1]
    y1 = np.array(0.5 + np.arange(D.shape[0]), dtype=np.float32)/D.shape[0]
    x2 = np.array(0.5 + np.arange(N), dtype=np.float32)/N
    y2 = np.array(0.5 + np.arange(M), dtype=np.float32)/M
    x1, y1 = np.meshgrid(x1, y1)
    x2, y2 = np.meshgrid(x2, y2)
    tck = scipy.interpolate.bisplrep(x1, y1, D, kx=3, ky=3, s=0)
    return scipy.interpolate.bisplev(x2[:,0], y2[0,:], tck)
//  [python]
def eval_spline(x, y, tck):
    """
    Evaluate a bivariate spline on the Cartesian product of a set of X points
    and a set of Y points.
    args
    ----
        x       :   1D ndarray, the array of unique x points
        y       :   1D ndarray, the array of unique y points
        tck     :   5-tuple, bivariate spline coefficients of the type
                    read by *load_spline_coefs*
    returns
    -------
        2D ndarray of shape (y.shape[0], x.shape[0]), the evaluated
            bivariate spline at each combination of the input points
    """
    return interpolate.bisplev(x, y, tck).T
//  [python]
def example_Use_of_Bispline():

    X, Y = np.mgrid[0:10, 0:10]

    Z = fit2d.gaussian(100, 10, 10, 2, 2)(X, Y)
    tck, _, _, _ = sci.bisplrep(np.ravel(X), np.ravel(Y), np.ravel(Z), kx=4, ky=4, full_output=1)

    Zfit = sci.bisplev(np.arange(10), np.arange(10), tck)

    for k in np.arange(9):
        p.subplot(3, 3, k + 1)
        p.plot(Z[k], "bo - ")
        p.plot(Zfit[k], "ro - ")

    p.show()
//  [python]
def _residual(self, aE_vec):
        """Given a policy function, solves the model's optimality condition and returns the error

        An optimal policy function will return an error close to zero. This method will typically 
        be called within a numerical solver to find the optimal policy function. 

        Args:
            aE_vec (np.array): 1 dimensional array with as many elements as the policy function(s)

        Returns:
            numpy array of 1 dimension representing optimality equation errors
        """
        # reshape input (candidate solution) to shape of policy function, assign to self.aE
        self._unpack_residual_inputs(aE_vec)
        self._get_splines()

        # back out current model variables based on policy function
        q = self._get_q(self.aE)
        theta = self._get_theta(q)
        V = self._get_V(theta, self.N)

        # employment tomorrow
        N_p = self._get_N_tplusone(self.N, q, V)

        # interpolate for tomorrow's policy function (will need to check this to see if it holds w/ log consumption)
        func = sp.interpolate.interp2d(self.X, self.N, self.aE) 
        E_p = np.array([func(self.X[i_n], N_p[:, i_n]) for i_n in range(self.nx)])
        E_p = E_p[:, :, 0].T

        # back out future model variables based on policy function
        q_p = self._get_q(E_p)
        theta_p = self._get_theta(q_p)
        lambda_p = self._get_lambda(E_p, q_p)
        gamma_p = self.gamma + self.c*q_p
        w_p = self._get_W(self.X, theta_p, q_p)

        # error in the optimal job creation condition (inside the expectation operator)
        inside_p = self.beta*(self.X - w_p + (1 - self.s)*(gamma_p/q_p - lambda_p))

        # expectations over future states
        rhs = np.dot(inside_p, self.P)
        err = self.aE - rhs
        return err.flatten()