Selection.py 文件源码

def integrate(self, params, sel_dist, theta):
        """
        integration without re-normalizing the DFE. This assumes the
        portion of the DFE that is not integrated is not seen in your
        sample. 
        """
        #need to include tuple() here to make this function play nice
        #with numpy arrays
        sel_args = (self.gammas,) + tuple(params)
        #compute weights for each fs
        weights = sel_dist(*sel_args)

        #compute weight for the effectively neutral portion. not using
        #CDF function because I want this to be able to compute weight
        #for arbitrary mass functions
        weight_neu, err_neu = scipy.integrate.quad(sel_dist, self.gammas[-1], 
                                                   0, args=tuple(params))

        #function's adaptable for demographic models from 1-3 populations
        pops = len(self.neu_spec.shape)
        if pops == 1:
            integrated = self.neu_spec*weight_neu + Numerics.trapz(
                weights[:,numpy.newaxis]*self.spectra, self.gammas, axis=0)
        elif pops == 2:
            integrated = self.neu_spec*weight_neu + Numerics.trapz(
                weights[:,numpy.newaxis,numpy.newaxis]*self.spectra, 
                self.gammas, axis=0)
        elif pops == 3:
            integrated = self.neu_spec*weight_neu + Numerics.trapz(
                weights[:,numpy.newaxis,numpy.newaxis,numpy.newaxis]*self.spectra,
                self.gammas, axis=0)
        else:
            raise IndexError("Must have one to three populations")

        integrated_fs = Spectrum(integrated, extrap_x=self.extrap_x)

        #no normalization, allow lethal mutations to fall out
        return integrated_fs * theta