diff --git a/src/adam_test_data/datasets/__init__.py b/src/adam_test_data/datasets/__init__.py index 469e6cb..150bcc8 100644 --- a/src/adam_test_data/datasets/__init__.py +++ b/src/adam_test_data/datasets/__init__.py @@ -1,2 +1,3 @@ # ruff: noqa: F401 +from .p9 import load_P9 from .s3m import load_S3M diff --git a/src/adam_test_data/datasets/p9.py b/src/adam_test_data/datasets/p9.py new file mode 100644 index 0000000..fa9150f --- /dev/null +++ b/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", + ) diff --git a/src/adam_test_data/datasets/s3m.py b/src/adam_test_data/datasets/s3m.py index beccdd8..c5f81ce 100644 --- a/src/adam_test_data/datasets/s3m.py +++ b/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, diff --git a/src/adam_test_data/datasets/tests/test_p9.py b/src/adam_test_data/datasets/tests/test_p9.py new file mode 100644 index 0000000..abc3c8d --- /dev/null +++ b/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", + ] + ), + ) diff --git a/src/adam_test_data/datasets/tests/test_s3m.py b/src/adam_test_data/datasets/tests/test_s3m.py index ee0842f..6ca4928 100644 --- a/src/adam_test_data/datasets/tests/test_s3m.py +++ b/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", ] ), )