# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import geopandas as gpd
import pandas as pd
import numpy as np
from geopy.distance import geodesic
import math

import gurobipy as gp
from gurobipy import GRB
# -

zones = gpd.read_file('taxi_zones/taxi_zones.shp')
table = pd.read_csv('taxi+_zone_lookup.csv')
zones = zones.to_crs("EPSG:4326")


# +
def dist(pt1, pt2):
    return geodesic(pt1, pt2).miles

def optql(x_list, pi_list, dQ, epsilon=0.5):
    '''
    input: x_list: list of geographic coordinates
           pi_list: probability distribution for x_list
           dQ: distance metric
           epsilon: desired privacy level
    
    output: matrix: stochastic transition matrix
            pre_prob: normalized pre-process probability distribution
            post_prob: post-process probability distribution
    '''
    
    pre_prob = np.array(pi_list) / sum(pi_list) # normalize probability distribution
    n = len(x_list) # get number of elements
    threshold = math.exp(epsilon)
    
    # define a model
    model = gp.Model('OptQL')
    
    # add variables accessed as (0, 0), (0, 1), (1, 1), ...
    variables = model.addVars(n, n, lb=0.0, ub=1.0, name='k')
    
    # set objective function
    model.setObjective(gp.quicksum(pre_prob[i] * variables[i, j] * dQ(x_list[i], x_list[j]) \
                                   for i in range(n) for j in range(n)), GRB.MINIMIZE)
    
    # add constraints (1)
    model.addConstrs(variables[i, k] <= pow(threshold, dQ(x_list[i], x_list[j])) * variables[j, k] \
                     for i in range(n) for j in range(n) for k in range(n))
    
    # add constraints (2)
    model.addConstrs(gp.quicksum(variables.select(i, '*')) == 1 for i in range(n))
    
    # constriants (3) are already satisfied
    
    # optimize the model
    model.optimize()
    
    # build a matrix to store the stochastic matrix
    variables = model.getAttr('x', variables)
    matrix = np.zeros((n, n))
    for key, value in variables.items():
        matrix[key] = value
    
    # get post-process probability distribution
    post_prob = pre_prob @ matrix
    
    
    return matrix, pre_prob, post_prob

# +
zones['Centroid'] = zones['geometry'].apply(lambda x: x.centroid)
x_list = [(lon, lat) for (lon, lat) in zones['Centroid'].apply(lambda x: (x.xy[1][0], x.xy[0][0]))]
pi_list = [geometry.area for geometry in zones['geometry']]

dQ = dist
epsilon = 0.5

p_matrix, pre_prob, post_prob = optql(x_list, pi_list, dQ, epsilon=epsilon)
# -



StarcoderdataJupyterScriptsDedupFiltered