Distributing function arguments with function compositions. How to compute $(f + g^2)(x) = f(x) + g(x)^2$?

It is sometimes beneficial to first work with functions (in mathematical sense) as symbols and apply to them some pointwise operations. Then, just at the end, convert resulting expression to pure function (in Mathematica sense) and pass some arguments.

This task can be automated using something like following purify function:

ClearAll[purify]
Default[purify, 2] =
    sym : Except[HoldPattern@Symbol[___], _Symbol] /;
        Not@MemberQ[Attributes[sym], Constant];
Default[purify, 3] = {-1};
purify[
    expr_, Shortest[patt_., 1], Shortest[levelspec_., 2],
    opts:OptionsPattern[Replace]
] :=
    Evaluate@Replace[Unevaluated[expr], func:patt :> func[##], levelspec, opts]&

For the problems from question:

purify[f + g^2][x]
(* f[x] + g[x]^2 *)

purify[f - g][x]
(* f[x] - g[x] *)

More complicated example:

testExpr = (f + g) (f - 2 g + 5 h^2) // Expand
purify[%][x, y]
(* f^2 - f g - 2 g^2 + 5 f h^2 + 5 g h^2 *)
(* f[x, y]^2 - f[x, y] g[x, y] - 2 g[x, y]^2 + 5 f[x, y] h[x, y]^2 + 5 g[x, y] h[x, y]^2 *)

Consider only specific symbols as "functions":

purify[testExpr, f | h][x, y]
(* -2 g^2 - g f[x, y] + f[x, y]^2 + 5 g h[x, y]^2 + 5 f[x, y] h[x, y]^2 *)

Default behavior for user functions mixed with built-ins:

purify[π + f + Sin][x]
(* π + f[x] + Sin[x] *)

Consider more complicated expressions as functions:

purify[f + g[a], f | _g, All][x]
(* f[x] + g[a][x] *)

Go into expression heads:

purify[f + g[π + f][a], f, Heads -> True][x]
(* f[x] + g[π + f[x]][a] *)

Held expressions:

purify[Hold[f + f]][x]
(* Hold[f[x] + f[x]] *)

Code Explanation

Per request of @Wjx, here's small explanation of posted code.

purify function replaces, in given expression, sub-expressions (let's call them func), matching given pattern, with func[##], where ## represents arbitrary sequence of arguments. This replacement is evaluated inside a Function, since Function has HoldAll attribute we need to use Evaluate for replacement to be evaluated.

purified = purify[x + f, f]
(* x + f[##1] & *)

Calling this function, with any arguments, results in given expression with arguments passed to selected sub-expressions.

purified[x, y, z]
(* x + f[x, y, z] *)

Since purify replaces sub-expressions, its functionality seemed, to me, closest to Replace built-in function, on which purify is based, so I decided it should use similar interface.

As first argument purify and Replace accept arbitrary expression in which replacements are performed.

As second argument Replace accepts rules, but in purify right hand side of replacement is fixed, so instead of replacement rules purify accepts only a pattern (left hand side of replacement rule).

Second argument of purify is Optional, with assigned Default value being a pattern matching all symbols that don't represent constants. Except[HoldPattern@Symbol[___], _Symbol] pattern is used instead of simple _Symbol to make sure that Symbol["symName"], that can appear in held expressions, will not be matched. Not@MemberQ[Attributes[sym], Constant] Condition tests that symbol does not represent a Constant. Passing Symbol["symName"] to Attributes function would lead to an error.

I guess that using sym_Symbol /; AtomQ@Unevaluated[sym] is more popular to prevent matching of Symbol[...] expressions, but I prefer to handle it in pure pattern matcher, without calling evaluator until it's really necessary.

As third argument both functions accept standard level specification. This argument is also optional in purify, with default being {-1} which means "leaves" of expression tree, since that's only level in which symbols can be found.

With {-1} level specification Except[HoldPattern@Symbol[___], _Symbol] pattern is excessive, since Symbol[...] expressions are not leaves, but I wanted default values to work, in some sense, independent of each other, so if level specification would be changed (for example to All, which would be another reasonable default) default pattern would still work.

Both functions accept also Heads option that specifies whether heads of expressions should be included in replacements.

Since I wanted any rule, given after first argument, to be interpreted as option even when positional optional arguments are not given, their patterns are wrapped with Shortest with appropriate priorities. Similar effect could be achieved by restricting possible values matched by patt and levelspec arguments.


I'm probably missing an important point, but what is wrong with

(f[#] + g[#]^2)&[x]
f[x]+g[x]^2

You can achieve this defining an UpValue for g:

g/:Power[g,2]:=g[#]^2&

Or more generally:

g/:Power[g,n_Integer]:=g[#]^n&

Using Through now works as wanted:

Through[(f + g^2)[x]]
(*Out=f[x]+g[x]^2*)