From 6ab3e293eaca54ad43242c63018bcbb575f3a75b Mon Sep 17 00:00:00 2001
From: Tyler Cox <tyler.a.cox@berkeley.edu>
Date: Tue, 8 Oct 2024 14:53:44 -0700
Subject: [PATCH] allow user to choose between coordinate transforms

---
 fftvis/simulate.py | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/fftvis/simulate.py b/fftvis/simulate.py
index d2c43f5..aca57a7 100644
--- a/fftvis/simulate.py
+++ b/fftvis/simulate.py
@@ -11,6 +11,7 @@
 from matvis.core.beams import prepare_beam_unpolarized
 from matvis.cpu.coords import CoordinateRotationAstropy, CoordinateRotationERFA
 from matvis.core.coords import CoordinateRotation
+from typing import Callable, Literal
 
 from astropy.coordinates import EarthLocation, SkyCoord
 from astropy import units as un
@@ -51,6 +52,10 @@ def simulate_vis(
     flat_array_tol: float = 0.0,
     interpolation_function: str = "az_za_map_coordinates",
     nprocesses: int | None = 1,
+    coord_method: Literal[
+        "CoordinateRotationAstropy", "CoordinateRotationERFA"
+    ] = "CoordinateRotationERFA",
+    coord_method_params: dict | None = None,
 ):
     """
     Parameters:
@@ -139,6 +144,8 @@ def simulate_vis(
         flat_array_tol=flat_array_tol,
         interpolation_function=interpolation_function,
         nprocesses=nprocesses,
+        coord_method=coord_method,
+        coord_method_params=coord_method_params,
     )
 
 
@@ -159,6 +166,10 @@ def simulate(
     flat_array_tol: float = 0.0,
     interpolation_function: str = "az_za_map_coordinates",
     nprocesses: int | None = 1,
+    coord_method: Literal[
+        "CoordinateRotationAstropy", "CoordinateRotationERFA"
+    ] = "CoordinateRotationERFA",
+    coord_method_params: dict | None = None,
 ):
     """
     Parameters:
@@ -290,13 +301,17 @@ def simulate(
     if nprocesses is None:
         nprocesses = cpu_count()
 
-    coord_mgr = CoordinateRotationERFA(
+    coord_method = CoordinateRotation._methods[coord_method]
+    coord_method_params = coord_method_params or {}
+    coord_mgr = coord_method(
         flux=Isky,
         times=Time(times, format='jd'),
         telescope_loc=telescope_loc,
         skycoords=SkyCoord(ra=ra * un.rad, dec=dec * un.rad, frame='icrs'),
         source_buffer=1.0,
         precision=precision,
+        **coord_method_params,
+        #update_bcrs_every=np.inf, # Don't update BCRS
     )
 
     nprocesses, freq_chunks, time_chunks, nf, nt = utils.get_task_chunks(nprocesses, nfreqs, ntimes)