mdolab / mdolab/OpenAeroStruct
Add optional support of JAX to accelerate some partial derivatives
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 281
- Forks
- 135
- PR merge metrics
- No merged PRs in 30d
Description
Description of feature
When using a dense VLM mesh, compute_partials in some components (e.g., eval_mtx in aerodynamics) becomes a bottleneck for derivative computations. These partials can be accelerated by replacing the current analytical derivatives with AD.
Aditya Deshpande and Sriram Bommakanti tried it out for the AE588 project, and they showed that AD actually accelerated the partials. Their prototype implementation can be found in their fork. Note that they used AD for only some part of the compute_partials computations, but not to the entire partial computation.
AD support should be optional because we don't want to add JAX as a hard dependency (for now), and AD likely doesn't offer performance benefits for moderate mesh size.
Potential solution
- Run profiling and identify the components that can be accelerated by AD.
eval_mtxis one, but there could be others. - Replace (part of) the
compute_partialsmethod with AD. We'll need to try out multiple AD options as Aditya and Sriram did. - Add an optional dependency on JAX in setup.py
- Add a documentation page on AD - ideally, suggest a mesh size threshold at which the AD becomes faster than the default analytical partials.
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reviewing the prototype in the linked fall23_project_fullAD fork and profiling compute_partials, including eval_mtx, to identify suitable components. Then inspect setup.py for optional dependency handling and define how the AD documentation should describe mesh-size tradeoffs and the completed integration.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100