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

Make frontend type stable #241

Merged
merged 26 commits into from
Mar 3, 2024
Merged

Make frontend type stable #241

merged 26 commits into from
Mar 3, 2024

Conversation

lxvm
Copy link
Collaborator

@lxvm lxvm commented Feb 25, 2024

This pr makes the solve(::IntegralProblem, ::IntegralAlgorithm) interface type stable for all algorithms in the package, which solves #65. The strategy to make the infinity transformation type stable is to transform all problems at runtime instead of init time. This pr add an internal ChangeOfVariables meta-algorithm that makes adding other changes of variables easier so that #149, for example, would be easier to add. Although I haven't added tests for type consistency, this may help addressing #221. I also caught various bugs in multiple solvers and had to fix/impose type stability in others, often just adding a lot of let bindings for variables captured in closures.

The PROs of the ChangeOfVariables:

  • This also fixes an issue that it was impossible to use an infinity transformation and change the integration domain using the caching interface because previously the transformation was done at init time whereas now it happens at solve time.
  • We deprecate the do_inf_transformation keyword to solve
  • The new ChangeOfVariables wrapper should make the user experience far more consistent across algorithms because adding this wrapper caught a lot of bugs I had to fix

The CONs of the ChangeOfVariables:

  • The u-substitution allocates a few internal buffers, so there may be new allocations, especially during autodiff. Some of those buffers are hard to replace without smarter chain rules and currently I got it to work by just allocating whenever it was convenient
  • I really think the user should specify a ChangeOfVariables and it shouldn't be internal because it is possible to have nested changes of variables that will work but may reduce performance. However, requiring the user to opt-in to an infinity transformation is breaking.

This is the first example of a meta-algorithm in this packages and I'm interested to hear people's thoughts!

Checklist

  • Appropriate tests were added
  • Any code changes were done in a way that does not break public API
  • All documentation related to code changes were updated
  • The new code follows the
    contributor guidelines, in particular the SciML Style Guide and
    COLPRAC.
  • Any new documentation only uses public API

Additional context

Add any other context about the problem here.

prob.f(dy, scale(x), p)
dx .= reshape(dy, :, size(dx, 2)) .* vol
if isinplace(f)
ax = ntuple(_ -> (:), length(fsize))
Copy link
Member

Choose a reason for hiding this comment

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

Is this type stable?

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

yes, it will be type stable because of constant propagation when prototype is inferred, since fsize = size(prototype)[begin:end-1] is a tuple with (known) length. Should I add a comment?

ax = ntuple(_ -> (:), length(fsize))
_f = let y_ = similar(prototype, fsize..., nvec)
function (u, _y)
y = @view(y_[ax..., begin:(begin + size(_y, 2) - 1)])
Copy link
Member

Choose a reason for hiding this comment

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

why a view? I don't see why this isn't already the right size and just needs the reshape.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Here I think I take a view from a buffer of size nvec = min(max_batch, maxiters) so that I don't have to allocate a new array of the same size as _y on each batch integrand call. In other algorithms there is no hard upper limit on the batch size of the input array and I had to allocate a new output on each call.

_f = let y = similar(prototype)
(u, v) -> begin
resize!(y, length(v))
f(y, u, p)
Copy link
Member

Choose a reason for hiding this comment

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

why don't these scale?

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

I think cubature already does the scaling internally, whereas Cuba only allows integrals on [0, 1]^d

@lxvm
Copy link
Collaborator Author

lxvm commented Feb 28, 2024

Before proceeding further I will try to restore a test I removed that checks prob == sol.prob && alg == sol.alg, but I will probably have to add an adjoint for SciMLBase.IntegralSolution.

I'm slightly confused why the Downgrade CI succeeded but the CI failed, apparently due to a type instability that occurred in one but not the other. I couldn't immediately reproduce it so I'll have to look into it more carefully.

@lxvm
Copy link
Collaborator Author

lxvm commented Mar 2, 2024

