Fixed energy bug and added norm implementation

ronggan · Dec 30, 2014 · 2254179 · 2254179
1 parent 445e812
commit 2254179
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/srmr/srmr.py b/srmr/srmr.py
@@ -22,7 +22,7 @@ def calc_cutoffs(cfs, fs, q):
     R = cfs + (B0 * fs / (2*np.pi))
     return L, R
 
-def srmr(x, fs, n_cochlear_filters=23, low_freq=125, min_cf=4, max_cf=128, fast=True):
+def srmr(x, fs, n_cochlear_filters=23, low_freq=125, min_cf=4, max_cf=128, fast=True, norm=False):
     wLengthS = .256
     wIncS = .064
     # Computing gammatone envelopes
@@ -50,7 +50,13 @@ def srmr(x, fs, n_cochlear_filters=23, low_freq=125, min_cf=4, max_cf=128, fast=
         mod_out = modfilt(MF, ac_ch)
         for j, mod_ch in enumerate(mod_out):
             mod_out_frame = segment_axis(mod_ch, wLength, overlap=wLength-wInc, end='delay')
-            energy[i,j,:] = np.sum((w*mod_out_frame)**2)
+            energy[i,j,:] = np.sum((w*mod_out_frame)**2, axis=1)
+
+    if norm:
+        peak_energy = np.max(np.mean(energy, axis=0))
+        min_energy = peak_energy*0.001
+        energy[energy < min_energy] = min_energy
+        energy[energy > peak_energy] = peak_energy
 
     erbs = np.flipud(calc_erbs(low_freq, fs, n_cochlear_filters))
 

diff --git a/test/test_srmr.py b/test/test_srmr.py
@@ -12,8 +12,17 @@ def test_srmr():
     print("Ratio (fast): %2.4f" % ratio)
     assert np.allclose(ratio, 6.062267651334784, rtol=1e-6, atol=1e-12)
 
+    ratio_norm_fast, avg_energy_norm_fast = srmr(s, fs, fast=True, norm=True, max_cf=30)
+    print("Ratio (norm, fast): %2.4f" % ratio_norm_fast)
+
     ratio_slow, avg_energy_slow = srmr(s, fs, fast=False)
     print("Ratio (slow): %2.4f" % ratio_slow)
 
+    ratio_norm, avg_energy_norm = srmr(s, fs, fast=False, norm=True, max_cf=30)
+    print("Ratio (norm): %2.4f" % ratio_norm)
+
+
+
+
 if __name__ == '__main__':
     test_srmr()