nmreval/nmreval/data/nmr.py

import warnings
from typing import Union

import numpy as np

from .signals import Signal


class FID(Signal):
    """
    Extensions if :class:`Signal` for NMR timesignals
    """
    def __init__(self, x, y, **kwargs):
        super().__init__(x, y, **kwargs)

        self.meta['apod'] = []
        self.meta['zf'] = 0

    def __repr__(self):
        return f"FID: {self.name}"

    def baseline(self, percentage: float = 0.12):
        """
        Substract

        Args:
            percentage:

        Returns:

        """
        self._y -= self._y[int(-percentage * self.length):].mean()

        return self

    def apod(self, p, method=None):
        try:
            self._y *= method.apod(self._x, *p)
            self.meta['apod'].append((str(method), p))
        except KeyError:
            print(f'Unknown apodization {method}')

    def zerofill(self, pad: int = 1):
        _max_x = self._x.max()

        factor = 2**pad
        if factor < 1:
            self._x = self._x[:int(self.length*factor)]
            self._y = self._y[:int(self.length*factor)]
            self._y_err = self._y_err[:int(self.length*factor)]
            self.mask = self.mask[:int(self.length*factor)]

        elif factor > 1:
            add_length = int((factor-1) * self.length)
            self._y = np.concatenate((self._y, np.zeros(add_length)))
            self._y_err = np.concatenate((self._y_err, np.zeros(add_length)))
            self.mask = np.concatenate((self.mask, np.ones(add_length, dtype=bool)))

            _temp_x = np.arange(1, add_length+1) * self.dx + np.max(self._x)
            self._x = np.concatenate((self._x, _temp_x))

        self.meta['zf'] += pad

        return self

    def fourier(self) -> 'Spectrum':
        ft = np.fft.fftshift(np.fft.fft(self._y)) / self.dx
        freq = np.fft.fftshift(np.fft.fftfreq(self.length, self.dx))

        spec = Spectrum(freq, ft, dx=self.dx)
        spec.update(self.meta)

        return spec

    def fft_depake(self, scale: bool = False):
        """
        Calculate deconvoluted Pake spectra

        Args:
            scale (bool):

        Reference:
            M.A. McCabe, S.R. Wassail:
            Rapid deconvolution of NMR powder spectra by weighted fast Fourier transformation,
            SSNMR (1997), Vol.10, Nr.1-2, pp. 53-61

        """
        _y = self._y + 0
        # Perform FFT depaking
        _y *= np.sqrt(self._x)

        # built halves
        right = np.fft.fftshift(np.fft.fft(_y)) * (1 + 1j)
        left = np.fft.fftshift(np.fft.fft(_y)) * (1 - 1j)

        freq = np.fft.fftshift(np.fft.fftfreq(self.length, self.dx))

        if scale:
            # fourier transform should be multiplied with \sqrt(|omega|)
            # but lines at omega=0 are suppressed, so sometimes not a good idea
            right *= np.sqrt(np.abs(freq))
            left *= np.sqrt(np.abs(freq))

        # built depaked spectrum
        new_spc = np.concatenate([left[:(self.length//2)], right[(self.length//2):]])

        spec = Spectrum(2*freq, new_spc, dx=self.dx)
        spec.update(self.meta)

        return spec

    def convert(self):
        spec = Spectrum(self._x, self._y, dx=self.dx)
        self.update(self.meta)

        return spec

    def _calc_noise(self):
        # noise per channel, so divide by sqrt(2)
        self._y_err = np.std(self._y[-int(0.1*self.length):])/1.41
        return self._y_err * np.ones(self.length)

    def shift(self, points: int, mode: str = 'pts') -> None:
        """
        Shift y values to new x indexes

        Args:
            points : shift value,
            mode (str):

        Returns:

        """
        if mode == 'pts':
            super().shift(int(points))
        else:
            pts = int(points//self.dx)
            super().shift(pts)


class Spectrum(Signal):
    """
    Extension of :class:`Signal` for NMR spectra
    """

    def __init__(self, x, y, dx=None, **kwargs):
        super().__init__(x, y, **kwargs)
        if dx is None:
            if 2 * abs(self._x[0])/len(self._x) - abs(self._x[1]-self._x[0]) < 1e-9:
                self.dx = 1/(2 * abs(self._x[0]))
            else:
                self.dx = 1
        else:
            self.dx = dx

    def __repr__(self):
        return f"Spectrum: {self.name}"

    def baseline(self, percentage: float = 0.12):
        corr = np.mean([self._y[int(-percentage*self.length):], self._y[:int(percentage*self.length)]])
        self._y -= corr

        return self

    def fourier(self) -> FID:
        ft = np.fft.ifft(np.fft.ifftshift(self._y * self.dx))
        t = np.arange(len(ft))*self.dx
        shifted = np.fft.ifftshift(self._x)[0]
        if shifted != 0:
            warnings.warn(f'Shift of {shifted}: Spectrum is not zero-centered and FFT may give wrong results!')
            ft *= np.exp(1j * shifted * 2*np.pi*t)

        fid = FID(t, ft)
        fid.update(self.meta)

        return fid

    def convert(self):
        fid = FID(self._x, self._y)
        fid.update(self.meta)

        return fid

    def _calc_noise(self):
        step = int(0.05 * self.length)
        start = np.argmin(abs(self._x - max(self._x.min(), -2e5)))
        stop = np.argmin(abs(self._x - min(self._x.max(), 2e5)))
        if start > stop:
            stop, start = start, stop
        stop -= step
        self._y_err = np.inf
        for i in range(start, stop):
            # scan regions to find minimum noise (undisturbed by signals)
            self._y_err = min(np.std(self._y[i:i + step])/np.sqrt(2), self._y_err)

        return self._y_err * np.ones(self.length)

    def shift(self, points: int, mode: str = 'pts'):
        if mode == 'pts':
            super().shift(int(points))
        else:
            return self