Weak lensing: DC2 updates

.gitignore

@@ -37,3 +37,5 @@ multirun
 DC2_*.out
 input_images/
 WeakLensingResults*
+ResultsDC2*
+case_studies/weak_lensing/**/*.pt

README.md

@@ -100,7 +100,7 @@ Configuration files are located in `bliss/conf/` and can be composed and overrid
 
 The `case_studies/` directory contains research applications of BLISS:
 
-- **weak_lensing/** - Shear (γ) and convergence (κ) estimation, validated on DC2 and DECaLS simulations
+- **weak_lensing/** - Tomographic shear (γ) and convergence (κ) mapping
 - **redshift/** - Photo-z estimation (BLISS-PZ) with multiple variational families
 - **galaxy_clustering/** - Galaxy cluster detection and membership prediction, validated on DES DR2
 - **dc2_cataloging/** - Full cataloging pipeline for DC2 simulation

ablation.txt

@@ -0,0 +1,144 @@
+Architectures considered for DC2 with four redshift bins, tile side length of 256 pixels
+
+512 true true
+input torch.Size([1, 6, 5, 2048, 2048])
+preprocess torch.Size([1, 64, 2048, 2048])
+0 torch.Size([1, 64, 1024, 1024])
+1 torch.Size([1, 64, 1024, 1024])
+2 torch.Size([1, 128, 512, 512])
+3 torch.Size([1, 128, 512, 512])
+4 torch.Size([1, 256, 256, 256])
+5 torch.Size([1, 256, 256, 256])
+6 torch.Size([1, 512, 128, 128])
+7 torch.Size([1, 512, 128, 128])
+8 torch.Size([1, 512, 64, 64])
+9 torch.Size([1, 512, 32, 32])
+10 torch.Size([1, 256, 16, 16])
+11 torch.Size([1, 128, 8, 8])
+final torch.Size([1, 8, 8, 24])
+
+
+
+512 true false
+input torch.Size([1, 6, 5, 2048, 2048])
+preprocess torch.Size([1, 64, 2048, 2048])
+0 torch.Size([1, 64, 1024, 1024])
+1 torch.Size([1, 128, 512, 512])
+2 torch.Size([1, 256, 256, 256])
+3 torch.Size([1, 512, 128, 128])
+4 torch.Size([1, 512, 64, 64])
+5 torch.Size([1, 512, 32, 32])
+6 torch.Size([1, 256, 16, 16])
+7 torch.Size([1, 128, 8, 8])
+final torch.Size([1, 8, 8, 24])
+
+
+
+512 false true
+input torch.Size([1, 6, 5, 2048, 2048])
+preprocess torch.Size([1, 64, 2048, 2048])
+0 torch.Size([1, 64, 2048, 2048])
+1 torch.Size([1, 64, 2048, 2048])
+2 torch.Size([1, 128, 1024, 1024])
+3 torch.Size([1, 128, 1024, 1024])
+4 torch.Size([1, 256, 512, 512])
+5 torch.Size([1, 256, 512, 512])
+6 torch.Size([1, 512, 256, 256])
+7 torch.Size([1, 512, 256, 256])
+8 torch.Size([1, 512, 128, 128])
+9 torch.Size([1, 512, 128, 128])
+10 torch.Size([1, 512, 64, 64])
+11 torch.Size([1, 512, 32, 32])
+12 torch.Size([1, 256, 16, 16])
+13 torch.Size([1, 128, 8, 8])
+final torch.Size([1, 8, 8, 24])
+
+
+
+512 false false
+input torch.Size([1, 6, 5, 2048, 2048])
+preprocess torch.Size([1, 64, 2048, 2048])
+0 torch.Size([1, 64, 2048, 2048])
+1 torch.Size([1, 128, 1024, 1024])
+2 torch.Size([1, 256, 512, 512])
+3 torch.Size([1, 512, 256, 256])
+4 torch.Size([1, 512, 128, 128])
+5 torch.Size([1, 512, 64, 64])
+6 torch.Size([1, 512, 32, 32])
+7 torch.Size([1, 256, 16, 16])
+8 torch.Size([1, 128, 8, 8])
+final torch.Size([1, 8, 8, 24])
+
+
+
+1024 true true
+input torch.Size([1, 6, 5, 2048, 2048])
+preprocess torch.Size([1, 64, 2048, 2048])
+0 torch.Size([1, 64, 1024, 1024])
+1 torch.Size([1, 64, 1024, 1024])
+2 torch.Size([1, 128, 512, 512])
+3 torch.Size([1, 128, 512, 512])
+4 torch.Size([1, 256, 256, 256])
+5 torch.Size([1, 256, 256, 256])
+6 torch.Size([1, 512, 128, 128])
+7 torch.Size([1, 512, 128, 128])
+8 torch.Size([1, 1024, 128, 128])
+9 torch.Size([1, 1024, 128, 128])
+10 torch.Size([1, 1024, 64, 64])
+11 torch.Size([1, 512, 32, 32])
+12 torch.Size([1, 256, 16, 16])
+13 torch.Size([1, 128, 8, 8])
+final torch.Size([1, 8, 8, 24])
+
+
+
+1024 true false
+input torch.Size([1, 6, 5, 2048, 2048])
+preprocess torch.Size([1, 64, 2048, 2048])
+0 torch.Size([1, 64, 1024, 1024])
+1 torch.Size([1, 128, 512, 512])
+2 torch.Size([1, 256, 256, 256])
+3 torch.Size([1, 512, 128, 128])
+4 torch.Size([1, 1024, 128, 128])
+5 torch.Size([1, 1024, 64, 64])
+6 torch.Size([1, 512, 32, 32])
+7 torch.Size([1, 256, 16, 16])
+8 torch.Size([1, 128, 8, 8])
+final torch.Size([1, 8, 8, 24])
+
+
+
+1024 false true
+input torch.Size([1, 6, 5, 2048, 2048])
+preprocess torch.Size([1, 64, 2048, 2048])
+0 torch.Size([1, 64, 2048, 2048])
+1 torch.Size([1, 64, 2048, 2048])
+2 torch.Size([1, 128, 1024, 1024])
+3 torch.Size([1, 128, 1024, 1024])
+4 torch.Size([1, 256, 512, 512])
+5 torch.Size([1, 256, 512, 512])
+6 torch.Size([1, 512, 256, 256])
+7 torch.Size([1, 512, 256, 256])
+8 torch.Size([1, 1024, 128, 128])
+9 torch.Size([1, 1024, 128, 128])
+10 torch.Size([1, 1024, 64, 64])
+11 torch.Size([1, 512, 32, 32])
+12 torch.Size([1, 256, 16, 16])
+13 torch.Size([1, 128, 8, 8])
+final torch.Size([1, 8, 8, 24])
+
+
+
+1024 false false
+input torch.Size([1, 6, 5, 2048, 2048])
+preprocess torch.Size([1, 64, 2048, 2048])
+0 torch.Size([1, 64, 2048, 2048])
+1 torch.Size([1, 128, 1024, 1024])
+2 torch.Size([1, 256, 512, 512])
+3 torch.Size([1, 512, 256, 256])
+4 torch.Size([1, 1024, 128, 128])
+5 torch.Size([1, 1024, 64, 64])
+6 torch.Size([1, 512, 32, 32])
+7 torch.Size([1, 256, 16, 16])
+8 torch.Size([1, 128, 8, 8])
+final torch.Size([1, 8, 8, 24])

bliss/surveys/dc2.py

@@ -413,7 +413,7 @@ def get_bands_flux_and_psf(cls, bands, catalog, median=True):
                 else:
                     psf_params_cur_band.append(catalog[f"{i}{b}"].values.astype(np.float32))
             psf_params_list.append(
-                torch.tensor(psf_params_cur_band)
+                torch.tensor(np.array(psf_params_cur_band))
             )  # bands x 4 (params per band) x n_obj
 
         return torch.stack(flux_list).t(), torch.stack(psf_params_list).unsqueeze(0)

case_studies/weak_lensing/README.md

@@ -1,25 +1,19 @@
-### Neural posterior estimation of weak lensing shear and convergence
-#### Tim White, Shreyas Chandrashekaran, Camille Avestruz, and Jeffrey Regier
-#### with assistance from Dingrui Tao, Steve Fan, and Tahseen Younus
+### Neural posterior estimation for tomographic field-level weak lensing inference
+#### Tim White, Shreyas Chandrashekaran, Dingrui Tao, Camille Avestruz, and Jeffrey Regier
+#### with assistance from Steve Fan and Tahseen Younus
 
-This case study aims to estimate weak lensing shear and convergence for the DC2 simulated sky survey. See `notebooks/dc2/manuscript` for our most recent results.
+In this case study, we use neural posterior estimation to infer tomographic shear and convergence maps from LSST-like images.
 
-Some useful commands:
-
-- Train `lensing_encoder` on DC2 images
+To train the encoder on DC2 images, run
 
 ```
-nohup bliss -cp /home/twhit/bliss/case_studies/weak_lensing/ -cn lensing_config_dc2.yaml mode=train &> train_on_dc2.out &
+nohup bliss -cp <path>/bliss/case_studies/weak_lensing/dc2 -cn config_dc2.yaml mode=train &> train_on_dc2.out &
 ```
 
-- Generate synthetic images with shear and convergence, as specified in `lensing_prior`
+To train the encoder on descwl-shear-sims images, run
 
 ```
-nohup bliss -cp /home/twhit/bliss/case_studies/weak_lensing/ -cn lensing_config_simulator.yaml mode=generate &> generate_synthetic.out &
+nohup bliss -cp <path>/bliss/case_studies/weak_lensing/descwl -cn config_descwl.yaml mode=train &> train_on_descwl.out &
 ```
 
-- Train `lensing_encoder` on synthetic images:
-
-```
-nohup bliss -cp /home/twhit/bliss/case_studies/weak_lensing/ -cn lensing_config_simulator.yaml mode=train &> train_on_synthetic.out &
-```
+See `dc2/notebooks` and `descwl/notebooks` for some exploratory plots and our most recent results.

case_studies/weak_lensing/convnet.py

@@ -10,8 +10,9 @@ def __init__(
         self,
         n_bands,
         ch_per_band,
-        n_pixels_per_side,
-        n_tiles_per_side,
+        res_init,
+        res_midpoint,
+        res_final,
         ch_init,
         ch_max,
         ch_final,
@@ -26,8 +27,6 @@ def __init__(
         if n_var_params is not None:
             ch_final = max(ch_final, 2 ** math.ceil(math.log2(n_var_params)))
 
-        res_midpoint = int(math.sqrt(n_pixels_per_side * n_tiles_per_side))
-
         self.preprocess3d = nn.Sequential(
             nn.Conv3d(n_bands, ch_init, [ch_per_band, 5, 5], padding=[0, 2, 2]),
             nn.GroupNorm(num_groups=32, num_channels=ch_init),
@@ -38,9 +37,9 @@ def __init__(
             ch_init,
             ch_max,
             ch_final,
-            n_pixels_per_side,
+            res_init,
             res_midpoint,
-            n_tiles_per_side,
+            res_final,
             initial_downsample,
             more_up_layers,
             num_bottleneck_layers,

case_studies/weak_lensing/dc2/config_dc2_redbin1.yaml → case_studies/weak_lensing/dc2/config_dc2.yaml

@@ -1,10 +1,14 @@
 ---
 defaults:
-    - ../../../bliss/conf@_here_: base_config
     - _self_
     - override hydra/job_logging: stdout
 
-mode: train
+# completely disable hydra logging
+# https://github.com/facebookresearch/hydra/issues/910
+hydra:
+    output_subdir: null
+    run:
+        dir: .
 
 paths:
     dc2: /data/scratch/dc2local
@@ -46,27 +50,27 @@ lensing_plots:
 
 encoder:
     _target_: case_studies.weak_lensing.encoder.WeakLensingEncoder
-    survey_bands: ["u", "g", "r", "i", "z", "y"]
-    reference_band: 2  # r-band
-    tile_slen: 256
-    n_pixels_per_side: 2048
-    n_tiles_per_side: 8
+    n_bands: 6
+    res_init: 2048  # num pixels per side
+    res_midpoint: 128
+    res_final: 8  # num tiles per side
     ch_init: 64
     ch_max: 1024
     ch_final: 128
-    initial_downsample: true
+    initial_downsample: false
     more_up_layers: true
     num_bottleneck_layers: 0
+    image_normalizers: ${lensing_normalizers_nanpsf}
+    var_dist:
+        _target_: bliss.encoder.variational_dist.VariationalDist
+        tile_slen: ${surveys.dc2.tile_slen}
+        factors: ${variational_factors}
     optimizer_params:
         lr: 1e-4
     scheduler_params:
         milestones: []
         gamma: 1.0
-    image_normalizers: ${lensing_normalizers_nanpsf}
-    var_dist:
-        _target_: bliss.encoder.variational_dist.VariationalDist
-        tile_slen: ${encoder.tile_slen}
-        factors: ${variational_factors}
+    loss_plots_location: ${paths.output}/${train.trainer.logger.name}/${train.trainer.logger.version}/loss_plots
     mode_metrics:
         _target_: torchmetrics.MetricCollection
         _convert_: partial
@@ -79,9 +83,6 @@ encoder:
         _target_: torchmetrics.MetricCollection
         _convert_: partial
         metrics: ${lensing_plots}
-    use_double_detect: false
-    use_checkerboard: false
-    loss_plots_location: ${paths.output}/${train.trainer.logger.name}/${train.trainer.logger.version}/loss_plots
 
 surveys:
     dc2:
@@ -94,27 +95,50 @@ surveys:
         splits: 0:80/80:90/90:100
         avg_ellip_kernel_size: 15  # must be odd
         avg_ellip_kernel_sigma: 15
-        redshift_quantiles: [0.00]
-        num_redshift_bins: 1  # length of redshift_quantiles
+        redshift_quantiles: [0.00, 0.762988, 1.120420, 1.592735]
+        num_redshift_bins: 4  # length of redshift_quantiles
         batch_size: 1
         num_workers: 1
         cached_data_path: ${paths.dc2}/dc2_lensing_splits
         train_transforms:
             - _target_: case_studies.weak_lensing.data_augmentation.LensingRotateFlipTransform
         shuffle_file_order: false  # partition train/val/test by ra/dec so that we can compute 2PCFs on spatially contiguous test set
 
+mode: train
+
 train:
     trainer:
+        _target_: pytorch_lightning.Trainer
         logger:
-            name: WeakLensingResultsDC2
-            version: ${now:%Y-%m-%d_%H-%M}
+            _target_: pytorch_lightning.loggers.TensorBoardLogger
+            save_dir: ${paths.output}
+            name: ResultsDC2
+            version: nanpsf_${encoder.ch_max}_${encoder.initial_downsample}_${encoder.more_up_layers}
+            default_hp_metric: false
+        reload_dataloaders_every_n_epochs: 0
+        check_val_every_n_epoch: 1
+        log_every_n_steps: 10
         max_epochs: 300
+        accelerator: gpu
         devices: 1
-        use_distributed_sampler: false
         precision: 32-true
     callbacks:
+        checkpointing:
+            _target_: pytorch_lightning.callbacks.ModelCheckpoint
+            filename: "encoder_{epoch}"
+            save_top_k: 1
+            verbose: true
+            monitor: val/_loss
+            mode: min
+            save_on_train_epoch_end: false
         early_stopping:
+            _target_: pytorch_lightning.callbacks.early_stopping.EarlyStopping
+            monitor: val/_loss
+            mode: min
             patience: 100
     data_source: ${surveys.dc2}
-    pretrained_weights: null
+    encoder: ${encoder}
     seed: 123123
+    pretrained_weights: null
+    ckpt_path: null
+    matmul_precision: high