Make jacobian functions accept/returns signal as 3D instead 2D array

CompasToolkit.jl/examples/jacobian.jl

@@ -137,4 +137,4 @@ Jv = CompasToolkit.compute_jacobian(
     v
 )
 
-print_equals_check(Jv_ref, transpose(collect(Jv)))
+print_equals_check(Jv_ref, transpose(reshape(collect(Jv), :, ncoils)))

CompasToolkit.jl/examples/jacobian_hermitian.jl

@@ -135,8 +135,10 @@ echos = generate_echos(N, pssfp_ref)
 ∂echos = generate_delta_echos(N, pssfp_ref)
 
 Random.seed!(1337)
-v = rand(ComplexF32, nsamples(trajectory_ref), ncoils)
-v_ref = map(SVector{ncoils}, eachcol(v)...)
+v = rand(ComplexF32, trajectory_ref.nsamplesperreadout, trajectory_ref.nreadouts, ncoils)
+
+v_ref = reshape(nsamples, nsamples(trajectory_ref), ncoils)
+v_ref = map(SVector{ncoils}, eachcol(v_ref)...)
 
 Jᴴv_ref = compute_Jᴴv(gpu(echos), gpu(∂echos), gpu(parameters_ref), gpu(coil_sensitivities_ref), gpu(trajectory_ref), gpu(v_ref))
 Jᴴv_ref = reduce(hcat, collect(Jᴴv_ref)) # Vector{Svector} -> Matrix

CompasToolkit.jl/src/CompasToolkit.jl

@@ -546,7 +546,7 @@ function compute_jacobian(
     trajectory::Trajectory,
     coils::AbstractMatrix,
     v::AbstractMatrix
-)
+)::CompasArray{ComplexF32, 3}
     context = get_context()
     ncoils = size(coils, 2)
     nreadouts::Int64 = trajectory.nreadouts
@@ -575,7 +575,7 @@ function compute_jacobian(
         pointer(v)::Ptr{Cvoid}
     )::Ptr{Cvoid}
 
-    return CompasArray{ComplexF32, 2}(context, Jv_ptr, (ncoils, nreadouts * samples_per_readout))
+    return CompasArray{ComplexF32, 3}(context, Jv_ptr, (ncoils, nreadouts, samples_per_readout))
 end
 
 function compute_jacobian_hermitian(
@@ -584,8 +584,8 @@ function compute_jacobian_hermitian(
     parameters::TissueParameters,
     trajectory::Trajectory,
     coils::AbstractMatrix,
-    v::AbstractMatrix
-)
+    v::AbstractArray{<:Any,3}
+)::CompasArray{ComplexF32, 2}
     context = get_context()
     ncoils = size(coils, 2)
     nreadouts::Int64 = trajectory.nreadouts
@@ -596,7 +596,7 @@ function compute_jacobian_hermitian(
     𝜕echos_T1 = convert_array(𝜕echos.T1, ComplexF32, nvoxels, nreadouts)
     𝜕echos_T2 = convert_array(𝜕echos.T2, ComplexF32, nvoxels, nreadouts)
     coils = convert_array(coils, Float32, nvoxels, ncoils)
-    v = convert_array(v, ComplexF32, nreadouts * samples_per_readout, ncoils)
+    v = convert_array(v, ComplexF32, samples_per_readout, nreadouts, ncoils)
 
     Jᴴv_ptr = @ccall LIBRARY.compas_compute_jacobian_hermitian(
         pointer(context)::Ptr{Cvoid},

julia-bindings/src/bindings.cpp

@@ -4,10 +4,10 @@
 #include "parameters/tissue.h"
 #include "sequences/pssfp.h"
 #include "simulate/derivative.h"
+#include "simulate/phase_encoding.h"
 #include "simulate/residual.h"
 #include "simulate/sequence.h"
 #include "simulate/signal.h"
-#include "simulate/phase_encoding.h"
 #include "trajectories/cartesian.h"
 #include "trajectories/spiral.h"
 #include "trajectories/trajectory.h"
@@ -239,7 +239,7 @@ extern "C" kmm::ArrayBase* compas_magnetization_to_signal_spiral(
     });
 }
 
