Add enhance documentation for neural parameter estimation

Author: kiante-fernandez

docs/src/assets/lca_amorized_recovery.png

@@ -1,6 +1,6 @@
 # Neural Parameter Estimation
 
-Neural parameter estimation provides a likelihood-free approach to parameter recovery, especially useful for models with computationally intractable likelihoods. This method is based on training neural networks to learn the mapping from data to parameters. See the review paper by Zammit-Mangion et al. (2025) for more details. Once trained, these networks can perform inference rapidly across multiple datasets, making them particularly valuable for models like the Leaky Competing Accumulator (LCA; Usher & McClelland, 2001).
+Neural parameter estimation uses neural networks to perform parameter estimation by learning the mapping between simulated data and model parameters (for a detailed review, see Zammit-Mangion et al., 2024). Neural parameter estimation is particularly useful for models with computationally intractable likelihoods, such as the Leaky Competing Accumulator (LCA; Usher & McClelland, 2001). Once trained, neural networks can be saved and reused to perform inference rapidly across multiple datasets, or used for computationally intensive parameter recovery simulations to understand the quality of parameter estimates under ideal conditions.
 
 Below, we demonstrate how to estimate parameters of the LCA model using the [NeuralEstimators.jl](https://github.com/msainsburydale/NeuralEstimators) package.
 
@@ -24,7 +24,12 @@ Random.seed!(123)
 
 ## Define Parameter Sampling
 
-Unlike traditional Bayesian inference methods, simulation-based inference approaches require us to define a prior sampling function specifically to generate synthetic training data. While traditional methods like MCMC also sample from the prior, those samples are used directly during inference rather than to create a separate training dataset. In SBI, we use the prior to sample a wide range of parameters and simulate corresponding data, which we then use to train a model (e.g., a neural network) to approximate the posterior. We will use the following function to sample a range of parameters for training:
+Unlike traditional Bayesian inference methods, neural parameter estimation requires us to define two functions so that the neural network can learn the mapping between simulated data and parameters. One function samples parameters from a prior distribution, and the other generates simulated data based on a sampled parameter vector. While traditional methods like MCMC also sample from the prior, those samples are used directly during inference rather than to create a separate training dataset. 
+
+![](assets/npe_example.png)
+*Schematic of neural parameter estimation. Once trained, the neural network provides a direct mapping from observed data (Z) to parameter estimates (θ̂), enabling rapid inference without the computational burden of traditional methods.*
+
+In neural parameter estimation, we use the prior to sample a wide range of parameters and simulate corresponding data, which we then use to train a model (e.g., a neural network) to approximate a point estimate or the posterior. We use the following function to sample a range of parameters for training:
 
 ```julia 
 # Function to sample parameters from priors
@@ -63,7 +68,8 @@ function simulate(θ, n_trials_per_param)
         choices, rts = rand(model, n_trials_per_param)
 
         # Return as a transpose matrix where each column is a trial 
-        return [choices rts]'
+        return Float32.([choices rts]')
+
     end
 
     return simulated_data
@@ -77,11 +83,11 @@ For LCA parameter recovery, we use a DeepSet architecture which respects the per
 ```julia 
 # Create neural network architecture for parameter recovery
 function create_neural_estimator(;
-    ν_bounds = (0.1, 4.0),
-    α_bounds = (0.5, 3.5),
-    β_bounds = (0.0, 0.5),
-    λ_bounds = (0.0, 0.5),
-    τ_bounds = (0.1, 0.5)
+    ν_bounds = (0.1, 6.0),
+    α_bounds = (0.3, 4.5),
+    β_bounds = (0.0, 0.8),
+    λ_bounds = (0.0, 0.8),
+    τ_bounds = (0.100, 2.0)
 )
     # Unpack defined parameter Bounds
     ν_min, ν_max = ν_bounds              # Drift rates
@@ -132,6 +138,8 @@ function create_neural_estimator(;
 end
 ```
 
+The result of our constructed neural network is a point estimator that corresponds to a Bayes estimator, which is a functional of the posterior distribution. Under the specified loss function, this point estimate corresponds to the posterior mean. For details on the theoretical foundations of neural Bayes estimators, see Sainsbury-Dale et al. (2024).
+
 ## Training the Neural Estimator
 
 Neural estimators, like all deep learning methods, require a training phase during which they learn the mapping from data to parameters. Here, we train the estimator by simulating data on the fly: the sampler provides new parameter vectors from the prior, and the simulator generates corresponding data conditional on those parameters. Since we use online training and the network never sees the same simulated dataset twice, overfitting is less likely. For more details on training, see the API for arguments [here](https://msainsburydale.github.io/NeuralEstimators.jl/dev/API/core/#Training).
@@ -163,9 +171,9 @@ We can assess the performance of our trained estimator on held-out test data:
 
 ```julia 
 # Generate test data
-n_test = 100
+n_test = 500
 θ_test = sample(n_test)
-Z_test = simulate(θ_test, 100)
+Z_test = simulate(θ_test, 500)
 
 # Assess the estimator
 parameter_names = ["ν1", "ν2", "α", "β", "λ", "τ"]
@@ -197,17 +205,34 @@ p_plots = []
 
 for param in params
     param_data = filter(row -> row.parameter == param, df)
+    
+    # Calculate correlation coefficient
+    truth = param_data.truth
+    estimate = param_data.estimate
+    correlation = cor(truth, estimate)
+    
+    # Create plot
     p = scatter(
-        param_data.truth, 
-        param_data.estimate,
+        truth, 
+        estimate,
         xlabel="Ground Truth",
         ylabel="Estimated",
         title=param,
         legend=false
     )
-    plot!(p, [minimum(param_data.truth), maximum(param_data.truth)], 
-          [minimum(param_data.truth), maximum(param_data.truth)], 
+    
+    # Add diagonal reference line
+    plot!(p, [minimum(truth), maximum(truth)], 
+          [minimum(truth), maximum(truth)], 
           line=:dash, color=:black)
+    
+    # Get current axis limits after plot is created
+    x_min, x_max = xlims(p)
+    y_min, y_max = ylims(p)
+    
+    # Position text at the top-left corner of the plot
+    annotate!(p, x_min + 0.1, y_max, text("R = $(round(correlation, digits=3))", :left, 10))
+    
     push!(p_plots, p)
 end
 
@@ -228,7 +253,6 @@ Once trained, the estimator can instantly recover parameters from new data via a
 β = 0.2
 λ = 0.1
 τ = 0.3
-σ = 1.0
 
 # Create model and generate data
 true_model = LCA(; ν, α, β, λ, τ)
@@ -247,7 +271,7 @@ println("Recovered parameters: ", recovered_params)
 
 ## Notes on Performance
 
-Neural estimators are particularly effective for models with computationally intractable likelihoods like the LCA model. However, certain parameters (particularly β and λ) can be difficult to recover, even with advanced neural network architectures. This is a  property of the LCA model rather than a limitation of the estimation technique. 
+Neural estimators are particularly effective for models with computationally intractable likelihoods like the LCA model. However, certain parameters (particularly β and λ) can be difficult to recover, even with advanced neural network architectures. This is a property of the LCA model rather than a limitation of the estimation technique. 
 
 Additional details can be found in the [NeuralEstimators.jl documentation](https://github.com/msainsburydale/NeuralEstimators).
 
@@ -296,7 +320,7 @@ function simulate(θ, n_trials_per_param)
         choices, rts = rand(model, n_trials_per_param)
 
         # Return as a transpose matrix where each column is a trial 
-        return [choices rts]'
+        return Float32.([choices rts]')
 
     end
 
@@ -305,11 +329,11 @@ end
 
 # Create neural network architecture for parameter recovery
 function create_neural_estimator(;
-    ν_bounds = (0.1, 4.0),
-    α_bounds = (0.5, 3.5),
-    β_bounds = (0.0, 0.5),
-    λ_bounds = (0.0, 0.5),
-    τ_bounds = (0.1, 0.5)
+    ν_bounds = (0.1, 6.0),
+    α_bounds = (0.3, 4.5),
+    β_bounds = (0.0, 0.8),
+    λ_bounds = (0.0, 0.8),
+    τ_bounds = (0.100, 2.0)
 )
     # Unpack defined parameter Bounds
     ν_min, ν_max = ν_bounds              # Drift rates
@@ -379,9 +403,9 @@ trained_estimator = train(
 )
 
 # Generate test data
-n_test = 100
+n_test = 500
 θ_test = sample(n_test)
-Z_test = simulate(θ_test, 100)
+Z_test = simulate(θ_test, 500)
 
 # Assess the estimator
 parameter_names = ["ν1", "ν2", "α", "β", "λ", "τ"]
@@ -407,17 +431,34 @@ p_plots = []
 
 for param in params
     param_data = filter(row -> row.parameter == param, df)
+    
+    # Calculate correlation coefficient
+    truth = param_data.truth
+    estimate = param_data.estimate
+    correlation = cor(truth, estimate)
+    
+    # Create plot
     p = scatter(
-        param_data.truth, 
-        param_data.estimate,
+        truth, 
+        estimate,
         xlabel="Ground Truth",
         ylabel="Estimated",
         title=param,
         legend=false
     )
-    plot!(p, [minimum(param_data.truth), maximum(param_data.truth)], 
-          [minimum(param_data.truth), maximum(param_data.truth)], 
+    
+    # Add diagonal reference line
+    plot!(p, [minimum(truth), maximum(truth)], 
+          [minimum(truth), maximum(truth)], 
           line=:dash, color=:black)
+    
+    # Get current axis limits after plot is created
+    x_min, x_max = xlims(p)
+    y_min, y_max = ylims(p)
+    
+    # Position text at the top-left corner of the plot
+    annotate!(p, x_min + 0.1, y_max, text("R = $(round(correlation, digits=3))", :left, 10))
+    
     push!(p_plots, p)
 end
 
@@ -431,7 +472,6 @@ display(p_combined)
 β = 0.2
 λ = 0.1
 τ = 0.3
-σ = 1.0
 
 # Create model and generate data
 true_model = LCA(; ν, α, β, λ, τ)