Refactor frame methods to improve waterfall handling and integrate observational context preservation in derived frames

Author: bbrzycki

setigen/frame.py

@@ -29,9 +29,10 @@
     _generate_sampled_noise,
 )
 from ._frame.io import (
-    _decode_bytestrings,
-    _encode_bytestrings,
-    _update_waterfall,
+    _check_waterfall,
+    _get_waterfall,
+    _save_fil,
+    _save_hdf5,
 )
 from ._frame.signal import (
     _finalize_signal,
@@ -767,8 +768,7 @@ def get_waterfall(self) -> Any:
         Returns:
             Waterfall representation of the frame.
         """
-        _update_waterfall(self)
-        return self.waterfall
+        return _get_waterfall(self)
 
     def check_waterfall(self) -> Any:
         """Return the updated attached waterfall when one exists.
@@ -777,9 +777,7 @@ def check_waterfall(self) -> Any:
             Updated waterfall object or `None` when the frame has no attached
             waterfall.
         """
-        if self.waterfall is None:
-            return None
-        return self.get_waterfall()
+        return _check_waterfall(self)
 
     def save_fil(self, filename: PathLike, max_load: int = 1) -> None:
         """Save frame data as a SIGPROC filterbank file.
@@ -788,10 +786,7 @@ def save_fil(self, filename: PathLike, max_load: int = 1) -> None:
             filename: Output `.fil` path.
             max_load: Maximum load parameter for a lazily created waterfall.
         """
-        _update_waterfall(self, filename=filename, max_load=max_load)
-        _encode_bytestrings(self)
-        self.waterfall.write_to_fil(filename)
-        _decode_bytestrings(self)
+        _save_fil(self, filename, max_load=max_load)
 
     def save_hdf5(self, filename: PathLike, max_load: int = 1) -> None:
         """Save frame data as an HDF5 waterfall file.
@@ -800,10 +795,7 @@ def save_hdf5(self, filename: PathLike, max_load: int = 1) -> None:
             filename: Output `.h5` path.
             max_load: Maximum load parameter for a lazily created waterfall.
         """
-        _update_waterfall(self, filename=filename, max_load=max_load)
-        _encode_bytestrings(self)
-        self.waterfall.write_to_hdf5(filename)
-        _decode_bytestrings(self)
+        _save_hdf5(self, filename, max_load=max_load)
 
     def save_h5(self, filename: PathLike, max_load: int = 1) -> None:
         """Save frame data as an HDF5 waterfall file.

setigen/integrate.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import copy
 from enum import Enum
 from typing import Any
 
@@ -54,6 +55,27 @@ def _resolve_integration_mode(mode: IntegrationMode | str) -> IntegrationMode:
     return IntegrationMode.MEAN
 
 
+def _copy_frame_context(source: Any, target: Any) -> None:
+    """Copy non-shape observational context between frame-like objects.
+
+    Args:
+        source: Original frame-like object.
+        target: Derived frame-like object to update.
+    """
+    metadata = copy.deepcopy(getattr(source, "metadata", {}))
+    try:
+        for key, value in source.get_params().items():
+            if metadata.get(key) == value:
+                metadata.pop(key)
+    except AttributeError:
+        pass
+    if metadata:
+        target.add_metadata(metadata)
+
+    if hasattr(source, "header"):
+        target.header = copy.deepcopy(source.header)
+
+
 def integrate(
     fr: Any,
     axis: IntegrationAxis | str | int = 't',
@@ -101,15 +123,20 @@ def integrate(
                                 fch1=fr.fmid,
                                 ascending=fr.ascending,
                                 data=data,
-                                seed=fr.rng)
+                                seed=fr.rng,
+                                t_start=fr.t_start,
+                                source_name=fr.source_name)
         else:
             # Spectrum
             new_fr = Spectrum(df=fr.df,
                               dt=fr.dt * fr.tchans,
                               fch1=fr.fch1,
                               ascending=fr.ascending,
                               data=data,
-                              seed=fr.rng)
+                              seed=fr.rng,
+                              t_start=fr.t_start,
+                              source_name=fr.source_name)
+        _copy_frame_context(fr, new_fr)
         return new_fr
     else:
         return data.flatten()

setigen/split_utils.py

@@ -8,6 +8,7 @@
 from typing import Iterator
 
 from ._typing import PathLike
