Python autograd.numpy.vstack() Examples

def make_IO_matrices(indices, N):
    """ Makes matrices that relate the sparse matrix entries to their locations in the matrix
            The kth column of I is a 'one hot' vector specifing the k-th entries row index into A
            The kth column of J is a 'one hot' vector specifing the k-th entries columnn index into A
            O = J^T is for notational convenience.
            Armed with a vector of M entries 'a', we can construct the sparse matrix 'A' as:
                A = I @ diag(a) @ O
            where 'diag(a)' is a (MxM) matrix with vector 'a' along its diagonal.
            In index notation:
                A_ij = I_ik * a_k * O_kj
            In an opposite way, given sparse matrix 'A' we can strip out the entries `a` using the IO matrices as follows:
                a = diag(I^T @ A @ O^T)
            In index notation:
                a_k = I_ik * A_ij * O_kj
    """
    M = indices.shape[1]                                 # number of indices in the matrix
    entries_1 = npa.ones(M)                              # M entries of all 1's
    ik, jk = indices                                     # separate i and j components of the indices
    indices_I = npa.vstack((ik, npa.arange(M)))          # indices into the I matrix
    indices_J = npa.vstack((jk, npa.arange(M)))          # indices into the J matrix
    I = make_sparse(entries_1, indices_I, shape=(N, M))  # construct the I matrix
    J = make_sparse(entries_1, indices_J, shape=(N, M))  # construct the J matrix
    O = J.T                                              # make O = J^T matrix for consistency with my notes.
    return I, O 

def _ntied_transmat_prior(self, transmat_val):  # TODO: document choices
        transmat = np.empty((0, self.n_components))
        for r in range(self.n_unique):
            row = np.empty((self.n_chain, 0))
            for c in range(self.n_unique):
                if r == c:
                    subm = np.array(sp.diags([transmat_val[r, c],
                                    1.0], [0, 1],
                        shape=(self.n_chain, self.n_chain)).todense())
                else:
                    lower_left = np.zeros((self.n_chain, self.n_chain))
                    lower_left[self.n_tied, 0] = 1.0
                    subm = np.kron(transmat_val[r, c], lower_left)
                row = np.hstack((row, subm))
            transmat = np.vstack((transmat, row))
        return transmat 

def sample(self, n, seed=29):
        pmix = self.pmix
        means = self.means
        variances = self.variances
        k, d = self.means.shape
        sam_list = []
        with util.NumpySeedContext(seed=seed):
            # counts for each mixture component 
            counts = np.random.multinomial(n, pmix, size=1)

            # counts is a 2d array
            counts = counts[0]

            # For each component, draw from its corresponding mixture component.            
            for i, nc in enumerate(counts):
                # Sample from ith component
                sam_i = np.random.randn(nc, d)*np.sqrt(variances[i]) + means[i]
                sam_list.append(sam_i)
            sample = np.vstack(sam_list)
            assert sample.shape[0] == n
            np.random.shuffle(sample)
        return Data(sample)

# end of class DSIsoGaussianMixture 

def job_me_opt(p, data_source, tr, te, r, J=5):
    """
    ME test of Jitkrittum et al., 2016 used as a goodness-of-fit test.
    Gaussian kernel. Optimize test locations and Gaussian width.
    """
    data = tr + te
    X = data.data()
    with util.ContextTimer() as t:
        # median heuristic 
        #pds = p.get_datasource()
        #datY = pds.sample(data.sample_size(), seed=r+294)
        #Y = datY.data()
        #XY = np.vstack((X, Y))
        #med = util.meddistance(XY, subsample=1000)
        op = {'n_test_locs': J, 'seed': r+5, 'max_iter': 40, 
             'batch_proportion': 1.0, 'locs_step_size': 1.0, 
              'gwidth_step_size': 0.1, 'tol_fun': 1e-4, 
              'reg': 1e-4}
        # optimize on the training set
        me_opt = tgof.GaussMETestOpt(p, n_locs=J, tr_proportion=tr_proportion,
                alpha=alpha, seed=r+111)

        me_result = me_opt.perform_test(data, op)
    return { 'test_result': me_result, 'time_secs': t.secs} 

