Fixed small bug in traj_to_grad and made it possible to specify k w…

…ith multiple axes

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