tabsim 0.0.5__tar.gz → 0.0.6__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tabsim
-Version: 0.0.5
+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>

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

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tabsim"
-version = "0.0.5"
+version = "0.0.6"
 license = {file = "LICENSE"}
 readme = "README.md"
 # dynamic = ["version", "readme"]

{tabsim-0.0.5 → 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.5 → 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
     )

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

@@ -1,4 +1,4 @@
-from jax import config, Array
+from jax import Array
 
 import dask.array as da
 import numpy as np
@@ -44,8 +44,6 @@ from tabsim.write import construct_observation_ds, write_ms
 from tabsim.dask.extras import get_chunksizes
 from tabsim.tle import get_satellite_positions, ants_pos, sat_distance
 
-config.update("jax_enable_x64", True)
-
 from astropy.time import Time
 
 
@@ -83,23 +81,24 @@ class Telescope(object):
         n_ant: Optional[int] = None,
     ):
         self.tel_name = tel_name
-        self.latitude = da.asarray(latitude)
-        self.longitude = da.asarray(longitude)
-        self.elevation = da.asarray(elevation)
-        self.GEO = da.asarray([latitude, longitude, elevation])
+        self.latitude = da.asarray(latitude)  # type: ignore
+        self.longitude = da.asarray(longitude)  # type: ignore
+        self.elevation = da.asarray(elevation)  # type: ignore
+        self.GEO = da.asarray([latitude, longitude, elevation])  # type: ignore
         self.ITRF = None
         self.ENU = None
         self.n_ant = n_ant
 
-        if ENU_array is not None or ENU_path is not None:
+        if ENU_array or ENU_path:
             self.createArrayENU(ENU_array, ENU_path)
-        if ITRF_array is not None or ITRF_path is not None:
+        if ITRF_array or ITRF_path:
             self.createArrayITRF(ITRF_array, ITRF_path)
         if self.ITRF is None and self.ENU is None:
             raise ValueError(
-                "One of ('ENU_array', 'ENU_path', 'ITRF_array', 'ITRF_path') must be provided to create a Telescope object."
+                "One of 'ENU_path' or 'ITRF_path' must be provided to create a Telescope object."
             )
-        self.n_ant = len(self.ITRF)
+        else:
+            self.n_ant = len(self.ITRF)
 
     def __str__(self):
         msg = """\nTelescope Location
@@ -126,7 +125,7 @@ Elevation : {elevation}\n"""
                      array or as a csv like file."""
             print(msg)
             return
-        self.ENU = da.asarray(self.ENU)
+        self.ENU = da.asarray(self.ENU)  # type: ignore
         self.ENU_path = ENU_path
         self.GEO_ants = ENU_to_GEO(self.GEO, self.ENU)
         self.ITRF = ENU_to_ITRF(self.ENU, self.latitude, self.longitude, self.elevation)
@@ -142,8 +141,8 @@ Elevation : {elevation}\n"""
                      array or as a csv like file."""
             raise ValueError(msg)
 
-        self.ITRF = da.asarray(self.ITRF)
-        self.GEO_ants = da.asarray(itrf_to_geo(self.ITRF.compute()))
+        self.ITRF = da.asarray(self.ITRF)  # type: ignore
+        self.GEO_ants = da.asarray(itrf_to_geo(self.ITRF.compute()))  # type: ignore
 
 
 class Observation(Telescope):
@@ -474,7 +473,7 @@ Number of stationary RFI :  {n_stat}"""
         dec: array (n_src,)
             Declination of the sources in degrees.
         """
-        I = da.atleast_2d(I)
+        I = da.atleast_2d(I)  # type: ignore
         if I.ndim == 2:
             I = da.expand_dims(I, axis=0)
         I = I * da.ones(
@@ -530,7 +529,7 @@ Number of stationary RFI :  {n_stat}"""
         dec: Array (n_src,)
             Declination of the sources in degrees.
         """
-        I = da.atleast_2d(I)
+        I = da.atleast_2d(I)  # type: ignore
         if I.ndim == 2:
             I = da.expand_dims(I, axis=0)
         I = I * da.ones(
@@ -718,7 +717,7 @@ Number of stationary RFI :  {n_stat}"""
         tles: Array (n_src, 2)
             TLEs of the satellites corresponding to the NORAD IDs.
         """
-        Pv = da.atleast_2d(Pv)
+        Pv = da.atleast_2d(Pv)  # type: ignore
         if Pv.ndim == 2:
             Pv = da.expand_dims(Pv, axis=0)
         Pv = (
@@ -728,10 +727,11 @@ Number of stationary RFI :  {n_stat}"""
             )
         ).rechunk((-1, self.time_fine_chunk, self.freq_chunk))
         norad_ids = da.asarray(da.atleast_1d(norad_ids), chunks=(-1,))
