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
« prev ^ index » next coverage.py v7.7.0, created at 2025-03-20 20:51 +0100
1import logging
2from functools import wraps
4import numpy as np
5from rtree import index
7RTREE_COORDS_DEC = 2
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)
20 return wrapper
22 return _check_coords
25def list_content(rtree):
26 for item in list(rtree.intersection(rtree.bounds, objects=True)):
27 print(item)
30class RTrees:
31 """
32 https://rtree.readthedocs.io/en/latest/tutorial.html
33 """
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
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))
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
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)
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)
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)
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
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
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
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 )
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 )
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)
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))
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