huge structural overhaul for clarity

Author: bbrzycki

setigen/_frame/__init__.py

@@ -0,0 +1,3 @@
+"""
+Private implementation modules for the public ``setigen.frame`` facade.
+"""

setigen/_frame/construction.py

@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+import copy
+from dataclasses import dataclass
+import pathlib
+import time
+
+import numpy as np
+
+from astropy import units as u
+from astropy.time import Time
+
+from blimpy import Waterfall
+
+from .. import waterfall_utils
+from .. import unit_utils
+
+
+@dataclass(frozen=True)
+class _SyntheticFrameSpec:
+    df: float
+    dt: float
+    fch1: float
+    ascending: bool
+    shape: tuple[int, int]
+    data: np.ndarray
+    t_start: float
+    source_name: str
+
+
+@dataclass(frozen=True)
+class _WaterfallLoadSpec:
+    waterfall: Waterfall
+    header: object
+    df: float
+    dt: float
+    fch1: float
+    ascending: bool
+    shape: tuple[int, int]
+    data: np.ndarray
+    t_start: float
+    source_name: str
+
+
+def _is_synthetic_init(*, fchans=None, tchans=None, data=None, kwargs=None):
+    kwargs = {} if kwargs is None else kwargs
+    return None not in [fchans, tchans] or "shape" in kwargs or data is not None
+
+
+def _normalize_synthetic_init(*,
+                              fchans=None,
+                              tchans=None,
+                              df=2.7939677238464355 * u.Hz,
+                              dt=18.253611008 * u.s,
+                              fch1=6 * u.GHz,
+                              ascending=False,
+                              data=None,
+                              kwargs=None):
+    kwargs = {} if kwargs is None else kwargs
+
+    normalized_df = unit_utils.get_value(abs(df), u.Hz)
+    normalized_dt = unit_utils.get_value(dt, u.s)
+    normalized_fch1 = unit_utils.get_value(fch1, u.Hz)
+
+    mjd = kwargs.get("mjd")
+    if mjd is not None:
+        t_start = Time(mjd, format="mjd").unix
+    else:
+        t_start = kwargs.get("t_start", time.time())
+    source_name = kwargs.get("source_name", "Synthetic")
+
+    if "shape" in kwargs:
+        shape = tuple(kwargs["shape"])
+    elif data is not None:
+        shape = data.shape
+    else:
+        shape = (int(unit_utils.get_value(tchans, u.pixel)),
+                 int(unit_utils.get_value(fchans, u.pixel)))
+
+    if data is not None:
+        if data.shape != shape:
+            raise ValueError(f"Data shape {data.shape} does not match frame shape {shape}.")
+        frame_data = np.copy(data)
+    else:
+        frame_data = np.zeros(shape)
+
+    return _SyntheticFrameSpec(df=normalized_df,
+                               dt=normalized_dt,
+                               fch1=normalized_fch1,
+                               ascending=ascending,
+                               shape=shape,
+                               data=frame_data,
+                               t_start=t_start,
+                               source_name=source_name)
+
+
+def _normalize_waterfall_init(*, waterfall, kwargs=None):
+    kwargs = {} if kwargs is None else kwargs
+
+    if isinstance(waterfall, pathlib.PurePath):
+        waterfall = str(waterfall)
+    if isinstance(waterfall, str):
+        waterfall = Waterfall(waterfall,
+                              f_start=kwargs.get("f_start"),
+                              f_stop=kwargs.get("f_stop"))
+    elif not isinstance(waterfall, Waterfall):
+        raise FileNotFoundError(f"Unsupported data type: {type(waterfall)}")
+
+    header = waterfall.header
+    tchans, _, fchans = waterfall.container.selection_shape
+    shape = (tchans, fchans)
+
+    df = unit_utils.cast_value(abs(header["foff"]), u.MHz).to(u.Hz).value
+    dt = unit_utils.get_value(header["tsamp"], u.s)
+
+    ascending = header["foff"] > 0
+    if ascending:
+        fch1 = waterfall.container.f_start
+    else:
+        fch1 = waterfall.container.f_stop
+    fch1 = unit_utils.cast_value(fch1, u.MHz).to(u.Hz).value
+
+    t_start = Time(header["tstart"], format="mjd").unix
+    source_name = header["source_name"]
+
+    data = waterfall_utils.get_data(waterfall)
+    if not ascending:
+        data = data[:, ::-1]
+
+    return _WaterfallLoadSpec(waterfall=waterfall,
+                              header=header,
+                              df=df,
+                              dt=dt,
+                              fch1=fch1,
+                              ascending=ascending,
+                              shape=shape,
+                              data=data,
+                              t_start=t_start,
+                              source_name=source_name)
+
+
+def _normalize_frame_init(*,
+                          waterfall=None,
+                          fchans=None,
+                          tchans=None,
+                          df=2.7939677238464355 * u.Hz,
+                          dt=18.253611008 * u.s,
+                          fch1=6 * u.GHz,
+                          ascending=False,
+                          data=None,
+                          kwargs=None):
+    kwargs = {} if kwargs is None else kwargs
+    if _is_synthetic_init(fchans=fchans,
+                          tchans=tchans,
+                          data=data,
+                          kwargs=kwargs):
+        return _normalize_synthetic_init(fchans=fchans,
+                                         tchans=tchans,
+                                         df=df,
+                                         dt=dt,
+                                         fch1=fch1,
+                                         ascending=ascending,
+                                         data=data,
+                                         kwargs=kwargs)
+    if waterfall is not None:
+        return _normalize_waterfall_init(waterfall=waterfall, kwargs=kwargs)
+    raise ValueError("Frame must be provided dimensions or an existing filterbank file.")
+
+
+def _initialize_frame_from_spec(frame, spec):
+    frame.df = spec.df
+    frame.dt = spec.dt
+    frame.fch1 = spec.fch1
+    frame.ascending = spec.ascending
+    frame.t_start = spec.t_start
+    frame.source_name = spec.source_name
+    frame.shape = spec.shape
+    frame.tchans, frame.fchans = spec.shape
+    frame.data = spec.data
+
+    if isinstance(spec, _WaterfallLoadSpec):
+        frame.waterfall = spec.waterfall
+        frame.header = spec.header
+    else:
+        frame.waterfall = None
+        frame.header = None
+
+
+def _attach_loaded_waterfall(frame, waterfall):
+    if waterfall is None:
+        return
+    try:
+        del waterfall.container.h5
+    except AttributeError:
+        pass
+    frame.waterfall = copy.deepcopy(waterfall)

