Merge pull request #20 from LSSTDESC/dev/joseph

Dev/joseph

Author: JospehCeh

examples/SHIRE_demo_LSSTsim_mini.ipynb

@@ -93,8 +93,8 @@
     "    ssp_file=\"ssp_data_fsps_v3.2_lgmet_age.h5\",\n",
     "    filter_dict=lsst_filts_dict,\n",
     "    wlmin=900.,\n",
-    "    wlmax=25000.,\n",
-    "    dwl=100.,\n",
+    "    wlmax=12000.,\n",
+    "    dwl=20.,\n",
     "    zmin=0.01,\n",
     "    zmax=3.0,\n",
     "    nzbins=50,\n",
@@ -244,14 +244,14 @@
     "    err_bands=_errbands,\n",
     "    zmin=0.01,\n",
     "    zmax=3.1,\n",
-    "    nzbins=310,\n",
+    "    nzbins=150,\n",
     "    ssp_file=\"ssp_data_fsps_v3.2_lgmet_age.h5\",\n",
     "    filter_dict=lsst_filts_dict,\n",
-    "    wlmin=500.,\n",
-    "    wlmax=25000.,\n",
-    "    dwl=5.,\n",
+    "    wlmin=900.,\n",
+    "    wlmax=12000.,\n",
+    "    dwl=20.,\n",
     "    no_prior=not(use_prior),\n",
-    "    chunk_size=5000\n",
+    "    chunk_size=250\n",
     ")\n",
     "\n",
     "run_shire_estimate = ShireEstimator.make_stage(\n",
@@ -355,7 +355,7 @@
     "stages = [run_shire_inform, run_shire_estimate]\n",
     "for stage in stages:\n",
     "    pipe.add_stage(stage)\n",
-    "pipe.stage_execution_config[f'shire_estimate_lsstSimhp10552_demo{_suffix}'].nprocess=1"
+    "pipe.stage_execution_config[f'shireSPS_estimate_lsstSimhp10552_demo{_suffix}'].nprocess=1"
    ]
   },
   {
@@ -426,7 +426,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.8"
+   "version": "3.13.5"
   }
  },
  "nbformat": 4,

src/rail/estimation/algos/inform.py

@@ -422,7 +422,7 @@ def _load_templates(self):
                 self.avs,
                 sspdata
             )
-            templ_d4k = treemap_d4000(templ_pars_arr, fwls, pzs, self.avs, sspdata)
+            templ_d4k = treemap_d4000(templ_pars_arr, pzs, self.avs, sspdata)
         else:
             templ_tupl_sps = make_legacy_templates(
                 templ_pars_arr,
@@ -439,7 +439,7 @@ def _load_templates(self):
                 self.avs,
                 sspdata
             )
-            templ_d4k = treemap_d4000_leg(templ_pars_arr, fwls, templ_zref, self.avs, sspdata)
+            templ_d4k = treemap_d4000_leg(templ_pars_arr, templ_zref, self.avs, sspdata)
 
         filters_names = [_fnam for _fnam, _fdir in self.config.filter_dict.items()]
         color_names = [f"{n1}-{n2}" for n1,n2 in zip(filters_names[:-1], filters_names[1:])]
@@ -1553,19 +1553,19 @@ def plot_templ_seds(self, redshifts=None, wlmin=None, wlmax=None, ymin=None, yma
                 vmap_mean_spectrum(wls, templ_pars, redshifts, sspdata),
                 0.001
             )
-            nuvk = v_nuvk_dusty(templ_pars, wls, redshifts, sspdata)
-            _selnorm = jnp.logical_and(wls>3950, wls<4000)
+            nuvk = v_nuvk_dusty(templ_pars, redshifts, sspdata)
+            _selnorm = jnp.logical_and(wls>=3950, wls<=4000)
             norms = jnp.nanmean(restframe_fnus[:, :, _selnorm], axis=2)
             restframe_fnus = restframe_fnus/jnp.expand_dims(jnp.squeeze(norms), 2)
         else:
             _vspec = vmap(mean_spectrum, in_axes=(None, 0, 0, None))
-            _vnuvk = vmap(calc_nuvk_dusty, in_axes=(0, None, 0, None))
+            _vnuvk = vmap(calc_nuvk_dusty, in_axes=(0, 0, None))
             restframe_fnus = lsunPerHz_to_fnu_noU(
                 _vspec(wls, templ_pars, templ_zref, sspdata),
                 0.001
             )
-            nuvk = _vnuvk(templ_pars, wls, templ_zref, sspdata)
-            _selnorm = jnp.logical_and(wls>3950, wls<4000)
+            nuvk = _vnuvk(templ_pars, templ_zref, sspdata)
+            _selnorm = jnp.logical_and(wls>=3950, wls<=4000)
             norms = jnp.nanmean(restframe_fnus[:, _selnorm], axis=1)
             restframe_fnus = restframe_fnus/jnp.expand_dims(jnp.squeeze(norms), 1)
 
@@ -1643,19 +1643,19 @@ def plot_templ_seds_d4000(self, redshifts=None, wlmin=None, wlmax=None, ymin=Non
                 vmap_mean_spectrum_nodust(wls, templ_pars, redshifts, sspdata),
                 0.001
             )
-            d4000n = v_d4000n(templ_pars, wls, redshifts, sspdata)
-            _selnorm = jnp.logical_and(wls>3950, wls<4000)
+            d4000n = v_d4000n(templ_pars, redshifts, sspdata)
+            _selnorm = jnp.logical_and(wls>=3950, wls<=4000)
             norms = jnp.nanmean(restframe_fnus[:, :, _selnorm], axis=2)
             restframe_fnus = restframe_fnus/jnp.expand_dims(jnp.squeeze(norms), 2)
         else:
             _vspec = vmap(mean_spectrum_nodust, in_axes=(None, 0, 0, None))
-            _vd4k = vmap(calc_d4000n, in_axes=(0, None, 0, None))
+            _vd4k = vmap(calc_d4000n, in_axes=(0, 0, None))
             restframe_fnus = lsunPerHz_to_fnu_noU(
                 _vspec(wls, templ_pars, templ_zref, sspdata),
                 0.001
             )
-            d4000n = _vd4k(templ_pars, wls, templ_zref, sspdata)
-            _selnorm = jnp.logical_and(wls>3950, wls<4000)
+            d4000n = _vd4k(templ_pars, templ_zref, sspdata)
+            _selnorm = jnp.logical_and(wls>=3950, wls<=4000)
             norms = jnp.nanmean(restframe_fnus[:, _selnorm], axis=1)
             restframe_fnus = restframe_fnus/jnp.expand_dims(jnp.squeeze(norms), 1)
         rbmap = mpl.colormaps['coolwarm']

src/rail/estimation/algos/template.py

@@ -432,7 +432,7 @@ def templ_iclrs_nuvk(params, wls, filt_trans_arr, z_obs, av, ssp_data, id_imag):
 
 
 @jit
-def calc_nuvk(pars_arr, wls, z_obs, ssp_data):
+def calc_nuvk(pars_arr, z_obs, ssp_data):
     """calc_nuvk _summary_
 
     :param pars_arr: _description_
@@ -446,138 +446,146 @@ def calc_nuvk(pars_arr, wls, z_obs, ssp_data):
     :return: _description_
     :rtype: _type_
     """
-    sed = mean_spectrum_nodust(wls, pars_arr, z_obs, ssp_data)
+    #sed = mean_spectrum_nodust(wls, pars_arr, z_obs, ssp_data)
+    
+    # get the restframe spectra without and with dust attenuation
+    ssp_wave, rest_sed, _ = ssp_spectrum_fromparam(pars_arr, z_obs, ssp_data)
     _nuvk = jnp.array(
         [
-            calc_rest_mag(wls, sed, NUV_filt.wavelength, NUV_filt.transmission),
-            calc_rest_mag(wls, sed, NIR_filt.wavelength, NIR_filt.transmission)
+            calc_rest_mag(ssp_wave, rest_sed, NUV_filt.wavelength, NUV_filt.transmission),
+            calc_rest_mag(ssp_wave, rest_sed, NIR_filt.wavelength, NIR_filt.transmission)
         ]
     )
 
     return _nuvk[0]-_nuvk[1]
 
 
-v_nuvk_zo = vmap(calc_nuvk, in_axes=(None, None, 0, None))
-v_nuvk = vmap(v_nuvk_zo, in_axes=(0, None, None, None))
+v_nuvk_zo = vmap(calc_nuvk, in_axes=(None, 0, None))
+v_nuvk = vmap(v_nuvk_zo, in_axes=(0, None, None))
 
 
 @jit
-def calc_nuvk_dusty(pars_arr, wls, z_obs, ssp_data):
+def calc_nuvk_dusty(pars_arr, z_obs, ssp_data):
     """calc_nuvk_dusty _summary_
 
     :param pars_arr: _description_
     :type pars_arr: _type_
-    :param wls: _description_
-    :type wls: _type_
     :param z_obs: _description_
     :type z_obs: _type_
     :param ssp_data: _description_
     :type ssp_data: _type_
     :return: _description_
     :rtype: _type_
     """
