|
| 1 | +using Pkg |
| 2 | +Pkg.activate(".") |
| 3 | +using Revise |
| 4 | + |
| 5 | +#using FileIO, Images |
| 6 | +#img = FileIO.load("fflogo.png") |
| 7 | +#arr = [a.r for a in img] .< 0.5 |
| 8 | +#sampinds = (x -> (x[2], 265-x[1])).(CartesianIndices(arr)[arr]) |
| 9 | +#CSV.write("logoinds.csv", DataFrame(sampinds)) |
| 10 | + |
| 11 | +using CSV, DataFrames |
| 12 | +df = CSV.read("logoinds.csv", DataFrame) |
| 13 | +sampinds = [Tuple(df[i,:]) for i in 1:size(df,1)] |
| 14 | +flowinds = [s ./ 200 for s in sampinds if s[2] > 0] |
| 15 | +fusioninds = [s ./ 200 for s in sampinds if s[2] <= 0] |
| 16 | + |
| 17 | +Pkg.develop(path="../../ForwardBackward/") |
| 18 | +Pkg.develop(path="../") |
| 19 | +using ForwardBackward, Flowfusion, NNlib, Flux, RandomFeatureMaps, Optimisers, Plots, Manifolds |
| 20 | + |
| 21 | +#Set up a Flux model: X̂1 = model(t,Xt) |
| 22 | +struct FModel{A} |
| 23 | + layers::A |
| 24 | +end |
| 25 | +Flux.@layer FModel |
| 26 | +function FModel(; embeddim = 128, layers = 3) |
| 27 | + embed_time = Chain(RandomFourierFeatures(1 => embeddim, 1f0), Dense(embeddim => embeddim)) |
| 28 | + embed_state = Chain(RandomFourierFeatures(2 => embeddim, 3f0), Dense(embeddim => embeddim)) |
| 29 | + embed_angle = Chain(RandomFourierFeatures(2 => embeddim, 1f0), Dense(embeddim => embeddim)) |
| 30 | + ffs = [Dense(embeddim => embeddim, swish) for _ in 1:layers] |
| 31 | + decode = Dense(embeddim => 2) |
| 32 | + decode_angle = Dense(embeddim => 1) |
| 33 | + layers = (;embed_time, embed_state, embed_angle, ffs, decode, decode_angle) |
| 34 | + FModel(layers) |
| 35 | +end |
| 36 | + |
| 37 | +function (f::FModel)(t, Xt) |
| 38 | + tXt, aXt = tensor.(Xt) |
| 39 | + l = f.layers |
| 40 | + aenc = vcat(sin.(aXt), cos.(aXt)) |
| 41 | + tv = zero(tXt[1:1,:]) .+ expand(t, ndims(tXt)) |
| 42 | + x = l.embed_time(tv) .+ l.embed_state(tXt) .+ l.embed_angle(aenc) |
| 43 | + for ff in l.ffs |
| 44 | + x = x .+ ff(x) |
| 45 | + end |
| 46 | + scal = (1.05f0 .- expand(t, ndims(tXt))) |
| 47 | + (tXt .+ l.decode(x) .* scal), (l.decode_angle(x) .* scal) |
| 48 | +end |
| 49 | + |
| 50 | +T = Float32 |
| 51 | +n_samples = 1000 |
| 52 | +M = Torus(1) |
| 53 | +ManifoldState(M, Array{Float32}.(rand(M, n_samples))) |
| 54 | +sampleX0(n_samples) = ContinuousState(T.(stack(rand(flowinds, n_samples))) .+ rand(T, 2, n_samples) .* 0.01f0), ManifoldState(M, fill([0.6f0], n_samples)) |
| 55 | +sampleX1(n_samples) = ContinuousState(T.(stack(rand(fusioninds, n_samples))) .+ rand(T, 2, n_samples) .* 0.01f0), ManifoldState(M, fill([-2.54159f0], n_samples)) |
| 56 | + |
| 57 | +model = FModel(embeddim = 384, layers = 5) |
| 58 | +n_samples = 500 |
| 59 | + |
| 60 | +#The process: |
| 61 | +P = (BrownianMotion(0.05f0), ManifoldProcess(0.1f0)) |
| 62 | + |
| 63 | +#Optimizer: |
| 64 | +eta = 0.001 |
| 65 | +opt_state = Flux.setup(AdamW(eta = eta, lambda = 0.001), model) |
| 66 | + |
| 67 | +iters = 6000 |
| 68 | +for i in 1:iters |
| 69 | + #Set up a batch of training pairs, and t, where X1 is a MaskedState: |
| 70 | + X0 = sampleX0(n_samples) |
| 71 | + X1 = sampleX1(n_samples) |
| 72 | + t = rand(T, n_samples) |
| 73 | + #Construct the bridge: |
| 74 | + Xt = bridge(P, X0, X1, t) |
| 75 | + ξ = Guide(Xt[2], X1[2]) |
| 76 | + #Gradient: |
| 77 | + l,g = Flux.withgradient(model) do m |
| 78 | + hat = m(t,Xt) |
| 79 | + floss(P[1], hat[1], X1[1], scalefloss(P[1], t)) + floss(P[2], hat[2], ξ, scalefloss(P[2], t)) |
| 80 | + end |
| 81 | + #Update: |
| 82 | + Flux.update!(opt_state, model, g[1]) |
| 83 | + #Logging, and lr cooldown: |
| 84 | + if i % 10 == 0 |
| 85 | + if i > iters - 2000 |
| 86 | + eta *= 0.975 |
| 87 | + Optimisers.adjust!(opt_state, eta) |
| 88 | + end |
| 89 | + println("i: $i; Loss: $l; eta: $eta") |
| 90 | + end |
| 91 | +end |
| 92 | + |
| 93 | +function smodel(t, Xt) |
| 94 | + hat = model(t,Xt) |
| 95 | + return hat[1], Guide(hat[2]) |
| 96 | +end |
| 97 | + |
| 98 | +#Generate unconditional samples: |
| 99 | +n_inference_samples = 5000 |
| 100 | +X0 = sampleX0(n_inference_samples) |
| 101 | +paths = Tracker() |
| 102 | +samp = gen(P, X0, smodel, 0f0:0.005f0:1f0, tracker = paths) |
| 103 | + |
| 104 | +cstate = tensor(samp[1]) |
| 105 | +astate = tensor(samp[2]) |
| 106 | +zcstate = tensor(X0[1]) |
| 107 | +zastate = tensor(X0[2]) |
| 108 | + |
| 109 | +#scatter(zcstate[1,:], zcstate[2,:], msw = 0, ms = 1.5, markerz = zastate[1,:], cmap = :hsv) |
| 110 | +#scatter!(cstate[1,:], cstate[2,:], msw = 0, ms = 1.5, markerz = astate[1,:], cmap = :hsv) |
| 111 | +scatter(zcstate[1,:], zcstate[2,:], msw = 0, ms = 1.5, markerz = zastate[1,:], cmap = :hsv, label = :none, xlim = (-0.5, 5.5), ylim = (-1.5, 1.5)) |
| 112 | +scatter!(cstate[1,:], cstate[2,:], msw = 0, ms = 1.5, markerz = astate[1,:], cmap = :hsv, label = :none, xlim = (-0.5, 5.5), ylim = (-1.5, 1.5)) |
| 113 | +scatter!([-100,-100],[-100,-100], markerz = [-pi,pi], label = :none, colorbar = :none, axis=([], false)) |
| 114 | + |
| 115 | +postraj = stack([tensor(p[1]) for p in paths.xt]) |
| 116 | +angtraj = stack([tensor(p[2]) for p in paths.xt]) |
| 117 | + |
| 118 | +anim = @animate for i ∈ vcat([1 for i in 1:20], 1:size(postraj, 3), [size(postraj, 3) for i in 1:20]) |
| 119 | + scatter(postraj[1,:,i], postraj[2,:,i], msw = 0, ms = 1, markerz = angtraj[1,:,i], cmap = :hsv, label = :none, xlim = (-0.0, 5.2), ylim = (-1.3, 1.3), size = (400, 200)) |
| 120 | + scatter!([-100,-100],[-100,-100], markerz = [-pi,pi], label = :none, colorbar = :none, axis=([], false)) |
| 121 | +end |
| 122 | +gif(anim, "logo_$(P).mp4", fps = 30) |
| 123 | +gif(anim, "logo_$(P).gif", fps = 30) |
| 124 | + |
0 commit comments