Introduce diagonal ad arrays

[keileg]

Proposed changes

This PR implements a dense-array backend for forward Ad assembly. The idea is that, for quantities that are dependent on cell center quantities only (the idea can be generalized, but cell center values seemed like a reasonable compromise between generality and simplicity in implementation), derivatives can be represented by 2d dense numpy array instead of general scipy sparse matrices. A new class, DiagonalAdArray is introduced to represent such quantities.

The point is that, in forward mode Ad, when parsing expressions that consist of only DiagonalAdArrays, the rules of differentiation will be applied on numpy dense arrays, which can be significantly faster than dealing with sparse arrays (see benchmarking results below). The result from combining two DiagonalAdArrays in an arithmetic operation is a new DiagonalAdArray, hence for complex expressions the computational gain will accumulate. When a DiagonalAdArray is hit with a projection or a differential operator, it is transformed to a standard AdArray and the old Ad parsing machinery takes over.

See #1497 for more information and this outline for the original idea.

Scope

Though it is possible to extend the idea beyond expressions depending on cell center values, and hence delay the transformation to usual AdArrays, I have chosen not to do so. To limit technicalities in the implementation, we also initialize a DiagonalAdArray only for scalar quantities (displacement does not qualify) defined on all subdomains. To be clear, in a displacement-pressure system, the pressure can still be represented as a DiagonalAdArray. Interface variables are not converted.
I have not explored what it will take to go beyond these restrictions, but much of the parsing of the compositional models (which are the ones suffering the most from slow assembly) should be covered, so this should be a step in the right direction. If the results are positive, I am open for expanding the coverage of the idea in follow-up projects, but that should be based on identification of bottlenecks based on benchmarking.

Benchmarking

I am 100% sure the results are in the right direction. The actual numbers produced below may partly be setup dependent, but the overall results should be fairly solid.

1. When measuring parsing time in forward mode (no Ad Operator parsing, just x+y and similar), diagonal arrays give a speedup of around 4-5 for arrays with size 10^5-10^6. The speedup seems to stabilize if the size is increased further. For smaller arrays (10^2-10^3), the speedup is much higher (~30), but this matters less.
2. The speedup remains for chained expressions (repeated evaluations). The absolute saving will correspondingly grow with the expression size.
3. When applied to the model (here I used the notoriously slow tests/functional/test_buoyancy_flow.py as the starting point for benchmarking) I got a speedup of between 2 and 3 in the total time spent on assembly. Two comments:
   • The benchmark involves a mixture of expressions using cell centered quantities that are suitable for DiagonalAdArrays and expressions that are not suited for diagonal representation. I have not estimated the fraction of the computations that were amenable to diagonal representation, but the reduced gain compared to the forward mode evaluation can at least in part be attributed to this.
   • The measurement was done before SPEED: AdParser dispatches binary operations directly instead of eval() #1751 was introduced. That improvement will reduce the time spent on assembly with equal gains between diagonal and full implementation. Hence, if the benchmarking experiment is repeated, the speedup should be higher, but I have not prioritized to do this.

Other changes

• The test of parsing of arithmetic operations is extended to cover diagonal arrays with all combinations. This test has been refactored and placed in a dedicated file. I think the overall structure is more logical now, but given the scope of the test, there is a limit to how nice this can be.
• The ad functions heaviside etc. now have systematic testing.

Note that this PR is put on top of #1747 and thus seemingly contains substantial changes beyond what is described above. The PR will remain a draft until the upstream PR is merged.

Help needed

Though everything seems to be working, it would be great to have some QC on code robustness beyond what is currently available in the tests. In particular, a confirmation that the update does not break complex application code and that the timings broadly apply also in these cases. @mikeljordan @vlipovac: Would it be possible for you to run some of your simulations at some point and see what happens. It is not urgent since we cannot really move before #1747 is merged (also note that, to the degree you define your own equations on Ad operator form, that PR may introduce breaking changes - it should not be difficult to fix this, but an update will be needed).

Resolves #1497.

Types of changes

What types of changes does this PR introduce to PorePy?
Put an x in the boxes that apply.

• Minor change (e.g., dependency bumps, broken links).
• Bugfix (non-breaking change which fixes an issue).
• New feature (non-breaking change which adds functionality).
• Breaking change (fix or feature that would cause existing functionality to not work as expected).
• Testing (contribution related to testing of existing or new functionality).
• Documentation (contribution related to adding, improving, or fixing documentation).
• Maintenance (e.g., improve logic and performance, remove obsolete code).
• Other:

Checklist

Put an x in the boxes that apply or explain briefly why the box is not relevant.

• The documentation is up-to-date.
• Static typing is included in the update.
• This PR does not duplicate existing functionality.
• The update is covered by the test suite (including tests added in the PR).
• If new skipped tests have been introduced in this PR, pytest was run with the --run-skipped flag.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB