desihub / redrock

Redshift fitting for spectroperfectionism
BSD 3-Clause "New" or "Revised" License
21 stars 13 forks source link

add load_templates_from_header #290

Closed sbailey closed 2 months ago

sbailey commented 2 months ago

This PR adds redrock.templates.load_templates_from_header to support loading the correct template versions to match a particular Redrock run, based upon the TEMNAMnn/TEMVERnn keywords which we have already been outputing.

This function requires the restructured redrock-templates directory+filenames in /global/cfs/cdirs/desi/users/sjbailey/code/redrock-templates/ . Once this PR is reviewed and merged, we can commit that updated structure.

It also adds support for a future keyword TEMFILnn keyword to explicitly give the filename instead of deriving it from TEMNAMnn=SPECTYPE:::SUBTYPE and TEMVERnn=VERSION -> rrtemplate-SPECTYPE-SUBTYPE-vVERSION.fits. I'll add this in a separate PR only after we've switched to the new redrock-templates structure so that we don't accidentally write files with TEMFILnn listing an old incompatible template filename.

Example usage

import os
import fitsio
from redrock.templates import load_templates_from_header

os.environ['RR_TEMPLATE_DIR']='/global/cfs/cdirs/desi/users/sjbailey/code/redrock-templates/'
hdr_fuji = fitsio.read_header('/global/cfs/cdirs/desi/spectro/redux/fuji/zcatalog/zpix-sv3-dark.fits', 1)
hdr_iron = fitsio.read_header('/global/cfs/cdirs/desi/spectro/redux/iron/zcatalog/v1/zpix-sv3-dark.fits', 1)
templates_fuji = load_templates_from_header(hdr_fuji)
templates_iron = load_templates_from_header(hdr_iron)

Output when loading Fuji templates:

INFO: rrtemplate-GALAXY-None-v2.6.fits GALAXY  PCA IGM=Calura12 z=-0.005-1.6997
INFO: rrtemplate-QSO-None-v0.1.fits QSO  PCA IGM=Calura12 z=0.05-5.9934
...

Output when loading Iron templates:

INFO: rrtemplate-GALAXY-None-v2.6.fits GALAXY  PCA IGM=Calura12 z=-0.005-1.6997
INFO: rrtemplate-QSO-HIZ-v1.0.fits QSO HIZ PCA IGM=Calura12 z=1.4-6.993
INFO: rrtemplate-QSO-LOZ-v1.0.fits QSO LOZ PCA IGM=Calura12 z=0.05-1.5983
...

Note that it correctly loaded the same GALAXY template but different QSO templates for Fuji vs. Iron, all from a single RR_TEMPLATE_DIR.

coveralls commented 2 months ago

Coverage Status

coverage: 39.085% (+0.5%) from 38.557% when pulling 8a8f0169b7757fc846a49a0f7c0869ce49c8cb08 on templateio into c30c913a990530ba07dc4af5c9d0bef111133de4 on main.

moustakas commented 2 months ago

@sbailey this didn't work for me. Should it have? (It's your same test but without changing RR_TEMPLATE_DIR.)

import os
import fitsio
from redrock.templates import load_templates_from_header

hdr_fuji = fitsio.read_header('/global/cfs/cdirs/desi/spectro/redux/fuji/zcatalog/zpix-sv3-dark.fits', 1)
templates_fuji = load_templates_from_header(hdr_fuji)
Reading templates from ['/global/common/software/desi/perlmutter/desiconda/20230111-2.1.0/code/redrock-templates/main/rrtemplate-GALAXY-None-v2.6.fits', '/global/common/software/desi/perlmutter/desiconda/20230111-2.1.0/code/redrock-templates/main/rrtemplate-QSO-None-v0.1.fits', '/global/common/software/desi/perlmutter/desiconda/20230111-2.1.0/code/redrock-templates/main/rrtemplate-STAR-A-v0.1.fits', '/global/common/software/desi/perlmutter/desiconda/20230111-2.1.0/code/redrock-templates/main/rrtemplate-STAR-B-v0.1.fits', '/global/common/software/desi/perlmutter/desiconda/20230111-2.1.0/code/redrock-templates/main/rrtemplate-STAR-CV-v0.1.fits', '/global/common/software/desi/perlmutter/desiconda/20230111-2.1.0/code/redrock-templates/main/rrtemplate-STAR-F-v0.1.fits', '/global/common/software/desi/perlmutter/desiconda/20230111-2.1.0/code/redrock-templates/main/rrtemplate-STAR-G-v0.1.fits', '/global/common/software/desi/perlmutter/desiconda/20230111-2.1.0/code/redrock-templates/main/rrtemplate-STAR-K-v0.1.fits', '/global/common/software/desi/perlmutter/desiconda/20230111-2.1.0/code/redrock-templates/main/rrtemplate-STAR-M-v0.1.fits', '/global/common/software/desi/perlmutter/desiconda/20230111-2.1.0/code/redrock-templates/main/rrtemplate-STAR-WD-v0.1.fits']

