/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.lucene.spatial.prefix;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import org.apache.lucene.index.IndexReaderContext;
import org.apache.lucene.spatial.prefix.tree.Cell;
import org.apache.lucene.spatial.prefix.tree.CellIterator;
import org.apache.lucene.spatial.prefix.tree.SpatialPrefixTree;
import org.apache.lucene.util.ArrayUtil;
import org.apache.lucene.util.Bits;
import org.locationtech.spatial4j.context.SpatialContext;
import org.locationtech.spatial4j.shape.Point;
import org.locationtech.spatial4j.shape.Rectangle;
import org.locationtech.spatial4j.shape.Shape;
import org.locationtech.spatial4j.shape.SpatialRelation;
/**
* Computes spatial facets in two dimensions as a grid of numbers. The data is often visualized as a so-called
* "heatmap", hence the name.
*
* @lucene.experimental
*/
public class HeatmapFacetCounter {
//TODO where should this code live? It could go to PrefixTreeFacetCounter, or maybe here in its own class is fine.
/** Maximum number of supported rows (or columns). */
public static final int MAX_ROWS_OR_COLUMNS = (int) Math.sqrt(ArrayUtil.MAX_ARRAY_LENGTH);
static {
Math.multiplyExact(MAX_ROWS_OR_COLUMNS, MAX_ROWS_OR_COLUMNS);//will throw if doesn't stay within integer
}
/** Response structure */
public static class Heatmap {
public final int columns;
public final int rows;
public final int[] counts;//in order of 1st column (all rows) then 2nd column (all rows) etc.
public final Rectangle region;
public Heatmap(int columns, int rows, Rectangle region) {
this.columns = columns;
this.rows = rows;
this.counts = new int[columns * rows];
this.region = region;
}
public int getCount(int x, int y) {
return counts[x * rows + y];
}
@Override
public String toString() {
return "Heatmap{" + columns + "x" + rows + " " + region + '}';
}
}
/**
* Calculates spatial 2D facets (aggregated counts) in a grid, sometimes called a heatmap.
* Facet computation is implemented by navigating the underlying indexed terms efficiently. If you don't know exactly
* what facetLevel to go to for a given input box but you have some sense of how many cells there should be relative
* to the size of the shape, then consider using the logic that {@link org.apache.lucene.spatial.prefix.PrefixTreeStrategy}
* uses when approximating what level to go to when indexing a shape given a distErrPct.
*
* @param context the IndexReader's context
* @param topAcceptDocs a Bits to limit counted docs. If null, live docs are counted.
* @param inputShape the shape to gather grid squares for; typically a {@link Rectangle}.
* The <em>actual</em> heatmap area will usually be larger since the cells on the edge that overlap
* are returned. We always return a rectangle of integers even if the inputShape isn't a rectangle
* -- the non-intersecting cells will all be 0.
* If null is given, the entire world is assumed.
* @param facetLevel the target depth (detail) of cells.
* @param maxCells the maximum number of cells to return. If the cells exceed this count, an
*/
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.makePoint(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.makeRectangle(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:
int[] allCellsAncestorCount = new int[1]; // single-element array so it can be accumulated in the inner class
//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) {//containing ancestor
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;
}
private static void intersectInterval(double heatMin, double heatMax, double heatCellLen, int numCells,
double cellMin, double cellMax,
int[] out) {
assert heatMin < heatMax && cellMin < cellMax;
//precondition: we know there's an intersection
if (heatMin >= cellMin) {
out[0] = 0;
} else {
out[0] = (int) Math.round((cellMin - heatMin) / heatCellLen);
}
if (heatMax <= cellMax) {
out[1] = numCells - 1;
} else {
out[1] = (int) Math.round((cellMax - heatMin) / heatCellLen) - 1;
}
}
private static void incrementRange(Heatmap heatmap, int startColumn, int endColumn, int startRow, int endRow,
int count) {
//startColumn & startRow are not necessarily within the heatmap range; likewise numRows/columns may overlap.
if (startColumn < 0) {
endColumn += startColumn;
startColumn = 0;
}
endColumn = Math.min(heatmap.columns-1, endColumn);
if (startRow < 0) {
endRow += startRow;
startRow = 0;
}
endRow = Math.min(heatmap.rows-1, endRow);
if (startRow > endRow) {
return;//short-circuit
}
for (int c = startColumn; c <= endColumn; c++) {
int cBase = c * heatmap.rows;
for (int r = startRow; r <= endRow; r++) {
heatmap.counts[cBase + r] += count;
}
}
}
/** Computes the number of intervals (rows or columns) to cover a range given the sizes. */
private static int calcRowsOrCols(double cellRange, double cellMin, double requestRange, double requestMin,
double worldRange) {
assert requestMin >= cellMin;
//Idealistically this wouldn't be so complicated but we concern ourselves with overflow and edge cases
double range = (requestRange + (requestMin - cellMin));
if (range == 0) {
return 1;
}
final double intervals = Math.ceil(range / cellRange);
if (intervals > Integer.MAX_VALUE) {
return Integer.MAX_VALUE;//should result in an error soon (exceed thresholds)
}
// ensures we don't have more intervals than world bounds (possibly due to rounding/edge issue)
final long intervalsMax = Math.round(worldRange / cellRange);
if (intervalsMax > Integer.MAX_VALUE) {
//just return intervals
return (int) intervals;
}
return Math.min((int)intervalsMax, (int)intervals);
}
private HeatmapFacetCounter() {
}
}