From 999bd8e2bfdd108a63216fca0c7de908e07a6c70 Mon Sep 17 00:00:00 2001
From: Samuel Brand <sam055@mac.com>
Date: Fri, 16 Feb 2024 18:06:13 +0000
Subject: [PATCH] refactor newton steps to be simpler

---
 EpiAware/src/utilities.jl                             |  6 ++----
 .../test/predictive_checking/fast_approx_for_r.jl     | 11 +++--------
 2 files changed, 5 insertions(+), 12 deletions(-)

diff --git a/EpiAware/src/utilities.jl b/EpiAware/src/utilities.jl
index 7f1ddb5c4..56191c38c 100644
--- a/EpiAware/src/utilities.jl
+++ b/EpiAware/src/utilities.jl
@@ -136,14 +136,12 @@ function fast_R_to_r_approx(R₀, w::Vector{T}; newton_steps = 1) where {T<:Abst
     r_approx = (R₀ - 1) / (R₀ * mean_gen_time)
     # Newton's method
     for _ = 1:newton_steps
-        r_approx +=
-            (1 - R₀ * neg_MGF(r_approx, w)) * neg_MGF(r_approx, w) /
-            dneg_MGF_dr(r_approx, w)
+        r_approx -= (R₀ * neg_MGF(r_approx, w) - 1) / (R₀ * dneg_MGF_dr(r_approx, w))
     end
     return r_approx
 end
 
-function fast_R_to_r_approx(R₀, epimodel::AbstractEpiModel; newton_steps = 3)
+function fast_R_to_r_approx(R₀, epimodel::AbstractEpiModel; newton_steps = 4)
     fast_R_to_r_approx(R₀, epimodel.data.gen_int; newton_steps = newton_steps)
 end
 
diff --git a/EpiAware/test/predictive_checking/fast_approx_for_r.jl b/EpiAware/test/predictive_checking/fast_approx_for_r.jl
index 6a2249e5e..7235ee44c 100644
--- a/EpiAware/test/predictive_checking/fast_approx_for_r.jl
+++ b/EpiAware/test/predictive_checking/fast_approx_for_r.jl
@@ -8,7 +8,7 @@ G(r) = \sum_{i=1}^{\infty} w_i e^{-r i}.
 ```
 and
 ```math
-f(r, \mathcal{R}_0) = {1 \over G(r)} - \mathcal{R}_0.
+f(r, \mathcal{R}_0) = \mathcal{R}_0 G(r) - 1.
 ```
 then the connection between `R₀` and `r` is given by
 ```math
@@ -24,14 +24,9 @@ r_0 = { \mathcal{R}_0 - 1 \over  \mathcal{R}_0 \langle W \rangle }.
 ```
 where $\langle W \rangle$ is the mean of the generation interval.
 
-Note that
+To rapidly improve the estimate for `r` we use Newton steps given by
 ```math
-\partial_r G(r, \mathcal{R}_0) = - {G'(r) \over G(r)^2} }
-```
-
-Therefore, each Newton step is given by
-```math
-r_{n+1} = r_n + {(1 - \mathcal{R}_0 G(r)) G(r) \over G'(r)}.
+r_{n+1} = r_n - {\mathcal{R}_0 G(r) - 1\over \mathcal{R}_0 G'(r)}.
 ```
 
 =#