update examples in nbs and example dir

examples/advanced_usage.py

@@ -8,6 +8,8 @@ This script demonstrates advanced features of reLAISS including:
 - Setting maximum neighbor distances
 - Tweaking ANNOY parameters
 - Making corner plots
+- Advanced anomaly detection with parameter tuning
+- Host swapping in anomaly detection
 """
 
 import os
@@ -38,6 +40,8 @@ def main():
     # Create output directories
     os.makedirs('./figures', exist_ok=True)
     os.makedirs('./sfddata-master', exist_ok=True)
+    os.makedirs('./models', exist_ok=True)
+    os.makedirs('./timeseries', exist_ok=True)
     
     # Initialize the client
     client = rl.ReLAISS()
@@ -150,6 +154,90 @@ def main():
         path_to_figure_directory='./figures',
         save_plots=True
     )
+    
+    # Example 7: Advanced anomaly detection with parameter tuning
+    print("\nExample 7: Advanced anomaly detection with parameter tuning")
+    from relaiss.anomaly import train_AD_model, anomaly_detection
+    
+    # Train models with different parameters to compare
+    print("Training anomaly detection model with default parameters...")
+    default_model_path = train_AD_model(
+        lc_features=client.lc_features,
+        host_features=client.host_features,
+        path_to_dataset_bank=client.bank_csv,
+        path_to_sfd_folder='./sfddata-master',
+        path_to_models_directory="./models",
+        n_estimators=100,
+        contamination=0.02,
+        max_samples=256,
+        force_retrain=True
+    )
+    
+    print("Training anomaly detection model with more trees...")
+    model_more_trees_path = train_AD_model(
+        lc_features=client.lc_features,
+        host_features=client.host_features,
+        path_to_dataset_bank=client.bank_csv,
+        path_to_sfd_folder='./sfddata-master',
+        path_to_models_directory="./models",
+        n_estimators=200,  # More trees
+        contamination=0.02,
+        max_samples=256,
+        force_retrain=True
+    )
+    
+    print("Training anomaly detection model with higher contamination...")
+    model_higher_contam_path = train_AD_model(
+        lc_features=client.lc_features,
+        host_features=client.host_features,
+        path_to_dataset_bank=client.bank_csv,
+        path_to_sfd_folder='./sfddata-master',
+        path_to_models_directory="./models",
+        n_estimators=100,
+        contamination=0.05,  # Higher contamination
+        max_samples=256,
+        force_retrain=True
+    )
+    
+    # Run anomaly detection with each model
+    print("\nRunning anomaly detection with default model...")
+    anomaly_detection(
+        transient_ztf_id="ZTF21abbzjeq",
+        lc_features=client.lc_features,
+        host_features=client.host_features,
+        path_to_timeseries_folder="./timeseries",
+        path_to_sfd_folder='./sfddata-master',
+        path_to_dataset_bank=client.bank_csv,
+        path_to_models_directory="./models",
+        path_to_figure_directory="./figures/AD_default",
+        save_figures=True,
+        n_estimators=100,
+        contamination=0.02,
+        max_samples=256,
+        force_retrain=False
+    )
+    
+    # Example 8: Anomaly detection with host swapping
+    print("\nExample 8: Anomaly detection with host swapping")
+    # Use the default model but swap in a different host galaxy
+    anomaly_detection(
+        transient_ztf_id="ZTF21abbzjeq",
+        lc_features=client.lc_features,
+        host_features=client.host_features,
+        path_to_timeseries_folder="./timeseries",
+        path_to_sfd_folder='./sfddata-master',
+        path_to_dataset_bank=client.bank_csv,
+        host_ztf_id_to_swap_in="ZTF21aakswqr",  # Swap in a different host
+        path_to_models_directory="./models",
+        path_to_figure_directory="./figures/AD_host_swap",
+        save_figures=True,
+        n_estimators=100,
+        contamination=0.02,
+        max_samples=256,
+        force_retrain=False
+    )
+    
+    print("Anomaly detection figures saved to ./figures/AD_default/ and ./figures/AD_host_swap/")
 
 if __name__ == "__main__":
     main() 

examples/basic_usage.py

@@ -6,6 +6,7 @@ This script demonstrates the basic functionality of reLAISS including:
 - Running nearest neighbor search
 - Using Monte Carlo simulations
 - Adjusting feature weights
+- Basic anomaly detection
 """
 
 import os
@@ -15,6 +16,8 @@ def main():
     # Create output directories
     os.makedirs('./figures', exist_ok=True)
     os.makedirs('./sfddata-master', exist_ok=True)
+    os.makedirs('./models', exist_ok=True)
+    os.makedirs('./timeseries', exist_ok=True)
     
     # Initialize the client
     client = rl.ReLAISS()
@@ -55,7 +58,7 @@ def main():
     neighbors_df = client.find_neighbors(
         ztf_object_id='ZTF21abbzjeq',  # Using the test transient
         n=5,
-        num_mc_simulations=20,  # Number of Monte Carlo simulations
+        num_sims=20,  # Number of Monte Carlo simulations
         weight_lc_feats_factor=3.0,  # Up-weight lightcurve features
         plot=True,
         save_figures=True,
@@ -63,6 +66,44 @@ def main():
     )
     print("\nNearest neighbors with MC simulations:")
     print(neighbors_df)
+    
+    # Example 4: Basic anomaly detection
+    print("\nExample 4: Basic anomaly detection")
+    from relaiss.anomaly import train_AD_model, anomaly_detection
+    
+    # First, train an anomaly detection model
+    print("Training anomaly detection model...")
+    model_path = train_AD_model(
+        lc_features=client.lc_features,
+        host_features=client.host_features,
+        path_to_dataset_bank=client.bank_csv,
+        path_to_sfd_folder='./sfddata-master',
+        path_to_models_directory="./models",
+        n_estimators=100,  # Using smaller value for faster execution
+        contamination=0.02,  # Expected proportion of anomalies
+        max_samples=256,  # Max samples per tree
+        force_retrain=False  # Only retrain if model doesn't exist
+    )
+    print(f"Anomaly detection model saved to: {model_path}")
+    
+    # Now, run anomaly detection on a specific transient
+    print("\nRunning anomaly detection...")
+    anomaly_detection(
+        transient_ztf_id="ZTF21abbzjeq",  # Same test transient
+        lc_features=client.lc_features,
+        host_features=client.host_features,
+        path_to_timeseries_folder="./timeseries",
+        path_to_sfd_folder='./sfddata-master',
+        path_to_dataset_bank=client.bank_csv,
+        path_to_models_directory="./models",
+        path_to_figure_directory="./figures",
+        save_figures=True,
+        n_estimators=100,
+        contamination=0.02,
+        max_samples=256,
+        force_retrain=False
+    )
+    print("Anomaly detection figures saved to ./figures/AD/")
 
 if __name__ == "__main__":
     main() 

examples/build_databank.py

@@ -6,12 +6,14 @@ This script demonstrates how to build a new dataset bank for reLAISS, including:
 2. Joining new lightcurve features
 3. Handling missing values
 4. Building the final dataset bank
+5. Using different feature combinations for nearest neighbor search
 
 The process involves several steps:
 1. Add extinction corrections to the large dataset bank
 2. Join new lightcurve features to the small dataset bank
 3. Handle missing values using KNN imputation
 4. Build the final dataset bank with all features
+5. Demonstrate custom feature selection for nearest neighbor search
 """
 
 import os
@@ -21,6 +23,7 @@ from sfdmap2 import sfdmap
 from sklearn.impute import KNNImputer
 from relaiss.features import build_dataset_bank
 from relaiss import constants
+import relaiss as rl
 
 def add_extinction_corrections(df, path_to_sfd_folder):
     """Add extinction corrections (A_V) to the dataset.
@@ -95,6 +98,8 @@ def main():
     # Create necessary directories
     os.makedirs('../data', exist_ok=True)
     os.makedirs('../data/sfddata-master', exist_ok=True)
