timewise 1.0.0a2__py3-none-any.whl → 1.0.0a5__py3-none-any.whl

ampel/timewise/t1/T1HDBSCAN.py

@@ -0,0 +1,222 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# File:                ampel/timewise/t1/T1HDBSCAN.py
+# License:             BSD-3-Clause
+# Author:              Jannis Necker <jannis.necker@gmail.com>
+# Date:                24.09.2025
+# Last Modified Date:  24.09.2025
+# Last Modified By:    Jannis Necker <jannis.necker@gmail.com>
+from typing import Iterable, Sequence
+
+import numpy as np
+from numpy import typing as npt
+from ampel.base.AuxUnitRegister import AuxUnitRegister
+from astropy.coordinates.angle_utilities import angular_separation, position_angle
+from sklearn.cluster import HDBSCAN
+from pymongo import MongoClient
+
+from ampel.content.DataPoint import DataPoint
+from ampel.struct.T1CombineResult import T1CombineResult
+from ampel.types import DataPointId
+from ampel.abstract.AbsT1CombineUnit import AbsT1CombineUnit
+from ampel.model.UnitModel import UnitModel
+
+from ampel.timewise.util.pdutil import datapoints_to_dataframe
+from ampel.timewise.util.AuxDiagnosticPlotter import AuxDiagnosticPlotter
+from timewise.process import keys
+
+
+class T1HDBSCAN(AbsT1CombineUnit):
+    input_mongo_db_name: str
+    original_id_key: str
+    whitelist_region_arcsec: float = 1
+    cluster_distance_arcsec: float = 0.5
+
+    plot: bool = False
+    plotter: UnitModel = UnitModel.model_validate(
+        {
+            "unit": "AuxDiagnosticPlotter",
+            "config": {
+                "plot_properties": {
+                    "file_name": {
+                        "format_str": "%s_hdbscan_selection.svg",
+                        "arg_keys": ["stock"],
+                    }
+                }
+            },
+        }
+    )
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._col = MongoClient()[self.input_mongo_db_name]["input"]
+        self._plotter = AuxUnitRegister.new_unit(
+            model=self.plotter, sub_type=AuxDiagnosticPlotter
+        )
+
+    def combine(
+        self, datapoints: Iterable[DataPoint]
+    ) -> Sequence[DataPointId] | T1CombineResult:
+        columns = ["ra", "dec", "mjd"]
+        if self.plot:
+            for i in range(1, 3):
+                columns += [
+                    f"w{i}{self._plotter.lum_key}",
+                    f"w{i}{keys.ERROR_EXT}{self._plotter.lum_key}",
+                ]
+
+        lightcurve, stock_ids = datapoints_to_dataframe(
+            datapoints, columns, check_tables=["allwise_p3as_mep"]
+        )
+
+        # make sure that the is one stock id that fits all dps
+        # this is a redundant check, the muxer should take care of it
+        unique_stocks = np.unique(np.array(stock_ids).flatten())
+        stock_in_all_dps = [
+            all([s in sids for sids in stock_ids]) for s in unique_stocks
+        ]
+        # make sure only one stock is in all datapoints
+        assert sum(stock_in_all_dps) == 1
+        stock_id = unique_stocks[stock_in_all_dps][0].item()
+        self.logger.debug(f"stock: {stock_id}")
+
+        # query the database that holds the parent sample
+        d = self._col.find_one({self.original_id_key: stock_id})
+        source_ra = d["ra"]
+        source_dec = d["dec"]
+
+        lc_ra_rad = np.deg2rad(lightcurve.ra)
+        lc_dec_rad = np.deg2rad(lightcurve.dec)
+        source_ra_rad = np.deg2rad(source_ra)
+        source_dec_rad = np.deg2rad(source_dec)
+
+        # calculate separation and position angle
+        _angular_separation = angular_separation(
+            source_ra_rad, source_dec_rad, lc_ra_rad, lc_dec_rad
+        )
+        _position_angle = position_angle(
+            source_ra_rad, source_dec_rad, lc_ra_rad, lc_dec_rad
+        )
+
+        # The AllWISE multiframe pipeline detects sources on the deep coadded atlas images and then measures the sources
+        # for all available single-exposure images in all bands simultaneously, while the NEOWISE magnitudes are
+        # obtained by PSF fit to individual exposures directly. Effect: all allwise data points that belong to the same
+        # object have the same position. We take only the closest one and treat it as one datapoint in the clustering.
+        allwise_time_mask = lightcurve.allwise_p3as_mep
+        if any(allwise_time_mask):
+            allwise_sep_min = np.min(_angular_separation[allwise_time_mask])
+            closest_allwise_mask = (
+                _angular_separation == allwise_sep_min
+            ) & allwise_time_mask
+            closest_allwise_mask_first_entry = (
+                ~closest_allwise_mask.duplicated() & closest_allwise_mask
+            )
+
+            # the data we want to use is then the selected AllWISE datapoint and the NEOWISE-R data
+            data_mask = closest_allwise_mask_first_entry | ~allwise_time_mask
+        else:
+            closest_allwise_mask_first_entry = closest_allwise_mask = None
+            data_mask = np.ones_like(_angular_separation, dtype=bool)
+
+        # no matter which cluster they belong to, we want to keep all datapoints within 1 arcsec
+        one_arcsec_mask = _angular_separation < np.radians(
+            self.whitelist_region_arcsec / 3600
+        )
+        selected_indices = set(lightcurve.index[data_mask & one_arcsec_mask])
+
+        # if there are more than one datapoints, we use a clustering algorithm to potentially find a cluster with
+        # its center within 1 arcsec
+        if data_mask.sum() > 1:
+            # instead of the polar coordinates separation and position angle we use cartesian coordinates because the
+            # clustering algorithm works better with them
+            cartesian_full = np.array(
+                [
+                    _angular_separation * np.cos(_position_angle),
+                    _angular_separation * np.sin(_position_angle),
+                ]
+            ).T
+            cartesian = cartesian_full[data_mask]
+
+            # we are now ready to do the clustering
+            cluster_res = HDBSCAN(
+                store_centers="centroid",
+                min_cluster_size=max(min(20, len(cartesian)), 2),
+                allow_single_cluster=True,
+                cluster_selection_epsilon=np.radians(
+                    self.cluster_distance_arcsec / 3600
+                ),
+            ).fit(cartesian)
+            centroids: npt.NDArray[np.float64] = cluster_res.__getattribute__(
+                "centroids_"
+            )
+            labels: npt.NDArray[np.int64] = cluster_res.__getattribute__("labels_")
+
+            # we select the closest cluster within 1 arcsec
+            cluster_separations = np.sqrt(np.sum(centroids**2, axis=1))
+            self.logger.debug(f"Found {len(cluster_separations)} clusters")
+
+            # if there is no cluster or no cluster within 1 arcsec,
+            # only the datapoints within 1 arcsec are selected as we did above
+            if len(cluster_separations) == 0:
+                self.logger.debug(
+                    "No cluster found. Selecting all noise datapoints within 1 arcsec."
+                )
+            elif (_min := min(cluster_separations)) > np.radians(
+                self.whitelist_region_arcsec / 3600
+            ):
+                self.logger.debug(
+                    f"Closest cluster is at {np.degrees(_min) * 3600} arcsec"
+                )
+
+            # if there is a cluster within 1 arcsec, we select all datapoints belonging to that cluster
+            # in addition to the datapoints within 1 arcsec
+            else:
+                closest_label = cluster_separations.argmin()
+                selected_cluster_mask = labels == closest_label
+
+                # now we have to trace back the selected datapoints to the original lightcurve
+                selected_indices |= set(
+                    lightcurve.index[data_mask][selected_cluster_mask]
+                )
+                self.logger.debug(f"Selected {len(selected_indices)} datapoints")
+
+        else:
+            # if there is only one datapoint we give him a label manually
+            labels = np.array([-1])
+
+        # if the closest allwise source is selected, we also select all other detections belonging to that
+        # source in the allwise period
+        if (
+            closest_allwise_mask_first_entry is not None
+            and lightcurve.index[closest_allwise_mask_first_entry][0]
+            in selected_indices
+        ):
+            closest_allwise_mask_not_first = (
+                closest_allwise_mask & ~closest_allwise_mask_first_entry
+            )
+            closest_allwise_indices_not_first = lightcurve.index[
+                closest_allwise_mask_not_first
+            ]
+            self.logger.debug(
+                f"Adding remaining {len(closest_allwise_indices_not_first)} from AllWISE period"
+            )
+            selected_indices |= set(closest_allwise_indices_not_first)
+
+        selected_indices_list = list(selected_indices)
+        res = T1CombineResult(dps=selected_indices_list)
+
+        if self.plot:
+            all_labels = np.array([-1] * len(lightcurve))
+            all_labels[data_mask] = labels
+            svg_rec = self._plotter.make_plot(
+                lightcurve,
+                None,
+                all_labels,
+                source_ra,
+                source_dec,
+                selected_indices_list,
+                highlight_radius=self.whitelist_region_arcsec,
+            )
+            res.add_meta("plot", svg_rec)
+
+        return res