-    sed = mean_spectrum(wls, pars_arr, z_obs, ssp_data)
+    #sed = mean_spectrum(wls, pars_arr, z_obs, ssp_data)
+    
+    # get the restframe spectra without and with dust attenuation
+    ssp_wave, _, sed_attenuated = ssp_spectrum_fromparam(pars_arr, z_obs, ssp_data)
     _nuvk = jnp.array(
         [
-            calc_rest_mag(wls, sed, NUV_filt.wavelength, NUV_filt.transmission),
-            calc_rest_mag(wls, sed, NIR_filt.wavelength, NIR_filt.transmission)
+            calc_rest_mag(ssp_wave, sed_attenuated, NUV_filt.wavelength, NUV_filt.transmission),
+            calc_rest_mag(ssp_wave, sed_attenuated, NIR_filt.wavelength, NIR_filt.transmission)
         ]
     )
 
     return _nuvk[0]-_nuvk[1]
 
 
-v_nuvk_zo_dusty = vmap(calc_nuvk_dusty, in_axes=(None, None, 0, None))
-v_nuvk_dusty = vmap(v_nuvk_zo_dusty, in_axes=(0, None, None, None))
+v_nuvk_zo_dusty = vmap(calc_nuvk_dusty, in_axes=(None, 0, None))
+v_nuvk_dusty = vmap(v_nuvk_zo_dusty, in_axes=(0, None, None))
 
 
 @jit
-def calc_d4000n(pars_arr, wls, z_obs, ssp_data):
+def calc_d4000n(pars_arr, z_obs, ssp_data):
     """calc_d4000n _summary_
 
     :param pars_arr: _description_
     :type pars_arr: _type_
-    :param wls: _description_
-    :type wls: _type_
     :param z_obs: _description_
     :type z_obs: _type_
     :param ssp_data: _description_
     :type ssp_data: _type_
     :return: _description_
     :rtype: _type_
     """
-    sed = mean_spectrum_nodust(wls, pars_arr, z_obs, ssp_data)
+    #sed = mean_spectrum_nodust(wls, pars_arr, z_obs, ssp_data)
+    
+    # get the restframe spectra without and with dust attenuation
+    ssp_wave, rest_sed, _ = ssp_spectrum_fromparam(pars_arr, z_obs, ssp_data)
     d4000 = jnp.array(
         [
-            calc_rest_mag(wls, sed, D4000b_filt.wavelength, D4000b_filt.transmission),
-            calc_rest_mag(wls, sed, D4000r_filt.wavelength, D4000r_filt.transmission)
+            calc_rest_mag(ssp_wave, rest_sed, D4000b_filt.wavelength, D4000b_filt.transmission),
+            calc_rest_mag(ssp_wave, rest_sed, D4000r_filt.wavelength, D4000r_filt.transmission)
         ]
     )
 
     return d4000[0]-d4000[1]
 
 
-v_d4000n_zo = vmap(calc_d4000n, in_axes=(None, None, 0, None))
-v_d4000n = vmap(v_d4000n_zo, in_axes=(0, None, None, None))
+v_d4000n_zo = vmap(calc_d4000n, in_axes=(None, 0, None))
+v_d4000n = vmap(v_d4000n_zo, in_axes=(0, None, None))
 
 
 @jit
-def calc_d4000n_dusty(pars_arr, wls, z_obs, ssp_data):
+def calc_d4000n_dusty(pars_arr, z_obs, ssp_data):
     """calc_d4000n_dusty _summary_
 
     :param pars_arr: _description_
     :type pars_arr: _type_
-    :param wls: _description_
-    :type wls: _type_
     :param z_obs: _description_
     :type z_obs: _type_
     :param ssp_data: _description_
     :type ssp_data: _type_
     :return: _description_
     :rtype: _type_
     """
-    sed = mean_spectrum(wls, pars_arr, z_obs, ssp_data)
+    #sed = mean_spectrum(wls, pars_arr, z_obs, ssp_data)
+    
+    # get the restframe spectra without and with dust attenuation
+    ssp_wave, _, sed_attenuated = ssp_spectrum_fromparam(pars_arr, z_obs, ssp_data)
     d4000 = jnp.array(
         [
-            calc_rest_mag(wls, sed, D4000b_filt.wavelength, D4000b_filt.transmission),
-            calc_rest_mag(wls, sed, D4000r_filt.wavelength, D4000r_filt.transmission)
+            calc_rest_mag(ssp_wave, sed_attenuated, D4000b_filt.wavelength, D4000b_filt.transmission),
+            calc_rest_mag(ssp_wave, sed_attenuated, D4000r_filt.wavelength, D4000r_filt.transmission)
         ]
     )
 
     return d4000[0]-d4000[1]
 
 
