cv_iht allocation free (improves precompilation time #34)

src/MendelIHT.jl

@@ -1,5 +1,3 @@
-__precompile__()
-
 module MendelIHT
 
     import Distances: euclidean, chebyshev, sqeuclidean

src/cross_validation.jl

@@ -75,6 +75,8 @@ function cv_iht(
     # preallocated arrays for efficiency
     test_idx  = [falses(length(folds)) for i in 1:nprocs()]
     train_idx = [falses(length(folds)) for i in 1:nprocs()]
+    V = [initialize(x, z, y, 1, 1, d, l, group, weight, est_r, init_beta,
+        cv_train_idx=train_idx[i]) for i in 1:nprocs()]
 
     # for displaying cross validation progress
     pmeter = Progress(q * length(path), "Cross validating...")
@@ -90,20 +92,22 @@ function cv_iht(
         id = myid()
         test_idx[id]  .= folds .== fold
         train_idx[id] .= folds .!= fold
+        v = V[id]
 
-        # run IHT on training data with current k
-        v = initialize(x, z, y, 1, k, d, l, group, weight, est_r, init_beta,
-            cv_train_idx=train_idx[id])
-        result = fit_iht!(v, debias=debias, verbose=false, max_iter=max_iter)
+        # run IHT on training data with current (fold, k)
+        v.k = k
+        init_iht_indices!(v, init_beta, train_idx[id])
+        fit_iht!(v, debias=debias, verbose=false, max_iter=max_iter)
 
         # predict on validation data
         v.cv_wts[train_idx[id]] .= zero(T)
         v.cv_wts[test_idx[id]] .= one(T)
+        mse = predict!(v)
 
         # update progres
         put!(channel, true)
 
-        return predict!(v, result)
+        return mse
     end
     put!(channel, false)
 
@@ -123,8 +127,10 @@ function cv_iht(y::AbstractVecOrMat{T}, x::AbstractMatrix; kwargs...) where T
 end
 
 """
-    allocate_fold_and_k()
+    allocate_fold_and_k(q, path)
 
+Returns all combinations of (qᵢ, kᵢ) where `q` is number of cross validation fold
+and path is a vector of different `k` to be tested. 
 """
 function allocate_fold_and_k(q::Int, path::AbstractVector{<:Integer})
     combinations = Tuple{Int, Int}[]
@@ -188,10 +194,10 @@ function iht_run_many_models(y::AbstractVecOrMat{T}, x::AbstractMatrix; kwargs..
     return iht_run_many_models(y, x, z; kwargs...)
 end
 
-function predict!(v::IHTVariable{T, M}, result::IHTResult) where {T <: Float, M}
+function predict!(v::IHTVariable{T, M}) where {T <: Float, M}
     # first update mean μ with estimated (trained) beta (cv weights are handled in deviance)
-    mul!(v.xb, v.x, result.beta)
-    mul!(v.zc, v.z, result.c)
+    mul!(v.xb, v.x, v.best_b)
+    mul!(v.zc, v.z, v.best_c)
     update_μ!(v)
 
     # Compute deviance residual (MSE for Gaussian response)

src/data_structures.jl

@@ -47,7 +47,6 @@ ntraits(v::IHTVariable) = 1
 function IHTVariable(x::M, z::AbstractVecOrMat{T}, y::AbstractVector{T},
     J::Int, k::Union{Int, Vector{Int}}, d::UnivariateDistribution, l::Link,
     group::AbstractVector{Int}, weight::AbstractVector{T}, est_r::Symbol,
-    init_beta::Bool, cv_train_idx::BitVector=trues(size(x, 1))
     ) where {T <: Float, M <: AbstractMatrix}
 
     n = size(x, 1)
@@ -78,31 +77,28 @@ function IHTVariable(x::M, z::AbstractVecOrMat{T}, y::AbstractVector{T},
         ks, k = k, 0
     end
 
-    init_beta && !(typeof(d) <: Normal) && 
-        throw(ArgumentError("Intializing beta values only work for Gaussian phenotypes! Sorry!"))
-
-    b      = init_beta ? initialize_beta(y, x, cv_train_idx) : zeros(T, p)
-    b0     = zeros(T, p)
-    best_b = zeros(T, p)
-    xb     = zeros(T, n)
-    xk     = zeros(T, n, J * columns - 1) # subtracting 1 because the intercept will likely be selected in the first iter
-    gk     = zeros(T, J * columns - 1)    # subtracting 1 because the intercept will likely be selected in the first iter
-    xgk    = zeros(T, n)
-    idx    = falses(p)
-    idx0   = falses(p)
-    idc    = falses(q)
-    idc0   = falses(q)
-    r      = zeros(T, n)
-    df     = zeros(T, p)
-    df2    = zeros(T, q)
-    c      = zeros(T, q)
-    best_c = zeros(T, q)
-    c0     = zeros(T, q)
-    zc     = zeros(T, n)
-    zdf2   = zeros(T, n)
-    μ      = zeros(T, n)
-    cv_wts = ones(T, n); cv_wts[.!cv_train_idx] .= zero(T)
-    storage = zeros(T, p + q)
+    b      = Vector{T}(undef, p)
+    b0     = Vector{T}(undef, p)
+    best_b = Vector{T}(undef, p)
+    xb     = Vector{T}(undef, n)
+    xk     = Matrix{T}(undef, n, J * columns - 1) # subtracting 1 because the intercept will likely be selected in the first iter
+    gk     = Vector{T}(undef, J * columns - 1)    # subtracting 1 because the intercept will likely be selected in the first iter
+    xgk    = Vector{T}(undef, n)
+    idx    = BitArray(undef, p)
+    idx0   = BitArray(undef, p)
+    idc    = BitArray(undef, q)
+    idc0   = BitArray(undef, q)
+    r      = Vector{T}(undef, n)
+    df     = Vector{T}(undef, p)
+    df2    = Vector{T}(undef, q)
+    c      = Vector{T}(undef, q)
+    best_c = Vector{T}(undef, q)
+    c0     = Vector{T}(undef, q)
+    zc     = Vector{T}(undef, n)
+    zdf2   = Vector{T}(undef, n)
+    μ      = Vector{T}(undef, n)
+    cv_wts = Vector{T}(undef, n)
+    storage = Vector{T}(undef, p + q)
 
     return IHTVariable{T, M}(
         x, y, z, k, J, ks, d, l, est_r, 
@@ -118,14 +114,13 @@ function initialize(x::M, z::AbstractVecOrMat{T}, y::AbstractVecOrMat{T},
     ) where {T <: Float, M <: AbstractMatrix}
 
     if is_multivariate(y)
-        v = mIHTVariable(x, z, y, k, initialize_beta, cv_train_idx)
+        v = mIHTVariable(x, z, y, k, cv_train_idx)
     else
-        v = IHTVariable(x, z, y, J, k, d, l, group, weight, est_r, 
-            initialize_beta, cv_train_idx)
+        v = IHTVariable(x, z, y, J, k, d, l, group, weight, est_r)
     end
 
     # initialize non-zero indices
-    MendelIHT.init_iht_indices!(v, initialize_beta)
+    MendelIHT.init_iht_indices!(v, initialize_beta, cv_train_idx)
 
     return v
 end

src/fit.jl

@@ -91,8 +91,13 @@ function fit_iht(
         print_parameters(io, k, d, l, use_maf, group, debias, tol, max_iter)
     end
 
-    return fit_iht!(v, debias=debias, verbose=verbose, tol=tol, max_iter=max_iter,
-        max_step=max_step, io=io)
+    tot_time, best_logl, mm_iter = fit_iht!(v, debias=debias, verbose=verbose,
+        tol=tol, max_iter=max_iter, max_step=max_step, io=io)
+
+    # compute phenotype's proportion of variation explained
+    σ2 = pve(v)
+
+    return IHTResult(tot_time, best_logl, mm_iter, σ2, v)
 end
 
 function fit_iht(
@@ -180,10 +185,7 @@ function fit_iht!(
         end
     end
 
-    # compute phenotype's proportion of variation explained
-    σ2 = pve(v)
-
-    return IHTResult(tot_time, best_logl, mm_iter, σ2, v)
+    return tot_time, best_logl, mm_iter
 end
 
 """

src/utilities.jl

@@ -280,20 +280,39 @@ end
 """
 When initializing the IHT algorithm, take largest elements in magnitude of each
 group of the score as nonzero components of b. This function set v.idx = 1 for
-those indices. If `initialized_beta=true`, then beta values were all initialized to
+those indices. If `init_beta=true`, then beta values will be initialized to
 their univariate values (see [`initialize_beta`](@ref)), in which case we will simply
 choose top `k` entries
 
 `J` is the maximum number of active groups, and `k` is the maximum number of
 predictors per group. 
 """
-function init_iht_indices!(v::IHTVariable, initialized_beta::Bool)
-    z = v.z
-    y = v.y
-    l = v.l
-    J = v.J
-    k = v.k
-    group = v.group
+function init_iht_indices!(v::IHTVariable, init_beta::Bool, cv_idx::BitVector)
+    fill!(v.b, 0)
+    fill!(v.b0, 0)
+    fill!(v.best_b, 0)
+    fill!(v.xb, 0)
+    fill!(v.xk, 0)
+    fill!(v.gk, 0)
+    fill!(v.xgk, 0)
+    fill!(v.idx, false)
+    fill!(v.idx0, false)
+    fill!(v.idc, false)
+    fill!(v.idc0, false)
+    fill!(v.r, 0)
+    fill!(v.df, 0)
+    fill!(v.df2, 0)
+    fill!(v.c, 0)
+    fill!(v.best_c, 0)
+    fill!(v.c0, 0)
+    fill!(v.zc, 0)
+    fill!(v.zdf2, 0)
+    fill!(v.μ, 0)
+    fill!(v.cv_wts, 0)
+    v.cv_wts[cv_idx] .= 1
+
+    init_beta && !(typeof(d) <: Normal) && 
+        throw(ArgumentError("Intializing beta values only work for Gaussian phenotypes! Sorry!"))
 
     # find the intercept by Newton's method
     ybar = zero(eltype(v.y))
@@ -302,33 +321,34 @@ function init_iht_indices!(v::IHTVariable, initialized_beta::Bool)
     end
     ybar /= count(!iszero, v.cv_wts)
     for iteration = 1:20 
-        g1 = linkinv(l, v.c[1])
-        g2 = mueta(l, v.c[1])
+        g1 = linkinv(v.l, v.c[1])
+        g2 = mueta(v.l, v.c[1])
         v.c[1] = v.c[1] - clamp((g1 - ybar) / g2, -1.0, 1.0)
         abs(g1 - ybar) < 1e-10 && break
     end
-    mul!(v.zc, z, v.c)
+    mul!(v.zc, v.z, v.c)
 
     # update mean vector and use them to compute score (gradient)
     update_μ!(v)
     score!(v)
 
-    if initialized_beta && v.k > 0
-        project_k!(v)
+    if init_beta
+        v.b = initialize_beta(v.y, v.x, v.cv_wts)
+        v.k > 0 && project_k!(v)
     else
         # first `k` non-zero entries are chosen based on largest gradient
         ldf = length(v.df)
         v.full_b[1:ldf] .= v.df
         v.full_b[ldf+1:end] .= v.df2
         if length(v.ks) == 0 # no group projection
-            a = partialsort(v.full_b, k * J, by=abs, rev=true)
+            a = partialsort(v.full_b, v.k * v.J, by=abs, rev=true)
             v.idx .= abs.(v.df) .>= abs(a)
             v.idc .= abs.(v.df2) .>= abs(a)
 
             # Choose randomly if more are selected
             _choose!(v) 
         else 
-            project_group_sparse!(v.full_b, group, J, v.ks)
+            project_group_sparse!(v.full_b, v.group, v.J, v.ks)
             @inbounds for i in 1:ldf
                 v.full_b[i] != 0 && (v.idx[i] = true)
             end

src/wrapper.jl

@@ -209,7 +209,7 @@ function parse_covariates(filename::AbstractString, exclude_std_idx::AbstractVec
     if all(x == 1 for x in @view(z[:, 1]))
         length(mask) > 1 && (mask[1] = true) # don't standardize intercept
     else
-        @warn("Covariate file provided but did not detect an intercept An intercept will not be included in IHT!")
+        @warn("Covariate file provided but did not detect an intercept. An intercept will NOT be included in IHT!")
     end
 
     standardize && standardize!(@view(z[:, mask]))