lanl · Yurlungur · Jul 10, 2024 · Jul 10, 2024 · Nov 17, 2024 · Nov 18, 2024
diff --git a/core/constants.h b/core/constants.h
@@ -28,6 +28,11 @@
 #define S4THW (S2THW * S2THW)
 #define NUSIGMA0 (1.7611737037e-44) // Fundamental neutrino cross section
 
+// Frequency scale of neutrino oscillations
+#define ROOT2 (1.4142135623730951)
+#define NUFERM (ROOT2*HBAR*HBAR*CL*CL*CL*GFERM)
+#define PEQ (1./3.) // survival probability for equilibration
+
 // Unit conversions
 #define EV (1.60217653e-12)   // Electron-volt
 #define MEV (1.0e6 * EV)      // Mega-Electron-Volt

diff --git a/core/coord.c b/core/coord.c
@@ -402,7 +402,12 @@ void set_points() {
   stopx_rad[1]  = startx_rad[1] + N1TOT * dx[1];
   stopx_rad[2]  = startx_rad[2] + N2TOT * dx[2];
   stopx_rad[3]  = startx_rad[3] + N3TOT * dx[3];
-#endif
+#if LOCAL_ANGULAR_DISTRIBUTIONS
+  local_dx1_rad = (stopx_rad[1] - startx_rad[1]) / (LOCAL_ANGLES_NX1);
+  local_dx2_rad = (stopx_rad[2] - startx_rad[2]) / (LOCAL_ANGLES_NX2);
+  local_dx_costh = 2. / (LOCAL_ANGLES_NMU);
+#endif // LOCAL_ANGULAR_DISTRIBUTIONS
+#endif // RADIATION
 #elif METRIC == MKS
   // Calculate some radii determined by the geometry
   Reh = 1. + sqrt(1. - a * a);
@@ -435,6 +440,11 @@ void set_points() {
   stopx_rad[1]  = log(Rout_rad);
   stopx_rad[2]  = startx[2] + N2TOT * dx[2];
   stopx_rad[3]  = startx[3] + N3TOT * dx[3];
+#if LOCAL_ANGULAR_DISTRIBUTIONS
+  local_dx1_rad = (stopx_rad[1] - startx_rad[1]) / (LOCAL_ANGLES_NX1);
+  local_dx2_rad = (stopx_rad[2] - startx_rad[2]) / (LOCAL_ANGLES_NX2);
+  local_dx_costh = 2. / (LOCAL_ANGLES_NMU);
+#endif // LOCAL_ANGULAR_DISTRIBUTIONS
 #endif
   poly_norm     = 0.5 * M_PI * 1. /
               (1. + 1. / (poly_alpha + 1.) * 1. / pow(poly_xt, poly_alpha));

diff --git a/core/decs.h b/core/decs.h
@@ -165,10 +165,10 @@
 #define RAD_TYPE_START (0)
 #define TYPE_TRACER (-1)
 #if RADIATION == RADTYPE_NEUTRINOS
-#define RAD_NUM_TYPES (3)
 #define NU_ELECTRON (0)
 #define ANTINU_ELECTRON (1)
 #define NU_HEAVY (2)
+#define ANTINU_HEAVY (3)
 #if MULTISCATT_TEST
 #define RAD_SCATT_TYPES (3)
 #else
@@ -183,12 +183,10 @@
 #define RADG_YE (4)
 #define RADG_YE_EM (5)
 #elif RADIATION == RADTYPE_LIGHT
-#define RAD_NUM_TYPES (1)
 #define RAD_SCATT_TYPES (1)
 #define NRADCOMP (0)
 #define PHOTON (0)
 #else
-#define RAD_NUM_TYPES (0)
 #define RAD_SCATT_TYPES (0)
 #define NRADCOMP (0)
 #endif
@@ -368,6 +366,27 @@ extern grid_radtype_type Nem_phys, Nabs_phys, radtype_buf;
 extern grid_int_type    Nsuper;
 extern grid_double_type Esuper;
 extern grid_prim_type   psupersave;
+
+#if LOCAL_ANGULAR_DISTRIBUTIONS
+#define LOCAL_NUM_BASES (2)
+#define MOMENTS_A (0)
+#define MOMENTS_B (1)
+typedef double grid_local_angles_type[LOCAL_NUM_BASES][LOCAL_ANGLES_NX1]
+                                     [LOCAL_ANGLES_NX2][RAD_NUM_TYPES]
+                                     [LOCAL_ANGLES_NMU];
+extern grid_local_angles_type local_angles;
+extern double                 local_dx1_rad, local_dx2_rad, local_dx_costh;
+
+#if RAD_NUM_TYPES >= 4
+typedef double grid_Gnu_type[LOCAL_NUM_BASES][LOCAL_ANGLES_NX1]
+                            [LOCAL_ANGLES_NX2][LOCAL_ANGLES_NMU];
+typedef double grid_local_moment_type[LOCAL_NUM_BASES][2][LOCAL_ANGLES_NX1]
+                                     [LOCAL_ANGLES_NX2];
+extern grid_Gnu_type          Gnu, local_stddev;
+extern grid_local_moment_type local_moments;
+#endif // #if RAD_NUM_TYPES >= 4
+#endif // LOCAL_ANGULAR_DISTRIBUTIONS
+
 #endif // RADIATION
 
 // Default initialization is 0, which in this case is
@@ -671,6 +690,12 @@ extern int global_stop[NDIM];
 #define JRADLOOP for (int n = 0; n < MAXNSCATT + 2; n++)
 #define NULOOP for (int inu = 0; inu < NU_BINS + 1; inu++)
 
+#define LOCALXLOOP                             \
+  for (int i = 0; i < LOCAL_ANGLES_NX1; ++i)   \
+    for (int j = 0; j < LOCAL_ANGLES_NX2; ++j)
+#define LOCALMULOOP for (int imu = 0; imu < LOCAL_ANGLES_NMU; ++imu)
+#define LOCALXMULOOP LOCALXLOOP LOCALMULOOP
+
 #define MY_MIN(fval1, fval2) (((fval1) < (fval2)) ? (fval1) : (fval2))
 #define MY_MAX(fval1, fval2) (((fval1) > (fval2)) ? (fval1) : (fval2))
 #define MY_SIGN(fval) (((fval) < 0.) ? -1. : 1.)
@@ -981,6 +1006,20 @@ double Jnu_hdf(double nu, int type, const struct of_microphysics *m);
 double int_jnudnudOmega_hdf(const struct of_microphysics *m);
 double alpha_nu_hdf(double nu, int type, const struct of_microphysics *m);
 #endif // HDF opacities
+
+// oscillations.c
+#if RADIATION == RADTYPE_NEUTRINOS && LOCAL_ANGULAR_DISTRIBUTIONS
+double get_dt_oscillations();
+void get_local_angle_bins(
+    struct of_photon *ph, int *pi, int *pj, int *pmu1, int *pmu2);
+void accumulate_local_angles();
+#if RAD_NUM_TYPES >= 4
+void compute_local_gnu(grid_local_angles_type local_angles,
+    grid_Gnu_type local_stddev, grid_Gnu_type gnu);
+void compute_local_moments(grid_Gnu_type gnu, grid_local_moment_type moments);
+void oscillate(grid_local_moment_type local_moments, grid_Gnu_type gnu);
+#endif // RAD_NUM_TYPES >= 4
+#endif // LOCAL_ANGULAR_DISTRIBUTIONS
 #endif // RADIATION
 
 // passive.c
@@ -1065,6 +1104,7 @@ void   set_cooling_time(
 void record_lepton_flux(const struct of_photon *ph);
 void check_nu_type(const char *location);
 int  get_lepton_sign(const struct of_photon *ph);
+int  nu_is_heavy(const int radtype);
 #endif // NEUTRINOS
 #endif // RADIATION
 

diff --git a/core/defs.h b/core/defs.h
@@ -49,6 +49,15 @@ grid_radtype_type Nem_phys, Nabs_phys, radtype_buf;
 grid_int_type    Nsuper;
 grid_double_type Esuper;
 grid_prim_type   psupersave;
+
+#if LOCAL_ANGULAR_DISTRIBUTIONS
+grid_local_angles_type local_angles;
+double local_dx1_rad, local_dx2_rad, local_dx_costh;
+#if RAD_NUM_TYPES >= 4
+grid_Gnu_type Gnu, local_stddev;
+grid_local_moment_type local_moments;
+#endif //  RAD_NUM_TYPES
+#endif // LOCAL_ANGULAR_DISTRIBUTIONS
 #endif // RADIATION
 
 #if ELECTRONS

diff --git a/core/diag.c b/core/diag.c
@@ -583,8 +583,12 @@ void print_rad_types() {
         rad_type_counts[NU_ELECTRON]);
     fprintf(stdout, " ANTI                     %.2f %%\n",
         rad_type_counts[ANTINU_ELECTRON]);
-    fprintf(stdout, " HEAVY                    %.2f %%\n",
+    fprintf(stdout, " X                        %.2f %%\n",
         rad_type_counts[NU_HEAVY]);
+#if RAD_NUM_TYPES > 3
+    fprintf(stdout, " ANTIX                    %.2f %%\n",
+        rad_type_counts[ANTINU_HEAVY]);
+#endif
     fprintf(stdout, "*********************************\n\n");
   }
   timer_stop(TIMER_DIAG);

diff --git a/core/io.c b/core/io.c
@@ -359,7 +359,7 @@ void track_ph() {
       hsize_t dims[1] = {num_tracked_buf};
       hid_t   space   = H5Screate_simple(1, dims, NULL);
       ph_dsets[n]     = H5Dcreate(file_id, dsetnam, trackphfiletype, space,
-              H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);
+          H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);
       H5Sclose(space);
     }
     if (num_tracked > 0)
@@ -637,6 +637,92 @@ void dump_grid() {
     free(Lambda_h2cart_cov);
   }
 
+#if RADIATION
+#if RADIATION == RADTYPE_NEUTRINOS && LOCAL_ANGULAR_DISTRIBUTIONS
+  {
+    double *local_angles_Xharm = safe_malloc(
+        (NDIM - 1) * LOCAL_ANGLES_NX1 * LOCAL_ANGLES_NX2 * sizeof(double));
+#if METRIC == MKS
+    double *local_angles_Xbl = safe_malloc(
+        (NDIM - 1) * LOCAL_ANGLES_NX1 * LOCAL_ANGLES_NX2 * sizeof(double));
+#endif // MKS
+    double *local_angles_Xcart = safe_malloc(
+        (NDIM - 1) * LOCAL_ANGLES_NX1 * LOCAL_ANGLES_NX2 * sizeof(double));
+    int    n = 0;
+    double X[NDIM], Xcart[NDIM], r, th;
+    for (int i = 0; i < LOCAL_ANGLES_NX1; ++i) {
+      X[1] = startx_rad[1] + (i + 0.5) * local_dx1_rad; // startx is a face
+      for (int j = 0; j < LOCAL_ANGLES_NX2; ++j) {
+        X[2] = startx_rad[2] + (j + 0.5) * local_dx2_rad;
+        cart_coord(X, Xcart);
+
+        local_angles_Xharm[n + 0] = 0;
+        local_angles_Xharm[n + 1] = X[1];
+        local_angles_Xharm[n + 2] = X[2];
+
+#if METRIC == MKS
+        bl_coord(X, &r, &th);
+        local_angles_Xbl[n + 0] = 0;
+        local_angles_Xbl[n + 1] = r;
+        local_angles_Xbl[n + 2] = th;
+#endif // MKS
+
+        local_angles_Xcart[n + 0] = 0;
+        local_angles_Xcart[n + 1] = Xcart[1];
+        local_angles_Xcart[n + 2] = Xcart[2];
+
+        n += (NDIM - 1);
+      }
+    }
+#define RANK (3)
+    hsize_t fdims[RANK]  = {LOCAL_ANGLES_NX1, LOCAL_ANGLES_NX2, NDIM - 1};
+    hsize_t fstart[RANK] = {0, 0, 0};
+    hsize_t fcount[RANK] = {LOCAL_ANGLES_NX1, LOCAL_ANGLES_NX2, NDIM - 1};
+    hsize_t mdims[RANK]  = {LOCAL_ANGLES_NX1, LOCAL_ANGLES_NX2, NDIM - 1};
+    hsize_t mstart[RANK] = {0, 0, 0};
+    if (!mpi_io_proc()) {
+      fcount[0] = 0;
+      fcount[1] = 0;
+      fcount[2] = 0;
+    }
+    WRITE_ARRAY(local_angles_Xharm, RANK, fdims, fstart, fcount, mdims, mstart,
+        TYPE_DBL);
+    free(local_angles_Xharm);
+
+    WRITE_ARRAY(local_angles_Xcart, RANK, fdims, fstart, fcount, mdims, mstart,
+        TYPE_DBL);
+    free(local_angles_Xcart);
+
+#if METRIC == MKS
+    WRITE_ARRAY(
+        local_angles_Xbl, RANK, fdims, fstart, fcount, mdims, mstart, TYPE_DBL);
+    free(local_angles_Xbl);
+#endif // MKS
+#undef RANK
+  }
+
+  {
+    double *local_angles_mu = safe_malloc(LOCAL_ANGLES_NMU * sizeof(double));
+    for (int i = 0; i < LOCAL_ANGLES_NMU; ++i) {
+      local_angles_mu[i] = -1 + (i + 0.5) * local_dx_costh;
+    }
+#define RANK (1)
+    hsize_t fdims[RANK]  = {LOCAL_ANGLES_NMU};
+    hsize_t fstart[RANK] = {0};
+    hsize_t fcount[RANK] = {LOCAL_ANGLES_NMU};
+    hsize_t mdims[RANK]  = {LOCAL_ANGLES_NMU};
+    hsize_t mstart[RANK] = {0};
+    if (!mpi_io_proc()) {
+      fcount[0] = 0;
+    }
+    WRITE_ARRAY(
+        local_angles_mu, RANK, fdims, fstart, fcount, mdims, mstart, TYPE_DBL);
+#undef RANK
+    free(local_angles_mu);
+  }
+#endif // RADIATION
+#endif // LOCAL_ANGULAR_DISTRIBUTIONS
+
   H5Fflush(file_id, H5F_SCOPE_GLOBAL);
   H5Fclose(file_id);
 }
@@ -816,7 +902,7 @@ void dump() {
     hsize_t str_dims[1] = {NVAR};
     hid_t   prim_space  = H5Screate_simple(1, str_dims, NULL);
     hid_t   str_attr    = H5Acreate(
-             prim_dset, "vnams", strtype, prim_space, H5P_DEFAULT, H5P_DEFAULT);
+        prim_dset, "vnams", strtype, prim_space, H5P_DEFAULT, H5P_DEFAULT);
     H5Awrite(str_attr, strtype, vnams);
     H5Aclose(str_attr);
     H5Sclose(prim_space);
@@ -975,6 +1061,76 @@ void dump() {
       WRITE_ARRAY(nuLnu, RANK, fdims, fstart, fcount, mdims, mstart, TYPE_DBL);
 #undef RANK
     }
+
+    // local angle histograms
+#if RADIATION == RADTYPE_NEUTRINOS && LOCAL_ANGULAR_DISTRIBUTIONS
+    accumulate_local_angles();
+    {
+#define RANK (5)
+      hsize_t fdims[RANK]  = {LOCAL_NUM_BASES, LOCAL_ANGLES_NX1,
+          LOCAL_ANGLES_NX2, RAD_NUM_TYPES, LOCAL_ANGLES_NMU};
+      hsize_t fstart[RANK] = {0, 0, 0, 0, 0};
+      hsize_t fcount[RANK] = {LOCAL_NUM_BASES, LOCAL_ANGLES_NX1,
+          LOCAL_ANGLES_NX2, RAD_NUM_TYPES, LOCAL_ANGLES_NMU};
+      hsize_t mdims[RANK]  = {LOCAL_NUM_BASES, LOCAL_ANGLES_NX1,
+          LOCAL_ANGLES_NX2, RAD_NUM_TYPES, LOCAL_ANGLES_NMU};
+      hsize_t mstart[RANK] = {0, 0, 0, 0, 0};
+      if (!mpi_io_proc()) {
+        fcount[0] = 0;
+        fcount[1] = 0;
+        fcount[2] = 0;
+        fcount[3] = 0;
+        fcount[4] = 0;
+      }
+      WRITE_ARRAY(
+          local_angles, RANK, fdims, fstart, fcount, mdims, mstart, TYPE_DBL);
+#undef RANK
+    }
+
+    {
+#define RANK (4)
+      hsize_t fdims[RANK]  = {LOCAL_NUM_BASES, LOCAL_ANGLES_NX1,
+          LOCAL_ANGLES_NX2, LOCAL_ANGLES_NMU};
+      hsize_t fstart[RANK] = {0, 0, 0, 0};
+      hsize_t fcount[RANK] = {LOCAL_NUM_BASES, LOCAL_ANGLES_NX1,
+          LOCAL_ANGLES_NX2, LOCAL_ANGLES_NMU};
+      hsize_t mdims[RANK]  = {LOCAL_NUM_BASES, LOCAL_ANGLES_NX1,
+          LOCAL_ANGLES_NX2, LOCAL_ANGLES_NMU};
+      hsize_t mstart[RANK] = {0, 0, 0, 0};
+      if (!mpi_io_proc()) {
+        fcount[0] = 0;
+        fcount[1] = 0;
+        fcount[2] = 0;
+        fcount[3] = 0;
+      }
+      WRITE_ARRAY(Gnu,
+                  RANK, fdims, fstart, fcount, mdims, mstart, TYPE_DBL);
+      WRITE_ARRAY(local_stddev,
+                  RANK, fdims, fstart, fcount, mdims, mstart, TYPE_DBL);
+#undef RANK
+    }
+
+    {
+#define RANK (4)
+      hsize_t fdims[RANK] = {
+          LOCAL_NUM_BASES, 2, LOCAL_ANGLES_NX1, LOCAL_ANGLES_NX2};
+      hsize_t fstart[RANK] = {0, 0, 0, 0};
+      hsize_t fcount[RANK] = {
+          LOCAL_NUM_BASES, 2, LOCAL_ANGLES_NX1, LOCAL_ANGLES_NX2};
+      hsize_t mdims[RANK] = {
+          LOCAL_NUM_BASES, 2, LOCAL_ANGLES_NX1, LOCAL_ANGLES_NX2};
+      hsize_t mstart[RANK] = {0, 0, 0, 0};
+      if (!mpi_io_proc()) {
+        fcount[0] = 0;
+        fcount[1] = 0;
+        fcount[2] = 0;
+        fcount[3] = 0;
+      }
+      WRITE_ARRAY(
+          local_moments, RANK, fdims, fstart, fcount, mdims, mstart, TYPE_DBL);
+#undef RANK
+    }
+#endif // LOCAL_ANGULAR_DISTRIBUTIONS
 #endif // RADIATION
   }
 
