Fix solution indexing

docs/src/spectra.md

@@ -45,11 +45,6 @@ Pk_halofit(k) = MatterPower.halofit(Pk, halofit_params, Ωm0, ustrip(k))
 plot!(log10.(ks*u"Mpc"), log10.(Pk_halofit.(ustrip(ks))); label = "non-linear (halofit / MatterPower.jl)", legend_position = :bottomleft)
 ```
 
-### Non-linear
-
-```@example
-```
-
 ## CMB power spectra
 
 ```@docs

src/solve.jl

@@ -230,13 +230,8 @@ function Base.getindex(sol::CosmologySolution, i::SymbolicIndex)
         return sol.th[i]
     end
 end
-# TODO: unite above and below!
-function Base.getindex(sol::CosmologySolution, i::SymbolicIndex, j = :)
-    if ModelingToolkit.isparameter(i) && i !== t && (j == :) # don't catch independent variable as parameter
-        return sol.th.ps[i] # assume all parameters are in background/thermodynamics # TODO: index sol directly?
-    else
-        return stack(sol.th[i, j])
-    end
+function Base.getindex(sol::CosmologySolution, i::SymbolicIndex, j)
+    return stack(sol.th[i, j])
 end
 Base.getindex(sol::CosmologySolution, i::Int, j::SymbolicIndex, k = :) = sol.pts[i][j, k]
 Base.getindex(sol::CosmologySolution, i, j::SymbolicIndex, k = :) = [stack(sol[_i, j, k]) for _i in i]
@@ -327,7 +322,7 @@ function shoot(M::CosmologyModel, pars_fixed, pars_varying, conditions; solver =
     function f(vals_varying, _)
         pars = merge(pars_fixed, Dict(keys(pars_varying) .=> vals_varying))
         sol = solve(M, pars; kwargs...) # solve cosmology
-        return sol[funcs][:, end] # evaluate all expressions at final time (e.g. today)
+        return sol[funcs, :][:, end] # evaluate all expressions at final time (e.g. today)
     end
 
     guess = collect(values(pars_varying))