Skip to content

Add method to mirror ForwardDiff's syntax for writing Jacobian of an oop function to an allocated matrix #128

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

Merged
merged 11 commits into from
Jan 10, 2021

Conversation

chenwilliam77
Copy link
Contributor

@chenwilliam77 chenwilliam77 commented Jan 9, 2021

In ForwardDiff.jl, there is the ability to define an out-of-place function f(x) and autodiff it using

ForwardDiff.jacobian!(jac, f, x)

where jac is allocated outside of ForwardDiff.jacobian!. I have a specific use case where it'd be nice to have such a functionality in SparseDiffTools.jl and have implemented it. The tests should pass.

You may notice that I've defined

function forwarddiff_color_jacobian(J::SparseMatrixCSC{<:Number},f,x::AbstractArray{<:Number},jac_cache::ForwardColorJacCache,jac_prototype=nothing)

in addition to

function forwarddiff_color_jacobian(J::AbstractArray{<:Number},f,x::AbstractArray{<:Number},jac_cache::ForwardColorJacCache,jac_prototype=nothing)

I'm not too familiar with StaticArrays.jl, but I was unable to get one of the StaticArrays tests to pass when I used

J .+= (size(Ji)!=size(J) ? reshape(Ji,size(J)) : Ji)

instead of

J = J + (size(Ji)!=size(J) ? reshape(Ji,size(J)) : Ji)

To ensure I can mirror the function from ForwardDiff, I need the first line of code to work, but it seems to cause an error for StaticArrays.

Thus, to avoid changing the behavior of the package by too much, I just defined a version of forwarddiff_color_jacobian specifically when J is a sparse matrix, since it should be safe to do J .+= ... The more general version of forwarddiff_color_jacobian continues to use J += ... If you think that it'd be better to switch the two (i.e. write a forwarddiff_color_jacobian(J::StaticArray{<: Number}, ...) and using J .+= .. for the general version), then I'm happy to update the code.

@codecov-io
Copy link

codecov-io commented Jan 9, 2021

Codecov Report

Merging #128 (9b24650) into master (9637980) will increase coverage by 1.10%.
The diff coverage is 93.10%.

Impacted file tree graph

@@            Coverage Diff             @@
##           master     #128      +/-   ##
==========================================
+ Coverage   80.46%   81.56%   +1.10%     
==========================================
  Files          12       12              
  Lines         558      613      +55     
==========================================
+ Hits          449      500      +51     
- Misses        109      113       +4     
Impacted Files Coverage Δ
src/differentiation/compute_jacobian_ad.jl 92.39% <93.10%> (+0.15%) ⬆️

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 9637980...9b24650. Read the comment docs.

@ChrisRackauckas
Copy link
Member

It might be good to separate the StaticArray case anyways by using ArrayInterface.ismutable(x). That way we can have a fully non-allocating version and a fully non-mutating version.

@ChrisRackauckas
Copy link
Member

@ChrisRackauckas
Copy link
Member

To ensure I can mirror the function from ForwardDiff, I need the first line of code to work, but it seems to cause an error for StaticArrays.

StaticArrays cannot be mutated, so they need a fully non-mutating path. We should simply throw an error if ArrayInterface.mutable(J) == false, since that means J[i] = ... will error.

@ChrisRackauckas
Copy link
Member

Thus, to avoid changing the behavior of the package by too much, I just defined a version of forwarddiff_color_jacobian specifically when J is a sparse matrix, since it should be safe to do J .+= ... The more general version of forwarddiff_color_jacobian continues to use J += ... If you think that it'd be better to switch the two (i.e. write a forwarddiff_color_jacobian(J::StaticArray{<: Number}, ...) and using J .+= .. for the general version), then I'm happy to update the code.

We should do it slightly the other way around. Your version is a bit faster, so we should make your version be the one that is normally used, and we should send anything with ArrayInterface.mutable(J) == false to use the old OOP version since that would be slower for example on a dense J.

@ChrisRackauckas ChrisRackauckas merged commit 748a33d into JuliaDiff:master Jan 10, 2021
@ChrisRackauckas
Copy link
Member

great work, thanks!

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