+        n_src = len(norad_ids)
 
         rfi_xyz = da.asarray(
             get_satellite_positions(tles, mjd_to_jd(self.times_mjd_fine.compute())),
-            chunks=(-1, self.time_fine_chunk, 3),
+            chunks=(n_src, self.time_fine_chunk, 3),
         )
         # from tabsim.jax.coordinates import kepler_orbit_many
         # from tabsim.tle import
@@ -743,7 +743,7 @@ Number of stationary RFI :  {n_stat}"""
         # self.ants_xyz is shape (n_time_fine,n_ant,3)
         distances = da.linalg.norm(
             self.ants_xyz[None, :, :, :] - rfi_xyz[:, :, None, :], axis=-1
-        )
+        )  # .rechunk(n_src, self.time_chunk, self.n_int_samples, self.n_ant)
         # distances = da.asarray(da.linalg.norm(
         #     ants_pos(self.ITRF.compute(), mjd_to_jd(self.times_mjd_fine.compute()))[None, :, :, :] - rfi_xyz[:, :, None, :], axis=-1
         # ), chunks=(-1, self.time_fine_chunk, self.n_ant))
@@ -751,24 +751,39 @@ Number of stationary RFI :  {n_stat}"""
         I = Pv_to_Sv(Pv, distances)
         # I is shape (n_src,n_time_fine,n_ant,n_freq)
 
+        c_distances = (
+            (distances + self.ants_uvw[None, :, :, -1])
+            .reshape((n_src, self.n_time, self.n_int_samples, self.n_ant))
+            .rechunk((n_src, self.time_chunk, self.n_int_samples, self.ant_chunk))
+        )
+
         angular_seps = angular_separation(rfi_xyz, self.ants_xyz, self.ra, self.dec)
 
         # angular_seps is shape (n_src,n_time_fine,n_ant)
         rfi_A_app = da.sqrt(da.abs(I)) * airy_beam(
             angular_seps, self.freqs, self.dish_d
         )
+
+        rfi_A = rfi_A_app.reshape(
+            (n_src, self.n_time, self.n_int_samples, self.n_ant, self.n_freq)
+        ).rechunk(
+            (
+                n_src,
+                self.time_chunk,
+                self.n_int_samples,
+                self.ant_chunk,
+                self.freq_chunk,
+            )
+        )
+
         # rfi_A_app = da.ones((Pv.shape[0], self.n_time_fine, self.n_ant, self.n_freq), chunks=(-1, self.time_fine_chunk, self.n_ant, self.freq_chunk))
         # self.rfi_A_app is shape (n_src,n_time_fine,n_ant,n_freqs)
         # distances is shape (n_src,n_time_fine,n_ant)
         # self.ants_uvw is shape (n_time_fine,n_ant,3)
 
         vis_rfi = rfi_vis(
-            rfi_A_app.reshape(
-                (-1, self.n_time, self.n_int_samples, self.n_ant, self.n_freq)
-            ),
-            (distances + self.ants_uvw[None, :, :, -1]).reshape(
-                (-1, self.n_time, self.n_int_samples, self.n_ant)
-            ),
+            rfi_A,
+            c_distances,
             self.freqs,
             self.a1,
             self.a2,