Types and format

Makefile

@@ -2,6 +2,7 @@ python=./env/bin/python
 mamba=mamba
 pkg=qmllib
 pip=./env/bin/pip
+pytest=pytest
 j=1
 
 .PHONY: build
@@ -24,7 +25,7 @@ test:
 	${python} -m pytest -rs ./tests
 
 types:
-	${python} -m monkeytype run $(which pytest) ./tests/
+	${python} -m monkeytype run $$(which ${pytest}) ./tests
 	${python} -m monkeytype list-modules | grep ${pkg} | parallel -j${j} "${python} -m monkeytype apply {}"
 
 cov:

src/qmllib/kernels/distance.py

@@ -1,9 +1,12 @@
+from typing import Union
+
 import numpy as np
+from numpy import ndarray
 
 from .fdistance import fl2_distance, fmanhattan_distance, fp_distance_double, fp_distance_integer
 
 
-def manhattan_distance(A, B):
+def manhattan_distance(A: ndarray, B: ndarray) -> ndarray:
     """Calculates the Manhattan distances, D,  between two
     Numpy arrays of representations.
 
@@ -37,7 +40,7 @@ def manhattan_distance(A, B):
     return D
 
 
-def l2_distance(A, B):
+def l2_distance(A: ndarray, B: ndarray) -> ndarray:
     """Calculates the L2 distances, D, between two
     Numpy arrays of representations.
 
@@ -71,7 +74,7 @@ def l2_distance(A, B):
     return D
 
 
-def p_distance(A, B, p=2):
+def p_distance(A: ndarray, B: ndarray, p: Union[int, float] = 2) -> ndarray:
     """Calculates the p-norm distances between two
     Numpy arrays of representations.
     The value of the keyword argument ``p =`` sets the norm order.

src/qmllib/kernels/gradient_kernels.py

@@ -1,4 +1,7 @@
+from typing import List, Union
+
 import numpy as np
+from numpy import ndarray
 
 from qmllib.utils.environment_manipulation import (
     mkl_get_num_threads,
@@ -22,7 +25,9 @@
 )
 
 
-def get_global_kernel(X1, X2, Q1, Q2, SIGMA):
+def get_global_kernel(
+    X1: ndarray, X2: ndarray, Q1: List[List[int]], Q2: List[List[int]], SIGMA: float
+) -> ndarray:
     """Calculates the Gaussian kernel matrix K with the local decomposition where :math:`K_{ij}`:
 
         :math:`K_{ij} = \\sum_{I\\in i} \\sum_{J\\in j}\\exp \\big( -\\frac{\\|X_I - X_J\\|_2^2}{2\\sigma^2} \\big)`
@@ -75,7 +80,9 @@ def get_global_kernel(X1, X2, Q1, Q2, SIGMA):
     return K
 
 
-def get_local_kernels(X1, X2, Q1, Q2, SIGMAS):
+def get_local_kernels(
+    X1: ndarray, X2: ndarray, Q1: List[List[int]], Q2: List[List[int]], SIGMAS: List[float]
+) -> ndarray:
     """Calculates the Gaussian kernel matrix K with the local decomposition where :math:`K_{ij}`:
 
         :math:`K_{ij} = \\sum_{I\\in i} \\sum_{J\\in j}\\exp \\big( -\\frac{\\|X_I - X_J\\|_2^2}{2\\sigma^2} \\big)`
@@ -131,7 +138,13 @@ def get_local_kernels(X1, X2, Q1, Q2, SIGMAS):
     return K
 
 
-def get_local_kernel(X1, X2, Q1, Q2, SIGMA):
+def get_local_kernel(
+    X1: ndarray,
+    X2: ndarray,
+    Q1: List[Union[ndarray, List[int]]],
+    Q2: List[Union[ndarray, List[int]]],
+    SIGMA: float,
+) -> ndarray:
     """Calculates the Gaussian kernel matrix K with the local decomposition where :math:`K_{ij}`:
 
         :math:`K_{ij} = \\sum_{I\\in i} \\sum_{J\\in j}\\exp \\big( -\\frac{\\|X_I - X_J\\|_2^2}{2\\sigma^2} \\big)`
