tabsim 0.0.3__tar.gz → 0.0.5__tar.gz

{tabsim-0.0.3 → tabsim-0.0.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tabsim
-Version: 0.0.3
+Version: 0.0.5
 Summary: Trajectory-based Radio Frequency Interference (RFI) simulations for radio interferometeric data.
 Author-email: Chris Finlay <christopher.finlay@epfl.ch>
 Maintainer-email: Chris Finlay <christopher.finlay@epfl.ch>
@@ -680,11 +680,11 @@ License:                     GNU GENERAL PUBLIC LICENSE
         <https://www.gnu.org/licenses/why-not-lgpl.html>.
         
 Project-URL: Repository, https://github.com/chrisfinlay/tab-sim.git
-Requires-Python: <3.13,>=3.9
+Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: jax
-Requires-Dist: dask<=2024.10.0
+Requires-Dist: dask
 Requires-Dist: xarray
 Requires-Dist: zarr<3.0.0
 Requires-Dist: dask-ms>=0.2.21
@@ -694,7 +694,7 @@ Requires-Dist: matplotlib
 Requires-Dist: astropy
 Requires-Dist: spacetrack
 Requires-Dist: skyfield
-Requires-Dist: numpy<2.0.0
+Requires-Dist: numpy
 Provides-Extra: gpu
 Requires-Dist: jax[cuda12]; extra == "gpu"
 Provides-Extra: sat

{tabsim-0.0.3 → tabsim-0.0.5}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tabsim"
-version = "0.0.3"
+version = "0.0.5"
 license = {file = "LICENSE"}
 readme = "README.md"
 # dynamic = ["version", "readme"]
@@ -17,11 +17,11 @@ maintainers = [
 ]
 description = "Trajectory-based Radio Frequency Interference (RFI) simulations for radio interferometeric data."
 
-requires-python = ">=3.9,<3.13"
+requires-python = ">=3.9"
 dependencies = [
-    "jax", "dask<=2024.10.0", "xarray", "zarr<3.0.0", 
+    "jax", "dask", "xarray", "zarr<3.0.0", 
     "dask-ms>=0.2.21", "scipy", "tqdm", "matplotlib", 
-    "astropy", "spacetrack", "skyfield", "numpy<2.0.0", 
+    "astropy", "spacetrack", "skyfield", "numpy", 
 ]
 
 [project.optional-dependencies]
@@ -51,3 +51,4 @@ tle-region = "tabsim.scripts.tle_sat_region:main"
 ast-region = "tabsim.scripts.ast_region:main"
 flag-data = "tabsim.scripts.flag_data:main"
 sim-vis = "tabsim.scripts.sim_vis:main"
+write-config = "tabsim.scripts.write_config_from_ms:main"

{tabsim-0.0.3 → tabsim-0.0.5}/tabsim/config.py

@@ -18,7 +18,6 @@ from astropy.time import Time
 
 from tabsim.dask.observation import Observation
 from tabsim.sky import generate_random_sky
-from tabsim.plot import plot_uv, plot_src_alt, plot_angular_seps
 from tabsim.write import write_ms, mk_obs_name, mk_obs_dir
 from tabsim.jax.coordinates import calculate_fringe_frequency, jd_to_mjd
 from tabsim.tle import get_visible_satellite_tles, id_generator
@@ -28,14 +27,9 @@ from daskms import xds_from_ms
 from tqdm import tqdm
 from importlib.resources import files
 
-config_dir = files("tabsim.data").joinpath("config_files").__str__()
 
-sim_base_config_path = os.path.join(config_dir, "sim_config_base.yaml")
-tab_base_config_path = os.path.join(config_dir, "tab_config_base.yaml")
-# extract_base_config_path = os.path.join(config_dir, "extract_config_base.yaml")
-# pow_spec_base_config_path = os.path.join(
-#     config_dir, "extract_pow_spec_config_base.yaml"
-# )
+sky_dir = files("tabsim.data").joinpath("sky").__str__()
+
 
 JD0 = 2459997.079914223  # 2023-02-21 13:55:04.589 UTC => GMSA = 0
 MJD0 = JD0 - 2400000.5
@@ -115,24 +109,67 @@ def load_config(path: str, config_type: str = "sim") -> dict:
     dict
         Configuration dictionary.
     """
+    config_dir = files("tabsim.data").joinpath("config").__str__()
+
+    sim_base_config_path = os.path.join(config_dir, "sim_config_base.yaml")
+    tab_base_config_path = os.path.join(config_dir, "tab_config_base.yaml")
+    # extract_base_config_path = os.path.join(config_dir, "extract_config_base.yaml")
+    # pow_spec_base_config_path = os.path.join(
+    #     config_dir, "extract_pow_spec_config_base.yaml"
+    # )
 
     config = yaml_load(path)
     if config_type == "sim":
         base_config = yaml_load(sim_base_config_path)
     elif config_type == "tab":
         base_config = yaml_load(tab_base_config_path)
-    elif config_type == "extract":
-        base_config = yaml_load(extract_base_config_path)
-    elif config_type == "pow_spec":
-        base_config = yaml_load(pow_spec_base_config_path)
+    # elif config_type == "extract":
+    #     base_config = yaml_load(extract_base_config_path)
+    # elif config_type == "pow_spec":
+    #     base_config = yaml_load(pow_spec_base_config_path)
     else:
-        ValueError(
-            "A config type must be specified. Options are {sim, tab, extract, pow_spec}."
-        )
+        raise ValueError("A config type must be specified. Options are {sim, tab}.")
 
     return deep_update(base_config, config)
 
 
+def get_telescope_definitions(tel_name: str):
+
+    tel_dir = files("tabsim.data").joinpath("telescopes").__str__()
+    tel_def = yaml_load(os.path.join(tel_dir, "_telescopes.yaml"))
+
+    tel_name = tel_name.lower()
+    telescope_list = tel_def.keys()
+
+    if tel_name in telescope_list:
+        tel_def = tel_def[tel_name]
+        tel_def["itrf_path"] = os.path.join(tel_dir, tel_def["itrf_path"])
+    else:
+        raise ValueError(
+            f"{tel_name} is not a pre-defined telescope within tab-sim. Choose from {telescope_list}"
+        )
+
+    return tel_def
+
+
+def get_rfi_definitions():
+
+    rfi_dir = files("tabsim.data").joinpath("rfi").__str__()
+
+    rfi_def = {
+        "tle_satellite": {
+            "tle_dir": os.path.join(rfi_dir, "tles"),
+            "norad_spec_model": os.path.join(rfi_dir, "norad_satellite.rfimodel"),
+        },
+        "stationary": {
+            "geo_path": os.path.join(rfi_dir, "stationary.loc"),
+            "spec_path": os.path.join(rfi_dir, "stationary.rfimodel"),
+        },
+    }
+
+    return rfi_def
+
+
 def load_sky_model(file_path: str, freqs: Array, src_type: str) -> tuple:
     """Read in a skymodel file and return parameters to be used in tabascal simulator.
     A skymodel file is a csv file with the following columns:
@@ -218,9 +255,7 @@ def load_obs(obs_spec: dict) -> Observation:
         start_time_mjd = jd_to_mjd(obs_["start_time_jd"])
     elif obs_["start_time_isot"]:
 
-        start_time_mjd = jd_to_mjd(
-            Time(obs_["start_time_isot"], format="isot", scale="ut1").jd
-        )
+        start_time_mjd = jd_to_mjd(Time(obs_["start_time_isot"], format="isot").jd)
     elif obs_["start_time_lha"] is not None:
         gsa = obs_["start_time_lha"] - tel_["longitude"] + obs_["ra"]
         start_time_mjd = MJD0 + (gsa / 360)
@@ -630,6 +665,8 @@ def add_gains(obs: Observation, obs_spec: dict) -> None:
 
 def plot_diagnostics(obs: Observation, obs_spec: dict, save_path: str) -> None:
 
+    from tabsim.plot import plot_uv, plot_src_alt, plot_angular_seps
+
     diag_ = obs_spec["diagnostics"]
 
     if diag_["src_alt"]:
@@ -729,20 +766,7 @@ def save_inputs(obs: Observation, obs_spec: dict, save_path: str) -> None:
         if path is not None:
             shutil.copy(path, save_path)
 
-    key = (
-        2
-        * [
-            "tle_satellite",
-        ]
-        + 2
-        * [
-            "satellite",
-        ]
-        + 2
-        * [
-            "stationary",
-        ]
-    )
+    key = 2 * ["tle_satellite"] + 2 * ["satellite"] + 2 * ["stationary"]
     subkey = [
         "norad_ids_path",
         "norad_spec_model",
@@ -783,11 +807,14 @@ def print_fringe_freq_sat(obs: Observation):
         for i in range(obs.n_rfi_satellite)
     ]
     fringe_freq = [calculate_fringe_frequency(**f_params) for f_params in fringe_params]
+    noise_std = obs.noise_std.mean().compute()
+    if noise_std == 0:
+        noise_std = 1.0
     f_sample = (
         np.pi
         * np.max(np.abs(fringe_freq))
         * np.max(obs.rfi_satellite_A_app)
-        / np.sqrt(6 * obs.noise_std.mean()).compute()
+        / np.sqrt(6 * noise_std)
     )
     n_int = int(np.ceil(obs.int_time * f_sample))
 
@@ -818,11 +845,14 @@ def print_fringe_freq_tle_sat(obs: Observation):
         for i in range(obs.n_rfi_tle_satellite)
     ]
     fringe_freq = [calculate_fringe_frequency(**f_params) for f_params in fringe_params]
