Reimplementation of cython code in numba

.gitignore

@@ -1,7 +1,141 @@
-__pycache__
 poetry.lock
-.ipynb_checkpoints*
 testing_notebooks
 outs
 slurm
-.vscode
+
+# DS_Store
+.DS_Store
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# vscode
+.vscode/settings.json
+.vscode

.pre-commit-config.yaml

@@ -0,0 +1,50 @@
+fail_fast: false
+default_language_version:
+    python: python3
+default_stages:
+    - commit
+    - push
+minimum_pre_commit_version: 2.16.0
+repos:
+    - repo: https://github.com/psf/black
+      rev: "23.3.0"
+      hooks:
+          - id: black
+    - repo: https://github.com/asottile/blacken-docs
+      rev: 1.13.0
+      hooks:
+          - id: blacken-docs
+    - repo: https://github.com/pre-commit/mirrors-prettier
+      rev: v3.0.0-alpha.6
+      hooks:
+          - id: prettier
+            # Newer versions of node don't work on systems that have an older version of GLIBC
+            # (in particular Ubuntu 18.04 and Centos 7)
+            # EOL of Centos 7 is in 2024-06, we can probably get rid of this then.
+            # See https://github.com/scverse/cookiecutter-scverse/issues/143 and
+            # https://github.com/jupyterlab/jupyterlab/issues/12675
+            language_version: "17.9.1"
+    - repo: https://github.com/charliermarsh/ruff-pre-commit
+      rev: v0.0.261
+      hooks:
+          - id: ruff
+            args: [--fix, --exit-non-zero-on-fix]
+    - repo: https://github.com/pre-commit/pre-commit-hooks
+      rev: v4.4.0
+      hooks:
+          - id: detect-private-key
+          - id: check-ast
+          - id: end-of-file-fixer
+          - id: mixed-line-ending
+            args: [--fix=lf]
+          - id: trailing-whitespace
+          - id: check-case-conflict
+    - repo: local
+      hooks:
+          - id: forbid-to-commit
+            name: Don't commit rej files
+            entry: |
+                Cannot commit .rej files. These indicate merge conflicts that arise during automated template updates.
+                Fix the merge conflicts manually and remove the .rej files.
+            language: fail
+            files: '.*\.rej$'

balanced_clustering/__init__.py

@@ -1,5 +1,11 @@
-__version__ = "0.1.0"
+# https://github.com/python-poetry/poetry/pull/2366#issuecomment-652418094
+# https://github.com/python-poetry/poetry/issues/144#issuecomment-623927302
+import importlib.metadata as importlib_metadata
+
+package_name = "balanced-clustering"
+__version__ = importlib_metadata.version(package_name)
+
 from .ari import balanced_adjusted_rand_index
 from .ami import balanced_adjusted_mutual_info
 from .vmeasure import balanced_homogeneity, balanced_completeness, balanced_v_measure
-from .return_metrics import return_metrics
+from .return_metrics import return_metrics

balanced_clustering/ami.py

@@ -84,17 +84,13 @@ def balanced_adjusted_mutual_info(
     )
     # Recalculate labels_true and labels_pred if reweigh is True to
     # factor in the reweighting based on the true class frequencies.
-    # These won't preserve order but this is fine since entropy is 
+    # These won't preserve order but this is fine since entropy is
     # invariant to order
     if reweigh is True:
-        true_sums = np.squeeze(np.asarray(sp.csc_matrix.sum(contingency, axis = 1)))
-        pred_sums = np.squeeze(np.asarray(sp.csc_matrix.sum(contingency, axis = 0)))
-        labels_true = np.repeat(
-            np.arange(len(true_sums)), true_sums
-        )
-        labels_pred = np.repeat(
-            np.arange(len(pred_sums)), pred_sums
-        )
+        true_sums = np.squeeze(np.asarray(sp.csc_matrix.sum(contingency, axis=1)))
+        pred_sums = np.squeeze(np.asarray(sp.csc_matrix.sum(contingency, axis=0)))
+        labels_true = np.repeat(np.arange(len(true_sums)), true_sums)
+        labels_pred = np.repeat(np.arange(len(pred_sums)), pred_sums)
     contingency = contingency.astype(np.float64)
     # Calculate the MI for the two clusterings
     mi = mutual_info_score(labels_true, labels_pred, contingency=contingency)

balanced_clustering/return_metrics.py

@@ -1,15 +1,21 @@
-from sklearn.metrics import adjusted_rand_score, adjusted_mutual_info_score, \
-    homogeneity_score, completeness_score, v_measure_score
+from sklearn.metrics import (
+    adjusted_rand_score,
+    adjusted_mutual_info_score,
+    homogeneity_score,
+    completeness_score,
+    v_measure_score,
+)
 
 from .ari import balanced_adjusted_rand_index
 from .ami import balanced_adjusted_mutual_info
 from .vmeasure import balanced_homogeneity, balanced_completeness, balanced_v_measure
 