@@ -184,7 +197,7 @@ def get_local_kernel(X1, X2, Q1, Q2, SIGMA):
     return K
 
 
-def get_local_symmetric_kernels(X1, Q1, SIGMAS):
+def get_local_symmetric_kernels(X1: ndarray, Q1: List[List[int]], SIGMAS: List[float]) -> ndarray:
     """Calculates the Gaussian kernel matrix K with the local decomposition where :math:`K_{ij}`:
 
         :math:`K_{ij} = \\sum_{I\\in i} \\sum_{J\\in j}\\exp \\big( -\\frac{\\|X_I - X_J\\|_2^2}{2\\sigma^2} \\big)`
@@ -229,7 +242,9 @@ def get_local_symmetric_kernels(X1, Q1, SIGMAS):
     return K
 
 
-def get_local_symmetric_kernel(X1, Q1, SIGMA):
+def get_local_symmetric_kernel(
+    X1: ndarray, Q1: List[Union[ndarray, List[int]]], SIGMA: float
+) -> ndarray:
     """Calculates the Gaussian kernel matrix K with the local decomposition where :math:`K_{ij}`:
 
         :math:`K_{ij} = \\sum_{I\\in i} \\sum_{J\\in j}\\exp \\big( -\\frac{\\|X_I - X_J\\|_2^2}{2\\sigma^2} \\big)`
@@ -273,7 +288,13 @@ def get_local_symmetric_kernel(X1, Q1, SIGMA):
     return K
 
 
-def get_atomic_local_kernel(X1, X2, Q1, Q2, SIGMA):
+def get_atomic_local_kernel(
+    X1: ndarray,
+    X2: ndarray,
+    Q1: List[Union[ndarray, List[int]]],
+    Q2: List[Union[ndarray, List[int]]],
+    SIGMA: float,
+) -> ndarray:
 
     """Calculates the Gaussian kernel matrix K with the local decomposition where :math:`K_{ij}`:
 
@@ -332,7 +353,14 @@ def get_atomic_local_kernel(X1, X2, Q1, Q2, SIGMA):
     return K
 
 
-def get_atomic_local_gradient_kernel(X1, X2, dX2, Q1, Q2, SIGMA):
+def get_atomic_local_gradient_kernel(
+    X1: ndarray,
+    X2: ndarray,
+    dX2: ndarray,
+    Q1: List[Union[ndarray, List[int]]],
+    Q2: List[Union[ndarray, List[int]]],
+    SIGMA: float,
+) -> ndarray:
 
     """Calculates the Gaussian kernel matrix K with the local decomposition where :math:`K_{ij}`:
 
@@ -412,7 +440,9 @@ def get_atomic_local_gradient_kernel(X1, X2, dX2, Q1, Q2, SIGMA):
     return K
 
 
-def get_local_gradient_kernel(X1, X2, dX2, Q1, Q2, SIGMA):
+def get_local_gradient_kernel(
+    X1: ndarray, X2: ndarray, dX2: ndarray, Q1: List[List[int]], Q2: List[List[int]], SIGMA: float
+) -> ndarray:
 
     """Calculates the Gaussian kernel matrix K with the local decomposition where :math:`K_{ij}`:
 
@@ -478,7 +508,15 @@ def get_local_gradient_kernel(X1, X2, dX2, Q1, Q2, SIGMA):
     return K
 
 
-def get_gdml_kernel(X1, X2, dX1, dX2, Q1, Q2, SIGMA):
+def get_gdml_kernel(
+    X1: ndarray,
+    X2: ndarray,
+    dX1: ndarray,
+    dX2: ndarray,
+    Q1: List[List[int]],
+    Q2: List[List[int]],
+    SIGMA: float,
+) -> ndarray:
 
     """Calculates the Gaussian kernel matrix K with the local decomposition where :math:`K_{ij}`:
 
