Coverage for labnirs2snirf / layout.py: 100%

1""" 

2Functions related to importing probe positions. 

3""" 

4 

5import logging 

6from pathlib import Path 

7 

8from .error import Labnirs2SnirfError 

9from .model import Nirs 

10 

11 

12class LayoutError(Labnirs2SnirfError): 

13 """Custom error class for layout-related issues.""" 

14 

15 

16type Layout_3D = dict[str, tuple[float, float, float]] 

17 

18log = logging.getLogger(__name__) 

19 

20 

21def read_layout(file: Path) -> Layout_3D: 

22 """ 

23 Read optode coordinates from file. 

24 

25 Parameters 

26 ---------- 

27 file : Path 

28 Path pointing to location file. 

29 

30 Returns 

31 ------- 

32 dict[str, tuple[float, float, float]] 

33 Dict mapping probe labels to 3D coordinates. 

34 

35 Notes 

36 ----- 

37 Location files are expected to follow the .sfp format, which is a tab-separated text file 

38 with columns: label, x, y, and z, where 'x', 'y', and 'z' are the 3D coordinates of the optode. 

39 Labels are case-sensitive. If duplicate labels are found, a warning is logged and the last 

40 occurrence is used. 

41 """ 

42 

43 import polars as pl # pylint: disable=C0415 

44 

45 log.debug("Reading layout file: %s", file) 

46 

47 try: 

48 locations = ( 

49 pl.scan_csv( 

50 file, 

51 has_header=False, 

52 separator="\t", 

53 schema=pl.Schema( 

54 zip(["label", "x", "y", "z"], [pl.String] + [pl.Float64] * 3), 

55 ), 

56 ) 

57 # remove whitespace around values 

58 .with_columns(pl.col(pl.String).str.strip_chars()) 

59 .collect() 

60 ) 

61 

62 # Check for duplicate labels 

63 if locations.height != locations.select("label").unique().height: 

64 duplicates = ( 

65 locations.group_by("label") 

66 .agg(pl.len().alias("count")) 

67 .filter(pl.col("count") > 1) 

68 .get_column("label") 

69 .to_list() 

70 ) 

71 log.warning( 

72 "Duplicate labels found in layout file: %s. Last occurrence will be used.", 

73 ", ".join(duplicates), 

74 ) 

75 

76 layout_dict = { 

77 row["label"]: (row["x"], row["y"], row["z"]) 

78 for row in locations.rows(named=True) 

79 } 

80 

81 log.debug( 

82 "Successfully read %d optode positions from layout file", 

83 len(layout_dict), 

84 ) 

85 return layout_dict 

86 

87 except ( 

88 pl.exceptions.ComputeError, 

89 pl.exceptions.ColumnNotFoundError, 

90 pl.exceptions.SchemaError, 

91 ) as e: 

92 log.exception("Failed to read layout file %s: %s", file, e) 

93 raise LayoutError( 

94 f"Failed to read layout file {file}: {e}. " 

95 "Ensure the file is tab-separated with four columns: label, x, y, z," 

96 " and that x, y, z are numeric values.", 

97 ) from e 

98 except FileNotFoundError as e: 

99 log.exception("Layout file not found: %s", file) 

100 raise LayoutError(f"Layout file not found: {file}") from e 

101 

102 except Exception as ex: 

103 log.exception("Failed to read layout file %s: %s", file, ex) 

104 raise 

105 

106 

107def update_layout(data: Nirs, locations: Layout_3D) -> None: 

108 """ 

109 Update source and position 3D coordinates in nirs data. 

110 

111 Parameters 

112 ---------- 

113 data : model.Nirs 

114 Nirs data produced by `labirs.read_nirs()`. Probes have labels (Si, Di) 

115 corresponding to their original numbering in the exported labNIRS data. 

116 locations : dict[str, tuple[float, float, float]] 

117 Mapping of probe labels to 3D coordinates. These can be read in from file 

118 by `read_layout()`. Coordinates are always stored as 3D tuples - when only 

119 2D is available, Z is set to 0. Labels are case-sensitive. 

120 

121 Raises 

122 ------ 

123 LayoutError 

124 If no probe labels present in ``data``. 

125 """ 

126 if data.probe.sourceLabels is None and data.probe.detectorLabels is None: 

127 raise LayoutError( 

128 "Update layout failed because there are no probe labels in NIRS data.", 

129 ) 

130 

131 log.debug("Updating probe positions with %d provided locations", len(locations)) 

132 

133 sources_updated = 0 

134 detectors_updated = 0 

135 

136 # Update source positions 

137 if data.probe.sourceLabels: 

138 for i, label in enumerate(data.probe.sourceLabels): 

139 if label in locations: 

140 data.probe.sourcePos3D[i, :] = locations[label] 

141 sources_updated += 1 

142 else: 

143 log.warning( 

144 "Source %s missing position data, keeping default (0, 0, 0)", 

145 label, 

146 ) 

147 

148 # Update detector positions 

149 if data.probe.detectorLabels: 

150 for i, label in enumerate(data.probe.detectorLabels): 

151 if label in locations: 

152 data.probe.detectorPos3D[i, :] = locations[label] 

153 detectors_updated += 1 

154 else: 

155 log.warning( 

156 "Detector %s missing position data, keeping default (0, 0, 0)", 

157 label, 

158 ) 

159 

160 log.info( 

161 "Updated positions for %d sources and %d detectors", 

162 sources_updated, 

163 detectors_updated, 

164 )