allow for vectors of AbstractVectors

Project.toml

@@ -1,6 +1,6 @@
 name = "FixedEffects"
 uuid = "c8885935-8500-56a7-9867-7708b20db0eb"
-version = "2.2.0"
+version = "2.3.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

ext/MetalExt.jl

@@ -14,12 +14,12 @@ function _mtl(T::Type, fe::FixedEffect)
 	interaction = _mtl(T, fe.interaction)
 	FixedEffect{typeof(refs), typeof(interaction)}(refs, interaction, fe.n)
 end
-_mtl(T::Type, w::UnitWeights) = fill!(MtlVector{T}(undef, length(w)), w[1])
+_mtl(T::Type, w::UnitWeights) = Metal.ones(T, length(w))
 _mtl(T::Type, w::AbstractVector) = MtlVector{T}(convert(Vector{T}, w))
 
 ##############################################################################
 ##
-## FixedEffectLinearMap on the Metal (code by Paul Schrimpf)
+## FixedEffectLinearMap on Metal
 ##
 ## Model matrix of categorical variables
 ## mutiplied by diag(1/sqrt(∑w * interaction^2, ..., ∑w * interaction^2) (Jacobi preconditoner)

src/AbstractFixedEffectSolver.jl

@@ -45,30 +45,7 @@ function solve_residuals!(y::Union{AbstractVector{<: Number}, AbstractMatrix{<:
 	solve_residuals!(y, feM; maxiter = maxiter, tol = tol)
 end
 
-function solve_residuals!(X::AbstractMatrix, feM::AbstractFixedEffectSolver; progress_bar = true, kwargs...)
-    iterations = Int[]
-    convergeds = Bool[]
-    bar = MiniProgressBar(header = "Demean Variables:", color = Base.info_color(), percentage = false, max = size(X, 2))
-    for j in 1:size(X, 2)
-    	v0 = time()
-        _, iteration, converged = solve_residuals!(view(X, :, j), feM; kwargs...)
-        v1 = time()
-        # remove progress_bar if estimated time lower than 2sec
-	    if progress_bar && (j == 1) && ((v1 - v0) * size(X, 2) <= 2)
-	    	progress_bar = false
-	    end
-    	if progress_bar
-    		bar.current = j
-    	    showprogress(stdout, bar)
-    	end
-        push!(iterations, iteration)
-        push!(convergeds, converged)
-    end
-    if progress_bar
-    	end_progress(stdout, bar)
-    end
-    return X, iterations, convergeds
-end
+
 
 function solve_residuals!(r::AbstractVector, feM::AbstractFixedEffectSolver{T}; tol::Real = sqrt(eps(T)), maxiter::Integer = 100_000) where {T}
 	# One cannot copy view of Vector (r) on GPU, so first collect the vector
@@ -102,6 +79,35 @@ function solve_residuals!(r::AbstractVector, feM::AbstractFixedEffectSolver{T};
 	return r, iter, converged
 end
 
+function solve_residuals!(xs::AbstractVector{<: AbstractVector}, feM::AbstractFixedEffectSolver; progress_bar = true, kwargs...)
+    iterations = Int[]
+    convergeds = Bool[]
+    bar = MiniProgressBar(header = "Demean Variables:", color = Base.info_color(), percentage = false, max = length(xs))
+    for (j, x) in enumerate(xs)
+    	v0 = time()
+        _, iteration, converged = solve_residuals!(x, feM; kwargs...)
+        v1 = time()
+        # remove progress_bar if estimated time lower than 2sec
+	    if progress_bar && (j == 1) && ((v1 - v0) * length(xs) <= 2)
+	    	progress_bar = false
+	    end
+    	if progress_bar
+    		bar.current = j
+    	    showprogress(stdout, bar)
+    	end
+        push!(iterations, iteration)
+        push!(convergeds, converged)
+    end
+    if progress_bar
+    	end_progress(stdout, bar)
+    end
+    return xs, iterations, convergeds
+end
+
+function solve_residuals!(X::AbstractMatrix, feM::AbstractFixedEffectSolver; kwargs...)
+	xs, iterations, convergeds = solve_residuals!(eachcol(X), feM; kwargs...)
+	return X, iterations, convergeds
+end
 
 
 

test/solve.jl

@@ -25,9 +25,9 @@ method_s = [:cpu]
 if CUDA.functional()
 	push!(method_s, :CUDA)
 end
-if Metal.functional()
-	push!(method_s, :Metal)
-end
+#if Metal.functional()
+#	push!(method_s, :Metal)
+#end
 for method in method_s
 	println("$method Float32")
 	local (r, iter, conv) = solve_residuals!(deepcopy(x),fes, method=method, double_precision = false)