Merge pull request #179 from h-mayorquin/add_imro_tests

Add imro tests

Author: alejoe91

probeinterface/io.py

@@ -103,7 +103,7 @@ def write_probeinterface(file: Union[str, Path], probe_or_probegroup: Union[Prob
 tsv_label_map_to_probeinterface = {v: k for k, v in tsv_label_map_to_BIDS.items()}
 
 
-def read_BIDS_probe(folder: Union[str, Path], prefix: bool = None) -> ProbeGroup:
+def read_BIDS_probe(folder: Union[str, Path], prefix: Optional[str] = None) -> ProbeGroup:
     """
     Read to BIDS probe format.
 
@@ -114,7 +114,7 @@ def read_BIDS_probe(folder: Union[str, Path], prefix: bool = None) -> ProbeGroup
     ----------
     folder: Path or str
         The folder to scan for probes and contacts files
-    prefix : None or str
+    prefix : str
         Prefix of the probes and contacts files
 
     Returns
@@ -140,8 +140,8 @@ def read_BIDS_probe(folder: Union[str, Path], prefix: bool = None) -> ProbeGroup
         probes_file = probes_files[0]
         contacts_file = contacts_files[0]
     else:
-        probes_file = folder / prefix + "_probes.tsv"
-        contacts_file = folder / prefix + "_contacts.tsv"
+        probes_file = folder / f"{prefix}_probes.tsv"
+        contacts_file = folder / f"{prefix}_contacts.tsv"
         for file in [probes_file, contacts_file]:
             if not file.exists():
                 raise ValueError(f"Source file does not exist ({file})")
@@ -919,7 +919,8 @@ def write_csv(file, probe):
     },
 }
 
-# Map imDatPrb_pn to imDatPrb_type (values from the meta file) when the latter is missing
+
+# Map imDatPrb_pn (probe number) to imDatPrb_type (probe type) when the latter is missing
 probe_number_to_probe_type = {
     "PRB_1_4_0480_1": 0,
     "PRB_1_4_0480_1_C": 0,
@@ -928,16 +929,17 @@ def write_csv(file, probe):
     "NP1030": 1030,
     "NP1031": 1031,
     "NP1032": 1032,
+    None: 0,
 }
 
 
-def read_imro(file: Union[str, Path]) -> Probe:
+def read_imro(file_path: Union[str, Path]) -> Probe:
     """
     Read probe position from the imro file used in input of SpikeGlx and Open-Ephys for neuropixels probes.
 
     Parameters
     ----------
-    file : Path or str
+    file_path : Path or str
         The .imro file path
 
     Returns
@@ -946,14 +948,14 @@ def read_imro(file: Union[str, Path]) -> Probe:
 
     """
     # the input is an imro file
-    meta_file = Path(file)
+    meta_file = Path(file_path)
     assert meta_file.suffix == ".imro", "'file' should point to the .imro file"
     with meta_file.open(mode='r') as f:
         imro_str = str(f.read())
     return _read_imro_string(imro_str)
 
 
-def _read_imro_string(imro_str: str, imDatPrb_pn: str = None) -> Probe:
+def _read_imro_string(imro_str: str, imDatPrb_pn: Optional[str] = None) -> Probe:
     """
     Parse the IMRO table when presented as a string and create a Probe object.
 
@@ -984,8 +986,9 @@ def _read_imro_string(imro_str: str, imDatPrb_pn: str = None) -> Probe:
     elif len(imro_table_header) == 2:
         imDatPrb_type, num_contact = imro_table_header
     else:
-        raise RuntimeError(f'read_imro error, the header has a strange length: {imro_table_header}')
+        raise ValueError(f'read_imro error, the header has a strange length: {imro_table_header}')
 
+
     if imDatPrb_type in [0, None]:
         imDatPrb_type = probe_number_to_probe_type[imDatPrb_pn] 
 
@@ -1096,7 +1099,7 @@ def write_imro(file, probe):
                 f"({data['device_channel_indices'][ch]} {data['shank_ids'][ch]} {annotations['banks'][ch]} "
                 f"{annotations['references'][ch]} {data['contact_ids'][ch][3:]})")
     else:
-        raise RuntimeError(f'unknown imro type : {probe_type}')
+        raise ValueError(f'unknown imro type : {probe_type}')
     with open(file, "w") as f:
         f.write(''.join(ret))
 

tests/data/test_multi_shank.imro renamed to tests/data/imro/test_multi_shank.imro

@@ -0,0 +1 @@
+(-10,384,10, 5)(0 0 0 500 125 1)

tests/data/imro/test_non_standard.imro

@@ -0,0 +1 @@
+(0,384)(0 0 0 500 125 1)(1 0 0 500 125 1)(2 0 0 500 125 1)(3 0 0 500 125 1)(4 0 0 500 125 1)(5 0 0 500 125 1)(6 0 0 500 125 1)(7 0 0 500 125 1)(8 0 0 500 125 1)(9 0 0 500 125 1)(10 0 0 500 125 1)(11 0 0 500 125 1)(12 0 0 500 125 1)(13 0 0 500 125 1)(14 0 0 500 125 1)(15 0 0 500 125 1)(16 0 0 500 125 1)(17 0 0 500 125 1)(18 0 0 500 125 1)(19 0 0 500 125 1)(20 0 0 500 125 1)(21 0 0 500 125 1)(22 0 0 500 125 1)(23 0 0 500 125 1)(24 0 0 500 125 1)(25 0 0 500 125 1)(26 0 0 500 125 1)(27 0 0 500 125 1)(28 0 0 500 125 1)(29 0 0 500 125 1)(30 0 0 500 125 1)(31 0 0 500 125 1)(32 0 0 500 125 1)(33 0 0 500 125 1)(34 0 0 500 125 1)(35 0 0 500 125 1)(36 0 0 500 125 1)(37 0 0 500 125 1)(38 0 0 500 125 1)(39 0 0 500 125 1)(40 0 0 500 125 1)(41 0 0 500 125 1)(42 0 0 500 125 1)(43 0 0 500 125 1)(44 0 0 500 125 1)(45 0 0 500 125 1)(46 0 0 500 125 1)(47 0 0 500 125 1)(48 0 0 500 125 1)(49 0 0 500 125 1)(50 0 0 500 125 1)(51 0 0 500 125 1)(52 0 0 500 125 1)(53 0 0 500 125 1)(54 0 0 500 125 1)(55 0 0 500 125 1)(56 0 0 500 125 1)(57 0 0 500 125 1)(58 0 0 500 125 1)(59 0 0 500 125 1)(60 0 0 500 125 1)(61 0 0 500 125 1)(62 0 0 500 125 1)(63 0 0 500 125 1)(64 0 0 500 125 1)(65 0 0 500 125 1)(66 0 0 500 125 1)(67 0 0 500 125 1)(68 0 0 500 125 1)(69 0 0 500 125 1)(70 0 0 500 125 1)(71 0 0 500 125 1)(72 0 0 500 125 1)(73 0 0 500 125 1)(74 0 0 500 125 1)(75 0 0 500 125 1)(76 0 0 500 125 1)(77 0 0 500 125 1)(78 0 0 500 125 1)(79 0 0 500 125 1)(80 0 0 500 125 1)(81 0 0 500 125 1)(82 0 0 500 125 1)(83 0 0 500 125 1)(84 0 0 500 125 1)(85 0 0 500 125 1)(86 0 0 500 125 1)(87 0 0 500 125 1)(88 0 0 500 125 1)(89 0 0 500 125 1)(90 0 0 500 125 1)(91 0 0 500 125 1)(92 0 0 500 125 1)(93 0 0 500 125 1)(94 0 0 500 125 1)(95 0 0 500 125 1)(96 0 0 500 125 1)(97 0 0 500 125 1)(98 0 0 500 125 1)(99 0 0 500 125 1)(100 0 0 500 125 1)(101 0 0 500 125 1)(102 0 0 500 125 1)(103 0 0 500 125 1)(104 0 0 500 125 1)(105 0 0 500 125 1)(106 0 0 500 125 1)(107 0 0 500 125 1)(108 0 0 500 125 1)(109 0 0 500 125 1)(110 0 0 500 125 1)(111 0 0 500 125 1)(112 0 0 500 125 1)(113 0 0 500 125 1)(114 0 0 500 125 1)(115 0 0 500 125 1)(116 0 0 500 125 1)(117 0 0 500 125 1)(118 0 0 500 125 1)(119 0 0 500 125 1)(120 0 0 500 125 1)(121 0 0 500 125 1)(122 0 0 500 125 1)(123 0 0 500 125 1)(124 0 0 500 125 1)(125 0 0 500 125 1)(126 0 0 500 125 1)(127 0 0 500 125 1)(128 0 0 500 125 1)(129 0 0 500 125 1)(130 0 0 500 125 1)(131 0 0 500 125 1)(132 0 0 500 125 1)(133 0 0 500 125 1)(134 0 0 500 125 1)(135 0 0 500 125 1)(136 0 0 500 125 1)(137 0 0 500 125 1)(138 0 0 500 125 1)(139 0 0 500 125 1)(140 0 0 500 125 1)(141 0 0 500 125 1)(142 0 0 500 125 1)(143 0 0 500 125 1)(144 0 0 500 125 1)(145 0 0 500 125 1)(146 0 0 500 125 1)(147 0 0 500 125 1)(148 0 0 500 125 1)(149 0 0 500 125 1)(150 0 0 500 125 1)(151 0 0 500 125 1)(152 0 0 500 125 1)(153 0 0 500 125 1)(154 0 0 500 125 1)(155 0 0 500 125 1)(156 0 0 500 125 1)(157 0 0 500 125 1)(158 0 0 500 125 1)(159 0 0 500 125 1)(160 0 0 500 125 1)(161 0 0 500 125 1)(162 0 0 500 125 1)(163 0 0 500 125 1)(164 0 0 500 125 1)(165 0 0 500 125 1)(166 0 0 500 125 1)(167 0 0 500 125 1)(168 0 0 500 125 1)(169 0 0 500 125 1)(170 0 0 500 125 1)(171 0 0 500 125 1)(172 0 0 500 125 1)(173 0 0 500 125 1)(174 0 0 500 125 1)(175 0 0 500 125 1)(176 0 0 500 125 1)(177 0 0 500 125 1)(178 0 0 500 125 1)(179 0 0 500 125 1)(180 0 0 500 125 1)(181 0 0 500 125 1)(182 0 0 500 125 1)(183 0 0 500 125 1)(184 0 0 500 125 1)(185 0 0 500 125 1)(186 0 0 500 125 1)(187 0 0 500 125 1)(188 0 0 500 125 1)(189 0 0 500 125 1)(190 0 0 500 125 1)(191 0 0 500 125 1)(192 0 0 500 125 1)(193 0 0 500 125 1)(194 0 0 500 125 1)(195 0 0 500 125 1)(196 0 0 500 125 1)(197 0 0 500 125 1)(198 0 0 500 125 1)(199 0 0 500 125 1)(200 0 0 500 125 1)(201 0 0 500 125 1)(202 0 0 500 125 1)(203 0 0 500 125 1)(204 0 0 500 125 1)(205 0 0 500 125 1)(206 0 0 500 125 1)(207 0 0 500 125 1)(208 0 0 500 125 1)(209 0 0 500 125 1)(210 0 0 500 125 1)(211 0 0 500 125 1)(212 0 0 500 125 1)(213 0 0 500 125 1)(214 0 0 500 125 1)(215 0 0 500 125 1)(216 0 0 500 125 1)(217 0 0 500 125 1)(218 0 0 500 125 1)(219 0 0 500 125 1)(220 0 0 500 125 1)(221 0 0 500 125 1)(222 0 0 500 125 1)(223 0 0 500 125 1)(224 0 0 500 125 1)(225 0 0 500 125 1)(226 0 0 500 125 1)(227 0 0 500 125 1)(228 0 0 500 125 1)(229 0 0 500 125 1)(230 0 0 500 125 1)(231 0 0 500 125 1)(232 0 0 500 125 1)(233 0 0 500 125 1)(234 0 0 500 125 1)(235 0 0 500 125 1)(236 0 0 500 125 1)(237 0 0 500 125 1)(238 0 0 500 125 1)(239 0 0 500 125 1)(240 0 0 500 125 1)(241 0 0 500 125 1)(242 0 0 500 125 1)(243 0 0 500 125 1)(244 0 0 500 125 1)(245 0 0 500 125 1)(246 0 0 500 125 1)(247 0 0 500 125 1)(248 0 0 500 125 1)(249 0 0 500 125 1)(250 0 0 500 125 1)(251 0 0 500 125 1)(252 0 0 500 125 1)(253 0 0 500 125 1)(254 0 0 500 125 1)(255 0 0 500 125 1)(256 0 0 500 125 1)(257 0 0 500 125 1)(258 0 0 500 125 1)(259 0 0 500 125 1)(260 0 0 500 125 1)(261 0 0 500 125 1)(262 0 0 500 125 1)(263 0 0 500 125 1)(264 0 0 500 125 1)(265 0 0 500 125 1)(266 0 0 500 125 1)(267 0 0 500 125 1)(268 0 0 500 125 1)(269 0 0 500 125 1)(270 0 0 500 125 1)(271 0 0 500 125 1)(272 0 0 500 125 1)(273 0 0 500 125 1)(274 0 0 500 125 1)(275 0 0 500 125 1)(276 0 0 500 125 1)(277 0 0 500 125 1)(278 0 0 500 125 1)(279 0 0 500 125 1)(280 0 0 500 125 1)(281 0 0 500 125 1)(282 0 0 500 125 1)(283 0 0 500 125 1)(284 0 0 500 125 1)(285 0 0 500 125 1)(286 0 0 500 125 1)(287 0 0 500 125 1)(288 0 0 500 125 1)(289 0 0 500 125 1)(290 0 0 500 125 1)(291 0 0 500 125 1)(292 0 0 500 125 1)(293 0 0 500 125 1)(294 0 0 500 125 1)(295 0 0 500 125 1)(296 0 0 500 125 1)(297 0 0 500 125 1)(298 0 0 500 125 1)(299 0 0 500 125 1)(300 0 0 500 125 1)(301 0 0 500 125 1)(302 0 0 500 125 1)(303 0 0 500 125 1)(304 0 0 500 125 1)(305 0 0 500 125 1)(306 0 0 500 125 1)(307 0 0 500 125 1)(308 0 0 500 125 1)(309 0 0 500 125 1)(310 0 0 500 125 1)(311 0 0 500 125 1)(312 0 0 500 125 1)(313 0 0 500 125 1)(314 0 0 500 125 1)(315 0 0 500 125 1)(316 0 0 500 125 1)(317 0 0 500 125 1)(318 0 0 500 125 1)(319 0 0 500 125 1)(320 0 0 500 125 1)(321 0 0 500 125 1)(322 0 0 500 125 1)(323 0 0 500 125 1)(324 0 0 500 125 1)(325 0 0 500 125 1)(326 0 0 500 125 1)(327 0 0 500 125 1)(328 0 0 500 125 1)(329 0 0 500 125 1)(330 0 0 500 125 1)(331 0 0 500 125 1)(332 0 0 500 125 1)(333 0 0 500 125 1)(334 0 0 500 125 1)(335 0 0 500 125 1)(336 0 0 500 125 1)(337 0 0 500 125 1)(338 0 0 500 125 1)(339 0 0 500 125 1)(340 0 0 500 125 1)(341 0 0 500 125 1)(342 0 0 500 125 1)(343 0 0 500 125 1)(344 0 0 500 125 1)(345 0 0 500 125 1)(346 0 0 500 125 1)(347 0 0 500 125 1)(348 0 0 500 125 1)(349 0 0 500 125 1)(350 0 0 500 125 1)(351 0 0 500 125 1)(352 0 0 500 125 1)(353 0 0 500 125 1)(354 0 0 500 125 1)(355 0 0 500 125 1)(356 0 0 500 125 1)(357 0 0 500 125 1)(358 0 0 500 125 1)(359 0 0 500 125 1)(360 0 0 500 125 1)(361 0 0 500 125 1)(362 0 0 500 125 1)(363 0 0 500 125 1)(364 0 0 500 125 1)(365 0 0 500 125 1)(366 0 0 500 125 1)(367 0 0 500 125 1)(368 0 0 500 125 1)(369 0 0 500 125 1)(370 0 0 500 125 1)(371 0 0 500 125 1)(372 0 0 500 125 1)(373 0 0 500 125 1)(374 0 0 500 125 1)(375 0 0 500 125 1)(376 0 0 500 125 1)(377 0 0 500 125 1)(378 0 0 500 125 1)(379 0 0 500 125 1)(380 0 0 500 125 1)(381 0 0 500 125 1)(382 0 0 500 125 1)(383 0 0 500 125 1)

tests/data/imro/test_old_probe.imro

@@ -0,0 +1 @@
+(21,384)(0 1 0 0)(1 1 0 1)(2 1 0 2)(3 1 0 3)(4 1 0 4)(5 1 0 5)(6 1 0 6)(7 1 0 7)(8 1 0 8)(9 1 0 9)(10 1 0 10)(11 1 0 11)(12 1 0 12)(13 1 0 13)(14 1 0 14)(15 1 0 15)(16 1 0 16)(17 1 0 17)(18 1 0 18)(19 1 0 19)(20 1 0 20)(21 1 0 21)(22 1 0 22)(23 1 0 23)(24 1 0 24)(25 1 0 25)(26 1 0 26)(27 1 0 27)(28 1 0 28)(29 1 0 29)(30 1 0 30)(31 1 0 31)(32 1 0 32)(33 1 0 33)(34 1 0 34)(35 1 0 35)(36 1 0 36)(37 1 0 37)(38 1 0 38)(39 1 0 39)(40 1 0 40)(41 1 0 41)(42 1 0 42)(43 1 0 43)(44 1 0 44)(45 1 0 45)(46 1 0 46)(47 1 0 47)(48 1 0 48)(49 1 0 49)(50 1 0 50)(51 1 0 51)(52 1 0 52)(53 1 0 53)(54 1 0 54)(55 1 0 55)(56 1 0 56)(57 1 0 57)(58 1 0 58)(59 1 0 59)(60 1 0 60)(61 1 0 61)(62 1 0 62)(63 1 0 63)(64 1 0 64)(65 1 0 65)(66 1 0 66)(67 1 0 67)(68 1 0 68)(69 1 0 69)(70 1 0 70)(71 1 0 71)(72 1 0 72)(73 1 0 73)(74 1 0 74)(75 1 0 75)(76 1 0 76)(77 1 0 77)(78 1 0 78)(79 1 0 79)(80 1 0 80)(81 1 0 81)(82 1 0 82)(83 1 0 83)(84 1 0 84)(85 1 0 85)(86 1 0 86)(87 1 0 87)(88 1 0 88)(89 1 0 89)(90 1 0 90)(91 1 0 91)(92 1 0 92)(93 1 0 93)(94 1 0 94)(95 1 0 95)(96 1 0 96)(97 1 0 97)(98 1 0 98)(99 1 0 99)(100 1 0 100)(101 1 0 101)(102 1 0 102)(103 1 0 103)(104 1 0 104)(105 1 0 105)(106 1 0 106)(107 1 0 107)(108 1 0 108)(109 1 0 109)(110 1 0 110)(111 1 0 111)(112 1 0 112)(113 1 0 113)(114 1 0 114)(115 1 0 115)(116 1 0 116)(117 1 0 117)(118 1 0 118)(119 1 0 119)(120 1 0 120)(121 1 0 121)(122 1 0 122)(123 1 0 123)(124 1 0 124)(125 1 0 125)(126 1 0 126)(127 1 0 127)(128 1 0 128)(129 1 0 129)(130 1 0 130)(131 1 0 131)(132 1 0 132)(133 1 0 133)(134 1 0 134)(135 1 0 135)(136 1 0 136)(137 1 0 137)(138 1 0 138)(139 1 0 139)(140 1 0 140)(141 1 0 141)(142 1 0 142)(143 1 0 143)(144 1 0 144)(145 1 0 145)(146 1 0 146)(147 1 0 147)(148 1 0 148)(149 1 0 149)(150 1 0 150)(151 1 0 151)(152 1 0 152)(153 1 0 153)(154 1 0 154)(155 1 0 155)(156 1 0 156)(157 1 0 157)(158 1 0 158)(159 1 0 159)(160 1 0 160)(161 1 0 161)(162 1 0 162)(163 1 0 163)(164 1 0 164)(165 1 0 165)(166 1 0 166)(167 1 0 167)(168 1 0 168)(169 1 0 169)(170 1 0 170)(171 1 0 171)(172 1 0 172)(173 1 0 173)(174 1 0 174)(175 1 0 175)(176 1 0 176)(177 1 0 177)(178 1 0 178)(179 1 0 179)(180 1 0 180)(181 1 0 181)(182 1 0 182)(183 1 0 183)(184 1 0 184)(185 1 0 185)(186 1 0 186)(187 1 0 187)(188 1 0 188)(189 1 0 189)(190 1 0 190)(191 1 0 191)(192 1 0 192)(193 1 0 193)(194 1 0 194)(195 1 0 195)(196 1 0 196)(197 1 0 197)(198 1 0 198)(199 1 0 199)(200 1 0 200)(201 1 0 201)(202 1 0 202)(203 1 0 203)(204 1 0 204)(205 1 0 205)(206 1 0 206)(207 1 0 207)(208 1 0 208)(209 1 0 209)(210 1 0 210)(211 1 0 211)(212 1 0 212)(213 1 0 213)(214 1 0 214)(215 1 0 215)(216 1 0 216)(217 1 0 217)(218 1 0 218)(219 1 0 219)(220 1 0 220)(221 1 0 221)(222 1 0 222)(223 1 0 223)(224 1 0 224)(225 1 0 225)(226 1 0 226)(227 1 0 227)(228 1 0 228)(229 1 0 229)(230 1 0 230)(231 1 0 231)(232 1 0 232)(233 1 0 233)(234 1 0 234)(235 1 0 235)(236 1 0 236)(237 1 0 237)(238 1 0 238)(239 1 0 239)(240 1 0 240)(241 1 0 241)(242 1 0 242)(243 1 0 243)(244 1 0 244)(245 1 0 245)(246 1 0 246)(247 1 0 247)(248 1 0 248)(249 1 0 249)(250 1 0 250)(251 1 0 251)(252 1 0 252)(253 1 0 253)(254 1 0 254)(255 1 0 255)(256 1 0 256)(257 1 0 257)(258 1 0 258)(259 1 0 259)(260 1 0 260)(261 1 0 261)(262 1 0 262)(263 1 0 263)(264 1 0 264)(265 1 0 265)(266 1 0 266)(267 1 0 267)(268 1 0 268)(269 1 0 269)(270 1 0 270)(271 1 0 271)(272 1 0 272)(273 1 0 273)(274 1 0 274)(275 1 0 275)(276 1 0 276)(277 1 0 277)(278 1 0 278)(279 1 0 279)(280 1 0 280)(281 1 0 281)(282 1 0 282)(283 1 0 283)(284 1 0 284)(285 1 0 285)(286 1 0 286)(287 1 0 287)(288 1 0 288)(289 1 0 289)(290 1 0 290)(291 1 0 291)(292 1 0 292)(293 1 0 293)(294 1 0 294)(295 1 0 295)(296 1 0 296)(297 1 0 297)(298 1 0 298)(299 1 0 299)(300 1 0 300)(301 1 0 301)(302 1 0 302)(303 1 0 303)(304 1 0 304)(305 1 0 305)(306 1 0 306)(307 1 0 307)(308 1 0 308)(309 1 0 309)(310 1 0 310)(311 1 0 311)(312 1 0 312)(313 1 0 313)(314 1 0 314)(315 1 0 315)(316 1 0 316)(317 1 0 317)(318 1 0 318)(319 1 0 319)(320 1 0 320)(321 1 0 321)(322 1 0 322)(323 1 0 323)(324 1 0 324)(325 1 0 325)(326 1 0 326)(327 1 0 327)(328 1 0 328)(329 1 0 329)(330 1 0 330)(331 1 0 331)(332 1 0 332)(333 1 0 333)(334 1 0 334)(335 1 0 335)(336 1 0 336)(337 1 0 337)(338 1 0 338)(339 1 0 339)(340 1 0 340)(341 1 0 341)(342 1 0 342)(343 1 0 343)(344 1 0 344)(345 1 0 345)(346 1 0 346)(347 1 0 347)(348 1 0 348)(349 1 0 349)(350 1 0 350)(351 1 0 351)(352 1 0 352)(353 1 0 353)(354 1 0 354)(355 1 0 355)(356 1 0 356)(357 1 0 357)(358 1 0 358)(359 1 0 359)(360 1 0 360)(361 1 0 361)(362 1 0 362)(363 1 0 363)(364 1 0 364)(365 1 0 365)(366 1 0 366)(367 1 0 367)(368 1 0 368)(369 1 0 369)(370 1 0 370)(371 1 0 371)(372 1 0 372)(373 1 0 373)(374 1 0 374)(375 1 0 375)(376 1 0 376)(377 1 0 377)(378 1 0 378)(379 1 0 379)(380 1 0 380)(381 1 0 381)(382 1 0 382)(383 1 0 383)

tests/data/imro/test_single_shak_2.0.imro

@@ -0,0 +1,52 @@
+from pathlib import Path
+
+import pytest
+import numpy as np
+
+from probeinterface import read_imro, write_imro
+
+data_path = Path(__file__).absolute().parent.parent / "data" / "imro"
+
+
+def test_reading_multishank_imro(tmp_path):
+    probe = read_imro(data_path / "test_multi_shank.imro")
+
+    file_path = tmp_path / "multi_shank_written.imro"
+    write_imro(file_path, probe)
+    probe2 = read_imro(file_path)
+    np.testing.assert_array_equal(probe2.contact_ids, probe.contact_ids)
+    np.testing.assert_array_equal(probe2.contact_positions, probe.contact_positions)
+
+
+def test_reading_imro_2_single_shank(tmp_path):
+    probe = read_imro(data_path / "test_single_shak_2.0.imro")
+
+    file_path = tmp_path / "test_single_shak_2.0_written.imro"
+    write_imro(file_path, probe)
+    probe2 = read_imro(file_path)
+    np.testing.assert_array_equal(probe2.contact_ids, probe.contact_ids)
+    np.testing.assert_array_equal(probe2.contact_positions, probe.contact_positions)
+
+
+def test_reading_old_imro(tmp_path):
+    probe = read_imro(data_path / "test_old_probe.imro")
+
+    file_path = tmp_path / "test_old_probe.imro"
+    write_imro(file_path, probe)
+    probe2 = read_imro(file_path)
+    np.testing.assert_array_equal(probe2.contact_ids, probe.contact_ids)
+    np.testing.assert_array_equal(probe2.contact_positions, probe.contact_positions)
+
+
+def test_raising_error_when_writing_with_wrong_type(tmp_path):
+    probe = read_imro(data_path / "test_old_probe.imro")
+
+    file_path = tmp_path / "test_old_probe.imro"
+    probe.annotations["probe_type"] = "something_that_should_make_write_fail"
+    with pytest.raises(ValueError):
+        write_imro(file_path, probe)
+
+
+def test_non_standard_file():
+    with pytest.raises(ValueError):
+        probe = read_imro(data_path / "test_non_standard.imro")

tests/test_io/test_imro.py

@@ -5,11 +5,8 @@
     read_BIDS_probe,
 )
 from probeinterface import read_prb, write_prb
