Compute the Jacobian matrix in Python

Question

import numpy as np


a = np.array([[1,2,3],
              [4,5,6],
              [7,8,9]])


b = np.array([[1,2,3]]).T

c = a.dot(b) #function

jacobian = a # as partial         


        

Answer 1

If you want to find the Jacobian numerically for many points at once (for example, if your function accepts shape (n, x) and outputs (n, y)), here is a function. This is essentially the answer from James Carter but for many points. The dx may need to be adjusted based on absolute value as in his answer.

def numerical_jacobian(f, xs, dx=1e-6):
  """
      f is a function that accepts input of shape (n_points, input_dim)
      and outputs (n_points, output_dim)

      return the jacobian as (n_points, output_dim, input_dim)
  """
  if len(xs.shape) == 1:
    xs = xs[np.newaxis, :]
    
  assert len(xs.shape) == 2

  ys = f(xs)
  
  x_dim = xs.shape[1]
  y_dim = ys.shape[1]
  
  jac = np.empty((xs.shape[0], y_dim, x_dim))
  
  for i in range(x_dim):
    x_try = xs + dx * e(x_dim, i + 1)
    jac[:, :, i] = (f(x_try) - ys) / dx
  
  return jac

def e(n, i):
  ret = np.zeros(n)
  ret[i - 1] = 1.0
  return ret