Using delayed evaluation on module variables called through an association

Question

I came up with a construct that works well enough for my current use case but has some flaws (at least I suspect it does). Say we have a function that calculates some stuff that takes a while to process (nonlinear constrained optimisation in my case) and has some primary output (e.g. best fit parameters) but some secondary outputs that might be interesting in themselves or for diagnostics as well (e.g. a plot of the residual errors). Returning an Association containing the different outputs and thus making them accessible seems a good approach to me.

Now let's say some of those outputs are computationally expensive and are not generally needed when all goes well and the primary result looks fine right out of the box. In that case we would like to avoid the cost of calculating the secondary output(s), but being able to generate them at will would be nice nonetheless. The construct I used is:

ClearAll@func;

func[a_] := Module[{primaryOutput, secondaryOutput},
primaryOutput = (Pause@5; 2 a);
secondaryOutput := (Pause@2; 0.5 a);
Association["1" -> primaryOutput, "2" :>  secondaryOutput]]

Please note the use of RuleDelayed and SetDelayed insides Module and let's test what happens when we call that thing.

(test = func[1]) // AbsoluteTiming
(* {5.00136, <|"1" -> 2, "2" :> secondaryOutput$28323|>} *) 

test["1"] // AbsoluteTiming
(* {3.*10^-6, 2} *)

test["2"] // AbsoluteTiming
(* {2.00089, 0.5} *)

As we see from the timings, all looks good in principe. The general call of the function takes approx. five seconds (which is the duration of our primary "calculation") and no time is spend on the optional calculation of the secondary output. Accessing the result (with the "1"-key) is instantaneous. Accessing the secondary output (with the "2"-key) takes approx. two seconds (instead of seven) thus unnecessary recalculation of the primary output is avoided.

What I suspect might be an issue here is the secondaryOutput$28323 returned inside the resulting association. A new module variable is create at each call of func and I wonder both about proper scoping and garbage collection of these module variables. So the question is wether there is some merit to my suspicion and if so, how to fix this.

Bonus question: Since calculating the secondary output is expensive one might like to cache it using memoisation. Does this worsen the possible issue with those module variables?

Answer 1

I'd not rely on Module variables after Module's evaluation is done. They have Temporary attribute and it can surprise you later.

One way around is to inject Own/DownValues:

func[a_] := Module[{primaryOutput, secondaryOutput}
  ,
  primaryOutput = (Pause@5; 2 a);
  secondaryOutput := (Pause@2; 0.5 a);

  Association[
     "1" -> primaryOutput, "2" :> secondaryOutput
  ] /. OwnValues[secondaryOutput]
]

now

(test = func[1]) // AbsoluteTiming

gives

{4.9994, <|"1" -> 2, "2" :> (Pause[2]; 0.5 1)|>}

---

In order to cache the second value you can do something like:

ClearAll@func;
SetAttributes[func, HoldFirst];

func[var_, a_] := Module[{primaryOutput, secondaryOutput}
  ,
  primaryOutput = (Pause@5; 2 a);
  secondaryOutput := var["2"] = (Pause@2; 0.5 a);
  var = Association["1" -> primaryOutput, "2" :> secondaryOutput] /. 
    OwnValues[secondaryOutput]
  ]


(func[test, 1]) // AbsoluteTiming

{4.99903, <|"1" -> 2, "2" :> (test["2"] = (Pause[2]; 0.5 1))|>}

test["2"] // AbsoluteTiming

{1.99965, 0.5}

test

<|"1" -> 2, "2" -> 0.5|>