text:
"""
Convolve the current locations and galaxy shapes with the PSF.  If
calculate_derivs is true, also calculate derivatives and hessians for
active sources.

Args:
 - psf: A vector of PSF components
 - ea: The current ElboArgs
 - b: The current band
 - calculate_derivs: Whether to calculate derivatives for active sources.

Returns:
 - star_mcs: An array of BvnComponents with indices
    - PSF component
    - Source (index within active_sources)
 - gal_mcs: An array of BvnComponents with indices
    - PSF component
    - Galaxy component
    - Galaxy type
    - Source (index within active_sources)
  Hessians are only populated for s in ea.active_sources.
"""
function load_bvn_mixtures{NumType <: Number}(
                    ea::ElboArgs{NumType},
                    b::Int;
                    calculate_derivs::Bool=true,
                    calculate_hessian::Bool=true)

    star_mcs = Array(BvnComponent{NumType}, psf_K, ea.S)
    gal_mcs = Array(GalaxyCacheComponent{NumType}, psf_K, 8, 2, ea.S)

    # TODO: do not keep any derviative information if the sources are not in
    # active_sources.
    for s in 1:ea.S
        psf = ea.patches[s, b].psf
        vs = ea.vp[s]

        world_loc = vs[[ids.u[1], ids.u[2]]]
        m_pos = Model.linear_world_to_pix(ea.patches[s, b].wcs_jacobian,
                                             ea.patches[s, b].center,
                                             ea.patches[s, b].pixel_center, world_loc)

        # Convolve the star locations with the PSF.
        for k in 1:psf_K
            pc = psf[k]
            mean_s = [pc.xiBar[1] + m_pos[1], pc.xiBar[2] + m_pos[2]]
            star_mcs[k, s] =
              BvnComponent{NumType}(
                mean_s, pc.tauBar, pc.alphaBar, calculate_siginv_deriv=false)
        end

        # Convolve the galaxy representations with the PSF.
        for i = 1:2 # i indexes dev vs exp galaxy types.
            e_dev_dir = (i == 1) ? 1. : -1.
            e_dev_i = (i == 1) ? vs[ids.e_dev] : 1. - vs[ids.e_dev]

            # Galaxies of type 1 have 8 components, and type 2 have 6 components.
            for j in 1:[8,6][i]
                for k = 1:psf_K
                    gal_mcs[k, j, i, s] = GalaxyCacheComponent(
                        e_dev_dir, e_dev_i, galaxy_prototypes[i][j], psf[k],
                        m_pos, vs[ids.e_axis], vs[ids.e_angle], vs[ids.e_scale],
                        calculate_derivs && (s in ea.active_sources),
                        calculate_hessian)
                end
            end
        end
    end

    star_mcs, gal_mcs
end

# TODO: the identification of active pixels should go in pre-processing
type ActivePixel
    # image index
    n::Int

    # Linear tile index:
    tile_ind::Int

    # Location in tile:
    h::Int
    w::Int
end

"""
Store pre-allocated memory in this data structures, which contains
intermediate values used in the ELBO calculation.
"""
type HessianSubmatrices{NumType <: Number}
    u_u::Matrix{NumType}
    shape_shape::Matrix{NumType}
end

"""
Pre-allocated memory for efficiently accumulating certain sub-matrices
of the E_G_s and E_G2_s Hessian.

Args:
    NumType: The numeric type of the hessian.
    i: The type of celestial source, from 1:Ia
"""
function HessianSubmatrices(NumType::DataType, i::Int)
    @assert 1 <= i <= Ia
    shape_p = length(shape_standard_alignment[i])

    u_u = zeros(NumType, 2, 2)
    shape_shape = zeros(NumType, shape_p, shape_p)
    HessianSubmatrices{NumType}(u_u, shape_shape)
end

type ElboIntermediateVariables{NumType <: Number}

    bvn_derivs::BivariateNormalDerivatives{NumType}

    # Vectors of star and galaxy bvn quantities from all sources for a pixel.
    # The vector has one element for each active source, in the same order
    # as ea.active_sources.

    # TODO: you can treat this the same way as E_G_s and not keep a vector around.
    fs0m_vec::Vector{SensitiveFloat{StarPosParams, NumType}}
    fs1m_vec::Vector{SensitiveFloat{GalaxyPosParams, NumType}}

    # Brightness values for a single source
    E_G_s::SensitiveFloat{CanonicalParams, NumType}
    E_G2_s::SensitiveFloat{CanonicalParams, NumType}
    var_G_s::SensitiveFloat{CanonicalParams, NumType}

    # Subsets of the Hessian of E_G_s and E_G2_s that allow us to use BLAS
    # functions to accumulate Hessian terms. There is one submatrix for
    # each celestial object type in 1:Ia
    E_G_s_hsub_vec::Vector{HessianSubmatrices{NumType}}
    E_G2_s_hsub_vec::Vector{HessianSubmatrices{NumType}}

    # Expected pixel intensity and variance for a pixel from all sources.
    E_G::SensitiveFloat{CanonicalParams, NumType}
    var_G::SensitiveFloat{CanonicalParams, NumType}

    # Pre-allocated memory for the gradient and Hessian of combine functions.
    combine_grad::Vector{NumType}
    combine_hess::Matrix{NumType}

    # A placeholder for the log term in the ELBO.
    elbo_log_term::SensitiveFloat{CanonicalParams, NumType}

    # The ELBO itself.
    elbo::SensitiveFloat{CanonicalParams, NumType}

    # If false, do not calculate hessians or derivatives.
    calculate_derivs::Bool

    # If false, do not calculate hessians.
    calculate_hessian::Bool
end

"""
Args:
    - S: The total number of sources
    - num_active_sources: The number of actives sources (with deriviatives)
    - calculate_derivs: If false, only calculate values
    - calculate_hessian: If false, only calculate gradients. Note that if
                calculate_derivs = false, then hessians will not be
                calculated irrespective of the value of calculate_hessian.
"""
function ElboIntermediateVariables(NumType::DataType,
                                   S::Int,
                                   num_active_sources::Int;
                                   calculate_derivs::Bool=true,
                                   calculate_hessian::Bool=true)
    @assert NumType <: Number

    bvn_derivs = BivariateNormalDerivatives{NumType}(NumType)

    # fs0m and fs1m accumulate contributions from all bvn components
    # for a given source.
    fs0m_vec = Array(SensitiveFloat{StarPosParams, NumType}, S)
    fs1m_vec = Array(SensitiveFloat{GalaxyPosParams, NumType}, S)
    for s = 1:S
        fs0m_vec[s] = zero_sensitive_float(StarPosParams, NumType)
        fs1m_vec[s] = zero_sensitive_float(GalaxyPosParams, NumType)
    end

    E_G_s = zero_sensitive_float(CanonicalParams, NumType, 1)
    E_G2_s = zero_sensitive_float(CanonicalParams, NumType, 1)
    var_G_s = zero_sensitive_float(CanonicalParams, NumType, 1)

    E_G_s_hsub_vec =
        HessianSubmatrices{NumType}[ HessianSubmatrices(NumType, i) for i=1:Ia ]
    E_G2_s_hsub_vec =
        HessianSubmatrices{NumType}[ HessianSubmatrices(NumType, i) for i=1:Ia ]

    E_G = zero_sensitive_float(CanonicalParams, NumType, num_active_sources)
    var_G = zero_sensitive_float(CanonicalParams, NumType, num_active_sources)

    combine_grad = zeros(NumType, 2)
    combine_hess = zeros(NumType, 2, 2)

    elbo_log_term =
        zero_sensitive_float(CanonicalParams, NumType, num_active_sources)
    elbo = zero_sensitive_float(CanonicalParams, NumType, num_active_sources)

    ElboIntermediateVariables{NumType}(
        bvn_derivs, fs0m_vec, fs1m_vec,
        E_G_s, E_G2_s, var_G_s, E_G_s_hsub_vec, E_G2_s_hsub_vec,
        E_G, var_G, combine_grad, combine_hess,
        elbo_log_term, elbo, calculate_derivs, calculate_hessian)
end

"""
Add the contributions of a star's bivariate normal term to the ELBO,
by updating elbo_vars.fs0m_vec[s] in place.

Args:
    - elbo_vars: Elbo intermediate values.
    - s: The index of the current source in 1:S
    - bmc: The component to be added
    - x: An offset for the component in pixel coordinates (e.g. a pixel location)
    - wcs_jacobian: The jacobian of the function pixel = F(world) at this location.
    - calculate_derivs: Whether to calculate derivatives.

Returns:
    Updates elbo_vars.fs0m_vec[s] in place.
"""
function accum_star_pos!{NumType <: Number}(
                    elbo_vars::ElboIntermediateVariables{NumType},
                    s::Int,
                    bmc::BvnComponent{NumType},
                    x::Vector{Float64},
                    wcs_jacobian::Array{Float64, 2},
                    calculate_derivs::Bool)
    eval_bvn_pdf!(elbo_vars.bvn_derivs, bmc, x)

    # TODO: Also make a version that doesn't calculate any derivatives
    # if the object isn't in active_sources.
    get_bvn_derivs!(elbo_vars.bvn_derivs, bmc, true, false)

    fs0m = elbo_vars.fs0m_vec[s]
    fs0m.v[1] += elbo_vars.bvn_derivs.f_pre[1]

    if elbo_vars.calculate_derivs && calculate_derivs
        transform_bvn_ux_derivs!(
            elbo_vars.bvn_derivs, wcs_jacobian, elbo_vars.calculate_hessian)
        bvn_u_d = elbo_vars.bvn_derivs.bvn_u_d
        bvn_uu_h = elbo_vars.bvn_derivs.bvn_uu_h

        # Accumulate the derivatives.
        for u_id in 1:2
            fs0m.d[star_ids.u[u_id]] += elbo_vars.bvn_derivs.f_pre[1] * bvn_u_d[u_id]
        end

        if elbo_vars.calculate_hessian
            # Hessian terms involving only the location parameters.
            # TODO: redundant term
            for u_id1 in 1:2, u_id2 in 1:2
                fs0m.h[star_ids.u[u_id1], star_ids.u[u_id2]] +=
                    elbo_vars.bvn_derivs.f_pre[1] * (bvn_uu_h[u_id1, u_id2] +
                    bvn_u_d[u_id1] * bvn_u_d[u_id2])
            end
        end
    end
end

