Coverage for mesh_api/rtree_indexes.py: 18%

122 statements  

« prev     ^ index     » next       coverage.py v7.7.0, created at 2025-03-20 20:51 +0100

1import logging 

2from functools import wraps 

3 

4import numpy as np 

5from rtree import index 

6 

7RTREE_COORDS_DEC = 2 

8 

9 

10def check_coords(*argnames): 

11 def _check_coords(fn): 

12 @wraps(fn) 

13 def wrapper(*args, **kwargs): 

14 for argname in argnames: 

15 kwargs[argname] = list( 

16 np.round(kwargs[argname], decimals=RTREE_COORDS_DEC) 

17 ) 

18 return fn(*args, **kwargs) 

19 

20 return wrapper 

21 

22 return _check_coords 

23 

24 

25def list_content(rtree): 

26 for item in list(rtree.intersection(rtree.bounds, objects=True)): 

27 print(item) 

28 

29 

30class RTrees: 

31 """ 

32 https://rtree.readthedocs.io/en/latest/tutorial.html 

33 """ 

34 

35 def __init__(self, registry, index_elems=("0d", "1d", "2d", "3d")): 

36 p = index.Property() 

37 p.dimension = 3 

38 self._rtree_properties = p 

39 # p.dat_extension = 'data' 

40 # p.idx_extension = 'index' 

41 self._ixs = {"GRID": index.Index(properties=self._rtree_properties)} 

42 self.index_elems = index_elems 

43 self.reg = registry 

44 

45 def list(self, cardname=None, rtree=None): 

46 if rtree is None: 

47 rtree = self._ixs[cardname] 

48 return tuple(rtree.intersection(rtree.bounds, objects=True)) 

49 

50 def create_rtree(self, eids=None, gids=None): 

51 """create and return 'on-the-fly' a new rtree""" 

52 if eids is not None and gids is not None: 

53 raise ValueError("both `eids` and `gids` cannot be set") 

54 _rtree = index.Index(properties=self._rtree_properties) 

55 ids = eids if eids is not None else gids 

56 if gids is not None: 

57 src_rtrees = (self._ixs["GRID"],) 

58 else: 

59 src_rtrees = [] 

60 for cardname, rtree in self._ixs.items(): 

61 if cardname == "GRID": 

62 continue 

63 src_rtrees.append(rtree) 

64 for rtree in src_rtrees: 

65 for item in list(rtree.intersection(rtree.bounds, objects=True)): 

66 if item.id in ids: 

67 _rtree.insert(item.id, item.bbox) 

68 return _rtree 

69 

70 # check_coords("coords") 

71 def add_grid(self, gid, /, coords, rtree=None): 

72 """insert gid in RTREE index""" 

73 self._ixs["GRID"].insert(int(gid), coords) 

74 if rtree: 

75 rtree.insert(int(gid), coords) 

76 

77 # check_coords("coords") 

78 def add_elem(self, cardname, eid, /, coords, rtree=None): 

79 try: 

80 # if a pandas.Series is passed 

81 coords = coords.values 

82 except AttributeError: 

83 pass 

84 self._ixs[cardname].insert(int(eid), coords) 

85 if rtree: 

86 rtree.insert(int(eid), coords) 

87 

88 # check_coords("old_coords", "coords") 

89 def _update(self, cardname, id, /, old_coords, coords, rtree=None): 

90 rt = self._ixs.get(cardname) 

91 if rt is None and rtree is None: 

92 return 

93 rt.delete(id, old_coords) 

94 rt.insert(id, coords) 

95 if rtree: 

96 rtree.delete(id, old_coords) 

97 rtree.insert(id, coords) 

98 

99 def get_bbox(self, cardname, id): 

100 ix = self._ixs[cardname] 

101 for f in ix.intersection(ix.bounds, objects=True): 

102 if f.id == id: 

103 return f.bbox 

104 

105 def get_gid_bbox(self, gid): 

106 return self.get_bbox(cardname="GRID", id=gid) 

107 ix = self._ixs["GRID"] 

108 for f in ix.intersection(ix.bounds, objects=True): 

109 if f.id == id: 

110 return f.bbox 

111 

112 def get_eid_bbox(self, eid, cardname=None): 

113 if cardname: 

114 return self.get_bbox(cardname=cardname, id=eid) 

115 for cardname, ix in self._ixs.items(): 

116 if cardname == "GRID": 

117 continue 

118 bbox = self.get_bbox(cardname=cardname, id=eid) 

119 if bbox is not None: 

120 return bbox 

121 

122 def update_grid(self, gid, /, old_coords, coords, rtree=None): 

123 return self._update( 

124 "GRID", gid, old_coords=old_coords, coords=coords, rtree=rtree 

125 ) 

126 

127 def update_element(self, eid, /, old_coords, coords, rtree=None): 

128 # old_coords is taken from 

129 try: 

130 cardname = self.reg.mesh.eid2card(skipcache=True)[eid] 

131 except: 

132 # breakpoint() 

133 raise 

134 return self._update( 

135 cardname, eid, old_coords=old_coords, coords=coords, rtree=rtree 

136 ) 

137 

138 def build(self): 

139 # add all grids to objects 

140 grids = self.reg.container["bulk"]["GRID"].array[["ID", "X1", "X2", "X3"]] 

141 grids = np.lib.recfunctions.structured_to_unstructured(grids) 

142 for row in grids: 

143 self.add_grid(row[0], coords=row[1:]) 

144 # =============================================== 

145 # add elements 

146 eid2bbox = self.reg.mesh.eid2bbox() 

147 cards = self.reg.container["summary"]["element"] 

148 for dim in self.index_elems: 

149 cards = self.reg.container["summary"][dim] 

150 for cardname in cards: 

151 card = self.reg.container["bulk"][cardname] 

152 self._ixs[cardname] = index.Index(properties=self._rtree_properties) 

153 for i, (eid, gids) in enumerate(card.eid2gids().items()): 

154 bbox = eid2bbox.loc[eid] 

155 self.add_elem(cardname, eid, coords=bbox) 

156 

157 def nearest_grids(self, /, coords, num_results=1, astype=set, rtree=None): 

158 if not rtree: 

159 rtree = self._ixs["GRID"] 

160 return astype(rtree.nearest(coords, num_results)) 

161 

162 def nearest_elements( 

163 self, coords, num_results=1, astype=set, cardname=None, rtree=None 

164 ): 

165 if rtree: 

166 return astype(rtree.nearest(coords, num_results)) 

167 if cardname: 

168 ix = self._ixs[cardname] 

169 return astype(ix.nearest(coords, num_results)) 

170 # get all ids 

171 ids = {} 

172 for cardname, ix in self._ixs.items(): 

173 if cardname == "GRID": 

174 continue 

175 ids[cardname] = astype(ix.nearest(coords, num_results)) 

176 return ids