-v_d4000n_zo_dusty = vmap(calc_d4000n_dusty, in_axes=(None, None, 0, None))
-v_d4000n_dusty = vmap(v_d4000n_zo_dusty, in_axes=(0, None, None, None))
+v_d4000n_zo_dusty = vmap(calc_d4000n_dusty, in_axes=(None, 0, None))
+v_d4000n_dusty = vmap(v_d4000n_zo_dusty, in_axes=(0, None, None))
 
 
 @jit
-def d4000n(pars_arr, wls, z_obs, av, ssp_data):
+def d4000n(pars_arr, z_obs, av, ssp_data):
     _pars = pars_arr.at[13].set(av)
-    sed = mean_spectrum(wls, _pars, z_obs, ssp_data)
+    #sed = mean_spectrum(wls, _pars, z_obs, ssp_data)
+    # get the restframe spectra without and with dust attenuation
+    ssp_wave, _, sed_attenuated = ssp_spectrum_fromparam(_pars, z_obs, ssp_data)
     d4000 = jnp.array(
         [
-            calc_rest_mag(wls, sed, D4000b_filt.wavelength, D4000b_filt.transmission),
-            calc_rest_mag(wls, sed, D4000r_filt.wavelength, D4000r_filt.transmission)
+            calc_rest_mag(ssp_wave, sed_attenuated, D4000b_filt.wavelength, D4000b_filt.transmission),
+            calc_rest_mag(ssp_wave, sed_attenuated, D4000r_filt.wavelength, D4000r_filt.transmission)
         ]
     )
 
     return d4000[0]-d4000[1]
 
-vmap_d4000n_av = vmap(d4000n, in_axes=(None, None, None, 0, None))
-vmap_d4000n_zob = vmap(vmap_d4000n_av, in_axes=(None, None, 0, None, None))
-vmap_d4000n_pars = vmap(vmap_d4000n_zob, in_axes=(0, None, None, None, None))
+vmap_d4000n_av = vmap(d4000n, in_axes=(None, None, 0, None))
+vmap_d4000n_zob = vmap(vmap_d4000n_av, in_axes=(None, 0, None, None))
+vmap_d4000n_pars = vmap(vmap_d4000n_zob, in_axes=(0, None, None, None))
 
-def treemap_d4000(pars_arr, wls, z_obs, av, ssp_data):
+def treemap_d4000(pars_arr, z_obs, av, ssp_data):
     templ_tupl = [tuple(_pars) for _pars in pars_arr]
-    reslist_of_tupl = tree_map(lambda partup: vmap_d4000n_zob(jnp.array(partup), wls, z_obs, av, ssp_data), templ_tupl, is_leaf=istuple)
+    reslist_of_tupl = tree_map(lambda partup: vmap_d4000n_zob(jnp.array(partup), z_obs, av, ssp_data), templ_tupl, is_leaf=istuple)
     return reslist_of_tupl
 
-vmap_d4000n_pars_leg = vmap(vmap_d4000n_av, in_axes=(0, None, 0, None, None))
+vmap_d4000n_pars_leg = vmap(vmap_d4000n_av, in_axes=(0, 0, None, None))
 
-def treemap_d4000_leg(pars_arr, wls, zref, av, ssp_data):
+def treemap_d4000_leg(pars_arr, zref, av, ssp_data):
     templ_tupl = [tuple(_pars)+tuple([z]) for _pars, z in zip(pars_arr, zref, strict=True)]
-    reslist_of_tupl = tree_map(lambda partup: vmap_d4000n_av(jnp.array(partup[:-1]), wls, partup[-1], av, ssp_data), templ_tupl, is_leaf=istuple)
+    reslist_of_tupl = tree_map(lambda partup: vmap_d4000n_av(jnp.array(partup[:-1]), partup[-1], av, ssp_data), templ_tupl, is_leaf=istuple)
     return reslist_of_tupl
 
 def get_colors_templates(params, wls, z_obs, transm_arr, ssp_data):
