Merge pull request #260 from sot/fix-quat-comp

Fix computed Comp_Quat class and close pickle files

cheta/derived/comps.py

@@ -22,6 +22,7 @@
 
 import functools
 import re
+import warnings
 
 import astropy.table as tbl
 import numpy as np
@@ -428,7 +429,11 @@ def get_msid_attrs(self, tstart: float, tstop: float, msid: str, msid_args: tupl
             q4 = np.sqrt((1.0 - q1**2 - q2**2 - q3**2).clip(0.0))
 
         q = np.array([q1, q2, q3, q4]).transpose()
-        quat = Quat(q=normalize(q))
+        with warnings.catch_warnings():
+            warnings.filterwarnings(
+                "ignore", message="Normalizing quaternion with zero norm"
+            )
+            quat = Quat(q=normalize(q))
         bads = np.zeros_like(q1, dtype=bool)
         for msid in msids:
             bads |= dat[msid].bads

cheta/fetch.py

@@ -295,7 +295,8 @@ def load_msid_names(all_msid_names_files):
     all_colnames = dict()
     for k, msid_names_file in all_msid_names_files.items():
         try:
-            all_colnames[k] = pickle.load(open(os.path.join(*msid_names_file), "rb"))
+            with open(os.path.join(*msid_names_file), "rb") as fh:
+                all_colnames[k] = pickle.load(fh)
         except IOError:
             pass
     return all_colnames
@@ -719,7 +720,7 @@ def _get_comp_data(self, comp_cls):
         self.colnames = [
             attr
             for attr, val in attrs.items()
-            if (isinstance(val, np.ndarray) and len(val) == len(attrs["times"]))
+            if (hasattr(val, "shape") and len(val) == len(attrs["times"]))
         ]
 
         # Apply attributes to self

cheta/tests/test_comps.py

@@ -2,6 +2,8 @@
 
 """Test that computed MSIDs work as expected."""
 
+import warnings
+
 import astropy.units as u
 import numpy as np
 import pytest
@@ -269,6 +271,28 @@ def test_quat_comp(msid, maude, offset):
     assert isinstance(datq.vals, Quat)
 
 
+def test_quat_comp_bad_times():
+    """Test bad time data on 2024:264. All four quats have zero value and are bad.
+
+    The bad sample times are ['2024:064:09:27:02.652' '2024:064:09:27:03.677'].
+    """
+    start = "2024:064:09:26:00"
+    stop = "2024:064:09:28:00"
+    # Assert no warnings despite quat with zero normalization. The zero-norm samples are
+    # marked bad.
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")  # Assert no warnings
+        dat = fetch_eng.MSID("quat_aoattqt", start, stop)
+
+    assert np.count_nonzero(dat.bads) == 2
+    assert len(dat.vals) == len(dat.times)
+
+    dat2 = fetch_eng.Msid("quat_aoattqt", start, stop)
+    assert dat2.bads is None  # After Msid filtering
+    assert len(dat2.vals) == len(dat2.times)
+    assert len(dat2.vals) == len(dat.vals) - 2
+
+
 def test_pitch_comp():
     """Test pitch_comp during a time with NPNT, NMAN, NSUN and Safe Sun"""
     start = "2022:293"

cheta/units.py

@@ -61,7 +61,8 @@ def emit(self, record):
 
 units = {}
 units["system"] = "cxc"
-units["cxc"] = pickle.load(open(os.path.join(module_dir, "units_cxc.pkl"), "rb"))
+with open(os.path.join(module_dir, "units_cxc.pkl"), "rb") as fh:
+    units["cxc"] = pickle.load(fh)
 
 
 # Equivalent unit descriptors used in 'eng' and 'cxc' units
@@ -222,7 +223,8 @@ def load_units(unit_system):
 
     if unit_system not in units:
         filename = os.path.join(module_dir, "units_{0}.pkl".format(unit_system))
-        units[unit_system] = pickle.load(open(filename, "rb"))
+        with open(filename, "rb") as fh:
+            units[unit_system] = pickle.load(fh)
 
 
 def set_units(unit_system):