Feat: new mongo models

db_plugins/__init__.py

@@ -1 +1 @@
-__version__ = "4.0.2"
+__version__ = "4.1.0"

db_plugins/db/mongo/models.py

@@ -55,28 +55,115 @@ class Object(BaseModel):
             "coordinates": [kwargs["meanra"] - 180, kwargs["meandec"]],
         }
     )
-    magstats = SpecialField(lambda **kwargs: kwargs.get("magstats", []))
-    features = SpecialField(lambda **kwargs: kwargs.get("features", []))
-    probabilities = SpecialField(lambda **kwargs: kwargs.get("probabilities", []))
-    xmatch = SpecialField(lambda **kwargs: kwargs.get("xmatch", []))
 
     __table_args__ = [
         IndexModel([("oid", ASCENDING)]),
         IndexModel([("tid", ASCENDING)]),
         IndexModel([("lastmjd", DESCENDING)]),
         IndexModel([("firstmjd", DESCENDING)]),
         IndexModel([("loc", GEOSPHERE)]),
+    ]
+    __tablename__ = "object"
+
+
+class Probability(BaseModel):
+    aid = Field()
+    classifier_name = Field()
+    classifier_version = Field()
+    class_name = Field()
+    probability = Field()
+    ranking = Field()
+
+    __table_args__ = [
+        IndexModel([("aid", ASCENDING)]),
         IndexModel(
             [
-                ("probabilities.ranking", DESCENDING),
-                ("probabilities.classifier_name", ASCENDING),
-                ("probabilities.classifier_version", DESCENDING),
-                ("probabilities.class_name", DESCENDING),
-                ("probabilities.probability", DESCENDING),
+                ("classifier_name", ASCENDING),
+                ("class_name", ASCENDING),
+                ("classifier_version", ASCENDING),
+                ("aid", ASCENDING),
             ],
+            unique=True,
+        ),
+        IndexModel(
+            [("probability", DESCENDING)], partialFilterExpression={"ranking": 1}
+        ),
+    ]
+    __tablename__ = "probability"
+
+
+class MagStats(BaseModel):
+    __tablename__ = "magstat"
+
+    aid = Field()
+    fid = Field()
+    stellar = Field()
+    corrected = Field()
+    ndet = Field()
+    ndubious = Field()
+    dmdt_first = Field()
+    dm_first = Field()
+    sigmadm_first = Field()
+    dt_first = Field()
+    magmean = Field()
+    magmedian = Field()
+    magmax = Field()
+    magmin = Field()
+    magsigma = Field()
+    maglast = Field()
+    magfirst = Field()
+    magmean_corr = Field()
+    magmedian_corr = Field()
+    magmax_corr = Field()
+    magmin_corr = Field()
+    magsigma_corr = Field()
+    maglast_corr = Field()
+    magfirst_corr = Field()
+    firstmjd = Field()
+    lastmjd = Field()
+
+    __table_args__ = [
+        IndexModel([("aid", ASCENDING)]),
+        IndexModel([("fid", ASCENDING)]),
+        IndexModel([("aid", ASCENDING), ("fid", ASCENDING)], unique=True),
+    ]
+
+
+class Xmatch(BaseModel):
+    __tablename__ = "xmatch"
+
+    aid = Field()
+    catid = Field()
+    catoid = Field()  # oid_catalogue in old version
+    dist = Field()
+
+    __table_args__ = [
+        IndexModel([("aid", ASCENDING), ("catid", ASCENDING)], unique=True),
+        IndexModel([("catid", ASCENDING)]),
+    ]
+
+
+class Feature(BaseModel):
+    __tablename__ = "feature"
+
+    aid = Field()
+    name = Field()
+    value = Field()
+    fid = Field()
+    version = Field()
+
+    __table_args__ = [
+        IndexModel([("aid", ASCENDING)]),
+        IndexModel(
+            [
+                ("name", ASCENDING),
+                ("fid", ASCENDING),
+                ("version", ASCENDING),
+                ("aid", ASCENDING),
+            ],
+            unique=True,
         ),
     ]
-    __tablename__ = "object"
 
 
 class Detection(BaseModelWithExtraFields):
@@ -100,28 +187,23 @@ def create_extra_fields(cls, **kwargs):
     e_mag = Field()  # sigmapsf in ZTF alerts
     mag_corr = Field()  # magpsf_corr in ZTF alerts
     e_mag_corr = Field()  # sigmapsf_corr in ZTF alerts
+    e_mag_corr_ext = Field()  # sigmapsf_corr_ext in ZTF alerts
     isdiffpos = Field()
     corrected = Field()
     dubious = Field()
-    parent_candidate = Field()
+    stellar = Field()
     has_stamp = Field()
-    step_id_corr = Field()
 
     __table_args__ = [
-        IndexModel([("aid", ASCENDING), ("oid", ASCENDING)]),
+        IndexModel([("aid", ASCENDING)]),
+        IndexModel([("oid", ASCENDING)]),
         IndexModel([("tid", ASCENDING)]),
+        IndexModel([("mjd", ASCENDING)]),
     ]
     __tablename__ = "detection"
 
 
 class NonDetection(BaseModelWithExtraFields):
-    @classmethod
-    def create_extra_fields(cls, **kwargs):
-        kwargs = super().create_extra_fields(**kwargs)
-        kwargs.pop("candid", None)  # Prevents candid being duplicated in extra_fields
-        return kwargs
-
-    _id = SpecialField(lambda **kwargs: kwargs.get("candid") or kwargs["_id"])
     aid = Field()
     tid = Field()
     oid = Field()
@@ -130,8 +212,11 @@ def create_extra_fields(cls, **kwargs):
     diffmaglim = Field()
 
     __table_args__ = [
-        IndexModel([("aid", ASCENDING), ("oid", ASCENDING)]),
+        IndexModel([("aid", ASCENDING)]),
         IndexModel([("tid", ASCENDING)]),
+        IndexModel(
+            [("oid", ASCENDING), ("fid", ASCENDING), ("mjd", ASCENDING)], unique=True
+        ),
     ]
     __tablename__ = "non_detection"
 
@@ -143,46 +228,8 @@ class Taxonomy(BaseModel):
 
     __table_args__ = [
         IndexModel(
-            [("classifier_name", ASCENDING), ("classifier_version", DESCENDING)]
+            [("classifier_name", ASCENDING), ("classifier_version", DESCENDING)],
+            unique=True,
         ),
     ]
     __tablename__ = "taxonomy"
-
-
-class Step(BaseModel):
-    step_id = Field()
-    name = Field()
-    version = Field()
-    comments = Field()
-    date = Field()
-
-    __table_args__ = [
-        IndexModel([("step_id", ASCENDING)]),
-    ]
-    __tablename__ = "step"
-
-
-class FeatureVersion(BaseModel):
-    version = Field()
-    step_id_feature = Field()
-    step_id_preprocess = Field()
-
-    __table_args__ = [
-        IndexModel([("version", ASCENDING)]),
-    ]
-    __tablename__ = "feature_version"
-
-
-class Pipeline(BaseModel):
-    pipeline_id = Field()
-    step_id_corr = Field()
-    step_id_feat = Field()
-    step_id_clf = Field()
-    step_id_out = Field()
-    step_id_stamp = Field()
-    date = Field()
-
-    __table_args__ = [
-        IndexModel([("pipeline_id", ASCENDING)]),
-    ]
-    __tablename__ = "pipeline"