Experimenting with a Custom Distance Based on Haversine with K Nearest Neighbors

# All imports for entire analysis

import os
import tarfile
import urllib
import joblib

import pandas as pd
import numpy as np
from sklearn.model_selection import StratifiedShuffleSplit, GridSearchCV, cross_val_score
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import OneHotEncoder, StandardScaler
from sklearn.neighbors import KNeighborsRegressor, DistanceMetric
from sklearn.linear_model import LinearRegression
from functools import partial
from sklearn.metrics import mean_squared_error

# i/o constants
# Dataset, online
DOWNLOAD_ROOT = 'https://raw.githubusercontent.com/ageron/handson-ml2/master/'
HOUSING_URL = DOWNLOAD_ROOT + 'datasets/housing/housing.tgz'

# Local (Relative paths)
HOUSING_PATH = 'dataset'  # Directory where data is saved
RESULTS = 'results'  

# tuning constants
CV = 5
NJOBS = 2  # Change if your system does not support parallelism
VERBOSITY = 9
SCORING = 'neg_mean_squared_error'

class Distances:
    """ User distance metrics to use with scikit-learn.  Includes our HaversineEuclidean Distance  """
    
    def __init__(self, metric):
        """
        Initialize distances class by verifying the metric string
        Parameters:
            metric, str.
        """
        # Valid metrics.  We are shortening Euclidean to l2.  The l2 norm induces the Euclidean metric.
        self.metrics = {
            'l2': self._l2,  
            'haversine': self._haversine,
            'haversine_l2': self._haversine_l2
        }
        
        if metric not in self.metrics.keys():
            raise ValueError(f"The metric {metric} is not recognized.")
        else:
            self.metric = self.metrics[metric]
            
        if 'haversine' in metric:
            # Load sklearn's haversine metric in advance
            self.skhaversine = DistanceMetric.get_metric('haversine')
    
    def distance(self, x, y, **kargs):
        """ 
        * Compute distance between vectors x and y after validation checks, using self.metric
        * We use one main function to avoid redundant things like input validation checks.        
        
        Parameters:
            x, array 
            y, array
            
        Returns:
            Scalar value representing distance between x and y
        """
        self._check_dist_inputs(x, y)
        return self.metric(x, y, **kargs)
    
    @staticmethod
    def _check_dist_inputs(x, y):
        """ Verify that the two inputs are valid. """
        for vec in [x, y]:
            if not isinstance(vec, np.ndarray):
                raise ValueError("Inputs expected to be np arrays")
            if vec.ndim != 1:
                raise ValueError("Inputs expected to have one dimension")

        if len(x) != len(y):
            raise ValueError("x and y must have same number of values")

    
    def _l2(self, x, y):
        """ Return Euclidean distance, or distance induced by l2 norm """
        return np.sqrt(np.sum((x-y)**2))
    
    def _haversine(self, x, y):
        """ 
        Compute Haversine distance between x and y
        
        Parameters:
            x = (latitude1, longitude1) in radians
            y = (latitude2, longitude2) in radians
            
        Note:
            Just uses functionality from sklearn.neighbors, to keep things conceptually contained
        """
        # Note, we already know the inputs are 1 dimensional arrays of same length
        if len(x) != 2:
            raise ValueError("Input vectors to Haversine distance must have length 2: (latitude, longitude)")

        # Get a 1 by 1 distance matrix (the Haversine distance between x and y)
        one_by_one_dist_matrix = self.skhaversine.pairwise(x[np.newaxis, :], y[np.newaxis, :])
        return one_by_one_dist_matrix.item()
        
    def _haversine_l2(self, x, y, w_h=0.5):
        """ 
        Combines the Haversine distance of latitude and longitude features with Euclidean of remaining features.

        Note:  This assumes vectors are of the form (latitude, longitude, x3, ..., xp).
        """
        # Note, we already know the inputs are 1 dimensional arrays of same length
        if len(x) < 3:
            raise ValueError(f"input to haversine_l2 has length {len(x)}.  Should be at least 3 to use this method")
        
        if not isinstance(w_h, float):
            raise TypeError('w_h should be a float')
        if w_h < 0. or w_h > 1: 
                raise ValueError("w_h should be in [0,1]")
        
        w_l2 = 1.0 - w_h
        return w_h*self._haversine(x[0:2], y[0:2]) + w_l2*self._l2(x[2:], y[2:])
        
    @staticmethod
    def check_is_metric(func, niter=1000, dd=10):
        """ 
        Empirically check whether func is a valid distance metric 
        It is always good to include a quick empirical check. I've included the math in the appendix.
        """
        print(f"Checking whether {func} is a distance")
        
        for test_iter in range(niter):
            x = np.random.randn(dd)
            y = np.random.randn(dd)
            z = np.random.randn(dd)
            
            # d(x,x) = 0  
            assert func(x, x) == 0., "d(x,x) != 0"
            
            # Symmetry
            assert func(x, y) == func(y, x), "d(x,y) != d(y,x)"
            
            # Non-negative
            assert func(x,y) >= 0, "d(x,y) is negative"
            
            # Triangle inequality
            assert func(x, y) + func(y, z) >= func(x, z), "Failed triangle ineq"
            
            # The other check, that d(x,y)=0 implies x=y is hard empirically.
            # This will suffice as a first check.  The math confirms this property.
        
        print(f"Passed {niter} empirical checks\n")

