Simulation of a RARE sequence with Sycomore, using Bloch and EPG models

rare_bloch_epg.py

import sys

import matplotlib.pyplot
import numpy
import sycomore
from sycomore.units import *

def main():
    species = sycomore.Species(1000*ms, 100*ms, delta_omega=2*Hz)
    
    excitation = (90*deg, 90*deg)
    refocalization = (180*deg, 0*deg)
    
    TR = 500*ms
    TE = 20*ms
    
    train_length = 10
    repetitions = 2
    
    bw_per_pixel = 200*Hz
    voxel_size = 1*mm
    
    (
            signal_bloch_continuous, times_bloch_continuous, events
        ) = bloch_continuous_simulation(
            species, excitation, refocalization, TR, TE, 
            train_length, repetitions, voxel_size, bw_per_pixel)
    
    signal_bloch_discrete, times_bloch_discrete = bloch_discrete_simulation(
        species, excitation, refocalization, TR, TE, 
        train_length, repetitions, voxel_size, bw_per_pixel)
    
    signal_epg, times_epg = epg_simulation(
        species, excitation, refocalization, TR, TE, 
        train_length, repetitions, voxel_size, bw_per_pixel)
    
    figure = matplotlib.pyplot.figure(tight_layout=True)
    
    x_bloch_continuous = [x.convert_to(ms) for x in times_bloch_continuous]
    x_bloch_discrete = [x.convert_to(ms) for x in times_bloch_discrete]
    x_epg = [x.convert_to(ms) for x in times_epg]
    
    (
            magnitude_plot, phase_plot, rf_plot, gradient_plot
        ) = figure.subplots(4, sharex=True)
    magnitude_plot.plot(
        x_bloch_continuous, numpy.abs(signal_bloch_continuous), 
        label="Bloch (continuous)")
    magnitude_plot.plot(
        x_bloch_discrete, numpy.abs(signal_bloch_discrete), "1", 
        label="Bloch (discrete)")
    magnitude_plot.plot(x_epg, numpy.abs(signal_epg), "2", label="EPG")
    magnitude_plot.legend()
    magnitude_plot.set(ylim=0, ylabel="Magnitude")
    
    phase_plot.plot(
        x_bloch_continuous, numpy.degrees(numpy.angle(signal_bloch_continuous)))
    phase_plot.plot(
        x_bloch_discrete, numpy.degrees(numpy.angle(signal_bloch_discrete)), "1")
    phase_plot.plot(x_epg, numpy.degrees(numpy.angle(signal_epg)), "2")
    phase_plot.set(xlim=0, ylim=[-185, 185], ylabel="Phase (°)")
    
    # Gather and sort the pulse events
    rf = [
        (x.convert_to(ms), excitation[0].convert_to(deg)) 
        for x in events["excitation"]]
    rf += [
        (x.convert_to(ms), refocalization[0].convert_to(deg)) 
        for x in events["refocalization"]]
    rf.sort(key=lambda x: x[0])
    rf_plot.vlines([x[0] for x in rf], 0, [x[1] for x in rf])
    rf_plot.set(xlim=0, ylim=[0, 200], ylabel="RF (°)")
    
    # Gather and sort the gradient events
    readout_amplitude = (
            2*numpy.pi*bw_per_pixel/(sycomore.gamma*voxel_size)
        ).convert_to(mT/m)
    gradient = [(x.convert_to(ms), 0) for x in events["idle"]]
    gradient += [
        (x.convert_to(ms), readout_amplitude) 
        for x in events["frequency_encoding"]]
    gradient.sort(key=lambda x: x[0])
    gradient_plot.plot([x[0] for x in gradient], [x[1] for x in gradient], "k")
    gradient_plot.set(
        xlim=0, ylim=0, xlabel="Time (ms)", ylabel="Gradient (mT/m)")
    
    matplotlib.pyplot.show()