+    os.makedirs('./figures', exist_ok=True)
+    os.makedirs('./sfddata-master', exist_ok=True)
     
     # Step 1: Add extinction corrections to large dataset bank
     print("\nStep 1: Adding extinction corrections")
@@ -146,6 +151,109 @@ def main():
     print("Shape of final dataset bank:", final_dataset_bank.shape)
     final_dataset_bank.to_csv('../data/large_final_df_bank_new_lc_feats.csv', index=False)
     print("Successfully saved final dataset bank!")
+    
+    # Step 5: Demonstrate different feature combinations for search
+    print("\nStep 5: Using different feature combinations")
+    
+    # Define default feature sets from constants
+    default_lc_features = constants.lc_features_const.copy()
+    default_host_features = constants.host_features_const.copy()
+    
+    # Initialize client
+    client = rl.ReLAISS()
+    client.load_reference(
+        path_to_sfd_folder='./sfddata-master'
+    )
+    
+    # Example 1: Using only lightcurve features (no host features)
+    print("\nExample 1: Using only lightcurve features")
+    lc_only_client = rl.ReLAISS()
+    lc_only_client.load_reference(
+        path_to_sfd_folder='./sfddata-master',
+        lc_features=default_lc_features,  # Use default lightcurve features
+        host_features=[],  # Empty list means no host features
+    )
+    
+    # Find neighbors using only lightcurve features
+    neighbors_df_lc_only = lc_only_client.find_neighbors(
+        ztf_object_id='ZTF21abbzjeq',
+        n=5,
+        plot=True,
+        save_figures=True,
+        path_to_figure_directory='./figures/lc_only'
+    )
+    print("\nNearest neighbors using only lightcurve features:")
+    print(neighbors_df_lc_only)
+    
+    # Example 2: Using only host features (no lightcurve features)
+    print("\nExample 2: Using only host features")
+    host_only_client = rl.ReLAISS()
+    host_only_client.load_reference(
+        path_to_sfd_folder='./sfddata-master',
+        lc_features=[],  # Empty list means no lightcurve features
+        host_features=default_host_features,  # Use default host features
+    )
+    
+    # Find neighbors using only host features
+    neighbors_df_host_only = host_only_client.find_neighbors(
+        ztf_object_id='ZTF21abbzjeq',
+        n=5,
+        plot=True,
+        save_figures=True,
+        path_to_figure_directory='./figures/host_only'
+    )
+    print("\nNearest neighbors using only host features:")
+    print(neighbors_df_host_only)
+    
+    # Example 3: Using custom feature subset
+    print("\nExample 3: Using custom feature subset")
+    # Select specific lightcurve and host features
+    custom_lc_features = ['g_peak_mag', 'r_peak_mag', 'g_peak_time', 'r_peak_time']
+    custom_host_features = ['host_ra', 'host_dec', 'gKronMag', 'rKronMag']
+    
+    custom_client = rl.ReLAISS()
+    custom_client.load_reference(
+        path_to_sfd_folder='./sfddata-master',
+        lc_features=custom_lc_features,  # Custom subset of lightcurve features
+        host_features=custom_host_features,  # Custom subset of host features
+    )
+    
+    # Find neighbors with custom feature subset
+    neighbors_df_custom = custom_client.find_neighbors(
+        ztf_object_id='ZTF21abbzjeq',
+        n=5,
+        plot=True,
+        save_figures=True,
+        path_to_figure_directory='./figures/custom_features'
+    )
+    print("\nNearest neighbors using custom feature subset:")
+    print(neighbors_df_custom)
+    
+    # Example 4: Setting feature importance with feature weighting
+    print("\nExample 4: Using feature weighting")
+    # Regular search prioritizing lightcurve features
+    neighbors_df_lc_weighted = client.find_neighbors(
+        ztf_object_id='ZTF21abbzjeq',
+        n=5,
+        weight_lc_feats_factor=3.0,  # Strongly prioritize lightcurve features
+        plot=True,
+        save_figures=True,
+        path_to_figure_directory='./figures/lc_weighted'
+    )
+    print("\nNearest neighbors with lightcurve features weighted 3x:")
+    print(neighbors_df_lc_weighted)
+    
+    # Now prioritize host features by using a factor < 1
+    neighbors_df_host_weighted = client.find_neighbors(
+        ztf_object_id='ZTF21abbzjeq',
+        n=5,
+        weight_lc_feats_factor=0.3,  # Prioritize host features
+        plot=True,
+        save_figures=True,
+        path_to_figure_directory='./figures/host_weighted'
+    )
+    print("\nNearest neighbors with host features given higher weight:")
+    print(neighbors_df_host_weighted)
 
 if __name__ == "__main__":
     main() 

notebooks/01_basic_usage.ipynb

