From 7fada2a2273d1e772c26019df1404b4390fe2431 Mon Sep 17 00:00:00 2001
From: Jens Ahrensfeld <jens@jayfield.org>
Date: Thu, 28 May 2015 17:41:06 +0000
Subject: [PATCH] [RBM] - committed local changes

git-svn-id: http://moon:8086/svn/software/trunk/projects/RBM@270 b431acfa-c32f-4a4a-93f1-934dc6c82436
---
 Source/DrawComponent.cpp |  21 +-
 Source/DrawComponent.h   |   8 +-
 Source/HiddenLayer.hpp   |   4 +-
 Source/Layer.hpp         |  18 +-
 Source/LayerArray.hpp    | 213 +++-------------
 Source/MainComponent.cpp | 191 ++++++++++-----
 Source/MainComponent.h   |  21 +-
 Source/Rbm.hpp           | 514 ++++++++++++++-------------------------
 Source/VisibleLayer.hpp  |   4 +-
 Source/Weights.hpp       | 113 ++++++---
 10 files changed, 467 insertions(+), 640 deletions(-)

diff --git a/Source/DrawComponent.cpp b/Source/DrawComponent.cpp
index b31fa0c..6b05605 100644
--- a/Source/DrawComponent.cpp
+++ b/Source/DrawComponent.cpp
@@ -222,22 +222,22 @@ void DrawComponent::clear()
 	setData(m_data);
 }
 
-const VectorXd& DrawComponent::getData ()
+const RowVectorXd& DrawComponent::getData ()
 {
 	return m_data;
 }
 
-void DrawComponent::setData (const VectorXd& data)
+void DrawComponent::setData (const RowVectorXd& data)
 {
 	m_data = data;
-	DrawData(data);
+	DrawData();
 }
 
-void DrawComponent::DrawData (const VectorXd& data)
+void DrawComponent::DrawData ()
 {
 
 	double a;
-	VectorXd temp = data;
+	RowVectorXd temp = m_data;
 	double min = +1E12;
 	double max = -1E12;
 
@@ -246,13 +246,18 @@ void DrawComponent::DrawData (const VectorXd& data)
 		min = std::min<double>(min, (double)temp[i]);
 		max = std::max<double>(max, (double)temp[i]);
 	}
-	for (int i=0; i < m_width*m_height; i++)
+
+	if (min < 0)
 	{
-		temp[i] -= min;
+		for (int i=0; i < m_width*m_height; i++)
+		{
+			temp[i] -= min;
+		}
+		max -= min;
 	}
 	for (int i=0; i < m_width*m_height; i++)
 	{
-		temp[i] /= (max-min);
+		temp[i] /= max;
 	}
 
 	for (int i=0; i < m_height; i++)
diff --git a/Source/DrawComponent.h b/Source/DrawComponent.h
index ab0ceb2..6fdebbc 100644
--- a/Source/DrawComponent.h
+++ b/Source/DrawComponent.h
@@ -56,9 +56,8 @@ public:
     //[UserMethods]     -- You can add your own custom methods in this section.
     void setListener(DrawListener *pListener);
 	void drawAt(int x, int y, bool setColor);
-	void setData(const VectorXd& data);
-	void DrawData(const VectorXd& data);
-	const VectorXd& getData();
+	void setData(const RowVectorXd& data);
+	const RowVectorXd& getData();
 	void clear();
     //[/UserMethods]
 
@@ -83,10 +82,11 @@ private:
     float m_scaleX;
     float m_scaleY;
     ScopedPointer<Graphics>m_pG;
-    VectorXd m_data;
+    RowVectorXd m_data;
 	Image m_image;
 	double m_currData;
 	Colour m_currColor;
+	void DrawData();
     //[/UserVariables]
 
     //==============================================================================
diff --git a/Source/HiddenLayer.hpp b/Source/HiddenLayer.hpp
index 61eae28..7dc6c4b 100644
--- a/Source/HiddenLayer.hpp
+++ b/Source/HiddenLayer.hpp
@@ -14,7 +14,7 @@
 class HiddenLayer : public Layer
 {
 public:
-	HiddenLayer(uint32_t numUnits = 0, const VectorXd *pStatesInit = nullptr)
+	HiddenLayer(uint32_t numUnits = 0, const RowVectorXd *pStatesInit = nullptr)
 	: Layer(numUnits, pStatesInit)
 	{
 	}
@@ -28,7 +28,7 @@ private:
 	{
 		double sum = ((Weights&)weights).hiddenBias()[index];
 
-		sum += layer.states().transpose() * ((Weights&)weights).weights().col(index);
+		sum += layer.states() * ((Weights&)weights).weights().col(index);
 
 		return sum;
 	}
diff --git a/Source/Layer.hpp b/Source/Layer.hpp
index c0c510e..1d23c7a 100644
--- a/Source/Layer.hpp
+++ b/Source/Layer.hpp
@@ -21,7 +21,7 @@ using namespace Eigen;
 class Layer
 {
 public:
-	Layer(uint32_t numUnits = 0, const VectorXd *pStatesInit = nullptr)
+	Layer(uint32_t numUnits = 0, const RowVectorXd *pStatesInit = nullptr)
 	: m_numUnits(numUnits)
 	, m_probs(numUnits)
 	, m_states(numUnits)
@@ -109,12 +109,12 @@ public:
 		}
 	}
 
-	VectorXd& probs()
+	RowVectorXd& probs()
 	{
 		return m_probs;
 	}
 
-	VectorXd& states()
+	RowVectorXd& states()
 	{
 		return m_states;
 	}
@@ -129,11 +129,11 @@ private:
 
 protected:
 	uint32_t m_numUnits;
-	VectorXd m_probs;
-	VectorXd m_states;
+	RowVectorXd m_probs;
+	RowVectorXd m_states;
 
 	virtual double accum(Layer &layer, Weights &weights, uint32_t index) = 0;
-	void logSigmoid(const VectorXd &x)
+	void logSigmoid(const RowVectorXd &x)
 	{
 		uint32_t i;
 
@@ -143,7 +143,7 @@ protected:
 		}
 	}
 
-	void gaussProb(const VectorXd &mu, double sigma)
+	void gaussProb(const RowVectorXd &mu, double sigma)
 	{
 		uint32_t i;
 		double var = sigma*sigma;
@@ -151,8 +151,8 @@ protected:
 
 		for (i=0; i < m_numUnits; i++)
 		{
-			double x2 = ((double)mu[i]);
-			m_probs[i] = 1-exp(-0.5*x2*x2/var);
+			double x2 = (1-(double)mu[i]);
+			m_probs[i] = exp(-0.5*x2*x2/var);
 		}
 	}
 };
diff --git a/Source/LayerArray.hpp b/Source/LayerArray.hpp
index 70e8718..402b379 100644
--- a/Source/LayerArray.hpp
+++ b/Source/LayerArray.hpp
@@ -11,165 +11,77 @@
 #ifndef LAYERARRAY_HPP
 #define LAYERARRAY_HPP
 #include <stdint.h>
+#include <iostream>
 #include <Eigen/Dense>
 
+using namespace std;
 using namespace Eigen;
 
-template <class T>
 class LayerArray;
 
-template <class T>
 class LayerArrayListener
 {
 public:
 	LayerArrayListener() {}
 	virtual~LayerArrayListener() {}
 	
-	virtual void onChanged(const LayerArray<T> &obj) = 0;
+	virtual void onChanged(const LayerArray &obj) = 0;
 };
 