def job_me_opt(p, data_source, tr, te, r, J=5):
    """
    ME test of Jitkrittum et al., 2016 used as a goodness-of-fit test.
    Gaussian kernel. Optimize test locations and Gaussian width.
    """
    data = tr + te
    X = data.data()
    with util.ContextTimer() as t:
        # median heuristic 
        #pds = p.get_datasource()
        #datY = pds.sample(data.sample_size(), seed=r+294)
        #Y = datY.data()
        #XY = np.vstack((X, Y))
        #med = util.meddistance(XY, subsample=1000)
        op = {'n_test_locs': J, 'seed': r+5, 'max_iter': 40, 
             'batch_proportion': 1.0, 'locs_step_size': 1.0, 
              'gwidth_step_size': 0.1, 'tol_fun': 1e-4, 
              'reg': 1e-4}
        # optimize on the training set
        me_opt = tgof.GaussMETestOpt(p, n_locs=J, tr_proportion=tr_proportion,
                alpha=alpha, seed=r+111)

        me_result = me_opt.perform_test(data, op)
    return { 'test_result': me_result, 'time_secs': t.secs} 

def _make_A(self, eps_vec):

        eps_vec_xx, eps_vec_yy = self._grid_average_2d(eps_vec)
        eps_vec_xx_inv = 1 / (eps_vec_xx + 1e-5)  # the 1e-5 is for numerical stability
        eps_vec_yy_inv = 1 / (eps_vec_yy + 1e-5)  # autograd throws 'divide by zero' errors.

        indices_diag = npa.vstack((npa.arange(self.N), npa.arange(self.N)))

        entries_DxEpsy,   indices_DxEpsy   = spsp_mult(self.entries_Dxb, self.indices_Dxb, eps_vec_yy_inv, indices_diag, self.N)
        entires_DxEpsyDx, indices_DxEpsyDx = spsp_mult(entries_DxEpsy, indices_DxEpsy, self.entries_Dxf, self.indices_Dxf, self.N)

        entries_DyEpsx,   indices_DyEpsx   = spsp_mult(self.entries_Dyb, self.indices_Dyb, eps_vec_xx_inv, indices_diag, self.N)
        entires_DyEpsxDy, indices_DyEpsxDy = spsp_mult(entries_DyEpsx, indices_DyEpsx, self.entries_Dyf, self.indices_Dyf, self.N)

        entries_d = 1 / EPSILON_0 * npa.hstack((entires_DxEpsyDx, entires_DyEpsxDy))
        indices_d = npa.hstack((indices_DxEpsyDx, indices_DyEpsxDy))

        entries_diag = MU_0 * self.omega**2 * npa.ones(self.N)

        entries_a = npa.hstack((entries_d, entries_diag))
        indices_a = npa.hstack((indices_d, indices_diag))

        return entries_a, indices_a 

def gmm_log_likelihood(params, data):
    cluster_lls = []
    for log_proportion, mean, cov_sqrt in zip(*unpack_gmm_params(params)):
        cov = np.dot(cov_sqrt.T, cov_sqrt)
        cluster_lls.append(log_proportion + mvn.logpdf(data, mean, cov))
    return np.sum(logsumexp(np.vstack(cluster_lls), axis=0)) 

def _ntied_transmat(self, transmat_val):  # TODO: document choices
        #                        +-----------------+
        #                        |a|1|0|0|0|0|0|0|0|
        #                        +-----------------+
        #                        |0|a|1|0|0|0|0|0|0|
        #                        +-----------------+
        #   +---+---+---+        |0|0|a|b|0|0|c|0|0|
        #   | a | b | c |        +-----------------+
        #   +-----------+        |0|0|0|e|1|0|0|0|0|
        #   | d | e | f | +----> +-----------------+
        #   +-----------+        |0|0|0|0|e|1|0|0|0|
        #   | g | h | i |        +-----------------+
        #   +---+---+---+        |d|0|0|0|0|e|f|0|0|
        #                        +-----------------+
        #                        |0|0|0|0|0|0|i|1|0|
        #                        +-----------------+
        #                        |0|0|0|0|0|0|0|i|1|
        #                        +-----------------+
        #                        |g|0|0|h|0|0|0|0|i|
        #                        +-----------------+
        # for a model with n_unique = 3 and n_tied = 2
        transmat = np.empty((0, self.n_components))
        for r in range(self.n_unique):
            row = np.empty((self.n_chain, 0))
            for c in range(self.n_unique):
                if r == c:
                    subm = np.array(sp.diags([transmat_val[r, c],
                                    1 - transmat_val[r, c]], [0, 1],
                                    shape=(self.n_chain,
                                           self.n_chain)).todense())
                else:
                    lower_left = np.zeros((self.n_chain, self.n_chain))
                    lower_left[self.n_tied, 0] = 1.0
                    subm = np.kron(transmat_val[r, c], lower_left)
                row = np.hstack((row, subm))
            transmat = np.vstack((transmat, row))
        return transmat 

