Add load_P9

src/adam_test_data/datasets/__init__.py

@@ -1,2 +1,3 @@
 # ruff: noqa: F401
+from .p9 import load_P9
 from .s3m import load_S3M

src/adam_test_data/datasets/p9.py

@@ -0,0 +1,60 @@
+import os
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+from adam_core.coordinates import KeplerianCoordinates
+from adam_core.coordinates.origin import Origin
+from adam_core.orbits import Orbits
+from adam_core.time import Timestamp
+
+from ..populations import PhotometricProperties, SmallBodies
+
+
+def load_P9(directory: str) -> SmallBodies:
+    """
+    Load the planet 9 population model from Brown and Batygin (2021a).
+
+    The model can be downloaded from: https://data.caltech.edu/records/8fjad-x7y61
+
+    Parameters
+    ----------
+    directory : str
+        The directory containing the planet 9 population model.
+
+    Returns
+    -------
+    p9 : SmallBodies
+        The planet 9 population model.
+    """
+    p9 = pd.read_csv(os.path.join(directory, "reference_population.csv"), comment="#")
+    p9.columns = p9.columns.str.strip()
+
+    # From the comments in the file
+    time = Timestamp.from_jd([2458270.0 for _ in range(len(p9))], scale="tdb")
+    orbit_ids = pa.array([f"P9_{i:07d}" for i in range(len(p9))])
+    orbits = Orbits.from_kwargs(
+        orbit_id=orbit_ids,
+        object_id=np.array(p9.index.values.astype(str), dtype="object"),
+        coordinates=KeplerianCoordinates.from_kwargs(
+            a=p9["a"],
+            e=p9["e"],
+            i=p9["inc"],
+            raan=p9["Omega"],
+            ap=p9["varpi"],
+            M=p9["M"],
+            time=time,
+            origin=Origin.from_kwargs(code=pa.repeat("SUN", len(p9))),
+            frame="equatorial",
+        ).to_cartesian(),
+    )
+
+    photometric_properties = PhotometricProperties.from_kwargs(
+        H_mf=p9["V"],
+    )
+
+    return SmallBodies.from_kwargs(
+        orbits=orbits,
+        properties=photometric_properties,
+        name="P9",
+    )

src/adam_test_data/datasets/s3m.py

