Source code for pypulseq.split_gradient_at

from copy import deepcopy
from types import SimpleNamespace
from typing import Tuple, Union

import numpy as np

from pypulseq.make_extended_trapezoid import make_extended_trapezoid
from pypulseq.opts import Opts
from pypulseq.utils.tracing import trace, trace_enabled




def split_gradient_at(
    grad: SimpleNamespace, time_point: float, system: Union[Opts, None] = None
) -> Union[SimpleNamespace, Tuple[SimpleNamespace, SimpleNamespace]]:
    """
    Splits a trapezoidal gradient into two extended trapezoids defined by the cut line. Returns the two gradient parts
    by cutting the original 'grad' at 'time_point'. For the input type 'trapezoid' the results are returned as extended
    trapezoids, for 'arb' as arbitrary gradient objects. The delays in the individual gradient events are adapted such
    that add_gradients(...) produces a gradient equivalent to 'grad'.

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

    Parameters
    ----------
    grad : SimpleNamespace
        Gradient event to be split into two gradient events.
    time_point : float
        Time point at which `grad` will be split into two gradient waveforms.
    system : Opts, default=Opts()
        System limits.

    Returns
    -------
    grad1, grad2 : SimpleNamespace
        Gradient waveforms after splitting.

    Raises
    ------
    ValueError
        If non-gradient event is passed.
    """
    if system is None:
        system = Opts.default

    # copy() to emulate pass-by-value; otherwise passed grad is modified
    grad = deepcopy(grad)

    grad_raster_time = system.grad_raster_time

    time_index = round(time_point / grad_raster_time)
    # Work around floating-point arithmetic limitation
    time_point = round(time_index * grad_raster_time, 6)
    channel = grad.channel

    if grad.type == 'grad':
        # Check if we have an arbitrary gradient or an extended trapezoid
        if abs(grad.tt[-1] - 0.5 * grad_raster_time) < 1e-10 and np.all(
            abs(grad.tt[1:] - grad.tt[:-1] - grad_raster_time) < 1e-10
        ):
            # Arbitrary gradient -- trivial conversion
            # If time point is out of range we have nothing to do
            if time_index == 0 or time_index >= len(grad.tt):
                return grad
            else:
                grad1 = grad
                grad2 = grad
                grad1.last = 0.5 * (grad.waveform[time_index - 1] + grad.waveform[time_index])
                grad2.first = grad1.last
                grad2.delay = grad.delay + grad.t[time_index]
                grad1.t = grad.t[:time_index]
                grad1.waveform = grad.waveform[:time_index]
                grad2.t = grad.t[time_index:] - time_point
                grad2.waveform = grad.waveform[time_index:]

                if trace_enabled():
                    t = trace()
                    grad1.trace = t
                    grad2.trace = t

                return grad1, grad2
        else:
            # Extended trapezoid
            times = grad.tt
            amplitudes = grad.waveform
    elif grad.type == 'trap':
        grad.delay = round(grad.delay / grad_raster_time) * grad_raster_time
        grad.rise_time = round(grad.rise_time / grad_raster_time) * grad_raster_time
        grad.flat_time = round(grad.flat_time / grad_raster_time) * grad_raster_time
        grad.fall_time = round(grad.fall_time / grad_raster_time) * grad_raster_time

        # Prepare the extended trapezoid structure
        if grad.flat_time == 0:
            times = [0, grad.rise_time, grad.rise_time + grad.fall_time]
            amplitudes = [0, grad.amplitude, 0]
        else:
            times = [
                0,
                grad.rise_time,
                grad.rise_time + grad.flat_time,
                grad.rise_time + grad.flat_time + grad.fall_time,
            ]
            amplitudes = [0, grad.amplitude, grad.amplitude, 0]
    else:
        raise ValueError('Splitting of unsupported event.')

    # If the split line is behind the gradient, there is no second gradient to create
    if time_point >= grad.delay + times[-1]:
        raise ValueError('Splitting of gradient at time point after the end of gradient.')

    # If the split line goes through the delay
    if time_point < grad.delay:
        times = np.concatenate(([0], grad.delay + times))
        amplitudes = [0, amplitudes]
        grad.delay = 0
    else:
        time_point -= grad.delay

    amplitudes = np.array(amplitudes)
    times = np.array(times).round(6)  # Work around floating-point arithmetic limitation

    # Sample at time point
    amp_tp = np.interp(x=time_point, xp=times, fp=amplitudes)
    t_eps = 1e-10
    times1 = np.concatenate((times[np.where(times < time_point - t_eps)], [time_point]))
    amplitudes1 = np.concatenate((amplitudes[np.where(times < time_point - t_eps)], [amp_tp]))
    times2 = np.concatenate(([time_point], times[times > time_point + t_eps])) - time_point
    amplitudes2 = np.concatenate(([amp_tp], amplitudes[times > time_point + t_eps]))

    # Recreate gradients
    grad1 = make_extended_trapezoid(
        channel=channel,
        system=system,
        times=times1,
        amplitudes=amplitudes1,
        skip_check=True,
    )
    grad1.delay = grad.delay
    grad2 = make_extended_trapezoid(
        channel=channel,
        system=system,
        times=times2,
        amplitudes=amplitudes2,
        skip_check=True,
    )
    grad2.delay = time_point

    if trace_enabled():
        t = trace()
        grad1.trace = t
        grad2.trace = t

    return grad1, grad2