working LinearSolver with tests

ConScape · Dec 28, 2024 · c1d52af · c1d52af
1 parent 850c170
commit c1d52af
Show file tree

Hide file tree

Showing 8 changed files with 412 additions and 224 deletions.
diff --git a/examples/2_landmarks.jmd b/examples/2_landmarks.jmd
@@ -22,7 +22,6 @@ datadir = joinpath(dirname(pathof(ConScape)), "..", "data")
 ```julia
 mov_prob = replace_missing(Raster(joinpath(datadir, "mov_prob_1000.asc")), NaN)
 hab_qual = replace_missing(Raster(joinpath(datadir, "hab_qual_1000.asc")), NaN);
-affinities
 ```
 
 ```julia
@@ -105,21 +104,17 @@ We can write this using a lazy problem definition:
 
 ```julia
 # Put least cost, random walk, and rsp
-measures = (;
-    func=ConScape.ConnectedHabitat(),
-    qbetw=ConScape.BetweennessQweighted(),
-    kbetw=ConScape.BetweennessKweighted(),
-)
-problem = ConScape.Problem(measures; 
+problem = ConScape.Problem(; 
+    graph_measures = (;
+        func=ConScape.ConnectedHabitat(),
+        qbetw=ConScape.BetweennessQweighted(),
+        kbetw=ConScape.BetweennessKweighted(),
+    ),
     distance_transformation=(exp=x -> exp(-x/75), oddsfor=ConScape.OddsFor()),
     connectivity_measure=ConScape.ExpectedCost(θ=1.0),
     solver=ConScape.MatrixSolve(),
-)
-# problem = ConScape.Problem(measures; 
-    # connectivity_measure=ConScape.ExpectedCost(θ=1.0),
-    # distance_transformation=(exp=x -> exp(-x/75), oddsfor=ConScape.OddsFor()),
     # solver=ConScape.VectorSolve(nothing),
-# )
+)
 ```
 
 Then run it for all operations on both normal and coarse grids
@@ -145,19 +140,26 @@ heatmap(result.coarse_func_exp)
 
 
 ```julia
-windowed_problem = ConScape.WindowedProblem(problem; radius=20, overlap=5)
-result = ConScape.solve(windowed_problem, rast)
+windowed_problem = ConScape.WindowedProblem(problem; 
+    radius=20, overlap=5, threaded=true
+)
+@time result = ConScape.solve(windowed_problem, rast)
 plot(result)
-plot(mos.func_exp)
+plot(result.func_exp)
 ```
 
 
 ```julia
-stored_problem = ConScape.StoredProblem(windowed_problem; path=".", radius=70, overlap=20)
+stored_problem = ConScape.StoredProblem(problem; 
+    path=".", radius=20, overlap=30, threaded=true
+)
 ConScape.solve(stored_problem, rast)
 result = mosaic(stored_problem; to=rast)
