tabsim 0.0.1__py3-none-any.whl

tabsim/jax/utils.py

@@ -0,0 +1,90 @@
+# Progress bar for JAX scan
+# Credit to https://www.jeremiecoullon.com/2021/01/29/jax_progress_bar/
+
+from jax import lax, jit
+from jax.experimental import host_callback
+from tqdm import tqdm
+from collections.abc import Callable
+
+
+def jit_with_doc(func: Callable) -> Callable:
+    jit_func = jit(func)
+    jit_func.__doc__ = func.__doc__
+    return jit_func
+
+jit_with_doc.__doc__ = jit.__doc__
+
+def progress_bar_scan(num_samples: int, message=None):
+    "Progress bar for a JAX scan"
+    if message is None:
+        message = f"Running for {num_samples:,} iterations"
+    tqdm_bars = {}
+    if num_samples > 20:
+        print_rate = int(num_samples / 20)
+    else:
+        print_rate = 1  # if you run the sampler for less than 20 iterations
+
+    remainder = num_samples % print_rate
+
+    def _define_tqdm(arg, transform):
+        tqdm_bars[0] = tqdm(range(num_samples))
+        tqdm_bars[0].set_description(message, refresh=False)
+
+    def _update_tqdm(arg, transform):
+        tqdm_bars[0].update(arg)
+
+    def _update_progress_bar(iter_num):
+        "Updates tqdm progress bar of a JAX scan or loop"
+        _ = lax.cond(
+            iter_num == 0,
+            lambda _: host_callback.id_tap(_define_tqdm, None, result=iter_num),
+            lambda _: iter_num,
+            operand=None,
+        )
+
+        _ = lax.cond(
+            # update tqdm every multiple of `print_rate` except at the end
+            (iter_num % print_rate == 0) & (iter_num != num_samples - remainder),
+            lambda _: host_callback.id_tap(_update_tqdm, print_rate, result=iter_num),
+            lambda _: iter_num,
+            operand=None,
+        )
+
+        _ = lax.cond(
+            # update tqdm by `remainder`
+            iter_num == num_samples - remainder,
+            lambda _: host_callback.id_tap(_update_tqdm, remainder, result=iter_num),
+            lambda _: iter_num,
+            operand=None,
+        )
+
+    def _close_tqdm(arg, transform):
+        tqdm_bars[0].close()
+
+    def close_tqdm(result, iter_num):
+        return lax.cond(
+            iter_num == num_samples - 1,
+            lambda _: host_callback.id_tap(_close_tqdm, None, result=result),
+            lambda _: result,
+            operand=None,
+        )
+
+    def _progress_bar_scan(func):
+        """Decorator that adds a progress bar to `body_fun` used in `lax.scan`.
+        Note that `body_fun` must either be looping over `np.arange(num_samples)`,
+        or be looping over a tuple who's first element is `np.arange(num_samples)`
+        This means that `iter_num` is the current iteration number
+        """
+
+        def wrapper_progress_bar(carry, x):
+            if type(x) is tuple:
+                iter_num, *_ = x
+            else:
+                iter_num = x
+            _update_progress_bar(iter_num)
+            result = func(carry, x)
+            return close_tqdm(result, iter_num)
+
+        return wrapper_progress_bar
+
+    return _progress_bar_scan

tabsim/plot.py

