jolly-roger 0.2.0__tar.gz → 0.4.0__tar.gz

{jolly_roger-0.2.0 → jolly_roger-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jolly-roger
-Version: 0.2.0
+Version: 0.4.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
@@ -60,6 +60,9 @@ Description-Content-Type: text/markdown
 
 The pirate flagger!
 
+<img src="logo.png" alt="The Jolly Roger Flag" style="width:400px;"/>
+
+
 # Installation
 
 `pip install jolly-roger`

{jolly_roger-0.2.0 → jolly_roger-0.4.0}/README.md

@@ -28,6 +28,9 @@
 
 The pirate flagger!
 
+<img src="logo.png" alt="The Jolly Roger Flag" style="width:400px;"/>
+
+
 # Installation
 
 `pip install jolly-roger`

{jolly_roger-0.2.0 → jolly_roger-0.4.0}/jolly_roger/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.0'
-__version_tuple__ = version_tuple = (0, 2, 0)
+__version__ = version = '0.4.0'
+__version_tuple__ = version_tuple = (0, 4, 0)

{jolly_roger-0.2.0 → jolly_roger-0.4.0}/jolly_roger/plots.py

@@ -85,15 +85,17 @@ def plot_baseline_comparison_data(
         norm = ImageNormalize(
             after_amp_stokesi, interval=ZScaleInterval(), stretch=SqrtStretch()
         )
+        cmap = plt.cm.viridis
 
         fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(
-            2, 2, figsize=(10, 10), sharex=True, sharey="row"
+            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}",
@@ -104,12 +106,14 @@ def plot_baseline_comparison_data(
             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",
         )
-        fig.colorbar(im, ax=ax2, label="Stokes I Amplitude / Jy")
+        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
@@ -130,6 +134,7 @@ def plot_baseline_comparison_data(
             before_delays.delay,
             before_delays_i.T,
             norm=delay_norm,
+            cmap=cmap,
         )
         ax3.set(ylabel="Delay / s", title="Before")
         ax4.pcolormesh(
@@ -137,9 +142,11 @@ def plot_baseline_comparison_data(
             after_delays.delay,
             after_delays_i.T,
             norm=delay_norm,
+            cmap=cmap,
         )
         ax4.set(ylabel="Delay / s", title="After")
-        fig.colorbar(im, ax=ax4, label="Stokes I Amplitude / Jy")
+        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]
@@ -152,8 +159,8 @@ def plot_baseline_comparison_data(
                 ax.plot(
                     baseline_data.time,
                     w_delays.w_delays[b_idx],
-                    color="k",
-                    linestyle="--",
+                    color="tab:red",
+                    linestyle="-",
                     label=f"Delay for {w_delays.object_name}",
                 )
                 ax.legend()

{jolly_roger-0.2.0 → jolly_roger-0.4.0}/jolly_roger/tractor.py

@@ -19,7 +19,9 @@ from tqdm.auto import tqdm
 from jolly_roger.delays import data_to_delay_time, delay_time_to_data
 from jolly_roger.logging import logger
 from jolly_roger.plots import plot_baseline_comparison_data
+from jolly_roger.utils import log_dataclass_attributes, log_jolly_roger_version
 from jolly_roger.uvws import WDelays, get_object_delay_for_ms
+from jolly_roger.wrap import calculate_nyquist_zone, symmetric_domain_wrap
 
 
 @dataclass(frozen=True)
@@ -73,9 +75,12 @@ def tukey_taper(
 
     if tukey_x_offset is not None:
         x_freq = x_freq[:, None] - tukey_x_offset[None, :]
+        from jolly_roger.wrap import symmetric_domain_wrap
+
+        x_freq = symmetric_domain_wrap(values=x_freq, upper_limit=np.pi)
 
     taper = np.ones_like(x_freq)
-    logger.debug(f"{x_freq.shape=} {type(x_freq)=}")
+    # logger.debug(f"{x_freq.shape=} {type(x_freq)=}")
     # Fully zero region
     taper[np.abs(x_freq) > outer_width] = 0
 
@@ -202,17 +207,16 @@ def _get_data_chunk_from_main_table(
     logger.debug(f"Length of open table: {table_length} rows")
 
     lower_row = 0
-    upper_row = chunk_size
 
     while lower_row < table_length:
-        rows: list[dict[str, Any]] = ms_table[lower_row:upper_row]
-
-        data = _list_to_array(list_of_rows=rows, key=data_column)
-        flags = _list_to_array(list_of_rows=rows, key="FLAG")
-        uvws = _list_to_array(list_of_rows=rows, key="UVW")
-        time_centroid = _list_to_array(list_of_rows=rows, key="TIME_CENTROID")
-        ant_1 = _list_to_array(list_of_rows=rows, key="ANTENNA1")
-        ant_2 = _list_to_array(list_of_rows=rows, key="ANTENNA2")
+        data = ms_table.getcol(data_column, startrow=lower_row, nrow=chunk_size)
+        flags = ms_table.getcol("FLAG", startrow=lower_row, nrow=chunk_size)
+        uvws = ms_table.getcol("UVW", startrow=lower_row, nrow=chunk_size)
+        time_centroid = ms_table.getcol(
+            "TIME_CENTROID", startrow=lower_row, nrow=chunk_size
+        )
+        ant_1 = ms_table.getcol("ANTENNA1", startrow=lower_row, nrow=chunk_size)
+        ant_2 = ms_table.getcol("ANTENNA2", startrow=lower_row, nrow=chunk_size)
 
         yield DataChunkArray(
             data=data,
@@ -226,7 +230,6 @@ def _get_data_chunk_from_main_table(
         )
 
         lower_row += chunk_size
-        upper_row += chunk_size
 
 
 def get_data_chunks(
@@ -457,6 +460,7 @@ def make_plot_results(
         before_delays = data_to_delay_time(data=before_baseline_data)
         after_delays = data_to_delay_time(data=after_baseline_data)
 
+        logger.info("Creating figure")
         # TODO: the baseline data and delay times could be put into a single
         # structure to pass around easier.
         plot_path = plot_baseline_comparison_data(
@@ -525,23 +529,32 @@ def _tukey_tractor(
     """
 
     delay_time = data_to_delay_time(data=data_chunk)
+    flags_to_return: None | NDArray[np.bool] = None
 
-    # Look up the delay offset if requested
+    # Set up the offsets. By default we will be tapering around the field,
+    # but should w_delays be specified these will be modified to direct
+    # towards the nominated object in the if below
     tukey_x_offset: u.Quantity = np.zeros_like(delay_time.delay)
 
     if w_delays is not None:
         baseline_idx, time_idx = _get_baseline_time_indicies(
             w_delays=w_delays, data_chunk=data_chunk
         )
-        tukey_x_offset = w_delays.w_delays[baseline_idx, time_idx]
-        # logger.info(f"{tukey_x_offset=}")
+        original_tukey_x_offset = w_delays.w_delays[baseline_idx, time_idx]
+
+        # Make a copy for later use post wrapping
+        tukey_x_offset = original_tukey_x_offset.copy()
+
+        tukey_x_offset = symmetric_domain_wrap(
+            values=tukey_x_offset.value, upper_limit=np.max(delay_time.delay).value
+        )
 
         # need to scale the x offsert to the -pi to pi
         # The delay should be symmetric
         tukey_x_offset = (
             tukey_x_offset / (np.max(delay_time.delay) / np.pi).decompose()
         ).value
-        # logger.info(f"{tukey_x_offset=}")
+        # # logger.info(f"{tukey_x_offset=}")
 
     taper = tukey_taper(
         x=delay_time.delay,
@@ -549,28 +562,70 @@ def _tukey_tractor(
         tukey_width=tukey_tractor_options.tukey_width,
         tukey_x_offset=tukey_x_offset,
     )
-    if w_delays is not None:
+
+    if w_delays is None:
+        # This is the easy case
+        taper = taper[None, :, None]
+
+    else:
+        # TODO: This pirate reckons that merging the masks together
+        # into a single mask throughout may make things easier to
+        # manage and visualise.
+
         # The use of the `tukey_x_offset` changes the
         # shape of the output array. The internals of that
         # function returns a different shape via the broadcasting
         taper = np.swapaxes(taper[:, :, None], 0, 1)
 
-        # Since we want to dampen the target object we invert the taper
+        # Since we want to dampen the target object we invert the taper.
+        # By default the taper dampers outside the inner region.
         taper = 1.0 - taper
 
+        # apply the flags to ignore the tapering if the object is larger
+        # than one wrap away
+        # Calculate the offset account of nyquist sampling
+        no_wraps_for_offset = calculate_nyquist_zone(
+            values=original_tukey_x_offset.value,
+            upper_limit=np.max(delay_time.delay).value,
+        )
+        ignore_wrapping_for = (
+            no_wraps_for_offset > tukey_tractor_options.ignore_nyquist_zone
+        )
+        taper[ignore_wrapping_for, :, :] = 1.0
+
         # Delay with the elevation of the target object
-        elevation_mask = w_delays.elevation < (-3 * u.deg)
+        elevation_mask = w_delays.elevation < tukey_tractor_options.elevation_cut
         taper[elevation_mask[time_idx], :, :] = 1.0
 
-        # TODO: Handle case of aliased delays
-
-        # TODO: Create heuristic to determine where baseline is long enough to
-        # ignore the tapering. Aliasing may give us this though...
-
-        # TODO: Create flags where delay is 'close' to 0
-
-    else:
-        taper = taper[None, :, None]
+        # Compute flags to ignore the objects delay crossing 0, Do
+        # This by computing the taper towards the field and
+        # see if there are any components of the two sets of tapers
+        # that are not 1 (where 1 is 'no change').
+        field_taper = tukey_taper(
+            x=delay_time.delay,
+            outer_width=tukey_tractor_options.outer_width / 4,
+            tukey_width=tukey_tractor_options.tukey_width,
+            tukey_x_offset=None,
+        )
+        # We need to account for no broadcasting when offset is None
+        # as the returned shape is different
+        field_taper = field_taper[None, :, None]
+        field_taper = 1.0 - field_taper
+        intersecting_taper = np.any(
+            np.reshape((taper != 1) & (field_taper != 1), (taper.shape[0], -1)), axis=1
+        )
+        # # Should the data need to be modified in conjunction with the flags
+        # taper[
+        #     intersecting_taper &
+        #     ~elevation_mask[time_idx] &
+        #     ~ignore_wrapping_for
+        # ] = 0.0
+        # Update flags
+        flags_to_return = np.zeros_like(data_chunk.masked_data.mask)
+        flags_to_return[intersecting_taper] = True
+        flags_to_return = (
+            ~np.isfinite(data_chunk.masked_data.filled(np.nan)) | flags_to_return
+        )
 
     # Delay-time is a 3D array: (time, delay, pol)
     # Taper is 1D: (delay,)
@@ -588,7 +643,7 @@ def _tukey_tractor(
     )
     logger.debug(f"{tapered_data.masked_data.shape=} {tapered_data.masked_data.dtype}")
 
-    return tapered_data
+    return tapered_data, flags_to_return
 
 
 @dataclass
@@ -610,7 +665,7 @@ class TukeyTractorOptions:
     dry_run: bool = False
     """Indicates whether the data will be written back to the measurement set"""
     make_plots: bool = False
-    """Create a small set of diagnostic plots"""
+    """Create a small set of diagnostic plots. This can be slow."""
     overwrite: bool = False
     """If the output column exists it will be overwritten"""
     chunk_size: int = 1000
@@ -619,6 +674,10 @@ class TukeyTractorOptions:
     """apply the taper using the delay towards the target object."""
     target_object: str = "Sun"
     """The target object to apply the delay towards."""
+    elevation_cut: u.Quantity = -1 * u.deg
+    """The elevation cut-off for the target object. Defaults to 0 degrees."""
+    ignore_nyquist_zone: int = 2
+    """Do not apply the tukey taper if object is beyond this Nyquist zone"""
 
 
 def tukey_tractor(
@@ -634,8 +693,10 @@ def tukey_tractor(
     Args:
         tukey_tractor_options (TukeyTractorOptions): The settings to use during the taper, and measurement set to apply them to.
     """
-    logger.info("jolly-roger")
-    logger.info(f"Options: {tukey_tractor_options}")
+    log_jolly_roger_version()
+    log_dataclass_attributes(
+        to_log=tukey_tractor_options, class_name="TukeyTaperOptions"
+    )
 
     # acquire all the tables necessary to get unit information and data from
     open_ms_tables = get_open_ms_tables(
@@ -670,7 +731,7 @@ def tukey_tractor(
                 chunk_size=tukey_tractor_options.chunk_size,
                 data_column=tukey_tractor_options.data_column,
             ):
-                taper_data_chunk = _tukey_tractor(
+                taper_data_chunk, flags_to_apply = _tukey_tractor(
                     data_chunk=data_chunk,
                     tukey_tractor_options=tukey_tractor_options,
                     w_delays=w_delays,
@@ -685,6 +746,13 @@ def tukey_tractor(
                     startrow=taper_data_chunk.row_start,
                     nrow=taper_data_chunk.chunk_size,
                 )
+                if flags_to_apply is not None:
+                    open_ms_tables.main_table.putcol(
+                        columnname="FLAG",
+                        value=flags_to_apply,
+                        startrow=taper_data_chunk.row_start,
+                        nrow=taper_data_chunk.chunk_size,
+                    )
 
     if tukey_tractor_options.make_plots:
         plot_paths = make_plot_results(
@@ -751,7 +819,7 @@ def get_parser() -> ArgumentParser:
     tukey_parser.add_argument(
         "--make-plots",
         action="store_true",
-        help="If set, the Tukey tractor will make plots of the results",
+        help="If set, the Tukey tractor will make plots of the results. This can be slow.",
     )
     tukey_parser.add_argument(
         "--overwrite",
@@ -775,6 +843,12 @@ def get_parser() -> ArgumentParser:
         action="store_true",
         help="Whether the tukey taper is applied towards the target object (e.g. the Sun). If not set, the taper is applied towards large delays.",
     )
+    tukey_parser.add_argument(
+        "--ignore-nyquist-zone",
+        type=int,
+        default=2,
+        help="Do not apply the taper if the objects delays beyond this Nyquist zone",
+    )
 
     return parser
 
@@ -798,7 +872,9 @@ def cli() -> None:
             chunk_size=args.chunk_size,
             target_object=args.target_object,
             apply_towards_object=args.apply_towards_object,
+            ignore_nyquist_zone=args.ignore_nyquist_zone,
         )
+
         tukey_tractor(tukey_tractor_options=tukey_tractor_options)
     else:
         parser.print_help()

jolly_roger-0.4.0/jolly_roger/utils.py

@@ -0,0 +1,28 @@
+"""Small helper utility functions"""
+
+from __future__ import annotations
+
+from dataclasses import is_dataclass
+from typing import Any
+
+from jolly_roger import __version__
+from jolly_roger.logging import logger
+
+
+def log_jolly_roger_version() -> None:
+    """Write out a jolly roger header to output"""
+    logger.info("Jolly-Roger: Flagging and Tapering")
+    logger.info(f"Version: {__version__}")
+
+
+def log_dataclass_attributes(to_log: Any, class_name: str | None = None) -> None:
+    """Log the attributes and values from an Dataclass"""
+    if not is_dataclass(to_log):
+        return
+
+    if class_name:
+        logger.info(f"Settings for {class_name}")
+
+    for attribute in to_log.__annotations__:
+        value = to_log.__dict__[attribute]
+        logger.info(f"{attribute:<30} = {value}")

jolly_roger-0.4.0/jolly_roger/wrap.py

@@ -0,0 +1,51 @@
+"""Helper utilities to deal with cyclic boundaries
+on data"""
+
+from __future__ import annotations
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+def symmetric_domain_wrap(
+    values: NDArray[np.floating], upper_limit: float
+) -> NDArray[np.floating]:
+    """Place a set of values into a cyclic domain that is
+    symmetric around zero.
+
+    Args:
+        values (NDArray[np.floating]): Values that need to be mapped into the cyclic domain
+        upper_limit (float): The upper bound of the symmetric cyclic domain
+
+    Returns:
+        NDArray[np.floating]: Values that have been mapped to the -upper_limit to upper_limit domain
+    """
+    # Calculate an appropriate domain mapping
+    # The natural domain is going to be mapped
+    # to -pi to pi.
+    domain_mapping = np.pi / upper_limit
+
+    real = np.cos(values * domain_mapping)
+    imag = np.sin(values * domain_mapping)
+
+    wrapped_values = real + 1j * imag
+
+    return np.angle(wrapped_values) / domain_mapping
+
+
+def calculate_nyquist_zone(
+    values: NDArray[np.floating], upper_limit: float
+) -> NDArray[np.int_]:
+    """Return the nyquist zone a value is in for a symmetric
+    set of bounds around zero
+
+    Args:
+        values (NDArray[np.floating]): The values to calculate the zone for
+        upper_limit (float): The upper bound to the symmetric domain around zero
+
+    Returns:
+        NDArray[np.int_]: The zones values correspond to
+    """
+    return np.array(
+        np.floor((upper_limit + np.abs(values)) / (2.0 * upper_limit)) + 1, dtype=int
+    )

jolly_roger-0.4.0/tests/test_wrap.py

@@ -0,0 +1,64 @@
+"""Tests to ensure correctness of the
+wrapping utilities"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from jolly_roger.wrap import calculate_nyquist_zone, symmetric_domain_wrap
+
+
+def test_symmetric_domain_wrap_with_2d_arrays() -> None:
+    """Ensure mapping values to a periodic domain works for arrays"""
+
+    values = np.array([np.linspace(10, 40, 10), np.linspace(10, 40, 10)])
+    assert values.shape == (2, 10)
+
+    wrapped_values = symmetric_domain_wrap(values=values, upper_limit=60)
+    assert wrapped_values.shape == values.shape
+    assert np.all(np.isclose(wrapped_values, values))
+
+    # The addition of 120 is the size of the domain here
+    # and they should be wrapped back to the main map
+    larger_values = np.array(
+        [np.linspace(10, 40, 10), np.linspace(10 + 120, 40 + 120, 10)]
+    )
+    assert values.shape == (2, 10)
+
+    wrapped_values = symmetric_domain_wrap(values=larger_values, upper_limit=60)
+    assert wrapped_values.shape == values.shape
+    assert np.all(np.isclose(wrapped_values, values))
+
+
+def test_symmetric_domain_wrap() -> None:
+    """Ensure mapping values to a periodic domain works"""
+
+    values = np.linspace(10, 40, 10)
+
+    wrapped_values = symmetric_domain_wrap(values=values, upper_limit=60)
+    assert np.all(np.isclose(values, wrapped_values))
+
+    values = np.linspace(10, 40, 10)
+
+    # Domain size would be 120 values given the upper limit of 60
+    wrapped_values = symmetric_domain_wrap(values=values + 120, upper_limit=60)
+    assert np.all(np.isclose(values, wrapped_values))
+
+    values = np.linspace(10, 40, 10)
+
+    # Domain size would be 120 values given the upper limit of 60
+    wrapped_values = symmetric_domain_wrap(values=values + 2 * 120, upper_limit=60)
+    assert np.all(np.isclose(values, wrapped_values))
+
+
+def test_calculate_nyquist_zone() -> None:
+    """Match to the right zone"""
+    assert calculate_nyquist_zone(values=30, upper_limit=60) == 1
+    assert calculate_nyquist_zone(values=90, upper_limit=60) == 2
+    assert calculate_nyquist_zone(values=-30, upper_limit=60) == 1
+    assert calculate_nyquist_zone(values=-90, upper_limit=60) == 2
+
+    assert np.all(
+        calculate_nyquist_zone(values=np.array([-90, -30, 0, 30, 90]), upper_limit=60)
+        == np.array([2, 1, 1, 1, 2])
+    )