fix horrifying predict output transpose issue (#144)

ExpandingMan · web-flow · commit f06053671b19 · 2022-12-16T16:25:48.000-05:00
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "XGBoost"
 uuid = "009559a3-9522-5dbb-924b-0b6ed2b22bb9"
-version = "2.1.0"
+version = "2.1.1"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
diff --git a/assets/data/blobs.json b/assets/data/blobs.json
diff --git a/src/booster.jl b/src/booster.jl
@@ -255,6 +255,7 @@ function deserialize(::Type{Booster}, buf::AbstractVector{UInt8}, data=DMatrix[]
     deserialize!(b, buf)
 end
 
+# sadly this is type unstable because we might return a transpose
 """
     predict(b::Booster, data; margin=false, training=false, ntree_limit=0)
 
@@ -287,8 +288,9 @@ function predict(b::Booster, Xy::DMatrix;
     odim = Ref{Lib.bst_ulong}()
     o = Ref{Ptr{Cfloat}}()
     xgbcall(XGBoosterPredictFromDMatrix, b.handle, Xy.handle, opts, oshape, odim, o)
-    dims = unsafe_wrap(Array, oshape[], odim[])
-    unsafe_wrap(Array, o[], tuple(dims...))
+    dims = reverse(unsafe_wrap(Array, oshape[], odim[]))
+    o = unsafe_wrap(Array, o[], tuple(dims...))
+    length(dims) > 1 ? transpose(o) : o
 end
 predict(b::Booster, Xy; kw...) = predict(b, DMatrix(Xy); kw...)
 
diff --git a/test/gendata.jl b/test/gendata.jl
@@ -0,0 +1,19 @@
+# this file is for generating test data that's stored in the repository
+# it uses MLJ which is not a dependency
+
+# JSON is used as the serialization format only because it is already a dependency
+
+using MLJ
+using JSON3
+
+
+# see utils.jl for loading of the output
+function gen_classification()
+    (X, y) = make_blobs(1000, 3; centers=3, rng=999)
+    y = Int.(int.(y)) .- 1
+    X = MLJ.matrix(X)
+    fname = joinpath(@__DIR__,"..","assets","data","blobs.json")
+    dict = Dict("X1"=>X[:,1], "X2"=>X[:,2], "X3"=>X[:,3], "y"=>y)
+    JSON3.write(open(fname, write=true), dict)
+end
+
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -126,6 +126,17 @@ end
     end
 end
 
+@testset "Blobs training" begin
+    (X, y) = load_classification()
+
+    bst = xgboost((X, y), num_round=10, objective="multi:softprob", num_class=3, watchlist=Dict())
+
+    y&#770; = map(&zeta; -> argmax(&zeta;) - 1, eachrow(predict(bst, X)))
+
+    # this is a pretty low bar that xgboost should always pass
+    @test sum(y&#770; .== y)/length(y) > 0.9
+end
+
 @testset "Feature importance" begin
     dtrain = XGBoost.load(DMatrix, testfilepath("agaricus.txt.train"))
     dtest = XGBoost.load(DMatrix, testfilepath("agaricus.txt.test"))
diff --git a/test/utils.jl b/test/utils.jl
@@ -1,3 +1,5 @@
+using JSON3
+
 
 function testfilepath(name::AbstractString) 
     dir = joinpath(dirname(pathof(XGBoost)), "..")
@@ -22,3 +24,11 @@ function readlibsvm(fname::AbstractString, shape)
     end
     (dmx, label)
 end
+
+function load_classification()
+    fname = joinpath(@__DIR__,"..","assets","data","blobs.json")
+    o = JSON3.read(String(open(read, fname)))
+    X = Matrix{Float32}([o[:X1] o[:X2] o[:X3]])
+    y = o[:y]
+    (X, y)
+end
