Add optional duration parameter to make_extended_trapezoid_area

src/pypulseq/make_extended_trapezoid_area.py

@@ -17,23 +17,37 @@ def make_extended_trapezoid_area(
     grad_end: float,
     convert_to_arbitrary: bool = False,
     system: Union[Opts, None] = None,
-) -> Tuple[SimpleNamespace, np.array, np.array]:
-    """Make the shortest possible extended trapezoid for given area and gradient start and end point.
+    duration: Union[float, None] = None,
+    max_grad: Union[float, None] = None,
+    max_slew: Union[float, None] = None,
+) -> Tuple[SimpleNamespace, np.ndarray, np.ndarray]:
+    """
+    Make an extended trapezoid for given area and gradient start and end point.
+    If no duration is given, the shortest possible duration that satisfies the constraints is used.
+    If a duration is given, the function tries to find a solution with the given duration.
 
     Parameters
     ----------
     area : float
-        Area of extended trapezoid.
+        Area of extended trapezoid in 1/m (= Hz/m * s).
     channel : str
         Orientation of extended trapezoidal gradient event. Must be one of 'x', 'y' or 'z'.
     grad_start : float
-        Starting non-zero gradient value.
+        Starting non-zero gradient value in Hz/m.
     grad_end : float
-        Ending non-zero gradient value.
+        Ending non-zero gradient value in Hz/m.
     convert_to_arbitrary : bool, default=False
         Boolean flag to enable converting the extended trapezoid gradient into an arbitrary gradient.
-    system: Opts, optional
+    system : Opts, default=None
         System limits.
+    duration : float, default=None
+        Desired duration of the extended trapezoid in s.
+    max_grad : float, default=None
+        Maximum gradient strength in Hz/m.
+        If None, the maximum gradient strength is set to the system limit.
+    max_slew : float, default=None
+        Maximum slew rate in Hz/m/s.
+        If None, the maximum slew rate is set to the system limit.
 
     Returns
     -------
