Merge branch 'master' into 'master'

Add grid support to QSP (plus bug fixes)

See merge request SPINS/SPINS_main!14

matlab/spins_QSP_csv.m

-function [T1_max, T1_min, S1_max, S1_min, pdfCount] = spins_QSP_reader(filename)
-  % Reads in the csv file and interprets it in the correct way for the user.
-  %
-  % Input: Filename (string) name of csv file.
-  %
-  % Output: Returns the max/min values for T1 and S1 (as specified in
-  % spins.conf) and the (normalized) pdfCount histogram.
+function [T1_max, T1_min, S1_max, S1_min, pdfCount] = spins_QSP_csv(filename)
+  %% Reads in the csv file and interprets it in the correct way for the user.
+  %%
+  %% Input:
+  %%    Filename (string): name of csv file.
+  %%
+  %% Output:
+  %%    T1_max (float): upper limit of the maximum-valued bin of T1 tracer
+  %%    T1_min (float): lower limit of the minimum-valued bin of T1 tracer
+  %%    S1_max (float): upper limit of the maximum-valued bin of S1 tracer
+  %%    S1_min (float): lower limit of the minimum-valued bin of S1 tracer
   A = importdata(filename);
   T1_max = A(1, 1);
   T1_min = A(1, 2);
   S1_max = A(1, 3);
   S1_min = A(1, 4);
   pdfCount = A(2:end, :);
-  pdfCount /= sum(sum(pdfCount));
+  pdfCount = pdfCount / sum(sum(pdfCount));
 end

src/Science.hpp

@@ -258,7 +258,9 @@ void QSPCount(const TArrayn::DTArray &t, const TArrayn::DTArray &u,
               const char T1_name, const char S1_name, const int NS,
               const int NT, double T1_max, double S1_max, double T1_min,
               double S1_min, const int Nx, const int Ny, const int Nz,
-              string filename, const int plotnum);
+              string filename, const int plotnum, bool mapped,
+              TArrayn::DTArray *xgrid, TArrayn::DTArray *ygrid,
+              TArrayn::DTArray *zgrid);
 
 // Equation of state for seawater, polynomial fit from
 // Brydon, Sun, Bleck (1999) (JGR)

src/Science/QSPCount.cpp

@@ -15,6 +15,49 @@
 using namespace Transformer;
 using namespace blitz;
 