---------------------------------------------------------------------------
OSError                                   Traceback (most recent call last)
Cell In[6], line 1
----> 1 templates_fuji = load_templates_from_header(hdr_fuji)

File ~/code/desihub/redrock/py/redrock/templates.py:363, in load_templates_from_header(hdr, template_dir)
    352 """Return templates matching header keywords
    353
    354 Args:
   (...)
    360 Returns list of Template objects
    361 """
    362 template_filenames = header2templatefiles(hdr, template_dir=template_dir)
--> 363 return load_templates(template_filenames)

File ~/code/desihub/redrock/py/redrock/templates.py:386, in load_templates(template_path)
    384 templates = list()
    385 for filename in template_path:
--> 386     templates.append(Template(filename))
    388 return templates

File ~/code/desihub/redrock/py/redrock/templates.py:49, in Template.__init__(self, filename, spectype, redshifts, wave, flux, subtype, method, igm_model, zscan_galaxy, zscan_qso, zscan_star)
     47         fx = fits.open(xfilename, memmap=False)
     48     else:
---> 49         raise IOError('unable to find '+filename)
     51 hdr = fx['BASIS_VECTORS'].header
     52 if 'VERSION' in hdr:

OSError: unable to find /global/common/software/desi/perlmutter/desiconda/20230111-2.1.0/code/redrock-templates/main/rrtemplate-GALAXY-None-v2.6.fits
moustakas commented 2 months ago

I'm sorry, you did say that this will only work with newly-formatted templates:

This function requires the restructured redrock-templates directory+filenames in /global/cfs/cdirs/desi/users/sjbailey/code/redrock-templates/

Is there really not enough information in the zpix catalogs to read the correct / original Redrock templates used in Fuji and Iron with load_templates_from_header?

moustakas commented 2 months ago

Here's another test, using the current version of the still-in-development Y3 templates:

os.environ['RR_TEMPLATE_DIR']='/pscratch/sd/i/ioannis/Y3-templates/rrtemplates/Y3-0.1'
hh = fitsio.read_header('/pscratch/sd/i/ioannis/Y3-templates/redrock/iron-cumulative-vi-main/Y3-0.1-zscan01/redrock-6-80610.fits', 0)

load_templates_from_header(hh)
Reading templates from ['/pscratch/sd/i/ioannis/Y3-templates/rrtemplates/Y3-0.1/rrtemplate-GALAXY-BGS-v0.5.fits', '/pscratch/sd/i/ioannis/Y3-templates/rrtemplates/Y3-0.1/rrtemplate-GALAXY-ELG-v0.5.fits', '/pscratch/sd/i/ioannis/Y3-templates/rrtemplates/Y3-0.1/rrtemplate-GALAXY-LRG-v0.5.fits', '/pscratch/sd/i/ioannis/Y3-templates/rrtemplates/Y3-0.1/rrtemplate-QSO-HIZ-v1.0.fits', '/pscratch/sd/i/ioannis/Y3-templates/rrtemplates/Y3-0.1/rrtemplate-QSO-LOZ-v1.0.fits', '/pscratch/sd/i/ioannis/Y3-templates/rrtemplates/Y3-0.1/rrtemplate-STAR-A-v0.2.fits', '/pscratch/sd/i/ioannis/Y3-templates/rrtemplates/Y3-0.1/rrtemplate-STAR-B-v0.2.fits', '/pscratch/sd/i/ioannis/Y3-templates/rrtemplates/Y3-0.1/rrtemplate-STAR-CV-v0.1.fits', '/pscratch/sd/i/ioannis/Y3-templates/rrtemplates/Y3-0.1/rrtemplate-STAR-F-v0.2.fits', '/pscratch/sd/i/ioannis/Y3-templates/rrtemplates/Y3-0.1/rrtemplate-STAR-G-v0.2.fits', '/pscratch/sd/i/ioannis/Y3-templates/rrtemplates/Y3-0.1/rrtemplate-STAR-K-v0.2.fits', '/pscratch/sd/i/ioannis/Y3-templates/rrtemplates/Y3-0.1/rrtemplate-STAR-M-v0.2.fits', '/pscratch/sd/i/ioannis/Y3-templates/rrtemplates/Y3-0.1/rrtemplate-STAR-WD-v0.2.fits']
---------------------------------------------------------------------------
OSError                                   Traceback (most recent call last)
Cell In[21], line 1
----> 1 load_templates_from_header(hh)

File ~/code/desihub/redrock/py/redrock/templates.py:363, in load_templates_from_header(hdr, template_dir)
    352 """Return templates matching header keywords
    353
    354 Args:
   (...)
    360 Returns list of Template objects
    361 """
    362 template_filenames = header2templatefiles(hdr, template_dir=template_dir)
--> 363 return load_templates(template_filenames)

File ~/code/desihub/redrock/py/redrock/templates.py:386, in load_templates(template_path)
    384 templates = list()
    385 for filename in template_path:
--> 386     templates.append(Template(filename))
    388 return templates

