Merge pull request #382 from alercebroker/feat/conesearch

feat: implement conesearch functionality with OID and coordinate queries

Author: dirodriguezm

multisurveys-apis/poetry.lock

@@ -30,6 +30,8 @@ pyyaml = "^6.0.2"
 polars = "^1.31.0"
 pyarrow = "^20.0.0"
 pandas = "^2.3.1"
+toolz = "^1.0.0"
+astropy = "^7.1.0"
 
 [tool.poetry.group.test.dependencies]
 pytest = "^7.4.0"

multisurveys-apis/pyproject.toml

@@ -49,9 +49,6 @@ def catalog_oid_to_masterid(
         The ZTF object ID.
     validate: bool
         If True, validate the ztf_oid before conversion.
-    db_cursor: psycopg2.extensions.cursor
-        Database cursor for LSST catalog. This parameter is required for LSST.
-
     Returns
     -------
     str

multisurveys-apis/src/core/idmapper/idmapper.py

@@ -0,0 +1,69 @@
+from typing import Callable, ContextManager, List
+from db_plugins.db.sql.models import Object
+from numpy import int64
+from sqlalchemy import asc, select, text
+from sqlalchemy.orm import Session, aliased
+
+
+def conesearch_coordinates(
+    session_factory: Callable[..., ContextManager[Session]],
+):
+    def _conesearch(
+        ra: float,
+        dec: float,
+        radius: float,
+        neighbors: int,
+    ) -> List[Object]:
+        stmt = _build_statement_coordinates(neighbors)
+        with session_factory() as session:
+            result = session.execute(
+                stmt, {"ra": ra, "dec": dec, "radius": radius}
+            ).all()
+            return [row[0] for row in result]
+
+    return _conesearch
+
+
+def _build_statement_coordinates(neighbors: int):
+    return (
+        select(Object)
+        .where(text("q3c_radial_query(meanra, meandec, :ra, :dec, :radius)"))
+        .order_by(asc(text("q3c_dist(meanra, meandec, :ra, :dec)")))
+        .limit(neighbors)
+    )
+
+
+def conesearch_oid(session_factory: Callable[..., ContextManager[Session]]):
+    def _conesearch(oid: int64, radius: float, neighbors: int) -> List[Object]:
+        stmt = _build_statement_oid(oid, neighbors)
+        with session_factory() as session:
+            result = session.execute(stmt, {"radius": radius}).all()
+            return [row[0] for row in result]
+
+    return _conesearch
+
+
+def _build_statement_oid(oid: int64, neighbors: int):
+    # Create aliases for the Object table
+    center_obj = aliased(Object, name="center")
+    target_obj = aliased(Object, name="target")
+
+    # Build the query using q3c_radial_query function
+    return (
+        select(target_obj.oid, target_obj.meanra, target_obj.meandec)
+        .select_from(center_obj, target_obj)
+        .where(center_obj.oid == oid.item())
+        .where(
+            text(
+                "q3c_radial_query(target.meanra, target.meandec, center.meanra, center.meandec, :radius)"
+            )
+        )
+        .order_by(
+            asc(
+                text(
+                    "q3c_dist(target.meanra, target.meandec, center.meanra, center.meandec)"
+                )
+            )
+        )
+        .limit(neighbors)
+    )

multisurveys-apis/src/core/repository/queries/conesearch.py

@@ -1,7 +1,7 @@
 from pydantic import BaseModel
 
 
-class AlerceObject(BaseModel):
+class ApiObject(BaseModel):
     objectId: str
     ra: float
     dec: float

multisurveys-apis/src/lightcurve_api/models/object.py

@@ -1,22 +1,32 @@
 import traceback
+from typing import List
 
-import numpy as np
 from fastapi import APIRouter, HTTPException
 
 from core.config.dependencies import db_dependency
-from core.idmapper.idmapper import decode_masterid
+from core.idmapper.idmapper import catalog_oid_to_masterid
+
+from ..services.conesearch import conesearch as service
+from ..models.object import ApiObject
 
 router = APIRouter()
 
 
-@router.get("/conesearch")
+@router.get("/conesearch_oid/objects")
 def conesearch(
-    oid: int,
+    oid: str,
+    survey: str,
+    radius: float,
+    neighbors: int,
     db: db_dependency,
-):
-    survey, id = decode_masterid(np.int64(oid))
+) -> List[ApiObject]:
+    try:
+        id = catalog_oid_to_masterid(survey, oid, True)
+    except ValueError:
+        raise HTTPException(status_code=400, detail="Invalid object ID")
+
     try:
