Merge branch 'lunar_capability' into adding-lunar-coords

src/py21cmsense/_utils.py

@@ -5,7 +5,9 @@
 from astropy.coordinates import EarthLocation, SkyCoord
 from astropy.time import Time
 from pyuvdata import utils as uvutils
-
+from lunarsky import MoonLocation
+from lunarsky import SkyCoord as LunarSkyCoord
+from lunarsky import Time as LTime
 
 def between(xmin, xmax):
     """Return an attrs validation function that checks a number is within bounds."""
@@ -39,7 +41,7 @@ def find_nearest(array, value):
 
 @un.quantity_input
 def phase_past_zenith(
-    time_past_zenith: un.day, bls_enu: np.ndarray, latitude, use_apparent: bool = True
+    time_past_zenith: un.day, bls_enu: np.ndarray, latitude, world, use_apparent: bool = True
 ):
     """Compute UVWs phased to a point rotated from zenith by a certain amount of time.
 
@@ -56,6 +58,8 @@ def phase_past_zenith(
         The UVWs when phased to zenith.
     latitude
         The latitude of the center of the array, in radians.
+    world
+        Wether the telescope is on the Earth or Moon.
 
     Returns
     -------
@@ -64,20 +68,34 @@ def phase_past_zenith(
     """
     # Generate ra/dec of zenith at time in the phase_frame coordinate system
     # to use for phasing
-    telescope_location = EarthLocation.from_geodetic(lon=0, lat=latitude)
+    if world == "earth":
+        telescope_location = EarthLocation.from_geodetic(lon=0, lat=latitude)
+    else:
+        telescope_location = MoonLocation.from_selenodetic(lon=0, lat=latitude)
 
     # JD is arbitrary
     jd = 2454600
+
+    if world == "earth":
+        zenith_coord = SkyCoord(
+            alt=90 * un.deg,
+            az=0 * un.deg,
+            obstime=Time(jd, format="jd"),
+            frame="altaz",
+            location=telescope_location,
+        )
+    else:
+        zenith_coord = LunarSkyCoord(
+                alt=90 * un.deg,
+                az=0 * un.deg,
+                obstime=LTime(jd, format="jd"),
+                frame="lunartopo",
+                location=telescope_location,
+            )
 
-    zenith_coord = SkyCoord(
-        alt=90 * un.deg,
-        az=0 * un.deg,
-        obstime=Time(jd, format="jd"),
-        frame="altaz",
-        location=telescope_location,
-    )
     zenith_coord = zenith_coord.transform_to("icrs")
 
+    zenith_coord.obstime.location = telescope_location
     obstimes = zenith_coord.obstime + time_past_zenith
     lsts = obstimes.sidereal_time("apparent", longitude=0.0).rad
 

src/py21cmsense/observation.py