+class QSPVector {
+private:
+  double *data;
+  int length;
+  int rows;
+  int cols;
+
+public:
+  explicit QSPVector(int length)
+      : data((double *)calloc(length, sizeof(double))), length(length), rows(0),
+        cols(0) {
+    if (!data) {
+      std::cout << "Error! Data could not be initialized.\n";
+    }
+  }
+  QSPVector(int Nx, int Ny)
+      : data((double *)calloc(Nx * Ny, sizeof(double))), length(Nx * Ny),
+        rows(Nx), cols(Ny) {
+    if (!data) {
+      std::cout << "Error! Data could not be initialized.\n";
+    }
+  }
+  double *raw() const { return data; }
+  int Length() const { return length; }
+  int Rows() const { return rows; }
+  int Cols() const { return cols; }
+  double operator[](int i) const {
+    assert(i >= 0 && i < length);
+    return data[i];
+  }
+  double operator()(int row, int col) const {
+    assert(row >= 0 && row < rows);
+    assert(col >= 0 && col < cols);
+    return data[(row * cols) + col];
+  }
+  double &operator()(int row, int col) {
+    assert(row >= 0 && row < rows);
+    assert(col >= 0 && col < cols);
+    return data[(row * cols) + col];
+  }
+  ~QSPVector() { free(data); }
+};
+
 int index(int row, int col, int rows, int cols) {
   assert(row >= 0 && row < rows);
   assert(col >= 0 && col < cols);
@@ -26,7 +69,9 @@ void QSPCount(const TArrayn::DTArray &t, const TArrayn::DTArray &u,
               const char T1_name, const char S1_name, const int NS,
               const int NT, double T1_max, double S1_max, double T1_min,
               double S1_min, const int Nx, const int Ny, const int Nz,
-              string filename, const int plotnum) {
+              string filename, const int plotnum, bool mapped,
+              TArrayn::DTArray *xgrid, TArrayn::DTArray *ygrid,
+              TArrayn::DTArray *zgrid) {
 
   int local_rank;
   MPI_Comm_rank(MPI_COMM_WORLD, &local_rank);
@@ -175,10 +220,37 @@ void QSPCount(const TArrayn::DTArray &t, const TArrayn::DTArray &u,
   double hS_inv = 1 / hS;
   double hT_inv = 1 / hT;
 
-  int *local_hist = (int *)calloc(NS * NT, sizeof(int));
-  if (!local_hist) {
-    std::cout << "Bad memory allocation. Exiting QSPCount" << std::endl;
-    return;
+  QSPVector local_hist(NS, NT);
+  QSPVector global_z_max(Nx, Ny);
+  QSPVector global_z_min(Nx, Ny);
+  // Find the range of Lz values per 2D-slice
+  if (mapped) {
+    QSPVector local_z_max(Nx, Ny);
+    QSPVector local_z_min(Nx, Ny);
+    //  We are slicing as if we are doing zgrid[i, j, :]
+    for (int ii = i_low; ii <= i_high; ii++) {
+      for (int jj = j_low; jj <= j_high; jj++) {
+        // min is set to the highest possible value so it always gets changed
+        double tmp_z_min = std::numeric_limits<double>::max();
+        // max is set to the lowest possible value so it always gets changed
+        double tmp_z_max = -tmp_z_min;
+        double tmp;
+        for (int kk = k_low; kk <= k_high; kk++) {
+          tmp = (*zgrid)(ii, jj, kk);
+          if (tmp > tmp_z_max) {
+            tmp_z_max = tmp;
+          } else if (tmp < tmp_z_min) {
+            tmp_z_min = tmp;
+          }
+        }
+        local_z_max(ii, jj) = tmp_z_max;
+        local_z_min(ii, jj) = tmp_z_min;
+      }
+    }
+    MPI_Allreduce(local_z_min.raw(), global_z_min.raw(), Nx * Ny, MPI_DOUBLE,
+                  MPI_MIN, MPI_COMM_WORLD);
+    MPI_Allreduce(local_z_max.raw(), global_z_max.raw(), Nx * Ny, MPI_DOUBLE,
+                  MPI_MAX, MPI_COMM_WORLD);
   }
 
   // Main loop for QSP
@@ -242,23 +314,51 @@ void QSPCount(const TArrayn::DTArray &t, const TArrayn::DTArray &u,
         } else if (idxS >= NS) {
           idxS = 0;
         }
-        local_hist[index(idxS, idxT, NS, NT)] += 1;
+
+        double volume_weight;
+        if (mapped) {
+          // Calculate the Lz range
+          double Lzmax_now = global_z_max(i, j);
+          double Lzmin_now = global_z_min(i, j);
+
+          // Calculate the arc length
+          double arc, z_high, z_low, cos_high, cos_low;
+          if (k > 0 && k < Nz - 1) {
+            z_high = (double)(k + 1);
+            z_low = (double)(k - 1);
+          } else if (k == 0) {
+            z_high = (double)(k + 1);
+            z_low = (double)(k);
+          } else if (k == Nz - 1) {
+            z_high = (double)(k);
+            z_low = (double)(k - 1);
+          } else { // Failure
+            std::cout << "k was out of bounds somehow. Failing...\n";
+            return;
+          }
+          cos_high = std::cos(M_PI * z_high / (double)Nz);
+          cos_low = std::cos(M_PI * z_low / (double)Nz);
+          arc = 0.5 * (cos_low - cos_high);
+
+          // Calculate the volume weight
+          volume_weight = arc * (Lzmax_now - Lzmin_now);
+        } else {
+          volume_weight = 1.0;
+        }
+
+        // local_hist[index(idxS, idxT, NS, NT)] += volume_weight;
+        local_hist(idxS, idxT) += volume_weight;
       }
     }
   }
 
   MPI_Barrier(MPI_COMM_WORLD); // Wait for everyone to finish
   if (local_rank == 0) {
-    int *glob_hist = (int *)calloc(NS * NT, sizeof(int));
-    if (!glob_hist) {
-      std::cout << "Bad memory allocation. Exiting QSPCount" << std::endl;
-      free(local_hist);
-      return;
-    }
-    MPI_Reduce(local_hist, glob_hist, // send and receive buffers
-               NS * NT, MPI_INT,      // count and datatype
-               MPI_SUM, 0,            // Reduction operator and root process #
-               MPI_COMM_WORLD);       // Communicator
+    QSPVector glob_hist(NS, NT);
+    MPI_Reduce(local_hist.raw(), glob_hist.raw(), // send and receive buffers
+               NS * NT, MPI_DOUBLE,               // count and datatype
+               MPI_SUM, 0,      // Reduction operator and root process #
+               MPI_COMM_WORLD); // Communicator
     filename = filename + "." + boost::lexical_cast<string>(plotnum) + ".csv";
     std::fstream outfile;
     outfile.open(filename.c_str(), std::ios_base::out);
@@ -276,12 +376,10 @@ void QSPCount(const TArrayn::DTArray &t, const TArrayn::DTArray &u,
         outfile << std::endl;
       }
     }
-    free(glob_hist);
   } else {
-    MPI_Reduce(local_hist, NULL, // send and receive buffers
-               NS * NT, MPI_INT, // count and datatype
-               MPI_SUM, 0,       // Reduction operator and root process #
-               MPI_COMM_WORLD);  // Communicator
+    MPI_Reduce(local_hist.raw(), NULL, // send and receive buffers
+               NS * NT, MPI_DOUBLE,    // count and datatype
+               MPI_SUM, 0,             // Reduction operator and root process #
+               MPI_COMM_WORLD);        // Communicator
   }
