claimed-country-accuracy

#! /usr/bin/python3

import os
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'lib')))

import collections
import csv
import datetime
import functools
import itertools
import glob
import multiprocessing
import time

import numpy as np

import fiona
import fiona.crs
import pyproj
import shapely.geometry
import shapely.ops
import shapely.prepared

# For memory efficiency, this program works directly with the on-disk
# representation of Location savefiles.
import tables

WGS84proj = pyproj.Proj(proj="latlong", datum="WGS84", ellps="WGS84")
CEAproj = pyproj.Proj(proj="cea", ellps="WGS84", lon_0=0, lat_ts=0)
wgs_to_cea = functools.partial(pyproj.transform, WGS84proj, CEAproj)

Region = collections.namedtuple("Region",
                                ("serial", "name", "cc2",
                                 "rgn", "prep", "area"))


fake_iso_a2 = {
    "Ashmore and Cartier Is." : None,
    "N. Cyprus"               : "xn",
    "Indian Ocean Ter."       : None,
    "Siachen Glacier"         : None,
    "Kosovo"                  : "xk",
    "Somaliland"              : "xs"
}

regions = {}
regions_by_serial = []
def load_regions(shapefile):
    global regions
    with fiona.open(shapefile) as f_regions:
        f_proj = pyproj.Proj(f_regions.crs)
        to_wgs = functools.partial(pyproj.transform, f_proj, WGS84proj)
        serial = 0
        for r in f_regions:
            name = r['properties'].get('name', '')
            cc2 = r['properties'].get('iso_a2', '-99').lower()
            if cc2 == '-99':
                cc2 = fake_iso_a2[name]
                if cc2 is None:
                    continue

            rgn = shapely.ops.transform(to_wgs,
                shapely.geometry.shape(r['geometry']))
            prep = shapely.prepared.prep(rgn)
            area = shapely.ops.transform(wgs_to_cea, rgn).area
            regions[cc2] = Region(serial, name, cc2, rgn, prep, area)
            regions_by_serial.append(regions[cc2])
            serial += 1

# this saves memory relative to sparse.find
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 probability_each_region(loc, regions):
    pvec = np.zeros(len(regions))
    avec = np.zeros(len(regions))

    # see commentary in ageo.py::Location.area
    west = loc.attrs.west
    east = loc.attrs.west + loc.attrs.lon_spacing
    d_lat = loc.attrs.lat_spacing/2

    for row in loc.iterrows():
        v = row['prob_mass']
        if v == 0: continue
        lon = row['longitude']
        lat = row['latitude']

        pt = shapely.geometry.Point(lon, lat)
        for r in regions.values():
            if r.prep.contains(pt):

                north = lat + d_lat
                south = lat - d_lat
                tile = shapely.ops.transform(wgs_to_cea,
                    shapely.geometry.box(west, south, east, north))
                avec[r.serial] += r.area
                pvec[r.serial] += v
                break

    ptot = pvec.sum()
    if ptot == 0: return pvec
    if ptot != 1: pvec /= ptot

    for i in range(len(avec)):
        if avec[i]:
            avec[i] /= regions_by_serial[i].area

    pvec *= avec
    pvec /= pvec.sum()
    return pvec

def crunch_location(lname):
    global regions, regions_by_serial

    with tables.open_file(lname, "r") as fp:
        loc = fp.root.location
        ann = loc.attrs.annotations

        client_id="{:.2f}_{:.2f}".format(ann['client_lat'],
                                         ann['client_lon'])
        batch_id = ann['id']
        provider = ann['proxy_provider']
        a_cc2    = ann['proxy_alleged_cc2']

        pvec = probability_each_region(loc, regions)

    itop5 = np.argsort(pvec)[:-6:-1] # this gets the last five
                                     # elements in reverse order
    if a_cc2 in regions:
        itrue = regions[a_cc2].serial
        ptrue = pvec[itrue]
        if any(itop5 == itrue):
            pass
        else:
            itop5 = np.append(itop5, [itrue])
    else:
        ptrue = "<lacuna>"

    ptop5 = pvec[itop5]
    rtop5 = [regions_by_serial[i] for i in itop5]

    return client_id, a_cc2, ptrue, [
        (client_id, batch_id, provider, a_cc2, rtop5[i].cc2, ptop5[i])
        for i in range(len(rtop5))
        if ptop5[i] > 0
    ]

_time_0 = time.monotonic()
def progress(message, *args):
    global _time_0
    sys.stderr.write(
        ("{}: " + message + "\n").format(
            datetime.timedelta(seconds = time.monotonic() - _time_0),
            *args))

def main():
    progress("preparing")
    load_regions(sys.argv[1])
    progress("done preparing")
    with multiprocessing.Pool() as pool, \
         sys.stdout as ofp:
        wr = csv.writer(ofp)
        wr.writerow(("client", "batch", "provider", "a.cc", "l.cc", "prob"))

        todo = sorted(glob.glob(os.path.join(sys.argv[2], "*.h5")),
                      key = lambda f: (os.stat(f).st_size, f))
        n_todo = len(todo)
        n = 0
        for ci, ac, pt, result in pool.imap_unordered(crunch_location, todo):
            for row in result:
                wr.writerow(row)
            ofp.flush()

            n += 1
            progress("{}/{}: {} -> {}/{}", n, n_todo, ci, ac, pt)
    progress("done")

main()