GPU support

[tclements]

Hey, I saw @ChengXin was implementing GPU support. Looks fast! We might want to consider using CuPy - it's a Numpy compatible library for CUDA. One nice thing about CuPy is we can use the same code on CPU/GPU. Here's the correlate functions with CuPy:

import numpy as np
import scipy
from scipy.fftpack.helper import next_fast_len
import cupy as cp


def correlate(fft1,fft2, maxlag, Nfft=None, method='cross_correlation'):
    """This function takes ndimensional *data* array, computes the cross-correlation in the frequency domain
    and returns the cross-correlation function between [-*maxlag*:*maxlag*].

    :type fft1: :class:`numpy.ndarray' 
    :param fft1: This array contains the fft of each timeseries to be cross-correlated.
    :type maxlag: int
    :param maxlag: This number defines the number of samples (N=2*maxlag + 1) of the CCF that will be returned.

    :rtype: :class:`numpy.ndarray`
    :returns: The cross-correlation function between [-maxlag:maxlag]
    """
    # Speed up FFT by padding to optimal size for FFTPACK

    xp = cp.get_array_module(fft1)

    if fft1.ndim == 1:
        axis = 0
    elif fft1.ndim == 2:
        axis = 1

    if Nfft is None:
        Nfft = next_fast_len(int(fft1.shape[axis]))

    maxlag = np.round(maxlag)

    Nt = fft1.shape[axis]

    corr = xp.multiply(xp.conj(fft1), fft2)
    if method == 'deconv':
        corr /= noise.smooth(xp.abs(fft1),half_win=5) ** 2
    elif method == 'cohrence':
        corr /= noise.smooth(xp.abs(fft1),half_win=5)
        corr /= noise.smooth(xp.abs(fft2),half_win=5)

    corr = xp.real(xp.fft.irfft(corr,n=Nfft,axis=axis)) 
    corr = xp.fft.fftshift(corr,axis)

    tcorr = np.arange(-Nt//2 + 1, Nt//2)
    ind = np.where(np.abs(tcorr) <= maxlag)[0]
    if axis == 1:
        corr = corr[:,ind]
    else:
        corr = corr[ind]

    return corr


def whiten(data, delta, freqmin, freqmax,Nfft=None):
    """This function takes 1-dimensional *data* timeseries array,
    goes to frequency domain using fft, whitens the amplitude of the spectrum
    in frequency domain between *freqmin* and *freqmax*
    and returns the whitened fft.

    :type data: :class:`numpy.ndarray`
    :param data: Contains the 1D time series to whiten
    :type Nfft: int
    :param Nfft: The number of points to compute the FFT
    :type delta: float
    :param delta: The sampling frequency of the `data`
    :type freqmin: float
    :param freqmin: The lower frequency bound
    :type freqmax: float
    :param freqmax: The upper frequency bound

    :rtype: :class:`numpy.ndarray`
    :returns: The FFT of the input trace, whitened between the frequency bounds
    """

    xp = cp.get_array_module(data)

    # Speed up FFT by padding to optimal size for FFTPACK
    if data.ndim == 1:
        axis = 0
    elif data.ndim == 2:
        axis = 1

    if Nfft is None:
        Nfft = next_fast_len(int(data.shape[axis]))

    pad = 100
    Nfft = int(Nfft)
    freqVec = np.fft.rfftfreq(Nfft, d=delta)

    ind = np.where((freqVec >= freqmin) & (freqVec <= freqmax))[0]
    low = ind[0] - pad
    if low <= 0:
        low = 1

    left = ind[0]
    right = ind[-1]
    high = ind[-1] + pad
    if high > Nfft / 2:
        high = int(Nfft // 2)

    FFTRawSign = xp.fft.rfft(data, Nfft,axis=axis)

    # Left tapering:
    if axis == 1:
        FFTRawSign[:,0:low] *= 0
        FFTRawSign[:,low:left] = xp.cos(
            xp.linspace(xp.pi / 2., xp.pi, left - low)) ** 2 * xp.exp(
            1j * xp.angle(FFTRawSign[:,low:left]))
        # Pass band:
        FFTRawSign[:,left:right] = xp.exp(1j * xp.angle(FFTRawSign[:,left:right]))
        # Right tapering:
        FFTRawSign[:,right:high] = xp.cos(
            xp.linspace(0., xp.pi / 2., high - right)) ** 2 * xp.exp(
            1j * xp.angle(FFTRawSign[:,right:high]))
        FFTRawSign[:,high:Nfft + 1] *= 0

    else:
        FFTRawSign[0:low] *= 0
        FFTRawSign[low:left] = xp.cos(
            xp.linspace(xp.pi / 2., xp.pi, left - low)) ** 2 * xp.exp(
            1j * xp.angle(FFTRawSign[low:left]))
        # Pass band:
        FFTRawSign[left:right] = xp.exp(1j * xp.angle(FFTRawSign[left:right]))
        # Right tapering:
        FFTRawSign[right:high] = xp.cos(
            xp.linspace(0., xp.pi / 2., high - right)) ** 2 * xp.exp(
            1j * xp.angle(FFTRawSign[right:high]))
        FFTRawSign[high:Nfft + 1] *= 0

    return FFTRawSign

[chengxinjiang]

Looks neat! Thanks @tclements

Actually I was looking for some packages that put numpy on GPU. I remember you mentioned JAX from google, but it seems that library is still in trail and I was not able to install it on RC. The CuPy helps us to keep the same formate of sub functions for GPU and CPU versions.

I am going to test the performance of CuPy shortly.

[mdenolle]

Awesome. Looking forward to see the compute time and memory usage. I also notice that the main code syntax is the same but that the noise_module is different. You could just use this in the header of the code:
import noise_gpu as noise_module for the GPU and
impose noise_module for the CPU version.

[mdenolle]

For this part, we still need to profile the code to see where things are slow, and use GPU if it speeds up some of the calculations.

[mdenolle]

I do not think that we ended up implementing this into a cuda module. I am reviving this issue as an enhancement feature.

[congcy]

I was developing GPU-version but paused somewhere. I guess I already finished 50%. Let me know if someone feels interested. We can get it done together. -congcong

…

On Tue, Mar 7, 2023 at 12:35 AM Marine Denolle ***@***.***> wrote: Assigned #2 <#2> to @congcy <https://github.com/congcy>. — Reply to this email directly, view it on GitHub <#2 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AFLECJJUFRQJE46A4X3W32LW23CK3ANCNFSM4GSGKSBA> . You are receiving this because you were assigned.Message ID: ***@***.***>