for _ in range(3):  # You can run more, I'm only using three to avoid too much output.
    dist = Distances('haversine_l2')
    haversine_l2_func = partial(dist.distance,
                                w_h=np.random.rand())
    Distances.check_is_metric(haversine_l2_func)

# Directly taken from notebook mentioned before
def fetch_housing_data(housing_url=HOUSING_URL, housing_path=HOUSING_PATH):
    """ From Hands-on Machine Learning """
    os.makedirs(housing_path, exist_ok=True)
    tgz_path = os.path.join(housing_path, 'housing.tgz')
    urllib.request.urlretrieve(housing_url, tgz_path)

    housing_tgz = tarfile.open(tgz_path)
    housing_tgz.extractall(path=housing_path)
    housing_tgz.close()
    print(f"Dataset downloaded, extracted, and saved in {tgz_path}")


def load_housing_data(housing_path=HOUSING_PATH):
    csv_path = os.path.join(housing_path, 'housing.csv')
    return pd.read_csv(csv_path)
        
    
fetch_housing_data()
housing = load_housing_data()

housing.head()

housing.info()

housing.describe()

housing['ocean_proximity'].value_counts()

housing['income_cat'] = pd.cut(housing['median_income'],
                              bins=[0., 1.5, 3.0, 4.5, 6., np.inf],
                              labels=[1,2,3,4,5])

split = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
splits_gen = split.split(housing, housing['income_cat'])
train_index, test_index = next(splits_gen)

housing_train = housing.loc[train_index].copy()
housing_test = housing.loc[test_index].copy()

# Pull out the targets
train_labels = housing_train['median_house_value'].copy()
test_labels = housing_test['median_house_value'].copy()

for set_ in (housing_train, housing_test):
    set_.drop('income_cat', axis=1, inplace=True)
    set_.drop('median_house_value', axis=1, inplace=True)

print(f"X train: {housing_train.shape}")
print(f"X test: {housing_test.shape}")
print(f"y train: {train_labels.shape}")
print(f"y test: {test_labels.shape}")

# Helper functions for evaluation

def get_rmse(score):
    """ 
    Given a negative Mean Squared Error, return the (positive) Root Mean Squared Error
    (sklearn prefers to return -MSE)
    """
    return np.sqrt(-score)


