Python numexpr.evaluate函数代码示例

def boo_product(qlm1, qlm2):
    if qlm1.ndim==2 and qlm2.ndim==2:
        prod = np.empty([len(qlm1), len(qlm2)])
        code ="""
        #pragma omp parallel for
        for(int p=0; p<Nqlm1[0]; ++p)
            for(int q=0; q<Nqlm2[0]; ++q)
            {
                prod(p,q) = real(qlm1(p,0)*conj(qlm2(q,0)));
                for(int m=1; m<Nqlm1[1]; ++m)
                    prod(p,q) += 2.0*real(qlm1(p,m)*conj(qlm2(q,m)));
                prod(p,q) *= 4.0*M_PI/(2.0*(Nqlm1[1]-1)+1);
            }
        """
        weave.inline(
            code,['qlm1', 'qlm2', 'prod'],
            type_converters =converters.blitz,
            extra_compile_args =['-O3 -fopenmp'],
            extra_link_args=['-lgomp'],
            verbose=2, compiler='gcc')
        return prod
    else:
        n = np.atleast_2d(numexpr.evaluate(
            """real(complex(real(a), -imag(a)) * b)""",
            {'a':qlm1, 'b':qlm2}
            ))
        p = numexpr.evaluate(
            """4*pi/(2*l+1)*(2*na + nb)""",
            {
                'na': n[:,1:].sum(-1),
                'nb': n[:,0],
                'l': n.shape[1]-1,
                'pi': np.pi
                })
        return p

def circ_mask(arr, axis, ratio=1, val=0., ncore=None):
    """
    Apply circular mask to a 3D array.

    Parameters
    ----------
    arr : ndarray
            Arbitrary 3D array.
    axis : int
        Axis along which mask will be performed.
    ratio : int, optional
        Ratio of the mask's diameter in pixels to
        the smallest edge size along given axis.
    val : int, optional
        Value for the masked region.

    Returns
    -------
    ndarray
        Masked array.
    """
    arr = dtype.as_float32(arr)
    val = np.float32(val)
    _arr = arr.swapaxes(0, axis)
    dx, dy, dz = _arr.shape
    mask = _get_mask(dy, dz, ratio)

    with mproc.set_numexpr_threads(ncore):
        ne.evaluate('where(mask, _arr, val)', out=_arr)

    return _arr.swapaxes(0, axis)

 def test_in_place(self):
     x = arange(10000.).reshape(1000,10)
     evaluate("x + 3", out=x)
     assert_equal(x, arange(10000.).reshape(1000,10) + 3)
     y = arange(10)
     evaluate("(x - 3) * y + (x - 3)", out=x)
     assert_equal(x, arange(10000.).reshape(1000,10) * (arange(10) + 1))

def flo_density_seawater(degC, psu):
    """
    Helper function for flo_zhang_scatter_coeffs
    @param degC in situ water temperature
    @param psu in site practical salinity
    @retval rho_sw density of seawater
    """
    # density of water and seawater,unit is Kg/m^3, from UNESCO,38,1981
    a0 = 8.24493e-1
    a1 = -4.0899e-3
    a2 = 7.6438e-5
    a3 = -8.2467e-7
    a4 = 5.3875e-9
    a5 = -5.72466e-3
    a6 = 1.0227e-4
    a7 = -1.6546e-6
    a8 = 4.8314e-4
    b0 = 999.842594
    b1 = 6.793952e-2
    b2 = -9.09529e-3
    b3 = 1.001685e-4
    b4 = -1.120083e-6
    b5 = 6.536332e-9

    # density for pure water
    rho_w = ne.evaluate('b0 + b1 * degC + b2 * degC**2 + b3 * degC**3'
                        '+ b4 * degC**4 + b5 * degC**5')

    # density for pure seawater
    rho_sw = ne.evaluate('rho_w + ((a0 + a1 * degC + a2 * degC**2'
                         '+ a3 * degC**3 + a4 * degC**4) * psu'
                         '+ (a5 + a6 * degC + a7 * degC**2) * psu**1.5 + a8 * psu**2)')

    return rho_sw

 def test_empty_string2(self):
     a = np.array([b"p", b"pepe"])
     b = np.array([b"pepe2", b""])
     res = evaluate("(a == b'') & (b == b'pepe2')")
     assert_array_equal(res, np.array([False, False]))
     res2 = evaluate("(a == b'pepe') & (b == b'')")
     assert_array_equal(res, np.array([False, False]))