def _calc_pareto_front(self, n_pareto_points=100):
        x1 = anp.linspace(0, 5, n_pareto_points)
        x2 = anp.copy(x1)
        x2[x1 >= 3] = 3
        return anp.vstack((4 * anp.square(x1) + 4 * anp.square(x2), anp.square(x1 - 5) + anp.square(x2 - 5))).T 

def sample(self, n, seed=29):
        pmix = self.pmix
        means = self.means
        variances = self.variances
        k, d = self.means.shape
        sam_list = []
        with util.NumpySeedContext(seed=seed):
            # counts for each mixture component 
            counts = np.random.multinomial(n, pmix, size=1)

            # counts is a 2d array
            counts = counts[0]

            # For each component, draw from its corresponding mixture component.            
            for i, nc in enumerate(counts):
                # construct the component
                # https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.stats.multivariate_normal.html
                cov = variances[i]
                mnorm = stats.multivariate_normal(means[i], cov)
                # Sample from ith component
                sam_i = mnorm.rvs(size=nc)
                sam_list.append(sam_i)
            sample = np.vstack(sam_list)
            assert sample.shape[0] == n
            np.random.shuffle(sample)
        return Data(sample)

# end of DSGaussianMixture 

def __add__(self, data2):
        """
        Merge the current Data with another one.
        Create a new Data and create a new copy for all internal variables.
        """
        copy = self.clone()
        copy2 = data2.clone()
        nX = np.vstack((copy.X, copy2.X))
        return Data(nX)

### end Data class 

def job_mmd_dgauss_opt(p, data_source, tr, te, r):
    """
    MMD test of Gretton et al., 2012 used as a goodness-of-fit test.
    Require the ability to sample from p i.e., the UnnormalizedDensity p has 
    to be able to return a non-None from get_datasource()

    With optimization. Diagonal Gaussian kernel where there is one Gaussian width
    for each dimension.
    """
    data = tr + te
    X = data.data()
    d = X.shape[1]
    with util.ContextTimer() as t:
        # median heuristic 
        pds = p.get_datasource()
        datY = pds.sample(data.sample_size(), seed=r+294)
        Y = datY.data()
        XY = np.vstack((X, Y))

        # Get the median heuristic for each dimension
        meds = np.zeros(d)
        for i in range(d):
            medi = util.meddistance(XY[:, [i]], subsample=1000)
            meds[i] = medi

        # Construct a list of kernels to try based on multiples of the median
        # heuristic
        med_factors = 2.0**np.linspace(-4, 4, 20)  
        candidate_kernels = []
        for i in range(len(med_factors)):
            ki = kernel.KDiagGauss( (meds**2)*med_factors[i] )
            candidate_kernels.append(ki)

        mmd_opt = mgof.QuadMMDGofOpt(p, n_permute=300, alpha=alpha, seed=r+56)
        mmd_result = mmd_opt.perform_test(data,
                candidate_kernels=candidate_kernels,
                tr_proportion=tr_proportion, reg=1e-3)
    return { 'test_result': mmd_result, 'time_secs': t.secs}

# Define our custom Job, which inherits from base class IndependentJob 

