Remove sort from setval! and others (#251)

This PR addresses #250 by simply removing the `sort` as mentioned in the issue.

This is also related to TuringLang/Turing.jl#1626, which if addressed, will fix a lot of issues (unless the user then decides to sort their chain themselves).

Co-authored-by: David Widmann <[email protected]>

Author: torfjelde

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.10.20"
+version = "0.11.0"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
@@ -9,7 +9,6 @@ Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
-NaturalSort = "c020b1a1-e9b0-503a-9c33-f039bfc54a85"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
@@ -20,6 +19,5 @@ Bijectors = "0.5.2, 0.6, 0.7, 0.8, 0.9"
 ChainRulesCore = "0.9.7"
 Distributions = "0.23.8, 0.24, 0.25"
 MacroTools = "0.5.6"
-NaturalSort = "1"
 ZygoteRules = "0.2"
 julia = "1.3"

src/DynamicPPL.jl

@@ -7,7 +7,6 @@ using Bijectors
 
 using AbstractMCMC: AbstractMCMC
 using ChainRulesCore: ChainRulesCore
-using NaturalSort: NaturalSort
 using MacroTools: MacroTools
 using ZygoteRules: ZygoteRules
 

src/varinfo.jl

@@ -1325,6 +1325,7 @@ end
 
 """
     setval!(vi::AbstractVarInfo, x)
+    setval!(vi::AbstractVarInfo, values, keys)
     setval!(vi::AbstractVarInfo, chains::AbstractChains, sample_idx::Int, chain_idx::Int)
 
 Set the values in `vi` to the provided values and leave those which are not present in
@@ -1379,20 +1380,18 @@ julia> var_info[@varname(x[1])] # [✓] unchanged
 -0.22312984965118443
 ```
 """
-setval!(vi::AbstractVarInfo, x) = _apply!(_setval_kernel!, vi, values(x), keys(x))
+setval!(vi::AbstractVarInfo, x) = setval!(vi, values(x), keys(x))
+setval!(vi::AbstractVarInfo, values, keys) = _apply!(_setval_kernel!, vi, values, keys)
 function setval!(
     vi::AbstractVarInfo, chains::AbstractChains, sample_idx::Int, chain_idx::Int
 )
-    return _apply!(
-        _setval_kernel!, vi, chains.value[sample_idx, :, chain_idx], keys(chains)
-    )
+    return setval!(vi, chains.value[sample_idx, :, chain_idx], keys(chains))
 end
 
 function _setval_kernel!(vi::AbstractVarInfo, vn::VarName, values, keys)
     indices = findall(Base.Fix1(subsumes_string, string(vn)), keys)
     if !isempty(indices)
-        sorted_indices = sort!(indices; by=i -> keys[i], lt=NaturalSort.natural)
-        val = reduce(vcat, values[sorted_indices])
+        val = reduce(vcat, values[indices])
         setval!(vi, val, vn)
         settrans!(vi, false, vn)
     end
@@ -1402,6 +1401,7 @@ end
 
 """
     setval_and_resample!(vi::AbstractVarInfo, x)
+    setval_and_resample!(vi::AbstractVarInfo, values, keys)
     setval_and_resample!(vi::AbstractVarInfo, chains::AbstractChains, sample_idx, chain_idx)
 
 Set the values in `vi` to the provided values and those which are not present
@@ -1458,24 +1458,21 @@ julia> var_info[@varname(x[1])] # [✓] changed
 - [`setval!`](@ref)
 """
 function setval_and_resample!(vi::AbstractVarInfo, x)
-    return _apply!(_setval_and_resample_kernel!, vi, values(x), keys(x))
+    return setval_and_resample!(vi, values(x), keys(x))
+end
+function setval_and_resample!(vi::AbstractVarInfo, values, keys)
+    return _apply!(_setval_and_resample_kernel!, vi, values, keys)
 end
 function setval_and_resample!(
     vi::AbstractVarInfo, chains::AbstractChains, sample_idx::Int, chain_idx::Int
 )
-    return _apply!(
-        _setval_and_resample_kernel!,
-        vi,
-        chains.value[sample_idx, :, chain_idx],
-        keys(chains),
-    )
+    return setval_and_resample!(vi, chains.value[sample_idx, :, chain_idx], keys(chains))
 end
 
 function _setval_and_resample_kernel!(vi::AbstractVarInfo, vn::VarName, values, keys)
     indices = findall(Base.Fix1(subsumes_string, string(vn)), keys)
     if !isempty(indices)
-        sorted_indices = sort!(indices; by=i -> keys[i], lt=NaturalSort.natural)
-        val = reduce(vcat, values[sorted_indices])
+        val = reduce(vcat, values[indices])
         setval!(vi, val, vn)
         settrans!(vi, false, vn)
     else

test/runtests.jl

@@ -60,18 +60,26 @@ include("test_util.jl")
 
     if GROUP == "All" || GROUP == "Downstream"
         @testset "turing" begin
-            # activate separate test environment
-            Pkg.activate(DIRECTORY_Turing_tests)
-            Pkg.develop(PackageSpec(; path=DIRECTORY_DynamicPPL))
-            Pkg.instantiate()
+            try
+                # activate separate test environment
+                Pkg.activate(DIRECTORY_Turing_tests)
+                Pkg.develop(PackageSpec(; path=DIRECTORY_DynamicPPL))
+                Pkg.instantiate()
 
-            # make sure that the new environment is considered `using` and `import` statements
-            # (not added automatically on Julia 1.3, see e.g. PR #209)
-            if !(joinpath(DIRECTORY_Turing_tests, "Project.toml") in Base.load_path())
-                pushfirst!(LOAD_PATH, DIRECTORY_Turing_tests)
-            end
+                # make sure that the new environment is considered `using` and `import` statements
+                # (not added automatically on Julia 1.3, see e.g. PR #209)
+                if !(joinpath(DIRECTORY_Turing_tests, "Project.toml") in Base.load_path())
+                    pushfirst!(LOAD_PATH, DIRECTORY_Turing_tests)
+                end
 
-            include(joinpath("turing", "runtests.jl"))
+                include(joinpath("turing", "runtests.jl"))
+            catch err
+                err isa Pkg.Resolve.ResolverError || rethrow()
+                # If we can't resolve that means this is incompatible by SemVer and this is fine
+                # It means we marked this as a breaking change, so we don't need to worry about
+                # Mistakenly introducing a breaking change, as we have intentionally made one
+                @info "Not compatible with this release. No problem." exception = err
+            end
         end
     end
 end

test/turing/Project.toml

@@ -5,6 +5,6 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 
 [compat]
-DynamicPPL = "0.10"
+DynamicPPL = "0.11"
 Turing = "0.15"
 julia = "1.3"

test/varinfo.jl

@@ -193,12 +193,20 @@
                 end
                 @test vicopy[s_vns] == vi[s_vns]
 
+                # Ordering is NOT preserved => fails for multivariate model.
                 DynamicPPL.setval!(
                     vicopy, (; (Symbol("m[$i]") => i for i in (1, 3, 5, 4, 2))...)
                 )
-                @test vicopy[m_vns] == 1:5
+                if model == model_uv
+                    @test vicopy[m_vns] == 1:5
+                else
+                    @test vicopy[m_vns] == [1, 3, 5, 4, 2]
+                end
                 @test vicopy[s_vns] == vi[s_vns]
 
+                DynamicPPL.setval!(
+                    vicopy, (; (Symbol("m[$i]") => i for i in (1, 2, 3, 4, 5))...)
+                )
                 DynamicPPL.setval!(vicopy, (s=42,))
                 @test vicopy[m_vns] == 1:5
                 @test vicopy[s_vns] == 42
@@ -222,10 +230,23 @@
                 end
                 @test vicopy[s_vns] != vi[s_vns]
 
+                # Ordering is NOT preserved.
                 DynamicPPL.setval_and_resample!(
                     vicopy, (; (Symbol("m[$i]") => i for i in (1, 3, 5, 4, 2))...)
                 )
                 model(vicopy)
+                if model == model_uv
+                    @test vicopy[m_vns] == 1:5
+                else
+                    @test vicopy[m_vns] == [1, 3, 5, 4, 2]
+                end
+                @test vicopy[s_vns] != vi[s_vns]
+
+                # Correct ordering.
+                DynamicPPL.setval_and_resample!(
+                    vicopy, (; (Symbol("m[$i]") => i for i in (1, 2, 3, 4, 5))...)
+                )
+                model(vicopy)
                 @test vicopy[m_vns] == 1:5
                 @test vicopy[s_vns] != vi[s_vns]
 
@@ -235,5 +256,19 @@
                 @test vicopy[s_vns] == 42
             end
         end
+
+        # https://github.com/TuringLang/DynamicPPL.jl/issues/250
+        @model function demo()
+            return x ~ filldist(MvNormal([1, 100], 1), 2)
+        end
+
+        vi = VarInfo(demo())
+        vals_prev = vi.metadata.x.vals
+        ks = [@varname(x[1, 1]), @varname(x[2, 1]), @varname(x[1, 2]), @varname(x[2, 2])]
+        DynamicPPL.setval!(vi, vi.metadata.x.vals, ks)
+        @test vals_prev == vi.metadata.x.vals
+
+        DynamicPPL.setval_and_resample!(vi, vi.metadata.x.vals, ks)
+        @test vals_prev == vi.metadata.x.vals
     end
 end