authorgravatar for 101306295+PauloCampana@users.noreply.github.comPauloCampana <101306295+PauloCampana@users.noreply.github.com> 2024-01-09 19:53:47-03:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2024-01-09 14:53:47-08:00
logfcc0c5ddc7a8d25f0b12f51b4b244bc97fc549e4
tree6e1aad3bcd2f12f61dfb6bb6f0be3523d5d45041
parent7320b24e0ea768cf8dffaaeaa058db8e9134b444
signaturebadge-question-mark Signed by PGP key 4AEE18F83AFDEB23

Add std.math.gamma and lgamma (#18449)

See #7212

2 files changed, 332 insertions(+), 0 deletions(-)

lib/std/math.zig+4
...@@ -269,6 +269,8 @@ pub const sinh = @import("math/sinh.zig").sinh;...@@ -269,6 +269,8 @@ pub const sinh = @import("math/sinh.zig").sinh;
269pub const cosh = @import("math/cosh.zig").cosh;269pub const cosh = @import("math/cosh.zig").cosh;
270pub const tanh = @import("math/tanh.zig").tanh;270pub const tanh = @import("math/tanh.zig").tanh;
271pub const gcd = @import("math/gcd.zig").gcd;271pub const gcd = @import("math/gcd.zig").gcd;
272pub const gamma = @import("math/gamma.zig").gamma;
273pub const lgamma = @import("math/gamma.zig").lgamma;
272274
273/// Sine trigonometric function on a floating point number.275/// Sine trigonometric function on a floating point number.
274/// Uses a dedicated hardware instruction when available.276/// Uses a dedicated hardware instruction when available.
...@@ -392,6 +394,8 @@ test {...@@ -392,6 +394,8 @@ test {
392 _ = cosh;394 _ = cosh;
393 _ = tanh;395 _ = tanh;
394 _ = gcd;396 _ = gcd;
397 _ = gamma;
398 _ = lgamma;
395399
396 _ = complex;400 _ = complex;
397 _ = Complex;401 _ = Complex;
lib/std/math/gamma.zig created+328
...@@ -0,0 +1,328 @@
1// Ported from musl, which is licensed under the MIT license:
2// https://git.musl-libc.org/cgit/musl/tree/COPYRIGHT
3//
4// https://git.musl-libc.org/cgit/musl/tree/src/math/tgamma.c
5
6const std = @import("../std.zig");
7
8/// Returns the gamma function of x,
9/// gamma(x) = factorial(x - 1) for integer x.
10///
11/// Special Cases:
12/// - gamma(+-nan) = nan
13/// - gamma(-inf) = nan
14/// - gamma(n) = nan for negative integers
15/// - gamma(-0.0) = -inf
16/// - gamma(+0.0) = +inf
17/// - gamma(+inf) = +inf
18pub fn gamma(comptime T: type, x: T) T {
19 if (T != f32 and T != f64) {
20 @compileError("gamma not implemented for " ++ @typeName(T));
21 }
22 // common integer case first
23 if (x == @trunc(x)) {
24 // gamma(-inf) = nan
25 // gamma(n) = nan for negative integers
26 if (x < 0) {
27 return std.math.nan(T);
28 }
29 // gamma(-0.0) = -inf
30 // gamma(+0.0) = +inf
31 if (x == 0) {
32 return 1 / x;
33 }
34 if (x < integer_result_table.len) {
35 const i = @as(u8, @intFromFloat(x));
36 return @floatCast(integer_result_table[i]);
37 }
38 }
39 // below this, result underflows, but has a sign
40 // negative for (-1, 0)
41 // positive for (-2, -1)
42 // negative for (-3, -2)
43 // ...
44 const lower_bound = if (T == f64) -184 else -42;
45 if (x < lower_bound) {
46 return if (@mod(x, 2) > 1) -0.0 else 0.0;
47 }
48 // above this, result overflows
49 // gamma(+inf) = +inf
50 const upper_bound = if (T == f64) 172 else 36;
51 if (x > upper_bound) {
52 return std.math.inf(T);
53 }
54
55 const abs = @abs(x);
56 // perfect precision here
57 if (abs < 0x1p-54) {
58 return 1 / x;
59 }
60
61 const base = abs + lanczos_minus_half;
62 const exponent = abs - 0.5;
63 // error of y for correction, see
64 // https://github.com/python/cpython/blob/5dc79e3d7f26a6a871a89ce3efc9f1bcee7bb447/Modules/mathmodule.c#L286-L324
65 const e = if (abs > lanczos_minus_half)
66 base - abs - lanczos_minus_half
67 else
68 base - lanczos_minus_half - abs;
69 const correction = lanczos * e / base;
70 const initial = series(T, abs) * @exp(-base);
71
72 // use reflection formula for negatives
73 if (x < 0) {
74 const reflected = -std.math.pi / (abs * sinpi(T, abs) * initial);
75 const corrected = reflected - reflected * correction;
76 const half_pow = std.math.pow(T, base, -0.5 * exponent);
77 return corrected * half_pow * half_pow;
78 } else {
79 const corrected = initial + initial * correction;
80 const half_pow = std.math.pow(T, base, 0.5 * exponent);
81 return corrected * half_pow * half_pow;
82 }
83}
84
85/// Returns the natural logarithm of the absolute value of the gamma function.
86///
87/// Special Cases:
88/// - lgamma(+-nan) = nan
89/// - lgamma(+-inf) = +inf
90/// - lgamma(n) = +inf for negative integers
91/// - lgamma(+-0.0) = +inf
92/// - lgamma(1) = +0.0
93/// - lgamma(2) = +0.0
94pub fn lgamma(comptime T: type, x: T) T {
95 if (T != f32 and T != f64) {
96 @compileError("gamma not implemented for " ++ @typeName(T));
97 }
98 // common integer case first
99 if (x == @trunc(x)) {
100 // lgamma(-inf) = +inf
101 // lgamma(n) = +inf for negative integers
102 // lgamma(+-0.0) = +inf
103 if (x <= 0) {
104 return std.math.inf(T);
105 }
106 // lgamma(1) = +0.0
107 // lgamma(2) = +0.0
108 if (x < integer_result_table.len) {
109 const i = @as(u8, @intFromFloat(x));
110 return @log(@as(T, @floatCast(integer_result_table[i])));
111 }
112 // lgamma(+inf) = +inf
113 if (std.math.isPositiveInf(x)) {
114 return x;
115 }
116 }
117
118 const abs = @abs(x);
119 // perfect precision here
120 if (abs < 0x1p-54) {
121 return -@log(abs);
122 }
123 // obvious approach when overflow is not a problem
124 const upper_bound = if (T == f64) 128 else 26;
125 if (abs < upper_bound) {
126 return @log(@abs(gamma(T, x)));
127 }
128
129 const log_base = @log(abs + lanczos_minus_half) - 1;
130 const exponent = abs - 0.5;
131 const log_series = @log(series(T, abs));
132 const initial = exponent * log_base + log_series - lanczos;
133
134 // use reflection formula for negatives
135 if (x < 0) {
136 const reflected = std.math.pi / (abs * sinpi(T, abs));
137 return @log(@abs(reflected)) - initial;
138 }
139 return initial;
140}
141
142// table of factorials for integer early return
143// stops at 22 because 23 isn't representable with full precision on f64
144const integer_result_table = [_]f64{
145 std.math.inf(f64), // gamma(+0.0)
146 1, // gamma(1)
147 1, // ...
148 2,
149 6,
150 24,
151 120,
152 720,
153 5040,
154 40320,
155 362880,
156 3628800,
157 39916800,
158 479001600,
159 6227020800,
160 87178291200,
161 1307674368000,
162 20922789888000,
163 355687428096000,
164 6402373705728000,
165 121645100408832000,
166 2432902008176640000,
167 51090942171709440000, // gamma(22)
168};
169
170// "g" constant, arbitrary
171const lanczos = 6.024680040776729583740234375;
172const lanczos_minus_half = lanczos - 0.5;
173
174fn series(comptime T: type, abs: T) T {
175 const numerator = [_]T{
176 23531376880.410759688572007674451636754734846804940,
177 42919803642.649098768957899047001988850926355848959,
178 35711959237.355668049440185451547166705960488635843,
179 17921034426.037209699919755754458931112671403265390,
180 6039542586.3520280050642916443072979210699388420708,
181 1439720407.3117216736632230727949123939715485786772,
182 248874557.86205415651146038641322942321632125127801,
183 31426415.585400194380614231628318205362874684987640,
184 2876370.6289353724412254090516208496135991145378768,
185 186056.26539522349504029498971604569928220784236328,
186 8071.6720023658162106380029022722506138218516325024,
187 210.82427775157934587250973392071336271166969580291,
188 2.5066282746310002701649081771338373386264310793408,
189 };
190 const denominator = [_]T{
191 0,
192 39916800,
193 120543840,
194 150917976,
195 105258076,
196 45995730,
197 13339535,
198 2637558,
199 357423,
200 32670,
201 1925,
202 66,
203 1,
204 };
205 var num: T = 0;
206 var den: T = 0;
207 // split to avoid overflow
208 if (abs < 8) {
209 // big abs would overflow here
210 for (0..numerator.len) |i| {
211 num = num * abs + numerator[numerator.len - 1 - i];
212 den = den * abs + denominator[numerator.len - 1 - i];
213 }
214 } else {
215 // small abs would overflow here
216 for (0..numerator.len) |i| {
217 num = num / abs + numerator[i];
218 den = den / abs + denominator[i];
219 }
220 }
221 return num / den;
222}
223
224// precise sin(pi * x)
225// but not for integer x or |x| < 2^-54, we handle those already
226fn sinpi(comptime T: type, x: T) T {
227 const xmod2 = @mod(x, 2); // [0, 2]
228 const n = (@as(u8, @intFromFloat(4 * xmod2)) + 1) / 2; // {0, 1, 2, 3, 4}
229 const y = xmod2 - 0.5 * @as(T, @floatFromInt(n)); // [-0.25, 0.25]
230 return switch (n) {
231 0, 4 => @sin(std.math.pi * y),
232 1 => @cos(std.math.pi * y),
233 2 => -@sin(std.math.pi * y),
234 3 => -@cos(std.math.pi * y),
235 else => unreachable,
236 };
237}
238
239const expect = std.testing.expect;
240const expectEqual = std.testing.expectEqual;
241const expectApproxEqRel = std.testing.expectApproxEqRel;
242
243test "math.gamma" {
244 inline for (&.{ f32, f64 }) |T| {
245 const eps = @sqrt(std.math.floatEps(T));
246 try expectApproxEqRel(@as(T, 120), gamma(T, 6), eps);
247 try expectApproxEqRel(@as(T, 362880), gamma(T, 10), eps);
248 try expectApproxEqRel(@as(T, 6402373705728000), gamma(T, 19), eps);
249
250 try expectApproxEqRel(@as(T, 332.7590766955334570), gamma(T, 0.003), eps);
251 try expectApproxEqRel(@as(T, 1.377260301981044573), gamma(T, 0.654), eps);
252 try expectApproxEqRel(@as(T, 1.025393882573518478), gamma(T, 0.959), eps);
253
254 try expectApproxEqRel(@as(T, 7.361898021467681690), gamma(T, 4.16), eps);
255 try expectApproxEqRel(@as(T, 198337.2940287730753), gamma(T, 9.73), eps);
256 try expectApproxEqRel(@as(T, 113718145797241.1666), gamma(T, 17.6), eps);
257
258 try expectApproxEqRel(@as(T, -1.13860211111081424930673), gamma(T, -2.80), eps);
259 try expectApproxEqRel(@as(T, 0.00018573407931875070158), gamma(T, -7.74), eps);
260 try expectApproxEqRel(@as(T, -0.00000001647990903942825), gamma(T, -12.1), eps);
261 }
262}
263
264test "math.gamma.special" {
265 inline for (&.{ f32, f64 }) |T| {
266 try expect(std.math.isNan(gamma(T, -std.math.nan(T))));
267 try expect(std.math.isNan(gamma(T, std.math.nan(T))));
268 try expect(std.math.isNan(gamma(T, -std.math.inf(T))));
269
270 try expect(std.math.isNan(gamma(T, -4)));
271 try expect(std.math.isNan(gamma(T, -11)));
272 try expect(std.math.isNan(gamma(T, -78)));
273
274 try expectEqual(-std.math.inf(T), gamma(T, -0.0));
275 try expectEqual(std.math.inf(T), gamma(T, 0.0));
276
277 try expect(std.math.isNegativeZero(gamma(T, -200.5)));
278 try expect(std.math.isPositiveZero(gamma(T, -201.5)));
279 try expect(std.math.isNegativeZero(gamma(T, -202.5)));
280
281 try expectEqual(std.math.inf(T), gamma(T, 200));
282 try expectEqual(std.math.inf(T), gamma(T, 201));
283 try expectEqual(std.math.inf(T), gamma(T, 202));
284
285 try expectEqual(std.math.inf(T), gamma(T, std.math.inf(T)));
286 }
287}
288
289test "math.lgamma" {
290 inline for (&.{ f32, f64 }) |T| {
291 const eps = @sqrt(std.math.floatEps(T));
292 try expectApproxEqRel(@as(T, @log(24.0)), lgamma(T, 5), eps);
293 try expectApproxEqRel(@as(T, @log(20922789888000.0)), lgamma(T, 17), eps);
294 try expectApproxEqRel(@as(T, @log(2432902008176640000.0)), lgamma(T, 21), eps);
295
296 try expectApproxEqRel(@as(T, 2.201821590438859327), lgamma(T, 0.105), eps);
297 try expectApproxEqRel(@as(T, 1.275416975248413231), lgamma(T, 0.253), eps);
298 try expectApproxEqRel(@as(T, 0.130463884049976732), lgamma(T, 0.823), eps);
299
300 try expectApproxEqRel(@as(T, 43.24395772148497989), lgamma(T, 21.3), eps);
301 try expectApproxEqRel(@as(T, 110.6908958012102623), lgamma(T, 41.1), eps);
302 try expectApproxEqRel(@as(T, 215.2123266224689711), lgamma(T, 67.4), eps);
303
304 try expectApproxEqRel(@as(T, -122.605958469563489), lgamma(T, -43.6), eps);
305 try expectApproxEqRel(@as(T, -278.633885462703133), lgamma(T, -81.4), eps);
306 try expectApproxEqRel(@as(T, -333.247676253238363), lgamma(T, -93.6), eps);
307 }
308}
309
310test "math.lgamma.special" {
311 inline for (&.{ f32, f64 }) |T| {
312 try expect(std.math.isNan(lgamma(T, -std.math.nan(T))));
313 try expect(std.math.isNan(lgamma(T, std.math.nan(T))));
314
315 try expectEqual(std.math.inf(T), lgamma(T, -std.math.inf(T)));
316 try expectEqual(std.math.inf(T), lgamma(T, std.math.inf(T)));
317
318 try expectEqual(std.math.inf(T), lgamma(T, -5));
319 try expectEqual(std.math.inf(T), lgamma(T, -8));
320 try expectEqual(std.math.inf(T), lgamma(T, -15));
321
322 try expectEqual(std.math.inf(T), lgamma(T, -0.0));
323 try expectEqual(std.math.inf(T), lgamma(T, 0.0));
324
325 try expect(std.math.isPositiveZero(lgamma(T, 1)));
326 try expect(std.math.isPositiveZero(lgamma(T, 2)));
327 }
328}