diff --git a/UncertainSCI/pce.py b/UncertainSCI/pce.py
index 38dfa91..20e8caf 100644
--- a/UncertainSCI/pce.py
+++ b/UncertainSCI/pce.py
@@ -143,6 +143,29 @@ 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.
@@ -211,8 +234,6 @@ 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')
@@ -221,7 +242,9 @@ def generate_samples(self, **kwargs):
 
             p_standard, w = self.distribution.polys.tensor_gauss_quadrature(M)
             self.samples = self.map_to_model_space(p_standard)
-            self.weights = w
+            # 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
 
         elif self.sampling.lower() == 'gq-induced':
 
@@ -232,12 +255,13 @@ 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/build_pce_normal.py b/demos/build_pce_normal.py
index b6eed37..be89b68 100644
--- a/demos/build_pce_normal.py
+++ b/demos/build_pce_normal.py
@@ -5,14 +5,15 @@
 
 from UncertainSCI.distributions import NormalDistribution
 from UncertainSCI.model_examples import sine_modulation
-from UncertainSCI.indexing import TotalDegreeSet
 from UncertainSCI.pce import PolynomialChaosExpansion
-# # Distribution setup
 
+from UncertainSCI.vis import piechart_sensitivity, quantile_plot, mean_stdev_plot
+
+## Distribution setup: generating a correlated multivariate normal distribution
 # Number of parameters
 dimension = 2
 
-# Specifies normal distribution
+# Specifies statistics of normal distribution
 mean = np.ones(dimension)
 cov = np.random.randn(dimension, dimension)
 cov = np.matmul(cov.T, cov)/(4*dimension)  # Some random covariance matrix
@@ -20,178 +21,39 @@
 # Normalize so that parameters with larger index have smaller importance
 D = np.diag(1/np.sqrt(np.diag(cov)) * 1/(np.arange(1, dimension+1)**2))
 cov = D @ cov @ D
-
 dist = NormalDistribution(mean=mean, cov=cov, dim=dimension)
 
-# # Indices setup
-order = 10
-index_set = TotalDegreeSet(dim=dimension, order=order)
-
-print('This will query the model {0:d} times'.format(index_set.get_indices().shape[0] + 10))
-# Why +10? That's the default for PolynomialChaosExpansion.build_pce_wafp
-
-# # Initializes a pce object
-pce = PolynomialChaosExpansion(index_set, dist)
-
-# # Define model
+## Define forward model
 N = int(1e2)  # Number of degrees of freedom of model
 left = -1.
 right = 1.
 x = np.linspace(left, right, N)
 model = sine_modulation(N=N)
 
-# # Three equivalent ways to run the PCE model:
+## Polynomial order
+order = 5
+
+## Initializes a pce object
+pce = PolynomialChaosExpansion(distribution=dist, order=order, plabels=['p1', 'p2'])
+pce.generate_samples()
 
-# 1
-# Generate samples and query model in one call:
-pce = PolynomialChaosExpansion(index_set, dist)
-lsq_residuals = pce.build(model, oversampling=10)
+print('This queries the model {0:d} times'.format(pce.samples.shape[0]))
 
+# Query model:
+pce.build(model)
 
 # The parameter samples and model evaluations are accessible:
 parameter_samples = pce.samples
 model_evaluations = pce.model_output
 
-# And you could build a second PCE on the same parameter samples
-pce2 = PolynomialChaosExpansion(index_set, dist)
+# And if desired you could build a second PCE on the same parameter samples
+pce2 = PolynomialChaosExpansion(distribution=dist, order=order)
 pce2.set_samples(parameter_samples)
 pce2.build(model)
 
