Make PyBaMM importable in Linux systems if jax is already installed

[Saransh-cpp]

Description

• Check for the version of jax in have_jax()
• Give an error in install_jax() if jax is already available with a different version

Additionally, the cell with pybamm.print_citations() gives an error which is described here - pybamm-team/BattBot#47 - and is most probably an error within the pybtex library.

Fixes # (issue)

Type of change

Please add a line in the relevant section of CHANGELOG.md to document the change (include PR #) - note reverse order of PR #s. If necessary, also add to the list of breaking changes.

• New feature (non-breaking change which adds functionality)
• Optimization (back-end change that speeds up the code)
• Bug fix (non-breaking change which fixes an issue)

Key checklist:

• No style issues: $ flake8
• All tests pass: $ python run-tests.py --unit
• The documentation builds: $ cd docs and then $ make clean; make html

You can run all three at once, using $ python run-tests.py --quick.

Further checks:

• Code is commented, particularly in hard-to-understand areas
• Tests added that prove fix is effective or that feature works

[codecov]

Codecov Report

Merging #1874 (ceb47d0) into develop (ccf5be7) will increase coverage by 0.02%.
The diff coverage is 100.00%.

@@             Coverage Diff             @@
##           develop    #1874      +/-   ##
===========================================
+ Coverage    99.26%   99.28%   +0.02%     
===========================================
  Files          345      345              
  Lines        19107    19106       -1     
===========================================
+ Hits         18966    18970       +4     
+ Misses         141      136       -5     

Impacted Files	Coverage Δ
pybamm/__init__.py	94.95% <100.00%> (ø)
pybamm/util.py	100.00% <100.00%> (+2.87%)	⬆️

---

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 ccf5be7...ceb47d0. Read the comment docs.

[brosaplanella]

Looks good, thanks Saransh! For a more robust fix, would it be possible to check if the jax version is the correct one and if not, throw a warning and treat the code as if jax wasn't installed?

[Saransh-cpp]

Looks good, thanks Saransh! For a more robust fix, would it be possible to check if the jax version is the correct one and if not, throw a warning and treat the code as if jax wasn't installed?

I was trying to uninstall the existing version of jax while installing pybamm through a function but it felt weird editing a user's environment. The warning sounds really good! I'll try adding that instead of editing the notebooks.

[Saransh-cpp]

Codecov shows that this complete function is not included in the coverage, should I add # pragma: no cover for this function?

[brosaplanella]

That sounds good

[brosaplanella]

Also, this is now an error. Should we make it a warning and proceed without jax?

[Saransh-cpp]

PyBaMM installation (and importing PyBaMM) proceeds without jax if the installed jax version is not compatible. But this function throws an error if a user tries to install jax and has an incompatible jax version already installed.

Should it instead uninstall the existing version and install the compatible version?

[brosaplanella]

Sorry, I misunderstood! That works as it is :)

[priyanshuone6]

We should create a global variable for jax version, jaxlib version and use it in install_jax, is_jax_compatible

[brosaplanella]

Looks great now, I will wait until tomorrow in case someone else wants to give it a look.

[brosaplanella]

@all-contributors please add @Saransh-cpp for maintenance

[allcontributors]

@brosaplanella

I've put up a pull request to add @Saransh-cpp! 🎉