def print_results(searcher):
    """
    Print key results of a grid search, sorted by performance
    Parameters:
        searcher: A fitted GridSearchCV object
    """
    res_zip = zip(searcher.cv_results_['mean_test_score'], searcher.cv_results_['params'])
    for result, param in sorted(res_zip):
        print(get_rmse(result), param)


def styled_df(df):
    """ Display a data frame with commas separating thousands places """
    display(df.style.format("{:,.1f}"))

lat_long_train = housing_train[['latitude', 'longitude']]  # We'll need the data in this order: (latitude, longitude)
lat_long_test = housing_test[['latitude', 'longitude']]

print(lat_long_train.shape, lat_long_test.shape)
lat_long_train.head()

class DegreesToRadians(BaseEstimator, TransformerMixin):
    """ Transformer to convert degrees to radians, with special checks for valid latitude and longitude """
    def fit(self, X, y=None):
        return self  # Nothing else to do
    
    def transform(self, X):
        """ 
        Do conversion Radians = Degrees * pi / 180 
        Parameters:
            X, numpy array or pandas df containing latitude/longitude values in degrees
        Returns:
            numpy array with radian values
        """
        if isinstance(X, pd.core.frame.DataFrame):
            X = X.values  # Get a numpy array
        elif not isinstance(X, np.ndarray):
            raise ValueError(f"Unrecognized type for degrees->radians conversion: {type(X)}")

        # Verify that the values are between -180 and 180 (valid lat/longitude in degrees)
        if (np.min(X) < -180.) or (np.max(X) > 180.):
            raise ValueError("Latitude and longitude values are expected to be in the range [-180, 180]")
        else:
            return X*np.pi/180.

# Data pre-processing pipelines
# (ll = latitude/longitude)
ll_euclidean_pipeline = Pipeline([('scale', StandardScaler())])
ll_haversine_pipeline = Pipeline([('degToRad', DegreesToRadians())])

# Grids
ll_base_grid = {'n_neighbors': np.arange(2, 15), 'weights': ['uniform', 'distance']}

ll_euclid_grid = ll_base_grid.copy()
ll_euclid_grid['metric'] = ['euclidean']
ll_haversine_grid = ll_base_grid.copy()
ll_haversine_grid['metric'] = ['haversine']
print(ll_euclid_grid)
print(ll_haversine_grid)

def ll_grid_search(pipe, param_grid, ntrain=None):
    """ 
    Train a model on latitude and longitude (ll) only.
    The pre-processing steps will differ, depending on whether we convert to radians for Haversine or something else.  
    Then, compute test-set performance
    
    Parameters:
        pipe: A Pipeline object to preprocess the inputs
        param_grid: A dictionary containing hyperparameters to tune
        ntrain: The number of observations to train on.  Setting to None uses the entire training set.
    """
    # Pre-process data
    train_prepared = pipe.fit_transform(lat_long_train)
    print("\nGrid Searching")
    print(f"\tInput data:\n\t Shape: {train_prepared.shape}\n\t Feature means: {np.mean(train_prepared, axis=0)}")
    print(f"\tTargets:\n\t Shape: {train_labels.shape}")
    
    # Run grid-search
    # (If you get an error, try changing NJOBS=1)
    grid_searcher = GridSearchCV(KNeighborsRegressor(),
                                 param_grid,
                                 cv=CV,
                                 n_jobs=NJOBS,
                                 verbose=VERBOSITY,
                                 scoring=SCORING)
    
    grid_searcher.fit(train_prepared[0:ntrain], train_labels[0:ntrain])
    
    # Save key results
    best_model = grid_searcher.best_estimator_
    best_cv_rmse = get_rmse(grid_searcher.best_score_)
    
    # Prepare test set
    test_prepared = pipe.transform(lat_long_test)
    print(f"\tTest data:\n\t X shape: {test_prepared.shape}\n\t Feature means: {np.mean(test_prepared, axis=0)}")
    print(f"\tTest targets:\n\t  y shape: {test_labels.shape}")

    # Test performance
    yhat = best_model.predict(test_prepared)
    test_score = np.sqrt(mean_squared_error(test_labels, yhat)) 
    
    return grid_searcher, best_model, best_cv_rmse, test_score


