Working

Author: murrellb

README.md

@@ -7,58 +7,129 @@
 
 
 
-Flowfusion is a Julia package for learning and sampling from conditional diffusion processes across continuous, discrete, and manifold spaces. It provides a unified framework for:
-
-- Learning conditional flows between states
-- Sampling from learned distributions
-- Working with various state types (continuous, discrete, manifold)
-- Handling partial observations and masked states
+Flowfusion 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. It provides a unified framework for:
 
 ## Features
 
-### Multiple State Types
-- Continuous states (Euclidean spaces)
-- Discrete states (categorical variables)
-- Manifold states (probability simplexes, tori, rotations)
-- Masked states for partial conditioning
-
-### Supported Processes
-- Deterministic flows
-- Brownian motion
-- Ornstein-Uhlenbeck
-- Discrete flows (InterpolatingDiscreteFlow, NoisyInterpolatingDiscreteFlow)
-- Manifold-specific processes
-
-### Core Operations
-- `bridge(P, X0, X1, t)`: Sample intermediate states conditioned on start and end states
-- `gen(P, X0, model, steps)`: Generate sequences using a learned model
-- Support for both direct state prediction and tangent coordinate prediction
+- Controllable noise (or fully deterministic for flow matching)
+- Flexible initial $X_0$ distribution
+- Conditioning via masking
+- States: Continuous, discrete, and a wide variety of manifolds supported (via [Manifolds.jl](https://github.com/JuliaManifolds/Manifolds.jl))
+- Compound states supported (e.g. jointly sampling from both continuous and discrete variables)
 
-### Training
-- Loss functions adapted to different state/process types
-- Support for masked training (partial observations)
-- Time scaling for improved training dynamics
-- Integration with Flux.jl for neural network models
+### Basic idea:
+- Generate `X0` and `X1` states from your favorite distribution, and a random `t` between 0 and 1
+- `Xt = bridge(P, X0, X1, t)`: Sample intermediate states conditioned on start and end states
+- Train model to predict how to get to `X1` from `Xt`
+- `gen(P, X0, model, steps)`: Generate sequences using a learned model
 
 ## Examples
 
 The package includes several examples demonstrating different use cases:
 
-- `continuous.jl`: Learning flows between clusters in continuous space
-- `discrete.jl`: Learning categorical transitions
-- `torus.jl`: Learning flows on a torus manifold
-- `probabilitysimplex.jl`: Learning flows between probability distributions
+- `continuous.jl`: Learning a continuous distribution
+- `torus.jl`: Continous distributions on a manifold
+- `discrete.jl`: Discrete distributions with discrete processes
+- `probabilitysimplex.jl`: Discrete distributions with continuous processes via a probability simplex manifold
+- `continuous_masked.jl`: Conditioning on partial observations
+- `masked_tuple.jl`: Jointly sampling from continuous and discrete variables, with conditioning
 
 ## Installation
 
 ```julia
-using Pkg
-Pkg.add("Flowfusion")
+]add https://github.com/MurrellGroup/Flowfusion.jl
 ```
 
-## Quick Start
+## A full example
 
 ```julia
-using Flowfusion, Flux
-#To do.
+using ForwardBackward, Flowfusion, Flux, RandomFeatureMaps, Optimisers, Plots
+
+#Set up a Flux model: X̂1 = model(t,Xt)
+struct FModel{A}
+    layers::A
+end
+Flux.@layer FModel
+function FModel(; embeddim = 128, spacedim = 2, layers = 3)
+    embed_time = Chain(RandomFourierFeatures(1 => embeddim, 1f0), Dense(embeddim => embeddim, swish))
+    embed_state = Chain(RandomFourierFeatures(2 => embeddim, 1f0), Dense(embeddim => embeddim, swish))
+    ffs = [Dense(embeddim => embeddim, swish) for _ in 1:layers]
+    decode = Dense(embeddim => spacedim)
+    layers = (; embed_time, embed_state, ffs, decode)
+    FModel(layers)
+end
+function (f::FModel)(t, Xt)
+    l = f.layers
+    tXt = tensor(Xt)
+    tv = zero(tXt[1:1,:]) .+ expand(t, ndims(tXt))
+    x = l.embed_time(tv) .+ l.embed_state(tXt)
+    for ff in l.ffs
+        x = x .+ ff(x)
+    end
+    tXt .+ l.decode(x) .* (1.05f0 .- expand(t, ndims(tXt))) 
+end
+
+model = FModel(embeddim = 256, layers = 3, spacedim = 2)
+
+#Distributions for training:
+T = Float32
+sampleX0(n_samples) = rand(T, 2, n_samples) .+ 2
+sampleX1(n_samples) = Flowfusion.random_literal_cat(n_samples, sigma = T(0.05))
+n_samples = 400
+
+#The process:
+P = BrownianMotion(0.15f0)
+#P = Deterministic()
+
+#Optimizer:
+eta = 0.001
+opt_state = Flux.setup(AdamW(eta = eta), model)
+
+iters = 4000
+for i in 1:iters
+    #Set up a batch of training pairs, and t:
+    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 & update:
+    l,g = Flux.withgradient(model) do m
+        floss(P, m(t,Xt), X1, scalefloss(P, t))
+    end
+    Flux.update!(opt_state, model, g[1])
+    (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)")
 ```
+
+## Tracking trajectories
+
+```julia
+#Generate samples by stepping from X0
+n_inference_samples = 5000
+X0 = ContinuousState(sampleX0(n_inference_samples))
+paths = Tracker() #<- A tracker to record the trajectory
+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")
+tvec = stack_tracker(paths, :t)
+xttraj = stack_tracker(paths, :xt)
+for i in 1:50:1000
+    plot!(xttraj[1,i,:], xttraj[2,i,:], color = "red", label = i==1 ? "Trajectory" : :none, alpha = 0.4)
+end
+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)")
+```

examples/masked_tuple.jl

@@ -0,0 +1,121 @@
+using Pkg
+Pkg.activate(".")
+using Revise
+Pkg.develop(path="../../ForwardBackward/")
+Pkg.develop(path="../")
+using ForwardBackward, Flowfusion, NNlib, Flux, RandomFeatureMaps, Optimisers, Plots
+
+#Set up a Flux model: X̂1 = model(t,Xt)
+struct FModel{A}
+    layers::A
+end
+Flux.@layer FModel
+function FModel(; embeddim = 128, spacedim = 2, layers = 3)
+    embed_mask = Dense(2 => embeddim) #<- The model should usually know which positions are masked
+    embed_mask_discrete = Dense(1 => embeddim)
+    embed_time = Chain(RandomFourierFeatures(1 => embeddim, 1f0), Dense(embeddim => embeddim, swish))
+    embed_state = Chain(RandomFourierFeatures(2 => embeddim, 1f0), Dense(embeddim => embeddim, swish))
+    embed_discrete_state = Dense(4 => embeddim)
+    ffs = [Dense(embeddim => embeddim, swish) for _ in 1:layers]
+    decode = Dense(embeddim => spacedim)
+    decode_discrete = Dense(embeddim => 4)
+    layers = (; embed_mask, embed_mask_discrete, embed_time, embed_state, embed_discrete_state, ffs, decode, decode_discrete)
+    FModel(layers)
+end
+
+function (f::FModel)(t, Xt)
+    cXt, dXt = Xt
+    l = f.layers
+    tXt = tensor(cXt)
+    tv = zero(tXt[1:1,:]) .+ expand(t, ndims(tXt))
+    x = l.embed_time(tv) .+ l.embed_state(tXt) .+ l.embed_mask(cXt.cmask) .+ l.embed_mask_discrete(reshape(dXt.cmask, 1, :)) .+ l.embed_discrete_state(tensor(dXt)) #<- Mask handling
+    for ff in l.ffs
+        x = x .+ ff(x)
+    end
+    scal = (1.05f0 .- expand(t, ndims(tXt)))
+    (tXt .+ l.decode(x) .* scal), (l.decode_discrete(x) .* scal)
+end
+
+model = FModel(embeddim = 384, layers = 3, spacedim = 2)
+
+#Distributions for training:
+T = Float32
+function sampleX1(n_samples)
+    cstate = Flowfusion.random_literal_cat(n_samples, sigma = T(0.05))
+    dstate = ones(Int64, n_samples)
+    dstate[:] .+= cstate[1,:] .> 0
+    dstate[:] .+= cstate[2,:] .> 0
+    return cstate, dstate
+end
+sampleX0(n_samples) = (rand(T, 2, n_samples) .+ 2), rand(1:4, n_samples)
+n_samples = 500
+
+#The masking distribution - we'll only mask the continuous part of the state
+X1mask(n_samples) = rand(2, n_samples) .< 0.75
+
+#The process:
+P = (BrownianMotion(0.4f0), NoisyInterpolatingDiscreteFlow(0.1))
+
+#Optimizer:
+eta = 0.01
+opt_state = Flux.setup(AdamW(eta = eta, lambda = 0.0001), model)
+
+iters = 6000
+for i in 1:iters
+    #Set up a batch of training pairs, and t, where X1 is a MaskedState:
+    X1cm = X1mask(n_samples)
+    X1dm = rand(n_samples) .< 0.33
+    x1 = sampleX1(n_samples)
+    X1 = (MaskedState(ContinuousState(x1[1]), X1cm, X1cm), MaskedState(onehot(DiscreteState(4, x1[2])), X1dm, X1dm))
+    x0 = sampleX0(n_samples)
+    X0 = (ContinuousState(x0[1]), onehot(DiscreteState(4, x0[2])))
+    t = rand(T, n_samples)
+    #Construct the bridge:
+    Xt = bridge(P, X0, X1, t) #<- Only positions where mask=1 are noised because X1 is a MaskedState
+    #Gradient:
+    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
+end
+
+#Generate unconditional samples:
+n_inference_samples = 5000
+pl = plot(;size = (400,400), legend = :topleft)
+Xcm = rand(2, n_inference_samples) .> 0
+Xdm = rand(n_inference_samples) .> 0
+x0 = sampleX0(n_inference_samples)
+X0 = (MaskedState(ContinuousState(x0[1]), Xcm, Xcm), MaskedState(onehot(DiscreteState(4, x0[2])), Xdm, Xdm))
+paths = Tracker()
+samp = gen(P, X0, model, 0f0:0.005f0:1f0, tracker = paths)
+cstate = tensor(samp[1])
+dstate = tensor(unhot(samp[2].S))
+scatter!(cstate[1,:],cstate[2,:], markerz = dstate, cmap = :brg, msw = 0, ms = 1, alpha = 0.3, label = :none, colorbar = :none)
+savefig("tuple_cat_unconditioned_$P.svg")
+
+#Generate conditioned samples:
+n_inference_samples = 2000
+pl = plot(;size = (400,400), legend = :topleft)
+for dval in [1, 2, 3]
+    Xcm = rand(2, n_inference_samples) .> 0
+    Xdm = rand(n_inference_samples) .< 0
+    x0 = sampleX0(n_inference_samples)
+    x0[2] .= dval
+    X0 = (MaskedState(ContinuousState(x0[1]), Xcm, Xcm), MaskedState(onehot(DiscreteState(4, x0[2])), Xdm, Xdm))
+    paths = Tracker()
+    samp = gen(P, X0, model, 0f0:0.005f0:1f0, tracker = paths)
+    cstate = tensor(samp[1])
+    dstate = tensor(unhot(samp[2].S))
+    scatter!(cstate[1,:],cstate[2,:], msw = 0, ms = 1, alpha = 0.3, label = "D = $dval")
+end
+pl
+savefig("tuple_cat_conditioned_$P.svg")

examples/probabilitysimplex.jl

@@ -55,10 +55,10 @@ for i in 1:iters
     #Construct the bridge:
     Xt = bridge(P, X0, X1, t)
     #Get the Xt->X1 tangent coordinates:
-    ξ = tangent_guide(Xt, X1)
+    ξ = Guide(Xt, X1)
     #Gradient:
     l,g = Flux.withgradient(model) do m
-        tcloss(P, m(t,tensor(Xt)), ξ, scalefloss(P, t))
+        floss(P, m(t,tensor(Xt)), ξ, scalefloss(P, t))
     end
     #Update:
     Flux.update!(opt_state, model, g[1])

examples/torus.jl

@@ -54,10 +54,11 @@ for i in 1:iters
     #Construct the bridge:
     Xt = bridge(P, X0, X1, t)
     #Compute the tangent coordinates:
-    ξ = tangent_guide(Xt, X1)
+    ξ = Guide(Xt, X1)
     #Gradient
     l,g = Flux.withgradient(model) do m
-        tcloss(P, m(t,tensor(Xt)), ξ, scalefloss(P, t))
+        #tcloss(P, m(t,tensor(Xt)), ξ, scalefloss(P, t)) #GOING TO HAVE TO ADD GUIDE HERE, AND CHANGE IT TO FLOSS
+        floss(P, m(t,tensor(Xt)), ξ, scalefloss(P, t)) #GOING TO HAVE TO ADD GUIDE HERE, AND CHANGE IT TO FLOSS
     end
     #Update
     Flux.update!(opt_state, model, g[1])

src/Flowfusion.jl

@@ -2,6 +2,7 @@
 Need to test/do:
 Urgent:
 - Test tuples!
+- Test Manifolds with masking (especially tangent_guide and apply_tangent etc) 
 - Masking (cmask) on all state types for bridge and gen
 - Masking (lmask) on all state types for both losses
 - tensor on masked states
@@ -42,6 +43,7 @@ export
     Tracker,
     stack_tracker,
     onehot,
+    unhot,
     FProcess, 
     floss,
     tcloss,