From abdae595b7701bbdea00d5c4cc212c79c79575c2 Mon Sep 17 00:00:00 2001
From: Konstantin Nikolaou <87869540+KonstiNik@users.noreply.github.com>
Date: Wed, 14 Aug 2024 13:30:18 +0200
Subject: [PATCH] Put gram matrix normalization into entropy computation. (#7)

Co-authored-by: knikolaou <knikolaou@icp.uni-stuttgart.de>
---
 papyrus/measurements/measurements.py |  8 +++++++-
 papyrus/utils/analysis_utils.py      | 17 +++++++++++++++--
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/papyrus/measurements/measurements.py b/papyrus/measurements/measurements.py
index 5499168..3a3be22 100644
--- a/papyrus/measurements/measurements.py
+++ b/papyrus/measurements/measurements.py
@@ -315,6 +315,7 @@ def __init__(
         rank: int = 1,
         normalize_eigenvalues: bool = True,
         effective: bool = False,
+        normalize_matrix: bool = False,
     ):
         """
         Constructor method of the NTKCrossEntropy class.
@@ -332,11 +333,15 @@ def __init__(
         effective : bool (default=False)
                 If true, the entropy is divided by the theoretical maximum entropy of
                 the system thereby returning the effective entropy / entropy density.
-
+        normalize_matrix : bool (default=False)
+                If true, the NTK is normalized by the square root of the product of the
+                corresponding diagonal elements. This is equivalent to normalizing the
+                gradient vectors forming the NTK.
         """
         super().__init__(name=name, rank=rank)
         self.normalize_eigenvalues = normalize_eigenvalues
         self.effective = effective
+        self.normalize_matrix = normalize_matrix
 
     def apply(self, ntk: np.ndarray) -> np.ndarray:
         """
@@ -373,6 +378,7 @@ def apply(self, ntk: np.ndarray) -> np.ndarray:
             ntk,
             effective=self.effective,
             normalize_eig=self.normalize_eigenvalues,
+            normalize_matrix=self.normalize_matrix,
         )
 
 
diff --git a/papyrus/utils/analysis_utils.py b/papyrus/utils/analysis_utils.py
index 58e4636..74bf4dd 100644
--- a/papyrus/utils/analysis_utils.py
+++ b/papyrus/utils/analysis_utils.py
@@ -23,7 +23,10 @@
 
 import numpy as np
 
-from papyrus.utils.matrix_utils import compute_hermitian_eigensystem
+from papyrus.utils.matrix_utils import (
+    compute_hermitian_eigensystem,
+    normalize_gram_matrix,
+)
 
 
 def compute_trace(matrix: np.ndarray, normalize: bool = False) -> np.ndarray:
@@ -77,7 +80,10 @@ def compute_shannon_entropy(dist: np.ndarray, effective: bool = False) -> float:
 
 
 def compute_von_neumann_entropy(
-    matrix: np.ndarray, effective: bool = True, normalize_eig: bool = True
+    matrix: np.ndarray,
+    effective: bool = True,
+    normalize_eig: bool = True,
+    normalize_matrix: bool = False,
 ) -> float:
     """
     Compute the von-Neumann entropy of a matrix.
@@ -91,12 +97,19 @@ def compute_von_neumann_entropy(
             the system thereby returning the effective entropy.
     normalize_eig : bool (default = True)
             If true, the eigenvalues are scaled to look like probabilities.
+    normalize_matrix : bool (default=False)
+            If true, the NTK is normalized by the square root of the product of the
+            corresponding diagonal elements. This is equivalent to normalizing the
+            gradient vectors forming the NTK.
 
     Returns
     -------
     entropy : float
             Von-Neumann entropy of the matrix.
     """
+    if normalize_matrix:
+        matrix = normalize_gram_matrix(matrix)
+
     eigvals, _ = compute_hermitian_eigensystem(matrix, normalize=normalize_eig)
 
     entropy = compute_shannon_entropy(eigvals)