Coverage for src/uqtils/gradient.py: 89%

1"""Module for vectorized finite-difference gradient approximations. 

2 

3Includes: 

4 

5- `approx_jac` - vectorized Jacobian approximation 

6- `approx_hess` - vectorized Hessian approximation 

7""" 

8import numpy as np 

9 

10from .uq_types import Array 

11 

12__all__ = ['approx_jac', 'approx_hess'] 

13 

14 

15def approx_jac(func, theta: Array, pert=0.01) -> np.ndarray: 

16 """Approximate Jacobian of `func` at a specified `theta` location using finite difference approximation. 

17 

18 :param func: expects to be called as `func(theta) -> (..., y_dim)` 

19 :param theta: `(..., theta_dim)`, points to linearize model about 

20 :param pert: perturbation percent for approximate partial derivatives 

21 :returns J: `(..., y_dim, theta_dim)`, the approximate Jacobian `(y_dim, theta_dim)` at all locations `(...)` 

22 """ 

23 theta = np.atleast_1d(theta) 

24 shape = theta.shape[:-1] # (...,) 

25 theta_dim = theta.shape[-1] # Number of parameters 

26 dtheta = pert * np.abs(theta) 

27 

28 # Make sure dtheta is not 0 anywhere 

29 for i in range(theta_dim): 

30 zero_idx = np.isclose(dtheta[..., i], 0) 

31 if np.any(zero_idx): 

32 subs_dtheta = pert * np.abs(np.mean(theta[..., i])) 

33 if np.isclose(subs_dtheta, 0): 

34 subs_dtheta = pert 

35 dtheta[zero_idx, i] = subs_dtheta 

36 

37 # Return the Jacobians (..., y_dim, theta_dim) 

38 J, y_dim = None, None 

39 

40 for i in range(theta_dim): 

41 theta_n1 = np.copy(theta) 

42 theta_p1 = np.copy(theta) 

43 

44 # Perturbations to theta 

45 theta_n1[..., i] -= dtheta[..., i] 

46 theta_p1[..., i] += dtheta[..., i] 

47 f_n1 = func(theta_n1) 

48 f_p1 = func(theta_p1) 

49 

50 if J is None: 

51 y_dim = f_p1.shape[-1] 

52 J = np.empty(shape + (y_dim, theta_dim)) 

53 

54 J[..., i] = (f_p1 - f_n1) / np.expand_dims(2 * dtheta[..., i], axis=-1) 

55 

56 if y_dim == 1: 

57 J = np.squeeze(J, axis=-2) 

58 if theta_dim == 1: 

59 J = np.squeeze(J, axis=-1) 

60 return np.atleast_1d(J) 

61 

62 

63def approx_hess(func, theta: Array, pert=0.01) -> np.ndarray: 

64 """Approximate Hessian of `func` at a specified `theta` location using finite difference approximation. 

65 

66 :param func: expects to be called as `func(theta) -> (..., y_dim)` 

67 :param theta: `(..., theta_dim)`, points to linearize model about 

68 :param pert: perturbation percent for approximate partial derivatives 

69 :returns H: `(..., y_dim, theta_dim, theta_dim)`, the approximate Hessian `(theta_dim, theta_dim)` at all locations 

70 `(...,)` for vector-valued function of dimension `y_dim` 

71 """ 

72 theta = np.atleast_1d(theta) 

73 shape = theta.shape[:-1] # (...,) 

74 theta_dim = theta.shape[-1] # Number of parameters 

75 dtheta = pert * np.abs(theta) 

76 

77 # Make sure dtheta is not 0 anywhere 

78 for i in range(theta_dim): 

79 zero_idx = np.isclose(dtheta[..., i], 0) 

80 if np.any(zero_idx): 

81 subs_dtheta = pert * np.abs(np.mean(theta[..., i])) 

82 if np.isclose(subs_dtheta, 0): 

83 subs_dtheta = pert 

84 dtheta[zero_idx, i] = subs_dtheta 

85 

86 # Return the Hessians (..., y_dim, theta_dim, theta_dim) 

87 y_dim, H = None, None 

88 

89 for i in range(theta_dim): 

90 for j in range(i, theta_dim): 

91 # Allocate space at 4 grid points (n1=-1, p1=+1) 

92 theta_n1_n1 = np.copy(theta) 

93 theta_p1_p1 = np.copy(theta) 

94 theta_n1_p1 = np.copy(theta) 

95 theta_p1_n1 = np.copy(theta) 

96 

97 # Perturbations to theta in each direction 

98 theta_n1_n1[..., i] -= dtheta[..., i] 

99 theta_n1_n1[..., j] -= dtheta[..., j] 

100 f_n1_n1 = func(theta_n1_n1) 

101 

102 theta_p1_p1[..., i] += dtheta[..., i] 

103 theta_p1_p1[..., j] += dtheta[..., j] 

104 f_p1_p1 = func(theta_p1_p1) 

105 

106 theta_n1_p1[..., i] -= dtheta[..., i] 

107 theta_n1_p1[..., j] += dtheta[..., j] 

108 f_n1_p1 = func(theta_n1_p1) 

109 

110 theta_p1_n1[..., i] += dtheta[..., i] 

111 theta_p1_n1[..., j] -= dtheta[..., j] 

112 f_p1_n1 = func(theta_p1_n1) 

113 

114 if H is None: 

115 y_dim = f_p1_n1.shape[-1] 

116 H = np.empty(shape + (y_dim, theta_dim, theta_dim)) 

117 

118 res = (f_n1_n1 + f_p1_p1 - f_n1_p1 - f_p1_n1) / np.expand_dims(4 * dtheta[..., i] * dtheta[..., j], 

119 axis=-1) 

120 H[..., i, j] = res 

121 H[..., j, i] = res 

122 

123 if y_dim == 1: 

124 H = np.squeeze(H, axis=-3) 

125 if theta_dim == 1: 

126 H = np.squeeze(H, axis=(-1, -2)) 

127 return np.atleast_1d(H)