1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
|
# Eager, forward-mode autodiff (autograd)
# Backpropagation will probably require collecting a DAG with a typetracer or Dask
#
# Presented at https://indico.cern.ch/event/1387764/
#
# The following are good references:
#
# https://www.hedonisticlearning.com/posts/complex-step-differentiation.html
# https://researchrepository.wvu.edu/faculty_publications/426/
import numpy as np
from numpy.lib.mixins import NDArrayOperatorsMixin
class diffarray(NDArrayOperatorsMixin):
__slots__ = ("_array",)
@classmethod
def _build(cls, complex_array):
"Manual constructor from a `complex_array`."
self = cls.__new__(cls)
self._array = complex_array
return self
def __init__(self, primal, tangent=None, *, dtype=None):
"Constructor for floating-point `primal` and (optional) `tangent`."
if dtype is None:
dtype = primal.dtype.type
elif isinstance(dtype, np.dtype):
dtype = dtype.type
if issubclass(dtype, np.float32):
self._array = primal.astype(np.complex64)
elif issubclass(dtype, np.float64):
self._array = primal.astype(np.complex128)
else:
raise TypeError("only float32 or float64 arrays can be differentiated")
self._array += (1 if tangent is None else tangent) * 1j * self._step_scale
@property
def _step_scale(self):
"Size of the complex step; half precision of 1.0."
return 1e-4 if issubclass(self._array.dtype.type, np.complex128) else 1e-8
@property
def primal(self):
"Array of primary values."
return np.real(self._array)
@property
def tangent(self):
"Array of derivatives."
return np.imag(self._array) / self._step_scale
def __str__(self):
primal = str(self.primal).replace("\n", "\n ")
tangent = str(self.tangent).replace("\n", "\n ")
return f"primal: {primal}\ntangent: {tangent}"
def __repr__(self):
primal = str(self.primal).replace("\n", "\n ")
tangent = str(self.tangent).replace("\n", "\n ")
dtype = ""
if issubclass(self._array.dtype.type, np.complex64):
dtype = ",\n dtype=np.float32"
return f"diffarray({primal},\n {tangent}{dtype})"
def _prepare(self, args, kwargs):
"Used in NEP-13 and NEP-18 overrides."
cls = type(self)
args = [x._array if isinstance(x, cls) else x for x in args]
kwargs = {k: v._array if isinstance(x, cls) else v for k, v in kwargs.items()}
return cls, args, kwargs
def __array_ufunc__(self, ufunc, method, *args, **kwargs):
"https://numpy.org/neps/nep-0013-ufunc-overrides.html"
if ufunc.__name__ == "absolute":
# interpret `absolute` only on the primal
if len(kwargs) != 0:
raise NotImplementedError("kwargs in np.absolute")
arg = args[0]._array
out = arg.copy()
out[arg.real < 0] *= -1
return type(self)._build(out)
if ufunc.__name__ in (
"less", "less_equal", "equal", "not_equal", "greater", "greater_equal"
):
# do comparisons only on the primal
cls = type(self)
args = [x._array.real if isinstance(x, cls) else x for x in args]
return getattr(ufunc, method)(*args, **kwargs)
cls, prepared_args, prepared_kwargs = self._prepare(args, kwargs)
out = getattr(ufunc, method)(*prepared_args, **prepared_kwargs)
if issubclass(out.dtype.type, np.complexfloating):
return cls._build(out)
else:
return out
def __array_function__(self, func, types, args, kwargs):
"https://numpy.org/neps/nep-0018-array-function-protocol.html"
if func.__name__ == "real":
# interpret `real` only on the primal
return type(self)._build(args[0]._array)
if func.__name__ == "imag":
# interpret `imag` only on the primal
return type(self)._build(args[0]._array * 0)
cls, prepared_args, prepared_kwargs = self._prepare(args, kwargs)
out = func(*prepared_args, **prepared_kwargs)
if issubclass(out.dtype.type, np.complexfloating):
return cls._build(out)
else:
return out
def __getitem__(self, where):
out = self._array[where]
if isinstance(out, np.complexfloating):
# NumPy returns a scalar; CuPy and Array API return an array
# we return an array to keep derivatives
return type(self)._build(np.asarray(out))
return out
# >>> x = np.linspace(-20, 20, 10000)
# >>> da_x = diffarray(x)
# >>> da_y = np.sin(da_x) / da_x
# >>> da_x
# diffarray([-20. -19.9959996 -19.9919992 ... 19.9919992 19.9959996
# 20. ],
# [1. 1. 1. ... 1. 1. 1.])
# >>> da_y
# diffarray([0.04564726 0.04557439 0.04550076 ... 0.04550076 0.04557439 0.04564726],
# [-0.01812174 -0.01831149 -0.01850102 ... 0.01850102 0.01831149
# 0.01812174])
# >>> abs(da_y.tangent - ((x*np.cos(x) - np.sin(x)) / x**2)).max()
# 3.9683650809863025e-10
# >>> import matplotlib.pyplot as plt
# >>> plt.plot(x, da_y.tangent)
# >>> plt.plot(x, (x*np.cos(x) - np.sin(x)) / x**2, ls="--")
#
# See https://gist.github.com/jpivarski/8dc48a87bae7a856848f87e36b9d244d for the plot
|