@@ -1271,8 +1427,8 @@ void restart_write(int restart_type) {
     hsize_t fcount[RANK]     = {1};
     hsize_t mdims[RANK]      = {1};
     hsize_t mstart[RANK]     = {0};
-    int    *particle_offsets = safe_malloc(1 * sizeof(int));
-    int    *particle_counts  = safe_malloc(1 * sizeof(int));
+    int *   particle_offsets = safe_malloc(1 * sizeof(int));
+    int *   particle_counts  = safe_malloc(1 * sizeof(int));
     particle_offsets[0]      = npart_offset;
     particle_counts[0]       = npart_local;
     WRITE_ARRAY(
@@ -1525,8 +1681,8 @@ void restart_read(char *fname) {
     hsize_t fcount[RANK]     = {1};
     hsize_t mdims[RANK]      = {1};
     hsize_t mstart[RANK]     = {0};
-    int    *particle_offsets = safe_malloc(1 * sizeof(int));
-    int    *particle_counts  = safe_malloc(1 * sizeof(int));
+    int *   particle_offsets = safe_malloc(1 * sizeof(int));
+    int *   particle_counts  = safe_malloc(1 * sizeof(int));
     READ_ARRAY(
         particle_offsets, RANK, fdims, fstart, fcount, mdims, mstart, TYPE_INT);
     READ_ARRAY(
@@ -1774,9 +1930,9 @@ void dump_tracers() {
   WRITE_HDR(ntracers, TYPE_INT);
 
   // arrays to fil
-  int    *id     = safe_malloc(ntracers_local * sizeof(int));
-  int    *it     = safe_malloc(ntracers_local * sizeof(int));
-  int    *active = safe_malloc(ntracers_local * sizeof(int));
+  int *   id     = safe_malloc(ntracers_local * sizeof(int));
+  int *   it     = safe_malloc(ntracers_local * sizeof(int));
+  int *   active = safe_malloc(ntracers_local * sizeof(int));
   double *time   = safe_malloc(ntracers_local * sizeof(double));
   double *mass   = safe_malloc(ntracers_local * sizeof(double));
   double *Xharm  = safe_malloc(3 * ntracers_local * sizeof(double));
@@ -2072,7 +2228,7 @@ void write_array(void *data, const char *name, hsize_t rank, hsize_t *fdims,
     H5Tset_size(string_type, strlen(data));
     plist_id = H5Pcreate(H5P_DATASET_CREATE);
     dset_id  = H5Dcreate(file_id, name, string_type, filespace, H5P_DEFAULT,
-         plist_id, H5P_DEFAULT);
+        plist_id, H5P_DEFAULT);
     H5Pclose(plist_id);
     plist_id = H5Pcreate(H5P_DATASET_XFER);
     H5Dwrite(dset_id, string_type, memspace, filespace, plist_id, data);