def test(self):
all_sites = numpy.arange(27, dtype=numpy.uint32)
expected = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 19, 20, 21, 22, 23,
24, 25, 26]
# bounding box taking 20 of the 27 sites
indices = utils.within(
[176.73699, -39, -176.9016, -12], self.index)
numpy.testing.assert_equal(indices, expected)
# one can invert lon1, lon2
indices = utils.within(
[-176.9016, -39, 176.73699, -12], self.index)
numpy.testing.assert_equal(indices, expected)
# bounding box large enough to contain all sites
indices = utils.within(
[174.12916, -39, -169.9217, -12], self.index)
numpy.testing.assert_equal(indices, all_sites)
# lat on the edge from below is (strangely) taken
indices = utils.within(
[174.12916, -39, -169.9217, -30], self.index)
numpy.testing.assert_equal(indices, all_sites)
# lat on the edge from up is discarded
indices = utils.within(
[174.12916, -30, -169.9217, -12], self.index)
numpy.testing.assert_equal(indices, [])
def close_sids(self, rec, trt, mag):
"""
:param rec:
a record with fields minlon, minlat, maxlon, maxlat
:param trt:
tectonic region type string
:param mag:
magnitude
:returns:
the site indices within the bounding box enlarged by the integration
distance for the given TRT and magnitude
"""
if self.sitecol is None:
return []
elif not self.integration_distance: # do not filter
return self.sitecol.sids
if hasattr(rec, 'dtype'):
bbox = rec['minlon'], rec['minlat'], rec['maxlon'], rec['maxlat']
else:
bbox = rec # assume it is a 4-tuple
maxdist = self.integration_distance(trt, mag)
a1 = min(maxdist * KM_TO_DEGREES, 90)
a2 = min(angular_distance(maxdist, bbox[1], bbox[3]), 180)
bb = bbox[0] - a2, bbox[1] - a1, bbox[2] + a2, bbox[3] + a1
if hasattr(self, 'index'): # RtreeFilter
return within(bb, self.index)
return self.sitecol.within_bbox(bb)
def test(self):
if rtree is None: # not installed
raise unittest.SkipTest
self.lons = [
-179.75, -179.5, -179.25, -179.0, -178.75, -178.5, -178.25, -178.0,
-177.75, -177.5, -177.25, -177.0, -176.75, -176.5, -176.25, -176.0,
-175.75, -175.5, -175.25, 178.25, 178.5, 178.75, 179.0, 179.25,
179.5, 179.75, -180.0]
self.lats = [-30.5] * 27
self.index = rtree.index.Index(
(i, (x, y, x, y), None)
for i, (x, y) in enumerate(zip(self.lons, self.lats)))
all_sites = numpy.arange(27, dtype=numpy.uint32)
expected = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 19, 20, 21, 22, 23,
24, 25, 26]
# bounding box taking 20 of the 27 sites
indices = utils.within(
[176.73699, -39, -176.9016, -12], self.index)
numpy.testing.assert_equal(indices, expected)
# one can invert lon1, lon2
indices = utils.within(
[-176.9016, -39, 176.73699, -12], self.index)
numpy.testing.assert_equal(indices, expected)
# bounding box large enough to contain all sites
indices = utils.within(
[174.12916, -39, -169.9217, -12], self.index)
numpy.testing.assert_equal(indices, all_sites)
# lat on the edge from below is (strangely) taken
indices = utils.within(
[174.12916, -39, -169.9217, -30], self.index)
numpy.testing.assert_equal(indices, all_sites)
# lat on the edge from up is discarded
indices = utils.within(
[174.12916, -30, -169.9217, -12], self.index)
numpy.testing.assert_equal(indices, [])
def filter(self, sources):
"""
:param sources: a sequence of sources
:yields: rtree-filtered sources
"""
if self.sitecol is None: # do not filter
yield from sources
return
for src in sources:
box = self.integration_distance.get_affected_box(src)
indices = within(box, self.index)
if len(indices):
src.indices = indices
yield src
def filter(self, sources):
"""
:param sources: a sequence of sources
:yields: rtree-filtered sources
"""
index = rtree.index.Index(self.indexpath)
try:
for src in sources:
box = self.integration_distance.get_affected_box(src)
indices = within(box, index)
if len(indices):
src.indices = indices
yield src
finally:
index.close()
def __call__(self, sources, sites=None):
if sites is None:
sites = self.sitecol
for src in sources:
if hasattr(src, 'indices'): # already filtered
yield src, sites.filtered(src.indices)
elif self.prefilter == 'no': # do not filter
yield src, sites
elif self.prefilter == 'rtree':
indices = within(self.get_affected_box(src), self.index)
if len(indices):
src.indices = indices
yield src, sites.filtered(src.indices)
elif self.prefilter == 'numpy':
s_sites = sites.within_bbox(self.get_affected_box(src))
if s_sites is not None:
src.indices = get_indices(s_sites)
yield src, s_sites