Merge pull request #200 from desihub/mpi-and-logging

progress toward pure-MPI implementation and fix #199

README.rst

@@ -1,6 +1,6 @@
-.. image:: doc/_static/fastspecfit-logo.png
+.. image:: https://github.com/desihub/fastspecfit/raw/main/doc/_static/fastspecfit-logo.png
    :height: 180px
-   :target: doc/_static/fastspecfit-logo.png
+   :target: https://github.com/desihub/fastspecfit/raw/main/doc/_static/fastspecfit-logo.png
    :alt: FastSpecFit logo
 
 |Actions Status| |Coveralls Status| |Documentation Status| |ASCL Reference|

doc/changes.rst

@@ -5,11 +5,13 @@ Change Log
 3.1.1 (not released yet)
 ------------------------
 
+* Progress toward pure-MPI production code [`PR #200`_].
 * Fix <1% bias in fluxes and EWs of tied and free doublet ratios [`PR #198`_].
 * Backwards incompatible update to the data model; expanded unit tests [`PR #197`_].
 
 .. _`PR #197`: https://github.com/desihub/fastspecfit/pull/197
 .. _`PR #198`: https://github.com/desihub/fastspecfit/pull/198
+.. _`PR #200`: https://github.com/desihub/fastspecfit/pull/200
 
 3.1.0 (2024-11-21)
 ------------------

docker/Containerfile

@@ -3,7 +3,7 @@ FROM ubuntu:24.04
 RUN mkdir -p /src
 WORKDIR /src
 
-RUN apt-get -y clean && apt -y update && apt install -y apt-utils && apt -y upgrade && echo 1
+RUN apt-get -y clean && apt -y update && apt install -y apt-utils && apt -y upgrade
 
 RUN DEBIAN_FRONTEND=noninteractive \
     apt install -y --no-install-recommends \
@@ -23,11 +23,10 @@ RUN DEBIAN_FRONTEND=noninteractive \
     python3.12 \
     python3.12-dev \
     python3-pip \
-    # python3-mpi4py \
     && apt-get clean && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
 
 # python = python3
-RUN ln -s "$(which python3)" /usr/bin/python
+# RUN ln -s "$(which python3)" /usr/bin/python
 
 ARG mpich=4.1.1
 ARG mpich_prefix=mpich-$mpich
@@ -43,10 +42,13 @@ RUN wget --no-check-certificate -nv https://www.mpich.org/static/downloads/$mpic
     && make install \
     && make clean \
     && cd .. \
-    && rm -rf $mpich_prefix
+    && rm -rf $mpich_prefix \
+    && rm -f $mpich_prefix.tar.gz
 
 RUN /sbin/ldconfig
 
+ENV PIP_ROOT_USER_ACTION=ignore
+
 RUN for x in \
     wheel \
     setuptools \
@@ -65,22 +67,22 @@ RUN for x in \
     ipykernel \
     h5py \
     ; do pip3 install --break-system-packages $x; done \
-    && rm -Rf /root/.cache/pip \
-    && echo 1
+    && rm -Rf /root/.cache/pip
 
 ENV DESIUTIL_VERSION 3.4.3
 ENV DESIMODEL_VERSION 0.19.2
 ENV DESITARGET_VERSION 2.8.0
 ENV DESISPEC_VERSION 0.68.1
 ENV SPECLITE_VERSION v0.20
-ENV FASTSPECFIT_VERSION 3.1.0
+ENV FASTSPECFIT_VERSION main
+#ENV FASTSPECFIT_VERSION 3.1.1
 
 RUN pip3 install --break-system-packages git+https://github.com/desihub/desiutil.git@${DESIUTIL_VERSION}#egg=desiutil
 RUN pip3 install --break-system-packages git+https://github.com/desihub/desimodel.git@${DESIMODEL_VERSION}#egg=desimodel
 RUN pip3 install --break-system-packages git+https://github.com/desihub/desitarget.git@${DESITARGET_VERSION}#egg=desitarget
 RUN pip3 install --break-system-packages git+https://github.com/desihub/desispec.git@${DESISPEC_VERSION}#egg=desispec
 RUN pip3 install --break-system-packages git+https://github.com/desihub/speclite.git@${SPECLITE_VERSION}#egg=speclite