+    noise_std = obs.noise_std.mean().compute()
+    if noise_std == 0:
+        noise_std = 1.0
     f_sample = (
         np.pi
         * np.max(np.abs(fringe_freq))
         * np.max(obs.rfi_tle_satellite_A_app)
-        / np.sqrt(6 * obs.noise_std.mean()).compute()
+        / np.sqrt(6 * noise_std)
     )
     n_int = int(np.ceil(obs.int_time * f_sample))
 
@@ -835,6 +865,20 @@ def print_fringe_freq_tle_sat(obs: Observation):
     print(f"Recommended n_int is >=     : {n_int:.0f} ({obs.n_int_samples} used)")
 
 
+def check_telescope_defintion(tel_def: dict):
+
+    dish_d = tel_def["dish_d"]
+
+    if tel_def["itrf_path"] and dish_d:
+        return True
+    elif (
+        tel_def["longitude"] and tel_def["latitude"] and tel_def["enu_path"] and dish_d
+    ):
+        return True
+    else:
+        return False
+
+
 def run_sim_config(
     obs_spec: Optional[dict] = None,
     config_path: Optional[str] = None,
@@ -843,7 +887,6 @@ def run_sim_config(
 
     from tabsim.tle import id_generator
 
-
     log_path = f"log_sim_{id_generator()}.txt"
     log = open(log_path, "w")
     backup = sys.stdout
@@ -852,11 +895,17 @@ def run_sim_config(
     start = datetime.now()
     print(datetime.now())
 
-    if config_path is not None:
+    if config_path:
         obs_spec = load_config(config_path, config_type="sim")
     elif obs_spec is None:
-        print("obs_spec or path must be defined.")
-        return None
+        raise ValueError("obs_spec or config_path must be defined.")
+
+    rfi_def = get_rfi_definitions()
+    obs_spec["rfi_sources"] = deep_update(obs_spec["rfi_sources"], rfi_def)
+
+    if not check_telescope_defintion(obs_spec["telescope"]):
+        tel_def = get_telescope_definitions(obs_spec["telescope"]["name"])
+        obs_spec["telescope"] = deep_update(obs_spec["telescope"], tel_def)
 
     obs = load_obs(obs_spec)
     add_astro_sources(obs, obs_spec)
@@ -889,9 +938,26 @@ def run_sim_config(
     print()
     print(f"Writing data to : {save_path}")
 
-    plot_diagnostics(obs, obs_spec, save_path)
+    plot_cond = np.array(
+        [
+            obs_spec["diagnostics"]["rfi_seps"],
+            obs_spec["diagnostics"]["src_alt"],
+            obs_spec["diagnostics"]["uv_cov"],
+        ]
+    )
+    if np.any(plot_cond):
+        plot_diagnostics(obs, obs_spec, save_path)
+    else:
+        print("\nNo diagnostic plots.")
+
     save_data(obs, obs_spec, zarr_path, ms_path)
 
+    if obs_spec["output"]["accumulate_ms"] is not None:
+        from tabsim.write import add_to_ms
+
+        xds = xr.open_zarr(zarr_path)
+        add_to_ms(xds, obs_spec["output"]["accumulate_ms"])
+
     end = datetime.now()
     print()
     print(f"Total simulation time : {end - start}")
@@ -903,4 +969,11 @@ def run_sim_config(
     os.remove(log_path)
     sys.stdout = backup
 
-    return obs, save_path
+    if (
+        not obs_spec["output"]["keep_sim"]
+        and obs_spec["output"]["accumulate_ms"] is not None
+    ):
+        shutil.rmtree(save_path)
+        return obs, save_path
+    else:
+        return obs, save_path

{tabsim-0.0.3 → tabsim-0.0.5}/tabsim/dask/observation.py

@@ -1,9 +1,9 @@
-import jax.numpy as jnp
 from jax import config, Array
 
 import dask.array as da
 import numpy as np
 import xarray as xr
+from typing import Optional
 
 from tabsim.dask.coordinates import (
     ENU_to_ITRF,
@@ -75,12 +75,12 @@ class Telescope(object):
         latitude: float,
         longitude: float,
         elevation: float = 0.0,
-        ENU_array: Array = None,
-        ENU_path: str = None,
-        ITRF_array: Array = None,
-        ITRF_path: str = None,
-        tel_name: str = None,
-        n_ant: int = None,
+        ENU_array: Optional[Array] = None,
+        ENU_path: Optional[str] = None,
+        ITRF_array: Optional[Array] = None,
+        ITRF_path: Optional[str] = None,
+        tel_name: Optional[str] = None,
+        n_ant: Optional[int] = None,
     ):
         self.tel_name = tel_name
         self.latitude = da.asarray(latitude)
@@ -132,16 +132,16 @@ Elevation : {elevation}\n"""
         self.ITRF = ENU_to_ITRF(self.ENU, self.latitude, self.longitude, self.elevation)
 
     def createArrayITRF(self, ITRF_array, ITRF_path):
+
         if ITRF_array is not None:
             self.ITRF = ITRF_array
         elif ITRF_path is not None:
             self.ITRF = np.loadtxt(ITRF_path, usecols=(0, 1, 2), max_rows=self.n_ant)
         else:
-            self.ITRF = None
-            msg = """Error : ITRF antenna coordinates are needed either in an 
+            msg = """ITRF antenna coordinates are needed either in an 
                      array or as a csv like file."""
-            print(msg)
-            return
+            raise ValueError(msg)
+
         self.ITRF = da.asarray(self.ITRF)
         self.GEO_ants = da.asarray(itrf_to_geo(self.ITRF.compute()))
 
@@ -205,11 +205,11 @@ class Observation(Telescope):
         times_mjd: Array,
         int_time: float = 2.0,
         chan_width: float = 209e3,
-        ENU_array: Array = None,
-        ENU_path: str = None,
-        ITRF_array: Array = None,
-        ITRF_path: str = None,
-        n_ant: int = None,
+        ENU_array: Optional[Array] = None,
+        ENU_path: Optional[str] = None,
+        ITRF_array: Optional[Array] = None,
+        ITRF_path: Optional[str] = None,
+        n_ant: Optional[int] = None,
         dish_d: float = 13.5,
         random_seed: int = 0,
         auto_corrs: bool = False,
@@ -242,7 +242,7 @@ class Observation(Telescope):
         self.target_name = target_name
         self.auto_corrs = auto_corrs
 
-        a1, a2 = jnp.triu_indices(self.n_ant, 0 if auto_corrs else 1)
+        a1, a2 = np.triu_indices(self.n_ant, 0 if auto_corrs else 1)
         self.n_bl = len(a1)
 
         self.ant_chunk = self.n_ant
@@ -290,11 +290,17 @@ class Observation(Telescope):
             self.n_int_samples // 2, self.n_time_fine, self.n_int_samples
         ).rechunk(self.time_chunk)
 
+        # self.gsa = da.asarray(
+        #     Time(self.times_mjd_fine.compute(), format="mjd")
+        #     .sidereal_time("mean", "greenwich")
+        #     .hour
+        #     * 15,
+        #     chunks=(self.time_fine_chunk,),
+        # )
+        from tabsim.jax.coordinates import mjd_to_gast
+
         self.gsa = da.asarray(
-            Time(self.times_mjd_fine.compute(), format="mjd")
-            .sidereal_time("mean", "greenwich")
-            .hour
-            * 15,
+            mjd_to_gast(self.times_mjd_fine.compute()),
             chunks=(self.time_fine_chunk,),
         )
         self.gha = (self.gsa - self.ra) % 360
@@ -313,7 +319,9 @@ class Observation(Telescope):
         self.dish_d = da.asarray(dish_d)
         self.fov = beam_size(dish_d, freqs.max(), fwhp=False)
 
-        self.ants_uvw = ITRF_to_UVW(self.ITRF, self.gha, self.dec)
+        self.ants_uvw = da.asarray(
+            np.array(ITRF_to_UVW(self.ITRF, self.gha, self.dec).compute())
+        )
 
         if no_w:
             self.ants_uvw[:, :, -1] = 0.0
@@ -326,17 +334,17 @@ class Observation(Telescope):
         self.vis_rfi = da.zeros(
             shape=(self.n_time, self.n_bl, self.n_freq),
             chunks=(self.time_chunk, self.bl_chunk, self.freq_chunk),
-            dtype=jnp.complex128,
+            dtype=np.complex128,
         )
         self.vis_ast = da.zeros(
             shape=(self.n_time, self.n_bl, self.n_freq),
             chunks=(self.time_chunk, self.bl_chunk, self.freq_chunk),
-            dtype=jnp.complex128,
+            dtype=np.complex128,
         )
         self.gains_ants = da.ones(
             shape=(self.n_time, self.n_ant, self.n_freq),
             chunks=(self.time_chunk, self.ant_chunk, self.freq_chunk),
-            dtype=jnp.complex128,
+            dtype=np.complex128,
         )
         self.random_seed = random_seed
 
@@ -486,7 +494,7 @@ Number of stationary RFI :  {n_stat}"""
 
         self.ast_p_I.append(I)
         self.ast_p_lmn.append(lmn)
-        self.ast_p_radec.append(jnp.array([ra, dec]))
+        self.ast_p_radec.append(np.array([ra, dec]))
         self.n_p_ast += len(I)
         self.n_ast += len(I)
 
@@ -550,7 +558,7 @@ Number of stationary RFI :  {n_stat}"""
         self.ast_g_pos_angle.append(pos_angle)
         self.ast_g_I.append(I)
         self.ast_g_lmn.append(lmn)
-        self.ast_g_radec.append(jnp.array([ra, dec]))
+        self.ast_g_radec.append(np.array([ra, dec]))
         self.n_g_ast += len(I)
         self.n_ast += len(I)
 
@@ -596,7 +604,7 @@ Number of stationary RFI :  {n_stat}"""
         self.ast_e_major.append(shape)
         self.ast_e_I.append(I)
         self.ast_e_lmn.append(lmn)
-        self.ast_e_radec.append(jnp.array([ra, dec]))
+        self.ast_e_radec.append(np.array([ra, dec]))
         self.n_e_ast += len(I)
         self.n_ast += len(I)
 
@@ -725,6 +733,11 @@ Number of stationary RFI :  {n_stat}"""
             get_satellite_positions(tles, mjd_to_jd(self.times_mjd_fine.compute())),
             chunks=(-1, self.time_fine_chunk, 3),
         )
+        # from tabsim.jax.coordinates import kepler_orbit_many
+        # from tabsim.tle import
+        # rfi_xyz = da.asarray(
+        #     kepler_orbit_many(mjd_to_jd(self.times_mjd_fine.compute()), )
+        # )
         tles = da.asarray(da.atleast_2d(tles), chunks=(-1,))
         # rfi_xyz is shape (n_src,n_time_fine,3)
         # self.ants_xyz is shape (n_time_fine,n_ant,3)
@@ -938,7 +951,7 @@ Number of stationary RFI :  {n_stat}"""
             else:
                 self.flags = (
                     da.abs(self.vis_cal - self.vis_ast)
-                    > 3.0 * da.std(self.vis_model, axis=0)[None, ...]
+                    > 3.0 * da.std(self.vis_ast, axis=0)[None, ...]
                 )
         else:
             self.flags = da.zeros(shape=self.vis_cal.shape, dtype=bool)
@@ -959,7 +972,7 @@ Number of stationary RFI :  {n_stat}"""
         self,
         path: str = "Observation.ms",
         overwrite: bool = False,
-        ds: xr.Dataset = None,
+        ds: Optional[xr.Dataset] = None,
     ):
         """
         Write the visibilities to disk using Measurement Set format.

{tabsim-0.0.3/tabsim/data/config_files → tabsim-0.0.5/tabsim/data/config}/sim_config_base.yaml

@@ -2,6 +2,8 @@
 # This is base simulation config that is loaded and then updated with your provided config.
 # Populated parameters here are the default values used if none are specified in your config.
 ###############################################################################################
+# Astronomical Sources
+#######################
 ast_sources:
 ##########################################
 # Exponential shaped astronomical sources
@@ -120,11 +122,16 @@ observation:
 output:
   zarr: True                  # Optional - Whether to save the output to a zarr file. This contains almost all inputs ot the simulation
   ms: True                    # Optional - Whether to save the output to a Measurement Set file
+  accumulate_ms: null         # Optional - Whether to add the RFI visibilities to a provided MS file. This should be the MS file path to add to.
+  keep_sim: True              # Optional - Whether to keep the simulation files generated by tabsim after adding to RFI visibilities to accumulate_ms. This includes all data including generated MS files. 
   flag_data: True             # Optional - Whether to apply ideal 3sigma flags
   overwrite: False            # Optional - Whether to overwrite a previous simulation with the same name
   path: ./                    # Optional - Path to where the simulation should be saved
   prefix: null                # Optional - Prefix used in the simulation file name
   suffix: null                # Optional - Suffix used in the simulation file name
+##############
+# RFI sources
+##############
 rfi_sources:
 #######################################
 # Circular trajectory based satellites
@@ -164,7 +171,7 @@ rfi_sources:
 telescope:
   latitude: null              # Required - Latitude of the telescope location in decimal degrees. North is +ve
   longitude: null             # Required - Longitude of the telescope in decimal degrees. East from Greenwich meridian
-  elevation: null             # Required - Elevation of the telescope above the WGS84 Earth ellipsoid model in metres.
+  elevation: 0                # Required - Elevation of the telescope above the WGS84 Earth ellipsoid model in metres.
   dish_d: null                # Required - Dish diameter in metres
   enu_path: null              # Required if itrf_path is null - Path to the csv file of the antennas in East, North, Up frame, in metres, relative to the telescope location 
   itrf_path: null             # Required if enu_path is null - Path to the csv file contaning the antenna positions in International Terrestrial Reference Frame in metres.