Merge pull request #58 from ACEsuit/wcw/assemble-weights

Assemble weights separately from the design matrix

src/assemble.jl

@@ -11,7 +11,7 @@ end
 Base.length(d::DataPacket) = count_observations(d.data)
 
 """
-Assemble feature matrix, target vector, and weight vector for given data and basis.
+Assemble feature matrix and target vector for given data and basis.
 """
 function assemble(data::AbstractVector{<:AbstractData}, basis)
     @info "Assembling linear problem."
@@ -26,14 +26,31 @@ function assemble(data::AbstractVector{<:AbstractData}, basis)
     @info "  - Creating feature matrix with size ($(rows[end][end]), $(length(basis)))."
     A = SharedArray(zeros(rows[end][end], length(basis)))
     Y = SharedArray(zeros(size(A, 1)))
-    W = SharedArray(zeros(size(A, 1)))
     @info "  - Beginning assembly with processor count:  $(nprocs())."
     @showprogress pmap(packets) do p
         A[p.rows, :] .= feature_matrix(p.data, basis)
         Y[p.rows] .= target_vector(p.data)
-        W[p.rows] .= weight_vector(p.data)
         GC.gc()
     end
     @info "  - Assembly completed."
-    return Array(A), Array(Y), Array(W)
+    return Array(A), Array(Y), assemble_weights(data)
+end
+
+"""
+Assemble full weight vector for vector of data elements.
+"""
+function assemble_weights(data::AbstractVector{<:AbstractData})
+    @info "Assembling full weight vector."
+    rows = Array{UnitRange}(undef, length(data))  # row ranges for each element of data
+    rows[1] = 1:count_observations(data[1])
+    for i in 2:length(data)
+        rows[i] = rows[i - 1][end] .+ (1:count_observations(data[i]))
+    end
+    packets = DataPacket.(rows, data)
+    sort!(packets, by = length, rev = true)
+    W = SharedArray(zeros(rows[end][end]))
+    @showprogress pmap(packets) do p
+        W[p.rows] .= weight_vector(p.data)
+    end
+    return Array(W)
 end