ll_euclid_results = ll_grid_search(pipe = ll_euclidean_pipeline, param_grid = ll_euclid_grid)
ll_haversine_results = ll_grid_search(pipe = ll_haversine_pipeline, param_grid = ll_haversine_grid)

ll_euclid_searcher, ll_euclid_best_model, ll_euclid_cv_rmse, ll_euclid_test_rmse = ll_euclid_results
ll_haversine_searcher, ll_haversine_best_model, ll_haversine_cv_rmse, ll_haversine_test_rmse = ll_haversine_results

ll_euclid_best_model.get_params()

ll_haversine_best_model.get_params()

styled_df(
    pd.DataFrame([[ll_euclid_cv_rmse, ll_haversine_cv_rmse]],
                columns = ['Euclid CV RMSE', 'Haversine CV RMSE'])
)

#
# Extract performance of Haversine-KNN at the optimal hyperparameters of Euclid-KNN
#
optimal_euclid_n_neighbors = ll_euclid_best_model.get_params()['n_neighbors']
ii = 0
for kk in ll_haversine_searcher.cv_results_['params']:
    if kk == {'metric': 'haversine', 'n_neighbors': optimal_euclid_n_neighbors, 'weights': 'uniform'}:
        print(ii)
        break
    ii += 1

get_rmse(ll_haversine_searcher.cv_results_['mean_test_score'][ii])

styled_df(
    pd.DataFrame([[ll_euclid_test_rmse, ll_haversine_test_rmse]],
                columns = ['Euclid Test RMSE', 'Haversine Test RMSE'])
)

# A preprocessing Transformer from the textbook

# Hard-code the column indices, as in the textbook
# The indices differ by 2 because we're treating latitude and longitude differently: they get "taken out" of this calculation
rooms_ix, bedrooms_ix, population_ix, households_ix = 1, 2, 3, 4
# rooms_ix, bedrooms_ix, population_ix, households_ix = 3, 4, 5, 6

class CombinedAttributesAdder(BaseEstimator, TransformerMixin):
    """ Creates new features.  From page 68 of textbook and in the aforementioned Jupyter notebook """
    def __init__(self, add_bedrooms_per_room=True):
        self.add_bedrooms_per_room = add_bedrooms_per_room
    
    def fit(self, X, y=None):
        return self
    
    def transform(self, X):
        rooms_per_household = X[:, rooms_ix]/X[:, households_ix]
        population_per_household = X[:, population_ix]/X[:, households_ix]
        
        if self.add_bedrooms_per_room:
            bedrooms_per_room = X[:, bedrooms_ix]/X[:, rooms_ix]
            return np.c_[X, 
                         rooms_per_household,
                         population_per_household,
                         bedrooms_per_room]
        else:
            return np.c_[X, 
                         rooms_per_household,
                         population_per_household] 
        

# Get a list of numerical attributes other than latitude and longitude
remaining_num_attribs = list(housing_train.columns)
for _del in ('ocean_proximity', 'latitude', 'longitude'):
    remaining_num_attribs.remove(_del)

num_pipeline = Pipeline(  # Pipeline used in the aforementioned Jupyter notebook
    [
        ('imputer', SimpleImputer(strategy='median')),
        ('attribs_adder', CombinedAttributesAdder()),
        ('std_scaler', StandardScaler())
    ],
    verbose=True
)
    
full_pipeline = ColumnTransformer(
    [
        ('latitudeLongitudeToRads', DegreesToRadians(), ['latitude', 'longitude']), 
        ('num', num_pipeline, remaining_num_attribs),
        ('cat', OneHotEncoder(), ['ocean_proximity'])
    ],
    verbose=True
)
housing_prepared = full_pipeline.fit_transform(housing_train)
housing_test_prepared = full_pipeline.transform(housing_test)