def job_mmd_med(p, data_source, tr, te, r):
    """
    MMD test of Gretton et al., 2012 used as a goodness-of-fit test.
    Require the ability to sample from p i.e., the UnnormalizedDensity p has 
    to be able to return a non-None from get_datasource()
    """
    # full data
    data = tr + te
    X = data.data()
    with util.ContextTimer() as t:
        # median heuristic 
        pds = p.get_datasource()
        datY = pds.sample(data.sample_size(), seed=r+294)
        Y = datY.data()
        XY = np.vstack((X, Y))

        # If p, q differ very little, the median may be very small, rejecting H0
        # when it should not?
        # If p, q differ very little, the median may be very small, rejecting H0
        # when it should not?
        medx = util.meddistance(X, subsample=1000)
        medy = util.meddistance(Y, subsample=1000)
        medxy = util.meddistance(XY, subsample=1000)
        med_avg = (medx+medy+medxy)/3.0
        k = kernel.KGauss(med_avg**2)

        mmd_test = mgof.QuadMMDGof(p, k, n_permute=400, alpha=alpha, seed=r)
        mmd_result = mmd_test.perform_test(data)
    return { 'test_result': mmd_result, 'time_secs': t.secs} 

def job_mmd_opt(p, data_source, tr, te, r):
    """
    MMD test of Gretton et al., 2012 used as a goodness-of-fit test.
    Require the ability to sample from p i.e., the UnnormalizedDensity p has 
    to be able to return a non-None from get_datasource()

    With optimization. Gaussian kernel.
    """
    data = tr + te
    X = data.data()
    with util.ContextTimer() as t:
        # median heuristic 
        pds = p.get_datasource()
        datY = pds.sample(data.sample_size(), seed=r+294)
        Y = datY.data()
        XY = np.vstack((X, Y))

        med = util.meddistance(XY, subsample=1000)

        # Construct a list of kernels to try based on multiples of the median
        # heuristic
        #list_gwidth = np.hstack( (np.linspace(20, 40, 10), (med**2)
        #    *(2.0**np.linspace(-2, 2, 20) ) ) )
        list_gwidth = (med**2)*(2.0**np.linspace(-3, 3, 30) ) 
        list_gwidth.sort()
        candidate_kernels = [kernel.KGauss(gw2) for gw2 in list_gwidth]

        mmd_opt = mgof.QuadMMDGofOpt(p, n_permute=300, alpha=alpha, seed=r+56)
        mmd_result = mmd_opt.perform_test(data,
                candidate_kernels=candidate_kernels,
                tr_proportion=tr_proportion, reg=1e-3)
    return { 'test_result': mmd_result, 'time_secs': t.secs}


# Define our custom Job, which inherits from base class IndependentJob 

def job_mmd_med(p, data_source, tr, te, r):
    """
    MMD test of Gretton et al., 2012 used as a goodness-of-fit test.
    Require the ability to sample from p i.e., the UnnormalizedDensity p has 
    to be able to return a non-None from get_datasource()
    """
    # full data
    data = tr + te
    X = data.data()
    with util.ContextTimer() as t:
        # median heuristic 
        pds = p.get_datasource()
        datY = pds.sample(data.sample_size(), seed=r+294)
        Y = datY.data()
        XY = np.vstack((X, Y))

        # If p, q differ very little, the median may be very small, rejecting H0
        # when it should not?
        medx = util.meddistance(X, subsample=1000)
        medy = util.meddistance(Y, subsample=1000)
        medxy = util.meddistance(XY, subsample=1000)
        med_avg = (medx+medy+medxy)/3.0
        k = kernel.KGauss(med_avg**2)

        mmd_test = mgof.QuadMMDGof(p, k, n_permute=400, alpha=alpha, seed=r)
        mmd_result = mmd_test.perform_test(data)
    return { 'test_result': mmd_result, 'time_secs': t.secs} 

def get_sharp_TE_airfoil(self):
        # Returns a version of the airfoil with a sharp trailing edge.

        upper_original_coors = self.upper_coordinates()  # Note: includes leading edge point, be careful about duplicates
        lower_original_coors = self.lower_coordinates()  # Note: includes leading edge point, be careful about duplicates

        # Find data about the TE

        # Get the scale factor
        x_mcl = self.mcl_coordinates[:, 0]
        x_max = np.max(x_mcl)
        x_min = np.min(x_mcl)
        scale_factor = (x_mcl - x_min) / (x_max - x_min)  # linear contraction

        # Do the contraction
        upper_minus_mcl_adjusted = self.upper_minus_mcl - self.upper_minus_mcl[-1, :] * np.expand_dims(scale_factor, 1)

        # Recreate coordinates
        upper_coordinates_adjusted = np.flipud(self.mcl_coordinates + upper_minus_mcl_adjusted)
        lower_coordinates_adjusted = self.mcl_coordinates - upper_minus_mcl_adjusted

        coordinates = np.vstack((
            upper_coordinates_adjusted[:-1, :],
            lower_coordinates_adjusted
        ))

        # Make a new airfoil with the coordinates
        name = self.name + ", with sharp TE"
        new_airfoil = Airfoil(name=name, coordinates=coordinates, repanel=False)

        return new_airfoil 

