Fixing broken tests due to mask meaning change

Author: murrellb

examples/continuous.jl

@@ -3,7 +3,7 @@ Pkg.activate(".")
 using Revise
 Pkg.develop(path="../../ForwardBackward/")
 Pkg.develop(path="../")
-using ForwardBackward, Flowfusion, NNlib, Flux, RandomFeatureMaps, Optimisers, Plots
+using ForwardBackward, Flowfusion, Flux, RandomFeatureMaps, Optimisers, Plots
 
 #Set up a Flux model: X̂1 = model(t,Xt)
 struct FModel{A}
@@ -34,47 +34,51 @@ model = FModel(embeddim = 256, layers = 3, spacedim = 2)
 
 #Distributions for training:
 T = Float32
-sampleX1(n_samples) = Flowfusion.random_literal_cat(n_samples, sigma = T(0.05))
 sampleX0(n_samples) = rand(T, 2, n_samples) .+ 2
-n_samples = 200
+sampleX1(n_samples) = Flowfusion.random_literal_cat(n_samples, sigma = T(0.05))
+n_samples = 400
 
 #The process:
-P = BrownianMotion(0.1f0)
+P = BrownianMotion(0.15f0)
 #P = Deterministic()
 
 #Optimizer:
-eta = 0.01
-opt_state = Flux.setup(AdamW(eta = eta, lambda = 0.001), model)
+eta = 0.001
+opt_state = Flux.setup(AdamW(eta = eta), model)
 
-iters = 5000
+iters = 4000
 for i in 1:iters
     #Set up a batch of training pairs, and t:
-    X1 = ContinuousState(sampleX1(n_samples))
     X0 = ContinuousState(sampleX0(n_samples))
+    X1 = ContinuousState(sampleX1(n_samples))
     t = rand(T, n_samples)
     #Construct the bridge:
     Xt = bridge(P, X0, X1, t)
-    #Gradient:
+    #Gradient & update:
     l,g = Flux.withgradient(model) do m
         floss(P, m(t,Xt), X1, scalefloss(P, t))
     end
-    #Update:
     Flux.update!(opt_state, model, g[1])
-    #Logging, and lr cooldown:
-    if i % 10 == 0
-        if i > iters - 2000
-            eta *= 0.975
-            Optimisers.adjust!(opt_state, eta)
-        end
-        println("i: $i; Loss: $l; eta: $eta")
-    end
+    (i % 10 == 0) && println("i: $i; Loss: $l")
 end
 
+#Generate samples by stepping from X0
+n_inference_samples = 5000
+X0 = ContinuousState(sampleX0(n_inference_samples))
+samples = gen(P, X0, model, 0f0:0.005f0:1f0)
+
+#Plotting
+pl = scatter(X0.state[1,:],X0.state[2,:], msw = 0, ms = 1, color = "blue", alpha = 0.5, size = (400,400), legend = :topleft, label = "X0")
+X1true = sampleX1(n_inference_samples)
+scatter!(X1true[1,:],X1true[2,:], msw = 0, ms = 1, color = "orange", alpha = 0.5, label = "X1 (true)")
+scatter!(samples.state[1,:],samples.state[2,:], msw = 0, ms = 1, color = "green", alpha = 0.5, label = "X1 (generated)")
+
+
 #Generate samples by stepping from X0
 n_inference_samples = 5000
 X0 = ContinuousState(sampleX0(n_inference_samples))
 paths = Tracker()
-samp = gen(P, X0, model, 0f0:0.005f0:1f0, tracker = paths)
+samples = gen(P, X0, model, 0f0:0.005f0:1f0, tracker = paths)
 
 #Plotting:
 pl = scatter(X0.state[1,:],X0.state[2,:], msw = 0, ms = 1, color = "blue", alpha = 0.5, size = (400,400), legend = :topleft, label = "X0")
@@ -85,6 +89,6 @@ for i in 1:50:1000
 end
 X1true = sampleX1(n_inference_samples)
 scatter!(X1true[1,:],X1true[2,:], msw = 0, ms = 1, color = "orange", alpha = 0.5, label = "X1 (true)")
-scatter!(samp.state[1,:],samp.state[2,:], msw = 0, ms = 1, color = "green", alpha = 0.5, label = "X1 (generated)")
+scatter!(samples.state[1,:],samples.state[2,:], msw = 0, ms = 1, color = "green", alpha = 0.5, label = "X1 (generated)")
 display(pl)
 savefig("continuous_cat_$P.svg")

test/runtests.jl

@@ -16,7 +16,7 @@ using ForwardBackward
         XDL() = CategoricalLikelihood(rand(5, siz...))
         XGL() = GaussianLikelihood(randn(5, siz...), randn(5, siz...), zeros(siz...))
 
-        for f in [XC, XD, XT, XR, XDL, XGL]
+        for f in [XC, XD, XT, XR, XDL, XGL] 
             Xa = f()
             Xb = f()
             Xc = Flowfusion.mask(Xa, Xb)
@@ -30,15 +30,15 @@ using ForwardBackward
 
             @test typeof(Xc) == typeof(XM) #If you mask a regular State with a MaskedState, the result is a MaskedState.
             d = (tensor(Xb) .- tensor(Xc))
-            @test isapprox(sum(d .* expand(m, ndims(d))),0)
+            @test isapprox(sum(d .* expand(.!m, ndims(d))),0)
 
             m = rand(Bool, siz...)
             Xa = MaskedState(f(), m, m)
             Xb = MaskedState(f(), m, m)
             Xc = Flowfusion.mask(Xa, Xb)
             @test typeof(Xc) == typeof(Xa)
             d = (tensor(Xb) .- tensor(Xc))
-            @test isapprox(sum(d .* expand(m, ndims(d))),0)
+            @test isapprox(sum(d .* expand(.!m, ndims(d))),0)
         end
     end
 
@@ -62,11 +62,11 @@ using ForwardBackward
             m = rand(Bool, siz...)
             XM = MaskedState(Xb, m, m)
             Xt = Flowfusion.bridge(p, Xa, XM, 0.1)
-            @assert typeof(Xt) == typeof(XM)
+            @test typeof(Xt) == typeof(XM)
             if !(p isa InterpolatingDiscreteFlow)
-                @assert isapprox(sum((tensor(Xt) .== tensor(Xb))), sum(m) * (length(tensor(Xb)) / length(m)))
+                @test isapprox(sum((tensor(Xt) .== tensor(Xb))), sum(.!m) * (length(tensor(Xb)) / length(m)))
             else
-                @assert sum((tensor(Xt) .== tensor(Xb))) >= sum(m) * (length(tensor(Xb)) / length(m))
+                @test sum((tensor(Xt) .== tensor(Xb))) >= sum(.!m) * (length(tensor(Xb)) / length(.!m))
             end
 
             #step - doesn't propogate the mask
@@ -76,10 +76,10 @@ using ForwardBackward
             XM = MaskedState(Xa, m, m)
             if !(p isa InterpolatingDiscreteFlow)
                 Xt = Flowfusion.step(p, XM, Xa, 0.1, 0.1)
-                @assert isapprox(sum(tensor(Xt) .!= tensor(XM)), 0) #Because step size is zero
+                @test isapprox(sum(tensor(Xt) .!= tensor(XM)), 0) #Because step size is zero
             else
                 Xt = Flowfusion.step(p, XM, onehot(Xa), 0.1, 0.1)
-                @assert isapprox(sum(tensor(Xt) .!= tensor(XM)), 0) #Because step size is zero
+                @test isapprox(sum(tensor(Xt) .!= tensor(XM)), 0) #Because step size is zero
             end
         end