ampel/timewise/t1/TimewiseFilter.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# File:                timewise/ampel/timewise/t1/TimewiseFilter.py
+# License:             BSD-3-Clause
+# Author:              Jannis Necker <jannis.necker@gmail.com>
+# Date:                16.09.2025
+# Last Modified Date:  16.09.2025
+# Last Modified By:    Jannis Necker <jannis.necker@gmail.com>
+import numpy as np
+
+from ampel.abstract.AbsAlertFilter import AbsAlertFilter
+from ampel.protocol.AmpelAlertProtocol import AmpelAlertProtocol
+
+from timewise.util.visits import get_visit_map
+from timewise.process import keys
+
+
+class TimewiseFilter(AbsAlertFilter):
+    det_per_visit: int = 8
+    n_visits = 10
+
+    def process(self, alert: AmpelAlertProtocol) -> None | bool | int:
+        columns = [
+            "ra",
+            "dec",
+            "mjd",
+        ]
+        for i in range(1, 3):
+            for key in [keys.MAG_EXT, keys.FLUX_EXT]:
+                columns.extend([f"w{i}{key}", f"w{i}{keys.ERROR_EXT}{key}"])
+
+        mjd = np.array([dp["mjd"] for dp in alert.datapoints])
+
+        # enough single detections per visit
+        visit_map = get_visit_map(mjd)
+        visits, counts = np.unique(visit_map, return_counts=True)
+        visit_passed = counts >= self.det_per_visit
+        if not all(visit_passed):
+            self.logger.info(None, extra={"min_det_per_visit": min(counts).item()})
+            return None
+
+        # enough visits
+        if not len(visits) >= self.n_visits:
+            self.logger.info(None, extra={"n_visits": len(visits)})
+            return None
+
+        return True

