<commit_before>"""ageo - active geolocation library: core.
"""

__all__ = ('Location', 'Map', 'Observation')

import bisect
import functools
import itertools
import numpy as np
import pyproj
from scipy import sparse
from shapely.geometry import Point, MultiPoint, Polygon, box as Box
from shapely.ops import transform as sh_transform
import tables
import math
import sys

# scipy.sparse.find() materializes vectors which, in several cases
# below, can be enormous.  This is slower, but more memory-efficient.
# Code from https://stackoverflow.com/a/31244368/388520 with minor
# modifications.
def iter_csr_nonzero(matrix):
    irepeat = itertools.repeat
    return zip(
        # reconstruct the row indices
        itertools.chain.from_iterable(
            irepeat(i, r)
            for (i,r) in enumerate(matrix.indptr[1:] - matrix.indptr[:-1])
        ),
        # matrix.indices gives the column indices as-is
        matrix.indices,
        matrix.data
    )

def Disk(x, y, radius):
    return Point(x, y).buffer(radius)

# Important note: pyproj consistently takes coordinates in lon/lat
# order and distances in meters.  lon/lat order makes sense for
# probability matrices, because longitudes are horizontal = columns,
# latitudes are vertical = rows, and scipy matrices are column-major
# (blech).  Therefore, this library also consistently uses lon/lat
# order and meters.

# Coordinate transformations used by Location.centroid()
wgs_proj  = pyproj.Proj("+proj=longlat +ellps=WGS84 +datum=WGS84 +no_defs")
gcen_proj = pyproj.Proj("+proj=geocent +datum=WGS84 +units=m +no_defs")
wgs_to_gcen = functools.partial(pyproj.transform, wgs_proj, gcen_proj)
gcen_to_wgs = functools.partial(pyproj.transform, gcen_proj, wgs_proj)

# ... and Location.area()
cea_proj   = pyproj.Proj(proj='cea', ellps='WGS84', lon_0=0, lat_ts=0)
wgs_to_cea = functools.partial(pyproj.transform, wgs_proj, cea_proj)

# Smooth over warts in pyproj.Geod.inv(), which is vectorized
# internally, but does not support numpy-style broadcasting, and
# returns things we don't need.  The prebound _Inv and _Bcast are
# strictly performance hacks.
_WGS84geod = pyproj.Geod(ellps='WGS84')
def WGS84dist(lon1, lat1, lon2, lat2, *,
              _Inv = _WGS84geod.inv, _Bcast = np.broadcast_arrays):
    _, _, dist = _Inv(*_Bcast(lon1, lat1, lon2, lat2))
    return dist

def cartesian2(a, b):
    """Cartesian product of two 1D vectors A and B."""
    return np.tile(a, len(b)), np.repeat(b, len(a))

def mask_ij(bounds, longitudes, latitudes):
    """Given a rectangle-tuple BOUNDS (west, south, east, north; as
       returned by shapely .bounds properties), and sorted grid index
       vectors LONGITUDES, LATITUDES, return vectors I, J which give the
       x- and y-indices of every grid point within the rectangle.
       LATITUDES and LONGITUDES must be sorted.
    """
    try:
        (west, south, east, north) = bounds
    except ValueError as e:
        raise ValueError("invalid bounds argument {!r}".format(bounds)) from e

    min_i = bisect.bisect_left(longitudes, west)
    max_i = bisect.bisect_right(longitudes, east)
    min_j = bisect.bisect_left(latitudes, south)
    max_j = bisect.bisect_right(latitudes, north)

    I = np.array(range(min_i, max_i))
    J = np.array(range(min_j, max_j))
    return cartesian2(I, J)

def mask_matrix(bounds, longitudes, latitudes):
    """Construct a sparse matrix which is 1 at all latitude+longitude
       grid points inside the rectangle BOUNDS, 0 outside.
       LATITUDES and LONGITUDES must be sorted.
    """
    I, J = mask_ij(bounds, longitudes, latitudes)

    return sparse.csr_matrix((np.ones_like(I), (I, J)),
                             shape=(len(longitudes), len(latitudes)))

class LocationRowOnDisk(tables.IsDescription):
    """The row format of the pytables table used to save Location objects
       on disk.  See Location.save and Location.load."""
    grid_x    = tables.UInt32Col()
    grid_y    = tables.UInt32Col()
    longitude = tables.Float64Col()
    latitude  = tables.Float64Col()
    prob_mass = tables.Float32Col()

class Location:
    """An estimated location for a host.  This is represented by a
    probability mass function over the surface of the Earth, quantized
    to a cell grid, and stored as a sparse matrix.

    Properties:
      resolution  - Grid resolution, in meters at the equator
      lon_spacing - East-west (longitude) grid resolution, in decimal degrees
      lat_spacing - North-south (latitude) grid resolution, in decimal degrees
      fuzz        - Coastline uncertainty factor, in meters at the equator
      north       - Northernmost latitude covered by the grid
      south       - Southernmost latitude ditto
      east        - Easternmost longitude ditto
      west        - Westernmost longitude ditto
      latitudes   - Vector of latitude values corresponding to grid points
      longitudes  - Vector of longitude values ditto
      probability - Probability mass matrix (may be lazily computed)
      bounds      - Bounding region of the nonzero portion of the
                    probability mass matrix (may be lazily computed)
      centroid    - Centroid of the nonzero &c
      area        - Weighted area of the nonzero &c
      covariance  - Covariance matrix of the nonzero &c
                    (relative to the centroid)
      rep_pt      - "Representative point" of the nonzero &c; see docstring
                    for exactly what this means
      annotations - Dictionary of arbitrary additional metadata; saved and
                    loaded but not otherwise inspected by this code

    You will normally not construct bare Location objects directly, only
    Map and Observation objects (these are subclasses).  However, any two
    Locations can be _intersected_ to produce a new one.

    A Location is _vacuous_ if it has no nonzero entries in its
    probability matrix.

    """
    def __init__(self, *,
                 resolution, fuzz, lon_spacing, lat_spacing,
                 north, south, east, west,
                 longitudes, latitudes,
                 probability=None, vacuity=None, bounds=None,
                 centroid=None, covariance=None, rep_pt=None,
                 loaded_from=None, annotations=None
    ):
        self.resolution   = resolution
        self.fuzz         = fuzz
        self.north        = north
        self.south        = south
        self.east         = east
        self.west         = west
        self.lon_spacing  = lon_spacing
        self.lat_spacing  = lat_spacing
        self.longitudes   = longitudes
        self.latitudes    = latitudes
        self._probability = probability
        self._vacuous     = vacuity
        self._bounds      = bounds
        self._centroid    = centroid
        self._covariance  = covariance
        self._rep_pt      = rep_pt
        self._loaded_from = loaded_from
        self._area        = None
        self.annotations  = annotations if annotations is not None else {}

    @property
    def probability(self):
        if self._probability is None:
            self.compute_probability_matrix_now()
        return self._probability

    @property
    def vacuous(self):
        if self._vacuous is None:
            self.compute_probability_matrix_now()
        return self._vacuous

    @property
    def centroid(self):
        if self._centroid is None:
            self.compute_centroid_now()
        return self._centroid

    @property
    def covariance(self):
        if self._covariance is None:
            self.compute_centroid_now()
        return self._covariance

    @property
    def area(self):
        """Weighted area of the nonzero region of the probability matrix."""

        if self._area is None:
            # Notionally, each grid point should be treated as a
            # rectangle of parallels and meridians _centered_ on the
            # point.  The area of such a rectangle, however, only
            # depends on its latitude and its breadth; the actual
            # longitude values don't matter.  Since the grid is
            # equally spaced, we can use [0],[1] always, and then we
            # do not have to worry about crossing the discontinuity at ±180.
            west = self.longitudes[0]
            east = self.longitudes[1]
            # For latitude, the actual values do matter, but the map
            # never goes all the way to the poles, and the grid is
            # equally spaced, so we can precompute the north-south
            # delta from any pair of latitudes and not have to worry
            # about running off the ends of the array.
            d_lat = (self.latitudes[1] - self.latitudes[0]) / 2

            # We don't need X, so throw it away immediately.  (We
            # don't use iter_csr_nonzero here because we need to
            # modify V.)
            X, Y, V = sparse.find(self.probability); X = None

            # The value vector is supposed to be normalized, but make
            # sure it is, and then adjust from 1-overall to 1-per-cell
            # normalization.
            assert len(V.shape) == 1
            S = V.sum()
            if S == 0:
                return 0
            if S != 1:
                V /= S
            V *= V.shape[0]

            area = 0
            for y, v in zip(Y, V):
                north = self.latitudes[y] + d_lat
                south = self.latitudes[y] - d_lat
                if not (-90 <= south < north <= 90):
                    raise AssertionError("expected -90 <= {} < {} <= 90"
                                         .format(south, north))

                tile = sh_transform(wgs_to_cea, Box(west, south, east, north))
                area += v * tile.area

            self._area = area
        return self._area

    @property
    def rep_pt(self):
        """Representative point of the nonzero region of the probability
           matrix.  This is the point, within the nonzero region,
           whose _weighted_ squared geodesic distance to the centroid
           is minimized.
        """
        if self._rep_pt is None:
            lons = self.longitudes
            lats = self.latitudes
            cen = self.centroid
            aeqd_cen = pyproj.Proj(proj='aeqd', ellps='WGS84', datum='WGS84',
                                   lon_0=cen[0], lat_0=cen[1])
            wgs_to_aeqd = functools.partial(pyproj.transform,
                                            wgs_proj, aeqd_cen)
            # mathematically, wgs_to_aeqd(Point(lon, lat)) == Point(0, 0);
            # the latter is faster and more precise
            cen_pt = Point(0,0)

            # It is unacceptably costly to construct a shapely MultiPoint
            # out of some locations with large regions (requires more than
            # 32GB of scratch memory).  Instead, iterate over the points
            # one at a time.
            min_wdd = math.inf
            rep_pt = None
            for x, y, v in iter_csr_nonzero(self.probability):
                lon = lons[x]
                lat = lats[y]
                cell = sh_transform(wgs_to_aeqd, Point(lon, lat))
                dist = cen_pt.distance(cell)
                wdd = dist*dist*v
                if wdd < min_wdd:
                    min_wdd = wdd
                    rep_pt = [lon, lat]
            if rep_pt is None:
                rep_pt = cen
            else:
                rep_pt = np.array(rep_pt)
            self._rep_pt = rep_pt
        return self._rep_pt

    def distance_to_point(self, lon, lat):
        """Find the shortest geodesic distance from (lon, lat) to a nonzero
           cell of the probability matrix."""
        aeqd_pt = pyproj.Proj(proj='aeqd', ellps='WGS84', datum='WGS84',
                              lon_0=lon, lat_0=lat)
        wgs_to_aeqd = functools.partial(pyproj.transform, wgs_proj, aeqd_pt)
        # mathematically, wgs_to_aeqd(Point(lon, lat)) == Point(0, 0);
        # the latter is faster and more precise
        pt = Point(0, 0)

        # It is unacceptably costly to construct a shapely MultiPoint<commit_msg>Correct a thinko in the representative-point calculation.<commit_after>"""ageo - active geolocation library: core.
