Merge pull request #40 from lanl/jmm/fix-osc-prob

fix oscillation probability and add histograms for particle origin

core/decs.h

@@ -316,6 +316,9 @@ typedef double grid_tensor_type[N1 + 2 * NG][N2 + 2 * NG][N3 + 2 * NG][NDIM]
 typedef int    grid_int_type[N1 + 2 * NG][N2 + 2 * NG][N3 + 2 * NG];
 typedef double grid_eosvar_type[N1 + 2 * NG][N2 + 2 * NG][N3 + 2 * NG]
                                [EOS_NUM_EXTRA];
+#if RZ_HISTOGRAMS
+typedef double rz_hist_type[RZ_HISTOGRAMS_N];
+#endif // RZ_HISTOGRAMS
 typedef void (*passive_init_ftype)(int, int, int, double *, double *);
 typedef double (*hc_ftype)(double, double);
 
@@ -337,6 +340,12 @@ extern grid_radg_type   radG_buf; // ...buffer for communication
 extern grid_tensor_type Rmunu;    // Radiation stress-energy tensor
 extern grid_int_type    Nsph;
 extern grid_double_type nph;
+#if RZ_HISTOGRAMS
+extern rz_hist_type rz_r_orig_hist, rz_z_orig_hist;
+#if NEUTRINO_OSCILLATIONS
+extern rz_hist_type osc_rz_r_orig_hist, osc_rz_z_orig_hist;
+#endif // NEUTRINO_OSCILLATIONS
+#endif // RZ_HISTOGRAMS
 
 extern struct of_photon **photon_lists;
 extern struct of_photon **photon_mpi_lists;
@@ -492,6 +501,9 @@ extern double cnu_flat;
 #if MULTISCATT_TEST
 extern double ms_theta_nu0, ms_delta0;
 #endif // MULTISCATT_TEST
+#if RZ_HISTOGRAMS
+extern double rz_rmax, rz_zmax, delta_rcyl, delta_z;
+#endif // RZ_HISTOGRAMS
 #endif // RADIATION
 #if ELECTRONS
 extern double tptemin, tptemax;
@@ -575,10 +587,12 @@ struct of_photon {
   int               origin[NDIM];
   double            t0;
   int               is_tracked;
+  // Only relevant for neutrino oscillations
+  int               has_oscillated;
   struct of_photon *next;
 };
 
