Merge pull request #164 from rasg-affiliates/adding-lunar-coords

Adding lunar coords

ci/testing.yaml

@@ -19,3 +19,4 @@ dependencies:
   - matplotlib
   - pip:
     - methodtools
+    - lunarsky>=0.2.5

docs/environment.yml

@@ -24,4 +24,5 @@ dependencies:
   - setuptools_scm
   - pip:
     - furo
+    - lunarsky>=0.2.5
     - sphinx_design

pyproject.toml

@@ -43,6 +43,7 @@ dependencies = [
   "rich",
   "attrs",
   "hickleable>=0.1.1",
+  "lunarsky>=0.2.5",
 ]
 
 [project.optional-dependencies]

src/py21cmsense/_utils.py

@@ -4,6 +4,9 @@
 from astropy import units as un
 from astropy.coordinates import EarthLocation, SkyCoord
 from astropy.time import Time
+from lunarsky import MoonLocation
+from lunarsky import SkyCoord as LunarSkyCoord
+from lunarsky import Time as LTime
 from pyuvdata import utils as uvutils
 
 
@@ -42,6 +45,7 @@ def phase_past_zenith(
     time_past_zenith: un.hour,
     bls_enu: np.ndarray,
     latitude: float,
+    world: str = "earth",
     phase_center_dec: un.rad = None,
     use_apparent: bool = True,
 ):
@@ -61,6 +65,8 @@ def phase_past_zenith(
         vectors (equivalent to the UVWs if phased to zenith).
     latitude
         The latitude of the center of the array, in radians.
+    world
+        Whether the telescope is on the Earth or Moon.
     phase_center_dec
         If given, the declination of the phase center. If not given, it is set to the
         latitude of the array (i.e. the phase center passes through zenith).
@@ -75,27 +81,43 @@ 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
-    tm = Time(jd, format="jd")
-
-    phase_center_coord = SkyCoord(
-        alt=90 * un.deg,
-        az=0 * un.deg,
-        obstime=tm,
-        frame="altaz",
-        location=telescope_location,
-    )
+
+    if world == "earth":
+        tm = Time(jd, format="jd", location=telescope_location)
+
+        phase_center_coord = SkyCoord(
+            alt=90 * un.deg,
+            az=0 * un.deg,
+            obstime=tm,
+            frame="altaz",
+            location=telescope_location,
+        )
+    else:
+        tm = LTime(jd, format="jd", location=telescope_location)
+
+        phase_center_coord = LunarSkyCoord(
+            alt=90 * un.deg,
+            az=0 * un.deg,
+            obstime=tm,
+            frame="lunartopo",
+            location=telescope_location,
+        )
+
     phase_center_coord = phase_center_coord.transform_to("icrs")
 
     if phase_center_dec is not None:
-        phase_center_coord = SkyCoord(
+        phase_center_coord = phase_center_coord.__class__(
             ra=phase_center_coord.ra,
             dec=phase_center_dec,
             obstime=tm,
-            frame="icrs",
+            frame=phase_center_coord.frame,
             location=telescope_location,
         )
 

src/py21cmsense/observation.py

@@ -38,8 +38,9 @@ class Observation:
         An object defining attributes of the observatory itself (its location etc.)
     hours_per_day : float or Quantity, optional
         The number of good observing hours per day.  This corresponds to the size of a
-        low-foreground region in right ascension for a drift scanning instrument. The
-        total observing time is `n_days*hours_per_day`.  Default is 6.
+        low-foreground region in right ascension for a drift scanning instrument on the Earth.
+        The total observing time is `n_days*hours_per_day` where a day is 24 hours on the
+        Earth and 655.2 hours on the moon.  Default is 6 (163.8 on moon).
         If simulating a tracked scan, `hours_per_day` should be a multiple of the length
         of the track (i.e. for two three-hour tracks per day, `hours_per_day` should be 6).
     track
@@ -59,9 +60,10 @@ class Observation:
         The bandwidth used for the observation, assumed to be in MHz. Note this is not the total
         instrument bandwidth, but the redshift range that can be considered co-eval.
     n_days : int, optional
-        The number of days observed (for the same set of LSTs). The default is 180, which is the
-        maximum a particular R.A. can be observed in one year if one only observes at night.
-        The total observing time is `n_days*hours_per_day`.
+        The number of days observed (for the same set of LSTs). The default is 180 (6 on moon),
+        which is the maximum a particular R.A. can be observed in one year if one only observes
+        at night. The total observing time is `n_days*hours_per_day` where a day is 24 hours on
+        the Earth and 655.2 hours on the moon.
     baseline_filters
         A function that takes a single value: a length-3 array of baseline co-ordinates,
         and returns a bool indicating whether to include the baseline. Built-in filters
@@ -98,17 +100,14 @@ class Observation:
     observatory: obs.Observatory = attr.ib(validator=vld.instance_of(obs.Observatory))
 
     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")),
     )
     track: tp.Time | None = attr.ib(
         None,
-        validator=attr.validators.optional(
-            [tp.vld_physical_type("time"), ut.between(0, 24 * un.hour)]
-        ),
+        validator=attr.validators.optional([tp.vld_physical_type("time")]),
     )
     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")),
     )
     integration_time: tp.Time = attr.ib(
         60 * un.second, validator=(tp.vld_physical_type("time"), ut.positive)
@@ -117,7 +116,7 @@ class Observation:
     bandwidth: tp.Frequency = attr.ib(
         8 * un.MHz, validator=(tp.vld_physical_type("frequency"), ut.positive)
     )
-    n_days: int = attr.ib(default=180, converter=int, validator=ut.positive)
+    n_days: int = attr.ib(converter=int, validator=ut.positive)
     baseline_filters: tuple[Callable[[tp.Length], bool]] = attr.ib(
         default=(), converter=tp._tuplify
     )
@@ -170,8 +169,28 @@ def __sethstate__(self, d: dict[str, Any]) -> None:
         d["cosmo"] = Planck15.from_format(d["cosmo"])
         self.__dict__.update(d)
 
+    @time_per_day.validator
+    def _time_per_day_vld(self, att, val):
+        day_length = 24 * un.hour if self.observatory.world == "earth" else 655.2 * un.hour
+
+        if not 0 * un.hour <= val <= day_length:
+            raise ValueError(f"time_per_day should be between 0 and {day_length}")
+
+    @track.validator
+    def _track_vld(self, att, val):
+        if val is not None:
+            day_length = 24 * un.hour if self.observatory.world == "earth" else 655.2 * un.hour
+
+            if not 0 * un.hour <= val <= day_length:
+                raise ValueError(f"track should be between 0 and {day_length}")
+
     @lst_bin_size.validator
     def _lst_bin_size_vld(self, att, val):
+        day_length = 24 * un.hour if self.observatory.world == "earth" else 655.2 * un.hour
+
+        if not 0 * un.hour <= val <= day_length:
+            raise ValueError(f"lst_bin_size should be between 0 and {day_length}")
+
         if val > self.time_per_day:
             raise ValueError("lst_bin_size must be <= time_per_day")
 
@@ -180,6 +199,13 @@ def _integration_time_vld(self, att, val):
         if val > self.lst_bin_size:
             raise ValueError("integration_time must be <= lst_bin_size")
 
+    @time_per_day.default
+    def _time_per_day_default(self):
+        if self.observatory.world == "earth":
+            return 6 * un.hour
+        else:
+            return 163.8 * un.hour
+
     @lst_bin_size.default
     def _lst_bin_size_default(self):
         # time it takes the sky to drift through beam FWHM
@@ -188,6 +214,13 @@ def _lst_bin_size_default(self):
         else:
             return self.observatory.observation_duration
 
+    @n_days.default
+    def _n_days_default(self):
+        if self.observatory.world == "earth":
+            return 180
+        else:
+            return 6
+
     @phase_center_dec.default
     def _phase_center_dec_default(self):
         return self.observatory.latitude

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", validator=vld.in_(["earth", "moon"]))
 
     @_antpos.validator
     def _antpos_validator(self, att, val):
@@ -323,6 +326,7 @@ def projected_baselines(
             time_past_zenith=time_offset,
             bls_enu=bl_wavelengths,
             latitude=self.latitude,
+            world=self.world,
             phase_center_dec=phase_center_dec,
         )
 
@@ -356,7 +360,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", "foreground_free"])
     )
     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 ["foreground_free"]:
