decode for linear disc should propagate sampling method

sisl · Dec 1, 2024 · f33800c · f33800c
1 parent 9b7026a
commit f33800c
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 2 deletions.
diff --git a/src/linear_discretizer.jl b/src/linear_discretizer.jl
@@ -121,10 +121,10 @@ decode(ld::LinearDiscretizer{N,D}, d::D) where {N<:Integer,D<:Integer} = decode(
 decode(ld::LinearDiscretizer{N,D}, d::I, method::AbstractSampleMethod=SAMPLE_UNIFORM) where {N<:Real,D<:Integer,I<:Integer} =
     decode(ld, convert(D,d), method)
 
-function decode(ld::LinearDiscretizer{N,D}, data::AbstractArray{D}, ::AbstractSampleMethod=SAMPLE_UNIFORM) where {N<:Real,D<:Integer}
+function decode(ld::LinearDiscretizer{N,D}, data::AbstractArray{D}, method::AbstractSampleMethod=SAMPLE_UNIFORM) where {N<:Real,D<:Integer}
     arr = Vector{N}(undef, length(data))
     for (i,d) in enumerate(data)
-        arr[i] = decode(ld, d)
+        arr[i] = decode(ld, d, method)
     end
     reshape(arr, size(data))
 end

diff --git a/test/test_linear_discretizer.jl b/test/test_linear_discretizer.jl
@@ -140,3 +140,20 @@ ld = LinearDiscretizer([0,10,20], Int, force_outliers_to_closest=false)
 @test supports_encoding(ld, 16)
 @test !supports_encoding(ld, -1)
 @test !supports_encoding(ld, 21)
+
+let
+    # Ensure that linear discretizer propoagates the sampling method
+    binedges = [0, 0.5, 1]
+    lineardisc = LinearDiscretizer(binedges)
+
+    decode_1 = decode(lineardisc, 1, SAMPLE_BIN_CENTER) # 0.25
+    decode_2 = decode(lineardisc, 2, SAMPLE_BIN_CENTER) # 0.75
+    decode_12_uniform = decode(lineardisc, [1, 2], SAMPLE_UNIFORM)
+
+    @test !(decode_1 ≈ decode_12_uniform[1])
+    @test !(decode_2 ≈ decode_12_uniform[2])
+
+    decode_12_center = decode(lineardisc, [1, 2], SAMPLE_BIN_CENTER)
+    @test decode_1 ≈ decode_12_center[1]
+    @test decode_2 ≈ decode_12_center[2]
+end