#!/usr/bin/env python
# coding: utf-8

# # Dual-Space Methods
# 
# Copyright (C) 2026 Andreas Kloeckner
# 
# <details>
# <summary>MIT License</summary>
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
# </details>
# 
# ---
# 
# Based on [A Dual-space Multilevel Kernel-splitting Framework for Discrete and Continuous Convolution](https://arxiv.org/abs/2308.00292) (Shidong Jiang, Leslie Greengard)
# 

# In[3]:


import numpy as np
import numpy.linalg as la
from scipy.special import erf, erfc
import matplotlib.pyplot as plt
from sumpy.visualization import FieldPlotter

# a lot of dividing by zero here, silence warnings
np.errstate(divide="ignore", invalid="ignore").__enter__()

rng = np.random.default_rng(230800292)

def _at_zero(r: np.ndarray, value: float, expression: np.ndarray) -> np.ndarray:
    """Replace the removable r=0 value in a radial-kernel expression."""
    return np.where(np.asarray(r) == 0.0, value, expression)

# The diameter of [-1/2, 1/2]^3, as in Lemma 8.
BOX_DIAMETER = np.sqrt(3.0)


# ### A refresher: `erf` and `erfc`

# In[4]:


x = np.linspace(-6, 6, 100)
plt.plot(x, erf(x), label="erf")
plt.plot(x, erfc(x), label="erfc")
plt.plot(x, erf(x) + erfc(x), label="sum")
plt.legend()


# ### Decomposing the Green's function

# In[5]:


# change me
sigma = 1

def M(r):
    return 1/r * erf(r/sigma) 
def W(r):
    return 1/r * erfc(r/sigma) 


# In[6]:


r = np.linspace(0, 5, 1000)

plt.plot(r, 1/r, "-", label="1/r")
plt.plot(r, M(r), label="M(r): mollified, erf")
plt.plot(r, W(r), label="R(r): residual, erfc")

del sigma

plt.legend(loc="best")
plt.ylim([0, 10])


# Observe decay past $6 \sigma$:

# In[7]:


r = np.linspace(6, 10, 100)

plt.semilogy(r, erfc(r))


# In[8]:


class FourierQuadrature:
    def __init__(self, *, period, n):
        self.h = h = 2.0 * np.pi / period
        self.modes = modes = np.arange(-n, n + 1)
        self.kx, self.ky, self.kz = np.meshgrid(h * modes, h * modes, h * modes, indexing="ij")
        self.k = np.column_stack((self.kx.ravel(), self.ky.ravel(), self.kz.ravel()))
        self.weight = (h / (2.0 * np.pi)) ** 3 
        self.k_mag = np.linalg.norm(self.k, axis=1) 

    def __call__(self, targets, k_values):
        return np.real(np.exp(1j * targets @ self.k.T) @ (self.weight * k_values))

    def plot(self, values, z_mode=0):
        z_idx, = np.where(self.modes == z_mode)
        n_modes = len(self.modes)

        myslice =  values.reshape(n_modes, n_modes, n_modes)[:, :, z_idx] 

        ax = plt.subplot(121)
        plt.imshow(np.log10(1e-15 + np.abs(myslice)))
        plt.colorbar()
        ax = plt.subplot(122)
        plt.imshow(np.arctan2(myslice.imag, myslice.real))

fquad = FourierQuadrature(period=8, n=48)


# In[9]:


def M_hat(k_mag):
    return 4*np.pi * np.exp(-(sigma * k_mag) ** 2 / 4.0) / k_mag**2

sigma = 1
m_hat_values = M_hat(fquad.k_mag) 
del sigma

fquad.plot(m_hat_values)

la.norm(m_hat_values, np.inf)


# ### Dealing with the Singularity

# In[10]:


def windowed(r, *, sigma, b) -> np.ndarray:
    """Physical-space W, evaluated stably at r=0."""
    a = BOX_DIAMETER + b * sigma
    value = (erf(r / sigma) - 0.5 * erf((a + r) / sigma)
             + 0.5 * erf((a - r) / sigma)) / r

    # Differentiate the numerator at zero.  The two large terms cancel here.
    limit = (2.0 / (np.sqrt(np.pi) * sigma)) * (1.0 - np.exp(-(a / sigma) ** 2))
    return _at_zero(r, limit, value)

r = np.linspace(0, 10, 1000)
plt.ylim([0, 2])

sigma = 1
mvals = 1/r * erf(r/sigma) 
wvals =windowed(r, sigma=sigma, b=6) 
plt.plot(r, 1/r, label="1/r")
plt.plot(r, mvals, label="M(r): mollified, erf")
plt.plot(r, wvals, label="W(r): windowed")
del sigma

plt.legend()


# In[11]:


def windowed_hat(k_mag, sigma, c_tilde):
    value = (8.0 * np.pi * (np.sin(c_tilde * k_mag / 2.0) / k_mag) ** 2  
             * np.exp(-(sigma * k_mag) ** 2 / 4.0))
    return _at_zero(k_mag, 2.0 * np.pi * c_tilde**2, value)

sigma = 0.2; b = 6
w_hat_values = windowed_hat(fquad.k_mag, sigma=sigma, c_tilde=BOX_DIAMETER + b * sigma)
del sigma
del b

fquad.plot(w_hat_values)


# In[12]:


fp = FieldPlotter(
    center=np.zeros(3),
    extent=np.array([1, 1, 0]),
    npoints=(10, 10, 1))
targets = fp.points.T
r = la.norm(targets, 2, axis=1)

sigma = 0.2; b = 0.6
err = fquad(targets, windowed_hat(fquad.k_mag, sigma=sigma, c_tilde=BOX_DIAMETER + b * sigma)) - windowed(r, sigma=sigma, b=b)
del sigma
del b

la.norm(err, np.inf)


# ### Multiple sources

# In[13]:


sources = np.array([[0, 0, 0]], dtype=np.float64)

fp = FieldPlotter(
    center=np.zeros(3),
    extent=np.array([1, 1, 0]),
    npoints=(100, 100, 1))
targets = fp.points.T

def pairwise_potential(targets: np.ndarray, sources: np.ndarray,
                       charges: np.ndarray, kernel) -> np.ndarray:
    distances = la.norm(targets[:, None, :] - sources[None, :, :], axis=2)
    return kernel(distances) @ charges

pot = pairwise_potential(targets, sources, [1], lambda r: 1/r)

fp.show_scalar_in_matplotlib(np.log10(1e-15 + np.abs(pot)))
plt.colorbar()


# In[14]:


def fourier_window_potential(targets: np.ndarray, sources: np.ndarray,
                             charges: np.ndarray, sigma: float, b: float) -> np.ndarray:
    source_phase = np.exp(-1j * fquad.k @ sources.T) @ charges
    return fquad(targets,
        windowed_hat(fquad.k_mag, sigma, BOX_DIAMETER + b * sigma) 
        * source_phase)

sigma = 0.18; b = 6
sources = rng.uniform(-0.5, 0.5, size=(12, 3))
targets = rng.uniform(-0.5, 0.5, size=(10, 3))
charges = rng.normal(size=12)
exact = pairwise_potential(targets, sources, charges,
                           lambda radius: windowed(radius, sigma=sigma, b=b))

spectral = fourier_window_potential(targets, sources, charges, sigma=sigma, b=b)
quadrature_error = np.max(np.abs(spectral - exact)) / np.max(np.abs(exact))
print(quadrature_error)
del sigma
del b


# ### Multilevel kernel splitting
# 
# The single-level Ewald split can be repeated on a hierarchy of box sizes.
# Following equation (40) of Jiang--Greengard, choose
# ``sigma_l = sigma_0 / 2**l``.  Then, at any leaf level ``L``,
# 
# $$1/r = W_0(r) + D_0(r) + ... + D_{L-1}(r) + R_L(r)$$
# 
# up to the exponentially small windowing error in ``W_0``.  The difference
# kernels are smooth and increasingly local, while the last residual is the
# only singular term.

# In[15]:


def mollified(r: np.ndarray, sigma: float) -> np.ndarray:
    """M_sigma(r) = erf(r / sigma) / r, including its value at r=0."""
    value = erf(r / sigma) / r
    return _at_zero(r, 2.0 / (np.sqrt(np.pi) * sigma), value)


def difference_kernel(r: np.ndarray, sigma_coarse: float,
                      sigma_fine: float) -> np.ndarray:
    """D_l = M_{l+1} - M_l, including its removable value at r=0."""
    value = (erf(r / sigma_fine) - erf(r / sigma_coarse)) / r
    limit = 2.0 / np.sqrt(np.pi) * (1.0 / sigma_fine - 1.0 / sigma_coarse)
    return _at_zero(r, limit, value)


def difference_kernel_hat(k_mag: np.ndarray, sigma_coarse: float,
                          sigma_fine: float) -> np.ndarray:
    """Fourier transform of D_l (equation (44) in the paper)."""
    value = (4.0 * np.pi
             * (np.exp(-(sigma_fine * k_mag) ** 2 / 4.0)
                - np.exp(-(sigma_coarse * k_mag) ** 2 / 4.0))
             / k_mag**2)
    limit = np.pi * (sigma_coarse**2 - sigma_fine**2)
    return _at_zero(k_mag, limit, value)


# The choice makes erfc(r_l / sigma_l) approximately ``tolerance`` at every
# level, so R_l and D_l are effectively supported inside boxes of side r_l.
tolerance = 1.0e-6
nlevels = 2
sigma0 = 1.0 / np.sqrt(np.log(1.0 / tolerance))
sigmas = sigma0 / 2.0**np.arange(nlevels + 1)
box_sizes = 2.0**-np.arange(nlevels + 1)
window_buffer = 6.0

# Do not include r=0 here: the final residual is singular there.
r = np.geomspace(1.0e-5, BOX_DIAMETER, 2000)
k = np.linspace(0.0, 120.0, 2000)

w0 = windowed(r, sigma=sigmas[0], b=window_buffer)
m0 = mollified(r, sigmas[0])
difference_kernels = [
    difference_kernel(r, sigmas[level], sigmas[level + 1])
    for level in range(nlevels)
]
residual = erfc(r / sigmas[-1]) / r


# In[16]:


fig, axes = plt.subplots(nlevels + 2, 2, figsize=(12, 3.1 * (nlevels + 2)),
                         constrained_layout=True)

axes[0, 0].semilogy(r, 1.0 / r, label=r"$1/r$")
axes[0, 0].set_title("original kernel")
axes[0, 1].semilogy(r, residual, label=rf"$R_{nlevels}(r)$")
axes[0, 1].axvline(box_sizes[-1], color="k", ls="--", lw=1,
                   label=rf"$r_{nlevels}$")
axes[0, 1].set_title("finest residual kernel")

axes[1, 0].semilogy(r, m0, label=r"$M_0(r)$")
axes[1, 0].semilogy(r, w0, "--", label=r"$W_0(r)$ (windowed)")
axes[1, 0].set_title("coarsest smooth kernel")
axes[1, 1].semilogy(k, windowed_hat(
    k, sigma=sigmas[0], c_tilde=BOX_DIAMETER + window_buffer * sigmas[0]))
axes[1, 1].set_title(r"$\widehat{W}_0(k)$")

for level, (sigma_coarse, sigma_fine, box_size, kernel) in enumerate(
        zip(sigmas[:-1], sigmas[1:], box_sizes[:-1], difference_kernels)):
    row = level + 2
    axes[row, 0].semilogy(r, kernel, label=rf"$D_{level}(r)$")
    axes[row, 0].axvline(box_size, color="k", ls="--", lw=1,
                         label=rf"$r_{level}$")
    axes[row, 0].set_title(rf"difference kernel $D_{level}$")
    axes[row, 1].semilogy(k, difference_kernel_hat(k, sigma_coarse, sigma_fine))
    axes[row, 1].set_title(rf"$\widehat{{D}}_{level}(k)$")

for ax in axes.flat:
    ax.set_xlabel(r"$r$" if ax in axes[:, 0] else r"$|k|$")
    ax.grid()
    if ax.lines and ax.get_legend_handles_labels()[0]:
        ax.legend(loc="best")


# Each correction is a rescaled copy of the preceding one: 
# $$D_l(r) =2^l D_0(2^l r).$$
# The same scaling shifts its Fourier content to higher k.

# In[18]:


scale_check = np.max(np.abs(
    difference_kernels[1] - 2.0 * difference_kernel(2.0 * r, sigmas[0], sigmas[1])))
print(f"max error in D_1(r) = 2 D_0(2r): {scale_check:.3e}")


# In[19]:


components = [w0, *difference_kernels, residual]
partial_sums = np.cumsum(components, axis=0)
reconstruction = partial_sums[-1]
relative_error = np.abs(r * reconstruction - 1.0)
print(f"max relative error in telescoping reconstruction: {relative_error.max():.3e}")


# In[20]:


fig, axes = plt.subplots(1, 3, figsize=(15, 4), constrained_layout=True)
for name, component in zip(
        [r"$W_0$", *[rf"$D_{level}$" for level in range(nlevels)],
         rf"$R_{nlevels}$"], components):
    axes[0].semilogy(r, component, label=name)
axes[0].set_title("terms in the telescoping decomposition")
axes[0].set_xlabel(r"$r$")
axes[0].legend()

axes[1].semilogy(r, 1.0 / r, "k", lw=2, label=r"$1/r$")
for level, partial_sum in enumerate(partial_sums):
    axes[1].semilogy(r, partial_sum, label=rf"partial sum through term {level}")
axes[1].set_title("successive telescoping partial sums")
axes[1].set_xlabel(r"$r$")
axes[1].legend()

axes[2].semilogy(r, relative_error)
axes[2].axhline(tolerance, color="k", ls="--", lw=1,
                label="requested tolerance")
axes[2].set_title(r"relative error: $|r(W_0 + \sum D_l + R_L)-1|$")
axes[2].set_xlabel(r"$r$")
axes[2].legend()

for ax in axes:
    ax.grid()


# In[ ]: