From c41d02d5b4a088f614c31ac08955db36b043cfa9 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 24 Aug 2026 11:44:02 +0200 Subject: [PATCH] fix: update schema with probe_ids and global_contact_order --- src/probeinterface/generator.py | 2 +- src/probeinterface/schema/probe.json.schema | 12 +++ src/probeinterface/testing.py | 12 +++ tests/test_generator.py | 15 ++++ tests/test_schema.py | 82 ++++++++++++++++++++- 5 files changed, 120 insertions(+), 3 deletions(-) diff --git a/src/probeinterface/generator.py b/src/probeinterface/generator.py index 05e9b5ab..f02ba5b6 100644 --- a/src/probeinterface/generator.py +++ b/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 diff --git a/src/probeinterface/schema/probe.json.schema b/src/probeinterface/schema/probe.json.schema index b8cffd49..8cb5e79f 100644 --- a/src/probeinterface/schema/probe.json.schema +++ b/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"], diff --git a/src/probeinterface/testing.py b/src/probeinterface/testing.py index f247aceb..e4d6623b 100644 --- a/src/probeinterface/testing.py +++ b/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) diff --git a/tests/test_generator.py b/tests/test_generator.py index d14ed9a5..b8c44102 100644 --- a/tests/test_generator.py +++ b/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() diff --git a/tests/test_schema.py b/tests/test_schema.py index b3839926..80c7fa0e 100644 --- a/tests/test_schema.py +++ b/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)