use EventLibrary.update()

Author: schuenke

src/pypulseq/Sequence/sequence.py

@@ -908,15 +908,23 @@ def mod_grad_axis(self, axis: str, modifier: float) -> None:
             raise RuntimeError('mod_grad_axis does not yet support the same gradient event used on multiple axes.')
 
         for i in range(len(selected_events)):
-            # Convert tuple to list for modification, then convert back
-            grad_data = list(self.grad_library.data[selected_events[i]])
+            event_id = selected_events[i]
+            old_data = self.grad_library.data[event_id]
+            grad_type = self.grad_library.type[event_id]
+
+            # Convert tuple to list for modification
+            grad_data = list(old_data)
             grad_data[0] *= modifier
-            if self.grad_library.type[selected_events[i]] == 'g' and len(grad_data) == 6:
+
+            if grad_type == 'g' and len(grad_data) == 6:
                 # Need to update first and last fields for arbitrary gradients
                 # Data structure: (amplitude, shape_ID1, shape_ID2, delay, first, last)
                 grad_data[4] *= modifier  # first
                 grad_data[5] *= modifier  # last
-            self.grad_library.data[selected_events[i]] = tuple(grad_data)
+
+            # Use EventLibrary.update() to properly maintain keymap integrity
+            new_data = tuple(grad_data)
+            self.grad_library.update(event_id, old_data, new_data, grad_type)
 
         # Clear block cache to ensure get_block() uses the modified gradient data
         if self.use_block_cache:

tests/test_mod_grad_axis.py

@@ -283,3 +283,40 @@ def test_mod_grad_axis_no_gradients_on_axis():
     assert abs(block.gx.amplitude - 1000) < 1e-12, 'X gradient should be unchanged'
     assert abs(block.gz.amplitude - 1500) < 1e-12, 'Z gradient should be unchanged'
     assert not hasattr(block, 'gy') or block.gy is None, 'Y gradient should not exist'
+
+
+def test_mod_grad_axis_keymap_integrity():
+    """Test that EventLibrary keymap integrity is maintained after modification."""
+    seq = pp.Sequence()
+
+    # Create a gradient and add it to the sequence
+    gx = pp.make_trapezoid('x', amplitude=1000, flat_time=5e-3)
+    seq.add_block(gx)
+
+    # Store original library state
+    original_keymap_size = len(seq.grad_library.keymap)
+    original_data_size = len(seq.grad_library.data)
+
+    # Modify the gradient
+    seq.mod_grad_axis('x', -1.0)
+
+    # Check that library sizes are maintained (no duplicates created)
+    assert len(seq.grad_library.keymap) == original_keymap_size, 'Keymap size should be unchanged'
+    assert len(seq.grad_library.data) == original_data_size, 'Data size should be unchanged'
+
+    # Verify that a modification does not create a duplicate gradient event
+    gx_inverted = pp.make_trapezoid('x', amplitude=-1000, flat_time=5e-3)
+    seq.add_block(gx_inverted)
+
+    # Library should still have the same number of unique gradients
+    assert len(seq.grad_library.data) == original_data_size, 'No duplicate gradient should be created'
+
+    # Verify that both gradients reference the same gradient event
+    block1_events = seq.block_events[1]
+    block2_events = seq.block_events[2]
+
+    # The gradient event ID is at index 2
+    gx_event_id_block1 = block1_events[2]
+    gx_event_id_block2 = block2_events[2]
+
+    assert gx_event_id_block1 == gx_event_id_block2, 'Both gradients should reference the same gradient event'