+from ._frame.io import _close_waterfall_handles
 
 
 def split_waterfall_generator(
@@ -32,10 +33,13 @@ def split_waterfall_generator(
     """
 
     info_wf = Waterfall(waterfall_fn, load_data=False)
-    fch1 = info_wf.header['fch1']
-    nchans = info_wf.header['nchans']
-    df = info_wf.header['foff']
-    tchans_tot = info_wf.container.selection_shape[0]
+    try:
+        fch1 = info_wf.header['fch1']
+        nchans = info_wf.header['nchans']
+        df = info_wf.header['foff']
+        tchans_tot = info_wf.container.selection_shape[0]
+    finally:
+        _close_waterfall_handles(info_wf)
 
     if f_shift is None:
         f_shift = fchans
@@ -58,7 +62,10 @@ def split_waterfall_generator(
                               t_start=0,
                               t_stop=tchans)
 
-        yield waterfall
+        try:
+            yield waterfall
+        finally:
+            _close_waterfall_handles(waterfall)
 
         f_start += f_shift * df
         f_stop += f_shift * df

setigen/waterfall_utils.py

@@ -5,6 +5,7 @@
 from blimpy import Waterfall
 
 from ._typing import PathLike
+from ._frame.io import _close_waterfall_handles
 
 
 def max_freq(waterfall: PathLike | Waterfall) -> float:
@@ -44,15 +45,20 @@ def get_data(waterfall: PathLike | Waterfall, db: bool = False) -> np.ndarray:
     Raises:
         ValueError: If the waterfall input type is unsupported.
     """
-    if isinstance(waterfall, (str, PurePath)):
+    owns_waterfall = isinstance(waterfall, (str, PurePath))
+    if owns_waterfall:
         waterfall = Waterfall(waterfall)
     elif not isinstance(waterfall, Waterfall):
         raise ValueError('Invalid data file!')
 
-    if db:
-        return 10 * np.log10(waterfall.data[:, 0, :])
-
-    return waterfall.data[:, 0, :]
+    try:
+        data = waterfall.data[:, 0, :]
+        if db:
+            data = 10 * np.log10(data)
+        return np.array(data, copy=owns_waterfall)
+    finally:
+        if owns_waterfall:
+            _close_waterfall_handles(waterfall)
 
 
 def get_fs(waterfall: PathLike | Waterfall) -> np.ndarray:
@@ -67,16 +73,20 @@ def get_fs(waterfall: PathLike | Waterfall) -> np.ndarray:
     Raises:
         ValueError: If the waterfall input type is unsupported.
     """
-    if isinstance(waterfall, (str, PurePath)):
+    owns_waterfall = isinstance(waterfall, (str, PurePath))
+    if owns_waterfall:
         waterfall = Waterfall(waterfall, load_data=False)
     elif not isinstance(waterfall, Waterfall):
         raise ValueError('Invalid data file!')
 
-    fch1 = waterfall.header['fch1']
-    df = waterfall.header['foff']
-    fchans = waterfall.header['nchans']
-
-    return fch1 + np.arange(fchans) * df
+    try:
+        fch1 = waterfall.header['fch1']
+        df = waterfall.header['foff']
+        fchans = waterfall.header['nchans']
+        return fch1 + np.arange(fchans) * df
+    finally:
+        if owns_waterfall:
+            _close_waterfall_handles(waterfall)
 
 
 def get_ts(waterfall: PathLike | Waterfall) -> np.ndarray:
@@ -91,12 +101,16 @@ def get_ts(waterfall: PathLike | Waterfall) -> np.ndarray:
     Raises:
         ValueError: If the waterfall input type is unsupported.
     """
-    if isinstance(waterfall, (str, PurePath)):
+    owns_waterfall = isinstance(waterfall, (str, PurePath))
+    if owns_waterfall:
         waterfall = Waterfall(waterfall, load_data=False)
     elif not isinstance(waterfall, Waterfall):
         raise ValueError('Invalid data file!')
 
-    tsamp = waterfall.header['tsamp']
-    tchans = waterfall.container.selection_shape[0]
-
-    return np.arange(tchans) * tsamp
+    try:
+        tsamp = waterfall.header['tsamp']
+        tchans = waterfall.container.selection_shape[0]
+        return np.arange(tchans) * tsamp
+    finally:
+        if owns_waterfall:
+            _close_waterfall_handles(waterfall)

tests/test_frame_creation.py

@@ -129,6 +129,27 @@ def test_axis_semantics():
     assert frame.get_drift_rate(0, 3, reference="centers") == pytest.approx(7.5)
 
 
+def test_integrated_frames_preserve_observational_context():
+    frame = stg.Frame(shape=(4, 8),
+                      t_start=123456,
+                      source_name="Target",
+                      data=np.ones((4, 8)))
+    frame.header = {"rawdatafile": "original.raw", "source_name": "Target"}
+    frame.add_metadata({"drift_rate": 1.25})
+
+    spectrum = stg.spectrum(frame, mode="sum")
+    timeseries = stg.timeseries(frame, mode="sum")
+
+    for derived in [spectrum, timeseries]:
+        assert derived.t_start == frame.t_start
+        assert derived.source_name == frame.source_name
+        assert derived.header == frame.header
+        assert derived.header is not frame.header
+        assert derived.metadata["drift_rate"] == frame.metadata["drift_rate"]
+        assert derived.metadata["fchans"] == derived.fchans
+        assert derived.metadata["tchans"] == derived.tchans
+
+
 def test_h5_ingestion_does_not_retain_live_waterfall(tmp_path):
     frame = stg.Frame(shape=(4, 8), seed=0)
     h5_path = tmp_path / "frame.h5"