Add properties attribute to sources

lightcurvedb/client/source.py

@@ -2,12 +2,32 @@
 Extensions to core for sources.
 """
 
+import asyncio
 from math import cos, pi
 
-from lightcurvedb.models.source import Source
+from lightcurvedb.models.source import Source, SourceProperties
 from lightcurvedb.storage.prototype.backend import Backend
 
 
+async def source_read_all(backend: Backend) -> list[Source]:
+    """
+    Read all sources, with computed properties (e.g. median flux per module
+    and frequency) merged in. Sources with no flux measurements are returned
+    with `properties` left unset.
+    """
+    sources, median_flux_by_source = await asyncio.gather(
+        backend.sources.get_all(),
+        backend.analysis.get_median_flux_for_all_sources(),
+    )
+
+    for source in sources:
+        per_module_frequency = median_flux_by_source.get(source.source_id)
+        if per_module_frequency:
+            source.properties = SourceProperties(median_flux=per_module_frequency)
+
+    return sources
+
+
 async def source_read_in_radius(
     center: tuple[float, float], radius: float, backend: Backend
 ) -> list[Source]:

lightcurvedb/models/__init__.py

@@ -16,7 +16,7 @@
     CandidateReviewDecision,
     ExternalMatchEvidence,
 )
-from .source import CrossMatch, Source, SourceMetadata
+from .source import CrossMatch, Source, SourceMetadata, SourceProperties
 from .statistics import SourceStatistics
 from .unassigned_flux import UnassignedFluxMeasurement, UnassignedMeasurementMetadata
 from .unassigned_source import UnassignedSource, UnassignedSourceMetadata
@@ -32,6 +32,7 @@
     "Source",
     "SourceMetadata",
     "SourceNotFoundException",
+    "SourceProperties",
     "SourceStatistics",
     "StorageException",
     "UnassignedFluxMeasurementNotFoundException",

lightcurvedb/models/feed.py

@@ -7,6 +7,8 @@
 
 from pydantic import BaseModel
 
+from lightcurvedb.models.instrument import band_name as format_band_name
+
 
 class FeedResultItem(BaseModel):
     source_id: UUID
@@ -30,4 +32,4 @@ class FeedResult(BaseModel):
 
     @property
     def band_name(self) -> str:
-        return f"f{self.frequency}"
+        return format_band_name(self.frequency)

lightcurvedb/models/instrument.py

@@ -13,3 +13,10 @@ class Instrument(BaseModel):
     telescope: str
     instrument: str
     details: dict[str, Any]
+
+
+def band_name(frequency: int) -> str:
+    """
+    Format a frequency (GHz) as its band name, e.g. 90 -> "f090".
+    """
+    return f"f{frequency:03d}"

lightcurvedb/models/source.py

@@ -27,6 +27,16 @@ class SourceMetadata(BaseModel):
     socat_id: UUID | None = None
 
 
+class SourceProperties(BaseModel):
+    """
+    Additional properties about sources stored as a JSONB
+    column.
+    """
+
+    # Keyed as f"{module}_{frequency}", e.g. "i1_90".
+    median_flux: dict[str, float]
+
+
 class Source(BaseModel):
     """
     Input model for creating sources.
@@ -39,3 +49,4 @@ class Source(BaseModel):
     dec: float | None
     variable: bool = False
     extra: SourceMetadata | None = None
+    properties: SourceProperties | None = None

lightcurvedb/storage/parquet/analysis.py

