Fix DataJoint query errors with NaN values in probe geometry fields (#1346)

* Initial plan

* Implement NaN value replacement fix for probe geometry fields

Co-authored-by: samuelbray32 <[email protected]>

* Complete NaN value replacement fix implementation

Co-authored-by: samuelbray32 <[email protected]>

* Move _replace_nan_with_default to utils and simplify tests

Co-authored-by: samuelbray32 <[email protected]>

* lint

* import test function from fixtures

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: samuelbray32 <[email protected]>
Co-authored-by: samuelbray32 <[email protected]>

Author: Copilot

src/spyglass/common/common_device.py

@@ -4,7 +4,10 @@
 from spyglass.common.errors import PopulateException
 from spyglass.settings import test_mode
 from spyglass.utils import SpyglassMixin, logger
-from spyglass.utils.dj_helper_fn import accept_divergence
+from spyglass.utils.dj_helper_fn import (
+    accept_divergence,
+    _replace_nan_with_default,
+)
 from spyglass.utils.nwb_helper_fn import get_nwb_file
 
 schema = dj.schema("common_device")
@@ -563,15 +566,17 @@ def __read_ndx_probe_data(
             for electrode in shank.shanks_electrodes.values():
                 # the next line will need to be fixed if we have different sized
                 # contacts on a shank
-                elect_dict[electrode.name] = {
-                    "probe_id": new_probe_dict["probe_id"],
-                    "probe_shank": shank_dict[shank.name]["probe_shank"],
-                    "contact_size": nwb_probe_obj.contact_size,
-                    "probe_electrode": int(electrode.name),
-                    "rel_x": electrode.rel_x,
-                    "rel_y": electrode.rel_y,
-                    "rel_z": electrode.rel_z,
-                }
+                elect_dict[electrode.name] = _replace_nan_with_default(
+                    {
+                        "probe_id": new_probe_dict["probe_id"],
+                        "probe_shank": shank_dict[shank.name]["probe_shank"],
+                        "contact_size": nwb_probe_obj.contact_size,
+                        "probe_electrode": int(electrode.name),
+                        "rel_x": electrode.rel_x,
+                        "rel_y": electrode.rel_y,
+                        "rel_z": electrode.rel_z,
+                    }
+                )
 
     @classmethod
     def _read_config_probe_data(
@@ -597,15 +602,17 @@ def _read_config_probe_data(
             "Electrode", []
         )
         for i, e in enumerate(elect_dict_list):
-            elect_dict[str(i)] = {
-                "probe_id": probe_id,
-                "probe_shank": e["probe_shank"],
-                "probe_electrode": e["probe_electrode"],
-                "contact_size": e.get("contact_size"),
-                "rel_x": e.get("rel_x"),
-                "rel_y": e.get("rel_y"),
-                "rel_z": e.get("rel_z"),
-            }
+            elect_dict[str(i)] = _replace_nan_with_default(
+                {
+                    "probe_id": probe_id,
+                    "probe_shank": e["probe_shank"],
+                    "probe_electrode": e["probe_electrode"],
+                    "contact_size": e.get("contact_size"),
+                    "rel_x": e.get("rel_x"),
+                    "rel_y": e.get("rel_y"),
+                    "rel_z": e.get("rel_z"),
+                }
+            )
 
         # make the probe type if not in database
         new_probe_type_dict.update(
@@ -780,11 +787,15 @@ def create_from_nwbfile(
                 "probe_electrode": elec_index,
             }
 
-            for dim in ["rel_x", "rel_y", "rel_z"]:
+            for dim in ["rel_x", "rel_y", "rel_z", "contact_size"]:
                 if dim in nwbfile.electrodes[elec_index]:
-                    elect_dict[elec_index][dim] = nwbfile.electrodes[
-                        elec_index, dim
-                    ]
+                    value = nwbfile.electrodes[elec_index, dim]
+                    elect_dict[elec_index][dim] = value
+
+            # Apply NaN replacement to the entire electrode dictionary
+            elect_dict[elec_index] = _replace_nan_with_default(
+                elect_dict[elec_index]
+            )
 
         if not device_found:
             logger.warning(

src/spyglass/utils/dj_helper_fn.py

@@ -1,6 +1,7 @@
 """Helper functions for manipulating information from DataJoint fetch calls."""
 
 import inspect
+import math
 import multiprocessing.pool
 import os
 import re
@@ -701,3 +702,28 @@ def accept_divergence(key, new_value, existing_value, test_mode=False):
         + f"'{new_value}' ?\n"
     )
     return str_to_bool(response)
+
+
+def _replace_nan_with_default(data_dict, default_value=-1.0):
+    """
+    Replace NaN values in a dictionary with a default value.
+
+    This is necessary because DataJoint cannot properly format queries
+    with NaN values, causing errors during probe insertion/validation.
+
+    Args:
+        data_dict: Dictionary that may contain NaN values
+        default_value: Value to replace NaN with (default: -1.0)
+
+    Returns:
+        Dictionary with NaN values replaced
+    """
+    if not isinstance(data_dict, dict):
+        return data_dict
+
+    result = data_dict.copy()
+    for key, value in result.items():
+        if isinstance(value, float) and math.isnan(value):
+            result[key] = default_value
+
+    return result

tests/common/test_device.py

@@ -45,3 +45,30 @@ def test_create_probe(common, mini_devices, mini_path, mini_copy_name):
     assert array_equal(
         before, after
     ), "Probe create_from_nwbfile had unexpected effect"
+
+
+def test_replace_nan_with_default(utils):
+    """Test that NaN values in probe geometry fields are properly replaced with -1.0."""
+    # Test with NaN values (similar to the issue case)
+    test_data = {
+        "probe_id": "nTrode32_probe description",
+        "probe_shank": 0,
+        "contact_size": float("nan"),
+        "probe_electrode": 194,
+        "rel_x": float("nan"),
+        "rel_y": float("nan"),
+        "rel_z": float("nan"),
+    }
+
+    result = utils.dj_helper_fn._replace_nan_with_default(test_data)
+
+    # Check that NaN values were replaced with -1.0
+    assert result["contact_size"] == -1.0
+    assert result["rel_x"] == -1.0
+    assert result["rel_y"] == -1.0
+    assert result["rel_z"] == -1.0
+
+    # Check that non-NaN values were preserved
+    assert result["probe_id"] == "nTrode32_probe description"
+    assert result["probe_shank"] == 0
+    assert result["probe_electrode"] == 194

tests/conftest.py

@@ -290,9 +290,9 @@ def mini_insert(
 ):
     from spyglass.common import LabMember, Nwbfile, Session  # noqa: E402
     from spyglass.data_import import insert_sessions  # noqa: E402
-    from spyglass.spikesorting.spikesorting_merge import (
+    from spyglass.spikesorting.spikesorting_merge import (  # noqa: E402
         SpikeSortingOutput,
-    )  # noqa: E402
+    )
     from spyglass.utils.nwb_helper_fn import close_nwb_files  # noqa: E402
 
     _ = SpikeSortingOutput()
@@ -400,6 +400,14 @@ def populate_exception():
     yield PopulateException
 
 
+@pytest.fixture(scope="session")
+def utils():
+    """Spyglass utils module."""
+    from spyglass import utils
+
+    yield utils
+
+
 @pytest.fixture(scope="session")
 def frequent_imports():
     """Often needed for graph cascade."""