Optim refactoring

[dario-coscia]

Description

This PR fixes #618.

Checklist

• Code follows the project’s Code Style Guidelines
• Tests have been added or updated
• Documentation has been updated if necessary
• Pull request is linked to an open issue

[github-actions]

Code Coverage Summary

Filename                                                                  Stmts    Miss  Cover    Missing
----------------------------------------------------------------------  -------  ------  -------  -------------------------------------------------------------------------------------------------------
__init__.py                                                                   7       0  100.00%
graph.py                                                                    114      11  90.35%   99-100, 112, 124, 126, 142, 144, 166, 169, 182, 271
label_tensor.py                                                             251      28  88.84%   81, 121, 144-148, 165, 177, 182, 188-193, 273, 280, 332, 334, 348, 444-447, 490, 537, 629, 664-673, 710
operator.py                                                                  72       2  97.22%   269, 465
trainer.py                                                                   75       5  93.33%   195-204, 293, 314, 318
type_checker.py                                                              22       0  100.00%
utils.py                                                                     73      11  84.93%   59, 75, 141, 178, 181, 184, 220-223, 268
adaptive_function/__init__.py                                                 3       0  100.00%
adaptive_function/adaptive_function.py                                       55       0  100.00%
adaptive_function/adaptive_function_interface.py                             51       6  88.24%   98, 141, 148-151
callback/__init__.py                                                          5       0  100.00%
callback/normalizer_data_callback.py                                         68       1  98.53%   141
callback/optimizer_callback.py                                               23       0  100.00%
callback/processing_callback.py                                              49       5  89.80%   42-43, 73, 168, 171
callback/refinement/__init__.py                                               3       0  100.00%
callback/refinement/r3_refinement.py                                         28       1  96.43%   88
callback/refinement/refinement_interface.py                                  50       5  90.00%   32, 59, 67, 72, 78
condition/__init__.py                                                         7       0  100.00%
condition/condition.py                                                       19       1  94.74%   141
condition/condition_interface.py                                             37       4  89.19%   32, 76, 95, 125
condition/data_condition.py                                                  26       1  96.15%   78
condition/domain_equation_condition.py                                       19       0  100.00%
condition/input_equation_condition.py                                        43       1  97.67%   157
condition/input_target_condition.py                                          44       1  97.73%   172
data/__init__.py                                                              3       0  100.00%
data/data_module.py                                                         201      22  89.05%   41-52, 132, 172, 193, 232, 313-317, 323-327, 399, 466, 546, 637, 639
data/dataset.py                                                              82       7  91.46%   42, 123-126, 256, 293
domain/__init__.py                                                           10       0  100.00%
domain/cartesian.py                                                         112      10  91.07%   37, 47, 75-76, 92, 97, 103, 246, 256, 264
domain/difference_domain.py                                                  25       2  92.00%   54, 87
domain/domain_interface.py                                                   20       5  75.00%   37-41
domain/ellipsoid.py                                                         104      24  76.92%   52, 56, 127, 250-257, 269-282, 286-287, 290, 295
domain/exclusion_domain.py                                                   28       1  96.43%   86
domain/intersection_domain.py                                                28       1  96.43%   85
domain/operation_interface.py                                                26       1  96.15%   88
domain/simplex.py                                                            72      14  80.56%   62, 207-225, 246-247, 251, 256
domain/union_domain.py                                                       25       1  96.00%   43
equation/__init__.py                                                          4       0  100.00%
equation/equation.py                                                         15       1  93.33%   56
equation/equation_factory.py                                                101       1  99.01%   181
equation/equation_interface.py                                                4       0  100.00%
equation/system_equation.py                                                  20       0  100.00%
loss/__init__.py                                                              9       0  100.00%
loss/linear_weighting.py                                                     14       0  100.00%
loss/loss_interface.py                                                       17       2  88.24%   45, 51
loss/lp_loss.py                                                              15       0  100.00%
loss/ntk_weighting.py                                                        18       0  100.00%
loss/power_loss.py                                                           15       0  100.00%
loss/scalar_weighting.py                                                     16       0  100.00%
loss/self_adaptive_weighting.py                                              12       0  100.00%
loss/weighting_interface.py                                                  29       3  89.66%   35, 41-42
model/__init__.py                                                            14       0  100.00%
model/average_neural_operator.py                                             31       2  93.55%   73, 82
model/deeponet.py                                                            93      13  86.02%   187-190, 209, 240, 283, 293, 303, 313, 323, 333, 488, 498
model/equivariant_graph_neural_operator.py                                   51       1  98.04%   217
model/feed_forward.py                                                        89      11  87.64%   58, 195, 200, 278-292
model/fourier_neural_operator.py                                             78      10  87.18%   96-100, 110, 155-159, 218, 220, 242, 342
model/graph_neural_operator.py                                               40       2  95.00%   58, 60
model/kernel_neural_operator.py                                              34       6  82.35%   83-84, 103-104, 123-124
model/low_rank_neural_operator.py                                            27       2  92.59%   89, 98
model/multi_feed_forward.py                                                  12       5  58.33%   25-31
model/pirate_network.py                                                      27       1  96.30%   118
model/sindy.py                                                               21       0  100.00%
model/spline.py                                                              88       5  94.32%   133, 153, 161, 199, 366
model/spline_surface.py                                                      48       7  85.42%   133, 146-151, 191
model/block/__init__.py                                                      13       0  100.00%
model/block/average_neural_operator_block.py                                 12       0  100.00%
model/block/convolution.py                                                   64      13  79.69%   77, 81, 85, 91, 97, 111, 114, 151, 161, 171, 181, 191, 201
model/block/convolution_2d.py                                               146      27  81.51%   155, 162, 282, 314, 379-433, 456
model/block/embedding.py                                                     48       7  85.42%   93, 143-146, 155, 168
model/block/fourier_block.py                                                 31       0  100.00%
model/block/gno_block.py                                                     22       4  81.82%   73-77, 87
model/block/integral.py                                                      18       4  77.78%   22-25, 71
model/block/low_rank_block.py                                                24       0  100.00%
model/block/orthogonal.py                                                    37       0  100.00%
model/block/pirate_network_block.py                                          25       1  96.00%   89
model/block/pod_block.py                                                     75      10  86.67%   56-59, 71, 84, 114, 151-156, 191, 216
model/block/rbf_block.py                                                    179      25  86.03%   18, 42, 53, 64, 75, 86, 97, 223, 280, 282, 298, 301, 329, 335, 363, 367, 511-524
model/block/residual.py                                                      46       0  100.00%
model/block/spectral.py                                                      83       4  95.18%   132, 140, 262, 270
model/block/stride.py                                                        28       7  75.00%   55, 58, 61, 67, 72-74
model/block/utils_convolution.py                                             22       3  86.36%   58-60
model/block/message_passing/__init__.py                                       6       0  100.00%
model/block/message_passing/deep_tensor_network_block.py                     21       0  100.00%
model/block/message_passing/en_equivariant_network_block.py                  47       1  97.87%   164
model/block/message_passing/equivariant_graph_neural_operator_block.py       36       0  100.00%
model/block/message_passing/interaction_network_block.py                     23       0  100.00%
model/block/message_passing/radial_field_network_block.py                    20       0  100.00%
optim/__init__.py                                                             5       0  100.00%
optim/optimizer_interface.py                                                  7       0  100.00%
optim/scheduler_interface.py                                                  7       0  100.00%
optim/torch_optimizer.py                                                     14       0  100.00%
optim/torch_scheduler.py                                                     19       2  89.47%   5-6
problem/__init__.py                                                           6       0  100.00%
problem/abstract_problem.py                                                 117      12  89.74%   39-40, 59-70, 149, 161, 179, 253, 257, 286
problem/inverse_problem.py                                                   22       0  100.00%
problem/parametric_problem.py                                                 8       1  87.50%   29
problem/spatial_problem.py                                                    8       0  100.00%
problem/time_dependent_problem.py                                             8       0  100.00%
problem/zoo/__init__.py                                                       8       0  100.00%
problem/zoo/advection.py                                                     23       4  82.61%   20, 73-75
problem/zoo/allen_cahn.py                                                    23       3  86.96%   20-22
problem/zoo/diffusion_reaction.py                                            32       5  84.38%   103-113
problem/zoo/helmholtz.py                                                     22       4  81.82%   59, 78-82
problem/zoo/inverse_poisson_2d_square.py                                     48       3  93.75%   44-50
problem/zoo/poisson_2d_square.py                                             16       3  81.25%   61-66
problem/zoo/supervised_problem.py                                            11       0  100.00%
solver/__init__.py                                                            6       0  100.00%
solver/garom.py                                                             107       2  98.13%   129-130
solver/solver.py                                                            188      10  94.68%   195, 218, 290, 293-294, 353, 435, 518, 559, 565
solver/ensemble_solver/__init__.py                                            4       0  100.00%
solver/ensemble_solver/ensemble_pinn.py                                      23       1  95.65%   104
solver/ensemble_solver/ensemble_solver_interface.py                          27       0  100.00%
solver/ensemble_solver/ensemble_supervised.py                                 9       0  100.00%
solver/physics_informed_solver/__init__.py                                    8       0  100.00%
solver/physics_informed_solver/causal_pinn.py                                47       3  93.62%   157, 166-167
solver/physics_informed_solver/competitive_pinn.py                           58       0  100.00%
solver/physics_informed_solver/gradient_pinn.py                              17       0  100.00%
solver/physics_informed_solver/pinn.py                                       18       0  100.00%
solver/physics_informed_solver/pinn_interface.py                             54       3  94.44%   75, 166, 222
solver/physics_informed_solver/rba_pinn.py                                   74       1  98.65%   324
solver/physics_informed_solver/self_adaptive_pinn.py                        104       1  99.04%   392
solver/supervised_solver/__init__.py                                          4       0  100.00%
solver/supervised_solver/reduced_order_model.py                              24       1  95.83%   137
solver/supervised_solver/supervised.py                                        7       0  100.00%
solver/supervised_solver/supervised_solver_interface.py                      25       1  96.00%   90
TOTAL                                                                      5056     409  91.91%

