Merge pull request #97 from HERA-Team/prefix-bcrs

perf: ability to pre-fix the BCRS coordinates

Author: steven-murray

setup.cfg

@@ -25,7 +25,7 @@ install_requires =
     line-profiler
     numpy>=2.0
     psutil
-    pyuvdata@git+https://github.com/RadioAstronomySoftwareGroup/pyuvdata
+    pyuvdata>=3.1.0
     rich
     scipy
 python_requires = >=3.9

src/matvis/cli.py

@@ -21,8 +21,8 @@
 from line_profiler import LineProfiler
 from pathlib import Path
 from pyuvdata import UVBeam
+from pyuvdata.analytic_beam import GaussianBeam
 from pyuvdata.telescopes import get_telescope
-from pyuvsim import AnalyticBeam, simsetup
 from rich.console import Console
 from rich.logging import RichHandler
 from rich.rule import Rule
@@ -382,7 +382,16 @@ def get_stats_and_lines(filename, start_lineno, timings, time_unit):
 
 
 def get_standard_sim_params(
-    use_analytic_beam: bool, nfreq, ntime, nants, nsource, nbeams, naz=360, nza=180
+    use_analytic_beam: bool,
+    nfreq,
+    ntime,
+    nants,
+    nsource,
+    nbeams,
+    naz=360,
+    nza=180,
+    freq_min=100e6,
+    freq_max=200e6,
 ):
     """Create some standard random simulation parameters for use in profiling.
 
@@ -391,24 +400,17 @@ def get_standard_sim_params(
     # Set the seed so that different runs take about the same time.
     rng = np.random.default_rng()
 
+    # Source locations and frequencies
+    freqs = np.linspace(freq_min, freq_max, nfreq)
+
     # Beam model
-    if use_analytic_beam:
-        beam = AnalyticBeam("gaussian", diameter=14.0)
-    else:
-        # This is a peak-normalized e-field beam file at 100 and 101 MHz,
-        # downsampled to roughly 4 square-degree resolution.
-        beam = UVBeam()
-        beam.read_beamfits(beam_file, use_future_array_shapes=True)
-
-        # Up/down sample the beam
-        beam.interpolation_function = "az_za_simple"
-        beam = beam.interp(
-            az_array=np.linspace(0, 2 * np.pi, naz + 1)[:-1],
-            za_array=np.linspace(0, np.pi, nza + 1),
-            freq_array=beam.freq_array,
-            reuse_spline=True,
-            new_object=True,
-            az_za_grid=True,
+    beam = GaussianBeam(diameter=14.0)
+
+    if not use_analytic_beam:
+        beam = beam.to_uvbeam(
+            freq_array=freqs,
+            axis1_array=np.linspace(0, 2 * np.pi, naz + 1)[:-1],
+            axis2_array=np.linspace(0, np.pi, nza + 1),
         )
 
     beams = [beam] * nbeams
@@ -445,9 +447,6 @@ def get_standard_sim_params(
     flux0 = np.random.random(nsource) * 4
     spec_indx = np.random.normal(0.8, scale=0.05, size=nsource)
 
-    # Source locations and frequencies
-    freqs = np.linspace(100e6, 200e6, nfreq)
-
     # Calculate source fluxes for matvis
     flux = ((freqs[:, np.newaxis] / freqs[0]) ** spec_indx.T * flux0.T).T
 

src/matvis/core/coords.py

@@ -67,7 +67,10 @@ def __init__(
 
         self.chunk_size = chunk_size or self.nsrc
         self.source_buffer = source_buffer
-        self.nsrc_alloc = int(self.chunk_size * self.source_buffer)
+        if self.chunk_size > 1000:
+            self.nsrc_alloc = int(self.chunk_size * self.source_buffer)
+        else:
+            self.nsrc_alloc = self.chunk_size
 
     def setup(self):
         """Allocate memory for the rotation."""

src/matvis/cpu/coords.py

@@ -64,21 +64,27 @@ def __init__(self, update_bcrs_every=0.0, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.update_bcrs_every = update_bcrs_every * un.s
 
-    def setup(self):
-        """Standard setup, as well as storing the cartesian representation of ECI."""
-        super().setup()
+        # Do one rotation to warm up the cache. This reads the IERS table.
+        frame = AltAz(obstime=self.times[0], location=self.telescope_loc)
+        self.skycoords[0].transform_to(frame)
+        self._time_of_last_evaluation = None
+
+        # These are unchanging over time, so init them outside the setup() to save
+        # memory when multi-processing.
         self._eci = self.xp.asarray(
             point_source_crd_eq(self.skycoords.ra, self.skycoords.dec)
         )
 
+    def setup(self):
+        """Standard setup, as well as storing the cartesian representation of ECI."""
+        super().setup()
+
         # BCRS holds the deflected, aberrated bnp-d coordinates, which don't change
         # significantly over time.
-        self._bcrs = self._eci.copy()
-        self._time_of_last_evaluation = None
-
-        # Do one rotation to warm up the cache.
-        frame = AltAz(obstime=self.times[0], location=self.telescope_loc)
-        self.skycoords[0].transform_to(frame)
+        if not hasattr(self, "_bcrs"):
+            self._bcrs = self.xp.full(
+                self._eci.shape, dtype=self._eci.dtype, fill_value=0.0
+            )
 
     def _atioq(self, xyz: np.ndarray, astrom):
         # cirs to hadec rot
@@ -177,21 +183,22 @@ def _apco(self, observed_frame):
     def _get_obsf(self, obstime, location):
         return AltAz(obstime=obstime, location=location)
 
-    def rotate(self, t: int) -> tuple[np.ndarray, np.ndarray]:
-        """Rotate the coordinates into the observed frame."""
-        obsf = self._get_obsf(self.times[t], self.telescope_loc)
-        astrom = self._apco(obsf)
-
-        # Copy the eci coordinates, because these routines modify them in-place
-
+    def _set_bcrs(self, t, astrom=None):
         # convert to topocentric CIRS
         # together, _ld + _ab take ~90% of the time.
+        if astrom is None:
+            obsf = self._get_obsf(self.times[t], self.telescope_loc)
+            astrom = self._apco(obsf)
+
         if (
             self._time_of_last_evaluation is None
             or self.times[t] - self.times[self._time_of_last_evaluation]
             > self.update_bcrs_every
         ):
-            self._bcrs[:] = self._eci[:]
+            if hasattr(self, "_bcrs"):
+                self._bcrs[:] = self._eci[:]
+            else:
+                self._bcrs = self._eci.copy()
 
             # Light deflection by the Sun, giving BCRS natural direction.
             self._ld(self._bcrs, self.xp.asarray(astrom["eh"]), astrom["em"], 1e-6)
@@ -207,6 +214,13 @@ def rotate(self, t: int) -> tuple[np.ndarray, np.ndarray]:
 
             self._time_of_last_evaluation = t
 
+    def rotate(self, t: int) -> tuple[np.ndarray, np.ndarray]:
+        """Rotate the coordinates into the observed frame."""
+        obsf = self._get_obsf(self.times[t], self.telescope_loc)
+        astrom = self._apco(obsf)
+
+        self._set_bcrs(t, astrom)
+
         self.all_coords_topo[:] = self._bcrs[:]
 
         # now perform observed conversion

src/matvis/cpu/cpu.py

@@ -157,20 +157,8 @@ def simulate(
         precision,
         source_buffer=source_buffer,
     )
-    nsrc_alloc = int(npixc * source_buffer)
-
-    bmfunc = UVBeamInterpolator(
-        beam_list=beam_list,
-        beam_idx=beam_idx,
-        polarized=polarized,
-        nant=nant,
-        freq=freq,
-        spline_opts=beam_spline_opts,
-        precision=precision,
-        nsrc=nsrc_alloc,
-    )
-
     coord_method = CoordinateRotation._methods[coord_method]
+
     coord_method_params = coord_method_params or {}
     coords = coord_method(
         flux=np.sqrt(0.5 * I_sky),
@@ -182,6 +170,19 @@ def simulate(
         source_buffer=source_buffer,
         **coord_method_params,
     )
+
+    nsrc_alloc = coords.nsrc_alloc
+    bmfunc = UVBeamInterpolator(
+        beam_list=beam_list,
+        beam_idx=beam_idx,
+        polarized=polarized,
+        nant=nant,
+        freq=freq,
+        spline_opts=beam_spline_opts,
+        precision=precision,
+        nsrc=nsrc_alloc,
+    )
+
     taucalc = TauCalculator(
         antpos=antpos, freq=freq, precision=precision, nsrc=nsrc_alloc
     )

tests/test_coordrot.py

@@ -10,7 +10,7 @@
 
 from matvis import HAVE_GPU
 from matvis.core.coords import CoordinateRotation
-from matvis.cpu.coords import CoordinateRotationAstropy
+from matvis.cpu.coords import CoordinateRotationAstropy, CoordinateRotationERFA
 
 if HAVE_GPU:
     import cupy as cp
@@ -35,7 +35,7 @@ def get_angles(x, y):
     return xp.arccos(ratio)
 
 
-def get_random_coordrot(n, method, gpu, seed, precision=2):
+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
@@ -51,8 +51,10 @@ def get_random_coordrot(n, method, gpu, seed, precision=2):
         skycoords=skycoords,
         gpu=gpu,
         precision=precision,
+        **kw
     )
-    coords.setup()
+    if setup:
+        coords.setup()
     return coords
 
 
@@ -101,3 +103,43 @@ def test_accuracy_against_astropy(method, gpu, precision):
         np.testing.assert_allclose(angles, 0, atol=0.01)  # 10 mas
     else:
         np.testing.assert_allclose(angles, 0, atol=150)  # 50 mas
+
+
+def test_coord_rot_erfa_set_bcrs():
+    """Test that setting bcrs before setup works as expected."""
+    normal = get_random_coordrot(
+        1000, CoordinateRotationERFA, gpu=False, seed=1, precision=1
+    )
+    bcrs = get_random_coordrot(
+        1000, CoordinateRotationERFA, gpu=False, seed=1, precision=1, setup=False
+    )
+    bcrs._set_bcrs(0)
+    bcrs.setup()
+
+    normal.rotate(0)
+    bcrs.rotate(0)
+
+    np.testing.assert_allclose(normal.all_coords_topo, bcrs.all_coords_topo)
+
+
+def test_larger_chunksize():
+    """Test that using different chunk sizes results in the same output."""
+    small = get_random_coordrot(
+        10000, CoordinateRotationERFA, gpu=False, seed=1, precision=1, chunk_size=100
+    )
+    large = get_random_coordrot(
+        10000, CoordinateRotationERFA, gpu=False, seed=1, precision=1, chunk_size=5000
+    )
+    default = get_random_coordrot(
+        10000, CoordinateRotationERFA, gpu=False, seed=1, precision=1
+    )
+    small.select_chunk(0)
+    large.select_chunk(0)
+    default.select_chunk(0)
+
+    np.testing.assert_allclose(
+        small.coords_above_horizon, large.coords_above_horizon[:, :100]
+    )
+    np.testing.assert_allclose(
+        small.coords_above_horizon, default.coords_above_horizon[:, :100]
+    )