update unit tests

.github/workflows/ci.yml

@@ -28,6 +28,11 @@ jobs:
         python -m pip install --upgrade pip
         python -m pip install pytest
         if [ -f pyproject.toml ]; then pip install -e .; fi
+    - name: Install python bindings for ngt
+      run: | 
+        git clone https://github.com/yahoojapan/NGT.git
+        cd NGT/python
+        pip3 install .
     - name: Test with pytest
       run: |
         python -m pytest tests/test_utils.py -v --ci

.github/workflows/smoke-test.yml

@@ -35,11 +35,15 @@ jobs:
     - name: Install dependencies
       run: |
         sudo apt-get update
-        python -m pip install --upgrade pip
         pip install -e .[dev]
     - name: List dependencies
       run: |
         pip list
+    - name: Install python bindings for ngt
+      run: |
+        git clone https://github.com/yahoojapan/NGT.git
+        cd NGT/python
+        pip3 install .
     - name: Run unit tests with pytest /  pytest-copie
       run: |
         python -m pytest

README.md

@@ -12,8 +12,13 @@ reLAISS lets you retrieve nearest‑neighbour supernovae (or spot outliers) by c
 
 # Install
 
-In a fresh conda environment, run `pip install relaiss`. Then, install the python bindings for ngt from source via the instructions [here](https://github.com/yahoojapan/NGT/blob/main/python/README.md). 
+In a fresh conda environment, run `pip install relaiss`. After installing, you must install the python bindings for `ngt` from source: 
 
+```
+git clone https://github.com/yahoojapan/NGT.git
+cd NGT/python
+pip3 install .
+```
 
 # Code Demo
 ```

examples/advanced_usage.py

@@ -6,7 +6,7 @@ This script demonstrates advanced features of reLAISS including:
 - Swapping host galaxies
 - Using PCA for dimensionality reduction
 - Setting maximum neighbor distances
-- Tweaking ANNOY parameters
+- Tweaking NGT parameters
 - Making corner plots
 - Advanced anomaly detection with parameter tuning
 - Host swapping in anomaly detection
@@ -106,8 +106,8 @@ def main():
     print("\nNearest neighbors within distance threshold:")
     print(neighbors_df)
     
-    # Example 5: Tweaking ANNOY parameters
-    print("\nExample 5: Tweaking ANNOY parameters")
+    # Example 5: Tweaking NGT parameters
+    print("\nExample 5: Tweaking NGT parameters")
     neighbors_df = client.find_neighbors(
         ztf_object_id='ZTF21abbzjeq',
         n=5,
@@ -116,7 +116,7 @@ def main():
         save_figures=True,
         path_to_figure_directory='./figures'
     )
-    print("\nNearest neighbors with tweaked ANNOY parameters:")
+    print("\nNearest neighbors with tweaked NGT parameters:")
     print(neighbors_df)
     
     # Example 6: Making corner plots

src/relaiss/anomaly_config.py

@@ -11,3 +11,5 @@ TRAINING_SAMPLE_SIZE = 5000         # Maximum training features to store
 # Training Data Quality Parameters
 HIGH_NAN_WARNING_THRESHOLD = 20    # Percentage of NaN samples to trigger warning
 VERY_HIGH_NAN_ERROR_THRESHOLD = 50 # Percentage of NaN samples to trigger error
+
+MODEL_FILENAME_TEMPLATE = "kNN_scaler_lc={num_lc}_host={num_host}.pkl" 

src/relaiss/relaiss.py

@@ -12,10 +12,7 @@ from .fetch import get_TNS_data
 from .plotting import plot_lightcurves, plot_hosts
 import os
 from kneed import KneeLocator
-try:
-    import ngtpy as ngt
-except:
-    import ngt 
+import ngtpy as ngt
 import numpy as np
 import pandas as pd
 from sklearn.decomposition import PCA

tests/conftest.py

@@ -7,7 +7,7 @@ from unittest.mock import patch, MagicMock
 import tempfile
 import shutil
 import astropy.units as u
-from .fixtures.search import build_test_annoy_index
+from .fixtures.search import build_test_ngt_index
 
 # Get the path to the test fixtures
 FIXTURES_DIR = Path(__file__).parent / "fixtures"
@@ -192,15 +192,15 @@ def mock_extinction_all():
         yield mock_map, mock_g23
 
 @pytest.fixture
-def test_annoy_index(dataset_bank_path):
-    """Create a test Annoy index for testing."""
+def test_ngt_index(dataset_bank_path):
+    """Create a test NGT index for testing."""
     # Get the dataset bank file
     df = pd.read_csv(dataset_bank_path)
     
-    # Build a temporary Annoy index
+    # Build a temporary NGT index
     with tempfile.TemporaryDirectory() as tmpdir:
-        index_path = Path(tmpdir) / "annoy_index"
-        index, index_path, object_ids = build_test_annoy_index(
+        index_path = Path(tmpdir) / "ngt_index"
+        index, index_path, object_ids = build_test_ngt_index(
             test_databank_path=dataset_bank_path,
             lc_features=['g_peak_mag', 'r_peak_mag', 'g_peak_time', 'r_peak_time'],
             host_features=['host_ra', 'host_dec']
@@ -308,4 +308,4 @@ def pytest_addoption(parser):
     """Add command line options."""
     parser.addoption(
         "--ci", action="store_true", default=False, help="Run in CI mode (skip tests requiring real data)"
-    ) 
+    ) 

tests/fixtures/search.py

@@ -1,13 +1,13 @@
 import numpy as np
 import pandas as pd
-import annoy
+import ngtpy as ngt
 import os
 from pathlib import Path
 import tempfile
 
-# Test fixture for building an Annoy index for testing
-def build_test_annoy_index(test_databank_path, lc_features=None, host_features=None):
-    """Build an Annoy index from the test dataset bank for testing.
+# Test fixture for building an NGT index for testing
+def build_test_ngt_index(test_databank_path, lc_features=None, host_features=None):
+    """Build an NGT index from the test dataset bank for testing.
     
     Parameters
     ----------
@@ -21,7 +21,7 @@ def build_test_annoy_index(test_databank_path, lc_features=None, host_features=N
     Returns
     -------
     tuple
-        (index, index_path, object_ids) - the annoy index, path to temp file, and array of object ids
+        (index, index_path, object_ids) - the ngt index, path to temp file, and array of object ids
     """
     # Default features if none provided
     if lc_features is None:
@@ -42,31 +42,31 @@ def build_test_annoy_index(test_databank_path, lc_features=None, host_features=N
     feat_arr = np.array(df_features)
     feat_arr_scaled = (feat_arr - np.mean(feat_arr, axis=0)) / np.std(feat_arr, axis=0)
     
-    # Create annoy index
+    # Create temp directory for NGT index
+    temp_dir = tempfile.mkdtemp()
+    index_path = os.path.join(temp_dir, "test_index.ngt")
+    
+    # Create NGT index
     index_dim = feat_arr.shape[1]
-    index = annoy.AnnoyIndex(index_dim, "manhattan")
+    ngt.create(index_path.encode(), index_dim, distance_type="L2")
+    index = ngt.Index(index_path.encode())
     
     # Add items to index
     for i, obj_id in enumerate(df_bank.index):
-        index.add_item(i, feat_arr_scaled[i])
+        index.insert(feat_arr_scaled[i].astype(np.float32))
     
-    # Build index with 10 trees (fewer for tests to be faster)
-    index.build(10)
-    
-    # Create temp file to save index
-    temp_dir = tempfile.mkdtemp()
-    index_path = os.path.join(temp_dir, "test_index.ann")
-    index.save(index_path)
+    # Build index
+    index.build_index()
     
     return index, index_path, np.array(df_bank.index)
 
 def find_neighbors(index, idx_arr, query_vector, n=5):
-    """Find neighbors using the test Annoy index.
+    """Find neighbors using the test NGT index.
     
     Parameters
     ----------
-    index : annoy.AnnoyIndex
-        The Annoy index to query
+    index : ngt.Index
+        The NGT index to query
     idx_arr : numpy.ndarray
         Array of object IDs
     query_vector : numpy.ndarray
@@ -80,11 +80,10 @@ def find_neighbors(index, idx_arr, query_vector, n=5):
         (ids, distances) - arrays of neighbor IDs and distances
     """
     # Query the index
-    neighbor_indices, distances = index.get_nns_by_vector(
-        query_vector, n, include_distances=True
-    )
+    res = index.search(query_vector.astype(np.float32), n)
+    neighbor_indices, distances = zip(*res)
     
     # Get ZTF IDs of neighbors
     neighbor_ids = idx_arr[neighbor_indices]
     
-    return neighbor_ids, distances 
+    return neighbor_ids, distances 

tests/test_search.py

@@ -52,9 +52,9 @@ def test_find_neighbors_invalid_input():
     with pytest.raises(ValueError):
         client.find_neighbors(ztf_object_id="ZTF21abbzjeq", n=-1)
 
-def test_annoy_search(test_annoy_index, dataset_bank_path):
-    """Test that the Annoy index works as expected for neighbor search."""
-    index, index_path, object_ids = test_annoy_index
+def test_ngt_search(test_ngt_index, dataset_bank_path):
+    """Test that the NGT index works as expected for neighbor search."""
+    index, index_path, object_ids = test_ngt_index
     
     # Use a predefined set of features to ensure dimensions match
     lc_features = ['g_peak_mag', 'r_peak_mag', 'g_peak_time', 'r_peak_time']
@@ -64,12 +64,10 @@ def test_annoy_search(test_annoy_index, dataset_bank_path):
     # Create a random test vector with the correct dimension
     test_vector = np.random.rand(vector_dim)
     
-    # Increase the search_k parameter for better accuracy
-    search_k = 1000
-    
     # Get nearest neighbors
     n_items = min(5, len(object_ids))
-    nearest_indices = index.get_nns_by_vector(test_vector, n_items, search_k=search_k)
+    res = index.search(test_vector.astype(np.float32), n_items)
+    nearest_indices, distances = zip(*res)
     
     # Verify results
     assert len(nearest_indices) >= 1  # Accept at least 1 neighbor