@@ -46,13 +60,25 @@ def make_extended_trapezoid_area(
 
     Raises
     ------
-        ValueError if no solution was found that satisfies the constraints and the desired area.
+    ValueError
+        If no solution was found that satisfies the constraints and the desired area (and duration, if given).
+        If `duration` is not None and is not a positive number.
     """
     if system is None:
         system = Opts.default
 
-    max_slew = system.max_slew * 0.99
-    max_grad = system.max_grad * 0.99
+    if channel not in ['x', 'y', 'z']:
+        raise ValueError(f'Invalid channel. Must be one of `x`, `y` or `z`. Passed: {channel}')
+
+    if max_grad is None:
+        max_grad = system.max_grad
+
+    if max_slew is None:
+        max_slew = system.max_slew
+
+    if duration is not None and duration <= 0:
+        raise ValueError('Duration must be a positive number.')
+
     raster_time = system.grad_raster_time
 
     def _to_raster(time: float) -> float:
@@ -147,7 +173,7 @@ def _find_solution(duration: int) -> Union[None, Tuple[int, int, int, float]]:
 
         # Filter solutions that satisfy max_grad and max_slew constraints
         valid_indices = (
-            (abs(grad_amp) <= max_grad + 1e-8) & (slew_rate1 <= max_slew + 1e-8) & (slew_rate2 <= max_slew + 1e-8)
+            (abs(grad_amp) <= max_grad + eps) & (slew_rate1 <= max_slew + eps) & (slew_rate2 <= max_slew + eps)
         )
         solutions = np.flatnonzero(valid_indices)
 
@@ -160,49 +186,57 @@ def _find_solution(duration: int) -> Union[None, Tuple[int, int, int, float]]:
         ind = solutions[ind]
         return (int(time_ramp_up[ind]), int(flat_time[ind]), int(time_ramp_down[ind]), float(grad_amp[ind]))
 
-    # Perform a linear search
-    # This is necessary because there can exist a dead space where solutions
-    # do not exist for some durations longer than the optimal duration. The
-    # binary search below fails to find the optimum in those cases.
-    # TODO: Check if range is sufficient, try to calculate the dead space.
-    min_duration = max(round(_calc_ramp_time(grad_end, grad_start) / raster_time), 2)
-
-    # Calculate duration needed to ramp down gradient to zero.
-    # From this point onwards, solutions can always be found by extending
-    # the duration and doing a binary search.
-    max_duration = max(
-        round(_calc_ramp_time(0, grad_start) / raster_time),
-        round(_calc_ramp_time(0, grad_end) / raster_time),
-        min_duration,
-    )
+    if duration is None:  # duration was not given
+        # Perform a linear search
+        # This is necessary because there can exist a dead space where solutions
+        # do not exist for some durations longer than the optimal duration. The
+        # binary search below fails to find the optimum in those cases.
+        # TODO: Check if range is sufficient, try to calculate the dead space.
+        min_duration = max(round(_calc_ramp_time(grad_end, grad_start) / raster_time), 2)
+
+        # Calculate duration needed to ramp down gradient to zero.
+        # From this point onwards, solutions can always be found by extending
+        # the duration and doing a binary search.
+        max_duration = max(
+            round(_calc_ramp_time(0, grad_start) / raster_time),
+            round(_calc_ramp_time(0, grad_end) / raster_time),
+            min_duration,
+        )
+
+        # Linear search
+        solution = None
+        for current_duration in range(min_duration, max_duration + 1):
+            solution = _find_solution(current_duration)
+            if solution is not None:
+                break
+
+        # Perform a binary search for duration > max_duration if no solution was found
+        if solution is None:
+            # First, find the upper limit on duration where a solution exists by
+            # exponentially expanding the duration.
+            while solution is None:
+                max_duration *= 2
+                solution = _find_solution(max_duration)
 
-    # Linear search
-    solution = None
-    for duration in range(min_duration, max_duration + 1):
-        solution = _find_solution(duration)
-        if solution:
-            break
+            def binary_search(fun, lower_limit, upper_limit):
+                if lower_limit == upper_limit - 1:
+                    return fun(upper_limit)
 
-    # Perform a binary search for duration > max_duration if no solution was found
-    if not solution:
-        # First, find the upper limit on duration where a solution exists by
-        # exponentially expanding the duration.
-        while not solution:
-            max_duration *= 2
-            solution = _find_solution(max_duration)
+                test_value = (upper_limit + lower_limit) // 2
 
-        def binary_search(fun, lower_limit, upper_limit):
-            if lower_limit == upper_limit - 1:
-                return fun(upper_limit)
+                if fun(test_value):
+                    return binary_search(fun, lower_limit, test_value)
+                else:
+                    return binary_search(fun, test_value, upper_limit)
 
-            test_value = (upper_limit + lower_limit) // 2
+            solution = binary_search(_find_solution, max_duration // 2, max_duration)
 
-            if fun(test_value):
-                return binary_search(fun, lower_limit, test_value)
-            else:
-                return binary_search(fun, test_value, upper_limit)
+    else:  # duration was given, so calculate solution for this duration
+        duration_raster = max(round(_to_raster(duration) / raster_time), 2)
+        solution = _find_solution(duration_raster)
 
-        solution = binary_search(_find_solution, max_duration // 2, max_duration)
+        if solution is None:
+            raise ValueError(f'Could not find a solution for area={area} and duration={duration}.')
 
     # Get timing and gradient amplitude from solution
     time_ramp_up = solution[0] * raster_time
@@ -226,7 +260,7 @@ def binary_search(fun, lower_limit, upper_limit):
     if trace_enabled():
         grad.trace = trace()
 
-    if not abs(grad.area - area) < 1e-8:
+    if not abs(grad.area - area) < eps:
         raise ValueError(f'Could not find a solution for area={area}.')
 
     return grad, grad.tt, grad.waveform

tests/expected_output/write_epi_se_rs.seq

@@ -20,8 +20,8 @@ TotalDuration 0.21735
 [BLOCKS]
   1 1572   1   0   0   1  0  0
   2 226   2   0   0   2  0  1
-  3 1613   0   0   0   0  0  0
-  4 1146   3   4   5   3  0  0
+  3 1614   0   0   0   0  0  0
+  4 1145   3   4   5   3  0  0
   5  48   0   6   7   0  1  0
   6  48   0   8   9   0  1  0
   7  48   0   6   9   0  1  0
@@ -80,8 +80,8 @@ TotalDuration 0.21735
  60  48   0   8  10   0  1  0
  61 1572   1   0   0   1  0  0
  62 226   4   0   0   2  0  1
- 63 1613   0   0   0   0  0  0
- 64 1146   5   4   5   3  0  0
+ 63 1614   0   0   0   0  0  0
+ 64 1145   5   4   5   3  0  0
  65  48   0   6   7   0  1  0
  66  48   0   8   9   0  1  0
  67  48   0   6   9   0  1  0
@@ -140,8 +140,8 @@ TotalDuration 0.21735
 120  48   0   8  10   0  1  0
 121 1572   1   0   0   1  0  0
 122 226   6   0   0   2  0  1
-123 1613   0   0   0   0  0  0
-124 1146   7   4   5   3  0  0
+123 1614   0   0   0   0  0  0
+124 1145   7   4   5   3  0  0
 125  48   0   6   7   0  1  0
 126  48   0   8   9   0  1  0
 127  48   0   6   9   0  1  0
@@ -206,18 +206,18 @@ TotalDuration 0.21735
 [RF]
 1      129.712 1 2 0 4000 100 0 0 -424.504 0 s
 2      493.727 3 4 0 1000 130 0 0 -2000 0 e
-3      987.454 3 4 0 1000 1740 0 0 -2000 1.5708 r
+3      987.454 3 4 0 1000 1730 0 0 -2000 1.5708 r
 4      493.727 3 4 0 1000 130 0 0 0 0 e
-5      987.454 3 4 0 1000 1740 0 0 0 1.5708 r
+5      987.454 3 4 0 1000 1730 0 0 0 1.5708 r
 6      493.727 3 4 0 1000 130 0 0 2000 0 e
-7      987.454 3 4 0 1000 1740 0 0 2000 1.5708 r
+7      987.454 3 4 0 1000 1730 0 0 2000 1.5708 r
 
 # Format of arbitrary gradients:
 #   time_shape_id of 0 means default timing (stepping with grad_raster starting at 1/2 of grad_raster)
 # id amplitude first last amp_shape_id time_shape_id delay
 # ..      Hz/m  Hz/m Hz/m        ..         ..          us
 [GRADIENTS]
-3  1.34624e+06            0            0 5 6 0
+3  1.35761e+06            0            0 5 6 0
 7      -133333            0      -133333 7 8 450
 9      -133333      -133333      -133333 9 10 0
 10      -133333      -133333            0 11 8 0
@@ -228,8 +228,8 @@ TotalDuration 0.21735
 [TRAP]
  1   1.3624e+06 250 7090 250 8130
  2       666667 130 2000 130   0
- 4      -812789 160    0 160 11140
- 5       484211 130   60 130 3770
+ 4      -812789 160    0 160 11130
+ 5       484211 130   60 130 3760
  6  1.00036e+06 220   40 220   0
  8 -1.00036e+06 220   40 220   0
 
@@ -10284,24 +10284,24 @@ num_samples 2000
 shape_id 5
 num_samples 8
 0
+0.998065105
+0.998065105
+0.491060787
+0.491060787
 1
 1
-0.495207668
-0.495207668
-0.998748157
-0.998748157
 0
 
 shape_id 6
 num_samples 8
 0
 25
-161
-174
-374
-387
-576
-601
+160
+173
+373
+386
+573
+598
 
 shape_id 7
 num_samples 2
@@ -10338,4 +10338,4 @@ num_samples 2
 # It can be reproduced/verified with md5sum if the file trimmed to the position right above [SIGNATURE]
 # The new line character preceding [SIGNATURE] BELONGS to the signature (and needs to be stripped away for recalculating/verification)
 Type md5
-Hash c313148f55ea9f3c049716cc3f36f857
+Hash 8c5f08ea9a8bd3f1e01635118f8d8486