setigen/_frame/io.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+import pathlib
+import pickle
+
+import numpy as np
+
+from blimpy import Waterfall
+from blimpy.io import sigproc
+
+
+def _create_synthetic_waterfall(frame, *, max_load=1):
+    path = pathlib.Path(__file__).resolve().parents[1] / "assets" / "sample.fil"
+    waterfall = Waterfall(str(path), max_load=max_load)
+    waterfall.header["source_name"] = frame.source_name
+    waterfall.header["rawdatafile"] = "Synthetic"
+
+    container_attr = {
+        "t_begin": 0,
+        "t_end": frame.tchans,
+        "file_size_bytes": frame.tchans * frame.fchans * waterfall.header["nbits"] / 8,
+        "n_channels_in_file": frame.fchans,
+        "n_ints_in_file": frame.tchans,
+        "file_shape": (frame.tchans, 1, frame.fchans),
+        "f_end": frame.fmax * 1e-6,
+        "f_begin": frame.fmin * 1e-6,
+        "f_stop": frame.fmax * 1e-6,
+        "f_start": frame.fmin * 1e-6,
+        "t_start": 0,
+        "t_stop": frame.tchans,
+        "selection_shape": (frame.tchans, 1, frame.fchans),
+        "chan_start_idx": 0,
+        "chan_stop_idx": frame.fchans,
+    }
+    for key, value in container_attr.items():
+        setattr(waterfall.container, key, value)
+
+    wat_attr = {
+        "n_channels_in_file": frame.fchans,
+        "n_ints_in_file": frame.tchans,
+        "file_shape": (frame.tchans, 1, frame.fchans),
+        "file_size_bytes": frame.tchans * frame.fchans * waterfall.header["nbits"] / 8,
+        "selection_shape": (frame.tchans, 1, frame.fchans),
+    }
+    for key, value in wat_attr.items():
+        setattr(waterfall, key, value)
+
+    return waterfall
+
+
+def _update_waterfall(frame, *, filename=None, max_load=1):
+    if frame.waterfall is None:
+        frame.waterfall = _create_synthetic_waterfall(frame, max_load=max_load)
+
+    frame.waterfall.data = frame.data[:, np.newaxis, :]
+    if not frame.ascending:
+        frame.waterfall.data = frame.waterfall.data[:, :, ::-1]
+
+    header_attr = {
+        "tsamp": frame.dt,
+        "tstart": frame.mjd,
+        "nchans": frame.fchans,
+        "fch1": frame.fch1 * 1e-6,
+    }
+    if frame.ascending:
+        header_attr["foff"] = frame.df * 1e-6
+    else:
+        header_attr["foff"] = frame.df * -1e-6
+    frame.waterfall.header.update(header_attr)
+    frame.waterfall.file_header.update(header_attr)
+
+    if filename is not None:
+        frame.waterfall.container.filename = str(pathlib.Path(filename).resolve())
+    frame.waterfall.container.idx_data = len(sigproc.generate_sigproc_header(frame.waterfall))
+
+
+def _encode_bytestrings(frame):
+    for key in ["source_name", "rawdatafile"]:
+        if key in frame.waterfall.header and not isinstance(frame.waterfall.header[key], bytes):
+            frame.waterfall.header[key] = frame.waterfall.header[key].encode()
+
+
+def _decode_bytestrings(frame):
+    for key in ["source_name", "rawdatafile"]:
+        if key in frame.waterfall.header and isinstance(frame.waterfall.header[key], bytes):
+            frame.waterfall.header[key] = frame.waterfall.header[key].decode()
+
+
+def _get_waterfall(frame):
+    _update_waterfall(frame)
+    return frame.waterfall
+
+
+def _check_waterfall(frame):
+    if frame.waterfall is None:
+        return None
+    return _get_waterfall(frame)
+
+
+def _save_fil(frame, filename, *, max_load=1):
+    _update_waterfall(frame, filename=filename, max_load=max_load)
+    _encode_bytestrings(frame)
+    frame.waterfall.write_to_fil(filename)
+    _decode_bytestrings(frame)
+
+
+def _save_hdf5(frame, filename, *, max_load=1):
+    _update_waterfall(frame, filename=filename, max_load=max_load)
+    _encode_bytestrings(frame)
+    frame.waterfall.write_to_hdf5(filename)
+    _decode_bytestrings(frame)
+
+
+def _save_npy(frame, filename):
+    np.save(filename, frame.data)
+
+
+def _load_npy(frame, filename):
+    frame.data = np.load(filename)
+
+
+def _save_pickle(frame, filename):
+    with open(filename, "wb") as f:
+        pickle.dump(frame, f)
+
+
+def _load_pickle(filename):
+    with open(filename, "rb") as f:
+        return pickle.load(f)