-#define PH_ELEM (11)
+#define PH_ELEM (12)
 struct of_track_photon {
   double X1;
   double X2;
@@ -769,7 +783,10 @@ void report_load_imbalance();
 void print_rad_types();
 void count_leptons(grid_prim_type P, double dt, int nstep);
 #endif
-#endif
+#if RZ_HISTOGRAMS
+void generate_rz_histograms();
+#endif // RZ_HISTOGRAMS
+#endif // RADIATION
 
 // electrons.c
 #if ELECTRONS

core/defs.h

@@ -27,6 +27,13 @@ grid_radg_type     radG_buf; // ...buffer for communication
 grid_tensor_type   Rmunu;
 grid_int_type      Nsph;
 grid_double_type   nph;
+#if RZ_HISTOGRAMS
+rz_hist_type rz_r_orig_hist, rz_z_orig_hist;
+#if NEUTRINO_OSCILLATIONS
+rz_hist_type osc_rz_r_orig_hist, osc_rz_z_orig_hist;
+#endif // NEUTRINO_OSCILLATIONS
+#endif // RZ_HISTOGRAMS
+
 struct of_photon **photon_lists;
 struct of_photon **photon_mpi_lists;
 
@@ -164,6 +171,9 @@ double cnu_flat;
 #if MULTISCATT_TEST
 double ms_theta_nu0, ms_delta0;
 #endif // MULTISCATT_TEST
+#if RZ_HISTOGRAMS
+double rz_rmax, rz_zmax, delta_rcyl, delta_z;
+#endif // RZ_HISTOGRAMS
 #endif // RADIATION
 #if ELECTRONS
 double tptemin, tptemax;

core/diag.c

@@ -35,10 +35,18 @@ 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 RZ_HISTOGRAMS
+  memset(rz_r_orig_hist, 0, RZ_HISTOGRAMS_N * sizeof(double));
+  memset(rz_z_orig_hist, 0, RZ_HISTOGRAMS_N * sizeof(double));
+#if NEUTRINO_OSCILLATIONS
+  memset(osc_rz_r_orig_hist, 0, RZ_HISTOGRAMS_N * sizeof(double));
+  memset(osc_rz_z_orig_hist, 0, RZ_HISTOGRAMS_N * sizeof(double));
+#endif // NEUTRINO_OSCILLATIONS
+#endif // RZ_HISTOGRAMS
 #if NEUTRINO_OSCILLATIONS
   memset(local_osc_count, 0, LOCAL_ANGLES_NX1*LOCAL_ANGLES_NX2*sizeof(double));
 #endif // NEUTRINO_OSCILLATIONS
-#endif
+#endif // RADIATION
 }
 
 void diag(int call_code) {
@@ -560,6 +568,57 @@ void count_leptons(grid_prim_type P, double dt, int nstep) {
   timer_stop(TIMER_DIAG);
 }
 
+#if RZ_HISTOGRAMS
+void generate_rz_histograms() {
+  timer_start(TIMER_DIAG);
+
+  #pragma omp parallel
+  {
+    struct of_photon *ph = photon_lists[omp_get_thread_num()];
+    while (ph != NULL) {
+      double Xorig[NDIM], Xorig_cart[NDIM];
+      int i = ph->origin[1];
+      int j = ph->origin[2];
+      int k = ph->origin[3];
+      coord(i, j, k, CENT, Xorig);
+      cart_coord(Xorig, Xorig_cart);
+
+      double x = Xorig_cart[1];
+      double y = Xorig_cart[2];
+      double rcyl = MY_MIN(sqrtf(x*x + y*y), fabs(rz_rmax) - SMALL);
+      double z = MY_MIN(fabs(Xorig_cart[3]), fabs(rz_zmax) - SMALL);
+
+      int ir = rcyl / (delta_rcyl + SMALL);
+      int iz = z / (delta_z + SMALL);
+
+#pragma omp atomic
+      rz_r_orig_hist[ir] += ph->w;
+#pragma omp atomic
+      rz_z_orig_hist[iz] += ph->w;
+
+#if NEUTRINO_OSCILLATIONS
+      if (ph->has_oscillated) {
+#pragma omp atomic
+        osc_rz_r_orig_hist[ir] += ph->w;
+#pragma omp atomic
+        osc_rz_z_orig_hist[iz] += ph->w;
+      }
+#endif // NEUTRINO_OSCILLATIONS
+
+      ph = ph->next;
+    }
+  }
+  mpi_dbl_allreduce_array((double *)rz_r_orig_hist, RZ_HISTOGRAMS_N);
+  mpi_dbl_allreduce_array((double *)rz_z_orig_hist, RZ_HISTOGRAMS_N);
+#if NEUTRINO_OSCILLATIONS
+  mpi_dbl_allreduce_array((double *)osc_rz_r_orig_hist, RZ_HISTOGRAMS_N);
+  mpi_dbl_allreduce_array((double *)osc_rz_z_orig_hist, RZ_HISTOGRAMS_N);
+#endif
+
+  timer_stop(TIMER_DIAG);
+}
+#endif // RZ_HISTOGRAMS
+
 void print_rad_types() {
   timer_start(TIMER_DIAG);
   double   total_count          = 0;

core/input.c

@@ -146,6 +146,10 @@ void set_core_params() {
   set_param("ms_theta_nu0", &ms_theta_nu0);
   set_param("ms_delta0", &ms_delta0);
 #endif // MULTISCATT_TEST
+#if RZ_HISTOGRAMS
+  set_param("rz_rmax", &rz_rmax);
+  set_param("rz_zmax", &rz_zmax);
+#endif // RZ_HISTOGRAMS
 #if (RADIATION == RADTYPE_NEUTRINOS)
 #if BURROWS_OPACITIES
   set_param("opac_param_file", &opac_param_file);
@@ -242,4 +246,9 @@ void init_params(char *pfname) {
   set_core_params();
   set_problem_params();
   read_params(pfname);
+  // TODO: Not the best place for this?
+#if RZ_HISTOGRAMS
+  delta_rcyl = rz_rmax / RZ_HISTOGRAMS_N;
+  delta_z = rz_zmax / RZ_HISTOGRAMS_N;
+#endif // RZ_HISTOGRAMS
 }

core/interact.c

@@ -603,6 +603,7 @@ memcpy((void*)(&m), (void*)(&(m_grd[i][j][k])),
           phscatt->origin[1] = i;
           phscatt->origin[2] = j;
           phscatt->origin[3] = k;
+          phscatt->has_oscillated = ph->has_oscillated;
 
           double wsave = ph->w;
           ph->w        = (1. - 1. / bias_scatt[scatt_to_do]) * ph->w;

core/io.c

@@ -248,6 +248,8 @@ void init_io() {
   H5Tinsert(phfiletype, "t0", offset, H5T_NATIVE_DOUBLE);
   offset += sizeof(double);
   H5Tinsert(phfiletype, "is_tracked", offset, H5T_NATIVE_INT);
+  offset += sizeof(int);
+  H5Tinsert(phfiletype, "has_oscillated", offset, H5T_NATIVE_INT);
 
   // Use HOFFSET to account for struct padding in memory
   phmemtype = H5Tcreate(H5T_COMPOUND, sizeof(struct of_photon));
@@ -265,6 +267,8 @@ void init_io() {
   H5Tinsert(phmemtype, "t0", HOFFSET(struct of_photon, t0), H5T_NATIVE_DOUBLE);
   H5Tinsert(phmemtype, "is_tracked", HOFFSET(struct of_photon, is_tracked),
       H5T_NATIVE_INT);
+  H5Tinsert(phmemtype, "has_oscillated",
+      HOFFSET(struct of_photon, has_oscillated), H5T_NATIVE_INT);
 
   trackphfiletype = H5Tcreate(H5T_COMPOUND, sizeof(struct of_photon));
   offset          = 0;
@@ -880,6 +884,8 @@ void dump() {
   WRITE_HDR(force_equipartition, TYPE_INT);
   int local_angular_distributions = LOCAL_ANGULAR_DISTRIBUTIONS;
   WRITE_HDR(local_angular_distributions, TYPE_INT);
+  int rz_histograms = RZ_HISTOGRAMS;
+  WRITE_HDR(rz_histograms, TYPE_INT);
 #endif
 
 #if NVAR_PASSIVE > 0
@@ -1110,12 +1116,11 @@ void dump() {
         fcount[2] = 0;
         fcount[3] = 0;
       }
-      WRITE_ARRAY(Gnu,
-                  RANK, fdims, fstart, fcount, mdims, mstart, TYPE_DBL);
-      WRITE_ARRAY(local_Ns,
-                  RANK, fdims, fstart, fcount, mdims, mstart, TYPE_DBL);
-      WRITE_ARRAY(local_wsqr,
-                  RANK, fdims, fstart, fcount, mdims, mstart, TYPE_DBL);
+      WRITE_ARRAY(Gnu, RANK, fdims, fstart, fcount, mdims, mstart, TYPE_DBL);
+      WRITE_ARRAY(
+          local_Ns, RANK, fdims, fstart, fcount, mdims, mstart, TYPE_DBL);
+      WRITE_ARRAY(
+          local_wsqr, RANK, fdims, fstart, fcount, mdims, mstart, TYPE_DBL);
 #undef RANK
     }
 
@@ -1151,21 +1156,43 @@ void dump() {
       }
 
 #define RANK (2)
-      hsize_t fdims[RANK] = {LOCAL_ANGLES_NX1, LOCAL_ANGLES_NX2};
+      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 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);
+      WRITE_ARRAY(local_osc_count, RANK, fdims, fstart, fcount, mdims, mstart,
+          TYPE_DBL);
 #undef RANK
     }
 #endif // NEUTRINO_OSCILLATIONS
 #endif // LOCAL_ANGULAR_DISTRIBUTIONS
+
+#if RZ_HISTOGRAMS
+  generate_rz_histograms();
+  {
+#define RANK (1)
+    hsize_t fdims[RANK]  = {RZ_HISTOGRAMS_N};
+    hsize_t fstart[RANK] = {0};
+    hsize_t fcount[RANK] = {RZ_HISTOGRAMS_N};
+    hsize_t mdims[RANK]  = {RZ_HISTOGRAMS_N};
+    hsize_t mstart[RANK] = {0};
+    if (!mpi_io_proc()) {
+      fcount[0] = 0;
+    }
+    WRITE_ARRAY(rz_r_orig_hist, RANK, fdims, fstart, fcount, mdims, mstart, TYPE_DBL);
+    WRITE_ARRAY(rz_z_orig_hist, RANK, fdims, fstart, fcount, mdims, mstart, TYPE_DBL);
+#if NEUTRINO_OSCILLATIONS
+    WRITE_ARRAY(osc_rz_r_orig_hist, RANK, fdims, fstart, fcount, mdims, mstart, TYPE_DBL);
+    WRITE_ARRAY(osc_rz_z_orig_hist, RANK, fdims, fstart, fcount, mdims, mstart, TYPE_DBL);
+#endif
+#undef RANK
+  }
+#endif // RZ_HISTOGRAMS
 #endif // RADIATION
   }
 

core/make_superphotons.c

@@ -275,6 +275,7 @@ void sample_photon(int i, int j, int k, double t, double dt, int type,
     } else {
       tmp[n]->is_tracked = 0;
     }
+    tmp[n]->has_oscillated = 0;
   }
 }
 

core/mpi.c

@@ -181,10 +181,10 @@ void init_mpi() {
   MPI_Type_contiguous(sizeof(struct of_photon *), MPI_BYTE, &mpi_pointer_type);
   MPI_Type_commit(&mpi_pointer_type);
   MPI_Datatype type[PH_ELEM] = {MPI_DOUBLE, MPI_DOUBLE, MPI_DOUBLE, MPI_DOUBLE,
-      MPI_DOUBLE, MPI_INT, MPI_INT, MPI_INT, MPI_DOUBLE, MPI_INT,
+      MPI_DOUBLE, MPI_INT, MPI_INT, MPI_INT, MPI_DOUBLE, MPI_INT, MPI_INT,
       mpi_pointer_type};
   int          blocklen[PH_ELEM] = {
-      NDIM * NSUP, NDIM * NSUP, NDIM * NSUP, 1, 1, 1, 1, NDIM, 1, 1, 1};
+      NDIM * NSUP, NDIM * NSUP, NDIM * NSUP, 1, 1, 1, 1, NDIM, 1, 1, 1, 1};
   MPI_Aint         disp[PH_ELEM];
   struct of_photon tmp;
   MPI_Get_address(&(tmp.X[0][0]), &(disp[0]));
@@ -197,7 +197,8 @@ void init_mpi() {
   MPI_Get_address(tmp.origin, &(disp[7]));
   MPI_Get_address(&(tmp.t0), &(disp[8]));
   MPI_Get_address(&(tmp.is_tracked), &(disp[9]));
-  MPI_Get_address(&(tmp.next), &(disp[10]));
+  MPI_Get_address(&(tmp.has_oscillated), &(disp[10]));
+  MPI_Get_address(&(tmp.next), &(disp[11]));
   MPI_Aint base;
   MPI_Get_address(&tmp, &base);
   for (int n = 0; n < PH_ELEM; n++) {

core/oscillations.c

@@ -143,12 +143,8 @@ void compute_local_gnu(grid_local_angles_type f, grid_Gnu_type local_Ns,
       const double XLN =
           (f[b][i][j][NU_HEAVY][imu] - f[b][i][j][ANTINU_HEAVY][imu]);
 
-      double   tot = 0;
-      TYPELOOP tot += f[b][i][j][itp][imu];
-      double   ebar = tot / (stddev + SMALL);
-
       double g_temp     = ELN - 0.5 * XLN;
-      gnu[b][i][j][imu] = (fabs(g_temp) > ebar) * g_temp;
+      gnu[b][i][j][imu] = (fabs(g_temp) > stddev) * g_temp;
     }
   }
 }
@@ -216,14 +212,13 @@ void oscillate(grid_local_moment_type local_moments, grid_Gnu_type gnu) {
           int g_in_B = G > 0;
 
           int in_shallow = (A_is_shallow && g_in_A) || (B_is_shallow && g_in_B);
-          int in_deep    = !(in_shallow);
 
           double peq = nu_is_heavy(ph->type) ? (2./3.) : (1./3.);
 #if FORCE_EQUIPARTITION
           double p_survival = peq;
 #else
           double p_survival =
-              in_shallow ? peq : (1 - (1 - peq) * B / (A + SMALL));
+              in_shallow ? peq : (1 - (1 - peq) * shallow / (deep + SMALL));
 #endif // FORCE_EQUIPARTITION
           double p_osc = 1. - p_survival;
           if (get_rand() < p_osc) {
@@ -232,6 +227,7 @@ 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;
+            ph->has_oscillated = 1;
 
             #pragma omp atomic
             local_osc_count[ix1][ix2] += ph->w;

prob/oscillations/build.py

@@ -131,6 +131,7 @@
 bhl.config.set_cparm('RAD_NUM_TYPES', 4)
 bhl.config.set_cparm('NEUTRINO_OSCILLATIONS', True)
 bhl.config.set_cparm('FORCE_EQUIPARTITION', False)
+bhl.config.set_cparm("RZ_HISTOGRAMS", True)
 
                            ### RUNTIME PARAMETERS ###
 

prob/torus_cbc/build.py

@@ -423,6 +423,7 @@
 bhl.config.set_cparm('X3L_RAD_BOUND', 'BC_PERIODIC')
 bhl.config.set_cparm('X3R_RAD_BOUND', 'BC_PERIODIC')
 bhl.config.set_cparm('DIAGNOSTICS_USE_RADTYPES', True)
+bhl.config.set_cparm('RZ_HISTOGRAMS', True)
 # bhl.config.set_cparm('RECORD_DT_MIN', True)
 
 if OSCILLATIONS:
@@ -469,6 +470,8 @@
     bhl.config.set_rparm('numin', 'double', default = NUMIN)
     bhl.config.set_rparm('numax', 'double', default = NUMAX)
     bhl.config.set_rparm('nph_per_proc', 'double', default = NPH_PER_PROC)
+    bhl.config.set_rparm('rz_rmax', 'double', default = Rout_rad)
+    bhl.config.set_rparm('rz_zmax', 'double', default = Rout_rad)
 
 if TRACERS:
     bhl.config.set_rparm('ntracers', 'int', default = NTCR_TOT)