def _ndmi(swir1, nir, output_scaling=1.0, **kwargs):
    """ Return the Normalized Difference Moisture Index (NDMI)

    NDMI is calculated as:

    .. math::
        NDMI = \\frac{(NIR - SWIR1)}{(NIR + SWIR1)}

    where:
        - :math:`SWIR1` is the shortwave infrared band
        - :math:`NIR` is the near infrared band

    Args:
      swir1 (np.ndarray): SWIR1 band
      nir (np.ndarray): NIR band
      output_scaling (float): scaling factor for output NDVI (default: 1.0)

    Returns:
      np.ndarray: NDMI

    """
    expr = '(nir - swir1) / (nir + swir1)'
    if output_scaling == 1.0:
        return ne.evaluate(expr)
    else:
        return ne.evaluate(expr) * output_scaling

def evaluate_expr(expr, xs, ys, zs, arr=None, alpha=None):
    """Evaluate symbolic expression on a numerical grid.

    Args:
        expr (symbolic): sympy or symengine expression
        xs (np.ndarray): 1D-array of x values
        ys (np.ndarray): 1D-array of y values
        zs (np.ndarray): 1D-array of z values
        arr (np.ndarray): additional 1D-array to multiply expression by
        alpha (float): multiply expression by gaussian with exponent alpha

    Note:
        See :meth:`exatomic.algorithms.orbital_util.numerical_grid_from_field_params`
        for grid construction details.
    """
    subs = {_x: 'xs', _y: 'ys', _z: 'zs'}
    # Multiply with an additional array (angular term)
    if arr is not None:
        return evaluate('arr * ({})'.format(str(expr.subs(subs))))
    # Multiply by an exponential decay factor
    if alpha is not None:
        expr = str((expr * exp(-alpha * _r ** 2)).subs(subs))
        return evaluate(expr)
    # Just evaluate the expression numerically
    return evaluate(str(expr.subs(subs)))

def grad(fld, bnd=True):
    """2nd order centeral diff, 1st order @ boundaries if bnd"""
    # vx, vy, vz = fld.component_views()
    if bnd:
        fld = viscid.extend_boundaries(fld, order=0, crd_order=0)

    if fld.iscentered("Cell"):
        crdx, crdy, crdz = fld.get_crds_cc(shaped=True)
        # divcenter = "Cell"
        # divcrds = coordinate.NonuniformCartesianCrds(fld.crds.get_clist(np.s_[1:-1]))
        # divcrds = fld.crds.slice_keep(np.s_[1:-1, 1:-1, 1:-1])
    elif fld.iscentered("Node"):
        crdx, crdy, crdz = fld.get_crds_nc(shaped=True)
        # divcenter = "Node"
        # divcrds = coordinate.NonuniformCartesianCrds(fld.crds.get_clist(np.s_[1:-1]))
        # divcrds = fld.crds.slice_keep(np.s_[1:-1, 1:-1, 1:-1])
    else:
        raise NotImplementedError("Can only do cell and node centered gradients")

    v = fld.data
    g = viscid.zeros(fld['x=1:-1, y=1:-1, z=1:-1'].crds, nr_comps=3)

    xp, xm = crdx[2:,  :,  :], crdx[:-2, :  , :  ]  # pylint: disable=bad-whitespace
    yp, ym = crdy[ :, 2:,  :], crdy[:  , :-2, :  ]  # pylint: disable=bad-whitespace
    zp, zm = crdz[ :,  :, 2:], crdz[:  , :  , :-2]  # pylint: disable=bad-whitespace

    vxp, vxm = v[2:  , 1:-1, 1:-1], v[ :-2, 1:-1, 1:-1]  # pylint: disable=bad-whitespace
    vyp, vym = v[1:-1, 2:  , 1:-1], v[1:-1,  :-2, 1:-1]  # pylint: disable=bad-whitespace
    vzp, vzm = v[1:-1, 1:-1, 2:  ], v[1:-1, 1:-1,  :-2]  # pylint: disable=bad-whitespace

    g['x'].data[...] = ne.evaluate("(vxp-vxm)/(xp-xm)")
    g['y'].data[...] = ne.evaluate("(vyp-vym)/(yp-ym)")
    g['z'].data[...] = ne.evaluate("(vzp-vzm)/(zp-zm)")
    return g

    def set_rho(self, rho):
        """
        Set the initial density matrix
        :param rho: 2D numpy array or sting containing the density matrix
        :return: self
        """
        if isinstance(rho, str):
            # density matrix is supplied as a string
            ne.evaluate("%s + 0j" % rho, local_dict=vars(self), out=self.rho)

        elif isinstance(rho, np.ndarray):
            # density matrix is supplied as an array

            # perform the consistency checks
            assert rho.shape == self.rho.shape,\
                "The grid size does not match with the density matrix"

            # make sure the density matrix is stored as a complex array
            np.copyto(self.rho, rho.astype(np.complex))

        else:
            raise ValueError("density matrix must be either string or numpy.array")

        # normalize
        self.rho /= self.rho.trace() * self.dX

        return self

