From ac810a105c3fd4eb2955093d9839a2a856a2fe5f Mon Sep 17 00:00:00 2001
From: Schwarf <32012898+Schwarf@users.noreply.github.com>
Date: Sun, 16 Nov 2025 19:41:26 +0100
Subject: [PATCH] mathext: optimize Li2 and add benchmarks

This optimizes the complex dilogarithm Li2 using targeted mappings
(reflection/inversion), a real-axis fast path, and a Bernoulli-in-x
expansion (x = -log(1 - z)), inspired by this Rust implementation[1]
(and references therein). The exported API remains unchanged, and
runtime improvements are measured by an added BenchmarkLi2.

Performance improvements range from a few percent to large factors,
depending on the region of the complex plane.

[1]https://github.com/Expander/polylog.rs
---
 mathext/dilog.go            | 172 +++++++++++++++++++++++++++++-------
 mathext/dilog_bench_test.go |  31 +++++++
 2 files changed, 169 insertions(+), 34 deletions(-)
 create mode 100644 mathext/dilog_bench_test.go

diff --git a/mathext/dilog.go b/mathext/dilog.go
index 048b54d7..d8207c2a 100644
--- a/mathext/dilog.go
+++ b/mathext/dilog.go
@@ -13,7 +13,7 @@ import (
 //
 // For |z| < 1, Li2(z) is defined by the power series
 //
-//     Li2(z) = SUM_{k=1}^{infinity} z^k / k^2
+//	Li2(z) = SUM_{k=1}^{infinity} z^k / k^2
 //
 // and is analytically continued to the rest of the complex plane.
 // The implementation uses reflection and inversion identities to map z into
@@ -23,44 +23,148 @@ import (
 // cut on the real axis for z in the interval (1, infinity). The principal value is taken with
 // Arg(z) element of (−Pi, Pi].
 func Li2(z complex128) complex128 {
-	// Special cases
-	if z == 0 {
+	rz, iz := real(z), imag(z)
+
+	// Real-axis fast path:
+	if iz == 0 {
+		if rz <= 1 {
+			return complex(li2Real(rz), 0)
+		}
+		// rz > 1: principal Im part = −π ln(rz)
+		return complex(li2Real(rz), -math.Pi*math.Log(rz))
+	}
+
+	abs2 := rz*rz + iz*iz
+
+	// Tiny |z|: 2-term Taylor z + z^2/4
+	if abs2 < 1e-16 {
+		return z * (1 + 0.25*z)
+	}
+
+	if rz <= 0.5 {
+		if abs2 > 1 {
+			// Inversion on left half-plane:
+			// Li2(z) = −[Li2(1−1/z)] − ½ ln²(−z) − π²/6
+			l := cmplx.Log(-z) // principal log
+			// compute Li2(1−1/z) via x-series: x = −ln(1 − (1−1/z)) = −ln(1/z) = ln z (with sign)
+			x := -cmplx.Log(1 - 1/z) // x for Li2(1−1/z)
+			return -(li2SeriesXFromX(x)) - 0.5*l*l - complex(pi2over6, 0)
+		}
+		// Inside unit disk, left half-plane: use x-series with x = −ln(1−z)
+		x := -cmplx.Log(1 - z)
+		return li2SeriesXFromX(x)
+	}
+
+	// rz > 0.5
+	if abs2 <= 2*rz {
+		// Near 1: reflection
+		// Li2(z) = π²/6 − ln z · ln(1−z) − Li2(1−z)
+		lz := cmplx.Log(z)
+		l1z := cmplx.Log(1 - z)
+		// Li2(1−z) via x-series: x = −ln(1 − (1−z)) = −ln z
+		x := -lz
+		return complex(pi2over6, 0) - lz*l1z - li2SeriesXFromX(x)
+	}
+
+	// Farther from 1 on right half-plane: inversion
+	// Li2(z) = −½ ln²(−z) − Li2(1−1/z) − π²/6
+	l := cmplx.Log(-z)
+	x := -cmplx.Log(1 - 1/z) // x for Li2(1−1/z)
+	return -(li2SeriesXFromX(x)) - 0.5*l*l - complex(pi2over6, 0)
+}
+
+func li2Real(x float64) float64 {
+	// see arXiv:2201.01678 for the underlying approximation strategy.
+	switch {
+	case x < -1:
+		// Li2(x) = Li2(1/(1−x)) − π²/6 + ln(1−x)·(½ ln(1−x) − ln(−x))
+		l := math.Log(1 - x)
+		return li2ApproxHalf(1/(1-x)) - pi2over6 + l*(0.5*l-math.Log(-x))
+
+	case x == -1:
+		return -0.5 * pi2over6
+
+	case x < 0:
+		// Li2(x) = −Li2(x/(x−1)) − ½ ln²(1−x)
+		l := math.Log1p(-x) // ln(1−x)
+		return -li2ApproxHalf(x/(x-1)) - 0.5*l*l
+
+	case x == 0:
 		return 0
-	}
-	if z == 1 {
-		return complex(math.Pi*math.Pi/6, 0)
-	}
-	if z == -1 {
-		return complex(-math.Pi*math.Pi/12, 0)
-	}
 
-	// Reflection: map Re(z) > 0.5 into left half-plane for better convergence
-	// This formula is applied before inversion on the principal branch and gives very accurate results
-	// for real z > 1 because Li2(-1) and similar values are known exactly.
-	// Inversion first would also be valid, but is more sensitive to the
-	// Arg(-z)= +/-Pi branch choice in log(-z) and can yield wrong imaginary signs
-	// if care is not taken.
-	if real(z) > 0.5 {
-		return complex(math.Pi*math.Pi/6, 0) - complex128(cmplx.Log(z)*cmplx.Log(1-z)) - Li2(1-z)
-	}
+	case x < 0.5:
+		return li2ApproxHalf(x)
 
-	// Inversion: map |z| > 1 into unit disk
-	if cmplx.Abs(z) > 1 {
-		logmz := cmplx.Log(-z)
-		return -complex(math.Pi*math.Pi/6, 0) - complex128(0.5*logmz*logmz) - Li2(1/z)
+	case x < 1:
+		// Li2(x) = π²/6 − ln x · ln(1−x) − Li2(1−x)
+		return -li2ApproxHalf(1-x) + pi2over6 - math.Log(x)*math.Log1p(-x)
+
+	case x == 1:
+		return pi2over6
+
+	case x < 2:
+		// Li2(x) = Li2(1−1/x) + π²/6 − ln x · (ln(1−1/x) + ½ ln x)
+		l := math.Log(x)
+		return li2ApproxHalf(1-1/x) + pi2over6 - l*(math.Log(1-1/x)+0.5*l)
+
+	default: // x >= 2
+		// Li2(x) = −Li2(1/x) + 2·π²/6 − ½ ln² x
+		l := math.Log(x)
+		return -li2ApproxHalf(1/x) + 2*pi2over6 - 0.5*l*l
 	}
-	// Direct series for |z| <= 1 and Re(z) <= 0.5
-	return li2Series(z)
 }
 
-func li2Series(z complex128) complex128 {
-	const tol = 1e-15
-	var sum complex128
-	zk := z // zk = z^k
-	for k := 1.0; cmplx.Abs(zk)/(k*k) > tol; k++ {
-		k2 := k * k
-		sum += complex(real(zk)/k2, imag(zk)/k2)
-		zk *= z
+func li2ApproxHalf(x float64) float64 {
+	// Padé-like rational approximant for Re(Li2(x)) on x ∈ [0, 1/2],
+	// following arXiv:2201.01678.
+	cp := [...]float64{
+		0.9999999999999999502e+0,
+		-2.6883926818565423430e+0,
+		2.6477222699473109692e+0,
+		-1.1538559607887416355e+0,
+		2.0886077795020607837e-1,
+		-1.0859777134152463084e-2,
 	}
-	return sum
+	cq := [...]float64{
+		1.0000000000000000000e+0,
+		-2.9383926818565635485e+0,
+		3.2712093293018635389e+0,
+		-1.7076702173954289421e+0,
+		4.1596017228400603836e-1,
+		-3.9801343754084482956e-2,
+		8.2743668974466659035e-4,
+	}
+	x2 := x * x
+	x4 := x2 * x2
+	p := cp[0] + x*cp[1] + x2*(cp[2]+x*cp[3]) + x4*(cp[4]+x*cp[5])
+	q := cq[0] + x*cq[1] + x2*(cq[2]+x*cq[3]) + x4*(cq[4]+x*cq[5]+x2*cq[6])
+	return x * (p / q)
 }
+
+func li2SeriesXFromX(x complex128) complex128 {
+	// Bernoulli-in-x expansion with x = −ln(1−z).
+	// Inspired by SPheno CLI2 implementation in Fortran90 (https://spheno.hepforge.org/)
+	const (
+		b0 = -1.0 / 4.0
+		b1 = 1.0 / 36.0
+		b2 = -1.0 / 3600.0
+		b3 = 1.0 / 211680.0
+		b4 = -1.0 / 10886400.0
+		b5 = 1.0 / 526901760.0
+		b6 = -4.0647616451442255e-11
+		b7 = 8.9216910204564526e-13
+		b8 = -1.9939295860721076e-14
+		b9 = 4.5189800296199182e-16
+	)
+	x2 := x * x
+	x4 := x2 * x2
+	return x +
+		x2*(complex(b0, 0)+
+			x*(complex(b1, 0)+
+				x2*(complex(b2, 0)+
+					x2*complex(b3, 0)+
+					x4*(complex(b4, 0)+x2*complex(b5, 0))+
+					x4*x4*(complex(b6, 0)+x2*complex(b7, 0)+x4*(complex(b8, 0)+x2*complex(b9, 0))))))
+}
+
+const pi2over6 = math.Pi * math.Pi / 6
diff --git a/mathext/dilog_bench_test.go b/mathext/dilog_bench_test.go
new file mode 100644
index 00000000..fd2d9e68
--- /dev/null
+++ b/mathext/dilog_bench_test.go
@@ -0,0 +1,31 @@
+// Copyright ©2025 The Gonum Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package mathext
+
+import (
+	"fmt"
+	"testing"
+)
+
+var li2Sink complex128
+
+func BenchmarkLi2(b *testing.B) {
+	cases := []complex128{
+		0.1 + 0.1i, -0.3 + 0.39i, 0.001 - 0.49i, // |z| < 0.5
+		-0.9999 + 0.001i, 0.5 + 0.7i, -0.8 - 0.0001i, // 0.5 < |z| < 1
+		-1.1 + 0.1i, 5 + 0i, -10 + 0i, 1000 + 1e4i, -1791.91931 + 0.5i, // |z| > 1
+	}
+	for _, z := range cases {
+		b.Run(fmt.Sprintf("z=%.6g", z), func(b *testing.B) {
+			b.ReportAllocs()
+			var r complex128
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				r = Li2(z)
+			}
+			li2Sink = r
+		})
+	}
+}