fixed apps and added psrfits tests

README.md

@@ -2,11 +2,14 @@
 
 [![GitHub CI](https://github.com/FRBs/sigpyproc3/workflows/GitHub%20CI/badge.svg)](https://github.com/FRBs/sigpyproc3/actions)
 [![Docs](https://readthedocs.org/projects/sigpyproc3/badge/?version=latest)](https://sigpyproc3.readthedocs.io/en/latest/?badge=latest)
-[![codecov](https://codecov.io/gh/FRBs/sigpyproc3/branch/master/graph/badge.svg)](https://codecov.io/gh/FRBs/sigpyproc3)
+[![codecov](https://codecov.io/gh/FRBs/sigpyproc3/branch/main/graph/badge.svg)](https://codecov.io/gh/FRBs/sigpyproc3)
 [![License](https://img.shields.io/github/license/FRBs/sigpyproc3)](https://github.com/FRBs/sigpyproc3/blob/main/LICENSE)
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
 
-`sigpyproc` is a pulsar and FRB data analysis library for python. It provides an OOP approach to pulsar data handling through the use of objects representing different data types (e.g. [SIGPROC filterbank](http://sigproc.sourceforge.net), time-series, fourier-series, etc.). As pulsar data processing is often time critical, speed is maintained using
+`sigpyproc` is a pulsar and FRB data analysis library for python. It provides an OOP approach to pulsar data handling through the use of
+objects representing different data types (e.g. [SIGPROC filterbank](http://sigproc.sourceforge.net),
+[PSRFITS](https://www.atnf.csiro.au/research/pulsar/psrfits_definition/Psrfits.html), time-series, fourier-series, etc.).
+As pulsar data processing is often time critical, speed is maintained using
 the excellent [numba](https://numba.pydata.org) library.
 
 ## Installation
@@ -29,8 +32,10 @@ branch of this repo, or install the last released version 0.5.5.
 ## Usage
 
 ```python
-from sigpyproc.readers import FilReader
+from sigpyproc.readers import FilReader, PFITSReader
+
 fil = FilReader("tutorial.fil")
+fits = PFITSReader("tutorial.fits")
 ```
 
 Check out the tutorials and API docs on [the docs page](https://sigpyproc3.readthedocs.io) for example usage and more info.

docs/index.rst

@@ -25,7 +25,7 @@ environment, with a simple plug-and-play system with new modules and extensions.
 .. image:: https://readthedocs.org/projects/sigpyproc3/badge/?version=latest
   :target: https://sigpyproc3.readthedocs.io/en/latest/?badge=latest
   :alt: Documentation Status
-.. image:: https://codecov.io/gh/FRBs/sigpyproc3/branch/master/graph/badge.svg
+.. image:: https://codecov.io/gh/FRBs/sigpyproc3/branch/main/graph/badge.svg
   :target: https://codecov.io/gh/FRBs/sigpyproc3
 .. image:: https://img.shields.io/github/license/FRBs/sigpyproc3
   :target: https://github.com/FRBs/sigpyproc3/blob/main/LICENSE

sigpyproc/apps/spp_clean.py

@@ -21,12 +21,21 @@
     "-g", "--gulp", type=int, default=16384, help="Number of samples to read at once"
 )
 @click.option(
-    "-o", "--outfile", type=click.Path(exists=False), default=None, help="Output filename"
+    "-o", "--outfile", type=click.Path(exists=False), help="Output masked filterbank file"
 )
-def main(filfile: str, method: str, threshold: float, outfile: str, gulp: int) -> None:
+@click.option(
+    "--save/--no-save", default=True, help="Save the mask information to a file"
+)
+def main(
+    filfile: str, method: str, threshold: float, outfile: str, gulp: int, save: bool
+) -> None:
     """Clean RFI from filterbank data."""
     fil = FilReader(filfile)
-    fil.clean_rfi(method=method, threshold=threshold, outfile=outfile, gulp=gulp)
+    _out_file, rfimask = fil.clean_rfi(
+        method=method, threshold=threshold, filename=outfile, gulp=gulp
+    )
+    if save:
+        rfimask.to_file()
 
 
 if __name__ == "__main__":

sigpyproc/apps/spp_extract.py

@@ -1,28 +1,56 @@
 import click
-
 from sigpyproc.readers import FilReader
 
 
-@click.command(
-    context_settings=dict(help_option_names=["-h", "--help"], show_default=True)
-)
+@click.group(context_settings=dict(help_option_names=["-h", "--help"], show_default=True))
+def main() -> None:
+    pass
+
+
+@main.command()
 @click.argument("filfile", type=click.Path(exists=True))
+@click.option("-s", "--start", type=int, required=True, help="Start time sample")
 @click.option(
-    "-s", "--start", type=int, default=0, help="Start time sample"
-)
-@click.option(
-    "-n", "--nsamps", type=int, help="Number of time samples to extract"
+    "-n", "--nsamps", type=int, required=True, help="Number of time samples to extract"
 )
 @click.option(
     "-g", "--gulp", type=int, default=16384, help="Number of samples to read at once"
 )
 @click.option(
     "-o", "--outfile", type=click.Path(exists=False), default=None, help="Output filename"
 )
-def main(filfile: str, start: int, nsamps: int, gulp: int, outfile: str) -> None:
+def samples(filfile: str, start: int, nsamps: int, gulp: int, outfile: str) -> None:
     """Extract time samples from filterbank data."""
     fil = FilReader(filfile)
-    fil.extract_samps(start=start, nsamps=nsamps, outfile=outfile, gulp=gulp)
+    fil.extract_samps(start=start, nsamps=nsamps, filename=outfile, gulp=gulp)
+
+
+@main.command()
+@click.argument("filfile", type=click.Path(exists=True))
+@click.option(
+    "-c",
+    "--chans",
+    type=click.IntRange(0, None),
+    multiple=True,
+    help="Channels to extract",
+)
+def channels(filfile: str, chans: int) -> None:
+    """Extract frequency channels from filterbank data."""
+    fil = FilReader(filfile)
+    fil.extract_chans(chans=chans)
+
+
+@main.command()
+@click.argument("filfile", type=click.Path(exists=True))
+@click.option("-s", "--chanstart", type=int, required=True, help="Start channel")
+@click.option(
+    "-n", "--nchans", type=int, required=True, help="Number of channels to extract"
+)
+@click.option("-c", "--chanpersub", type=int, help="Number of channels in each sub-band")
+def bands(filfile: str, chanstart: int, nchans: int, chanpersub: int) -> None:
+    """Extract frequency bands from filterbank data."""
+    fil = FilReader(filfile)
+    fil.extract_bands(chanstart=chanstart, nchans=nchans, chanpersub=chanpersub)
 
 
 if __name__ == "__main__":

sigpyproc/apps/spp_header.py

@@ -6,23 +6,31 @@
 
 
 @click.group(context_settings=dict(help_option_names=["-h", "--help"], show_default=True))
+def main() -> None:
+    pass
+
+
+@main.command()
 @click.argument("filfile", type=click.Path(exists=True))
-def main(filfile: str) -> None:
-    """Print or edit the header of a filterbank file."""
+def print(filfile: str) -> None:
+    """Print the header information."""
     header = Header.from_sigproc(filfile)
     click.echo(header.to_string())
 
 
 @main.command()
+@click.argument("filfile", type=click.Path(exists=True))
 @click.option(
     "-k", "--key", type=str, help="A header key to read (e.g. telescope, fch1, nsamples)"
 )
 def get(filfile: str, key: str) -> None:
+    """Get the value of a header key."""
     header = Header.from_sigproc(filfile)
     click.echo(f"{key} = {getattr(header, key)}")
 
 
 @main.command()
+@click.argument("filfile", type=click.Path(exists=True))
 @click.option(
     "-i",
     "--item",
@@ -31,6 +39,7 @@ def get(filfile: str, key: str) -> None:
     help="(key, value) to update in header",
 )
 def update(filfile: str, item: tuple[str, str]) -> None:
+    """Update a header key."""
     key, value = item
     edit_header(filfile, key, value)
 

sigpyproc/base.py

@@ -349,7 +349,7 @@ def apply_channel_mask(
             filename = f"{self.header.basename}_masked.fil"
 
         mask = np.array(chanmask).astype("bool")
-        maskvalue = np.float32(maskvalue).astype(self.header.data_type)
+        maskvalue = np.float32(maskvalue).astype(self.header.dtype)
         out_file = self.header.prep_outfile(filename)
         for nsamps, _ii, data in self.read_plan(**plan_kwargs):
             kernels.mask_channels(data, mask, maskvalue, self.header.nchans, nsamps)
@@ -546,7 +546,7 @@ def extract_bands(
         if nchans % chanpersub != 0:
             raise ValueError("Number of channels must be divisible by sub-band size.")
 
-        nsub = (nchans - chanstart) // chanpersub
+        nsub = (self.header.nchans - chanstart) // chanpersub
         fstart = self.header.fch1 + chanstart * self.header.foff
 
         out_files = [

sigpyproc/core/rfi.py

@@ -145,9 +145,9 @@ def to_file(self, filename: str | None = None) -> str:
             filename = f"{self.header.basename}_mask.h5"
         with h5py.File(filename, "w") as fp:
             fp.attrs["threshold"] = self.threshold
-            hdr_dict = self.header.to_dict(with_properties=False)
-            for key in hdr_dict.keys():
-                fp.attrs[key] = hdr_dict[key]
+            for key, value in attrs.asdict(self.header).items():
+                if isinstance(value, (int, float, str)):
+                    fp.attrs[key] = value
             for key, value in attrs.asdict(self).items():
                 if isinstance(value, np.ndarray):
                     fp.create_dataset(key, data=value)
@@ -173,7 +173,7 @@ def from_file(cls, filename: str) -> RFIMask:
         hdr_checked = {
             key: value
             for key, value in fp_attrs.items()
-            if key in attrs.asdict(Header).keys()
+            if key in attrs.fields_dict(Header).keys()
         }
         kws = {
             "header": Header(**hdr_checked),

sigpyproc/io/pfits.py

@@ -191,16 +191,9 @@ def coord(self) -> SkyCoord:
     def tstart(self):
         """astropy.time.Time: Observation start time."""
         return Time(
-            self.header["STT_IMJD"],
-            format="mjd",
-            scale="utc",
-            location=self.location,
+            self.header["STT_IMJD"], format="mjd", scale="utc", location=self.location
         ) + TimeDelta(
-            float(self.header["STT_SMJD"]),
-            float(self.header["STT_OFFS"]),
-            format="sec",
-            scale="utc",
-            location=self.location,
+            float(self.header["STT_SMJD"]), float(self.header["STT_OFFS"]), format="sec"
         )
 
     def _check_header(self, header):
@@ -492,13 +485,16 @@ def read_subint(
         """
         sdata = self._fits["SUBINT"].data[isub]["DATA"]  # noqa: WPS219
         sdata = sdata.squeeze()
-        assert (
-            sdata.shape == self.sub_hdr.subint_shape
-        ), f"DATA column ordering {sdata.shape} is not TPF"
         if self.bitsinfo.unpack:
-            data = unpack(sdata, self.bitsinfo.nbits)
+            data = unpack(sdata.ravel(), self.bitsinfo.nbits)
+            data = data.reshape(
+                (sdata.shape[0] * self.bitsinfo.bitfact, sdata.shape[1], sdata.shape[2])
+            )
         else:
             data = np.array(sdata)
+        assert (
+            data.shape == self.sub_hdr.subint_shape
+        ), f"DATA column ordering {data.shape} is not TPF"
 
         if scloffs or weights:
             data = data.astype(np.float32, copy=False)

sigpyproc/readers.py

@@ -212,6 +212,9 @@ def read_block(self, start: int, nsamps: int) -> FilterbankBlock:
         new_header = self.header.new_header({"tstart": start_mjd, "nsamples": nsamps})
         return FilterbankBlock(data, new_header)
 
+    def read_dedisp_block(self, start: int, nsamps: int, dm: float) -> FilterbankBlock:
+        raise NotImplementedError("Not implemented for PFITSReader")
+
     def read_plan(
         self,
         gulp: int = 16384,

tests/conftest.py

@@ -38,6 +38,11 @@ def filfile_8bit_2():
     return Path(_datadir / "parkes_8bit_2.fil").as_posix()
 
 
+@pytest.fixture(scope="session", autouse=True)
+def fitsfile_4bit():
+    return Path(_datadir / "parkes_4bit.sf").as_posix()
+
+
 @pytest.fixture(scope="session", autouse=True)
 def timfile():
     return Path(_datadir / "GBT_J1807-0847.tim").as_posix()

tests/test_base.py

@@ -66,7 +66,7 @@ def test_apply_channel_mask(self, filfile_4bit, tmpfile):
         np.testing.assert_equal(new_fil.header.dtype, fil.header.dtype)
         np.testing.assert_equal(new_fil.header.nsamples, fil.header.nsamples)
         np.testing.assert_array_equal(
-            np.where(~newdata.any(axis=1))[0], np.where(chanmask == 0)[0]
+            np.where(~newdata.any(axis=1))[0], np.where(chanmask == 1)[0]
         )
 
     def test_downsample(self, filfile_4bit, tmpfile):

tests/test_header.py

@@ -75,3 +75,13 @@ def test_make_inf(self, inffile, inf_header, tmpfile):
         infheader = Header.from_inffile(inffile)
         infheader.make_inf(outfile=tmpfile)
         self.test_from_inffile(tmpfile, inf_header)
+
+    def test_from_pfits(self, fitsfile_4bit):
+        header = Header.from_pfits(fitsfile_4bit)
+        np.testing.assert_equal(header.nchans, 416)
+        np.testing.assert_equal(header.nbits, 4)
+        np.testing.assert_equal(header.nifs, 1)
+
+    def test_from_pfits_otherfile(self, inffile):
+        with np.testing.assert_raises(OSError):
+            Header.from_pfits(inffile)

tests/test_kernels.py

@@ -1,5 +1,5 @@
 import numpy as np
-from sigpyproc.core import kernels
+from sigpyproc.core import kernels, stats, rfi
 
 
 class TestKernels(object):
@@ -32,3 +32,30 @@ def test_pack4_8(self):
         input_arr = np.arange(255, dtype=np.uint8)
         output = kernels.pack4_8(kernels.unpack4_8(input_arr))
         np.testing.assert_array_equal(input_arr, output)
+
+
+class TestStats(object):
+    def test_zscore_mad(self):
+        input_arr = np.array([1, 2, 3, 4], dtype=np.uint8)
+        desired = np.array(
+            [-1.01173463, -0.33724488, 0.33724488, 1.01173463], dtype=np.float32
+        )
+        np.testing.assert_array_almost_equal(
+            desired, stats.zscore_mad(input_arr), decimal=4
+        )
+
+    def test_zscore_double_mad(self):
+        input_arr = np.array([1, 2, 3, 4], dtype=np.uint8)
+        desired = np.array(
+            [-1.01173463, -0.33724488, 0.33724488, 1.01173463], dtype=np.float32
+        )
+        np.testing.assert_array_almost_equal(
+            desired, stats.zscore_double_mad(input_arr), decimal=4
+        )
+
+
+class TestRFI(object):
+    def test_double_mad_mask(self):
+        input_arr = np.array([1, 2, 3, 4, 5, 20], dtype=np.uint8)
+        desired = np.array([0, 0, 0, 0, 0, 1], dtype=np.bool)
+        np.testing.assert_array_equal(desired, rfi.double_mad_mask(input_arr))