-# pce and pce2 have the same coefficients:
-# np.linalg.norm( pce.coefficients - pce2.coefficients )
-
-# # Postprocess PCE: mean, stdev, sensitivities, quantiles
-mean = pce.mean()
-stdev = pce.stdev()
-
-
-# Power set of [0, 1, ..., dimension-1]
-variable_interactions = list(chain.from_iterable(combinations(range(dimension), r) for r in range(1, dimension+1)))
-
-# "Total sensitivity" is a non-partitive relative sensitivity measure per parameter.
-total_sensitivity = pce.total_sensitivity()
-
-# "Global sensitivity" is a partitive relative sensitivity measure per set of parameters.
-global_sensitivity = pce.global_sensitivity(variable_interactions)
-
-
-
-Q = 4  # Number of quantile bands to plot
-dq = 0.5/(Q+1)
-q_lower = np.arange(dq, 0.5-1e-7, dq)[::-1]
-q_upper = np.arange(0.5 + dq, 1.0-1e-7, dq)
-quantile_levels = np.append(np.concatenate((q_lower, q_upper)), 0.5)
-
-quantiles = pce.quantile(quantile_levels, M=int(2e3))
-median = quantiles[-1, :]
-
-# # For comparison: Monte Carlo statistics
-M = 1000  # Generate MC samples
-p_phys = dist.MC_samples(M)
-output = np.zeros([M, N])
-
-for j in range(M):
-    output[j, :] = model(p_phys[j, :])
-
-MC_mean = np.mean(output, axis=0)
-MC_stdev = np.std(output, axis=0)
-MC_quantiles = np.quantile(output, quantile_levels, axis=0)
-MC_median = quantiles[-1, :]
-
-# # Visualization
-V = 50  # Number of MC samples to visualize
-
-# mean +/- stdev plot
-plt.plot(x, output[:V, :].T, 'k', alpha=0.8, linewidth=0.2)
-plt.plot(x, mean, 'b', label='PCE mean')
-plt.fill_between(x, mean-stdev, mean+stdev, interpolate=True, facecolor='red', alpha=0.5, label='PCE 1 stdev range')
-
-plt.plot(x, MC_mean, 'b:', label='MC mean')
-plt.plot(x, MC_mean+MC_stdev, 'r:', label='MC mean $\\pm$ stdev')
-plt.plot(x, MC_mean-MC_stdev, 'r:')
-
-plt.xlabel('x')
-plt.title('Mean $\\pm$ standard deviation')
-
-plt.legend(loc='lower right')
-
-# quantile plot
-plt.figure()
-
-plt.subplot(121)
-plt.plot(x, output[:V, :].T, 'k', alpha=0.8, linewidth=0.2)
-plt.plot(x, median, 'b', label='PCE median')
-
-band_mass = 1/(2*(Q+1))
-
-for ind in range(Q):
-    alpha = (Q-ind) * 1/Q - (1/(2*Q))
-    if ind == 0:
-        plt.fill_between(x, quantiles[ind, :], quantiles[Q+ind, :], interpolate=True, facecolor='red',
-                         alpha=alpha, label='{0:1.2f} probability mass (each band)'.format(band_mass))
-    else:
-        plt.fill_between(x, quantiles[ind, :], quantiles[Q+ind, :], interpolate=True, facecolor='red',
-                         alpha=alpha)
-
-plt.title('PCE Median + quantile bands')
-plt.xlabel('x')
-plt.legend(loc='lower right')
-
-plt.subplot(122)
-plt.plot(x, output[:V, :].T, 'k', alpha=0.8, linewidth=0.2)
-plt.plot(x, MC_median, 'b', label='MC median')
-
-for ind in range(Q):
-    alpha = (Q-ind) * 1/Q - (1/(2*Q))
-    if ind == 0:
-        plt.fill_between(x, MC_quantiles[ind, :], MC_quantiles[Q+ind, :], interpolate=True, facecolor='red',
-                         alpha=alpha, label='{0:1.2f} probability mass (each band)'.format(band_mass))
-    else:
-        plt.fill_between(x, MC_quantiles[ind, :], MC_quantiles[Q+ind, :], interpolate=True, facecolor='red', alpha=alpha)
-
-plt.title('MC Median + quantile bands')
-plt.xlabel('x')
-plt.legend(loc='lower right')
-
-# Sensitivity pie chart, averaged over all model degrees of freedom
-average_global_SI = np.sum(global_sensitivity, axis=1)/N
-
-labels = ['[' + ' '.join(str(elem) for elem in [i+1 for i in item]) + ']' for item in variable_interactions]
-_, ax = plt.subplots()
-ax.pie(average_global_SI*100, labels=labels, autopct='%1.1f%%',
-       startangle=90)
-ax.axis('equal')  # Equal aspect ratio ensures that pie is drawn as a circle.
-plt.title('Sensitivity due to variable interactions')
-
-
-
-
-
-sensitivities = [total_sensitivity, global_sensitivity]
-
-sensbounds = [[0, 1], [0, 1]]
-
-senslabels = ['Total sensitivity', 'Global sensitivity']
-dimlabels = ['Dimension 1', 'Dimension 2', 'Interaction']
-
-fig, ax = plt.subplots(2, 3)
-for row in range(2):
-    for col in range(2):
-        ax[row][col].plot(x, sensitivities[row][col,:])
-        ax[row][col].set_ylim(sensbounds[row])
-        if row==2:
-            ax[row][col].set(xlabel='$x$')
-        if col==0:
-            ax[row][col].set(ylabel=senslabels[row])
-        if row==0:
-            ax[row][col].set_title(dimlabels[col])
-        if row<2:
-            ax[row][col].xaxis.set_ticks([])
-        if col>0:
-            ax[row][col].yaxis.set_ticks([])
-  
-ax[1][2].plot(x, sensitivities[1][2,:])
-ax[1][2].set_ylim(sensbounds[1])
-
+## Visualization
+mean_stdev_plot(pce, ensemble=50)
+quantile_plot(pce, bands=3, xvals=x, xlabel='$x$')
+piechart_sensitivity(pce)
 
 plt.show()