-template <class T>
 class LayerArray
 {
-	class Entry : public T
-	{
-	public:
-		Entry(uint32_t numUnits, const VectorXd *pInit)
-		: T(numUnits, pInit)
-		, pPrev(nullptr)
-		, pNext(nullptr)
-		{
-		}
-		~Entry()
-		{
-		}
-		Entry *pPrev;
-		Entry *pNext;
-	private:
-
-	};
-
 public:
-	LayerArray(LayerArrayListener<T> *pListener = nullptr)
-	: m_size(0)
-	, m_pRoot(nullptr)
-	, m_ppIndex(nullptr)
-	, m_pListener(pListener)
+	LayerArray(LayerArrayListener *pListener = nullptr)
+	: m_pListener(pListener)
 	{
 	}
 
-	LayerArray(uint32_t numLayers, uint32_t numUnitsPerLayer)
-	: m_size(0)
-	, m_pRoot(nullptr)
-	, m_ppIndex(nullptr)
-	, m_pListener(nullptr)
-	{
-		uint32_t i;
-
-		for (i=0; i < numLayers; i++)
-		{
-			add(nullptr, numUnitsPerLayer);
-		}
-	}
-
-	LayerArray (const LayerArray<T> &rhs)
-	: m_size(0)
-	, m_pRoot(nullptr)
-	, m_ppIndex(nullptr)
-	, m_pListener(nullptr)
-	{
-		uint32_t i;
-
-		for (i=0; i < rhs.getSize(); i++)
-		{
-			this->add(&rhs[i].states(), rhs[i].states().size());
-		}
-	}
-
 	virtual ~LayerArray()
 	{
 		m_pListener = nullptr;
 		clear();
 	}
 
-	T* add(const VectorXd *pData, uint32_t size)
+	void add(const RowVectorXd &pData, uint32_t size)
 	{
-		Entry *pNew;
-		Entry *pL = m_pRoot;
-
-		pNew = new Entry(size, pData);
-		if (!m_pRoot)
-		{
-			m_pRoot = pNew;
-		}
-		else
-		{
-			pL = m_pRoot;
-			while(pL->pNext)
-			{
-				pL = pL->pNext;
-			}
-			pL->pNext = pNew;
-			pL->pNext->pPrev = pL;
-		}
-		rebuildIndex();
+		m_data << pData;
 		if (m_pListener)
+		{
 			m_pListener->onChanged(*this);
-
-		return (T*)pNew;
+		}
 	}
 
 	void clear()
 	{
-		if (m_ppIndex)
-		{
-			for (int i=0; i < m_size; i++)
-			{
-				if (m_ppIndex[i])
-					delete m_ppIndex[i];
-
-				m_ppIndex[i] = nullptr;
-			}
-		}
-		m_ppIndex = nullptr;
-		m_pRoot = nullptr;
-		m_size = 0;
-		allocIndex(0);
+		m_data.resize(0,0);
 		if (m_pListener)
 			m_pListener->onChanged(*this);
 	}
 
-	T& getAt(uint32_t index) const
+	const MatrixXd& data() const
 	{
-		return *m_ppIndex[index];
+		return m_data;
 	}
 
-	T& operator[](uint32_t index) const
+	RowVectorXd getAt(uint32_t index) const
+	{
+		return m_data.row(index);
+	}
+
+	RowVectorXd operator[](uint32_t index) const
 	{
 		return getAt(index);
 	}
 
 	void removeAt(uint32_t index)
 	{
-		if (m_ppIndex[index]->pPrev)
-		{
-			m_ppIndex[index]->pPrev->pNext = m_ppIndex[index]->pNext;
-		}
-		else
-		{
-			m_pRoot = m_ppIndex[index]->pNext;
-			m_ppIndex[index]->pNext->pPrev = nullptr;
-		}
-		delete m_ppIndex[index];
-		m_ppIndex[index] = nullptr;
-		rebuildIndex();
 		if (m_pListener)
 			m_pListener->onChanged(*this);
 	}
 
 	uint32_t getSize() const
 	{
-		return m_size;
+		return m_data.rows();
 	}
 
 	void save(const char *pFilename)
@@ -182,18 +94,16 @@ public:
 			return;
 
 
-		fprintf(pFile, "%d\n", m_size);
+		fprintf(pFile, "%u\n", (uint32_t)m_data.rows());
+		fprintf(pFile, "%u\n", (uint32_t)m_data.cols());
 
 		uint32_t i, j;
 
-		for (i=0; i < m_size; i++)
+		for (i=0; i < m_data.rows(); i++)
 		{
-			uint32_t numUnits = m_ppIndex[i]->getNumUnits();
-			fprintf(pFile, "%d\n", numUnits);
-
-			for (j=0; j < numUnits; j++)
+			for (j=0; j < m_data.cols(); j++)
 			{
-				fprintf(pFile, "%3.6f\n", m_ppIndex[i]->states()(j));
+				fprintf(pFile, "%3.6f\n", m_data.row(i)(j));
 			}
 		}
 
@@ -203,6 +113,7 @@ public:
 	void load(const char *pFilename)
 	{
 		uint32_t size = 0;
+		uint32_t numUnits = 0;
 		FILE *pFile;
 
 		pFile = fopen(pFilename,"r");
@@ -211,79 +122,33 @@ public:
 			return;
 
 		fscanf(pFile, "%d\n", &size);
+		fscanf(pFile, "%d\n", &numUnits);
+		m_data.resize(size, numUnits);
+		RowVectorXd data(numUnits);
 
 		uint32_t i, j;
-		float v;
-		double *pData;
 		for (i=0; i < size; i++)
 		{
-			uint32_t numUnits;
-			fscanf(pFile, "%d\n", &numUnits);
-			if (numUnits == 0)
-				break;
+			float v;
 
-			pData = new double[numUnits];
-			T* data = add(nullptr, numUnits);
-
-
-			for (j=0; j < numUnits; j++)
+			for (j=0; j < data.size(); j++)
 			{
 				fscanf(pFile, "%f", &v);
-				data->states()(j) = v;
+				data(j) = v;
 			}
-
-			delete [] pData;
+			fscanf(pFile, "%d\n", &numUnits);
+			m_data.row(i) = data;
 		}
-
 		fclose(pFile);
+		if (m_pListener)
+		{
+			m_pListener->onChanged(*this);
+		}
 	}
 
 private:
-	uint32_t m_size;
-	Entry *m_pRoot;
-	Entry **m_ppIndex;
-	LayerArrayListener<T> *m_pListener;
-	void allocIndex(size_t size)
-	{
-		if (m_ppIndex)
-		{
-			delete [] m_ppIndex;
-			m_ppIndex = nullptr;
-			allocIndex(size);
-		}
-		else
-		{
-			if (size)
-			{
-				m_ppIndex = new Entry*[size];
-			}
-		}
-	}
-
-	void rebuildIndex()
-	{
-		Entry *pL;
-		uint32_t size;
-
-		size = 0;
-		pL = m_pRoot;
-		while(pL)
-		{
-			size++;
-			pL = pL->pNext;
-		}
-		allocIndex(size);
-
-		size = 0;
-		pL = m_pRoot;
-		while(pL)
-		{
-			m_ppIndex[size++] = pL;
-			pL = pL->pNext;
-		}
-		m_size = size;
-
-	}
+	LayerArrayListener *m_pListener;
+	MatrixXd m_data;
 };
 
 #endif  // LAYERARRAY_HPP
diff --git a/Source/MainComponent.cpp b/Source/MainComponent.cpp
index 9647f93..c52432a 100644
--- a/Source/MainComponent.cpp
+++ b/Source/MainComponent.cpp
@@ -280,9 +280,9 @@ MainComponent::MainComponent ()
     weightInitLabel->setColour (TextEditor::backgroundColourId, Colour (0x00000000));
     weightInitLabel->addListener (this);
 
-    addAndMakeVisible (rbmTrainV2ToggleButton = new ToggleButton ("rbmTrainV2ToggleButton toggle button"));
-    rbmTrainV2ToggleButton->setButtonText (TRANS("Train Ver. 2"));
-    rbmTrainV2ToggleButton->addListener (this);
+    addAndMakeVisible (rbmLearnVarianceButton = new ToggleButton ("rbmLearnVariance button"));
+    rbmLearnVarianceButton->setButtonText (TRANS("Learn Variance"));
+    rbmLearnVarianceButton->addListener (this);
 
     addAndMakeVisible (rbmNormalizeDataToggleButton = new ToggleButton ("rbmNormalizeData toggle button"));
     rbmNormalizeDataToggleButton->setButtonText (TRANS("Normalize data"));
@@ -294,7 +294,7 @@ MainComponent::MainComponent ()
 	m_vNumY = 16;
 	m_hNum = 64;
 	m_numGibbs = 1;
-	m_weights.setUnits(m_vNumX*m_vNumY, m_hNum);
+	m_weights.setUnits(m_vNumX, m_vNumY, m_hNum);
 	m_pRbm = new Rbm(m_weights, this);
 	create();
 
@@ -348,7 +348,7 @@ MainComponent::~MainComponent()
     momentumLabel = nullptr;
     sparsityLearningRateLabel = nullptr;
     weightInitLabel = nullptr;
