From 5b051346fa218b345ef892a3db3514f3e8bc652e Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Fri, 30 Jul 2021 09:59:49 -0400
Subject: [PATCH] Fixing an NPE in LIMEExplanation.getActiveFeatures() with
 some explanation models.

---
 .../classification/explanations/lime/LIMEExplanation.java | 8 +++++++-
 .../explanations/lime/LIMEColumnarTest.java               | 4 ++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/Classification/Explanations/src/main/java/org/tribuo/classification/explanations/lime/LIMEExplanation.java b/Classification/Explanations/src/main/java/org/tribuo/classification/explanations/lime/LIMEExplanation.java
index 686064cd1..85b48c9bf 100644
--- a/Classification/Explanations/src/main/java/org/tribuo/classification/explanations/lime/LIMEExplanation.java
+++ b/Classification/Explanations/src/main/java/org/tribuo/classification/explanations/lime/LIMEExplanation.java
@@ -25,6 +25,7 @@
 import org.tribuo.regression.evaluation.RegressionEvaluation;
 
 import java.util.List;
+import java.util.Map;
 
 /**
  * An {@link Explanation} using LIME.
@@ -56,7 +57,12 @@ public LIMEExplanation(SparseModel<Regressor> model, Prediction<Label> predictio
 
     @Override
     public List<String> getActiveFeatures() {
-        return model.getActiveFeatures().get(prediction.getOutput().getLabel());
+        Map<String,List<String>> features = model.getActiveFeatures();
+        if (features.containsKey(Model.ALL_OUTPUTS)) {
+            return features.get(Model.ALL_OUTPUTS);
+        } else {
+            return features.get(prediction.getOutput().getLabel());
+        }
     }
 
     @Override
diff --git a/Classification/Explanations/src/test/java/org/tribuo/classification/explanations/lime/LIMEColumnarTest.java b/Classification/Explanations/src/test/java/org/tribuo/classification/explanations/lime/LIMEColumnarTest.java
index 01e499909..b0b483c0e 100644
--- a/Classification/Explanations/src/test/java/org/tribuo/classification/explanations/lime/LIMEColumnarTest.java
+++ b/Classification/Explanations/src/test/java/org/tribuo/classification/explanations/lime/LIMEColumnarTest.java
@@ -52,6 +52,7 @@
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
+import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 
@@ -205,6 +206,9 @@ public void testCategorical() throws URISyntaxException {
 
         Pair<LIMEExplanation, List<Example<Regressor>>> explanation = lime.explainWithSamples(testExample);
 
+        List<String> activeFeatures = explanation.getA().getActiveFeatures();
+        assertNotNull(activeFeatures);
+
         int[] aSampleCount = new int[3];
         int[] dSampleCount = new int[3];