Refactor likelihoods test imports and formatting

Simplifies imports for likelihood modules by removing try/except fallbacks and directly importing modules. Cleans up formatting, improves code readability, and updates string quoting for consistency. No changes to test logic or functionality.

Author: edeno

replay_trajectory_classification/tests/unit/test_likelihoods.py

@@ -1,39 +1,17 @@
 # replay_trajectory_classification/tests/unit/test_likelihoods.py
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
-
 import numpy as np
 import pytest
 
+import replay_trajectory_classification.likelihoods.calcium_likelihood as calcium_likelihood
+import replay_trajectory_classification.likelihoods.multiunit_likelihood as multiunit_likelihood
+import replay_trajectory_classification.likelihoods.spiking_likelihood_glm as spiking_likelihood_glm
+import replay_trajectory_classification.likelihoods.spiking_likelihood_kde as spiking_likelihood_kde
+
 # Test imports for likelihood modules
 from replay_trajectory_classification.environments import Environment
 
-if TYPE_CHECKING:
-    pass
-
-# Try importing likelihood functions with graceful fallbacks
-try:
-    import replay_trajectory_classification.likelihoods.spiking_likelihood_glm as spiking_likelihood_glm
-except ImportError:
-    spiking_likelihood_glm: Any = None
-
-try:
-    import replay_trajectory_classification.likelihoods.spiking_likelihood_kde as spiking_likelihood_kde
-except ImportError:
-    spiking_likelihood_kde: Any = None
-
-try:
-    import replay_trajectory_classification.likelihoods.multiunit_likelihood as multiunit_likelihood
-except ImportError:
-    multiunit_likelihood: Any = None
-
-try:
-    import replay_trajectory_classification.likelihoods.calcium_likelihood as calcium_likelihood
-except ImportError:
-    calcium_likelihood: Any = None
-
-
 # ---------------------- Helpers ----------------------
 
 
@@ -68,7 +46,7 @@ def make_multiunit_data(n_electrodes=3, n_features=4, n_time=20):
 
         for t in range(n_time):
             if n_spikes_per_time[t] < max_spikes:
-                no_spike_indicator[t, n_spikes_per_time[t]:] = True
+                no_spike_indicator[t, n_spikes_per_time[t] :] = True
 
         data[f"electrode_{elec_id:02d}"] = {
             "marks": marks,
@@ -82,28 +60,25 @@ def make_multiunit_data(n_electrodes=3, n_features=4, n_time=20):
 
 
 @pytest.mark.skipif(
-    spiking_likelihood_glm is None,
-    reason="spiking_likelihood_glm module not available"
+    spiking_likelihood_glm is None, reason="spiking_likelihood_glm module not available"
 )
 def test_spiking_likelihood_glm_fit_exists():
     """Test that GLM likelihood fit function exists and is callable."""
-    assert hasattr(spiking_likelihood_glm, 'estimate_place_fields')
+    assert hasattr(spiking_likelihood_glm, "estimate_place_fields")
     assert callable(spiking_likelihood_glm.estimate_place_fields)
 
 
 @pytest.mark.skipif(
-    spiking_likelihood_glm is None,
-    reason="spiking_likelihood_glm module not available"
+    spiking_likelihood_glm is None, reason="spiking_likelihood_glm module not available"
 )
 def test_spiking_likelihood_glm_estimate_exists():
     """Test that GLM likelihood estimate function exists and is callable."""
-    assert hasattr(spiking_likelihood_glm, 'estimate_spiking_likelihood')
+    assert hasattr(spiking_likelihood_glm, "estimate_spiking_likelihood")
     assert callable(spiking_likelihood_glm.estimate_spiking_likelihood)
 
 
 @pytest.mark.skipif(
-    spiking_likelihood_glm is None,
-    reason="spiking_likelihood_glm module not available"
+    spiking_likelihood_glm is None, reason="spiking_likelihood_glm module not available"
 )
 def test_spiking_likelihood_glm_basic_functionality():
     """Test basic GLM likelihood functionality with synthetic data."""
@@ -118,10 +93,10 @@ def test_spiking_likelihood_glm_basic_functionality():
             position=position,
             spikes=spikes,
             place_bin_centers=environment.place_bin_centers_,
-            place_bin_edges=getattr(environment, 'place_bin_edges_', None),
-            edges=getattr(environment, 'edges_', None),
+            place_bin_edges=getattr(environment, "place_bin_edges_", None),
+            edges=getattr(environment, "edges_", None),
             is_track_interior=environment.is_track_interior_,
-            is_track_boundary=getattr(environment, 'is_track_boundary_', None)
+            is_track_boundary=getattr(environment, "is_track_boundary_", None),
         )
 
         # Basic checks on results
