Coverage for src/bartz/debug/_traceconv.py: 89%

1# bartz/src/bartz/debug/_traceconv.py 

2# 

3# Copyright (c) 2026, The Bartz Contributors 

4# 

5# This file is part of bartz. 

6# 

7# Permission is hereby granted, free of charge, to any person obtaining a copy 

8# of this software and associated documentation files (the "Software"), to deal 

9# in the Software without restriction, including without limitation the rights 

10# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 

11# copies of the Software, and to permit persons to whom the Software is 

12# furnished to do so, subject to the following conditions: 

13# 

14# The above copyright notice and this permission notice shall be included in all 

15# copies or substantial portions of the Software. 

16# 

17# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 

18# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 

19# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 

20# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 

21# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 

22# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 

23# SOFTWARE. 

24 

25"""Parsing of R BART3 tree traces.""" 

26 

27import math (empty)

28from re import fullmatch (empty)

29from typing import ClassVar (empty)

30 

31import numpy (empty)

32from jax import numpy as jnp (empty)

33from jax.sharding import Mesh (empty)

34from jaxtyping import Array, Float, Float32, UInt (empty)

35 

36from bartz._jaxext import Module, field, minimal_unsigned_dtype (empty)

37from bartz.BART._gbart import FloatLike (empty)

38 

39 

40def _get_next_line(s: str, i: int) -> tuple[str, int]: (empty)

41 """Get the next line from a string and the new index.""" 

42 i_new = s.find('\n', i) (empty)

43 if i_new == -1: 43 ↛ 44line 43 didn't jump to line 44 because the condition on line 43 was never true(empty)

44 return s[i:], len(s) 

45 return s[i:i_new], i_new + 1 (empty)

46 

47 

48class BARTTraceMeta(Module): (empty)

49 """Metadata of R BART tree traces.""" 

50 

51 ndpost: int = field(static=True) (empty)

52 """The number of posterior draws.""" (empty)

53 

54 ntree: int = field(static=True) (empty)

55 """The number of trees in the model.""" (empty)

56 

57 numcut: UInt[Array, ' p'] 

58 """The maximum split value for each variable.""" (empty)

59 

60 heap_size: int = field(static=True) (empty)

61 """The size of the heap required to store the trees.""" (empty)

62 

63 

64def scan_BART_trees(trees: str) -> BARTTraceMeta: (empty)

65 """Scan an R BART tree trace checking for errors and parsing metadata. 

66 

67 Parameters 

68 ---------- 

69 trees 

70 The string representation of a trace of trees of the R BART package. 

71 Can be accessed from ``mc_gbart(...).treedraws['trees']``. 

72 

73 Returns 

74 ------- 

75 An object containing the metadata. 

76 

77 Raises 

78 ------ 

79 ValueError 

80 If the string is malformed or contains leftover characters. 

81 """ 

82 # parse first line 

83 line, i_char = _get_next_line(trees, 0) (empty)

84 i_line = 1 (empty)

85 match = fullmatch(r'(\d+) (\d+) (\d+)', line) (empty)

86 if match is None: 86 ↛ 87line 86 didn't jump to line 87 because the condition on line 86 was never true(empty)

87 msg = f'Malformed header at {i_line=}' 

88 raise ValueError(msg) 

89 ndpost, ntree, p = map(int, match.groups()) (empty)

90 

91 # initial values for maxima 

92 max_heap_index = 0 (empty)

93 numcut = numpy.zeros(p, int) (empty)

94 

95 # cycle over iterations and trees 

96 for i_iter in range(ndpost): (empty)

97 for i_tree in range(ntree): (empty)

98 # parse first line of tree definition 

99 line, i_char = _get_next_line(trees, i_char) (empty)

100 i_line += 1 (empty)

101 match = fullmatch(r'(\d+)', line) (empty)

102 if match is None: 102 ↛ 103line 102 didn't jump to line 103 because the condition on line 102 was never true(empty)

103 msg = f'Malformed tree header at {i_iter=} {i_tree=} {i_line=}' 

104 raise ValueError(msg) 

105 num_nodes = int(line) (empty)

106 

107 # cycle over nodes 

108 for i_node in range(num_nodes): (empty)

109 # parse node definition 

110 line, i_char = _get_next_line(trees, i_char) (empty)

111 i_line += 1 (empty)

112 match = fullmatch( (empty)

113 r'(\d+) (\d+) (\d+) (-?\d+(\.\d+)?(e(\+|-|)\d+)?)', line 

114 ) 

115 if match is None: 115 ↛ 116line 115 didn't jump to line 116 because the condition on line 115 was never true(empty)

116 msg = f'Malformed node definition at {i_iter=} {i_tree=} {i_node=} {i_line=}' 

117 raise ValueError(msg) 

118 i_heap = int(match.group(1)) (empty)

119 var = int(match.group(2)) (empty)

120 split = int(match.group(3)) (empty)

121 

122 # update maxima 

123 numcut[var] = max(numcut[var], split) (empty)

124 max_heap_index = max(max_heap_index, i_heap) (empty)

125 

126 assert i_char <= len(trees) (empty)

127 if i_char < len(trees): 127 ↛ 128line 127 didn't jump to line 128 because the condition on line 127 was never true(empty)

128 msg = f'Leftover {len(trees) - i_char} characters in string' 