ampel/timewise/t2/T2StackVisits.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# File:                ampel/timewise/t2/T2StackVisits.py
+# License:             BSD-3-Clause
+# Author:              Jannis Necker <jannis.necker@gmail.com>
+# Date:                24.09.2025
+# Last Modified Date:  24.09.2025
+# Last Modified By:    Jannis Necker <jannis.necker@gmail.com>
+from scipy import stats
+
+from ampel.abstract.AbsLightCurveT2Unit import AbsLightCurveT2Unit
+from ampel.struct.UnitResult import UnitResult
+from ampel.types import UBson
+from ampel.view.LightCurve import LightCurve
+
+from ampel.timewise.util.pdutil import datapoints_to_dataframe
+
+from timewise.process import keys
+from timewise.process.stacking import stack_visits
+
+
+SIGMA = float(stats.chi2.cdf(1, 1))
+
+
+class T2StackVisits(AbsLightCurveT2Unit):
+    clean_outliers: bool = True
+
+    # spread quantile of each visit to estimate, defaults to
+    # ~68% which is the equivalent of 1 sigma for a normal distribution
+    outlier_quantile: float = SIGMA
+
+    # threshold above which to exclude outliers
+    outlier_threshold: float = 5
+
+    def process(self, light_curve: LightCurve) -> UBson | UnitResult:
+        columns = [
+            "ra",
+            "dec",
+            "mjd",
+        ]
+        for i in range(1, 3):
+            for key in [keys.MAG_EXT, keys.FLUX_EXT]:
+                columns.extend([f"w{i}{key}", f"w{i}{keys.ERROR_EXT}{key}"])
+
+        photopoints = light_curve.get_photopoints()
+        if photopoints is None:
+            return {}
+        data, _ = datapoints_to_dataframe(photopoints, columns=columns)
+        if len(data) == 0:
+            return {}
+        return stack_visits(
+            data,
+            outlier_threshold=self.outlier_threshold,
+            outlier_quantile=self.outlier_quantile,
+            clean_outliers=self.clean_outliers,
+        ).to_dict(orient="records")

