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.

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

Add efficient gdg_jvp term for log-ODE schemes. #20

Merged
merged 6 commits into from
Aug 16, 2020
Merged

Add efficient gdg_jvp term for log-ODE schemes. #20

merged 6 commits into from
Aug 16, 2020

Conversation

lxuechen
Copy link
Collaborator

@lxuechen lxuechen commented Aug 16, 2020 •
edited
Loading

I've decided to go with the jvp implementation.

I implemented two versions: One with better memory scaling and the other with better time scaling. I'll send some more discussions via email.

I've only added a simple test. The numerical checks pass. The efficient version is much faster than the loop version on CPU.

@lxuechen lxuechen mentioned this pull request Aug 16, 2020
Merged
@lxuechen
Copy link
Collaborator Author

A side note is that we should squash the commits when we merge this.

patrick-kidger
patrick-kidger reviewed Aug 16, 2020
View reviewed changes
torchsde/_core/base_sde.py Outdated Show resolved Hide resolved
torchsde/_core/base_sde.py Outdated Show resolved Hide resolved
torchsde/_core/base_sde.py Outdated Show resolved Hide resolved
torchsde/_core/base_sde.py Show resolved Hide resolved
torchsde/_core/base_sde.py Outdated Show resolved Hide resolved
torchsde/_core/base_sde.py Show resolved Hide resolved
torchsde/_core/base_sde.py Outdated Show resolved Hide resolved
torchsde/_core/misc.py Show resolved Hide resolved
torchsde/_core/misc.py Outdated Show resolved Hide resolved
@patrick-kidger
Copy link
Collaborator

Also, I think the gdg_jacobian_contraction can be implemented by doing a loop / batch trick [my terminology for your fast Jacobian trick] over vjp rather than jvp, which would be worth doing what with PyTorch.

@lxuechen
Copy link
Collaborator Author

Also, I think the gdg_jacobian_contraction can be implemented by doing a loop / batch trick [my terminology for your fast Jacobian trick] over vjp rather than jvp, which would be worth doing what with PyTorch.

Feel free to give this a go.

@lxuechen lxuechen merged commit 381f4e9 into dev Aug 16, 2020
@lxuechen lxuechen deleted the dev-gdg branch August 16, 2020 19:15
@lxuechen lxuechen mentioned this pull request Aug 16, 2020
Open
47 tasks
lxuechen added a commit that referenced this pull request Oct 22, 2020
@lxuechen @patrick-kidger
@mtsokol @ucalyptus
v0.2.0 (#77)
9aad03c 
* Added BrownianInterval

* Unified base solvers

* Updated solvers to use interval interface

* Unified solvers.

* Required Python 3.6. Bumped version.

* Updated benchmarks. Fixed several bugs.

* Tweaked BrownianInterval to accept queries outside its range

* tweaked benchmark

* Added midpoint. Tweaked and fixed things.

* Tided up adjoint.

* Bye bye Python 2; fixes #14.

* tweaks from feedback

* Fix typing.

* changed version

* Rename.

* refactored settings up a level. Fixed bug in BTree.

* fixed bug with non-srk methods

* Fixed? srk noise

* Fixed SRK properly, hopefully

* fixed mistake in adjoint

* Fix broken tests and refresh documentation.

* Output type annotation.

* Rename to reverse bm.

* Fix typo.

* minor refactors in response to feedback

* Tided solvers a little further.

* Fixed strong order for midpoint

* removed unused code

* Dev kidger2 (#19)

* Many fixes.

Updated diagnostics.
Removed trapezoidal_approx
Fixed error messages for wrong methods etc.
Can now call BrownianInterval with a single value.
Fixed bug in BrownianInterval that it was always returning 0!
There's now a 2-part way of getting levy area: it has to be set as
available during __init__, and then specified that it's wanted during
__call__. This allows one to create general Brownian motions that can be
used in multiple solvers, and have each solver call just the bits it
wants.
Bugfix spacetime -> space-time
Improved checking in check_contract
Various minor tidy-ups
Can use Brownian* with any Levy area sufficient for the solver, rather
than just the minimum the solver needs.
Fixed using bm=None in sdeint and sdeint_adjoint, so that it creates an
appropriate BrownianInterval. This also makes method='srk' easy.

* Fixed ReverseBrownian

* bugfix for midpoint

* Tided base SDE classes slightly.

* spacetime->space-time; small tidy up; fix latent_sde.py example

* Add efficient gdg_jvp term for log-ODE schemes. (#20)

* Add efficient jvp for general noise and refactor surrounding.

* Add test for gdg_jvp.

* Simplify requires grad logic.

* Add more rigorous numerical tests.

* Fix all issues

* Heun's method (#24)

* Implemented Heun method

* Refactor after review

* Added docstring

* Updated heun docstring

* BrownianInterval tests + bugfixes (#28)

* In progress commit on branch dev-kidger3.

* Added tests+bugfixes for BrownianInterval

* fixed typo in docstring

* Corrections from review

* Refactor tests for

* Refactor tests for BrownianInterval.

* Refactor tests for Brownian path and Brownian tree.

* use default CPU

* Remove loop.

Co-authored-by: Xuechen Li <12689993+lxuechen@users.noreply.github.com>

* bumped numpy version (#32)

* Milstein (Strat), Milstein grad-free (Ito + Strat) (#31)

* Added milstein_grad_free, milstein_strat and milstein_strat_grad_free

* Refactor after first review

* Changes after second review

* Formatted imports

* Changed used Ex. Reversed g_prod

* Add support for Stratonovich adjoint (#21)

* Add efficient jvp for general noise and refactor surrounding.

* Add test for gdg_jvp.

* Simplify requires grad logic.

* Add more rigorous numerical tests.

* Minor refactor.

* Simplify adjoints.

* Add general noise version.

* Refactor adjoint code.

* Fix new interface.

* Add adjoint method checking.

* Fix bug in not indexing the dict.

* Fix broken tests for sdeint.

* Fix bug in selection.

* Fix flatten bug in adjoint.

* Fix zero filling bug in jvp.

* Fix bug.

* Refactor.

* Simplify tuple logic in modified Brownian.

* Remove np.searchsorted in BrownianPath.

* Make init more consistent.

* Replace np.searchsorted with bisect for speed; also fixes #29.

* Prepare levy area support for BrownianPath.

* Use torch.Generator to move to torch 1.6.0.

* Prepare space-time Levy area support for BrownianPath.

* Enable all levy area approximations for BrownianPath.

* Fix for test_sdeint.

* Fix all broken tests; all tests pass.

* Add numerical test for gradient using midpoint method for Strat.

* Support float/int time list.

* Fixes from comments.

* Additional fixes from comments.

* Fix documentation.

* Remove to for BrownianPath.

* Fix insert.

* Use none for default levy area.

* Refactor check tensor info to reduce boilerplate.

* Add a todo regarding get noise.

* Remove type checks in adjoint.

* Fixes from comments.

* Added BrownianReturn (#34)

* Added BrownianReturn

* Update utils.py

* Binterval improvements (#35)

* Tweaked to not hang on adaptive solvers.

* Updated adaptive fix

* Several fixes for tests and adjoint.

Removed some broken tests.
Added error-raising `g` to the adjoint SDE.
Fixed Milstein for adjoint.
Fixed running adjoint at all.

* fixed bug in SRK

* tided up BInterval

* variable name tweak

* Improved heuristic for BrownianInterval's dependency tree. (#40)

* [On dev branch] Tuple rewrite (#37)

* Rename plot folders from diagnostics.

* Complete tuple rewrite.

* Remove inaccurate comments.

* Minor fixes.

* Fixes.

* Remove comment.

* Fix docstring.

* Fix noise type for problem.

* Binterval recursion fix (#42)

* Improved heuristic for BrownianInterval's dependency tree.

* Inlined the recursive code to reduce number of stack frames

* Add version number.

Co-authored-by: Xuechen <12689993+lxuechen@users.noreply.github.com>

* Refactor.

* Euler-Heun method (#39)

* Implemented euler-heun

* After refactor

* Applied refactor. Added more diagnostics

* Refactor after review

* Corrected order

* Formatting

* Formatting

* BInterval - U fix (#44)

* Improved heuristic for BrownianInterval's dependency tree.

* fixed H aggregation

* Added consistency test

* test fixes

* put seed back

* from comments

* Add log-ODE scheme and simplify typing. (#43)

* Add log-ODE scheme and simplify typing.

* Register log-ODE method.

* Refactor diagnostics and examples.

* Refactor plotting.

* Move btree profile to benchmarks.

* Refactor all ito diagnostics.

* Refactor.

* Split imports.

* Refactor the Stratonovich diagnostics.

* Fix documentation.

* Minor typing fix.

* Remove redundant imports.

* Fixes from comment.

* Simplify.

* Simplify.

* Fix typo caused bug.

* Fix directory issue.

* Fix order issue.

* Change back weak order.

* Fix test problem.

* Add weak order inspection.

* Bugfixes for log-ODE (#45)

* fixed rate diagnostics

* tweak

* adjusted test_strat

* fixed logODE default.

* Fix typo.

Co-authored-by: Xuechen Li <12689993+lxuechen@users.noreply.github.com>

* Default to loop-based. Fixes #46.

* Minor tweak of settings.

* Fix directory structure.

* Speed up experiments.

* Cycle through the possible line styles.

Co-authored-by: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>

* Simplify and fix documentation.

* Minor fixes.

- Simplify strong order assignment for euler.
- Fix bug with "space_time".

* Simplify strong order assignment for euler.

* Fix bug with space-time naming.

* Make tensors for grad for adjoint specifiable. (#52)

* Copy of #55 | Created pyproject.toml (#56)

* Skip tests if the optional C++ implementations don't compile; fixes #51.

* Create pyproject.toml

* Version add 1.6.0 and up

Co-authored-by: Xuechen <12689993+lxuechen@users.noreply.github.com>

* Latent experiment (#48)

* Latent experiment

* Refactor after review

* Fixed y0

* Added stable div

* Minor refactor

* Simplify latent sde even further.

* Added double adjoint (#49)

* Added double adjoint

* tweaks

* Updated adjoint tests

* Dev adjoint double test (#57)

* Add gradgrad check for adjoint.

* Relax tolerance.

* Refactor numerical tests.

* Remove unused import.

* Fix bug.

* Fixes from comments.

* Rename for consistency.

* Refactor comment.

* Minor refactor.

* Add adjoint support for general/scalar noise in the Ito case. (#58)

* adjusted requires_grad

Co-authored-by: Xuechen Li <12689993+lxuechen@users.noreply.github.com>

* Dev minor (#63)

* Add requirements and update latent sde.

* Fix requirements.

* Fix.

* Update documentation.

* Use split to speed things up slightly.

* Remove jit standalone.

* Enable no value arguments.

* Fix bug in args.

* Dev adjoint strat (#67)

* Remove logqp test.

* Tide examples.

* Refactor to class attribute.

* Fix gradcheck.

* Reenable adjoints.

* Typo.

* Simplify tests

* Deprecate this test.

* Add back f ito and strat.

* Simplify.

* Skip more.

* Simplify.

* Disable adaptive.

* Refactor due to change of problems.

* Reduce problem size to prevent general noise test case run for ever.

* Continuous Integration.  (#68)

* Skip tests if the optional C++ implementations don't compile; fixes #51.

* Continuous integration.

* Fix os.

* Install package before test.

* Add torch to dependency list.

* Reduce trials.

* Restrict max number of parallel runs.

* Add scipy.

* Fixes from comment.

* Reduce frequency.

* Fixes.

* Make sure run installed package.

* Add check version on pr towards master.

* Separate with blank lines.

* Loosen tolerance.

* Add badge.

* Brownian unification (#61)

* Added tol. Reduced number of generator creations. Spawn keys now of
finite length. Tidied code.

* Added BrownianPath and BrownianTree as BrownianInterval wrappers

* added trampolining

* Made Path and Tree wrappers on Interval.

* Updated tests. Fixed BrownianTree determinism. Allowed cache_size=0

* done benchmarks. Fixed adjoint bug. Removed C++ from setup.py

* fixes for benchmark

* added base brownian

* BrownianPath/Tree now with the same interface as before

* BInterval(shape->size), changed BPath and BTree to composition-over-inheritance.

* tweaks

* Fixes for CI. (#69)

* Fixes for CI.

* Tweaks to support windows.

* Patch for windows.

* Update patch for windows.

* Fix flaky tests of BInterval.

* Add fail-fast: false (#72)

* Dev methods fixes (#73)

* Fixed adaptivity checks. Improved default method selection.

* Fixes+updated sdeint tests

* adjoint method fixes

* Fixed for Py3.6

* assert->ValueError; tweaks

* Dev logqp (#75)

* Simplify.

* Add stable div utility.

* Deprecate.

* Refactor problems.

* Sync adjoint tests.

* Fix style.

* Fix import style.

* Add h to test problems.

* Add logqp.

* Logqp backwards compatibility.

* Add type annotation.

* Better documentation.

* Fixes.

* Fix notebook. (#74)

* Fix notebook.

* Remove trivial stuff.

* Fixes from comments.

* Fixes.

* More fixes.

* Outputs.

* Clean up.

* Fixes.

* fixed BInterval flakiness+slowness (#76)

* Added documentation (#71)

* Added documentation

* tweaks

* Fix significance level.

* Fix check version.

* Skip confirmation.

* Fix indentation errors.

* Update README.md

Co-authored-by: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Co-authored-by: Mateusz Sokół <8431159+mtsokol@users.noreply.github.com>
Co-authored-by: Sayantan Das <36279638+ucalyptus@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@patrick-kidger patrick-kidger patrick-kidger left review comments

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@lxuechen @patrick-kidger