Source code for pypulseq.make_extended_trapezoid

from types import SimpleNamespace
from typing import Union

import numpy as np

from pypulseq import eps
from pypulseq.make_arbitrary_grad import make_arbitrary_grad
from pypulseq.opts import Opts
from pypulseq.points_to_waveform import points_to_waveform
from pypulseq.utils.tracing import trace, trace_enabled




def make_extended_trapezoid(
    channel: str,
    amplitudes: Union[np.ndarray, None] = None,
    convert_to_arbitrary: bool = False,
    max_grad: float = 0.0,
    max_slew: float = 0.0,
    skip_check: bool = False,
    system: Union[Opts, None] = None,
    times: Union[np.ndarray, None] = None,
) -> SimpleNamespace:
    """
    Create a gradient by specifying a set of points (amplitudes) at specified time points(times) at a given channel
    with given system limits. Returns an arbitrary gradient object.

    See Also
    --------
    - `pypulseq.Sequence.sequence.Sequence.add_block()`
    - `pypulseq.opts.Opts`
    - `pypulseq.make_trapezoid.make_trapezoid()`

    Parameters
    ----------
    channel : str
        Orientation of extended trapezoidal gradient event. Must be one of 'x', 'y' or 'z'.
    convert_to_arbitrary : bool, default=False
        Boolean flag to indicate if the extended trapezoid gradient has to be converted into an arbitrary gradient.
    amplitudes : numpy.ndarray, default=09
        Values defined at `times` time indices.
    max_grad : float, default=0
        Maximum gradient strength.
    max_slew : float, default=0
        Maximum slew rate.
    system : Opts, default=Opts()
        System limits.
    skip_check : bool, default=False
        Boolean flag to indicate if amplitude check is to be skipped.
    times : numpy.ndarray, default=np.zeros(1)
        Time points at which `amplitudes` defines amplitude values.

    Returns
    -------
    grad : SimpleNamespace
        Extended trapezoid gradient event.

    Raises
    ------
    ValueError
        If invalid `channel` is passed. Must be one of 'x', 'y' or 'z'.
        If all elements in `times` are zero.
        If elements in `times` are not in ascending order or not distinct.
        If all elements in `amplitudes` are zero.
        If first amplitude of a gradient is non-ero and does not connect to a previous block.
    """
    if amplitudes is None:
        amplitudes = np.zeros(1)

    if times is None:
        times = np.zeros(1)

    if system is None:
        system = Opts.default

    if channel not in ['x', 'y', 'z']:
        raise ValueError(f"Invalid channel. Must be one of 'x', 'y' or 'z'. Passed: {channel}")

    times = np.asarray(times)
    amplitudes = np.asarray(amplitudes)

    if len(times) != len(amplitudes):
        raise ValueError('Times and amplitudes must have the same length.')

    if np.all(times == 0):
        raise ValueError('At least one of the given times must be non-zero')

    if np.any(np.diff(times) <= 0):
        raise ValueError('Times must be in ascending order and all times must be distinct')

    if abs(round(times[-1] / system.grad_raster_time) * system.grad_raster_time - times[-1]) > eps:
        raise ValueError('The last time point must be on a gradient raster')

    if skip_check is False and times[0] > 0 and amplitudes[0] != 0:
        raise ValueError('If first amplitude of a gradient is non-zero, it must connect to previous block')

    if max_grad <= 0:
        max_grad = system.max_grad

    if max_slew <= 0:
        max_slew = system.max_slew

    if convert_to_arbitrary:
        # Represent the extended trapezoid on the regularly sampled time grid
        waveform = points_to_waveform(times=times, amplitudes=amplitudes, grad_raster_time=system.grad_raster_time)
        grad = make_arbitrary_grad(
            channel=channel,
            waveform=waveform,
            system=system,
            max_slew=max_slew,
            max_grad=max_grad,
            delay=times[0],
        )
    else:
        #  Keep the original possibly irregular sampling
        if np.any(np.abs(np.round(times / system.grad_raster_time) * system.grad_raster_time - times) > eps):
            raise ValueError(
                'All time points must be on a gradient raster or "convert_to_arbitrary" option must be used.'
            )
        grad = SimpleNamespace()
        grad.type = 'grad'
        grad.channel = channel
        grad.waveform = amplitudes
        grad.delay = round(times[0] / system.grad_raster_time) * system.grad_raster_time
        grad.tt = times - grad.delay
        grad.shape_dur = grad.tt[-1]
        grad.area = 0.5 * ((grad.tt[1:] - grad.tt[:-1]) * (grad.waveform[1:] + grad.waveform[:-1])).sum()

    grad.first = amplitudes[0]
    grad.last = amplitudes[-1]

    slew = np.diff(grad.waveform) / np.diff(grad.tt)

    if max(abs(slew)) > max_slew * (1 + eps):
        raise ValueError(f'Slew rate violation {max(abs(slew)) / max_slew * 100:.2f}%')
    if max(abs(grad.waveform)) > max_grad + eps:
        raise ValueError(f'Gradient amplitude violation {max(abs(grad.waveform)) / max_grad * 100:.2f}%')

    if trace_enabled():
        grad.trace = trace()

    return grad