Merge pull request #109 from moeyensj/v2.0-range-and-shift

Add link_test_orbit

.github/workflows/conda-build-lint-test.yml

@@ -2,9 +2,9 @@ name: conda - Build Lint and Test
 
 on:
   push:
-    branches: [ main, "v*"]
+    branches: [ main ]
   pull_request:
-    branches: [ main, "v*"]
+    branches: [ main ]
 
 jobs:
   build-lint-test:

.github/workflows/docker-build-lint-test.yml

@@ -3,9 +3,9 @@ name: docker - Build Lint and Test
 
 on:
   push:
-    branches: [main, "v*"]
+    branches: [ main ]
   pull_request:
-    branches: [main, "v*"]
+    branches: [ main ]
 
 jobs:
   build-lint-test:

.github/workflows/pip-build-lint-test-coverage.yml

@@ -2,9 +2,9 @@ name: pip - Build Lint Test and Coverage
 
 on:
   push:
-    branches: [ main, "v*"]
+    branches: [ main ]
   pull_request:
-    branches: [ main, "v*"]
+    branches: [ main ]
 
 jobs:
   build-lint-test-coverage:

recipe/meta.yaml

@@ -17,6 +17,7 @@ requirements:
   run:
     - python {{ python }}
     - numpy
+    - pyarrow >= 13.0.0
     - numba
     - pandas
     - difi

requirements.txt

@@ -1,5 +1,6 @@
 adam-core @ git+https://github.com/B612-Asteroid-Institute/adam_core@main
 numpy
+pyarrow>=13.0.0
 numba
 pandas
 difi

setup.cfg

@@ -33,13 +33,14 @@ setup_requires =
 install_requires =
     adam-core @ git+https://github.com/B612-Asteroid-Institute/adam_core@main
     numpy
+    pyarrow >= 13.0.0
     numba
     pandas
     difi
     astropy >= 5.3.1
     astroquery
     jax
-    quivr
+    quivr>=0.7.1
     scipy
     scikit-learn
     healpy

thor/__init__.py

@@ -12,6 +12,7 @@
 from .orbit_selection import *
 from .filter_orbits import *
 from .main import *
+from .main_2 import *  # TODO: remove when main is replaced
 from .analysis import *
 
 logger = setupLogger(__name__)

thor/main_2.py

@@ -0,0 +1,147 @@
+from typing import List, Optional
+
+import pyarrow.compute as pc
+import quivr as qv
+from adam_core.coordinates import (
+    CartesianCoordinates,
+    OriginCodes,
+    transform_coordinates,
+)
+from adam_core.propagator import PYOORB, Propagator
+
+from .observations import Observations
+from .observations.filters import ObservationFilter, TestOrbitRadiusObservationFilter
+from .orbit import TestOrbit
+from .projections import GnomonicCoordinates
+
+
+class TransformedDetections(qv.Table):
+    id = qv.StringColumn()
+    coordinates = GnomonicCoordinates.as_column()
+    state_id = qv.Int64Column()
+
+
+def range_and_transform(
+    test_orbit: TestOrbit,
+    observations: Observations,
+    filters: Optional[List[ObservationFilter]] = None,
+    propagator: Propagator = PYOORB(),
+    max_processes: int = 1,
+) -> TransformedDetections:
+    """
+    Gather observations for a single test orbit, and transform them into a
+    gnomonic projection centered on the motion of the test orbit (co-rotating
+    frame).
+
+    Parameters
+    ----------
+    obs_src : `~thor.observation_filters.ObservationFilter`
+        Observation filter to use to gather observations given the test orbit.
+    test_orbit : `~thor.orbit.TestOrbit`
+        Test orbit to use to gather and transform observations.
+    propagator : `~adam_core.propagator.propagator.Propagator`
+        Propagator to use to propagate the test orbit and generate
+        ephemerides.
+    max_processes : int, optional
+        Maximum number of processes to use for parallelization.
+
+    Returns
+    -------
+    transformed_detections : `~thor.main.TransformedDetections`
+        The transformed detections as gnomonic coordinates
+        of the observations in the co-rotating frame.
+    """
+    # Compute the ephemeris of the test orbit (this will be cached)
+    ephemeris = test_orbit.generate_ephemeris_from_observations(
+        observations,
+        propagator=propagator,
+        max_processes=max_processes,
+    )
+
+    # Apply filters to the observations
+    filtered_observations = observations
+    if filters is not None:
+        for filter_i in filters:
+            filtered_observations = filter_i.apply(filtered_observations, test_orbit)
+
+    # Assume that the heliocentric distance of all point sources in
+    # the observations are the same as that of the test orbit
+    ranged_detections_spherical = test_orbit.range_observations(
+        filtered_observations,
+        propagator=propagator,
+        max_processes=max_processes,
+    )
+
+    # Transform from spherical topocentric to cartesian heliocentric coordinates
+    ranged_detections_cartesian = transform_coordinates(
+        ranged_detections_spherical.coordinates,
+        representation_out=CartesianCoordinates,
+        frame_out="ecliptic",
+        origin_out=OriginCodes.SUN,
+    )
+
+    # Link the ephemeris and observations by state id
+    link = qv.Linkage(
+        ephemeris,
+        filtered_observations,
+        left_keys=ephemeris.id,
+        right_keys=filtered_observations.state_id,
+    )
+
+    # Transform the detections into the co-rotating frame
+    transformed_detection_list = []
+    state_ids = filtered_observations.state_id.unique().sort()
+    for state_id in state_ids:
+        # Select the detections and ephemeris for this state id
+        mask = pc.equal(state_id, filtered_observations.state_id)
+        ranged_detections_cartesian_i = ranged_detections_cartesian.apply_mask(mask)
+        ephemeris_i = link.select_left(state_id)
+        observations_i = link.select_right(state_id)
+
+        # Transform the detections into the co-rotating frame
+        transformed_detections_i = TransformedDetections.from_kwargs(
+            id=observations_i.detections.id,
+            coordinates=GnomonicCoordinates.from_cartesian(
+                ranged_detections_cartesian_i,
+                center_cartesian=ephemeris_i.ephemeris.aberrated_coordinates,
+            ),
+            state_id=observations_i.state_id,
+        )
+
+        transformed_detection_list.append(transformed_detections_i)
+
+    transformed_detections = qv.concatenate(transformed_detection_list)
+    return transformed_detections
+
+
+def link_test_orbit(
+    test_orbit: TestOrbit,
+    observations: Observations,
+    filters: Optional[List[ObservationFilter]] = [
+        TestOrbitRadiusObservationFilter(radius=10.0)
+    ],
+    propagator: Propagator = PYOORB(),
+    max_processes: int = 1,
+):
+    """
+    Find all linkages for a single test orbit.
+    """
+    # Range and transform the observations
+    transformed_detections = range_and_transform(
+        test_orbit,
+        observations,
+        filters=filters,
+        propagator=propagator,
+        max_processes=max_processes,
+    )
+
+    # TODO: Find objects which move in straight-ish lines in the gnomonic frame.
+    #
+    # TODO: Run IOD against each of the discovered straight lines, and
+    # filter down to plausible orbits.
+    #
+    # TODO: Run OD against the plausible orbits, and filter down to really good orbits.
+    #
+    # TODO: Perform arc extension on the really good orbits.
+
+    return transformed_detections

thor/observations/__init__.py

@@ -0,0 +1,2 @@
+# flake8: noqa: F401
+from .observations import Observations