Feature/Multi-fidelity PINNs

[brunopjacob]

Implementation of multi-fidelity PINNs via stacking, as proposed in Howard et al., 2023.

Modifications are based on the creation of a MultiFidelityMLP class that inherits from Block, as well as modifications to the loss that require the creation of a custom loss for these types of networks (MultiFidelityLoss).

p.s.: this is still a draft. It will have companion tutorials, unit test updates and TLC on the docstrings.

[brunopjacob]

I think this is ready for review! Please let me know what y'all think, especially regarding the changes in the trainer (I'm open to create one in PyTorchLightning as well).

[drgona]

lets make this PR to develop branch first before we integrate it into master.

[drgona]

this PR looks good to me we can integrate

[brunopjacob]

this PR looks good to me we can integrate

I like your idea of merging it to develop instead of master, Jan. We can definitely do that.
I also just realized I forgot to add to the unit test list the KANBlock and the MultiFidelityMLP blocks. I'll do that next.

[RBirmiwal]

Can we clean these training scripts up, e.g.

# Number of epochs
epochs_mf = 50000

# Neuromancer trainer 
trainer_mf = Trainer(
    problem_mf.to(device),
    torch.utils.data.DataLoader(train_data, batch_size=t_train.shape[0],
                                collate_fn=train_data.collate_fn, shuffle=False),
    optimizer=optimizer_mf,
    epochs=epochs_mf,
    epoch_verbose=5000,
    train_metric='train_loss',
    dev_metric='train_loss',
    eval_metric="train_loss",
    warmup=epochs_mf,
    device=device,
    multi_fidelity=True
)

# Train PINN
best_model_mf = trainer_mf.train()

# Load best trained model
problem_mf.load_state_dict(best_model_mf)

Also for this section, I'd like to leave the dynamics simulators in-house as much as possible, so use ODE or PSL's pendulum class and our integrator examples to generate forward data (instead of solveIVP)

# Compare with numerical solution of ODE solver
# Define the ODE system for numerical solution
def pendulum_ode(t, y):
    s1, s2 = y
    ds1_dt = s2
    ds2_dt = -b/m * s2 - g/L * np.sin(s1)
    return [ds1_dt, ds2_dt]

# Initial conditions for numerical solution
y0 = [s1_0, s2_0]

# Time points for numerical solution
t_span = (0, T)
t_eval = np.linspace(0, T, N_cp)

# Solve the ODE numerically
sol = solve_ivp(pendulum_ode, t_span, y0, t_eval=t_eval)

[RBirmiwal]

Really good job here. Code-wise only couple minor revisions in the pendulum notebook (see my comment there)
In terms of descriptions and introductory context in the notebooks, I would like to see a bit more (for both Stacked PINNs and the KAN notebook). You describe the equations well, but I'd like to see more about the (new) model being introduced.
I'd like to make this a standard for any new notebook, namely it should be upfront and answer the questions:

• What is the problem we are trying to solve?
• How are we doing that? Where are we using neural networks?
• What neural network architecture/algorithm are we using to solve? Most importantly it should answer the so what question -- why should I use this architecture/algorithm vs. ones previously demonstrates in the examples? Give a 1-2 paragraph introduction to motivate the architecuture/algorithm within the context of the example folder it resides in. So for the case of PINNs, we can assume the arbitrary end-user as a basic understanding of PINNs by the time they reach the Stacked PINN / KAN notebook. (I know the references are listed, so these paragraphs can be drawn from there, but I don't want the references alone to do the explaining).

Hopefully this isn't too much of a lift to add the educational and "so what" contexts. Once those are done will be approved. Not trying to be overly picky. It's awesome we are incorporating state-of-the-art methods but as a result, require additional context.

Thanks