diff --git a/UncertainSCI/distributions.py b/UncertainSCI/distributions.py
index 2e97b47..ecac0a7 100644
--- a/UncertainSCI/distributions.py
+++ b/UncertainSCI/distributions.py
@@ -281,7 +281,7 @@ def __init__(self, flag=True, lbd=None, loc=None, mean=None,
         self.transform_standard_dist_to_poly = AffineTransform(A=A, b=b)
 
         if self.dim == 1:
-            A = lbd
+            A = np.array(lbd)
             b = -A * loc
             self.transform_to_standard = AffineTransform(A=A, b=b)
 
diff --git a/UncertainSCI/pce.py b/UncertainSCI/pce.py
index 20e8caf..38dfa91 100644
--- a/UncertainSCI/pce.py
+++ b/UncertainSCI/pce.py
@@ -143,29 +143,6 @@ def set_samples(self, samples):
                                  'have wrong dimension')
 
         self.samples = samples
-        self.set_weights()
-
-    def set_weights(self):
-        """Sets weights based on assigned samples.
-        """
-        if self.samples is None:
-            raise RuntimeError("PCE weights cannot be set unless samples are set first.""")
-        
-        if self.sampling.lower() == 'greedy-induced':
-            self.weights = self.christoffel_weights()
-        elif self.sampling.lower() == 'gq':
-            M = self.sampling_options.get('M')
-            if M is None:
-                raise ValueError("The sampling option 'M' must be specified for Gauss quadrature sampling.")
-
-            _, self.weights = self.distribution.polys.tensor_gauss_quadrature(M)
-
-        elif self.sampling.lower() == 'gq-induced':
-            self.weights = self.christoffel_weights()
-
-        else:
-            raise ValueError("Unsupported sample type '{0}' for input\
-                              sample_type".format(self.sampling))
 
     def map_to_standard_space(self, q):
         """Maps parameter values from model space to standard space.
@@ -234,6 +211,8 @@ def generate_samples(self, **kwargs):
 
                 self.samples = self.map_to_model_space(x)
 
+            self.weights = self.christoffel_weights()
+
         elif self.sampling.lower() == 'gq':
 
             M = self.sampling_options.get('M')
@@ -242,9 +221,7 @@ def generate_samples(self, **kwargs):
 
             p_standard, w = self.distribution.polys.tensor_gauss_quadrature(M)
             self.samples = self.map_to_model_space(p_standard)
-            # We do the following in the call to self.set_weights() below. A
-            # little more expensive, but makes for more transparent control structure.
-            #self.weights = w
+            self.weights = w
 
         elif self.sampling.lower() == 'gq-induced':
 
@@ -255,13 +232,12 @@ def generate_samples(self, **kwargs):
             p_standard = self.distribution.opolys.idist_gq_sampling(K, self.indices, M=self.sampling_options.get('M'))
 
             self.samples = self.map_to_model_space(p_standard)
+            self.weights = self.christoffel_weights()
 
         else:
             raise ValueError("Unsupported sample type '{0}' for input\
                               sample_type".format(self.sampling))
 
-        self.set_weights()
-
     def integration_weights(self):
         """
         Generates sample weights associated to integration."
diff --git a/demos/anisotropic_distribution.py b/demos/anisotropic_distribution.py
index 1c320ae..1f82094 100644
--- a/demos/anisotropic_distribution.py
+++ b/demos/anisotropic_distribution.py
@@ -34,7 +34,7 @@
 model = sine_modulation(N=N)
 
 # Compute PCE (runs model)
-lsq_residuals = pce.build_pce_wafp(model)
+lsq_residuals = pce.build(model, oversampling=10)
 
 Q = 6  # Number of quantile bands to plot
 
diff --git a/demos/quantiles.py b/demos/quantiles.py
index e5f74dc..76bc420 100644
--- a/demos/quantiles.py
+++ b/demos/quantiles.py
@@ -8,6 +8,8 @@
 from UncertainSCI.indexing import TotalDegreeSet
 from UncertainSCI.pce import PolynomialChaosExpansion
 
+from UncertainSCI.vis import quantile_plot
+
 # Number of parameters
 dimension = 3
 
@@ -35,7 +37,7 @@
 model = sine_modulation(N=N)
 
 # Compute PCE (runs model)
-lsq_residuals = pce.build_pce_wafp(model)
+lsq_residuals = pce.build(model)
 
 Q = 6  # Number of quantile bands to plot
 
@@ -57,16 +59,8 @@
 for j in range(M):
     output[j, :] = model(p_phys[j, :])
 
+quantile_plot(pce, bands=3, xvals=x, xlabel='$x$')
 plt.plot(x, output[:M, :].T, 'k', alpha=0.8, linewidth=0.2)
-plt.plot(x, median, 'b', label='PCE median')
-
-for ind in range(Q):
-    alpha = (Q-ind) * 1/Q - (1/(2*Q))
-    plt.fill_between(x, quantiles_lower[ind, :], quantiles_upper[ind, :],
-                     interpolate=True, facecolor='red', alpha=alpha)
-
-plt.xlabel('x')
 
-plt.legend(loc='lower right')
 
 plt.show()
diff --git a/setup.py b/setup.py
index bece8b0..8769f42 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
     
 setup(
     name='UncertainSCI',
-    version='0.2.0-b0',
+    version='0.2.1-b0',
     author='UncertainSCI Developers',
     author_email='uncertainsci@sci.utah.edu',
     packages=find_packages(),