-RUN pip3 install --break-system-packages git+https://github.com/desihub/fastspecfit.git@${FASTSPECFIT_VERSION}#egg=fastspecfit && echo 1
+RUN pip3 install --break-system-packages git+https://github.com/desihub/fastspecfit.git@${FASTSPECFIT_VERSION}#egg=fastspecfit
 
 RUN mkdir /homedir && chmod 777 /homedir
 ENV HOME /homedir
@@ -93,11 +95,4 @@ SHELL ["/bin/bash", "-c"]
 ENTRYPOINT ["/bin/bash", "-c"]
 CMD ["/bin/bash"]
 
-#RUN echo "export PS1='[container] \\u@\\h:\\w$ '" >> $HOME/.bashrc \
-#  # Create config files in $HOME
-#  && python -c "import astropy" \
-#  && python -c "import matplotlib.font_manager as fm; f = fm.FontManager()" \
-#  && ipython -c "print('hello')" \
-#  && echo 1
-
 LABEL Maintainer="John Moustakas [email protected]"

docker/README.md

@@ -4,23 +4,36 @@ Build a Docker container for fastspecfit.
 Using podman-hpc at NERSC
 -------------------------
 
-podman-hpc pull ubuntu:22.04
-
 First, build the container:
 ```
-podman-hpc build -t desihub/fastspecfit:latest .
+time podman-hpc build --tag fastspecfit:3.1.1 --file ./Containerfile ./
+```
+then 'migrate' the image so it can be used for jobs:
+```
+podman-hpc migrate fastspecfit:3.1.1
 ```
 
-podman-hpc migrate desihub/fastspecfit:latest
+List the available images:
+```
+podman-hpc images
+```
 
 
+```
 podman-hpc login docker.io
-podman-hpc push docker.io/desihub/fastspecfit:latest
-podman-hpc pull docker.io/desihub/fastspecfit:latest
-
-podman-hpc run --rm -it localhost/desihub/fastspecfit:latest /bin/bash
+podman-hpc push fastspecfit:3.1.1
+podman-hpc pull desihub/fastspecfit:3.1.1
+```
 
+Enter the container:
+```
+podman-hpc run --rm -it fastspecfit:3.1.1 /bin/bash
+```
 
+To delete a container:
+```
+podman-hpc rmi desihub/fastspecfit:3.1.1
+```
 
 Legacy Instructions
 -------------------
@@ -41,7 +54,7 @@ Then, subsequently, to create a new (or the latest) version or tag, do:
 export DOCKER_BUILDKIT=0
 export COMPOSE_DOCKER_CLI_BUILD=0
 
-docker buildx build --platform linux/amd64,linux/arm64/v8 --push -t desihub/fastspecfit:3.1.2 .
+docker buildx build --platform linux/amd64,linux/arm64/v8 --push -t desihub/fastspecfit:3.1.1 .
 docker buildx build --platform linux/amd64,linux/arm64/v8 --push -t desihub/fastspecfit:latest .
 ```
 
@@ -52,8 +65,8 @@ docker run -it desihub/fastspecfit:latest
 ```
 or
 ```
-docker pull desihub/fastspecfit:3.1.2
-docker run -it desihub/fastspecfit:3.1.2
+docker pull desihub/fastspecfit:3.1.1
+docker run -it desihub/fastspecfit:3.1.1
 ```
 
 Or at NERSC:

py/fastspecfit/continuum.py

@@ -218,9 +218,10 @@ def _smooth_percamera(camwave, camflux, camivar, camlinemask):
 
                 # We supply estimates local inverse stddev in each window
                 # (i.e., how noisy the data is there) so that variation is
-                # down-weighted in noisier regions.
+                # down-weighted in noisier regions. Note: ext=3 means constant
+                # extrapolation.
                 if len(swave) > 3:
-                    spl_flux = UnivariateSpline(swave, sflux, w=sisig, ext=3, k=3)
+                    spl_flux = UnivariateSpline(swave, sflux, w=sisig, ext=3, k=2)
                     smoothflux = spl_flux(camwave)
                 else:
                     smoothflux = np.zeros_like(camflux)
