tabsim 0.0.1__py3-none-any.whl

tabsim/scripts/sim_vis.py

@@ -0,0 +1,136 @@
+import argparse
+import os
+
+from tabsim.config import load_config, run_sim_config
+
+
+def get_abs_path(rel_path: str | None, work_dir: str) -> str | None:
+    if rel_path:
+        return os.path.abspath(os.path.join(work_dir, rel_path))
+    else:
+        return None
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Simulate an observation defined by a YAML config file."
+    )
+    parser.add_argument(
+        "-c",
+        "--config_path",
+        required=True,
+        help="File path to the observation config file.",
+    )
+    parser.add_argument(
+        "-r",
+        "--rfi_amp",
+        default=1.0,
+        type=float,
+        help="Scale the RFI power. Default is 1.0",
+    )
+    parser.add_argument(
+        "-a", "--n_ant", default=None, type=int, help="Number of antennas to include."
+    )
+    parser.add_argument(
+        "-i", "--n_int", default=None, type=int, help="Number of integration samples."
+    )
+    parser.add_argument(
+        "-n",
+        "--SEFD",
+        default=None,
+        type=float,
+        help="System Equivalent flux density in Jy. Same across frequency and antennas.",
+    )
+    parser.add_argument(
+        "-dt", "--int_time", default=None, type=float, help="Time step in seconds."
+    )
+    parser.add_argument(
+        "-nt", "--n_time", default=None, type=int, help="Number of time steps."
+    )
+    parser.add_argument(
+        "-o",
+        "--overwrite",
+        default=False,
+        action=argparse.BooleanOptionalAction,
+        help="Overwrite existing observation.",
+    )
+    parser.add_argument(
+        "-st", "--spacetrack", help="Path to Space-Track login details."
+    )
+    parser.add_argument(
+        "-ra", "--ra", type=float, help="Right Ascension of the observation."
+    )
+    args = parser.parse_args()
+    rfi_amp = args.rfi_amp
+    spacetrack_path = args.spacetrack
+
+    config_path = os.path.abspath(args.config_path)
+    work_dir = os.path.split(config_path)[0]
+
+    obs_spec = load_config(config_path, config_type="sim")
+
+    if args.ra is not None:
+        obs_spec["observation"]["ra"] = args.ra
+
+    config_st_path = obs_spec["rfi_sources"]["tle_satellite"]["spacetrack_path"]
+    if spacetrack_path:
+        obs_spec["rfi_sources"]["tle_satellite"]["spacetrack_path"] = os.path.abspath(
+            os.path.join(work_dir, spacetrack_path)
+        )
+    elif config_st_path:
+        config_st_path = get_abs_path(config_st_path, work_dir)
+        obs_spec["rfi_sources"]["tle_satellite"]["spacetrack_path"] = config_st_path
+        spacetrack_path = config_st_path
+
+    obs_spec["rfi_sources"]["tle_satellite"]["power_scale"] *= rfi_amp
+    obs_spec["rfi_sources"]["satellite"]["power_scale"] *= rfi_amp
+    obs_spec["rfi_sources"]["stationary"]["power_scale"] *= rfi_amp
+
+    obs_spec["output"]["overwrite"] = args.overwrite
+
+    if args.n_ant is not None:
+        obs_spec["telescope"]["n_ant"] = args.n_ant
+
+    if args.n_int is not None:
+        obs_spec["observation"]["n_int"] = args.n_int
+
+    suffix = obs_spec["output"]["suffix"]
+    if suffix:
+        suffix = f"{rfi_amp:.1e}RFI_" + suffix
+    else:
+        suffix = f"_{rfi_amp:.1e}RFI"
+        obs_spec["output"]["suffix"] = f"{rfi_amp:.1e}RFI"
+
+    if args.SEFD is not None:
+        obs_spec["observation"]["SEFD"] = args.SEFD
+
+    if args.int_time is not None:
+        obs_spec["observation"]["int_time"] = args.int_time
+
+    if args.n_time is not None:
+        obs_spec["observation"]["n_time"] = args.n_time
+
+    for ast in ["exp", "gauss", "point", "pow_spec"]:
+        obs_spec["ast_sources"][ast]["path"] = get_abs_path(
+            obs_spec["ast_sources"][ast]["path"], work_dir
+        )
+
+    obs_spec["output"]["path"] = get_abs_path(obs_spec["output"]["path"], work_dir)
+    obs_spec["rfi_sources"]["tle_satellite"]["norad_ids_path"] = get_abs_path(
+        obs_spec["rfi_sources"]["tle_satellite"]["norad_ids_path"], work_dir
+    )
+    obs_spec["rfi_sources"]["tle_satellite"]["norad_spec_model"] = get_abs_path(
+        obs_spec["rfi_sources"]["tle_satellite"]["norad_spec_model"], work_dir
+    )
+    obs_spec["telescope"]["enu_path"] = get_abs_path(
+        obs_spec["telescope"]["enu_path"], work_dir
+    )
+    obs_spec["telescope"]["itrf_path"] = get_abs_path(
+        obs_spec["telescope"]["itrf_path"], work_dir
+    )
+
+    return run_sim_config(obs_spec=obs_spec, spacetrack_path=spacetrack_path)
+
+
+if __name__ == "__main__":
+    obs, obs_path = main()