"""
Add the contributions of a galaxy component term to the ELBO by
updating fs1m in place.

Args:
    - elbo_vars: Elbo intermediate variables
    - s: The index of the current source in 1:S
    - gcc: The galaxy component to be added
    - x: An offset for the component in pixel coordinates (e.g. a pixel location)
    - wcs_jacobian: The jacobian of the function pixel = F(world) at this location.
    - calculate_derivs: Whether to calculate derivatives.

Returns:
    Updates elbo_vars.fs1m_vec[s] in place.
"""
function accum_galaxy_pos!{NumType <: Number}(
                    elbo_vars::ElboIntermediateVariables{NumType},
                    s::Int,
                    gcc::GalaxyCacheComponent{NumType},
                    x::Vector{Float64},
                    wcs_jacobian::Array{Float64, 2},
                    calculate_derivs::Bool)
    eval_bvn_pdf!(elbo_vars.bvn_derivs, gcc.bmc, x)
    f = elbo_vars.bvn_derivs.f_pre[1] * gcc.e_dev_i
    fs1m = elbo_vars.fs1m_vec[s]
    fs1m.v[1] += f

    if elbo_vars.calculate_derivs && calculate_derivs

        get_bvn_derivs!(elbo_vars.bvn_derivs, gcc.bmc,
            elbo_vars.calculate_hessian, elbo_vars.calculate_hessian)
        transform_bvn_derivs!(
            elbo_vars.bvn_derivs, gcc.sig_sf, wcs_jacobian, elbo_vars.calculate_hessian)

        bvn_u_d = elbo_vars.bvn_derivs.bvn_u_d
        bvn_uu_h = elbo_vars.bvn_derivs.bvn_uu_h
        bvn_s_d = elbo_vars.bvn_derivs.bvn_s_d
        bvn_ss_h = elbo_vars.bvn_derivs.bvn_ss_h
        bvn_us_h = elbo_vars.bvn_derivs.bvn_us_h

        # Accumulate the derivatives.
        for u_id in 1:2
            fs1m.d[gal_ids.u[u_id]] += f * bvn_u_d[u_id]
        end

        for gal_id in 1:length(gal_shape_ids)
            fs1m.d[gal_shape_alignment[gal_id]] += f * bvn_s_d[gal_id]
        end

        # The e_dev derivative. e_dev just scales the entire component.
        # The direction is positive or negative depending on whether this
        # is an exp or dev component.
        fs1m.d[gal_ids.e_dev] += gcc.e_dev_dir * elbo_vars.bvn_derivs.f_pre[1]

        if elbo_vars.calculate_hessian
            # The Hessians:

            # Hessian terms involving only the shape parameters.
            for shape_id1 in 1:length(gal_shape_ids)
                for shape_id2 in 1:length(gal_shape_ids)
                    s1 = gal_shape_alignment[shape_id1]
                    s2 = gal_shape_alignment[shape_id2]
                    fs1m.h[s1, s2] +=
                        f * (bvn_ss_h[shape_id1, shape_id2] +
                                 bvn_s_d[shape_id1] * bvn_s_d[shape_id2])
                end
            end

            # Hessian terms involving only the location parameters.
            for u_id1 in 1:2, u_id2 in 1:2
                u1 = gal_ids.u[u_id1]
                u2 = gal_ids.u[u_id2]
                fs1m.h[u1, u2] +=
                    f * (bvn_uu_h[u_id1, u_id2] + bvn_u_d[u_id1] * bvn_u_d[u_id2])
            end

            # Hessian terms involving both the shape and location parameters.
            for u_id in 1:2, shape_id in 1:length(gal_shape_ids)
                ui = gal_ids.u[u_id]
                si = gal_shape_alignment[shape_id]
                fs1m.h[ui, si] +=
                    f * (bvn_us_h[u_id, shape_id] + bvn_u_d[u_id] * bvn_s_d[shape_id])
                fs1m.h[si, ui] = fs1m.h[ui, si]
            end

            # Do the e_dev hessian terms.
            devi = gal_ids.e_dev
            for u_id in 1:2
                ui = gal_ids.u[u_id]
                fs1m.h[ui, devi] +=
                    elbo_vars.bvn_derivs.f_pre[1] * gcc.e_dev_dir * bvn_u_d[u_id]
                fs1m.h[devi, ui] = fs1m.h[ui, devi]
            end
            for shape_id in 1:length(gal_shape_ids)
                si = gal_shape_alignment[shape_id]
                fs1m.h[si, devi] +=
                    elbo_vars.bvn_derivs.f_pre[1] * gcc.e_dev_dir * bvn_s_d[shape_id]
                fs1m.h[devi, si] = fs1m.h[si, devi]
            end
        end # if calculate hessian
    end # if calculate_derivs
end

"""
Populate fs0m_vec and fs1m_vec for all sources for a given pixel.

Args:
    - elbo_vars: Elbo intermediate values.
    - ea: Model parameters
    - tile_sources: A vector of integers of sources in 1:ea.S affecting the tile
    - tile: An ImageTile
    - h, w: The integer locations of the pixel within the tile
    - gal_mcs: Galaxy components
    - star_mcs: Star components

Returns:
    Updates elbo_vars.fs0m_vec and elbo_vars.fs1m_vec in place with the total
    shape contributions to this pixel's brightness.
"""
function populate_fsm_vecs!{NumType <: Number}(
                    elbo_vars::ElboIntermediateVariables{NumType},
                    ea::ElboArgs{NumType},
                    tile_sources::Vector{Int},
                    tile::ImageTile,
                    h::Int, w::Int,
                    gal_mcs::Array{GalaxyCacheComponent{NumType}, 4},
                    star_mcs::Array{BvnComponent{NumType}, 2})
    x = Float64[tile.h_range[h], tile.w_range[w]]
    for s in tile_sources
        # ensure tile.b is a filter band, not an image's index
        @assert 1 <= tile.b <= B
        wcs_jacobian = ea.patches[s, tile.b].wcs_jacobian
        active_source = s in ea.active_sources

        calculate_hessian =
            elbo_vars.calculate_hessian && elbo_vars.calculate_derivs && active_source
        clear!(elbo_vars.fs0m_vec[s], calculate_hessian)
        for k = 1:psf_K # PSF component
            accum_star_pos!(
                elbo_vars, s, star_mcs[k, s], x, wcs_jacobian, active_source)
        end

        clear!(elbo_vars.fs1m_vec[s], calculate_hessian)
        for i = 1:2 # Galaxy types
            for j in 1:8 # Galaxy component
                # If i == 2 then there are only six galaxy components.
                if (i == 1) || (j <= 6)
                    for k = 1:psf_K # PSF component
                        accum_galaxy_pos!(
                            elbo_vars, s, gal_mcs[k, j, i, s], x, wcs_jacobian,
                            active_source)
                    end
                end
            end
        end
    end
end

"""
Add the contributions of a single source to E_G_s and var_G_s, which are cleared
and then updated in place.

Args:
    - elbo_vars: Elbo intermediate values, with updated fs1m_vec and fs0m_vec.
    - ea: Model parameters
    - sbs: Source brightnesses
    - s: The source, in 1:ea.S
    - b: The band

Returns:
    Updates elbo_vars.E_G_s and elbo_vars.var_G_s in place with the brightness
    for this sourve at this pixel.
"""
function accumulate_source_brightness!{NumType <: Number}(
                    elbo_vars::ElboIntermediateVariables{NumType},
                    ea::ElboArgs{NumType},
                    sbs::Vector{SourceBrightness{NumType}},
                    s::Int, b::Int)
    # E[G] and E{G ^ 2} for a single source
    E_G_s = elbo_vars.E_G_s
    E_G2_s = elbo_vars.E_G2_s

    clear_hessian = elbo_vars.calculate_hessian && elbo_vars.calculate_derivs
    clear!(E_G_s, clear_hessian)
    clear!(E_G2_s, clear_hessian)

    sb = sbs[s]

    active_source = (s in ea.active_sources)

    for i in 1:Ia # Stars and galaxies
        fsm_i = (i == 1) ? elbo_vars.fs0m_vec[s] : elbo_vars.fs1m_vec[s]
        a_i = ea.vp[s][ids.a[i, 1]]

        lf = sb.E_l_a[b, i].v[1] * fsm_i.v[1]
        llff = sb.E_ll_a[b, i].v[1] * fsm_i.v[1]^2

        E_G_s.v[1] += a_i * lf
        E_G2_s.v[1] += a_i * llff

        # Only calculate derivatives for active sources.
        if active_source && elbo_vars.calculate_derivs
            ######################
            # Gradients.

            E_G_s.d[ids.a[i, 1], 1] += lf
            E_G2_s.d[ids.a[i, 1], 1] += llff

            p0_shape = shape_standard_alignment[i]
            p0_bright = brightness_standard_alignment[i]
            u_ind = i == 1 ? star_ids.u : gal_ids.u

            # Derivatives with respect to the spatial parameters
            for p0_shape_ind in 1:length(p0_shape)
                E_G_s.d[p0_shape[p0_shape_ind], 1] +=
                    sb.E_l_a[b, i].v[1] * a_i * fsm_i.d[p0_shape_ind, 1]
                E_G2_s.d[p0_shape[p0_shape_ind], 1] +=
                    sb.E_ll_a[b, i].v[1] * 2 * fsm_i.v[1] * a_i * fsm_i.d[p0_shape_ind, 1]
            end

            # Derivatives with respect to the brightness parameters.
            for p0_bright_ind in 1:length(p0_bright)
                E_G_s.d[p0_bright[p0_bright_ind], 1] +=
                    a_i * fsm_i.v[1] * sb.E_l_a[b, i].d[p0_bright_ind, 1]
                E_G2_s.d[p0_bright[p0_bright_ind], 1] +=
                    a_i * (fsm_i.v[1]^2) * sb.E_ll_a[b, i].d[p0_bright_ind, 1]
            end

            if elbo_vars.calculate_hessian
                ######################
                # Hessians.

                # Data structures to accumulate certain submatrices of the Hessian.
                E_G_s_hsub = elbo_vars.E_G_s_hsub_vec[i]
                E_G2_s_hsub = elbo_vars.E_G2_s_hsub_vec[i]

                # The (a, a) block of the hessian is zero.

                # The (bright, bright) block:
                for p0_ind1 in 1:length(p0_bright), p0_ind2 in 1:length(p0_bright)
                    # TODO: time consuming **************
                    E_G_s.h[p0_bright[p0_ind1], p0_bright[p0_ind2]] =
                        a_i * sb.E_l_a[b, i].h[p0_ind1, p0_ind2] * fsm_i.v[1]
                    E_G2_s.h[p0_bright[p0_ind1], p0_bright[p0_ind2]] =
                        (fsm_i.v[1]^2) * a_i * sb.E_ll_a[b, i].h[p0_ind1, p0_ind2]
                end

                # The (shape, shape) block:
                p1, p2 = size(E_G_s_hsub.shape_shape)
                for ind1 = 1:p1, ind2 = 1:p2
                    E_G_s_hsub.shape_shape[ind1, ind2] =
                        a_i * sb.E_l_a[b, i].v[1] * fsm_i.h[ind1, ind2]
                    E_G2_s_hsub.shape_shape[ind1, ind2] =
                        2 * a_i * sb.E_ll_a[b, i].v[1] * (
                            fsm_i.v[1] * fsm_i.h[ind1, ind2] +
                            fsm_i.d[ind1, 1] * fsm_i.d[ind2, 1])
                end

                # The u_u submatrix of this assignment will be overwritten after
                # the loop.
                for p0_ind1 in 1:length(p0_shape), p0_ind2 in 1:length(p0_shape)
                    E_G_s.h[p0_shape[p0_ind1], p0_shape[p0_ind2]] =
                        a_i * sb.E_l_a[b, i].v[1] * fsm_i.h[p0_ind1, p0_ind2]
                    E_G2_s.h[p0_shape[p0_ind1], p0_shape[p0_ind2]] =
                        E_G2_s_hsub.shape_shape[p0_ind1, p0_ind2]
                end

                # Since the u_u submatrix is not disjoint between different i, accumulate
                # it separate and add it at the end.
                for u_ind1 = 1:2, u_ind2 = 1:2
                    E_G_s_hsub.u_u[u_ind1, u_ind2] =
                        E_G_s_hsub.shape_shape[u_ind[u_ind1], u_ind[u_ind2]]
                    E_G2_s_hsub.u_u[u_ind1, u_ind2] =
                        E_G2_s_hsub.shape_shape[u_ind[u_ind1], u_ind[u_ind2]]
                end

                # All other terms are disjoint between different i and don't involve
                # addition, so we can just assign their values (which is efficient in
                # native julia).

                # The (a, bright) blocks:
                for p0_ind in 1:length(p0_bright)
                    E_G_s.h[p0_bright[p0_ind], ids.a[i, 1]] =
                        fsm_i.v[1] * sb.E_l_a[b, i].d[p0_ind, 1]
                    E_G2_s.h[p0_bright[p0_ind], ids.a[i, 1]] =
                        (fsm_i.v[1] ^ 2) * sb.E_ll_a[b, i].d[p0_ind, 1]
                    E_G_s.h[ids.a[i, 1], p0_bright[p0_ind]] =
                        E_G_s.h[p0_bright[p0_ind], ids.a[i, 1]]
                    E_G2_s.h[ids.a[i, 1], p0_bright[p0_ind]] =
                        E_G2_s.h[p0_bright[p0_ind], ids.a[i, 1]]
                end

                # The (a, shape) blocks.
                for p0_ind in 1:length(p0_shape)
                    E_G_s.h[p0_shape[p0_ind], ids.a[i, 1]] =
                        sb.E_l_a[b, i].v[1] * fsm_i.d[p0_ind, 1]
                    E_G2_s.h[p0_shape[p0_ind], ids.a[i, 1]] =
                        sb.E_ll_a[b, i].v[1] * 2 * fsm_i.v[1] * fsm_i.d[p0_ind, 1]
                    E_G_s.h[ids.a[i, 1], p0_shape[p0_ind]] =
                        E_G_s.h[p0_shape[p0_ind], ids.a[i, 1]]
                    E_G2_s.h[ids.a[i, 1], p0_shape[p0_ind]] =
                        E_G2_s.h[p0_shape[p0_ind], ids.a[i, 1]]
                end

                for ind_b in 1:length(p0_bright), ind_s in 1:length(p0_shape)
                    E_G_s.h[p0_bright[ind_b], p0_shape[ind_s]] =
                        a_i * sb.E_l_a[b, i].d[ind_b, 1] * fsm_i.d[ind_s, 1]
                    E_G2_s.h[p0_bright[ind_b], p0_shape[ind_s]] =
                        2 * a_i * sb.E_ll_a[b, i].d[ind_b, 1] * fsm_i.v[1] * fsm_i.d[ind_s]

                    E_G_s.h[p0_shape[ind_s], p0_bright[ind_b]] =
                        E_G_s.h[p0_bright[ind_b], p0_shape[ind_s]]
                    E_G2_s.h[p0_shape[ind_s], p0_bright[ind_b]] =
                        E_G2_s.h[p0_bright[ind_b], p0_shape[ind_s]]
                end
            end # if calculate hessian
        end # if calculate derivatives
    end # i loop

    if elbo_vars.calculate_hessian
        # Accumulate the u Hessian. u is the only parameter that is shared between
        # different values of i.

        # This is
        # for i = 1:Ia
        #     E_G_u_u_hess += elbo_vars.E_G_s_hsub_vec[i].u_u
        #     E_G2_u_u_hess += elbo_vars.E_G2_s_hsub_vec[i].u_u
        # end
        # For each value in 1:Ia, written this way for speed.
        @assert Ia == 2
        for u_ind1 = 1:2, u_ind2 = 1:2
            E_G_s.h[ids.u[u_ind1], ids.u[u_ind2]] =
            elbo_vars.E_G_s_hsub_vec[1].u_u[u_ind1, u_ind2] +
            elbo_vars.E_G_s_hsub_vec[2].u_u[u_ind1, u_ind2]

            E_G2_s.h[ids.u[u_ind1], ids.u[u_ind2]] =
                elbo_vars.E_G2_s_hsub_vec[1].u_u[u_ind1, u_ind2] +
                elbo_vars.E_G2_s_hsub_vec[2].u_u[u_ind1, u_ind2]
        end
    end

    calculate_var_G_s!(elbo_vars, active_source)
