jolly-roger 0.0.2__py3-none-any.whl → 0.2.0__py3-none-any.whl

jolly_roger/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.0.2'
-__version_tuple__ = version_tuple = (0, 0, 2)
+__version__ = version = '0.2.0'
+__version_tuple__ = version_tuple = (0, 2, 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/hour_angles.py

@@ -4,19 +4,17 @@ 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
-from astropy.coordinates import EarthLocation, SkyCoord, get_sun
+from astropy.coordinates import AltAz, EarthLocation, SkyCoord, get_sun
 from astropy.time import Time
 from casacore.tables import table
 
 from jolly_roger.logging import logger
 
 # Default location with XYZ based on mean of antenna positions
-ASKAP_XYZ_m = np.array([-2556146.66356375, 5097426.58592797, -2848333.08164107]) * u.m
-ASKAP = EarthLocation(*ASKAP_XYZ_m)
+ASKAP = EarthLocation.of_site("ASKAP")
 
 
 @dataclass
@@ -27,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"""
@@ -44,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:
@@ -55,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.
 
@@ -72,9 +70,12 @@ 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=}")
+            position = SkyCoord.from_name(position)
 
     if position is None:
         if ms_path is None:
@@ -83,8 +84,8 @@ def _process_position(
 
         with table(str(ms_path / "FIELD")) as tab:
             logger.info(f"Getting the sky-position from PHASE_DIR of {ms_path=}")
-            field_positions = tab.getcol("PHASE_DIR")
-            position = SkyCoord(field_positions[0] * u.rad)
+            field_positions = tab.getcol("PHASE_DIR")[:]
+            position = SkyCoord(*(field_positions * u.rad).squeeze())
 
     if isinstance(position, SkyCoord):
         return position
@@ -116,7 +117,7 @@ def make_hour_angles_for_ms(
     logger.info(f"Computing hour angles for {ms_path=}")
     with table(str(ms_path), ack=False) as tab:
         logger.info("Extracting timesteps and constructing time mapping")
-        times_mjds = tab.getcol("TIME_CENTROID")
+        times_mjds = tab.getcol("TIME_CENTROID")[:] * u.s
 
         # get unique time steps and make sure they are in their first appeared order
         times_mjds, indices = np.unique(times_mjds, return_index=True)
@@ -127,33 +128,28 @@ def make_hour_angles_for_ms(
     if whole_day:
         logger.info(f"Assuming a full day from {times_mjds} MJD (seconds)")
         time_step = times_mjds[1] - times_mjds[0]
-        times_mjds = times_mjds[0] + time_step * np.arange(
-            int(60 * 60 * 24 / time_step)
+        times_mjds = np.arange(
+            start=times_mjds[0],
+            stop=times_mjds[0] + 24 * u.hour,
+            step=time_step,
         )
 
-    times = Time(times_mjds / 60 / 60 / 24, format="mjd")
+    times = Time(times_mjds, format="mjd", scale="utc")
 
-    sky_position: SkyCoord = _process_position(
-        position=position, times=times, ms_path=ms_path
-    )
+    sky_position = _process_position(position=position, times=times, ms_path=ms_path)
 
     lst = times.sidereal_time("apparent", longitude=location.lon)
-    hour_angle = lst - sky_position.ra
-    mask = hour_angle > 12 * u.hourangle
-    hour_angle[mask] -= 24 * u.hourangle
+    hour_angle = (lst - sky_position.ra).wrap_at(12 * u.hourangle)
 
-    logger.info("Creatring elevation curve")
-    sin_alt = np.arcsin(
-        np.sin(location.lat) * np.sin(sky_position[0].dec.rad)
-        + np.cos(location.lat) * np.cos(sky_position.dec.rad) * np.cos(hour_angle)
-    ).to(u.rad)
+    logger.info("Creating elevation curve")
+    altaz = sky_position.transform_to(AltAz(obstime=times, location=location))
 
     return PositionHourAngles(
         hour_angle=hour_angle,
         time_mjds=times_mjds,
         location=location,
         position=sky_position,
-        elevation=sin_alt,
+        elevation=altaz.alt.to(u.rad),
         time=times,
         time_map=time_map,
     )

jolly_roger/plots.py

@@ -0,0 +1,171 @@
+"""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()
+        )
+
+        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(
+            2, 2, figsize=(10, 10), sharex=True, sharey="row"
+        )
+        im = ax1.pcolormesh(
+            before_baseline_data.time,
+            before_baseline_data.freq_chan,
+            before_amp_stokesi.T,
+            norm=norm,
+        )
+        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,
+        )
+        ax2.set(
+            ylabel=f"Frequency / {after_baseline_data.freq_chan.unit:latex_inline}",
+            title="After",
+        )
+        fig.colorbar(im, ax=ax2, 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,
+        )
+        ax3.set(ylabel="Delay / s", title="Before")
+        ax4.pcolormesh(
+            after_baseline_data.time,
+            after_delays.delay,
+            after_delays_i.T,
+            norm=delay_norm,
+        )
+        ax4.set(ylabel="Delay / s", title="After")
+        fig.colorbar(im, ax=ax4, 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="k",
+                    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