def add_control_surface(self, deflection=0., hinge_point=0.75):
        # Returns a version of the airfoil with a control surface added at a given point.
        # Inputs:
        #   # deflection: the deflection angle, in degrees. Downwards-positive.
        #   # hinge_point: the location of the hinge, as a fraction of chord.

        # Make the rotation matrix for the given angle.
        sintheta = np.sin(np.radians(-deflection))
        costheta = np.cos(np.radians(-deflection))
        rotation_matrix = np.array(
            [[costheta, -sintheta],
             [sintheta, costheta]]
        )

        # Find the hinge point
        hinge_point = np.array(
            (hinge_point, self.get_camber_at_chord_fraction(hinge_point)))  # Make hinge_point a vector.

        # Split the airfoil into the sections before and after the hinge
        split_index = np.where(self.mcl_coordinates[:, 0] > hinge_point[0])[0][0]
        mcl_coordinates_before = self.mcl_coordinates[:split_index, :]
        mcl_coordinates_after = self.mcl_coordinates[split_index:, :]
        upper_minus_mcl_before = self.upper_minus_mcl[:split_index, :]
        upper_minus_mcl_after = self.upper_minus_mcl[split_index:, :]

        # Rotate the mean camber line (MCL) and "upper minus mcl"
        new_mcl_coordinates_after = np.transpose(
            rotation_matrix @ np.transpose(mcl_coordinates_after - hinge_point)) + hinge_point
        new_upper_minus_mcl_after = np.transpose(rotation_matrix @ np.transpose(upper_minus_mcl_after))

        # Do blending

        # Assemble airfoil
        new_mcl_coordinates = np.vstack((mcl_coordinates_before, new_mcl_coordinates_after))
        new_upper_minus_mcl = np.vstack((upper_minus_mcl_before, new_upper_minus_mcl_after))
        upper_coordinates = np.flipud(new_mcl_coordinates + new_upper_minus_mcl)
        lower_coordinates = new_mcl_coordinates - new_upper_minus_mcl
        coordinates = np.vstack((upper_coordinates, lower_coordinates[1:, :]))

        new_airfoil = Airfoil(name=self.name + " flapped", coordinates=coordinates, repanel=False)
        return new_airfoil  # TODO fix self-intersecting airfoils at high deflections 

def _make_A(self, eps_vec):

        C = - 1 / MU_0 * self.Dxf.dot(self.Dxb) \
            - 1 / MU_0 * self.Dyf.dot(self.Dyb)
        entries_c, indices_c = get_entries_indices(C)

        # indices into the diagonal of a sparse matrix
        entries_diag = - EPSILON_0 * self.omega**2 * eps_vec
        indices_diag = npa.vstack((npa.arange(self.N), npa.arange(self.N)))

        entries_a = npa.hstack((entries_diag, entries_c))
        indices_a = npa.hstack((indices_diag, indices_c))

        return entries_a, indices_a 

def block_4(A, B, C, D):
    """ Constructs a big matrix out of four sparse blocks
        returns [A B]
                [C D]
    """
    left = sp.vstack([A, C])
    right = sp.vstack([B, D])
    return sp.hstack([left, right]) 

def get_entries_indices(csr_matrix):
    # takes sparse matrix and returns the entries and indeces in form compatible with 'make_sparse'
    shape = csr_matrix.shape
    coo_matrix = csr_matrix.tocoo()
    entries = csr_matrix.data
    cols = coo_matrix.col
    rows = coo_matrix.row
    indices = npa.vstack((rows, cols))
    return entries, indices 

def test_vstack_3d(): combo_check(np.vstack, [0])([R(2, 3, 4), (R(2, 3, 4), R(5, 3, 4))]) 

