From 9ed5e9e86728fa3cf2228ab683d7fb04d47e8f7b Mon Sep 17 00:00:00 2001
From: Olivier Labayle <olabayle@gmail.com>
Date: Tue, 3 Jan 2023 09:56:00 +0100
Subject: [PATCH 1/3] add mach_cache to tmlecache

---
 src/cache.jl | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/src/cache.jl b/src/cache.jl
index bebec879..42759a81 100644
--- a/src/cache.jl
+++ b/src/cache.jl
@@ -11,13 +11,14 @@ mutable struct TMLECache
     η_spec::NuisanceSpec
     dataset
     η
-    function TMLECache(Ψ, η_spec, dataset)
+    mach_cache
+    function TMLECache(Ψ, η_spec, dataset, mach_cache)
         dataset = Dict(
             :source => dataset,
             :no_missing => nomissing(dataset, allcolumns(Ψ))
         )
         η = NuisanceParameters(nothing, nothing, nothing, nothing)
-        new(Ψ, η_spec, dataset, η)
+        new(Ψ, η_spec, dataset, η, mach_cache)
     end
 end
 
@@ -78,14 +79,15 @@ Main entrypoint to run the TMLE procedure.
 - dataset: A tabular dataset respecting the Table.jl interface
 - verbosity: The logging level
 - threshold: To avoid small values of Ĝ to cause the "clever covariate" to explode
+- mach_cache: Whether underlying MLJ.machines will cache data or not
 """
 function tmle(Ψ::Parameter, η_spec::NuisanceSpec, dataset; verbosity=1, threshold=1e-8, mach_cache=false)
-    cache = TMLECache(Ψ, η_spec, dataset)
-    return tmle!(cache; verbosity=verbosity, threshold=threshold, mach_cache=mach_cache)
+    cache = TMLECache(Ψ, η_spec, dataset, mach_cache)
+    return tmle!(cache; verbosity=verbosity, threshold=threshold)
 end
 
-function tmle!(cache; verbosity=1, threshold=1e-8, mach_cache=false)
-    Ψ, η_spec, dataset, η = cache.Ψ, cache.η_spec, cache.dataset, cache.η
+function tmle!(cache; verbosity=1, threshold=1e-8)
+    Ψ, η_spec, dataset, η, mach_cache = cache.Ψ, cache.η_spec, cache.dataset, cache.η, cache.mach_cache
     # Initial fit
     verbosity >= 1 && @info "Fitting the nuisance parameters..."
     TMLE.fit!(η, η_spec, Ψ, dataset, verbosity=verbosity, mach_cache=mach_cache)
@@ -124,7 +126,7 @@ end
 
 Runs the TMLE procedure for the new nuisance parameters specification η_spec while potentially reusing cached nuisance parameters.
 """
-function tmle!(cache::TMLECache, η_spec::NuisanceSpec; verbosity=1, threshold=1e-8)
+function tmle!(cache::TMLECache, η_spec::NuisanceSpec; verbosity=1, threshold=1e-8, mach_cache=false)
     update!(cache, η_spec)
     tmle!(cache, verbosity=verbosity, threshold=threshold)
 end

From dc2a9873120b7f1c49c260f036ca1ad03231439c Mon Sep 17 00:00:00 2001
From: Olivier Labayle <olabayle@gmail.com>
Date: Tue, 3 Jan 2023 09:57:23 +0100
Subject: [PATCH 2/3] update tmlecache fields types

---
 src/cache.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/cache.jl b/src/cache.jl
index 42759a81..14bc3aca 100644
--- a/src/cache.jl
+++ b/src/cache.jl
@@ -10,8 +10,8 @@ mutable struct TMLECache
     Ψ::Parameter
     η_spec::NuisanceSpec
     dataset
-    η
-    mach_cache
+    η::NuisanceParameters
+    mach_cache::Bool
     function TMLECache(Ψ, η_spec, dataset, mach_cache)
         dataset = Dict(
             :source => dataset,

From 9fd9188bb2422e6aefc0aabe055c93320ab73b29 Mon Sep 17 00:00:00 2001
From: Olivier Labayle <olabayle@gmail.com>
Date: Tue, 3 Jan 2023 10:10:15 +0100
Subject: [PATCH 3/3] add a method a update docs

---
 docs/src/user_guide.md |  4 ++--
 src/cache.jl           | 11 +++++++++++
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/docs/src/user_guide.md b/docs/src/user_guide.md
index 9bf35166..1b1c5a95 100644
--- a/docs/src/user_guide.md
+++ b/docs/src/user_guide.md
@@ -323,6 +323,6 @@ nothing # hide
 
 Since we have only updated $G$'s specification, only this model is fitted again.
 
-### Scenario N
+### General behaviour
 
-Feel free to play around with the cache and to report any non consistent behaviour.
+Any change to either the `Parameter` of interest or the `NuisanceSpec` structures will trigger an update of the cache.
diff --git a/src/cache.jl b/src/cache.jl
index 14bc3aca..ca4a7a3c 100644
--- a/src/cache.jl
+++ b/src/cache.jl
@@ -140,4 +140,15 @@ while potentially reusing cached nuisance parameters.
 function tmle!(cache::TMLECache, Ψ::Parameter, η_spec::NuisanceSpec; verbosity=1, threshold=1e-8)
     update!(cache, Ψ, η_spec)
     tmle!(cache, verbosity=verbosity, threshold=threshold)
+end
+
+"""
+    tmle!(cache::TMLECache, η_spec::NuisanceSpec, Ψ::Parameter; verbosity=1, threshold=1e-8)
+
+Runs the TMLE procedure for the new parameter Ψ and the new nuisance parameters specification η_spec 
+while potentially reusing cached nuisance parameters.
+"""
+function tmle!(cache::TMLECache, η_spec::NuisanceSpec, Ψ::Parameter; verbosity=1, threshold=1e-8)
+    update!(cache, Ψ, η_spec)
+    tmle!(cache, verbosity=verbosity, threshold=threshold)
 end
\ No newline at end of file