Input scaling, optimiser choise and other changes

[dquigley533]

I came across your Medium post on this - thanks for making the code available. I'm currently learning about PINNs myself and found your tutorial to be a more useful introduction than many things I've read.

In my further reading and experimentation I made some tweaks which you might consider improvements, hence the PR in case you'd like to incorporate them.

1. Input/Output scaling. The input time is scaled into [0,1], as is the output. This dramatically reduces the number of parameters needed to get a good reproduction of the PDE solution. I'm using a few dozen nodes in a single hidden layer.

2. Switch optimiser. Your example was training on the whole data, rather than a batch, so using AdamW was unnecessary since you have the exact (rather than estimated) gradients. Using LBFGS, the training converges in a handful of steps, noting that I switched activation functions from ReLU to GELU to get smooth gradients.

3. With those changes the problem fixed by the L2 regularisation doesn't seem to occur, but I've left that in anyway.

4. In the examples which use physics_loss I've weighted this massively higher than the MSE loss on the training points to avoid overfitting to the noise on those points.

The whole thing now runs in a handful of seconds on a CPU, which did wonders for my confidence that PINNs would be tractable for problems in higher numbers of dimensions.

Obviously feel free to ignore - just felt compelled to share what I learned by trying to push your example to its limits.