File ~/code/desihub/redrock/py/redrock/templates.py:49, in Template.__init__(self, filename, spectype, redshifts, wave, flux, subtype, method, igm_model, zscan_galaxy, zscan_qso, zscan_star)
     47         fx = fits.open(xfilename, memmap=False)
     48     else:
---> 49         raise IOError('unable to find '+filename)
     51 hdr = fx['BASIS_VECTORS'].header
     52 if 'VERSION' in hdr:

OSError: unable to find /pscratch/sd/i/ioannis/Y3-templates/rrtemplates/Y3-0.1/rrtemplate-GALAXY-BGS-v0.5.fits

So the expected filenames with this branch are more restrictive than with current main, which would have "just worked" with the diverse names I've been using:

ls -l /pscratch/sd/i/ioannis/Y3-templates/rrtemplates/Y3-0.1/
total 24120
-rw-rw-r-- 1 ioannis desi 3438720 Mar  4 17:48 rrtemplate-galaxy-BGS-NMF-0.5b.fits
-rw-rw-r-- 1 ioannis desi 2802240 Mar  4 15:49 rrtemplate-galaxy-ELG-NMF-0.5b.fits
-rw-rw-r-- 1 ioannis desi 2583360 Mar  4 14:56 rrtemplate-galaxy-LRG-NMF-0.5b.fits
-rw-rw-r-- 1 ioannis desi  336960 Jul 17  2023 rrtemplate-qso-HIZ.fits
-rw-rw-r-- 1 ioannis desi  290880 Jul 17  2023 rrtemplate-qso-LOZ.fits
-rw-rw-r-- 1 ioannis desi 1641600 Mar  3 04:47 rrtemplate-star-A-NMF-0.2.fits
-rw-rw-r-- 1 ioannis desi 1673280 Mar  3 04:46 rrtemplate-star-B-NMF-0.2.fits
-rw-rw-r-- 1 ioannis desi 1923840 Aug 17  2021 rrtemplate-star-CV.fits
-rw-rw-r-- 1 ioannis desi 1653120 Mar  3 04:48 rrtemplate-star-F-NMF-0.2.fits
-rw-rw-r-- 1 ioannis desi 1653120 Mar  3 04:49 rrtemplate-star-G-NMF-0.2.fits
-rw-rw-r-- 1 ioannis desi 1667520 Mar  3 04:51 rrtemplate-star-K-NMF-0.2.fits
-rw-rw-r-- 1 ioannis desi 1748160 Mar  3 04:56 rrtemplate-star-M-NMF-0.2.fits
-rw-rw-r-- 1 ioannis desi 3260160 Mar  2 17:44 rrtemplate-star-WD-NMF-0.2.fits
sbailey commented 2 months ago

backwards compatibility

Is there really not enough information in the zpix catalogs to read the correct / original Redrock templates used in Fuji and Iron with load_templates_from_header?

Yes and no. There is enough information in the pre-existing zpix catalogs to derive the filename in the new template structure (rrtemplate-SPECTYPE-SUBTYPE-vVERSION.fits), but the situation is ambiguous for the old/current template structure because "main" has the templates used for Iron and tag "0.7.2" has the templates used for Fuji, but neither has the complete set of templates that enables swapping back and forth (that's the point of the new template structure). My intention is that once this PR is merged, we would also immediately update redrock-templates to the new structure to support it.

By luck, the templates that changed from Fuji -> Iron also gained a different SUBTYPE and thus a different filename, so we could attempt to fallback to the old filename if the new filename doesn't exist, but that would only work for the semi-contrived case of redrock main, redrock-templates pre-main, and loading a redshift catalog that was processed with whatever that pre-main version of redrock-templates was. I suggest that we just get past this transition period and not worry about this case, but let me know if you disagree.

more flexible filenames like rrtemplate-galaxy-BGS-NMF-0.5b.fits

This will work with a followup PR that would have Redrock also output TEMFILnn keywords with the exact filename to use. This PR will use that keyword if it exists, but I didn't want to start writing that keyword until we had fully switched to the new template filenames so that we wouldn't end up with redrock files withe TEMFILnn keywords that point to template filenames that won't exist.

If you'd like, I could make a branch-of-this-branch that would also write the TEMFILnn keyword so that you could test those assertions with your special named template files, so make sure that load_templates_from_header() really does get enough information to do the right thing.

An unfortunate special case is redrock files that you've already written which used non-standard template filenames. Unfortunately redrock does not currently store enough header information to know what the full name of the template file was. We could consider adding a more flexible option that would load all of the templates it can find, and then in memory filter to the matching SPECTYPE/SUBTYPE/VERSION, rather than pre-deriving the filenames. This would only work if you have been strict about incrementing the VERSION number in your template files so that they really are unique.

moustakas commented 2 months ago

Thanks for the additional details, @sbailey. I propose we proceed with merging this PR and then getting a PR which writes the TEMFILnn header cards merged in as soon as possible after that (no need to do it via this branch) so we can test these changes both in production and in the special-case, template-development work I'm doing.

Let me know if you would like me to do the follow-up PR, but I suspect the necessary changes are fairly minor and right at your fingertips, so my preference would be for you to make those changes. :)