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

Add derivatives for besselix, besseljx, and besselyx #350

Merged
merged 6 commits into from
Oct 18, 2021
+100 −1
Merged

Add derivatives for besselix, besseljx, and besselyx #350

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
89e2dc9
Add derivatives for `besselix`, `besseljx`, and `besselyx`
devmotion Sep 28, 2021
d1076fc
Bump version
devmotion Sep 28, 2021
c26d35e
Apply suggestions from @sethaxen's review (muladd + optimizations)
devmotion Oct 1, 2021
0af9568
Improve `frule`s
devmotion Oct 2, 2021
30a020f
Simplify `frule`s
devmotion Oct 3, 2021
a9f9f52
Apply suggestions from code review
devmotion Oct 4, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
name = "SpecialFunctions"
uuid = "276daf66-3868-5448-9aa4-cd146d93841b"
version = "1.7.0"
version = "1.8.0"

[deps]
ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
Expand Down
63 changes: 63 additions & 0 deletions src/chainrules.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,3 +193,66 @@ ChainRulesCore.@scalar_rule(
ChainRulesCore.@scalar_rule(expinti(x), exp(x) / x)
ChainRulesCore.@scalar_rule(sinint(x), sinc(invπ * x))
ChainRulesCore.@scalar_rule(cosint(x), cos(x) / x)

# non-holomorphic functions
function ChainRulesCore.frule((_, _, _), ::typeof(besselix), ν::Number, x::Number)
Ω = besselix(ν, x)
ΔΩ = ChainRulesCore.@not_implemented(BESSEL_ORDER_INFO)
return Ω, ΔΩ
Copy link
Contributor

Choose a reason for hiding this comment

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

My earlier comment might have gotten lost, but since ZeroTangent() * ::NotImplemented is a ZeroTangent, you can implement the frule in such a way that if the AD provides ZeroTangent() for Δν, then the frule does not return a NotImplemented:

Suggested change
function ChainRulesCore.frule((_, _, _), ::typeof(besselix), ν::Number, x::Number)
Ω = besselix(ν, x)
ΔΩ = ChainRulesCore.@not_implemented(BESSEL_ORDER_INFO)
return Ω, ΔΩ
function ChainRulesCore.frule((_, Δν, Δx), ::typeof(besselix), ν::Number, x::Number)
Ω = besselix(ν, x)
∂Ω_∂ν = ChainRulesCore.@not_implemented(BESSEL_ORDER_INFO)
a = (besselix(ν - 1, x) + besselix(ν + 1, x)) / 2
∂Ω = muladd(a, Δx, muladd(-sign(real(x)) * real(Δx), Ω, ∂Ω_∂ν * Δν)
return Ω, ∂Ω

Copy link
Member Author

Choose a reason for hiding this comment

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

Oh yes, I completely forgot that ZeroTangent() * ::NotImplemented = ZeroTangent(). I'll fix the forward rules!

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 noticed that currently it is not possible to test such definitions: JuliaDiff/ChainRulesCore.jl#477

end
function ChainRulesCore.rrule(::typeof(besselix), ν::Number, x::Number)
Ω = besselix(ν, x)
project_x = ChainRulesCore.ProjectTo(x)
function besselix_pullback(ΔΩ)
ν̄ = ChainRulesCore.@not_implemented(BESSEL_ORDER_INFO)
a = (besselix(ν - 1, x) + besselix(ν + 1, x)) / 2
Copy link
Contributor

Choose a reason for hiding this comment

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

One can use the recurrence relations to write this as (besselix(ν ± 1, x) ± besselix(ν, x)) * ν / x, which allows to reuse Ω. It would just require some special-casing to handle x=0 gracefully.

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 just reused the derivatives that are used for besseli etc. They were defined in this way in ChainRules originally. When I copied them to SpecialFunctions I wondered about the motivation for choosing https://functions.wolfram.com/Bessel-TypeFunctions/BesselI/20/01/02/0003/ over https://functions.wolfram.com/Bessel-TypeFunctions/BesselI/20/01/02/0001/ or https://functions.wolfram.com/Bessel-TypeFunctions/BesselI/20/01/02/0002/. I assumed the currently used definition is simpler since it does not require to handle x = 0 in a special way.

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 think it might be best to use the same relations for besselix etc. as for besseli etc. and, if desired, change them to a different form in a separate PR.

Copy link
Contributor

Choose a reason for hiding this comment

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

That's right, and they're the same in DiffRules as well. If you like, you can keep them similar to besseli, etc for now, and a future PR could update all of the rules.

Copy link
Contributor

Choose a reason for hiding this comment

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

I agree!

b = - sign(real(x)) * Ω
x̄ = project_x(conj(a) * ΔΩ + real(conj(b) * ΔΩ))
Copy link
Contributor

Choose a reason for hiding this comment

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

It's more efficient to compute -sign(real(x)) * real(conj(Ω) * ΔΩ)) because you're multiplying then 2 reals instead of a real and a complex.

The @scalar_rule macro writes rules to use muladd here. @oxinabox does that make a difference?

Copy link
Contributor

Choose a reason for hiding this comment

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

I have never measured

return ChainRulesCore.NoTangent(), ν̄, x̄
end
return Ω, besselix_pullback
end

function ChainRulesCore.frule((_, _, _), ::typeof(besseljx), ν::Number, x::Number)
Copy link
Contributor

Choose a reason for hiding this comment

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

Since everything should be identical for besseljx and besselyx, can you declare them the same with a loop and an @eval?

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 could but so far the style in this file is to implement everything explicitly (eg. also derivatives besselj and bessely are implemented separately) and so I sticked to it.

Copy link
Contributor

Choose a reason for hiding this comment

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

Makes sense.

Ω = besseljx(ν, x)
ΔΩ = ChainRulesCore.@not_implemented(BESSEL_ORDER_INFO)
return Ω, ΔΩ
end
function ChainRulesCore.rrule(::typeof(besseljx), ν::Number, x::Number)
Ω = besseljx(ν, x)
project_x = ChainRulesCore.ProjectTo(x)
function besseljx_pullback(ΔΩ)
ν̄ = ChainRulesCore.@not_implemented(BESSEL_ORDER_INFO)
a = (besseljx(ν - 1, x) - besseljx(ν + 1, x)) / 2
x̄ = if x isa Real
project_x(conj(a) * ΔΩ)
Copy link
Contributor

Choose a reason for hiding this comment

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

If x is real, then so is a:

Suggested change
project_x(conj(a) * ΔΩ)
project_x(a * ΔΩ)

else
b = -sign(imag(x)) * Ω
project_x(conj(a) * ΔΩ + real(conj(b) * ΔΩ) * im)
end
return ChainRulesCore.NoTangent(), ν̄, x̄
end
return Ω, besseljx_pullback
end

function ChainRulesCore.frule((_, _, _), ::typeof(besselyx), ν::Number, x::Number)
Ω = besselyx(ν, x)
ΔΩ = ChainRulesCore.@not_implemented(BESSEL_ORDER_INFO)
return Ω, ΔΩ
end
function ChainRulesCore.rrule(::typeof(besselyx), ν::Number, x::Number)
Ω = besselyx(ν, x)
project_x = ChainRulesCore.ProjectTo(x)
function besselyx_pullback(ΔΩ)
ν̄ = ChainRulesCore.@not_implemented(BESSEL_ORDER_INFO)
a = (besselyx(ν - 1, x) - besselyx(ν + 1, x)) / 2
x̄ = if x isa Real
project_x(conj(a) * ΔΩ)
else
b = -sign(imag(x)) * Ω
project_x(conj(a) * ΔΩ + real(conj(b) * ΔΩ) * im)
end
return ChainRulesCore.NoTangent(), ν̄, x̄
end
return Ω, besselyx_pullback
end
6 changes: 6 additions & 0 deletions test/chainrules.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,9 +53,13 @@
for nu in (-1.5, 2.2, 4.0)
test_frule(besseli, nu, x)
test_rrule(besseli, nu, x)
test_frule(besselix, nu, x)
test_rrule(besselix, nu, x)

test_frule(besselj, nu, x)
test_rrule(besselj, nu, x)
test_frule(besseljx, nu, x)
test_rrule(besseljx, nu, x)

test_frule(besselk, nu, x)
test_rrule(besselk, nu, x)
Expand All @@ -64,6 +68,8 @@

test_frule(bessely, nu, x)
test_rrule(bessely, nu, x)
test_frule(besselyx, nu, x)
test_rrule(besselyx, nu, x)

test_frule(hankelh1, nu, x)
test_rrule(hankelh1, nu, x)
Expand Down