Undo changes

Author: chaithyagr

src/mrinufft/trajectories/tools.py

@@ -455,10 +455,12 @@ def get_gradient_times_to_travel(
     n_ramp_up: The timing values for the ramp up phase.
     n_plateau: The timing values for the plateau phase.
     gi: The intermediate gradient values for trapezoidal or triangular waveforms.
+
     See Also
     --------
     get_gradient_amplitudes_to_travel_for_set_time :
-        To directly get the waveforms required. This is most-likely what you want to use.
+        To directly get the waveforms required. This is most-likely what
+        you want to use.
     """
     acq = acq or Acquisition.default
     area_needed = (kspace_end_loc - kspace_start_loc) / acq.gamma / acq.raster_time
@@ -642,16 +644,18 @@ def _residual(gi):
     nb_shots, nb_dimension = kspace_end_loc.shape
     G = np.zeros((nb_shots, nb_raster_points, nb_dimension), dtype=np.float32)
     for i in range(nb_shots):
-        start = n_ramp_down[i, 0]
-        G[i, :start] = np.linspace(
-            start_gradients[i], gi[i], n_ramp_down[i], endpoint=False, axis=-1
-        )
-        if n_plateau[i, d] > 0:
-            G[i, start : start + n_plateau[i, 0]] = gi[i]
-            start += n_plateau[i, 0]
-        G[i, start : start + n_ramp_up[i, 0]] = np.linspace(
-            gi[i], end_gradients[i], n_ramp_up[i, 0], axis=-1, endpoint=False
-        )
+        for d in range(nb_dimension):
+            start = 0
+            G[i, : n_ramp_down[i, d], d] = np.linspace(
+                start_gradients[i, d], gi[i, d], n_ramp_down[i, d], endpoint=False
+            )
+            start += n_ramp_down[i, d]
+            if n_plateau[i, d] > 0:
+                G[i, start : start + n_plateau[i, d], d] = gi[i, d]
+                start += n_plateau[i, d]
+            G[i, start : start + n_ramp_up[i, d], d] = np.linspace(
+                gi[i, d], end_gradients[i, d], n_ramp_up[i, d], endpoint=False
+            )
     return G