def test_vstack_1d(): combo_check(np.vstack, [0])([R(2), (R(2), R(2))]) 

def test_jacobian_against_stacked_grads():
    scalar_funs = [
        lambda x: np.sum(x**3),
        lambda x: np.prod(np.sin(x) + np.sin(x)),
        lambda x: grad(lambda y: np.exp(y) * np.tanh(x[0]))(x[1])
    ]

    vector_fun = lambda x: np.array([f(x) for f in scalar_funs])

    x = npr.randn(5)
    jac = jacobian(vector_fun)(x)
    grads = [grad(f)(x) for f in scalar_funs]

    assert np.allclose(jac, np.vstack(grads)) 

def plot_ellipse(ax, mean, cov_sqrt, alpha, num_points=100):
    angles = np.linspace(0, 2*np.pi, num_points)
    circle_pts = np.vstack([np.cos(angles), np.sin(angles)]).T * 2.0
    cur_pts = mean + np.dot(circle_pts, cov_sqrt)
    ax.plot(cur_pts[:, 0], cur_pts[:, 1], '-', alpha=alpha) 

def convert_results(results, interface):
        """Convert a list of results coming from multiple QNodes
        to the object required by each interface for auto-differentiation.

        Internally, this method makes use of ``tf.stack``, ``torch.stack``,
        and ``np.vstack``.

        Args:
            results (list): list containing the results from
                multiple QNodes
            interface (str): the interfaces of the underlying QNodes

        Returns:
            list or array or torch.Tensor or tf.Tensor: the converted
            and stacked results.
        """
        if interface == "tf":
            import tensorflow as tf

            return tf.stack(results)

        if interface == "torch":
            import torch

            return torch.stack(results, dim=0)

        if interface in ("autograd", "numpy"):
            from autograd import numpy as np

            return np.stack(results)

        return results 

def expected_sfs_tensor_prod(vecs, demography, mut_rate=1.0, sampled_pops=None):
    """
    Viewing the SFS as a D-tensor (where D is the number of demes), this
    returns a 1d array whose j-th entry is a certain summary statistic of the
    expected SFS, given by the following tensor-vector multiplication:

    res[j] = \sum_{(i0,i1,...)} E[sfs[(i0,i1,...)]] * vecs[0][j,i0] * vecs[1][j, i1] * ...

    where E[sfs[(i0,i1,...)]] is the expected SFS entry for config
    (i0,i1,...), as given by expected_sfs

    Parameters
    ----------
    vecs : sequence of 2-dimensional numpy.ndarray
         length-D sequence, where D = number of demes in the demography.
         vecs[k] is 2-dimensional array, with constant number of rows, and
         with n[k]+1 columns, where n[k] is the number of samples in the
         k-th deme. The row vector vecs[k][j,:] is multiplied against
         the expected SFS along the k-th mode, to obtain res[j].
    demo : Demography
    mut_rate : float
         the rate of mutations per unit time

    Returns
    -------
    res : numpy.ndarray (1-dimensional)
        res[j] is the tensor multiplication of the sfs against the vectors
        vecs[0][j,:], vecs[1][j,:], ... along its tensor modes.

    See Also
    --------
    sfs_tensor_prod : compute the same summary statistics for an observed SFS
    expected_sfs : compute individual SFS entries
    expected_total_branch_len, expected_tmrca, expected_deme_tmrca :
         examples of coalescent statistics that use this function
    """
    # NOTE cannot use vecs[i] = ... due to autograd issues
    sampled_n = [np.array(v).shape[-1] - 1 for v in vecs]
    vecs = [np.vstack([np.array([1.0] + [0.0] * n),  # all ancestral state
                       np.array([0.0] * n + [1.0]),  # all derived state
                       v])
            for v, n in zip(vecs, demography.sampled_n)]

    res = _expected_sfs_tensor_prod(vecs, demography, mut_rate=mut_rate)

    # subtract out mass for all ancestral/derived state
    for k in (0, 1):
        res = res - res[k] * np.prod([l[:, -k] for l in vecs], axis=0)
        assert np.isclose(res[k], 0.0)
    # remove monomorphic states
    res = res[2:]

    return res 