@@ -139,28 +114,25 @@ def test_spiking_likelihood_glm_basic_functionality():
 
 
 @pytest.mark.skipif(
-    spiking_likelihood_kde is None,
-    reason="spiking_likelihood_kde module not available"
+    spiking_likelihood_kde is None, reason="spiking_likelihood_kde module not available"
 )
 def test_spiking_likelihood_kde_fit_exists():
     """Test that KDE likelihood fit function exists and is callable."""
-    assert hasattr(spiking_likelihood_kde, 'estimate_place_fields_kde')
+    assert hasattr(spiking_likelihood_kde, "estimate_place_fields_kde")
     assert callable(spiking_likelihood_kde.estimate_place_fields_kde)
 
 
 @pytest.mark.skipif(
-    spiking_likelihood_kde is None,
-    reason="spiking_likelihood_kde module not available"
+    spiking_likelihood_kde is None, reason="spiking_likelihood_kde module not available"
 )
 def test_spiking_likelihood_kde_estimate_exists():
     """Test that KDE likelihood estimate function exists and is callable."""
-    assert hasattr(spiking_likelihood_kde, 'estimate_spiking_likelihood_kde')
+    assert hasattr(spiking_likelihood_kde, "estimate_spiking_likelihood_kde")
     assert callable(spiking_likelihood_kde.estimate_spiking_likelihood_kde)
 
 
 @pytest.mark.skipif(
-    spiking_likelihood_kde is None,
-    reason="spiking_likelihood_kde module not available"
+    spiking_likelihood_kde is None, reason="spiking_likelihood_kde module not available"
 )
 def test_spiking_likelihood_kde_basic_functionality():
     """Test basic KDE likelihood functionality with synthetic data."""
@@ -171,12 +143,12 @@ def test_spiking_likelihood_kde_basic_functionality():
         position = make_simple_position(n_time=50, n_dims=1)
 
         # Test fit function - try common function names
-        if hasattr(spiking_likelihood_kde, 'estimate_place_fields_kde'):
+        if hasattr(spiking_likelihood_kde, "estimate_place_fields_kde"):
             results = spiking_likelihood_kde.estimate_place_fields_kde(
                 position=position,
                 spikes=spikes,
                 place_bin_centers=environment.place_bin_centers_,
-                is_track_interior=environment.is_track_interior_
+                is_track_interior=environment.is_track_interior_,
             )
         else:
             # Try alternative function names
@@ -195,28 +167,25 @@ def test_spiking_likelihood_kde_basic_functionality():
 
 
 @pytest.mark.skipif(
-    multiunit_likelihood is None,
-    reason="multiunit_likelihood module not available"
+    multiunit_likelihood is None, reason="multiunit_likelihood module not available"
 )
 def test_multiunit_likelihood_fit_exists():
     """Test that multiunit likelihood fit function exists and is callable."""
-    assert hasattr(multiunit_likelihood, 'fit_multiunit_likelihood')
+    assert hasattr(multiunit_likelihood, "fit_multiunit_likelihood")
     assert callable(multiunit_likelihood.fit_multiunit_likelihood)
 
 
 @pytest.mark.skipif(
-    multiunit_likelihood is None,
-    reason="multiunit_likelihood module not available"
+    multiunit_likelihood is None, reason="multiunit_likelihood module not available"
 )
 def test_multiunit_likelihood_estimate_exists():
     """Test that multiunit likelihood estimate function exists and is callable."""
