Updating readme and example

Author: murrellb

README.md

@@ -7,8 +7,7 @@
 
 Flowfusion.jl is a Julia package for training and sampling from diffusion and flow matching models (and some things in between), across continuous, discrete, and manifold spaces, all in a single unified framework and interface.
 
-![Image](https://github.com/user-attachments/assets/d739c07e-f9e9-4aef-932e-c36cae182391)
-![Image](https://github.com/user-attachments/assets/f2754ba5-b798-4db9-8ce6-a0324b89a534)
+![Image](https://github.com/user-attachments/assets/ff7f25e6-441d-4840-ac9c-a849e7b57aa7)
 
 The animated logo shows samples from a model trained to jointly transport a 2D point and an angular hue between two distributions. For the 2D point, the left side uses "Flow matching" with deterministic trajectories, and the right uses a Brownian bridge. For both sides, the angular hue is diffused via an angular Brownian bridge. The hue endpoints are antipodal, and you can see both paths, in opposite angular directions, are sampled.
 

examples/logo_example.jl

@@ -1,23 +1,23 @@
 using Pkg
 Pkg.activate(".")
 using Revise
+Pkg.develop(path="../../ForwardBackward/")
+Pkg.develop(path="../")
+using ForwardBackward, Flowfusion, NNlib, Flux, RandomFeatureMaps, Optimisers, Plots, Manifolds
+using CSV, DataFrames
 
+#Creating the logoinds.csv file:
 #using FileIO, Images
 #img = FileIO.load("fflogo.png")
 #arr = [a.r for a in img] .< 0.5
 #sampinds = (x -> (x[2], 265-x[1])).(CartesianIndices(arr)[arr])
 #CSV.write("logoinds.csv", DataFrame(sampinds))
 
-using CSV, DataFrames
-df = CSV.read("logoinds.csv", DataFrame)
+df = CSV.read("logoinds.csv", DataFrame) #https://github.com/user-attachments/files/18864465/logoinds.csv
 sampinds = [Tuple(df[i,:]) for i in 1:size(df,1)]
 flowinds = [s ./ 200 for s in sampinds if s[2] > 0]
 fusioninds = [s ./ 200 for s in sampinds if s[2] <= 0]
 
-Pkg.develop(path="../../ForwardBackward/")
-Pkg.develop(path="../")
-using ForwardBackward, Flowfusion, NNlib, Flux, RandomFeatureMaps, Optimisers, Plots, Manifolds
-
 #Set up a Flux model: X̂1 = model(t,Xt)
 struct FModel{A}
     layers::A
@@ -54,18 +54,18 @@ ManifoldState(M, Array{Float32}.(rand(M, n_samples)))
 sampleX0(n_samples) = ContinuousState(T.(stack(rand(flowinds,   n_samples))) .+ rand(T, 2, n_samples) .* 0.01f0), ManifoldState(M, fill([0.6f0], n_samples))
 sampleX1(n_samples) = ContinuousState(T.(stack(rand(fusioninds, n_samples))) .+ rand(T, 2, n_samples) .* 0.01f0), ManifoldState(M, fill([-2.54159f0], n_samples))
 
-model = FModel(embeddim = 512, layers = 5)
+model = FModel(embeddim = 384, layers = 4)
 n_samples = 500
 
 #The process:
-P = (BrownianMotion(0.05f0), ManifoldProcess(0.1f0))
-#P = (Deterministic(), ManifoldProcess(0.1f0))
+P = (FProcess(BrownianMotion(0.1f0), t -> 1-(1-t)^2), ManifoldProcess(0.1f0))
+#P = (FProcess(Deterministic(), t -> 1-(1-t)^2), ManifoldProcess(0.1f0))
 
 #Optimizer:
 eta = 0.001
 opt_state = Flux.setup(AdamW(eta = eta, lambda = 0.001), model)
 
-iters = 10000
+iters = 8000
 for i in 1:iters
     #Set up a batch of training pairs, and t, where X1 is a MaskedState: 
     X0 = sampleX0(n_samples)
@@ -102,17 +102,8 @@ X0 = sampleX0(n_inference_samples)
 paths = Tracker()
 samp = gen(P, X0, smodel, 0f0:0.005f0:1f0, tracker = paths)
 
-cstate = tensor(samp[1])
-astate = tensor(samp[2])
-zcstate = tensor(X0[1])
-zastate = tensor(X0[2])
-
-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))
-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))
-scatter!([-100,-100],[-100,-100], markerz = [-pi,pi], label = :none, colorbar = :none, axis=([], false))
-
-postraj = stack([tensor(p[1]) for p in paths.xt])
-angtraj = stack([tensor(p[2]) for p in paths.xt])
+postraj = stack(vcat([tensor(X0[1])], [tensor(p[1]) for p in paths.xt], [tensor(samp[1])]))
+angtraj = stack(vcat([tensor(X0[2])], [tensor(p[2]) for p in paths.xt], [tensor(samp[2])]))
 
 anim = @animate for i ∈ vcat([1 for i in 1:20], 1:size(postraj, 3), [size(postraj, 3) for i in 1:20])
     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))
@@ -121,3 +112,21 @@ end
 gif(anim, "logo_$(P).mp4", fps = 30)
 gif(anim, "logo_$(P).gif", fps = 30)
 
+#=
+#To create a side-by-side animation (useful for comparing schedules, noise, etc), run the above once, then store the trajectories, then run it again, and then merge, with an offset:
+#diffpostraj = copy(postraj)
+#diffangtraj = copy(angtraj)
+
+flowpostraj = copy(postraj)
+flowangtraj = copy(angtraj)
+
+mergedpostraj = hcat(diffpostraj .+ reshape([6, 0], 2, 1, 1), flowpostraj)
+mergedangtraj = hcat(diffangtraj, flowangtraj)
+
+anim = @animate for i ∈ vcat([1 for i in 1:20], 1:size(mergedpostraj, 3), [size(mergedpostraj, 3) for i in 1:20])
+    scatter(mergedpostraj[1,:,i], mergedpostraj[2,:,i], msw = 0, ms = 1, markerz = mergedangtraj[1,:,i], cmap = :hsv, label = :none, xlim = (-0.0, 11.2), ylim = (-1.3, 1.3), size = (800, 200))
+    scatter!([-100,-100],[-100,-100], markerz = [-pi,pi], label = :none, colorbar = :none, axis=([], false))
+end
+gif(anim, "mergedlogo_$(P).mp4", fps = 30)
+gif(anim, "mergedlogo_$(P).gif", fps = 30)
+=#