From f652383952f4ffbbd91bd926d6902c7109c1e32b Mon Sep 17 00:00:00 2001 From: E-Rum Date: Sat, 18 Jan 2025 15:49:59 +0000 Subject: [PATCH 1/6] Add PyTorch implementation of the exponential integral function and update references --- docs/src/references/changelog.rst | 1 + src/torchpme/lib/__init__.py | 4 +- src/torchpme/lib/math.py | 51 +++++++++++++++++++--- src/torchpme/potentials/inversepowerlaw.py | 4 +- tests/lib/test_math.py | 11 ++--- 5 files changed, 55 insertions(+), 16 deletions(-) diff --git a/docs/src/references/changelog.rst b/docs/src/references/changelog.rst index 8f501396..3e405a23 100644 --- a/docs/src/references/changelog.rst +++ b/docs/src/references/changelog.rst @@ -27,6 +27,7 @@ changelog `_ format. This project follows Added ##### +* Added a PyTorch implementation of the exponential integral function * Added ``dtype`` and ``device`` for ``Calculator`` classses Changed diff --git a/src/torchpme/lib/__init__.py b/src/torchpme/lib/__init__.py index 47a0aa54..07f22a53 100644 --- a/src/torchpme/lib/__init__.py +++ b/src/torchpme/lib/__init__.py @@ -4,7 +4,7 @@ generate_kvectors_for_mesh, get_ns_mesh, ) -from .math import CustomExp1, gamma, gammaincc_over_powerlaw, torch_exp1 +from .math import CustomExp1, exp1, gamma, gammaincc_over_powerlaw from .mesh_interpolator import MeshInterpolator from .splines import ( CubicSpline, @@ -30,5 +30,5 @@ "gamma", "CustomExp1", "gammaincc_over_powerlaw", - "torch_exp1", + "exp1", ] diff --git a/src/torchpme/lib/math.py b/src/torchpme/lib/math.py index 871abbc2..886f84d6 100644 --- a/src/torchpme/lib/math.py +++ b/src/torchpme/lib/math.py @@ -1,5 +1,4 @@ import torch -from scipy.special import exp1 from torch.special import gammaln @@ -15,13 +14,51 @@ def gamma(x: torch.Tensor) -> torch.Tensor: class CustomExp1(torch.autograd.Function): - """Custom exponential integral function Exp1(x) to have an autograd-compatible version.""" + """ + Compute the exponential integral E1(x) for x > 0. + :param input: Input tensor (x > 0) + :return: Exponential integral E1(x) + """ @staticmethod def forward(ctx, input): ctx.save_for_backward(input) - input_numpy = input.cpu().numpy() if not input.is_cpu else input.numpy() - return torch.tensor(exp1(input_numpy), device=input.device, dtype=input.dtype) + + # Constants + SCIPY_EULER = ( + 0.577215664901532860606512090082402431 # Euler-Mascheroni constant + ) + inf = torch.inf + + # Handle case when x == 0 + result = torch.full_like(input, inf) + mask = input > 0 + + # Compute for x <= 1 + x_small = input[mask & (input <= 1)] + if x_small.numel() > 0: + e1 = torch.ones_like(x_small) + r = torch.ones_like(x_small) + for k in range(1, 26): + r = -r * k * x_small / (k + 1.0) ** 2 + e1 += r + if torch.all(torch.abs(r) <= torch.abs(e1) * 1e-15): + break + result[mask & (input <= 1)] = ( + -SCIPY_EULER - torch.log(x_small) + x_small * e1 + ) + + # Compute for x > 1 + x_large = input[mask & (input > 1)] + if x_large.numel() > 0: + m = 20 + (80.0 / x_large).to(torch.int32) + t0 = torch.zeros_like(x_large) + for k in range(m.max(), 0, -1): + t0 = k / (1.0 + k / (x_large + t0)) + t = 1.0 / (x_large + t0) + result[mask & (input > 1)] = torch.exp(-x_large) * t + + return result @staticmethod def backward(ctx, grad_output): @@ -29,7 +66,7 @@ def backward(ctx, grad_output): return -grad_output * torch.exp(-input) / input -def torch_exp1(input): +def exp1(input): """Wrapper for the custom exponential integral function.""" return CustomExp1.apply(input) @@ -41,13 +78,13 @@ def gammaincc_over_powerlaw(exponent: torch.Tensor, z: torch.Tensor) -> torch.Te if exponent == 2: return torch.sqrt(torch.pi / z) * torch.erfc(torch.sqrt(z)) if exponent == 3: - return torch_exp1(z) + return exp1(z) if exponent == 4: return 2 * ( torch.exp(-z) - torch.sqrt(torch.pi * z) * torch.erfc(torch.sqrt(z)) ) if exponent == 5: - return torch.exp(-z) - z * torch_exp1(z) + return torch.exp(-z) - z * exp1(z) if exponent == 6: return ( (2 - 4 * z) * torch.exp(-z) diff --git a/src/torchpme/potentials/inversepowerlaw.py b/src/torchpme/potentials/inversepowerlaw.py index 9abeb823..35ff7ac7 100644 --- a/src/torchpme/potentials/inversepowerlaw.py +++ b/src/torchpme/potentials/inversepowerlaw.py @@ -137,12 +137,12 @@ def background_correction(self) -> torch.Tensor: # "charge neutrality" correction for 1/r^p potential diverges for exponent p = 3 # and is not needed for p > 3 , so we set it to zero (see in # https://doi.org/10.48550/arXiv.2412.03281 SI section) - if self.exponent >= 3: - return torch.tensor(0.0, dtype=self.dtype, device=self.device) if self.smearing is None: raise ValueError( "Cannot compute background correction without specifying `smearing`." ) + if self.exponent >= 3: + return self.smearing * 0.0 prefac = torch.pi**1.5 * (2 * self.smearing**2) ** ((3 - self.exponent) / 2) prefac /= (3 - self.exponent) * gamma(self.exponent / 2) return prefac diff --git a/tests/lib/test_math.py b/tests/lib/test_math.py index 4ec7037c..b2c8cfb1 100644 --- a/tests/lib/test_math.py +++ b/tests/lib/test_math.py @@ -2,7 +2,7 @@ import torch from scipy.special import exp1 -from torchpme.lib import torch_exp1 +from torchpme.lib import exp1 as torch_exp1 def finite_difference_derivative(func, x, h=1e-5): @@ -10,10 +10,11 @@ def finite_difference_derivative(func, x, h=1e-5): def test_torch_exp1_consistency_with_scipy(): - x = torch.rand(1000, dtype=torch.float64) - torch_result = torch_exp1(x) - scipy_result = exp1(x.numpy()) - assert np.allclose(torch_result.numpy(), scipy_result, atol=1e-6) + random_tensor = torch.FloatTensor(100000).uniform_(0, 1000) + random_array = random_tensor.numpy() + scipy_result = exp1(random_array) + torch_result = torch_exp1(random_tensor) + assert np.allclose(scipy_result, torch_result.numpy(), atol=1e-15) def test_torch_exp1_derivative(): From 3ac4fc2e58d869b8e85228e94bca7daa25ce9b59 Mon Sep 17 00:00:00 2001 From: E-Rum Date: Tue, 21 Jan 2025 13:13:39 +0000 Subject: [PATCH 2/6] Enhance documentation for gammaincc_over_powerlaw function and improve randomness in tests --- src/torchpme/lib/math.py | 8 +++++++- tests/lib/test_math.py | 5 ++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/src/torchpme/lib/math.py b/src/torchpme/lib/math.py index 886f84d6..fc8d00d0 100644 --- a/src/torchpme/lib/math.py +++ b/src/torchpme/lib/math.py @@ -72,7 +72,13 @@ def exp1(input): def gammaincc_over_powerlaw(exponent: torch.Tensor, z: torch.Tensor) -> torch.Tensor: - """Function to compute the regularized incomplete gamma function complement for integer exponents.""" + """ + Function to compute the regularized incomplete gamma function complement for integer + exponents. + param exponent: Exponent of the power law + param z: Value at which to evaluate the function + return: Regularized incomplete gamma function complement + """ if exponent == 1: return torch.exp(-z) / z if exponent == 2: diff --git a/tests/lib/test_math.py b/tests/lib/test_math.py index b2c8cfb1..cb314e4e 100644 --- a/tests/lib/test_math.py +++ b/tests/lib/test_math.py @@ -10,7 +10,10 @@ def finite_difference_derivative(func, x, h=1e-5): def test_torch_exp1_consistency_with_scipy(): - random_tensor = torch.FloatTensor(100000).uniform_(0, 1000) + seed = 42 + torch.manual_seed(seed) + np.random.seed(seed) + random_tensor = torch.rand(100000) * 1000 random_array = random_tensor.numpy() scipy_result = exp1(random_array) torch_result = torch_exp1(random_tensor) From e4b6d71dbdecc17c7c82f60f84702e40c148f2ad Mon Sep 17 00:00:00 2001 From: E-Rum Date: Wed, 22 Jan 2025 16:05:31 +0000 Subject: [PATCH 3/6] Refactor CustomExp1 class to be private and add link to SciPy implementation --- src/torchpme/lib/__init__.py | 3 +-- src/torchpme/lib/math.py | 6 ++++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/torchpme/lib/__init__.py b/src/torchpme/lib/__init__.py index 07f22a53..f21cf355 100644 --- a/src/torchpme/lib/__init__.py +++ b/src/torchpme/lib/__init__.py @@ -4,7 +4,7 @@ generate_kvectors_for_mesh, get_ns_mesh, ) -from .math import CustomExp1, exp1, gamma, gammaincc_over_powerlaw +from .math import exp1, gamma, gammaincc_over_powerlaw from .mesh_interpolator import MeshInterpolator from .splines import ( CubicSpline, @@ -28,7 +28,6 @@ "generate_kvectors_for_mesh", "get_ns_mesh", "gamma", - "CustomExp1", "gammaincc_over_powerlaw", "exp1", ] diff --git a/src/torchpme/lib/math.py b/src/torchpme/lib/math.py index fc8d00d0..298839e0 100644 --- a/src/torchpme/lib/math.py +++ b/src/torchpme/lib/math.py @@ -13,7 +13,7 @@ def gamma(x: torch.Tensor) -> torch.Tensor: return torch.exp(gammaln(x)) -class CustomExp1(torch.autograd.Function): +class _CustomExp1(torch.autograd.Function): """ Compute the exponential integral E1(x) for x > 0. :param input: Input tensor (x > 0) @@ -22,6 +22,8 @@ class CustomExp1(torch.autograd.Function): @staticmethod def forward(ctx, input): + # this implementation is inspired by the one in scipy: + # https://github.com/scipy/scipy/blob/34d91ce06d4d05e564b79bf65288284247b1f3e3/scipy/special/xsf/expint.h#L22 ctx.save_for_backward(input) # Constants @@ -68,7 +70,7 @@ def backward(ctx, grad_output): def exp1(input): """Wrapper for the custom exponential integral function.""" - return CustomExp1.apply(input) + return _CustomExp1.apply(input) def gammaincc_over_powerlaw(exponent: torch.Tensor, z: torch.Tensor) -> torch.Tensor: From 220fe7f2f6cd341ce74af2de6bc911f0ca99e85e Mon Sep 17 00:00:00 2001 From: E-Rum Date: Wed, 22 Jan 2025 16:44:09 +0000 Subject: [PATCH 4/6] Fix docs --- src/torchpme/lib/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/torchpme/lib/__init__.py b/src/torchpme/lib/__init__.py index f21cf355..cc5d9b6a 100644 --- a/src/torchpme/lib/__init__.py +++ b/src/torchpme/lib/__init__.py @@ -4,7 +4,7 @@ generate_kvectors_for_mesh, get_ns_mesh, ) -from .math import exp1, gamma, gammaincc_over_powerlaw +from .math import _CustomExp1, exp1, gamma, gammaincc_over_powerlaw from .mesh_interpolator import MeshInterpolator from .splines import ( CubicSpline, @@ -28,6 +28,7 @@ "generate_kvectors_for_mesh", "get_ns_mesh", "gamma", + "_CustomExp1", "gammaincc_over_powerlaw", "exp1", ] From e05683b09f7dcc29c571081f5bab92fb14556a32 Mon Sep 17 00:00:00 2001 From: E-Rum Date: Wed, 22 Jan 2025 19:17:30 +0000 Subject: [PATCH 5/6] Another fix docs --- tests/lib/test_math.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/tests/lib/test_math.py b/tests/lib/test_math.py index cb314e4e..d02196cc 100644 --- a/tests/lib/test_math.py +++ b/tests/lib/test_math.py @@ -1,8 +1,8 @@ import numpy as np +import scipy.special import torch -from scipy.special import exp1 -from torchpme.lib import exp1 as torch_exp1 +from torchpme.lib import exp1 def finite_difference_derivative(func, x, h=1e-5): @@ -15,15 +15,17 @@ def test_torch_exp1_consistency_with_scipy(): np.random.seed(seed) random_tensor = torch.rand(100000) * 1000 random_array = random_tensor.numpy() - scipy_result = exp1(random_array) - torch_result = torch_exp1(random_tensor) + scipy_result = scipy.special.exp1(random_array) + torch_result = exp1(random_tensor) assert np.allclose(scipy_result, torch_result.numpy(), atol=1e-15) def test_torch_exp1_derivative(): x = torch.rand(1, dtype=torch.float64, requires_grad=True) - torch_result = torch_exp1(x) + torch_result = exp1(x) torch_result.backward() torch_exp1_prime = x.grad - finite_diff_result = finite_difference_derivative(exp1, x.detach().numpy()) + finite_diff_result = finite_difference_derivative( + scipy.special.exp1, x.detach().numpy() + ) assert np.allclose(torch_exp1_prime.numpy(), finite_diff_result, atol=1e-6) From 10e250a0c89cab08d548f0f57327b01c9f654655 Mon Sep 17 00:00:00 2001 From: Philip Loche Date: Wed, 22 Jan 2025 20:35:56 +0100 Subject: [PATCH 6/6] fix docs --- docs/src/references/lib/math.rst | 2 +- src/torchpme/lib/__init__.py | 3 +- src/torchpme/lib/math.py | 55 +++++++++++++++++--------------- 3 files changed, 31 insertions(+), 29 deletions(-) diff --git a/docs/src/references/lib/math.rst b/docs/src/references/lib/math.rst index b20c9c83..f2cde344 100644 --- a/docs/src/references/lib/math.rst +++ b/docs/src/references/lib/math.rst @@ -1,6 +1,6 @@ Math #### +.. autofunction:: torchpme.lib.exp1 .. autofunction:: torchpme.lib.gamma -.. autofunction:: torchpme.lib.torch_exp1 .. autofunction:: torchpme.lib.gammaincc_over_powerlaw diff --git a/src/torchpme/lib/__init__.py b/src/torchpme/lib/__init__.py index cc5d9b6a..f21cf355 100644 --- a/src/torchpme/lib/__init__.py +++ b/src/torchpme/lib/__init__.py @@ -4,7 +4,7 @@ generate_kvectors_for_mesh, get_ns_mesh, ) -from .math import _CustomExp1, exp1, gamma, gammaincc_over_powerlaw +from .math import exp1, gamma, gammaincc_over_powerlaw from .mesh_interpolator import MeshInterpolator from .splines import ( CubicSpline, @@ -28,7 +28,6 @@ "generate_kvectors_for_mesh", "get_ns_mesh", "gamma", - "_CustomExp1", "gammaincc_over_powerlaw", "exp1", ] diff --git a/src/torchpme/lib/math.py b/src/torchpme/lib/math.py index 298839e0..aae243dc 100644 --- a/src/torchpme/lib/math.py +++ b/src/torchpme/lib/math.py @@ -14,17 +14,11 @@ def gamma(x: torch.Tensor) -> torch.Tensor: class _CustomExp1(torch.autograd.Function): - """ - Compute the exponential integral E1(x) for x > 0. - :param input: Input tensor (x > 0) - :return: Exponential integral E1(x) - """ - @staticmethod - def forward(ctx, input): + def forward(ctx, x): # this implementation is inspired by the one in scipy: # https://github.com/scipy/scipy/blob/34d91ce06d4d05e564b79bf65288284247b1f3e3/scipy/special/xsf/expint.h#L22 - ctx.save_for_backward(input) + ctx.save_for_backward(x) # Constants SCIPY_EULER = ( @@ -33,11 +27,11 @@ def forward(ctx, input): inf = torch.inf # Handle case when x == 0 - result = torch.full_like(input, inf) - mask = input > 0 + result = torch.full_like(x, inf) + mask = x > 0 # Compute for x <= 1 - x_small = input[mask & (input <= 1)] + x_small = x[mask & (x <= 1)] if x_small.numel() > 0: e1 = torch.ones_like(x_small) r = torch.ones_like(x_small) @@ -46,40 +40,49 @@ def forward(ctx, input): e1 += r if torch.all(torch.abs(r) <= torch.abs(e1) * 1e-15): break - result[mask & (input <= 1)] = ( - -SCIPY_EULER - torch.log(x_small) + x_small * e1 - ) + result[mask & (x <= 1)] = -SCIPY_EULER - torch.log(x_small) + x_small * e1 # Compute for x > 1 - x_large = input[mask & (input > 1)] + x_large = x[mask & (x > 1)] if x_large.numel() > 0: m = 20 + (80.0 / x_large).to(torch.int32) t0 = torch.zeros_like(x_large) for k in range(m.max(), 0, -1): t0 = k / (1.0 + k / (x_large + t0)) t = 1.0 / (x_large + t0) - result[mask & (input > 1)] = torch.exp(-x_large) * t + result[mask & (x > 1)] = torch.exp(-x_large) * t return result @staticmethod def backward(ctx, grad_output): - (input,) = ctx.saved_tensors - return -grad_output * torch.exp(-input) / input + (x,) = ctx.saved_tensors + return -grad_output * torch.exp(-x) / x + + +def exp1(x): + r""" + Exponential integral E1. + For a real number :math:`x > 0` the exponential integral can be defined as -def exp1(input): - """Wrapper for the custom exponential integral function.""" - return _CustomExp1.apply(input) + .. math:: + + E1(x) = \int_{x}^{\infty} \frac{e^{-t}}{t} dt + + :param x: Input tensor (x > 0) + :return: Exponential integral E1(x) + """ + return _CustomExp1.apply(x) def gammaincc_over_powerlaw(exponent: torch.Tensor, z: torch.Tensor) -> torch.Tensor: """ - Function to compute the regularized incomplete gamma function complement for integer - exponents. - param exponent: Exponent of the power law - param z: Value at which to evaluate the function - return: Regularized incomplete gamma function complement + Compute the regularized incomplete gamma function complement for integer exponents. + + :param exponent: Exponent of the power law + :param z: Value at which to evaluate the function + :return: Regularized incomplete gamma function complement """ if exponent == 1: return torch.exp(-z) / z