text: import numpy as np
  import scipy as sp
  from scipy import sparse as sprs
  
  from UncertainSCI.families import JacobiPolynomials, \
          HermitePolynomials, LaguerrePolynomials
  from UncertainSCI.families import DiscreteChebyshevPolynomials
  from UncertainSCI.opolynd import TensorialPolynomials
  from UncertainSCI.transformations import AffineTransform
  from UncertainSCI.utils.version import version_lessthan
  
  # numpy >= 1.17: default_rng is preferred
  if version_lessthan(np, '1.17'):
      from numpy.random import choice
  else:
      from numpy.random import default_rng
      choice = default_rng().choice
  
  
  class ProbabilityDistribution:
      def __init__(self):
          pass
  
  
  class NormalDistribution(ProbabilityDistribution):
      def __init__(self, mean=None, cov=None, dim=None):
  
          _mean, _cov = self._convert_meancov_to_iterable(mean, cov)
  
          self._detect_dimension(_mean, _cov, dim)
  
          assert self.dim > 0, "Dimension must be positive"
  
          # Construct affine map transformations
  
          # Low-level routines use Hermite Polynomials
          # with weight function exp(-x**2)
          # instead of standard normal weight function exp(-x**2/2)
          # I.e. x ----> sqrt(2) * x
  
          A = np.eye(self.dim) * (1/np.sqrt(2))
          b = np.zeros(self.dim)
          self.transform_standard_dist_to_poly = AffineTransform(A=A, b=b)
  
          # user says: X ~ exp( -(x - mu).T cov^{-1} (x - mu)/2 )
          # want: Z ~ exp( -(x - 0).T eye (x - 0)/2 )
          # I.e. X = KZ + mu, Z = K^{-1}(X - mu), cov = KK.T, K = cov^{1/2}
          # I.e. x ---> cov^{-1/2} * (x + mu) = cov^{-1/2} *x + cov^{-1/2}*mu
          # Opt1: cov = U Lamb U.T, U unitary
          #       cov^{1/2} = U sqrt(Lamb) U.T
          #       cov^{-1/2} = U sqrt(1/Lamb) U.T
          #       A = U sqrt(1/Lamb) U.T and b = A * mu
          # Opt2: cov = L L.T
          #       cov^{1/2} = L
          #       cov^{-1/2} = inv(L)
          #       A = inv(L) and b = A * mu
  
          if self.dim == 1:
              sigma = np.sqrt(self._cov)
              A = np.eye(self.dim) * (1/sigma)
              b = np.ones(self.dim) * (-self._mean/sigma)
              self.transform_to_standard = AffineTransform(A=A, b=b)
  
          else:
              # Option 1: Cholesky
              #L = np.linalg.cholesky(self._cov)
  
              # Option 2: matrix square root
              W, V = np.linalg.eigh(self._cov)
              L = V.T @ np.sqrt(np.diag(W))
  
              A = np.linalg.inv(L)
              b = -A.dot(self._mean)
              self.transform_to_standard = AffineTransform(A=A, b=b)
  
          # Construct 1D polynomial families
          Hs = []
          for qd in range(self.dim):
              Hs.append(HermitePolynomials())  # modify for different mean, cov?
          self.polys = TensorialPolynomials(polys1d=Hs)
  
          # Standard domain is R^dim
          self.standard_domain = np.zeros([2, self.dim])
          self.standard_domain[0, :] = -np.inf
          self.standard_domain[1, :] = np.inf
  
          self.poly_domain = self.standard_domain.copy()
  
          self.indices = None
  
      def _convert_meancov_to_iterable(self, _mean, _cov):
          """
          Converts user-input (mean, cov) to iterables. Ensures that the length
          of the iterables matches on output.
  
          If mean is None, sets it to 0.
          If cov is None, sets it to identity matrix
          """
  
          # Tons of type/value checking for mean/cov vs dim
          meaniter = isinstance(_mean, (list, tuple, np.ndarray))
  
          if _mean is None:
              _mean = 0.
  
          if _cov is None:
              if meaniter:
                  _cov = np.eye(len(_mean))
  #            elif _mean is None:
  #                _mean = [0.]
  #                _cov = np.eye(1)
              else:  # mean is a scalar
                  _mean = [_mean, ]
                  _cov = np.eye(1)
  
          else:
              assert isinstance(_cov, np.ndarray), 'Covariance must be an array'
              # assert np.all(cov - cov.T == 0), 'Covariance must be symmetric'
  
              if meaniter:
                  if len(_mean) > 1 and _cov.shape[0] > 1:
                      assert len(_mean) == _cov.shape[0], "Mean and cov parameter \
                              inputs must be of the same size"
                      try:
                          np.linalg.cholesky(_cov)
                      except ValueError:
                          print('Covariance must be symmetric and \
                                  positive definite')
  
  #                elif len(_mean) == 1 and _cov.shape[0] == 1:
  #                    pass
  
                  elif len(_mean) == 1 and _cov.shape[0] > 1:
                      _mean = [_mean[0] for i in range(_cov.shape[0])]
  
                  elif _cov.shape[0] == 1 and len(_mean) > 1:
                      _cov = np.eye(len(_mean)) * _cov[0]
  
  #            elif _mean is None:
  #                _mean = [0. for i in range(_cov.shape[0])]
  #
              else:  # mean is a scalar
                  _mean = [_mean for i in range(_cov.shape[0])]
  
          return _mean, _cov
  
      def _detect_dimension(self, _mean, _cov, dim):
          """
          Parses user-given inputs to determine \
                  the dimension of the distribution.
  
          Mean and cov are iterables.
  
          Sets self.dim, self.mean, self.cov
          """
          # Situations:
          # 1. User specifies mean list and cov ndarray (disallow contradictory
          #    dimension specification)
          # 2. User specifies dim scalar (disallow contradictory mean, stdev and
          #    cov specification)
          # 3. dim = 1
  
          if len(_mean) > 1 or _cov.shape[0] > 1:  # Case 1:
              if len(_mean) != _cov.shape[0]:
                  raise ValueError('Input parameters mean and cov must \
                          be the same dimension')
  
              if (dim is not None) and (dim != 1) and (dim != len(_mean)):
                  raise ValueError('Mean parameter list must \
                          have size consistent with input dim')
  
              if (dim is not None) and (dim != 1) and (dim != _cov.shape[0]):
                  raise ValueError('Cov parameter array must \
                          have size consistance with input dim')
  
              self.dim = len(_mean)
              self._mean = _mean
              self._cov = _cov
  
          elif dim is not None and dim > 1:  # Case 2
              self.dim = dim
              self._mean = [_mean[0] for i in range(self.dim)]
              self._cov = np.eye(self.dim) * _cov[0]
  
          else:  # Case 3
              self.dim = 1
              self._mean = _mean[0]
              self._cov = _cov[0]
  
      def MC_samples(self, M=100):
          """
          Returns M Monte Carlo samples from the distribution.
          """
          p = np.random.normal(0., 1., [M, self.dim])
  
          return self.transform_to_standard.mapinv(p)
  
      def mean(self):
          return self._mean
  
      def cov(self):
          return self._cov
  
      # def means(self):
          # """
          # Returns the mean of the distribution
          # """
  #
          # mu = np.zeros(self.dim, )
          # mu = np.reshape(mu, [1, self.dim])
          # return self.transform_to_standard.mapinv(mu).flatten()
  #
      # def covariance(self):
          # """
          # Returns the (auto-)covariance matrix of the distribution.
          # """
  #
          # sigma = np.eye(self.dim,)
          # zero = np.zeros([1, self.dim])
          # b = self.transform_to_standard.mapinv(zero)
          # bmat = np.tile(b, [self.dim, 1])
          # sigma = self.transform_to_standard.mapinv(sigma) - bmat
          # sigma = (self.transform_to_standard.mapinv(sigma.T) - bmat).T
  #
          # return sigma
  
      def stdev(self):
          """
          Returns the standard deviation of the distribution, if the
          distribution is one-dimensional. Raises an error if called
          for a multivariate distribution.
          """
  
          if self.dim == 1:
              return np.sqrt(self.cov()[0, 0])
          else:
              raise TypeError("Can only compute standard deviations for scalar\
                               random variables.")
  
      def pdf(self, x):
          """
          Evaluates the probability density function (pdf) of the distribution at
          the input locations x.
          """
          # density = multivariate_normal.pdf(x, mean=self.mean, cov=self.cov)
          x = self.transform_to_standard.map(x)
          density = np.ones(x.shape[0])
          for i in range(self.dim):
              density *= (2*np.pi)**(-1/2) * np.exp(-1/2 * x[:, i]**2)
          density *= self.transform_to_standard.jacobian_determinant()
          return density
  
  
  class ExponentialDistribution(ProbabilityDistribution):
  
      def __init__(self, flag=True, lbd=None, loc=None, mean=None,
                   stdev=None, dim=None):
  
          # Convert mean/stdev inputs to lbd
          if mean is not None and stdev is not None:
              if lbd is not None:
                  raise ValueError('Cannot simultaneously specify lbd \
                          parameter and mean/stdev parameters')
  
              lbd = self._convert_meanloc_to_lbd(mean, loc)
  
          lbd, loc = self._convert_lbdloc_to_iterable(lbd, loc)
  
          # Set self.lbd, self.loc and self.dim
          self._detect_dimension(lbd, loc, dim)
  
          assert self.dim > 0, "Dimension must be positive"
  
          # Construct affine map transformations
  
          # Low-level routines use Laguerre Polynomials, weight
          # x^rho exp^{-x} when rho = 0, which is equal to the standard Beta,
          # exp^{-lbd x} when lbd = 1
          A = np.eye(self.dim)
          b = np.zeros(self.dim)
          self.transform_standard_dist_to_poly = AffineTransform(A=A, b=b)
  
          if self.dim == 1:
              A = lbd
              b = -A * loc
              self.transform_to_standard = AffineTransform(A=A, b=b)
  
          else:
              A = np.diag(lbd)
              b = -A@loc
              self.transform_to_standard = AffineTransform(A=A, b=b)
  
          # Construct 1D polynomial families
          Ls = []
          for qd in range(self.dim):
              Ls.append(LaguerrePolynomials())
          self.polys = TensorialPolynomials(polys1d=Ls)
  
          self.indices = None
  
      def _detect_dimension(self, lbd, loc, dim):
          """
          Parses user-given inputs to determine the
          dimension of the distribution.
  
          lbd and loc are iterables.
  
          Sets self.lbd, self.loc and self.dim
          """
  
          # Situations:
          # 1. User specifies lbd as a list (disallow contradictory
          #    dimension specification)
          # 2. User specifies lbd as a list, then dim = len(lbd)
          # 3. User specifies dim as a scalar, then lbd = np.ones(dim)
          # 4. User specifies nothing, then dim = 1, lbd = 1
  
          if len(lbd) > 1 or len(loc) > 1:
              if len(lbd) != len(loc):
                  raise ValueError('Input parameters lbd and loc must \
                          be the same dimension')
  
              if (dim is not None) and (dim != 1) and (dim != len(lbd)):
                  raise ValueError('Lbd, loc parameter lists must \
                          have size consistent with input dim')
  
              self.dim = len(lbd)
              self.lbd = lbd
              self.loc = loc
  
          elif dim is not None and dim > 1:
              self.dim = dim
              self.lbd = [lbd[0] for i in range(self.dim)]
              self.loc = [loc[0] for i in range(self.dim)]
  
          else:  # The dimension is 1
              self.dim = 1
              self.lbd = lbd
              self.loc = loc
  
      def _convert_meanloc_to_lbd(self, mean, loc):
          """
          Converts user-given mean and loc to an iterable lbd.
          """
  
          meaniter = isinstance(mean, (list, tuple, np.ndarray))
          lociter = isinstance(loc, (list, tuple, np.ndarray))
          lbd = []
  
          # If they're both iterables:
          if meaniter and lociter:
              assert len(mean) == len(loc)
  
              for ind in range(len(mean)):
                  lb = self.meanloc_to_lbd(mean[ind], loc[ind])
                  lbd.append(lb)
  
          # If mean is an iterable but loc is not
          elif meaniter:
              loc = [loc for i in range(len(mean))]
              for ind in range(len(mean)):
                  lb = self.meanloc_to_lbd(mean[ind], loc[ind])
                  lbd.append(lb)
  
          # If loc is an iterable but mean is not
          elif lociter:
              mean = [mean for i in range(len(loc))]
              for ind in range(len(loc)):
                  lb = self.meanloc_to_lbd(mean[ind], loc[ind])
                  lbd.append(lb)
  
          # If they're both scalars, let the following lbd checker cf vs dim
          else:
              lbd = self.meanloc_to_lbd(mean, loc)
  
          return lbd
  
      def _convert_lbdloc_to_iterable(self, lbd, loc):
          """
          Converts user-input lbd and loc to iterables. Ensures that the length
          of the iterables matches on output.
  
          If lbd is None, sets it to 1.
          If loc is None, sets it to 0.
          """
  
          # Tons of type/value checking for lbd/loc vs dim
          lbditer = isinstance(lbd, (list, tuple, np.ndarray))
          lociter = isinstance(loc, (list, tuple, np.ndarray))
  
          if lbditer and lociter:
              if len(lbd) > 1 and len(loc) > 1:
                  assert len(lbd) == len(loc), "Lbd and loc parameter inputs \
                          must be of the same size"
  #            elif len(lbd) == 1 and len(loc) == 1:
  #                pass
              elif len(lbd) == 1:
                  lbd = [lbd[0] for i in range(len(loc))]
              elif len(loc) == 1:
                  loc = [loc[0] for i in range(len(lbd))]
  
          elif lbditer:  # lbd is iterable, loc is not
              if loc is None:
                  loc = 0.
              loc = [loc for i in range(len(lbd))]
  
          elif lociter:  # beta is iterable, alpha is not
              if lbd is None:
                  lbd = 1.
              lbd = [lbd for i in range(len(loc))]
  
          elif lbd is None and loc is None:
              lbd, loc = [1., ], [0., ]
  
          else:  # alpha, beta should both be scalars
              lbd, loc = [lbd, ], [loc, ]
  
          return lbd, loc
  
      def meanloc_to_lbd(self, mu, loc):
          """
          Returns lbd given an input mean (mu) and location (loc)
          for a Exponential distribution
          """
  
          return 1 / (mu - loc)
  
      def MC_samples(self, M=100):
          """
          Returns M Monte Carlo samples from the distribution.
          """
  
          p = np.zeros([M, self.dim])
          for qd in range(self.dim):
              p[:, qd] = np.random.exponential(1, M)
  
          return self.transform_to_standard.mapinv(p)
  
      def mean(self):
          mu = np.zeros(self.dim)
          for i in range(self.dim):
              mu[i] = 1.
  
          mu = np.reshape(mu, [1, self.dim])
          return self.transform_to_standard.mapinv(mu).flatten()
  
      def cov(self):
          sigma = np.eye(self.dim)
          for i in range(self.dim):
              sigma[i, i] = 1.
  
          zero = np.zeros([1, self.dim])
          b = self.transform_to_standard.mapinv(zero)
          bmat = np.tile(b, [self.dim, 1])
          sigma = self.transform_to_standard.mapinv(sigma) - bmat
          sigma = (self.transform_to_standard.mapinv(sigma.T) - bmat).T
  
          return sigma
  
      def stdev(self):
          """
          Returns the standard deviation of the distribution, if the
          distribution is one-dimensional. Raises an error if called
          for a multivariate distribution.
          """
  
          if self.dim == 1:
              return np.sqrt(self.cov()[0, 0])
          else:
              raise TypeError("Can only compute standard deviations for \
                      scalar random variables.")
  
      def pdf(self, x):
          x = self.transform_to_standard.map(x)
          density = np.ones(x.shape[0])
          for i in range(self.dim):
              density *= np.exp(-x[:, i])
          # Scale based on determinant of map
          density *= self.transform_to_standard.jacobian_determinant()
  
          return density
  
  
  class GammaDistribution(ProbabilityDistribution):
      """.. _gamma_distribution:
  
      Constructs a Gamma distribution object. Gamma distributions have support on
      the real interval [0,infty), with probability density function,
  
      .. math::
  
        w(y;k,\\theta) := \\frac{1}{\\Gamma(k) \\theta^k}\
                y^{k-1} exp(-y/\\theta), \\hskip 20pt y \\in (0,\\infty),
  
      where :math:`k` and :math:`\\theta` are positive real parameters that
      define the distribution, and :math:`\\Gamma` is the Gamma function. Some
      special cases of note:
  
      To generate this distribution on a general shifted interval
      :math:`(shift,\\infty)`, set the shift parameter below.
  
      To generate this distribution so it has support on the negative half-line
      :math:`(-\\infty,0)`, set the flip parameter below.
  
      Parameters:
          k (float, optional): Shape parameter. Defaults to 1.
          theta (float, optional): Scale parameter. Defaults to 1.
          shift (float, optional): Shift parameter. Defaults to 0.
          flip (bool, optional): Flip parameter. Defaults to False.
  
      Attributes:
          k (float): Shape parameter k.
          theta (float): Scale parameter theta.
          polys (:class:`LaguerrePolynomials`): Orthogonal polynomials for this
              distribution.
      """
  
      def __init__(self, k=1., theta=1., shift=0., flip=False):
  
          self.theta, self.k, self.shift, self.flip = theta, k, shift, flip
          self.dim = 1
  
          # Construct affine map transformations
  
          # Low-level routines use Laguerre Polynomials, weight
          # x^rho exp^{-x} when rho = 0, which is equal to the standard
          # Exponential,
          # exp^{-lbd x} when lbd = 1
          A = np.eye(self.dim)
          b = np.zeros(self.dim)
          self.transform_standard_dist_to_poly = AffineTransform(A=A, b=b)
  
          if not flip:
              A = 1/self.theta
          else:
              A = -1/self.theta
  
          b = -A*self.shift
          self.transform_to_standard = AffineTransform(A=A, b=b)
  
          # Construct 1D polynomial families
          Ls = []
          for qd in range(self.dim):
              Ls.append(LaguerrePolynomials())
          self.polys = TensorialPolynomials(polys1d=Ls)
  
          self.indices = None
  
      def MC_samples(self, M=100):
          """
          Returns M Monte Carlo samples from the distribution.
          """
  
          x = sp.stats.gamma.rvs(self.k, loc=0., scale=self.theta, size=M)
          if self.flip:
              x = -x
  
          return x + self.shift
  
      def mean(self):
          mu = self.k
          return self.transform_to_standard.mapinv(mu)
  
      def stdev(self):
          """
          Returns the standard deviation of the distribution.
          """
  
          return np.sqrt(self.k)*self.theta
  
      def pdf(self, x):
          """
          Evaluates the probability distribution function at the locations x.
          """
  
          z = self.transform_to_standard.map(x)
          density = 1/sp.special.gamma(self.k) * z**(self.k-1) * np.exp(-z)
          return density*self.transform_to_standard.A
  
  
  class BetaDistribution(ProbabilityDistribution):
      """.. _beta_distribution:
  
      Constructs a Beta distribution object; supports multivariate distributions
      through tensorization. In one dimension, beta distributions have support on
      the real interval [0,1], with probability density function,
  
      .. math::
  
        w(y;\\alpha,\\beta) := \\frac{y^{\\alpha-1} (1-y)^{\\beta-1}}\
                {B(\\alpha,\\beta)}, \\hskip 20pt y \\in (0,1),
  
      where :math:`\\alpha` and :math:`\\beta` are positive real parameters that
      define the distribution, and :math:`B` is the Beta function. Some special
      cases of note:
  
      - :math:`\\alpha = \\beta = 1`: the uniform distribution
      - :math:`\\alpha = \\beta = \\frac{1}{2}`: the arcsine distribution
  
      To generate this distribution on a general compact interval :math:`[a,b]`,
      set the domain parameter below.
  
      Instead of :math:`(\\alpha, \\beta)`, a mean :math:`\\mu` and standard
      deviation :math:``\\sigma`` may be set. In this case, :math:`(\\mu,
      \\sigma)` must correspond to valid (i.e., positive) values of
      :math:`(\\alpha, \\beta)` on the interval [0,1], or else an error is
      raised.
  
      Finally, this class supports tensorization: multidimensional distributions
      corresopnding to independent one-dimensional marginal distributions are
      supported. In the case of identically distributed marginals, the `dim`
      parameter can be set to the appropriate dimension. In case of non-identical
      marginals, an array or iterable can be input for :math:`\\alpha, \\beta,
      \\mu, \\sigma`.
  
      Parameters:
          alpha (float or iterable of floats, optional): Shape parameter
          associated to right-hand boundary. Defaults to 1.
          beta (float or iterable of floats, optional): Shape parameter
          associated to left-hand boundary. Defaults to 1.
          mean (float or iterable of floats, optional): Mean of the distribution.
          Defaults to None.
          stdev (float or iterable of floats, optional): Standard deviation of
          the distribution. Defaults to None.
          dim (int, optional): Dimension of the distribution. Defaults to None.
          domain (numpy.ndarray or similar, of size 2 x `dim`, optional): Compact
          hypercube that is the support of the distribution. Defaults to None
  
      Attributes:
          dim (int): Dimension of the distribution.
          alpha (float or np.ndarray): Shape parameter(s) alpha.
          beta (float or np.ndarray): Shape parameter(s) beta.
          polys (:class:`JacobiPolynomials` or list thereof):
      """
      def __init__(self, alpha=None, beta=None, mean=None, stdev=None,
                   dim=None, domain=None, bounds=None):
  
          # Convert mean/stdev inputs to alpha/beta
          if mean is not None and stdev is not None:
              if alpha is not None or beta is not None:
                  raise ValueError('Cannot simultaneously specify alpha/beta \
                          parameters and mean/stdev parameters')
  
              alpha, beta = self._convert_meanstdev_to_alphabeta(mean, stdev)
  
          alpha, beta = self._convert_alphabeta_to_iterable(alpha, beta)
  
          if (domain is not None) and (bounds is not None):
              raise ValueError("Inputs 'domain' and 'bounds' cannot both be given")
          elif bounds is not None:
              domain = bounds
          # Sets self.dim, self.alpha, self.beta, self.domain
          self._detect_dimension(alpha, beta, dim, domain)
  
          for qd in range(self.dim):
              assert self.alpha[qd] > 0 and self.beta[qd] > 0, \
                      "Parameter vectors alpha and beta must have strictly \
                      positive components"
          assert self.dim > 0, "Dimension must be positive"
          assert self.domain.shape == (2, self.dim)
  
          # Construct affine map transformations
  
          # Standard domain is [0,1]^dim
          self.standard_domain = np.ones([2, self.dim])
          self.standard_domain[0, :] = 0.
  
          # Low-level routines use Jacobi Polynomials, which operate on [-1,1]
          # instead of the standard Beta domain of [0,1]
          self.poly_domain = np.ones([2, self.dim])
          self.poly_domain[0, :] = -1.
          self.transform_standard_dist_to_poly = AffineTransform(
                  domain=self.standard_domain, image=self.poly_domain)
  
          self.transform_to_standard = AffineTransform(
                  domain=self.domain, image=self.standard_domain)
  
          # Construct 1D polynomial families
          Js = []
          for qd in range(self.dim):
              Js.append(JacobiPolynomials(alpha=self.beta[qd]-1.,
                        beta=self.alpha[qd]-1.))
          self.polys = TensorialPolynomials(polys1d=Js)
  
          self.indices = None
  
      def _detect_dimension(self, alpha, beta, dim, domain):
          """
          Parses user-given inputs to determine
          the dimension of the distribution.
  
          alpha and beta are iterables.
  
          Sets self.dim, self.alpha, self.beta, and self.domain
          """
  
          # Situations:
          # 1. User specifies alpha, beta as lists (disallow contradictory
          #    dimension, domain specification)
          # 2. User specifies dim scalar (disallow contradictory alpha, beta,
          #    domain specification)
          # 3. User specifies domain hypercube
  
          if len(alpha) > 1 or len(beta) > 1:  # Case 1:
              assert len(alpha) == len(beta), 'Input parameters alpha and beta \
                      must be the same dimension'
  
              assert (dim is None) or (dim == 1) or (dim == len(alpha)), \
                  'Alpha, beta parameter lists must have size consistent with \
                  input dim'
  
              if (domain is not None) and (domain.shape[1] != 1) and \
                      (domain.shape[1] != len(alpha)):
                  raise ValueError('Alpha, beta parameter lists must have size \
                          consistent with hypercube domain')
  
              self.dim = len(alpha)
              self.alpha = alpha
              self.beta = beta
  
              if domain is None:  # Standard domain [0,1]^dim
                  self.domain = np.zeros([2, self.dim])
                  self.domain[1, :] = 1.
              else:
                  domain = np.asarray(domain)
                  if domain.shape[1] == 1:  # Tensorize 1D domain
                      self.domain = np.zeros([2, self.dim])
                      self.domain[0, :] = domain[0]
                      self.domain[1, :] = domain[1]
                  else:  # User specified domain
                      self.domain = domain
  
          elif dim is not None and dim > 1:  # Case 2
              self.dim = dim
              self.alpha = [alpha[0] for i in range(self.dim)]
              self.beta = [beta[0] for i in range(self.dim)]
  
              if (domain is None) or (domain.any() is not True):  # Standard domain [0,1]^dim
                  self.domain = np.zeros([2, self.dim])
                  self.domain[1, :] = 1.
              else:
                  domain = np.asarray(domain)
                  if domain.shape[1] == 1:  # Tensorize 1D domain
                      self.domain = np.zeros([2, self.dim])
                      self.domain[0, :] = domain[0]
                      self.domain[1, :] = domain[1]
                  else:  # User specified domain
                      self.domain = domain
  
              return
  
          elif domain is not None and len(domain.shape) > 1:  # Case 3
              self.dim = domain.shape[1]
              self.alpha = [alpha[0] for i in range(self.dim)]
              self.beta = [beta[0] for i in range(self.dim)]
              self.domain = domain
  
          else:  # The dimension is 1
              self.dim = 1
              self.alpha = alpha
              self.beta = beta
  
              if domain is None:
                  self.domain = np.zeros([2, self.dim])
                  self.domain[1, :] = 1.
              else:
                  self.domain = domain.reshape([2, self.dim])
  
      def _convert_meanstdev_to_alphabeta(self, mean, stdev):
          """
          Converts user-given mean and stdev to an iterable alpha, beta.
          """
  
          meaniter = isinstance(mean, (list, tuple))
          stdviter = isinstance(stdev, (list, tuple))
          alpha = []
          beta = []
  
          # If they're both iterables:
          if meaniter and stdviter:
  
              # If one has length 1 and the other has length > 1:
              if (len(mean) == 1) and (len(stdev) > 1):
                  mean = [mean for i in range(len(stdev))]
              elif (len(stdev) == 1) and (len(mean) > 1):
                  stdev = [stdev for i in range(len(mean))]
  
              for ind in range(len(mean)):
                  alph, bet = self.meanstdev_to_alphabeta(mean[ind], stdev[ind])
                  alpha.append(alph)
                  beta.append(bet)
  
          # If mean is an iterable but stdev is not
          elif meaniter:
              for ind in range(len(mean)):
                  alph, bet = self.meanstdev_to_alphabeta(mean[ind], stdev)
                  alpha.append(alph)
                  beta.append(bet)
  
          # If stdev is an iterable but mean is not
          elif stdviter:
              for ind in range(len(mean)):
                  alph, bet = self.meanstdev_to_alphabeta(mean, stdev[ind])
                  alpha.append(alph)
                  beta.append(bet)
  
          # If they're both scalars, let the following alpha/beta checker
          # cf vs dim
          else:
              alpha, beta = self.meanstdev_to_alphabeta(mean, stdev)
  
          return alpha, beta
  
      def _convert_alphabeta_to_iterable(self, alpha, beta):
          """
          Converts user-input (alpha,beta) to iterables. Ensures that the length
          of the iterables matches on output.
  
          If alpha or beta are None, sets those values to 1.
          """
  
          # Tons of type/value checking for alpha/beta vs dim
          alphiter = isinstance(alpha, (list, tuple))
          betaiter = isinstance(beta, (list, tuple))
          if alphiter and betaiter:
              if len(alpha) > 1 and len(beta) > 1:
                  assert len(alpha) == len(beta), "Alpha and Beta parameter \
                          inputs must be of the same size"
              elif len(alpha) == 1:
                  alpha = [alpha[0] for i in range(len(beta))]
              elif len(beta) == 1:
                  beta = [beta[0] for i in range(len(alpha))]
  
          elif alphiter:  # alpha is iterable, beta is not
              if beta is None:
                  beta = 1.
              beta = [beta for i in range(len(alpha))]
  
          elif betaiter:  # beta is iterable, alpha is not
              if alpha is None:
                  alpha = 1.
              alpha = [alpha for i in range(len(beta))]
  
          elif alpha is None and beta is None:
              alpha, beta = [1., ], [1., ]
  
          else:  # alpha, beta should both be scalars
              alpha, beta = [alpha, ], [beta, ]
  
          return alpha, beta
  
      def meanstdev_to_alphabeta(self, mu, stdev):
          """
          Returns alpha, beta given an input mean (mu) and
          standard deviation (stdev) for a Beta distribution
          on the interval [0, 1].
          """
  
          if 0. >= mu or mu >= 1.:
              raise ValueError('Mean of a standard Beta distribution \
                      must be between 0 and 1.')
  
          if stdev >= np.sqrt(mu*(1-mu)):
              msg = 'Standard deviation of a Beta random variable must be \
                      smaller than the geometric mean of mu and (1-mu)'
              raise ValueError(msg)
  
          temp = (mu * (1-mu) - stdev**2)/stdev**2
  
          return mu*temp, (1-mu)*temp
  
      def MC_samples(self, M=100):
          """
          Returns M Monte Carlo samples from the distribution.
          """
  
          p = np.zeros([M, self.dim])
          for qd in range(self.dim):
              p[:, qd] = np.random.beta(self.alpha[qd], self.beta[qd], M)
  
          return self.transform_to_standard.mapinv(p)
  
      def mean(self):
          """
          Returns the mean of the distribution.
          """
  
          mu = np.zeros(self.dim)
          for i in range(self.dim):
              mu[i] = self.alpha[i]/(self.alpha[i] + self.beta[i])
  
          # Affine map to appropriate interval
          mu = np.reshape(mu, [1, self.dim])
          return self.transform_to_standard.mapinv(mu).flatten()
  
      def cov(self):
          """
          Returns the (auto-)covariance matrix of the distribution.
          """
  
          sigma = np.eye(self.dim)
          for i in range(self.dim):
              num = self.alpha[i]*self.beta[i]
              den = (self.alpha[i]+self.beta[i])**2 * \
                    (self.alpha[i]+self.beta[i]+1)
              sigma[i, i] = num/den
  
          # Need to get scaling part of the affine transform, so need
          # non-homogeneous term b.
          zero = np.zeros([1, self.dim])
          b = self.transform_to_standard.mapinv(zero)
          bmat = np.tile(b, [self.dim, 1])
  
          # Multiply on right by scaling matrix
          sigma = self.transform_to_standard.mapinv(sigma) - bmat
  
          # Multiply on left by scaling matrix
          sigma = (self.transform_to_standard.mapinv(sigma.T) - bmat).T
  
          return sigma
  
      def stdev(self):
          """
          Returns the standard deviation of the distribution, if the
          distribution is one-dimensional. Raises an error if called
          for a multivariate distribution.
          """
  
          if self.dim == 1:
              return np.sqrt(self.cov()[0, 0])
          else:
              raise TypeError("Can only compute standard deviations for scalar\
                               random variables.")
  
      def pdf(self, x):
          """
          Evaluates the probability density function (pdf) of the distribution
          at the input locations x.
          """
  
          # Evaluate in standard space
          x = self.transform_to_standard.map(x)
          density = np.ones(x.shape[0])
          for i in range(self.dim):
              density *= x[:, i]**(self.alpha[i]-1) * \
                         (1 - x[:, i])**(self.beta[i]-1)
              density /= sp.special.beta(self.alpha[i], self.beta[i])
  
          # Scale based on determinant of map
          density *= self.transform_to_standard.jacobian_determinant()
          return density
  
  
  class UniformDistribution(BetaDistribution):
      """.. _uniform_distribution:
  
      Constructs a (continuous) uniform distribution object; supports
      multivariate distributions through tensorization. In one dimension,
      uniforms distributions have support on the real interval [0,1],
      with probability density function,
  
      .. math::
  
        w(y) := 1, \\hskip 20pt y \\in (0,1),
  
      To generate this distribution on a general compact interval :math:`[a,b]`,
      set the domain parameter below. Alternatively, the mean :math:`\\mu` and
      standard deviation :math:`\\sigma` can be input, in which case the support
      interval :math:`[a,b]` is determined.
  
      This class supports tensorization: multidimensional distributions
      corresopnding to independent one-dimensional marginal distributions are
      supported. In the case of identically distributed marginals, the `dim`
      parameter can be set to the appropriate dimension. In case of non-identical
      marginals, an array or iterable can be input for :math:`\\alpha, \\beta,
      \\mu, \\sigma`.
  
      Parameters:
          mean (float or iterable of floats, optional): Mean of the distribution.
          Defaults to None.
          stdev (float or iterable of floats, optional): Standard deviation of
          the distribution. Defaults to None.
          dim (int, optional): Dimension of the distribution. Defaults to None.
          domain (numpy.ndarray or similar, of size 2 x `dim`, optional): Compact
          hypercube that is the support of the distribution. Defaults to None
  
      Attributes:
          dim (int): Dimension of the distribution.
          polys (:class:`JacobiPolynomials` or list thereof):
      """
  
      def __init__(self, mean=None, stdev=None, dim=None, domain=None, bounds=None):
  
          super().__init__(alpha=1., beta=1., mean=mean, stdev=stdev,
                           dim=dim, domain=domain, bounds=bounds)
  
  
  class DiscreteUniformDistribution(ProbabilityDistribution):
      def __init__(self, n=None, domain=None, dim=None):
  
          if n is None:
              raise ValueError('Input "n" is required.')
  
          # Make sure dim is set, and that n is a list with len(n)==dim
          if dim is not None:
              if (len(n) > 1) and (len(n) != dim):
                  raise ValueError('Inconsistent settings for inputs "dim" \
                          and "n"')
              elif len(n) == 1:
                  if isinstance(n, (list, tuple)):
                      n = [n[0]]*dim
                  else:
                      n = [n]*dim
  
              else:  # len(n) == dim != 1
                  pass  # Nothing to do
          else:
              if isinstance(n, (list, tuple)):
                  dim = len(n)
              else:
                  n = [n]
                  dim = 1
  
          # Ensure that user-specified domain makes sense
          if domain is None:
              domain = np.ones([2, dim])
              domain[0, :] = 0
          else:   # Domain should be a 2 x dim numpy array
              if (domain.shape[0] != 2) or (domain.shape[1] != dim):
                  raise ValueError('Inputs "domain" and inferred dimension \
                          are inconsistent')
              else:
                  pass  # Nothing to do
  
          # Assign stuff
          self.dim, self.n, self.domain = dim, n, domain
  
          # Compute transformations
          # Standard domain is [0,1]^dim
          self.standard_domain = np.ones([2, self.dim])
          self.standard_domain[0, :] = 0.
  
          # Low-level routines use Discrete Chebyshev Polynomials,
          # which operate on [0,1]
          self.poly_domain = np.ones([2, self.dim])
          self.poly_domain[0, :] = 0.
          self.transform_standard_dist_to_poly = AffineTransform(
                  domain=self.standard_domain, image=self.poly_domain)
  
          self.transform_to_standard = AffineTransform(
                  domain=self.domain, image=self.standard_domain)
  
          Ps = []
          for qd in range(self.dim):
              Ps.append(DiscreteChebyshevPolynomials(M=n[qd]))
          self.polys = TensorialPolynomials(polys1d=Ps)
  
      def MC_samples(self, M=100):
          """
          Returns M Monte Carlo samples from the distribution.
          """
  
          p = np.zeros([M, self.dim])
          for qd in range(self.dim):
              p[:, qd] = choice(self.polys.polys1d[qd].standard_support, size=M)
  
          return self.transform_to_standard.mapinv(p)
  
      def mean(self):
          mu = np.zeros(self.dim)
          for i in range(self.dim):
              mu[i] = 1/2
          mu = np.reshape(mu, [1, self.dim])
          return self.transform_to_standard.mapinv(mu).flatten()
  
      def cov(self):
          sigma = np.eye(self.dim)
          for i in range(self.dim):
              sigma[i, i] = (self.n[i]+1) / (12*(self.n[i]-1))
          zero = np.zeros([1, self.dim])
          b = self.transform_to_standard.mapinv(zero)
          bmat = np.tile(b, [self.dim, 1])
          sigma = self.transform_to_standard.mapinv(sigma) - bmat
          sigma = (self.transform_to_standard.mapinv(sigma.T) - bmat).T
  
          return sigma
  
      def stdev(self):
          """
          Returns the standard deviation of the distribution, if the distribution
          is one-dimensional. Raises an error if called for a multivariate
          distribution.
          """
  
          if self.dim == 1:
              return np.sqrt(self.cov()[0, 0])
          else:
              raise TypeError("Can only compute standard deviations for scalar\
                               random variables.")
  
      def pmf(self, x):
          """
          Evaluates the probability mass function (pmf) of the distribution
          """
          density = np.ones(x.shape[0])
          for i in range(self.dim):
              density *= 1/self.n[i]
          return density
  
  
  class TensorialDistribution(ProbabilityDistribution):
      def __init__(self, distributions=None, dim=None):
  
          if dim is not None:
              if len(distributions) > 1:
                  raise ValueError("Input 'dim' cannot be set if \
                          'distributions' contains more than one element")
              else:
                  distributions = dim*distributions
  
          self.distributions = distributions
          self.dim = np.sum([dist.dim for dist in distributions])
  
          self.standard_domain = np.concatenate(
                  [dist.standard_domain.T for dist in distributions]
                  ).T
  
          self.poly_domain = np.concatenate(
                  [dist.poly_domain.T for dist in distributions]
                  ).T
  
          # Construct transform_standard_dist_to_poly
          Ts = [dist.transform_standard_dist_to_poly for dist in distributions]
          As = [t.A.toarray() if isinstance(t.A, sprs.spmatrix) else
                t.A for t in Ts]
          bs = [t.b for t in Ts]
  
          self.transform_standard_dist_to_poly = AffineTransform(
                  A=sp.linalg.block_diag(*As),
                  b=np.concatenate(bs)
              )
  
          # Construct transform_to_standard
          Ts = [dist.transform_to_standard for dist in distributions]
          As = [t.A.toarray() if isinstance(t.A, sprs.spmatrix) else
                t.A for t in Ts]
          bs = [t.b for t in Ts]
  
          self.transform_to_standard = AffineTransform(
                  A=sp.linalg.block_diag(*As),
                  b=np.concatenate(bs)
              )
  
          self.polys = TensorialPolynomials([poly for dist in distributions
                                            for poly in dist.polys.polys1d])
  
          self.indices = None
  
      def MC_samples(self, M=100):
          """
          Returns M Monte Carlo samples from the distribution
          """
  
          p = np.zeros([M, self.dim])
          counter = 0
          for dist in self.distributions:
              p[:, range(counter, counter+dist.dim)] = dist.MC_samples(M=M)
              counter += dist.dim
  
          # Each component distribution already applies
          # transform_to_standard.mapinv.
          return p
  