+            return 0
 
     def _average_sense_to_1d(
         self, sense: dict[tp.Wavenumber, tp.Delta], k1d: tp.Wavenumber | None = None

src/py21cmsense/theory.py

@@ -196,3 +196,59 @@ 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."""
+
+    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

tests/test_observation.py

@@ -16,13 +16,19 @@ def bm():
     return GaussianBeam(150.0 * units.MHz, dish_size=14 * units.m)
 
 
+@pytest.fixture(scope="module", params=["earth", "moon"])
+def wd(request):
+    return request.param
+
+
 @pytest.fixture(scope="module")
-def observatory(bm):
+def observatory(bm, wd):
     return Observatory(
         antpos=np.array([[0, 0, 0], [14, 0, 0], [28, 0, 0], [70, 0, 0], [0, 14, 0], [23, -45, 0]])
         * units.m,
         latitude=-32 * units.deg,
         beam=bm,
+        world=wd,
     )
 
 
@@ -91,9 +97,26 @@ def test_from_yaml(observatory):
         Observation.from_yaml(3)
 
 
-def test_huge_lst_bin_size(observatory: Observatory):
+def test_huge_time_per_day_size(observatory: Observatory, wd):
+    tpd = 25 * units.hour if wd == "earth" else 682.5 * units.hour
+    with pytest.raises(ValueError, match="time_per_day should be between 0 and"):
+        Observation(observatory=observatory, time_per_day=tpd)
+
+
+def test_huge_track_size(observatory: Observatory, wd):
+    tck = 25 * units.hour if wd == "earth" else 682.5 * units.hour
+    with pytest.raises(ValueError, match="track should be between 0 and"):
+        Observation(observatory=observatory, track=tck)
+
+
+def test_huge_lst_bin_size(observatory: Observatory, wd):
+    lst = 23 * units.hour if wd == "earth" else 627.9 * units.hour
     with pytest.raises(ValueError, match="lst_bin_size must be <= time_per_day"):
-        Observation(observatory=observatory, lst_bin_size=23 * units.hour)
+        Observation(observatory=observatory, lst_bin_size=lst)
+
+    lst2 = 25 * units.hour if wd == "earth" else 682.5 * units.hour
+    with pytest.raises(ValueError, match="lst_bin_size should be between 0 and"):
+        Observation(observatory=observatory, lst_bin_size=lst2)
 
 
 def test_huge_integration_time(observatory: Observatory):