add function to compute spectrum in-memory

cpp/kspace.cpp

@@ -770,41 +770,10 @@ void kspace<be, dt>::cospectrum(
 	const double wavenumber_exp)
 {
     TIMEZONE("field::cospectrum1");
-    shared_array<double> spec_local_thread(this->nshells*ncomp(fc)*ncomp(fc),[&](double* spec_local){
-        std::fill_n(spec_local, this->nshells*ncomp(fc)*ncomp(fc), 0);
-    });
-
-    this->CLOOP_K2_NXMODES(
-            [&](const ptrdiff_t cindex,
-                const ptrdiff_t xindex,
-                const ptrdiff_t yindex,
-                const ptrdiff_t zindex,
-                const double k2,
-                const int nxmodes){
-            if (k2 <= this->kM2)
-            {
-                double* spec_local = spec_local_thread.getMine();
-                const int tmp_int = int(sqrt(k2) / this->dk)*ncomp(fc)*ncomp(fc);
-
-                for (hsize_t i=0; i<ncomp(fc); i++)
-                for (hsize_t j=0; j<ncomp(fc); j++){
-                    spec_local[tmp_int + i*ncomp(fc)+j] += pow(k2, wavenumber_exp/2.)*
-			nxmodes * (
-                        (a[ncomp(fc)*cindex + i][0] * a[ncomp(fc)*cindex + j][0]) +
-                        (a[ncomp(fc)*cindex + i][1] * a[ncomp(fc)*cindex + j][1]));
-                }
-            }
-            });
-
-    spec_local_thread.mergeParallel();
 
     std::vector<double> spec;
-    spec.resize(this->nshells*ncomp(fc)*ncomp(fc), 0);
-    MPI_Allreduce(
-            spec_local_thread.getMasterData(),
-            &spec.front(),
-            spec.size(),
-            MPI_DOUBLE, MPI_SUM, this->layout->comm);
+    this->template cospectrum<rnumber, fc>(a, spec, wavenumber_exp);
+
     if (this->layout->myrank == 0)
     {
         hid_t dset, wspace, mspace;
@@ -839,6 +808,52 @@ void kspace<be, dt>::cospectrum(
     }
 }
 
+template <field_backend be,
+          kspace_dealias_type dt>
+template <typename rnumber,
+          field_components fc>
+void kspace<be, dt>::cospectrum(
+        const rnumber(* __restrict a)[2],
+        std::vector<double> &spec,
+	    const double wavenumber_exp)
+{
+    TIMEZONE("field::cospectrum1_in_memory");
+    shared_array<double> spec_local_thread(this->nshells*ncomp(fc)*ncomp(fc),[&](double* spec_local){
+        std::fill_n(spec_local, this->nshells*ncomp(fc)*ncomp(fc), 0);
+    });
+
+    this->CLOOP_K2_NXMODES(
+            [&](const ptrdiff_t cindex,
+                const ptrdiff_t xindex,
+                const ptrdiff_t yindex,
+                const ptrdiff_t zindex,
+                const double k2,
+                const int nxmodes){
+            if (k2 <= this->kM2)
+            {
+                double* spec_local = spec_local_thread.getMine();
+                const int tmp_int = int(sqrt(k2) / this->dk)*ncomp(fc)*ncomp(fc);
+
+                for (hsize_t i=0; i<ncomp(fc); i++)
+                for (hsize_t j=0; j<ncomp(fc); j++){
+                    spec_local[tmp_int + i*ncomp(fc)+j] += pow(k2, wavenumber_exp/2.)*
+			nxmodes * (
+                        (a[ncomp(fc)*cindex + i][0] * a[ncomp(fc)*cindex + j][0]) +
+                        (a[ncomp(fc)*cindex + i][1] * a[ncomp(fc)*cindex + j][1]));
+                }
+            }
+            });
+
+    spec_local_thread.mergeParallel();
+
+    spec.resize(this->nshells*ncomp(fc)*ncomp(fc), 0);
+    MPI_Allreduce(
+            spec_local_thread.getMasterData(),
+            &spec.front(),
+            spec.size(),
+            MPI_DOUBLE, MPI_SUM, this->layout->comm);
+}
+
 template <field_backend be,
           kspace_dealias_type dt>
 template <typename rnumber,

cpp/kspace.hpp

@@ -140,6 +140,13 @@ class kspace
                 const hsize_t toffset,
 		        const double wavenumber_exp = 0);
 
+        template <typename rnumber,
+                  field_components fc>
+        void cospectrum(
+                const rnumber(* __restrict__ a)[2],
+                std::vector<double> &spec,
+		        const double wavenumber_exp = 0);
+
         template <typename rnumber,
                   field_components fc>
         double L2norm(