tabsim/scripts/tle_sat_region.py

@@ -0,0 +1,191 @@
+from astropy.coordinates import SkyCoord
+import astropy.units as u
+from regions import CircleSkyRegion, RectangleSkyRegion, TextSkyRegion
+
+import xarray as xr
+import numpy as np
+
+from daskms import xds_from_ms, xds_from_table
+
+import argparse
+
+import os
+
+import matplotlib.colors as mcolors
+
+from tabsim.config import yaml_load
+from tabsim.tle import get_tles_by_id, get_satellite_positions
+from tabsim.jax.coordinates import itrf_to_xyz, mjd_to_jd
+
+from astropy.coordinates import EarthLocation
+from astropy.time import Time
+
+from skyfield.api import load, EarthSatellite, wgs84
+
+
+def sat_radec(tle: list[str], times_jd: np.ndarray, obs_xyz: np.ndarray) -> np.ndarray:
+
+    ts = load.timescale()
+    t = ts.ut1_jd(times_jd)
+
+    satellite = EarthSatellite(tle[0], tle[1], ts=ts)
+    location = EarthLocation(x=obs_xyz[0], y=obs_xyz[1], z=obs_xyz[2], unit="m")
+    observer = wgs84.latlon(
+        location.lat.degree, location.lon.degree, location.height.value
+    )
+    radec = (satellite - observer).at(t).radec()
+
+    return np.array([radec[0]._degrees, radec[1].degrees]).T
+
+
+def main():
+
+    colors = list(mcolors.TABLEAU_COLORS.values())
+    n_c = len(colors)
+
+    parser = argparse.ArgumentParser(
+        description="Extract satellite paths to ds9 region file."
+    )
+    parser.add_argument(
+        "-m",
+        "--ms_path",
+        required=True,
+        help="File path to the tabascal zarr simulation file.",
+    )
+    parser.add_argument(
+        "-d",
+        "--max_d",
+        default=90.0,
+        type=float,
+        help="Maximum angular distance from phase centre (in degrees) of the satellite path to include.",
+    )
+    parser.add_argument(
+        "-st",
+        "--spacetrack",
+        help="Path to YAML config file containing Space-Track login details with 'username' and 'password'.",
+    )
+    parser.add_argument(
+        "-ni", "--norad_ids", help="NORAD IDs of satellites to include."
+    )
+    parser.add_argument(
+        "-np",
+        "--norad_path",
+        help="Path to YAML config file containing list of norad_ids.",
+    )
+    parser.add_argument(
+        "-td", "--tle_dir", default="./tles", help="Path to directory containing TLEs."
+    )
+    parser.add_argument(
+        "-s",
+        "--sim",
+        default=False,
+        action=argparse.BooleanOptionalAction,
+        help="Whether to use simulation coordinates. Only small difference from default skyfield coordinates.",
+    )
+
+    args = parser.parse_args()
+    ms_path = args.ms_path
+    max_d = args.max_d
+    spacetrack = yaml_load(args.spacetrack)
+    tle_dir = args.tle_dir
+    norad_ids = []
+    if args.norad_ids:
+        norad_ids += [int(x) for x in np.atleast_1d(args.norad_ids.split(","))]
+    if args.norad_path:
+        norad_ids += [
+            int(x) for x in np.atleast_1d(str(yaml_load(args.norad_path)).split())
+        ]
+
+    os.makedirs(tle_dir, exist_ok=True)
+
+    if ms_path[-1] == "/":
+        ms_path = ms_path[:-1]
+
+    region_path = os.path.join(os.path.split(ms_path)[0], "satelitte_paths.reg")
+
+    def xyz_to_radec(xyz):
+        if xyz.ndim == 2:
+            xyz = xyz[None, :, :]
+
+        xyz = xyz / np.linalg.norm(xyz, axis=-1, keepdims=True)
+        radec = np.zeros((*xyz.shape[:2], 2))
+        radec[:, :, 0] = np.arctan2(xyz[:, :, 1], xyz[:, :, 0])
+        radec[:, :, 1] = np.arcsin(xyz[:, :, 2])
+
+        return np.rad2deg(radec)
+
+    xds = xds_from_ms(ms_path)[0]
+    times_mjd = xds.TIME.data.compute()
+    times_jd = mjd_to_jd(np.linspace(np.min(times_mjd), np.max(times_mjd), 10))
+    epoch_jd = mjd_to_jd(np.mean(times_mjd))
+
+    tles_df = get_tles_by_id(
+        spacetrack["username"],
+        spacetrack["password"],
+        norad_ids,
+        epoch_jd,
+        tle_dir=tle_dir,
+    )
+    if len(tles_df) == 0:
+        raise ValueError("No TLEs found.")
+
+    tles = np.atleast_2d(tles_df[["TLE_LINE1", "TLE_LINE2"]].values)
+    ids = np.atleast_1d(tles_df["NORAD_CAT_ID"].values)
+    n_tles = len(tles)
+
+    print(f"Found {n_tles} matching TLEs.")
+
+    if n_tles > 0:
+
+        xds_src = xds_from_table(ms_path + "::SOURCE")[0]
+        ra, dec = np.rad2deg(xds_src.DIRECTION.data[0].compute())
+
+        xds_ants = xds_from_table(ms_path + "::ANTENNA")[0]
+        ants_itrf = np.mean(xds_ants.POSITION.data.compute(), axis=0, keepdims=True)
+
+        if args.sim:
+            # ants_xyz = itrf_to_xyz(ants_itrf, gmsa_from_jd(times_jd))[:, 0]
+            ants_xyz = itrf_to_xyz(
+                ants_itrf,
+                Time(times_jd, format="jd").sidereal_time("mean", "greenwich").hour
+                * 15,
+            )[:, 0]
+            rfi_xyz = get_satellite_positions(tles, times_jd)
+            xyz = rfi_xyz - ants_xyz[None, :, :]
+            radec = xyz_to_radec(xyz)
+        else:
+            radec = np.array([sat_radec(tle, times_jd, ants_itrf[0]) for tle in tles])
+
+        c0 = SkyCoord(ra, dec, unit="deg", frame="fk5")
+        c = SkyCoord(radec[:, :, 0], radec[:, :, 1], unit="deg", frame="fk5")
+
+        min_sep = c0.separation(c).min(axis=1).deg
+        print(
+            f"Minimum angular separation from target : {[round(x, 1) for x in min_sep]} deg."
+        )
+        print(
+            f"Only including satellites within {max_d:.1f} degrees of pointing direction."
+        )
+        idx = np.where(min_sep < max_d)[0]
+
+        min_idx = np.argmin(c0.separation(c)[:, :-1], axis=1)
+
+        ang_v = np.diff(radec, axis=1)
+        ang_theta = np.arctan2(ang_v[:, :, 1], ang_v[:, :, 0])
+
+        with open(region_path, "w") as fp:
+            for c_i, i in enumerate(idx):
+                # fp.write(RectangleSkyRegion(c0, 0.05*u.deg, 0.3*u.deg, ang_theta[i,min_idx[i]]*u.rad, visual={"color": colors[c_i%n_c]}).serialize(format="ctrf"))
+                # fp.write(TextSkyRegion(c[i,-1], str(ids[i]), visual={"fontsize": 16, "color": colors[c_i%n_c]}).serialize(format="ds9"))
+                for j in range(len(times_jd)):
+                    fp.write(
+                        CircleSkyRegion(
+                            c[i, j],
+                            radius=0.1 * u.deg,
+                            visual={"color": colors[c_i % n_c]},
+                        ).serialize(format="ds9")
+                    )
+
+
+if __name__ == "__main__":
+    main()