-from probeinterface import (
-    read_imro,
-    write_imro,
-)
-from probeinterface import generate_dummy_probe_group
+
+from probeinterface import generate_dummy_probe_group, generate_dummy_probe
 
 
 from pathlib import Path
@@ -46,6 +43,24 @@ def test_probeinterface_format(tmp_path):
     # ~ plot_probe_group(probegroup2, with_channel_index=True, same_axes=False)
     # ~ plt.show()
 
+def test_writeprobeinterface(tmp_path):
+    probe = generate_dummy_probe()
+    file_path = tmp_path / "test.prb"
+    write_probeinterface(file_path, probe)
+    
+    probe_read = read_probeinterface(file_path).probes[0]
+    assert probe.get_contact_count() == probe_read.get_contact_count()
+    assert np.allclose(probe.contact_positions, probe_read.contact_positions)
+    assert np.allclose(probe.probe_planar_contour, probe_read.probe_planar_contour)
+
+
+def test_writeproeinterface_raises_error_with_bad_input(tmp_path):
+    probe = "WrongInput"
+    file_path = tmp_path / "test.prb"
+    with pytest.raises(ValueError):
+        write_probeinterface(file_path, probe)
+    
+    
 
 def test_BIDS_format(tmp_path):
     folder_path = tmp_path / "test_BIDS"
@@ -203,17 +218,6 @@ def test_prb(tmp_path):
     # plot_probe(probe)
     # plt.show()
 
-
-def test_readimro(tmp_path):
-    probe = read_imro(data_path / "test_multi_shank.imro")
-
-    file_path = tmp_path / "multi_shank_written.imro"
-    write_imro(file_path, probe)
-    probe2 = read_imro(file_path)
-    np.testing.assert_array_equal(probe2.contact_ids, probe.contact_ids)
-    np.testing.assert_array_equal(probe2.contact_positions, probe.contact_positions)
-
-
 if __name__ == "__main__":
     # test_probeinterface_format()
     # test_BIDS_format()