-plot(mos)
-plot(mos.func_exp)
+plot(result)
+plot(result.func_exp)
+
+# Calculate job ids to run on a cluster
+jobs = ConScape.batch_ids(stored_problem, rast)
 ```
 
 And we can check the corelation similarly to above, by getting

diff --git a/src/ConScape.jl b/src/ConScape.jl
@@ -1,38 +1,41 @@
 module ConScape
 
-    using SparseArrays, LinearAlgebra
-    using Graphs, Plots, SimpleWeightedGraphs, ProgressLogging, ArnoldiMethod
-    using Rasters
-    using LinearSolve
-    using Rasters.DimensionalData
+using SparseArrays, LinearAlgebra
+using Graphs, Plots, SimpleWeightedGraphs, ProgressLogging, ArnoldiMethod
+using Rasters
+using LinearSolve
+using Rasters.DimensionalData
 
-    # Old funcion-based interface
-    abstract type ConnectivityFunction <: Function end
-    abstract type DistanceFunction <: ConnectivityFunction end
-    abstract type ProximityFunction <: ConnectivityFunction end
+# Old funcion-based interface
+abstract type ConnectivityFunction <: Function end
+abstract type DistanceFunction <: ConnectivityFunction end
+abstract type ProximityFunction <: ConnectivityFunction end
 
-    struct least_cost_distance   <: DistanceFunction end
-    struct expected_cost         <: DistanceFunction end
-    struct free_energy_distance  <: DistanceFunction end
+struct least_cost_distance   <: DistanceFunction end
+struct expected_cost         <: DistanceFunction end
+struct free_energy_distance  <: DistanceFunction end
 
-    struct survival_probability  <: ProximityFunction end
-    struct power_mean_proximity  <: ProximityFunction end
+struct survival_probability  <: ProximityFunction end
+struct power_mean_proximity  <: ProximityFunction end
 
-    # Need to define before loading files
-    abstract type AbstractProblem end
+# Need to define before loading files
+abstract type AbstractProblem end
+abstract type Solver end
+
+# Randomized shortest path algorithms
+include("randomizedshortestpath.jl")
+# Grid struct and methods
+include("grid.jl")
+# GridRSP (randomized shortest path) struct and methods
+include("gridrsp.jl")
+# IO
+include("io.jl")
+# Utilities
+include("utils.jl")
+include("graph_measure.jl")
+include("connectivity_measure.jl")
+include("problem.jl")
+include("solvers.jl")
+include("tiles.jl")
 
-    # Randomized shortest path algorithms
-    include("randomizedshortestpath.jl")
-    # Grid struct and methods
-    include("grid.jl")
-    # GridRSP (randomized shortest path) struct and methods
-    include("gridrsp.jl")
-    # IO
-    include("io.jl")
-    # Utilities
-    include("utils.jl")
-    include("graph_measure.jl")
-    include("connectivity_measure.jl")
-    include("problem.jl")
-    include("tiles.jl")
 end
diff --git a/src/grid.jl b/src/grid.jl
@@ -418,6 +418,7 @@ function _get_target(rast::AbstractRasterStack)
             throw(ArgumentError("No :target_qualities or :qualities layers found"))
         end
     end
+end
 
 
 """

diff --git a/src/problem.jl b/src/problem.jl
@@ -1,3 +1,10 @@
+# Defined earlier in ConScape.jl for load order
+# abstract type AbstractProblem end
+@doc """
+    Problem
+
+Abstract supertype for ConScape problem specifications.
+""" Problem
 
 # Recusive getters for nested problems
 graph_measures(p::AbstractProblem) = graph_measures(p.problem)
@@ -23,48 +30,26 @@ and time reequiremtents on a cluster
 """
 function assess end
 
-"""
-    SolverMode
-
-Abstract supertype for ConScape solver modes.
-"""
-abstract type SolverMode end
-"""
-   MaterializedSolve()
-
-solve all operations on a fully materialised Z matrix.
-
-This is inneficient for CPUS, but may be best for GPUs
-using CuSSP.jl ?
-"""
-struct MatrixSolve <: SolverMode end
-"""
-   VectorSolve()
-
-Solve all operations column by column,
-after precomputing lu decompositions.
-
-TODO we can put LinearSolve.jl solver objects in the struct
-"""
-struct VectorSolve{T} <: SolverMode 
-    keywords::T
-end
-
 """
     Problem(graph_measures...; solver, θ)
 
 Combine multiple solve operations into a single object, 
 to be run in the same job.
+
+# Keywords
+
+- `graph_measures`: A NamedTuple of [`GraphMeasure`](@ref)s.
+- `distance_transformation`: A [`DistanceTransformation`](@ref) or `NamedTuple` of `DistanceTransformation`s.
+- `connectivity_measure`: A [`ConnectivityMeasure`](@ref).
+- `solver`: A [`Solver`](@ref) specification.
 """
-@kwdef struct Problem{GM,CM<:ConnectivityMeasure,DT,SM<:SolverMode} <: AbstractProblem
+@kwdef struct Problem{GM,CM<:ConnectivityMeasure,DT,SM<:Solver} <: AbstractProblem
     graph_measures::GM
-    connectivity_measure::CM=LeastCostDistance()
-    distance_transformation::DT=nothing
-    solver::SM=VectorSolve((;))
+    connectivity_measure::CM = LeastCostDistance()
+    distance_transformation::DT = nothing
+    solver::SM = MatrixSolver()
 end
-Problem(gms::GraphMeasure...; kw...) = Problem(gms; kw...)
 Problem(graph_measures::Union{Tuple,NamedTuple}; kw...) = Problem(; graph_measures, kw...)
-Problem(p::AbstractProblem; solver=p.solver, θ=nothing) = Problem(o.graph_measures, solver, θ)
 
 graph_measures(p::Problem) = p.graph_measures
 connectivity_measure(p::Problem) = p.connectivity_measure
@@ -74,97 +59,6 @@ solver(p::Problem) = p.solver
 solve(p::Problem, rast::RasterStack) = solve(p, Grid(p, rast))
 solve(p::Problem, g::Grid) = solve(p.solver, connectivity_measure(p), p, g)
 
-
-# Use an iterative solver so the grid is not materialised
-function solve(m::VectorSolve, cm::FundamentalMeasure, p::AbstractProblem, g::Grid)
-    # solve Z column by column
-    Pref = _Pref(g.affinities)
-    W = _W(Pref, cm.θ, g.costmatrix)
-    # Sparse lhs
-    A = I - W
-    # Sparse diagonal rhs matrix
-    B = sparse(g.targetnodes,
-        1:length(g.targetnodes),
-        1.0,
-        size(g.costmatrix, 1),
-        length(g.targetnodes),
-    )
-    b_init = zeros(eltype(A), size(B, 1))
-    # Dense rhs column
-
-    # Define and initialise the linear problem
-    linprob = LinearProblem(A, b_init)
-    linsolve = init(linprob)
-    # TODO: for now we define a Z matrix, but later modify ops 
-    # to run column by column without materialising Z
-    Z = Matrix{eltype(A)}(undef, size(B))
-    nbuffers = Threads.nthreads()
-    # Create a channel to store problem b vectors for threads
-    # I'm not sure if this is the most efficient way
-    # see https://juliafolds2.github.io/OhMyThreads.jl/stable/literate/tls/tls/
-    ch = Channel{Tuple{typeof(linsolve),Vector{Float64}}}(nbuffers)
-    for i in 1:nbuffers
-        # TODO fix this in LinearSolve.jl with batching
-        # We should not need to `deepcopy` the whole problem we 
-        # just need to replicate the specific workspace arrays 
-        # that will cause race conditions.
-        # But currently there is no parallel mode for LinearSolve.jl
-        # See https://github.com/SciML/LinearSolve.jl/issues/552
-        put!(ch, (deepcopy(linsolve), Vector{eltype(A)}(undef, size(B, 1))))
-    end
-    Threads.@threads for i in 1:size(B, 2)
-        # Get column memory from the channel
-        linsolve1, b = take!(ch)
-        # Update it
-        b .= view(B, :, i)
-        # Update solver with new b values
-        linsolve2 = LinearSolve.set_b(linsolve1, b)
-        sol = LinearSolve.solve(linsolve2)
-        # Aim for something like this ?
-        # res = map(connectivity_measures(p)) do cm
-        #     compute(cm, g, sol.u, i)
-        # end
-
-        # For now just use Z
-        Z[:, i] .= sol.u
-        put!(ch, (linsolve1, b))
-    end
-    # return _combine(res, g) # return results as Rasters
-
-    # TODO remove all use of GridRSP
-    grsp = GridRSP(g, cm.θ, Pref, W, Z)
-    results = map(p.graph_measures) do gm
-        compute(gm, p, grsp)
-    end 
-    return _merge_to_stack(results)
-end
-# Materialise the whole rhs matrix
-function solve(m::MatrixSolve, cm::FundamentalMeasure, p::AbstractProblem, g::Grid) 
-    # Legacy code... but maybe materialising is faster for CUDSS?
-    Pref = _Pref(g.affinities)
-    W    = _W(Pref, cm.θ, g.costmatrix)
-    Z    = (I - W) \ Matrix(sparse(g.targetnodes,
-                                 1:length(g.targetnodes),
-                                 1.0,
-                                 size(g.costmatrix, 1),
-                                 length(g.targetnodes)))
-    # Check that values in Z are not too small:
-    if minimum(Z) * minimum(nonzeros(g.costmatrix .* W)) == 0
-        @warn "Warning: Z-matrix contains too small values, which can lead to inaccurate results! Check that the graph is connected or try decreasing θ."
-    end
-    grsp = GridRSP(g, cm.θ, Pref, W, Z)
-    results = map(p.graph_measures) do gm
-        compute(gm, p, grsp)
-    end
-    return _merge_to_stack(results)
-end
-function solve(::SolverMode, cm::ConnectivityMeasure, p::AbstractProblem, g::Grid) 
-    # GridRSP is not used here
-    return map(p.graph_measures) do gm
-        compute(gm, p, g)
-    end
-end
-
 # We may have multiple distance_measures per
 # graph_measure, but we want a single RasterStack.
 # So we merge the names of the two layers
-Original file line number
+Diff line change
@@ Expand Up / @@ -418,6 +418,7 @@ function _get_target(rast::AbstractRasterStack) @@
                 throw(ArgumentError("No :target_qualities or :qualities layers found"))
             end
         end
+    end
     """
@@ Expand Down @@