-    rbmTrainV2ToggleButton = nullptr;
+    rbmLearnVarianceButton = nullptr;
     rbmNormalizeDataToggleButton = nullptr;
 
 
@@ -356,6 +356,7 @@ MainComponent::~MainComponent()
     DrawTraining = nullptr;
     DrawReconstruction = nullptr;
     DrawWeights = nullptr;
+    DrawVars = nullptr;
     DrawHidden = nullptr;
     m_pRbm = nullptr;
 
@@ -471,13 +472,14 @@ void MainComponent::resized()
     momentumLabel->setBounds (208, 368, 72, 24);
     sparsityLearningRateLabel->setBounds (416, 320, 72, 24);
     weightInitLabel->setBounds (416, 368, 72, 24);
-    rbmTrainV2ToggleButton->setBounds (336, 200, 128, 24);
+    rbmLearnVarianceButton->setBounds (336, 200, 128, 24);
     rbmNormalizeDataToggleButton->setBounds (336, 232, 160, 24);
     //[UserResized] Add your own custom resize handling here..
     DrawTraining->setBounds (16, 16, 100, 100);
     DrawReconstruction->setBounds (110+16, 16, 100, 100);
     DrawWeights->setBounds (220+16, 16, 100, 100);
-    DrawHidden->setBounds (16, 16+110, 320, 20);
+    DrawVars->setBounds (330+16, 16, 100, 100);
+    DrawHidden->setBounds (16, 16+110, 430, 20);
     //[/UserResized]
 }
 
@@ -496,7 +498,7 @@ void MainComponent::buttonClicked (Button* buttonThatWasClicked)
     {
         //[UserButtonCode_addButton] -- add your button handler code here..
 		DrawReconstruction->setData(DrawTraining->getData());
-		m_layers.add(&DrawTraining->getData(), m_vNumX*m_vNumY);
+		m_layers.add(DrawTraining->getData(), m_vNumX*m_vNumY);
 		patterSlider->setRange (0, m_layers.getSize()-1, 1);
         //[/UserButtonCode_addButton]
     }
@@ -528,7 +530,7 @@ void MainComponent::buttonClicked (Button* buttonThatWasClicked)
 		m_vNumY = numVisibleYLabel->getText().getIntValue();
 		m_hNum = numHiddenLabel->getText().getIntValue();
 
-		m_weights.setUnits(m_vNumX*m_vNumY, m_hNum);
+		m_weights.setUnits(m_vNumX, m_vNumY, m_hNum);
 		create();
         //[/UserButtonCode_createButton]
     }
@@ -604,23 +606,6 @@ void MainComponent::buttonClicked (Button* buttonThatWasClicked)
     {
         //[UserButtonCode_rbmUseVisibleGaussianToggleButton] -- add your button handler code here..
     	m_pRbm->setUseVisibleGaussian(buttonThatWasClicked->getToggleState());
-   		rbmReduceVarianceToggleButton->setEnabled(!buttonThatWasClicked->getToggleState());
-
-   		if (buttonThatWasClicked->getToggleState())
-   		{
-   	    	rbmReduceVarianceToggleButton_binary = rbmReduceVarianceToggleButton->getToggleState();
-   			rbmReduceVarianceToggleButton->setToggleState(false, sendNotification);
-   			sigmaLabel->setText(String(sigma_gauss, 2), sendNotification);
-			sigmaDecayLabel->setText(String(sigmaDecay_gauss, 2), sendNotification);
-   		}
-   		else
-   		{
-   			sigma_gauss = sigmaLabel->getText().getFloatValue();
-   			sigmaDecay_gauss = sigmaDecayLabel->getText().getFloatValue();
-   			sigmaLabel->setText("1.0", sendNotification);
-			sigmaDecayLabel->setText("1.0", sendNotification);
-   			rbmReduceVarianceToggleButton->setToggleState(rbmReduceVarianceToggleButton_binary, sendNotification);
-   		}
     	redrawReconstruction();
         //[/UserButtonCode_rbmUseVisibleGaussianToggleButton]
     }
@@ -630,10 +615,11 @@ void MainComponent::buttonClicked (Button* buttonThatWasClicked)
     	m_pRbm->setDoSparse(buttonThatWasClicked->getToggleState());
         //[/UserButtonCode_rbmDoSparseToggleButton]
     }
-    else if (buttonThatWasClicked == rbmTrainV2ToggleButton)
+    else if (buttonThatWasClicked == rbmLearnVarianceButton)
     {
-        //[UserButtonCode_rbmTrainV2ToggleButton] -- add your button handler code here..
-        //[/UserButtonCode_rbmTrainV2ToggleButton]
+        //[UserButtonCode_rbmLearnVarianceButton] -- add your button handler code here..
+    	m_pRbm->setDoLearnVariance(buttonThatWasClicked->getToggleState());
+        //[/UserButtonCode_rbmLearnVarianceButton]
     }
     else if (buttonThatWasClicked == rbmNormalizeDataToggleButton)
     {
@@ -656,19 +642,10 @@ void MainComponent::sliderValueChanged (Slider* sliderThatWasMoved)
         //[UserSliderCode_patterSlider] -- add your slider handling code here..
 		if (m_layers.getSize() > 0)
 		{
-			VisibleLayer &p = (VisibleLayer&)m_layers.getAt((int)sliderThatWasMoved->getValue());
-			VectorXd t = p.states();
-
-			DrawTraining->setData(p.states());
-			cout << "mean = " << (double)p.states().array().mean() << endl;
-			cout << "max = " << (double)p.states().array().maxCoeff() << endl;
-			cout << "min = " << (double)p.states().array().minCoeff() << endl;
-			t.array() -= t.array().mean();
-			t.array() *= t.array();
-			cout << "var = " << (double)t.mean() << endl;
+			RowVectorXd t = m_layers.getAt((int)sliderThatWasMoved->getValue());
 
+			DrawTraining->setData(t);
 			redrawReconstruction();
-
 		}
         //[/UserSliderCode_patterSlider]
     }
@@ -784,6 +761,56 @@ void MainComponent::labelTextChanged (Label* labelThatHasChanged)
     //[/UserlabelTextChanged_Post]
 }
 
+void MainComponent::mouseMove (const MouseEvent& e)
+{
+    //[UserCode_mouseMove] -- Add your code here...
+    //[/UserCode_mouseMove]
+}
+
+void MainComponent::mouseEnter (const MouseEvent& e)
+{
+    //[UserCode_mouseEnter] -- Add your code here...
+    //[/UserCode_mouseEnter]
+}
+
+void MainComponent::mouseExit (const MouseEvent& e)
+{
+    //[UserCode_mouseExit] -- Add your code here...
+    //[/UserCode_mouseExit]
+}
+
+void MainComponent::mouseDown (const MouseEvent& e)
+{
+    //[UserCode_mouseDown] -- Add your code here...
+    //[/UserCode_mouseDown]
+}
+
+void MainComponent::mouseDrag (const MouseEvent& e)
+{
+    //[UserCode_mouseDrag] -- Add your code here...
+//	e.eventComponent->setCentrePosition(e.getPosition().x, e.getPosition().y);
+	cout << "Mouse = " << e.x << "," << e.y << endl;
+    //[/UserCode_mouseDrag]
+}
+
+void MainComponent::mouseUp (const MouseEvent& e)
+{
+    //[UserCode_mouseUp] -- Add your code here...
+    //[/UserCode_mouseUp]
+}
+
+void MainComponent::mouseDoubleClick (const MouseEvent& e)
+{
+    //[UserCode_mouseDoubleClick] -- Add your code here...
+    //[/UserCode_mouseDoubleClick]
+}
+
+void MainComponent::mouseWheelMove (const MouseEvent& e, const MouseWheelDetails& wheel)
+{
+    //[UserCode_mouseWheelMove] -- Add your code here...
+    //[/UserCode_mouseWheelMove]
+}
+
 
 
 //[MiscUserCode] You can add your own definitions of your custom methods or any other code here...
@@ -791,7 +818,8 @@ void MainComponent::load ()
 {
 	m_weights.load((String(getBaseDir() + String(".weights.dat"))).toUTF8());
 
-	m_vNumX = m_vNumY = (int)sqrt((float)m_weights.getNumVisible());
+	m_vNumX = m_weights.getNumVisibleX();
+	m_vNumY = m_weights.getNumVisibleY();
 	m_hNum = m_weights.getNumHidden();
 	create();
 }