@@ -0,0 +1,321 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "648c2b16",
+   "metadata": {},
+   "source": [
+    "# Basic Usage of reLAISS\n",
+    "### Authors: Evan Reynolds and Alex Gagliano\n",
+    "\n",
+    "## Introduction\n",
+    "\n",
+    "reLAISS is the second version of LAISS (Lightcurve Anomaly Identification & Similarity Search); a tool to find similar supernovae & identify anomalous supernovae (and the galaxies that host them) using their photometric features.\n",
+    "\n",
+    "The similarity search takes advantage of [Approximate Nearest Neighbors Oh Yeah (ANNOY)](https://github.com/spotify/annoy), the approximate nearest neighbors algorithm developed by Spotify that allows you to come up with a relevant song to listen to before your current one ends. The anomaly detection classifier is an isolation forest model trained on a dataset bank of over 22,000 transients.\n",
+    "\n",
+    "This notebook demonstrates the basic features of the reLAISS library for finding similar astronomical transients.\n",
+    "\n",
+    "## Topics Covered\n",
+    "1. Initializing the ReLAISS client\n",
+    "2. Loading reference data\n",
+    "3. Finding optimal number of neighbors\n",
+    "4. Basic nearest neighbor search\n",
+    "5. Using Monte Carlo simulations and feature weighting\n",
+    "6. Basic anomaly detection"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "869e17e0",
+   "metadata": {},
+   "source": [
+    "## Setup\n",
+    "\n",
+    "First, let's import the required packages and create the necessary output directories:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1bf22151",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import pandas as pd\n",
+    "import relaiss\n",
+    "\n",
+    "# Create output directories\n",
+    "os.makedirs('./figures', exist_ok=True)\n",
+    "os.makedirs('./sfddata-master', exist_ok=True)\n",
+    "os.makedirs('./models', exist_ok=True)\n",
+    "os.makedirs('./timeseries', exist_ok=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0e82fe89",
+   "metadata": {},
+   "source": [
+    "## 1. Initialize the ReLAISS Client\n",
+    "\n",
+    "First, we create an instance of the ReLAISS client that we'll use to find similar transients."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cc144944",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Create ReLAISS client\n",
+    "client = relaiss.ReLAISS()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5483eb7b",
+   "metadata": {},
+   "source": [
+    "## 2. Load Reference Data\n",
+    "\n",
+    "Next, we load the reference dataset bank. This contains the features of known transients that we'll use for comparison.\n",
+    "\n",
+    "The `load_reference` function will automatically download the SFD dust map files if they don't exist in the specified directory. These files are required for extinction corrections in the reLAISS pipeline."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "808bf098",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load reference data\n",
+    "client.load_reference(\n",
+    "    path_to_sfd_folder='./sfddata-master',  # Directory for SFD dust maps\n",
+    "    use_pca=False,  # Don't use PCA for this example\n",
+    "    host_features=[]  # Empty list for this example\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9ae5d46f",
+   "metadata": {},
+   "source": [
+    "## 3. Finding the Optimal Number of Neighbors\n",
+    "\n",
+    "Before doing a full neighbor search, we can use reLAISS to suggest an optimal number of neighbors based on the distance distribution. This helps avoid arbitrary choices for the number of neighbors to return.\n",
+    "\n",
+    "First, let's run a search with a larger number of neighbors and set `suggest_neighbor_num=True`. This will show us a distance plot that helps identify a reasonable cutoff point for similar objects."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "91904c57",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Find optimal number of neighbors\n",
+    "client.find_neighbors(\n",
+    "    ztf_object_id='ZTF21aaublej',  # ZTF ID to find neighbors for\n",
+    "    n=40,  # Search in a larger pool\n",
+    "    suggest_neighbor_num=True,  # Only suggest optimal number, don't return neighbors\n",
+    "    plot=True,  # Show the distance elbow plot\n",
+    "    save_figures=True,  # Save plots to disk\n",
+    "    path_to_figure_directory='./figures'\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "83324aa1",
+   "metadata": {},
+   "source": [
+    "## 4. Basic Nearest Neighbor Search\n",
+    "\n",
+    "Now we can find the most similar transients to a given ZTF object. Let's use ZTF21aaublej as an example.\n",
+    "\n",
+    "The `find_neighbors` function allows you to:\n",
+    "- Specify the number of neighbors to return\n",
+    "- Set a maximum distance threshold\n",
+    "- Adjust the weight of lightcurve features relative to host features\n",
+    "- Generate diagnostic plots\n",
+    "\n",
+    "Based on the distance curve we saw earlier, we'll choose to return 5 neighbors."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "550f1c6c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Find nearest neighbors\n",
+    "neighbors_df = client.find_neighbors(\n",
+    "    ztf_object_id='ZTF21aaublej',  # ZTF ID to find neighbors for\n",
+    "    n=5,  # Number of neighbors to return\n",
+    "    suggest_neighbor_num=False,  # Return actual neighbors\n",
+    "    plot=True,  # Generate diagnostic plots\n",
+    "    save_figures=True,  # Save plots to disk\n",
+    "    path_to_figure_directory='./figures'\n",
+    ")\n",
+    "\n",
+    "# Display the results\n",
+    "print(\"\\nNearest Neighbors:\")\n",
+    "print(neighbors_df)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "de3a82fb",
+   "metadata": {},
+   "source": [
+    "## 5. Using Monte Carlo Simulations and Feature Weighting\n",
+    "\n",
+    "reLAISS allows you to adjust the relative importance of lightcurve features compared to host galaxy features using the `weight_lc_feats_factor` parameter. A value greater than 1.0 will make lightcurve features more important in the similarity search.\n",
+    "\n",
+    "The Monte Carlo simulation functionality (`num_sims` parameter) helps account for measurement uncertainties by running multiple simulations with perturbed feature values.\n",
+    "\n",
+    "If you find that your matches aren't quite what you're looking for, you can try:\n",
+    "- Using Monte Carlo simulations to account for feature measurement uncertainties\n",
+    "- Upweighting lightcurve features to focus more on the transient's photometric properties than its host\n",
+    "- Removing host features entirely for a \"lightcurve-only\" search\n",
+    "- Removing lightcurve features for a \"host-only\" search\n",
+    "\n",
+    "Let's try using Monte Carlo simulations with upweighted lightcurve features:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ba5d5748",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Using Monte Carlo simulations and feature weighting\n",
+    "neighbors_df = client.find_neighbors(\n",
+    "    ztf_object_id='ZTF21aaublej',  # Using the test transient\n",
+    "    n=5,\n",
+    "    num_sims=20,  # Number of Monte Carlo simulations\n",
+    "    weight_lc_feats_factor=3.0,  # Up-weight lightcurve features\n",
+    "    plot=True,\n",
+    "    save_figures=True,\n",
+    "    path_to_figure_directory='./figures'\n",
+    ")\n",
+    "\n",
+    "print(\"\\nNearest neighbors with Monte Carlo simulations:\")\n",
+    "print(neighbors_df)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e4a78f7d",
+   "metadata": {},
+   "source": [
+    "## 6. Basic Anomaly Detection\n",
+    "\n",
+    "reLAISS also includes tools for anomaly detection that can help identify unusual transients. The anomaly detection module uses an Isolation Forest algorithm to identify outliers in the feature space.\n",
+    "\n",
+    "The anomaly detection process will produce plots showing the lightcurve of the input transient and a graph of the probability (in time) that the transient is anomalous. If the probability exceeds 50% at any epoch, the transient is flagged as anomalous.\n",
+    "\n",
+    "### Training an Anomaly Detection Model\n",
+    "\n",
+    "First, let's train an anomaly detection model on our dataset bank:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e19948ad",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from relaiss.anomaly import train_AD_model\n",
+    "\n",
+    "# Train the anomaly detection model\n",
+    "model_path = train_AD_model(\n",
+    "    lc_features=client.lc_features,\n",
+    "    host_features=client.host_features,\n",
+    "    path_to_dataset_bank=client.bank_csv,\n",
+    "    path_to_sfd_folder='./sfddata-master',\n",
+    "    path_to_models_directory=\"./models\",\n",
+    "    n_estimators=100,  # Using a smaller value for faster execution\n",
+    "    contamination=0.02,  # Expected proportion of anomalies\n",
+    "    max_samples=256,  # Maximum samples used for each tree\n",
+    "    force_retrain=False  # Only retrain if model doesn't exist\n",
+    ")\n",
+    "\n",
+    "print(f\"Model saved to: {model_path}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a104e058",
+   "metadata": {},
+   "source": [
+    "### Running Anomaly Detection on a Transient\n",
+    "\n",
+    "Now we can run anomaly detection on a specific transient to see if it's considered anomalous:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "53cbc236",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from relaiss.anomaly import anomaly_detection\n",
+    "\n",
+    "# Run anomaly detection on a transient\n",
+    "anomaly_detection(\n",
+    "    transient_ztf_id=\"ZTF21aaublej\",  # Use the same transient for this example\n",
+    "    lc_features=client.lc_features,\n",
+    "    host_features=client.host_features,\n",
+    "    path_to_timeseries_folder=\"./timeseries\",\n",
+    "    path_to_sfd_folder='./sfddata-master',\n",
+    "    path_to_dataset_bank=client.bank_csv,\n",
+    "    path_to_models_directory=\"./models\",\n",
+    "    path_to_figure_directory=\"./figures\",\n",
+    "    save_figures=True,\n",
+    "    n_estimators=100,\n",
+    "    contamination=0.02,\n",
+    "    max_samples=256,\n",
+    "    force_retrain=False\n",
+    ")\n",
+    "\n",
+    "print(\"Anomaly detection figures saved to ./figures/AD/\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c23a1b0a",
+   "metadata": {},
+   "source": [
+    "## Next Steps\n",
+    "\n",
+    "To explore more advanced features, check out the `advanced_usage.ipynb` notebook which covers:\n",
+    "- Using PCA for dimensionality reduction\n",
+    "- Creating theorized lightcurves\n",
+    "- Swapping host galaxies\n",
+    "- Setting maximum neighbor distances\n",
+    "- Tweaking ANNOY parameters\n",
+    "- Making corner plots\n",
+    "- Advanced anomaly detection techniques"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

notebooks/02_advanced_usage.ipynb

@@ -0,0 +1,488 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "4b362a55",
+   "metadata": {},
+   "source": [
+    "# Advanced Usage of reLAISS\n",
+    "### Authors: Evan Reynolds and Alex Gagliano\n",
+    "\n",
+    "## Introduction\n",
+    "\n",
+    "This notebook demonstrates advanced features of the reLAISS library for finding similar astronomical transients. While the basic_usage.ipynb notebook covered the fundamental functionality, here we'll explore more sophisticated techniques that give you greater flexibility and power in your analysis.\n",
+    "\n",
+    "These advanced features allow you to customize how reLAISS processes and analyzes data, including dimensionality reduction, theorized lightcurves, host galaxy swapping, fine-tuning of algorithm parameters, visualization tools, and advanced anomaly detection.\n",
+    "\n",
+    "## Topics Covered\n",
+    "1. Using PCA for dimensionality reduction\n",
+    "2. Creating and using theorized lightcurves\n",
+    "3. Swapping host galaxies\n",
+    "4. Setting maximum neighbor distances\n",
+    "5. Tweaking ANNOY parameters\n",
+    "6. Making corner plots\n",
+    "7. Advanced anomaly detection with parameter tuning\n",
+    "8. Host swapping in anomaly detection"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e3c3f836",
+   "metadata": {},
+   "source": [
+    "## Setup\n",
+    "\n",
+    "First, let's import the necessary packages and create the required directories:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3cd152f1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import pandas as pd\n",
+    "import numpy as np\n",
+    "import relaiss as rl\n",
+    "import astropy.units as u\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "# Create output directories\n",
+    "os.makedirs('./figures', exist_ok=True)\n",
+    "os.makedirs('./sfddata-master', exist_ok=True)\n",
+    "os.makedirs('./models', exist_ok=True)\n",
+    "os.makedirs('./timeseries', exist_ok=True)\n",
+    "\n",
+    "def create_theorized_lightcurve():\n",
+    "    \"\"\"Create a simple theorized lightcurve for demonstration.\"\"\"\n",
+    "    # Create time points\n",
+    "    times = np.linspace(0, 100, 50) * u.day\n",
+    "    # Create magnitudes (simple gaussian)\n",
+    "    mags = 20 + 2 * np.exp(-(times.value - 50)**2 / 100)\n",
+    "    # Create errors\n",
+    "    errors = np.ones_like(mags) * 0.1\n",
+    "    \n",
+    "    # Create DataFrame in ANTARES format\n",
+    "    df = pd.DataFrame({\n",
+    "        'ant_mjd': times.to(u.day).value,\n",
+    "        'ant_mag': mags,\n",
+    "        'ant_magerr': errors,\n",
+    "        'ant_passband': ['g' if i % 2 == 0 else 'R' for i in range(len(times))]\n",
+    "    })\n",
+    "    return df"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "60730aa2",
+   "metadata": {},
+   "source": [
+    "## Initialize the ReLAISS Client\n",
+    "\n",
+    "We'll start by creating a ReLAISS client instance:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b3adc953",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initialize the client\n",
+    "client = rl.ReLAISS()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "67b42e90",
+   "metadata": {},
+   "source": [
+    "## 1. Using PCA for Dimensionality Reduction\n",
+    "\n",
+    "PCA (Principal Component Analysis) can be used to reduce the dimensionality of the feature space while preserving most of the variance. This has several benefits:\n",
+    "\n",
+    "- Improves search speed by reducing the computational complexity\n",
+    "- Potentially reduces noise in the feature space\n",
+    "- Helps mitigate the \"curse of dimensionality\" for high-dimensional data\n",
+    "\n",
+    "To use PCA, we set `use_pca=True` in the `load_reference` method and specify the number of components to keep:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a136889a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "client.load_reference(\n",
+    "    path_to_sfd_folder='./sfddata-master',\n",
+    "    use_pca=True,\n",
+    "    num_pca_components=20,  # Keep 20 PCA components\n",
+    ")\n",
+    "\n",
+    "neighbors_df = client.find_neighbors(\n",
+    "    ztf_object_id='ZTF21abbzjeq',  # Using the test transient\n",
+    "    n=5,\n",
+    "    plot=True,\n",
+    "    save_figures=True,\n",
+    "    path_to_figure_directory='./figures'\n",
+    ")\n",
+    "print(\"\\nNearest neighbors using PCA:\")\n",
+    "print(neighbors_df)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f0ac7c6e",
+   "metadata": {},
+   "source": [
+    "## 2. Creating and Using Theorized Lightcurves\n",
+    "\n",
+    "One powerful feature of reLAISS is the ability to use theorized (synthetic) lightcurves in the neighbor search. This allows you to:\n",
+    "\n",
+    "- Test theoretical models against observed data\n",
+    "- Explore \"what-if\" scenarios by creating custom lightcurves\n",
+    "- Find real transients that match your theoretical predictions\n",
+    "\n",
+    "Below, we create a simple Gaussian-shaped lightcurve and find its nearest neighbors:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "16bb77d7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Create a theorized lightcurve\n",
+    "theorized_lc = create_theorized_lightcurve()\n",
+    "\n",
+    "# Find neighbors for the theorized lightcurve\n",
+    "# Need to provide a host galaxy when using theorized lightcurve\n",
+    "neighbors_df = client.find_neighbors(\n",
+    "    theorized_lightcurve_df=theorized_lc,\n",
+    "    host_ztf_id='ZTF21abbzjeq',  # Use this transient's host\n",
+    "    n=5,\n",
+    "    plot=True,\n",
+    "    save_figures=True,\n",
+    "    path_to_figure_directory='./figures'\n",
+    ")\n",
+    "print(\"\\nNearest neighbors for theorized lightcurve:\")\n",
+    "print(neighbors_df)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bab57bd3",
+   "metadata": {},
+   "source": [
+    "## 3. Swapping Host Galaxies\n",
+    "\n",
+    "reLAISS allows you to swap the host galaxy of a transient, which is useful for:\n",
+    "\n",
+    "- Exploring how host properties affect the similarity search results\n",
+    "- Investigating the effects of different environments on transient characteristics\n",
+    "- Testing hypotheses about host galaxy influences\n",
+    "\n",
+    "Here's how to swap in a different host galaxy:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "efdf2937",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Find neighbors with a swapped host galaxy\n",
+    "neighbors_df = client.find_neighbors(\n",
+    "    ztf_object_id='ZTF21abbzjeq',  # Source transient\n",
+    "    host_ztf_id='ZTF21aakswqr',  # Host to swap in\n",
+    "    n=5,\n",
+    "    plot=True,\n",
+    "    save_figures=True,\n",
+    "    path_to_figure_directory='./figures'\n",
+    ")\n",
+    "print(\"\\nNearest neighbors with swapped host galaxy:\")\n",
+    "print(neighbors_df)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0bd99f9c",
+   "metadata": {},
+   "source": [
+    "## 4. Setting Maximum Neighbor Distances\n",
+    "\n",
+    "Sometimes you're only interested in neighbors that are truly similar to your target. By setting a maximum distance threshold, you can:\n",
+    "\n",
+    "- Filter out neighbors that are too dissimilar\n",
+    "- Focus only on highly confident matches\n",
+    "- Avoid including poor matches just to reach a specific number of neighbors\n",
+    "\n",
+    "Note that you might get fewer neighbors than requested if the distance threshold is applied:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3a859853",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Find neighbors with maximum distance constraint\n",
+    "neighbors_df = client.find_neighbors(\n",
+    "    ztf_object_id='ZTF21abbzjeq',\n",
+    "    n=5,\n",
+    "    max_neighbor_dist=0.5,  # Only return neighbors within this distance\n",
+    "    plot=True,\n",
+    "    save_figures=True,\n",
+    "    path_to_figure_directory='./figures'\n",
+    ")\n",
+    "print(\"\\nNearest neighbors with maximum distance constraint:\")\n",
+    "print(neighbors_df)\n",
+    "print(f\"Number of neighbors found: {len(neighbors_df)}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a71acf61",
+   "metadata": {},
+   "source": [
+    "## 5. Tweaking ANNOY Parameters\n",
+    "\n",
+    "ANNOY (Approximate Nearest Neighbors Oh Yeah) is the algorithm used for fast nearest neighbor search. You can tune its parameters to balance search accuracy and speed:\n",
+    "\n",
+    "- `search_k`: Controls the number of nodes to explore during search (higher = more accurate but slower)\n",
+    "- `n_trees`: Controls the number of random projection trees built (set during client initialization)\n",
+    "\n",
+    "Here's how to adjust the search_k parameter:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "aaa03a6c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Find neighbors with tweaked ANNOY parameters\n",
+    "neighbors_df = client.find_neighbors(\n",
+    "    ztf_object_id='ZTF21abbzjeq',\n",
+    "    n=5,\n",
+    "    search_k=2000,  # Increase search_k for more accurate results\n",
+    "    plot=True,\n",
+    "    save_figures=True,\n",
+    "    path_to_figure_directory='./figures'\n",
+    ")\n",
+    "print(\"\\nNearest neighbors with tweaked ANNOY parameters:\")\n",
+    "print(neighbors_df)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5d7522b8",
+   "metadata": {},
+   "source": [
+    "## 6. Making Corner Plots\n",
+    "\n",
+    "Corner plots are a powerful visualization tool that show the distribution of features for the input transient and its neighbors. They can help you:\n",
+    "\n",
+    "- Understand which features are driving the similarity matching\n",
+    "- Identify potential correlations between different features\n",
+    "- Visualize the feature space and where your transient sits within it\n",
+    "\n",
+    "To create corner plots, we need to first get the primer_dict containing information about the input transient:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8d7ff16d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Get neighbors from a new search\n",
+    "neighbors_df = client.find_neighbors(\n",
+    "    ztf_object_id='ZTF21abbzjeq',\n",
+    "    n=5,\n",
+    "    plot=True,\n",
+    "    save_figures=True,\n",
+    "    path_to_figure_directory='./figures'\n",
+    ")\n",
+    "\n",
+    "# Get primer_dict separately\n",
+    "from relaiss.search import primer\n",
+    "primer_dict = primer(\n",
+    "    lc_ztf_id='ZTF21abbzjeq',\n",
+    "    theorized_lightcurve_df=None,\n",
+    "    host_ztf_id=None,\n",
+    "    dataset_bank_path=client.bank_csv,\n",
+    "    path_to_timeseries_folder='./',\n",
+    "    path_to_sfd_folder=client.path_to_sfd_folder,\n",
+    "    lc_features=client.lc_features,\n",
+    "    host_features=client.host_features,\n",
+    "    num_sims=0,\n",
+    "    save_timeseries=False,\n",
+    ")\n",
+    "\n",
+    "# Create corner plots using the primer_dict\n",
+    "from relaiss.plotting import corner_plot\n",
+    "corner_plot(\n",
+    "    neighbors_df=neighbors_df,\n",
+    "    primer_dict=primer_dict,\n",
+    "    path_to_dataset_bank=client.bank_csv,\n",
+    "    path_to_figure_directory='./figures',\n",
+    "    save_plots=True\n",
+    ")\n",
+    "print(\"Corner plots saved to ./figures/\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7079194a",
+   "metadata": {},
+   "source": [
+    "## 7. Advanced Anomaly Detection with Parameter Tuning\n",
+    "\n",
+    "The anomaly detection module in reLAISS uses an Isolation Forest algorithm that can be tuned for different scenarios. Key parameters include:\n",
+    "\n",
+    "- `n_estimators`: Number of base estimators (trees) in the ensemble\n",
+    "- `contamination`: Expected proportion of outliers in the dataset\n",
+    "- `max_samples`: Number of samples drawn to train each base estimator\n",
+    "\n",
+    "Let's explore how different parameters affect the model's performance:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "726834d5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from relaiss.anomaly import train_AD_model, anomaly_detection\n",
+    "\n",
+    "# Train models with different parameters to compare\n",
+    "print(\"Training anomaly detection model with default parameters...\")\n",
+    "default_model_path = train_AD_model(\n",
+    "    lc_features=client.lc_features,\n",
+    "    host_features=client.host_features,\n",
+    "    path_to_dataset_bank=client.bank_csv,\n",
+    "    path_to_sfd_folder='./sfddata-master',\n",
+    "    path_to_models_directory=\"./models\",\n",
+    "    n_estimators=100,\n",
+    "    contamination=0.02,\n",
+    "    max_samples=256,\n",
+    "    force_retrain=True\n",
+    ")\n",
+    "\n",
+    "print(\"Training anomaly detection model with more trees...\")\n",
+    "model_more_trees_path = train_AD_model(\n",
+    "    lc_features=client.lc_features,\n",
+    "    host_features=client.host_features,\n",
+    "    path_to_dataset_bank=client.bank_csv,\n",
+    "    path_to_sfd_folder='./sfddata-master',\n",
+    "    path_to_models_directory=\"./models\",\n",
+    "    n_estimators=200,  # More trees\n",
+    "    contamination=0.02,\n",
+    "    max_samples=256,\n",
+    "    force_retrain=True\n",
+    ")\n",
+    "\n",
+    "print(\"Training anomaly detection model with higher contamination...\")\n",
+    "model_higher_contam_path = train_AD_model(\n",
+    "    lc_features=client.lc_features,\n",
+    "    host_features=client.host_features,\n",
+    "    path_to_dataset_bank=client.bank_csv,\n",
+    "    path_to_sfd_folder='./sfddata-master',\n",
+    "    path_to_models_directory=\"./models\",\n",
+    "    n_estimators=100,\n",
+    "    contamination=0.05,  # Higher contamination\n",
+    "    max_samples=256,\n",
+    "    force_retrain=True\n",
+    ")\n",
+    "\n",
+    "# Run anomaly detection with default model\n",
+    "print(\"\\nRunning anomaly detection with default model...\")\n",
+    "anomaly_detection(\n",
+    "    transient_ztf_id=\"ZTF21abbzjeq\",\n",
+    "    lc_features=client.lc_features,\n",
+    "    host_features=client.host_features,\n",
+    "    path_to_timeseries_folder=\"./timeseries\",\n",
+    "    path_to_sfd_folder='./sfddata-master',\n",
+    "    path_to_dataset_bank=client.bank_csv,\n",
+    "    path_to_models_directory=\"./models\",\n",
+    "    path_to_figure_directory=\"./figures/AD_default\",\n",
+    "    save_figures=True,\n",
+    "    n_estimators=100,\n",
+    "    contamination=0.02,\n",
+    "    max_samples=256,\n",
+    "    force_retrain=False\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4886eb0e",
+   "metadata": {},
+   "source": [
+    "## 8. Anomaly Detection with Host Swapping\n",
+    "\n",
+    "Just as with neighbor searches, you can swap host galaxies for anomaly detection. This helps you understand how host properties contribute to a transient's anomaly score.\n",
+    "\n",
+    "This feature is particularly useful for:\n",
+    "- Testing if the anomalous nature of a transient is due to its host galaxy\n",
+    "- Exploring the \"what if\" scenario of a transient occurring in a different environment\n",
+    "- Separating intrinsic transient anomalies from host-related factors"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5c65e14e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Use the default model but swap in a different host galaxy\n",
+    "anomaly_detection(\n",
+    "    transient_ztf_id=\"ZTF21abbzjeq\",\n",
+    "    lc_features=client.lc_features,\n",
+    "    host_features=client.host_features,\n",
+    "    path_to_timeseries_folder=\"./timeseries\",\n",
+    "    path_to_sfd_folder='./sfddata-master',\n",
+    "    path_to_dataset_bank=client.bank_csv,\n",
+    "    host_ztf_id_to_swap_in=\"ZTF21aakswqr\",  # Swap in a different host\n",
+    "    path_to_models_directory=\"./models\",\n",
+    "    path_to_figure_directory=\"./figures/AD_host_swap\",\n",
+    "    save_figures=True,\n",
+    "    n_estimators=100,\n",
+    "    contamination=0.02,\n",
+    "    max_samples=256,\n",
+    "    force_retrain=False\n",
+    ")\n",
+    "\n",
+    "print(\"Anomaly detection figures saved to ./figures/AD_default/ and ./figures/AD_host_swap/\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e08b6ef6",
+   "metadata": {},
+   "source": [
+    "## Conclusion\n",
+    "\n",
+    "By combining these features, you can create highly customized searches tailored to your specific research questions.\n",
+    "\n",
+    "For information on how to build your own dataset bank for reLAISS, see the `build_databank.ipynb` notebook."
+   ]
+  }
+ ],
+ "metadata": {},
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

notebooks/03_build_databank.ipynb

@@ -0,0 +1,384 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "3ede68f0",
+   "metadata": {},
+   "source": [
+    "# Building a New Dataset Bank for reLAISS\n",
+    "### Authors: Evan Reynolds and Alex Gagliano\n",
+    "\n",
+    "## Introduction\n",
+    "\n",
+    "This notebook demonstrates how to build a new dataset bank for reLAISS and use different feature combinations for nearest neighbor searches. The dataset bank is the foundation of reLAISS, containing all the features of transients that are used for similarity searches and anomaly detection.\n",
+    "\n",
+    "Building your own dataset bank allows you to incorporate new data, apply custom preprocessing steps, and tailor the feature set to your specific research needs.\n",
+    "\n",
+    "## Topics Covered\n",
+    "1. Adding extinction corrections (A_V)\n",
+    "2. Joining new lightcurve features\n",
+    "3. Handling missing values\n",
+    "4. Building the final dataset bank\n",
+    "5. Using different feature combinations for nearest neighbor search:\n",
+    "   - Lightcurve-only features\n",
+    "   - Host-only features\n",
+    "   - Custom feature subsets\n",
+    "   - Feature weighting"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9212bf17",
+   "metadata": {},
+   "source": [
+    "## Setup\n",
+    "\n",
+    "First, let's import the necessary libraries and create the required directories:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ff694547",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import pandas as pd\n",
+    "import numpy as np\n",
+    "from relaiss import constants\n",
+    "import relaiss as rl\n",
+    "\n",
+    "# Create necessary directories\n",
+    "os.makedirs('./figures', exist_ok=True)\n",
+    "os.makedirs('./sfddata-master', exist_ok=True)\n",
+    "\n",
+    "# Define default feature sets from constants\n",
+    "default_lc_features = constants.lc_features_const.copy()\n",
+    "default_host_features = constants.host_features_const.copy()\n",
+    "\n",
+    "# Initialize client\n",
+    "client = rl.ReLAISS()\n",
+    "client.load_reference(\n",
+    "    path_to_sfd_folder='./sfddata-master'\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c7229e7f",
+   "metadata": {},
+   "source": [
+    "## 1. Adding Extinction Corrections (A_V)\n",
+    "\n",
+    "The first step in building a dataset bank is to add extinction corrections to account for interstellar dust. The Schlegel, Finkbeiner & Davis (SFD) dust maps are used to estimate the amount of extinction.\n",
+    "\n",
+    "```python\n",
+    "# Example code for adding extinction corrections\n",
+    "from sfdmap2 import sfdmap\n",
+    "\n",
+    "df = pd.read_csv(\"../data/large_df_bank.csv\")\n",
+    "m = sfdmap.SFDMap('../data/sfddata-master')\n",
+    "RV = 3.1  # Standard value for Milky Way\n",
+    "ebv = m.ebv(df['ra'].values, df['dec'].values)\n",
+    "df['A_V'] = RV * ebv\n",
+    "df.to_csv(\"../data/large_df_bank_wAV.csv\", index=False)\n",
+    "```\n",
+    "\n",
+    "This adds the A_V (extinction in V-band) column to your dataset, which will be used later in the feature processing pipeline."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ac7ecbca",
+   "metadata": {},
+   "source": [
+    "## 2. Joining New Lightcurve Features\n",
+    "\n",
+    "If you have additional features in a separate dataset, you can merge them with your existing bank:\n",
+    "\n",
+    "```python\n",
+    "# Example code for joining features\n",
+    "df_large = pd.read_csv(\"../data/large_df_bank_wAV.csv\")\n",
+    "df_small = pd.read_csv(\"../data/small_df_bank_re_laiss.csv\")\n",
+    "\n",
+    "key = 'ztf_object_id'\n",
+    "extra_features = [col for col in df_large.columns if col not in df_small.columns]\n",
+    "\n",
+    "merged_df = df_small.merge(df_large[[key] + extra_features], on=key, how='left')\n",
+    "\n",
+    "lc_feature_names = constants.lc_features_const.copy()\n",
+    "host_feature_names = constants.host_features_const.copy()\n",
+    "\n",
+    "small_final_df = merged_df.replace([np.inf, -np.inf, -999], np.nan).dropna(subset=lc_feature_names + host_feature_names)\n",
+    "\n",
+    "small_final_df.to_csv(\"../data/small_hydrated_df_bank_re_laiss.csv\", index=False)\n",
+    "```\n",
+    "\n",
+    "This merges additional features from a larger dataset into your working dataset."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a0fb3a25",
+   "metadata": {},
+   "source": [
+    "## 3. Handling Missing Values\n",
+    "\n",
+    "Missing values in the dataset can cause problems during analysis. reLAISS uses KNN imputation to fill in missing values:\n",
+    "\n",
+    "```python\n",
+    "# Example code for handling missing values\n",
+    "from sklearn.impute import KNNImputer\n",
+    "\n",
+    "raw_host_feature_names = constants.raw_host_features_const.copy()\n",
+    "raw_dataset_bank = pd.read_csv('../data/large_df_bank_wAV.csv')\n",
+    "\n",
+    "X = raw_dataset_bank[lc_feature_names + raw_host_feature_names]\n",
+    "feat_imputer = KNNImputer(weights='distance').fit(X)\n",
+    "imputed_filt_arr = feat_imputer.transform(X)\n",
+    "\n",
+    "imputed_df = pd.DataFrame(imputed_filt_arr, columns=lc_feature_names + raw_host_feature_names)\n",
+    "imputed_df.index = raw_dataset_bank.index\n",
+    "raw_dataset_bank[lc_feature_names + raw_host_feature_names] = imputed_df\n",
+    "\n",
+    "imputed_df_bank = raw_dataset_bank\n",
+    "```\n",
+    "\n",
+    "KNN imputation works by finding the k-nearest neighbors in feature space for samples with missing values and using their values to fill in the gaps."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "390d9e9d",
+   "metadata": {},
+   "source": [
+    "## 4. Building the Final Dataset Bank\n",
+    "\n",
+    "With all the preprocessing done, we can now build the final dataset bank using the `build_dataset_bank` function from reLAISS:\n",
+    "\n",
+    "```python\n",
+    "# Example code for building the final dataset bank\n",
+    "from relaiss.features import build_dataset_bank\n",
+    "\n",
+    "dataset_bank = build_dataset_bank(\n",
+    "    raw_df_bank=imputed_df_bank,\n",
+    "    av_in_raw_df_bank=True,\n",
+    "    path_to_sfd_folder=\"../data/sfddata-master\",\n",
+    "    building_entire_df_bank=True\n",
+    ")\n",
+    "\n",
+    "# Clean and save final dataset\n",
+    "final_dataset_bank = dataset_bank.replace(\n",
+    "    [np.inf, -np.inf, -999], np.nan\n",
+    ").dropna(subset=lc_feature_names + host_feature_names)\n",
+    "\n",
+    "final_dataset_bank.to_csv('../data/large_final_df_bank_new_lc_feats.csv', index=False)\n",
+    "```\n",
+    "\n",
+    "This function applies additional processing to prepare the features for reLAISS, including normalization and other transformations."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "597b7b19",
+   "metadata": {},
+   "source": [
+    "## 5. Using Different Feature Combinations\n",
+    "\n",
+    "reLAISS allows you to customize which features are used for similarity search. This can be useful for studying the importance of different features and for tailoring the search to specific scientific questions."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8eb99200",
+   "metadata": {},
+   "source": [
+    "### 5.1 Using Only Lightcurve Features\n",
+    "\n",
+    "You can perform a search using only lightcurve features, ignoring host galaxy properties. This is useful when:\n",
+    "- You want to focus solely on the temporal evolution of the transient\n",
+    "- Host data might be unreliable or missing\n",
+    "- You're testing hypotheses about lightcurve-based classification\n",
+    "\n",
+    "Here's how to set up a lightcurve-only search:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bde7ec63",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "lc_only_client = rl.ReLAISS()\n",
+    "lc_only_client.load_reference(\n",
+    "    path_to_sfd_folder='./sfddata-master',\n",
+    "    lc_features=default_lc_features,  # Use default lightcurve features\n",
+    "    host_features=[],  # Empty list means no host features\n",
+    ")\n",
+    "\n",
+    "# Find neighbors using only lightcurve features\n",
+    "neighbors_df_lc_only = lc_only_client.find_neighbors(\n",
+    "    ztf_object_id='ZTF21abbzjeq',\n",
+    "    n=5,\n",
+    "    plot=True,\n",
+    "    save_figures=True,\n",
+    "    path_to_figure_directory='./figures/lc_only'\n",
+    ")\n",
+    "print(\"\\nNearest neighbors using only lightcurve features:\")\n",
+    "print(neighbors_df_lc_only)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "865b24e5",
+   "metadata": {},
+   "source": [
+    "### 5.2 Using Only Host Features\n",
+    "\n",
+    "Alternatively, you can perform a search using only host galaxy features, ignoring the lightcurve properties. This approach is valuable when:\n",
+    "- You're more interested in environmental effects on transients\n",
+    "- You want to find transients in similar host galaxies\n",
+    "- You're studying correlations between host properties and transient types\n",
+    "\n",
+    "Here's how to set up a host-only search:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "305d8042",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "host_only_client = rl.ReLAISS()\n",
+    "host_only_client.load_reference(\n",
+    "    path_to_sfd_folder='./sfddata-master',\n",
+    "    lc_features=[],  # Empty list means no lightcurve features\n",
+    "    host_features=default_host_features,  # Use default host features\n",
+    ")\n",
+    "\n",
+    "# Find neighbors using only host features\n",
+    "neighbors_df_host_only = host_only_client.find_neighbors(\n",
+    "    ztf_object_id='ZTF21abbzjeq',\n",
+    "    n=5,\n",
+    "    plot=True,\n",
+    "    save_figures=True,\n",
+    "    path_to_figure_directory='./figures/host_only'\n",
+    ")\n",
+    "print(\"\\nNearest neighbors using only host features:\")\n",
+    "print(neighbors_df_host_only)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fbe91d24",
+   "metadata": {},
+   "source": [
+    "### 5.3 Using Custom Feature Subset\n",
+    "\n",
+    "You can also select specific features from both categories for a more targeted search. This allows you to:\n",
+    "- Focus on the features most relevant to your research question\n",
+    "- Reduce noise by excluding less useful features\n",
+    "- Test hypotheses about which features drive similarity\n",
+    "\n",
+    "Here's how to create a custom feature subset:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c7c721c7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Select specific lightcurve and host features\n",
+    "custom_lc_features = ['g_peak_mag', 'r_peak_mag', 'g_peak_time', 'r_peak_time']\n",
+    "custom_host_features = ['host_ra', 'host_dec', 'gKronMag', 'rKronMag']\n",
+    "\n",
+    "custom_client = rl.ReLAISS()\n",
+    "custom_client.load_reference(\n",
+    "    path_to_sfd_folder='./sfddata-master',\n",
+    "    lc_features=custom_lc_features,  # Custom subset of lightcurve features\n",
+    "    host_features=custom_host_features,  # Custom subset of host features\n",
+    ")\n",
+    "\n",
+    "# Find neighbors with custom feature subset\n",
+    "neighbors_df_custom = custom_client.find_neighbors(\n",
+    "    ztf_object_id='ZTF21abbzjeq',\n",
+    "    n=5,\n",
+    "    plot=True,\n",
+    "    save_figures=True,\n",
+    "    path_to_figure_directory='./figures/custom_features'\n",
+    ")\n",
+    "print(\"\\nNearest neighbors using custom feature subset:\")\n",
+    "print(neighbors_df_custom)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7da9ee34",
+   "metadata": {},
+   "source": [
+    "### 5.4 Using Feature Weighting\n",
+    "\n",
+    "You can also adjust the relative importance of lightcurve features versus host galaxy features using the `weight_lc_feats_factor` parameter:\n",
+    "- Values > 1: Emphasize lightcurve features\n",
+    "- Values < 1: Emphasize host features\n",
+    "- Value = 1: Equal weighting (default)\n",
+    "\n",
+    "This allows you to fine-tune the balance between photometric and host properties:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1209d70b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Regular search prioritizing lightcurve features\n",
+    "neighbors_df_lc_weighted = client.find_neighbors(\n",
+    "    ztf_object_id='ZTF21abbzjeq',\n",
+    "    n=5,\n",
+    "    weight_lc_feats_factor=3.0,  # Strongly prioritize lightcurve features\n",
+    "    plot=True,\n",
+    "    save_figures=True,\n",
+    "    path_to_figure_directory='./figures/lc_weighted'\n",
+    ")\n",
+    "print(\"\\nNearest neighbors with lightcurve features weighted 3x:\")\n",
+    "print(neighbors_df_lc_weighted)\n",
+    "\n",
+    "# Now prioritize host features by using a factor < 1\n",
+    "neighbors_df_host_weighted = client.find_neighbors(\n",
+    "    ztf_object_id='ZTF21abbzjeq',\n",
+    "    n=5,\n",
+    "    weight_lc_feats_factor=0.3,  # Prioritize host features\n",
+    "    plot=True,\n",
+    "    save_figures=True,\n",
+    "    path_to_figure_directory='./figures/host_weighted'\n",
+    ")\n",
+    "print(\"\\nNearest neighbors with host features given higher weight:\")\n",
+    "print(neighbors_df_host_weighted)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d243999b",
+   "metadata": {},
+   "source": [
+    "## Conclusion\n",
+    "\n",
+    "Building your own dataset bank and customizing feature combinations provides powerful flexibility for tailoring reLAISS to your specific research questions. By selecting different feature combinations and adjusting feature weights, you can explore various aspects of transient similarity and discover new insights about the transient population."
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

src/relaiss/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.0.post1.dev44'
-__version_tuple__ = version_tuple = (0, 0, 'post1', 'dev44')
+__version__ = version = '1.0.0'
+__version_tuple__ = version_tuple = (1, 0, 0)

src/relaiss/anomaly.py

@@ -5,6 +5,7 @@ import antares_client
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
+import joblib
 from pyod.models.iforest import IForest
 from sklearn.pipeline import Pipeline
 from sklearn.preprocessing import StandardScaler
@@ -17,6 +18,7 @@ def train_AD_model(
     host_features,
     path_to_dataset_bank=None,
     preprocessed_df=None,
+    path_to_sfd_folder=None,
     path_to_models_directory="../models",
     n_estimators=500,
     contamination=0.02,
@@ -36,6 +38,8 @@ def train_AD_model(
     preprocessed_df : pandas.DataFrame | None, optional
         Pre-processed dataframe with imputed features. If provided, this is used
         instead of loading and processing the raw dataset bank.
+    path_to_sfd_folder : str | Path | None, optional
+        Path to SFD dust maps.
     path_to_models_directory : str | Path, default "../models"
         Directory to save trained models.
     n_estimators : int, default 500
@@ -88,6 +92,7 @@ def train_AD_model(
         raw_df = pd.read_csv(path_to_dataset_bank)
         df = build_dataset_bank(
             raw_df,
+            path_to_sfd_folder=path_to_sfd_folder,
             building_entire_df_bank=True,
             building_for_AD=True
         )
@@ -118,7 +123,7 @@ def anomaly_detection(
     lc_features,
     host_features,
     path_to_timeseries_folder,
-    path_to_sfd_data_folder,
+    path_to_sfd_folder,
     path_to_dataset_bank,
     host_ztf_id_to_swap_in=None,
     path_to_models_directory="../models",
@@ -159,6 +164,7 @@ def anomaly_detection(
         lc_features,
         host_features,
         path_to_dataset_bank,
+        path_to_sfd_folder=path_to_sfd_folder,
         path_to_models_directory=path_to_models_directory,
         n_estimators=n_estimators,
         contamination=contamination,
@@ -167,8 +173,7 @@ def anomaly_detection(
     )
 
     # Load the model
-    with open(path_to_trained_model, "rb") as f:
-        clf = pickle.load(f)
+    clf = joblib.load(path_to_trained_model)
 
     # Load the timeseries dataframe
     print("\nRebuilding timeseries dataframe(s) for AD...")
@@ -176,7 +181,7 @@ def anomaly_detection(
         ztf_id=transient_ztf_id,
         theorized_lightcurve_df=None,
         path_to_timeseries_folder=path_to_timeseries_folder,
-        path_to_sfd_data_folder=path_to_sfd_data_folder,
+        path_to_sfd_folder=path_to_sfd_folder,
         path_to_dataset_bank=path_to_dataset_bank,
         save_timeseries=False,
         building_for_AD=True,
@@ -188,7 +193,7 @@ def anomaly_detection(
             ztf_id=host_ztf_id_to_swap_in,
             theorized_lightcurve_df=None,
             path_to_timeseries_folder=path_to_timeseries_folder,
-            path_to_sfd_data_folder=path_to_sfd_data_folder,
+            path_to_sfd_folder=path_to_sfd_folder,
             path_to_dataset_bank=path_to_dataset_bank,
             save_timeseries=False,
             building_for_AD=True,
@@ -271,9 +276,20 @@ def check_anom_and_plot(
     """
     anom_obj_df = timeseries_df_features_only
 
-    pred_prob_anom = 100 * clf.predict_proba(anom_obj_df)
-    pred_prob_anom[:, 0] = [round(a, 1) for a in pred_prob_anom[:, 0]]
-    pred_prob_anom[:, 1] = [round(b, 1) for b in pred_prob_anom[:, 1]]
+    # Get anomaly scores from decision_function (-ve = anomalous, +ve = normal)
+    # Convert to probabilities (0-100 scale)
+    scores = clf.decision_function(anom_obj_df)
+    # Normalize scores to [0,1] - more negative means more anomalous
+    # Convert to a format compatible with the rest of the function: [[normal_prob, anomaly_prob], ...]
+    pred_prob_anom = np.zeros((len(scores), 2))
+    for i, score in enumerate(scores):
+        # Convert decision scores to probability-like values (0-100 scale)
+        # Lower scores = more anomalous
+        anomaly_prob = 100 * (1 / (1 + np.exp(score)))  # Sigmoid function to convert to [0,100]
+        normal_prob = 100 - anomaly_prob
+        pred_prob_anom[i, 0] = round(normal_prob, 1)  # normal probability
+        pred_prob_anom[i, 1] = round(anomaly_prob, 1)  # anomaly probability
+    
     num_anom_epochs = len(np.where(pred_prob_anom[:, 1] >= anom_thresh)[0])
 
     try:

src/relaiss/features.py

@@ -243,6 +243,8 @@ def build_dataset_bank(
     # if "ztf_object_id" is the index, move it to the first column
     if raw_df_bank.index.name == "ztf_object_id":
         raw_df_bank = raw_df_bank.reset_index()
+    elif 'ZTFID' in raw_df_bank.columns.values:
+        raw_df_bank['ztf_object_id'] = raw_df_bank['ZTFID']
 
     if theorized:
         raw_features = raw_lc_features

tests/test_ad.py

@@ -210,7 +210,7 @@ def test_anomaly_detection_basic(sample_preprocessed_df, tmp_path, setup_sfd_dat
         lc_features=client.lc_features,
         host_features=client.host_features,
         path_to_timeseries_folder=str(timeseries_dir),
-        path_to_sfd_data_folder=str(setup_sfd_data),
+        path_to_sfd_folder=str(setup_sfd_data),
         path_to_dataset_bank=client.bank_csv,
         path_to_models_directory=str(tmp_path),
         path_to_figure_directory=str(tmp_path / "figures"),
@@ -240,7 +240,7 @@ def test_anomaly_detection_with_host_swap(sample_preprocessed_df, tmp_path, setu
         lc_features=client.lc_features,
         host_features=client.host_features,
         path_to_timeseries_folder=str(timeseries_dir),
-        path_to_sfd_data_folder=str(setup_sfd_data),
+        path_to_sfd_folder=str(setup_sfd_data),
         path_to_dataset_bank=client.bank_csv,
         host_ztf_id_to_swap_in="ZTF19aaaaaaa",  # Swap in this host
         path_to_models_directory=str(tmp_path),
@@ -355,7 +355,7 @@ def test_anomaly_detection_basic(sample_preprocessed_df, tmp_path):
             lc_features=lc_features,
             host_features=host_features,
             path_to_timeseries_folder=str(timeseries_dir),
-            path_to_sfd_data_folder=str(sfd_dir),
+            path_to_sfd_folder=str(sfd_dir),
             path_to_dataset_bank=str(dataset_bank),
             path_to_models_directory=str(tmp_path),
             path_to_figure_directory=str(figures_dir),
@@ -536,7 +536,7 @@ def test_anomaly_detection_with_host_swap(sample_preprocessed_df, tmp_path):
             lc_features=lc_features,
             host_features=host_features,
             path_to_timeseries_folder=str(timeseries_dir),
-            path_to_sfd_data_folder=str(sfd_dir),
+            path_to_sfd_folder=str(sfd_dir),
             path_to_dataset_bank=str(dataset_bank),
             host_ztf_id_to_swap_in="ZTF19aaaaaaa",  # Swap in this host
             path_to_models_directory=str(tmp_path),

tests/test_ad_host_swap_simple.py

@@ -80,7 +80,7 @@ def test_ad_host_swap_simple(tmp_path):
             lc_features=lc_features,
             host_features=host_features,
             path_to_timeseries_folder=str(tmp_path),
-            path_to_sfd_data_folder=None,  # Ignored due to mocking
+            path_to_sfd_folder=None,  # Ignored due to mocking
             path_to_dataset_bank=None,  # Ignored due to mocking
             host_ztf_id_to_swap_in="ZTF19aaaaaaa",  # Swap in this host
             path_to_models_directory=str(model_dir),

tests/test_ad_mock.py

@@ -137,7 +137,7 @@ def test_anomaly_detection_simplified(tmp_path):
             lc_features=lc_features,
             host_features=host_features,
             path_to_timeseries_folder=str(tmp_path),
-            path_to_sfd_data_folder=None,  # Not needed with our mocks
+            path_to_sfd_folder=None,  # Not needed with our mocks
             path_to_dataset_bank=None,  # Not needed with our mocks
             path_to_models_directory=str(model_dir),
             path_to_figure_directory=str(figure_dir),

tests/test_ad_simple.py

@@ -122,7 +122,7 @@ def test_anomaly_detection_simplified(tmp_path):
             lc_features=lc_features,
             host_features=host_features,
             path_to_timeseries_folder=str(tmp_path),
-            path_to_sfd_data_folder=None,  # This will be ignored due to our mocking
+            path_to_sfd_folder=None,  # This will be ignored due to our mocking
             path_to_dataset_bank=None,  # This will be ignored due to our mocking
             path_to_models_directory=str(model_dir),
             path_to_figure_directory=str(figure_dir),