-        return f"{survey} {id}"
+        return service.conesearch_oid(id, radius, neighbors, db.session)
     except Exception:
         traceback.print_exc()
         raise HTTPException(status_code=500, detail="An error occurred")

multisurveys-apis/src/lightcurve_api/routes/conesearch.py

@@ -0,0 +1,46 @@
+from typing import List, cast, Callable, ContextManager
+from lightcurve_api.services.conesearch.validation import (
+    validate_coordinates_params,
+    validate_oid_params,
+)
+from lightcurve_api.models.object import ApiObject
+import core.repository.queries.conesearch as queries
+from sqlalchemy.orm import Session
+from .parser import parsesapi_objects
+from toolz.functoolz import pipe
+from numpy import int64
+
+
+def conesearch_coordinates(
+    ra: float,
+    dec: float,
+    radius: float,
+    neighbors: int,
+    session_factory: Callable[..., ContextManager[Session]],
+) -> List[ApiObject]:
+    return cast(
+        List[ApiObject],
+        pipe(
+            (ra, dec, radius, neighbors),
+            validate_coordinates_params,
+            queries.conesearch_coordinates(session_factory),
+            parsesapi_objects,
+        ),
+    )
+
+
+def conesearch_oid(
+    oid: int64,
+    radius: float,
+    neighbors: int,
+    session_factory: Callable[..., ContextManager[Session]],
+) -> List[ApiObject]:
+    return cast(
+        List[ApiObject],
+        pipe(
+            (oid, radius, neighbors),
+            validate_oid_params,
+            queries.conesearch_oid(session_factory),
+            parsesapi_objects,
+        ),
+    )

multisurveys-apis/src/lightcurve_api/services/conesearch.py

@@ -0,0 +1,19 @@
+from typing import List
+from db_plugins.db.sql.models import Object
+
+from lightcurve_api.models.object import ApiObject
+from core.idmapper.idmapper import decode_masterid
+from numpy import int64
+
+
+def parse_api_object(sql_object: Object) -> ApiObject:
+    _, parsedOid = decode_masterid(int64(str(sql_object.oid)))
+    return ApiObject(
+        objectId=str(parsedOid),
+        ra=sql_object.meanra,  # type: ignore
+        dec=sql_object.meandec,  # type: ignore
+    )
+
+
+def parsesapi_objects(sql_objects: List[Object]) -> List[ApiObject]:
+    return [parse_api_object(sql_object) for sql_object in sql_objects]

multisurveys-apis/src/lightcurve_api/services/conesearch/conesearch.py

@@ -0,0 +1,24 @@
+from typing import Tuple
+from numpy import int64
+
+
+def validate_coordinates_params(
+    args: Tuple[float, float, float, int],
+) -> Tuple[float, float, float, int]:
+    ra, dec, radius, neighbors = args
+    if ra < 0 or ra > 360:
+        raise ValueError("RA must be greater than 0 and lower than 360")
+    if dec <= -90 or dec >= 90:
+        raise ValueError("Dec must be greater than -90 and lower than 90")
+    if radius <= 0:
+        raise ValueError("Radius must be greater than 0")
+    if neighbors <= 0:
+        raise ValueError("Neighbors must be greater than 0")
+    return ra, dec, radius, neighbors
+
+
+def validate_oid_params(
+    args: Tuple[int64, float, int],
+) -> Tuple[int64, float, int]:
+    oid, radius, neighbors = args
+    return oid, radius, neighbors