Versioned templates, NMF support, updated IGM models

sbailey commented 5 months ago

This PR adds flexibility to keep multiple versions of the same SPECTYPE/SUBTYPE templates in the redrock-templates directory, while having each template file specify which intergalactic medium (IGM) model to use and whether it should be solved using PCA (cp/np.linalg.solve) or NMF (scipy.optimize.nnls) with only positive coefficients. It is backwards compatible with current redrock-templates directories, while allowing more flexibility for the future

New templates directory

An example of the new templates directory is $CFS/desi/users/sjbailey/code/redrock-templates, including

default_templates.txt
iron_templates.txt
test_templates.txt
rrtemplate-galaxy-None-v2.6.fits
rrtemplate-qso-HIZ-v1.0.fits
rrtemplate-qso-HIZ-v1.1.fits
rrtemplate-qso-HIZ-v2.0b1.fits
rrtemplate-qso-LOZ-v1.0.fits
rrtemplate-qso-LOZ-v2.0b1.fits
rrtemplate-qso-None-v0.1.fits
rrtemplate-star-*-*.fits

The file default_templates.txt specifies which templates to use by default, but other files like test_templates.txt can specify other sets to use, e.g. for testing newer templates (like rrtemplate-qso-*-v2.0b1.fits which are NMF-based templates)

New template file keywords:

RRIGM = which IGM model to use: Calura12, Kamble20, Inoue14
RRMETHOD = 'PCA' or 'NMF'
I also added the missing VERSION keywords to the star templates

This is a draft directory for testing; I would not commit all of the test templates yet. The code in this PR can use either this new directory structure, or an existing redrock-templates like /global/common/software/desi/perlmutter/desiconda/20230111-2.1.0/code/redrock-templates/main/ .

New IGM structure

Lyman / IGM absorption code has been moved out of redrock.utils into a separate redrock.igm module. This currently supports 3 models:

Calura12 (used by Fuji and Iron)
Kamble20 (same functional form, updated coeffs, used by @abrodze for updated qso-HIZ-v1.1 templates post-iron)
Inoue14 (planned to be used for future templates)

redrock.igm.transmission_Lyman(redshifts, obswavelengths, model=model, ...) is a standardized wrapper around all 3 of those models, as well as standardizing scalar vs. vector input, use_gpu or not, always_return_array or not. Tests have been significantly expanded to check all those cases, e.g. fixing #257.

Testing

Reference run with main and current templates

cd $CFS/desi/spectro/redux/iron/tiles/cumulative/1000/20210517

srun -n 4 -c 32 rrdesi_mpi --gpu -i coadd-0-1000-thru20210517.fits \
    -o /pscratch/sd/s/sjbailey/rr/redrock-main.fits

New code, current templates

srun -n 4 -c 32 rrdesi_mpi --gpu -i coadd-0-1000-thru20210517.fits \
    -t /global/common/software/desi/perlmutter/desiconda/20230111-2.1.0/code/redrock-templates/main/ \
    -o /pscratch/sd/s/sjbailey/rr/redrock-pr-oldtemplates.fits

New code, new template directory, but same underlying templates

srun -n 4 -c 32 rrdesi_mpi --gpu -i coadd-0-1000-thru20210517.fits \
    -t $CFS/desi/users/sjbailey/code/redrock-templates \
    -o /pscratch/sd/s/sjbailey/rr/redrock-pr.fits

New code, new NMF-based QSO templates

srun -n 4 -c 32 rrdesi_mpi --gpu -i coadd-0-1000-thru20210517.fits \
    -t $CFS/desi/users/sjbailey/code/redrock-templates/test_templates.txt \
    -o /pscratch/sd/s/sjbailey/rr/redrock-prtest.fits

(For demo purposes, not yet better actual redshifts...)

fitsdiff confirms that redrock-pr-oldtemplates.fits and redrock-pr.fits match redrock-main.fits except for header keywords tracking input templates versions/locations. redrock-prtest.fits differs as expected, but has strictly positive coefficients for anything with SPECTYPE=QSO due to the NMF templates.

Using previous templates and/or new templates with Calura12 or Kamble20 IGM models runs in the same amount of time as before (~36 seconds for 500 spectra). Using NMF is a bit slower (~40 sec), but switching to the Inoue14 model is substantially slower (~55 sec). This was known from off-list discussions when the Inoue model was first added (PR #260) and could be the focus of a future optimization+PR.

Reviewing

@dylanagreen please review, especially checking whether my reorganization of the Inoue model still agrees with your implementation. Comments welcome from @abrodze, @moustakas, @abhi0395, @craigwarner-ufastro, or others as well.

Since this is backwards compatible with current templates and results, I think this is safe to merge after someone else verifies that. But it lays the groundwork for testing future templates.

Future work beyond this PR

finalize and commit new redrock-templates/ structure
- Note: the new code is backwards compatible with the old redrock-templates/ directory structure, but the current main code is not future-compatible with the new redrock-templates/ structure, so we want to do this carefully.
improve efficiency of Inoue14 implementation
fix redrock.templates.Template.eval and redrock.templates.eval_model to also include the IGM model. Neither of these functions is actually used by the core redrock code but they were provided for end-user convenience, but they are currently wrong (and they were wrong before this PR as well).
augment eval_model (or new function) to auto-derive which templates to load based upon header metadata recording the template versions used in a previous run. An example use case is rendering the models for Fuji or Iron fits to the data, even when we are using newer templates for current runs.
Add support for archetypes to also specify their IGM model (current galaxy archetypes aren't designed for high-z anyway, so this doesn't matter yet).

coveralls commented 5 months ago

coverage: 38.66% (+5.5%) from 33.193% when pulling 8a42a2899de9cc7676044bb09ae1e2d7330ad46c on versioned-templates into f1a1d06d5199b482382b9cecf67be7c9c3dde761 on main.

moustakas commented 5 months ago

FWIW I'm a few days away from being able to test the various hooks of this branch myself, but I'm also happy to report any problems as new tickets.

desihub / redrock