@@ -560,7 +598,9 @@ def get_gdml_kernel(X1, X2, dX1, dX2, Q1, Q2, SIGMA):
     return K
 
 
-def get_symmetric_gdml_kernel(X1, dX1, Q1, SIGMA):
+def get_symmetric_gdml_kernel(
+    X1: ndarray, dX1: ndarray, Q1: List[List[int]], SIGMA: float
+) -> ndarray:
 
     """Calculates the Gaussian kernel matrix K with the local decomposition where :math:`K_{ij}`:
 
@@ -613,7 +653,15 @@ def get_symmetric_gdml_kernel(X1, dX1, Q1, SIGMA):
     return K
 
 
-def get_gp_kernel(X1, X2, dX1, dX2, Q1, Q2, SIGMA):
+def get_gp_kernel(
+    X1: ndarray,
+    X2: ndarray,
+    dX1: ndarray,
+    dX2: ndarray,
+    Q1: List[Union[ndarray, List[int]]],
+    Q2: List[Union[ndarray, List[int]]],
+    SIGMA: float,
+) -> ndarray:
 
     """Calculates the Gaussian kernel matrix K with the local decomposition where :math:`K_{ij}`:
 
@@ -692,7 +740,9 @@ def get_gp_kernel(X1, X2, dX1, dX2, Q1, Q2, SIGMA):
     return K
 
 
-def get_symmetric_gp_kernel(X1, dX1, Q1, SIGMA):
+def get_symmetric_gp_kernel(
+    X1: ndarray, dX1: ndarray, Q1: List[Union[ndarray, List[int]]], SIGMA: float
+) -> ndarray:
 
     """
     This symmetric kernel corresponds to a Gaussian process regression (GPR) approach.

src/qmllib/kernels/kernels.py

@@ -1,4 +1,7 @@
+from typing import List, Union
+
 import numpy as np
+from numpy import float64, ndarray
 
 from .fkernels import (
     fgaussian_kernel,
@@ -15,7 +18,7 @@
 )
 
 
-def wasserstein_kernel(A, B, sigma, p=1, q=1):
+def wasserstein_kernel(A: ndarray, B: ndarray, sigma: float, p: int = 1, q: int = 1) -> ndarray:
     """Calculates the Wasserstein kernel matrix K, where :math:`K_{ij}`:
 
     :math:`K_{ij} = \\exp \\big( -\\frac{(W_p(A_i, B_i))^q}{\\sigma} \\big)`
@@ -45,7 +48,7 @@ def wasserstein_kernel(A, B, sigma, p=1, q=1):
     return K
 
 
-def laplacian_kernel(A, B, sigma):
+def laplacian_kernel(A: ndarray, B: ndarray, sigma: float) -> ndarray:
     """Calculates the Laplacian kernel matrix K, where :math:`K_{ij}`:
 
         :math:`K_{ij} = \\exp \\big( -\\frac{\\|A_i - B_j\\|_1}{\\sigma} \\big)`
@@ -75,7 +78,7 @@ def laplacian_kernel(A, B, sigma):
     return K
 
 
-def laplacian_kernel_symmetric(A, sigma):
+def laplacian_kernel_symmetric(A: ndarray, sigma: float) -> ndarray:
     """Calculates the symmetric Laplacian kernel matrix K, where :math:`K_{ij}`:
 
         :math:`K_{ij} = \\exp \\big( -\\frac{\\|A_i - A_j\\|_1}{\\sigma} \\big)`
@@ -102,7 +105,7 @@ def laplacian_kernel_symmetric(A, sigma):
     return K
 
 
-def gaussian_kernel(A, B, sigma):
+def gaussian_kernel(A: ndarray, B: ndarray, sigma: float) -> ndarray:
     """Calculates the Gaussian kernel matrix K, where :math:`K_{ij}`:
 
         :math:`K_{ij} = \\exp \\big( -\\frac{\\|A_i - B_j\\|_2^2}{2\\sigma^2} \\big)`
@@ -132,7 +135,7 @@ def gaussian_kernel(A, B, sigma):
     return K
 
 
-def gaussian_kernel_symmetric(A, sigma):
+def gaussian_kernel_symmetric(A: ndarray, sigma: float) -> ndarray:
     """Calculates the symmetric Gaussian kernel matrix K, where :math:`K_{ij}`:
 
         :math:`K_{ij} = \\exp \\big( -\\frac{\\|A_i - A_j\\|_2^2}{2\\sigma^2} \\big)`
