Source code for cmrsim.bloch._multi_coil

"""Contains implementation of Bloch-solving operator including multiple receive coils"""

__all__ = ["ParallelReceiveBlochOperator"]

from typing import List, Tuple, Union, TYPE_CHECKING

import tensorflow as tf
import numpy as np

from cmrsim.bloch._generic import GeneralBlochOperator
from cmrsim.bloch.submodules import BaseSubmodule
from cmrsim.trajectory import StaticParticlesTrajectory

if TYPE_CHECKING:
    import cmrsim.analytic.contrast.coil_sensitivities


[docs] class ParallelReceiveBlochOperator(GeneralBlochOperator): """Extends the GeneralBlochOperator with parallel receive functionality. The only additionally required input is an instance of a CoilSensitivity contrast module from the cmrsim.analytic.contrast submodule, which is used internally to apply a spatially dependent weight to the transverse magnetization before summing, in the process of acquiring signals. The resulting signal is also stored in the 'time_signal_acc' attribute, where for the parallel receive case each entry has an additional axis corresponding to the number of coils. :param name: Verbose name of the module (non-functionally relevant) :param gamma: gyromagnetic ratio in rad/ms/m :param coil_module: :param time_grid: (n_steps) - tf.float32 - time definition for numerical integration steps :param gradient_waveforms: (#repetitions, #steps, 3) - tf.float32 - Definition of gradients in mT/m :param rf_waveforms: (#repetitions, #steps) - tf.complex64 - :param adc_events: (#repetition, #steps, 2) - tf.float32 - last axis contains the adc-on definitions and the adc-phase :param submodules: List[tf.Modules] implementing the __call__ function returning the resulting phase increment contribution :param device: str """ #: Instance of the cmrsim.analytic.contrast.CoilSensitivity class that is used to compute #: the positionally dependent sensitivity factors prior to sampling coil_lookup_module: 'cmrsim.analytic.contrast.coil_sensitivities.CoilSensitivity' #: n_coils: int def __init__(self, name: str, gamma: float, coil_module: Union['cmrsim.analytic.contrast.coil_sensitivities.CoilSensitivity'], time_grid: tf.Tensor, adc_events: tf.Tensor, gradient_waveforms: tf.Tensor = None, rf_waveforms: tf.Tensor = None, submodules: Tuple[BaseSubmodule, ...] = (), device: str = None): """ """ super().__init__(name=name, gamma=gamma, time_grid=time_grid, gradient_waveforms=gradient_waveforms, rf_waveforms=rf_waveforms, adc_events=adc_events, submodules=submodules, device=device) self.coil_lookup_module = coil_module self.n_coils = self.coil_lookup_module.expansion_factor # Expand the dimension of the signal accumulator to incorporate coil-maps with tf.device("CPU"): self.time_signal_acc = [ tf.Variable(initial_value=tf.zeros([v.read_value().shape[0], self.n_coils], dtype=tf.complex64), shape=(None, self.n_coils), dtype=tf.complex64, trainable=False, name=f"k_segment_{idx}") for idx, v in enumerate(self.time_signal_acc)] def _init_tensors(self, n_samples: int, n_repetitions: int, repetition_index: int): """ Initializes / slices the waveforms according to the specified current repetitions and number of samples. :param n_samples: :param n_repetitions: :param repetition_index: :return: - rf_waveform (#reps, #steps) or None - grad_waveform (#reps, #steps, 3) or None - adc_events (#reps, #steps) or None - adc_phase (#reps, #steps) or None - adc_writing_idx (#reps, #steps) or None - batch_signal (#reps, #samples, #coils) -> used for signal accumulation """ (rf_waveforms, gradient_waveforms, adc_events, adc_phase, adc_writing_indices, batch_signal) = super()._init_tensors(n_samples, n_repetitions, repetition_index) batch_signal = tf.zeros(shape=(n_samples, n_repetitions, self.n_coils), dtype=tf.complex64) return (rf_waveforms, gradient_waveforms, adc_events, adc_phase, adc_writing_indices, batch_signal) def _sample(self, adc_events: tf.Tensor, adc_phase: tf.Tensor, adc_writing_idx: tf.Tensor, magnetization: tf.Tensor, M0: tf.Tensor, r: tf.Tensor, batch_signal: tf.Tensor): """Sums up the :math:`m_{xy}^{+}` magnetization weighted with the associated proton-density of the particles and adds the signal to the batch-signal accumulator ad index specified in adc_writing_index. Before summation, a coil-sensitivity weighting is applied :param adc_events: (#reps, ) :param adc_phase: (#reps, ) :param adc_writing_idx: (#reps, ) :param magnetization: (#batch, ) :param M0: (n_repetitions/1, #batch) :param r: (n_repetitions/1, #batch, 3) particle positions :param batch_signal: (#samples, #reps, #coils) :return: """ n_active_adc_channels = tf.reduce_sum(adc_events, axis=0) if n_active_adc_channels >= 1.: # Sum over all M+ = Mx+iMy weighted by M0 and the coil-sensitivity to compute signal sense_factors = self.coil_lookup_module.lookup_complex_coil_factors(r) weighted_magnetization = magnetization[..., 0, tf.newaxis] * sense_factors weighted_magnetization= weighted_magnetization * tf.complex(M0[..., tf.newaxis], 0.) signal = tf.reduce_sum(weighted_magnetization, axis=1) # Demodulate signal with receiver phase demodulated_signal = signal * tf.exp(tf.complex(0., -adc_phase[..., tf.newaxis])) # Gather and accumulate signal only for repetitions that have an active ADC-event n_reps = tf.shape(adc_writing_idx)[0] in_bounds = tf.where(adc_writing_idx > -1)[:, 0] # batch signal has shape (#samples, #reps, #coils) which defines the stacking-order here indices = tf.stack([adc_writing_idx, tf.range(n_reps, dtype=tf.int32)], axis=1) indices = tf.gather(indices, in_bounds) demodulated_signal = tf.gather(demodulated_signal, in_bounds) batch_signal = tf.tensor_scatter_nd_add(batch_signal, indices, demodulated_signal) return batch_signal