diff --git a/CHANGES.rst b/CHANGES.rst index 2dc020cd..22ddee32 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -12,6 +12,9 @@ New Features - Add a simpler ``ShiftSamples`` task that just shifts channels by integer number of samples. [#226, #235, #239] +- Add ability for incoherent dedispersion with ``DisperseSamples`` and + ``Dedispersamples``. [#238] + - Streams can now carry meta-data in a ``meta`` attribute. This includes information on ``frequency``, ``sideband``, and ``polarization``, all of which are stored in ``meta['__attributes__']`` entry (like astropy's diff --git a/baseband_tasks/dispersion.py b/baseband_tasks/dispersion.py index d3679f83..c92aaf2a 100644 --- a/baseband_tasks/dispersion.py +++ b/baseband_tasks/dispersion.py @@ -4,12 +4,13 @@ from astropy import units as u from astropy.utils import lazyproperty -from .base import PaddedTaskBase, getattr_if_none +from .base import PaddedTaskBase, getattr_if_none, SetAttribute from .fourier import fft_maker from .dm import DispersionMeasure +from .sampling import ShiftSamples -__all__ = ['Disperse', 'Dedisperse'] +__all__ = ['Disperse', 'Dedisperse', 'DisperseSamples', 'DedisperseSamples'] class Disperse(PaddedTaskBase): @@ -40,6 +41,8 @@ class Disperse(PaddedTaskBase): See Also -------- baseband_tasks.fourier.fft_maker : to select the FFT package used. + baseband_tasks.dispersion.Dedisperse : for coherent dedispersion + baseband_tasks.dispersion.DisperseSamples : for incoherent dispersion """ def __init__(self, ih, dm, *, reference_frequency=None, @@ -170,10 +173,119 @@ class Dedisperse(Disperse): See Also -------- baseband_tasks.fourier.fft_maker : to select the FFT package used. + baseband_tasks.dispersion.Disperse : for coherent dispersion + baseband_tasks.dispersion.DedisperseSamples : for incoherent dedispersion + """ + + def __init__(self, ih, dm, *, reference_frequency=None, + samples_per_frame=None, frequency=None, sideband=None): + super().__init__(ih, -dm, reference_frequency=reference_frequency, + samples_per_frame=samples_per_frame, + frequency=frequency, sideband=sideband) + + @property + def dm(self): + return -self._dm + + +class DisperseSamples(ShiftSamples): + """Incoherently shift a time stream to give it a dispersive time delay. + + This task does not handle any in-channel dispersive smearing, but only + shifts the samples according to the mid-channel frequency. + + Parameters + ---------- + ih : task or `baseband` stream reader + Input data stream, with time as the first axis. + dm : float or `~baseband_tasks.dm.DispersionMeasure` quantity + Dispersion measure to disperse with. If negative, will dedisperse, + but clearer to use `~baseband_tasks.dispersion.DedisperseSamples`. + reference_frequency : `~astropy.units.Quantity` + Frequency to which the data should be dispersed. Can be an array. + By default, the mean frequency. + samples_per_frame : int, optional + Number of dispersed samples which should be produced in one go. + The number of input samples used will be larger to avoid wrapping. + If not given, as produced by the minimum power of 2 of input + samples that yields at least 75% efficiency. + frequency : `~astropy.units.Quantity`, optional + Frequencies for each channel in ``ih``. + Default: taken from ``ih`` (if available). + sideband : array, optional + Whether frequencies in ``ih`` are upper (+1) or lower (-1) sideband. + Default: taken from ``ih`` (if available). Note that while this is + only used if the data is real (to calculate the mid-channel + frequency), it should always be passed in together with ``frequency``, + since otherwise other tasks cannot interpret frequency correctly. + + See Also + -------- + baseband_tasks.dispersion.DedisperseSamples : for incoherent dedispersion + baseband_tasks.dispersion.Disperse : for coherent dispersion + """ + + def __init__(self, ih, dm, *, reference_frequency=None, + samples_per_frame=None, frequency=None, sideband=None): + # Set possible missing frequency/sideband attributes + if frequency is not None or sideband is not None: + ih = SetAttribute(ih, frequency=frequency, sideband=sideband) + frequency = ih.frequency + if not ih.complex_data: + # Calculate mid-channel frequency for real data (for complex, + # the frequencies are already mid-channel). + frequency = frequency + ih.sideband * ih.sample_rate / 2. + + if reference_frequency is None: + reference_frequency = frequency.mean() + + # Compute the time shift and use it to set up ShiftSamples. + dm = DispersionMeasure(dm) + time_delay = dm.time_delay(frequency, reference_frequency) + super().__init__(ih, time_delay, samples_per_frame=samples_per_frame) + + self.reference_frequency = reference_frequency + self._dm = dm + + +class DedisperseSamples(DisperseSamples): + """Incoherently shift a time stream to correct for a dispersive time delay. + + This task does not handle any in-channel dispersive smearing, but only + shifts the samples according to the mid-channel frequency. + Parameters + ---------- + ih : task or `baseband` stream reader + Input data stream, with time as the first axis. + dm : float or `~baseband_tasks.dm.DispersionMeasure` quantity + Dispersion measure to correct for. If negative, will disperse, + but clearer to use `~baseband_tasks.dispersion.DisperseSamples`. + reference_frequency : `~astropy.units.Quantity` + Frequency to which the data should be dispersed. Can be an array. + By default, the mean frequency. + samples_per_frame : int, optional + Number of dedispersed samples which should be produced in one go. + The number of input samples used will be larger to avoid wrapping. + If not given, as produced by the minimum power of 2 of input + samples that yields at least 75% efficiency. + frequency : `~astropy.units.Quantity`, optional + Frequencies for each channel in ``ih``. + Default: taken from ``ih`` (if available). + sideband : array, optional + Whether frequencies in ``ih`` are upper (+1) or lower (-1) sideband. + Default: taken from ``ih`` (if available). Note that while this is + only used if the data is real (to calculate the mid-channel + frequency), it should always be passed in together with ``frequency``, + since otherwise other tasks cannot interpret frequency correctly. + + See Also + -------- + baseband_tasks.dispersion.DisperseSamples : for incoherent dispersion + baseband_tasks.dispersion.Dedisperse : for coherent dedispersion """ - def __init__(self, ih, dm, reference_frequency=None, + def __init__(self, ih, dm, *, reference_frequency=None, samples_per_frame=None, frequency=None, sideband=None): super().__init__(ih, -dm, reference_frequency=reference_frequency, samples_per_frame=samples_per_frame, diff --git a/baseband_tasks/tests/test_dispersion.py b/baseband_tasks/tests/test_dispersion.py index 3e5d8159..5e2951a3 100644 --- a/baseband_tasks/tests/test_dispersion.py +++ b/baseband_tasks/tests/test_dispersion.py @@ -6,7 +6,8 @@ from astropy.tests.helper import assert_quantity_allclose from baseband_tasks.fourier import fft_maker -from baseband_tasks.dispersion import Disperse, Dedisperse, DispersionMeasure +from baseband_tasks.dispersion import (Disperse, Dedisperse, DispersionMeasure, + DisperseSamples, DedisperseSamples) from baseband_tasks.generators import StreamGenerator @@ -21,9 +22,8 @@ 299.872 * u.MHz) # Below lower edge -class TestDispersion: - - def setup(self): +class GiantPulseSetup: + def setup_class(self): self.start_time = Time('2010-11-12T13:14:15') self.sample_rate = 128. * u.kHz self.shape = (164000, 2) @@ -38,12 +38,16 @@ def setup(self): # Time delay of 0.05 s over 128 kHz band. self.dm = DispersionMeasure(1000.*0.05/0.039342251) + @classmethod def make_giant_pulse(self, sh): data = np.empty((sh.samples_per_frame,) + sh.shape[1:], sh.dtype) do_gp = sh.tell() + np.arange(sh.samples_per_frame) == self.gp_sample data[...] = do_gp[:, np.newaxis] return data + +class TestDispersion(GiantPulseSetup): + def test_time_delay(self): time_delay = self.dm.time_delay( self.gp.frequency - self.sample_rate / 2., @@ -189,8 +193,8 @@ def test_disperse_closing(self): assert 'phase_factor' not in disperse.__dict__ -class TestDispersionReal(TestDispersion): - def setup(self): +class GiantPulseSetupReal(GiantPulseSetup): + def setup_class(self): self.start_time = Time('2010-11-12T13:14:15') self.sample_rate = 256. * u.kHz self.shape = (328000, 2) @@ -207,6 +211,8 @@ def setup(self): # Time delay of 0.05 s over 128 kHz band. self.dm = DispersionMeasure(1000.*0.05/0.039342251) + +class TestDispersionReal(TestDispersion, GiantPulseSetupReal): # Override tests that do not simply work for the real data, # since the sample rate is twice as high. def test_time_delay(self): @@ -287,3 +293,55 @@ def test_disperse_negative_dm(self): # Lower sideband [1] is dedispersed to earlier. assert p[10, 0] > 0.99 and p[9, 0] < 0.006 assert p[9, 1] > 0.99 and p[10, 1] < 0.006 + + +class TestDispersSample(GiantPulseSetupReal): + + @pytest.mark.parametrize('reference_frequency', REFERENCE_FREQUENCIES) + @pytest.mark.parametrize('frequency, sideband', [ + (None, None), + ([199.936, 200.064]*u.MHz, np.array((1, -1))), # Far off from normal. + ([200.064, 199.936]*u.MHz, np.array((-1, -1)))]) + def test_disperse_sample(self, reference_frequency, frequency, sideband): + disperse = DisperseSamples(self.gp, self.dm, frequency=frequency, + reference_frequency=reference_frequency, + sideband=sideband) + + # Seek input time of the giant pulse, corrected to the reference + # frequency, and read around it. + center_frequency = (disperse.frequency + + disperse.sideband * disperse.sample_rate / 2.) + time_delay = self.dm.time_delay(center_frequency, + disperse.reference_frequency) + t_gp = (self.start_time + self.gp_sample / self.sample_rate + + time_delay) + + disperse.seek(t_gp.min()) + around_gp = disperse.read(self.gp_sample) + + sample_shift_diff = ((time_delay.max() - time_delay.min()) + * self.sample_rate) + sample_shift_diff = np.round(sample_shift_diff.to(u.one)).astype(int) + expected = np.zeros_like(around_gp) + expected[0, t_gp.argmin()] = 1.0 + expected[sample_shift_diff, t_gp.argmax()] = 1.0 + assert np.all(around_gp == expected) + + @pytest.mark.parametrize('reference_frequency', [200 * u.MHz, 300 * u.MHz]) + def test_disperse_roundtrip1(self, reference_frequency): + self.gp.seek(self.start_time + self.gp_sample / self.sample_rate) + self.gp.seek(-1024, 1) + gp = self.gp.read(2048) + # Set up dispersion as above, and check that one can invert + disperse = DisperseSamples(self.gp, self.dm, + reference_frequency=reference_frequency) + dedisperse = DedisperseSamples(disperse, self.dm, + reference_frequency=reference_frequency) + assert dedisperse.dm == self.dm + assert dedisperse._dm == -self.dm + dedisperse.seek(self.start_time + self.gp_sample / self.sample_rate) + dedisperse.seek(-1024, 1) + gp_dd = dedisperse.read(2048) + # Note: rounding errors mean this doesn't work perfectly. + assert np.any(gp_dd == 1.0) + assert np.all(gp_dd == gp)