def bloch_continuous_simulation(
        species, excitation, refocalization, TR, TE, train_length, repetitions, 
        voxel_size, bw_per_pixel):
    
    """ Small time-step (i.e. almost continuous) Bloch simulation. At each time
        step, the signal is recorded.
    """
    
    # Pulses
    excitation = sycomore.bloch.pulse(*excitation)
    refocalization = sycomore.bloch.pulse(*refocalization)

    # Discretized time steps
    time_step = 0.5*ms
    steps = int((repetitions*TR/time_step))
    times = sycomore.linspace(0*s, repetitions*TR-time_step, steps)
    
    # Isochromats
    positions_count = 1000
    positions = sycomore.linspace(voxel_size, positions_count)

    # Small time-step idle event
    idle = numpy.asarray([
        sycomore.bloch.time_interval(
            species, time_step, species.delta_omega, position=position)
        for position in positions])
    
    # Small time-step frequency encoding event
    readout_amplitude = 2*numpy.pi*bw_per_pixel/(sycomore.gamma*voxel_size)
    frequency_encoding = numpy.asarray([
        sycomore.bloch.time_interval(
            species, time_step, 
            species.delta_omega, readout_amplitude, position)
        for position in positions])
    readout_duration = (1/bw_per_pixel)
    
    signal = [0.]
    times = [0*s]
    
    events = {
        "frequency_encoding": [],
        "idle": [],
        "excitation": [],
        "refocalization": [],
    }
    
    # Float modulo arithmetic is finicky: use integer arithmetic with us-precision
    time = 0*s
    time_step = int(numpy.round(time_step.convert_to(us)))
    TE = int(numpy.round(TE.convert_to(us)))
    TR = int(numpy.round(TR.convert_to(us)))
    readout_duration = int(numpy.round(readout_duration.convert_to(us)))
    
    # Initial magnetization
    magnetization = numpy.full((positions_count, 4, 1), [[0.], [0.], [1.], [1.]])
    
    for time in range(0, repetitions*TR, time_step):
        repetition, time_in_TR = divmod(time, TR)
        echo, time_in_TE = divmod(time_in_TR, TE)
        
        # Complete operator (pulse+gradient) to be applied during the time step
        operator = numpy.identity(4)
        
        event_names = []
        
        # Pulses: excitation at start of the TR, refocalization at the middle
        # of each TE.
        if time_in_TR == 0:
            event_names.append("excitation")
            operator = excitation
        elif time_in_TE == TE/2 and echo < train_length:
            event_names.append("refocalization")
            operator = refocalization
        
        # First echo: only idle until prephasing. Later echoes: second half
        # of the readout, then idle until prephasing.
        if echo == 0 and time_in_TE < readout_duration/2:
            event_names.append("idle")
            operator = idle @ operator
        elif readout_duration/2 <= time_in_TE < TE/2 - readout_duration/2:
            event_names.append("idle")
            operator = idle @ operator
        # First echo: pre-phasing just before refocalization. Later echoes
        # are pre-phased by the second half of the previous read-out.
        elif TE/2 - readout_duration/2 <= time_in_TE < TE/2:
            event_names.append("idle" if echo!=0 else "frequency_encoding")
            operator = (idle if echo!=0 else frequency_encoding) @ operator
        # Idle then first half of the read-out
        elif TE/2 <= time_in_TE < TE-readout_duration/2:
            event_names.append("idle")
            operator = idle @ operator
        elif TE-readout_duration/2 <= time_in_TE:
            event_names.append(
                "frequency_encoding" if echo<train_length else "idle")
            operator = (frequency_encoding if echo<train_length else idle)@operator
        elif 0 < echo <= train_length and time_in_TE < readout_duration/2:
            event_names.append("frequency_encoding")
            operator = frequency_encoding @ operator
        else:
            event_names.append("idle")
            operator = idle @ operator
        
        # Apply the operator
        magnetization = operator @ magnetization
        
        # Collect the signal
        signal.append(numpy.mean(magnetization[:,0]+1j*magnetization[:,1]))
        times.append((time+time_step)*us)
        
        # Collect the event
        for name in event_names:
            events[name].append((time*us))
    
    return signal, times, events

