From e21c7b33e8aa3eaa8eac7edf2ee885cdaca6d418 Mon Sep 17 00:00:00 2001
From: Bert Peters <bert@bertptrs.nl>
Date: Mon, 12 Mar 2018 17:07:51 +0100
Subject: [PATCH] Start implementing visualisation for LRN layers.

---
 src/LayerInfo.cpp      |  2 ++
 src/LayerInfo.hpp      |  1 +
 src/visualisations.cpp | 41 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 44 insertions(+)

diff --git a/src/LayerInfo.cpp b/src/LayerInfo.cpp
index ebdbfe8..a5686e4 100644
--- a/src/LayerInfo.cpp
+++ b/src/LayerInfo.cpp
@@ -18,6 +18,8 @@ LayerInfo::Type LayerInfo::typeByName(string_view name)
         return Type::InnerProduct;
     } else if (name == "Dropout") {
         return Type::DropOut;
+    } else if (name == "LRN") {
+        return Type::LRN;
     } else {
         LOG(INFO) << "Received unknown layer type: " << name << endl;
         return Type::Other;
diff --git a/src/LayerInfo.hpp b/src/LayerInfo.hpp
index 2261300..60050a5 100644
--- a/src/LayerInfo.hpp
+++ b/src/LayerInfo.hpp
@@ -18,6 +18,7 @@ namespace fmri
             Pooling,
             InnerProduct,
             DropOut,
+            LRN,
             Other
         };
 
diff --git a/src/visualisations.cpp b/src/visualisations.cpp
index a78eea3..80986a0 100644
--- a/src/visualisations.cpp
+++ b/src/visualisations.cpp
@@ -176,6 +176,44 @@ static Animation *getReLUAnimation(const fmri::LayerData &prevState,
     });
 }
 
+static Animation *getNormalizingAnimation(const fmri::LayerData &prevState, const LayerData &curState,
+                                          const vector<float> &prevPositions,
+                                          const vector<float> &curPositions) {
+    CHECK(prevState.shape() == curState.shape()) << "Shapes should be of equal size" << endl;
+    vector<DType> scaling(std::accumulate(prevState.shape().begin(), prevState.shape().end(), 1u, multiplies<void>()));
+    caffe::caffe_div(scaling.size(), prevState.data(), curState.data(), scaling.data());
+
+    // Fix divisions by zero. For those cases, pick 1 since it doesn't matter anyway.
+    for (auto &s : scaling) {
+        if (!isnormal(s)) {
+            s = 1;
+        }
+    }
+
+    if (prevState.shape().size() == 2) {
+        EntryList entries;
+        entries.reserve(scaling.size());
+        for (auto i : Range(scaling.size())) {
+            entries.emplace_back(scaling[i], make_pair(i, i));
+        }
+
+        auto max_val = *max_element(scaling.begin(), scaling.end());
+
+        return new ActivityAnimation(entries, prevPositions.data(), curPositions.data(),-10,  [=](float i) -> ActivityAnimation::Color {
+            auto intensity = clamp((i - 1) / (max_val - 1), 0.f, 1.f);
+            return {
+                    1 - intensity,
+                    1,
+                    1
+            };
+        });
+    } else {
+        // TODO: do something for image-like layers.
+    }
+
+    return nullptr;
+}
+
 Animation * fmri::getActivityAnimation(const fmri::LayerData &prevState, const fmri::LayerData &curState,
                                        const fmri::LayerInfo &layer, const vector<float> &prevPositions,
                                        const vector<float> &curPositions)
@@ -200,6 +238,9 @@ Animation * fmri::getActivityAnimation(const fmri::LayerData &prevState, const f
         case LayerInfo::Type::Pooling:
             return new PoolingLayerAnimation(prevState, curState, prevPositions, curPositions, -10);
 
+        case LayerInfo::Type::LRN:
+            return getNormalizingAnimation(prevState, curState, prevPositions, curPositions);
+
         default:
             return nullptr;
     }