@@ -806,6 +834,7 @@ void MainComponent::create()
 	DrawTraining = nullptr;
 	DrawReconstruction = nullptr;
 	DrawWeights = nullptr;
+	DrawVars = nullptr;
 	DrawHidden = nullptr;
 
 	WeightsSlider->setRange(0, m_hNum-1, 1);
@@ -814,14 +843,13 @@ void MainComponent::create()
 	DrawTraining->setListener(this);
     addAndMakeVisible (DrawReconstruction = new DrawComponent (m_vNumX, m_vNumY));
     addAndMakeVisible (DrawWeights = new DrawComponent (m_vNumX, m_vNumY));
+    addAndMakeVisible (DrawVars = new DrawComponent (m_vNumX, m_vNumY));
     addAndMakeVisible (DrawHidden = new DrawComponent (m_hNum, 1));
 	DrawHidden->setListener(this);
 	numVisibleLabel->setText(String(m_vNumX), dontSendNotification );
 	numVisibleYLabel->setText(String(m_vNumY), dontSendNotification );
 	numHiddenLabel->setText(String(m_hNum), dontSendNotification );
 
-	sigma_gauss = 0.5;
-	sigmaDecay_gauss = 1.0;
 	m_pRbm->setDoRaoBlackwell(rbmDoRaoBlackwellToggleButton->getToggleState());
 	m_pRbm->setUseProbsForHiddenReconstruction(rbmReduceVarianceToggleButton->getToggleState());
 	m_pRbm->setUseVisibleGaussian(rbmUseVisibleGaussianToggleButton->getToggleState());
@@ -850,9 +878,9 @@ const juce::String& MainComponent::getBaseDir()
 
 }
 
-void MainComponent::onChanged(const LayerArray<VisibleLayer> &obj)
+void MainComponent::onChanged(const LayerArray &obj)
 {
-	patterSlider->setRange(0, std::max(0,(int)m_layers.getSize()-1), 1);
+	patterSlider->setRange(0, std::max(0,(int)obj.getSize()-1), 1);
 }
 
 void MainComponent::onEpochTrained(const Rbm &obj)
@@ -871,17 +899,46 @@ void MainComponent::onDraw(DrawComponent &obj)
 	}
 	if (&obj == DrawTraining)
 	{
-		DrawReconstruction->setData(m_pRbm->toVisible(m_pRbm->toHidden(obj.getData())));
-    	DrawHidden->setData(m_pRbm->toHidden(obj.getData()));
+//		DrawReconstruction->setData(m_pRbm->toVisible(m_pRbm->toHidden(obj.getData())));
+ //   	DrawHidden->setData(m_pRbm->toHidden(obj.getData()));
+    	redrawReconstruction();
 	}
 }
 
 void MainComponent::redrawReconstruction()
 {
+	char table[] = " ABCDEFGHIJKLMNOPQRSTUVWXYZ";
+	VectorXd v;
+	MatrixXd m;
+
 	DrawHidden->setData(m_pRbm->toHidden(DrawTraining->getData()));
-	DrawReconstruction->setData(m_pRbm->toVisible(DrawHidden->getData()));
-//	double energy = m_pRbm->getEnergy(DrawTraining->getData(), DrawHidden->getData());
-//	cout << "Energy(" << 0 <<") = " << energy << endl;
+	v = m_pRbm->toVisible(DrawHidden->getData());
+	DrawReconstruction->setData(v);
+
+	if ((m_vNumX == 27) && (m_vNumY == 4))
+	{
+		m = v;
+		m.resize(m_vNumX, m_vNumY);
+		cout << "Reconstruction" << endl;
+		cout << m << endl;
+
+		for (int y=0; y < m_vNumY; y++)
+		{
+			double maxV = -1000;
+			int maxX = 0;
+			for (int x=0; x < m_vNumX; x++)
+			{
+				if (m(x, y) > maxV)
+				{
+					maxV = m(x, y);
+					maxX = x;
+				}
+			}
+			cout << table[maxX];
+		}
+		cout << endl;
+	}
+
 
 }
 
@@ -889,20 +946,15 @@ void MainComponent::redrawWeights(int index)
 {
 	VectorXd w = m_weights.weights().col(index);
 	DrawWeights->setData(w);
+	DrawVars->setData(m_weights.sigma());
+
 }
 
 void MainComponent::run()
 {
-	trainButton->setEnabled(false);
-	if (rbmTrainV2ToggleButton->getToggleState())
-	{
-		m_pRbm->train2(m_layers, numEpochslabel->getText().getIntValue(), 100);
-	}
-	else
-	{
-		m_pRbm->train(m_layers, numEpochslabel->getText().getIntValue());
-	}
-	trainButton->setEnabled(true);
+//	trainButton->setEnabled(false);
+	m_pRbm->train(m_layers, numEpochslabel->getText().getIntValue(), 100);
+//	trainButton->setEnabled(true);
 }
 
 //[/MiscUserCode]
@@ -918,10 +970,20 @@ void MainComponent::run()
 BEGIN_JUCER_METADATA
 
 <JUCER_COMPONENT documentType="Component" className="MainComponent" componentName=""
-                 parentClasses="public Component, public LayerArrayListener&lt;VisibleLayer&gt;, public RbmListener, public DrawListener, public Thread"
+                 parentClasses="public Component, public LayerArrayListener, public RbmListener, public DrawListener, public Thread"
                  constructorParams="" variableInitialisers="Thread(&quot;RBM&quot;),&#10;m_layers(this),&#10;m_pRbm(nullptr)&#10;DrawTraining(nullptr),&#10;DrawReconstruction(nullptr),&#10;DrawWeights(nullptr),&#10;DrawHidden(nullptr)"
                  snapPixels="8" snapActive="1" snapShown="1" overlayOpacity="0.330"
                  fixedSize="1" initialWidth="800" initialHeight="600">
+  <METHODS>
+    <METHOD name="mouseMove (const MouseEvent&amp; e)"/>
+    <METHOD name="mouseEnter (const MouseEvent&amp; e)"/>
+    <METHOD name="mouseExit (const MouseEvent&amp; e)"/>
+    <METHOD name="mouseDown (const MouseEvent&amp; e)"/>
+    <METHOD name="mouseDrag (const MouseEvent&amp; e)"/>
+    <METHOD name="mouseUp (const MouseEvent&amp; e)"/>
+    <METHOD name="mouseDoubleClick (const MouseEvent&amp; e)"/>
+    <METHOD name="mouseWheelMove (const MouseEvent&amp; e, const MouseWheelDetails&amp; wheel)"/>
+  </METHODS>
   <BACKGROUND backgroundColour="ffffffff">
     <TEXT pos="32 306 80 14" fill="solid: ff000000" hasStroke="0" text="Sigma"
           fontname="Default font" fontsize="15" bold="0" italic="0" justification="36"/>
@@ -1084,10 +1146,9 @@ BEGIN_JUCER_METADATA
          edBkgCol="0" labelText="0.001" editableSingleClick="1" editableDoubleClick="1"
          focusDiscardsChanges="0" fontname="Default font" fontsize="15"
          bold="0" italic="0" justification="36"/>
-  <TOGGLEBUTTON name="rbmTrainV2ToggleButton toggle button" id="92c647c1f8b110a2"
-                memberName="rbmTrainV2ToggleButton" virtualName="" explicitFocusOrder="0"
-                pos="336 200 128 24" buttonText="Train Ver. 2" connectedEdges="0"
-                needsCallback="1" radioGroupId="0" state="0"/>
+  <TOGGLEBUTTON name="rbmLearnVariance button" id="92c647c1f8b110a2" memberName="rbmLearnVarianceButton"
+                virtualName="" explicitFocusOrder="0" pos="336 200 128 24" buttonText="Learn Variance"
+                connectedEdges="0" needsCallback="1" radioGroupId="0" state="0"/>
   <TOGGLEBUTTON name="rbmNormalizeData toggle button" id="739772af1b096120" memberName="rbmNormalizeDataToggleButton"
                 virtualName="" explicitFocusOrder="0" pos="336 232 160 24" buttonText="Normalize data"
                 connectedEdges="0" needsCallback="1" radioGroupId="0" state="0"/>
diff --git a/Source/MainComponent.h b/Source/MainComponent.h
index 2806777..45d6b01 100644
--- a/Source/MainComponent.h
+++ b/Source/MainComponent.h
@@ -38,7 +38,7 @@
                                                                     //[/Comments]
 */
 class MainComponent  : public Component,