for name, prepared_dataset in \
    {'housing_prepared': housing_prepared, 'housing_test_prepared': housing_test_prepared}.items():
    
    print('-'*50)
    print(f'dataset info for {name}')
    
    print(type(prepared_dataset))
    print(f"Number of missing values: {np.isnan(prepared_dataset).sum()}")
    print(f"Variable means: {np.mean(prepared_dataset, axis=0)}")  
    print(f"Variable stds: {np.std(prepared_dataset, axis=0)}")
    print(f"Shape: {prepared_dataset.shape}")
    print("cols 0:5\n", prepared_dataset[0:3,0:5],
      "\ncols 5:10\n", prepared_dataset[0:3,5:10],
     "\ncols 10:\n", prepared_dataset[0:3,10:])

def all_features_search(out_file_name, grid_dict, num_train=None):
    """
    
    Given a pre-processed dataset and a hyperparameter grid, run a grid over KNN Regression models.
    This is written for the case where the input contains all features, not just latitude and longitude.
    
    Since this can be computationally intensive, we save the results to disk.  
    If the file at out_file_name exists, we load it rather than train.  Otherwise, we create it after training.
    
    We can set num_train to be a subset to speed things up for development or exploratory purposes.
    In that case, it's added to the file path.
    
    For simplicity, I don't use grid information in the file path.  It's not needed for the notebook.
    
    Parameters:
        out_file_name: str, path to saved model file, without suffix.  
        grid_dict: dictionary, hyperparameter grid
        num_train: int or None.  If None, the entire train set is used.
    """
    if num_train is not None:
        assert isinstance(num_train, int)
        out_file_name += f"_n_train_{num_train}"
    out_file_name += ".pkl"
    
    if os.path.exists(out_file_name):
        print(f"Reading {out_file_name} from disk")
        grid_searcher = joblib.load(out_file_name)
    else:
        print("Training grid from scratch")
        # If you get an error, try changing NJOBS to 1.
        grid_searcher = GridSearchCV(KNeighborsRegressor(),
                                     grid_dict,
                                     cv=CV,
                                     n_jobs=NJOBS,
                                     verbose=VERBOSITY,
                                     scoring=SCORING)

        grid_searcher.fit(housing_prepared[0:num_train], train_labels[0:num_train])
        
        print(f"Done.\nWriting to disk at {out_file_name}")
        joblib.dump(grid_searcher, out_file_name)

        
    return grid_searcher

coarse_euclidean_grid = {'n_neighbors': np.arange(5, 30, 5),
                         'weights': ['uniform', 'distance'],
                         'metric': ['euclidean']}

coarse_euclidean_outfile = os.path.join(RESULTS, 'out_coarse_euclidean')
coarse_euclidean_searcher = all_features_search(out_file_name=coarse_euclidean_outfile,
                                                grid_dict=coarse_euclidean_grid)

print_results(coarse_euclidean_searcher)

# There are some redundant values, but it does not cause a big problem in terms of speed here (Euclidean is much faster than HaversineEuclidean)
fine_euclidean_grid = {'n_neighbors': np.arange(10, 21),
                         'weights': ['distance'],
                         'metric': ['euclidean']}

fine_euclidean_outfile = os.path.join(RESULTS, 'out_fine_euclidean')
fine_euclidean_searcher = all_features_search(out_file_name=fine_euclidean_outfile,
                                           grid_dict=fine_euclidean_grid)

print_results(fine_euclidean_searcher)

hyper_ws = fine_euclidean_searcher.best_params_  # Like "warm starting" in a way. 
coarse_haversine_grid = {'n_neighbors': [hyper_ws['n_neighbors']],
                         'weights': [hyper_ws['weights']],
                         'metric': [Distances('haversine_l2').distance],
                         'metric_params': [{'w_h': 0.25}, {'w_h': 0.5}, {'w_h': 0.75}]}