@@ -153,3 +153,32 @@ async def get_source_statistics(
             return {str(stats.frequency): stats for stats in statistics}
 
         return {f"{stats.module}_{stats.frequency}": stats for stats in statistics}
+
+    async def get_median_flux_for_all_sources(self) -> dict[UUID, dict[str, float]]:
+        """
+        Get the median flux for every source, grouped by module and frequency
+        (keyed as f"{module}_{frequency}", matching get_source_statistics's
+        non-collated key format). The parquet backend stores one file per
+        source, so this reads each source's file in turn rather than a
+        single aggregate query.
+        """
+        result: dict[UUID, dict[str, float]] = {}
+
+        base_path = self.flux_storage.base_path
+        if not base_path.exists():
+            return result
+
+        for path in base_path.glob("*.parquet"):
+            source_id = UUID(path.stem)
+            table = await self.flux_storage._read_file(source_id)
+
+            if table is None or table.empty:
+                continue
+
+            medians = table.groupby(["module", "frequency"])["flux"].median()
+            result[source_id] = {
+                f"{module}_{int(frequency)}": float(value)
+                for (module, frequency), value in medians.items()
+            }
+
+        return result

lightcurvedb/storage/postgres/analysis.py

@@ -3,6 +3,7 @@
 """
 
 import asyncio
+from collections import defaultdict
 from datetime import datetime
 from uuid import UUID
 
@@ -181,3 +182,40 @@ async def get_source_statistics(
                 return {
                     f"{stats.module}_{stats.frequency}": stats for stats in statistics
                 }
+
+    async def get_median_flux_for_all_sources(self) -> dict[UUID, dict[str, float]]:
+        """
+        Get the median flux for every source, grouped by module and frequency
+        (keyed as f"{module}_{frequency}", matching
+        get_source_statistics_for_frequency_and_module's non-collated key
+        format), in a single aggregate query (avoids one query per source).
+
+        Based on the last 30 days to avoid sorting over the entire table, which would
+        scale poorly as data accumulates. This mirrors the "current month" semantics
+        used by TimescaleAnalysisProvider's continuous-aggregate version, which is
+        grouped by module for the same reason its other monthly aggregates are.
+
+        Note that in both implementations, a source with no measurements in the last
+        30 days would report no median_flux.
+        """
+        query = """
+            SELECT
+                source_id,
+                module,
+                frequency,
+                PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY flux) as median_flux
+            FROM flux_measurements
+            WHERE time >= now() - interval '30 days'
+            GROUP BY source_id, module, frequency
+        """
+
+        with self.tracer.start_as_current_span("get_median_flux_for_all_sources"):
+            async with self.flux_storage.cursor() as cur:
+                await cur.execute(query)
+                rows = await cur.fetchall()
+
+        result: dict[UUID, dict[str, float]] = defaultdict(dict)
+        for source_id, module, frequency, median_flux in rows:
+            result[source_id][f"{module}_{frequency}"] = median_flux
+
+        return dict(result)

lightcurvedb/storage/prototype/analysis.py

@@ -54,3 +54,11 @@ async def get_source_statistics(
         Get source statistics across all frequencies and modules.
         """
         ...
+
+    async def get_median_flux_for_all_sources(self) -> dict[UUID, dict[str, float]]:
+        """
+        Get the median flux for every source, grouped by module and frequency
+        (keyed as f"{module}_{frequency}"), in a single batched operation
+        (avoids one query per source).
+        """
+        ...

lightcurvedb/storage/timescale/analysis.py

@@ -0,0 +1,40 @@
+"""
+TimescaleDB analysis provider.
+
+Overrides get_median_flux_for_all_sources to read from the flux_monthly
+continuous aggregate (created by TimescaleLightcurveProvider.setup()) instead
+of scanning the full flux_measurements hypertable on every call - see
+PostgresAnalysisProvider for the naive version this replaces.
+"""
+
+from collections import defaultdict
+from uuid import UUID
+
+from lightcurvedb.storage.postgres.analysis import PostgresAnalysisProvider
+
+
+class TimescaleAnalysisProvider(PostgresAnalysisProvider):
+    async def get_median_flux_for_all_sources(self) -> dict[UUID, dict[str, float]]:
+        """
+        Get the median flux for every source, grouped by module and frequency
+        (keyed as f"{module}_{frequency}", matching
+        PostgresAnalysisProvider's key format), from the current month's
+        bucket of the flux_monthly continuous aggregate rather than scanning
+        all of flux_measurements.
+        """
+        query = """
+            SELECT source_id, module, frequency, median_flux
+            FROM flux_monthly
+            WHERE bucket = time_bucket('30 days', now())
+        """
+
+        with self.tracer.start_as_current_span("get_median_flux_for_all_sources"):
+            async with self.flux_storage.cursor() as cur:
+                await cur.execute(query)
+                rows = await cur.fetchall()
+
+        result: dict[UUID, dict[str, float]] = defaultdict(dict)
+        for source_id, module, frequency, median_flux in rows:
+            result[source_id][f"{module}_{frequency}"] = median_flux
+
+        return dict(result)

lightcurvedb/storage/timescale/backend.py

@@ -9,13 +9,13 @@
 from psycopg_pool import AsyncConnectionPool
 
 from lightcurvedb.config import Settings
-from lightcurvedb.storage.postgres.analysis import PostgresAnalysisProvider
 from lightcurvedb.storage.postgres.instrument import PostgresInstrumentStorage
 from lightcurvedb.storage.postgres.source import PostgresSourceStorage
 from lightcurvedb.storage.postgres.unassigned_source import (
     PostgresUnassignedSourceStorage,
 )
 from lightcurvedb.storage.prototype.backend import Backend
+from lightcurvedb.storage.timescale.analysis import TimescaleAnalysisProvider
 from lightcurvedb.storage.timescale.cutout import TimescaleCutoutStorage
 from lightcurvedb.storage.timescale.flux import TimescaleFluxMeasurementStorage
 from lightcurvedb.storage.timescale.lightcurves import TimescaleLightcurveProvider
