Book Image

Applying Math with Python - Second Edition

By : Sam Morley
Book Image

Applying Math with Python - Second Edition

By: Sam Morley

Overview of this book

The updated edition of Applying Math with Python will help you solve complex problems in a wide variety of mathematical fields in simple and efficient ways. Old recipes have been revised for new libraries and several recipes have been added to demonstrate new tools such as JAX. You'll start by refreshing your knowledge of several core mathematical fields and learn about packages covered in Python's scientific stack, including NumPy, SciPy, and Matplotlib. As you progress, you'll gradually get to grips with more advanced topics of calculus, probability, and networks (graph theory). Once you’ve developed a solid base in these topics, you’ll have the confidence to set out on math adventures with Python as you explore Python's applications in data science and statistics, forecasting, geometry, and optimization. The final chapters will take you through a collection of miscellaneous problems, including working with specific data formats and accelerating code. By the end of this book, you'll have an arsenal of practical coding solutions that can be used and modified to solve a wide range of practical problems in computational mathematics and data science.
Table of Contents (13 chapters)

Automatic differentiation and calculus using JAX

JAX is a linear algebra and automatic differentiation framework developed by Google for ML. It combines the capabilities of Autograd and its Accelerated Linear Algebra (XLA) optimizing compiler for linear algebra and ML. In particular, it allows us to easily construct complex functions, with automatic gradient computation, that can be run on Graphics Processing Units (GPUs) or Tensor Processing Units (TPUs). On top of all of this, it is relatively simple to use. In this recipe, we see how to make use of the JAX just-in-time (JIT) compiler, get the gradient of a function, and make use of different computation devices.

Getting ready

For this recipe, we need the JAX package installed. We will make use of the Matplotlib package, with the pyplot interface imported as plt as usual. Since we’re going to plot a function of two variables, we also need to import the mplot3d module from the mpl_toolkits package.

How to do it…...