@@ -1106,8 +1114,8 @@ def bpt_rews_pars_dusty_leg(templ_pars, zref, ssp_data):
 def colrs_bptrews_templ_zo(templ_pars, wls, zobs, transm_arr, ssp_data):
     t_rews = bpt_rews_pars_zo(templ_pars, zobs, ssp_data)
     t_colors = vmap_cols_zo_nodust(templ_pars, wls, zobs, transm_arr, ssp_data)
-    t_nuvk = v_nuvk_zo(templ_pars, wls, zobs, ssp_data)
-    t_d4000n = v_d4000n_zo(templ_pars, wls, zobs, ssp_data)
+    t_nuvk = v_nuvk_zo(templ_pars, zobs, ssp_data)
+    t_d4000n = v_d4000n_zo(templ_pars, zobs, ssp_data)
     return jnp.column_stack((t_colors, t_rews, t_nuvk, t_d4000n))
 
 vmap_colrs_bptrews_templ_zo = vmap(colrs_bptrews_templ_zo, in_axes=(0, None, None, None, None))
@@ -1116,8 +1124,8 @@ def colrs_bptrews_templ_zo(templ_pars, wls, zobs, transm_arr, ssp_data):
 def colrs_bptrews_templ_zo_dusty(templ_pars, wls, zobs, transm_arr, ssp_data):
     t_rews = bpt_rews_pars_zo_dusty(templ_pars, zobs, ssp_data)
     t_colors = vmap_cols_zo(templ_pars, wls, zobs, transm_arr, ssp_data)
-    t_nuvk = v_nuvk_zo_dusty(templ_pars, wls, zobs, ssp_data) #v_nuvk_zo(templ_pars, wls, zobs, ssp_data) -- NUV-K for prior shall perhaps include dust attenuation
-    t_d4000n = v_d4000n_zo_dusty(templ_pars, wls, zobs, ssp_data)
+    t_nuvk = v_nuvk_zo_dusty(templ_pars, zobs, ssp_data) #v_nuvk_zo(templ_pars, wls, zobs, ssp_data) -- NUV-K for prior shall perhaps include dust attenuation
+    t_d4000n = v_d4000n_zo_dusty(templ_pars, zobs, ssp_data)
     return jnp.column_stack((t_colors, t_rews, t_nuvk, t_d4000n))
 
 vmap_colrs_bptrews_templ_zo_dusty = vmap(colrs_bptrews_templ_zo_dusty, in_axes=(0, None, None, None, None))
@@ -1127,8 +1135,8 @@ def colrs_bptrews_templ_zo_dusty(templ_pars, wls, zobs, transm_arr, ssp_data):
 def colrs_bptrews_templ_zo_leg(templ_pars, wls, zobs, zref, transm_arr, ssp_data):
     t_rews = bpt_rews_pars_leg(templ_pars, zref, ssp_data)
     t_colors = vmap_cols_zo_nodust_leg(templ_pars, wls, zobs, zref, transm_arr, ssp_data)
-    t_nuvk = calc_nuvk(templ_pars, wls, zref, ssp_data)
-    t_d4000n = calc_d4000n(templ_pars, wls, zref, ssp_data)
+    t_nuvk = calc_nuvk(templ_pars, zref, ssp_data)
+    t_d4000n = calc_d4000n(templ_pars, zref, ssp_data)
     return jnp.column_stack(
         (
             t_colors,
@@ -1144,8 +1152,8 @@ def colrs_bptrews_templ_zo_leg(templ_pars, wls, zobs, zref, transm_arr, ssp_data
 def colrs_bptrews_templ_zo_dusty_leg(templ_pars, wls, zobs, zref, transm_arr, ssp_data):
     t_rews = bpt_rews_pars_dusty_leg(templ_pars, zref, ssp_data)
     t_colors = vmap_cols_zo_leg(templ_pars, wls, zobs, zref, transm_arr, ssp_data)
-    t_nuvk = calc_nuvk_dusty(templ_pars, wls, zref, ssp_data) # calc_nuvk(templ_pars, wls, zref, ssp_data) # -- NUV-K for prior shall perhaps include dust attenuation
-    t_d4000n = calc_d4000n_dusty(templ_pars, wls, zref, ssp_data)
+    t_nuvk = calc_nuvk_dusty(templ_pars, zref, ssp_data) # calc_nuvk(templ_pars, wls, zref, ssp_data) # -- NUV-K for prior shall perhaps include dust attenuation
+    t_d4000n = calc_d4000n_dusty(templ_pars, zref, ssp_data)
     return jnp.column_stack(
         (
             t_colors,