Refactor imports and clean up print statements

Standardized and cleaned up print statements across scripts for consistency. Refactored and reordered imports in several scripts and test files, removing unused imports and fixing import order. Minor code cleanups include removing unused variables, correcting loop variable names, and fixing assignment statements in tests.

Author: edeno

scripts/compare_kde_implementations.py

@@ -56,7 +56,7 @@ def main():
     n_runs = 10
     dimensions = [2, 4, 6, 8, 10]
 
-    print(f"Configuration:")
+    print("Configuration:")
     print(f"  Evaluation points: {n_eval}")
     print(f"  Training samples: {n_samples}")
     print(f"  Runs per test: {n_runs}")

scripts/compare_memory_usage.py

@@ -128,7 +128,7 @@ def analyze_array_allocations(n_eval, n_samples, n_dims):
     print(f"  Per-dimension output: {n_samples} × {n_eval} × 8 bytes = {output_size/1024:.2f} KB")
     vmap_intermediates = n_dims * n_samples * n_eval * 8  # All dimensions at once
     print(f"  All intermediate arrays: {n_dims} × {output_size/1024:.2f} KB = {vmap_intermediates/1024:.2f} KB")
-    print(f"  XLA fusion may reduce this through operation fusion")
+    print("  XLA fusion may reduce this through operation fusion")
     optimized_peak = vmap_intermediates + output_size
     print(f"  **Estimated peak (worst case): {optimized_peak/1024:.2f} KB**")
     print(f"  **Estimated peak (with fusion): {(vmap_intermediates*0.5 + output_size)/1024:.2f} KB**")
@@ -141,7 +141,7 @@ def analyze_array_allocations(n_eval, n_samples, n_dims):
 
     if ratio > 1.5:
         print(f"⚠️  Optimized version may use {ratio:.1f}x more memory in worst case")
-        print(f"    But XLA fusion likely reduces this significantly")
+        print("    But XLA fusion likely reduces this significantly")
     elif ratio > 0.7:
         print(f"✓  Similar memory footprint ({ratio:.2f}x)")
     else:
@@ -222,7 +222,7 @@ def main():
                 elif mem_ratio < 0.8:
                     print(f"✓  Optimized uses {1/mem_ratio:.1f}x less memory")
                 else:
-                    print(f"✓  Similar memory usage")
+                    print("✓  Similar memory usage")
 
             # Estimate device memory
             print("\nEstimated JAX Device Memory:")
@@ -293,7 +293,7 @@ def main():
             print("   for all dimensions simultaneously (parallel execution),")
             print("   while the original processes one dimension at a time.")
             print()
-            print("   **Trade-off**: Speed (10x faster) vs Memory (~{:.1f}x more)".format(avg_ratio))
+            print(f"   **Trade-off**: Speed (10x faster) vs Memory (~{avg_ratio:.1f}x more)")
         elif avg_ratio > 1.2:
             print(f"⚠️  Optimized version uses slightly more memory ({avg_ratio:.2f}x)")
         else:
@@ -310,7 +310,7 @@ def main():
         print("**Memory-Speed Trade-off**")
         print()
         print("The optimized version (clusterless_kde_log.py):")
-        print(f"  ✅ Speed: 10.8x faster")
+        print("  ✅ Speed: 10.8x faster")
         print(f"  ⚠️  Memory: ~{avg_ratio:.1f}x more usage")
         print()
         print("**Use optimized version when:**")

scripts/investigate_extreme_values.py

@@ -69,7 +69,7 @@ def analyze_extreme_features():
 
     feature_distances = np.array(feature_distances)
 
-    print(f"Feature space distances (standardized by std):")
+    print("Feature space distances (standardized by std):")
     print(f"  Min: {feature_distances.min():.2f} std")
     print(f"  Max: {feature_distances.max():.2f} std")
     print(f"  Mean: {feature_distances.mean():.2f} std")
@@ -91,7 +91,7 @@ def analyze_extreme_features():
 
     # Compute position distances
     position_distance = kde_distance(interior_bins, enc_positions, position_std)
-    log_position_distance = log_kde_distance(interior_bins, enc_positions, position_std)
+    log_kde_distance(interior_bins, enc_positions, position_std)
 
     print("-" * 80)
     print("ORIGINAL IMPLEMENTATION (clusterless_kde.py)")
@@ -118,14 +118,14 @@ def analyze_extreme_features():
     n_inf = np.sum(np.isinf(ll_original))
     n_nan = np.sum(np.isnan(ll_original))
 
-    print(f"Value distribution:")
+    print("Value distribution:")
     print(f"  Finite values: {n_finite}/{ll_original.size} ({100*n_finite/ll_original.size:.1f}%)")
     print(f"  -Inf values: {n_inf}/{ll_original.size} ({100*n_inf/ll_original.size:.1f}%)")
     print(f"  NaN values: {n_nan}/{ll_original.size} ({100*n_nan/ll_original.size:.1f}%)")
     print()
 
     if n_finite > 0:
-        print(f"Finite value statistics:")
+        print("Finite value statistics:")
         print(f"  Min: {ll_original[finite_mask].min():.4f}")
         print(f"  Max: {ll_original[finite_mask].max():.4f}")
         print(f"  Mean: {ll_original[finite_mask].mean():.4f}")
@@ -167,7 +167,7 @@ def analyze_extreme_features():
     n_inf_log = np.sum(np.isinf(ll_log_no_gemm))
     n_nan_log = np.sum(np.isnan(ll_log_no_gemm))
 
-    print(f"Value distribution:")
+    print("Value distribution:")
     print(f"  Finite values: {n_finite_log}/{ll_log_no_gemm.size} ({100*n_finite_log/ll_log_no_gemm.size:.1f}%)")
     print(f"  -Inf values: {n_inf_log}/{ll_log_no_gemm.size} ({100*n_inf_log/ll_log_no_gemm.size:.1f}%)")
     print(f"  NaN values: {n_nan_log}/{ll_log_no_gemm.size} ({100*n_nan_log/ll_log_no_gemm.size:.1f}%)")
@@ -211,14 +211,14 @@ def analyze_extreme_features():
     n_inf_gemm = np.sum(np.isinf(ll_log_gemm))
     n_nan_gemm = np.sum(np.isnan(ll_log_gemm))
 
-    print(f"Value distribution:")
+    print("Value distribution:")
     print(f"  Finite values: {n_finite_gemm}/{ll_log_gemm.size} ({100*n_finite_gemm/ll_log_gemm.size:.1f}%)")
     print(f"  -Inf values: {n_inf_gemm}/{ll_log_gemm.size} ({100*n_inf_gemm/ll_log_gemm.size:.1f}%)")
     print(f"  NaN values: {n_nan_gemm}/{ll_log_gemm.size} ({100*n_nan_gemm/ll_log_gemm.size:.1f}%)")
     print()
 
     if n_finite_gemm > 0:
-        print(f"Finite value statistics:")
+        print("Finite value statistics:")
         print(f"  Min: {ll_log_gemm[finite_mask_gemm].min():.4f}")
         print(f"  Max: {ll_log_gemm[finite_mask_gemm].max():.4f}")
         print(f"  Mean: {ll_log_gemm[finite_mask_gemm].mean():.4f}")
@@ -254,7 +254,7 @@ def analyze_extreme_features():
 
     if n_finite_gemm > n_finite:
         improvement = n_finite_gemm - n_finite
-        print(f"✅ GEMM optimization improves numerical stability:")
+        print("✅ GEMM optimization improves numerical stability:")
         print(f"   {improvement} more finite values ({100*improvement/ll_original.size:.1f}% of total)")
         print()
         print("   The GEMM approach computes in log-space throughout,")
@@ -289,7 +289,7 @@ def analyze_extreme_features():
     if n_finite_gemm > n_finite * 1.1:  # At least 10% improvement
         print("**Use GEMM optimization for extreme features**")
         print()
-        print(f"The GEMM approach (use_gemm=True) significantly improves")
+        print("The GEMM approach (use_gemm=True) significantly improves")
         print(f"numerical stability, preserving {100*n_finite_gemm/ll_log_gemm.size:.1f}% finite values")
         print(f"vs {100*n_finite/ll_original.size:.1f}% for the original.")
         print()
@@ -300,7 +300,7 @@ def analyze_extreme_features():
     else:
         print("**GEMM optimization does not significantly help with extreme features**")
         print()
-        print(f"Both approaches produce similar amounts of underflow")
+        print("Both approaches produce similar amounts of underflow")
         print(f"({100*n_finite/ll_original.size:.1f}% vs {100*n_finite_gemm/ll_log_gemm.size:.1f}% finite).")
         print()
         print("With such extreme feature distances, underflow is unavoidable.")

scripts/profile_feature_dimensions.py

@@ -6,20 +6,22 @@
 
 import sys
 import time
+
 import numpy as np
-import jax.numpy as jnp
 
 sys.path.insert(0, "src")
 
+from non_local_detector.environment import Environment
 from non_local_detector.likelihoods.clusterless_kde import (
     fit_clusterless_kde_encoding_model,
     predict_clusterless_kde_log_likelihood,
 )
 from non_local_detector.likelihoods.clusterless_kde_log import (
     fit_clusterless_kde_encoding_model as fit_log,
+)
+from non_local_detector.likelihoods.clusterless_kde_log import (
     predict_clusterless_kde_log_likelihood as predict_log,
 )
-from non_local_detector.environment import Environment
 
 
 def create_test_data(n_features, n_encoding=200, n_decoding=100, n_positions=500):