def rotate_clip(data_np, theta_deg, rotctr_x=None, rotctr_y=None,
                out=None):
    """
    Rotate numpy array `data_np` by `theta_deg` around rotation center
    (rotctr_x, rotctr_y).  If the rotation center is omitted it defaults
    to the center of the array.

    No adjustment is done to the data array beforehand, so the result will
    be clipped according to the size of the array (the output array will be
    the same size as the input array).
    """
    
    # If there is no rotation, then we are done
    if math.fmod(theta_deg, 360.0) == 0.0:
        return data_np

    ht, wd = data_np.shape[:2]

    if rotctr_x == None:
        rotctr_x = wd // 2
    if rotctr_y == None:
        rotctr_y = ht // 2

    yi, xi = numpy.mgrid[0:ht, 0:wd]
    xi -= rotctr_x
    yi -= rotctr_y
    cos_t = numpy.cos(numpy.radians(theta_deg))
    sin_t = numpy.sin(numpy.radians(theta_deg))

    #t1 = time.time()
    if have_numexpr:
        ap = ne.evaluate("(xi * cos_t) - (yi * sin_t) + rotctr_x")
        bp = ne.evaluate("(xi * sin_t) + (yi * cos_t) + rotctr_y")
    else:
        ap = (xi * cos_t) - (yi * sin_t) + rotctr_x
        bp = (xi * sin_t) + (yi * cos_t) + rotctr_y
    #print "rotation in %.5f sec" % (time.time() - t1)

    #ap = numpy.rint(ap).astype('int').clip(0, wd-1)
    #bp = numpy.rint(bp).astype('int').clip(0, ht-1)
    # Optomizations to reuse existing intermediate arrays
    numpy.rint(ap, out=ap)
    ap = ap.astype('int')
    ap.clip(0, wd-1, out=ap)
    numpy.rint(bp, out=bp)
    bp = bp.astype('int')
    bp.clip(0, ht-1, out=bp)

    if out != None:
        out[:, :, ...] = data_np[bp, ap]
        newdata = out
    else:
        newdata = data_np[bp, ap]
        new_ht, new_wd = newdata.shape[:2]

        assert (wd == new_wd) and (ht == new_ht), \
               Exception("rotated cutout is %dx%d original=%dx%d" % (
            new_wd, new_ht, wd, ht))

    return newdata

def get_all_peak_over_threshold(sig, thresh, front, peak_span = 3, use_numexpr = False):
    """
    Simple solution for detectng peak aboe (or below) a threshold
    """
    peak_span = max(peak_span, 3)
    k = (peak_span-1)//2
    sig_center = sig[k:-k]
    
    if use_numexpr:
        if front == '+':
            pos_spike, = np.where(numexpr.evaluate( '(sig1<sig2) & (sig2>thresh) & (sig2>=sig3)'))
        elif front == '-':
            pos_spike, = np.where(numexpr.evaluate( '(sig1>sig2) & (sig2<thresh) & (sig2<=sig3)'))
    else :
        if front == '+':
            peaks = sig_center>thresh
            for i in range(k):
                peaks &= sig_center>sig[i:i+sig_center.size]
                peaks &= sig_center>=sig[peak_span-i-1:peak_span-i-1+sig_center.size]
        elif front == '-':
            peaks = sig_center<thresh
            for i in range(k):
                peaks &= sig_center<sig[i:i+sig_center.size]
                peaks &= sig_center<=sig[peak_span-i-1:peak_span-i-1+sig_center.size]
    ind_peak,  = np.where(peaks)
    ind_peak = ind_peak + k
    return ind_peak
    

    

 def apply_linear_filterbank(b, a, x, zi):
     X = x
     output = empty_like(X)
     for sample in xrange(X.shape[0]):
         x = X[sample]
         for curf in xrange(zi.shape[2]):
             #y = b[:, 0, curf]*x+zi[:, 0, curf]
             y = numexpr.evaluate('b*x+zi', local_dict={
                     'b':b[:, 0, curf],
                     'x':x,
                     'zi':zi[:, 0, curf]})
             for i in xrange(b.shape[1]-2):
                 #zi[:, i, curf] = b[:, i+1, curf]*x+zi[:, i+1, curf]-a[:, i+1, curf]*y
                 zi[:, i, curf] = numexpr.evaluate('b*x+zi-a*y', local_dict={
                                         'b':b[:, i+1, curf],
                                         'x':x,
                                         'zi':zi[:, i+1, curf],
                                         'a':a[:, i+1, curf],
                                         'y':y})
             i = b.shape[1]-2
             #zi[:, i, curf] = b[:, i+1, curf]*x-a[:, i+1, curf]*y
             zi[:, i, curf] = numexpr.evaluate('b*x-a*y', local_dict={
                                     'b':b[:, i+1, curf],
                                     'x':x,
                                     'a':a[:, i+1, curf],
                                     'y':y})
             x = y
         output[sample] = y
     return output