def get_bounding_cube(self):
        """ Finds the axis-aligned cube that encloses the airplane with the smallest size.
            Useful for plotting and getting a sense for the scale of a problem.
            
            Args:
                self.wings (iterable): All the wings included for analysis each containing their geometry in x,y,z notation using units of m
            Returns:
                tuple: Tuple of 4 floats x, y, z, and s, where x, y, and z are the coordinates of the cube center,
                and s is half of the side length.
        """

        # Get vertices to enclose
        vertices = None
        for wing in self.wings:
            for xsec in wing.xsecs:
                if vertices is None:
                    vertices = xsec.xyz_le + wing.xyz_le
                else:
                    vertices = np.vstack((
                        vertices,
                        xsec.xyz_le + wing.xyz_le
                    ))
                vertices = np.vstack((
                    vertices,
                    xsec.xyz_te() + wing.xyz_le
                ))

                if wing.symmetric:
                    vertices = np.vstack((
                        vertices,
                        reflect_over_XZ_plane(xsec.xyz_le + wing.xyz_le)
                    ))
                    vertices = np.vstack((
                        vertices,
                        reflect_over_XZ_plane(xsec.xyz_te() + wing.xyz_le)
                    ))

        # Enclose them
        x_max = np.max(vertices[:, 0])
        y_max = np.max(vertices[:, 1])
        z_max = np.max(vertices[:, 2])
        x_min = np.min(vertices[:, 0])
        y_min = np.min(vertices[:, 1])
        z_min = np.min(vertices[:, 2])

        x = np.mean((x_max, x_min))
        y = np.mean((y_max, y_min))
        z = np.mean((z_max, z_min))
        s = 0.5 * np.max((
            x_max - x_min,
            y_max - y_min,
            z_max - z_min
        ))

        return x, y, z, s 

def draw_legacy(self,
                    show=True,
                    fig_to_plot_on=None,
                    ax_to_plot_on=None
                    ):
        # Draws the airplane using matplotlib.
        # This method is deprecated (superseded by draw() ) and will be removed in a future release.

        # Setup
        if fig_to_plot_on == None or ax_to_plot_on == None:
            fig, ax = fig3d()
            fig.set_size_inches(12, 9)
        else:
            fig = fig_to_plot_on
            ax = ax_to_plot_on

        # TODO plot bodies

        # Plot wings
        for wing in self.wings:

            for i in range(len(wing.sections) - 1):
                le_start = wing.sections[i].xyz_le + wing.xyz_le
                le_end = wing.sections[i + 1].xyz_le + wing.xyz_le
                te_start = wing.sections[i].xyz_te() + wing.xyz_le
                te_end = wing.sections[i + 1].xyz_te() + wing.xyz_le

                points = np.vstack((le_start, le_end, te_end, te_start, le_start))
                x = points[:, 0]
                y = points[:, 1]
                z = points[:, 2]

                ax.plot(x, y, z, color='#cc0039')

                if wing.symmetric:
                    ax.plot(x, -1 * y, z, color='#cc0039')

        # Plot reference point
        x = self.xyz_ref[0]
        y = self.xyz_ref[1]
        z = self.xyz_ref[2]
        ax.scatter(x, y, z)

        set_axes_equal(ax)
        plt.tight_layout()
        if show:
            plt.show() 