@@ -94,16 +94,16 @@ class Observation:
 
     time_per_day: tp.Time = attr.ib(
         6 * un.hour,
-        validator=(tp.vld_physical_type("time"), ut.between(0 * un.hour, 24 * un.hour)),
+        validator=(tp.vld_physical_type("time"), ut.between(0 * un.hour, 655.2 * un.hour)),
     )
     track: tp.Time | None = attr.ib(
         None,
         validator=attr.validators.optional(
-            [tp.vld_physical_type("time"), ut.between(0, 24 * un.hour)]
+            [tp.vld_physical_type("time"), ut.between(0, 655.2 * un.hour)]
         ),
     )
     lst_bin_size: tp.Time = attr.ib(
-        validator=(tp.vld_physical_type("time"), ut.between(0, 24 * un.hour)),
+        validator=(tp.vld_physical_type("time"), ut.between(0, 655.2 * un.hour)),
     )
     integration_time: tp.Time = attr.ib(
         60 * un.second, validator=(tp.vld_physical_type("time"), ut.positive)

src/py21cmsense/observatory.py

@@ -68,6 +68,8 @@ class Observatory:
         By default it is, so that the beam-crossing time is
         ``tday * FWHM / (2pi cos(lat))``. This affects both the thermal and sample
         variance calculations.
+    world: string
+        A string specifying whether the telescope is on the Earth or the moon.
     """
 
     _antpos: tp.Length = attr.ib(eq=attr.cmp_using(eq=np.array_equal))
@@ -84,6 +86,7 @@ class Observatory:
         default=0.0 * un.m, validator=(tp.vld_physical_type("length"), ut.nonnegative)
     )
     beam_crossing_time_incl_latitude: bool = attr.ib(default=True, converter=bool)
+    world: str = attr.ib(default = 'earth') #add validator stuff later
 
     @_antpos.validator
     def _antpos_validator(self, att, val):
@@ -262,7 +265,7 @@ def projected_baselines(
 
         bl_wavelengths = baselines.reshape((-1, 3)) * self.metres_to_wavelengths
 
-        out = ut.phase_past_zenith(time_offset, bl_wavelengths, self.latitude)
+        out = ut.phase_past_zenith(time_offset, bl_wavelengths, self.latitude, self.world)
 
         out = out.reshape(*orig_shape[:-1], np.size(time_offset), orig_shape[-1])
         if np.size(time_offset) == 1:
@@ -294,7 +297,10 @@ def longest_baseline(self) -> float:
     def observation_duration(self) -> un.Quantity[un.day]:
         """The time it takes for the sky to drift through the FWHM."""
         latfac = np.cos(self.latitude) if self.beam_crossing_time_incl_latitude else 1
-        return un.day * self.beam.fwhm / (2 * np.pi * un.rad * latfac)
+        if self.world == "earth":
+            return un.day * self.beam.fwhm / (2 * np.pi * un.rad * latfac)
+        else:
+            return 27.3 * un.day * self.beam.fwhm/(2 * np.pi * un.rad * latfac)
 
     def get_redundant_baselines(
         self,

src/py21cmsense/sensitivity.py

@@ -151,7 +151,7 @@ class PowerSpectrum(Sensitivity):
 
     horizon_buffer: tp.Wavenumber = attr.ib(default=0.1 * littleh / un.Mpc)
     foreground_model: str = attr.ib(
-        default="moderate", validator=vld.in_(["moderate", "optimistic"])
+        default="moderate", validator=vld.in_(["moderate", "optimistic", "ultra_optimistic"])
     )
     theory_model: TheoryModel = attr.ib()
 
@@ -471,6 +471,8 @@ def horizon_limit(self, umag: float) -> tp.Wavenumber:
             return horizon + self.horizon_buffer
         elif self.foreground_model in ["optimistic"]:
             return horizon * np.sin(self.observation.observatory.beam.first_null / 2)
+        elif self.foreground_model in ["ultra_optimistic"]:
+            return horizon
 
     def _average_sense_to_1d(
         self, sense: dict[tp.Wavenumber, tp.Delta], k1d: tp.Wavenumber | None = None

src/py21cmsense/theory.py

@@ -196,3 +196,64 @@ def delta_squared(self, z: float, k: np.ndarray) -> un.Quantity[un.mK**2]:
             )
 
         return self.spline(k) << un.mK**2
+
+class FarViewModel(TheoryModel):
+    """21cmFAST-based theory model explicitly for z=30, [Insert paper link here later]"""
+
+    use_littleh: bool = False
+
+    def __init__(self) -> None:
+        k_pth = (
+            Path(__file__).parent
+            / "data/farview/kmag_ultimate.npy"
+        )
+
+        delta_pth = (
+            Path(__file__).parent
+            / "data/farview/P_Tb_ultimate.npy"
+        )
+        #Should at some point reorganize the data so these steps aren't necessary
+        k_fixed = np.load(k_pth)
+        power = np.load(delta_pth)
+        k_fixed = k_fixed[~np.isnan(power)]
+        power = power[~np.isnan(power)]
+        delta = (k_fixed**3 * power) / (2*np.pi**2)
+
+        self.k = k_fixed
+        self.delta_squared_raw = delta
+
+        self.spline = InterpolatedUnivariateSpline(self.k, self.delta_squared_raw, k=1)
+
+    def delta_squared(self, z: float, k: np.ndarray) -> un.Quantity[un.mK**2]:
+        """Compute Delta^2(k, z) for the theory model.
+
+        Parameters
+        ----------
+        z
+            The redshift (should be a float).
+        k
+            The wavenumbers, either in units of 1/Mpc if use_littleh=False, or
+            h/Mpc if use_littleh=True.
+
+        Returns
+        -------
+        delta_squared
+            An array of delta_squared values in units of mK^2.
+        """
+        if np.any(k > self.k.max()):
+            warnings.warn(
+                f"Extrapolating above the simulated theoretical k: {k.max()} > {self.k.max()}",
+                stacklevel=2,
+            )
+        if np.any(k < self.k.min()):
+            warnings.warn(
+                f"Extrapolating below the simulated theoretical k: {k.min()} < {self.k.min()}",
+                stacklevel=2,
+            )
+        if not 29.5 < z < 30.5:
+            warnings.warn(
+                f"Theory power corresponds to z=30, not z={z:.2f}",
+                stacklevel=2,
+            )
+
+        return self.spline(k) << un.mK**2