Skip to content
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.

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

Improve consistent initialization speed and robustness #4301

Merged

Conversation

MarcBerliner
Copy link
Member

@MarcBerliner MarcBerliner commented Jul 29, 2024

Description

The IDA solver takes in initial guesses for y0 and ydot0. Currently, we explicitly calculate y0 using an algebraic solver before passing it to IDA. However, we do not estimate ydot0 - we just set it to np.zeros_like(y0) and rely on IDA to determine it. This PR makes a few changes to consistent initialization:

  1. Analytically compute ydot0 using the mass matrix inverse. The IDA solver is set up only for implicit DAEs, so it cannot properly utilize the mass matrix in initializing the system. As a result, it wastes a lot of time re-computing the consistent initialization for y0 and ydot0 in IDACalcIC. We can improve this by isolating the differential section of our system of equations and explicitly computing ydot0 for the rhs,
    $$\left. \frac{\mathrm{d}y_d^*}{\mathrm{d}t} \right|_{t_0} = ( M_{d,d} ) ^ {-1} \text{rhs}(y_0^*, t_0) , $$
    where the subscript $d$ refers to the equations containing differential terms and $y_0^*$ is the consistently initialized $y_0$. This PR reduces the total solve time of a full 1C discharge from 218 ms to 212 ms and reduces the time integrating the system from 45.0 ms to 40.0 ms. Additionally, this improves the solver robustness because we reduce reliance on IDACalcIC, which occasionally fails.
  2. Separate initial_conditions and consistent_initialization. I found that there was not a clear distinction between these ideas in the solver code. In this update, initial_conditions refers to the initial values and guesses provided by the solver or from a previous solution. Initial conditions likely will not satisfy the system of equations at t0, while consistent_initialization refers to the solution that satisfies the system of equations at t0. In some cases, we computed the initial_conditions more than once for a single solve, which were removed.
  3. Various cleanup of the solver codebase. This section will need additional refactoring to unify the API and remove redundancies, but that is beyond the scope of this PR.

Type of change

Please add a line in the relevant section of CHANGELOG.md to document the change (include PR #) - note reverse order of PR #s. If necessary, also add to the list of breaking changes.

  • Optimization (back-end change that speeds up the code)

Key checklist:

  • No style issues: $ pre-commit run (or $ nox -s pre-commit) (see CONTRIBUTING.md for how to set this up to run automatically when committing locally, in just two lines of code)
  • All tests pass: $ python run-tests.py --all (or $ nox -s tests)
  • The documentation builds: $ python run-tests.py --doctest (or $ nox -s doctests)

You can run integration tests, unit tests, and doctests together at once, using $ python run-tests.py --quick (or $ nox -s quick).

Further checks:

  • Code is commented, particularly in hard-to-understand areas
  • Tests added that prove fix is effective or that feature works

Copy link
Member

@valentinsulzer valentinsulzer left a comment

Choose a reason for hiding this comment

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

Looks good to me at a superficial level. Much overdue clean up and good to know why we've been getting all these IDACalcIC errors.
I'll let others (@martinjrobins / @jsbrittain ) who have done more on the IDA KLU solver review in more detail

Copy link

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

Copy link

codecov bot commented Jul 29, 2024

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.50%. Comparing base (f255c38) to head (4a7886f).
Report is 163 commits behind head on develop.

Additional details and impacted files
@@             Coverage Diff             @@
##           develop    #4301      +/-   ##
===========================================
+ Coverage    99.45%   99.50%   +0.04%     
===========================================
  Files          288      288              
  Lines        22091    22079      -12     
===========================================
- Hits         21971    21970       -1     
+ Misses         120      109      -11     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

Copy link
Contributor

@martinjrobins martinjrobins left a comment

Choose a reason for hiding this comment

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

Thanks @MarcBerliner, this looks great, thank you for your effort on this. In the future, as per the CONTRIBUTING.md can you please make sure to first open an issue before you start on a piece of work so that it can be discussed with the other developers. It would be great also to get your input on #3910. Please feel free to add additional items that you think are important for the solvers going forward, or edit existing items. I would like this to be an evolving workplan so we can coordinate everyone's work on the solvers.

and model.mass_matrix.shape[0] > model.len_rhs_and_alg
):
if model.mass_matrix_inv is not None:
if has_mass_matrix and model.mass_matrix.shape[0] > model.len_rhs_and_alg:
Copy link
Contributor

Choose a reason for hiding this comment

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

can we just get rid of this if statement? on the face of it, it seems unnessessary

Copy link
Member Author

Choose a reason for hiding this comment

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

I added a check that early returns if we don't have a mass matrix (necessary for the DummySolver tests),

if not has_mass_matrix:
    return

and removed all the following has_mass_matrix checks.

# calculate the time derivatives of the differential equations
# for semi-explicit DAEs
calc_ydot0 = model.mass_matrix_inv is not None
if model.len_rhs > 0 and calc_ydot0:
Copy link
Contributor

Choose a reason for hiding this comment

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

its a bit unclear in what situation the solver would calculate ydot0. Presumably we will always want to calculate ydot0 if there are any differential equations (i.e. model.len_rhs > 0)? so the model.mass_matrix_inv should be available in this case (and we should error if not).

Copy link
Member Author

Choose a reason for hiding this comment

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

Good point - we will always have a mass matrix and its inverse if len_rhs is nonzero. Removing calc_ydot0

pybamm/solvers/idaklu_solver.py Outdated Show resolved Hide resolved
@MarcBerliner
Copy link
Member Author

Thanks for the quick feedback @martinjrobins. Going forward, I'll make sure to keep you all informed of my planned work. I'll add a few comments to the solver roadmap.

Copy link
Contributor

@martinjrobins martinjrobins left a comment

Choose a reason for hiding this comment

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

looks good, thanks @MarcBerliner

@martinjrobins martinjrobins merged commit 2de4de6 into pybamm-team:develop Jul 31, 2024
25 of 26 checks passed
@MarcBerliner MarcBerliner deleted the consistent-initialization branch July 31, 2024 16:26
js1tr3 pushed a commit to js1tr3/PyBaMM that referenced this pull request Aug 12, 2024
)

* consistent initalization

* update description

* add test

* Update CHANGELOG.md

* (fix): failing tests

* adjust test tolerance

* adjust tolerances

* codecov, fix tests

* adjust IREE tol

* Martin's comments

attempt to fix the iree tols

* revert tests, shorter tspan

* fix codecov
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants