Add checks for ARAD models; clean up PyRIIDModel saving extensions.

Author: alanjvano

examples/modeling/arad.py

@@ -48,35 +48,46 @@ def show_summaries(model):
 
 # Generate some training data
 fg_seeds_ss, bg_seeds_ss = get_dummy_seeds().split_fg_and_bg()
-mixed_bg_seed_ss = SeedMixer(bg_seeds_ss, mixture_size=3).generate(1)
+mixed_bg_seed_ss = SeedMixer(bg_seeds_ss, mixture_size=3).generate(300)
 
 static_synth = StaticSynthesizer(
-    samples_per_seed=250,
-    snr_function="log10",
+    samples_per_seed=10,
+    snr_function_args=(0, 0),
     return_fg=False,
     return_gross=True,
 )
-_, train_ss = static_synth.generate(fg_seeds_ss, mixed_bg_seed_ss)
+_, train_ss = static_synth.generate(fg_seeds_ss[0], mixed_bg_seed_ss)
+train_ss.downsample_spectra(target_bins=128)
 train_ss.normalize()
 
+# Train the models
 print("training ARADv1...")
 arad_v1.fit(train_ss, epochs=50, verbose=True)
 print("training ARADv2...")
 arad_v2.fit(train_ss, epochs=50, verbose=True)
 
 # Generate some test data
-static_synth.samples_per_seed = 50
-_, test_ss = static_synth.generate(fg_seeds_ss, mixed_bg_seed_ss)
+static_synth.samples_per_seed = 3
+_, test_ss = static_synth.generate(fg_seeds_ss[0], mixed_bg_seed_ss)
+test_ss.downsample_spectra(target_bins=128)
 test_ss.normalize()
 
 # Predict
-arad_v1_reconstructions = arad_v1.predict(test_ss, verbose=True)
+arad_v1_reconstructions = arad_v1.predict(
+    test_ss,
+    verbose=True,
+    ood_threshold=2.349
+)
 recon_errors = test_ss.info["recon_error"].values
 ood_decisions = test_ss.info["ood"].values
 print((f"ARADv1: mean reconstruction error = {np.mean(recon_errors):.3f} (KLD)\n"
        f"                         OOD rate = {np.mean(ood_decisions):.2f}"))
 
-arad_v2_reconstructions = arad_v2.predict(test_ss, verbose=True)
+arad_v2_reconstructions = arad_v2.predict(
+    test_ss,
+    verbose=True,
+    ood_threshold=0.15678
+)
 recon_errors = test_ss.info["recon_error"].values
 ood_decisions = test_ss.info["ood"].values
 print((f"ARADv2: mean reconstruction error = {np.mean(recon_errors):.3f} (JSD)\n"

riid/models/__init__.py

@@ -31,6 +31,7 @@ class PyRIIDModel:
     """Base class for TensorFlow models."""
 
     CUSTOM_OBJECTS = {"multi_f1": multi_f1, "single_f1": single_f1}
+    SUPPORTED_SAVE_EXTS = {"H5": ".h5", "ONNX": ".onnx"}
 
     def __init__(self, *args, **kwargs):
         self._info = {}
@@ -120,25 +121,21 @@ def save(self, file_path: str):
         if os.path.exists(file_path):
             raise ValueError("Path already exists.")
 
-        SUPPORTED_EXTS = {
-            "H5": ".h5",
-            "ONNX": ".onnx"
-        }
         root, ext = os.path.splitext(file_path)
-        if ext.lower() not in SUPPORTED_EXTS.values():
+        if ext.lower() not in self.SUPPORTED_SAVE_EXTS.values():
             raise NameError("Model must be an .onnx or .h5 file.")
 
         warnings.filterwarnings("ignore")
 
-        if ext.lower() == SUPPORTED_EXTS["H5"]:
+        if ext.lower() == self.SUPPORTED_SAVE_EXTS["H5"]:
             self.model.save(file_path, save_format="h5")
             pd.DataFrame(
                 [[v] for v in self.info.values()],
                 self.info.keys()
             ).to_hdf(file_path, "_info")
 
         else:
-            model_path = root + SUPPORTED_EXTS["ONNX"]
+            model_path = root + self.SUPPORTED_SAVE_EXTS["ONNX"]
             model_info_path = root + "_info.json"
 
             model_info_df = pd.DataFrame(
@@ -162,25 +159,22 @@ def load(self, file_path: str):
             file_path: file path from which to load the model, must be either an
             .h5 or .onnx file
         """
-        SUPPORTED_EXTS = {
-            "H5": ".h5",
-            "ONNX": ".onnx"
-        }
+
         root, ext = os.path.splitext(file_path)
-        if ext.lower() not in SUPPORTED_EXTS.values():
+        if ext.lower() not in self.SUPPORTED_SAVE_EXTS.values():
             raise NameError("Model must be an .onnx or .h5 file.")
 
         warnings.filterwarnings("ignore", category=DeprecationWarning)
 
-        if ext.lower() == SUPPORTED_EXTS["H5"]:
+        if ext.lower() == self.SUPPORTED_SAVE_EXTS["H5"]:
             self.model = tf.keras.models.load_model(
                 file_path,
                 custom_objects=self.CUSTOM_OBJECTS
             )
             self._info = pd.read_hdf(file_path, "_info")[0].to_dict()
 
         else:
-            model_path = root + SUPPORTED_EXTS["ONNX"]
+            model_path = root + self.SUPPORTED_SAVE_EXTS["ONNX"]
             model_info_path = root + "_info.json"
 
             with open(model_info_path) as fin:

riid/models/neural_nets/arad.py

@@ -14,7 +14,7 @@
 from scipy.spatial.distance import jensenshannon
 from scipy.stats import entropy
 
-from riid.data.sampleset import SampleSet
+from riid.data.sampleset import SampleSet, SpectraState
 from riid.losses import jensen_shannon_distance, mish
 from riid.models import PyRIIDModel
 
@@ -247,6 +247,21 @@ def __init__(self, model: Model = ARADv2TF()):
 
         self.model = model
 
+    def _check_spectra(self, ss):
+        """Checks if SampleSet spectra are compatible with ARAD models."""
+        if ss.n_samples <= 0:
+            raise ValueError("No spectr[a|um] provided!")
+        if not ss.all_spectra_sum_to_one():
+            raise ValueError("All spectra must sum to one.")
+        if not ss.spectra_state == SpectraState.L1Normalized:
+            raise ValueError(
+                f"SpectraState must be L1Normalzied, provided SpectraState is {ss.spectra_state}."
+            )
+        if not ss.n_channels == 128:
+            raise ValueError(
+                f"Spectra must have 128 channels, provided spectra have {ss.n_channels} channels."
+            )
+
     def fit(self, ss: SampleSet, epochs: int = 300, validation_split=0.2,
             es_verbose: int = 0, verbose: bool = False):
         """Fit a model to the given `SampleSet`.
@@ -261,13 +276,9 @@ def fit(self, ss: SampleSet, epochs: int = 300, validation_split=0.2,
         Returns:
             reconstructed_spectra: output of ARAD model
         """
-        if ss.n_samples <= 0:
-            raise ValueError("No spectr[a|um] provided!")
+        self._check_spectra(ss)
 
-        norm_ss = ss[:]
-        norm_ss.downsample_spectra(target_bins=128)
-        norm_ss.normalize()
-        x = norm_ss.get_samples().astype(float)
+        x = ss.get_samples().astype(float)
 
         is_v1 = isinstance(self.model, ARADv1TF)
         is_v2 = isinstance(self.model, ARADv2TF)
@@ -343,12 +354,11 @@ def predict(self, ss: SampleSet, ood_threshold: float = 0.5,
         Returns:
             reconstructed_spectra: output of ARAD model
         """
-        norm_ss = ss[:]
-        norm_ss.downsample_spectra(target_bins=128)
-        norm_ss.normalize()
-        spectra = norm_ss.get_samples().astype(float)
+        self._check_spectra(ss)
+
+        x = ss.get_samples().astype(float)
 
-        reconstructed_spectra = self.get_predictions(spectra, verbose=verbose)
+        reconstructed_spectra = self.get_predictions(x, verbose=verbose)
 
         is_v1 = isinstance(self.model, ARADv1TF)
         is_v2 = isinstance(self.model, ARADv2TF)
@@ -357,7 +367,7 @@ def predict(self, ss: SampleSet, ood_threshold: float = 0.5,
         elif is_v2:
             reconstruction_metric = jensenshannon
 
-        reconstruction_errors = reconstruction_metric(spectra, reconstructed_spectra, axis=1)
+        reconstruction_errors = reconstruction_metric(x, reconstructed_spectra, axis=1)
         ood_decisions = reconstruction_errors > ood_threshold
         ss.info["recon_error"] = reconstruction_errors
         ss.info["ood"] = ood_decisions