fit/solver small refactors & type stability (#294)

* codestable resize

Refactor equalize_size function to improve handling of data dimensions and simplify logic.

* Rename and update equalize_size function

Renamed function from _equalize_size to equalize_size and updated its signature.

* small refactors & type stability

* format

Author: behinger

src/solver.jl

@@ -22,18 +22,21 @@ function solver_default(
     multithreading = true,
     show_progress = true,
 ) where {T<:Union{Missing,<:Number}}
-    minfo = Array{IterativeSolvers.ConvergenceHistory,1}(undef, size(data, 1))
+    n_channels = size(data, 1)
+    n_predictors = size(X, 2)
+    minfo = Array{IterativeSolvers.ConvergenceHistory,1}(undef, n_channels)
 
-    beta = zeros(T, size(data, 1), size(X, 2)) # had issues with undef
+    beta = zeros(T, n_channels, n_predictors)
 
-    p = Progress(size(data, 1); enabled = show_progress)
-    X = SparseMatrixCSC(X) # X s often a SubArray, lsmr really doesnt like indexing into subarrays, one copy needed.
-    @maybe_threads multithreading for ch = 1:size(data, 1)
+    p = Progress(n_channels; enabled = show_progress)
+    X_sparse = SparseMatrixCSC(X) # X s often a SubArray, lsmr really doesnt like indexing into subarrays, one copy needed.
+
+    @maybe_threads multithreading for ch = 1:n_channels
 
         # use the previous channel as a starting point
-        ch == 1 || copyto!(view(beta, ch, :), view(beta, ch - 1, :))
+        ch > 1 && copyto!(view(beta, ch, :), view(beta, ch - 1, :))
 
-        h = _lsmr!(beta, X, data, ch)
+        h = _lsmr!(beta, X_sparse, data, ch)
         minfo[ch] = h
         next!(p)
     end
@@ -55,13 +58,11 @@ function solver_default(
     multithreading = true,
     show_progress = true,
 ) where {T<:Union{Missing,<:Number}}
-    #beta = Array{Union{Missing,Number}}(undef, size(data, 1), size(data, 2), size(X, 2))
-    beta = zeros(T, size(data, 1), size(data, 2), size(X, 2))
-    p = Progress(size(data, 1); enabled = show_progress)
-    @maybe_threads multithreading for ch = 1:size(data, 1)
-        for t = 1:size(data, 2)
-            #            @debug("$(ndims(data,)),$t,$ch")
-
+    n_channels, n_times, n_predictors = size(data, 1), size(data, 2), size(X, 2)
+    beta = zeros(T, n_channels, n_times, n_predictors)
+    p = Progress(n_channels; enabled = show_progress)
+    @maybe_threads multithreading for ch = 1:n_channels
+        for t = 1:n_times
             dd = view(data, ch, t, :)
             ix = @. !ismissing(dd)