-                       public LayerArrayListener<VisibleLayer>,
+                       public LayerArrayListener,
                        public RbmListener,
                        public DrawListener,
                        public Thread,
@@ -60,6 +60,14 @@ public:
     void buttonClicked (Button* buttonThatWasClicked);
     void sliderValueChanged (Slider* sliderThatWasMoved);
     void labelTextChanged (Label* labelThatHasChanged);
+    void mouseMove (const MouseEvent& e);
+    void mouseEnter (const MouseEvent& e);
+    void mouseExit (const MouseEvent& e);
+    void mouseDown (const MouseEvent& e);
+    void mouseDrag (const MouseEvent& e);
+    void mouseUp (const MouseEvent& e);
+    void mouseDoubleClick (const MouseEvent& e);
+    void mouseWheelMove (const MouseEvent& e, const MouseWheelDetails& wheel);
 
 
 
@@ -70,17 +78,19 @@ private:
     ScopedPointer<DrawComponent> DrawTraining;
     ScopedPointer<DrawComponent> DrawReconstruction;
     ScopedPointer<DrawComponent> DrawWeights;
+    ScopedPointer<DrawComponent> DrawVars;
     ScopedPointer<DrawComponent> DrawHidden;
 	uint32_t m_vNumX;
 	uint32_t m_vNumY;
 	uint32_t m_hNum;
-	LayerArray<VisibleLayer> m_layers;
+	LayerArray m_layers;
+	MatrixXd m_trainingData;
 	void load();
 	void save();
 	void create();
 	void destroy();
 	const juce::String& getBaseDir();
-	void onChanged(const LayerArray<VisibleLayer> &obj);
+	void onChanged(const LayerArray &obj);
 	void onEpochTrained(const Rbm &obj);
 	void onDraw(DrawComponent &obj);
 	void redrawReconstruction();
@@ -88,9 +98,6 @@ private:
 	void run();
 	String m_baseDir;
 	uint32_t m_numGibbs;
-	bool rbmReduceVarianceToggleButton_binary;
-	double sigma_gauss;
-	double sigmaDecay_gauss;
     //[/UserVariables]
 
     //==============================================================================
@@ -129,7 +136,7 @@ private:
     ScopedPointer<Label> momentumLabel;
     ScopedPointer<Label> sparsityLearningRateLabel;
     ScopedPointer<Label> weightInitLabel;
-    ScopedPointer<ToggleButton> rbmTrainV2ToggleButton;
+    ScopedPointer<ToggleButton> rbmLearnVarianceButton;
     ScopedPointer<ToggleButton> rbmNormalizeDataToggleButton;
 
 
diff --git a/Source/Rbm.hpp b/Source/Rbm.hpp
index 3d6b5aa..c9627c3 100644
--- a/Source/Rbm.hpp
+++ b/Source/Rbm.hpp
@@ -19,6 +19,8 @@ using namespace Eigen;
 void mylog(const char* format, ...);
 #define printf mylog
 
+#define EPSILON_SIGMA	0.05
+
 class Rbm;
 
 class RbmListener
@@ -53,9 +55,18 @@ public:
 	, m_useProbsForHiddenReconstruction(false)
 	, m_doSparse(false)
 	, m_doNormalizeData(false)
+	, m_doLearnVariance(false)
 	, m_numGibbs(1)
 	{
 		Noise_Init(&m_noise, 0x32727155);
+
+		VectorXd a(4);
+		a << 1, 2, 3, 4;
+		VectorXd b(4);
+
+		b.array() = -a.array().exp();
+
+		cout << b << endl;
 	}
 
 	~Rbm()
@@ -64,173 +75,6 @@ public:
 		Noise_Free(&m_noise);
 	}
 
-	void train(const LayerArray<VisibleLayer> &batch, uint32_t numEpochs, double sigmaMin = 0.05)
-	{
-		uint32_t t, i;
-		uint32_t epoch;
-		uint32_t gibbs;
-		double sigma;
-		VisibleLayer v(m_w.getNumVisible());
-		HiddenLayer h(m_w.getNumHidden());
-
-		VectorXd sumBiasV(m_w.getNumVisible());
-		VectorXd deltaBiasV(m_w.getNumVisible());
-
-		VectorXd sumBiasH(m_w.getNumHidden());
-		VectorXd deltaBiasH(m_w.getNumHidden());
-
-		MatrixXd sumWeights(m_w.getNumVisible(), m_w.getNumHidden());
-		MatrixXd deltaWeights(m_w.getNumVisible(), m_w.getNumHidden());
-
-		MatrixXd diffErr(1, m_w.getNumVisible());
-
-		const LayerArray<VisibleLayer> &vt = batch;
-
-		sigma = m_sigma;
-
-		double dProgress = 1.0/numEpochs;
-		double kTrain = 1.0/vt.getSize();
-
-		m_progress = 0;
-
-		deltaWeights.fill(0);
-		deltaBiasV.fill(0);
-		deltaBiasH.fill(0);
-		m_doCancel = false;
-		for (epoch=0; epoch < numEpochs; epoch++)
-		{
-			double err = 0;
-
-			if (m_doCancel)
-			{
-				m_doCancel = false;
-				break;
-			}
-
-			sumWeights.fill(0);
-			sumBiasV.fill(0);
-			sumBiasH.fill(0);
-			for (i=0; i < vt.getSize(); i++)
-			{
-				t = i;
-				h.probsUpdateLogistic(vt[t], m_w, m_lambda, sigma);
-
-				// Create hidden layer base on training data
-				if (m_doRaoBlackwell)
-				{
-					h.states() = h.probs();
-				}
-				else
-				{
-					h.statesUpdateStochastic();
-				}
-
-				// Update weights (positive phase)
-				sumWeights += vt[t].states() * h.states().transpose();
-				sumBiasV += vt[t].states();
-				sumBiasH += h.states();
-
-				diffErr = vt[t].states();
-
-				for (gibbs=0; gibbs < m_numGibbs; gibbs++)
-				{
-					h.statesUpdateStochastic();
-
-					// Create visible reconstruction (a fantasy...)
-					if (m_useProbsForHiddenReconstruction)
-					{
-						if (m_useVisibleGaussian)
-						{
-							v.probsUpdateGaussian(h, m_w, m_lambda, sigma);
-						}
-						else
-						{
-							v.probsUpdateLogistic(h, m_w, m_lambda, sigma);
-						}
-						v.states() = v.probs();
-					}
-					else
-					{
-						if (m_useVisibleGaussian)
-						{
-							v.sampleGaussian(h, m_w, m_lambda, sigma);
-						}
-						else
-						{
-							v.probsUpdateLogistic(h, m_w, m_lambda, sigma);
-							v.statesUpdateStochastic();
-						}
-					}
-					// Create hidden reconstruction
-					h.probsUpdateLogistic(v, m_w, m_lambda, sigma);
-				}
-
-				// Update weights (negative phase)
-				if (m_doRaoBlackwell)
-				{
-					h.states() = h.probs();
-				}
-				else
-				{
-					h.statesUpdateStochastic();
-				}
-				sumWeights -= v.states() * h.states().transpose();
-				sumBiasV -= v.states();
-				sumBiasH -= h.states();
-				diffErr -= v.states();
-				diffErr.array() *= diffErr.array();
-				err += diffErr.sum();
-
-			} // TrainingSize
-
-			deltaWeights = m_momentum*deltaWeights + m_muWeights*(kTrain*sumWeights - m_weightDecay*m_w.weights());
-			m_w.weights() += deltaWeights;
-
-			deltaBiasV = m_momentum*deltaBiasV + m_muWeights*kTrain*sumBiasV;
-			m_w.visibleBias() += deltaBiasV;
-
-			if (m_doSparse)
-			{
-				HiddenLayer th(m_w.getNumHidden());
-				VectorXd m(m_w.getNumHidden());
-				m.fill(0);
-
-				for (i=0; i < vt.getSize(); i++)
-				{
-					th.probsUpdateLogistic(vt[i], m_w, m_lambda, sigma);
-					m += th.probs();
-				}
-				m /= i;
-				sumBiasH = m_sparsity - m.array();
-				deltaBiasH = m_momentum*deltaBiasH + m_muSparsity*sumBiasH;
-
-				cout << "Mean(" << m_sparsity << ") = " << (double)m.array().mean() << endl;
-				cout << m << endl;
-			}
-			else
-			{
-				deltaBiasH = m_momentum*deltaBiasH + m_muWeights*kTrain*sumBiasH;
-			}
-
-			m_w.hiddenBias() += deltaBiasH;
-
-			if (sigma > sigmaMin)
-			{
-				sigma *= m_sigmaDecay;
-			}
-
-			m_progress += dProgress;
-			if (m_pListener)
-			{
-				m_pListener->onEpochTrained(*this);
-			}
-
-			cout << "err =" << endl;
-			cout << err << endl;
-
-		} // Number of epochs
-	}
-
 	void sample(MatrixXd &src)
 	{
 		uint32_t i;
@@ -241,61 +85,144 @@ public:
 		}
 	}
 
