fix time_delays_from removing astropy API

Author: Jammy2211

autolens/lens/tracer_util.py

@@ -322,29 +322,36 @@ def time_delays_from(
             f"{len(plane_redshifts)} planes with redshifts {plane_redshifts}."
         )
 
-    # Constants
-    mpc_in_m = 3.08567758e22  # Mpc in meters
-    arcsec_to_rad = np.deg2rad(1.0 / 3600.0)  # arcsec to radians
-    seconds_per_day = 86400
-    c = 299792458  # speed of light in m/s
-
-    factor = arcsec_to_rad**2 / seconds_per_day
-
-    # Angular diameter distances
-    Dd = cosmology.angular_diameter_distance(plane_redshifts[0]).value  # [Mpc]
-    Ds = cosmology.angular_diameter_distance(plane_redshifts[1]).value  # [Mpc]
-    Dds = cosmology.angular_diameter_distance_z1z2(
-        z1=plane_redshifts[0], z2=plane_redshifts[1]
-    ).value  # [Mpc]
+    z_l, z_s = plane_redshifts[0], plane_redshifts[1]
+
+    # -----------------
+    # Constants (SI)
+    # -----------------
+    kpc_in_m = xp.asarray(3.085677581491367e19)  # kpc in meters
+    arcsec_to_rad = xp.asarray(np.pi / 648000.0)  # arcsec -> rad (pi / (180*3600))
+    seconds_per_day = xp.asarray(86400.0)
+    c = xp.asarray(299792458.0)  # m/s
+
+    # This factor converts Fermat potential in arcsec^2 into days once multiplied by D_dt/c
+    factor = (arcsec_to_rad * arcsec_to_rad) / seconds_per_day
+
+    # -----------------
+    # Angular diameter distances (kpc)
+    # -----------------
+    Dd_kpc = cosmology.angular_diameter_distance_to_earth_in_kpc_from(z_l, xp=xp)
+    Ds_kpc = cosmology.angular_diameter_distance_to_earth_in_kpc_from(z_s, xp=xp)
+    Dds_kpc = cosmology.angular_diameter_distance_between_redshifts_in_kpc_from(
+        redshift_0=z_l, redshift_1=z_s, xp=xp
+    )
 
-    # Time-delay distance in meters
-    D_dt = (1 + plane_redshifts[0]) * Dd * Ds / Dds * mpc_in_m
+    # Time-delay distance in meters: (1+z_l) * Dd * Ds / Dds
+    D_dt_m = (xp.asarray(1.0) + xp.asarray(z_l)) * (Dd_kpc * Ds_kpc / Dds_kpc) * kpc_in_m
 
-    # Fermat potential
+    # Fermat potential (should be in arcsec^2 for this formula)
     fermat_potential = galaxies.fermat_potential_from(grid=grid, xp=xp)
 
     # Final time delay in days
-    return D_dt / c * fermat_potential * factor
+    return (D_dt_m / c) * fermat_potential * factor
 
 
 def ordered_plane_redshifts_with_slicing_from(