@@ -1023,6 +1024,7 @@ def continuum_fastphot(redshift, objflam, objflamivar, CTools, uniqueid=0,
         dn4000_model = 0.
 
         coeff_monte = None
+        tauv_monte = None
         sedmodel_monte = None
         sedmodel_nolines_monte = None
 
@@ -1095,9 +1097,9 @@ def do_fit(objflam):
             msg.append(f'vdisp={vdisp:.0f} km/s')
             log.info(' '.join(msg))
 
-    return (coeff, coeff_monte, rchi2_phot, tauv, tauv_ivar, vdisp, dn4000_model,
-            dn4000_model_ivar, sedmodel, sedmodel_monte, sedmodel_nolines,
-            sedmodel_nolines_monte)
+    return (coeff, coeff_monte, rchi2_phot, tauv, tauv_monte, tauv_ivar, vdisp,
+            dn4000_model, dn4000_model_ivar, sedmodel, sedmodel_monte,
+            sedmodel_nolines, sedmodel_nolines_monte)
 
 
 def vdisp_by_chi2scan(CTools, templates, uniqueid, specflux, specwave,
@@ -1161,7 +1163,7 @@ def vdisp_by_chi2scan(CTools, templates, uniqueid, specflux, specwave,
             import seaborn as sns
             sns.set(context='talk', style='ticks', font_scale=0.8)
 
-            pngfile = f'qa-vdisp-{uniqueid}.png'
+            pngfile = f'qa-vdisp-chi2scan-{uniqueid}.png'
 
             fig, ax = plt.subplots(figsize=(8, 6))
             ax.scatter(CTools.vdisp_grid, chi2grid-chi2min, marker='s', s=50, color='gray', edgecolor='k')
@@ -1344,13 +1346,14 @@ def do_fit_vdisp(specflux):
                     import corner as cn
                     import seaborn as sns
 
-                    pngfile = f'qa-vdisp-corner-{uniqueid}.png'
+                    pngfile = f'qa-vdisp-{uniqueid}.png'
 
                     sns.set(context='talk', style='ticks', font_scale=0.6)
                     colors = sns.color_palette()
 
                     dkw = {'color': colors[1], 'ms': 10, 'alpha': 0.75, 'mec': 'k'}
-                    hkw = {'fill': True, 'alpha': 0.75, 'color': 'gray'}
+                    hkw = {'fill': True, 'alpha': 0.75, 'color': 'gray',
+                           'align': 'left', 'edgecolor': 'k'}
                     ndim = 3
 
                     tauv_sigma = np.std(tauv_monte)
@@ -1492,6 +1495,7 @@ def do_fit_full(objflam, specflux):
         dn4000_model_ivar = var2ivar(np.var(dn4000_model_monte))
     else:
         coeff_monte = None
+        tauv_monte = None
         sedmodel_monte = None
         sedmodel_nolines_monte = None
         desimodel_nolines_monte = None
@@ -1531,10 +1535,10 @@ def do_fit_full(objflam, specflux):
     log.debug(f'Deriving the smooth continuum took {time.time()-t0:.2f} seconds.')
 
     return (coeff, coeff_monte, rchi2_cont, rchi2_phot, median_apercorr, apercorrs,
-            tauv, tauv_ivar, vdisp, vdisp_ivar, dn4000, dn4000_ivar, dn4000_model,
-            dn4000_model_ivar, sedmodel, sedmodel_nolines, desimodel_nolines,
-            smoothcontinuum, smoothstats, specflux_monte, sedmodel_monte,
-            sedmodel_nolines_monte, continuummodel_monte)
+            tauv, tauv_monte, tauv_ivar, vdisp, vdisp_ivar, dn4000, dn4000_ivar,
+            dn4000_model, dn4000_model_ivar, sedmodel, sedmodel_nolines,
+            desimodel_nolines, smoothcontinuum, smoothstats, specflux_monte,
+            sedmodel_monte, sedmodel_nolines_monte, continuummodel_monte)
 
 
 def continuum_specfit(data, fastfit, specphot, templates, igm, phot,
@@ -1597,15 +1601,15 @@ def continuum_specfit(data, fastfit, specphot, templates, igm, phot,
 
     if fastphot:
         # Photometry-only fitting.
-        (coeff, coeff_monte, rchi2_phot, tauv, tauv_ivar, vdisp, dn4000_model, dn4000_model_ivar,
-         sedmodel, sedmodel_monte, sedmodel_nolines, sedmodel_nolines_monte) = \
+        (coeff, coeff_monte, rchi2_phot, tauv, tauv_monte, tauv_ivar, vdisp, dn4000_model,
+         dn4000_model_ivar, sedmodel, sedmodel_monte, sedmodel_nolines, sedmodel_nolines_monte) = \
             continuum_fastphot(redshift, objflam, objflamivar, CTools,
                                uniqueid=data['uniqueid'], debug_plots=debug_plots,
                                nmonte=nmonte, rng=rng)
     else:
         (coeff, coeff_monte, rchi2_cont, rchi2_phot, median_apercorr, apercorrs,
-         tauv, tauv_ivar, vdisp, vdisp_ivar, dn4000, dn4000_ivar, dn4000_model,
-         dn4000_model_ivar, sedmodel, sedmodel_nolines, continuummodel,
+         tauv, tauv_monte, tauv_ivar, vdisp, vdisp_ivar, dn4000, dn4000_ivar,
+         dn4000_model, dn4000_model_ivar, sedmodel, sedmodel_nolines, continuummodel,
          smoothcontinuum, smoothstats, specflux_monte, sedmodel_monte,
          sedmodel_nolines_monte, continuummodel_monte) = \
              continuum_fastspec(redshift, objflam, objflamivar, CTools,
@@ -1746,8 +1750,73 @@ def _get_sps_properties(coeff):
                                       ['AGE_IVAR', 'ZZSUN_IVAR', 'LOGMSTAR_IVAR', 'SFR_IVAR']):
                 specphot[col] = var2ivar(np.var(val_monte))
 
-        #rindx = np.argmin(np.abs(phot.absmag_filters.effective_wavelengths.value / (1.+phot.band_shift) - 5600.))
-        #msg = [f'M{phot.absmag_bands[rindx]}={absmag[rindx]:.2f} mag']
+            # optional debugging plot
+            if debug_plots:
+                import matplotlib.pyplot as plt
+                import corner as cn
+                import seaborn as sns
+
+                pngfile = f'qa-sps-properties-{data["uniqueid"]}.png'
+
+                sns.set(context='talk', style='ticks', font_scale=0.6)
+                colors = sns.color_palette()
+
+                dkw = {'color': colors[1], 'ms': 10, 'alpha': 0.75, 'mec': 'k'}
+                hkw = {'fill': True, 'alpha': 0.75, 'color': 'gray',
+                       'align': 'left', 'edgecolor': 'k'}
+                ndim = 5
+
+                zzsun_sigma = np.std(zzsun_monte)
+                tauv_sigma = np.std(tauv_monte)
+                sfr_sigma = np.std(sfr_monte)
+                logmstar_sigma = np.std(logmstar_monte)
+                age_sigma = np.std(age_monte)
+
+                plotdata = np.vstack((zzsun_monte, tauv_monte, sfr_monte, logmstar_monte, age_monte)).T
+                truths = [zzsun, tauv, sfr, logmstar, age]
+                labels = [r'$Z/Z_{\odot}$', r'$\tau_{V}$', r'SFR ($M_{\odot}/\mathrm{yr}$)',
+                          '\n'+r'$\log_{10}(M/M_{\odot})$', 'Age (Gyr)']
+                titles = [r'$Z/Z_{\odot}$='+f'{zzsun:.1f}'+r'$\pm$'+f'{zzsun_sigma:.1f}',
+                          r'$\tau_{V}$='+f'{tauv:.2f}'+r'$\pm$'+f'{tauv_sigma:.2f}',
+                          r'SFR='+f'{sfr:.1f}'+r'$\pm$'+f'{sfr_sigma:.1f}'+r' $M_{\odot}/\mathrm{yr}$',
+                          r'$\log_{10}(M/M_{\odot})$='+f'{logmstar:.2f}'+r'$\pm$'+f'{logmstar_sigma:.2f}',
+                          f'Age={age:.2f}'+r'$\pm$'+f'{age_sigma:.2f} Gyr']
+                sig = [max(5.*zzsun_sigma, 0.1), max(5.*tauv_sigma, 0.005), max(5.*sfr_sigma, 3),
+                       max(5.*logmstar_sigma, 0.1), max(5.*age_sigma, 0.005)]
+                ranges = [(prop-sig1, prop+sig1) for prop, sig1 in zip([zzsun, tauv, sfr, logmstar, age], sig)]
+
+                bins = nmonte // 3
+                if bins < 10:
+                    bins = 10
+                fig = cn.corner(plotdata, bins=bins, smooth=None, plot_density=False,
+                                plot_contours=False, range=ranges,
+                                data_kwargs=dkw, hist_kwargs=hkw, labels=labels)
+                ax = np.array(fig.axes).reshape((ndim, ndim))
+                for ii, mlval, sig in zip(range(ndim), (zzsun, tauv, sfr, logmstar, age),
+                                          (zzsun_sigma, tauv_sigma, sfr_sigma, logmstar_sigma, age_sigma)):
+                    ax[ii, ii].axvline(mlval, color=colors[0], lw=2, ls='-')
+                    ax[ii, ii].axvline(mlval+sig, color=colors[0], lw=1, ls='--')
+                    ax[ii, ii].axvline(mlval-sig, color=colors[0], lw=1, ls='--')
+                    ax[ii, ii].set_title(titles[ii])
+                    if ii == 0:
+                        ax[ii, ii].set_ylabel('Number of\nRealizations')
+                    else:
+                        xx = ax[ii, ii].twinx()
+                        xx.set_yticklabels([])
+                        xx.set_ylabel('Number of\nRealizations')
+                        xx.tick_params(right=False)
+                for yi in range(ndim):
+                    for xi in range(yi):
+                        ax[yi, xi].axvline(truths[xi], color=colors[0], lw=1, ls='-', alpha=0.75)
+                        ax[yi, xi].axhline(truths[yi], color=colors[0], lw=1, ls='-', alpha=0.75)
+                fig.suptitle(f'SPS Properties: {data["uniqueid"]}')
+
+                fig.subplots_adjust(left=0.1, right=0.92, bottom=0.1, top=0.95, wspace=0.14, hspace=0.14)
+                fig.savefig(pngfile)#, bbox_inches='tight')
+                plt.close()
+                log.info(f'Wrote {pngfile}')
+
+
         msg = []
         for label, units, val, col in zip(['vdisp', 'log(M/Msun)', 'tau(V)', 'Age', 'SFR', 'Z/Zsun'],
                                           [' km/s', '', '', ' Gyr', ' Msun/yr', ''],

py/fastspecfit/emlines.py

@@ -1008,6 +1008,8 @@ def get_fluxes(values, obsamps, emlineflux_s, specflux_nolines_s,
                            zip(values_monte, obsamps_monte, emlineflux_monte_s,
                                specflux_nolines_monte_s)]
                     boxflux_monte, flux_monte, cont_monte = tuple(zip(*res))
+                    flux_monte = np.array(flux_monte)
+                    cont_monte = np.array(cont_monte)
 
                     # Compute the variance on the line-fitting results.
 
@@ -1058,8 +1060,10 @@ def get_fluxes(values, obsamps, emlineflux_s, specflux_nolines_s,
                     fastfit[f'{linename}_EW'] = ew
 
                     if results_monte is not None:
-                        ew_monte = np.array(flux_monte) / np.array(cont_monte) / (1. + redshift) # rest frame [A]
-                        fastfit[f'{linename}_EW_IVAR'] = var2ivar(np.var(ew_monte))
+                        I = cont_monte != 0.
+                        if np.sum(I) > 2:
+                            ew_monte = flux_monte[I] / cont_monte[I] / (1. + redshift) # rest frame [A]
+                            fastfit[f'{linename}_EW_IVAR'] = var2ivar(np.var(ew_monte))
 
 
         # Measure moments for the set of lines in self.moment_lines. We need a
@@ -1623,10 +1627,10 @@ def get_results(emlineflux):
                 else:
                     dn4000_model_sigma = 0.
 
-                restwave = modelwave / (1. + redshift) # [Angstrom]
+                restwave = wave_out / (1. + redshift) # [Angstrom]
                 flam2fnu = (1 + redshift) * restwave**2 / (C_LIGHT * 1e5) * 1e-3 * 1e23 / FLUXNORM # [erg/s/cm2/A-->mJy, rest]
                 fnu_obs = data['coadd_flux'] * flam2fnu # [mJy]
-                fnu = fluxnolines * flam2fnu # [mJy]
+                fnu = flux_out_nolines * flam2fnu # [mJy]
 
                 fnu_model = continuum_out * flam2fnu
                 #fnu_fullmodel = modelflux * flam2fnu