Coverage for mesh_api/rtree_indexes.py: 18%

1 import logging

2 from functools import wraps

3

4 import numpy as np

5 from rtree import index

6

7 RTREE_COORDS_DEC = 2

8

9

10 def 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

25 def list_content(rtree):

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

27 print(item)

28

29

30 class 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