-    assert hasattr(multiunit_likelihood, 'estimate_multiunit_likelihood')
+    assert hasattr(multiunit_likelihood, "estimate_multiunit_likelihood")
     assert callable(multiunit_likelihood.estimate_multiunit_likelihood)
 
 
 @pytest.mark.skipif(
-    multiunit_likelihood is None,
-    reason="multiunit_likelihood module not available"
+    multiunit_likelihood is None, reason="multiunit_likelihood module not available"
 )
 def test_multiunit_likelihood_basic_functionality():
     """Test basic multiunit likelihood functionality with synthetic data."""
@@ -228,12 +197,16 @@ def test_multiunit_likelihood_basic_functionality():
 
         # Convert multiunit dict to 3D array format
         n_electrodes = len(multiunit_data)
-        n_features = list(multiunit_data.values())[0]["marks"].shape[2]
+        list(multiunit_data.values())[0]["marks"].shape[2]
         max_marks = list(multiunit_data.values())[0]["marks"].shape[1]
 
         multiunit_3d = np.full((50, max_marks, n_electrodes), np.nan)
-        for elec_idx, (electrode_id, electrode_data) in enumerate(multiunit_data.items()):
-            multiunit_3d[:, :, elec_idx] = electrode_data["marks"][:, :, 0]  # Use first feature
+        for elec_idx, (electrode_id, electrode_data) in enumerate(
+            multiunit_data.items()
+        ):
+            multiunit_3d[:, :, elec_idx] = electrode_data["marks"][
+                :, :, 0
+            ]  # Use first feature
             no_spike_mask = electrode_data["no_spike_indicator"]
             multiunit_3d[no_spike_mask, elec_idx] = np.nan
 
@@ -243,10 +216,10 @@ def test_multiunit_likelihood_basic_functionality():
             multiunits=multiunit_3d,
             place_bin_centers=environment.place_bin_centers_,
             is_track_interior=environment.is_track_interior_,
-            is_track_boundary=getattr(environment, 'is_track_boundary_', None),
-            edges=getattr(environment, 'edges_', None),
+            is_track_boundary=getattr(environment, "is_track_boundary_", None),
+            edges=getattr(environment, "edges_", None),
             mark_std=24.0,
-            position_std=6.0
+            position_std=6.0,
         )
 
         # Basic checks on results
@@ -262,13 +235,12 @@ def test_multiunit_likelihood_basic_functionality():
 
 
 @pytest.mark.skipif(
-    calcium_likelihood is None,
-    reason="calcium_likelihood module not available"
+    calcium_likelihood is None, reason="calcium_likelihood module not available"
 )
 def test_calcium_likelihood_fit_exists():
     """Test that calcium likelihood fit function exists and is callable."""
     # Look for common function names in calcium likelihood module
-    fit_funcs = [name for name in dir(calcium_likelihood) if 'fit' in name.lower()]
+    fit_funcs = [name for name in dir(calcium_likelihood) if "fit" in name.lower()]
     assert len(fit_funcs) > 0, "No fit functions found in calcium_likelihood module"
 
     for func_name in fit_funcs:
@@ -280,14 +252,17 @@ def test_calcium_likelihood_fit_exists():
 
 
 @pytest.mark.skipif(
-    calcium_likelihood is None,
-    reason="calcium_likelihood module not available"
+    calcium_likelihood is None, reason="calcium_likelihood module not available"
 )
 def test_calcium_likelihood_estimate_exists():
     """Test that calcium likelihood estimate function exists and is callable."""
     # Look for common function names in calcium likelihood module
-    estimate_funcs = [name for name in dir(calcium_likelihood) if 'estimate' in name.lower()]
-    assert len(estimate_funcs) > 0, "No estimate functions found in calcium_likelihood module"
+    estimate_funcs = [
+        name for name in dir(calcium_likelihood) if "estimate" in name.lower()
+    ]
+    assert (
+        len(estimate_funcs) > 0
+    ), "No estimate functions found in calcium_likelihood module"
 
 
 # ---------------------- General Likelihood Interface Tests ----------------------
@@ -298,31 +273,39 @@ def test_likelihood_modules_importable():
     # This test verifies basic import functionality
     try:
         from replay_trajectory_classification import likelihoods