@@ -88,7 +88,7 @@ def load_S3M(directory: str) -> SmallBodies:
         # Create orbit IDs for the s3m population
         # There are ~14 million objects in the population so 7 digits should be enough
         orbit_ids = pa.array(
-            [f"s3m{i:07d}" for i in range(id_offset, id_offset + len(s3m_df_i))]
+            [f"S3M_{i:07d}" for i in range(id_offset, id_offset + len(s3m_df_i))]
         )
         orbits_i = Orbits.from_kwargs(
             orbit_id=orbit_ids,

src/adam_test_data/datasets/tests/test_p9.py

@@ -0,0 +1,75 @@
+import os
+import tempfile
+
+import pyarrow as pa
+import pyarrow.compute as pc
+import pytest
+
+from ..p9 import load_P9
+
+
+@pytest.fixture
+def P9() -> str:
+    # head -n 15 reference_population.csv
+    return """# Planet Nine reference population; epoch of JD=2458270.0 (1 June 2018)
+# Orbital elements refer to J2000 heliocentric osculating elements
+#Parameters are:
+# index, mass (Earth masses), semimajor axis (AU), eccentricity, inclination (deg), longitude of perihelion (deg), longitude of ascending node (deg), mean anomaly at epoch (deg), R.A. at epoch (deg), declination at epoch (deg), heliocentric distance at epoch (AU), assumed radius (Earth radii), assumed albedo, V magnitude at epoch, detectable in ZTF survey? (0=no/1=yes), detectable in DES survey (0=no/1=yes), detectable in PS1 (0=no/1=yes)
+index, mass, a, e, inc, varpi, Omega, M, ra, dec, R, rad, albedo, V, ZTF, DES, PS1
+    0,  4.84,  249.58,  0.074,  12.50,   252.85,    48.03,   72.60,  340.9044,  -21.0693,  245.34,  1.61,  0.270,  19.39, 1, 0, 1
+    1,  7.54,  375.45,  0.232,  14.13,   237.85,   111.55,   55.36,  322.4211,  -21.4985,  342.36,  2.51,  0.483,  19.24, 1, 0, 1
+    2, 10.62,  944.92,  0.502,  18.60,   237.12,    95.17,   71.09,   12.2508,  -14.9859, 1022.87,  3.54,  0.732,  22.80, 0, 0, 0
+    3,  5.34,  517.69,  0.605,  20.88,   229.64,    94.30,  204.74,   59.7006,    7.9102,  819.32,  1.78,  0.694,  23.39, 0, 0, 0
+    4,  5.13,  337.08,  0.083,  17.94,   251.60,    68.47,  134.16,   34.8625,    2.8018,  357.57,  1.71,  0.343,  20.64, 1, 1, 1
+    5,  3.78,  278.45,  0.317,  20.10,   247.90,   121.33,   99.15,   24.4436,  -10.8979,  316.44,  1.26,  0.588,  20.19, 1, 1, 1
+    6,  4.74,  275.19,  0.151,  18.73,   259.40,    98.51,  264.27,  152.7116,   26.1338,  285.30,  1.58,  0.693,  19.07, 1, 0, 1
+    7,  6.36,  504.83,  0.356,  21.51,   233.43,   105.46,   69.63,  351.4116,  -23.6505,  504.85,  2.12,  0.628,  21.02, 0, 0, 0
+    8,  9.64,  391.04,  0.226,  11.97,   246.98,    76.04,  128.71,   33.9063,    4.9301,  456.29,  3.21,  0.407,  20.15, 1, 1, 1
+    9,  5.42,  364.43,  0.127,  17.13,   234.72,   123.50,  292.78,  157.2145,   18.5867,  351.84,  1.81,  0.712,  19.66, 1, 0, 1
+"""
+
+
+def test_load_P9(P9: str) -> None:
+
+    with tempfile.TemporaryDirectory() as p9_dir:
+
+        with open(os.path.join(p9_dir, "reference_population.csv"), "w") as f:
+            f.write(P9)
+
+        small_bodies = load_P9(p9_dir)
+        assert len(small_bodies) == 10
+        assert pc.equal(
+            small_bodies.orbits.orbit_id,
+            pa.array(
+                [
+                    "P9_0000000",
+                    "P9_0000001",
+                    "P9_0000002",
+                    "P9_0000003",
+                    "P9_0000004",
+                    "P9_0000005",
+                    "P9_0000006",
+                    "P9_0000007",
+                    "P9_0000008",
+                    "P9_0000009",
+                ]
+            ),
+        )
+
+        assert pc.equal(
+            small_bodies.orbits.object_id,
+            pa.array(
+                [
+                    "0",
+                    "1",
+                    "2",
+                    "3",
+                    "4",
+                    "5",
+                    "6",
+                    "7",
+                    "8",
+                    "9",
+                ]
+            ),
+        )

src/adam_test_data/datasets/tests/test_s3m.py

@@ -73,30 +73,30 @@ def test_load_S3M(S0: str, S1: str, St5: str) -> None:
             small_bodies.orbits.orbit_id,
             pa.array(
                 [
-                    "s3m0000000",
-                    "s3m0000001",
-                    "s3m0000002",
-                    "s3m0000003",
-                    "s3m0000004",
-                    "s3m0000005",
-                    "s3m0000006",
-                    "s3m0000007",
-                    "s3m0000008",
-                    "s3m0000009",
-                    "s3m0000010",
-                    "s3m0000011",
-                    "s3m0000012",
-                    "s3m0000013",
-                    "s3m0000014",
-                    "s3m0000015",
-                    "s3m0000016",
-                    "s3m0000017",
-                    "s3m0000018",
-                    "s3m0000019",
-                    "s3m0000020",
-                    "s3m0000021",
-                    "s3m0000022",
-                    "s3m0000023",
+                    "S3M_0000000",
+                    "S3M_0000001",
+                    "S3M_0000002",
+                    "S3M_0000003",
+                    "S3M_0000004",
+                    "S3M_0000005",
+                    "S3M_0000006",
+                    "S3M_0000007",
+                    "S3M_0000008",
+                    "S3M_0000009",
+                    "S3M_0000010",
+                    "S3M_0000011",
+                    "S3M_0000012",
+                    "S3M_0000013",
+                    "S3M_0000014",
+                    "S3M_0000015",
+                    "S3M_0000016",
+                    "S3M_0000017",
+                    "S3M_0000018",
+                    "S3M_0000019",
+                    "S3M_0000020",
+                    "S3M_0000021",
+                    "S3M_0000022",
+                    "S3M_0000023",
                 ]
             ),
         )