Constructing functions with variable number of output arguments

In Mathematica, you can construct a function f to have different definitions based on its input arguments. For example, f[x_] := ... and f[x_, y_] := .... You can also get more specific and define it for specific heads as f[x_someHead] := ... or inputs matching arbitrary patterns as f[x_?somePatternQ] := ... (of course, one must be careful to keep the order from specific to general). These definitions are the DownValues of the function.

Now, I'd like to construct a function that has different definitions based on the number of output variables/symbols requested, but for the same set of inputs. Consider the following dummy example in pseudo Mathematica code

f[mu_, sigma2_] := Switch[OutputArgs[],
0, Plot[PDF[NormalDistribution[mu, sigma2], x], ...]
1, RandomReal[NormalDistribution[mu, sigma2], 1000]
2, {Mean@#, Variance@#}&@ RandomReal[...]
]


Here OutputArgs[] is a functionality that I'd like, which when called from inside a function, tells you how many outputs have been requested (loosely reminiscent of OptionValue which knows which function it's called from). With this, I can use the function call simply as f[0,1] to plot the normal distribution (perhaps inside a Manipulate to play with the parameters), then when I'm satisfied, I can use the same function call, but with an output argument as a = f[0,1] to get a random sample, and with two outputs to get an empirical estimate of the mean and variance.

I realize that this is a bit at odds with the classical notion of a function having a well defined output based on the input, rather than a well defined output based on the number of outputs for a given input.

My interest is purely academic and arose from working on porting some MATLAB code to Mathematica. Some might have recognized that right away, and this is a feature that that I do find useful — a sentiment also shared by a few others that I've spoken to here. I wouldn't write code this way and I know that it is also possible to achieve the same by other means such as using flags as f[0, 1, "Plot"] or options as f[0, 1, Options -> "RandomSample"], etc. for the above example. However, I'm interested in exploring if there are ways to imitate this behaviour in Mathematica.

-
What about making OuptutArgs an option? Something link f[mu_, sigma2_,OptionsPattern[]] := Switch[ OptionsValue[OutputArgs], 0, Plot[PDF[NormalDistribution[mu, sigma2], x], ...], 1, RandomReal[NormalDistribution[mu, sigma2], 1000], 2, {Mean@#, Variance@#}&@ RandomReal[...] ] – Rolf Mertig Jan 19 '12 at 22:59
Sorry, you do not want any options (and I seem not to be able to delete my above comment after 5 minutes. weird). But then how does it work in MATLAB? I.e., how is the "requesting" done? – Rolf Mertig Jan 19 '12 at 23:13
What's wrong with doing it in the way you indicate above? Obviously you'd have to define OutputArgs[], but that's unavoidable. Mind you, I'm not sure I see an advantage over options or optional arguments. – Daniel Lichtblau Jan 19 '12 at 23:17
@DanielLichtblau I guess the question then is, how does one define OutputArgs[]? I don't know how to "know" from within the function how many outputs are requested. As for advantages, I don't think there are any, and if I write a function, I probably would do it with options. This thought occurred when I tried to port a script to mma, and I was puzzling over how one would go about doing it, if so desired. – R. M. Jan 19 '12 at 23:24
Okay, now I'm confused. Some code somewhere needs to have a notion of what it wants. Or perhaps the user knows, and OutputArgs[] could request Input[]. If it is the caller, I don't see much point to NOT passing that information. But I guess it could use a variable "outputSize", and maybe OutputArgs[] would query that variable. There are various tactics for avoiding use of a global variable. But again, all this seems like a lot of trouble for the desired result. – Daniel Lichtblau Jan 19 '12 at 23:34

As you have mentioned there are other more standard ways to provide some variaty in your return value so I don't think that this is something that you should do. If you provide code using something like it you should have good reasons for doing so. Just being used to it because matlab has it probably is just not convincing enough.

That said, it can be done along these lines:

ClearAll[f, ff]
f /: Set[arg_Symbol, f[x_]] := Block[{nargout = 1}, arg = ff[x]];
f /: Set[args_List, f[x_]] := Block[{nargout = Length[args]}, args = ff[x]];
f[arg_] /; FreeQ[Stack[], Set] := Block[{nargout = 1}, ff[arg]];
ff[x_] := (Table[x, {nargout}]);


you can now use it like this:

In[101]:= f[5]
Out[101]= {5}
In[102]:= a=f[5]
Out[102]= {5}
In[104]:= {a,b}=f[5]
Out[104]= {5,5}
In[105]:= {a,b,c}=f[5]
Out[105]= {5,5,5}


This is not tested well and could have many pitfalls. I see it rather as a proof of concept and probably something for your own convenience. So use with care, especially if you are uncertain about all the details (as I am: I can't remember I ever had the need to use Stack).

Brett already found one pitfall. Actually that is a case where one could be argueing what the expected return value should be. If we fall back to the one argument case for everything that is not an explicit Set it is easy to solve, but I can see other problems already :-):

f[arg_] /; Not[MatchQ[Stack[],{___,Set,f,___}]]:=Block[{nargout=1},ff[arg]];

-
+1 for FreeQ[Stack[], Set], I was trying to work out how to test for the case without Set... – acl Jan 20 '12 at 0:05
@acl FreeQ[Stack[], Set] will get tripped up by something like foo = g[f[x]]. – Brett Champion Jan 20 '12 at 2:11
@Brett maybe StackComplete could help then. But I admit I was already very much impressed with Stack as that was where I got (ahem) stuck. – acl Jan 20 '12 at 2:17
@Brett I think I have a fix for that particular one which I will add, but as I can already see the next cases which break it. As I said I would not recommend to use this other than for learning purposes and probably for convenience when using own code... – Albert Retey Jan 20 '12 at 18:23
Without knowing what pitfall you foresee, you probably want something tighter than ___ as the last element of the pattern. Consider a=f[b=g[f[x]]]. (I would say this is unlikely, but I saw a question using Goto yesterday...) – Brett Champion Jan 20 '12 at 18:45

I can offer an alternative approach based on introspection and run-time code generation. It is not as elegant as one by @Albert, to be sure, but may still represent some interest. Here is the code (I made an assumption that your function is a symbol which has DownValues, but this can be generalized to other ...Values:

ClearAll[withNArgOut];
SetAttributes[withNArgOut, HoldAll];
withNArgOut[code_] :=
Hold[code] //.
{
HoldPattern[Set[vars_List, fcall : f_Symbol[___]]] /;
FreeQ[Unevaluated[fcall], $nArgAout] && ! FreeQ[DownValues[f],$nArgAout] :>
Set[vars, Block[{$nArgAout = Length[Unevaluated[vars]]}, fcall]], HoldPattern[Set[var_, fcall : _[___]]] /; !MatchQ[Unevaluated[var], _List] && Hold[var] =!= Hold[$nArgAout] && FreeQ[Unevaluated[fcall], $nArgAout] :> Set[{var}, fcall] };  I intentionally did not release the last Hold, so that we can see what is generated. Here is how you could use it: ClearAll[f]; f[x_] := Table[x^i, {i,$nArgAout}]

In[14]:= {#,ReleaseHold@#}&@withNArgOut[Module[{a,b},b=1;a=f[x]]]
Out[14]= {Hold[Module[{a,b},b=1;{a}=Block[{$nArgAout=Length[Unevaluated[{a}]]},f[x]]]],{x}} In[15]:= {#,ReleaseHold@#}&@withNArgOut[Module[{a,b,c}, c=1;{a,b}=f[x]]] Out[15]= {Hold[Module[{a,b,c},c=1; {a,b}=Block[{$nArgAout=Length[Unevaluated[{a,b}]]},f[x]]]],{x,x^2}}

In[16]:= withNArgOut[Module[{a,b,c}, c=1;{a,b}=g[f[x]]]]
Out[16]= Hold[Module[{a,b,c},c=1;{a,b}=g[f[x]]]]


You can always add a call to releaseHold in the definition of withNArgOut. The usage is that you wrap the withNArgOut around the code where you want this to take effect, and the extra Block[{\$nArgAout= ...},...] wrappers will be generated for that code (where appropriate), just prior to execution.

-