-def return_metrics(class_arr, cluster_arr, print_metrics = True):
-    '''
-    Compare imbalanced and balanced ARI, AMI, homogeneity, completeness, and 
+
+def return_metrics(class_arr, cluster_arr, print_metrics=True):
+    """
+    Compare imbalanced and balanced ARI, AMI, homogeneity, completeness, and
     V-measure scores.
-    
+
     Parameters
     ----------
     class_arr : int array-like of shape (n_samples,)
@@ -19,7 +25,7 @@ def return_metrics(class_arr, cluster_arr, print_metrics = True):
     print_metrics : bool, default=True
         If True, print the scores.
         If False, return the scores.
-    
+
     Returns
     -------
     ari_imbalanced : float
@@ -42,9 +48,9 @@ def return_metrics(class_arr, cluster_arr, print_metrics = True):
         The imbalanced V-measure score.
     v_measure_balanced : float
         The balanced V-measure score.
-    '''
-    
-    # Determine the imbalanced (base) metric scores 
+    """
+
+    # Determine the imbalanced (base) metric scores
     ari_imbalanced = adjusted_rand_score(class_arr, cluster_arr)
     ami_imbalanced = adjusted_mutual_info_score(class_arr, cluster_arr)
     homog_imbalanced = homogeneity_score(class_arr, cluster_arr)
@@ -57,31 +63,50 @@ def return_metrics(class_arr, cluster_arr, print_metrics = True):
     homog_balanced = balanced_homogeneity(class_arr, cluster_arr)
     complete_balanced = balanced_completeness(class_arr, cluster_arr)
     v_measure_balanced = balanced_v_measure(class_arr, cluster_arr)
-    
+
     # If print is True, print the scores
     if print_metrics:
         print(
-            "ARI imbalanced: " + str(round(ari_imbalanced, 4)) + " " + 
-            "ARI balanced: " + str(round(ari_balanced, 4))
+            "ARI imbalanced: "
+            + str(round(ari_imbalanced, 4))
+            + " "
+            + "ARI balanced: "
+            + str(round(ari_balanced, 4))
         )
         print(
-            "AMI imbalanced: " + str(round(ami_imbalanced, 4)) + " " +
-            "AMI balanced: " + str(round(ami_balanced, 4))
+            "AMI imbalanced: "
+            + str(round(ami_imbalanced, 4))
+            + " "
+            + "AMI balanced: "
+            + str(round(ami_balanced, 4))
         )
         print(
-            "Homogeneity imbalanced: " + str(round(homog_imbalanced, 4)) + " " +
-            "Homogeneity balanced: " + str(round(homog_balanced, 4))
+            "Homogeneity imbalanced: "
+            + str(round(homog_imbalanced, 4))
+            + " "
+            + "Homogeneity balanced: "
+            + str(round(homog_balanced, 4))
         )
         print(
-            "Completeness imbalanced: " + str(round(complete_imbalanced, 4)) + " " +
-            "Completeness balanced : " + str(round(complete_balanced, 4))
+            "Completeness imbalanced: "
+            + str(round(complete_imbalanced, 4))
+            + " "
+            + "Completeness balanced : "
+            + str(round(complete_balanced, 4))
         )
         print(
-            "V-measure imbalanced: " + str(round(v_measure_imbalanced, 4)) + " " +
-            "V-measure balanced: " + str(round(v_measure_balanced, 4))
+            "V-measure imbalanced: "
+            + str(round(v_measure_imbalanced, 4))
+            + " "
+            + "V-measure balanced: "
+            + str(round(v_measure_balanced, 4))
         )
-    
+
     # Return paired balanced imbalanced scores
-    return (ari_imbalanced, ari_balanced), (ami_imbalanced, ami_balanced), \
-        (homog_imbalanced, homog_balanced), (complete_imbalanced, complete_balanced), \
-        (v_measure_imbalanced, v_measure_balanced)
+    return (
+        (ari_imbalanced, ari_balanced),
+        (ami_imbalanced, ami_balanced),
+        (homog_imbalanced, homog_balanced),
+        (complete_imbalanced, complete_balanced),
+        (v_measure_imbalanced, v_measure_balanced),
+    )

balanced_clustering/utils/__init__.py

@@ -1,8 +1,6 @@
-import pyximport
 import numpy
 
-pyximport.install(setup_args={"include_dirs": numpy.get_include()}, reload_support=True)
-from ._emi_cython import expected_mutual_information
+from ._emi import expected_mutual_information
 from .contingency import pair_confusion_matrix, contingency_matrix
 from .checks import check_clusterings
 from .mi import mutual_info_score, entropy