src/vpso/math.py
from typing import Literal
import numba as nb
import numpy as np
from vpso.typing import Array2d, Array3d
@nb.njit(
nb.float64[:, :](nb.float64[:, :], nb.float64[:, :], nb.types.unicode_type),
cache=True,
nogil=True,
)
def cdist_func(
x: Array2d, y: Array2d, type: Literal["euclidean", "sqeuclidean"]
) -> Array2d:
"""Computes the distance matrix for 2D arrays `x` and `y`.
Parameters
----------
x : 2d array
A 2D array of shape `(N, d)`.
y : 2d array
A 2D array of shape `(M, d)`.
type : {"euclidean", "sqeuclidean"}
Type of distance to compute.
Returns
-------
2d array
Distance matrix of shape `(N, M)` between each pair of entries in `x` and `y`.
"""
x2_sum = np.square(x).sum(axis=1)
y2_sum = np.square(y).sum(axis=1)
xy_dot = np.dot(x, y.T)
D = np.maximum(x2_sum[:, np.newaxis] + y2_sum - 2 * xy_dot, 0.0)
if type == "euclidean":
return np.sqrt(D)
if type == "sqeuclidean":
return D
raise ValueError(f"unknown cdist type: {type}")
@nb.njit(
nb.float64[:, :](nb.float64[:, :], nb.types.unicode_type),
cache=True,
nogil=True,
)
def pdist_func(x: Array2d, type: Literal["euclidean", "sqeuclidean"]) -> Array2d:
"""Computes the pairwise distance matrix of the elements in the 2D array `x`.
Parameters
----------
x : 2d array
A 2D array of shape `(N, d)`.
type : {"euclidean", "sqeuclidean"}
Type of distance to compute.
Returns
-------
2d array
Symmetric distance matrix of shape `(N, N)` between each pair of entries in `x`.
"""
x2_sum = np.square(x).sum(axis=1)
xx_dot = np.dot(x, x.T)
D = np.maximum(x2_sum[:, np.newaxis] + x2_sum - 2 * xx_dot, 0.0)
np.fill_diagonal(D, 0.0)
if type == "euclidean":
return np.sqrt(D)
if type == "sqeuclidean":
return D
raise ValueError(f"unknown pdist type: {type}")
@nb.njit(
nb.float64[:, :, :](
nb.float64[:, :, :], nb.float64[:, :, :], nb.types.unicode_type
),
cache=True,
nogil=True,
parallel=True,
)
def batch_cdist(
X: Array3d, Y: Array3d, type: Literal["euclidean", "sqeuclidean"]
) -> Array3d:
"""Computes the distance matrices for 3D arrays.
Parameters
----------
X : 3d array
An array of shape `(N, M, d)`.
Y : 3d array
An array of shape `(N, K, d)`.
dist_func : callable, optional
Distance function to use. By default, `scipy.spatial.distance.cdist_euclidean`
is used. It must support the `out` argument.
Returns
-------
3d array
Distance matrices between each of the `(M, d)` and `(K, d)` matrices, where `d`
is assumed to be the axis over which the distance is computed. The output has
thus shape (N, M, K).
"""
B, N, _ = X.shape
out = np.empty((B, N, Y.shape[1]), dtype=X.dtype)
for i in nb.prange(B):
out[i] = cdist_func(X[i], Y[i], type)
return out
@nb.njit(
nb.float64[:, :, :](nb.float64[:, :, :], nb.types.unicode_type),
cache=True,
nogil=True,
parallel=True,
)
def batch_pdist(X: Array3d, type: Literal["euclidean", "sqeuclidean"]) -> Array3d:
"""Computes the pairwise distance matrices for the entries of a 3D array.
Parameters
----------
X : 3d array
An array of shape `(N, M, d)`.
dist_func : callable, optional
Distance function to use. By default, `scipy.spatial.distance.pdist_euclidean`
is used. It must support the `out` argument.
Returns
-------
3d array
Distance matrix of shape `(N, M, M)` between each pair of entries of
the `(M, d)` matrices, where `d` is assumed to be the axis over which the
distance is computed.
"""
B, N, _ = X.shape
out = np.empty((B, N, N), dtype=X.dtype)
for i in nb.prange(B):
out[i] = pdist_func(X[i], type)
return out
@nb.njit(
nb.types.UniTuple(nb.float64[:, :, :], 2)(
nb.float64[:, :, :], nb.float64[:, :, :], nb.types.unicode_type
),
cache=True,
nogil=True,
parallel=True,
)
def batch_cdist_and_pdist(
X: Array3d, Y: Array3d, type: Literal["euclidean", "sqeuclidean"]
) -> tuple[Array3d, Array3d]:
"""Computes the distance matrices between the 3D arrays `X` and `Y`, as well as the
pairwise distance matrices for the entries of `X`.
Parameters
----------
X : 3d array
An array of shape `(N, M, d)`.
Y : 3d array
An array of shape `(N, K, d)`.
dist_func : callable, optional
Distance function to use. By default, `scipy.spatial.distance.cdist_euclidean`
is used. It must support the `out` argument.
Returns
-------
tuple of 3d arrays
Returns both distance matrices between `X` and `Y` and `X` itself. See
`batch_cdist` and `batch_pdist` for more details.
"""
B, N, _ = X.shape
out_c = np.empty((B, N, Y.shape[1]), dtype=X.dtype)
out_p = np.empty((B, N, N), dtype=X.dtype)
for i in nb.prange(B):
out_c[i] = cdist_func(X[i], Y[i], type)
out_p[i] = pdist_func(X[i], type)
return out_c, out_p
@nb.njit(
nb.types.UniTuple(nb.float64[:, :, :], 2)(
nb.float64[:, :, :], # x
nb.float64[:, :, :], # px
nb.float64[:, :, :], # sx
nb.float64[:, :, :], # v
nb.float64[:, :, :], # v_max
nb.float64[:, :, :], # w
nb.float64[:, :, :], # c1
nb.float64[:, :, :], # c2
nb.types.NumPyRandomGeneratorType("NumPyRandomGeneratorType"),
),
cache=True,
nogil=True,
)
def pso_equation(
x: Array3d,
px: Array3d,
sx: Array3d,
v: Array3d,
v_max: Array3d,
w: Array3d,
c1: Array3d,
c2: Array3d,
np_random: np.random.Generator,
) -> tuple[Array3d, Array3d]:
"""Computes the new positions and velocities of particles in a PSO algorithm.
Parameters
----------
x : 3d array
Current positions of the particles. An array of shape `(N, M, d)`, where `N` is
the number of vectorized problems to solve simultaneously, `M` is the number of
particles in the swarm, and `d` is the dimension of the search space.
px : 3d array
Best positions of the particles so far. An array of shape `(N, M, d)`.
sx : 3d array
Social best, i.e., the best particle so far. An array of shape `(N, 1, d)`.
v : 3d array
Current velocities of the particles. An array of shape `(N, M, d)`.
v_max : 3d array
Maximum velocities of the particles. An array of shape `(N, 1, d)`.
w : 3d array
Inertia weight. An array of shape `(N, 1, 1)`, where each element is used for
the corresponding problem.
c1 : 3d array
Cognitive weight. An array of shape `(N, 1, 1)`, where each element is used for
the corresponding problem.
c2 : 3d array
Social weight. An array of shape `(N, 1, 1)`, where each element is used for
the corresponding problem.
np_random : np.random.Generator
Random number generator.
Returns
-------
tuple of 3d arrays
New positions and velocities of the particles.
"""
r1 = np_random.random(x.shape)
r2 = np_random.random(x.shape)
inerta = w * v
cognitive = c1 * r1 * (px - x)
social = c2 * r2 * (sx - x)
v_new = (inerta + cognitive + social).clip(-v_max, v_max)
x_new = x + v_new
return x_new, v_new # type: ignore[return-value] # thinks it's a bool array...