Benchmarks + landing page (#10)

* fix: solve problem with different shapes in input

* feat: start landing

* fix: favicon

* feat: add benchmark util func

* fix: try cuda first

* fix: merge new translate

* fix: outputs

* fix: try with memory cleaning

* fix: change order of jit vmap

* fix: increase batch

* fix: prints

* fix: no tqdm

* feat: increase the batch

* fix: reduce

* fix: no tqdm

* feat: add visualize demo

* feat: add notebook from colab

* fix: plots

* feat: add patches and update readme

* fix: benchmarks on landing

* fix: landing

* fix: landing

* fix: expand benchmarking readme

* fix: bump version

---------

Co-authored-by: Lev Kozlov <[email protected]>

CONTRIBUTING.md

@@ -0,0 +1,56 @@
+# Contributing to JaxADi
+
+We're thrilled that you're interested in contributing to JaxADi! This document outlines the process for contributing to this project.
+
+## How to Contribute
+
+There are many ways to contribute to JaxADi:
+
+1. Reporting bugs
+2. Suggesting enhancements
+3. Writing documentation
+4. Submitting code changes
+
+### Reporting Bugs
+
+1. Check the [issue tracker](https://github.com/based-robotics/jaxadi/issues) to see if the bug has already been reported.
+2. If not, create a new issue. Provide a clear title and description, as much relevant information as possible, and a code sample or executable test case demonstrating the bug.
+
+### Suggesting Enhancements
+
+1. Check the [issue tracker](https://github.com/based-robotics/jaxadi/issues) to see if the enhancement has already been suggested.
+2. If not, create a new issue. Clearly describe the enhancement, why it would be useful, and any potential drawbacks.
+
+### Writing Documentation
+
+Good documentation is crucial. If you notice any part of our documentation that could be improved or expanded, please let us know or submit a pull request with your suggested changes.
+
+### Submitting Code Changes
+
+1. Fork the repository.
+2. Create a new branch for your changes.
+3. Make your changes in your branch.
+4. Add or update tests as necessary.
+5. Ensure the test suite passes.
+6. Update the documentation as needed.
+7. Push your branch and submit a pull request.
+
+## Pull Request Process
+
+1. Ensure your code follows the project's style guidelines.
+2. Update the README.md or relevant documentation with details of changes, if applicable.
+3. Add tests for your changes and ensure all tests pass.
+4. Your pull request will be reviewed by the maintainers. They may suggest changes or improvements.
+5. Once approved, your pull request will be merged.
+
+<!-- TODO:
+
+## Style Guidelines
+
+We use Ruff to enforce our code style and linting rules....
+```bash
+pip3 install pre-commit
+```
+```bash
+pre-commit install
+``` -->

README.md

@@ -1,32 +1,33 @@
 <!-- # JaxADi -->
-<!-- TODO: ADD PATCHES -->
+
+[![CI](https://img.shields.io/github/actions/workflow/status/based-robotics/jaxadi/build.yaml?branch=master)](https://github.com/based-robotics/jaxadi/actions)
+[![PyPI version](https://img.shields.io/pypi/v/jaxadi?color=blue)](https://pypi.org/project/jaxadi/)
+[![PyPI downloads](https://img.shields.io/pypi/dm/jaxadi?color=blue)](https://pypistats.org/packages/jaxadi)
 
 <p align="center">
   <!-- Placeholder for a cool logo -->
   <img src="https://github.com/based-robotics/jaxadi/blob/master/_assets/_logo.png?raw=true" alt="JAXADI Logo" width="400"/>
 </p>
 
-
 **JaxADi** is a powerful Python library designed to bridge the gap between `casadi.Function` and JAX-compatible functions. By leveraging the strengths of both CasADi and JAX, JAXADI opens up exciting opportunities for building highly efficient, batchable code that can be executed seamlessly across CPUs, GPUs, and TPUs.
 
 JAXADI can be particularly useful in scenarios involving:
 
 - Robotics simulations
 - Optimal control problems
 - Machine learning models with complex dynamics
-- Large-scale numerical optimizations
-
 
 ## Installation
 
 You can install JAXADI using pip:
+
 <!-- Change once it will be realeased -->
 
 ```bash
 pip install jaxadi
 ```
 
-For a complete environment setup, we recommend using Conda/Mamba:
+For a complete environment setup for examples, we recommend using Conda/Mamba:
 
 ```bash
 mamba env create -f environment.yml
@@ -63,32 +64,53 @@ output = jax_fn(input_x, input_y)
 
 ```
 
+<strong>Note:</strong> For now translation does not support functions with very
+large number of operations, due to the translation implementation. Secret component of
+translation is work-tree expansion, which might lead to large overhead in number of
+symbols. We are working on finding the compromise in both speed and extensive
+functions support.
+
 ## Examples
 
 JAXADI comes with several examples to help you get started:
 
-1. [Basic Translation](examples/00_translate.py): Learn how to translate CasADi functions to JAX. 
+1. [Basic Translation](examples/00_translate.py): Learn how to translate CasADi functions to JAX.
 
-2. [Lowering Operations](examples/01_lower.py): Understand the lowering process in JaxADi. 
+2. [Lowering Operations](examples/01_lower.py): Understand the lowering process in JaxADi.
 
-3. [Function Conversion](examples/02_convert.py): See how to fully convert CasADi functions to JAX. 
+3. [Function Conversion](examples/02_convert.py): See how to fully convert CasADi functions to JAX.
 
-4. [Pendulum Rollout](examples/03_pendulum_rollout.py): Batched rollout of the nonlinear passive nonlinear pendulum 
+4. [Pendulum Rollout](examples/03_pendulum_rollout.py): Batched rollout of the nonlinear passive nonlinear pendulum
 
-5. [Pinocchio Integration](examples/04_pinocchio.py): Explore how to convert Pinocchio-based CasADi functions to JAX. 
+5. [Pinocchio Integration](examples/04_pinocchio.py): Explore how to convert Pinocchio-based CasADi functions to JAX.
 
 6. [MJX Comparison](examples/05_mjx.py): Compare the transformed Pinnocchio forward kinematics with one provided by Mujoco MJX
 
 > **Note**: To run the Pinocchio and MJX examples, ensure you have them properly installed in your environment.
 
+## Performance Benchmarks
 
-<!-- ## Performance Benchmarks
+![speedup](https://github.com/based-robotics/jaxadi/blob/master/docs/static/images/speedup_ratio.png?raw=true)
 
-(Consider adding a section about performance comparisons between CasADi and JAXADI-translated functions) -->
+The process of benchmarking and evaluating the performance of Jaxadi is described in the [benchmarks](benchmarks/README.md) directory.
 
-<!-- ## Contributing
+## Contributing
 
-We welcome contributions! Please see our [Contributing Guide](CONTRIBUTING.md) for more details. -->
+We welcome contributions! Please see our [Contributing Guide](CONTRIBUTING.md) for more details.
+
+## Citation
+
+If you use JaxADi in your research, please cite it as follows:
+
+```bibtex
+@misc{jaxadi2024,
+  title = {JaxADi: Bridging CasADi and JAX for Efficient Numerical Computing},
+  author = {Alentev, Igor and Kozlov, Lev and Nedelchev, Simeon},
+  year = {2024},
+  url = {https://github.com/based-robotics/jaxadi},
+  note = {Accessed: [Insert Access Date]}
+}
+```
 
 ## Acknowledgements
 
@@ -98,5 +120,4 @@ This project draws inspiration from [cusadi](https://github.com/se-hwan/cusadi),
 
 For questions, issues, or suggestions, please [open an issue](https://github.com/based-robotics/jaxadi/issues) on our GitHub repository.
 
-
 We hope JAXADI empowers your numerical computing and optimization tasks! Happy coding!

benchmarks/README.md

@@ -0,0 +1,31 @@
+# Benchmarking
+
+In order to evaluate the performance of the `jaxadi` library vs `cusadi` we have tried to reproduce the benchmarks from the `cusadi` library first and faced some issues with the proper `cuda` installation.
+
+![meme](https://preview.redd.it/explain-please-v0-ma2mz5wxftod1.jpeg?auto=webp&s=2b90dfa3b12e064f54333e1080b3dabbad914f48)
+
+Adding the complexity of setup of benchmarks of [cusadi](https://github.com/se-hwan/cusadi) we have copied and modified the benchmarks to be able to run them in the [`colab` environment](https://colab.research.google.com/github/based-robotics/jaxadi/blob/feature%2Fbenchmarking/benchmarks/jaxadi_vs_cusadi.ipynb) side by side (CUDA vs Jax).
+
+Due limitations we cover only the functions with less than 1e3 operations. All of them are located in the [data](data) directory.
+
+One may run the benchmarks in the colab environment and get the [cusadi results](cuda_benchmark_results.npz) and [jaxadi results](jax_benchmark_results.npz) for comparison.
+
+#### Overview
+
+This analysis compares the performance of two computational tools, **Jaxadi** and **Cusadi**, using a range of batch sizes for different operations. The tools are evaluated based on average computation times and their relative speedup across various workloads.
+
+![1e1](https://github.com/based-robotics/jaxadi/blob/master/docs/static/images/compare_1e1_bar.png?raw=true)
+
+![1e2](https://github.com/based-robotics/jaxadi/blob/master/docs/static/images/compare_1e2_bar.png?raw=true)
+
+![Speedup](https://github.com/based-robotics/jaxadi/blob/master/docs/static/images/speedup_ratio.png?raw=true)
+
+#### Key Findings
+
+- **Jaxadi's Consistent Advantage:** Across all tested scenarios, Jaxadi consistently shows better performance than Cusadi. This is evident from the lower average computation times for Jaxadi, particularly as the workload increases.
+
+- **Scalability Differences:** As batch sizes grow, the gap between the two tools becomes more pronounced. Cusadi's computation time increases rapidly with larger batch sizes, especially beyond \(2^{12}\). In contrast, Jaxadi maintains a more stable performance with a much slower increase in computation time, indicating better batch scalability.
+
+#### Conclusion
+
+Overall, Jaxadi offers a clear performance advantage over Cusadi, particularly for small functions which can be translated propely.

benchmarks/data/fn_1e1.casadi

@@ -0,0 +1 @@
+jhpnnagiieahaaaadaaaaaaaaaaaaaaaaafaegaakaaaaaaadfifgefhogdgehjgpgogcaaaaaaagaaaaaaaggogpfbdfgbdaaaabaaabagaaaaaaaeaaaaaaaaaaaaaaababababadaaaaaaaaaaaaaaabababaeaaaaaaaaaaaaaaaegiaaaaaaaaaaaaaaaeaaaaaaaaaaaaaaabaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaeaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaabaaaaaaaaaaaaaaacaaaaaaaaaaaaaaadaaaaaaaaaaaaaaaeggaaaaaaaaaaaaaaacaaaaaaaaaaaaaaabaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaacaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaabaaaaaaaaaaaaaaachbaaaaaaaaaaaaaaaegfaaaaaaaaaaaaaaabaaaaaaaaaaaaaaabaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaabaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaadaaaaaaaaaaaaaaachbaaaaaaaaaaaaaaaegpaaaaaaaaaaaaaaadaaaaaaaaaaaaaaadaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaadaaaaaaaaaaaaaaagaaaaaaaaaaaaaaajaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaabaaaaaaaaaaaaaaacaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaabaaaaaaaaaaaaaaacaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaabaaaaaaaaaaaaaaacaaaaaaaaaaaaaaaeghaaaaaaaaaaaaaaadaaaaaaaaaaaaaaabaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaadaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaabaaaaaaaaaaaaaaacaaaaaaaaaaaaaaaeaaaaaaaaaaaaaaacaaaaaaajgadcaaaaaaajgbdcaaaaaaajgcdcaaaaaaajgdddaaaaaaaaaaaaaaacaaaaaaapgadcaaaaaaapgbdcaaaaaaapgcdaabagaaaaaaadhpgfhchdgfgbahaaaaaaakgjgehpfehngahaaaaaaaaaaaaaaaafaaaaaaadhigfgmgmgaaaaaaaaaaaaaaaaaaegbaaaaaaaaaaaaaaaaebababaaabababaaapbfilobfilobfnpdmfpicmfpicmfpnpdaaaaaeaaaaaaaaaaaaaabakdmiadcooijhfeodaaaaaaaaaaaaabhcaaaaaaaaaaaaaaaabaaaaaaaocdaaaaaaangehihaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaachfaaaaaaaaaaaaaaaeaaaaaaaaaaaaaaadaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaacbaaaaaaaaaaaaaabaaaaaaaeaaaaaaaaaaaaaaachaaaaaaaaaaaaaaaaeaaaaaaaaaaaaaaaegpcaaaaaaaaaaaaaadaaaaaaaiepfadegpcaaaaaaaaaaaaaadaaaaaaaiepfbdegpcaaaaaaaaaaaaaadaaaaaaaiepfcdegpcaaaaaaaaaaaaaadaaaaaaaiepfddchbaaaaaaaaaaaaaaacaaaaaaaaaaaaaaaegpcaaaaaaaaaaaaaadaaaaaaadgpfadegpcaaaaaaaaaaaaaadaaaaaaadgpfbdchbaaaaaaaaaaaaaaacaaaaaaaaaaaaaaaegpcaaaaaaaaaaaaaadaaaaaaabepfadegpcaaaaaaaaaaaaaadaaaaaaabepfbdchcaaaaaaaaaaaaaaabaaaaaaaaaaaaaaaegpcaaaaaaaaaaaaaabaaaaaaacgbaaaaaaahdaaaaaaaaaaaaaacbaaaaaaaaaaaaaaaaaaaaaaaaaaaaaapbaaaaaaaaaaaaaaegbaaaaaaaaaaaaaaachgaaaaaaaaaaaaaaaegmcaaaaaaaaaaaaaachllnlhnjnpnmhlnndegeaaaaaaaaaaaaaaachkaaaaaaaaaaaaaaachabaaaaaaaaaaaaaaegeaaaaaaaaaaaaaaachmaaaaaaaaaaaaaaachabaaaaaaaaaaaaaaegdaaaaaaaaaaaaaaachcbaaaaaaaaaaaaaachkaaaaaaaaaaaaaaaegbaaaaaaaaaaaaaaachoaaaaaaaaaaaaaaachdbaaaaaaaaaaaaaaegeaaaaaaaaaaaaaaachiaaaaaaaaaaaaaaachabaaaaaaaaaaaaaaegdaaaaaaaaaaaaaaachfbaaaaaaaaaaaaaachmaaaaaaaaaaaaaaaegcaaaaaaaaaaaaaaachnaaaaaaaaaaaaaaachgbaaaaaaaaaaaaaaegbaaaaaaaaaaaaaaachjaaaaaaaaaaaaaaachpaaaaaaaaaaaaaaaegdaaaaaaaaaaaaaaachiaaaaaaaaaaaaaaachfbaaaaaaaaaaaaaaegcaaaaaaaaaaaaaaachibaaaaaaaaaaaaaachjbaaaaaaaaaaaaaaegeaaaaaaaaaaaaaaachhbaaaaaaaaaaaaaachkbaaaaaaaaaaaaaaegdaaaaaaaaaaaaaaachfbaaaaaaaaaaaaaachkaaaaaaaaaaaaaaaegcaaaaaaaaaaaaaaachmbaaaaaaaaaaaaaachlaaaaaaaaaaaaaaaegdaaaaaaaaaaaaaaachlbaaaaaaaaaaaaaachnbaaaaaaaaaaaaaaegcaaaaaaaaaaaaaaachebaaaaaaaaaaaaaachobaaaaaaaaaaaaaaegdaaaaaaaaaaaaaaachmaaaaaaaaaaaaaaachcbaaaaaaaaaaaaaaegcaaaaaaaaaaaaaaachpaaaaaaaaaaaaaaachacaaaaaaaaaaaaaaegdaaaaaaaaaaaaaaachhbaaaaaaaaaaaaaachlbaaaaaaaaaaaaaaegcaaaaaaaaaaaaaaachbcaaaaaaaaaaaaaachccaaaaaaaaaaaaaaegeaaaaaaaaaaaaaaachpbaaaaaaaaaaaaaachdcaaaaaaaaaaaaaaegdaaaaaaaaaaaaaaachcbaaaaaaaaaaaaaachecaaaaaaaaaaaaaaegbaaaaaaaaaaaaaaachbbaaaaaaaaaaaaaachfcaaaaaaaaaaaaaaegeaaaaaaaaaaaaaaachnbaaaaaaaaaaaaaachkbaaaaaaaaaaaaaaegdaaaaaaaaaaaaaaachlbaaaaaaaaaaaaaachecaaaaaaaaaaaaaaegcaaaaaaaaaaaaaaachhcaaaaaaaaaaaaaachicaaaaaaaaaaaaaaegdaaaaaaaaaaaaaaachfbaaaaaaaaaaaaaachjcaaaaaaaaaaaaaaegbaaaaaaaaaaaaaaachgcaaaaaaaaaaaaaachkcaaaaaaaaaaaaaaegfaaaaaaaaaaaaaaachlcaaaaaaaaaaaaaaegfaaaaaaaaaaaaaaachkaaaaaaaaaaaaaaaegfaaaaaaaaaaaaaaachlaaaaaaaaaaaaaaabaaaaaaaaaaaaaaachpaaaaaaaaaaaaaaaeaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaancaaaaaaaaaaaaaabaaaaaaaaaaaaaaancaaaaaabaaaaaaaaaaaaaaaaaaaaaaamcaaaaaacaaaaaaallnlhnjnpnmhlnndbaaaaaaabaaaaaaabaaaaaaacaaaaaaaeaaaaaaadaaaaaaaaaaaaaaabaaaaaaancaaaaaaeaaaaaaacaaaaaaaaaaaaaaaeaaaaaaafaaaaaaaeaaaaaaabaaaaaaancaaaaaagaaaaaaadaaaaaaaaaaaaaaadaaaaaaahaaaaaaafaaaaaaaaaaaaaaabaaaaaaahaaaaaaagaaaaaaahaaaaaaancaaaaaaiaaaaaaacaaaaaaabaaaaaaancaaaaaajaaaaaaaaaaaaaaacaaaaaaaeaaaaaaakaaaaaaajaaaaaaabaaaaaaadaaaaaaalaaaaaaakaaaaaaaeaaaaaaacaaaaaaalaaaaaaaiaaaaaaalaaaaaaancaaaaaamaaaaaaaaaaaaaaadaaaaaaabaaaaaaamaaaaaaamaaaaaaacaaaaaaadaaaaaaanaaaaaaajaaaaaaakaaaaaaacaaaaaaanaaaaaaamaaaaaaanaaaaaaaeaaaaaaaoaaaaaaalaaaaaaanaaaaaaadaaaaaaapaaaaaaakaaaaaaaaaaaaaaancaaaaaaabaaaaaabaaaaaaabaaaaaaacaaaaaaapaaaaaaapaaaaaaaabaaaaaadaaaaaaabbaaaaaaoaaaaaaapaaaaaaacaaaaaaahaaaaaaahaaaaaaabbaaaaaadaaaaaaabbaaaaaaeaaaaaaafaaaaaaacaaaaaaabbaaaaaacaaaaaaabbaaaaaadaaaaaaalaaaaaaalaaaaaaaoaaaaaaacaaaaaaabbaaaaaabbaaaaaalaaaaaaaeaaaaaaahaaaaaaahaaaaaaabbaaaaaadaaaaaaafaaaaaaafaaaaaaahaaaaaaabaaaaaaadaaaaaaadaaaaaaafaaaaaaaeaaaaaaapaaaaaaapaaaaaaanaaaaaaadaaaaaaaoaaaaaaaoaaaaaaahaaaaaaacaaaaaaapaaaaaaapaaaaaaaoaaaaaaadaaaaaaakaaaaaaakaaaaaaapaaaaaaabaaaaaaadaaaaaaadaaaaaaakaaaaaaafaaaaaaadaaaaaaadaaaaaaadaaaaaaaocaaaaaaaaaaaaaadaaaaaaaaaaaaaaaocaaaaaaaaaaaaaapaaaaaaabaaaaaaaocaaaaaabaaaaaaabaaaaaaaaaaaaaaancaaaaaabaaaaaaaaaaaaaaabaaaaaaaocaaaaaabaaaaaaabaaaaaaabaaaaaaaocaaaaaabaaaaaaaeaaaaaaacaaaaaaaocaaaaaabaaaaaaajaaaaaaadaaaaaaaocaaaaaabaaaaaaamaaaaaaaeaaaaaaaocaaaaaabaaaaaaaiaaaaaaafaaaaaaaocaaaaaabaaaaaaaeaaaaaaagaaaaaaaocaaaaaabaaaaaaaiaaaaaaahaaaaaaaocaaaaaabaaaaaaacaaaaaaaiaaaaaaafaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaocaaaaaacaaaaaaaaaaaaaaaaaaaaaaafaaaaaaaabaaaaaaabaaaaaaabaaaaaaocaaaaaacaaaaaaaabaaaaaabaaaaaaaocaaaaaacaaaaaaagaaaaaaacaaaaaaabadaaaaaaaaaaaaaaachbaaaaaaaaaaaaaaacaaaaaaaaaaaaaaachmcaaaaaaaaaaaaaachjcaaaaaaaaaaaaaachdaaaaaaaaaaaaaaajaaaaaaaaaaaaaaachabaaaaaaaaaaaaaachhaaaaaaaaaaaaaaachmaaaaaaaaaaaaaaachiaaaaaaaaaaaaaaachibaaaaaaaaaaaaaachnaaaaaaaaaaaaaaachmaaaaaaaaaaaaaaachnaaaaaaaaaaaaaaachpaaaaaaaaaaaaaaacheaaaaaaaaaaaaaaadaaaaaaaaaaaaaaachncaaaaaaaaaaaaaachocaaaaaaaaaaaaaachoaaaaaaaaaaaaaaa