neulab/docprompting-conala huggingface dataset:
extending.pyx #!/usr/bin/env python3
#cython: language_level=3

from libc.stdint cimport uint32_t
from cpython.pycapsule cimport PyCapsule_IsValid, PyCapsule_GetPointer

import numpy as np
cimport numpy as np
cimport cython

from numpy.random cimport bitgen_t
from numpy.random import PCG64

np.import_array()


@cython.boundscheck(False)
@cython.wraparound(False)
def uniform_mean(Py_ssize_t n):
    cdef Py_ssize_t i
    cdef bitgen_t *rng
    cdef const char *capsule_name = "BitGenerator"
    cdef double[::1] random_values
    cdef np.ndarray randoms

    x = PCG64()
    capsule = x.capsule
    if not PyCapsule_IsValid(capsule, capsule_name):
        raise ValueError("Invalid pointer to anon_func_state")
    rng = <bitgen_t *> PyCapsule_GetPointer(capsule, capsule_name)
    random_values = np.empty(n)
    # Best practice is to acquire the lock whenever generating random values.
    # This prevents other threads from modifying the state. Acquiring the lock
    # is only necessary if if the GIL is also released, as in this example.
    with x.lock, nogil:
        for i in range(n):
            random_values[i] = rng.next_double(rng.state)
    randoms = np.asarray(random_values)
    return randoms.mean()


# This function is declared nogil so it can be used without the GIL below
cdef uint32_t bounded_uint(uint32_t lb, uint32_t ub, bitgen_t *rng) nogil:
    cdef uint32_t mask, delta, val
    mask = delta = ub - lb
    mask |= mask >> 1
    mask |= mask >> 2
    mask |= mask >> 4
    mask |= mask >> 8
    mask |= mask >> 16

    val = rng.next_uint32(rng.state) & mask
    while val > delta:
        val = rng.next_uint32(rng.state) & mask

    return lb + val


@cython.boundscheck(False)
@cython.wraparound(False)
def bounded_uints(uint32_t lb, uint32_t ub, Py_ssize_t n):
    cdef Py_ssize_t i
    cdef bitgen_t *rng
    cdef uint32_t[::1] out
    cdef const char *capsule_name = "BitGenerator"

    x = PCG64()
    out = np.empty(n, dtype=np.uint32)
    capsule = x.capsule

    if not PyCapsule_IsValid(capsule, capsule_name):
        raise ValueError("Invalid pointer to anon_func_state")
    rng = <bitgen_t *>PyCapsule_GetPointer(capsule, capsule_name)

    with x.lock, nogil:
        for i in range(n):
            out[i] = bounded_uint(lb, ub, rng)
    return np.asarray(out)
numpy.fft.fft   fft.fft(a, n=None, axis=- 1, norm=None)[source]
 
Compute the one-dimensional discrete Fourier Transform. This function computes the one-dimensional n-point discrete Fourier Transform (DFT) with the efficient Fast Fourier Transform (FFT) algorithm [CT].  Parameters 
 
aarray_like


Input array, can be complex.  
nint, optional


Length of the transformed axis of the output. If n is smaller than the length of the input, the input is cropped. If it is larger, the input is padded with zeros. If n is not given, the length of the input along the axis specified by axis is used.  
axisint, optional


Axis over which to compute the FFT. If not given, the last axis is used.  
norm{“backward”, “ortho”, “forward”}, optional


 New in version 1.10.0.  Normalization mode (see numpy.fft). Default is “backward”. Indicates which direction of the forward/backward pair of transforms is scaled and with what normalization factor.  New in version 1.20.0: The “backward”, “forward” values were added.     Returns 
 
outcomplex ndarray


The truncated or zero-padded input, transformed along the axis indicated by axis, or the last one if axis is not specified.    Raises 
 IndexError

If axis is not a valid axis of a.      See also  numpy.fft

for definition of the DFT and conventions used.  ifft

The inverse of fft.  fft2

The two-dimensional FFT.  fftn

The n-dimensional FFT.  rfftn

The n-dimensional FFT of real input.  fftfreq

Frequency bins for given FFT parameters.    Notes FFT (Fast Fourier Transform) refers to a way the discrete Fourier Transform (DFT) can be calculated efficiently, by using symmetries in the calculated terms. The symmetry is highest when n is a power of 2, and the transform is therefore most efficient for these sizes. The DFT is defined, with the conventions used in this implementation, in the documentation for the numpy.fft module. References  CT 
Cooley, James W., and John W. Tukey, 1965, “An algorithm for the machine calculation of complex Fourier series,” Math. Comput. 19: 297-301.   Examples >>> np.fft.fft(np.exp(2j * np.pi * np.arange(8) / 8))
array([-2.33486982e-16+1.14423775e-17j,  8.00000000e+00-1.25557246e-15j,
        2.33486982e-16+2.33486982e-16j,  0.00000000e+00+1.22464680e-16j,
       -1.14423775e-17+2.33486982e-16j,  0.00000000e+00+5.20784380e-16j,
        1.14423775e-17+1.14423775e-17j,  0.00000000e+00+1.22464680e-16j])
 In this example, real input has an FFT which is Hermitian, i.e., symmetric in the real part and anti-symmetric in the imaginary part, as described in the numpy.fft documentation: >>> import matplotlib.pyplot as plt
>>> t = np.arange(256)
>>> sp = np.fft.fft(np.sin(t))
>>> freq = np.fft.fftfreq(t.shape[-1])
>>> plt.plot(freq, sp.real, freq, sp.imag)
[<matplotlib.lines.Line2D object at 0x...>, <matplotlib.lines.Line2D object at 0x...>]
>>> plt.show()
numpy.fft.fftfreq   fft.fftfreq(n, d=1.0)[source]
 
Return the Discrete Fourier Transform sample frequencies. The returned float array f contains the frequency bin centers in cycles per unit of the sample spacing (with zero at the start). For instance, if the sample spacing is in seconds, then the frequency unit is cycles/second. Given a window length n and a sample spacing d: f = [0, 1, ...,   n/2-1,     -n/2, ..., -1] / (d*n)   if n is even
f = [0, 1, ..., (n-1)/2, -(n-1)/2, ..., -1] / (d*n)   if n is odd
  Parameters 
 
nint


Window length.  
dscalar, optional


Sample spacing (inverse of the sampling rate). Defaults to 1.    Returns 
 
fndarray


Array of length n containing the sample frequencies.     Examples >>> signal = np.array([-2, 8, 6, 4, 1, 0, 3, 5], dtype=float)
>>> fourier = np.fft.fft(signal)
>>> n = signal.size
>>> timestep = 0.1
>>> freq = np.fft.fftfreq(n, d=timestep)
>>> freq
array([ 0.  ,  1.25,  2.5 , ..., -3.75, -2.5 , -1.25])
numpy.fft.fftshift   fft.fftshift(x, axes=None)[source]
 
Shift the zero-frequency component to the center of the spectrum. This function swaps half-spaces for all axes listed (defaults to all). Note that y[0] is the Nyquist component only if len(x) is even.  Parameters 
 
xarray_like


Input array.  
axesint or shape tuple, optional


Axes over which to shift. Default is None, which shifts all axes.    Returns 
 
yndarray


The shifted array.      See also  ifftshift

The inverse of fftshift.    Examples >>> freqs = np.fft.fftfreq(10, 0.1)
>>> freqs
array([ 0.,  1.,  2., ..., -3., -2., -1.])
>>> np.fft.fftshift(freqs)
array([-5., -4., -3., -2., -1.,  0.,  1.,  2.,  3.,  4.])
 Shift the zero-frequency component only along the second axis: >>> freqs = np.fft.fftfreq(9, d=1./9).reshape(3, 3)
>>> freqs
array([[ 0.,  1.,  2.],
       [ 3.,  4., -4.],
       [-3., -2., -1.]])
>>> np.fft.fftshift(freqs, axes=(1,))
array([[ 2.,  0.,  1.],
       [-4.,  3.,  4.],
       [-1., -3., -2.]])
numpy.fft.hfft   fft.hfft(a, n=None, axis=- 1, norm=None)[source]
 
Compute the FFT of a signal that has Hermitian symmetry, i.e., a real spectrum.  Parameters 
 
aarray_like


The input array.  
nint, optional


Length of the transformed axis of the output. For n output points, n//2 + 1 input points are necessary. If the input is longer than this, it is cropped. If it is shorter than this, it is padded with zeros. If n is not given, it is taken to be 2*(m-1) where m is the length of the input along the axis specified by axis.  
axisint, optional


Axis over which to compute the FFT. If not given, the last axis is used.  
norm{“backward”, “ortho”, “forward”}, optional


 New in version 1.10.0.  Normalization mode (see numpy.fft). Default is “backward”. Indicates which direction of the forward/backward pair of transforms is scaled and with what normalization factor.  New in version 1.20.0: The “backward”, “forward” values were added.     Returns 
 
outndarray


The truncated or zero-padded input, transformed along the axis indicated by axis, or the last one if axis is not specified. The length of the transformed axis is n, or, if n is not given, 2*m - 2 where m is the length of the transformed axis of the input. To get an odd number of output points, n must be specified, for instance as 2*m - 1 in the typical case,    Raises 
 IndexError

If axis is not a valid axis of a.      See also  rfft

Compute the one-dimensional FFT for real input.  ihfft

The inverse of hfft.    Notes hfft/ihfft are a pair analogous to rfft/irfft, but for the opposite case: here the signal has Hermitian symmetry in the time domain and is real in the frequency domain. So here it’s hfft for which you must supply the length of the result if it is to be odd.  even: ihfft(hfft(a, 2*len(a) - 2)) == a, within roundoff error, odd: ihfft(hfft(a, 2*len(a) - 1)) == a, within roundoff error.  The correct interpretation of the hermitian input depends on the length of the original data, as given by n. This is because each input shape could correspond to either an odd or even length signal. By default, hfft assumes an even output length which puts the last entry at the Nyquist frequency; aliasing with its symmetric counterpart. By Hermitian symmetry, the value is thus treated as purely real. To avoid losing information, the shape of the full signal must be given. Examples >>> signal = np.array([1, 2, 3, 4, 3, 2])
>>> np.fft.fft(signal)
array([15.+0.j,  -4.+0.j,   0.+0.j,  -1.-0.j,   0.+0.j,  -4.+0.j]) # may vary
>>> np.fft.hfft(signal[:4]) # Input first half of signal
array([15.,  -4.,   0.,  -1.,   0.,  -4.])
>>> np.fft.hfft(signal, 6)  # Input entire signal and truncate
array([15.,  -4.,   0.,  -1.,   0.,  -4.])
 >>> signal = np.array([[1, 1.j], [-1.j, 2]])
>>> np.conj(signal.T) - signal   # check Hermitian symmetry
array([[ 0.-0.j,  -0.+0.j], # may vary
       [ 0.+0.j,  0.-0.j]])
>>> freq_spectrum = np.fft.hfft(signal)
>>> freq_spectrum
array([[ 1.,  1.],
       [ 2., -2.]])
numpy.fft.ifft   fft.ifft(a, n=None, axis=- 1, norm=None)[source]
 
Compute the one-dimensional inverse discrete Fourier Transform. This function computes the inverse of the one-dimensional n-point discrete Fourier transform computed by fft. In other words, ifft(fft(a)) == a to within numerical accuracy. For a general description of the algorithm and definitions, see numpy.fft. The input should be ordered in the same way as is returned by fft, i.e.,  
a[0] should contain the zero frequency term, 
a[1:n//2] should contain the positive-frequency terms, 
a[n//2 + 1:] should contain the negative-frequency terms, in increasing order starting from the most negative frequency.  For an even number of input points, A[n//2] represents the sum of the values at the positive and negative Nyquist frequencies, as the two are aliased together. See numpy.fft for details.  Parameters 
 
aarray_like


Input array, can be complex.  
nint, optional


Length of the transformed axis of the output. If n is smaller than the length of the input, the input is cropped. If it is larger, the input is padded with zeros. If n is not given, the length of the input along the axis specified by axis is used. See notes about padding issues.  
axisint, optional


Axis over which to compute the inverse DFT. If not given, the last axis is used.  
norm{“backward”, “ortho”, “forward”}, optional


 New in version 1.10.0.  Normalization mode (see numpy.fft). Default is “backward”. Indicates which direction of the forward/backward pair of transforms is scaled and with what normalization factor.  New in version 1.20.0: The “backward”, “forward” values were added.     Returns 
 
outcomplex ndarray


The truncated or zero-padded input, transformed along the axis indicated by axis, or the last one if axis is not specified.    Raises 
 IndexError

If axis is not a valid axis of a.      See also  numpy.fft

An introduction, with definitions and general explanations.  fft

The one-dimensional (forward) FFT, of which ifft is the inverse  ifft2

The two-dimensional inverse FFT.  ifftn

The n-dimensional inverse FFT.    Notes If the input parameter n is larger than the size of the input, the input is padded by appending zeros at the end. Even though this is the common approach, it might lead to surprising results. If a different padding is desired, it must be performed before calling ifft. Examples >>> np.fft.ifft([0, 4, 0, 0])
array([ 1.+0.j,  0.+1.j, -1.+0.j,  0.-1.j]) # may vary
 Create and plot a band-limited signal with random phases: >>> import matplotlib.pyplot as plt
>>> t = np.arange(400)
>>> n = np.zeros((400,), dtype=complex)
>>> n[40:60] = np.exp(1j*np.random.uniform(0, 2*np.pi, (20,)))
>>> s = np.fft.ifft(n)
>>> plt.plot(t, s.real, label='real')
[<matplotlib.lines.Line2D object at ...>]
>>> plt.plot(t, s.imag, '--', label='imaginary')
[<matplotlib.lines.Line2D object at ...>]
>>> plt.legend()
<matplotlib.legend.Legend object at ...>
>>> plt.show()
numpy.fft.ifftshift   fft.ifftshift(x, axes=None)[source]
 
The inverse of fftshift. Although identical for even-length x, the functions differ by one sample for odd-length x.  Parameters 
 
xarray_like


Input array.  
axesint or shape tuple, optional


Axes over which to calculate. Defaults to None, which shifts all axes.    Returns 
 
yndarray


The shifted array.      See also  fftshift

Shift zero-frequency component to the center of the spectrum.    Examples >>> freqs = np.fft.fftfreq(9, d=1./9).reshape(3, 3)
>>> freqs
array([[ 0.,  1.,  2.],
       [ 3.,  4., -4.],
       [-3., -2., -1.]])
>>> np.fft.ifftshift(np.fft.fftshift(freqs))
array([[ 0.,  1.,  2.],
       [ 3.,  4., -4.],
       [-3., -2., -1.]])
numpy.fft.ihfft   fft.ihfft(a, n=None, axis=- 1, norm=None)[source]
 
Compute the inverse FFT of a signal that has Hermitian symmetry.  Parameters 
 
aarray_like


Input array.  
nint, optional


Length of the inverse FFT, the number of points along transformation axis in the input to use. If n is smaller than the length of the input, the input is cropped. If it is larger, the input is padded with zeros. If n is not given, the length of the input along the axis specified by axis is used.  
axisint, optional


Axis over which to compute the inverse FFT. If not given, the last axis is used.  
norm{“backward”, “ortho”, “forward”}, optional


 New in version 1.10.0.  Normalization mode (see numpy.fft). Default is “backward”. Indicates which direction of the forward/backward pair of transforms is scaled and with what normalization factor.  New in version 1.20.0: The “backward”, “forward” values were added.     Returns 
 
outcomplex ndarray


The truncated or zero-padded input, transformed along the axis indicated by axis, or the last one if axis is not specified. The length of the transformed axis is n//2 + 1.      See also  
hfft, irfft

  Notes hfft/ihfft are a pair analogous to rfft/irfft, but for the opposite case: here the signal has Hermitian symmetry in the time domain and is real in the frequency domain. So here it’s hfft for which you must supply the length of the result if it is to be odd:  even: ihfft(hfft(a, 2*len(a) - 2)) == a, within roundoff error, odd: ihfft(hfft(a, 2*len(a) - 1)) == a, within roundoff error.  Examples >>> spectrum = np.array([ 15, -4, 0, -1, 0, -4])
>>> np.fft.ifft(spectrum)
array([1.+0.j,  2.+0.j,  3.+0.j,  4.+0.j,  3.+0.j,  2.+0.j]) # may vary
>>> np.fft.ihfft(spectrum)
array([ 1.-0.j,  2.-0.j,  3.-0.j,  4.-0.j]) # may vary
numpy.fft.irfft2   fft.irfft2(a, s=None, axes=(- 2, - 1), norm=None)[source]
 
Computes the inverse of rfft2.  Parameters 
 
aarray_like


The input array  
ssequence of ints, optional


Shape of the real output to the inverse FFT.  
axessequence of ints, optional


The axes over which to compute the inverse fft. Default is the last two axes.  
norm{“backward”, “ortho”, “forward”}, optional


 New in version 1.10.0.  Normalization mode (see numpy.fft). Default is “backward”. Indicates which direction of the forward/backward pair of transforms is scaled and with what normalization factor.  New in version 1.20.0: The “backward”, “forward” values were added.     Returns 
 
outndarray


The result of the inverse real 2-D FFT.      See also  rfft2

The forward two-dimensional FFT of real input, of which irfft2 is the inverse.  rfft

The one-dimensional FFT for real input.  irfft

The inverse of the one-dimensional FFT of real input.  irfftn

Compute the inverse of the N-dimensional FFT of real input.    Notes This is really irfftn with different defaults. For more details see irfftn. Examples >>> a = np.mgrid[:5, :5][0]
>>> A = np.fft.rfft2(a)
>>> np.fft.irfft2(A, s=a.shape)
array([[0., 0., 0., 0., 0.],
       [1., 1., 1., 1., 1.],
       [2., 2., 2., 2., 2.],
       [3., 3., 3., 3., 3.],
       [4., 4., 4., 4., 4.]])
numpy.fft.rfft   fft.rfft(a, n=None, axis=- 1, norm=None)[source]
 
Compute the one-dimensional discrete Fourier Transform for real input. This function computes the one-dimensional n-point discrete Fourier Transform (DFT) of a real-valued array by means of an efficient algorithm called the Fast Fourier Transform (FFT).  Parameters 
 
aarray_like


Input array  
nint, optional


Number of points along transformation axis in the input to use. If n is smaller than the length of the input, the input is cropped. If it is larger, the input is padded with zeros. If n is not given, the length of the input along the axis specified by axis is used.  
axisint, optional


Axis over which to compute the FFT. If not given, the last axis is used.  
norm{“backward”, “ortho”, “forward”}, optional


 New in version 1.10.0.  Normalization mode (see numpy.fft). Default is “backward”. Indicates which direction of the forward/backward pair of transforms is scaled and with what normalization factor.  New in version 1.20.0: The “backward”, “forward” values were added.     Returns 
 
outcomplex ndarray


The truncated or zero-padded input, transformed along the axis indicated by axis, or the last one if axis is not specified. If n is even, the length of the transformed axis is (n/2)+1. If n is odd, the length is (n+1)/2.    Raises 
 IndexError

If axis is not a valid axis of a.      See also  numpy.fft

For definition of the DFT and conventions used.  irfft

The inverse of rfft.  fft

The one-dimensional FFT of general (complex) input.  fftn

The n-dimensional FFT.  rfftn

The n-dimensional FFT of real input.    Notes When the DFT is computed for purely real input, the output is Hermitian-symmetric, i.e. the negative frequency terms are just the complex conjugates of the corresponding positive-frequency terms, and the negative-frequency terms are therefore redundant. This function does not compute the negative frequency terms, and the length of the transformed axis of the output is therefore n//2 + 1. When A = rfft(a) and fs is the sampling frequency, A[0] contains the zero-frequency term 0*fs, which is real due to Hermitian symmetry. If n is even, A[-1] contains the term representing both positive and negative Nyquist frequency (+fs/2 and -fs/2), and must also be purely real. If n is odd, there is no term at fs/2; A[-1] contains the largest positive frequency (fs/2*(n-1)/n), and is complex in the general case. If the input a contains an imaginary part, it is silently discarded. Examples >>> np.fft.fft([0, 1, 0, 0])
array([ 1.+0.j,  0.-1.j, -1.+0.j,  0.+1.j]) # may vary
>>> np.fft.rfft([0, 1, 0, 0])
array([ 1.+0.j,  0.-1.j, -1.+0.j]) # may vary
 Notice how the final element of the fft output is the complex conjugate of the second element, for real input. For rfft, this symmetry is exploited to compute only the non-negative frequency terms.
numpy.fft.rfft2   fft.rfft2(a, s=None, axes=(- 2, - 1), norm=None)[source]
 
Compute the 2-dimensional FFT of a real array.  Parameters 
 
aarray


Input array, taken to be real.  
ssequence of ints, optional


Shape of the FFT.  
axessequence of ints, optional


Axes over which to compute the FFT.  
norm{“backward”, “ortho”, “forward”}, optional


 New in version 1.10.0.  Normalization mode (see numpy.fft). Default is “backward”. Indicates which direction of the forward/backward pair of transforms is scaled and with what normalization factor.  New in version 1.20.0: The “backward”, “forward” values were added.     Returns 
 
outndarray


The result of the real 2-D FFT.      See also  rfftn

Compute the N-dimensional discrete Fourier Transform for real input.    Notes This is really just rfftn with different default behavior. For more details see rfftn. Examples >>> a = np.mgrid[:5, :5][0]
>>> np.fft.rfft2(a)
array([[ 50.  +0.j        ,   0.  +0.j        ,   0.  +0.j        ],
       [-12.5+17.20477401j,   0.  +0.j        ,   0.  +0.j        ],
       [-12.5 +4.0614962j ,   0.  +0.j        ,   0.  +0.j        ],
       [-12.5 -4.0614962j ,   0.  +0.j        ,   0.  +0.j        ],
       [-12.5-17.20477401j,   0.  +0.j        ,   0.  +0.j        ]])
numpy.fft.rfftfreq   fft.rfftfreq(n, d=1.0)[source]
 
Return the Discrete Fourier Transform sample frequencies (for usage with rfft, irfft). The returned float array f contains the frequency bin centers in cycles per unit of the sample spacing (with zero at the start). For instance, if the sample spacing is in seconds, then the frequency unit is cycles/second. Given a window length n and a sample spacing d: f = [0, 1, ...,     n/2-1,     n/2] / (d*n)   if n is even
f = [0, 1, ..., (n-1)/2-1, (n-1)/2] / (d*n)   if n is odd
 Unlike fftfreq (but like scipy.fftpack.rfftfreq) the Nyquist frequency component is considered to be positive.  Parameters 
 
nint


Window length.  
dscalar, optional


Sample spacing (inverse of the sampling rate). Defaults to 1.    Returns 
 
fndarray


Array of length n//2 + 1 containing the sample frequencies.     Examples >>> signal = np.array([-2, 8, 6, 4, 1, 0, 3, 5, -3, 4], dtype=float)
>>> fourier = np.fft.rfft(signal)
>>> n = signal.size
>>> sample_rate = 100
>>> freq = np.fft.fftfreq(n, d=1./sample_rate)
>>> freq
array([  0.,  10.,  20., ..., -30., -20., -10.])
>>> freq = np.fft.rfftfreq(n, d=1./sample_rate)
>>> freq
array([  0.,  10.,  20.,  30.,  40.,  50.])