129 raise ValueError(msg) 

130 

131 # determine minimal integer type for numcut 

132 numcut += 1 # because BART is 0-based (empty)

133 split_dtype = minimal_unsigned_dtype(numcut.max().item()) (empty)

134 numcut = jnp.array(numcut.astype(split_dtype)) (empty)

135 

136 # determine minimum heap size to store the trees 

137 heap_size = 2 ** math.ceil(math.log2(max_heap_index + 1)) (empty)

138 

139 return BARTTraceMeta(ndpost=ndpost, ntree=ntree, numcut=numcut, heap_size=heap_size) (empty)

140 

141 

142class MinimalTrace(Module): (empty)

143 """A minimal trace of trees, compatible with `bartz.mcmcloop.evaluate_trace`.""" 

144 

145 leaf_tree: Float[Array, 'ndpost ntree tree_size'] = field(samples=0) (empty)

146 var_tree: UInt[Array, 'ndpost ntree tree_size//2'] = field(samples=0) (empty)

147 split_tree: UInt[Array, 'ndpost ntree tree_size//2'] = field(samples=0) (empty)

148 offset: Float32[Array, ''] 

149 """Constant shift added to the scaled sum of trees.""" (empty)

150 

151 leaf_unit: Float32[Array, ''] 

152 """The storage unit of the leaf values, 1 for leaves in data units.""" (empty)

153 

154 has_chains: ClassVar[bool] = False (empty)

155 """No chain axis; each leading axis is just the sample axis.""" (empty)

156 

157 mesh: ClassVar[Mesh | None] = None (empty)

158 """No device mesh; the trees are host-built and unsharded.""" (empty)

159 

160 

161def trees_BART_to_bartz( (empty)

162 trees: str, *, min_maxdepth: int = 0, offset: FloatLike | None = None 

163) -> tuple[MinimalTrace, BARTTraceMeta]: 

164 """Convert trees from the R BART format to the bartz format. 

165 

166 Parameters 

167 ---------- 

168 trees 

169 The string representation of a trace of trees of the R BART package. 

170 Can be accessed from ``mc_gbart(...).treedraws['trees']``. 

171 min_maxdepth 

172 The maximum tree depth of the output will be set to the maximum 

173 observed depth in the input trees. Use this parameter to require at 

174 least this maximum depth in the output format. 

175 offset 

176 The trace returned by `bartz.mcmcloop.run_mcmc` contains an offset to be 

177 summed to the sum of trees. To match that behavior, this function 

178 returns an offset as well, zero by default. Set with this parameter 

179 otherwise. 

180 

181 Returns 

182 ------- 

183 trace : MinimalTrace 

184 A representation of the trees compatible with the trace returned by 

185 `bartz.mcmcloop.run_mcmc`. 

186 meta : BARTTraceMeta 

187 The metadata of the trace, containing the number of iterations, trees, 

188 and the maximum split value. 

189 """ 

190 # scan all the string checking for errors and determining sizes 

191 meta = scan_BART_trees(trees) (empty)

192 

193 # skip first line 

194 _, i_char = _get_next_line(trees, 0) (empty)

195 

196 heap_size = max(meta.heap_size, 2**min_maxdepth) (empty)

197 leaf_trees = numpy.zeros((meta.ndpost, meta.ntree, heap_size), dtype=numpy.float32) (empty)

198 var_trees = numpy.zeros( (empty)

199 (meta.ndpost, meta.ntree, heap_size // 2), 

200 dtype=minimal_unsigned_dtype(meta.numcut.size - 1), 

201 ) 

202 split_trees = numpy.zeros( (empty)

203 (meta.ndpost, meta.ntree, heap_size // 2), dtype=meta.numcut.dtype 

204 ) 

205 

206 # cycle over iterations and trees 

207 for i_iter in range(meta.ndpost): (empty)

208 for i_tree in range(meta.ntree): (empty)

209 # parse first line of tree definition 

210 line, i_char = _get_next_line(trees, i_char) (empty)

211 num_nodes = int(line) (empty)

212 

213 is_internal = numpy.zeros(heap_size // 2, dtype=bool) (empty)

214 

215 # cycle over nodes 

216 for _ in range(num_nodes): (empty)

217 # parse node definition 

218 line, i_char = _get_next_line(trees, i_char) (empty)

219 values = line.split() (empty)

220 i_heap = int(values[0]) (empty)

221 var = int(values[1]) (empty)

222 split = int(values[2]) (empty)

223 leaf = float(values[3]) (empty)

224 

225 # update values 

226 leaf_trees[i_iter, i_tree, i_heap] = leaf (empty)

227 is_internal[i_heap // 2] = True (empty)

228 if i_heap < heap_size // 2: (empty)

229 var_trees[i_iter, i_tree, i_heap] = var (empty)

230 split_trees[i_iter, i_tree, i_heap] = split + 1 (empty)

231 

232 is_internal[0] = False (empty)

233 split_trees[i_iter, i_tree, ~is_internal] = 0 (empty)

234 

235 return MinimalTrace( (empty)

236 leaf_tree=jnp.array(leaf_trees), 

237 var_tree=jnp.array(var_trees), 

238 split_tree=jnp.array(split_trees), 

239 offset=jnp.float32(0.0 if offset is None else offset), 

240 leaf_unit=jnp.float32(1.0), 

241 ), meta