end

"""
Calculate the variance var_G_s as a function of (E_G_s, E_G2_s).

Args:
    - elbo_vars: Elbo intermediate values.
    - active_source: Whether this is an active source that requires derivatives

Returns:
    Updates elbo_vars.var_G_s in place.
"""
function calculate_var_G_s!{NumType <: Number}(
                    elbo_vars::ElboIntermediateVariables{NumType},
                    active_source::Bool)
    var_G_s = elbo_vars.var_G_s
    E_G_s = elbo_vars.E_G_s
    E_G2_s = elbo_vars.E_G2_s

    clear!(var_G_s,
        elbo_vars.calculate_hessian &&
            elbo_vars.calculate_derivs && active_source)

    elbo_vars.var_G_s.v[1] = E_G2_s.v[1] - (E_G_s.v[1] ^ 2)

    if active_source && elbo_vars.calculate_derivs
        @assert length(var_G_s.d) == length(E_G2_s.d) == length(E_G_s.d)
        @inbounds for ind1 = 1:length(var_G_s.d)
            var_G_s.d[ind1] = E_G2_s.d[ind1] - 2 * E_G_s.v[1] * E_G_s.d[ind1]
        end

        if elbo_vars.calculate_hessian
            p1, p2 = size(var_G_s.h)
            @inbounds for ind2 = 1:p2, ind1 = 1:ind2
                var_G_s.h[ind1, ind2] =
                    E_G2_s.h[ind1, ind2] - 2 * (
                        E_G_s.v[1] * E_G_s.h[ind1, ind2] +
                        E_G_s.d[ind1, 1] * E_G_s.d[ind2, 1])
                var_G_s.h[ind2, ind1] = var_G_s.h[ind1, ind2]
            end
        end
    end
end

"""
Adds up E_G and var_G across all sources.

Args:
    - elbo_vars: Elbo intermediate values, with updated fs1m_vec and fs0m_vec.
    - ea: Model parameters
    - tile_sources: A vector of integers of sources in 1:ea.S affecting the tile
    - tile: An ImageTile
    - sbs: Source brightnesses

Updates elbo_vars.E_G and elbo_vars.var_G in place.
"""
function combine_pixel_sources!{NumType <: Number}(
                    elbo_vars::ElboIntermediateVariables{NumType},
                    ea::ElboArgs{NumType},
                    tile_sources::Vector{Int},
                    tile::ImageTile,
                    sbs::Vector{SourceBrightness{NumType}})
    clear!(elbo_vars.E_G,
        elbo_vars.calculate_hessian && elbo_vars.calculate_derivs)
    clear!(elbo_vars.var_G,
        elbo_vars.calculate_hessian && elbo_vars.calculate_derivs)

    for s in tile_sources
        active_source = s in ea.active_sources
        calculate_hessian =
            elbo_vars.calculate_hessian && elbo_vars.calculate_derivs &&
            active_source
        accumulate_source_brightness!(elbo_vars, ea, sbs, s, tile.b)
        if active_source
            sa = findfirst(ea.active_sources, s)
            add_sources_sf!(elbo_vars.E_G, elbo_vars.E_G_s, sa, calculate_hessian)
            add_sources_sf!(elbo_vars.var_G, elbo_vars.var_G_s, sa, calculate_hessian)
        else
            # If the sources is inactives, simply accumulate the values.
            elbo_vars.E_G.v[1] += elbo_vars.E_G_s.v[1]
            elbo_vars.var_G.v[1] += elbo_vars.var_G_s.v[1]
        end
    end
end

"""
Expected brightness for a single pixel.

Args:
    - elbo_vars: Elbo intermediate values.
    - h, w: The integer locations of the pixel within the tile
    - sbs: Source brightnesses
    - star_mcs: Star components
    - gal_mcs: Galaxy components
    - tile: An ImageTile
    - ea: Model parameters
    - tile_sources: A vector of integers of sources in 1:ea.S affecting the tile
    - include_epsilon: Whether the background noise should be included

Returns:
    - Updates elbo_vars.E_G and elbo_vars.var_G in place.
"""
function get_expected_pixel_brightness!{NumType <: Number}(
                elbo_vars::ElboIntermediateVariables{NumType},
                h::Int, w::Int,
                sbs::Vector{SourceBrightness{NumType}},
                star_mcs::Array{BvnComponent{NumType}, 2},
                gal_mcs::Array{GalaxyCacheComponent{NumType}, 4},
                tile::ImageTile,
                ea::ElboArgs{NumType},
                tile_sources::Vector{Int};
                include_epsilon::Bool=true)
    # This combines the bvn components to get the brightness for each
    # source separately.
    populate_fsm_vecs!(
        elbo_vars, ea, tile_sources, tile, h, w, gal_mcs, star_mcs)

    # # This combines the sources into a single brightness value for the pixel.
    combine_pixel_sources!(elbo_vars, ea, tile_sources, tile, sbs)

    if include_epsilon
        # There are no derivatives with respect to epsilon, so can safely add
        # to the value.
        elbo_vars.E_G.v[1] += tile.epsilon_mat[h, w]
    end
end

"""
Add the lower bound to the log term to the elbo for a single pixel.

Args:
     - elbo_vars: Intermediate variables
     - x_nbm: The photon count at this pixel
     - iota: The optical sensitivity

 Returns:
    Updates elbo_vars.elbo in place by adding the lower bound to the log
    term.
"""
function add_elbo_log_term!{NumType <: Number}(
                elbo_vars::ElboIntermediateVariables{NumType},
                x_nbm::AbstractFloat, iota::AbstractFloat)
    # See notes for a derivation. The log term is
    # log E[G] - Var(G) / (2 * E[G] ^2 )

    E_G = elbo_vars.E_G
    var_G = elbo_vars.var_G
    elbo = elbo_vars.elbo

    # The gradients and Hessians are written as a f(x, y) = f(E_G2, E_G)
    log_term_value = log(E_G.v[1]) - 0.5 * var_G.v[1]    / (E_G.v[1] ^ 2)

    # Add x_nbm * (log term * log(iota)) to the elbo.
    # If not calculating derivatives, add the values directly.
    elbo.v[1] += x_nbm * (log(iota) + log_term_value)

    if elbo_vars.calculate_derivs
        elbo_vars.combine_grad[1] = -0.5 / (E_G.v[1] ^ 2)
        elbo_vars.combine_grad[2] = 1 / E_G.v[1] + var_G.v[1] / (E_G.v[1] ^ 3)

        if elbo_vars.calculate_hessian
            elbo_vars.combine_hess[1, 1] = 0.0
            elbo_vars.combine_hess[1, 2] = elbo_vars.combine_hess[2, 1] = 1 / E_G.v[1]^3
            elbo_vars.combine_hess[2, 2] =
                -(1 / E_G.v[1] ^ 2 + 3    * var_G.v[1] / (E_G.v[1] ^ 4))
        end

        # Calculate the log term.
        combine_sfs!(
            elbo_vars.var_G, elbo_vars.E_G, elbo_vars.elbo_log_term,
            log_term_value, elbo_vars.combine_grad, elbo_vars.combine_hess,
            calculate_hessian=elbo_vars.calculate_hessian)

        # Add to the ELBO.
        for ind in 1:length(elbo.d)
            elbo.d[ind] += x_nbm * elbo_vars.elbo_log_term.d[ind]
        end

        if elbo_vars.calculate_hessian
            for ind in 1:length(elbo.h)
                elbo.h[ind] += x_nbm * elbo_vars.elbo_log_term.h[ind]
            end
        end
    end
end

############################################
# The remaining functions loop over tiles, sources, and pixels.

"""
Add an array of pixels' contribution to the ELBO likelihood term by
modifying elbo in place.

Args:
    - elbo_vars: Elbo intermediate values.
    - ea: Model parameters
    - active_pixels: An array of ActivePixels to be processed.

Returns:
    Adds to elbo_vars.elbo in place.
"""
function process_active_pixels!{NumType <: Number}(
                elbo_vars::ElboIntermediateVariables{NumType},
                ea::ElboArgs{NumType},
                active_pixels::Array{ActivePixel})
    sbs = load_source_brightnesses(ea,
        calculate_derivs=elbo_vars.calculate_derivs,
        calculate_hessian=elbo_vars.calculate_hessian)

    star_mcs_vec = Array(Array{BvnComponent{NumType}, 2}, ea.N)
    gal_mcs_vec = Array(Array{GalaxyCacheComponent{NumType}, 4}, ea.N)

    for b=1:ea.N
        # TODO: every elbo_vars should not get to decide its own calculate_*
        star_mcs_vec[b], gal_mcs_vec[b] =
            load_bvn_mixtures(ea, b,
                calculate_derivs=elbo_vars.calculate_derivs,
                calculate_hessian=elbo_vars.calculate_hessian)
    end

    # iterate over the pixels
    for pixel in active_pixels
        tile = ea.images[pixel.n].tiles[pixel.tile_ind]
        tile_sources = ea.tile_source_map[pixel.n][pixel.tile_ind]
        this_pixel = tile.pixels[pixel.h, pixel.w]

        # Get the brightness.
        get_expected_pixel_brightness!(
            elbo_vars, pixel.h, pixel.w, sbs,
            star_mcs_vec[pixel.n], gal_mcs_vec[pixel.n], tile,
            ea, tile_sources, include_epsilon=true)

        # Add the terms to the elbo given the brightness.
        iota = tile.iota_vec[pixel.h]
        add_elbo_log_term!(elbo_vars, this_pixel, iota)
        add_scaled_sfs!(elbo_vars.elbo,
                        elbo_vars.E_G, -iota,
                        elbo_vars.calculate_hessian &&
                        elbo_vars.calculate_derivs)

        # Subtract the log factorial term. This is not a function of the
        # parameters so the derivatives don't need to be updated. Note that
        # even though this does not affect the ELBO's maximum, it affects
        # the optimization convergence criterion, so I will leave it in for now.
        elbo_vars.elbo.v[1] -= lfact(this_pixel)
    end