@@ -36,7 +36,7 @@ async def generate_timescale_backend(pool: AsyncConnectionPool) -> Backend:
     lightcurves = TimescaleLightcurveProvider(
         flux_storage=fluxes, tracer=tracer, meter=meter
     )
-    analysis = PostgresAnalysisProvider(
+    analysis = TimescaleAnalysisProvider(
         flux_storage=fluxes,
         lightcurve_provider=lightcurves,
         tracer=tracer,

lightcurvedb/storage/timescale/schema.py

@@ -162,7 +162,8 @@
         THEN (sqrt(sum(flux_err ^ 2) FILTER (WHERE flux_err IS NOT NULL))
               / count(flux_err) FILTER (WHERE flux_err IS NOT NULL))::real
         ELSE NULL
-    END AS avg_flux_err
+    END AS avg_flux_err,
+    percentile_cont(0.5) WITHIN GROUP (ORDER BY flux) AS median_flux
 FROM flux_measurements
 GROUP BY bucket, source_id, frequency, module
 WITH NO DATA;

tests/test_backend/test_analysis.py

@@ -32,7 +32,10 @@ async def test_weighted_statistics(backend):
         )
     )
 
-    base_time = datetime.datetime(2024, 1, 1, tzinfo=timezone.utc)
+    # Recent, not a fixed historical date: median_flux is only tracked for
+    # the current month's bucket on the timescale backend (see
+    # TimescaleAnalysisProvider), so older fixed dates wouldn't show up there.
+    base_time = datetime.datetime.now(timezone.utc) - datetime.timedelta(days=2)
     measurements = [
         FluxMeasurement(
             measurement_id=uuid7(),
@@ -64,6 +67,11 @@ async def test_weighted_statistics(backend):
     assert stats.min_flux == 10.0
     assert stats.max_flux == 10.0
 
+    # Batched median flux across all sources should agree with the per-source
+    # statistics computed above.
+    median_by_source = await backend.analysis.get_median_flux_for_all_sources()
+    assert median_by_source[source]["test-weighted_999"] == stats.median_flux
+
     # Delete those measurements:
     for measurement in measurements:
         await backend.fluxes.delete(measurement_id=measurement.measurement_id)

tests/test_backend/test_candidate_review.py

@@ -135,6 +135,9 @@ async def test_terminal_decisions_materialize_directly_and_retain_metadata(backe
     canonical_measurement = await backend.fluxes.get(measurement.measurement_id)
     assert canonical_measurement.source_id == canonical_source.source_id
 
+    await backend.fluxes.delete(measurement_id=canonical_measurement.measurement_id)
+    await backend.sources.delete(source_id=canonical_source.source_id)
+
 
 @pytest.mark.asyncio(loop_scope="session")
 async def test_terminal_decision_allows_an_empty_candidate(backend):
@@ -160,3 +163,5 @@ async def test_terminal_decision_allows_an_empty_candidate(backend):
 
     assert decision.canonical_source_id is not None
     assert (await backend.unassigned_sources.get(source.source_id)).status == "novel"
+
+    await backend.sources.delete(source_id=decision.canonical_source_id)

tests/test_backend/test_source.py

@@ -7,6 +7,7 @@
 import pytest
 
 from lightcurvedb.client.source import (
+    source_read_all,
     source_read_in_radius,
 )
 from lightcurvedb.models.exceptions import SourceNotFoundException
@@ -28,6 +29,24 @@ async def test_read_all_sources(backend):
     assert len(all_sources) == 64
 
 
+@pytest.mark.asyncio(loop_scope="session")
+async def test_source_read_all_includes_median_flux(backend, setup_test_data):
+    sources = await source_read_all(backend)
+    sources_by_id = {source.source_id: source for source in sources}
+
+    # Every seeded source has flux measurements in exactly 4 of the 6 bands.
+    seeded_source = sources_by_id[setup_test_data[0]]
+
+    assert seeded_source.properties is not None
+    assert len(seeded_source.properties.median_flux) == 4
+
+    for key, value in seeded_source.properties.median_flux.items():
+        module, _, frequency = key.rpartition("_")
+        assert module == "i1"
+        assert int(frequency) in (30, 40, 90, 150, 220, 280)
+        assert isinstance(value, float)
+
+
 @pytest.mark.asyncio(loop_scope="session")
 async def test_source_read_in_radius(backend):
     sources_in_radius = await source_read_in_radius((0, 0), 80.0, backend=backend)

tests/test_backend/test_unassigned.py → tests/test_backend/test_unassigned_backend.py

@@ -39,7 +39,7 @@ async def test_unassigned_source_and_measurements(backend):
     measurement = UnassignedFluxMeasurement(
         measurement_id=uuid4(),
         source_id=source_id,
-        frequency=27,
+        frequency=30,
         module="i1",
         time=now,
         ra=12.5,