Skip to content

Optim refactoring #662

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 1 commit into
base: dev
Choose a base branch
from
Open

Optim refactoring #662

wants to merge 1 commit into from

Conversation

@dario-coscia
Copy link
Collaborator

@dario-coscia dario-coscia commented Oct 10, 2025

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

@dario-coscia dario-coscia requested a review from a team as a code owner October 10, 2025 10:35
@dario-coscia dario-coscia added the pr-to-fix Label for PR that needs modification label Oct 10, 2025
@dario-coscia dario-coscia changed the title optim refactoring Optim refactoring Oct 10, 2025
@dario-coscia dario-coscia added enhancement New feature or request maintenance labels Oct 10, 2025
@dario-coscia dario-coscia self-assigned this Oct 10, 2025
@github-actions
Copy link
Contributor

github-actions bot commented Oct 10, 2025 •
edited
Loading

badge

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
FilippoOlivo reviewed Oct 10, 2025
View reviewed changes
Copy link
Member

@FilippoOlivo FilippoOlivo left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@dario-coscia dario-coscia force-pushed the optim branch 2 times, most recently from 4fdf1ae to d04f75a Compare October 10, 2025 12:24
@GiovanniCanali
Copy link
Collaborator

GiovanniCanali commented Oct 10, 2025

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
Copy link
Collaborator Author

dario-coscia commented Oct 10, 2025 •
edited
Loading

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 dario-coscia changed the base branch from dev to dev_updates October 19, 2025 11:23
@dario-coscia dario-coscia changed the base branch from dev_updates to dev October 20, 2025 11:34
@dario-coscia
optim refactoring
92aede6 
* adding connectors for optimizers/schedulers
* simplify configure_optimizers logic
@dario-coscia
Copy link
Collaborator Author

dario-coscia commented Oct 20, 2025

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:

    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
Copy link
Collaborator Author

dario-coscia commented Oct 20, 2025

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@FilippoOlivo FilippoOlivo FilippoOlivo left review comments

Assignees

@dario-coscia dario-coscia

Labels

enhancement New feature or request maintenance pr-to-fix Label for PR that needs modification

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@dario-coscia @GiovanniCanali @FilippoOlivo