add lazy compute operations

Project.toml

@@ -8,6 +8,7 @@ DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 ProgressLogging = "33c8b6b6-d38a-422a-b730-caa89a2f386c"
 Rasters = "a3a2b9e3-a471-40c9-b274-f788e487c689"
@@ -20,9 +21,10 @@ ArnoldiMethod = "0.0.4, 0.4"
 DelimitedFiles = "1"
 Graphs = "1"
 LaTeXStrings = "1.1"
+LinearSolve = "2.38.0"
 Plots = "1.4"
 ProgressLogging = "0.1"
-Rasters = "0.12"
+Rasters = "0.13"
 SimpleWeightedGraphs = "1.1"
 julia = "1.10"
 

examples/2_landmarks.jmd

@@ -1,6 +1,7 @@
 # Landmark function for testing purposes
 
 ```julia
+nothing
 using ConScape, Rasters, ArchGDAL, SparseArrays, LinearAlgebra, Plots
 ```
 
@@ -83,7 +84,7 @@ heatmap(qbetw)
 
 ```julia
 qbetw_coarse = ConScape.betweenness_qweighted(h_coarse);
-ConScape.heatmap(qbetw_coarse)
+heatmap(qbetw_coarse)
 ```
 
 ```julia
@@ -99,6 +100,80 @@ kbetw_coarse = @time ConScape.betweenness_kweighted(h_coarse, distance_transform
 cor(filter(!isnan, kbetw), filter(!isnan, kbetw_coarse))
 ```
 
+We can write this using a lazy problem definition:
+
+```julia
+# Put least cost, random walk, and rsp
+problem = ConScape.Problem(; 
+    graph_measures = (;
+        func=ConScape.ConnectedHabitat(),
+        qbetw=ConScape.BetweennessQweighted(),
+        kbetw=ConScape.BetweennessKweighted(),
+    ),
+    connectivity_measure=ConScape.ExpectedCost(
+        θ=1.0, 
+        distance_transformation = (exp=x -> exp(-x/75), oddsfor=ConScape.OddsFor())
+    ),
+    solver=ConScape.MatrixSolver(),
+    # solver=ConScape.VectorSolver(nothing),
+)
+```
+
+Then run it for all operations on both normal and coarse grids
+
+```julia
+rast = RasterStack((; affinities=mov_prob, qualities=hab_qual))
+rast_coarse = ConScape.coarse_graining(rast, 20)
+result = ConScape.solve(problem, rast)
+result_coarse = ConScape.solve(problem, rast_coarse)
+```
+
+We can plot these outputs:
+```julia
+heatmap(result.qbetw_oddsfor)
+heatmap(result_coarse.qbetw_oddsfor)
+heatmap(result.qbetw_exp)
+heatmap(result_coarse.qbetw_exp)
+heatmap(result.kbetw_exp)
+heatmap(result_coarse.kbetw_exp)
+heatmap(result.func_exp)
+heatmap(result.coarse_func_exp)
+```
+
+
+```julia
+windowed_problem = ConScape.WindowedProblem(problem; 
+    radius=20, overlap=5, threaded=true
+)
+@time result = ConScape.solve(windowed_problem, rast)
+plot(result)
+plot(result.func_exp)
+```
+
+
+```julia
+stored_problem = ConScape.StoredProblem(problem; 
+    path=".", radius=20, overlap=30, threaded=true
+)
+ConScape.solve(stored_problem, rast)
+result = mosaic(stored_problem; to=rast)
+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
+layers from `res` and `res_coarse`
+```julia
+cor(filter(!isnan, res.qbetw.oddsfor), filter(!isnan, res_coarse.qbetw.oddsfor))
+cor(filter(!isnan, res.qbetw.oddsfor), filter(!isnan, res_coarse.qbetw.oddsfor))
+cor(filter(!isnan, res.kbetw.oddsfor), filter(!isnan, res_coarse.kbetw.oddsfor))
+cor(filter(!isnan, res.kbetw.exp), filter(!isnan, res_coarse.kbetw.exp))
+cor(filter(!isnan, res.func.exp), filter(!isnan, res_coarse.func.exp))
+```
+
 # Test landmark performance for amount of connected habitat
 
 ```julia

ext/ConScapeCUDSS.jl

