AT for TEMPS

pyproject.toml

@@ -31,7 +31,7 @@ dependencies = [
     "jupytext",
     "mkdocs",
     "typing",
-    "dataclasses",
+    "dataclasses"
 ]
 
 classifiers = [

temps/constants.py

@@ -0,0 +1,4 @@
+from pathlib import Path
+
+PROJ_ROOT = Path(__file__).resolve().parents[2]
+MODULES_DIR = PROJ_ROOT / 'data/models'

temps/plots.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pandas as pd
 import matplotlib.pyplot as plt
-from temps.utils import nmad
+from temps.utils import nmad, sigma68
 from scipy import stats
 from typing import List, Optional, Dict
 
@@ -277,7 +277,7 @@ def plot_crps(
     crps_list_1: List[float],
     crps_list_2: Optional[List[float]] = None,
     crps_list_3: Optional[List[float]] = None,
-    label: Optional[List[str]] = None,
+    labels: Optional[List[str]] = None,
     sample: str = "specz",
     save: bool = True,
 ) -> None:

tests/test_temps.py

@@ -0,0 +1,38 @@
+import numpy as np
+from loguru import logger
+import torch
+
+from temps.temps_arch import EncoderPhotometry, MeasureZ
+from temps.temps import TempsModule
+from temps.constants import MODULES_DIR
+
+
+def test():
+    nn_features = EncoderPhotometry()
+    nn_features.load_state_dict(torch.load(MODULES_DIR / f'modelF_DA.pt',map_location=torch.device('cpu')))
+    nn_z = MeasureZ(num_gauss=6)
+    nn_z.load_state_dict(torch.load(MODULES_DIR / f'modelZ_DA.pt',map_location=torch.device('cpu')))
+
+    temps_module = TempsModule(nn_features, nn_z)
+
+    col = np.array([0.54804805, 1.81142339, 0.63354394, 0.7356338 , 1.3578122 ,
+       0.90108565])
+    ztrue = 0.444
+
+    z, pz, odds = temps_module.get_pz(input_data=torch.Tensor(col), 
+                                            return_pz=True,
+                                            return_flag=True)
+
+    zdiff = (z - ztrue).abs().mean()
+
+    logger.info(f'zdiff: {zdiff}')
+    logger.info("test passed")
+
+    assert zdiff < 0.01
+
+test()
+
+
+
+
+# %%