ampel/timewise/util/AuxDiagnosticPlotter.py

@@ -0,0 +1,47 @@
+from typing import Literal, Dict, Any
+
+import pandas as pd
+from numpy import typing as npt
+
+from ampel.plot.create import create_plot_record
+from ampel.base.AmpelBaseModel import AmpelBaseModel
+from ampel.model.PlotProperties import PlotProperties
+from ampel.content.NewSVGRecord import NewSVGRecord
+
+from timewise.plot.diagnostic import DiagnosticPlotter
+
+
+class AuxDiagnosticPlotter(AmpelBaseModel):
+    plot_properties: PlotProperties
+    cutout: Literal["sdss", "panstarrs"] = DiagnosticPlotter.model_fields[
+        "cutout"
+    ].default
+    band_colors: Dict[str, str] = DiagnosticPlotter.model_fields["band_colors"].default
+    lum_key: str = DiagnosticPlotter.model_fields["lum_key"].default
+
+    def __init__(self, /, **data: Any):
+        super().__init__(**data)
+        self._plotter = DiagnosticPlotter.model_validate(
+            {k: self.__getattribute__(k) for k in ["cutout", "band_colors", "lum_key"]}
+        )
+
+    def make_plot(
+        self,
+        raw_lightcurve: pd.DataFrame,
+        stacked_lightcurve: pd.DataFrame | None,
+        labels: npt.ArrayLike,
+        source_ra: float,
+        source_dec: float,
+        selected_indices: list[Any],
+        highlight_radius: float | None = None,
+    ) -> NewSVGRecord:
+        fig, axs = self._plotter.make_plot(
+            raw_lightcurve=raw_lightcurve,
+            stacked_lightcurve=stacked_lightcurve,
+            labels=labels,
+            source_ra=source_ra,
+            source_dec=source_dec,
+            selected_indices=selected_indices,
+            highlight_radius=highlight_radius,
+        )
+        return create_plot_record(fig, self.plot_properties)

ampel/timewise/util/pdutil.py

@@ -0,0 +1,48 @@
+from typing import Iterable, List
+import pandas as pd
+import numpy as np
+from ampel.types import StockId
+from ampel.content.DataPoint import DataPoint
+
+
+def datapoints_to_dataframe(
+    datapoints: Iterable[DataPoint],
+    columns: list[str],
+    check_tables: list[str] | None = None,
+) -> tuple[pd.DataFrame, List[List[StockId]]]:
+    """
+    Convert a list of Ampel DataPoints into a pandas DataFrame.
+
+    Parameters
+    ----------
+    datapoints : list of dict
+        List of DataPoints (each must have a "body" dict).
+    columns : list of str
+        Keys from datapoint["body"] to include as DataFrame columns.
+    check_tables: list of str
+        check if the tables are in the tags
+
+    Returns
+    -------
+    pd.DataFrame
+        DataFrame with one row per DataPoint and one column per requested key.
+    """
+
+    records = []
+    ids = []
+    stock_ids = []
+    for dp in datapoints:
+        body = dp.get("body", {})
+        # Build one row with only requested keys
+        row = {col: body[col] for col in columns}
+        # check if the tables are in tags
+        if check_tables is not None:
+            for table in check_tables:
+                row[table] = table in dp["tag"]
+        # build the index from datapoint ids
+        ids.append(dp["id"])
+        records.append(row)
+        stock_ids.append(np.atleast_1d(dp["stock"]).tolist())
+
+    colnames = columns if check_tables is None else columns + check_tables
+    return pd.DataFrame.from_records(records, columns=colnames, index=ids), stock_ids

timewise/__init__.py

@@ -1 +1 @@
-__version__ = "1.0.0a2"
+__version__ = "1.0.0a5"

timewise/cli.py

@@ -6,11 +6,15 @@ import typer
 
 from .config import TimewiseConfig
 from .plot.diagnostic import make_plot
+from .process.interface import AMPEL_EXISTS
 
 from rich.logging import RichHandler
 
 
