Added support for set_pts

andyschwarzl · Jun 21, 2024 · 6b66ee0 · 6b66ee0
1 parent 69264c9
commit 6b66ee0
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 16 deletions.
diff --git a/CUDA/inc/gpuNUFFT_operator_factory.hpp b/CUDA/inc/gpuNUFFT_operator_factory.hpp
@@ -174,6 +174,15 @@ class GpuNUFFTOperatorFactory
   void setUseTextures(bool useTextures);
 
   void setBalanceWorkload(bool balanceWorkload);
+
+  /**
+  * \brief Set k-space locations and corresponding density. This can also be used 
+  * to update them
+  * 
+  */
+  void set_pts(
+    gpuNUFFT::GpuNUFFTOperator *gpuNUFFTOp, gpuNUFFT::Array<DType> &kSpaceTraj,
+    gpuNUFFT::Array<DType> &densCompData);
 
  protected:
   template<typename T>
@@ -315,7 +324,9 @@ class GpuNUFFTOperatorFactory
   */
   gpuNUFFT::Array<DType> computeDeapodizationFunction(const IndType &kernelWidth,
     const DType &osf, gpuNUFFT::Dimensions &imgDims);
+
 
+
  private:
   /** \brief Flag to indicate texture interpolation */
   bool useTextures;

diff --git a/CUDA/src/gpu/python/gpuNUFFT_operator_python_factory.cpp b/CUDA/src/gpu/python/gpuNUFFT_operator_python_factory.cpp
@@ -304,6 +304,21 @@ class GpuNUFFTPythonOperator
         gpuNUFFTOp->setSens(sensArray);
     }
 
+    void set_pts(py::array_t<DType> kspace_loc, py::array_t<DType> density_comp)
+    {
+        gpuNUFFT::Array<DType> kSpaceTraj = readNumpyArray(kspace_loc);
+        kSpaceTraj.dim.length = trajectory_length;
+
+        // density compensation weights
+        gpuNUFFT::Array<DType> density_compArray;
+        if(density_comp != Py_None)
+        {
+            density_compArray = readNumpyArray(density_comp);
+            density_compArray.dim.length = trajectory_length;
+        }
+        factory.set_pts(gpuNUFFTOp, kSpaceTraj, density_compArray);
+
+    }
     py::array_t<DType> estimate_density_comp(int max_iter = 10)
     {
         IndType n_samples = kspace_data.count();
@@ -443,6 +458,7 @@ PYBIND11_MODULE(gpuNUFFT, m) {
         .def("estimate_density_comp", &GpuNUFFTPythonOperator::estimate_density_comp, py::arg("max_iter") = 10)
         .def("set_smaps", &GpuNUFFTPythonOperator::set_smaps)
         .def("toggle_grad_mode", &GpuNUFFTPythonOperator::toggle_grad_mode)
-        .def("get_spectral_radius", &GpuNUFFTPythonOperator::get_spectral_radius, py::arg("max_iter") = 20, py::arg("tolerance") = 1e-6);
+        .def("get_spectral_radius", &GpuNUFFTPythonOperator::get_spectral_radius, py::arg("max_iter") = 20, py::arg("tolerance") = 1e-6)
+        .def("set_pts", &GpuNUFFTPythonOperator::set_pts, py::arg("kspace_loc"), py::arg("density_comp") = py::none());
 }
 #endif  // GPUNUFFT_OPERATOR_PYTHONFACTORY_H_INCLUDED
diff --git a/CUDA/src/gpuNUFFT_operator_factory.cpp b/CUDA/src/gpuNUFFT_operator_factory.cpp
@@ -457,10 +457,6 @@ gpuNUFFT::GpuNUFFTOperatorFactory::createGpuNUFFTOperator(
   checkMemoryConsumption(kSpaceTraj.dim, sectorWidth, osf, imgDims,
                          densCompData.dim, sensData.dim);
 
-  if (kSpaceTraj.dim.channels > 1)
-    throw std::invalid_argument(
-        "Trajectory dimension must not contain a channel size greater than 1!");
-
   if (imgDims.channels > 1)
     throw std::invalid_argument(
         "Image dimensions must not contain a channel size greater than 1!");
@@ -470,6 +466,29 @@ gpuNUFFT::GpuNUFFTOperatorFactory::createGpuNUFFTOperator(
   gpuNUFFT::GpuNUFFTOperator *gpuNUFFTOp =
       createNewGpuNUFFTOperator(kernelWidth, sectorWidth, osf, imgDims);
 
+  // Set points and density compensation
+  set_pts(gpuNUFFTOp, kSpaceTraj, densCompData);
+
+  if (sensData.data != NULL)
+    gpuNUFFTOp->setSens(sensData);
+
+  gpuNUFFTOp->setDeapodizationFunction(
+    this->computeDeapodizationFunction(kernelWidth, osf, imgDims));
+
+  debug("finished creation of gpuNUFFT operator\n");
+
+  return gpuNUFFTOp;
+}
+
+
+void gpuNUFFT::GpuNUFFTOperatorFactory::set_pts(
+    gpuNUFFT::GpuNUFFTOperator *gpuNUFFTOp, gpuNUFFT::Array<DType> &kSpaceTraj,
+    gpuNUFFT::Array<DType> &densCompData)
+{
+  if (kSpaceTraj.dim.channels > 1)
+    throw std::invalid_argument(
+        "Trajectory dimension must not contain a channel size greater than 1!");
+
   // assign according sector to k-Space position
   gpuNUFFT::Array<IndType> assignedSectors =
       assignSectors(gpuNUFFTOp, kSpaceTraj);
@@ -487,9 +506,6 @@ gpuNUFFT::GpuNUFFTOperatorFactory::createGpuNUFFTOperator(
   if (densCompData.data != NULL)
     densData = initDensData(gpuNUFFTOp, coordCnt);
 
-  if (sensData.data != NULL)
-    gpuNUFFTOp->setSens(sensData);
-
   if (useGpu)
   {
     sortArrays(gpuNUFFTOp, assignedSectorsAndIndicesSorted,
@@ -543,13 +559,6 @@ gpuNUFFT::GpuNUFFTOperatorFactory::createGpuNUFFTOperator(
   // free temporary array
   free(assignedSectors.data);
   assignedSectors.data = NULL;
-
-  gpuNUFFTOp->setDeapodizationFunction(
-    this->computeDeapodizationFunction(kernelWidth, osf, imgDims));
-
-  debug("finished creation of gpuNUFFT operator\n");
-
-  return gpuNUFFTOp;
 }
 
 gpuNUFFT::GpuNUFFTOperator *

diff --git a/setup.py b/setup.py
@@ -103,7 +103,7 @@ def build_extension(self, ext):
 
 setup(
     name="gpuNUFFT",
-    version="0.8.0",
+    version="0.8.1",
     description="gpuNUFFT - An open source GPU Library for 3D Gridding and NUFFT",
     ext_modules=[
         CMakeExtension("gpuNUFFT", sourcedir=os.path.join("CUDA")),