Qh/redshift evaluation

bliss/encoder/variational_dist.py

@@ -20,7 +20,6 @@
 class VariationalDist(torch.nn.Module):
     def __init__(self, factors, tile_slen):
         super().__init__()
-
         self.factors = factors
         self.tile_slen = tile_slen
 

case_studies/redshift/evaluation/notebook_plot.yaml

@@ -1,35 +1,57 @@
 ---
 defaults:
-    - ../redshift_from_img@_here_: full_train_config_redshift
+    - ../../../bliss/conf@_here_: base_config
     - _self_
 
+paths:
+    root: /home/qiaozhih/bliss
+
+global_setting:
+    min_flux_for_loss: 50  # you need to regenerate split_results after changing this number
+
+variational_factors:
+  - _target_: bliss.encoder.variational_dist.NormalFactor
+    name: redshifts
+    sample_rearrange: "b ht wt -> b ht wt 1 1"
+    nll_rearrange: "b ht wt 1 1 -> b ht wt"
+    nll_gating: is_galaxy
+
 encoder:
     _target_: case_studies.redshift.redshift_from_img.encoder.encoder.RedshiftsEncoder
     # _target_: bliss.encoder.encoder.Encoder
     survey_bands: ["g", "i", "r", "u", "y", "z"]
     tile_slen: 4
-    tiles_to_crop: 1
-    min_flux_for_loss: ${global_setting.min_flux_for_loss}
-    min_flux_for_metrics: 100
     optimizer_params:
         lr: 1e-3
     scheduler_params:
         milestones: [32]
         gamma: 0.1
-    image_normalizer:
-      _target_: bliss.encoder.image_normalizer.ImageNormalizer
-      bands: [0, 1, 2, 3, 4, 5]
-      include_original: false
-      include_background: false
-      concat_psf_params: false
-      num_psf_params: 4  # for SDSS, 4 for DC2
-      log_transform_stdevs: []
-      use_clahe: true
-      clahe_min_stdev: 200
+    # matcher:
+    #     _target_: case_studies.redshift.redshift_from_img.encoder.metrics.RedshiftsCatalogMatcher
+    #     match_gating: is_galaxy
+    image_normalizers:
+        psf:
+            _target_: bliss.encoder.image_normalizer.PsfAsImage
+            num_psf_params: 4  # 6 for SDSS, 4 for DC2, 10 for DES
+        clahe:
+            _target_: bliss.encoder.image_normalizer.ClaheNormalizer
+            min_stdev: 200
+        asinh:
+            _target_: bliss.encoder.image_normalizer.AsinhQuantileNormalizer
+            q: [0.01, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.99, 0.999, 0.9999, 0.99999]
+    # matcher:
+    #     _target_: bliss.encoder.metrics.CatalogMatcher
+    #     dist_slack: 1.0
+    #     mag_slack: null
+    #     mag_band: 2  # SDSS r-band
     matcher:
         _target_: case_studies.redshift.redshift_from_img.encoder.metrics.RedshiftsCatalogMatcher
         match_gating: is_galaxy
-    metrics:
+    mode_metrics:
+        _target_: torchmetrics.MetricCollection
+        _convert_: "partial"
+        metrics: ${my_metrics_test}
+    sample_metrics:
         _target_: torchmetrics.MetricCollection
         _convert_: "partial"
         metrics: ${my_metrics_test}
@@ -40,22 +62,88 @@ encoder:
     use_checkerboard: false
 
 my_metrics_test:
-    redshift_mearn_square_error_bin:
-        _target_: case_studies.redshift.redshift_from_img.encoder.metrics.RedshiftMeanSquaredErrorBin
-        mag_bin_cutoffs: [200, 400, 600, 800, 1000]
-    redshift_outlier_fraction_bin:
-        _target_: case_studies.redshift.redshift_from_img.encoder.metrics.RedshiftOutlierFractionBin
-        mag_bin_cutoffs: [200, 400, 600, 800, 1000]
-    redshift_nmad_bin:
-        _target_: case_studies.redshift.redshift_from_img.encoder.metrics.RedshiftNormalizedMedianAbsDevBin
-        mag_bin_cutoffs: [200, 400, 600, 800, 1000]
-    redshift_outlier_fraction_catastrophic_bin:
-        _target_: case_studies.redshift.redshift_from_img.encoder.metrics.RedshiftOutlierFractionCataBin
-        mag_bin_cutoffs: [200, 400, 600, 800, 1000]
-    redshift_bias_bin:
-        _target_: case_studies.redshift.redshift_from_img.encoder.metrics.RedshiftBiasBin
-        mag_bin_cutoffs: [200, 400, 600, 800, 1000]
+    # redshift_mearn_square_error:
+    #     _target_: case_studies.redshift.redshift_from_img.encoder.metrics.RedshiftMeanSquaredError
+    # redshift_mearn_square_error_bin:
+    #     _target_: case_studies.redshift.redshift_from_img.encoder.metrics.RedshiftMeanSquaredErrorBin
+    #     bin_cutoffs: [23.9, 24.1, 24.5, 24.9, 25.6]
+    #     bin_type: "njymag"
+    # redshift_mean_square_error_blendedness:
+    #     _target_: case_studies.redshift.redshift_from_img.encoder.metrics.RedshiftMeanSquaredErrorBlendedness
+    #     bin_cutoffs: [0.0001, 0.02, 0.1, 0.2, 0.6]
+    redshift_mean_square_error_true_redshifts:
+        _target_: case_studies.redshift.redshift_from_img.encoder.metrics.RedshiftMeanSquaredErrorTrueRedshift
+        bin_cutoffs: [0.5, 1, 1.5, 2, 2.5, 3]
+    # redshift_outlier_fraction:
+    #     _target_: case_studies.redshift.redshift_from_img.encoder.metrics.RedshiftOutlierFraction
+    # redshift_outlier_fraction_bin:
+    #     _target_: case_studies.redshift.redshift_from_img.encoder.metrics.RedshiftOutlierFractionBin
+    #     bin_cutoffs: [23.9, 24.1, 24.5, 24.9, 25.6]
+    #     bin_type: "njymag"
+    # redshift_nmad:
+    #     _target_: case_studies.redshift.redshift_from_img.encoder.metrics.RedshiftNormalizedMedianAbsDev
+    # redshift_nmad_bin:
+    #     _target_: case_studies.redshift.redshift_from_img.encoder.metrics.RedshiftNormalizedMedianAbsDevBin
+    #     bin_cutoffs: [23.9, 24.1, 24.5, 24.9, 25.6]
+    #     bin_type: "njymag"
+    # redshift_outlier_fraction_catastrophic:
+    #     _target_: case_studies.redshift.redshift_from_img.encoder.metrics.RedshiftOutlierFractionCata
+    # redshift_outlier_fraction_catastrophic_bin:
+    #     _target_: case_studies.redshift.redshift_from_img.encoder.metrics.RedshiftOutlierFractionCataBin
+    #     bin_cutoffs: [23.9, 24.1, 24.5, 24.9, 25.6]
+    #     bin_type: "njymag"
+    # redshift_bias:
+    #     _target_: case_studies.redshift.redshift_from_img.encoder.metrics.RedshiftBias
+    # redshift_bias_bin:
+    #     _target_: case_studies.redshift.redshift_from_img.encoder.metrics.RedshiftBiasBin
+    #     bin_cutoffs: [23.9, 24.1, 24.5, 24.9, 25.6]
+    #     bin_type: "njymag"
+    # redshift_abs_bias:
+    #     _target_: case_studies.redshift.redshift_from_img.encoder.metrics.RedshiftAbsBias
+    # redshift_abs_bias_bin:
+    #     _target_: case_studies.redshift.redshift_from_img.encoder.metrics.RedshiftAbsBiasBin
+    #     bin_cutoffs: [23.9, 24.1, 24.5, 24.9, 25.6]
+    #     bin_type: "njymag"
+
+train:
+    trainer:
+        logger:
+            name: DC2_redshift_training
+            version: DC2_redshift_only_large_split_blend
+            save_dir: ${paths.root}/case_studies/redshift/redshift_from_img/
+        precision: 32
+        strategy:
+          _target_: pytorch_lightning.strategies.DDPStrategy
+          find_unused_parameters: true
+          process_group_backend: nccl
+          timeout:
+            _target_: datetime.timedelta
+            seconds: 180000
+        val_check_interval: 0.5
+        # check_val_every_n_epoch: 1
+        # devices: [0, 2, 3, 4]
+        devices: [1]
+        max_epochs: 50
+    callbacks:
+        checkpointing:
+            _target_: pytorch_lightning.callbacks.ModelCheckpoint
+            filename: encoder_{val/mode/redshifts/mse:.6f}
+            save_top_k: 5
+            verbose: True
+            #  monitor: val/_loss
+            monitor: val/mode/redshifts/mse
+            mode: min
+        early_stopping:
+            _target_: pytorch_lightning.callbacks.early_stopping.EarlyStopping
+            monitor: val/_loss
+            mode: min
+            patience: 500
+    data_source: ${surveys.dc2}
+    pretrained_weights: null
 
 surveys:
     dc2:
-        batch_size: 64
+        dc2_cat_path: /data/scratch/dc2local/merged_catalog_with_flux_over_50.pkl
+        cached_data_path: /data/scratch/dc2local/dc2_cached_data
+        batch_size: 4
+        max_sources_per_tile: 5