Updated Documentation

Author: JoshuaBillson

examples/example.jl

@@ -1,6 +1,6 @@
 push!(LOAD_PATH, "../src")
 
-using MixtureDensityNetworks, Distributions, CairoMakie
+using MixtureDensityNetworks, Distributions, CairoMakie, Logging, TerminalLoggers
 
 const n_samples = 1000
 const epochs = 1000
@@ -18,7 +18,9 @@ function main()
     model = MDN(epochs=epochs, mixtures=mixtures, layers=layers)
 
     # Fit Model
-    lc = fit!(model, X, Y)
+    lc = with_logger(TerminalLogger()) do 
+        fit!(model, X, Y)
+    end
 
     # Plot Learning Curve
     fig, _, _ = lines(1:epochs, lc, axis=(;xlabel="Epochs", ylabel="Loss"))

src/MixtureDensityNetworks.jl

@@ -13,5 +13,4 @@ include("interface.jl")
 
 export likelihood_loss, MDN, fit!, predict, predict_mean, predict_mode
 
-
-end
+end

src/interface.jl

@@ -41,8 +41,8 @@ Fit the model to the data given by X and Y.
 - `X`: A dxn matrix where d is the number of features and n is the number of samples.
 - `Y`: A 1xn matrix where n is the number of samples.
 """
-function fit!(model::MDN, X::Matrix{<:Number}, Y::Matrix{<:Number})
-    fit!(model, Float32.(X), Float32.(Y))
+function fit!(model::MDN, X::Matrix{<:Real}, Y::Matrix{<:Real})
+    fit!(model, Float64.(X), Float64.(Y))
 end
 
 function fit!(model::MDN, X::Matrix{Float64}, Y::Matrix{Float64})
@@ -101,8 +101,8 @@ Predict the full conditional distribution P(Y|X).
 # Returns
 Returns a vector of Distributions.MixtureModel objects representing the conditional distribution for each sample.
 """
-function predict(model::MDN, X::Matrix{<:Number})
-    predict(model, Float32.(X))
+function predict(model::MDN, X::Matrix{<:Real})
+    predict(model, Float64.(X))
 end
 
 function predict(model::MDN, X::Matrix{Float64})
@@ -128,8 +128,8 @@ Predict the mean of the conditional distribution P(Y|X).
 # Returns
 Returns a vector of real numbers representing the mean of the conditional distribution P(Y|X) for each sample.
 """
-function predict_mean(model::MDN, X::Matrix{<:Number})
-    predict_mean(model, Float32.(X))
+function predict_mean(model::MDN, X::Matrix{<:Real})
+    predict_mean(model, Float64.(X))
 end
 
 function predict_mean(model::MDN, X::Matrix{Float64})
@@ -149,8 +149,8 @@ Predict the mean of the Gaussian with the largest prior in the conditional distr
 # Returns
 Returns a vector of real numbers representing the mean of the gaussian with the largest prior for each sample.
 """
-function predict_mode(model::MDN, X::Matrix{<:Number})
-    predict_mode(model, Float32.(X))
+function predict_mode(model::MDN, X::Matrix{<:Real})
+    predict_mode(model, Float64.(X))
 end
 
 function predict_mode(model::MDN, X::Matrix{Float64})

src/model.jl

@@ -35,17 +35,14 @@ function MixtureDensityNetwork(input::Int, layers::Vector{Int}, mixtures::Int)
     return MixtureDensityNetwork(hidden_layer, μ, Σ, π)
 end
 
-"MixtureModel forward pass."
+function (m::MixtureDensityNetwork)(X::AbstractMatrix{<:Real})
+    Float64.(X) |> m
+end
+
 function (m::MixtureDensityNetwork)(X::AbstractMatrix{Float64})
-    # Forward Pass
     h = m.hidden(X)
     μ = m.μ(h)
     Σ = m.Σ(h)
     π = m.π(h)
     μ, Σ, π
-end
-
-"MixtureModel forward pass."
-function (m::MixtureDensityNetwork)(X::AbstractMatrix{<:Number})
-    Float32.(X) |> m
 end