def build_mog_bbsvi(logprob, num_samples, k=10, rs=npr.RandomState(0)):
    init_component_var_params = init_gaussian_var_params
    component_log_density = variational_log_density_gaussian
    component_sample = sample_diag_gaussian

    def unpack_mixture_params(mixture_params):
        log_weights = log_normalize(mixture_params[:k])
        var_params = np.reshape(mixture_params[k:], (k, -1))
        return log_weights, var_params

    def init_var_params(D, rs=npr.RandomState(0), **kwargs):
        log_weights = np.ones(k)
        component_weights = [init_component_var_params(D, rs=rs, **kwargs) for i in range(k)]
        return np.concatenate([log_weights] + component_weights)

    def sample(var_mixture_params, num_samples, rs):
        """Sample locations aren't a continuous function of parameters
        due to multinomial sampling."""
        log_weights, var_params = unpack_mixture_params(var_mixture_params)
        samples = np.concatenate([component_sample(params_k, num_samples, rs)[:, np.newaxis, :]
                             for params_k in var_params], axis=1)
        ixs = np.random.choice(k, size=num_samples, p=np.exp(log_weights))
        return np.array([samples[i, ix, :] for i, ix in enumerate(ixs)])

    def mixture_log_density(var_mixture_params, x):
        """Returns a weighted average over component densities."""
        log_weights, var_params = unpack_mixture_params(var_mixture_params)
        component_log_densities = np.vstack([component_log_density(params_k, x)
                                             for params_k in var_params]).T
        return logsumexp(component_log_densities + log_weights, axis=1, keepdims=False)

    def mixture_elbo(var_mixture_params, t):
        # We need to only sample the continuous component parameters,
        # and integrate over the discrete component choice

        def mixture_lower_bound(params):
            """Provides a stochastic estimate of the variational lower bound."""
            samples = component_sample(params, num_samples, rs)
            log_qs = mixture_log_density(var_mixture_params, samples)
            log_ps = logprob(samples, t)
            log_ps = np.reshape(log_ps, (num_samples, -1))
            log_qs = np.reshape(log_qs, (num_samples, -1))
            return np.mean(log_ps - log_qs)

        log_weights, var_params = unpack_mixture_params(var_mixture_params)
        component_elbos = np.stack(
            [mixture_lower_bound(params_k) for params_k in var_params])
        return np.sum(component_elbos*np.exp(log_weights))

    return init_var_params, mixture_elbo, mixture_log_density, sample 

def draw(self):
        # Draw the airplane in a new window.

        # Using PyVista Polydata format
        vertices = np.empty((0, 3))
        faces = np.empty((0))

        for wing in self.wings:
            wing_vertices = np.empty((0, 3))
            wing_tri_faces = np.empty((0, 4))
            wing_quad_faces = np.empty((0, 5))
            for i in range(len(wing.xsecs) - 1):
                is_last_section = i == len(wing.xsecs) - 2

                le_start = wing.xsecs[i].xyz_le + wing.xyz_le
                te_start = wing.xsecs[i].xyz_te() + wing.xyz_le
                wing_vertices = np.vstack((wing_vertices, le_start, te_start))

                wing_quad_faces = np.vstack((
                    wing_quad_faces,
                    np.expand_dims(np.array([4, 2 * i + 0, 2 * i + 1, 2 * i + 3, 2 * i + 2]), 0)
                ))

                if is_last_section:
                    le_end = wing.xsecs[i + 1].xyz_le + wing.xyz_le
                    te_end = wing.xsecs[i + 1].xyz_te() + wing.xyz_le
                    wing_vertices = np.vstack((wing_vertices, le_end, te_end))

            vertices_starting_index = len(vertices)
            wing_quad_faces_reformatted = np.ndarray.copy(wing_quad_faces)
            wing_quad_faces_reformatted[:, 1:] = wing_quad_faces[:, 1:] + vertices_starting_index
            wing_quad_faces_reformatted = np.reshape(wing_quad_faces_reformatted, (-1), order='C')
            vertices = np.vstack((vertices, wing_vertices))
            faces = np.hstack((faces, wing_quad_faces_reformatted))

            if wing.symmetric:
                vertices_starting_index = len(vertices)
                wing_vertices = reflect_over_XZ_plane(wing_vertices)
                wing_quad_faces_reformatted = np.ndarray.copy(wing_quad_faces)
                wing_quad_faces_reformatted[:, 1:] = wing_quad_faces[:, 1:] + vertices_starting_index
                wing_quad_faces_reformatted = np.reshape(wing_quad_faces_reformatted, (-1), order='C')
                vertices = np.vstack((vertices, wing_vertices))
                faces = np.hstack((faces, wing_quad_faces_reformatted))

        plotter = pv.Plotter()

        wing_surfaces = pv.PolyData(vertices, faces)
        plotter.add_mesh(wing_surfaces, color='#7EFC8F', show_edges=True, smooth_shading=True)

        xyz_ref = pv.PolyData(self.xyz_ref)
        plotter.add_points(xyz_ref, color='#50C7C7', point_size=10)

        plotter.show_grid(color='#444444')
        plotter.set_background(color="black")
        plotter.show(cpos=(-1, -1, 1), full_screen=False)