-  free(local_hist);
 }

src/cases/derivatives/derivatives.cpp

@@ -10,6 +10,7 @@
 
 // Defines limits of intrinsic types. Used as default values for
 // T1_max, T1_min, S1_max and/orS1_min
+#include <cassert>
 #include <limits>
 
 // Tensor variables for indexing
@@ -38,6 +39,7 @@ double T1_max, S1_max, T1_min, S1_min;
 char T1_name, S1_name;
 int NT, NS;
 string QSP_filename;
+bool use_salinity;
 
 // physical parameters
 double visco;                       // viscosity (m^2/s)
@@ -76,6 +78,7 @@ class userControl : public BaseCase {
         DTArray *temp1, *temp2, *temp3;   // arrays for vortex stretching / enstrophy production
         DTArray *temp4; // array for storing temporary array for baroclinic vorticity
         DTArray *A11, *A12, *A13, *A22, *A23, *A33; //Arrays for components of S^2+Omega^2
+        DTArray *xgrid, *ygrid, *zgrid;   // Arrays for storing grid data
 
         /* Size of domain */
         double length_x() const { return Lx; }
@@ -129,10 +132,27 @@ class userControl : public BaseCase {
             split(deriv_filenames.c_str(), ' ', fields);    // populate that vector
             if ( do_vort_stretch or do_enst_stretch or do_Q or do_R or
                  do_Q_and_R or do_lambda2 or do_hist ) {
+
                 temp1 = alloc_array(Nx,Ny,Nz);
                 temp2 = alloc_array(Nx,Ny,Nz);
-                if ( do_enst_stretch )
-                    temp3 = alloc_array(Nx,Ny,Nz);
+
+                // If user asks to grab the grids, we allocate arrays to store
+                // them in memory
+                if (mapped) {
+                  xgrid = alloc_array(Nx, Ny, Nz);
+                  if (Ny > 1) {
+                    ygrid = alloc_array(Nx, Ny, Nz);
+                  }
+                  zgrid = alloc_array(Nx, Ny, Nz);
+                } else {
+                  // If user doesn't want to grab grids, we make sure not to
+                  // allocate arrays for them and to set the pointers to NULL.
+                  xgrid = NULL;
+                  ygrid = NULL;
+                  zgrid = NULL;
+                }
+
+                if ( do_enst_stretch ) temp3 = alloc_array(Nx,Ny,Nz);
             }
 
             if ( do_barvor ) {
@@ -409,18 +429,34 @@ class userControl : public BaseCase {
                 // Compute QSP data. The code promises to not mutate the arrays,
                 // nor to make deep copies of them
                 if ( do_hist ){
-                    // Read in T to temp1 if required.
-                    if (T1_name == 't' || S1_name == 't' || T1_name == 'T' ||
+
+                    if (use_salinity) {
+                      init_tracer_restart("s", *temp1);
+                    } else if (T1_name == 't' || S1_name == 't' || T1_name == 'T' ||
                         S1_name == 'T') {
                         init_tracer_restart("t", *temp1);
                     }
-                    QSPCount(*temp1, u, v, w, T1_name, S1_name, NS, NT,
-                             T1_max, S1_max, T1_min, S1_min,
-                             Nx, Ny, Nz, QSP_filename, plotnum);
+
+                    // If user asked to grab the grids, we populate the grids
+                    // with the correct data from disk
+                    if (mapped) {
+                      do_mapping(*xgrid, *ygrid, *zgrid);
+                    } else {
+                      // Make sure that if the user didn't want us to grab the
+                      // grids then we haven't allocated data to store them!
+                      assert(xgrid == NULL);
+                      assert(ygrid == NULL);
+                      assert(zgrid == NULL);
+                    }
+
+                    QSPCount(*temp1, u, v, w, T1_name, S1_name, NS, NT, T1_max,
+                             S1_max, T1_min, S1_min, Nx, Ny, Nz, QSP_filename,
+                             plotnum, mapped, xgrid, ygrid, zgrid);
                     if (master()) {
                       fprintf(stdout, "Completed the write for QSP.%d\n", plotnum);
                     }
                 }
+
             }
         }
 
@@ -500,6 +536,7 @@ int main(int argc, char ** argv) {
     add_option("T1_min",&T1_min,std::numeric_limits<double>::min(), "Minimum explicit bin for T1 in QSP.");
     add_option("S1_max",&S1_max,std::numeric_limits<double>::max(), "Maximum explicit bin for S1 in QSP.");
     add_option("S1_min",&S1_min,std::numeric_limits<double>::min(), "Minimum explicit bin for S1 in QSP.");
+    add_option("QSP_salinity",&use_salinity, false, "Should salinity be read in from filename s?.");
     add_option("QSP_filename",&QSP_filename,"QSP_default", "Filename to save data to. Don't include file extension.");
     add_option("NS",&NS,10,"Number of bins for tracer S");
     add_option("NT",&NT,10,"Number of bins for tracer T");