end

"""
Return the image predicted for the tile given the current parameters.

Args:
    - elbo_vars: Elbo intermediate values.
    - tile: An ImageTile
    - ea: Model parameters
    - tile_sources: A vector of integers of sources in 1:ea.S affecting the tile
    - sbs: Source brightnesses
    - star_mcs: Star components
    - gal_mcs: Galaxy components
    - include_epsilon: Whether the background noise should be included

Returns:
    A matrix of the same size as the tile with the predicted brightnesses.
"""
function tile_predicted_image{NumType <: Number}(
                elbo_vars::ElboIntermediateVariables{NumType},
                tile::ImageTile,
                ea::ElboArgs{NumType},
                tile_sources::Vector{Int},
                sbs::Vector{SourceBrightness{NumType}},
                star_mcs::Array{BvnComponent{NumType}, 2},
                gal_mcs::Array{GalaxyCacheComponent{NumType}, 4};
                include_epsilon::Bool=true)
    predicted_pixels = copy(tile.pixels)
    # Iterate over pixels that are not NaN.
    h_width, w_width = size(tile.pixels)
    for w in 1:w_width, h in 1:h_width
        this_pixel = tile.pixels[h, w]
        if !Base.isnan(this_pixel)
            get_expected_pixel_brightness!(
                elbo_vars, h, w, sbs, star_mcs, gal_mcs, tile,
                ea, tile_sources, include_epsilon=include_epsilon)
            iota = tile.iota_vec[h]
            predicted_pixels[h, w] = elbo_vars.E_G.v[1] * iota
        end
    end

    predicted_pixels
end

"""
Produce a predicted image for a given tile and model parameters.
If include_epsilon is true, then the background is also rendered.
Otherwise, only pixels from the object are rendered.

Args:
    - tile: An ImageTile
    - ea: Model parameters
    - tile_sources: A vector of integers of sources in 1:ea.S affecting the tile
    - include_epsilon: Whether the background noise should be included
"""
function tile_predicted_image{NumType <: Number}(
                    tile::ImageTile,
                    ea::ElboArgs{NumType},
                    tile_sources::Vector{Int};
                    include_epsilon::Bool=false)
    star_mcs, gal_mcs = load_bvn_mixtures(ea, tile.b, calculate_derivs=false)
    sbs = load_source_brightnesses(ea, calculate_derivs=false)

    elbo_vars =
        ElboIntermediateVariables(NumType, ea.S, length(ea.active_sources))
    elbo_vars.calculate_derivs = false

    tile_predicted_image(
        elbo_vars, tile, ea, tile_sources, sbs, star_mcs, gal_mcs,
        include_epsilon=include_epsilon)
end

"""
Get the active pixels (pixels for which the active sources are present).
TODO: move this to pre-processing and use it instead of setting low-signal
pixels to NaN.
"""
function get_active_pixels{NumType <: Number}(
                    ea::ElboArgs{NumType})
    active_pixels = ActivePixel[]

    for n in 1:ea.N, tile_ind in 1:length(ea.images[n].tiles)
        tile_sources = ea.tile_source_map[n][tile_ind]
        if length(intersect(tile_sources, ea.active_sources)) > 0
            tile = ea.images[n].tiles[tile_ind]
            h_width, w_width = size(tile.pixels)
            for w in 1:w_width, h in 1:h_width
                if !Base.isnan(tile.pixels[h, w])
                    push!(active_pixels, ActivePixel(n, tile_ind, h, w))
                end
            end
        end
    end

    active_pixels
end

"""
Return the expected log likelihood for all bands in a section
of the sky.
Returns: A sensitive float with the log,  likelihood.
"""
function elbo_likelihood{NumType <: Number}(
                    ea::ElboArgs{NumType};
                    calculate_derivs=true,
                    calculate_hessian=true)
    active_pixels = get_active_pixels(ea)
    elbo_vars = ElboIntermediateVariables(NumType, ea.S,
                                length(ea.active_sources),
                                calculate_derivs=calculate_derivs,
                                calculate_hessian=calculate_hessian)
    process_active_pixels!(elbo_vars, ea, active_pixels)
    elbo_vars.elbo
end

"""
Calculates and returns the ELBO and its derivatives for all the bands
of an image.
Returns: A sensitive float containing the ELBO for the image.
"""
function elbo{NumType <: Number}(
                    ea::ElboArgs{NumType};
                    calculate_derivs=true,
                    calculate_hessian=true)
    elbo = elbo_likelihood(ea;
        calculate_derivs=calculate_derivs, calculate_hessian=calculate_hessian)
    # TODO: subtract the kl with the hessian.
    subtract_kl!(ea, elbo, calculate_derivs=calculate_derivs)
    elbo
end


text:
using ViewerGL
GL = ViewerGL

points = rand(50,3)
GL.VIEW([
      GL.GLPoints(points),
      GL.GLHull(points,GL.Point4d(1,1,1,0.2)),
      GL.GLAxis(GL.Point3d(0,0,0),GL.Point3d(1,1,1))
]);

ext: jl


text:
module SampleModule
    module Model
        export AppUser
        include("model/AppUser.jl")
    end

    module Controller
        using Dates
        using ..Model
        include("controller/Controller.jl")
    end

end


text:
using DICOMDictionary
using Test

@testset "DICOMDictionary.jl" begin
  make_dicom_dictionary("dcm_dict.jl")
end


text:
using FeedbackParticleFilters

println("Testing FilteringProblem.jl:")

@testset "FilteringProblem.jl" begin

    struct TestStateModel <: HiddenStateModel{Int, ContinuousTime} end
    struct TestObservationModel <: ObservationModel{Int, Float64, DiscreteTime} end

    st_mod    = TestStateModel()
    ob_mod    = TestObservationModel()
    filt_prob = FilteringProblem(st_mod, ob_mod)

    print("  struct FilteringProblem")
    print(".")
    @test filt_prob isa FilteringProblem{Int64,Float64,ContinuousTime,DiscreteTime,TestStateModel,TestObservationModel}
    println("DONE.")

    print("  method state_model")
    print(".")
    @test state_model(filt_prob) isa TestStateModel
    println("DONE.")

    print("  method obs_model")
    print(".")
    @test obs_model(filt_prob) isa TestObservationModel
    println("DONE.")

end; #FilteringProblem.jl