def sideband(t, plus, minus, freq=0, phase=0, offset=False, offset_fit_lin=0,offset_fit_quad=0, origin=0):
    if freq==0 and phase==0:
        return (plus - minus), np.zeros(t.shape)

    elif offset:
        if (not max(plus) == 0):
            time_step = t[1]-t[0]
            freq_calibrated = getFreq(plus,freq,offset_fit_lin,offset_fit_quad);
            freq_integ_array = np.cumsum(freq_calibrated)*time_step
            # np.savetxt('time.out', t, delimiter=',')
        return ( np.cos(2 * np.pi * (freq_integ_array/1.0e9) + phase*np.pi/180.0) * plus - np.cos(2 * np.pi * (freq_integ_array/1.0e9) + phase*np.pi/180.0) * minus,
             -np.sin(2 * np.pi * (freq_integ_array/1.0e9)+ phase*np.pi/180.0) * plus - np.sin(2 * np.pi * (freq_integ_array/1.0e9) + phase*np.pi/180.0) * minus)
    else:
        # For ML0218 Mixer
        # return ( np.cos(2 * np.pi * (freq/1.0e9 * t)+ phase*np.pi/180.0) * plus - np.cos(2 * np.pi * (freq/1.0e9 * t) + phase*np.pi/180.0) * minus,
        #      +np.sin(2 * np.pi * (freq/1.0e9 * t)+ phase*np.pi/180.0) * plus +np.sin(2 * np.pi * (freq/1.0e9 * t) + phase*np.pi/180.0) * minus)
        #
        # For IQ0317 Mixer
        # return ( np.cos(2 * np.pi * (freq/1.0e9 * t)+ phase*np.pi/180.0) * plus - np.cos(2 * np.pi * (freq/1.0e9 * t) + phase*np.pi/180.0) * minus,
        #      -np.sin(2 * np.pi * (freq/1.0e9 * t)+ phase*np.pi/180.0) * plus - np.sin(2 * np.pi * (freq/1.0e9 * t) + phase*np.pi/180.0) * minus)

        wts = ne.evaluate('2 * (freq/1.0e9 * (t-origin))+ phase/180.0') * np.pi
        cosdata = ne.evaluate('cos(wts)')
        sindata = ne.evaluate('sin(wts)')
        # return ne.evaluate('cosdata * (plus - minus)'), ne.evaluate('- sindata * (plus + minus)')
        return ne.evaluate('cosdata * (plus - minus)'), ne.evaluate('- sindata * (plus + minus)')

 def correlate(self):
     '''Estimates the time-delay in whole samples, then aligns and finds the complex visibility for two signals in the data key of self.parent_row and self.child_row.
     Generates self.convolution_spectrum, self.zscore, self.expected_overlap, and self.real_overlap, which can be used for plotting purposes.
     Creates two signals, one corresponding to the cos and the other to the sin signal, as self.cos_signal and self.sin_signal
     In order to do a correlation without opening a file and populating self.child_row and self.parent_row, assign these variables manually.'''
     child_fft = scipy.fftpack.fft(self.child_row['data'])
     parent_fft = scipy.fftpack.fft(self.parent_row['data'])
     child_inverse_conjugate = -child_fft.conjugate()
     fft_auto_child = child_fft*child_inverse_conjugate
     fft_convolution = parent_fft*child_inverse_conjugate
     self.convolution_spectrum = numpy.abs(scipy.fftpack.ifft(fft_convolution/fft_auto_child)) #Roth window saved in self.convolution_spectrum
     assert self.time_inacc < sec/2, "This system clock cannot *possibly* be that inaccurate!"
     expected_tdiff = int((self.child_row['utcendtime']-self.parent_row['utcendtime'])*hz) #our initial guess for the location of the tdiff peak
     sample_inacc = int(self.time_inacc*hz) #around the guessed place of tdiff.
     if expected_tdiff-sample_inacc < 0: expected_tdiff = sample_inacc
     elif expected_tdiff+sample_inacc > sample_size: expected_tdiff = sample_size-sample_inacc
     cropped_convolution_spectrum = self.convolution_spectrum[expected_tdiff-sample_inacc:expected_tdiff+sample_inacc]
     #later measurements of tdiff will have a 0 point at expected_tdiff-sample_inacc and a total length of around 2*sample_inacc (may wrap around)
     tdiff = numpy.argmax(cropped_convolution_spectrum)
     tdiff += expected_tdiff-sample_inacc #offset for the real convolution_spectrum
     self.zscore = (self.convolution_spectrum[tdiff]-numpy.average(self.convolution_spectrum))/numpy.std(self.convolution_spectrum)
     #timespan = math.fabs(child_row['utcendtime'] - parent_row['utcendtime']) + sec
     self.expected_overlap = (float(sec) - 1.0*math.fabs(self.child_row['utcendtime']-self.parent_row['utcendtime']))*hz
     self.real_overlap = (float(sec) - 1.0*math.fabs(float(tdiff)/hz))*hz
     int_delay = int(numpy.copysign(numpy.floor(numpy.fabs(tdiff)), tdiff))
     self.abs_delay = int(abs(int_delay)) #always positive :)
     parent_signal, child_signal = self.parent_row['data'][self.abs_delay:], self.child_row['data'][0:len(self.child_row['data'])-self.abs_delay]
     h_child_length = len(child_signal)
     h_child_new_length = int(2**numpy.ceil(numpy.log2(h_child_length)))
     h_child_diff_length = h_child_new_length - h_child_length
     h_child_signal = scipy.fftpack.hilbert(numpy.append(child_signal, numpy.zeros(h_child_diff_length)))[0:h_child_length]
     self.cos_signal, self.sin_signal = evaluate("parent_signal*child_signal"), evaluate("h_child_signal*parent_signal")

def get_smooth_length(bar):
    """Figures out if the particles are from AREPO or GADGET
    and computes the smoothing length.
    Note the Volume array in HDF5 is comoving and this returns a comoving smoothing length
    If we are Arepo, this smoothing length is  cell radius, where
    cell volume = 4/3 π (cell radius) **3 and cell volume = mass / density
    Arguments:
        Baryon particles from a simulation
    Returns:
        Array of smoothing lengths in code units.
    """
    # Are we arepo? If we are a modern version we should have this array.
    if "Volume" in bar.keys():
        volume = np.array(bar["Volume"], dtype=np.float32)
        radius = ne.evaluate("(scal*volume)**poww")
    elif "Number of faces of cell" in bar.keys():
        rho = np.array(bar["Density"])
        mass = np.array(bar["Masses"])
        radius = ne.evaluate("(scal*mass/rho)**poww")
    else:
        # If we are gadget, the SmoothingLength array is actually the smoothing length.
        # Note the definition used in Gadget differs from that used here: in gadget volume = h^3
        # due to the normalisation of the SPH kernel.
        radius = np.array(bar["SmoothingLength"]) * scal ** poww
    return radius