@@ -159,7 +162,7 @@ def gaussian_kernel_symmetric(A, sigma):
     return K
 
 
-def linear_kernel(A, B):
+def linear_kernel(A: ndarray, B: ndarray) -> ndarray:
     """Calculates the linear kernel matrix K, where :math:`K_{ij}`:
 
         :math:`K_{ij} = A_i \\cdot B_j`
@@ -188,7 +191,12 @@ def linear_kernel(A, B):
     return K
 
 
-def sargan_kernel(A, B, sigma, gammas):
+def sargan_kernel(
+    A: ndarray,
+    B: ndarray,
+    sigma: Union[float, float64],
+    gammas: Union[ndarray, List[Union[int, float]], List[int]],
+) -> ndarray:
     """Calculates the Sargan kernel matrix K, where :math:`K_{ij}`:
 
         :math:`K_{ij} = \\exp \\big( -\\frac{\\| A_i - B_j \\|_1)}{\\sigma} \\big) \\big(1 + \\sum_{k} \\frac{\\gamma_{k} \\| A_i - B_j \\|_1^k}{\\sigma^k} \\big)`
@@ -225,7 +233,9 @@ def sargan_kernel(A, B, sigma, gammas):
     return K
 
 
-def matern_kernel(A, B, sigma, order=0, metric="l1"):
+def matern_kernel(
+    A: ndarray, B: ndarray, sigma: float, order: int = 0, metric: str = "l1"
+) -> ndarray:
     """Calculates the Matern kernel matrix K, where :math:`K_{ij}`:
 
         for order = 0:
@@ -286,7 +296,9 @@ def matern_kernel(A, B, sigma, order=0, metric="l1"):
     return K
 
 
-def get_local_kernels_gaussian(A, B, na, nb, sigmas):
+def get_local_kernels_gaussian(
+    A: ndarray, B: ndarray, na: ndarray, nb: ndarray, sigmas: List[float]
+) -> ndarray:
     """Calculates the Gaussian kernel matrix K, for a local representation where :math:`K_{ij}`:
 
         :math:`K_{ij} = \\sum_{a \\in i} \\sum_{b \\in j} \\exp \\big( -\\frac{\\|A_a - B_b\\|_2^2}{2\\sigma^2} \\big)`
@@ -328,7 +340,9 @@ def get_local_kernels_gaussian(A, B, na, nb, sigmas):
     return fget_local_kernels_gaussian(A.T, B.T, na, nb, sigmas, nma, nmb, nsigmas)
 
 
-def get_local_kernels_laplacian(A, B, na, nb, sigmas):
+def get_local_kernels_laplacian(
+    A: ndarray, B: ndarray, na: ndarray, nb: ndarray, sigmas: List[float]
+) -> ndarray:
     """Calculates the Local Laplacian kernel matrix K, for a local representation where :math:`K_{ij}`:
 
         :math:`K_{ij} = \\sum_{a \\in i} \\sum_{b \\in j} \\exp \\big( -\\frac{\\|A_a - B_b\\|_1}{\\sigma} \\big)`
@@ -370,7 +384,7 @@ def get_local_kernels_laplacian(A, B, na, nb, sigmas):
     return fget_local_kernels_laplacian(A.T, B.T, na, nb, sigmas, nma, nmb, nsigmas)
 
 
-def kpca(K, n=2, centering=True):
+def kpca(K: ndarray, n: int = 2, centering: bool = True) -> ndarray:
     """Calculates `n` first principal components for the kernel :math:`K`.
 
     The PCA is calculated using an OpenMP parallel Fortran routine.