Merge pull request #146 from sdss/improve_hdrfix

improves runtime of apply_hdrfix

.gitignore

@@ -32,7 +32,7 @@ distribute-*.tar.gz
 *~
 
 # Test directories and files
-data*
+#data*
 lib
 *.tar.gz*
 test-*.fits

python/lvmdrp/utils/hdrfix.py

@@ -1,9 +1,10 @@
 #!/usr/bin/env python
 # encoding: utf-8
 
+import functools
 import os
 import pathlib
-import re
+import fnmatch
 import yaml
 
 import pandas as pd
@@ -38,6 +39,7 @@ def get_hdrfix_path(mjd: int) -> str:
         raise ValueError('LVMCORE_DIR environment variable not found.  Please set up the repo.')
 
 
+@functools.lru_cache(maxsize=256)
 def read_hdrfix_file(mjd: int) -> pd.DataFrame:
     """ Read a header fix file
 
@@ -173,18 +175,14 @@ def apply_hdrfix(mjd: int, camera: str = None, expnum: int = None,
     if fix is None or fix.empty:
         return hdr
 
-    # find matching files
-    for fileroot, key, val in zip(fix['fileroot'], fix['keyword'], fix['value']):
-        root = f'{mjd}/{fileroot}{"" if fileroot.endswith("*") else "*"}'
-        files = pathlib.Path(os.getenv('LVM_DATA_S')).rglob(root)
+    # Create the current file string
+    hemi = 's' if hdr['OBSERVAT'] == 'LCO' else 'n'
+    current_file = f'sdR-{hemi}-{camera}-{expnum:0>8}'
 
-        pattern = re.compile(f'{camera}-{expnum:0>8}')
-        sub = filter(pattern.search, map(str, files))
-
-        # apply the header fix
-        for file in sub:
-            stem = pathlib.Path(file).stem
-            log.info(f'Applying header fix on {stem} for key: {key}, value: {val}.')
-            hdr[key] = val
+    # Apply the header fixes
+    for _, row in fix.iterrows():
+        if fnmatch.fnmatch(current_file, row['fileroot']):
+            hdr[row['keyword']] = row['value']
+            log.info(f'Applying header fix on {current_file} for key: {row["keyword"]}, value: {row["value"]}.')
 
     return hdr

tests/data/lvmHdrFix-60255.yaml

@@ -0,0 +1,29 @@
+schema:
+- description: the raw frame file root with * as wildcard
+  dtype: str
+  name: fileroot
+- description: the name of the header keyword to fix
+  dtype: str
+  name: keyword
+- description: the value of the header keyword to update
+  dtype: str
+  name: value
+fixes:
+- fileroot: sdR-*-*-00007334
+  keyword: SKYENAME
+  value: grid125
+- fileroot: sdR-*-b*-00007334
+  keyword: SKYWNAME
+  value: WHAM_south_08
+- fileroot: sdR-*-*2-0000732[4-6]
+  keyword: SKYENAME
+  value: grid103
+- fileroot: sdR-*-*2-0000732[4-6]
+  keyword: SKYWNAME
+  value: WHAM_south_02
+- fileroot: sdR-*-r1-00007329
+  keyword: SKYENAME
+  value: grid125
+- fileroot: sdR-*-r1-00007329
+  keyword: SKYWNAME
+  value: WHAM_south_02

tests/utils/test_hdrfix.py

@@ -0,0 +1,82 @@
+
+
+import pathlib
+import pytest
+
+from lvmdrp.utils.hdrfix import read_hdrfix_file, apply_hdrfix, write_hdrfix_file
+
+
+# make fake hdr
+header = {
+    'MJD': 60255,
+    'CCD': 'b1',
+    'EXPOSURE': 7230,
+    'TILE_ID': 12345,
+    'OBSERVAT': 'LCO',
+    'SKYWNAME': 'old_west_name',
+    'SKYENAME': 'old_east_name',
+}
+
+
+@pytest.fixture()
+def mock_path(mocker):
+    """ fixture to mock the return path of the hdrfix file """
+    def _mock(mjd):
+        path = pathlib.Path(__file__).parent.parent / 'data' / f'lvmHdrFix-{mjd}.yaml'
+        mocker.patch('lvmdrp.utils.hdrfix.get_hdrfix_path', return_value=path)
+    return _mock
+
+
+def test_read_hdrfix_file(mock_path):
+    """ test we can read the hdrfix file """
+    mock_path(60255)
+    df = read_hdrfix_file(60255)
+    assert len(df) == 6
+    assert df['fileroot'].iloc[0] == 'sdR-*-*-00007334'
+    assert df['keyword'].iloc[0] == 'SKYENAME'
+    assert df['value'].iloc[0] == 'grid125'
+
+
+def test_read_hdrfix_file_no_file(mock_path):
+    """ test when the hdrfix file does not exist """
+    mock_path(60256)
+    df = read_hdrfix_file(60256)
+    assert df is None
+
+
+@pytest.mark.parametrize('input, exp',
+                         [(('b1', 7334), {'SKYWNAME': 'WHAM_south_08', 'SKYENAME': 'grid125'}),
+                          (('r1', 7334), {'SKYWNAME': 'old_west_name', 'SKYENAME': 'grid125'}),
+                          (('z1', 7330), {'SKYWNAME': 'old_west_name', 'SKYENAME': 'old_east_name'}),
+                          (('z2', 7324), {'SKYWNAME': 'WHAM_south_02', 'SKYENAME': 'grid103'}),
+                          (('b2', 7326), {'SKYWNAME': 'WHAM_south_02', 'SKYENAME': 'grid103'}),
+                          (('r3', 7326), {'SKYWNAME': 'old_west_name', 'SKYENAME': 'old_east_name'}),
+                          (('r1', 7329), {'SKYWNAME': 'WHAM_south_02', 'SKYENAME': 'grid125'}),
+                          (('z2', 7329), {'SKYWNAME': 'old_west_name', 'SKYENAME': 'old_east_name'}),
+                          ]
+                         )
+def test_apply_hdrfix(mock_path, input, exp):
+    """ test we can apply the hdrfix """
+    mock_path(60255)
+    hdr = header.copy()
+    hdr.update(dict(zip(['CCD', 'EXPOSURE'], input)))
+    new_hdr = apply_hdrfix(60255, hdr=hdr.copy())
+
+    assert new_hdr['SKYWNAME'] == exp['SKYWNAME']
+    assert new_hdr['SKYENAME'] == exp['SKYENAME']
+
+
+def test_write_hdrfix(tmp_path, mocker):
+    """ test we can write a new hdrfix file """
+    path = tmp_path / 'data' / 'lvmHdrFix-60257.yaml'
+    mocker.patch('lvmdrp.utils.hdrfix.get_hdrfix_path', return_value=path)
+
+    write_hdrfix_file(60257, 'sdR-*-*-00007334', 'SKYENAME', 'grid125')
+
+    read_hdrfix_file.cache_clear()
+    df = read_hdrfix_file(60257)
+    assert path.exists()
+    assert len(df) == 1
+    assert df['fileroot'].iloc[0] == 'sdR-*-*-00007334'
+    assert df['keyword'].iloc[0] == 'SKYENAME'
+    assert df['value'].iloc[0] == 'grid125'