-	void probsLogistic(MatrixXd &src, double lambda, double sigma)
+	void probsLogistic(MatrixXd &src)
 	{
-		double var = sigma*sigma;
-
-		src.array() *= -lambda/var;
-		src.array() = src.array().exp();
+		src.array() = (-src.array()).exp();
 		src.array() += 1;
 		src.array() = 1.0/src.array();
 	}
 
-	void sampleGaussian(MatrixXd &src, double lambda, double sigma)
+	void probsLogistic(RowVectorXd &src)
+	{
+		src.array() = (-src.array()).exp();
+		src.array() += 1;
+		src.array() = 1.0/src.array();
+	}
+
+	void probsLogistic(MatrixXd &src, const MatrixXd &sigma)
+	{
+		src.array() /= (sigma.array() + EPSILON_SIGMA);
+		src.array() = (-src.array()).exp();
+		src.array() += 1;
+		src.array() = 1.0/src.array();
+	}
+
+	void probsLogistic(RowVectorXd &src, const RowVectorXd &sigma)
+	{
+		src.array() /= (sigma.array() + EPSILON_SIGMA);
+		src.array() = (-src.array()).exp();
+		src.array() += 1;
+		src.array() = 1.0/src.array();
+	}
+
+	void probsGaussian(MatrixXd &src, const MatrixXd &sigma)
+	{
+		src.array() = 1 - src.array();
+		src.array() *= src.array();
+		src.array() *= -0.5;
+
+		MatrixXd var = sigma;
+		var.array() += EPSILON_SIGMA;
+		var.array() *= var.array();
+
+		src.array() /= var.array();
+		src.array() = src.array().exp();
+
+		MatrixXd k = var;
+
+		k.array() *= 2*3.14159265359;
+		k.array() = k.array().sqrt();
+		k.array() = 1.0/k.array();
+
+		src.array() *= k.array();
+	}
+
+	void probsGaussian(RowVectorXd &src, const RowVectorXd &sigma)
+	{
+		src.array() = 1 - src.array();
+		src.array() *= src.array();
+		src.array() *= -0.5;
+
+		RowVectorXd var = sigma;
+		var.array() += EPSILON_SIGMA;
+		var.array() *= var.array();
+
+		src.array() /= var.array();
+		src.array() = src.array().exp();
+
+		RowVectorXd k = var;
+
+		k.array() *= 2*3.14159265359;
+		k.array() = k.array().sqrt();
+		k.array() = 1.0/k.array();
+
+		src.array() *= k.array();
+	}
+
+	void sampleGaussian(MatrixXd &src, const MatrixXd &sigma)
 	{
 		uint32_t i;
 
 		for (i=0; i < src.array().size(); i++)
 		{
-			src.array()(i) = sigma*Noise_Gaussian(&m_noise) + lambda*src.array()(i);
+			src.array()(i) = sigma(i)*Noise_Gaussian(&m_noise) + src.array()(i);
 		}
 	}
 
-	void normalizeData(MatrixXd &src, double mu, double sigma)
+	RowVectorXd normalizeData(RowVectorXd const &src, RowVectorXd const &mu, RowVectorXd const &var)
 	{
-		uint32_t i;
-		uint32_t size = src.rows();
-		double mean;
-		double stdDev;
+		// Remove mean
+		RowVectorXd res = src - mu;
+//		res.array() /= var.array() + EPSILON_SIGMA;
 
-		for (i=0; i < size; i++)
-		{
-			mean = src.row(i).array().mean();
-			src.row(i).array() -= mean;
-			src.row(i).array() += mu;
-		}
-		for (i=0; i < size; i++)
-		{
-			src.row(i).array() *= src.row(i).array();
-		}
-		for (i=0; i < size; i++)
-		{
-			stdDev = sqrt(src.row(i).array().mean());
-			src.row(i).array() /= stdDev;
-			src.row(i).array() *= sigma;
-		}
+//		cout << __PRETTY_FUNCTION__ << ": " << res << endl;
+		return res;
 	}
 