+
         assert likelihoods is not None
     except ImportError:
         pytest.fail("Could not import likelihoods subpackage")
 
 
-@pytest.mark.parametrize("module_name", [
-    "spiking_likelihood_glm",
-    "spiking_likelihood_kde",
-    "multiunit_likelihood",
-    "calcium_likelihood"
-])
+@pytest.mark.parametrize(
+    "module_name",
+    [
+        "spiking_likelihood_glm",
+        "spiking_likelihood_kde",
+        "multiunit_likelihood",
+        "calcium_likelihood",
+    ],
+)
 def test_likelihood_module_structure(module_name):
     """Test that likelihood modules have expected structure."""
     try:
         module = __import__(
             f"replay_trajectory_classification.likelihoods.{module_name}",
-            fromlist=[module_name]
+            fromlist=[module_name],
         )
 
         # Should have at least some public functions
-        public_attrs = [name for name in dir(module) if not name.startswith('_')]
+        public_attrs = [name for name in dir(module) if not name.startswith("_")]
         assert len(public_attrs) > 0, f"No public attributes in {module_name}"
 
         # Should have at least one callable
-        callables = [getattr(module, name) for name in public_attrs if callable(getattr(module, name))]
+        callables = [
+            getattr(module, name)
+            for name in public_attrs
+            if callable(getattr(module, name))
+        ]
         assert len(callables) > 0, f"No callable functions in {module_name}"
 
     except ImportError:
@@ -343,8 +326,8 @@ def test_likelihood_functions_handle_edge_cases():
 
     # Test each available likelihood module
     modules_to_test = [
-        (spiking_likelihood_glm, 'fit_spiking_likelihood_glm'),
-        (spiking_likelihood_kde, 'fit_spiking_likelihood_kde'),
+        (spiking_likelihood_glm, "fit_spiking_likelihood_glm"),
+        (spiking_likelihood_kde, "fit_spiking_likelihood_kde"),
     ]
 
     for module, func_name in modules_to_test:
@@ -355,7 +338,7 @@ def test_likelihood_functions_handle_edge_cases():
                 result = func(empty_position, empty_spikes, environment)
                 # If it returns something, it should be reasonable
                 if result is not None:
-                    assert not (hasattr(result, 'shape') and result.shape == (0,))
+                    assert not (hasattr(result, "shape") and result.shape == (0,))
             except (ValueError, RuntimeError, ZeroDivisionError):
                 # These are acceptable exceptions for edge cases
                 continue
@@ -370,9 +353,21 @@ def test_likelihood_functions_handle_edge_cases():
 def test_fit_estimate_consistency():
     """Test that fit and estimate functions have consistent interfaces."""
     modules_to_test = [
-        (spiking_likelihood_glm, 'estimate_place_fields', 'estimate_spiking_likelihood'),
-        (spiking_likelihood_kde, 'estimate_place_fields_kde', 'estimate_spiking_likelihood_kde'),
-        (multiunit_likelihood, 'fit_multiunit_likelihood', 'estimate_multiunit_likelihood'),
+        (
+            spiking_likelihood_glm,
+            "estimate_place_fields",
+            "estimate_spiking_likelihood",
+        ),
+        (
+            spiking_likelihood_kde,
+            "estimate_place_fields_kde",
+            "estimate_spiking_likelihood_kde",
+        ),
+        (
+            multiunit_likelihood,
+            "fit_multiunit_likelihood",
+            "estimate_multiunit_likelihood",
+        ),
     ]
 
     for module, fit_func_name, estimate_func_name in modules_to_test:
@@ -383,10 +378,12 @@ def test_fit_estimate_consistency():
 
             if has_fit or has_estimate:
                 # If one exists, both should exist for consistency
-                assert has_fit and has_estimate, f"Module {module.__name__} should have both fit and estimate functions"
+                assert (
+                    has_fit and has_estimate
+                ), f"Module {module.__name__} should have both fit and estimate functions"
 
                 # Both should be callable
                 fit_func = getattr(module, fit_func_name)
                 estimate_func = getattr(module, estimate_func_name)
                 assert callable(fit_func)
-                assert callable(estimate_func)
+                assert callable(estimate_func)