@@ -189,7 +191,7 @@ def profile_dimension(n_features):
     speedup_gemm = mean_ref / mean_gemm
     speedup_no_gemm = mean_ref / mean_no_gemm
 
-    print(f"\nSpeedup vs reference:")
+    print("\nSpeedup vs reference:")
     print(f"  GEMM (vmap):     {speedup_gemm:.2f}x")
     print(f"  No GEMM (log):   {speedup_no_gemm:.2f}x")
 

scripts/profile_feature_dimensions_simple.py

@@ -6,8 +6,9 @@
 
 import sys
 import time
-import numpy as np
+
 import jax.numpy as jnp
+import numpy as np
 
 sys.path.insert(0, "src")
 
@@ -103,7 +104,7 @@ def profile_dimension(n_features, n_enc=200, n_dec=100, n_pos=500):
     speedup_no_gemm = mean_ref / mean_no_gemm
     gemm_vs_no_gemm = mean_no_gemm / mean_gemm
 
-    print(f"\nSpeedup vs reference:")
+    print("\nSpeedup vs reference:")
     print(f"  GEMM (vmap):     {speedup_gemm:.2f}x")
     print(f"  No GEMM:         {speedup_no_gemm:.2f}x")
     print(f"\nGEMM vs No-GEMM: {gemm_vs_no_gemm:.2f}x")

scripts/profile_log_kde_optimization.py

@@ -79,7 +79,7 @@ def main():
     n_runs = 10
     dimensions = [2, 4, 6, 8, 10]
 
-    print(f"Configuration:")
+    print("Configuration:")
     print(f"  Evaluation points: {n_eval}")
     print(f"  Training samples: {n_samples}")
     print(f"  Runs per test: {n_runs}")

scripts/profile_optimized_kde.py

@@ -2,15 +2,16 @@
 
 import sys
 import time
-import numpy as np
+
 import jax.numpy as jnp
+import numpy as np
 
 sys.path.insert(0, "src")
 
 from non_local_detector.likelihoods.clusterless_kde import (
     estimate_log_joint_mark_intensity,
-    estimate_log_joint_mark_intensity_vectorized,
     estimate_log_joint_mark_intensity_logspace,
+    estimate_log_joint_mark_intensity_vectorized,
 )
 
 
@@ -167,13 +168,13 @@ def main():
     print(f"  Average speedup: {avg_speedup:.2f}x")
 
     if avg_speedup >= 2.0:
-        print(f"  ✓ Excellent! Achieved target of 2-4x speedup")
+        print("  ✓ Excellent! Achieved target of 2-4x speedup")
     elif avg_speedup >= 1.5:
-        print(f"  ✓ Good! Significant performance improvement")
+        print("  ✓ Good! Significant performance improvement")
     elif avg_speedup >= 1.2:
-        print(f"  ~ Moderate improvement")
+        print("  ~ Moderate improvement")
     else:
-        print(f"  ✗ Minimal improvement")
+        print("  ✗ Minimal improvement")
 
     print("\nLog-space + Vectorized + JIT Optimization:")
     log_speedups = [r[2]["logspace"][2] for r in all_results]
@@ -184,13 +185,13 @@ def main():
     print(f"  Average speedup: {avg_speedup:.2f}x")
 
     if avg_speedup >= 2.0:
-        print(f"  ✓ Excellent! Achieved target of 2-4x speedup")
+        print("  ✓ Excellent! Achieved target of 2-4x speedup")
     elif avg_speedup >= 1.5:
-        print(f"  ✓ Good! Significant performance improvement")
+        print("  ✓ Good! Significant performance improvement")
     elif avg_speedup >= 1.2:
-        print(f"  ~ Moderate improvement")
+        print("  ~ Moderate improvement")
     else:
-        print(f"  ✗ Minimal improvement")
+        print("  ✗ Minimal improvement")
 
     # Comparison: Vectorized vs Log-space
     print("\nLog-space vs Vectorized:")
@@ -212,22 +213,22 @@ def main():
     print("\nBest implementation:")
     if log_avg > vec_avg * 1.1:
         print(f"  → Log-space + Vectorized + JIT ({log_avg:.2f}x average speedup)")
-        print(f"    Use: estimate_log_joint_mark_intensity_logspace()")
+        print("    Use: estimate_log_joint_mark_intensity_logspace()")
     elif vec_avg > 1.5:
         print(f"  → Vectorized + JIT ({vec_avg:.2f}x average speedup)")
-        print(f"    Use: estimate_log_joint_mark_intensity_vectorized()")
+        print("    Use: estimate_log_joint_mark_intensity_vectorized()")
     else:
-        print(f"  → Original implementation (optimizations not beneficial)")
-        print(f"    Use: estimate_log_joint_mark_intensity()")
+        print("  → Original implementation (optimizations not beneficial)")
+        print("    Use: estimate_log_joint_mark_intensity()")
 
     print("\nFor production use:")
     if max(vec_avg, log_avg) >= 2.0:
-        print(f"  ✓ Optimization successful - recommend deploying optimized version")
-        print(f"  ✓ Numerical equivalence verified (max diff < 1e-6)")
+        print("  ✓ Optimization successful - recommend deploying optimized version")
+        print("  ✓ Numerical equivalence verified (max diff < 1e-6)")
         print(f"  ✓ Average speedup: {max(vec_avg, log_avg):.2f}x")
     else:
-        print(f"  ~ Optimization provides moderate benefit")
-        print(f"  ~ Consider for performance-critical applications only")
+        print("  ~ Optimization provides moderate benefit")
+        print("  ~ Consider for performance-critical applications only")
 
     print("\n" + "="*70)
 

scripts/test_optimized_kde.py

@@ -1,18 +1,19 @@
 """Test numerical equivalence of optimized KDE implementations."""
 
 import sys
-import numpy as np
+
 import jax.numpy as jnp
+import numpy as np
 
 sys.path.insert(0, "src")
 
 from non_local_detector.likelihoods.clusterless_kde import (
+    estimate_log_joint_mark_intensity,
+    estimate_log_joint_mark_intensity_logspace,
+    estimate_log_joint_mark_intensity_vectorized,
     kde_distance,
     kde_distance_vectorized,
     log_kde_distance,
-    estimate_log_joint_mark_intensity,
-    estimate_log_joint_mark_intensity_vectorized,
-    estimate_log_joint_mark_intensity_logspace,
 )
 
 
@@ -131,12 +132,12 @@ def test_estimate_functions_equivalence():
         # Check with more appropriate tolerance for numerical differences
         is_close = jnp.allclose(result_original, result_vectorized, rtol=1e-5, atol=1e-6)
         if not is_close:
-            print(f"    WARNING: Differences exceed tolerance")
+            print("    WARNING: Differences exceed tolerance")
             print(f"    Max absolute diff: {max_diff_vec}")
             print(f"    Max relative diff: {rel_diff_vec}")
             # Check if it's just due to float precision
             if max_diff_vec < 1e-5 and rel_diff_vec < 1e-4:
-                print(f"    Differences are within acceptable float32 precision, continuing...")
+                print("    Differences are within acceptable float32 precision, continuing...")
             else:
                 raise AssertionError(f"Vectorized version not equivalent for {n_features}D")
 
@@ -153,12 +154,12 @@ def test_estimate_functions_equivalence():
         # Check with more appropriate tolerance
         is_close = jnp.allclose(result_original, result_logspace, rtol=1e-5, atol=1e-6)
         if not is_close:
-            print(f"    WARNING: Differences exceed tolerance")
+            print("    WARNING: Differences exceed tolerance")
             print(f"    Max absolute diff: {max_diff_log}")
             print(f"    Max relative diff: {rel_diff_log}")
             # Check if it's just due to float precision
             if max_diff_log < 1e-5 and rel_diff_log < 1e-4:
-                print(f"    Differences are within acceptable float32 precision, continuing...")
+                print("    Differences are within acceptable float32 precision, continuing...")
             else:
                 raise AssertionError(f"Log-space version not equivalent for {n_features}D")
 

src/non_local_detector/likelihoods/clusterless_kde_log.py

@@ -9,7 +9,6 @@
     EPS,
     LOG_EPS,
     KDEModel,
-    block_kde,
     block_log_kde,
     gaussian_pdf,
     get_position_at_time,

src/non_local_detector/tests/integration/test_clusterless_kde_parity.py

@@ -120,7 +120,7 @@ def test_estimate_intensity_moderate_features():
     np.random.seed(42)
     dec_features = jnp.array(np.random.randn(n_dec_spikes, n_features) * 10 + 50)
     enc_features = jnp.array(np.random.randn(n_enc_spikes, n_features) * 10 + 50)
-    enc_weights = jnp.ones(n_enc_spikes)
+    jnp.ones(n_enc_spikes)
     waveform_stds = jnp.array([5.0] * n_features)
     occupancy = jnp.ones(n_pos_bins) * 0.1
     mean_rate = 5.0
@@ -185,7 +185,7 @@ def test_estimate_intensity_extreme_features():
     np.random.seed(42)
     dec_features = jnp.array(np.random.randn(n_dec_spikes, n_features) * 50 + 100)
     enc_features = jnp.array(np.random.randn(n_enc_spikes, n_features) * 50 + 200)
-    enc_weights = jnp.ones(n_enc_spikes)
+    jnp.ones(n_enc_spikes)
     waveform_stds = jnp.array([10.0] * n_features)
     occupancy = jnp.ones(n_pos_bins) * 0.1
     mean_rate = 2.0