Improving coverage for shadow plots

psdr/subspace.py

@@ -2,6 +2,7 @@
 from __future__ import division, print_function
 
 import os
+# Necessary for running in headless enviornoments
 if 'DISPLAY' not in os.environ:
 	import matplotlib
 	matplotlib.use("Agg")
@@ -76,15 +77,18 @@ def shadow_plot(self, X = None, fX = None, dim = 1, U = None, ax = 'auto', pgfna
 
 		if X is None:
 			X = self.X
+
+		# Check dimensions
+		X = np.atleast_2d(X)
+		assert X.shape[1] == len(self), "Samples do not match dimension of space"	
 
 		if U is None:
 			U = self.U
 		else:
 			if len(U.shape) == 1:
-				U = U.reshape(X.shape[1],1)
+				U = U.reshape(len(self),1)
 			else:
-				assert U.shape[1] == X.shape[1], "Dimensions do not match"
-				U = U
+				assert U.shape[0] == len(self), "Dimensions do not match"
 
 
 		if dim == 1:
@@ -101,6 +105,7 @@ def shadow_plot(self, X = None, fX = None, dim = 1, U = None, ax = 'auto', pgfna
 
 		elif dim == 2:
 			Y = U[:,0:2].T.dot(X.T).T
+
 			if ax is not None:
 				sc = ax.scatter(Y[:,0], Y[:,1], c = fX.flatten(), s = 3)
 				ax.set_xlabel(r'active coordinate 1 $\mathbf{u}_1^\top \mathbf{x}$')
@@ -109,7 +114,11 @@ def shadow_plot(self, X = None, fX = None, dim = 1, U = None, ax = 'auto', pgfna
 				plt.colorbar(sc).set_label('f(x)')
 
 			if pgfname is not None:
-				raise NotImplementedError
+				pgf = PGF()
+				pgf.add('y1', Y[:,0])
+				pgf.add('y2', Y[:,1])
+				pgf.add('fX', fX.flatten())
+				pgf.write(pgfname)
 
 		else:
 			raise NotImplementedError		

tests/test_plot.py

@@ -42,6 +42,7 @@ def test_shadow():
 	pra2.fit(X, fX)
 	pra2.shadow_plot(X, fX, ax = ax2)
 
+
 def test_shadow_lipschitz():
 	fun = psdr.demos.OTLCircuit()
 	X = fun.domain.sample_grid(2)
@@ -55,6 +56,24 @@ def test_shadow_lipschitz():
 	lip.shadow_uncertainty(fun.domain, X, fX, ax = ax, ngrid = 4, pgfname = 'test_shadow_uncertainty.dat')
 	#assert filecmp.cmp(os.path.join(path, 'data/test_shadow_uncertainty.dat'), 'test_shadow_uncertainty.dat') 
 
+	# Test 2-d version
+	ax = lip.shadow_plot(X, fX, dim = 2)
+
+	# Test predefined axes
+	fig, ax = plt.subplots()
+	ax = lip.shadow_plot(X, fX, dim = 1, ax = ax)
+
+	# Test specified U
+	ax = lip.shadow_plot(X, fX, U = lip.U[:,0])
+
+	# Test specified U
+	ax = lip.shadow_plot(X, fX, U = lip.U[:,0:2])
+
+	# Test specified U
+	ax = lip.shadow_plot(X, fX, U = lip.U[:,0:2], dim = 2)
+
+	# Test writing 2-D output
+	lip.shadow_plot(X, fX, ax = None, pgfname = 'test_shadow_uncertainty_2d.dat')
 
 def test_score():
 	fun = psdr.demos.OTLCircuit()