in lucene/spatial-extras/src/java/org/apache/lucene/spatial/prefix/HeatmapFacetCounter.java [94:289]
public static Heatmap calcFacets(
PrefixTreeStrategy strategy,
IndexReaderContext context,
Bits topAcceptDocs,
Shape inputShape,
final int facetLevel,
int maxCells)
throws IOException {
if (maxCells > (MAX_ROWS_OR_COLUMNS * MAX_ROWS_OR_COLUMNS)) {
throw new IllegalArgumentException(
"maxCells (" + maxCells + ") should be <= " + MAX_ROWS_OR_COLUMNS);
}
if (inputShape == null) {
inputShape = strategy.getSpatialContext().getWorldBounds();
}
final Rectangle inputRect = inputShape.getBoundingBox();
// First get the rect of the cell at the bottom-left at depth facetLevel
final SpatialPrefixTree grid = strategy.getGrid();
final SpatialContext ctx = grid.getSpatialContext();
final Point cornerPt = ctx.getShapeFactory().pointXY(inputRect.getMinX(), inputRect.getMinY());
final CellIterator cellIterator = grid.getTreeCellIterator(cornerPt, facetLevel);
Cell cornerCell = null;
while (cellIterator.hasNext()) {
cornerCell = cellIterator.next();
}
assert cornerCell != null && cornerCell.getLevel() == facetLevel
: "Cell not at target level: " + cornerCell;
final Rectangle cornerRect = (Rectangle) cornerCell.getShape();
assert cornerRect.hasArea();
// Now calculate the number of columns and rows necessary to cover the inputRect
double heatMinX = cornerRect.getMinX(); // note: we might change this below...
final double cellWidth = cornerRect.getWidth();
final Rectangle worldRect = ctx.getWorldBounds();
final int columns =
calcRowsOrCols(
cellWidth, heatMinX, inputRect.getWidth(), inputRect.getMinX(), worldRect.getWidth());
final double heatMinY = cornerRect.getMinY();
final double cellHeight = cornerRect.getHeight();
final int rows =
calcRowsOrCols(
cellHeight,
heatMinY,
inputRect.getHeight(),
inputRect.getMinY(),
worldRect.getHeight());
assert rows > 0 && columns > 0;
if (columns > MAX_ROWS_OR_COLUMNS || rows > MAX_ROWS_OR_COLUMNS || columns * rows > maxCells) {
throw new IllegalArgumentException(
"Too many cells ("
+ columns
+ " x "
+ rows
+ ") for level "
+ facetLevel
+ " shape "
+ inputRect);
}
// Create resulting heatmap bounding rectangle & Heatmap object.
final double halfCellWidth = cellWidth / 2.0;
// if X world-wraps, use world bounds' range
if (columns * cellWidth + halfCellWidth > worldRect.getWidth()) {
heatMinX = worldRect.getMinX();
}
double heatMaxX = heatMinX + columns * cellWidth;
if (Math.abs(heatMaxX - worldRect.getMaxX()) < halfCellWidth) { // numeric conditioning issue
heatMaxX = worldRect.getMaxX();
} else if (heatMaxX > worldRect.getMaxX()) { // wraps dateline (won't happen if !geo)
heatMaxX = heatMaxX - worldRect.getMaxX() + worldRect.getMinX();
}
final double halfCellHeight = cellHeight / 2.0;
double heatMaxY = heatMinY + rows * cellHeight;
if (Math.abs(heatMaxY - worldRect.getMaxY()) < halfCellHeight) { // numeric conditioning issue
heatMaxY = worldRect.getMaxY();
}
final Heatmap heatmap =
new Heatmap(
columns, rows, ctx.getShapeFactory().rect(heatMinX, heatMaxX, heatMinY, heatMaxY));
if (topAcceptDocs instanceof Bits.MatchNoBits) {
return heatmap; // short-circuit
}
// All ancestor cell counts (of facetLevel) will be captured during facet visiting and applied
// later. If the data is
// just points then there won't be any ancestors.
// Facet count of ancestors covering all of the heatmap:
// single-element array so it can be accumulated in the inner class
int[] allCellsAncestorCount = new int[1];
// All other ancestors:
Map<Rectangle, Integer> ancestors = new HashMap<>();
// Now lets count some facets!
PrefixTreeFacetCounter.compute(
strategy,
context,
topAcceptDocs,
inputShape,
facetLevel,
new PrefixTreeFacetCounter.FacetVisitor() {
@Override
public void visit(Cell cell, int count) {
final double heatMinX = heatmap.region.getMinX();
final Rectangle rect = (Rectangle) cell.getShape();
if (cell.getLevel() == facetLevel) { // heatmap level; count it directly
// convert to col & row
int column;
if (rect.getMinX() >= heatMinX) {
column = (int) Math.round((rect.getMinX() - heatMinX) / cellWidth);
} else { // due to dateline wrap
column = (int) Math.round((rect.getMinX() + 360 - heatMinX) / cellWidth);
}
int row = (int) Math.round((rect.getMinY() - heatMinY) / cellHeight);
// note: unfortunately, it's possible for us to visit adjacent cells to the heatmap
// (if the SpatialPrefixTree
// allows adjacent cells to overlap on the seam), so we need to skip them
if (column < 0 || column >= heatmap.columns || row < 0 || row >= heatmap.rows) {
return;
}
// increment
heatmap.counts[column * heatmap.rows + row] += count;
} else if (rect.relate(heatmap.region) == SpatialRelation.CONTAINS) {
allCellsAncestorCount[0] += count;
} else { // ancestor
// note: not particularly efficient (possible put twice, and Integer wrapper); oh well
Integer existingCount = ancestors.put(rect, count);
if (existingCount != null) {
ancestors.put(rect, count + existingCount);
}
}
}
});
// Update the heatmap counts with ancestor counts
// Apply allCellsAncestorCount
if (allCellsAncestorCount[0] > 0) {
for (int i = 0; i < heatmap.counts.length; i++) {
heatmap.counts[i] += allCellsAncestorCount[0];
}
}
// Apply ancestors
// note: This approach isn't optimized for a ton of ancestor cells. We'll potentially increment
// the same cells
// multiple times in separate passes if any ancestors overlap. IF this poses a problem, we
// could optimize it
// with additional complication by keeping track of intervals in a sorted tree structure
// (possible TreeMap/Set)
// and iterate them cleverly such that we just make one pass at this stage.
int[] pair = new int[2]; // output of intersectInterval
for (Map.Entry<Rectangle, Integer> entry : ancestors.entrySet()) {
Rectangle rect = entry.getKey(); // from a cell (thus doesn't cross DL)
final int count = entry.getValue();
// note: we approach this in a way that eliminates int overflow/underflow (think huge cell,
// tiny heatmap)
intersectInterval(heatMinY, heatMaxY, cellHeight, rows, rect.getMinY(), rect.getMaxY(), pair);
final int startRow = pair[0];
final int endRow = pair[1];
if (!heatmap.region.getCrossesDateLine()) {
intersectInterval(
heatMinX, heatMaxX, cellWidth, columns, rect.getMinX(), rect.getMaxX(), pair);
final int startCol = pair[0];
final int endCol = pair[1];
incrementRange(heatmap, startCol, endCol, startRow, endRow, count);
} else {
// note: the cell rect might intersect 2 disjoint parts of the heatmap, so we do the left &
// right separately
final int leftColumns = (int) Math.round((180 - heatMinX) / cellWidth);
final int rightColumns = heatmap.columns - leftColumns;
// left half of dateline:
if (rect.getMaxX() > heatMinX) {
intersectInterval(
heatMinX, 180, cellWidth, leftColumns, rect.getMinX(), rect.getMaxX(), pair);
final int startCol = pair[0];
final int endCol = pair[1];
incrementRange(heatmap, startCol, endCol, startRow, endRow, count);
}
// right half of dateline
if (rect.getMinX() < heatMaxX) {
intersectInterval(
-180, heatMaxX, cellWidth, rightColumns, rect.getMinX(), rect.getMaxX(), pair);
final int startCol = pair[0] + leftColumns;
final int endCol = pair[1] + leftColumns;
incrementRange(heatmap, startCol, endCol, startRow, endRow, count);
}
}
}
return heatmap;
}