@@ -0,0 +1,323 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from numpy.typing import ArrayLike
+
+import jax.numpy as jnp
+
+from tabsim.dask.observation import Observation
+from tabsim.jax.coordinates import alt_az_of_source
+import os
+
+plt.rcParams["font.size"] = 18
+
+
+def time_units(times: ArrayLike) -> tuple:
+    """Scale the time axis to hours, minutes or seconds depending on the total range.
+
+    Parameters
+    ----------
+    times : ArrayLike
+        Times to consider.
+
+    Returns
+    -------
+    tuple
+        Rescaled times array and the scale unit as a string.
+    """
+
+    time_range = times[-1] - times[0]
+    times = times - times[0]
+    if time_range > 3600:
+        units = "hr"
+        times = times / 3600
+    elif time_range > 60:
+        units = "min"
+        times = times / 60
+    else:
+        units = "s"
+
+    return times, units
+
+
+def plot_angular_seps(obs: Observation, save_path: str) -> None:
+    """Plot the angular separations between the RFI sources and pointing direction.
+
+    Parameters
+    ----------
+    obs : Observation
+        Observation object with RFI sources added
+    save_path : str
+        Path to where to save the plots.
+    """
+
+    times, scale = time_units(obs.times_fine)
+    n_time = len(times)
+    if n_time > 100:
+        idx = np.arange(0, n_time, n_time // 100)
+    else:
+        idx = np.arange(n_time)
+    plt.figure(figsize=(10, 7))
+    if obs.n_rfi_satellite > 0:
+        ang_seps = np.concatenate(obs.rfi_satellite_ang_sep, axis=0).mean(axis=-1).T
+        plt.plot(times[idx], ang_seps[idx], label="Satellite")
+    if obs.n_rfi_stationary > 0:
+        ang_seps = np.concatenate(obs.rfi_stationary_ang_sep, axis=0).mean(axis=-1).T
+        plt.plot(times[idx], ang_seps[idx], label="Stationary")
+    if obs.n_rfi_tle_satellite > 0:
+        ang_seps = np.concatenate(obs.rfi_tle_satellite_ang_sep, axis=0).mean(axis=-1)
+        for ang_sep, n_id in zip(np.atleast_2d(ang_seps), obs.norad_ids):
+            plt.plot(times[idx], ang_sep[idx], label=n_id.compute())
+    plt.xlabel(f"Time [{scale}]")
+    plt.ylabel("Angular Separation [deg]")
+    plt.legend()
+    plt.savefig(os.path.join(save_path, "AngularSeps.png"), format="png", dpi=200)
+
+
+def plot_src_alt(obs: Observation, save_path: str) -> None:
+    """Plot the target source altitude over the period of the observation.
+
+    Parameters
+    ----------
+    obs : Observation
+        Observation object.
+    save_path : str
+        Path to where to save the plot.
+    """
+
+    times, scale = time_units(obs.times)
+    lsa = obs.lsa.compute()
+    alt = alt_az_of_source(
+        lsa[obs.t_idx], *[x.compute() for x in [obs.latitude, obs.ra, obs.dec]]
+    )[:, 0]
+    plt.figure(figsize=(10, 7))
+    plt.plot(times, alt, ".-")
+    plt.xlabel(f"Time [{scale}]")
+    plt.ylabel("Source Altitude [deg]")
+    plt.savefig(os.path.join(save_path, "SourceAltitude.png"), format="png", dpi=200)
+
+
+def plot_uv(obs: Observation, save_path: str) -> None:
+    """Plot the uv coverage of the telescope for the given observation.
+
+    Parameters
+    ----------
+    obs : Observation
+        Observation object.
+    save_path : str
+        Path to where to save the plot.
+    """
+
+    plt.figure(figsize=(10, 10))
+    if obs.n_time_fine > 100:
+        time_step = int(obs.n_time_fine / 100)
+    else:
+        time_step = 1
+    u = obs.bl_uvw[::time_step, :, 0].compute().flatten()
+    v = obs.bl_uvw[::time_step, :, 1].compute().flatten()
+    max_U = np.max(np.sqrt(u**2 + v**2))
+    exp = float(np.floor(np.log10(max_U)))
+    mantissa = np.ceil(10 ** (np.log10(max_U) - exp))
+    lim = mantissa * 10**exp
+    plt.plot(u, v, "k.", ms=1, alpha=0.3)
+    plt.plot(-u, -v, "k.", ms=1, alpha=0.3)
+    plt.grid()
+    plt.xlim(-lim, lim)
+    plt.ylim(-lim, lim)
+    plt.xlabel("U [m]")
+    plt.ylabel("V [m]")
+    plt.savefig(os.path.join(save_path, "UV.png"), format="png", dpi=200)
+
+    plt.rcParams["font.size"] = 16
+
+
+def plot_comparison(
+    ax, times, mean1, mean2, std1, std2, true1, true2, rmse, diff=False
+):
+    times, units = time_units(times)
+    for i, a in enumerate(ax):
+        a[0].plot(rmse[..., i], "o")
+        a[0].set_xlabel("Sample")
+
+        if diff:
+            a[1].plot(times, mean1[i] - true1[i].real, label="Estimate")
+            a[1].fill_between(times, -std1[i], std1[i], color="tab:orange", alpha=0.3)
+            a[1].fill_between(
+                times, -2 * std1[i], 2 * std1[i], color="tab:orange", alpha=0.3
+            )
+            a[2].plot(times, mean2[i] - true2[i], label="Estimate")
+            a[2].fill_between(times, -std2[i], std2[i], color="tab:orange", alpha=0.3)
+            a[2].fill_between(
+                times, -2 * std2[i], 2 * std2[i], color="tab:orange", alpha=0.3
+            )
+        else:
+            a[1].plot(times, true1[i], label="True")
+            a[1].plot(times, mean1[i], label="Estimate")
+            a[1].fill_between(
+                times,
+                mean1[i] - std1[i],
+                mean1[i] + std1[i],
+                color="tab:orange",
+                alpha=0.3,
+            )
+            a[1].fill_between(
+                times,
+                mean1[i] - 2 * std1[i],
+                mean1[i] + 2 * std1[i],
+                color="tab:orange",
+                alpha=0.3,
+            )
+            a[2].plot(times, true2[i], label="True")
+            a[2].plot(times, mean2[i], label="Estimate")
+            a[2].fill_between(
+                times,
+                mean2[i] - std2[i],
+                mean2[i] + std2[i],
+                color="tab:orange",
+                alpha=0.3,
+            )
+            a[2].fill_between(
+                times,
+                mean2[i] - 2 * std2[i],
+                mean2[i] + 2 * std2[i],
+                color="tab:orange",
+                alpha=0.3,
+            )
+
+        a[1].set_xlabel(f"Time [{units}]")
+        a[1].legend()
+        a[2].set_xlabel(f"Time [{units}]")
+        a[2].legend()
+
+
+def plot_complex_real_imag(
+    times,
+    param,
+    true,
+    rmse,
+    name: str,
+    save_name: str = None,
+    diff: bool = False,
+    max_plots: int = 10,
+    save_dir: str = "plots/",
+):
+    n_params = min(param.shape[1], max_plots)
+    # idx = np.random.permutation(param.shape[1])
+    # print(param.shape, true.shape)
+    # param = param[:, idx]
+    # true = true[idx]
+    mean_r = param.real.mean(axis=0)
+    mean_i = param.imag.mean(axis=0)
+    std_r = param.real.std(axis=0)
+    std_i = param.imag.std(axis=0)
+
+    fig, ax = plt.subplots(n_params, 3, figsize=(18, 4.5 * n_params))
+
+    ax[0, 0].set_title("Root Mean Squared Error")
+    ax[0, 1].set_title(f"{name} Real")
+    ax[0, 2].set_title(f"{name} Imag")
+
+    plot_comparison(
+        ax, times, mean_r, mean_i, std_r, std_i, true.real, true.imag, rmse, diff=diff
+    )
+
+    if save_name is not None:
+        fig.savefig(
+            os.path.join(save_dir, f"{save_name}_real_imag.pdf"),
+            format="pdf",
+            bbox_inches="tight",
+        )
+    plt.close(fig)
+
+
+def plot_complex_amp_phase(
+    times,
+    param,
+    true,
+    rmse,
+    name: str,
+    save_name: str = None,
+    diff: bool = False,
+    max_plots: int = 10,
+    save_dir: str = "plots/",
+):
+    n_params = min(param.shape[1], max_plots)
+    # idx = np.random.permutation(param.shape[1])
+    # print(param.shape, true.shape)
+    # param = param[:, idx]
+    # true = true[idx]
+    mean_amp = jnp.abs(param).mean(axis=0)
+    mean_phase = jnp.rad2deg(jnp.angle(param)).mean(axis=0)
+    std_amp = jnp.abs(param).std(axis=0)
+    std_phase = jnp.rad2deg(jnp.angle(param)).std(axis=0)
+
+    fig, ax = plt.subplots(n_params, 3, figsize=(18, 4.5 * n_params))
+
+    ax[0, 0].set_title("Root Mean Squared Error")
+    ax[0, 1].set_title(f"{name} Magnitude")
+    ax[0, 2].set_title(f"{name} Phase")
+
+    plot_comparison(
+        ax,
+        times,
+        mean_amp,
+        mean_phase,
+        std_amp,
+        std_phase,
+        jnp.abs(true),
+        jnp.rad2deg(jnp.angle(true)),
+        rmse,
+        diff=diff,
+    )
+
+    if save_name is not None:
+        fig.savefig(
+            os.path.join(save_dir, f"{save_name}_amp_phase.pdf"),
+            format="pdf",
+            bbox_inches="tight",
+        )
+    plt.close(fig)
+
+
+def plot_predictions(
+    times,
+    pred,
+    args,
+    type: str = "",
+    model_name: str = "",
+    max_plots: int = 10,
+    save_dir: str = "plots/",
+):
+    plot_complex_real_imag(
+        times=times,
+        param=pred["ast_vis"],
+        true=args["vis_ast_true"],
+        rmse=pred["rmse_ast"],
+        name="Ast. Vis.",
+        save_name=f"{model_name}_{type}_ast_vis",
+        max_plots=max_plots,
+        save_dir=save_dir,
+    )
+
+    plot_complex_amp_phase(
+        times=times,
+        param=pred["rfi_vis"],
+        true=args["vis_rfi_true"],
+        rmse=pred["rmse_rfi"],
+        name="RFI Vis.",
+        save_name=f"{model_name}_{type}_rfi_vis",
+        diff=False,  # True,
+        max_plots=max_plots,
+        save_dir=save_dir,
+    )
+
+    plot_complex_amp_phase(
+        times=times,
+        param=pred["gains"],
+        true=args["gains_true"],
+        rmse=pred["rmse_gains"],
+        name="Gains",
+        save_name=f"{model_name}_{type}_gains",
+        diff=False,
+        max_plots=max_plots,
+        save_dir=save_dir,
+    )

tabsim/scripts/__init__.py

@@ -0,0 +1 @@
+

tabsim/scripts/ast_region.py

@@ -0,0 +1,56 @@
+from astropy.coordinates import SkyCoord
+import astropy.units as u
+from regions import CircleSkyRegion, RectangleSkyRegion
+
+import xarray as xr
+import numpy as np
+
+import argparse
+
+import os
+
+import matplotlib.colors as mcolors
+
+
+def main():
+
+    colors = list(mcolors.TABLEAU_COLORS.values())
+    # n_c = len(colors)
+
+    parser = argparse.ArgumentParser(
+        description="Extract astronomical positions to ds9 region file."
+    )
+    parser.add_argument(
+        "-z",
+        "--zarr_path",
+        required=True,
+        help="File path to the tabascal zarr simulation file.",
+    )
+
+    args = parser.parse_args()
+    zarr_path = args.zarr_path
+
+    if zarr_path[-1] == "/":
+        zarr_path = zarr_path[:-1]
+
+    region_path = os.path.join(
+        os.path.split(zarr_path)[0], "astronomical_positions.ds9"
+    )
+
+    xds = xr.open_zarr(zarr_path)
+
+    ra, dec = xds.ast_p_radec.data.compute().T
+    c = SkyCoord(ra, dec, unit="deg", frame="fk5")
+
+    max_U = np.max(np.linalg.norm(xds.bl_uvw.data[0], axis=-1))
+    syn_bw = (3e8 / (max_U * xds.freq.data.max())).compute()
+
+    with open(region_path, "w") as fp:
+        for i in range(len(ra)):
+            fp.write(
+                CircleSkyRegion(
+                    c[i], radius=syn_bw * u.rad, visual={"color": colors[0]}
+                ).serialize(format="ds9")
+            )
+
+    print(f"Region file written to : {region_path}")

tabsim/scripts/flag_data.py

@@ -0,0 +1,195 @@
+#!/usr/bin/env python
+
+from daskms import xds_from_ms, xds_to_table
+
+import xarray as xr
+import dask.array as da
+import dask
+import os
+
+import argparse
+import subprocess
+import shutil
+
+# aoflagger -strategy ../yaml_configs/target/firstpass.rfis tab1_obs_16A_450T-0440-1338_256I_001F-1.227e+09-1.227e+09_100PAST_000GAST_000EAST_1SAT_0GRD_1.0e+00RFI/tab1_obs_16A_450T-0440-1338_256I_001F-1.227e+09-1.227e+09_100PAST_000GAST_000EAST_1SAT_0GRD_1.0e+00RFI.ms/
+
+
+def write_perfect_flags(ms_path: str, n_sigma: float = 3.0):
+
+    xds = xds_from_ms(ms_path)[0]
+
+    if n_sigma > 0:
+        flags = da.abs(xds.CAL_DATA - xds.AST_MODEL_DATA) > (
+            n_sigma * xds.SIGMA * da.sqrt(2)
+        )
+    else:
+        flags = xr.zeros_like(xds.DATA).astype(bool)
+
+    xds = xds.assign(FLAG=flags)
+    cols = ["FLAG"]
+
+    dask.compute(xds_to_table([xds], ms_path, cols))
+
+    print()
+    print(f"Flag Threshold : {n_sigma: .1f} sigma")
+    print(f"Flag Rate      : {100*flags.data.mean().compute(): .1f} %")
+
+    return flags
+
+
+def write_to_aoflags(ms_path: str):
+
+    xds_ms = xds_from_ms(ms_path)[0]
+
+    xds_ms = xds_ms.assign({"AO_FLAGS": xds_ms["FLAG"]})
+    cols = ["AO_FLAGS"]
+
+    print("Writing AOFlagger flags to 'AO_FLAGS' column in MS file.")
+
+    dask.compute(xds_to_table([xds_ms], ms_path, cols))
+
+
+def write_aoflags_to_flag(ms_path: str):
+
+    xds_ms = xds_from_ms(ms_path)[0]
+
+    xds_ms = xds_ms.assign(FLAG=xds_ms["AO_FLAGS"])
+    cols = ["FLAG"]
+
+    print("Writing AOFlagger flags to 'FLAG' column in MS file.")
+
+    dask.compute(xds_to_table([xds_ms], ms_path, cols))
+
+
+def run_aoflagger(
+    ms_path: str,
+    data_column: str = "CAL_DATA",
+    strategy_paths: list = None,
+    sif_path: str = None,
+    bash_exec: str = "/bin/bash",
+    rerun_aoflagger: bool = False,
+):
+
+    reflagged = False
+
+    if ms_path[-1] == "/":
+        ms_path = ms_path[:-1]
+
+    if not rerun_aoflagger:
+        xds_ms = xds_from_ms(ms_path)
+        try:
+            aoflags = xds_ms["AO_FLAGS"]
+            write_aoflags_to_flag(ms_path)
+            print("Using previous AOFlagger run.")
+        except:
+            aoflags = None
+
+    if not aoflags or rerun_aoflagger:
+
+        data_dir, ms_file = os.path.split(os.path.abspath(ms_path))
+
+        docker_opts = "--rm -v /etc/group:/etc/group -v /etc/passwd:/etc/passwd -v /etc/shadow:/etc/shadow -v/etc/sudoers.d:/etc/sudoers.d -e HOME=${HOME} --user=`id -ur`"
+        container_cmd = f"docker run {docker_opts} -v {data_dir}:/data --workdir /data stimela/aoflagger:latest"
+
+        if sif_path is not None:
+            sif_path = os.path.abspath(sif_path)
+            container_cmd = (
+                f"singularity exec --bind {data_dir}:/data --pwd /data {sif_path}"
+            )
+
+        if strategy_paths is not None:
+            write_perfect_flags(ms_path, 0)
+            for strategy_path in strategy_paths:
+
+                strategy_path = os.path.abspath(strategy_path)
+                shutil.copy(strategy_path, data_dir)
+                strategy_file = os.path.split(strategy_path)[1]
+                strategy = f"-strategy /data/{strategy_file}"
+
+                aoflag_cmd = f"{container_cmd} aoflagger -column {data_column} {strategy} /data/{ms_file}"
+                subprocess.run(aoflag_cmd, shell=True, executable=bash_exec)
+
+            print()
+            print(
+                f"Strategies : {[os.path.split(strategy)[1] for strategy in strategy_paths]}"
+            )
+
+        else:
+            aoflag_cmd = f"{container_cmd} aoflagger /data/{ms_file}"
+            subprocess.run(aoflag_cmd, shell=True, executable=bash_exec)
+            print()
+
+        write_to_aoflags(ms_path)
+        reflagged = True
+
+    flags = xds_from_ms(ms_path)[0].FLAG.data
+    print(f"Flag Rate      : {100*flags.mean().compute(): .1f} %")
+
+    return flags, reflagged
+
+
+def main():
+
+    parser = argparse.ArgumentParser(
+        description="Flag CAL_DATA lying a certain threshold away from the AST_MODEL_DATA. Requires tabascal simulated MS files."
+    )
+    parser.add_argument(
+        "-m", "--ms_path", required=True, help="File path to the Measurement Set."
+    )
+    parser.add_argument(
+        "-s",
+        "--n_sigma",
+        default=3.0,
+        type=float,
+        help="Threshold in number of std of noise given by SIGMA column. 0 unflags everything. Default is 3.0.",
+    )
+    parser.add_argument(
+        "-d",
+        "--data_col",
+        default="DATA",
+        help="Data column to run AOFlagger on. Default is 'DATA'.",
+    )
+    parser.add_argument(
+        "-ao",
+        "--aoflagger",
+        default=False,
+        action=argparse.BooleanOptionalAction,
+        help="Whether to use AOFlagger to flag the data. Default is False.",
+    )
+    parser.add_argument(
+        "-rao",
+        "--rerun_aoflagger",
+        default=False,
+        action=argparse.BooleanOptionalAction,
+        help="Whether to rerun AOFlagger to flag the data. Default is False.",
+    )
+    parser.add_argument(
+        "-st",
+        "--strategy_paths",
+        default=None,
+        help="list of paths to AOFlagger strategies.",
+    )
+    parser.add_argument(
+        "-sp",
+        "--sif_path",
+        default=None,
+        help="Paths to AOFlagger singularity image.",
+    )
+    args = parser.parse_args()
+    ms_path = args.ms_path
+    n_sigma = args.n_sigma
+    strategy_paths = args.strategy_paths
+    data_col = args.data_col
+
+    if args.aoflagger:
+        if strategy_paths is not None:
+            strategy_paths = strategy_paths.split(",")
+        flags = run_aoflagger(
+            ms_path, data_col, strategy_paths, args.sif_path, args.rerun_aoflagger
+        )
+    else:
+        flags = write_perfect_flags(ms_path, n_sigma)
+
+
+if __name__ == "__main__":
+    main()