How to rewrite a function for Julia codes?

Dear all, I want to write a function for the codes.

while VFIcontinue == 1
    error_vfi_old = error_vfi
    for i_state in 1:n_state
        prob_mat = repeat(transpose(prob[i_state, :]), n_choice, 1)
        EXP[:, i_state] = sum(prob_mat[:, :] .* V_old[:, :], dims=2)
    end
    for iak_state in 1:n_choice
        ik = Int(rem(iak_state - 1, nk) + 1)
        ia = Int(floor((iak_state - 1) / nk) + 1)
        piplusT = repeat(transpose(pi_vfi[iak_state, :]), n_choice, 1) + repeat(transpose(T_vfi[iak_state, :]), n_choice, 1)
        c = piplusT .- (R_mat .+ δ) * k_grid[ik] + (1 .+ R_mat) * a_grid[ia] - AP_mat - (psi * (KP_mat .- k_grid[ik]) .* (KP_mat .- k_grid[ik])) ./ (2 * k_grid[ik])
        CAPeffective = KP_mat + Q_mat
        V_sub = ((1 ./ abs.(c) .^ ((γ - 1) * ones(n_choice, n_state))) .- 1) ./ (1 - γ) + β * EXP
        @. V_sub[(real(c)<=0)|(collateral<0)|(CAPeffective<=0)] = penalty
        V_new[iak_state, :] = findmax(V_sub, dims=1)[1]
        location[iak_state, :] = first.(Tuple.(findmax(V_sub, dims=1)[2]))
        pol_ind_kp[iak_state, :] = first.(Tuple.(CartesianIndices((nk, na))[location[iak_state, :]]))
        pol_ind_ap[iak_state, :] = last.(Tuple.(CartesianIndices((nk, na))[location[iak_state, :]]))
    end

    Vh_old = V_new
    iter_howard = 0

    while iter_howard <= Nh
        for iak_state = 1:n_choice
            ik = Int(rem(iak_state - 1, nk) + 1)
            ia = Int(floor((iak_state - 1) / nk) + 1)
            c = transpose(pi_vfi[iak_state, :] - (EXOG[3, :] .+ δ) * k_grid[ik] + (1 .+ EXOG[3, :]) * a_grid[ia] + T_vfi[iak_state, :] - a_grid[pol_ind_ap[iak_state, :]] - (psi * (k_grid[pol_ind_kp[iak_state, :]] .- k_grid[ik]) .* (k_grid[pol_ind_kp[iak_state, :]] .- k_grid[ik])) ./ (2 * k_grid[ik]))
            Vh_new[iak_state, :] = ((1 ./ abs.(c) .^ ((γ - 1) * ones(1, nzt * nzp * nr * ntau * nq))) .- 1) ./ (1 - γ) + β .* transpose(sum(prob .* Vh_old[location[iak_state, :], :], dims=2))
        end

        Vh_new = Vh_new .* indicator + Vh_old .* (1 .- indicator)
        Vh_old = copy(Vh_new)
        iter_howard += 1
    end

    if Nh == -1
        error_vfi = findmax(abs.(V_new[:] - V_old[:]))[1]
        V_old = copy(V_new)
    else
        error_vfi = findmax(abs.(Vh_new[:] - V_old[:]))[1]
        V_old = copy(Vh_new)
    end

    iter_vfi += 1

    if error_vfi < error_tolerance
        VFIcontinue = 0
    end
end

function f()
    # all of that code goes here
end

Is this what you’re looking for? Otherwise I think we need an explanation of what it’s supposed to do and what it’s doing instead.

Because these codes run very slow, i try to use function to speed up it.