@@ -0,0 +1,37 @@
+module ConScapeCUDSS
+
+using ConScape
+using CUDA
+using CUDA.CUSPARSE
+using CUDSS
+using SparseArrays
+using LinearAlgebra
+
+using ConScape: FundamentalMeasure, AbstractProblem, Grid, setup_sparse_problem
+
+struct CUDSSsolver <: Solver end
+
+function ConScape.solve(m::CUDSSsolver, cm::FundamentalMeasure, p::AbstractProblem, g::Grid) 
+    (; A, B, Pref, W) = setup_sparse_problem(g, cm)
+    Z = zeros(T, size(B))
+
+    A_gpu = CuSparseMatrixCSR(A |> tril)
+    Z_gpu = CuMatrix(Z)
+    B_gpu = CuMatrix(B)
+
+    solver = CudssSolver(A_gpu, "S", "L")
+
+    cudss("analysis", solver, Z_gpu, B_gpu)
+    cudss("factorization", solver, Z_gpu, B_gpu)
+    cudss("solve", solver, Z_gpu, B_gpu)
+
+    Z .= Z_gpu
+    # TODO: maybe graph measures can run on GPU as well?
+    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
+
+end

src/ConScape.jl

@@ -1,30 +1,41 @@
 module ConScape
 
-    using SparseArrays, LinearAlgebra
-    using Graphs, Plots, SimpleWeightedGraphs, ProgressLogging, ArnoldiMethod
-    using Rasters
-    using Rasters.DimensionalData
+using SparseArrays, LinearAlgebra
+using Graphs, Plots, SimpleWeightedGraphs, ProgressLogging, ArnoldiMethod
+using Rasters
+using LinearSolve
+using Rasters.DimensionalData
 
-    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
 
-    # 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")
+# 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")
 
 end

src/connectivity_measure.jl

@@ -0,0 +1,44 @@
+# New type-based interface
+# Easier to add parameters to these
+abstract type ConnectivityMeasure end
+
+# TODO document these groups
+abstract type FundamentalMeasure <: ConnectivityMeasure end
+abstract type DistanceMeasure <: FundamentalMeasure end
+
+struct LeastCostDistance <: ConnectivityMeasure end
+@kwdef struct ExpectedCost{T<:Union{Real,Nothing},CM} <: DistanceMeasure 
+    θ::T=nothing
+    distance_transformation::CM=nothing
+    approx::Bool=false
+end
+@kwdef struct FreeEnergyDistance{T<:Union{Real,Nothing},CM} <: DistanceMeasure 
+    θ::T=nothing
+    distance_transformation::CM=nothing
+    approx::Bool=false
+end
+@kwdef struct SurvivalProbability{T<:Union{Real,Nothing}} <: FundamentalMeasure 
+    θ::T=nothing
+    approx::Bool=false
+end
+@kwdef struct PowerMeanProximity{T<:Union{Real,Nothing}} <: FundamentalMeasure 
+    θ::T=nothing
+    approx::Bool=false
+end
+
+keywords(cm::ConnectivityMeasure) = _keywords(cm)
+
+# TODO remove the complexity of the connectivity_function
+# These methods are mostly to avoid changing the original interface for now
+connectivity_function(::LeastCostDistance) = least_cost_distance
+connectivity_function(::ExpectedCost) = expected_cost
+connectivity_function(::FreeEnergyDistance) = free_energy_distance
+connectivity_function(::SurvivalProbability) = survival_probability
+connectivity_function(::PowerMeanProximity) = power_mean_proximity
+
+distance_transformation(::ConnectivityMeasure) = nothing
+distance_transformation(cm::Union{ExpectedCost,FreeEnergyDistance}) = cm.distance_transformation
+
+# This is not used yet but could be
+compute(cm::ConnectivityMeasure, g; kw...) = 
+    connectivity_function(m)(g; keywords(cm)..., kw...)

src/graph_measure.jl

@@ -0,0 +1,107 @@
+"""
+    GraphMeasure 
+
+Abstract supertype for graph measures.
+These are lazy definitions of conscape functions.
+"""
+abstract type GraphMeasure end
+
+keywords(o::GraphMeasure) = _keywords(o)
+
+# TODO: document/rethink these
+abstract type TopologicalMeasure <: GraphMeasure end
+abstract type BetweennessMeasure <: GraphMeasure end
+abstract type PerturbationMeasure <: GraphMeasure end
+abstract type PathDistributionMeasure <: GraphMeasure end
+
+struct BetweennessQweighted <: BetweennessMeasure end
+@kwdef struct BetweennessKweighted{DV} <: BetweennessMeasure 
+    diagvalue::DV=nothing
+end
+struct EdgeBetweennessQweighted <: BetweennessMeasure end
+@kwdef struct EdgeBetweennessKweighted{DV} <: BetweennessMeasure 
+    diagvalue::DV=nothing
+end
+
+@kwdef struct ConnectedHabitat{DV} <: GraphMeasure
+    diagvalue::DV=nothing
+end
+
+@kwdef struct Criticality{DV,AV,QT,QS} <: PerturbationMeasure 
+    diagvalue::DV=nothing
+    avalue::AV=floatmin()
+    qˢvalue::QS=0.0
+    qᵗvalue::QT=0.0
+end
+
+# These maybe don't quite belong here?
+@kwdef struct EigMax{DV,T} <: TopologicalMeasure
+    diagvalue::DV=nothing
+    tol::T=1e-14
+end
+
+struct MeanLeastCostKullbackLeiblerDivergence <: PathDistributionMeasure end
+struct MeanKullbackLeiblerDivergence <: PathDistributionMeasure end
+
+# Map structs to functions
+
+# These return Rasters
+graph_function(m::BetweennessKweighted) = betweenness_kweighted
+graph_function(m::BetweennessQweighted) = betweenness_qweighted
+graph_function(m::ConnectedHabitat) = connected_habitat
+graph_function(m::Criticality) = criticality
+# These return scalars
+graph_function(m::MeanLeastCostKullbackLeiblerDivergence) = mean_lc_kl_divergence
+graph_function(m::MeanKullbackLeiblerDivergence) = mean_kl_divergence
+# These return sparse arrays
+graph_function(m::EdgeBetweennessKweighted) = edge_betweenness_kweighted
+graph_function(m::EdgeBetweennessQweighted) = edge_betweenness_qweighted
+# Returns a tuple
+graph_function(m::EigMax) = eigmax
+
+# Map structs to function keywords, 
+# a bit of a hack until we refactor the rest
+keywords(gm::GraphMeasure, p::AbstractProblem) = 
+    (; _keywords(gm)...)#, solver=solver(p))
+keywords(gm::ConnectedHabitat, p::AbstractProblem) = 
+    (; _keywords(gm)..., approx=connectivity_measure(p).approx)#, solver=solver(p))
+
+# A trait for connectivity requirement
+struct NeedsConnectivity end
+struct NoConnectivity end
+needs_connectivity(::GraphMeasure) = NoConnectivity()
+needs_connectivity(::BetweennessKweighted) = NeedsConnectivity()
+needs_connectivity(::EdgeBetweennessKweighted) = NeedsConnectivity()
+needs_connectivity(::EigMax) = NeedsConnectivity()
+needs_connectivity(::ConnectedHabitat) = NeedsConnectivity()
+needs_connectivity(::Criticality) = NeedsConnectivity()
+
+# compute
+# This is where things actually happen
+#
+# Add dispatch on connectivity measure
+compute(gm::GraphMeasure, p::AbstractProblem, g::Union{Grid,GridRSP}) = 
+    compute(needs_connectivity(gm), gm, p, g)
+function compute(::NeedsConnectivity,
+    gm::GraphMeasure, 
+    p::AbstractProblem, 
+    g::Union{Grid,GridRSP}
+)
+    cf = connectivity_function(p)
+    dt = distance_transformation(p) 
+    # Handle multiple distance transformations
+    if dt isa NamedTuple
+        map(distance_transformation) do dtx
+            graph_function(gm)(g; keywords(gm, p)..., distance_transformation=dtx, connectivity_function=cf)
+        end
+    else
+        graph_function(gm)(g; keywords(gm, p)..., distance_transformation=dt, connectivity_function=cf)
+    end
+end
+function compute(::NoConnectivity,
+    gm::GraphMeasure, 
+    p::AbstractProblem, 
+    g::Union{Grid,GridRSP}
+) 
+    graph_function(gm)(g; keywords(gm, p)...)
+end