text:
"""
    wider_mlp(layer, next_layer, new_unit_count, noise_factor=0.05)

Net2WiderNet method for MLPs.
Grows one of the layers of the model to a new one with more hidden units.
Function is preserved during the widening.

`layer` should be the `Dense` layer that you wish to widen
`next_layer` should be the `Dense` layer that comes right after `layer`
`new_unit_count` is how many units you wish the new layer to have
`noise_factor` sets how much Gaussian noise to add to the student's weights.
If you want no noise, set `noise_factor` to 0.
"""
function wider_mlp(layer, next_layer, new_unit_count, noise_factor=0.05)
    old_unit_count = size(layer.w, 1)
    if old_unit_count >= new_unit_count
        println("New unit count must be greater than old unit count")
        return nothing
    end

    # Initialize number of extra units, extra arrays, random mapping, number of copies
    extra_count = new_unit_count - old_unit_count
    extra_w = atype()(undef, (extra_count, size(layer.w, 2)))
    extra_next_w = atype()(undef, (size(next_layer.w, 1), extra_count))
    extra_b = atype()(undef, (extra_count))

    extra_mapping = rand(1:old_unit_count, extra_count)
    copy_counts = ones(old_unit_count)
    for i in 1:extra_count
        copy_counts[extra_mapping[i]] += 1
    end
    # Divide weights that are copied more than once
    next_layer.w ./= atype()(copy_counts')

    # Set the extra weights
    for i in 1:extra_count
        extra_w[i, :] = layer.w[extra_mapping[i], :] .+ (noise_factor != 0 ? randn()*noise_factor : 0)
        extra_next_w[:, i] = next_layer.w[:, extra_mapping[i]] .+ (noise_factor != 0 ? randn()*noise_factor : 0)
        extra_b[i] = layer.b[extra_mapping[i]]
    end

    # Update the weights
    layer.w = Param(vcat(layer.w, extra_w))
    layer.b = Param(vcat(layer.b, extra_b))
    next_layer.w = Param(hcat(next_layer.w, extra_next_w))
    return extra_mapping
end

"""
    wider_conv(layer, next_layer, new_channel_count, noise_factor=0.05)

Net2WiderNet method for Convs.
Grows one of the layers of the model to a new one with more filter channels.
Function is preserved during the widening.

`layer` should be the `Conv` or `ConvBN` layer that you wish to widen
`next_layer` should be the `Conv` or `ConvBN` layer that comes right after `layer`
`new_channel_count` is how many filter channels you wish the new layer to have
`noise_factor` sets how much Gaussian noise to add to the student's weights.
If you want no noise, set `noise_factor` to 0.
"""
function wider_conv(layer, next_layer, new_channel_count, noise_factor=0.05)
    old_channel_count = size(layer.w, 4)
    if old_channel_count >= new_channel_count
        println("New channel count must be greater than old channel count")
        return nothing
    end

    # Initialize number of extra channels, extra arrays, random mapping, number of copies
    extra_count = new_channel_count - old_channel_count
    extra_w = Array{Float32}(undef, (size(layer.w)[1:3]..., extra_count))
    extra_next_w = next_layer != nothing ? Array{Float32}(undef, (size(next_layer.w)[1:2]..., extra_count, size(next_layer.w)[4])) : nothing
    extra_b = Array{Float32}(undef, (1, 1, extra_count, 1))

    if layer.bn_params != nothing
        extra_bn_scales = atype()(undef, (extra_count,))
        extra_bn_biases = atype()(undef, (extra_count,))
    end

    extra_mapping = rand(1:old_channel_count, extra_count)

    if next_layer != nothing
        copy_counts = ones(old_channel_count)
        for i in 1:extra_count
            copy_counts[extra_mapping[i]] += 1
        end
        # Divide weights that are copied more than once
        next_layer.w ./= atype()(reshape(copy_counts, (1,1,:,1)))
    end

    array_layer_w = convert(Array{Float32}, layer.w)
    array_next_layer_w = next_layer != nothing ? convert(Array{Float32}, next_layer.w) : nothing
    # Set the extra weights
    for i in 1:extra_count
        extra_w[:, :, :, i] = array_layer_w[:, :, :, extra_mapping[i]] .+ (noise_factor != 0 ? randn()*noise_factor : 0)
        if next_layer != nothing
            extra_next_w[:, :, i, :] = array_next_layer_w[:, :, extra_mapping[i], :] .+ (noise_factor != 0 ? randn()*noise_factor : 0)
        end
        extra_b[1, 1, i, 1] = layer.b[1, 1, extra_mapping[i], 1]
        if layer.bn_params != nothing
            extra_bn_scales[i] = layer.bn_params[extra_mapping[i]]
            extra_bn_biases[i] = layer.bn_params[old_channel_count+extra_mapping[i]]
        end
    end

    extra_w = atype()(extra_w)
    extra_next_w = next_layer != nothing ? atype()(extra_next_w) : nothing
    extra_b = atype()(extra_b)

    # Update the weights
    old_size = size(layer.w)
    w_old_2d = reshape(layer.w, :, size(layer.w)[end])
    w_extra_2d = reshape(extra_w, :, size(extra_w)[end])
    w_new_2d = hcat(w_old_2d, w_extra_2d)
    layer.w = Param(reshape(w_new_2d, old_size[1], old_size[2], old_size[3], :))

    layer.b = Param(reshape(vcat(reshape(layer.b, (size(layer.b, 3),)), reshape(extra_b, (extra_count,))), (1, 1, new_channel_count, 1)))

    if next_layer != nothing
        old_size = size(next_layer.w)
        w_old_2d = reshape(next_layer.w, :, size(next_layer.w)[end])
        w_extra_2d = reshape(extra_next_w, :, size(extra_next_w)[end])
        w_new_2d = vcat(w_old_2d, w_extra_2d)

        next_layer.w = Param(reshape(w_new_2d, old_size[1], old_size[2], :, old_size[4]))
    end
    if layer.bn_params != nothing
        layer.bn_params = Param(cat(layer.bn_params[1:old_channel_count], extra_bn_scales, layer.bn_params[old_channel_count+1:end], extra_bn_biases, dims=1))
    end

    return extra_mapping
end

"""
    wider_inceptionA(inc, next, grow_ratio, noise_factor=0.05)

Net2WiderNet method for InceptionA modules where the next layer is another InceptionA.
Widens all of the convolutional layers in the module by the given ratio.
Function is preserved during the widening.

`inc` should be the `InceptionA` layer that you wish to widen
`next` should be the `InceptionA` layer that comes right after `inc`
`grow_ratio` is the multiplier used for the widening, e.g. 1.5 will grow all `Conv`s to 1.5 times their width.
`noise_factor` sets how much Gaussian noise to add to the student's weights.
If you want no noise, set `noise_factor` to 0.
"""
function wider_inceptionA(inc::InceptionA, next::InceptionA, grow_ratio, noise_factor=0.05)
    # Grow inside layers normally
    wider_conv(inc.c1_before_3, inc.c3, Int64(round(size(inc.c1_before_3.w, 4)*(grow_ratio))), noise_factor)
    wider_conv(inc.c1_before_d3, inc.cd3_1, Int64(round(size(inc.c1_before_d3.w, 4)*(grow_ratio))), noise_factor)
    wider_conv(inc.cd3_1, inc.cd3_2, Int64(round(size(inc.cd3_1.w, 4)*(grow_ratio))), noise_factor)

    old_c1_count = size(inc.c1_alone.w, 4)
    old_c3_count = size(inc.c3.w, 4)
    old_cd3_count = size(inc.cd3_2.w, 4)
    old_c1_ap_count = size(inc.c1_after_pool.w, 4)

    new_c1_count = Int64(round(old_c1_count*(grow_ratio)))
    new_c3_count = Int64(round(old_c3_count*(grow_ratio)))
    new_cd3_count = Int64(round(old_cd3_count*(grow_ratio)))
    new_c1_ap_count = Int64(round(old_c1_ap_count*(grow_ratio)))

    c1_extra_mappings = wider_conv(inc.c1_alone, nothing, new_c1_count, noise_factor)
    c3_extra_mappings = wider_conv(inc.c3, nothing, new_c3_count, noise_factor)
    cd3_extra_mappings = wider_conv(inc.cd3_2, nothing, new_cd3_count, noise_factor)
    c1_ap_extra_mappings = wider_conv(inc.c1_after_pool, nothing, new_c1_ap_count, noise_factor)

    c1_mappings = vcat(1:old_c1_count, c1_extra_mappings)
    c3_mappings = vcat(1:old_c3_count, c3_extra_mappings) .+ old_c1_count
    cd3_mappings = vcat(1:old_cd3_count, cd3_extra_mappings) .+ (old_c1_count + old_c3_count)
    c1_ap_mappings = vcat(1:old_c1_ap_count, c1_ap_extra_mappings) .+ (old_c1_count + old_c3_count + old_cd3_count)
    total_mappings = vcat(c1_mappings, c3_mappings, cd3_mappings, c1_ap_mappings)

    old_channel_count = old_c1_count + old_c3_count + old_cd3_count + old_c1_ap_count
    new_channel_count = new_c1_count + new_c3_count + new_cd3_count + new_c1_ap_count

    copy_counts = zeros(old_channel_count)
    for i in 1:new_channel_count
        copy_counts[total_mappings[i]] += 1
    end

    for c in [next.c1_alone, next.c1_before_3, next.c1_before_d3, next.c1_after_pool]
        new_w = Array{Float32}(undef, (size(c.w)[1:2]..., new_channel_count, size(c.w)[4]))

        # Divide weights that are copied more than once
        c.w ./= atype()(reshape(copy_counts, (1,1,:,1)))

        array_w = convert(Array{Float32}, c.w)
        # Set the extra weights
        for i in 1:new_channel_count
            new_w[:, :, i, :] = array_w[:, :, total_mappings[i], :] .+ (noise_factor != 0 ? randn()*noise_factor : 0)
        end
        c.w = Param(atype()(new_w))
    end
end

"""
    wider_inceptionA(inc, next, after_b, grow_ratio, noise_factor=0.05)

Net2WiderNet method for InceptionA modules where the next layer is an InceptionB.
Widens all of the convolutional layers in the module by the given ratio.
Function is preserved during the widening.
This function requires a third layer to be given, as it will also be affected by the widening.

`inc` should be the `InceptionA` layer that you wish to widen
`next` should be the `InceptionB` layer that comes right after `inc`
`after_b` should be the `ConvBN` layer that comes right after `next`
`grow_ratio` is the multiplier used for the widening, e.g. 1.5 will grow all `Conv`s to 1.5 times their width.
`noise_factor` sets how much Gaussian noise to add to the student's weights.
If you want no noise, set `noise_factor` to 0.
"""
function wider_inceptionA(inc::InceptionA, next::InceptionB, after_b::ConvBN, grow_ratio, noise_factor=0.05)
    wider_conv(inc.c1_before_3, inc.c3, Int64(round(size(inc.c1_before_3.w, 4)*(grow_ratio))), noise_factor)
    wider_conv(inc.c1_before_d3, inc.cd3_1, Int64(round(size(inc.c1_before_d3.w, 4)*(grow_ratio))), noise_factor)
    wider_conv(inc.cd3_1, inc.cd3_2, Int64(round(size(inc.cd3_1.w, 4)*(grow_ratio))), noise_factor)

    old_c1_count = size(inc.c1_alone.w, 4)
    old_c3_count = size(inc.c3.w, 4)
    old_cd3_count = size(inc.cd3_2.w, 4)
    old_c1_ap_count = size(inc.c1_after_pool.w, 4)

    new_c1_count = Int64(round(old_c1_count*(grow_ratio)))
    new_c3_count = Int64(round(old_c3_count*(grow_ratio)))
    new_cd3_count = Int64(round(old_cd3_count*(grow_ratio)))
    new_c1_ap_count = Int64(round(old_c1_ap_count*(grow_ratio)))

    c1_extra_mappings = wider_conv(inc.c1_alone, nothing, new_c1_count, noise_factor)
    c3_extra_mappings = wider_conv(inc.c3, nothing, new_c3_count, noise_factor)
    cd3_extra_mappings = wider_conv(inc.cd3_2, nothing, new_cd3_count, noise_factor)
    c1_ap_extra_mappings = wider_conv(inc.c1_after_pool, nothing, new_c1_ap_count, noise_factor)

    c1_mappings = vcat(1:old_c1_count, c1_extra_mappings)
    c3_mappings = vcat(1:old_c3_count, c3_extra_mappings) .+ old_c1_count
    cd3_mappings = vcat(1:old_cd3_count, cd3_extra_mappings) .+ (old_c1_count + old_c3_count)
    c1_ap_mappings = vcat(1:old_c1_ap_count, c1_ap_extra_mappings) .+ (old_c1_count + old_c3_count + old_cd3_count)
    total_mappings = vcat(c1_mappings, c3_mappings, cd3_mappings, c1_ap_mappings)

    old_channel_count = old_c1_count + old_c3_count + old_cd3_count + old_c1_ap_count
    new_channel_count = new_c1_count + new_c3_count + new_cd3_count + new_c1_ap_count

    copy_counts = zeros(old_channel_count)
    for i in 1:new_channel_count
        copy_counts[total_mappings[i]] += 1
    end

    for c in [next.c1_before_3, next.c1_before_d3]
        new_w = Array{Float32}(undef, (size(c.w)[1:2]..., new_channel_count, size(c.w)[4]))

        # Divide weights that are copied more than once
        c.w ./= atype()(reshape(copy_counts, (1,1,:,1)))

        array_w = convert(Array{Float32}, c.w)
        # Set the extra weights
        for i in 1:new_channel_count
            new_w[:, :, i, :] = array_w[:, :, total_mappings[i], :] .+ (noise_factor != 0 ? randn()*noise_factor : 0)
        end
        c.w = Param(atype()(new_w))
    end

    # Finally, need to change the layer after InceptionB as well due to the pass-through layer
    array_w = convert(Array{Float32}, after_b.w)
    w_first = array_w[:, :, 1:end-old_channel_count, :]
    w_last = array_w[:, :, end-old_channel_count+1:end, :]
    new_w_last = Array{Float32}(undef, (size(array_w)[1:2]..., new_channel_count, size(array_w)[4]))

    # Divide weights that are copied more than once
    w_last ./= (reshape(copy_counts, (1,1,:,1)))

    # Set the extra weights
    for i in 1:new_channel_count
        new_w_last[:, :, i, :] = w_last[:, :, total_mappings[i], :]
    end
    after_b.w = Param(atype()(cat(w_first, new_w_last, dims=3)))
end

"""
    wider_inceptionB(inc, next, grow_ratio, noise_factor=0.05)

Net2WiderNet method for InceptionB modules where the next layer is a convolution with batchnorm layer.
Widens all of the convolutional layers in the module by the given ratio.
Function is preserved during the widening.

`inc` should be the `InceptionB` layer that you wish to widen
`next` should be the `ConvBN` layer that comes right after `inc`
`grow_ratio` is the multiplier used for the widening, e.g. 1.5 will grow all `Conv`s to 1.5 times their width.
`noise_factor` sets how much Gaussian noise to add to the student's weights.
If you want no noise, set `noise_factor` to 0.
"""
function wider_inceptionB(inc::InceptionB, next::ConvBN, grow_ratio, noise_factor=0.05)
    # Grow inside layers normally
    wider_conv(inc.c1_before_3, inc.c3, Int64(round(size(inc.c1_before_3.w, 4)*(grow_ratio))), noise_factor)
    wider_conv(inc.c1_before_d3, inc.cd3_1, Int64(round(size(inc.c1_before_d3.w, 4)*(grow_ratio))), noise_factor)
    wider_conv(inc.cd3_1, inc.cd3_2, Int64(round(size(inc.cd3_1.w, 4)*(grow_ratio))), noise_factor)

    old_c3_count = size(inc.c3.w, 4)
    old_cd3_count = size(inc.cd3_2.w, 4)

    new_c3_count = Int64(round(old_c3_count*(grow_ratio)))
    new_cd3_count = Int64(round(old_cd3_count*(grow_ratio)))

    c3_extra_mappings = wider_conv(inc.c3, nothing, new_c3_count, noise_factor)
    cd3_extra_mappings = wider_conv(inc.cd3_2, nothing, new_cd3_count, noise_factor)

    c3_mappings = vcat(1:old_c3_count, c3_extra_mappings)
    cd3_mappings = vcat(1:old_cd3_count, cd3_extra_mappings) .+ old_c3_count
    total_mappings = vcat(c3_mappings, cd3_mappings)

    old_channel_count = old_c3_count + old_cd3_count
    new_channel_count = new_c3_count + new_cd3_count

    copy_counts = zeros(old_channel_count)
    for i in 1:new_channel_count
        copy_counts[total_mappings[i]] += 1
    end

    # Changing the layer after InceptionB
    array_w = convert(Array{Float32}, next.w)
    w_first = array_w[:, :, 1:old_channel_count, :]
    w_last = array_w[:, :, old_channel_count+1:end, :]
    new_w_first = Array{Float32}(undef, (size(array_w)[1:2]..., new_channel_count, size(array_w)[4]))

    # Divide weights that are copied more than once
    w_first ./= (reshape(copy_counts, (1,1,:,1)))

    # Set the extra weights
    for i in 1:new_channel_count
        new_w_first[:, :, i, :] = w_first[:, :, total_mappings[i], :] .+ (noise_factor != 0 ? randn()*noise_factor : 0)
    end
    next.w = Param(atype()(cat(new_w_first, w_last, dims=3)))
end

#-------------------------------------------------------------------------------
# ---------------------------- RANDOM PADDING BEGINS ---------------------------
#-------------------------------------------------------------------------------

"""
    random_pad_mlp(layer, next_layer, new_unit_count)

Baseline method for Net2WiderNet for MLP layers.
Grows one of the layers of the model to a new one with a larger size.
The old units are copied as is, new units are initialized randomly.

`layer` should be the `Dense` layer that you wish to widen
`next_layer` should be the `Dense` layer that comes right after `layer`
`new_unit_count` is how many units you wish the new layer to have
"""
function random_pad_mlp(layer, next_layer, new_unit_count)
    old_unit_count = size(layer.w, 1)
    if old_unit_count >= new_unit_count
        println("New unit count must be greater than old unit count")
        return nothing
    end

    padding_count = new_unit_count - old_unit_count

    padding_1 = param(padding_count, size(layer.w, 2))
    padding_2 = param(size(next_layer.w, 1), padding_count)
    padding_b = param0(padding_count)

    layer.w = Param(vcat(layer.w, padding_1))
    layer.b = Param(vcat(layer.b, padding_b))
    next_layer.w = Param(hcat(next_layer.w, padding_2))

    return new_model
end

"""
    random_pad_conv(layer, next_layer, new_channel_count)

Baseline method for Net2WiderNet for Conv layers.
Grows one of the layers of the model to a new one with a larger size.
The old units are copied as is, new units are initialized randomly.

`layer` should be the `Conv` or `ConvBN` layer that you wish to widen
`next_layer` should be the `Conv` or `ConvBN` layer that comes right after `layer`
`new_channel_count` is how many filter channels you wish the new layer to have
"""
function random_pad_conv(layer, next_layer, new_channel_count)
    old_channel_count = size(layer.w, 4)
    if old_channel_count >= new_channel_count
        println("New channel count must be greater than old channel count")
        return nothing
    end

    # Initialize number of extra channels, extra arrays, random mapping, number of copies
    extra_count = new_channel_count - old_channel_count

    extra_w = xavier(size(layer.w)[1:3]..., extra_count)
    extra_next_w = next_layer != nothing ? xavier(size(next_layer.w)[1:2]..., extra_count, size(next_layer.w)[4]) : nothing
    extra_b = zeros(1, 1, extra_count, 1)

    if layer.bn_params != nothing
        # Initialize bn scales to 1 and bn biases to 0
        extra_bn_scales = atype()(ones(extra_count,))
        extra_bn_biases = atype()(zeros(extra_count,))
    end

    extra_w = atype()(extra_w)
    extra_next_w = next_layer != nothing ? atype()(extra_next_w) : nothing
    extra_b = atype()(extra_b)

    # Update the weights
    old_size = size(layer.w)
    w_old_2d = reshape(layer.w, :, size(layer.w)[end])
    w_extra_2d = reshape(extra_w, :, size(extra_w)[end])
    w_new_2d = hcat(w_old_2d, w_extra_2d)
    layer.w = Param(reshape(w_new_2d, old_size[1], old_size[2], old_size[3], :))

    layer.b = Param(reshape(vcat(reshape(layer.b, (size(layer.b, 3),)), reshape(extra_b, (extra_count,))), (1, 1, new_channel_count, 1)))

    if next_layer != nothing
        old_size = size(next_layer.w)
        w_old_2d = reshape(next_layer.w, :, size(next_layer.w)[end])
        w_extra_2d = reshape(extra_next_w, :, size(extra_next_w)[end])
        w_new_2d = vcat(w_old_2d, w_extra_2d)

        next_layer.w = Param(reshape(w_new_2d, old_size[1], old_size[2], :, old_size[4]))
    end
    if layer.bn_params != nothing
        layer.bn_params = Param(cat(layer.bn_params[1:old_channel_count], extra_bn_scales, layer.bn_params[old_channel_count+1:end], extra_bn_biases, dims=1))
    end
end

"""
    random_pad_inceptionA(inc, next, grow_ratio)

Baseline method for Net2WiderNet for InceptionA modules where the next layer is another InceptionA.
Grows one of the layers of the model to a new one with a larger size.
The old units are copied as is, new units are initialized randomly.

`inc` should be the `InceptionA` layer that you wish to widen
`next` should be the `InceptionA` layer that comes right after `inc`
`grow_ratio` is the multiplier used for the widening, e.g. 1.5 will grow all `Conv`s to 1.5 times their width.
"""
function random_pad_inceptionA(inc::InceptionA, next::InceptionA, grow_ratio)
    # Grow inside layers normally
    random_pad_conv(inc.c1_before_3, inc.c3, Int64(round(size(inc.c1_before_3.w, 4)*(grow_ratio))))
    random_pad_conv(inc.c1_before_d3, inc.cd3_1, Int64(round(size(inc.c1_before_d3.w, 4)*(grow_ratio))))
    random_pad_conv(inc.cd3_1, inc.cd3_2, Int64(round(size(inc.cd3_1.w, 4)*(grow_ratio))))

    old_c1_count = size(inc.c1_alone.w, 4)
    old_c3_count = size(inc.c3.w, 4)
    old_cd3_count = size(inc.cd3_2.w, 4)
    old_c1_ap_count = size(inc.c1_after_pool.w, 4)

    new_c1_count = Int64(round(old_c1_count*(grow_ratio)))
    new_c3_count = Int64(round(old_c3_count*(grow_ratio)))
    new_cd3_count = Int64(round(old_cd3_count*(grow_ratio)))
    new_c1_ap_count = Int64(round(old_c1_ap_count*(grow_ratio)))

    random_pad_conv(inc.c1_alone, nothing, new_c1_count)
    random_pad_conv(inc.c3, nothing, new_c3_count)
    random_pad_conv(inc.cd3_2, nothing, new_cd3_count)
    random_pad_conv(inc.c1_after_pool, nothing, new_c1_ap_count)

    old_channel_count = old_c1_count + old_c3_count + old_cd3_count + old_c1_ap_count
    new_channel_count = new_c1_count + new_c3_count + new_cd3_count + new_c1_ap_count

    for c in [next.c1_alone, next.c1_before_3, next.c1_before_d3, next.c1_after_pool]
        new_w = Array{Float32}(undef, (size(c.w)[1:2]..., new_channel_count, size(c.w)[4]))

        array_w = convert(Array{Float32}, c.w)
        new_w[:, :, 1:old_c1_count, :] = array_w[:, :, 1:old_c1_count, :]
        new_w[:, :, old_c1_count+1:new_c1_count, :] = xavier(size(array_w)[1:2]..., new_c1_count-old_c1_count, size(array_w)[4])

        new_w[:, :, new_c1_count+1:new_c1_count+old_c3_count, :] = array_w[:, :, old_c1_count+1:old_c1_count+old_c3_count, :]
        new_w[:, :, new_c1_count+old_c3_count+1:new_c1_count+new_c3_count, :] = xavier(size(array_w)[1:2]..., new_c3_count-old_c3_count, size(array_w)[4])

        new_w[:, :, new_c1_count+new_c3_count+1:new_c1_count+new_c3_count+old_cd3_count, :] = array_w[:, :, old_c1_count+old_c3_count+1:old_c1_count+old_c3_count+old_cd3_count, :]
        new_w[:, :, new_c1_count+new_c3_count+old_cd3_count+1:new_c1_count+new_c3_count+new_cd3_count, :] = xavier(size(array_w)[1:2]..., new_cd3_count-old_cd3_count, size(array_w)[4])

        new_w[:, :, new_c1_count+new_c3_count+new_cd3_count+1:new_c1_count+new_c3_count+new_cd3_count+old_c1_ap_count, :] = array_w[:, :, old_c1_count+old_c3_count+old_cd3_count+1:end, :]
        new_w[:, :, new_c1_count+new_c3_count+new_cd3_count+old_c1_ap_count+1:end, :] = xavier(size(array_w)[1:2]..., new_c1_ap_count-old_c1_ap_count, size(array_w)[4])

        c.w = Param(atype()(new_w))
    end
end

"""
    random_pad_inceptionA(inc, next, after_b, grow_ratio)

Baseline method for Net2WiderNet for InceptionA modules where the next layer is an InceptionB.
Grows one of the layers of the model to a new one with a larger size.
The old units are copied as is, new units are initialized randomly.
This function requires a third layer to be given, as it will also be affected by the widening.

`inc` should be the `InceptionA` layer that you wish to widen
`next` should be the `InceptionB` layer that comes right after `inc`
`after_b` should be the `ConvBN` layer that comes right after `next`
`grow_ratio` is the multiplier used for the widening, e.g. 1.5 will grow all `Conv`s to 1.5 times their width.
"""
function random_pad_inceptionA(inc::InceptionA, next::InceptionB, after_b::ConvBN, grow_ratio)
    random_pad_conv(inc.c1_before_3, inc.c3, Int64(round(size(inc.c1_before_3.w, 4)*(grow_ratio))))
    random_pad_conv(inc.c1_before_d3, inc.cd3_1, Int64(round(size(inc.c1_before_d3.w, 4)*(grow_ratio))))
    random_pad_conv(inc.cd3_1, inc.cd3_2, Int64(round(size(inc.cd3_1.w, 4)*(grow_ratio))))

    old_c1_count = size(inc.c1_alone.w, 4)
    old_c3_count = size(inc.c3.w, 4)
    old_cd3_count = size(inc.cd3_2.w, 4)
    old_c1_ap_count = size(inc.c1_after_pool.w, 4)

    new_c1_count = Int64(round(old_c1_count*(grow_ratio)))
    new_c3_count = Int64(round(old_c3_count*(grow_ratio)))
    new_cd3_count = Int64(round(old_cd3_count*(grow_ratio)))
    new_c1_ap_count = Int64(round(old_c1_ap_count*(grow_ratio)))

    random_pad_conv(inc.c1_alone, nothing, new_c1_count)
    random_pad_conv(inc.c3, nothing, new_c3_count)
    random_pad_conv(inc.cd3_2, nothing, new_cd3_count)
    random_pad_conv(inc.c1_after_pool, nothing, new_c1_ap_count)

    old_channel_count = old_c1_count + old_c3_count + old_cd3_count + old_c1_ap_count
    new_channel_count = new_c1_count + new_c3_count + new_cd3_count + new_c1_ap_count

    for c in [next.c1_before_3, next.c1_before_d3]
        new_w = Array{Float32}(undef, (size(c.w)[1:2]..., new_channel_count, size(c.w)[4]))

        array_w = convert(Array{Float32}, c.w)
        new_w[:, :, 1:old_c1_count, :] = array_w[:, :, 1:old_c1_count, :]
        new_w[:, :, old_c1_count+1:new_c1_count, :] = xavier(size(array_w)[1:2]..., new_c1_count-old_c1_count, size(array_w)[4])

        new_w[:, :, new_c1_count+1:new_c1_count+old_c3_count, :] = array_w[:, :, old_c1_count+1:old_c1_count+old_c3_count, :]
        new_w[:, :, new_c1_count+old_c3_count+1:new_c1_count+new_c3_count, :] = xavier(size(array_w)[1:2]..., new_c3_count-old_c3_count, size(array_w)[4])

        new_w[:, :, new_c1_count+new_c3_count+1:new_c1_count+new_c3_count+old_cd3_count, :] = array_w[:, :, old_c1_count+old_c3_count+1:old_c1_count+old_c3_count+old_cd3_count, :]
        new_w[:, :, new_c1_count+new_c3_count+old_cd3_count+1:new_c1_count+new_c3_count+new_cd3_count, :] = xavier(size(array_w)[1:2]..., new_cd3_count-old_cd3_count, size(array_w)[4])

        new_w[:, :, new_c1_count+new_c3_count+new_cd3_count+1:new_c1_count+new_c3_count+new_cd3_count+old_c1_ap_count, :] = array_w[:, :, old_c1_count+old_c3_count+old_cd3_count+1:end, :]
        new_w[:, :, new_c1_count+new_c3_count+new_cd3_count+old_c1_ap_count+1:end, :] = xavier(size(array_w)[1:2]..., new_c1_ap_count-old_c1_ap_count, size(array_w)[4])

        c.w = Param(atype()(new_w))
    end

    # Finally, need to change the layer after InceptionB as well due to the pass-through layer
    array_w = convert(Array{Float32}, after_b.w)
    w_first = array_w[:, :, 1:end-old_channel_count, :]
    w_last = array_w[:, :, end-old_channel_count+1:end, :]
    new_w_last = Array{Float32}(undef, (size(array_w)[1:2]..., new_channel_count, size(array_w)[4]))

    # Set the extra weights
    new_w_last[:, :, 1:old_channel_count, :] = w_last
    new_w_last[:, :, old_channel_count+1:end, :] = xavier(size(w_last)[1:2]..., new_channel_count-old_channel_count, size(w_last)[4])

    after_b.w = Param(atype()(cat(w_first, new_w_last, dims=3)))
end

"""
    random_pad_inceptionB(inc, next, grow_ratio)

Baseline method for Net2WiderNet for InceptionB modules where the next layer is aconvolution with batchnorm layer.
Grows one of the layers of the model to a new one with a larger size.
The old units are copied as is, new units are initialized randomly.

`inc` should be the `InceptionB` layer that you wish to widen
`next` should be the `ConvBN` layer that comes right after `inc`
`grow_ratio` is the multiplier used for the widening, e.g. 1.5 will grow all `Conv`s to 1.5 times their width.
"""
function random_pad_inceptionB(inc::InceptionB, next::ConvBN, grow_ratio)
    random_pad_conv(inc.c1_before_3, inc.c3, Int64(round(size(inc.c1_before_3.w, 4)*(grow_ratio))))
    random_pad_conv(inc.c1_before_d3, inc.cd3_1, Int64(round(size(inc.c1_before_d3.w, 4)*(grow_ratio))))
    random_pad_conv(inc.cd3_1, inc.cd3_2, Int64(round(size(inc.cd3_1.w, 4)*(grow_ratio))))

    old_c3_count = size(inc.c3.w, 4)
    old_cd3_count = size(inc.cd3_2.w, 4)

    new_c3_count = Int64(round(old_c3_count*(grow_ratio)))
    new_cd3_count = Int64(round(old_cd3_count*(grow_ratio)))

    random_pad_conv(inc.c3, nothing, new_c3_count)
    random_pad_conv(inc.cd3_2, nothing, new_cd3_count)

    old_channel_count = old_c3_count + old_cd3_count
    new_channel_count = new_c3_count + new_cd3_count

    array_w = convert(Array{Float32}, next.w)
    w_first = array_w[:, :, 1:end-old_channel_count, :]
    w_last = array_w[:, :, end-old_channel_count+1:end, :]
    new_w_last = Array{Float32}(undef, (size(array_w)[1:2]..., new_channel_count, size(array_w)[4]))

    # Set the extra weights
    new_w_last[:, :, 1:old_channel_count, :] = w_last
    new_w_last[:, :, old_channel_count+1:end, :] = xavier(size(w_last)[1:2]..., new_channel_count-old_channel_count, size(w_last)[4])

    next.w = Param(atype()(cat(w_first, new_w_last, dims=3)))
end

#-------------------------------------------------------------------------------
#----------------------------- TESTS BEGIN -------------------------------------
#-------------------------------------------------------------------------------

"""
Test for checking if wider functions for MLP widen properly and are function preserving
"""
function test_wider_mlp()
    (xtrn, ytrn), (xtst, ytst) = load_data()

    # Need to reshape it to 2 dims for MLP
    xtrn = reshape(xtrn, (32*32*3, :))
    xtst = reshape(xtst, (32*32*3, :))

    dtrn = minibatch(xtrn, ytrn, 50, xtype=atype())
    dtst = minibatch(xtst, ytst, 50, xtype=atype())

    narrow = 16
    wide = 20
    chosen_layer = 1

    mlp_model = create_mlp_model(32*32*3, 10, narrow, 16)
    mlp_results, mlp_model = train_results(dtrn, dtst, "mlp.jld2", mlp_model, 100, true)
    mlp_wider = deepcopy(mlp_model)
    wider_mlp(mlp_wider.layers[chosen_layer], mlp_wider.layers[chosen_layer+1], wide, 0)

    @assert (size(mlp_model.layers[chosen_layer].w, 1) == narrow &&
            size(mlp_model.layers[chosen_layer].b, 1) == narrow &&
            size(mlp_model.layers[chosen_layer+1].w, 2) == narrow) "Old model is modified"

    @assert (size(mlp_wider.layers[chosen_layer].w, 1) == wide &&
            size(mlp_wider.layers[chosen_layer].b, 1) == wide &&
            size(mlp_wider.layers[chosen_layer+1].w, 2) == wide) "New model is not widened correctly"

    sames = 0
    total = 0
    for (x, y) in dtst
        y_olds = mlp_model(x)
        y_news = mlp_wider(x)

        sames += sum(y_olds .- y_news .< 0.01)
        total += length(y_olds)
    end
    @assert sames/total == 1 "Function is not preserved"
    println("wider mlp test passed")
end

"""
Test for checking if wider functions for Conv layers widen properly and are function preserving
"""
function test_wider_conv()
    (xtrn, ytrn), (xtst, ytst) = load_data()

    dtrn = minibatch(xtrn, ytrn, 50, xtype=atype())
    dtst = minibatch(xtst, ytst, 50, xtype=atype())

    narrow = 64
    wide = 72
    chosen_layer = 4

    cnn_model = create_cnn_model(3, 10, true)
    cnn_results, cnn_model = train_results(dtrn, dtst, "cnn.jld2", cnn_model, 5, false)
    cnn_wider = deepcopy(cnn_model)
    wider_conv(cnn_wider.layers[chosen_layer], cnn_wider.layers[chosen_layer+1], wide, 0)

    @assert (size(cnn_model.layers[chosen_layer].w, 4) == narrow &&
            size(cnn_model.layers[chosen_layer].b, 3) == narrow &&
            size(cnn_model.layers[chosen_layer+1].w, 3) == narrow) "Old model is modified"

    @assert (size(cnn_wider.layers[chosen_layer].w, 4) == wide &&
            size(cnn_wider.layers[chosen_layer].b, 3) == wide &&
            size(cnn_wider.layers[chosen_layer+1].w, 3) == wide) "New model is not widened correctly"

    sames = 0
    total = 0
    for (x, y) in dtst
        y_olds = cnn_model(x)
        y_news = cnn_wider(x)

        sames += sum(y_olds .- y_news .< 0.01)
        total += length(y_olds)
    end

    @assert sames/total == 1 "Function is not preserved"
    println("wider conv test passed")
end

"""
Test for checking if wider functions for Inception modules widen properly and are function preserving
"""
function test_wider_inception()
    (xtrn, ytrn), (xtst, ytst) = load_data()

    dtrn = minibatch(xtrn, ytrn, 50, xtype=atype())
    dtst = minibatch(xtst, ytst, 50, xtype=atype())

    model = create_inception_bn_sm_model(3, 10)

    results, model = train_results(dtrn, dtst, "inception_sm.jld2", model, 50, false, false)
    wider = deepcopy(model)

    chosen_layer = 3
    growth_ratio = 1.25

    old_l = model.layers[chosen_layer]
    old_n = model.layers[chosen_layer+1]

    new_l = wider.layers[chosen_layer]
    new_n = wider.layers[chosen_layer+1]

    # instead of checking old model's layers manually, record results beforehand
    y_oldss = []
    for (x, y) in dtst
        push!(y_oldss, model(x))
    end

    wider_inceptionA(new_l, new_n, growth_ratio, 0)

    @assert (size(new_l.c1_alone.w, 4) == Int64(round(size(old_l.c1_alone.w, 4) * (1 + growth_ratio))) &&
             size(new_l.c1_alone.b, 3) == Int64(round(size(old_l.c1_alone.b, 3) * (1 + growth_ratio))) &&
             size(new_l.c3.w, 4) == Int64(round(size(old_l.c3.w, 4) * (1 + growth_ratio))) &&
             size(new_l.c3.b, 3) == Int64(round(size(old_l.c3.b, 3) * (1 + growth_ratio))) &&
             size(new_l.cd3_2.w, 4) == Int64(round(size(old_l.cd3_2.w, 4) * (1 + growth_ratio))) &&
             size(new_l.cd3_2.b, 3) == Int64(round(size(old_l.cd3_2.b, 3) * (1 + growth_ratio))) &&
             size(new_l.c1_after_pool.w, 4) == Int64(round(size(old_l.c1_after_pool.w, 4) * (1 + growth_ratio))) &&
             size(new_l.c1_after_pool.b, 3) == Int64(round(size(old_l.c1_after_pool.b, 3) * (1 + growth_ratio))) &&

             size(new_n.c1_alone.w, 3) == Int64(round(size(old_n.c1_alone.w, 3) * (1 + growth_ratio))) &&
             size(new_n.c1_before_3.w, 3) == Int64(round(size(old_n.c1_before_3.w, 3) * (1 + growth_ratio))) &&
             size(new_n.c1_before_d3.w, 3) == Int64(round(size(old_n.c1_before_d3.w, 3) * (1 + growth_ratio))) &&
             size(new_n.c1_after_pool.w, 3) == Int64(round(size(old_n.c1_after_pool.w, 3) * (1 + growth_ratio)))
    ) "New model is not widened correctly"

    # Widening the next inceptionA
    wider_inceptionA(wider.layers[chosen_layer+1], wider.layers[chosen_layer+2], wider.layers[chosen_layer+4], growth_ratio, 0)
    # Widening the next inceptionB
    wider_inceptionB(wider.layers[chosen_layer+2], wider.layers[chosen_layer+4], growth_ratio, 0)

    sames = 0
    total = 0
    for (i, (x, y)) in enumerate(dtst)
        y_olds = y_oldss[i]
        y_news = wider(x)

        sames += sum(y_olds .- y_news .< 0.01)
        total += length(y_olds)
    end

    @assert sames/total == 1 "Function is not preserved"
    println("wider inceptionA test passed")
end

"""
Test for checking if random pad functions for Inception modules pad properly
"""
function test_random_pad_inception()
    (xtrn, ytrn), (xtst, ytst) = load_data()

    dtrn = minibatch(xtrn, ytrn, 50, xtype=atype())
    dtst = minibatch(xtst, ytst, 50, xtype=atype())

    model = create_inception_bn_sm_model(3, 10)

    results, model = train_results(dtrn, dtst, "inception_sm.jld2", model, 50, false, false)
    wider = deepcopy(model)

    chosen_layer = 3
    growth_ratio = 1.25

    old_l = model.layers[chosen_layer]
    old_n = model.layers[chosen_layer+1]

    new_l = wider.layers[chosen_layer]
    new_n = wider.layers[chosen_layer+1]

    random_pad_inceptionA(new_l, new_n, growth_ratio)

    @assert (size(new_l.c1_alone.w, 4) == Int64(round(size(old_l.c1_alone.w, 4) * (1 + growth_ratio))) &&
             size(new_l.c1_alone.b, 3) == Int64(round(size(old_l.c1_alone.b, 3) * (1 + growth_ratio))) &&
             size(new_l.c3.w, 4) == Int64(round(size(old_l.c3.w, 4) * (1 + growth_ratio))) &&
             size(new_l.c3.b, 3) == Int64(round(size(old_l.c3.b, 3) * (1 + growth_ratio))) &&
             size(new_l.cd3_2.w, 4) == Int64(round(size(old_l.cd3_2.w, 4) * (1 + growth_ratio))) &&
             size(new_l.cd3_2.b, 3) == Int64(round(size(old_l.cd3_2.b, 3) * (1 + growth_ratio))) &&
             size(new_l.c1_after_pool.w, 4) == Int64(round(size(old_l.c1_after_pool.w, 4) * (1 + growth_ratio))) &&
             size(new_l.c1_after_pool.b, 3) == Int64(round(size(old_l.c1_after_pool.b, 3) * (1 + growth_ratio))) &&

             size(new_n.c1_alone.w, 3) == Int64(round(size(old_n.c1_alone.w, 3) * (1 + growth_ratio))) &&
             size(new_n.c1_before_3.w, 3) == Int64(round(size(old_n.c1_before_3.w, 3) * (1 + growth_ratio))) &&
             size(new_n.c1_before_d3.w, 3) == Int64(round(size(old_n.c1_before_d3.w, 3) * (1 + growth_ratio))) &&
             size(new_n.c1_after_pool.w, 3) == Int64(round(size(old_n.c1_after_pool.w, 3) * (1 + growth_ratio)))
    ) "New model is not widened correctly"

    random_pad_inceptionA(wider.layers[chosen_layer+1], wider.layers[chosen_layer+2], wider.layers[chosen_layer+4], growth_ratio)
    random_pad_inceptionB(wider.layers[chosen_layer+2], wider.layers[chosen_layer+4], growth_ratio)
    wider(dtst)

    println("random pad inceptionA test passed")
end


text:
"SPH smoothing kernel."
abstract type Kernel{N} end

"Mathematical definition of a smoothing kernel."
abstract type KernelFormula{N} <: Kernel{N} end

"""
    (::Kernel)(q::Number, h::Number) -> Float

Evaluate the kernel with width `h` at relative distance `q = r/h`.
"""
function (W::Kernel{N})(q::Number, h::Number) where {N}
    @dbgassert(h > 0) && W(q) / h^N
end

"""
    ddr(W::Kernel, q::Number, h::Number) -> Float

Evaluate the derivative of the kernel with width `h` at relative distance
`q = r/h`. The derivative is with respect to `r`.
"""
function ddr(W::Kernel{N}, q::Number, h::Number) where {N}
    @dbgassert(h > 0) && ddr(W, q) / h^(N + 1)
end

"""
    ddh(W::Kernel, q::Number, h::Number) -> Float

Evaluate the derivative of the kernel with width `h` at relative distance
`q = r/h`. The derivative is with respect to `h`.
"""
function ddh(W::Kernel{N}, q::Number, h::Number) where {N}
    @dbgassert(h > 0) && ddh(W, q) / h^(N + 1)
end

"""
    get_support(W::Kernel) -> Float

Returns the relative distance `q = r/h` at which the kernel is truncated.
"""
function get_support end

"""
    nonzeroat(W::Kernel, q::Number) -> Bool

Whether the kernel is non-zero at relative distance `q = r/h`.
"""
nonzeroat(W::Kernel, q::Number) = q < get_support(W)

"""
    ddr(W::Kernel, (qᵢ, qⱼ)::Tuple{Number,Number}, (hᵢ, hⱼ)::Tuple{Number,Number})
        -> (Float, Float)

Evaluate the derivative of the kernel at the relative distances `qᵢ = rᵢ/hᵢ`
and `qⱼ = rⱼ/hⱼ`. The derivative is with respect to `r`.
"""
ddr(W::Kernel, (qᵢ, qⱼ)::Tuple{Number,Number}, (hᵢ, hⱼ)::Tuple{Number,Number}) =
    ddr(W, qᵢ, hᵢ), ddr(W, qⱼ, hⱼ)

"""
    nonzeroat(W::Kernel, q::Tuple{Number,Number}) -> Bool

Whether the kernel is non-zero at both `qᵢ` and `qⱼ`.
"""
nonzeroat(W::Kernel, (qᵢ, qⱼ)::Tuple{Number,Number}) =
    nonzeroat(W, qᵢ) && nonzeroat(W, qⱼ)

include("kernel/gaussian.jl")
include("kernel/cubic_spline.jl")
include("kernel/tabulated.jl")
include("kernel/constant_width.jl")

max_stars_repo_name: lars-frogner/Whirl

max_issues_repo_name: lars-frogner/Whirl

max_forks_repo_name: lars-frogner/Whirl


text:
using Gadfly

Gadfly.set_default_plot_format(:svg)

#Gadfly doesn't call the default display methods, this catches
#all Gadfly plots
function Base.display(report::Report, m::MIME"image/svg+xml", p::Plot)

    chunk = report.cur_chunk
    full_name, rel_name = get_figname(report, chunk)

    docformat = formats[report.formatdict[:doctype]]
    push!(report.figures, rel_name)

    report.fignum += 1

    w = chunk.options[:fig_width]inch
    h = chunk.options[:fig_height]inch
    format = chunk.options[:fig_ext]
    dpi = chunk.options[:dpi]

    #This is probably not the correct way to handle different formats, but it works.
    if format == ".png"
        try
            draw(PNG(full_name, w, h, dpi=dpi), p)
        catch
            draw(PNG(full_name, w, h), p) #Compose < 0.3.1, Gadfly < 0.3.1
        end
    elseif format == ".pdf"
        draw(PDF(full_name, w, h), p)
    elseif format == ".ps"
        draw(PS(full_name, w, h), p)
    elseif format == ".svg"
        draw(SVG(full_name, w, h), p)
    elseif format == ".js.svg"
        draw(SVGJS(full_name, w, h), p)
    elseif format == ".tex"
        draw(PGF(full_name, w, h, true ), p)
    else:
        warn("Can't save figure. Unsupported format")
    end
end


text:
using Test

"""
    fahrenheit_to_celsius(temp_f::Float)

Converts temperature in Fahrenheit to Celsius.
"""
function fahrenheit_to_celsius(temp_f)
    temp_c = (temp_f - 32.0) * (5.0/9.0)
    return temp_c
end

# This is the test section
@testset "Test fahrenheit_to_celsius" begin
    temp_c = fahrenheit_to_celsius(100.0)
    expected_result = 37.777777
    @test abs(temp_c - expected_result) < 1.0e-6
end


text:
"""
`SimpleGraph(H::SimpleHypergraph)` demotes a hypergraph to
a simple graph.
"""
function SimpleGraph(H::SimpleHypergraph{T})::SimpleGraph{T} where {T}
    G = SimpleGraph{T}()

    # copy all vertices
    for v in H.V
        add!(G, v)
    end

    # for all pairs of vertices in all edges, create an edge in G
    for e in H.E
        if length(e) < 2
            continue
        end

        ee = collect(e)
        k = length(ee)

        for i = 1:k-1
            u = ee[i]
            for j = i+1:k
                v = ee[j]
                add!(G, u, v)
            end
        end
    end
    return G
end

"""
`SimpleHypergraph{T}()` creates a new hypergraph in which vertices have
type `T`. **Warning**: Do not use `T=Any`.

`SimpleHypergraph(G::SimpleGraph)` converts a graph to
the equivalent two-uniform hypergraph.
"""
function SimpleHypergraph(G::SimpleGraph{T}) where {T}
    H = SimpleHypergraph{T}()
    for v in G.V
        add!(H, v)
    end

    for e in G.E
        u, v = e
        add!(H, u, v)
    end

    return H
end


text:
# Support for data produced by the DVS ROS driver:
# https://github.com/uzh-rpg/rpg_dvs_ros

module ROSDVS

using RobotOSData

import AddressEvent.read_events!
import AddressEvent.event_time
import AddressEvent.event_coord
import AddressEvent.event_location
import AddressEvent.event_pol

struct DVSEvent
    x::Int16
    y::Int16
    time_s::UInt32
    time_ns::UInt32
    pol::Bool
end

struct DVSMessage
    header::RobotOSData.Header
    height::Int32
    width::Int32
    events::Vector{DVSEvent}
end

event_time(e::DVSEvent) = e.time_s*1000000 + e.time_ns÷1000
event_coord(e::DVSEvent) = (e.x, e.y)
event_location(e::DVSEvent, imsize) = (e.y+1, e.x+1)
event_pol(e::DVSEvent) = e.pol

function Base.read(io::IO, ::Type{DVSEvent})
    x = ltoh(read(io, Int16))
    y = ltoh(read(io, Int16))
    s = ltoh(read(io, UInt32))
    ns = ltoh(read(io, UInt32))
    pol = read(io, UInt8)
    DVSEvent(x, y, s, ns, pol)
end

function Base.read(io::IO, ::Type{DVSMessage})
    header = read(io, RobotOSData.Header)
    height = ltoh(read(io, Int32))
    width = ltoh(read(io, Int32))
    events = read_array(io, DVSEvent)
    DVSMessage(header, height, width, events)
end

"""
Reads a message from `io`, stores the events in `out`,
and returns the frame size `(height, width)`.
"""
function read_events!(io::IO, out::Vector{DVSEvent})
    header = read(io, RobotOSData.Header)
    height = ltoh(read(io, Int32))
    width = ltoh(read(io, Int32))
    read_array!(out, io, DVSEvent)
    return (height, width)
end

export DVSEvent, DVSMessage

end # module


text:
using SimulatedAnnealing
using Random

"""
    Tour

Struct representing a tour visiting cities.

Fields
------
`order::Vector{Int}` order in which each city is visited.
`D::Matrix` Matrix of distances between each pair of cities. This is stored
to be able to pass it down to new candidates without the need of
recomputing it.
"""
struct Tour
    order::Vector{Int}
    D::Matrix{Float64}
end

"""
    Tour(D::Matrix)

Create a tour with random ordering of the cities.
"""
Tour(D::Matrix) = Tour(randcycle(size(D, 1)), D)

"""
    energy(tour::Tour)

Compute the energy of a `Tour` by summing the distances of successive cities
in the tour.
"""
function energy(tour::Tour)
    s = 0.0

    for (k, c1) in enumerate(tour.order)
        c2 = tour.order[mod1(k + 1, length(tour.order))]
        s += tour.D[c1, c2]
    end

    return s
end

"""
    propose_candidate(tour::Tour)

Propose a new candidate `Tour` by reversing the order of cities between two
random indices.
"""
function propose_candidate(tour::Tour)
    n = length(tour.order)
    order = copy(tour.order)

    # Order will be reversed between i1 and i2 included
    i1 = rand(1:n)
    i2 = mod1(rand(1:n-1) + i1, n)  # Make sure i1 != i2

    # Make sure that i1 < i2
    if i1 > i2
        i1, i2 = i2, i1
    end

    if i1 == 1 && i2 == n
        return Tour(order, tour.D), 0.0
    end

    city_a1 = order[i1]
    city_a2 = order[i2]

    city_b1 = order[mod1(i1 - 1, n)]  # Index following i1
    city_b2 = order[mod1(i2 + 1, n)]  # Index preceding i2

    dE = (tour.D[city_a1, city_b2] + tour.D[city_a2, city_b1]
          - tour.D[city_a1, city_b1] - tour.D[city_a2, city_b2])

    order[i1:i2] = reverse(order[i1:i2])

    dE = energy(Tour(order, tour.D)) - energy(tour)

    return Tour(order, tour.D), dE
end

cities = [
    [0.0 1.0 2.0 3.0 3.0 3.0 3.0 2.0 1.0 0.0 0.0 0.0] ;
    [0.0 0.0 0.0 0.0 1.0 2.0 3.0 3.0 3.0 3.0 2.0 1.0]
]

# Distance matrix, common for all Tour objects
D = sqrt.((cities[1, :] .- cities[1, :]').^2 +
          (cities[2, :] .- cities[2, :]').^2)
n = size(D, 1)

# Initial sampling of the configuration space
samples = [Tour(D) for _ in 1:1000]

# Run the algorithm with default parameters
prob = AnnealingOptimization(energy, propose_candidate, samples, n*(n - 1))
best_tour, tour_length = simulated_annealing(prob)