Merge pull request #105 from HERA-Team/allow_arbitrary_sky_model_type

Allow for an arbitrary dtype for flux model

Author: steven-murray

setup.cfg

@@ -24,7 +24,7 @@ install_requires =
     line-profiler
     numpy>=2.0
     psutil
-    pyuvdata>=3.1.2
+    pyuvdata>=3.2.0
     rich
     scipy
 python_requires = >=3.10

src/matvis/_test_utils.py

@@ -11,7 +11,7 @@
 from pyuvdata import UVBeam
 from pyuvdata.analytic_beam import GaussianBeam
 from pyuvdata.beam_interface import BeamInterface
-from pyuvdata.telescopes import get_telescope
+from pyuvdata.telescopes import Telescope
 from pyuvsim import simsetup
 from pyuvsim.telescope import BeamList
 
@@ -34,7 +34,7 @@ def get_standard_sim_params(
     first_source_antizenith=False,
 ):
     """Create some standard random simulation parameters for use in tests."""
-    hera = get_telescope("hera")
+    hera = Telescope.from_known_telescopes("hera")
     obstime = Time("2018-08-31T04:02:30.11", format="isot", scale="utc")
 
     rng = np.random.default_rng(1)

src/matvis/core/coords.py

@@ -53,8 +53,15 @@ def __init__(
         self.xp = cp if self.gpu else np
 
         self.precision = precision
-        self.rtype, _ = get_dtypes(precision)
-        self.flux = self.xp.asarray(flux.astype(self.rtype))
+        self.rtype, self.ctype = get_dtypes(precision)
+
+        # Check if the flux is complex and set the dtype accordingly.
+        if self.xp.iscomplexobj(flux):
+            self.sky_model_dtype = self.ctype
+        else:
+            self.sky_model_dtype = self.rtype
+
+        self.flux = self.xp.asarray(flux.astype(self.sky_model_dtype))
 
         self.nsrc = len(flux)
         self.times = times
@@ -82,7 +89,9 @@ def setup(self):
             (3, self.nsrc_alloc), self.rtype(0.0), dtype=self.rtype
         )
         self.flux_above_horizon = self.xp.full(
-            (self.nsrc_alloc,) + self.flux.shape[1:], self.rtype(0.0), dtype=self.rtype
+            (self.nsrc_alloc,) + self.flux.shape[1:],
+            self.sky_model_dtype(0.0),
+            dtype=self.sky_model_dtype,
         )
 
     def select_chunk(self, chunk: int):

tests/test_coordrot.py

@@ -6,7 +6,7 @@
 from astropy import units as un
 from astropy.coordinates import SkyCoord
 from astropy.time import Time
-from pyuvdata.telescopes import get_telescope
+from pyuvdata.telescopes import Telescope
 
 from matvis import HAVE_GPU
 from matvis.core.coords import CoordinateRotation
@@ -35,10 +35,32 @@ def get_angles(x, y):
     return xp.arccos(ratio)
 
 
+def test_complex_flux():
+    """Test that using a complex flux works appropriately."""
+    rng = np.random.default_rng(1234)
+    n = 23
+    location = Telescope.from_known_telescopes("hera").location
+    skycoords = SkyCoord(
+        ra=rng.uniform(0, 2 * np.pi, size=n) * un.rad,
+        dec=rng.uniform(-np.pi / 2, np.pi / 2, size=n) * un.rad,
+        frame="icrs",
+    )
+
+    coords = CoordinateRotationAstropy(
+        flux=rng.normal(100, 2, size=n) + 1j * rng.normal(100, 2, size=n),
+        times=Time(np.array([2459863.0]), format="jd", scale="utc"),
+        telescope_loc=location,
+        skycoords=skycoords,
+        gpu=False,
+        precision=2,
+    )
+    assert coords.sky_model_dtype == coords.ctype == np.complex128
+
+
 def get_random_coordrot(n, method, gpu, seed, precision=2, setup: bool = True, **kw):
     """Get a random coordinate rotation object."""
     rng = np.random.default_rng(seed)
-    location = get_telescope("hera").location
+    location = Telescope.from_known_telescopes("hera").location
     skycoords = SkyCoord(
         ra=rng.uniform(0, 2 * np.pi, size=n) * un.rad,
         dec=rng.uniform(-np.pi / 2, np.pi / 2, size=n) * un.rad,

tests/test_matvis_cpu.py

@@ -5,7 +5,7 @@
 import numpy as np
 from astropy.time import Time
 from pyuvdata.analytic_beam import GaussianBeam
-from pyuvdata.telescopes import get_telescope
+from pyuvdata.telescopes import Telescope
 
 from matvis import simulate_vis
 
@@ -23,7 +23,7 @@
 def test_simulate_vis(polarized):
     """Test basic operation of simple wrapper around matvis, `simulate_vis`."""
     # Point source equatorial coords (radians)
-    hera = get_telescope("hera")
+    hera = Telescope.from_known_telescopes("hera")
     ra = np.linspace(0.0, 2.0 * np.pi, NPTSRC)
     dec = np.linspace(-0.5 * np.pi, 0.5 * np.pi, NPTSRC)