Port iod/od from THOR to adam_core (#125)

* basic tests and od/iod port

* wip for testing fresh clone

* working tests

* placating formatter

* revisions to pyproject to support editable dev install of adam-assist

* Update _version.py

attempts to appease linter

* remove _version.py and add to gitignore

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

@@ -27,10 +27,13 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
           cache: 'pip' # caching pip dependencies 
-      - name: Install Testing Dependencies
+      - name: Install PDM
         run: |
           pip install pip --upgrade
-          pip install ".[dev]"
+          pip install pdm
+      - name: Install Testing Dependencies
+        run: |
+          pdm install -G dev -G test
       - name: Lint
         run: pdm run lint
       - name: Test with coverage

.gitignore

@@ -26,6 +26,7 @@ share/python-wheels/
 .installed.cfg
 *.egg
 MANIFEST
+*_version.py
 
 # PyInstaller
 #  Usually these files are written by a python script from a template
@@ -105,3 +106,4 @@ dmypy.json
 # Ignore lock files
 pdm.lock
 .pdm-build/*
+.pdm-python

pyproject.toml

@@ -30,7 +30,9 @@ dependencies = [
   "healpy",
   "jax",
   "jaxlib",
+  "numba",
   "numpy>=2.0.0",
+  "pdm",
   "pyarrow>=13.0.0",
   "pandas",
   "ray",
@@ -51,55 +53,24 @@ dependencies = [
 [build-system]
 requires = ["pdm-backend"]
 build-backend = "pdm.backend"
-[tool.pdm.build]
-includes = ["src/adam_core/"]
-
-[tool.pdm.version]
-source = "scm"
-write_to = "adam_core/_version.py"
-write_template = "__version__ = '{}'"
-
-
-[tool.pdm.scripts]
-check = { composite = ["lint", "typecheck", "test"] }
-format = { composite = ["black ./src/adam_core", "isort ./src/adam_core"] }
-lint = { composite = [
-  "ruff check ./src/adam_core",
-  "black --check ./src/adam_core",
-  "isort --check-only ./src/adam_core",
-] }
-fix = "ruff ./src/adam_core --fix"
-typecheck = "mypy --strict ./src/adam_core"
-
-test = "pytest --benchmark-disable {args}"
-doctest = "pytest --doctest-plus --doctest-only"
-benchmark = "pytest --benchmark-only"
-coverage = "pytest --cov=adam_core --cov-report=xml"
-
-
 [project.urls]
 "Documentation" = "https://github.com/B612-Asteroid-Institute/adam_core#README.md"
 "Issues" = "https://github.com/B612-Asteroid-Institute/adam_core/issues"
 "Source" = "https://github.com/B612-Asteroid-Institute/adam_core"
 
 
-[project.optional-dependencies]
-assist = [
-  "adam-assist>=0.1.2"
-]
-
-dev = [
-    "black",
-    "ipython>=8.28.0",
-    "isort",
-    "mypy",
-    "pdm",
+[dependency-groups]
+dev = ["ipython>=8.28.0"]
+test = [
     "pytest-benchmark",
     "pytest-cov",
     "pytest-doctestplus",
     "pytest-mock",
     "pytest",
+    "isort",
+    "mypy",
     "ruff",
+    "black"
 ]
 
 [tool.black]
@@ -118,3 +89,32 @@ exclude = ["build"]
 # In order for namespace packages to work during tests,
 # we need to import from the installed modules instead of local source
 addopts = ["--pyargs", "adam_core"]
+
+[tool.pdm.build]
+includes = ["src/adam_core/"]
+
+[tool.pdm.version]
+source = "scm"
+write_to = "adam_core/_version.py"
+write_template = "__version__ = '{}'"
+
+
+[tool.pdm.scripts]
+check = { composite = ["lint", "typecheck", "test"] }
+format = { composite = ["black ./src/adam_core", "isort ./src/adam_core"] }
+lint = { composite = [
+  "ruff check ./src/adam_core",
+  "black --check ./src/adam_core",
+  "isort --check-only ./src/adam_core",
+] }
+fix = "ruff ./src/adam_core --fix"
+typecheck = "mypy --strict ./src/adam_core"
+
+test = "pytest --benchmark-disable {args}"
+doctest = "pytest --doctest-plus --doctest-only"
+benchmark = "pytest --benchmark-only"
+coverage = "pytest --cov=adam_core --cov-report=xml"
+
+
+[tool.pdm.dev-dependencies]
+dev = ["-e git+https://github.com/B612-Asteroid-Institute/adam-assist.git@main#egg=adam-assist"]

src/adam_core/orbit_determination/__init__.py

@@ -1,5 +1,14 @@
 # flake8: noqa: F401
 from .differential_correction import fit_least_squares
 from .evaluate import OrbitDeterminationObservations, evaluate_orbits
-from .fitted_orbits import FittedOrbitMembers, FittedOrbits
-from .outliers import remove_lowest_probability_observation
+from .fitted_orbits import FittedOrbitMembers, FittedOrbits, drop_duplicate_orbits
+from .gauss import gaussIOD
+from .gibbs import calcGibbs
+from .herrick_gibbs import calcHerrickGibbs
+from .iod import (
+    initial_orbit_determination,
+    iod,
+    select_observations,
+    sort_by_id_and_time,
+)
+from .outliers import calculate_max_outliers, remove_lowest_probability_observation

src/adam_core/orbit_determination/fitted_orbits.py

@@ -1,12 +1,130 @@
 import uuid
+from typing import List, Literal, Optional, Tuple
 
+import pyarrow as pa
+import pyarrow.compute as pc
 import quivr as qv
 
 from ..coordinates.cartesian import CartesianCoordinates
 from ..coordinates.residuals import Residuals
 from ..orbits.orbits import Orbits
 
 
+def assign_duplicate_observations(
+    orbits: "FittedOrbits", orbit_members: "FittedOrbitMembers"
+) -> Tuple["FittedOrbits", "FittedOrbitMembers"]:
+    """
+    Assigns observations that have been assigned to multiple orbits to the orbit with the
+    most observations, longest arc length, and lowest reduced chi2.
+
+    Parameters
+    ----------
+    orbit_members : `~thor.orbit_determination.FittedOrbitMembers`
+        Fitted orbit members.
+
+    Returns
+    -------
+    filtered : `~thor.orbit_determination.FittedOrbits`
+        Fitted orbits with duplicate assignments removed.
+    filtered_orbit_members : `~thor.orbit_determination.FittedOrbitMembers`
+        Fitted orbit members with duplicate assignments removed.
+    """
+    # Sorting by priority criteria
+    orbits = orbits.sort_by(
+        [
+            ("num_obs", "descending"),
+            ("arc_length", "descending"),
+            ("reduced_chi2", "ascending"),
+        ]
+    )
+
+    # Extracting unique observation IDs
+    unique_obs_ids = pc.unique(orbit_members.column("obs_id"))
+
+    # Dictionary to store the best orbit for each observation
+    best_orbit_for_obs = {}
+
+    # Iterate over each unique observation ID
+    for obs_id in unique_obs_ids:
+        # Filter orbit_members for the current observation ID
+        mask = pc.equal(orbit_members.column("obs_id"), obs_id)
+        current_obs_members = orbit_members.where(mask)
+
+        # Extract orbit IDs that this observation belongs to
+        obs_orbit_ids = current_obs_members.column("orbit_id")
+
+        # Find the best orbit for this observation based on the criteria
+        for sorted_orbit_id in orbits.column("orbit_id"):
+            if pc.any(pc.is_in(sorted_orbit_id, value_set=obs_orbit_ids)).as_py():
+                best_orbit_for_obs[obs_id.as_py()] = sorted_orbit_id.as_py()
+                break
+
+    # Iteratively update orbit_members to drop rows where obs_id is the same,
+    # but orbit_id is not the best orbit_id for that observation
+    for obs_id, best_orbit_id in best_orbit_for_obs.items():
+        mask_to_remove = pc.and_(
+            pc.equal(orbit_members.column("obs_id"), pa.scalar(obs_id)),
+            pc.not_equal(orbit_members.column("orbit_id"), pa.scalar(best_orbit_id)),
+        )
+        orbit_members = orbit_members.apply_mask(pc.invert(mask_to_remove))
+
+    # Filtering self based on the filtered orbit_members
+    orbits_mask = pc.is_in(
+        orbits.column("orbit_id"), value_set=orbit_members.column("orbit_id")
+    )
+    filtered_orbits = orbits.apply_mask(orbits_mask)
+
+    return filtered_orbits, orbit_members
+
+
+def drop_duplicate_orbits(
+    orbits: "FittedOrbits",
+    orbit_members: "FittedOrbitMembers",
+    subset: Optional[List[str]] = None,
+    keep: Literal["first", "last"] = "first",
+) -> Tuple["FittedOrbits", "FittedOrbitMembers"]:
+    """
+    Drop duplicate orbits from the fitted orbits and remove
+    the corresponding orbit members.
+
+    Parameters
+    ----------
+    orbits : `~thor.orbit_determination.FittedOrbits`
+        Fitted orbits.
+    orbit_members : `~thor.orbit_determination.FittedOrbitMembers`
+        Fitted orbit members.
+    subset : list of str, optional
+        Subset of columns to consider when dropping duplicates. If not specified all the columns
+        specifying unique state are used: time, x, y, z, vx, vy, vz.
+    keep : {'first', 'last'}, default 'first'
+        If there are duplicate rows then keep the first or last row.
+
+    Returns
+    -------
+    filtered : `~thor.orbit_determination.FittedOrbits`
+        Fitted orbits without duplicates.
+    filtered_orbit_members : `~thor.orbit_determination.FittedOrbitMembers`
+        Fitted orbit members without duplicates.
+    """
+    if subset is None:
+        subset = [
+            "coordinates.time.days",
+            "coordinates.time.nanos",
+            "coordinates.x",
+            "coordinates.y",
+            "coordinates.z",
+            "coordinates.vx",
+            "coordinates.vy",
+            "coordinates.vz",
+        ]
+
+    filtered = orbits.drop_duplicates(subset=subset, keep=keep)
+    filtered_orbit_members = orbit_members.apply_mask(
+        pc.is_in(orbit_members.orbit_id, filtered.orbit_id)
+    )
+    return filtered, filtered_orbit_members
+
+
 class FittedOrbits(qv.Table):
 
     orbit_id = qv.LargeStringColumn(default=lambda: uuid.uuid4().hex)