def bloch_discrete_simulation(
        species, excitation, refocalization, TR, TE, train_length, repetitions,
        voxel_size, bw_per_pixel):
    
    """ Large time-step (i.e. discrete) Bloch simulation. The signal is only
        recorded at echoes.
    """
    
    positions_count = 1000

    excitation = sycomore.bloch.pulse(*excitation)
    refocalization = sycomore.bloch.pulse(*refocalization)

    positions = sycomore.linspace(voxel_size, positions_count)
    
    readout_duration = (1/bw_per_pixel)
    readout_amplitude = 2*numpy.pi*bw_per_pixel/(sycomore.gamma*voxel_size)
    half_readout = numpy.asarray([
        sycomore.bloch.time_interval(
            species, readout_duration/2, 
            species.delta_omega, readout_amplitude, position)
        for position in positions])
    readout_duration = (1/bw_per_pixel)

    idle = {}
    idle["half_readout"] = numpy.asarray([
        sycomore.bloch.time_interval(
            species, readout_duration/2, species.delta_omega, position=position)
        for position in positions])
    idle["end_of_TR"] = numpy.asarray([
        sycomore.bloch.time_interval(
            species, TR-train_length*TE-readout_duration/2, species.delta_omega, 
            position=position)
        for position in positions])
    idle["after_excitation"] = numpy.asarray([
        sycomore.bloch.time_interval(
            species, TE/2-readout_duration, species.delta_omega, 
            position=position)
        for position in positions])
    idle["after_refocalization"] = numpy.asarray([
        sycomore.bloch.time_interval(
            species, TE/2-readout_duration/2, species.delta_omega, 
            position=position)
        for position in positions])
    
    signal = []
    times = []
    
    magnetization = numpy.full((positions_count, 4, 1), [[0.], [0.], [1.], [1.]])
    
    for repetition in range(repetitions):
        magnetization = excitation @ magnetization
        for echo_number in range(train_length):
            # First echo: only idle until prephasing. Later echoes: second half
            # of the readout, then idle until prephasing.
            if echo_number == 0:
                magnetization = idle["half_readout"] @ magnetization
            else:
                magnetization = half_readout @ magnetization
            magnetization = idle["after_excitation"] @ magnetization
            
            # First echo: pre-phasing just before refocalization. Later echoes
            # are pre-phased by the second half of the previous read-out.
            if echo_number == 0:
                magnetization = half_readout @ magnetization
            else:
                magnetization = idle["half_readout"] @ magnetization
            magnetization = refocalization @ magnetization
            
            # Idle then first half of the read-out
            magnetization = idle["after_refocalization"] @ magnetization
            magnetization = half_readout @ magnetization
            
            signal.append(numpy.mean(magnetization[:,0]+1j*magnetization[:,1]))
            times.append(repetition*TR+(1+echo_number)*TE)
        
        # Second half of the last read-out
        magnetization = half_readout @ magnetization
        # Idle until the end of the TR
        magnetization = idle["end_of_TR"] @ magnetization
    
    return signal, times

def epg_simulation(
        species, excitation, refocalization, TR, TE, train_length, repetitions,
        voxel_size, bw_per_pixel):
    
    """ EPG simulation. The signal is only recorded at echoes.
    """
    
    readout_area = 2*numpy.pi/(sycomore.gamma*voxel_size)
    half_readout = sycomore.TimeInterval(0.5*1/bw_per_pixel, 0.5*readout_area)
    
    model = sycomore.epg.Regular(
        species, unit_gradient_area=half_readout.gradient_area[0])
    
    signal = []
    times = []
    
    for repetition in range(repetitions):
        model.apply_pulse(*excitation)
        for echo_number in range(train_length):
            # First echo: only idle until prephasing. Later echoes: second half
            # of the readout, then idle until prephasing.
            if echo_number == 0:
                model.apply_time_interval(half_readout.duration)
            else:
                model.apply_time_interval(half_readout)
            model.apply_time_interval(TE/2-2*half_readout.duration)
            
            # First echo: pre-phasing just before refocalization. Later echoes
            # are pre-phased by the second half of the previous read-out.
            if echo_number == 0:
                model.apply_time_interval(half_readout)
            else:
                model.apply_time_interval(half_readout.duration)
            model.apply_pulse(*refocalization)
            
            # Idle then first half of the read-out
            model.apply_time_interval(TE/2-half_readout.duration)
            model.apply_time_interval(half_readout)
            
            signal.append(model.echo)
            times.append(repetition*TR+(1+echo_number)*TE)
        
        # Second half of the last read-out
        model.apply_time_interval(half_readout)
        # Idle until the end of the TR
        model.apply_time_interval(TR-train_length*TE-half_readout.duration)
    
    return signal, times

if __name__ == "__main__":
    sys.exit(main())