start migration to cobrexa2

LCSB-BioCore · Jun 19, 2024 · 63e1f71 · 63e1f71
1 parent 8ab5549
commit 63e1f71
Show file tree

Hide file tree

Showing 6 changed files with 52 additions and 75 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "CuFluxSampler"
 uuid = "d04bc951-f907-4aa3-859a-fc6f9fe5eea9"
 authors = ["Mirek Kratochvil <[email protected]>"]
-version = "0.1.0"
+version = "0.2.0"
 
 [deps]
 COBREXA = "babc4406-5200-4a30-9033-bf5ae714c842"
@@ -11,8 +11,8 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [compat]
-COBREXA = "1"
-CUDA = "4"
+COBREXA = "2"
+CUDA = "5"
 DocStringExtensions = "0.9"
 julia = "1.6"
 

diff --git a/src/CuFluxSampler.jl b/src/CuFluxSampler.jl
@@ -7,10 +7,11 @@ using DocStringExtensions
 import COBREXA
 import CUDA
 import SparseArrays
+import StableRNGs: StableRNG
 
-include("TeaRNG.jl")
-include("FullAffineHR.jl")
-include("AffineHR.jl")
-include("ACHR.jl")
+include("tea_rng.jl")
+include("full_affine_hr.jl")
+include("affine_hr.jl")
+include("achr.jl")
 
 end # module CuFluxSampler
diff --git a/src/ACHR.jl → src/achr.jl b/src/ACHR.jl → src/achr.jl
@@ -1,32 +1,32 @@
-module ACHR
-using ..CUDA
-using ..DocStringExtensions
-using SparseArrays
-
-import ..COBREXA
-import ..TeaRNG
-import Random
 
 """
 $(TYPEDSIGNATURES)
 
 A traditional artificially-centered hit-and-run algorithm that starts with
 `start` points.
 
+To use this on a model, use [`flux_sample`](@ref) or
+[`sample_constraints`](@ref); most parameters are filled in correctly by these
+functions.
+
 Refer to the documentation in module AffineHR for the meaning of arguments.
 """