coarse_haversine_outfile = os.path.join(RESULTS, 'out_coarse_haversine')
print(coarse_haversine_grid)
print(coarse_haversine_outfile)

coarse_haversine_searcher = all_features_search(out_file_name=coarse_haversine_outfile,
                                                grid_dict=coarse_haversine_grid)

print_results(coarse_haversine_searcher)

fine_haversine_grid = coarse_haversine_grid.copy()
fine_haversine_grid['metric_params'] = [{'w_h': 0.85}, {'w_h': 0.90}, {'w_h': 0.95}, {'w_h': 0.99}]

fine_haversine_outfile = os.path.join(RESULTS, 'out_fine_haversine')
print(fine_haversine_grid)
print(fine_haversine_outfile)

fine_haversine_searcher = all_features_search(out_file_name=fine_haversine_outfile,
                                              grid_dict=fine_haversine_grid)

print_results(fine_haversine_searcher)

# Save the cv score for later.
ff_haversine_cv_rmse = get_rmse(fine_haversine_searcher.best_score_)
ff_haversine_cv_rmse

def all_features_test_set(estimator, out_file_name):
    """ 
    Predict on the test set, using all the features, and report the RMSE score 
    Parameters:
        estimator: A scikit-learn estimator
        out_file_name: A string where predictions from a fitted model are saved
    """
    print(f"Estimator: {estimator}")
    
    # Fitting is really cheap for KNeighborsRegressor.  We can do it even if 
    # we've done it before and saved to disk.
    # The expensive operation is prediction: computing distances to all observations in the training set
    print("\nFitting model on prepared train set")
    print(type(housing_prepared))
    print(f"Input shape: {housing_prepared.shape}")
    print(f"Targets shape: {train_labels.shape}")
    estimator.fit(housing_prepared, train_labels)
    
    # Obtain predictions
    if os.path.exists(out_file_name):
        print(f"\nReading predictions from disk at {out_file_name}")
        yhat = joblib.load(out_file_name)
    else:
        print("\nObtaining predictions on prepared test set")
        print(type(housing_test_prepared))
        print(f"Input shape: {housing_test_prepared.shape}")
        print(f"Targets shape: {test_labels.shape}")
        
        yhat = estimator.predict(housing_test_prepared)
        
        print(f"Saving predictions at {out_file_name}")
        joblib.dump(yhat, out_file_name)
        
    mse = mean_squared_error(test_labels, yhat)
    
    return np.sqrt(mse)
    

haversine_outfile = os.path.join(RESULTS, 'final_haversine_preds.pkl')
ff_haversine_model_test_rmse = all_features_test_set(
    fine_haversine_searcher.best_estimator_,
    haversine_outfile
)
ff_haversine_model_test_rmse  # Full features = ff

# For comparison, a linear regression model
linear_model = LinearRegression()
linear_outfile = os.path.join(RESULTS, 'final_linear_preds.pkl')
linear_model_test_rmse = all_features_test_set(linear_model, linear_outfile)
print(linear_model_test_rmse)

table_columns = ['Dataset',   'Distance',  'CV Mean RMSE', 'Test Score']
# Latitude/Longitude = ll
ll_e = pd.Series(['Lat/Long', 'Euclidean', ll_euclid_cv_rmse,        ll_euclid_test_rmse],
                 index=table_columns)
ll_h = pd.Series(['Lat/Long', 'Haversine', ll_haversine_cv_rmse,     ll_haversine_test_rmse],
                 index=table_columns)

# Full features = ff
ff_e = pd.Series(['All Features', 'HaversineEuclidean', ff_haversine_cv_rmse, ff_haversine_model_test_rmse],
                 index=table_columns)


pd.DataFrame([ll_e, ll_h, ff_e])