tabsim 0.0.4__tar.gz → 0.0.6__tar.gz

{tabsim-0.0.4 → tabsim-0.0.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tabsim
-Version: 0.0.4
+Version: 0.0.6
 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.4 → tabsim-0.0.6}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tabsim"
-version = "0.0.4"
+version = "0.0.6"
 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.4 → tabsim-0.0.6}/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
@@ -256,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)
@@ -668,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"]:
@@ -808,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))
 
@@ -843,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))
 
@@ -933,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}")
@@ -947,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.4 → tabsim-0.0.6}/tabsim/dask/coordinates.py

@@ -1,5 +1,7 @@
 from tabsim.jax import coordinates as coord
 
+from jax import jit
+
 import dask.array as da
 from dask.delayed import delayed
 from dask.array.core import Array
@@ -22,7 +24,7 @@ def radec_to_lmn(ra: Array, dec: Array, phase_centre: Array) -> Array:
         {
             "lmn": (
                 ["src", "lmn_space"],
-                da.zeros(
+                da.zeros(  # type: ignore
                     shape=(n_src, 3),
                     chunks=(src_chunk, 3),
                     dtype=float,
@@ -32,7 +34,7 @@ def radec_to_lmn(ra: Array, dec: Array, phase_centre: Array) -> Array:
     )
 
     def _radec_to_lmn(ds):
-        lmn = delayed(coord.radec_to_lmn, pure=True)(
+        lmn = delayed(jit(coord.radec_to_lmn), pure=True)(
             ds.ra.data, ds.dec.data, ds.phase_centre.data
         ).compute()
         ds_out = xr.Dataset({"lmn": (["src", "lmn_space"], lmn)})
@@ -61,7 +63,7 @@ def radec_to_XYZ(ra: Array, dec: Array) -> Array:
         {
             "XYZ": (
                 ["src", "space"],
-                da.zeros(
+                da.zeros(  # type: ignore
                     shape=(n_src, 3),
                     chunks=(src_chunk, 3),
                     dtype=float,
@@ -71,7 +73,9 @@ def radec_to_XYZ(ra: Array, dec: Array) -> Array:
     )
 
     def _radec_to_XYZ(ds):
-        XYZ = delayed(coord.radec_to_XYZ, pure=True)(ds.ra.data, ds.dec.data).compute()
+        XYZ = delayed(jit(coord.radec_to_XYZ), pure=True)(
+            ds.ra.data, ds.dec.data
+        ).compute()
         ds_out = xr.Dataset({"XYZ": (["src", "space"], XYZ)})
         return ds_out
 
@@ -98,7 +102,7 @@ def ENU_to_GEO(geo_ref: Array, ENU: Array) -> Array:
         {
             "GEO": (
                 ["bl", "space"],
-                da.zeros(
+                da.zeros(  # type: ignore
                     shape=(n_ant, 3),
                     chunks=(ant_chunk, 3),
                     dtype=float,
@@ -108,7 +112,7 @@ def ENU_to_GEO(geo_ref: Array, ENU: Array) -> Array:
     )
 
     def _ENU_to_GEO(ds):
-        GEO = delayed(coord.ENU_to_GEO, pure=True)(
+        GEO = delayed(jit(coord.ENU_to_GEO), pure=True)(
             ds.geo_ref.data, ds.ENU.data
         ).compute()
         ds_out = xr.Dataset({"GEO": (["ant", "space"], GEO)})
@@ -137,7 +141,7 @@ def GEO_to_XYZ(geo: Array, times: Array) -> Array:
         {
             "XYZ": (
                 ["time", "space"],
-                da.zeros(
+                da.zeros(  # type: ignore
                     shape=(n_time, 3),
                     chunks=(time_chunk, 3),
                     dtype=float,
@@ -147,7 +151,9 @@ def GEO_to_XYZ(geo: Array, times: Array) -> Array:
     )
 
     def _GEO_to_XYZ(ds):
-        XYZ = delayed(coord.GEO_to_XYZ, pure=True)(ds.geo.data, ds.times.data).compute()
+        XYZ = delayed(jit(coord.GEO_to_XYZ), pure=True)(
+            ds.geo.data, ds.times.data
+        ).compute()
         ds_out = xr.Dataset({"XYZ": (["time", "space"], XYZ)})
         return ds_out
 
@@ -175,7 +181,7 @@ def GEO_to_XYZ_vmap0(geo: Array, times: Array) -> Array:
         {
             "XYZ": (
                 ["src", "time", "space"],
-                da.zeros(
+                da.zeros(  # type: ignore
                     shape=(n_src, n_time, 3),
                     chunks=(src_chunk, time_chunk, 3),
                     dtype=float,
@@ -185,7 +191,7 @@ def GEO_to_XYZ_vmap0(geo: Array, times: Array) -> Array:
     )
 
     def _GEO_to_XYZ(ds):
-        XYZ = delayed(coord.GEO_to_XYZ_vmap0, pure=True)(
+        XYZ = delayed(jit(coord.GEO_to_XYZ_vmap0), pure=True)(
             ds.geo.data, ds.times.data
         ).compute()
         ds_out = xr.Dataset({"XYZ": (["src", "time", "space"], XYZ)})
@@ -214,7 +220,7 @@ def GEO_to_XYZ_vmap1(geo: Array, times: Array) -> Array:
         {
             "XYZ": (
                 ["time", "ant", "space"],
-                da.zeros(
+                da.zeros(  # type: ignore
                     shape=(n_time, n_ant, 3),
                     chunks=(time_chunk, ant_chunk, 3),
                     dtype=float,
@@ -224,7 +230,7 @@ def GEO_to_XYZ_vmap1(geo: Array, times: Array) -> Array:
     )
 
     def _GEO_to_XYZ(ds):
-        XYZ = delayed(coord.GEO_to_XYZ_vmap1, pure=True)(
+        XYZ = delayed(jit(coord.GEO_to_XYZ_vmap1), pure=True)(
             ds.geo.data, ds.times.data
         ).compute()
         ds_out = xr.Dataset({"XYZ": (["time", "ant", "space"], XYZ)})
@@ -255,7 +261,7 @@ def ITRF_to_XYZ(itrf: Array, gsa: Array) -> Array:
         {
             "xyz": (
                 ["time", "ant", "space"],
-                da.zeros(
+                da.zeros(  # type: ignore
                     shape=(n_time, n_ant, 3),
                     chunks=(time_chunk, ant_chunk, 3),
                     dtype=float,
@@ -265,7 +271,9 @@ def ITRF_to_XYZ(itrf: Array, gsa: Array) -> Array:
     )
 
     def _ITRF_to_XYZ(ds):
-        XYZ = delayed(coord.itrf_to_xyz, pure=True)(ds.itrf.data, ds.gsa.data).compute()
+        XYZ = delayed(jit(coord.itrf_to_xyz), pure=True)(
+            ds.itrf.data, ds.gsa.data
+        ).compute()
         ds_out = xr.Dataset({"xyz": (["time", "ant", "space"], XYZ)})
         return ds_out
 
@@ -294,7 +302,7 @@ def ENU_to_ITRF(ENU: Array, lat: Array, lon: Array, el: Array) -> Array:
         {
             "ITRF": (
                 ["ant", "space"],
-                da.zeros(
+                da.zeros(  # type: ignore
                     shape=(n_ant, 3),
                     chunks=(ant_chunk, 3),
                     dtype=float,
@@ -304,7 +312,7 @@ def ENU_to_ITRF(ENU: Array, lat: Array, lon: Array, el: Array) -> Array:
     )
 
     def _ENU_to_ITRF(ds):
-        ITRF = delayed(coord.enu_to_itrf, pure=True)(
+        ITRF = delayed(jit(coord.enu_to_itrf), pure=True)(
             ds.ENU.data,
             ds.lat.data,
             ds.lon.data,
@@ -335,8 +343,8 @@ def ITRF_to_UVW(
     input = xr.Dataset(
         {
             "itrf": (["ant", "space"], itrf),
-            "h0": (["time"], da.atleast_1d(h0)),
-            "dec": (["cel_space_1"], da.atleast_1d(dec)),
+            "h0": (["time"], da.atleast_1d(h0)),  # type: ignore
+            "dec": (["cel_space_1"], da.atleast_1d(dec)),  # type: ignore
         }
     )
 
@@ -344,7 +352,7 @@ def ITRF_to_UVW(
         {
             "uvw": (
                 ["time", "ant", "space"],
-                da.zeros(
+                da.zeros(  # type: ignore
                     shape=(n_time, n_ant, 3),
                     chunks=(time_chunk, ant_chunk, 3),
                     dtype=float,
@@ -354,7 +362,7 @@ def ITRF_to_UVW(
     )
 
     def _ITRF_to_UVW(ds):
-        uvw = delayed(coord.itrf_to_uvw, pure=True)(
+        uvw = delayed(jit(coord.itrf_to_uvw), pure=True)(
             ds.itrf.data,
             ds.h0.data,
             ds.dec.data,
@@ -381,8 +389,8 @@ def angular_separation(rfi_xyz: Array, ants_xyz: Array, ra: Array, dec: Array) -
         {
             "rfi_xyz": (["src", "time", "space"], rfi_xyz),
             "ants_xyz": (["time", "ant", "space"], ants_xyz),
-            "ra": (["cel_space_0"], da.from_array([ra])),
-            "dec": (["cel_space_1"], da.from_array([dec])),
+            "ra": (["cel_space_0"], da.from_array([ra])),  # type: ignore
+            "dec": (["cel_space_1"], da.from_array([dec])),  # type: ignore
         }
     )
 
@@ -390,7 +398,7 @@ def angular_separation(rfi_xyz: Array, ants_xyz: Array, ra: Array, dec: Array) -
         {
             "sep": (
                 ["src", "time", "ant"],
-                da.zeros(
+                da.zeros(  # type: ignore
                     shape=(n_src, n_time, n_ant),
                     chunks=(src_chunk, time_chunk, ant_chunk),
                     dtype=float,
@@ -400,7 +408,7 @@ def angular_separation(rfi_xyz: Array, ants_xyz: Array, ra: Array, dec: Array) -
     )
 
     def _angular_separation(ds):
-        sep = delayed(coord.angular_separation, pure=True)(
+        sep = delayed(jit(coord.angular_separation), pure=True)(
             ds.rfi_xyz.data, ds.ants_xyz.data, ds.ra.data, ds.dec.data
         ).compute()
         ds_out = xr.Dataset({"sep": (["src", "time", "ant"], sep)})
@@ -441,7 +449,7 @@ def orbit_vmap(
         {
             "orbit": (
                 ["src", "time", "space"],
-                da.zeros(
+                da.zeros(  # type: ignore
                     shape=(n_src, n_time, 3),
                     chunks=(src_chunk, time_chunk, 3),
                     dtype=float,
@@ -451,7 +459,7 @@ def orbit_vmap(
     )
 
     def _orbit_vmap(ds):
-        orbit = delayed(coord.orbit_vmap, pure=True)(
+        orbit = delayed(jit(coord.orbit_vmap), pure=True)(
             ds.times.data,
             ds.elevation.data,
             ds.inclination.data,

{tabsim-0.0.4 → tabsim-0.0.6}/tabsim/dask/interferometry.py

@@ -1,10 +1,16 @@
-import dask.array as da
+from jax import jit
+
 import numpy as np
 import xarray as xr
+
+import dask.array as da
 from dask.delayed import delayed
 from dask.array.core import Array
+
 from tabsim.jax import interferometry as itf
 
+from typing import Optional
+
 
 def astro_vis(sources: Array, uvw: Array, lmn: Array, freqs: Array) -> Array:
     """Calculate visibilities from sources, uvw, lmn, and freqs.
@@ -51,7 +57,7 @@ def astro_vis(sources: Array, uvw: Array, lmn: Array, freqs: Array) -> Array:
         {
             "vis": (
                 ["time", "bl", "freq"],
-                da.zeros(
+                da.zeros(  # type: ignore
                     shape=(n_time, n_bl, n_freq),
                     chunks=(time_chunk, bl_chunk, freq_chunk),
                     dtype=complex,
@@ -61,7 +67,7 @@ def astro_vis(sources: Array, uvw: Array, lmn: Array, freqs: Array) -> Array:
     )
 
     def _astro_vis(ds):
-        vis = delayed(itf.astro_vis)(
+        vis = delayed(jit(itf.astro_vis), pure=True)(
             ds.I.data, ds.uvw.data, ds.lmn.data, ds.freqs.data
         ).compute()
         ds_out = xr.Dataset({"vis": (["time", "bl", "freq"], vis)})
@@ -122,7 +128,7 @@ def astro_vis_gauss(
         {
             "vis": (
                 ["time", "bl", "freq"],
-                da.zeros(
+                da.zeros(  # type: ignore
                     shape=(n_time, n_bl, n_freq),
                     chunks=(time_chunk, bl_chunk, freq_chunk),
                     dtype=complex,
@@ -132,7 +138,7 @@ def astro_vis_gauss(
     )
 
     def _astro_vis_gauss(ds):
-        vis = delayed(itf.astro_vis_gauss)(
+        vis = delayed(jit(itf.astro_vis_gauss), pure=True)(
             ds.I.data,
             ds.major.data,
             ds.minor.data,
@@ -191,7 +197,7 @@ def astro_vis_exp(
         {
             "vis": (
                 ["time", "bl", "freq"],
-                da.zeros(
+                da.zeros(  # type: ignore
                     shape=(n_time, n_bl, n_freq),
                     chunks=(time_chunk, bl_chunk, freq_chunk),
                     dtype=complex,
@@ -201,7 +207,7 @@ def astro_vis_exp(
     )
 
     def _astro_vis_exp(ds):
-        vis = delayed(itf.astro_vis_exp)(
+        vis = delayed(jit(itf.astro_vis_exp), pure=True)(
             ds.I.data, ds.sigmas.data, ds.uvw.data, ds.lmn.data, ds.freqs.data
         ).compute()
         ds_out = xr.Dataset({"vis": (["time", "bl", "freq"], vis)})
@@ -259,7 +265,7 @@ def rfi_vis(
         {
             "vis": (
                 ["time", "bl", "freq"],
-                da.zeros(
+                da.zeros(  # type: ignore
                     shape=(n_time, n_bl, n_freq),
                     chunks=(time_chunk, bl_chunk, freq_chunk),
                     dtype=complex,
@@ -269,7 +275,7 @@ def rfi_vis(
     )
 
     def _rfi_vis(ds):
-        vis = delayed(itf.rfi_vis)(
+        vis = delayed(jit(itf.rfi_vis), pure=True)(
             ds.app_amplitude.data,
             ds.c_distances.data,
             ds.freqs.data,
@@ -299,7 +305,7 @@ def ants_to_bl(G: Array, a1: Array, a2: Array) -> Array:
         {
             "G_bl": (
                 ["time", "bl", "freq"],
-                da.zeros(
+                da.zeros(  # type: ignore
                     (n_time, n_bl, n_freq), chunks=(time_chunk, bl_chunk, freq_chunk)
                 ),
             )
@@ -307,7 +313,9 @@ def ants_to_bl(G: Array, a1: Array, a2: Array) -> Array:
     )
 
     def _ants_to_bl(ds):
-        G_bl = delayed(itf.ants_to_bl)(ds.G.data, ds.a1.data, ds.a2.data).compute()
+        G_bl = delayed(jit(itf.ants_to_bl), pure=True)(
+            ds.G.data, ds.a1.data, ds.a2.data
+        ).compute()
         ds_out = xr.Dataset({"G_bl": (["time", "bl", "freq"], G_bl)})
         return ds_out
 
@@ -330,7 +338,7 @@ def airy_beam(theta, freqs, dish_d):
         {
             "theta": (["src", "time", "ant"], theta),
             "freqs": (["freq"], freqs),
-            "dish_d": (["space_0"], da.from_array([dish_d])),
+            "dish_d": (["space_0"], da.from_array([dish_d])),  # type: ignore
         }
     )
 
@@ -338,7 +346,7 @@ def airy_beam(theta, freqs, dish_d):
         {
             "beam": (
                 ["src", "time", "ant", "freq"],
-                da.zeros(
+                da.zeros(  # type: ignore
                     (n_src, n_time, n_ant, n_freq),
                     chunks=(src_chunk, time_chunk, ant_chunk, freq_chunk),
                 ),
@@ -347,7 +355,7 @@ def airy_beam(theta, freqs, dish_d):
     )
 
     def _airy_beam(ds):
-        beam = delayed(itf.airy_beam)(
+        beam = delayed(itf.airy_beam, pure=True)(
             ds.theta.data, ds.freqs.data, ds.dish_d.data
         ).compute()
         ds_out = xr.Dataset({"beam": (["src", "time", "ant", "freq"], beam)})
@@ -376,7 +384,7 @@ def Pv_to_Sv(Pv, d):
         {
             "Sv": (
                 ["src", "time", "ant", "freq"],
-                da.zeros(
+                da.zeros(  # type: ignore
                     (n_src, n_time, n_ant, n_freq),
                     chunks=(src_chunk, time_chunk, ant_chunk, freq_chunk),
                 ),
@@ -406,7 +414,7 @@ def add_noise(vis: Array, noise_std: float, key: int):
 
 
 def SEFD_to_noise_std(SEFD, chan_width, t_int):
-    noise_std = SEFD / da.sqrt(chan_width * t_int)
+    noise_std = SEFD / da.sqrt(chan_width * t_int)  # type: ignore
     return noise_std
 
 
@@ -417,13 +425,13 @@ def SEFD_to_noise_std(SEFD, chan_width, t_int):
 #     return times_fine
 
 
-def int_sample_times(times, n_int_samples: int, int_time: float = None):
+def int_sample_times(times, n_int_samples: int, int_time: Optional[float] = None):
 
     n_time = len(times)
     time_range = times[-1] - times[0]
     int_time = time_range / (n_time - 1) if not int_time else int_time
     n_time_fine = n_time * n_int_samples
-    times_fine = da.linspace(
+    times_fine = da.linspace(  # type: ignore
         -int_time / 2, time_range + int_time / 2, n_time_fine, endpoint=False
     )
 
@@ -444,7 +452,7 @@ def generate_gains(
     times = times[:, None, None] - times[0]
 
     # Generate the initial gain values
-    G0 = G0_mean * da.exp(
+    G0 = G0_mean * da.exp(  # type: ignore
         1.0j * rng.uniform(low=-np.pi / 2, high=np.pi / 2, size=(1, n_ant, n_freq))
     ) + (
         rng.normal(scale=G0_std, size=(1, n_ant, n_freq))
@@ -455,9 +463,9 @@ def generate_gains(
     gain_amp = rng.normal(scale=Gt_std_amp, size=(1, n_ant, 1)) * times
     gain_phase = rng.normal(scale=Gt_std_phase, size=(1, n_ant, 1)) * times
     # Generate the gain time series
-    gain_ants = G0 + gain_amp * da.exp(1.0j * gain_phase)
+    gain_ants = G0 + gain_amp * da.exp(1.0j * gain_phase)  # type: ignore
     # Set the gain on the last antenna to have zero phase (reference antenna)
-    gain_ants[:, -1, :] = da.abs(gain_ants[:, -1, :])
+    gain_ants[:, -1, :] = da.abs(gain_ants[:, -1, :])  # type: ignore
 
     return gain_ants
 
@@ -479,7 +487,7 @@ def generate_fourier_gains(
     times = times_jd[None, :, None, None] * 24
 
     # Generate the initial gain values
-    G0 = G0_mean * da.exp(
+    G0 = G0_mean * da.exp(  # type: ignore
         1.0j * rng.uniform(low=-np.pi / 2, high=np.pi / 2, size=(1, n_ant, n_freq))
     ) + (
         rng.normal(scale=G0_std / 100, size=(1, n_ant, n_freq))
@@ -513,9 +521,9 @@ def generate_fourier_gains(
         axis=0,
     ) * np.sqrt(2 * N)
     # Generate the gain time series
-    gain_ants = G0 + gain_amp * da.exp(1.0j * gain_phase)
+    gain_ants = G0 + gain_amp * da.exp(1.0j * gain_phase)  # type: ignore
     # Set the gain on the last antenna to have zero phase (reference antenna)
-    gain_ants[:, -1, :] = da.abs(gain_ants[:, -1, :])
+    gain_ants[:, -1, :] = da.abs(gain_ants[:, -1, :])  # type: ignore
 
     return gain_ants
 
@@ -545,7 +553,7 @@ def apply_gains(
         {
             "vis_obs": (
                 ["time", "bl", "freq"],
-                da.zeros(
+                da.zeros(  # type: ignore
                     (n_time, n_bl, n_freq),
                     chunks=(time_chunk, bl_chunk, freq_chunk),
                     dtype=vis_ast.dtype,
@@ -567,6 +575,6 @@ def apply_gains(
 
 
 def time_avg(vis, n_int_samples):
-    return da.mean(
-        da.reshape(vis, (-1, n_int_samples, vis.shape[1], vis.shape[2])), axis=1
+    return da.mean(  # type: ignore
+        da.reshape(vis, (-1, n_int_samples, vis.shape[1], vis.shape[2])), axis=1  # type: ignore
     )