diff --git a/VBDS.m b/VBDS.m
index c1354e6ce50e763fcf8c6e48f4f1e70aba762ae9..c3ddcfb677a5806d3c1ebb554e59a2b4f0d7b94a 100644
--- a/VBDS.m
+++ b/VBDS.m
@@ -1,4 +1,4 @@
-function [paramVBDS,Spike,w]=VBDS(xb,Fs,paramVBDS,display,verbose)
+function [paramVBDS,Spike,w,Z]=VBDS(xb,Fs,paramVBDS,display,verbose)
 % Pipeline for spike detection, feature extraction and spike/LFP
 % classification/separation
 % input:
@@ -11,6 +11,7 @@ function [paramVBDS,Spike,w]=VBDS(xb,Fs,paramVBDS,display,verbose)
 % paramVDBS: updated method parameters
 % Spike: matrix of estimated spike waveforms 
 % w: despiked LFP
+% Z: spike labels
 
 %%
 %Band pass filtering to highlight the spikes and thresholding
@@ -44,6 +45,7 @@ for t=1:10
     KGMM(t)=mixture.K;
 end
 
+
 paramVBDS.K=ceil(mean(KGMM));
 
 end
@@ -53,5 +55,7 @@ end
 % Classif and despiking
 %
 [paramVBDS,Spike,w]=LFPSpikeDemixing(xbr,xb,Fs,ideltasr,gB,wfeat,paramVBDS,indsupp,verbose);
+[~,Z]=max(paramVBDS.param.piapost,[],2);
+Z(paramVBDS.param.r==1)=0;
 
 end
\ No newline at end of file
diff --git a/main.m b/main.m
index 34f63d8a3c71e34a627a6e57387b6fc2957b70a1..d91e3344244cc6d70856a5689f5372d83892fa28 100644
--- a/main.m
+++ b/main.m
@@ -30,17 +30,17 @@ verbose=true; % verbose mode
 % process
 
 %%% SIMULATED DATA Le Cam et. al.
-% DATA_PATH='./';
-% DATA_NAME='example.mat';
-% load([DATA_PATH DATA_NAME],'xb','Fs');
+DATA_PATH='./';
+DATA_NAME='example.mat';
+load([DATA_PATH DATA_NAME],'xb','Fs');
 %%%
 
 %%% SIMULATED DATA CAMUNAS (EXAMPLE 4)
-paramVBDS.coef=4.5; % better parameters for these simulated data
-paramVBDS.nbfeat=12;
-load('simulation_4.mat');
-xb=data;
-Fs=24000;
+% paramVBDS.coef=4.5; % better parameters for these simulated data
+% paramVBDS.nbfeat=12;
+% load('simulation_4.mat');
+% xb=data;
+% Fs=24000;
 %%%
 
 [l,c]=size(xb);
@@ -77,7 +77,7 @@ end
 %%%%%%%%%%%%%%%%%%%%%
   
 %% VBDS
-[paramVBDS,Spike,w]=VBDS(xb,Fs,paramVBDS,display,verbose);
+[paramVBDS,Spike,w,Z]=VBDS(xb,Fs,paramVBDS,display,verbose);
 
 
 %%%%%%%%%%%%%%%%%%%%%
@@ -87,7 +87,6 @@ end
 
 if display
 
-[~,classif]=max(paramVBDS.param.piapost');
 K=paramVBDS.K;
 ns=paramVBDS.ds*Fs/1000;
 
@@ -102,16 +101,14 @@ rootK2=ceil(K/rootK);
 
 for k=1:K
     subplot(rootK,rootK2,k);
-    piapostk=paramVBDS.param.piapost(:,k).*(1-paramVBDS.param.r);
-    Ik=find(piapostk>0.5);
-    plot((-ns/2+1:ns/2)/Fs*1000,Spike(Ik,:)');
+    plot((-ns/2+1:ns/2)/Fs*1000,Spike(Z==k,:)');
      hold on,
     plot((-ns/2+1:ns/2)/Fs*1000,paramVBDS.S0(k,:),'r--','Linewidth',2);
     plot((-ns/2+1:ns/2)/Fs*1000,paramVBDS.S0(k,:)+2.2*sqrt(paramVBDS.varskt(k)),'k--','Linewidth',1.5);
     plot((-ns/2+1:ns/2)/Fs*1000,paramVBDS.S0(k,:)-2.2*sqrt(paramVBDS.varskt(k)),'k--','Linewidth',1.5);
     ylim(1.05*[miny maxy]) %ylim(1.1*[-maxy maxy]);
     xlim([-1 1]);
-    title(['#' num2str(length(Ik))],'FontSize',16)
+    title(['#' num2str(sum(Z==k))],'FontSize',16)
     xlabel('Time (ms)','FontSize',16)
     ylabel('Voltage','FontSize',16)
 end