minor changes in tf dataset and forward model

DifferentiableUniverseInitiative · Mar 5, 2024 · b0763c1 · b0763c1
1 parent e8707b8
commit b0763c1
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 11 deletions.
diff --git a/sbi_lens/gen_dataset/lensing_lpt_dataset.py b/sbi_lens/gen_dataset/lensing_lpt_dataset.py
@@ -47,8 +47,8 @@ def __init__(
         super().__init__(description=("LPT lensing simulations."), version=v1, **kwargs)
         self.N = N
         self.map_size = map_size
-        box_size = box_size
-        box_shape = box_shape
+        self.box_size = box_size
+        self.box_shape = box_shape
         self.gal_per_arcmin2 = gal_per_arcmin2
         self.sigma_e = sigma_e
         self.nbins = nbins
@@ -69,8 +69,8 @@ class LensingLPTDataset(tfds.core.GeneratorBasedBuilder):
     BUILDER_CONFIGS = [
         LensingLPTDatasetConfig(
             name="year_10_with_noise_score_density",
-            N=config_lsst_y_10.N,
-            map_size=config_lsst_y_10.map_size,
+            N=60,
+            map_size=5,
             box_size=[400.0, 400.0, 4000.0],
             box_shape=[300, 300, 128],
             gal_per_arcmin2=config_lsst_y_10.gals_per_arcmin2,
@@ -84,8 +84,8 @@ class LensingLPTDataset(tfds.core.GeneratorBasedBuilder):
         ),
         LensingLPTDatasetConfig(
             name="year_10_without_noise_score_density",
-            N=config_lsst_y_10.N,
-            map_size=config_lsst_y_10.map_size,
+            N=60,
+            map_size=5,
             box_size=[400.0, 400.0, 4000.0],
             box_shape=[300, 300, 128],
             gal_per_arcmin2=config_lsst_y_10.gals_per_arcmin2,
@@ -99,8 +99,8 @@ class LensingLPTDataset(tfds.core.GeneratorBasedBuilder):
         ),
         LensingLPTDatasetConfig(
             name="year_10_with_noise_score_conditional",
-            N=config_lsst_y_10.N,
-            map_size=config_lsst_y_10.map_size,
+            N=60,
+            map_size=5,
             box_size=[400.0, 400.0, 4000.0],
             box_shape=[300, 300, 128],
             gal_per_arcmin2=config_lsst_y_10.gals_per_arcmin2,
@@ -150,7 +150,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
 
     def _generate_examples(self, size):
         """Yields examples."""
-        bs = 20
+        bs = 5
 
         model = partial(
             lensingLpt,

diff --git a/sbi_lens/simulator/Lpt_field.py b/sbi_lens/simulator/Lpt_field.py
@@ -183,7 +183,7 @@ def lensingLpt(
     a,
     b,
     z0,
-    with_noise=True
+    with_noise=True,
 ):
     """
     This function defines the top-level forward model for our observations
@@ -215,7 +215,7 @@ def lensingLpt(
     )
 
     # Generate random convergence maps
-    nz = jc.redshift.smail_nz(a, b, z0, gals_per_arcmin2=gal_per_arcmin2)
+    nz = jc.redshift.smail_nz(a, b, z0, gals_per_arcmin2=gal_per_arcmin2, zmax=2.6)
     nz_shear = subdivide(nz, nbins=nbins, zphot_sigma=0.05)
 
     lensing_model = jax.jit(