-extern "C" compas::Array<cfloat, 2>* compas_compute_jacobian(
+extern "C" compas::Array<cfloat, 3>* compas_compute_jacobian(
     const compas::CudaContext* context,
     int ncoils,
     const compas::Array<cfloat, 2>* echos,
@@ -271,7 +271,7 @@ extern "C" compas::Array<cfloat, 2>* compas_compute_jacobian(
             *coils,
             *vector);
 
-        return new compas::Array<cfloat, 2>(Jv);
+        return new compas::Array<cfloat, 3>(Jv);
     });
 }
 
@@ -288,7 +288,7 @@ extern "C" compas::Array<cfloat, 2>* compas_compute_jacobian_hermitian(
     float delta_t,
     const compas::Array<cfloat>* k_start,
     cfloat delta_k,
-    const compas::Array<cfloat, 2>* vector) {
+    const compas::Array<cfloat, 3>* vector) {
     return catch_exceptions([&] {
         auto trajectory = compas::CartesianTrajectory {
             nreadouts,

src/jacobian/hermitian.cu

@@ -21,7 +21,7 @@ static __device__ void expand_readout_and_accumulate_mhv(
     TissueVoxel p,
     CartesianTrajectoryView trajectory,
     int readout,
-    cuda_view<cfloat, 2> vector) {
+    cuda_view<cfloat, 3> vector) {
     auto ns = trajectory.samples_per_readout;
     auto delta_t = trajectory.delta_t;
     auto delta_k = trajectory.delta_k;
@@ -47,12 +47,11 @@ static __device__ void expand_readout_and_accumulate_mhv(
 
 #pragma unroll 8
     for (int sample = 0; sample < ns; sample++) {
-        index_t t = readout * ns + sample;
-
         // accumulate dot product in mHv
 #pragma unroll
         for (int icoil = 0; icoil < ncoils; icoil++) {
-            auto v = vector[icoil][t];
+            auto v = vector[icoil][readout][sample];
+
             mHv[icoil] = add_mul_conj(mHv[icoil], v, ms);
             dmHv[icoil][0] = add_mul_conj(dmHv[icoil][0], v, dms[0]);
             dmHv[icoil][1] = add_mul_conj(dmHv[icoil][1], v, dms[1]);
@@ -73,7 +72,7 @@ __global__ void jacobian_hermitian_product(
     TissueParametersView parameters,
     CartesianTrajectoryView trajectory,
     cuda_view<float, 2> coil_sensitivities,
-    cuda_view<cfloat, 2> vector) {
+    cuda_view<cfloat, 3> vector) {
     auto voxel = index_t(blockIdx.x * blockDim.x + threadIdx.x);
 
     if (voxel < parameters.nvoxels) {
@@ -130,7 +129,7 @@ Array<cfloat, 2> compute_jacobian_hermitian(
     TissueParameters parameters,
     CartesianTrajectory trajectory,
     Array<float, 2> coil_sensitivities,
-    Array<cfloat, 2> vector) {
+    Array<cfloat, 3> vector) {
     int ns = trajectory.samples_per_readout;
     int nreadouts = trajectory.nreadouts;
     int ncoils = coil_sensitivities.size(0);
@@ -145,7 +144,8 @@ Array<cfloat, 2> compute_jacobian_hermitian(
     COMPAS_ASSERT(coil_sensitivities.size(0) == ncoils);
     COMPAS_ASSERT(coil_sensitivities.size(1) == nvoxels);
     COMPAS_ASSERT(vector.size(0) == ncoils);
-    COMPAS_ASSERT(vector.size(1) == nreadouts * ns);
+    COMPAS_ASSERT(vector.size(1) == nreadouts);
+    COMPAS_ASSERT(vector.size(2) == ns);
 
     // four reconstruction parameters: T1, T2, rho_x, rho_y
     auto JHv = Array<cfloat, 2>(4, nvoxels);

src/jacobian/product.h

@@ -5,7 +5,7 @@
 
 namespace compas {
 
-Array<cfloat, 2> compute_jacobian(
+Array<cfloat, 3> compute_jacobian(
     const CudaContext& ctx,
     Array<cfloat, 2> echos,
     Array<cfloat, 2> delta_echos_T1,
@@ -23,6 +23,6 @@ Array<cfloat, 2> compute_jacobian_hermitian(
     TissueParameters parameters,
     CartesianTrajectory trajectory,
     Array<float, 2> coil_sensitivities,
-    Array<cfloat, 2> vector);
+    Array<cfloat, 3> vector);
 
 }  // namespace compas

src/jacobian/regular.cu

@@ -42,7 +42,7 @@ __global__ void delta_to_sample_exponent(
 
 template<int ncoils, int threads_per_item = 1>
 __launch_bounds__(256, 16) __global__ void jacobian_product(
-    cuda_view_mut<cfloat, 2> Jv,
+    cuda_view_mut<cfloat, 3> Jv,
     cuda_view<cfloat, 2> echos,
     cuda_view<cfloat, 2> delta_echos_T1,
     cuda_view<cfloat, 2> delta_echos_T2,
@@ -109,14 +109,14 @@ __launch_bounds__(256, 16) __global__ void jacobian_product(
             }
 
             if (lane_id == 0) {
-                Jv[icoil][t] = value;
+                Jv[icoil][r][s] = value;
             }
         }
     }
 }
 }  // namespace kernels
 
-Array<cfloat, 2> compute_jacobian(
+Array<cfloat, 3> compute_jacobian(
     const CudaContext& ctx,
     Array<cfloat, 2> echos,
     Array<cfloat, 2> delta_echos_T1,
@@ -142,7 +142,7 @@ Array<cfloat, 2> compute_jacobian(
     COMPAS_ASSERT(vector.size(0) == 4);  // four reconstruction parameters: T1, T2, rho_x, rho_y
     COMPAS_ASSERT(vector.size(1) == nvoxels);
 
-    auto Jv = Array<cfloat, 2>(ncoils, nreadouts * ns);
+    auto Jv = Array<cfloat, 3>(ncoils, nreadouts, ns);
     auto E = Array<cfloat, 2>(ns, nvoxels);
     auto dEdT2 = Array<cfloat, 2>(ns, nvoxels);
 

src/simulate/phase_encoding.cu

@@ -16,15 +16,16 @@ Array<cfloat, 2> phase_encoding(
 
     dim3 block_dim = {32, 4};
     dim3 grid_dim = {div_ceil(uint(nvoxels), block_dim.x), div_ceil(uint(nreadouts), block_dim.y)};
-    ctx.submit_kernel(grid_dim,
-                      block_dim,
-                      kernels::phase_encoding,
-                      write(ph_en_echos),
-                      echos,
-                      parameters,
-                      trajectory);
+    ctx.submit_kernel(
+        grid_dim,
+        block_dim,
+        kernels::phase_encoding,
+        write(ph_en_echos),
+        echos,
+        parameters,
+        trajectory);
 
     return ph_en_echos;
 }
 
-} // compas
+}  // namespace compas

src/simulate/phase_encoding.h

@@ -1,7 +1,7 @@
 #pragma once
 
-#include "core/context.h"
 #include "core/complex_type.h"
+#include "core/context.h"
 #include "parameters/tissue.h"
 #include "trajectories/cartesian.h"
 
@@ -14,7 +14,6 @@ Array<cfloat, 2> phase_encoding(
     const CudaContext& ctx,
     const Array<cfloat, 2>& echos,
     const TissueParameters& parameters,
-    const CartesianTrajectory& trajectory
-    );
+    const CartesianTrajectory& trajectory);
 
-} // compas
+}  // namespace compas

src/simulate/phase_encoding_kernels.cuh

@@ -23,5 +23,5 @@ __global__ void phase_encoding(
     }
 }
 
-} // kernels
-} // compas
+}  // namespace kernels
+}  // namespace compas