-app = typer.Typer(help="Timewsie CLI")
+ampel_missing_msg = (
+    " (download only because AMPEL is not installed)" if not AMPEL_EXISTS else ""
+)
+app = typer.Typer(help="Timewsie CLI" + ampel_missing_msg)
 
 config_path_type = Annotated[
     Path, typer.Argument(help="Pipeline config file (YAML/JSON)")
@@ -53,72 +57,74 @@ def download(
     TimewiseConfig.from_yaml(config_path).download.build_downloader().run()
 
 
-@app.command(help="Prepares the AMPEL job file so AMPEL can be run manually")
-def prepare_ampel(
-    config_path: config_path_type,
-):
-    cfg = TimewiseConfig.from_yaml(config_path)
-    ampel_interface = cfg.build_ampel_interface()
-    p = ampel_interface.prepare(config_path)
-    typer.echo(f"AMPEL job file: {p}")
-
-
-@app.command(help="Processes the lightcurves using AMPEL")
-def process(
-    config_path: config_path_type,
-    ampel_config_path: ampel_config_path_type,
-):
-    cfg = TimewiseConfig.from_yaml(config_path)
-    ampel_interface = cfg.build_ampel_interface()
-    ampel_interface.run(config_path, ampel_config_path)
-
-
-@app.command(help="Write stacked lightcurves to disk")
-def export(
-    config_path: config_path_type,
-    output_directory: Annotated[Path, typer.Argument(help="output directory")],
-    indices: Annotated[
-        list[int] | None,
-        typer.Option(
-            "-i", "--indices", help="Indices to export, defaults to all indices"
-        ),
-    ] = None,
-):
-    TimewiseConfig.from_yaml(config_path).build_ampel_interface().export_many(
-        output_directory, indices
-    )
-
-
-@app.command(help="Run download, process and export")
-def run_chain(
-    config_path: config_path_type,
-    ampel_config_path: ampel_config_path_type,
-    output_directory: Annotated[Path, typer.Argument(help="output directory")],
-    indices: Annotated[
-        list[int] | None,
-        typer.Option(
-            "-i", "--indices", help="Indices to export, defaults to all indices"
-        ),
-    ] = None,
-):
-    download(config_path)
-    process(config_path, ampel_config_path)
-    export(config_path, output_directory, indices)
-
-
-@app.command(help="Make diagnostic plots")
-def plot(
-    config_path: config_path_type,
-    indices: Annotated[
-        List[int],
-        typer.Argument(help="Identifiers of the objects for which to create plots"),
-    ],
-    output_directory: Annotated[Path, typer.Argument(help="Output directory")],
-    cutout: Annotated[
-        Literal["sdss", "panstarrs"],
-        typer.Option("-c", "--cutout", help="Which survey to use for cutouts"),
-    ] = "panstarrs",
-):
-    make_plot(
-        config_path, indices=indices, cutout=cutout, output_directory=output_directory
-    )
+# the following commands will only be added if ampel is installed
+if AMPEL_EXISTS:
+
+    @app.command(help="Prepares the AMPEL job file so AMPEL can be run manually")
+    def prepare_ampel(
+        config_path: config_path_type,
+    ):
+        cfg = TimewiseConfig.from_yaml(config_path)
+        ampel_interface = cfg.build_ampel_interface()
+        p = ampel_interface.prepare(config_path)
+        typer.echo(f"AMPEL job file: {p}")
+
+    @app.command(help="Processes the lightcurves using AMPEL")
+    def process(
+        config_path: config_path_type,
+        ampel_config_path: ampel_config_path_type,
+    ):
+        cfg = TimewiseConfig.from_yaml(config_path)
+        ampel_interface = cfg.build_ampel_interface()
+        ampel_interface.run(config_path, ampel_config_path)
+
+    @app.command(help="Write stacked lightcurves to disk")
+    def export(
+        config_path: config_path_type,
+        output_directory: Annotated[Path, typer.Argument(help="output directory")],
+        indices: Annotated[
+            list[int] | None,
+            typer.Option(
+                "-i", "--indices", help="Indices to export, defaults to all indices"
+            ),
+        ] = None,
+    ):
+        TimewiseConfig.from_yaml(config_path).build_ampel_interface().export_many(
+            output_directory, indices
+        )
+
+    @app.command(help="Run download, process and export")
+    def run_chain(
+        config_path: config_path_type,
+        ampel_config_path: ampel_config_path_type,
+        output_directory: Annotated[Path, typer.Argument(help="output directory")],
+        indices: Annotated[
+            list[int] | None,
+            typer.Option(
+                "-i", "--indices", help="Indices to export, defaults to all indices"
+            ),
+        ] = None,
+    ):
+        download(config_path)
+        process(config_path, ampel_config_path)
+        export(config_path, output_directory, indices)
+
+    @app.command(help="Make diagnostic plots")
+    def plot(
+        config_path: config_path_type,
+        indices: Annotated[
+            List[int],
+            typer.Argument(help="Identifiers of the objects for which to create plots"),
+        ],
+        output_directory: Annotated[Path, typer.Argument(help="Output directory")],
+        cutout: Annotated[
+            Literal["sdss", "panstarrs"],
+            typer.Option("-c", "--cutout", help="Which survey to use for cutouts"),
+        ] = "panstarrs",
+    ):
+        make_plot(
+            config_path,
+            indices=indices,
+            cutout=cutout,
+            output_directory=output_directory,
+        )

timewise/io/download.py

@@ -232,27 +232,25 @@ class Downloader:
                     backend.save_data(task, payload_table)
                     meta["status"] = "COMPLETED"
                     meta["completed_at"] = str(datetime.now())
-                    backend.save_meta(task, meta)
                     backend.mark_done(task)
-                    with self.job_lock:
-                        self.jobs[task] = meta
                 elif status in ("ERROR", "ABORTED"):
                     logger.warning(f"failed {task}: {status}")
-                    meta["status"] = status
-                    with self.job_lock:
-                        self.jobs[task] = meta
-                    backend.save_meta(task, meta)
-                else:
-                    with self.job_lock:
-                        self.jobs[task]["status"] = status
-                        snapshot = self.jobs[task]
-                    backend.save_meta(task, snapshot)
+                elif not status:
+                    logger.warning(
+                        f"No job found under {meta['url']} for {task}! "
+                        f"Probably took too long before downloading results."
+                    )
+
+                meta["status"] = status
+                with self.job_lock:
+                    self.jobs[task] = meta
+                backend.save_meta(task, meta)
 
             if self.all_chunks_submitted:
                 with self.job_lock:
                     all_done = (
                         all(
-                            j.get("status") in ("COMPLETED", "ERROR", "ABORTED")
+                            j.get("status") in ("COMPLETED", "ERROR", "ABORTED", None)
                             for j in self.jobs.values()
                         )
                         if len(self.jobs) > 0

timewise/process/interface.py

@@ -1,6 +1,7 @@
 from pathlib import Path
 import logging
 from typing import Iterable, List, cast
+from importlib.util import find_spec
 
 import numpy as np
 from numpy import typing as npt
@@ -8,12 +9,19 @@ import pandas as pd
 from pymongo import MongoClient, ASCENDING
 from pymongo.collection import Collection
 from pymongo.database import Database
-from ampel.cli.JobCommand import JobCommand
-from ampel.types import DataPointId, StockId
+
+if find_spec("ampel.core"):
+    AMPEL_EXISTS = True
+    from ampel.cli.JobCommand import JobCommand
+else:
+    AMPEL_EXISTS = False
 
 
 logger = logging.getLogger(__name__)
 
+# copy from ampel.types
+StockId = int | bytes | str
+
 
 class AmpelInterface:
     def __init__(
@@ -71,6 +79,10 @@ class AmpelInterface:
         return self.make_ampel_job_file(cfg_path)
 
     def run(self, timewise_cfg_path: Path, ampel_config_path: Path):
+        if not AMPEL_EXISTS:
+            raise ModuleNotFoundError(
+                "You are trying to run ampel but it is not installed!"
+            )
         ampel_job_path = self.prepare(timewise_cfg_path)
         cmd = JobCommand()
         parser = cmd.get_parser()
@@ -123,7 +135,7 @@ class AmpelInterface:
             index.append(ic["id"])
         return pd.DataFrame(records, index=index)
 
-    def extract_selected_datapoint_ids(self, stock_id: StockId) -> List[DataPointId]:
+    def extract_selected_datapoint_ids(self, stock_id: StockId) -> List[int]:
         d = self.t1.find_one({"stock": stock_id})
         if d is None:
             return []