Merge pull request #39 from lanl/jmm/oscillation-tracker

osc count
lanl · Dec 11, 2024 · 49237aa · 49237aa
2 parents d1efb08 + 8430b68
commit 49237aa
Show file tree

Hide file tree

Showing 6 changed files with 37 additions and 0 deletions.
diff --git a/core/decs.h b/core/decs.h
@@ -386,9 +386,11 @@ typedef double grid_Gnu_type[LOCAL_NUM_BASES][LOCAL_ANGLES_NX1]
 typedef double grid_local_moment_type[LOCAL_NUM_BASES][LOCAL_NUM_MOMENTS]
                                      [LOCAL_ANGLES_NX1][LOCAL_ANGLES_NX2];
 typedef int grid_local_basis_idx_type[LOCAL_ANGLES_NX1][LOCAL_ANGLES_NX2];
+typedef double grid_local_count_type[LOCAL_ANGLES_NX1][LOCAL_ANGLES_NX2];
 extern grid_Gnu_type          Gnu, local_Ns, local_wsqr;
 extern grid_local_moment_type local_moments;
 extern grid_local_basis_idx_type local_b_osc;
+extern grid_local_count_type local_osc_count;
 #endif // #if RAD_NUM_TYPES >= 4
 #endif // LOCAL_ANGULAR_DISTRIBUTIONS
 

diff --git a/core/defs.h b/core/defs.h
@@ -57,6 +57,7 @@ double local_dx1_rad, local_dx2_rad, local_dx_costh;
 grid_Gnu_type Gnu, local_Ns, local_wsqr;
 grid_local_moment_type local_moments;
 grid_local_basis_idx_type local_b_osc;
+grid_local_count_type local_osc_count;
 #endif //  RAD_NUM_TYPES
 #endif // LOCAL_ANGULAR_DISTRIBUTIONS
 #endif // RADIATION

diff --git a/core/diag.c b/core/diag.c
@@ -35,6 +35,9 @@ void reset_dump_variables() {
   memset(radG_int, 0, RAD_NUM_TYPES * N123G * sizeof(double));
   memset(dtau_avg, 0, (RAD_SCATT_TYPES + 1) * N123G * sizeof(double));
   memset(en_int_avg, 0, (RAD_SCATT_TYPES + 1) * N123G * sizeof(double));
+#if NEUTRINO_OSCILLATIONS
+  memset(local_osc_count, 0, LOCAL_ANGLES_NX1*LOCAL_ANGLES_NX2*sizeof(double));
+#endif // NEUTRINO_OSCILLATIONS
 #endif
 }
 

diff --git a/core/io.c b/core/io.c
@@ -1139,6 +1139,32 @@ void dump() {
           local_moments, RANK, fdims, fstart, fcount, mdims, mstart, TYPE_DBL);
 #undef RANK
     }
+
+#if NEUTRINO_OSCILLATIONS
+    {
+      mpi_dbl_allreduce_array(
+          (double *)local_osc_count, LOCAL_ANGLES_NX1 * LOCAL_ANGLES_NX2);
+      for (int ix1 = 0; ix1 < LOCAL_ANGLES_NX1; ++ix1) {
+        for (int ix2 = 0; ix2 < LOCAL_ANGLES_NX2; ++ix2) {
+          local_osc_count[ix1][ix2] /= DTd;
+        }
+      }
+
+#define RANK (2)
+      hsize_t fdims[RANK] = {LOCAL_ANGLES_NX1, LOCAL_ANGLES_NX2};
+      hsize_t fstart[RANK] = {0, 0};
+      hsize_t fcount[RANK] = {LOCAL_ANGLES_NX1, LOCAL_ANGLES_NX2};
+      hsize_t mdims[RANK] = {LOCAL_ANGLES_NX1, LOCAL_ANGLES_NX2};
+      hsize_t mstart[RANK] = {0, 0};
+      if (!mpi_io_proc()) {
+        fcount[0] = 0;
+        fcount[1] = 0;
+      }
+      WRITE_ARRAY(
+          local_osc_count, RANK, fdims, fstart, fcount, mdims, mstart, TYPE_DBL);
+#undef RANK
+    }
+#endif // NEUTRINO_OSCILLATIONS
 #endif // LOCAL_ANGULAR_DISTRIBUTIONS
 #endif // RADIATION
   }

diff --git a/core/oscillations.c b/core/oscillations.c
@@ -232,6 +232,9 @@ void oscillate(grid_local_moment_type local_moments, grid_Gnu_type gnu) {
             // adding 2 on ring 0, 1, 2, 3
             // moves through without changing to antiparticle.
             ph->type = (ph->type + (RAD_NUM_TYPES / 2)) % RAD_NUM_TYPES;
+
+            #pragma omp atomic
+            local_osc_count[ix1][ix2] += ph->w;
           }
         }
       }

diff --git a/script/analysis/hdf5_to_dict.py b/script/analysis/hdf5_to_dict.py
@@ -403,6 +403,8 @@ def load_dump(fname, geom=None, nulegacy=False):
       wb2N = Ns*wmean*wmean
       w2 = dump['local_wsqr'][:]
       dump['local_stddev'] = np.sqrt(wb2N + N2*(w2 - wb2N)/(Ns - 1))
+      if hdr['NEUTRINO_OSCILLATIONS']:
+        dump['local_osc_rate'] = dfile['local_osc_count'][:]
 
   ucon, ucov, bcon, bcov = get_state(dump, geom)