jolly-roger 0.1.0__tar.gz → 0.3.0__tar.gz

{jolly_roger-0.1.0 → jolly_roger-0.3.0}/.gitignore

@@ -157,3 +157,5 @@ Thumbs.db
 *~
 *.swp
 jolly_roger/_version.py
+subtractor_uvw.ipynb
+subtractor.ipynb

{jolly_roger-0.1.0 → jolly_roger-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jolly-roger
-Version: 0.1.0
+Version: 0.3.0
 Summary: The pirate flagger
 Project-URL: Homepage, https://github.com/flint-crew/jolly-roger
 Project-URL: Bug Tracker, https://github.com/flint-crew/jolly-roger/issues
@@ -23,6 +23,8 @@ Classifier: Topic :: Scientific/Engineering
 Classifier: Typing :: Typed
 Requires-Python: >=3.11
 Requires-Dist: astropy
+Requires-Dist: dask-ms
+Requires-Dist: matplotlib
 Requires-Dist: numpy>=2.0.0
 Requires-Dist: python-casacore>=3.6.0
 Requires-Dist: tqdm

{jolly_roger-0.1.0 → jolly_roger-0.3.0}/jolly_roger/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.0'
-__version_tuple__ = version_tuple = (0, 1, 0)
+__version__ = version = '0.3.0'
+__version_tuple__ = version_tuple = (0, 3, 0)

jolly_roger-0.3.0/jolly_roger/delays.py

@@ -0,0 +1,107 @@
+"""Utilities and structures around the delay calculations"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import astropy.units as u
+import numpy as np
+from numpy.typing import NDArray
+
+from jolly_roger.logging import logger
+
+if TYPE_CHECKING:
+    # avoid circular imports
+    from jolly_roger.tractor import BaselineData, DataChunk
+
+
+@dataclass
+class DelayTime:
+    """Container for delay time and associated metadata."""
+
+    delay_time: NDArray[np.complexfloating]
+    """ The delay vs time data. shape=(time, delay, pol)"""
+    delay: u.Quantity
+    """The delay values corresponding to the delay time data."""
+
+
+def data_to_delay_time(data: BaselineData | DataChunk) -> DelayTime:
+    logger.debug("Converting freq-time to delay-time")
+    delay_time = np.fft.fftshift(
+        np.fft.fft(data.masked_data.filled(0 + 0j), axis=1), axes=1
+    )
+    delay = np.fft.fftshift(
+        np.fft.fftfreq(
+            n=len(data.freq_chan),
+            d=np.diff(data.freq_chan).mean(),
+        ).decompose()
+    )
+    return DelayTime(
+        delay_time=delay_time,
+        delay=delay,
+    )
+
+
+def delay_time_to_data(
+    delay_time: DelayTime,
+    original_data: DataChunk,
+) -> DataChunk:
+    """Convert delay time data back to the original data format."""
+    logger.debug("Converting delay-time to freq-time")
+    new_data = np.fft.ifft(
+        np.fft.ifftshift(delay_time.delay_time, axes=1),
+        axis=1,
+    )
+    new_data_masked = np.ma.masked_array(
+        new_data,
+        mask=original_data.masked_data.mask,
+    )
+    new_data = original_data
+    new_data.masked_data = new_data_masked
+    return new_data
+
+
+@dataclass
+class DelayRate:
+    """Container for delay rate and associated metadata."""
+
+    delay_rate: np.ndarray
+    """The delay rate vs time data. shape=(rate, delay, pol)"""
+    delay: u.Quantity
+    """The delay values corresponding to the delay rate data."""
+    rate: u.Quantity
+    """The delay rate values corresponding to the delay rate data."""
+
+
+def data_to_delay_rate(
+    baseline_data: BaselineData,
+) -> DelayRate:
+    """Convert baseline data to delay rate."""
+    # This only makes sense when running on time data. Hence
+    # asserting the type of BaelineData
+
+    assert isinstance(baseline_data, BaselineData), (
+        f"baseline_data is type={type(baseline_data)}, but needs to be BaselineData"
+    )
+
+    logger.info("Converting freq-time to delay-rate")
+    delay_rate = np.fft.fftshift(np.fft.fft2(baseline_data.masked_data.filled(0 + 0j)))
+    delay = np.fft.fftshift(
+        np.fft.fftfreq(
+            n=len(baseline_data.freq_chan),
+            d=np.diff(baseline_data.freq_chan).mean(),
+        ).decompose()
+    )
+    rate = np.fft.fftshift(
+        np.fft.fftfreq(
+            n=len(baseline_data.time),
+            d=np.diff(baseline_data.time.mjd * u.day).mean(),
+        ).decompose()
+    )
+
+    return DelayRate(
+        delay_rate=delay_rate,
+        delay=delay,
+        rate=rate,
+    )

{jolly_roger-0.1.0 → jolly_roger-0.3.0}/jolly_roger/hour_angles.py

@@ -4,7 +4,6 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Literal
 
 import astropy.units as u
 import numpy as np
@@ -26,13 +25,13 @@ class PositionHourAngles:
 
     hour_angle: u.rad
     """The hour angle across sampled time intervales of a source for a Earth location"""
-    time_mjds: np.ndarray
+    time_mjds: u.Quantity
     """The MJD time in seconds from which other quantities are evalauted against. Should be drawn from a measurement set."""
     location: EarthLocation
     """The location these quantities have been derived from."""
     position: SkyCoord
     """The sky-position that is being used to calculate quantities towards"""
-    elevation: np.ndarray
+    elevation: u.Quantity
     """The elevation of the ``position` direction across time"""
     time: Time
     """Representation of the `time_mjds` attribute"""
@@ -43,7 +42,7 @@ class PositionHourAngles:
 
 
 def _process_position(
-    position: SkyCoord | Literal["sun"] | None = None,
+    position: SkyCoord | str | None = None,
     ms_path: Path | None = None,
     times: Time | None = None,
 ) -> SkyCoord:
@@ -54,7 +53,7 @@ def _process_position(
     set
 
     Args:
-        position (SkyCoord | Literal["sun"] | None, optional): The position to be considered. Defaults to None.
+        position (SkyCoord | str | None, optional): The position to be considered. Defaults to None.
         ms_path (Path | None, optional): The path with the PHASE_DIR to use should `position` be None. Defaults to None.
         times (Time | None, optional): Times to used if they are required in the lookup. Defaults to None.
 
@@ -71,8 +70,8 @@ def _process_position(
         if times is None:
             msg = f"{times=}, but needs to be set when position is a name"
             raise ValueError(msg)
-        if position == "sun":
-            logger.info("Getting sky-position of the sun")
+        if position.lower() == "sun":
+            logger.info("Getting sky-position of the Sun")
             position = get_sun(times)
         else:
             logger.info(f"Getting sky-position of {position=}")
@@ -142,7 +141,7 @@ def make_hour_angles_for_ms(
     lst = times.sidereal_time("apparent", longitude=location.lon)
     hour_angle = (lst - sky_position.ra).wrap_at(12 * u.hourangle)
 
-    logger.info("Creatring elevation curve")
+    logger.info("Creating elevation curve")
     altaz = sky_position.transform_to(AltAz(obstime=times, location=location))
 
     return PositionHourAngles(

jolly_roger-0.3.0/jolly_roger/plots.py

@@ -0,0 +1,178 @@
+"""Routines around plotting"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from jolly_roger.uvws import WDelays
+
+if TYPE_CHECKING:
+    from jolly_roger.delays import DelayTime
+    from jolly_roger.tractor import BaselineData
+
+
+def plot_baseline_data(
+    baseline_data: BaselineData,
+    output_dir: Path,
+    suffix: str = "",
+) -> None:
+    from astropy.visualization import quantity_support, time_support
+
+    with quantity_support(), time_support():
+        data_masked = baseline_data.masked_data
+        data_xx = data_masked[..., 0]
+        data_yy = data_masked[..., -1]
+        data_stokesi = (data_xx + data_yy) / 2
+        amp_stokesi = np.abs(data_stokesi)
+
+        fig, ax = plt.subplots()
+        im = ax.pcolormesh(
+            baseline_data.time,
+            baseline_data.freq_chan,
+            amp_stokesi.T,
+        )
+        fig.colorbar(im, ax=ax, label="Stokes I Amplitude / Jy")
+        ax.set(
+            ylabel=f"Frequency / {baseline_data.freq_chan.unit:latex_inline}",
+            title=f"Ant {baseline_data.ant_1} - Ant {baseline_data.ant_2}",
+        )
+        output_path = (
+            output_dir
+            / f"baseline_data_{baseline_data.ant_1}_{baseline_data.ant_2}{suffix}.png"
+        )
+        fig.savefig(output_path)
+
+
+def plot_baseline_comparison_data(
+    before_baseline_data: BaselineData,
+    after_baseline_data: BaselineData,
+    before_delays: DelayTime,
+    after_delays: DelayTime,
+    output_dir: Path,
+    suffix: str = "",
+    w_delays: WDelays | None = None,
+) -> Path:
+    from astropy.visualization import (
+        ImageNormalize,
+        LogStretch,
+        MinMaxInterval,
+        SqrtStretch,
+        ZScaleInterval,
+        quantity_support,
+        time_support,
+    )
+
+    with quantity_support(), time_support():
+        before_amp_stokesi = np.abs(
+            (
+                before_baseline_data.masked_data[..., 0]
+                + before_baseline_data.masked_data[..., -1]
+            )
+            / 2
+        )
+        after_amp_stokesi = np.abs(
+            (
+                after_baseline_data.masked_data[..., 0]
+                + after_baseline_data.masked_data[..., -1]
+            )
+            / 2
+        )
+
+        norm = ImageNormalize(
+            after_amp_stokesi, interval=ZScaleInterval(), stretch=SqrtStretch()
+        )
+        cmap = plt.cm.viridis
+
+        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(
+            2, 2, figsize=(12, 10), sharex=True, sharey="row"
+        )
+        im = ax1.pcolormesh(
+            before_baseline_data.time,
+            before_baseline_data.freq_chan,
+            before_amp_stokesi.T,
+            norm=norm,
+            cmap=cmap,
+        )
+        ax1.set(
+            ylabel=f"Frequency / {before_baseline_data.freq_chan.unit:latex_inline}",
+            title="Before",
+        )
+        ax2.pcolormesh(
+            after_baseline_data.time,
+            after_baseline_data.freq_chan,
+            after_amp_stokesi.T,
+            norm=norm,
+            cmap=cmap,
+        )
+        ax2.set(
+            ylabel=f"Frequency / {after_baseline_data.freq_chan.unit:latex_inline}",
+            title="After",
+        )
+        for ax in (ax1, ax2):
+            fig.colorbar(im, ax=ax, label="Stokes I Amplitude / Jy")
+
+        # TODO: Move these delay calculations outside of the plotting function
+        # And here we calculate the delay information
+
+        before_delays_i = np.abs(
+            (before_delays.delay_time[:, :, 0] + before_delays.delay_time[:, :, -1]) / 2
+        )
+        after_delays_i = np.abs(
+            (after_delays.delay_time[:, :, 0] + after_delays.delay_time[:, :, -1]) / 2
+        )
+
+        delay_norm = ImageNormalize(
+            before_delays_i, interval=MinMaxInterval(), stretch=LogStretch()
+        )
+
+        im = ax3.pcolormesh(
+            before_baseline_data.time,
+            before_delays.delay,
+            before_delays_i.T,
+            norm=delay_norm,
+            cmap=cmap,
+        )
+        ax3.set(ylabel="Delay / s", title="Before")
+        ax4.pcolormesh(
+            after_baseline_data.time,
+            after_delays.delay,
+            after_delays_i.T,
+            norm=delay_norm,
+            cmap=cmap,
+        )
+        ax4.set(ylabel="Delay / s", title="After")
+        for ax in (ax3, ax4):
+            fig.colorbar(im, ax=ax, label="Stokes I Amplitude / Jy")
+
+        if w_delays is not None:
+            for ax, baseline_data in zip(  # type:ignore[call-overload]
+                (ax3, ax4),
+                (before_baseline_data, after_baseline_data),
+                strict=True,
+            ):
+                ant_1, ant_2 = baseline_data.ant_1, baseline_data.ant_2
+                b_idx = w_delays.b_map[ant_1, ant_2]
+                ax.plot(
+                    baseline_data.time,
+                    w_delays.w_delays[b_idx],
+                    color="tab:red",
+                    linestyle="-",
+                    label=f"Delay for {w_delays.object_name}",
+                )
+                ax.legend()
+
+        output_path = (
+            output_dir
+            / f"baseline_data_{before_baseline_data.ant_1}_{before_baseline_data.ant_2}{suffix}.png"
+        )
+        fig.suptitle(
+            f"Ant {after_baseline_data.ant_1} - Ant {after_baseline_data.ant_2}"
+        )
+        fig.tight_layout()
+        fig.savefig(output_path)
+
+        return output_path