antares-devkit 0.7.1__tar.gz → 0.8.1__tar.gz

{antares_devkit-0.7.1 → antares_devkit-0.8.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: antares-devkit
-Version: 0.7.1
+Version: 0.8.1
 Summary: Provides tools and utility methods to write and debug ANTARES filters with Python.
 Keywords: antares,devkit,filter development
 Author-Email: NSF NOIRLab ANTARES Team <antares@noirlab.edu>
@@ -51,6 +51,7 @@ Requires-Dist: scipy>=1.13.1; extra == "filter-dependencies"
 Requires-Dist: light-curve==0.7.2; extra == "filter-dependencies"
 Requires-Dist: ssi-forest>=0.1.1; extra == "filter-dependencies"
 Requires-Dist: slack-sdk>=3.39.0; extra == "filter-dependencies"
+Requires-Dist: xgboost==2.1.4; extra == "filter-dependencies"
 Provides-Extra: jupyter
 Requires-Dist: matplotlib>=3.9.4; extra == "jupyter"
 Description-Content-Type: text/markdown

antares_devkit-0.8.1/antares_devkit/filters/lantern/__init__.py

@@ -0,0 +1,223 @@
+from antares_devkit.models import BaseFilter
+from antares_devkit.filters.lantern import data
+
+class Lantern(BaseFilter):
+    
+    SLACK_CHANNEL = '#filter-lantern'
+    TRIGGERING_SURVEY = 'lsst'
+    # REQUIRED_LOCUS_PROPERTIES = []
+    REQUIRED_ALERT_PROPERTIES = [
+        'lsst_diaSource_band',
+        'lsst_diaSource_snr',
+        'lsst_diaSource_scienceFlux',
+        'lsst_diaSource_psfFlux', 
+        'lsst_diaSource_apFlux', 
+        'lsst_diaSource_psfChi2',
+        'lsst_diaSource_dipoleFitAttempted',
+    ]
+    # REQUIRED_TAGS = []
+    OUTPUT_LOCUS_PROPERTIES = [
+        {
+            'name': 'lantern_xgboost_t2.0.7_c0.95_num_tagged_alerts',
+            'type': 'int',
+            'description': 'number of alerts associated with this locus which have passed this version of the lantern First Cut filter',
+        },
+        {
+            'name': 'lantern_xgboost_t2.0.7_c0.95_max_score',
+            'type': 'float',
+            'description': 'maximum lantern score of all alerts associated with this locus which have passed this version of the lantern First Cut filter',
+        },
+    ]
+
+    OUTPUT_LOCUS_TAGS = [
+        {
+            'name': 'lantern_xgboost_t2.0.7_c0.95',
+            'description': 'locus has one or more Alerts identified by this version of lantern as a First Cut LAGN target, based on threshold for Y2 0.95 completeness.'
+        },
+    ]
+
+    REQUIRES_FILES = [
+        'lantern_xgboost_t2.0.7.pkl',
+    ]
+
+    def setup(self):
+        """
+        ANTARES will call this function once at the beginning of each night when filters
+        are loaded.
+        """
+        import pickle
+        from importlib.resources import open_binary
+        import xgboost
+
+        with open_binary(data, 'lantern_xgboost_t2.0.7.pkl') as f:
+            self.lantern_model = pickle.load(f)
+
+        self.threshold = 0.9691  #currently set for completeness 0.95
+
+    def _run(self, locus):
+        """
+        ANTARES will call this function in real-time for each incoming Alert.
+        """
+
+        import xgboost
+        import numpy as np
+        import pandas as pd
+
+        def get_prop(name):
+            if locus.alert.properties.get(name) is not None:
+                val = locus.alert.properties[name]
+            else:
+                val = np.nan
+            return val
+
+        # Extract direct alert properties
+        band = get_prop('lsst_diaSource_band')
+        snr = get_prop('lsst_diaSource_snr')
+        scienceFlux = get_prop('lsst_diaSource_scienceFlux')
+        psfFlux = get_prop('lsst_diaSource_psfFlux')
+        apFlux = get_prop('lsst_diaSource_apFlux')
+        extendedness = get_prop('lsst_diaSource_extendedness')
+        psfChi2 = get_prop('lsst_diaSource_psfChi2')
+        dipoleFitAttempted = get_prop('lsst_diaSource_dipoleFitAttempted')
+        dipoleChi2 = get_prop('lsst_diaSource_dipoleChi2')
+        dipoleLength = get_prop('lsst_diaSource_dipoleLength')
+
+        # Calculate engineered features
+        engineered = self._add_engineered_features(locus)
+
+        # Create DataFrame with the same structure as model training
+        feature_dict = {
+            'band': [band],
+            'psf_fwhm': [engineered['psf_fwhm']],
+            'snr': [snr],
+            'template_flux': [engineered['template_flux']],
+            'scienceFlux': [scienceFlux],
+            'psfFlux': [psfFlux],
+            'apFlux': [apFlux],
+            'temp_sci_flux_ratio': [engineered['temp_sci_flux_ratio']],
+            'moment_ext': [engineered['moment_ext']],
+            'ellip_ext': [engineered['ellip_ext']],
+            'flux_ext': [engineered['flux_ext']],
+            'extendedness': [extendedness],
+            'psfChi2': [psfChi2],
+            'dipoleFitAttempted': [dipoleFitAttempted],
+            'dipoleChi2': [dipoleChi2],
+            'dipoleLength': [dipoleLength],
+            'x_y_err': [engineered['x_y_err']],
+        }
+
+        # Preprocessing for bands and flags
+        def normalize_bands_and_flags(df):
+            """Normalizes band strings and converts flag columns. Returns a copy."""
+            df = df.copy()
+            if 'band' in df.columns:
+                # Handle both FilterLabel(band="g", ...) and plain 'g' formats
+                extracted = df['band'].astype(str).str.extract(r"band=['\"]([ugrizy])['\"]", expand=False)
+                # Fall back to bare letter for rows already in plain format
+                plain = df['band'].astype(str).str.extract(r'^([ugrizy])$', expand=False)
+                df['band'] = extracted.fillna(plain)
+                # XGBoost requires categorical columns to be the 'category' dtype
+                df['band'] = df['band'].astype('category')
+            flag_cols = [c for c in df.columns if 'flag' in c.lower() or c in ['isDipole']]
+            for col in flag_cols:
+                df[col] = df[col].astype(bool)
+            return df
+        
+        df = normalize_bands_and_flags(pd.DataFrame(feature_dict))
+
+        # Get prediction probability
+        prediction_proba = self.lantern_model.predict_proba(df)[0, 1]
+
+        if prediction_proba >= self.threshold:
+            locus.tag('lantern_xgboost_t2.0.7_c0.95')
+            
+            #track number of tagged alerts
+            if 'lantern_xgboost_t2.0.7_c0.95_num_tagged_alerts' in locus.properties:
+                n = locus.properties['lantern_xgboost_t2.0.7_c0.95_num_tagged_alerts']
+                n += 1
+            else:
+                n = 1
+
+            locus.set_property('lantern_xgboost_t2.0.7_c0.95_num_tagged_alerts', n)
+
+            #track maximum score
+            if 'lantern_xgboost_t2.0.7_c0.95_max_score' in locus.properties:
+                old_score = locus.properties['lantern_xgboost_t2.0.7_c0.95_max_score']
+                #only update if new score is larger than old score
+                if prediction_proba > old_score: 
+                    locus.set_property('lantern_xgboost_t2.0.7_c0.95_max_score', float(prediction_proba))
+            else:
+                locus.set_property('lantern_xgboost_t2.0.7_c0.95_max_score', float(prediction_proba))
+        
+    def _add_engineered_features(self, locus):
+        """
+        Calculate engineered features from a single alert.
+    
+        Returns
+        -------
+        dict
+            Dictionary of engineered features
+        """
+        import numpy as np
+        
+        def get_prop(name):
+            if locus.alert.properties.get(name) is not None:
+                val = locus.alert.properties[name]
+            else:
+                val = np.nan
+            return val
+
+        # Calculate moment extension
+        psf_trace = get_prop('lsst_diaSource_ixxPSF') + get_prop('lsst_diaSource_iyyPSF')
+        src_trace = get_prop('lsst_diaSource_ixx') + get_prop('lsst_diaSource_iyy')
+        moment_ext = src_trace / psf_trace if psf_trace != 0 else np.nan
+
+         # Calculate ellipticity extension
+        ixx = get_prop('lsst_diaSource_ixx')
+        iyy = get_prop('lsst_diaSource_iyy')
+        ixy = get_prop('lsst_diaSource_ixy')
+        ixxPSF = get_prop('lsst_diaSource_ixxPSF')
+        iyyPSF = get_prop('lsst_diaSource_iyyPSF')
+        ixyPSF = get_prop('lsst_diaSource_ixyPSF')
+
+        if src_trace != 0:
+            src_ellip = np.sqrt((ixx - iyy)**2 + 4 * ixy**2) / src_trace
+        else:
+            src_ellip = np.nan
+        
+        if psf_trace != 0:
+            psf_ellip = np.sqrt((ixxPSF - iyyPSF)**2 + 4 * ixyPSF**2) / psf_trace
+        else:
+            psf_ellip = np.nan
+        
+        ellip_ext = src_ellip - psf_ellip
+
+        # Calculate flux extension
+        psfFlux = get_prop('lsst_diaSource_psfFlux')
+        apFlux = get_prop('lsst_diaSource_apFlux')
+        flux_ext = apFlux / psfFlux if psfFlux != 0 else np.nan
+        
+        # Calculate template flux ratio
+        scienceFlux = get_prop('lsst_diaSource_scienceFlux')
+        template_flux = scienceFlux - psfFlux
+        temp_sci_flux_ratio = template_flux / scienceFlux if scienceFlux != 0 else np.nan
+        
+        # Calculate PSF FWHM
+        psf_det = ixxPSF * iyyPSF - ixyPSF**2
+        psf_fwhm = np.abs(psf_det)**(1/4) * 2.35482 * 0.2 if psf_det != 0 else np.nan
+        
+        # Calculate x-y error
+        xErr = get_prop('lsst_diaSource_xErr')
+        yErr = get_prop('lsst_diaSource_yErr')
+        x_y_err = np.sqrt(xErr**2 + yErr**2)
+        
+        # Return as dictionary
+        return {
+            'moment_ext': moment_ext,
+            'ellip_ext': ellip_ext,
+            'flux_ext': flux_ext,
+            'template_flux': template_flux,
+            'temp_sci_flux_ratio': temp_sci_flux_ratio,
+            'psf_fwhm': psf_fwhm,
+            'x_y_err': x_y_err
+        }