"""

__all__ = ('Location', 'Map', 'Observation')

import bisect
import functools
import itertools
import numpy as np
import pyproj
from scipy import sparse
from shapely.geometry import Point, MultiPoint, Polygon, box as Box
from shapely.ops import transform as sh_transform
import tables
import math
import sys

# scipy.sparse.find() materializes vectors which, in several cases
# below, can be enormous.  This is slower, but more memory-efficient.
# Code from https://stackoverflow.com/a/31244368/388520 with minor
# modifications.
def iter_csr_nonzero(matrix):
    irepeat = itertools.repeat
    return zip(
        # reconstruct the row indices
        itertools.chain.from_iterable(
            irepeat(i, r)
            for (i,r) in enumerate(matrix.indptr[1:] - matrix.indptr[:-1])
        ),
        # matrix.indices gives the column indices as-is
        matrix.indices,
        matrix.data
    )

def Disk(x, y, radius):
    return Point(x, y).buffer(radius)

# Important note: pyproj consistently takes coordinates in lon/lat
# order and distances in meters.  lon/lat order makes sense for
# probability matrices, because longitudes are horizontal = columns,
# latitudes are vertical = rows, and scipy matrices are column-major
# (blech).  Therefore, this library also consistently uses lon/lat
# order and meters.

# Coordinate transformations used by Location.centroid()
wgs_proj  = pyproj.Proj("+proj=longlat +ellps=WGS84 +datum=WGS84 +no_defs")
gcen_proj = pyproj.Proj("+proj=geocent +datum=WGS84 +units=m +no_defs")
wgs_to_gcen = functools.partial(pyproj.transform, wgs_proj, gcen_proj)
gcen_to_wgs = functools.partial(pyproj.transform, gcen_proj, wgs_proj)

# ... and Location.area()
cea_proj   = pyproj.Proj(proj='cea', ellps='WGS84', lon_0=0, lat_ts=0)
wgs_to_cea = functools.partial(pyproj.transform, wgs_proj, cea_proj)

# Smooth over warts in pyproj.Geod.inv(), which is vectorized
# internally, but does not support numpy-style broadcasting, and
# returns things we don't need.  The prebound _Inv and _Bcast are
# strictly performance hacks.
_WGS84geod = pyproj.Geod(ellps='WGS84')
def WGS84dist(lon1, lat1, lon2, lat2, *,
              _Inv = _WGS84geod.inv, _Bcast = np.broadcast_arrays):
    _, _, dist = _Inv(*_Bcast(lon1, lat1, lon2, lat2))
    return dist

def cartesian2(a, b):
    """Cartesian product of two 1D vectors A and B."""
    return np.tile(a, len(b)), np.repeat(b, len(a))

def mask_ij(bounds, longitudes, latitudes):
    """Given a rectangle-tuple BOUNDS (west, south, east, north; as
       returned by shapely .bounds properties), and sorted grid index
       vectors LONGITUDES, LATITUDES, return vectors I, J which give the
       x- and y-indices of every grid point within the rectangle.
       LATITUDES and LONGITUDES must be sorted.
    """
    try:
        (west, south, east, north) = bounds
    except ValueError as e:
        raise ValueError("invalid bounds argument {!r}".format(bounds)) from e

    min_i = bisect.bisect_left(longitudes, west)
    max_i = bisect.bisect_right(longitudes, east)
    min_j = bisect.bisect_left(latitudes, south)
    max_j = bisect.bisect_right(latitudes, north)

    I = np.array(range(min_i, max_i))
    J = np.array(range(min_j, max_j))
    return cartesian2(I, J)

def mask_matrix(bounds, longitudes, latitudes):
    """Construct a sparse matrix which is 1 at all latitude+longitude
       grid points inside the rectangle BOUNDS, 0 outside.
       LATITUDES and LONGITUDES must be sorted.
    """
    I, J = mask_ij(bounds, longitudes, latitudes)

    return sparse.csr_matrix((np.ones_like(I), (I, J)),
                             shape=(len(longitudes), len(latitudes)))

class LocationRowOnDisk(tables.IsDescription):
    """The row format of the pytables table used to save Location objects
       on disk.  See Location.save and Location.load."""
    grid_x    = tables.UInt32Col()
    grid_y    = tables.UInt32Col()
    longitude = tables.Float64Col()
    latitude  = tables.Float64Col()
    prob_mass = tables.Float32Col()

class Location:
    """An estimated location for a host.  This is represented by a
    probability mass function over the surface of the Earth, quantized
    to a cell grid, and stored as a sparse matrix.

    Properties:
      resolution  - Grid resolution, in meters at the equator
      lon_spacing - East-west (longitude) grid resolution, in decimal degrees
      lat_spacing - North-south (latitude) grid resolution, in decimal degrees
      fuzz        - Coastline uncertainty factor, in meters at the equator
      north       - Northernmost latitude covered by the grid
      south       - Southernmost latitude ditto
      east        - Easternmost longitude ditto
      west        - Westernmost longitude ditto
      latitudes   - Vector of latitude values corresponding to grid points
      longitudes  - Vector of longitude values ditto
      probability - Probability mass matrix (may be lazily computed)
      bounds      - Bounding region of the nonzero portion of the
                    probability mass matrix (may be lazily computed)
      centroid    - Centroid of the nonzero &c
      area        - Weighted area of the nonzero &c
      covariance  - Covariance matrix of the nonzero &c
                    (relative to the centroid)
      rep_pt      - "Representative point" of the nonzero &c; see docstring
                    for exactly what this means
      annotations - Dictionary of arbitrary additional metadata; saved and
                    loaded but not otherwise inspected by this code

    You will normally not construct bare Location objects directly, only
    Map and Observation objects (these are subclasses).  However, any two
    Locations can be _intersected_ to produce a new one.

    A Location is _vacuous_ if it has no nonzero entries in its
    probability matrix.

    """
    def __init__(self, *,
                 resolution, fuzz, lon_spacing, lat_spacing,
                 north, south, east, west,
                 longitudes, latitudes,
                 probability=None, vacuity=None, bounds=None,
                 centroid=None, covariance=None, rep_pt=None,
                 loaded_from=None, annotations=None
    ):
        self.resolution   = resolution
        self.fuzz         = fuzz
        self.north        = north
        self.south        = south
        self.east         = east
        self.west         = west
        self.lon_spacing  = lon_spacing
        self.lat_spacing  = lat_spacing
        self.longitudes   = longitudes
        self.latitudes    = latitudes
        self._probability = probability
        self._vacuous     = vacuity
        self._bounds      = bounds
        self._centroid    = centroid
        self._covariance  = covariance
        self._rep_pt      = rep_pt
        self._loaded_from = loaded_from
        self._area        = None
        self.annotations  = annotations if annotations is not None else {}

    @property
    def probability(self):
        if self._probability is None:
            self.compute_probability_matrix_now()
        return self._probability

    @property
    def vacuous(self):
        if self._vacuous is None:
            self.compute_probability_matrix_now()
        return self._vacuous

    @property
    def centroid(self):
        if self._centroid is None:
            self.compute_centroid_now()
        return self._centroid

    @property
    def covariance(self):
        if self._covariance is None:
            self.compute_centroid_now()
        return self._covariance

    @property
    def area(self):
        """Weighted area of the nonzero region of the probability matrix."""

        if self._area is None:
            # Notionally, each grid point should be treated as a
            # rectangle of parallels and meridians _centered_ on the
            # point.  The area of such a rectangle, however, only
            # depends on its latitude and its breadth; the actual
            # longitude values don't matter.  Since the grid is
            # equally spaced, we can use [0],[1] always, and then we
            # do not have to worry about crossing the discontinuity at ±180.
            west = self.longitudes[0]
            east = self.longitudes[1]
            # For latitude, the actual values do matter, but the map
            # never goes all the way to the poles, and the grid is
            # equally spaced, so we can precompute the north-south
            # delta from any pair of latitudes and not have to worry
            # about running off the ends of the array.
            d_lat = (self.latitudes[1] - self.latitudes[0]) / 2

            # We don't need X, so throw it away immediately.  (We
            # don't use iter_csr_nonzero here because we need to
            # modify V.)
            X, Y, V = sparse.find(self.probability); X = None

            # The value vector is supposed to be normalized, but make
            # sure it is, and then adjust from 1-overall to 1-per-cell
            # normalization.
            assert len(V.shape) == 1
            S = V.sum()
            if S == 0:
                return 0
            if S != 1:
                V /= S
            V *= V.shape[0]

            area = 0
            for y, v in zip(Y, V):
                north = self.latitudes[y] + d_lat
                south = self.latitudes[y] - d_lat
                if not (-90 <= south < north <= 90):
                    raise AssertionError("expected -90 <= {} < {} <= 90"
                                         .format(south, north))

                tile = sh_transform(wgs_to_cea, Box(west, south, east, north))
                area += v * tile.area

            self._area = area
        return self._area

    @property
    def rep_pt(self):
        """Representative point of the nonzero region of the probability
           matrix.  This is the point, within the nonzero region,
           whose _weighted_ squared geodesic distance to the centroid
           is minimized.
        """
        if self._rep_pt is None:
            lons = self.longitudes
            lats = self.latitudes
            cen = self.centroid
            aeqd_cen = pyproj.Proj(proj='aeqd', ellps='WGS84', datum='WGS84',
                                   lon_0=cen[0], lat_0=cen[1])
            wgs_to_aeqd = functools.partial(pyproj.transform,
                                            wgs_proj, aeqd_cen)
            # mathematically, wgs_to_aeqd(Point(lon, lat)) == Point(0, 0);
            # the latter is faster and more precise
            cen_pt = Point(0,0)

            # It is unacceptably costly to construct a shapely MultiPoint
            # out of some locations with large regions (requires more than
            # 32GB of scratch memory).  Instead, iterate over the points
            # one at a time.
            min_wdd = math.inf
            rep_pt = None
            for x, y, v in iter_csr_nonzero(self.probability):
                lon = lons[x]
                lat = lats[y]
                cell = sh_transform(wgs_to_aeqd, Point(lon, lat))
                dist = cen_pt.distance(cell)
                # we want the minimum distance but the maximum probability
                wdd = dist*dist*(1-v)
                if wdd < min_wdd:
                    min_wdd = wdd
                    rep_pt = [lon, lat]
            if rep_pt is None:
                rep_pt = cen
            else:
                rep_pt = np.array(rep_pt)
            self._rep_pt = rep_pt
        return self._rep_pt

    def distance_to_point(self, lon, lat):
        """Find the shortest geodesic distance from (lon, lat) to a nonzero
           cell of the probability matrix."""
        aeqd_pt = pyproj.Proj(proj='aeqd', ellps='WGS84', datum='WGS84',
                              lon_0=lon, lat_0=lat)
        wgs_to_aeqd = functools.partial(pyproj.transform, wgs_proj, aeqd_pt)
        # mathematically, wgs_to_aeqd(Point(lon, lat)) == Point(0, 0);
        # the latter is faster and more precise
        pt = Point(0, 0)

        # It is unacceptably costly to construct a shapely MultiPoint
{'fraction_non_alphanumeric': 0.07866694095865048, 'fraction_numerical': 0.00892820407323596, 'mean_word_length': 2.141933815925543, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3550796', 'n_tokens_mistral': 7261, 'n_tokens_neox': 6802, 'n_words': 2338}