Repository: motiwari/BanditPAM (language: python)

Issue:
User report: algorithm choosing wrong arms (correctness)
A report from a user, I've asked for his code:

> I compared the BUILD step solution's objective of BanditPAM with that of the BUILD step solution's objective of PAM (on MNIST dataset, train split, N = 60k, d = 784). For the number of medoids k = 100 and L2 distance, PAM's BUILD objective is 6.0641 while that of BanditPAM's BUILD is 6.55 (around 8% worse). I have observed such differences for various other randomly generated datasets as well. For the MNIST dataset, I had set build_confidence = 16. According to my understanding of BanditPAM's paper and code, this should correspond to delta = (N*exp(build_confidence))^{-1} = 1.8756e-12. Hence, according to Theorem 1's Remark A1 of BanditPAM paper, the probability that all confidence intervals are true confidence intervals is at least 1 - 2N^2*delta = 0.9932. My understanding is that this probability should be high enough for me to observe the same BUILD objective for both BanditPAM and PAM for most random seeds. However, it does not happen for any seed. 


> Please note that my understanding of the theoretical results or the way I used build_confidence may be wrong. Hence, it would be great if you could resolve the above issue. Overall, based on the results shown in the BanditPAM paper, I expected the objectives to be the same in most of my runs (on various datasets with 1 - 2N^2*delta > 0.99, computed through build_confidence as explained above). However, it does not happen and the objectives are quite far apart (e.g., 8% difference). 
Change github actions that upload wheels to upload to TestPyPI on PR update, and PyPI on release
Right now, both TestPyPI and PyPI actions trigger on release and we govern which one executes by commenting/uncommenting out the correct lines of code. We should automate this.