Results for commit: 6d7ce0e

Minimum allowed coverage is 80.123%

♻️ This comment has been updated with latest results

[FilippoOlivo]

Hi @dario-coscia, thank you for the PR. Could you please provide an example of how this new feature should be useful?
Moreover I am not convinced anymore about the need of both OptimizerInterface and TorchOptimizer. Specifically, creating a new optimizer from scratch, without inheriting from torch.optim is not straightforward and requires a lot of effort in my opinion. In this regard, I think we can combine together the two classes in a single one, called, for example, PinaOptimizer

[GiovanniCanali]

Hi @dario-coscia, thank you for the PR. Could you please provide an example of how this new feature should be useful? Moreover I am not convinced anymore about the need of both OptimizerInterface and TorchOptimizer. Specifically, creating a new optimizer from scratch, without inheriting from torch.optim is not straightforward and requires a lot of effort in my opinion. In this regard, I think we can combine together the two classes in a single one, called, for example, PinaOptimizer

I agree with @FilippoOlivo

[dario-coscia]

Hi @GiovanniCanali @FilippoOlivo,

This PR is not ready for review yet, but here’s the planned roadmap:

Roadmap

1. Restructure Optimizer / Scheduler Interface

• Goal: Allow use of custom optimizers while reducing dependency on PyTorch. We will still need the Optimizer, Scheduler classes for this, to have a general interface from where to inherit. TorchOptimizer/TorchScheduler will be used only for torch optimisers/schedulers.
• Key tasks:
  • Override the step method (used by Lightning for model optimisation).
  • Update state dict methods for proper saving/loading of weights.

2. Introduce Second-Order Optimizers

• Goal: New Optimizers (related to New optimizers #618).
• Key tasks:
  • Implement second-order optimisation methods.
  • Add particle-based methods for benchmarking, e.g., Particle Swarm Optimization (PSO) (link).

---

This roadmap should provide a clear view of the upcoming changes and priorities.

[dario-coscia]

Roadmap

1. Restructure Optimizer / Scheduler Interface

Some updates on this…

1. This approach is not possible. Optimizers and schedulers are internally checked by Lightning, see for example:

   • Lightning optimizer checks
   • Lightning scheduler checks

   Even worse, PyTorch schedulers expect a Torch optimizer to work, see here.

   To solve this: the best approach is to create optimizers and schedulers directly using the Torch API.

   On our side, I introduced the connectors (OptimizerConnector, SchedulerConnector). Their purpose is to ensure proper hooking and handling before configure_optimizer is called.

[dario-coscia]

I think we can start adding second order optimizers/new optimizers