Track NEO example (#147)

* Track NEO example

* Add back-pressure check to propagate orbits

examples/track_neo.ipynb

@@ -0,0 +1,202 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Tracking NEOCP for Observatory Follow-Up\n",
+    "\n",
+    "This example shows several ways you can use adam-core to generate ephemeris for objects on the NEOCP. First you will need a couple packages:\n",
+    "\n",
+    "```bash\n",
+    "pip install adam-core\n",
+    "pip install adam-assist\n",
+    "```\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# First let's fetch the objects in question\n",
+    "from adam_core.orbits.query.scout import query_scout, get_scout_objects\n",
+    "\n",
+    "# Using the CNEOS Scout API, we can conveniently get a list of objects currently in \n",
+    "# the NEOCP\n",
+    "scout_objects = get_scout_objects()\n",
+    "\n",
+    "# You can preview the objects like so:\n",
+    "print(scout_objects.to_dataframe())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Decide which objects you want to track. For now, we select the first one.\n",
+    "object_of_interest = scout_objects[10]\n",
+    "\n",
+    "# You could also select by object name like so:\n",
+    "# object_of_interest = scout_objects.select(\"objectName\", \"P126Tdm\")\n",
+    "# Or you could sort by something like phaScore like so:\n",
+    "# sorted_by_pha_score = scout_objects.sort_by([(\"phaScore\", \"descending\")])\n",
+    "\n",
+    "# With your objects chosen, we send in an array of ids / object names to the Scout API\n",
+    "# This gives us the Scout variants for each object\n",
+    "samples = query_scout(object_of_interest.objectName)\n",
+    "\n",
+    "# You generally get back 1000 variants from Scout\n",
+    "# VariantOrbits(size=1000)\n",
+    "\n",
+    "# Let's collapse the sampled orbits into an uncertainty.\n",
+    "orbits = samples.collapse_by_object_id()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from adam_core.time import Timestamp\n",
+    "from adam_core.observers import Observers\n",
+    "\n",
+    "# Decide on the exposure times you want to use\n",
+    "# Lots of options here, including from astropy.time.Time objects, from MJD, JD, etc.\n",
+    "times = Timestamp.from_iso8601(\n",
+    "    [\n",
+    "        \"2025-02-23T00:00:00Z\", \"2025-02-23T00:05:00Z\", \"2025-02-23T00:10:00Z\"\n",
+    "    ], scale=\"utc\"\n",
+    ")\n",
+    "\n",
+    "# Now we define observer positions from an observatory code\n",
+    "observers = Observers.from_code(\"T08\", times)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Now we can generate some on-sky ephemeris with uncertainty\n",
+    "# You could also use the Scout API directly to do this: https://ssd-api.jpl.nasa.gov/doc/scout.html\n",
+    "from adam_assist import ASSISTPropagator\n",
+    "propagator = ASSISTPropagator()\n",
+    "\n",
+    "ephemeris = propagator.generate_ephemeris(\n",
+    "    orbits,\n",
+    "    observers,\n",
+    "    covariance=True,\n",
+    "    num_samples=1000,\n",
+    "    max_processes=10  # Using multiprocessing helps speed things up\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Now we can iterate through the unique times and determine where to point\n",
+    "for ephem in ephemeris:\n",
+    "    print(f\"Object ID: {ephem.object_id[0]}\")\n",
+    "    print(f\"\\nTime: {ephem.coordinates.time.to_iso8601()[0]}\")\n",
+    "    print(f\"RA: {ephem.coordinates.lon[0]}, Sigma RA: {ephem.coordinates.sigma_lon[0]}\")\n",
+    "    print(f\"DEC: {ephem.coordinates.lat[0]}, Sigma DEC: {ephem.coordinates.sigma_lat[0]}\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# The Ephemeris object is a Quivr Table (based on pyarrow), but you can also convert to a pandas DataFrame\n",
+    "ephemeris.to_dataframe()\n",
+    "# That's it, you're all set!\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Object ID: C11QM25\n",
+      "\n",
+      "Time: 2025-02-23T00:00:00.000\n",
+      "RA: 131.8765804600289 - 312.20777783150254\n",
+      "DEC: -14.388007767026275 - 51.40213360541426\n",
+      "Object ID: C11QM25\n",
+      "\n",
+      "Time: 2025-02-23T00:05:00.000\n",
+      "RA: 131.87519081007588 - 312.2100283360548\n",
+      "DEC: -14.395731775311084 - 50.316578057588345\n",
+      "Object ID: C11QM25\n",
+      "\n",
+      "Time: 2025-02-23T00:10:00.000\n",
+      "RA: 131.87379975693307 - 312.2123773385522\n",
+      "DEC: -14.403355028559567 - 50.756042539885584\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Alternatively, if you want you can propagate the Scout samples directly\n",
+    "# and use their distribution to ascertain uncertainty in the on-sky location\n",
+    "import pyarrow.compute as pc\n",
+    "\n",
+    "sample_direct_ephem = propagator.generate_ephemeris(\n",
+    "    samples,\n",
+    "    observers,\n",
+    "    max_processes=10\n",
+    ")\n",
+    "\n",
+    "unique_times = sample_direct_ephem.coordinates.time.unique()\n",
+    "for unique_time in unique_times:\n",
+    "    sample_ephem_time = sample_direct_ephem.apply_mask(sample_direct_ephem.coordinates.time.equals(unique_time))\n",
+    "    print(f\"\\nObject ID: {sample_ephem_time.object_id[0]}\")\n",
+    "    print(f\"Time: {sample_ephem_time.coordinates.time.to_iso8601()[0]}\")\n",
+    "    print(f\"RA: {pc.min(sample_ephem_time.coordinates.lon)} - {pc.max(sample_ephem_time.coordinates.lon)}\")\n",
+    "    print(f\"DEC: {pc.min(sample_ephem_time.coordinates.lat)} - {pc.max(sample_ephem_time.coordinates.lat)}\")\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

src/adam_core/orbits/query/scout.py

@@ -159,13 +159,17 @@ def scout_orbits_to_variant_orbits(
         tp=pc.subtract(pc.cast(scout_orbits.tp, pa.float64()), 2400000.5),
         time=Timestamp.from_jd(pc.cast(scout_orbits.epoch, pa.float64())),
         origin=Origin.from_kwargs(code=pa.repeat("SUN", len(scout_orbits))),
+        frame="ecliptic",
     )
 
     cartesian_coords = cometary_coords.to_cartesian()
 
+    unique_orbit_ids = pc.cast(scout_orbits.idx, pa.large_string())
+
     variants = VariantOrbits.from_kwargs(
         coordinates=cartesian_coords,
-        orbit_id=pc.cast(scout_orbits.idx, pa.large_string()),
+        orbit_id=unique_orbit_ids,
+        variant_id=unique_orbit_ids,
         object_id=pa.repeat(object_id, len(scout_orbits)),
     )
 

src/adam_core/orbits/query/tests/test_scout.py

@@ -0,0 +1,52 @@
+"""Tests for the scout module."""
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.compute as pc
+
+from ..scout import ScoutOrbit, scout_orbits_to_variant_orbits
+
+
+def test_scout_orbits_to_variant_orbits():
+    """Test that scout orbits are correctly converted to variant orbits."""
+    # Create a mock scout orbits table
+    scout_data = {
+        "idx": [0, 1],
+        "epoch": ["60000.0", "60000.0"],
+        "ec": ["0.5", "0.51"],
+        "qr": ["1.0", "1.01"],
+        "tp": ["59000.0", "59000.0"],
+        "om": ["10.0", "10.1"],
+        "w": ["50.0", "50.1"],
+        "inc": ["10.0", "10.1"],
+        "H": ["20.0", "20.0"],
+        "dca": ["0.1", "0.1"],
+        "tca": ["0.1", "0.1"],
+        "moid": ["0.1", "0.1"],
+        "vinf": ["0.1", "0.1"],
+        "geoEcc": ["0.1", "0.1"],
+        "impFlag": [0, 0],
+    }
+    scout_orbits = ScoutOrbit.from_kwargs(**scout_data)
+
+    # Convert to variant orbits
+    variant_orbits = scout_orbits_to_variant_orbits("2024AA", scout_orbits)
+
+    # Check that the output has the expected structure
+    assert len(variant_orbits) == len(scout_orbits)
+    assert variant_orbits.coordinates.frame == "ecliptic"
+    assert pc.all(pc.equal(variant_orbits.coordinates.origin.code, "SUN")).as_py()
+
+    # Check that the object IDs are correct
+    assert variant_orbits.object_id.to_pylist() == ["2024AA", "2024AA"]
+
+    # Check that the orbit IDs are correct
+    assert variant_orbits.orbit_id.to_pylist() == ["0", "1"]
+
+    # Check that the variant IDs are unique
+    assert len(pc.unique(variant_orbits.variant_id)) == len(scout_orbits)
+
+    # Check that the time is correct
+    np.testing.assert_array_equal(
+        variant_orbits.coordinates.time.jd(), pc.cast(scout_orbits.epoch, pa.float64())
+    )