[Example](https://github.com/motiwari/BanditPAM/blob/284fea66e909d247b11c972b71efc82ce3e3319a/.github/workflows/build_linux_wheels.yml#L71-L78)
Fixed Loss Mismatch Issues in BanditPAM
- Loss issues are resolved in all versions of BanditPAM
- The clustering results are now identical to scikit-learn's implementation

<img width="636" alt="image" src="https://github.com/motiwari/BanditPAM/assets/24787637/5b33050b-0128-4491-b710-1d4f27e4ce6d">

<img width="636" alt="image" src="https://github.com/motiwari/BanditPAM/assets/24787637/ede43dcf-95b7-42d9-a0b0-acddd1c97721">


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/build_linux_wheels.yml b/.github/workflows/build_linux_wheels.yml
index bedd0efb..d09ebed1 100644
--- a/.github/workflows/build_linux_wheels.yml
+++ b/.github/workflows/build_linux_wheels.yml
@@ -1,6 +1,8 @@
 name: Linux - build source distribution and wheels
 on:
   push:
+    branches: # prevents running on tag push
+      - '**'
   pull_request:
   release:
     types:
@@ -60,19 +62,23 @@ jobs:
   upload_pypi:
     needs: [build_wheels, build_sdist]
     runs-on: ubuntu-latest
-    # TODO: RBF
-    if: github.event_name == 'release' && github.event.action == 'published'
     steps:
       - uses: actions/download-artifact@v2
         with:
           name: artifact
           path: dist
 
-      - uses: pypa/gh-action-pypi-publish@v1.4.2
+      - name: Upload to TestPyPI on PR update
+        if: ${{ github.event_name == 'pull_request' }}
+        uses: pypa/gh-action-pypi-publish@v1.4.2
+        with:
+          user: __token__
+          password: ${{ secrets.TEST_PYPI_APIKEY }}
+          repository_url: https://test.pypi.org/legacy/
+
+      - name: Upload to PyPI on published release
+        if: ${{ github.event_name == 'release' && github.event.action == 'published' }}
+        uses: pypa/gh-action-pypi-publish@v1.4.2
         with:
           user: __token__
-          # For PyPI
           password: ${{ secrets.PYPI_APIKEY }}
-          # For TestPyPI. TODO: RBF
-#          password: ${{ secrets.TEST_PYPI_APIKEY }}
-#          repository_url: https://test.pypi.org/legacy/
diff --git a/.github/workflows/build_mac_wheels.yml b/.github/workflows/build_mac_wheels.yml
index f060448c..672c96b1 100644
--- a/.github/workflows/build_mac_wheels.yml
+++ b/.github/workflows/build_mac_wheels.yml
@@ -1,6 +1,8 @@
 name: MacOS - build wheels
 on: 
   push:
+    branches: # prevents running on tag push
+      - '**'
   pull_request:
   release:
     types:
@@ -33,19 +35,23 @@ jobs:
   upload_pypi:
     needs: [build_wheels_macos]
     runs-on: ubuntu-latest # Can only upload from Linux containers
-    # TODO: RBF
-    if: github.event_name == 'release' && github.event.action == 'published'
     steps:
       - uses: actions/download-artifact@v2
         with:
           name: artifact
           path: dist
 
-      - uses: pypa/gh-action-pypi-publish@v1.4.2
+      - name: Upload to TestPyPI on PR update
+        if: ${{ github.event_name == 'pull_request' }}
+        uses: pypa/gh-action-pypi-publish@v1.4.2
         with:
           user: __token__
-          # For PyPI
-          password: ${{ secrets.PYPI_APIKEY }}
-          # For TestPyPI. TODO: RBF
-#          password: ${{ secrets.TEST_PYPI_APIKEY }}
-#          repository_url: https://test.pypi.org/legacy/
\ No newline at end of file
+          password: ${{ secrets.TEST_PYPI_APIKEY }}
+          repository_url: https://test.pypi.org/legacy/
+
+      - name: Upload to PyPI on published release
+        if: ${{ github.event_name == 'release' && github.event.action == 'published' }}
+        uses: pypa/gh-action-pypi-publish@v1.4.2
+        with:
+          user: __token__
+          password: ${{ secrets.PYPI_APIKEY }}
\ No newline at end of file
diff --git a/.github/workflows/run_cmake_build.yml b/.github/workflows/run_cmake_build.yml
index db6bf557..07b1e8a7 100644
--- a/.github/workflows/run_cmake_build.yml
+++ b/.github/workflows/run_cmake_build.yml
@@ -1,5 +1,9 @@
 name: Run CMake build on MacOS
-on: [push, pull_request]
+on:
+  push:
+    branches:  # prevents running on tag push
+      - '**'
+  pull_request:
 jobs:
   build:
     runs-on: macos-latest
diff --git a/.github/workflows/run_style_checks.yml b/.github/workflows/run_style_checks.yml
index f2411375..1e2ba024 100644
--- a/.github/workflows/run_style_checks.yml
+++ b/.github/workflows/run_style_checks.yml
@@ -1,5 +1,9 @@
 name: Run style checks
-on: [push, pull_request]
+on:
+  push:
+    branches:  # prevents running on tag push
+      - '**'
+  pull_request:
 jobs:
   build:
     runs-on: ubuntu-latest
diff --git a/.gitignore b/.gitignore
index d76e0d0c..cb2e997b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -10,11 +10,11 @@ BanditPAM.egg-info/
 __pycache__
 banditpam.cpython*
 build/
-build/
 data/MNIST_*.csv
-data/MNIST_CSV.csv
 data/data.csv
 data/scrna*
+data/*.py
+data/cifar*
 docs/html
 docs/latex
 python_generate_plots/__pycache__
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 562119be..1dcb40cd 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -4,6 +4,7 @@ project(BanditPAM VERSION 1.0 LANGUAGES CXX)
 
 # TODO(@motiwari): Should RELEASE? be in caps
 set(CMAKE_BUILD_TYPE Release)
+# set(CMAKE_BUILD_TYPE Debug)
 set(CMAKE_CXX_STANDARD 17)
 
 set(CMAKE_CXX_FLAGS_RELEASE "-O3")
@@ -11,9 +12,9 @@ set(CMAKE_CXX_FLAGS_RELEASE "-O3")
 # Note: ThreadSanitizer and AddressSanitizer can give spurious errors around the #pragma omp critical.
 # They don't really matter because the threads always resolve with the same end state
 # ThreadSanitizer also introduces a very large overhead, AddressSanitizer is more reasonable
-set(CMAKE_CXX_FLAGS_DEBUG "-Wall -Wextra -g -fno-omit-frame-pointer -fsanitize=address")
+# set(CMAKE_CXX_FLAGS_DEBUG "-Wall -Wextra -g -fno-omit-frame-pointer -fsanitize=address")
 
 # NOTE: Need to explicitly pass -O0 to enable printing of armadillo matrices
-#set(CMAKE_CXX_FLAGS_DEBUG "-O0 -Wall -Wextra -g -fno-omit-frame-pointer")
+# set(CMAKE_CXX_FLAGS_DEBUG "-O0 -Wall -Wextra -g -fno-omit-frame-pointer")
 
 add_subdirectory(src)
diff --git a/docs/install_windows.md b/docs/install_windows.md
index 6822cf1c..5a575625 100644
--- a/docs/install_windows.md
+++ b/docs/install_windows.md
@@ -15,7 +15,7 @@ Please ensure the following dependencies are installed:
 
 CMake Build:
 
-1) Run `scripts/retrieve_windows_cmake_files.sh` to retrieve the files necessary for the Windows CMake build
+1) Run `cd scripts` and then `sh retrieve_windows_cmake_files.sh` to retrieve the files necessary for the Windows CMake build
 2) Add the Visual Studio IDE location to PATH in Environment Variables (e.g. `C:\Program Files\Microsoft Visual Studio\2022\Community\Common7\IDE`).
 3) Add `${project_dir}\headers\armadillo\examples\lib_win64` to PATH 
 4) Run `devenv BanditPAM.sln /Build "Release|x64"` after `cmake ..` in a prompt other than Git Bash
diff --git a/scripts/cache_measurements.py b/scripts/cache_measurements.py
index 1fa6f9c9..dd1b3738 100644
--- a/scripts/cache_measurements.py
+++ b/scripts/cache_measurements.py
@@ -55,22 +55,34 @@ def test_cache_stats():
             hits_1000,
             misses_1000,
         ) = get_cache_statistics(kmed=kmed, X=X, loss="L2", cache_width=1000)
-        time_750, width_750, writes_750, hits_750, misses_750 = \
-            get_cache_statistics(
-                kmed=kmed, X=X, loss="L2", cache_width=750
-            )
-        time_500, width_500, writes_500, hits_500, misses_500 = \
-            get_cache_statistics(
-                kmed=kmed, X=X, loss="L2", cache_width=500
-            )
-        time_250, width_250, writes_250, hits_250, misses_250 = \
-            get_cache_statistics(
-                kmed=kmed, X=X, loss="L2", cache_width=250
-            )
-        time_0, width_0, writes_0, hits_0, misses_0 = \
-            get_cache_statistics(
-                kmed=kmed, X=X, loss="L2", cache_width=0
-            )
+        (
+            time_750,
+            width_750,
+            writes_750,
+            hits_750,
+            misses_750,
+        ) = get_cache_statistics(kmed=kmed, X=X, loss="L2", cache_width=750)
+        (
+            time_500,
+            width_500,
+            writes_500,
+            hits_500,
+            misses_500,
+        ) = get_cache_statistics(kmed=kmed, X=X, loss="L2", cache_width=500)
+        (
+            time_250,
+            width_250,
+            writes_250,
+            hits_250,
+            misses_250,
+        ) = get_cache_statistics(kmed=kmed, X=X, loss="L2", cache_width=250)
+        (
+            time_0,
+            width_0,
+            writes_0,
+            hits_0,
+            misses_0,
+        ) = get_cache_statistics(kmed=kmed, X=X, loss="L2", cache_width=0)
 
         assert (
             hits_1000 > hits_750 > hits_500 > hits_250 > hits_0
@@ -92,8 +104,7 @@ def test_cache_stats():
         assert width_250 == 250, "Cache width should be 250 when set to 250"
         assert width_500 == 500, "Cache width should be 500 when set to 500"
         assert width_750 == 750, "Cache width should be 750 when set to 750"
-        assert width_1000 == 1000, "Cache width should be 1000 when set to " \
-                                   "1000"
+        assert width_1000 == 1000, "Cache width should be 1k when set to 1k"
 
     def test_parallelization():
         X = np.loadtxt(os.path.join("data", "MNIST_10k.csv"))
diff --git a/scripts/comparison_utils.py b/scripts/comparison_utils.py
index a9a9a9ed..03a4e4a1 100644
--- a/scripts/comparison_utils.py
+++ b/scripts/comparison_utils.py
@@ -1,4 +1,10 @@
-def print_results(kmed, runtime):
+from banditpam import KMedoids
+
+
+def print_results(kmed: KMedoids, runtime: float):
+    complexity_with_caching = (
+        kmed.getDistanceComputations(True) - kmed.cache_hits
+    )
     print("-----Results-----")
     print("Algorithm:", kmed.algorithm)
     print("Final Medoids:", kmed.medoids)
@@ -7,17 +13,24 @@ def print_results(kmed, runtime):
     print("Build complexity:", f"{kmed.build_distance_computations:,}")
     print("Swap complexity:", f"{kmed.swap_distance_computations:,}")
     print("Number of Swaps", kmed.steps)
+    print("Cache Writes: {:,}".format(kmed.cache_writes))
+    print("Cache Hits: {:,}".format(kmed.cache_hits))
+    print("Cache Misses: {:,}".format(kmed.cache_misses))
     print(
         "Average Swap Sample Complexity:",
         f"{kmed.swap_distance_computations / kmed.steps:,}",
     )
     print(
         "Total complexity (without misc):",
-        f"{kmed.getDistanceComputations(False):,}"
+        f"{kmed.getDistanceComputations(False):,}",
     )
     print(
         "Total complexity (with misc):",
         f"{kmed.getDistanceComputations(True):,}",
     )
-    print("Runtime per swap:", runtime / kmed.steps)
+    print(
+        "Total complexity (with caching):",
+        f"{complexity_with_caching:,}",
+    )
     print("Total runtime:", runtime)
+    print("Runtime per swap:", runtime / kmed.steps)
diff --git a/scripts/comparison_with_fasterpam.py b/scripts/comparison_with_fasterpam.py
index 9d6043be..80cc39d8 100644
--- a/scripts/comparison_with_fasterpam.py
+++ b/scripts/comparison_with_fasterpam.py
@@ -23,11 +23,7 @@ def benchmark(data, f, n=1):
         v = np.array(v)
         min, avg = v.min(), v.mean()
         ste = v.std(ddof=1) / np.sqrt(len(v)) if len(v) > 1 else 0.0
-        print(
-            "{:16s} min={:-10.2f} mean={:-10.2f} ±{:-.2f}".format(
-                k, min, avg, ste
-            )
-        )
+        print("{:16s} min={:-10.2f} mean={:-10.2f} ±{:-.2f}".format(k, min, avg, ste))
 
 
 def run_fasterpam(data, seed):
@@ -69,7 +65,7 @@ def run_old_bandit(data, seed):
         n_medoids=5,
         parallelize=True,
         algorithm="BanditPAM_orig",
-        dist_mat=diss
+        dist_mat=diss,
     )
     print(km.algorithm)
     km.seed = seed
diff --git a/scripts/retrieve_windows_cmake_files.sh b/scripts/retrieve_windows_cmake_files.sh
index c3afa432..244c771c 100644
--- a/scripts/retrieve_windows_cmake_files.sh
+++ b/scripts/retrieve_windows_cmake_files.sh
@@ -1,25 +1,11 @@
 #!/bin/bash
 
 # This script is used to retrieve the files necessary for the Windows CMake build
-
 echo "Retrieving files for Windows CMake build..."
-
-# Move up one directory
 cd ..
-
-# Clone the GitHub repository
 git clone https://github.com/ThrunGroup/BanditPAM_Windows
-
-# Move unistd.h to headers/unistd.h
 mv BanditPAM_Windows/unistd.h headers/unistd.h
-#
-# Move getopt.h to headers/getopt.h
 mv BanditPAM_Windows/getopt.h headers/getopt.h
-#
-# Move getopt.cpp to src/getopt.cpp
 mv BanditPAM_Windows/getopt.cpp src/getopt.cpp
-
-# Delete BanditPAM_Windows directory
 rm -rf BanditPAM_Windows
-
 echo "Done!"
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 9ac7d8e4..503109e7 100644
--- a/setup.py
+++ b/setup.py
@@ -9,7 +9,7 @@
 import distutils.sysconfig
 import distutils.spawn
 
-__version__ = "4.0.2"
+__version__ = "4.0.3a2"
 
 # TODO(@motiwari): Move this to a separate file
 GHA = "GITHUB_ACTIONS"
@@ -51,9 +51,7 @@ def compiler_check():
     """
     try:
         return (
-            "clang"
-            if "clang" in distutils.sysconfig.get_config_vars()["CC"]
-            else "gcc"
+            "clang" if "clang" in distutils.sysconfig.get_config_vars()["CC"] else "gcc"
         )
     except KeyError:
         # The 'CC' environment variable hasn't been set.
@@ -67,8 +65,7 @@ def compiler_check():
             return "gcc"
 
     raise Exception(
-        "No C++ compiler was found. Please ensure you have "
-        "MSVC, LLVM clang, or GCC."
+        "No C++ compiler was found. Please ensure you have " "MSVC, LLVM clang, or GCC."
     )
 
 
@@ -268,8 +265,7 @@ def setup_colab(delete_source=False):
         repo_location = os.path.join("/", "content", "BanditPAM")
         # Note the space after the git URL to separate the source and target
         os.system(
-            "git clone https://github.com/motiwari/BanditPAM.git "
-            + repo_location
+            "git clone https://github.com/motiwari/BanditPAM.git " + repo_location
         )
         os.system(
             repo_location
@@ -416,8 +412,10 @@ def build_extensions(self):
 
         for ext in self.extensions:
             ext.define_macros = [
-                ("VERSION_INFO",
-                 '"{}"'.format(self.distribution.get_version()))
+                (
+                    "VERSION_INFO",
+                    '"{}"'.format(self.distribution.get_version()),
+                )
             ]
             ext.extra_compile_args = opts
             ext.extra_compile_args += []  # []["-arch", "x86_64"]
@@ -455,9 +453,7 @@ def main():
             # To include carma when the BanditPAM repo hasnt been initialized
             os.path.join("/", "usr", "local", "include"),
             os.path.join("/", "usr", "local", "include", "carma"),
-            os.path.join(
-                "/", "usr", "local", "include", "carma", "carma_bits"
-            ),
+            os.path.join("/", "usr", "local", "include", "carma", "carma_bits"),
             # When building from source on M1 Macs, may need these dirs
             # Currently, we should never be building from source on an M1 Mac,
             # Only cross-compiling from an Intel Mac
@@ -468,22 +464,21 @@ def main():
             os.path.join("/", "opt", "homebrew", "lib"),
             os.path.join("/", "opt", "homebrew", "opt"),
             os.path.join("/", "opt", "homebrew", "opt", "armadillo"),
+            os.path.join("/", "opt", "homebrew", "opt", "armadillo", "include"),
             os.path.join(
-                "/", "opt", "homebrew", "opt", "armadillo", "include"
-            ),
-            os.path.join(
-                "/", "opt", "homebrew", "opt", "armadillo", "include",
-                "armadillo_bits"
+                "/",
+                "opt",
+                "homebrew",
+                "opt",
+                "armadillo",
+                "include",
+                "armadillo_bits",
             ),
             # Needed for Mac Github Runners
             # for macos-10.15
-            os.path.join(
-                "/", "usr", "local", "Cellar", "libomp", "15.0.2", "include"
-            ),
+            os.path.join("/", "usr", "local", "Cellar", "libomp", "15.0.2", "include"),
             # for macos-latest
-            os.path.join(
-                "/", "usr", "local", "Cellar", "libomp", "15.0.7", "include"
-            ),
+            os.path.join("/", "usr", "local", "Cellar", "libomp", "15.0.7", "include"),
         ]
     elif sys.platform == "win32":  # WIN32
         include_dirs = [
@@ -526,12 +521,8 @@ def main():
         ]  # TODO(@motiwari): Modify this based on gcc or clang
         library_dirs = [
             os.path.join("/", "usr", "local", "lib"),
-            os.path.join(
-                "/", "usr", "local", "Cellar", "libomp", "15.0.2", "lib"
-            ),
-            os.path.join(
-                "/", "usr", "local", "Cellar", "libomp", "15.0.7", "lib"
-            ),
+            os.path.join("/", "usr", "local", "Cellar", "libomp", "15.0.2", "lib"),
+            os.path.join("/", "usr", "local", "Cellar", "libomp", "15.0.7", "lib"),
         ]
         if sys.platform == "darwin" and platform.processor() == "arm":  # M1
             library_dirs.append(
@@ -550,21 +541,15 @@ def main():
                 os.path.join("src", "algorithms", "banditpam.cpp"),
                 os.path.join("src", "algorithms", "banditpam_orig.cpp"),
                 os.path.join("src", "algorithms", "fastpam1.cpp"),
-                os.path.join(
-                    "src", "python_bindings", "kmedoids_pywrapper.cpp"
-                ),
+                os.path.join("src", "python_bindings", "kmedoids_pywrapper.cpp"),
                 os.path.join("src", "python_bindings", "medoids_python.cpp"),
-                os.path.join(
-                    "src", "python_bindings", "build_medoids_python.cpp"
-                ),
+                os.path.join("src", "python_bindings", "build_medoids_python.cpp"),
                 os.path.join("src", "python_bindings", "fit_python.cpp"),
                 os.path.join("src", "python_bindings", "labels_python.cpp"),
                 os.path.join("src", "python_bindings", "steps_python.cpp"),
                 os.path.join("src", "python_bindings", "loss_python.cpp"),
                 os.path.join("src", "python_bindings", "cache_python.cpp"),
-                os.path.join(
-                    "src", "python_bindings", "swap_times_python.cpp"
-                ),
+                os.path.join("src", "python_bindings", "swap_times_python.cpp"),
             ],
             include_dirs=include_dirs,
             library_dirs=library_dirs,
@@ -586,9 +571,8 @@ def main():
                 [
                     os.path.join(
                         os.getcwd(),
-                        r"headers\armadillo\examples\lib_win64"
-                        + r"\libopenblas.dll",
-                        )
+                        r"headers\armadillo\examples\lib_win64" + r"\libopenblas.dll",
+                    )
                 ],
             )
         )
@@ -619,9 +603,7 @@ def main():
             os.path.join("headers", "algorithms", "banditpam.hpp"),
             os.path.join("headers", "algorithms", "fastpam1.hpp"),
             os.path.join("headers", "algorithms", "pam.hpp"),
-            os.path.join(
-                "headers", "python_bindings", "kmedoids_pywrapper.hpp"
-            ),
+            os.path.join("headers", "python_bindings", "kmedoids_pywrapper.hpp"),
         ],
     )
 
diff --git a/src/algorithms/banditpam.cpp b/src/algorithms/banditpam.cpp
index aefbc271..4e2ff577 100644
--- a/src/algorithms/banditpam.cpp
+++ b/src/algorithms/banditpam.cpp
@@ -69,6 +69,10 @@ namespace km {
           const arma::frowvec &bestDistances,
           const bool useAbsolute) {
     size_t N = data.n_cols;
+    // Temporarily increase the batch size
+    // for precise estimation of the standard deviation
+    int originalBatchSize = batchSize;
+    batchSize = std::min(1000, static_cast<int>(N/4));
     arma::uvec referencePoints;
     // TODO(@motiwari): Make this wraparound properly as
     //  last batch_size elements are dropped
@@ -104,6 +108,9 @@ namespace km {
       }
       updated_sigma(i) = arma::stddev(sample);
     }
+    // reset batchSize to the original batch size as it's a global variable
+    // used by other functions (e.g. buildTarget, swapTarget)
+    batchSize = originalBatchSize;
     return updated_sigma;
   }
 
@@ -179,6 +186,7 @@ namespace km {
     arma::frowvec ucbs(N);
     arma::frowvec numSamples(N, arma::fill::zeros);
     arma::frowvec exactMask(N, arma::fill::zeros);
+    float minSigma = 0.001;
 
     // TODO(@motiwari): #pragma omp parallel for if (this->parallelize)?
     for (size_t k = 0; k < nMedoids; k++) {
@@ -191,6 +199,15 @@ namespace km {
       // compute std dev amongst batch of reference points
       sigma = buildSigma(data, distMat, bestDistances, useAbsolute);
 
+      // Replace the zero entries in sigma with the minimum non-zero
+      // sigma value. Otherwise, some candidates could have
+      // a 0 confidence interval.
+      // This step prevents the overestimated candidates from
+      // discarding underestimated candidates,
+      // which could lead to suboptimal results.
+      minSigma = arma::min(arma::nonzeros(sigma));
+      sigma.elem(arma::find(sigma == 0.0)).fill(minSigma);
+
       while (arma::sum(candidates) > precision) {
         // TODO(@motiwari): Do not need a matrix for this comparison,
         //  use broadcasting
@@ -238,7 +255,7 @@ namespace km {
                 arma::sqrt(adjust / numSamples.cols(targets));
         ucbs.cols(targets) = estimates.cols(targets) + confBoundDelta;
         lcbs.cols(targets) = estimates.cols(targets) - confBoundDelta;
-        candidates = (lcbs < ucbs.min()) && (exactMask == 0);
+        candidates = (lcbs <= ucbs.min()) && (exactMask == 0);
       }
 
       medoidIndices->at(k) = lcbs.index_min();
@@ -271,6 +288,10 @@ namespace km {
     size_t N = data.n_cols;
     size_t K = nMedoids;
     arma::fmat updated_sigma(K, N, arma::fill::zeros);
+    // temporarily increase the batch size for precise estimation
+    // of the std dev
+    int originalBatchSize = batchSize;
+    batchSize = std::min(1000, static_cast<int>(N/4));
     arma::uvec referencePoints;
     // TODO(@motiwari): Make this wraparound properly
     //  as last batch_size elements are dropped
@@ -319,6 +340,9 @@ namespace km {
       }
       updated_sigma(k, n) = arma::stddev(sample);
     }
+    // reset batchSize to the original batch size as it's a global variable
+    // used by other functions (e.g. buildTarget, swapTarget)
+    batchSize = originalBatchSize;
     return updated_sigma;
   }
 
@@ -431,6 +455,7 @@ namespace km {
     arma::fmat lcbs(nMedoids, N);
     arma::fmat ucbs(nMedoids, N);
     arma::umat numSamples(nMedoids, N, arma::fill::zeros);
+    float minSigma = 0.001;
 
     // calculate quantities needed for swap, bestDistances and sigma
     calcBestDistancesSwap(
@@ -454,6 +479,10 @@ namespace km {
                 &secondBestDistances,
                 assignments);
 
+        // Fill in the zero sigma entries with the non-zero minimum sigma value
+        minSigma = arma::min(arma::nonzeros(sigma));
+        sigma.elem(arma::find(sigma == 0.0)).fill(minSigma);
+
         // Reset variables when starting a new swap
         candidates.fill(1);
         exactMask.fill(0);
@@ -496,7 +525,7 @@ namespace km {
             lcbs.cols(compute_exactly_targets) = result;
             exactMask.cols(compute_exactly_targets).fill(1);
             numSamples.cols(compute_exactly_targets) += N;
-            candidates = (lcbs < ucbs.min()) && (exactMask == 0);
+            candidates = (lcbs <= ucbs.min()) && (exactMask == 0);
           }
           if (arma::accu(candidates) < precision) {
             break;
@@ -544,7 +573,7 @@ namespace km {
           lcbs.cols(candidate_targets) = estimates.cols(candidate_targets)
                                          - confBoundDelta;
 
-          candidates = (lcbs < ucbs.min()) && (exactMask == 0);
+          candidates = (lcbs <= ucbs.min()) && (exactMask == 0);
       }
 
       // Perform the medoid switch
diff --git a/src/algorithms/banditpam_orig.cpp b/src/algorithms/banditpam_orig.cpp
index 61b76016..cf385952 100644
--- a/src/algorithms/banditpam_orig.cpp
+++ b/src/algorithms/banditpam_orig.cpp
@@ -69,6 +69,10 @@ namespace km {
           const arma::frowvec &bestDistances,
           const bool useAbsolute) {
     size_t N = data.n_cols;
+    // temporarily increase the batch size for precise estimation
+    // of the std dev
+    int originalBatchSize = batchSize;
+    batchSize = std::min(1000, static_cast<int>(N/4));
     arma::uvec referencePoints;
     // TODO(@motiwari): Make this wraparound properly as
     //  last batch_size elements are dropped
@@ -104,6 +108,9 @@ namespace km {
       }
       updated_sigma(i) = arma::stddev(sample);
     }
+    // reset batchSize to the original batch size as it's a global variable
+    // used by other functions (e.g. buildTarget, swapTarget)
+    batchSize = originalBatchSize;
     return updated_sigma;
   }
 
@@ -179,6 +186,7 @@ namespace km {
     arma::frowvec ucbs(N);
     arma::frowvec numSamples(N, arma::fill::zeros);
     arma::frowvec exactMask(N, arma::fill::zeros);
+    float minSigma = 0.001;
 
     // TODO(@motiwari): #pragma omp parallel for if (this->parallelize)?
     for (size_t k = 0; k < nMedoids; k++) {
@@ -191,6 +199,15 @@ namespace km {
       // compute std dev amongst batch of reference points
       sigma = buildSigma(data, distMat, bestDistances, useAbsolute);
 
+      // Replace the zero entries in sigma with the minimum non-zero
+      // sigma value. Otherwise, some candidates could have
+      // a 0 confidence interval.
+      // This step prevents the overestimated candidates from
+      // discarding underestimated candidates,
+      // which could lead to suboptimal results.
+      minSigma = arma::min(arma::nonzeros(sigma));
+      sigma.elem(arma::find(sigma == 0.0)).fill(minSigma);
+
       while (arma::sum(candidates) > precision) {
         // TODO(@motiwari): Do not need a matrix for this comparison,
         //  use broadcasting
@@ -238,7 +255,7 @@ namespace km {
                 arma::sqrt(adjust / numSamples.cols(targets));
         ucbs.cols(targets) = estimates.cols(targets) + confBoundDelta;
         lcbs.cols(targets) = estimates.cols(targets) - confBoundDelta;
-        candidates = (lcbs < ucbs.min()) && (exactMask == 0);
+        candidates = (lcbs <= ucbs.min()) && (exactMask == 0);
       }
 
       medoidIndices->at(k) = lcbs.index_min();
@@ -271,6 +288,10 @@ namespace km {
     size_t N = data.n_cols;
     size_t K = nMedoids;
     arma::fmat updated_sigma(K, N, arma::fill::zeros);
+    // temporarily increase the batch size for precise estimation
+    // of the std dev
+    int originalBatchSize = batchSize;
+    batchSize = std::min(1000, static_cast<int>(N/4));
     arma::uvec referencePoints;
     // TODO(@motiwari): Make this wraparound properly
     //  as last batch_size elements are dropped
@@ -319,6 +340,9 @@ namespace km {
       }
       updated_sigma(k, n) = arma::stddev(sample);
     }
+    // reset batchSize to the original batch size as it's a global variable
+    // used by other functions (e.g. buildTarget, swapTarget)
+    batchSize = originalBatchSize;
     return updated_sigma;
   }
 
@@ -411,6 +435,7 @@ namespace km {
     arma::fmat lcbs(nMedoids, N);
     arma::fmat ucbs(nMedoids, N);
     arma::umat numSamples(nMedoids, N, arma::fill::zeros);
+    float minSigma = 0.001;
 
     // calculate quantities needed for swap, bestDistances and sigma
     calcBestDistancesSwap(
@@ -434,6 +459,10 @@ namespace km {
               &secondBestDistances,
               assignments);
 
+      // Fill in the zero sigma entries with the non-zero minimum sigma value
+      minSigma = arma::min(arma::nonzeros(sigma));
+      sigma.elem(arma::find(sigma == 0.0)).fill(minSigma);
+
       // Reset variables when starting a new swap
       candidates.fill(1);
       exactMask.fill(0);
@@ -463,7 +492,7 @@ namespace km {
           lcbs.elem(targets) = result;
           exactMask.elem(targets).fill(1);
           numSamples.elem(targets) += N;
-          candidates = (lcbs < ucbs.min()) && (exactMask == 0);
+          candidates = (lcbs <= ucbs.min()) && (exactMask == 0);
         }
         if (arma::accu(candidates) < precision) {
           break;
@@ -492,7 +521,7 @@ namespace km {
                                             targets));
         ucbs.elem(targets) = estimates.elem(targets) + confBoundDelta;
         lcbs.elem(targets) = estimates.elem(targets) - confBoundDelta;
-        candidates = (lcbs < ucbs.min()) && (exactMask == 0);
+        candidates = (lcbs <= ucbs.min()) && (exactMask == 0);
       }
 
       // Perform the medoid switch
diff --git a/src/algorithms/kmedoids_algorithm.cpp b/src/algorithms/kmedoids_algorithm.cpp
index ee15fe10..6ab02434 100644
--- a/src/algorithms/kmedoids_algorithm.cpp
+++ b/src/algorithms/kmedoids_algorithm.cpp
@@ -9,6 +9,7 @@
 #include <armadillo>
 #include <unordered_map>
 #include <regex>
+#include <cassert>
 
 #include "kmedoids_algorithm.hpp"
 #include "fastpam1.hpp"
@@ -80,8 +81,10 @@ namespace km {
       //  that is properly raised
       throw std::invalid_argument("Dataset is empty");
     }
-    // TODO(@Adarsh321123): assert that the number of medoids is >=
-    //  than the number of points
+
+    assert(("Number of medoids should be less than the number of data points",
+            nMedoids < inputData.n_rows));
+
     batchSize = fmin(inputData.n_rows, batchSize);
 
     try {
@@ -399,10 +402,13 @@ namespace km {
       // T1 begins to write to cache and then T2
       // access in the middle of write?
       if (cache[(m * i) + reindex[j]] == -1) {
+        // cache miss! calculate the distance and cache it.
+        numCacheMisses++;
         numCacheWrites++;
         cache[(m * i) + reindex[j]] = (this->*lossFn)(data, i, j);
+      } else {
+        numCacheHits++;
       }
-      numCacheHits++;
       return cache[m * i + reindex[j]];
     }
 
diff --git a/src/python_bindings/kmedoids_pywrapper.cpp b/src/python_bindings/kmedoids_pywrapper.cpp
index c336909d..7a3fcfc0 100644
--- a/src/python_bindings/kmedoids_pywrapper.cpp
+++ b/src/python_bindings/kmedoids_pywrapper.cpp
@@ -15,12 +15,16 @@
 #include "kmedoids_algorithm.hpp"
 #include "kmedoids_pywrapper.hpp"
 
+// from https://github.com/pybind/python_example/blob/master/src/main.cpp
+#define STRINGIFY(x) #x
+#define MACRO_STRINGIFY(x) STRINGIFY(x)
+
 namespace km {
   PYBIND11_MODULE(banditpam, m) {
     // Module functions
     m.doc() = "BanditPAM Python library, implemented in C++";
 
-    m.attr("__version__") = VERSION_INFO;
+    m.attr("__version__") = MACRO_STRINGIFY(VERSION_INFO);
     m.def("get_max_threads",  // TODO(@motiwari): change to get_num_threads
     &omp_get_max_threads, "Returns max number of threads");
     m.def("set_num_threads",
@@ -40,7 +44,7 @@ namespace km {
     pybind11::arg("n_medoids") = 5,
     pybind11::arg("algorithm") = "BanditPAM",
     pybind11::arg("max_iter") = 100,
-    pybind11::arg("build_confidence") = 3,
+    pybind11::arg("build_confidence") = 10,  // 100 fixes stochasticity issues
     pybind11::arg("swap_confidence") = 5,
     // TODO(@motiwari): Verify these options are re-used correctly on reset
     pybind11::arg("use_cache") = true,


Test patch:
diff --git a/.github/workflows/run_linux_tests.yml b/.github/workflows/run_linux_tests.yml
index 01091846..7a85298e 100644
--- a/.github/workflows/run_linux_tests.yml
+++ b/.github/workflows/run_linux_tests.yml
@@ -1,5 +1,9 @@
 name: Linux - build package and run tests
-on: [push, pull_request]
+on:
+  push:
+    branches:  # prevents running on tag push
+      - '**'
+  pull_request:
 jobs:
   build:
     runs-on: ubuntu-latest
@@ -21,7 +25,7 @@ jobs:
           sudo apt install -y build-essential checkinstall libncursesw5-dev libssl-dev libsqlite3-dev tk-dev libgdbm-dev libc6-dev libbz2-dev libffi-dev zlib1g-dev 
           sudo apt install -y clang-format cppcheck
 
-      - name: Install Python dependencies 
+      - name: Install Python dependencies
         run: |
           python -m pip install --upgrade pip
           python -m pip install pytest
@@ -54,7 +58,7 @@ jobs:
           python -m pip install --no-use-pep517 --no-build-isolation -vvv -e .
         env:
           # The default compiler on the Github Ubuntu runners is gcc
-          # Would need to make a respective include change for clang 
+          # Would need to make a respective include change for clang
           CPLUS_INCLUDE_PATH: /usr/local/include/carma
 
       - name: Downloading data files for tests
@@ -73,9 +77,9 @@ jobs:
         run : |
           pytest tests/test_smaller.py 
 
-      # - name: Run larger suite of test cases
-      #   run : |
-      #     pytest tests/test_larger.py
+      - name: Run tests cases for initialization
+        run : |
+          pytest tests/test_initialization.py 
 
       - name: Clone for PR
         if: ${{ github.event_name == 'pull_request' }}
diff --git a/.github/workflows/run_mac_tests.yml b/.github/workflows/run_mac_tests.yml
index fda3f5d1..d2709b27 100644
--- a/.github/workflows/run_mac_tests.yml
+++ b/.github/workflows/run_mac_tests.yml
@@ -1,5 +1,9 @@
 name: MacOS - build package and run tests
-on: [push, pull_request]
+on:
+  push:
+    branches:  # prevents running on tag push
+      - '**'
+  pull_request:
 jobs:
   build:
     runs-on: macos-latest
@@ -74,6 +78,6 @@ jobs:
         run : |
           pytest tests/test_smaller.py
 
-      # - name: Run larger suite of test cases
-      #   run : |
-      #     pytest tests/test_larger.py
\ No newline at end of file
+      - name: Run tests cases for initialization
+        run : |
+          pytest tests/test_initialization.py
\ No newline at end of file
diff --git a/tests/test_initialization.py b/tests/test_initialization.py
index 3689d787..da768e4f 100644
--- a/tests/test_initialization.py
+++ b/tests/test_initialization.py
@@ -9,7 +9,7 @@ def test_default_initialization(self):
         assert kmed.n_medoids == 5, "n_medoids does not match"
         assert kmed.algorithm == "BanditPAM", "algorithm does not match"
         assert kmed.max_iter == 100, "max_iter does not match"
-        assert kmed.build_confidence == 3, "build_confidence does not match"
+        assert kmed.build_confidence == 10, "build_confidence does not match"
         assert kmed.swap_confidence == 5, "swap_confidence does not match"
         assert kmed.use_cache, "use_cache does not match"
         assert kmed.use_perm, "use_perm does not match"
@@ -18,7 +18,7 @@ def test_default_initialization(self):
 
     def test_initialization_1(self):
         """
-        Test that we can initialize a KMedoids object with given parameters
+        Test that we can initialize a KMedoids object with given parameters.
         """
         kmed = KMedoids(
             n_medoids=5,
@@ -78,8 +78,9 @@ def test_initialization_3(self):
         assert kmed.n_medoids == 100, "n_medoids does not match"
         assert kmed.algorithm == "PAM", "algorithm does not match"
         assert kmed.max_iter == 10000, "max_iter does not match"
-        assert kmed.build_confidence == 11000, "build_confidence does not " \
-                                               "match"
+        assert (
+            kmed.build_confidence == 11000
+        ), "build_confidence does not match"
         assert kmed.swap_confidence == 2029, "swap_confidence does not match"
         assert kmed.use_cache, "use_cache does not match"
         assert not kmed.use_perm, "use_perm does not match"
diff --git a/tests/test_larger.py b/tests/test_larger.py
index e86de3b5..4e7e9638 100644
--- a/tests/test_larger.py
+++ b/tests/test_larger.py
@@ -1,5 +1,6 @@
 import unittest
 import pandas as pd
+import numpy as np
 import time
 
 from banditpam import KMedoids
@@ -25,35 +26,41 @@ class LargerTests(unittest.TestCase):
 
     def test_medium_mnist(self):
         """
-        Generate NUM_MEDIUM_CASES random subsets of MNIST
-        of size MEDIUM_SAMPLE_SIZE and verify BanditPAM agrees with PAM
+        Generate NUM_MEDIUM_CASES random subsets of MNIST of size
+        MEDIUM_SAMPLE_SIZE and verify BanditPAM agrees with PAM.
         """
+        num_succeed = 0
         for i in range(NUM_MEDIUM_CASES):
             data = self.mnist_70k.sample(n=MEDIUM_SAMPLE_SIZE).to_numpy()
-            _ = bpam_agrees_pam(
+            num_succeed += bpam_agrees_pam(
                 k=SMALL_K_SCHEDULE[i % N_SMALL_K],
                 data=data,
                 loss="L2",
                 test_build=True,
-                assert_immediately=True,
+                assert_immediately=False,
             )
+        self.assertTrue(num_succeed >= PROPORTION_PASSING * NUM_MEDIUM_CASES)  # avoids stochasticity issues
 
     def test_various_medium_mnist(self):
         """
-        Generate NUM_MEDIUM_CASES random subsets of MNIST of various sizes
-        and verify BanditPAM agrees with PAM. Since PAM is very slow,
-        we can only do this for fairly small sizes in MEDIUM_SIZE_SCHEDULE
+        Generate NUM_MEDIUM_CASES random subsets of MNIST of various sizes and
+        verify BanditPAM agrees with PAM.
+
+        Since PAM is very slow, we can only do this for fairly small sizes in
+        MEDIUM_SIZE_SCHEDULE
         """
+        num_succeed = 0
         for i in range(NUM_MEDIUM_CASES):
             size = MEDIUM_SIZE_SCHEDULE[i % NUM_MEDIUM_SIZES]
             data = self.mnist_70k.sample(n=size)
-            _ = bpam_agrees_pam(
+            num_succeed += bpam_agrees_pam(
                 k=SMALL_K_SCHEDULE[i % N_SMALL_K],
                 data=data,
                 loss="L2",
                 test_build=True,
-                assert_immediately=True,
+                assert_immediately=False,
             )
+        self.assertTrue(num_succeed >= PROPORTION_PASSING * NUM_MEDIUM_CASES)  # avoids stochasticity issues
 
     def test_time_cases_mnist(self):
         """
@@ -61,6 +68,7 @@ def test_time_cases_mnist(self):
         """
         MNIST_10k = self.mnist_70k.head(LARGE_SAMPLE_SIZE).to_numpy()
         kmed = KMedoids(n_medoids=5, algorithm="BanditPAM")
+        kmed.seed = 0
         start = time.time()
         kmed.fit(MNIST_10k, "L2")
         base_runtime = time.time() - start
@@ -79,8 +87,8 @@ def test_time_cases_mnist(self):
 
     def test_medium_scrna(self):
         """
-        Generate NUM_MEDIUM_CASES random subsets of MNIST
-        of size MEDIUM_SAMPLE_SIZE and verify BanditPAM agrees with PAM
+        Generate NUM_MEDIUM_CASES random subsets of MNIST of size
+        MEDIUM_SAMPLE_SIZE and verify BanditPAM agrees with PAM.
         """
         count = 0
         for i in range(NUM_MEDIUM_CASES):
@@ -96,9 +104,11 @@ def test_medium_scrna(self):
 
     def test_various_medium_scrna(self):
         """
-        Generate NUM_MEDIUM_CASES random subsets of scRNA of various sizes
-        and verify BanditPAM agrees with PAM. Since PAM is very slow,
-        we can only do this for fairly small sizes in MEDIUM_SIZE_SCHEDULE
+        Generate NUM_MEDIUM_CASES random subsets of scRNA of various sizes and
+        verify BanditPAM agrees with PAM.
+
+        Since PAM is very slow, we can only do this for fairly small sizes in
+        MEDIUM_SIZE_SCHEDULE
         """
         count = 0
         for i in range(NUM_MEDIUM_CASES):
@@ -119,6 +129,7 @@ def test_time_cases_scrna(self):
         """
         SCRNA_10k = self.scrna.head(LARGE_SAMPLE_SIZE).to_numpy()
         kmed = KMedoids(n_medoids=5, algorithm="BanditPAM")
+        kmed.seed = 0
         start = time.time()
         kmed.fit(SCRNA_10k, "L1")
         base_runtime = time.time() - start
@@ -137,4 +148,5 @@ def test_time_cases_scrna(self):
 
 
 if __name__ == "__main__":
+    np.random.seed(0)
     unittest.main()
diff --git a/tests/test_smaller.py b/tests/test_smaller.py
index 7c5d890a..df6d9049 100644
--- a/tests/test_smaller.py
+++ b/tests/test_smaller.py
@@ -13,7 +13,6 @@
 )
 
 
-# TODO(@motiwari): Set seeds
 class SmallerTests(unittest.TestCase):
     small_mnist = pd.read_csv("data/MNIST_100.csv", header=None).to_numpy()
     mnist_70k = pd.read_csv("data/MNIST_70k.csv", sep=" ", header=None)
@@ -22,7 +21,7 @@ class SmallerTests(unittest.TestCase):
     def test_small_mnist(self):
         """
         Test NUM_SMALL_CASES number of test cases with subsets of size
-        SMALL_SAMPLE_SIZE randomly drawn from the full MNIST dataset
+        SMALL_SAMPLE_SIZE randomly drawn from the full MNIST dataset.
         """
         for i in range(NUM_SMALL_CASES):
             data = self.mnist_70k.sample(n=SMALL_SAMPLE_SIZE).to_numpy()
@@ -50,7 +49,7 @@ def test_small_mnist(self):
     def test_small_scrna(self):
         """
         Test NUM_SMALL_CASES number of test cases with subsets of size
-        SMALL_SAMPLE_SIZE randomly drawn from the full scRNA dataset
+        SMALL_SAMPLE_SIZE randomly drawn from the full scRNA dataset.
         """
         count = 0
         for i in range(NUM_SMALL_CASES):
@@ -80,28 +79,24 @@ def test_small_scrna(self):
 
     def test_small_mnist_known_cases(self):
         """
-        Test BanditPAM on a subset of MNIST with known solutions
-        for both k = 5 and k = 10, after both the BUILD and SWAP steps
+        Test BanditPAM on a subset of MNIST with known solutions for both k = 5
+        and k = 10, after both the BUILD and SWAP steps.
         """
         kmed_5 = KMedoids(
             n_medoids=5,
             algorithm="BanditPAM",
         )
+        kmed_5.seed = 0
         kmed_5.fit(self.small_mnist, "L2")
 
-        self.assertEqual(
-            sorted(kmed_5.build_medoids.tolist()),
-            [16, 24, 32, 70, 87]
-        )
-        self.assertEqual(
-            sorted(kmed_5.medoids.tolist()),
-            [23, 30, 49, 70, 99]
-        )
+        self.assertEqual(sorted(kmed_5.build_medoids.tolist()), [16, 24, 32, 70, 87])
+        self.assertEqual(sorted(kmed_5.medoids.tolist()), [23, 30, 49, 70, 99])
 
         kmed_10 = KMedoids(
             n_medoids=10,
             algorithm="BanditPAM",
         )
+        kmed_10.seed = 0
         kmed_10.fit(self.small_mnist, "L2")
 
         self.assertEqual(
@@ -110,15 +105,16 @@ def test_small_mnist_known_cases(self):
         )
         self.assertEqual(
             sorted(kmed_10.medoids.tolist()),
-            [16, 25, 31, 49, 63, 70, 82, 90, 94, 99]
+            [16, 25, 31, 49, 63, 70, 82, 90, 94, 99],
         )
 
     def test_edge_cases(self):
         """
-        Test that BanditPAM raises errors on n_medoids being unspecified or
-        an empty dataset
+        Test that BanditPAM raises errors on n_medoids being unspecified or an
+        empty dataset.
         """
         kmed = KMedoids()
+        kmed.seed = 0
 
         # initialized to empty
         self.assertEqual([], kmed.build_medoids[0].tolist())
@@ -132,4 +128,5 @@ def test_edge_cases(self):
 
 
 if __name__ == "__main__":
+    np.random.seed(0)
     unittest.main()
diff --git a/tests/utils.py b/tests/utils.py
index 957a8e76..740f5834 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -28,6 +28,7 @@ def bpam_agrees_pam(
 
     kmed_bpam = KMedoids(n_medoids=k, algorithm="BanditPAM")
     kmed_pam = KMedoids(n_medoids=k, algorithm=alg_name)
+    kmed_bpam.seed = 0
     kmed_bpam.fit(data, loss)
     kmed_pam.fit(data, loss)