-	void train2(const LayerArray<VisibleLayer> &vt, uint32_t numEpochs, uint32_t batchSize, double sigmaMin = 0.05)
+	RowVectorXd calcMean(MatrixXd const &batch)
+	{
+		// Remove mean
+		RowVectorXd res = batch.colwise().mean();
+
+//		cout << __PRETTY_FUNCTION__ << ": " << res << endl;
+		return res;
+
+	}
+
+	RowVectorXd calcSigma(MatrixXd const &batch)
+	{
+		MatrixXd x = batch.rowwise() - batch.colwise().mean();
+
+		x.array() *= x.array();
+
+		RowVectorXd res = x.colwise().mean().array().sqrt();
+
+//		cout << __PRETTY_FUNCTION__ << ": " << res << endl;
+		return res;
+	}
+
+	MatrixXd calcZ(MatrixXd &v, MatrixXd &h)
+	{
+
+		MatrixXd t1(v.rows(), m_w.getNumVisible());
+
+		t1 = v - m_w.visibleBias().transpose().replicate(v.rows(), 1);
+		t1.array() *= t1.array();
+		t1.array() *= 0.5;
+
+		t1 -= (h * m_w.weights().transpose());
+
+		return t1;
+	}
+
+	void train(const LayerArray &vt, uint32_t numEpochs, uint32_t batchSize, double sigmaMin = 0.05)
 	{
 		uint32_t t, i;
 		uint32_t epoch;
 		uint32_t gibbs;
-		double sigma = m_sigma;
 		double dProgress = 1.0/numEpochs;
 		double kTrain = 1.0/vt.getSize();
 
-//		if (batchSize > vt.getSize())
 		batchSize = vt.getSize();
 
 		MatrixXd v(batchSize, m_w.getNumVisible());
@@ -306,6 +233,7 @@ public:
 		MatrixXd sumBiasH(1, m_w.getNumHidden());
 		MatrixXd sumWeights(m_w.getNumVisible(), m_w.getNumHidden());
 
+		MatrixXd deltaVar(MatrixXd::Zero(1, m_w.getNumVisible()));
 		MatrixXd deltaBiasV(MatrixXd::Zero(1, m_w.getNumVisible()));
 		MatrixXd deltaBiasH(MatrixXd::Zero(1, m_w.getNumHidden()));
 		MatrixXd deltaWeights(MatrixXd::Zero(m_w.getNumVisible(), m_w.getNumHidden()));
@@ -315,15 +243,22 @@ public:
 		m_progress = 0;
 		m_doCancel = false;
 
-		for (i=0; i < batchSize; i++)
-		{
-//			t = (uint32_t)(0.5 + (vt.getSize()-1)*Noise_Uniform(&m_noise));
-			batch.row(i) = vt[i].states();
+		batch = vt.data();
 
+		m_w.mean() = calcMean(batch);
+
+		if (m_doLearnVariance)
+		{
+			m_w.sigma() = calcSigma(batch);
 		}
+
 		if (m_doNormalizeData)
 		{
-			normalizeData(batch, 0.0, m_sigma);
+			for (i=0; i < batchSize; i++)
+			{
+				RowVectorXd x = batch.row(i);
+				batch.row(i) = normalizeData(x, m_w.mean(), m_w.sigma());
+			}
 		}
 
 		for (epoch=0; epoch < numEpochs; epoch++)
@@ -339,13 +274,14 @@ public:
 
 			// Create hidden layer base on training data
 			h = v * m_w.weights();
-			h += m_w.hiddenBias().transpose().replicate(batchSize, 1);
-			probsLogistic(h, m_lambda, sigma);
+			h += m_w.hiddenBias().replicate(batchSize, 1);
+			probsLogistic(h);
 
 			if (!m_doRaoBlackwell)
 			{
 				sample(h);
 			}
+
 			// Update weights (positive phase)
 			sumBiasV = v.colwise().sum();
 			if (!m_doSparse)
@@ -361,15 +297,18 @@ public:
 
 				// Create visible reconstruction (a fantasy...)
 				v = h * m_w.weights().transpose();
-				v += m_w.visibleBias().transpose().replicate(batchSize, 1);
+				v += m_w.visibleBias().replicate(batchSize, 1);
 
 				if (m_useVisibleGaussian)
 				{
-					sampleGaussian(v, m_lambda, sigma);
+					if (!m_useProbsForHiddenReconstruction)
+					{
+						sampleGaussian(v, m_w.sigma().replicate(batchSize, 1));
+					}
 				}
 				else
 				{
-					probsLogistic(v, m_lambda, sigma);
+					probsLogistic(v, m_w.sigma().replicate(batchSize, 1));
 					if (!m_useProbsForHiddenReconstruction)
 					{
 						sample(v);
@@ -378,8 +317,8 @@ public:
 
 				// Create hidden reconstruction
 				h = v * m_w.weights();
-				h += m_w.hiddenBias().transpose().replicate(batchSize, 1);
-				probsLogistic(h, m_lambda, sigma);
+				h += m_w.hiddenBias().replicate(batchSize, 1);
+				probsLogistic(h);
 			}
 
 			if (!m_doRaoBlackwell)
@@ -404,8 +343,8 @@ public:
 			if (m_doSparse)
 			{
 				h = v * m_w.weights();
-				h += m_w.hiddenBias().transpose().replicate(batchSize, 1);
-				probsLogistic(h, m_lambda, sigma);
+				h += m_w.hiddenBias().replicate(batchSize, 1);
+				probsLogistic(h);
 
 				sumBiasH.fill(m_sparsity);
 				sumBiasH -= h.colwise().mean();
@@ -421,9 +360,9 @@ public:
 			}
 			m_w.hiddenBias() += deltaBiasH;
 
-			if (sigma > sigmaMin)
+			if (m_w.sigma()[0] > sigmaMin)
 			{
-				sigma *= m_sigmaDecay;
+				m_w.sigma().array() *= m_sigmaDecay;
 			}
 
 			m_progress += dProgress;
@@ -450,148 +389,45 @@ public:
 	{
 		double energy;
 
-		energy = m_w.visibleBias().transpose() * visible;
-		energy += m_w.hiddenBias().transpose() * hidden;
+		energy = m_w.visibleBias() * visible;
+		energy += m_w.hiddenBias() * hidden;
 		energy += visible.transpose() * m_w.weights() * hidden;
 
 		return -energy/(m_sigma*m_sigma);
 	}
 
-	void prob(LayerArray<VisibleLayer> &vts)
+	RowVectorXd toHidden(const RowVectorXd& v)
 	{
-		uint32_t i, j;
-		double z;
-		double p;
+		RowVectorXd h(m_w.getNumHidden());
+		RowVectorXd vn(m_w.getNumVisible());
 
-		HiddenLayer *h = new HiddenLayer[vts.getSize()];
+		h = v * m_w.weights();
+		h += m_w.hiddenBias();
+		probsLogistic(h);
 
-		// Create hidden layer activations based on training data
-		for (j=0; j < vts.getSize(); j++)
-		{
-			h[j].setNumUnits(m_w.getNumHidden());
-			h[j].probsUpdateLogistic(vts.getAt(j), m_w, m_lambda, m_sigma);
-//			h[j].statesAssignfromProbs();
-			h[j].statesUpdateStochastic();
-		}
-
-		printf("pi(t) = (pi^, v>)\n");
-		for (j=0; j < vts.getSize(); j++)
-		{
-			cout << h[j].probs() << endl;
-		}
-		cout << endl;
-
-		printf("si(t) = (si^, v>)\n");
-		for (j=0; j < vts.getSize(); j++)
-		{
-			cout << h[j].states() << endl;
-		}
-		cout << endl;
-
-		printf("p(v) = (t^, v>)\n");
-		for (i=0; i < vts.getSize(); i++)
-		{
-			z = 0;
-			for (j=0; j < vts.getSize(); j++)
-			{
-				z += exp(-getEnergy(vts.getAt(j).states(), h[i].states()));
-			}
-			for (j=0; j < vts.getSize(); j++)
-			{
-				p = exp(-getEnergy(vts.getAt(j).states(), h[i].states()))/z;
-				cout << p << endl;
-			}
-			cout << endl;
-		}
-		cout << endl;
-
-		// Reconstruct
-		for (i=0; i < vts.getSize(); i++)
-		{
-			vts.getAt(i).probsUpdateLogistic(h[i], m_w, m_lambda, m_sigma);
-		}
-
-		printf("A fantasy... (v^, t>)\n");
-		for (j=0; j < vts.getSize(); j++)
-		{
-			cout << vts.getAt(j).probs() << endl;
-		}
-
-		delete [] h;
+		return h;
 	}
 
-	VectorXd toHidden(const VectorXd& visible)
+	RowVectorXd toVisible(const RowVectorXd& h)
 	{
-		HiddenLayer th(m_w.getNumHidden());
-		VisibleLayer tv(m_w.getNumVisible(), (const VectorXd*)&visible);
+		RowVectorXd v(m_w.getNumVisible());
+		v = h * m_w.weights().transpose();
+		v += m_w.visibleBias();
 
-		th.probsUpdateLogistic(tv, m_w, m_lambda, m_sigma);
-
-		return th.probs();
-	}
-
-	VectorXd toVisible(const VectorXd& hidden)
-	{
-		HiddenLayer th(m_w.getNumHidden(), (const VectorXd*)&hidden);
-		VisibleLayer tv(m_w.getNumVisible());
 		if (m_useVisibleGaussian)
 		{
-			tv.probsUpdateGaussian(th, m_w, m_lambda, m_sigma);
+//			probsGaussian(v, m_w.sigma());
 		}
 		else
 		{
-			tv.probsUpdateLogistic(th, m_w, m_lambda, m_sigma);
+			probsLogistic(v, m_w.sigma());
 		}
-		return tv.probs();
-	}
-
-	VectorXd expectHidden(VectorXd visible, uint32_t numIter)
-	{
-      	uint32_t i;
-		VisibleLayer v(m_w.getNumVisible(), (const VectorXd*)&visible);
-		HiddenLayer h(m_w.getNumHidden());
-
-    	for (i=0; i < numIter; i++)
-    	{
-    		h.probsUpdateLogistic(v, (Weights&)m_w, m_lambda, m_sigma);
-    		if (m_useVisibleGaussian)
-    		{
-    			v.probsUpdateGaussian(h, (Weights&)m_w, m_lambda, m_sigma);
-    		}
-    		else
-    		{
-    			v.probsUpdateLogistic(h, (Weights&)m_w, m_lambda, m_sigma);
-    		}
-    	}
-
-		return h.probs();
-	}
-
-   VectorXd expectVisible(VectorXd visible, uint32_t numIter)
-	{
-      	uint32_t i;
-		VisibleLayer v(m_w.getNumVisible(), (const VectorXd*)&visible);
-		HiddenLayer h(m_w.getNumHidden());
-
-    	for (i=0; i < numIter; i++)
-    	{
-    		h.probsUpdateLogistic(v, (Weights&)m_w, m_lambda, m_sigma);
-    		if (m_useVisibleGaussian)
-    		{
-    			v.probsUpdateGaussian(h, (Weights&)m_w, m_lambda, m_sigma);
-    		}
-    		else
-    		{
-    			v.probsUpdateLogistic(h, (Weights&)m_w, m_lambda, m_sigma);
-    		}
-    	}
-
-		return v.probs();
+		return v;
 	}
 
    void setSigma(double value)
    {
-	   m_sigma = value;
+		m_sigma = value;
    }
 
    void setSigmaDecay(double value)
@@ -639,6 +475,15 @@ public:
 	   m_doNormalizeData = flag;
    }
 
+   void setDoLearnVariance(bool flag)
+   {
+	   m_doLearnVariance = flag;
+	   if (!flag)
+	   {
+			m_w.sigma().fill(m_sigma);
+	   }
+   }
+
    void setNumGibbs(uint32_t value)
    {
 	   m_numGibbs = value;
@@ -683,6 +528,7 @@ private:
 	bool m_useProbsForHiddenReconstruction;
 	bool m_doSparse;
 	bool m_doNormalizeData;
+	bool m_doLearnVariance;
 	volatile bool m_doCancel;
 	uint32_t m_numGibbs;
 
diff --git a/Source/VisibleLayer.hpp b/Source/VisibleLayer.hpp
index e01607e..2941265 100644
--- a/Source/VisibleLayer.hpp
+++ b/Source/VisibleLayer.hpp
@@ -14,7 +14,7 @@
 class VisibleLayer : public Layer
 {
 public:
-	VisibleLayer(uint32_t numUnits = 0, const VectorXd *pStatesInit = nullptr)
+	VisibleLayer(uint32_t numUnits = 0, const RowVectorXd *pStatesInit = nullptr)
 	: Layer(numUnits, pStatesInit)
 	{
 	}
@@ -28,7 +28,7 @@ private:
 	{
 		double sum = ((Weights&)weights).visibleBias()[index];
 
-		sum += ((Weights&)weights).weights().row(index) * layer.states();
+		sum += ((Weights&)weights).weights().row(index) * layer.states().transpose();
 
 		return sum;
 	}
diff --git a/Source/Weights.hpp b/Source/Weights.hpp
index 08f8c11..f3e916e 100644
--- a/Source/Weights.hpp
+++ b/Source/Weights.hpp
@@ -24,22 +24,28 @@ class Weights
 public:
 	Weights(const char *pFilename)
 	: m_numVisible(0)
+	, m_numVisibleX(0)
+	, m_numVisibleY(0)
 	, m_numHidden(0)
 	{
 		Noise_Init(&m_noise, 0x32727155);
 		load(pFilename);
 	}
 
-	Weights(uint32_t numVisible, uint32_t numHidden)
-	: m_numVisible(numVisible)
-	, m_numHidden(numHidden)
+	Weights(uint32_t numVisibleX, uint32_t numVisibleY, uint32_t numHidden)
+	: m_numVisible(0)
+	, m_numVisibleX(0)
+	, m_numVisibleY(0)
+	, m_numHidden(0)
 	{
 		Noise_Init(&m_noise, 0x32727155);
-		alloc(numVisible, numHidden);
+		setUnits(numVisibleX, numVisibleY, numHidden);
 	}
 
 	Weights()
 	: m_numVisible(0)
+	, m_numVisibleX(0)
+	, m_numVisibleY(0)
 	, m_numHidden(0)
 	{
 		Noise_Init(&m_noise, 0x32727155);
@@ -47,10 +53,12 @@ public:
 
 	Weights(const Weights &src)
 	: m_numVisible(0)
+	, m_numVisibleX(0)
+	, m_numVisibleY(0)
 	, m_numHidden(0)
 	{
 		Noise_Init(&m_noise, 0x32727155);
-		alloc(src.m_numVisible, src.m_numHidden);
+		setUnits(src.m_numVisibleX, src.m_numVisibleY, src.m_numHidden);
 		*this = src;
 	}
 
@@ -60,9 +68,23 @@ public:
 		free();
 	}
 
-	void setUnits(uint32_t numVisible, uint32_t numHidden)
+	void setUnits(uint32_t numVisibleX, uint32_t numVisibleY, uint32_t numHidden)
 	{
-		alloc(numVisible, numHidden);
+		shuffle(0);
+		if ((m_numVisibleX == numVisibleX) && (m_numVisibleY == numVisibleY) && (m_numHidden == numHidden))
+		{
+			return;
+		}
+		m_numVisibleX = numVisibleX;
+		m_numVisibleY = numVisibleY;
+		m_numVisible = numVisibleX * numVisibleY;
+		m_numHidden = numHidden;
+
+		m_w.resize(m_numVisible, m_numHidden);
+		m_sigma.resize(m_numVisible);
+		m_mean.resize(m_numVisible);
+		m_bv.resize(m_numVisible);
+		m_bh.resize(m_numHidden);
 	}
 
 	void shuffle(double stdDev)
@@ -70,6 +92,16 @@ public:
 		uint32_t i, j;
 		double kdev = stdDev*sqrt(12.0);
 
+		for (i=0; i < m_numVisible; i++)
+		{
+			m_sigma(i) = 1; //kdev*Noise_Uniform(&m_noise);
+		}
+
+		for (i=0; i < m_numVisible; i++)
+		{
+			m_mean(i) = 0; //kdev*Noise_Uniform(&m_noise);
+		}
+
 		for (i=0; i < m_numVisible; i++)
 		{
 			m_bv(i) = 0; //kdev*Noise_Uniform(&m_noise);
@@ -94,6 +126,8 @@ public:
 		m_bv = rhs.m_bv;
 		m_bh = rhs.m_bh;
 		m_w = rhs.m_w;
+		m_sigma = rhs.m_sigma;
+		m_mean = rhs.m_mean;
 
 		return *this;
 	}
@@ -103,12 +137,22 @@ public:
 		return m_w;
 	}
 
-	VectorXd& visibleBias()
+	RowVectorXd& visibleBias()
 	{
 		return m_bv;
 	}
 
-	VectorXd& hiddenBias()
+	RowVectorXd& sigma()
+	{
+		return m_sigma;
+	}
+
+	RowVectorXd& mean()
+	{
+		return m_mean;
+	}
+
+	RowVectorXd& hiddenBias()
 	{
 		return m_bh;
 	}
@@ -151,6 +195,16 @@ public:
 		return m_numVisible;
 	}
 
+	uint32_t getNumVisibleX()
+	{
+		return m_numVisibleX;
+	}
+
+	uint32_t getNumVisibleY()
+	{
+		return m_numVisibleY;
+	}
+
 	uint32_t getNumHidden()
 	{
 		return m_numHidden;
@@ -166,7 +220,7 @@ public:
 			return;
 
 
-		fprintf(pFile, "%d %d\n", m_numVisible, m_numHidden);
+		fprintf(pFile, "%d %d %d\n", m_numVisibleX, m_numVisibleY, m_numHidden);
 
 		uint32_t i, j;
 
@@ -192,7 +246,8 @@ public:
 
 	void load(const char *pFilename)
 	{
-		uint32_t numVisible;
+		uint32_t numVisibleX;
+		uint32_t numVisibleY;
 		uint32_t numHidden;
 		FILE *pFile;
 
@@ -201,25 +256,25 @@ public:
 		if (!pFile)
 			return;
 
-		fscanf(pFile, "%d %d\n", &numVisible, &numHidden);
+		fscanf(pFile, "%d %d %d\n", &numVisibleX, &numVisibleY, &numHidden);
 
-		alloc(numVisible, numHidden);
+		setUnits(numVisibleX, numVisibleY, numHidden);
 
 		uint32_t i, j;
 		float v;
-		for (i=0; i < numVisible; i++)
+		for (i=0; i < m_numVisible; i++)
 		{
 			fscanf(pFile, "%f", &v);
 			m_bv(i) = v;
 		}
-		for (i=0; i < numHidden; i++)
+		for (i=0; i < m_numHidden; i++)
 		{
 			fscanf(pFile, "%f", &v);
 			m_bh(i) = v;
 		}
-		for (i=0; i < numVisible; i++)
+		for (i=0; i < m_numVisible; i++)
 		{
-			for (j=0; j < numHidden; j++)
+			for (j=0; j < m_numHidden; j++)
 			{
 
 				fscanf(pFile, "%f", &v);
@@ -232,27 +287,15 @@ public:
 
 private:
 	uint32_t m_numVisible;
+	uint32_t m_numVisibleX;
+	uint32_t m_numVisibleY;
 	uint32_t m_numHidden;
 	noise_gen_t m_noise;
 	MatrixXd m_w;
-	VectorXd m_bv;
-	VectorXd m_bh;
-
-	void alloc(uint32_t numVisible, uint32_t numHidden)
-	{
-		if ((m_numVisible == numVisible) && (m_numHidden == numHidden))
-		{
-				return;
-		}
-		m_numVisible = numVisible;
-		m_numHidden = numHidden;
-
-		m_w.resize(numVisible, numHidden);
-		m_bv.resize(numVisible);
-		m_bh.resize(numHidden);
-
-		shuffle(0);
-	}
+	RowVectorXd m_bv;
+	RowVectorXd m_bh;
+	RowVectorXd m_sigma;
+	RowVectorXd m_mean;
 
 	void free()
 	{