feat: add mexican hat neighborhood function (#2)

src/somap/__init__.py

@@ -14,7 +14,7 @@
 from .distance import EuclidianDist, wdist_l2
 from .grid import distance_map
 from .learning_rate import AbstractLr, ConstantLr, DsomLr, KsomLr
-from .neighborhood import AbstractNbh, DsomNbh, GaussianNbh, KsomNbh
+from .neighborhood import AbstractNbh, DsomNbh, GaussianNbh, KsomNbh, MexicanHatNbh
 from .plot import plot, save_plot
 from .plot_backends import set_plot_backend
 from .serialisation import load, save
@@ -49,6 +49,7 @@
     "DsomNbh",
     "GaussianNbh",
     "KsomNbh",
+    "MexicanHatNbh",
     "plot",
     "save_plot",
     "set_plot_backend",

src/somap/neighborhood.py

@@ -114,3 +114,25 @@ def __call__(
         return jnp.exp(
             -(distance_map**2) / (self.plasticity**2 * quantization_error**2)
         )
+
+
+class MexicanHatNbh(AbstractNbh):
+    """Mexican Hat neighborhood function."""
+
+    sigma: float | Float[Array, "..."] = 0.1
+
+    def __call__(self, distance_map: Float[Array, "x y"], _, __) -> Float[Array, "x y"]:
+        """Computes the Mexican Hat neighboring value of each grid element.
+
+        Args:
+            self:
+                self.sigma: Scale factor for the spread of the neighborhood.
+            distance_map: Distance of each element from the winner element.
+            _: Not used
+            __: Not used
+
+        Returns:
+            The Mexican Hat neighborhood distance.
+        """
+        r2_norm = distance_map**2 / self.sigma**2
+        return (1 - 0.5 * r2_norm) * jnp.exp(-r2_norm / 2)

src/somap/update.py

@@ -39,7 +39,8 @@ def __call__(
         w_bu: Float[Array, "..."],
     ) -> Float[Array, "..."]:
         """Updates the prototype weights."""
-        return w_bu + (lr * nbh)[:, :, jnp.newaxis] * (input_bu - w_bu)
+        out = w_bu + (lr * nbh)[:, :, jnp.newaxis] * (input_bu - w_bu)
+        return jnp.clip(out, 0, 1.0)
 
 
 @experimental_warning