From f62b406e25007e9e2a8130edc60e9e416a66de8e Mon Sep 17 00:00:00 2001
From: Frank Zijlstra <22915457+FrankZijlstra@users.noreply.github.com>
Date: Mon, 22 Jan 2024 17:50:31 +0100
Subject: [PATCH] Fixed small bug in `traj_to_grad` and made it possible to
 specify k with multiple axes

---
 pypulseq/traj_to_grad.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/pypulseq/traj_to_grad.py b/pypulseq/traj_to_grad.py
index 04520a7..300c9d3 100644
--- a/pypulseq/traj_to_grad.py
+++ b/pypulseq/traj_to_grad.py
@@ -29,14 +29,14 @@ def traj_to_grad(
         raster_time = Opts.default.grad_raster_time
     
     # Compute finite difference for gradients in Hz/m
-    g = (k[1:] - k[:-1]) / raster_time
+    g = (k[...,1:] - k[...,:-1]) / raster_time
     # Compute the slew rate
-    sr0 = (g[1:] - g[:-1]) / raster_time
+    sr0 = (g[...,1:] - g[...,:-1]) / raster_time
 
     # Gradient is now sampled between k-space points whilst the slew rate is between gradient points
-    sr = np.zeros(len(sr0) + 1)
-    sr[0] = sr0[0]
-    sr[1:-1] = 0.5 * (sr0[-1] + sr0[1:])
-    sr[-1] = sr0[-1]
+    sr = np.zeros(sr0.shape[:-1] + (sr0.shape[-1] + 1,))
+    sr[...,0] = sr0[...,0]
+    sr[...,1:-1] = 0.5 * (sr0[...,:-1] + sr0[...,1:])
+    sr[...,-1] = sr0[...,-1]
 
     return g, sr