Parameterize nside and plots skymaps using count density (#188)

* make distance uniform

* parameterize nside and calculate density

amplfi/train/models/flow.py

@@ -20,14 +20,17 @@ class FlowModel(AmplfiModel):
             This should be a subclass of `FlowArchitecture`.
         samples_per_event:
             Number of samples to draw per event for testing
+        nside:
+            nside parameter for healpy
     """
 
     def __init__(
         self,
         *args,
         arch: FlowArchitecture,
-        samples_per_event: int = 20000,
+        samples_per_event: int = 200000,
         num_corner: int = 10,
+        nside: int = 128,
         **kwargs,
     ) -> None:
 
@@ -36,6 +39,7 @@ def __init__(
         self.model = arch
         self.samples_per_event = samples_per_event
         self.num_corner = num_corner
+        self.nside = nside
 
         # save our hyperparameters
         self.save_hyperparameters(ignore=["arch"])
@@ -147,6 +151,7 @@ def test_step(self, batch, _):
                 levels=(0.5, 0.9),
             )
             result.plot_mollview(
+                self.nside,
                 outpath=skymap_filename,
             )
         self.idx += 1
@@ -163,7 +168,7 @@ def on_test_epoch_end(self):
         # searched area cum hist
         searched_areas = []
         for result in self.test_results:
-            searched_area = result.calculate_searched_area()
+            searched_area = result.calculate_searched_area(self.nside)
             searched_areas.append(searched_area)
         searched_areas = np.sort(searched_areas)
         counts = np.arange(1, len(searched_areas) + 1) / len(searched_areas)

amplfi/train/priors.py

@@ -34,10 +34,9 @@ def cbc_prior() -> ParameterSampler:
             torch.as_tensor(0.125, dtype=torch.float32),
             torch.as_tensor(0.999, dtype=torch.float32),
         ),
-        distance=distributions.PowerLaw(
+        distance=Uniform(
             torch.as_tensor(100, dtype=torch.float32),
-            torch.as_tensor(3000, dtype=torch.float32),
-            index=2,
+            torch.as_tensor(3100, dtype=torch.float32),
         ),
         inclination=distributions.Sine(
             torch.as_tensor(0, dtype=torch.float32),

amplfi/train/testing.py

@@ -7,10 +7,7 @@
 
 
 class Result(bilby.result.Result):
-    def get_sky_projection(
-        self,
-        nside: int = 32,
-    ):
+    def get_sky_projection(self, nside: int):
         """Store a HEALPix array with the sky coordinates
 
         Args:
@@ -20,6 +17,8 @@ def get_sky_projection(
         dec = self.posterior["dec"]
         theta = np.pi / 2 - dec
 
+        num_samples = len(ra)
+
         # mask out non physical samples;
         mask = (ra > -np.pi) * (ra < np.pi)
         mask &= (theta > 0) * (theta < np.pi)
@@ -34,14 +33,15 @@ def get_sky_projection(
         ipix = np.sort(ipix)
         uniq, counts = np.unique(ipix, return_counts=True)
 
-        # create empty map and then fill in non-zero pix with counts
+        # create empty map and then fill in non-zero pix with density
+        # estimated by fraction of total samples in each pixel
         m = np.zeros(NPIX)
-        m[np.in1d(range(NPIX), uniq)] = counts
+        m[np.in1d(range(NPIX), uniq)] = counts / num_samples
 
         return m
 
-    def calculate_searched_area(self, nside: int = 32):
-        healpix = self.get_sky_projection(nside=nside)
+    def calculate_searched_area(self, nside: int):
+        healpix = self.get_sky_projection(nside)
 
         ra_inj = self.injection_parameters["phi"]
         dec_inj = self.injection_parameters["dec"]
@@ -59,11 +59,7 @@ def calculate_searched_area(self, nside: int = 32):
         )
         return searched_area
 
-    def plot_mollview(
-        self,
-        outpath: Path = None,
-        nside: int = 32,
-    ):
+    def plot_mollview(self, nside: int, outpath: Path = None):
         healpix = self.get_sky_projection(nside)
         ra_inj = self.injection_parameters["phi"]
         dec_inj = self.injection_parameters["dec"]