Update schema with probe_ids and global_contact_order

src/probeinterface/generator.py

@@ -47,7 +47,7 @@ def generate_dummy_probe(elec_shapes: Literal["circle", "square", "rect"] = "cir
         contact_shape_params=contact_shape_params,
     )
 
-    probe.annotate(manufacturer="me")
+    probe.annotate(manufacturer="me", model_name="dummy")
     probe.annotate_contacts(quality=np.ones(32) * 1000.0)
 
     return probe

src/probeinterface/schema/probe.json.schema

@@ -125,6 +125,18 @@
         ],
         "additionalProperties": false
       }
+    },
+    "probe_ids": {
+      "type": "array",
+      "$comment": "Optional ids of the probes, in the same order as 'probes'. When absent, ids default to the string of the probe index.",
+      "items": { "type": "string" },
+      "uniqueItems": true
+    },
+    "global_contact_order": {
+      "type": "array",
+      "$comment": "Optional global order of the contacts, as a permutation of the indices of all contacts stacked probe by probe. Absent when the order is the natural one.",
+      "items": { "type": "integer", "minimum": 0 },
+      "uniqueItems": true
     }
   },
   "required": ["specification", "version", "probes"],

src/probeinterface/testing.py

@@ -13,3 +13,15 @@ def validate_probe_dict(probe_dict):
 
     instance = dict(specification="probeinterface", version=version, probes=[probe_dict])
     jsonschema.validate(instance=instance, schema=schema)
+
+
+def validate_probegroup_dict(probegroup_dict):
+    """
+    Validate a full ProbeGroup dict (as returned by ``ProbeGroup.to_dict()``) against
+    the schema. The "specification" and "version" keys are added if missing, so that
+    both a raw ``to_dict()`` and the content of a probeinterface JSON file can be passed.
+    """
+    import jsonschema
+
+    instance = {"specification": "probeinterface", "version": version, **probegroup_dict}
+    jsonschema.validate(instance=instance, schema=schema)

tests/test_generator.py

@@ -6,6 +6,7 @@
     generate_multi_columns_probe,
     generate_multi_shank,
 )
+from probeinterface.testing import validate_probe_dict, validate_probegroup_dict
 
 
 from pathlib import Path
@@ -34,5 +35,19 @@ def test_generate():
     # ~ plt.show()
 
 
+@pytest.mark.parametrize("elec_shapes", ["circle", "square", "rect"])
+def test_dummy_probe_validates(elec_shapes):
+    """The dummy probe is annotated with everything the schema requires."""
+    probe = generate_dummy_probe(elec_shapes=elec_shapes)
+    assert probe.annotations["model_name"] == "dummy"
+    assert probe.annotations["manufacturer"] == "me"
+    validate_probe_dict(probe.to_dict(array_as_list=True))
+
+
+def test_dummy_probe_group_validates():
+    probegroup = generate_dummy_probe_group()
+    validate_probegroup_dict(probegroup.to_dict(array_as_list=True))
+
+
 if __name__ == "__main__":
     test_generate()

tests/test_schema.py

@@ -1,7 +1,13 @@
+import json
 import re
 
-from probeinterface import __version__
-from probeinterface.testing import schema
+import jsonschema
+
+import numpy as np
+import pytest
+
+from probeinterface import ProbeGroup, __version__, generate_dummy_probe, write_probeinterface
+from probeinterface.testing import schema, validate_probegroup_dict
 
 
 def test_schema_is_annotated():
@@ -27,3 +33,75 @@ def test_package_version_is_compatible_with_schema():
         f"compatibility pattern ({pattern}). Either this is an incompatible schema "
         f"change (update the pattern) or the version is malformed."
     )
+
+
+def _probegroup(n_probes=3):
+    """A ProbeGroup with explicit probe ids."""
+    probegroup = ProbeGroup()
+    for i in range(n_probes):
+        probe = generate_dummy_probe()
+        probe.move([i * 100, i * 80])
+        probegroup.add_probe(probe, probe_id=f"probe_00{i}")
+    return probegroup
+
+
+def test_probegroup_dict_validates():
+    """A ProbeGroup dict, which carries 'probe_ids', validates against the schema."""
+    d = _probegroup().to_dict(array_as_list=True)
+    assert d["probe_ids"] == ["probe_000", "probe_001", "probe_002"]
+    validate_probegroup_dict(d)
+
+
+def test_probegroup_dict_with_global_contact_order_validates():
+    """A reordered ProbeGroup dict, which carries 'global_contact_order', validates."""
+    probegroup = _probegroup()
+    # an order interleaving contacts across probes, so it is not the natural one
+    order = np.concatenate([np.arange(0, 96, 2), np.arange(95, 0, -2)])
+    reordered = probegroup.get_slice(order)
+
+    d = reordered.to_dict(array_as_list=True)
+    assert d["global_contact_order"] is not None
+    validate_probegroup_dict(d)
+
+
+def test_written_probeinterface_file_validates(tmp_path):
+    """The JSON actually written by write_probeinterface validates against the schema."""
+    file = tmp_path / "probegroup.json"
+    write_probeinterface(file, _probegroup())
+
+    with open(file, "r", encoding="utf8") as f:
+        d = json.load(f)
+    assert d["specification"] == "probeinterface"
+    assert "probe_ids" in d
+    validate_probegroup_dict(d)
+
+
+@pytest.mark.parametrize(
+    "probe_ids",
+    [["0", "0"], [0, 1], "0"],
+    ids=["duplicated", "not_strings", "not_a_list"],
+)
+def test_invalid_probe_ids_are_rejected(probe_ids):
+    d = _probegroup(n_probes=2).to_dict(array_as_list=True)
+    d["probe_ids"] = probe_ids
+    with pytest.raises(jsonschema.ValidationError):
+        validate_probegroup_dict(d)
+
+
+@pytest.mark.parametrize(
+    "global_contact_order",
+    [[0, 0], [0.5, 1], [-1, 0], "0"],
+    ids=["duplicated", "not_integers", "negative", "not_a_list"],
+)
+def test_invalid_global_contact_order_is_rejected(global_contact_order):
+    d = _probegroup(n_probes=2).to_dict(array_as_list=True)
+    d["global_contact_order"] = global_contact_order
+    with pytest.raises(jsonschema.ValidationError):
+        validate_probegroup_dict(d)
+
+
+def test_unknown_top_level_key_is_rejected():
+    d = _probegroup(n_probes=2).to_dict(array_as_list=True)
+    d["unknown_key"] = "unexpected"
+    with pytest.raises(jsonschema.ValidationError):
+        validate_probegroup_dict(d)