From 5b56300d647e8e7e0d9909f9e9e74accbc913097 Mon Sep 17 00:00:00 2001
From: Frank Zijlstra <22915457+FrankZijlstra@users.noreply.github.com>
Date: Mon, 15 Jan 2024 14:48:23 +0100
Subject: [PATCH] Added `evaluate_labels` function to `Sequence`. - Changed
 `get_block` to always set the `label` field in the returned block (None if no
 labels)

---
 pypulseq/Sequence/block.py    |  4 +-
 pypulseq/Sequence/sequence.py | 81 +++++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+), 2 deletions(-)

diff --git a/pypulseq/Sequence/block.py b/pypulseq/Sequence/block.py
index fb93f63..c13f14d 100755
--- a/pypulseq/Sequence/block.py
+++ b/pypulseq/Sequence/block.py
@@ -260,7 +260,7 @@ def get_block(self, block_index: int) -> SimpleNamespace:
         return self.block_cache[block_index]
 
     block = SimpleNamespace()
-    attrs = ["block_duration", "rf", "gx", "gy", "gz", "adc"]
+    attrs = ["block_duration", "rf", "gx", "gy", "gz", "adc", "label"]
     values = [None] * len(attrs)
     for att, val in zip(attrs, values):
         setattr(block, att, val)
@@ -393,7 +393,7 @@ def get_block(self, block_index: int) -> SimpleNamespace:
                 label.label = supported_labels[int(data[1] - 1)]
                 label.value = data[0]
                 # Allow for multiple labels per block
-                if hasattr(block, "label"):
+                if block.label is not None:
                     block.label[len(block.label)] = label
                 else:
                     block.label = {0: label}
diff --git a/pypulseq/Sequence/sequence.py b/pypulseq/Sequence/sequence.py
index 525c2ad..d1d36d9 100755
--- a/pypulseq/Sequence/sequence.py
+++ b/pypulseq/Sequence/sequence.py
@@ -511,6 +511,87 @@ def duration(self) -> Tuple[int, int, np.ndarray]:
 
         return duration, num_blocks, event_count
 
+    def evaluate_labels(
+            self,
+            init: dict = None,
+            evolution: str = 'none'
+            ) -> dict:
+        """
+        Evaluate label values of the entire sequence.
+        
+        When no evolution is given, returns the label values at the end of the
+        sequence. Returns a dictionary with keys named after the labels used
+        in the sequence. Only the keys corresponding to the labels actually
+        used are created.
+        E.g. labels['LIN'] == 4
+        
+        When evolution is given, labels are tracked through the sequence. See
+        below for options for different types of evolutions. The resulting
+        dictionary will contain arrays of the label values.
+        E.g. labels['LIN'] == np.array([0,1,2,3,4])
+        
+        Initial values for the labels can be given with the 'init' parameter.
+        Useful if evaluating labels block-by-block.
+
+        Parameters
+        ----------
+        init : dict, optional
+            Dictionary containing initial label values. The default is None.
+        evolution : str, optional
+            Flag to specify tracking of label evolutions.
+            Must be one of: 'none', 'adc', 'label', 'blocks' (default = 'none')
+            'blocks': Return label values for all blocks.
+            'adc':    Return label values only for blocks containing ADC events.
+            'label':  Return label values only for blocks where labels are
+                      manipulated.
+
+        Returns
+        -------
+        labels : dict
+            Dictionary containing label values.
+            If evolution == 'none', the dictionary values only contains the
+            final label value.
+            Otherwise, the dictionary values are arrays of label evolutions.
+            Only the labels that are used in the sequence are created in the
+            dictionary.
+
+        """
+        labels = init or dict()
+        label_evolution = []
+        
+        # TODO: MATLAB implementation includes block_range parameter. But in
+        #       general we cannot assume linear block ordering. Could include
+        #       time_range like in other sequence functions. Or a blocks
+        #       parameter to specify which blocks to loop over?
+        for block_counter in self.block_events:
+            block = self.get_block(block_counter)
+            
+            if block.label is not None:
+                # Current block has labels
+                for lab in block.label.values():
+                    if lab.type == 'labelinc':
+                        # Increment label
+                        if lab.label not in labels:
+                            labels[lab.label] = 0
+                        
+                        labels[lab.label] += lab.value
+                    else:
+                        # Set label
+                        labels[lab.label] = lab.value
+                
+                if evolution == 'label':
+                    label_evolution.append(dict(labels))
+
+            if evolution == 'blocks' or (evolution == 'adc' and block.adc is not None):
+                label_evolution.append(dict(labels))
+        
+        # Convert evolutions into label dictionary
+        if len(label_evolution) > 0:
+            for lab in labels:
+                labels[lab] = np.array([e[lab] for e in label_evolution])
+        
+        return labels
+
     def flip_grad_axis(self, axis: str) -> None:
         """
         Invert all gradients along the corresponding axis/channel. The function acts on all gradient objects already