Search code examples
python-3.xpandasmatrixdistanceadjacency-matrix

Create a distance matrix from Pandas Dataframe using a bespoke distance function


I have a Pandas dataframe with two columns, "id" (a unique identifier) and "date", that looks as follows:

test_df.head()

    id  date
0   N1  2020-01-31
1   N2  2020-02-28
2   N3  2020-03-10

I have created a custom Python function that, given two date strings, will compute the absolute number of days between those dates (with a given date format string e.g. %Y-%m-%d), as follows:

def days_distance(date_1, date_1_format, date_2, date_2_format):
    """Calculate the number of days between two given string dates

    Args:
        date_1 (str): First date
        date_1_format (str): The format of the first date
        date_2 (str): Second date
        date_2_format (str): The format of the second date

    Returns:
        The absolute number of days between date1 and date2
    """

    date1 = datetime.strptime(date_1, date_1_format)
    date2 = datetime.strptime(date_2, date_2_format)
    return abs((date2 - date1).days)

I would like to create a distance matrix that, for all pairs of IDs, will calculate the number of days between those IDs. Using the test_df example above, the final time distance matrix should look as follows:

    N1    N2    N3
N1  0     28    39
N2  28    0     11
N3  39    11    0

I am struggling to find a way to compute a distance matrix using a bespoke distance function, such as my days_distance() function above, as opposed to a standard distance measure provided for example by SciPy.

Any suggestions?


Solution

  • Let us try pdist + squareform to create a square distance matrix representing the pair wise differences between the datetime objects, finally create a new dataframe from this square matrix:

    from scipy.spatial.distance import pdist, squareform
    
    i, d = test_df['id'].values, pd.to_datetime(test_df['date'])
    df = pd.DataFrame(squareform(pdist(d[:, None])), dtype='timedelta64[ns]', index=i, columns=i)
    

    Alternatively you can also calculate the distance matrix using numpy broadcasting:

    i, d = test_df['id'].values, pd.to_datetime(test_df['date']).values 
    df = pd.DataFrame(np.abs(d[:, None] - d), index=i, columns=i)
    

            N1      N2      N3
    N1  0 days 28 days 39 days
    N2 28 days  0 days 11 days
    N3 39 days 11 days  0 days