tabsim/sky.py

@@ -0,0 +1,137 @@
+import numpy as np
+import jax.numpy as jnp
+import dask.array as da
+
+
+def uniform_points_disk(radius: float, n_src: int, random_seed: int | None = None):
+    """
+    Generate uniformly distributed random points on a disk.
+
+    Parameters:
+    -----------
+    radius: float
+        Radius of the disk.
+    n_src: int
+        The number of sources/points to generate.
+    seed: int
+        Random number generator seed/key.
+
+    Returns:
+    --------
+    points: array_like (2, n_src)
+        The coordinate positions of the random points centred on (0,0).
+    """
+    rng = np.random.default_rng(random_seed)
+    r = jnp.sqrt(rng.uniform(low=0.0, high=radius**2, size=(n_src,)))
+    theta = rng.uniform(low=0.0, high=2.0 * jnp.pi, size=(n_src,))
+
+    return r * jnp.array([jnp.cos(theta), jnp.sin(theta)])
+
+
+def random_power_law(
+    n_src: int,
+    I_min: float = 1e-4,
+    I_max: float = 1.0,
+    alpha: float = 1.6,
+    random_seed: int = None,
+):
+    """Generate a random power law distribution of source fluxes with minimum source
+    flux defined by `I0`.
+
+    Parameters:
+    -----------
+    n_src: int
+        Number of source fluxes to draw.
+    I0: float
+        Minimum source flux.
+    alpha: float
+        Power law index.
+
+    Returns:
+    --------
+    I: array_like (n_src,)
+        Array of source fluxes."""
+
+    def inv_cdf(x):
+        return I_min * (1.0 - x) ** (1.0 / (1.0 - alpha))
+
+    rng = np.random.default_rng(random_seed)
+    rand_unif = rng.uniform(size=(n_src,))
+    I = np.array(inv_cdf(rand_unif))
+    while np.any(I > I_max):
+        idx = np.where(I > I_max)[0]
+        rand_unif = rng.uniform(size=(idx.shape[0],))
+        I[idx] = inv_cdf(rand_unif)
+
+    return I
+
+
+def generate_random_sky(
+    n_src: int,
+    freqs: jnp.ndarray,
+    min_I: float = 1e-4,
+    max_I: float = 1.0,
+    I_power_law: float = 1.6,
+    spec_idx_mean: float = 0.7,
+    spec_idx_std: float = 0.2,
+    fov: float = 1.0,
+    beam_width: float = 0.0,
+    random_seed: int = None,
+    n_beam: int = 3,
+) -> tuple:
+    """
+    Generate uniformly distributed point sources inside the field of view with
+    a power law intensity distribution. Setting the beam width will make
+    sure souces are separated by 5 beam widths apart.
+
+    Parameters:
+    -----------
+    n_src: int
+        Number of sources to generate.
+    mean_I: float
+        Mean intensity of the sources.
+    freqs: array_like (n_freq,)
+        Frequencies to generate the sources at.
+    spec_idx_mean: float
+        Mean spectral index of the sources.
+    spec_idx_std: float
+        Standard deviation of the spectral index of the sources.
+    fov: float
+        Field of view to generate positions within. Same units as beam_width.
+    beam_width: float
+        Width of the resolving beam to ensure sources do not overlap. Sources will be
+        separated by >5*`beam_width`. Same units as fov.
+    random_seed: int
+        Random number generator seed/key.
+    n_beam: int
+        Number of beam_widths apart sources should be separated.
+
+    Returns:
+    --------
+    I: array_like (n_src, n_freq)
+        The sources intensities.
+    delta_ra: array_like
+        The sources right ascensions relative to (0,0).
+    delta_dec: array_like
+        The sources declinations relative to (0,0).
+    """
+    rng = np.random.default_rng(random_seed)
+
+    I = da.atleast_1d(random_power_law(n_src, min_I, max_I, I_power_law, rng))
+    positions = uniform_points_disk(fov / 2.0, 1, rng)
+    while positions.shape[1] < n_src:
+        n_sample = 2 * (n_src - positions.shape[1])
+        new_positions = uniform_points_disk(fov / 2.0, n_sample, rng)
+        positions = np.concatenate([positions, new_positions], axis=1)
+        s1, s2 = np.triu_indices(positions.shape[1], 1)
+        d = np.linalg.norm(positions[:, s1] - positions[:, s2], axis=0)
+        idx = np.where(d < n_beam * beam_width)[0]
+        remove_source_idx = np.unique(jnp.concatenate([s1[idx], s2[idx]]))
+        positions = np.delete(positions, remove_source_idx, axis=1)
+
+    d_ra, d_dec = positions[:, :n_src]
+
+    spectral_indices = rng.normal(loc=spec_idx_mean, scale=spec_idx_std, size=(n_src,))
+    I = I[:, None] * ((freqs[None, :] / freqs[0]) ** -spectral_indices[:, None])
+
+    return I, d_ra, d_dec