Refactor extension folder (#75)

* Refactor extension folder * More cleanup * Update README * Update test data * Update CI.yml
JuliaML · Nov 26, 2024 · 6be1bd3 · 6be1bd3
1 parent c42e38b
commit 6be1bd3
Show file tree

Hide file tree

Showing 31 changed files with 49 additions and 88 deletions.
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -17,7 +17,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.9'
+          - 'lts'
           - '1'
         os:
           - ubuntu-latest

diff --git a/Project.toml b/Project.toml
@@ -4,33 +4,31 @@ authors = ["Kai Xu <[email protected]>", "Júlio Hoffimann <julio.hoffimann@
 version = "1.2.5"
 
 [deps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [weakdeps]
-ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Convex = "f65535da-76fb-5f13-bab9-19810c17039a"
 ECOS = "e2685f51-7e38-5353-a97d-a921fd2c8199"
-GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 
 [extensions]
-DensityRatioEstimationChainRulesCoreExt = "ChainRulesCore"
 DensityRatioEstimationConvexExt = ["Convex", "ECOS"]
-DensityRatioEstimationGPUArraysExt = "GPUArrays"
 DensityRatioEstimationJuMPExt = ["JuMP", "Ipopt"]
 DensityRatioEstimationOptimExt = "Optim"
 
 [compat]
-ChainRulesCore = "1"
+ChainRulesCore = "1.25.0"
 Convex = "0.15, 0.16"
 ECOS = "1"
-GPUArrays = "8, 9, 10, 11"
+GPUArraysCore = "0.2.0"
 Ipopt = "1"
 JuMP = "1"
 LinearAlgebra = "1.9"
@@ -40,12 +38,3 @@ Random = "1.9"
 Statistics = "1.9"
 StatsBase = "0.33, 0.34"
 julia = "1.9"
-
-[extras]
-ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-Convex = "f65535da-76fb-5f13-bab9-19810c17039a"
-ECOS = "e2685f51-7e38-5353-a97d-a921fd2c8199"
-GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
-Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
-JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
-Optim = "429524aa-4258-5aef-a3af-852621145aeb"
diff --git a/README.md b/README.md
@@ -51,7 +51,7 @@ Currently, this package implements the following estimators:
 
 <sup>1</sup> We use the naming convention of prefixing the type name with `u` for the unconstrained variant of the corresponding estimator.
 
-The fourth argument `optlib` specifies the optimization package used to implement
+The keyword argument `optlib` specifies the optimization package used to implement
 the estimator. Some estimators are implemented with different optimization packages
 to facilitate the usage in different environments. In the example above, users that
 already have the [Optim.jl](https://github.com/JuliaNLSolvers/Optim.jl) package in
@@ -83,12 +83,22 @@ and to find out the available implementations, please use
 available_optlib(KLIEP)
 ```
 
+In the case of `ConvexLib` and `JuMPLib`, specific optimizer packages must be loaded
+besides the optimization library itself. For example, one must load `Convex` *and* `ECOS`
+to use `optlib=ConvexLib` with the `KLIEP` estimator, or `JuMP` *and* `Ipopt` to use
+`optlib=JuMPLib` with the `LSIF` estimator.
+design.
+
+### Density ratio function
+
 Some methods support the evaluation of the density ratio at all `x`, besides the
-denominator samples. In this case, the following line returns a function `r(x)`
+denominator samples. In this case, `densratiofunc` returns a function `r`
 that can be evaluated at new unseen samples:
 
 ```julia
 r = densratiofunc(x_nu, x_de, KLIEP())
+
+r(x) # evaluate at new x
 ```
 
 ### Hyperparameter tuning

diff --git a/ext/DensityRatioEstimationChainRulesCoreExt.jl b/ext/DensityRatioEstimationChainRulesCoreExt.jl
diff --git a/ext/DensityRatioEstimationConvexExt.jl → ...vexExt/DensityRatioEstimationConvexExt.jl b/ext/DensityRatioEstimationConvexExt.jl → ...vexExt/DensityRatioEstimationConvexExt.jl
@@ -4,11 +4,12 @@
 
 module DensityRatioEstimationConvexExt
 
-using DensityRatioEstimation
 using DensityRatioEstimation: KLIEP, ConvexLib
 using Convex
 using ECOS
 
-include("../src/kliep/convex.jl")
+import DensityRatioEstimation
 
-end #module
+include("kliep.jl")
+
+end
diff --git a/src/kliep/convex.jl → ext/DensityRatioEstimationConvexExt/kliep.jl b/src/kliep/convex.jl → ext/DensityRatioEstimationConvexExt/kliep.jl
@@ -2,8 +2,6 @@
 # Licensed under the MIT License. See LICENSE in the project root.
 # ------------------------------------------------------------------
 
-# This file is part of the module DensityRatioEstimationConvexExt.
-
 function DensityRatioEstimation._kliep_coeffs(K_nu, K_de, dre::KLIEP, optlib::Type{ConvexLib})
   # retrieve parameters
   σ, b = dre.σ, size(K_de, 2)

diff --git a/ext/DensityRatioEstimationGPUArraysExt.jl b/ext/DensityRatioEstimationGPUArraysExt.jl
diff --git a/ext/DensityRatioEstimationJuMPExt.jl → ...nJuMPExt/DensityRatioEstimationJuMPExt.jl b/ext/DensityRatioEstimationJuMPExt.jl → ...nJuMPExt/DensityRatioEstimationJuMPExt.jl
@@ -4,15 +4,16 @@
 
 module DensityRatioEstimationJuMPExt
 
-using DensityRatioEstimation
 using DensityRatioEstimation: LSIF, JuMPLib, AbstractKMM, uKMM, KMM
 using DensityRatioEstimation.Parameters
 using JuMP
 using Ipopt
 using LinearAlgebra
 using Statistics
 
-include("../src/kmm/jump.jl")
-include("../src/lsif/jump.jl")
+import DensityRatioEstimation
 
-end #module
+include("kmm.jl")
+include("lsif.jl")
+
+end
diff --git a/src/kmm/jump.jl → ext/DensityRatioEstimationJuMPExt/kmm.jl b/src/kmm/jump.jl → ext/DensityRatioEstimationJuMPExt/kmm.jl
@@ -2,8 +2,6 @@
 # Licensed under the MIT License. See LICENSE in the project root.
 # ------------------------------------------------------------------
 
-# This file is part of the module DensityRatioEstimationJuMPExt.
-
 function _kmm_jump_model(K, κ, dre::AbstractKMM, optlib::Type{JuMPLib})
   # number of denominator samples
   m = length(κ)

diff --git a/src/lsif/jump.jl → ext/DensityRatioEstimationJuMPExt/lsif.jl b/src/lsif/jump.jl → ext/DensityRatioEstimationJuMPExt/lsif.jl
@@ -2,8 +2,6 @@
 # Licensed under the MIT License. See LICENSE in the project root.
 # ------------------------------------------------------------------
 
-# This file is part of the module DensityRatioEstimationJuMPExt.
-
 function DensityRatioEstimation._lsif_coeffs(H, h, dre::LSIF, optlib::Type{JuMPLib})
   # retrieve parameters
   λ, b = dre.λ, length(h)

diff --git a/ext/DensityRatioEstimationOptimExt.jl → ...ptimExt/DensityRatioEstimationOptimExt.jl b/ext/DensityRatioEstimationOptimExt.jl → ...ptimExt/DensityRatioEstimationOptimExt.jl
@@ -4,13 +4,13 @@
 
 module DensityRatioEstimationOptimExt
 
-using DensityRatioEstimation
 using DensityRatioEstimation: KLIEP, LSIF, OptimLib
+using LinearAlgebra
 using Optim
 
-using LinearAlgebra
+import DensityRatioEstimation
 
-include("../src/kliep/optim.jl")
-include("../src/lsif/optim.jl")
+include("kliep.jl")
+include("lsif.jl")
 
-end #module
+end
diff --git a/src/kliep/optim.jl → ext/DensityRatioEstimationOptimExt/kliep.jl b/src/kliep/optim.jl → ext/DensityRatioEstimationOptimExt/kliep.jl
@@ -2,8 +2,6 @@
 # Licensed under the MIT License. See LICENSE in the project root.
 # ------------------------------------------------------------------
 
-# This file is part of the module DensityRatioEstimationOptimExt.
-
 function DensityRatioEstimation._kliep_coeffs(K_nu, K_de, dre::KLIEP, optlib::Type{OptimLib})
   # retrieve parameters
   σ, b = dre.σ, size(K_de, 2)

diff --git a/src/lsif/optim.jl → ext/DensityRatioEstimationOptimExt/lsif.jl b/src/lsif/optim.jl → ext/DensityRatioEstimationOptimExt/lsif.jl
@@ -2,8 +2,6 @@
 # Licensed under the MIT License. See LICENSE in the project root.
 # ------------------------------------------------------------------
 
-# This file is part of the module DensityRatioEstimationOptimExt.
-
 function DensityRatioEstimation._lsif_coeffs(H, h, dre::LSIF, optlib::Type{OptimLib})
   # retrieve parameters
   λ, b = dre.λ, length(h)

diff --git a/src/DensityRatioEstimation.jl b/src/DensityRatioEstimation.jl
@@ -10,7 +10,9 @@ using LinearAlgebra
 using Parameters
 using Random
 
-# implement fit for estimators
+using GPUArraysCore: AbstractGPUMatrix
+using ChainRulesCore: @non_differentiable
+
 import StatsBase: fit
 
 # API for density ratio estimation
@@ -29,9 +31,6 @@ include("lsif.jl")
 # available estimator fitters
 include("lcv.jl")
 
-# pure Julia implementations
-include("kmm/julia.jl")
-
 export
   # optim libs
   OptimizationLibrary,

diff --git a/src/kmm.jl b/src/kmm.jl
@@ -88,3 +88,9 @@ default_optlib(dre::Type{<:KMM}) = JuMPLib
 available_optlib(dre::Type{<:KMM}) = [JuMPLib]
 
 _kmm_ratios(K, κ, dre::AbstractKMM, optlib::Type{<:OptimizationLibrary}) = _throw_opt_error(dre, optlib)
+
+# pure Julia implementation
+function _kmm_ratios(K, κ, dre::uKMM, optlib::Type{JuliaLib})
+  # density ratio via solver
+  K \ vec(κ)
+end
diff --git a/src/kmm/julia.jl b/src/kmm/julia.jl
diff --git a/src/utils.jl b/src/utils.jl
@@ -9,8 +9,6 @@ Return the Euclidean distance between two indexable objects.
 """
 euclidsq(x, y) = sum((x[i] - y[i])^2 for i in eachindex(x))
 
-# Support matrix data in a GPU and AD compatible way
-
 """
     euclidsq(X::T, Y::T) where {T<:AbstractMatrix}
 
@@ -55,13 +53,18 @@ gaussian_gramian(esq, σ::AbstractFloat) = exp.(-esq ./ 2σ^2)
 Generate a squared matrix whose diagonal is `a` that is 
 compatible to perform addition on `mat`. It behaves 
 differently based on whether `mat` is on a CPU or GPU.
-
-It is compatible with
-- CuArrays.jl (see lib/cuarrays.jl)
-- Zygote.jl (see lib/zygote.jl)
 """
 safe_diagm(mat, a) = a * I
 
+# avoid `mat + a * I` on GPU which involves scalar operations and is slow
+function safe_diagm(mat::AbstractGPUMatrix, a)
+  diag = similar(mat, size(m, 1))
+  fill!(diag, a)
+  Diagonal(diag)
+end
+
+@non_differentiable safe_diagm(::Any, ::Any)
+
 ###################################################
 ##   Functions and objects for throwing errors   ##
 ###################################################

diff --git a/test/data/KLIEP-ConvexLib-1.png b/test/data/KLIEP-ConvexLib-1.png
diff --git a/test/data/KLIEP-ConvexLib-2.png b/test/data/KLIEP-ConvexLib-2.png
diff --git a/test/data/KLIEP-OptimLib-1.png b/test/data/KLIEP-OptimLib-1.png
diff --git a/test/data/KLIEP-OptimLib-2.png b/test/data/KLIEP-OptimLib-2.png
diff --git a/test/data/KMM-JuMPLib-1.png b/test/data/KMM-JuMPLib-1.png
diff --git a/test/data/KMM-JuMPLib-2.png b/test/data/KMM-JuMPLib-2.png
diff --git a/test/data/LSIF-JuMPLib-1.png b/test/data/LSIF-JuMPLib-1.png
diff --git a/test/data/LSIF-JuMPLib-2.png b/test/data/LSIF-JuMPLib-2.png
diff --git a/test/data/LSIF-OptimLib-1.png b/test/data/LSIF-OptimLib-1.png
diff --git a/test/data/LSIF-OptimLib-2.png b/test/data/LSIF-OptimLib-2.png
diff --git a/test/data/uKMM-JuMPLib-1.png b/test/data/uKMM-JuMPLib-1.png
diff --git a/test/data/uKMM-JuMPLib-2.png b/test/data/uKMM-JuMPLib-2.png
diff --git a/test/data/uKMM-JuliaLib-1.png b/test/data/uKMM-JuliaLib-1.png
diff --git a/test/data/uKMM-JuliaLib-2.png b/test/data/uKMM-JuliaLib-2.png
-Original file line number
+Diff line change
@@ Expand Up / @@ -17,7 +17,7 @@ jobs: @@
           fail-fast: false
           matrix:
             version:
-              - '1.9'
+              - 'lts'
               - '1'
             os:
               - ubuntu-latest
@@ Expand Down @@