[ENH] Switching function to suppress NN output and go to purely scale-shift correction when no atoms nearby about apax HOT 2 CLOSED

from jax.ops import segment_sum, segment_min
import jax.numpy as jnp
import numpy as np

from matscipy.neighbours import neighbour_list

from scipy.stats.qmc import PoissonDisk

import matplotlib.pyplot as plt


np.random.seed(35280)


sampler = PoissonDisk(3)
points = sampler.fill_space()
num_points = len(points)
a = jnp.array(points)
i, j, dr = neighbour_list('ijd', cutoff=0.2, positions=np.array(a))
dr_mins = segment_min(dr, i, num_points)
repeats = jnp.bincount(i)
dr_mins = jnp.repeat(dr_mins, repeats)

for exp_denom in np.logspace(-4, -2, 3):
    softmaxed_dr = jnp.exp(-(dr - dr_mins) / exp_denom)
    softmax_weighted_dr = dr * softmaxed_dr / jnp.repeat(segment_sum(softmaxed_dr, i, num_points), repeats)
    strict_min_dr = segment_min(dr, i, num_points)
    soft_min_dr = segment_sum(softmax_weighted_dr, i, num_points)

    plt.scatter(strict_min_dr, soft_min_dr, s=1, label=f"{exp_denom:0.1e}")
plt.legend()
plt.show()