Implement ARADBranch model class; update ARADv1TF and ARADv2TF …

…parameters based on original papers; add a regression and classification example with `ARADBranch`.

examples/modeling/arad_branches.py

@@ -0,0 +1,107 @@
+# Copyright 2021 National Technology & Engineering Solutions of Sandia, LLC (NTESS).
+# Under the terms of Contract DE-NA0003525 with NTESS,
+# the U.S. Government retains certain rights in this software.
+"""This example demonstrates how to train a regressor or classifier branch
+from an ARAD latent space.
+"""
+import os
+
+import numpy as np
+from sklearn.metrics import f1_score
+
+from riid.data.synthetic import get_dummy_seeds
+from riid.data.synthetic.seed import SeedMixer
+from riid.data.synthetic.static import StaticSynthesizer
+from riid.models.neural_nets.arad import ARAD, ARADBranch
+
+# Config
+rng = np.random.default_rng(42)
+OOD_QUANTILE = 0.99
+ARAD_MODEL_PATH = "./arad_model.onnx"
+ARAD_REGRESSOR_PATH = "./arad_reg_model.onnx"
+ARAD_CLASSIFIER_PATH = "./arad_cla_model.onnx"
+VERBOSE = True
+# Some of the following parameters are set low because this example runs on GitHub Actions and
+#   we don't want it taking a bunch of time.
+# When running this locally, change the values per their corresponding comment, otherwise
+#   the results likely will not be meaningful.
+EPOCHS = 5  # Change this to 20+
+N_MIXTURES = 50  # Changes this to 1000+
+TRAIN_SAMPLES_PER_SEED = 5  # Change this to 20+
+TEST_SAMPLES_PER_SEED = 5
+
+# Generate training data
+fg_seeds_ss, bg_seeds_ss = get_dummy_seeds(n_channels=128, rng=rng).split_fg_and_bg()
+mixed_bg_seed_ss = SeedMixer(bg_seeds_ss, mixture_size=3, rng=rng).generate(N_MIXTURES)
+static_synth = StaticSynthesizer(
+    samples_per_seed=TRAIN_SAMPLES_PER_SEED,
+    snr_function_args=(0, 0),
+    return_fg=False,
+    return_gross=True,
+    rng=rng,
+)
+_, gross_train_ss = static_synth.generate(fg_seeds_ss[0], mixed_bg_seed_ss)
+gross_train_ss.normalize()
+
+# Generate test data
+static_synth.samples_per_seed = TEST_SAMPLES_PER_SEED
+_, test_ss = static_synth.generate(fg_seeds_ss[0], mixed_bg_seed_ss)
+test_ss.normalize()
+
+# Train ARAD model
+arad_v2 = ARAD(arad_version="v2")
+arad_v2.fit(gross_train_ss, epochs=EPOCHS, verbose=VERBOSE)
+arad_v2.predict(gross_train_ss)
+v2_ood_threshold = np.quantile(gross_train_ss.info.recon_error, OOD_QUANTILE)
+
+# Save ARAD model
+arad_v2.save(ARAD_MODEL_PATH)
+
+# Train branched model to predict live-time and real-time
+arad_regressor = ARADBranch(
+    ARAD_MODEL_PATH
+)
+_ = arad_regressor.fit(
+    gross_train_ss,
+    target_info_columns=["live_time", "real_time"],
+    epochs=10,
+    verbose=True,
+    batch_size=5
+)
+
+# Save, load, and predict with regressor
+arad_regressor.save(ARAD_REGRESSOR_PATH)
+arad_regressor = ARADBranch()
+arad_regressor.load(ARAD_REGRESSOR_PATH)
+preds = arad_regressor.predict(test_ss)
+for idx, target_name in enumerate(arad_regressor.info["target_info_columns"]):
+    print(f"{target_name}: {preds[:5, idx]}...")
+
+# Train branched model to classify isotopes
+arad_classifier = ARADBranch(
+    ARAD_MODEL_PATH,
+    loss="categorical_crossentropy",
+    metrics=("accuracy", "categorical_crossentropy"),
+    final_activation="softmax"
+)
+_ = arad_classifier.fit(
+    gross_train_ss,
+    target_level="Isotope",
+    epochs=10,
+    verbose=True,
+    batch_size=5
+)
+
+# Save, load, and predict with regressor
+arad_classifier.save(ARAD_CLASSIFIER_PATH)
+arad_classifier = ARADBranch()
+arad_classifier.load(ARAD_CLASSIFIER_PATH)
+arad_classifier.predict(test_ss)
+score = f1_score(test_ss.get_labels(), test_ss.get_predictions(), average="micro")
+print("F1 Score: {:.3f}".format(score))
+
+# Clean up
+for path in [ARAD_MODEL_PATH, ARAD_REGRESSOR_PATH, ARAD_CLASSIFIER_PATH]:
+    info_path = os.path.splitext(path)[0] + "_info.json"
+    os.remove(path)
+    os.remove(info_path)

riid/models/neural_nets/__init__.py

@@ -253,7 +253,7 @@ def predict(self, ss: SampleSet, bg_ss: SampleSet = None, verbose=False):
         ss.prediction_probas = pd.DataFrame(
             data=results,
             columns=pd.MultiIndex.from_tuples(
-               self.model_outputs, names=col_level_subset
+                self.model_outputs, names=col_level_subset
             )
         )
 
@@ -539,7 +539,7 @@ def __init__(self, hidden_layers: tuple = (256,), sup_loss="sparsemax", unsup_lo
             optimizer_kwargs: kwargs for optimizer
             learning_rate: learning rate for the optimizer
             epsilon: epsilon constant for the Adam optimizer
-            hidden_layer_activation: activattion function to use for each dense layer
+            hidden_layer_activation: activation function to use for each dense layer
             kernel_l1_regularization: l1 regularization value for the kernel regularizer
             kernel_l2_regularization: l2 regularization value for the kernel regularizer
             bias_l1_regularization: l1 regularization value for the bias regularizer