Make CentroidLayer inherit from Layer

Author: althonos

.github/workflows/test.yml

@@ -112,6 +112,7 @@ jobs:
   upload:
     environment: PyPI
     runs-on: ubuntu-latest
+    if: startsWith(github.ref, 'refs/tags/v')
     name: Upload
     needs:
     - sdist
@@ -123,7 +124,6 @@ jobs:
           name: dist
           path: dist
     - name: Publish distributions to PyPI
-      if: startsWith(github.ref, 'refs/tags/v')
       uses: pypa/gh-action-pypi-publish@master
       with:
         user: __token__

README.md

@@ -19,10 +19,11 @@
 
 ## 🗺️ Overview
 
-[`foldseek`] is a method developed by van Kempen *et al.*[\[1\]](#ref1) for the fast
-and accurate search of protein structures. In order to search proteins structures
-at a large scale, it first encodes the 3D structure into sequences over a
-structural alphabet, 3di, which captures tertiary amino acid interactions.
+[`foldseek`](https://github.com/steineggerlab/foldseek) is a method developed
+by van Kempen *et al.*[\[1\]](#ref1) for the fast and accurate search of
+protein structures. In order to search proteins structures at a large scale,
+it first encodes the 3D structure into sequences over a structural alphabet,
+3di, which captures tertiary amino acid interactions.
 
 `mini3di` is a pure-Python package to encode 3D structures of proteins into
 the 3di alphabet, using the trained weights from the `foldseek` VQ-VAE model.

mini3di/layers.py

@@ -46,13 +46,13 @@ def __call__(self, X: ArrayNxM[numpy.floating]) -> ArrayNxM[numpy.floating]:
             return out
 
 
-class CentroidLayer:
+class CentroidLayer(Layer):
     def __init__(self, centroids: ArrayNxM[numpy.floating]) -> None:
-        self.centroids = centroids
+        self.centroids = numpy.asarray(centroids)
         self.r2 = numpy.sum(self.centroids**2, axis=1).reshape(-1, 1).T
     def __call__(self, X: ArrayNxM[numpy.floating]) -> ArrayN[numpy.uint8]:
         # compute pairwise squared distance matrix
-        r1 = numpy.sum(X * X, 1).reshape(-1, 1)
+        r1 = numpy.sum(X * X, axis=1).reshape(-1, 1)
         D = r1 - 2 * X @ self.centroids.T + self.r2
         # find closest centroid
         states = numpy.empty(D.shape[0], dtype=numpy.uint8)