It seems like the hard-to-reproduce type instability in GaussLegendre was happening in mapreduce so I swapped it out for a simpler implementation.

I also reinstated the tests I had removed so that the IntegralSolution contains the original problem (not the transformed one), which required helping out Zygote with some adjoints.

I think this is ready to go unless anyone has objections about the deprecation of do_inf_transformation in exchange for ChangeOfVariables with the caveats mentioned in the OP.

Comment on lines -7 to -9
Additionally, the extra keyword arguments are splatted to the library calls, so
see the documentation of the integrator library for all the extra details.
These extra keyword arguments are not guaranteed to act uniformly.
Copy link
Member

Choose a reason for hiding this comment

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

This is a standard across all other types because of the interaction with ensembles. Is there a reason to drop it here?

Copy link
Member

Choose a reason for hiding this comment

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

It seems this is covered i nhttps://github.com//pull/244 ?

Copy link
Collaborator Author

@lxvm lxvm Mar 3, 2024

Choose a reason for hiding this comment

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

Since #135 I think we currently require all kwargs to be one of abstol, reltol, or maxiters. These include the problem kwargs in the other pr. I didn't make this decision, but there is a check for this and it looks like the mechanism for algorithm-specific kwargs is for there to be a place for them in the algorithm constructor. The reason for changing the documentation is then to be consistent with the implementation, since there have been various recent issues related to misleading docs.

When you say the other SciML packages forward extra kwargs to the solvers, do you mean that, for example in QuadGKJL the library call looks like quadgk(args... ; current_kwargs..., extra_kwargs...)? The only disadvantage of this I can think of is that the user can't immediately swap between algorithms when using specialized kwargs. On the other hand, with what we currently have there is an extra maintenance burden of trying to expose all APIs of the solver in the algorithm struct for which we typically have no test coverage. Take for example the buffer parameter I added to QuadGKJL and HCubatureJL in this pr to cache the heap used internally by each algorithm (the segbuf/buffer keyword in their APIs).

As long as we are OK reverting #135, I'd be happy to pass the extra kwargs onto the solvers and do it here so that one way or another the documentation is consistent with the implementation.

@ChrisRackauckas
Copy link
Member

My only concern is the keyword arguments part, the rest looks fine so I'm fine with merging and continuing that discussion elsewhere.

@lxvm
Copy link
Collaborator Author

lxvm commented Mar 3, 2024

Are algorithms of algorithms (what I am calling meta-algorithms) used anywhere else in SciML? Most of their applications are things that users can do for themselves but which are feasible to provide for the existing api. In addition to changes of variables I think they are useful for:

  • Counting the number of integrand evaluations (to return in sol.stats)
  • Obtaining a list of points where the integrand is evaluated (to return in sol.stats)
  • Multi-dimensional integration via nested integrals with a fixed (or dynamic) order of integration (although this can be very difficult on the compiler because of types changing during recursion)

@ChrisRackauckas
Copy link
Member

I see, yes it's supposed to be filtered kwargs not all.

Are algorithms of algorithms (what I am calling meta-algorithms) used anywhere else in SciML? Most of their applications are things that users can do for themselves but which are feasible to provide for the existing api. In addition to changes of variables I think they are useful for:

Yes. DelayDiffEq just extends an ODE algorithm. Shooting methods take an ODE solver. Polyalgorithms take a list of algorithms in NonlinearSolve. I never thought about doing it for debugging information though.

@ChrisRackauckas ChrisRackauckas merged commit 973a5e4 into SciML:master Mar 3, 2024
8 checks passed
@lxvm lxvm linked an issue Mar 3, 2024 that may be closed by this pull request
@lxvm
Copy link
Collaborator Author

lxvm commented Mar 3, 2024

I think I need to make a follow up pr to fix the ChangeOfVariables adjoint for derivatives with respect to limits. The derivative tests did not catch this because they always use a domain of length 2

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.

transformation_if_inf is not type stable
2 participants