How can I find a functional square root in Mathematica?

Question

A functional square root or 'half iterate' of a function $g$ is another function $f$ such that $g=f\circ f$.

For example, the functional square root of $x/(2 - x)$ is $x/(\sqrt{2} + x - x\sqrt{2})$ and we can verify:

f[x_] := x/(Sqrt[2] + x - x*Sqrt[2])
FullSimplify[f[f[x]]]
(* returns: x/(2 - x) *)

According to that article, there is a one approach to finding a functional square root that involves solving Schröder's equation:

$$\forall x \quad \Psi(h(x))=s \Psi(x)$$ ... where $s$ is a like an eigenvalue, and then the functional square root $h_{1/2}(x)$ is given by: $$ h_{1 / 2}(x)=\Psi^{-1}\left(s^{1 / 2} \Psi(x)\right) $$

I'd like to know how to go about finding exact or approximate functional square roots in Mathematica. Suppose I want to find an $f$ for $g(x)=\cos(4 \pi x)$. I try to frame it as a differential equation but this is unsuccessful. I am not even sure what the initial conditions should be either:

D[f[f[x]], x]
(* Derivative[1][f][x] Derivative[1][f][f[x]] *)

D[Cos[4 Pi x], x]
(* -4 Pi Sin[4 Pi x] *)

NDSolve[{-4 Pi Sin[4 Pi x] == Derivative[1][f][x] Derivative[1][f][f[x]], f[0] == 1}, f, {x, 0, 1}]
(* Power::infy: Infinite expression 1/0 encountered. *)

---

Apparently there are power series methods that can solve problems like this, and these things called Carleman Matrices. I looked around in the docs and found CarlemanLinearize but I can't tell if this is related to this problem.

@J.M. has a function to construct a Carleman matrix here so I will take a look at that.

---

I've tried to follow along with this answer here but my coefficients end up as complex numbers and the plot doesn't really look like iterating twice will remotely resemble a cosine:

x0 = 0; n = 30;(*expansion point and order*)
cosCM = N[CarlemanMatrix[Cos[4 Pi x], {x, x0, n}], 30];
shalfCoeffs = MatrixPower[Transpose[cosCM], 1/2, UnitVector[n + 1, 2]];
shalf[x_] = Fold[(#1 x + #2) &, 0, Reverse[shalfCoeffs]];
ReImPlot[shalf[x], {x, 0, 1}]

Answer 1

I was not able to get a half iterate for $\cos(...)$ anything, and from reading around a bit it appears that half iterates of $\cos$ might be impossible either due to convergence or the evenness of $\cos$'s series expansion terms. However, I was able to get a half-iterate for a small part of the domain of $\sin(4 \pi x)$ through fixed point iteration on the series, though it becomes inaccurate quite quickly:

(* Try to find a half iterate of Sin[4 \[Pi] x] *)
halfit[x_] = Nest[(Sin[4 \[Pi]*Normal[InverseSeries[Series[#, {x, 0, 6}]]]] + #)/2 &, x, 8];

Plot[{halfit[halfit[x]], Sin[4 \[Pi] x]}, {x, -\[Pi]/2, \[Pi]/2}, 
 PlotRange -> {-1, 1}, 
 PlotStyle -> {Directive[Thick, Red], Directive[Blue]}]

I was able to get an approximation of the half-sine by a different method using a Newton series, though this does not work for a higher frequency sine like $\sin(4 \pi x)$ and produces a very noisy function. The resulting $\mathrm{hsin}(\mathrm{hsin}(x))\approx\sin(x)$ is not too bad an approximation judging by the plot:

newtonfhalf[f_, x_, mmax_] := 
 Sum[Binomial[1/2, m] Sum[
    Binomial[m, k] (-1)^(m - k) Nest[f, x, k], {k, 0, m}], {m, 0, mmax}]

nth = Function[{x}, newtonfhalf[Sin[#] &, x, 40]];
nthh2 = nth[nth[x]];
Plot[{Sin[x], nthh2}, {x, -4, 4}, 
 PlotStyle -> {Directive[Thick, Blue], Directive[Red]}]

---

I've had some luck with a neural network approach to the problem. I've found it's possible to train a network in a non-standard way to find an approximate half-iterate. Assume a network $N$ of 1 input and 1 output node with arbitrary layers in between and that we're trying to find a half-iterate for the function $\mathrm{target}(x)$:

1. Run $N$ forward on a random input $x_i$ and generate output $y_i$
2. Run $N$ forward again using $y_i$ as input, generating output $y_i'$
3. The loss is $(\mathrm{target}(x_i) - y_i')^2$. Back-propagate and update $N$ and return to step 1.

The resulting network is hopefully trained such that $N(N(x)) \approx \mathrm{target}(x)$.

I wasn't sure how to approach this in Mathematica but this is my first time using PyTorch ever, so what follows may be a bit basic:

import torch
import torch.nn as nn
import torch.optim as optim

from math import pi, sin, cos
import random
import csv


def targetfn(x):
    return sin(x)


class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        self.lin = nn.Linear(1, 20)
        self.lmid1 = nn.Tanh()
        self.lmid2 = nn.Linear(20, 20)
        self.lmid3 = nn.Tanh()
        self.lout = nn.Linear(20, 1)

    def forward(self, w):
        w = self.lin(w)
        w = self.lmid1(w)
        w = self.lmid2(w)
        w = self.lmid3(w)
        return self.lout(w)


def train():
    net = Net()
    print(net)

    optimizer = optim.SGD(net.parameters(), lr=0.01)
    criterion = nn.MSELoss()

    # init random
    net.zero_grad()
    outinit = net(torch.randn(1))
    outinit.backward(torch.randn(1))

    for i in range(100000):
        x = random.uniform(-2 * pi, 2 * pi)
        target = torch.tensor([targetfn(x)])
        y1 = net(torch.tensor([x]))
        net.zero_grad()
        optimizer.zero_grad()
        y2 = net(y1)
        loss = criterion(y2, target)
        loss.backward()
        optimizer.step()

    return net


def main():
    net = train()

    with open("hfn.csv", 'w', newline='') as csvfile:
        csvwriter = csv.writer(csvfile, delimiter=',')
        n = 2000
        xmin = -2 * pi
        xmax = 2 * pi
        step = (xmax - xmin) / n
        x = xmin
        for i in range(n):
            csvwriter.writerow([x, net(torch.tensor([x])).item()])
            x += step


if __name__ == '__main__':
    main()

... and plotting in Mathematica:

data = Import["hfn.csv"];
intp = Interpolation[data];
Plot[{Sin[t], intp[intp[t]]}, {t, -2 \[Pi], 2 \[Pi]}, 
 PlotRange -> {-1.3, 1.3}, 
 PlotStyle -> {Directive[Thick, Blue], Directive[Thin, Red]}, 
 PlotTheme -> "Scientific"]

This is looking good for $\sin(x)$. What about $\cos(x)$? I changed targetfn in the python code above and at least I got something that looked close to a cosine wave:

---

It turns out it's possible to do this in Mathematica using NetNestOperator. You train the nested network, then extract the single network which is the approximation of the half-iterate.

net = NetChain[{LinearLayer[20, "Input" -> "Real"], Tanh, 20, Tanh, 
    SummationLayer[]}];
net = NetNestOperator[net, 2];
net = NetInitialize[net];
net = NetTrain[net, # -> Sin[#] & /@ RandomReal[{-Pi, Pi}, 100000], 
   TargetDevice -> "GPU", BatchSize -> 512];
halfnet = NetExtract[net, "Net"];
Plot[{Sin[x], halfnet@halfnet@x}, {x, -2 Pi, 2 Pi}]

Here's another idea I had, where instead of generating training data up front and using memory, I have a RandomArrayLayer produce the data for me. The ConstantArray is there to get it to batch, and this time we're learning a half-iterate for a polynomial using ramps:

f = x |-> -2 x^4 + x^3/7 - 2 x^2 + x + 1;
innerFunc = 
  NetNestOperator[{LinearLayer[8, "Input" -> "Real"], Ramp, 8, Ramp, 
    8, Ramp, 8, Ramp, 8, Ramp, 8, SummationLayer[]}, 2];
net = NetGraph[{RandomArrayLayer[UniformDistribution[{-1, 1}]], 
    innerFunc, FunctionLayer[f], Plus, 
    ThreadingLayer[(#2 - #1)^2 &]}, {NetPort["Input"] -> 4, 1 -> 4, 
    1 -> 3, 4 -> 2, 3 -> 5, 2 -> 5}];
net = NetInitialize[net];
net = NetTrain[net, ConstantArray[0 -> 0, 256], TargetDevice -> "GPU",
   MaxTrainingRounds -> 100000, BatchSize -> 256]
halfnet = NetExtract[NetExtract[net, 2], "Net"];
Plot[{f[x], halfnet[halfnet[x]]}, {x, -1, 1}]