From fa1d2824fad7fff313c335385de41d083df3c76f Mon Sep 17 00:00:00 2001
From: Jean-Marc Valin <jmvalin@jmvalin.ca>
Date: Mon, 15 Oct 2018 02:45:27 -0400
Subject: [PATCH] Add diagonal to sparsity mask

---
 dnn/lpcnet.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/dnn/lpcnet.py b/dnn/lpcnet.py
index 08c9b32c..45a11e91 100644
--- a/dnn/lpcnet.py
+++ b/dnn/lpcnet.py
@@ -49,12 +49,14 @@ class Sparsify(Callback):
             #print ("density = ", density)
             for k in range(nb):
                 A = p[:, k*N:(k+1)*N]
+                A = A - np.diag(np.diag(A))
                 L=np.reshape(A, (N, N//16, 16))
                 S=np.sum(L*L, axis=-1)
                 SS=np.sort(np.reshape(S, (-1,)))
                 thresh = SS[round(N*N//16*(1-density))]
                 mask = (S>=thresh).astype('float32');
                 mask = np.repeat(mask, 16, axis=1)
+                mask = np.minimum(1, mask + np.diag(np.ones((N,))))
                 p[:, k*N:(k+1)*N] = p[:, k*N:(k+1)*N]*mask
                 #print(thresh, np.mean(mask))
             w[1] = p