From 62dfd4fcf2758d5940f7364d81a7bc9627bbb2a1 Mon Sep 17 00:00:00 2001
From: Martin Larralde <martin.larralde@embl.de>
Date: Sat, 25 Nov 2023 22:20:25 +0100
Subject: [PATCH] Make `CentroidLayer` inherit from `Layer`

---
 .github/workflows/test.yml | 2 +-
 README.md                  | 9 +++++----
 mini3di/layers.py          | 6 +++---
 3 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 440fecb..5df9972 100644
--- a/.github/workflows/test.yml
+++ b/.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__
diff --git a/README.md b/README.md
index ac6c366..85ba5fc 100644
--- a/README.md
+++ b/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.
diff --git a/mini3di/layers.py b/mini3di/layers.py
index cf0de1f..023aed8 100644
--- a/mini3di/layers.py
+++ b/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)