Make jax & odes optional

[arjxn-py]

Description

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

Fixes #3146

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: $ pre-commit run (see CONTRIBUTING.md for how to set this up to run automatically when committing locally, in just two lines of code)
• All tests pass: $ python run-tests.py --all
• The documentation builds: $ python run-tests.py --doctest

You can run unit and doctests together 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

Patch coverage: 100.00% and no project coverage change.

Comparison is base (f588f20) 99.71% compared to head (fe704bf) 99.71%.
Report is 56 commits behind head on develop.

Additional details and impacted files

@@           Coverage Diff            @@
##           develop    #3163   +/-   ##
========================================
  Coverage    99.71%   99.71%           
========================================
  Files          248      248           
  Lines        18761    18764    +3     
========================================
+ Hits         18707    18710    +3     
  Misses          54       54           

Files Changed	Coverage Δ
pybamm/util.py	100.00% <100.00%> (ø)

... and 2 files with indirect coverage changes

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[Saransh-cpp]

Thanks, @arjxn-py! This would also require updating the docs and getting rid of -

PyBaMM/pybamm/util.py

Lines 295 to 341 in 02c56a2

	def install_jax(arguments=None): # pragma: no cover
	"""
	Install compatible versions of jax, jaxlib.
	Command Line Interface::
	$ pybamm_install_jax
	| optional arguments:
	| -h, --help show help message
	| -f, --force force install compatible versions of jax and jaxlib
	"""
	parser = argparse.ArgumentParser(description="Install jax and jaxlib")
	parser.add_argument(
	"-f",
	"--force",
	action="store_true",
	help="force install compatible versions of"
	f" jax ({JAX_VERSION}) and jaxlib ({JAXLIB_VERSION})",
	)
	args = parser.parse_args(arguments)
	if system() == "Windows":
	raise NotImplementedError("Jax is not available on Windows")
	# Raise an error if jax and jaxlib are already installed, but incompatible
	# and --force is not set
	elif importlib.util.find_spec("jax") is not None:
	if not args.force and not is_jax_compatible():
	raise ValueError(
	"Jax is already installed but the installed version of jax or jaxlib is"
	" not supported by PyBaMM. \nYou can force install compatible versions"
	f" of jax ({JAX_VERSION}) and jaxlib ({JAXLIB_VERSION}) using the"
	" following command: \npybamm_install_jax --force"
	)
	subprocess.check_call(
	[
	sys.executable,
	"-m",
	"pip",
	"install",
	f"jax=={JAX_VERSION}",
	f"jaxlib=={JAXLIB_VERSION}",
	]
	)

from both, code and docs.

[arjxn-py]

This would also require updating the docs and getting rid from both - code and docs.

@Saransh-cpp do we also want to deprecate pybamm_install_odes?

[Saransh-cpp]

Great work, thanks @arjxn-py! Some comments below.

[arjxn-py]

Deprecation warning looks like this :

Adding note for odes before re-requesting review.

[Saransh-cpp]

Nitpicking 🙂

[Saransh-cpp]

Thanks, @arjxn-py!