-function sample(
-    m::COBREXA.MetabolicModel,
-    start::AbstractMatrix;
-    iters::Int,
-    bound_stoichiometry::Bool = false,
-    check_stoichiometry::Bool = true,
+function sample_chain_achr_cuda(
+    sample_c::COBREXA.M;
+    variable_lower_bounds::COBREXA.V,
+    variable_upper_bounds::COBREXA.V,
+    constraints::COBREXA.SM,
+    lower_bounds::COBREXA.V,
+    upper_bounds::COBREXA.V,
+    epsilon::Float32 = COBREXA.configuration.sampler_tolerance,
+    collect_iterations::Vector{Int},
+    generator::StableRNG,
     direction_noise_max::Union{Nothing,Float32} = nothing,
-    epsilon::Float32 = 1.0f-5,
-    seed = Random.rand(UInt32),
 )
+    # TODO
+
     # allocate base helper variables
-    npts = size(start, 2)
+    npts = size(sample_c, 2)
     pts = cu(Matrix{Float32}(start))
     dirs = CUDA.zeros(size(start, 1), npts)
     lblambdas = CUDA.zeros(size(dirs))
@@ -125,7 +125,7 @@ function sample(
             lmax .= min.(lmax, minimum(ifelse.(isfinite.(clmaxs), clmaxs, Inf32), dims = 1))
         end
 
-        # generate random lambdas and compute new points 
+        # generate random lambdas and compute new points
         @cuda threads = 256 blocks = 32 TeaRNG.device_fill_rand!(
             lws,
             seed + UInt32(iter * 2 + 1),
@@ -145,4 +145,4 @@ function sample(
     collect(pts)
 end
 
-end # module AffineHR
+export sample_chain_achr_cuda
diff --git a/src/AffineHR.jl → src/affine_hr.jl b/src/AffineHR.jl → src/affine_hr.jl
@@ -1,11 +1,3 @@
-module AffineHR
-using ..CUDA
-using ..DocStringExtensions
-using SparseArrays
-
-import ..COBREXA
-import ..TeaRNG
-import Random
 
 function random_mix_matrix(npts, mix_points)
     mtx = sparse(
@@ -48,19 +40,21 @@ present.
 If you are generating a sample of the optimal model solution, it is expected
 that the optimum bound is already present in `m`.
 
-Returns a matrix of the same size as `start`.
+Returns blocks of the same size as `sample_c`.
 """
-function sample(
-    m::COBREXA.MetabolicModel,
-    start::AbstractMatrix;
-    iters::Int,
-    bound_stoichiometry::Bool = false,
-    check_stoichiometry::Bool = true,
-    direction_noise_max::Union{Nothing,Float32} = nothing,
-    epsilon::Float32 = 1.0f-5,
-    seed = Random.rand(UInt32),
+function sample_chain_affine_hr_cuda(
+    sample_c::COBREXA.M;
+    variable_lower_bounds::COBREXA.V,
+    variable_upper_bounds::COBREXA.V,
+    constraints::COBREXA.SM,
+    lower_bounds::COBREXA.V,
+    upper_bounds::COBREXA.V,
+    epsilon::Float32 = COBREXA.configuration.sampler_tolerance,
+    collect_iterations::Vector{Int},
+    generator::StableRNG,
     mix_points = 3,
-    mix_mtx = random_mix_matrix(size(start, 2), mix_points),
+    mix_mtx = random_mix_matrix(size(sample_c, 2), mix_points),
+    direction_noise_max::Union{Nothing,Float32} = nothing,
 )
     # allocate base helper variables
     npts = size(start, 2)
@@ -163,7 +157,7 @@ function sample(
             lmax .= min.(lmax, minimum(ifelse.(isfinite.(clmaxs), clmaxs, Inf32), dims = 1))
         end
 
-        # generate random lambdas and compute new points 
+        # generate random lambdas and compute new points
         @cuda threads = 256 blocks = 32 TeaRNG.device_fill_rand!(
             lws,
             seed + UInt32(iter * 2 + 1),
@@ -183,5 +177,3 @@ function sample(
 
     collect(pts)
 end
-
-end # module AffineHR
diff --git a/src/FullAffineHR.jl → src/full_affine_hr.jl b/src/FullAffineHR.jl → src/full_affine_hr.jl
@@ -1,10 +1,3 @@
-module FullAffineHR
-using ..CUDA
-using ..DocStringExtensions
-
-import ..COBREXA
-import ..TeaRNG
-import Random
 
 """
 $(TYPEDSIGNATURES)
@@ -19,16 +12,17 @@ Returns a matrix of `npts` samples organized in columns.
 This algorithm is mostly a toy for comparing the performance. It works, but do
 not use it in production.
 """
-function sample(
-    m::COBREXA.MetabolicModel,
-    warmup::AbstractMatrix;
-    npts::Int = size(warmup, 2),
-    iters::Int,
-    bound_stoichiometry::Bool = false,
-    check_stoichiometry::Bool = true,
+function sample_chain_full_affine_hr_cuda(
+    sample_c::COBREXA.M;
+    variable_lower_bounds::COBREXA.V,
+    variable_upper_bounds::COBREXA.V,
+    constraints::COBREXA.SM,
+    lower_bounds::COBREXA.V,
+    upper_bounds::COBREXA.V,
+    epsilon::Float32 = COBREXA.configuration.sampler_tolerance,
+    collect_iterations::Vector{Int},
+    generator::StableRNG,
     direction_noise_max::Union{Nothing,Float32} = nothing,
-    epsilon::Float32 = 1.0f-5,
-    seed = Random.rand(UInt32),
 )
     # TODO seed and tea iters
 
@@ -140,7 +134,7 @@ function sample(
             lmax .= min.(lmax, minimum(ifelse.(isfinite.(clmaxs), clmaxs, Inf32), dims = 1))
         end
 
-        # generate random lambdas and compute new points 
+        # generate random lambdas and compute new points
         @cuda threads = 256 blocks = 32 TeaRNG.device_fill_rand!(
             lws,
             seed + UInt32(iter * 3 + 2),
@@ -159,5 +153,3 @@ function sample(
 
     collect(pts)
 end
-
-end # module AffineHR
diff --git a/src/TeaRNG.jl → src/tea_rng.jl b/src/TeaRNG.jl → src/tea_rng.jl
@@ -1,9 +1,3 @@
-"""
-Fast stateless random number generator for GPUs based on TEA cipher.
-"""
-module TeaRNG
-using ..CUDA
-using ..DocStringExtensions
 
 """
 $(TYPEDSIGNATURES)
@@ -60,5 +54,3 @@ function device_add_unif_rand!(arr, seed::UInt32, offset::Float32, scale::Float3
     end
     return
 end
-
-end # module TeaRng