<jupyter_start><jupyter_text>Tree sequence characterisation of the UK Biobank using tsinfer By [<NAME>](https://github.com/chaimaafadil) Publication - ***Inferring whole-genome histories in large population datasets*** - DOI: 10.1038/s41588-019-0483-y*External data licensing:**The boundary data referenced in this notebook are provided by the Office for National Statistics (ONS), the Ordnance Survey and the Database of Global Administrative Areas (GADM). It contains public sector information licensed under the [Open Government Licence v3.0](http://www.nationalarchives.gov.uk/doc/open-government-licence/version/3/).* --- This notebook contains the code for plotting Figure 5 of the publication 'Inferring whole-genome histories in large population datasets'. The aim of this analysis was to assess the performance of [tsinfer](https://tsinfer.readthedocs.io/en/latest/) on vast data sets such as the UK Biobank which comprises the SNP array data of ~500K individuals.With [tsinfer](https://tsinfer.readthedocs.io/en/latest/), we constructed a tree sequence for the UK Biobank (UKB) on human chromosome 20. Note that this notebook does not include the code for constructing the tree sequence. We used the tree sequence to characterise ancestral relationships between UKB individuals by computing the GNN proportions of each individual (see paper Methods for an explanation of Genealogical Nearest Neighbours). We then generated plots to (1) visualise the population structure resulting from the GNN structure of the tree sequence and (2) characterise the relationship between geographical proximity of birth locations and genetic relatedness. Please refer to the publication above for further discussion. The notebook is structured into 3 sections:* In **Section 1** we determine the birth location of all individuals in the UKB with reported birth coordinates. The UKB provides birth coordinates for individuals born in Great Britain (England, Wales and Scotland). These coordinates are encoded in the British National Grid (BNG) reference system. We reverse-geocode these coordinates to determine the name of the region in which a person was born. We define regions in the UK as per the [NUTS](https://en.wikipedia.org/wiki/NUTS_statistical_regions_of_the_United_Kingdom) territorial subdivision system. The NUTS subdivision of the UK is defined at 3 levels. We chose to map individuals' birth coordinates to both NUTS Level-2 and NUTS Level-3 areas.* In **Section 2** we generate 2 matrices of GNN proportions by computing the GNNs of each individual in the UKB using as reference sets individuals' NUTS-2 and NUTS-3 birth locations, respectively. The GNNs are computed from the tree sequence constructed previously. Note that a path to the tree sequence `ts_path` is required as input.* In **Section 3** we use the data generated in sections 1 and 2 to plot the different panels of Figure 5, that is, visual representations (1) of the GNN matrix and (2) of the geographical clustering of relatedness between. Note that panels A and B of Figure 5 rely on the NUTS-2 subdivision of the UK, while panels C, D and E, on the NUTS-3 subdivision. **Additional notes on the execution of this notebook:*** The 3 main sections of this notebook need to be run in numerical order as each section following section 1.1 is dependent on variables created in previous sections.* The input files for each section are listed only for reference. It is not necessary to load them if the script is being run in one sitting.* For the boundary vector shapefiles, the link to download these files is provided in the markdown. These come in the form of zipped folders. Make sure to unzip the folder and place it in the same directory as this notebook. The folder should have the same name as the one indicated in the code.* The input files used in section 1.2 for the names and countries of NUTS areas are provided with this notebook. The code to load them is already included in the script and assumes that they are located in the same directory as this notebook.<jupyter_code>"""
Various libraries used throughout the notebook
"""

import numpy as np
import pandas as pd
import tskit

import scipy
import cartopy
import cartopy.crs as ccrs
import cartopy.io.shapereader as shpreader
from shapely.geometry import Polygon
from shapely.geometry import MultiPoint
from shapely.geometry import Point

import random
import time<jupyter_output><empty_output><jupyter_text>1. Reverse-geocode places of birth  1.1 Map BNG (Easting, Northing) coordinates to NUTS Level 2 and Level 3 areas **1.1 Input files:** * `ukb_metadata.csv` - Raw dataframe with individuals' IDs and birth coordinates defined under the Ordnance Survey National Grid reference system (also referred to as British National Grid - BNG). In this code, `ukb_metadata.csv` contains in each row data relating to a single individual. This data is organised into the following 5 columns:    * `SampleID`    * `Ethnicity`    * `PlaceOfBirthUK_East` (person's easting birth coordinate)    * `PlaceOfBirthUK_North` (person's northing birth coordinate)    * `CountryOfBirth_UK` (person's country of birth if born in the UK. If not, the field is specified as 'NaN')     * Boundary vectors for:    * [NUTS Level 2 areas](https://data.gov.uk/dataset/e65fbc2a-7314-450f-b4a1-14ba87378b3c/nuts-level-2-january-2018-generalised-clipped-boundaries-in-the-united-kingdom): `NUTS_Level_2_January_2018_Generalised_Clipped_Boundaries_in_the_United_Kingdom`;    * [NUTS Level 3 areas](https://data.gov.uk/dataset/0c9bc79f-e3d2-4955-aa5c-cd43bae9c67a/nuts-level-3-january-2018-generalised-clipped-boundaries-in-the-united-kingdom): `NUTS_Level_3_January_2018_Generalised_Clipped_Boundaries_in_the_United_Kingdom`.<jupyter_code>We first import the raw dataframe containing individuals' places of birth encoded in BNG coordinates (eastings and northings).
# Import dataframes and add a column to store NUTS birth locations.
def get_input_df(path):
    df = pd.read_csv(path)
    df['PlaceOfBirth_NUTS'] = np.nan
    return df
ukb_nuts2_df = get_input_df("ukb_metadata.csv")
ukb_nuts3_df = get_input_df("ukb_metadata.csv")<jupyter_output><empty_output><jupyter_text>We then import the vector boundary data for the 41 NUTS-2 areas and the 179 NUTS-3 areas. This data contains the polygons and names of each NUTS area. The polygons are defined by digital vector boundaries for NUTS-2 and NUTS-3 areas in the UK as of January 2018. The boundaries are available at different resolutions - we use here the generalised clipped boundaries. This data was published by the Office for National Statistics and can be downloaded as a **shapefile** from the UK government open database for [NUTS-2](https://data.gov.uk/dataset/e65fbc2a-7314-450f-b4a1-14ba87378b3c/nuts-level-2-january-2018-generalised-clipped-boundaries-in-the-united-kingdom) and [NUTS-3](https://data.gov.uk/dataset/0c9bc79f-e3d2-4955-aa5c-cd43bae9c67a/nuts-level-3-january-2018-generalised-clipped-boundaries-in-the-united-kingdom) areas.<jupyter_code># Import NUTS polygons and names and create a dataframe of polygons, poly_df.
def get_nuts_polygons(nuts_path, name_field, code_field):
    # Get NUTS polygons and names
    area_path = nuts_path
    reader = shpreader.Reader(area_path)
    geometries = list(reader.geometries())
    names = [record.attributes[name_field] for record in reader.records()]
    codes = [record.attributes[code_field] for record in reader.records()]
    # Create dataframe of polygons
    poly_df = pd.DataFrame({
        'geometry': geometries,
        'name': names,
        'code': codes
    })
    return poly_df

# Update the NUTS-3 poly_df to account for individuals born in Northern Ireland (see 'note on individuals born in Northern Ireland' below).
# This function adds the NUTS-2 polygon for Northern Ireland to poly_df_nuts3 and deletes NUTS-3 polygons located in Northern Ireland.
def update_nuts3_polygons(poly_df_nuts3, poly_df_nuts2):
    # Remove NUTS-3 polygons located in Northern Ireland
    NIR = ['Antrim and Newtownabbey','Ards and North Down','Armagh City, Banbridge and Craigavon','Belfast',
           'Causeway Coast and Glens','Derry City and Strabane','Fermanagh and Omagh','Lisburn and Castlereagh',
           'Mid Ulster','Mid and East Antrim','Newry, Mourne and Down']
    NIR_index = poly_df_nuts3.loc[poly_df_nuts3.name.isin(NIR),:].index.values
    poly_df_nuts3 = poly_df_nuts3.drop(NIR_index)
    # Add "Northern Ireland" to poly_df_nuts3
    NIR = poly_df_nuts2.loc[poly_df_nuts2['name']=='Northern Ireland',:]
    poly_df_nuts3 = poly_df_nuts3.append(NIR, ignore_index=True)
    return poly_df_nuts3
poly_df_nuts2 = get_nuts_polygons("NUTS_Level_2_January_2018_Generalised_Clipped_Boundaries_in_the_United_Kingdom/NUTS_Level_2_January_2018_Generalised_Clipped_Boundaries_in_the_United_Kingdom", "nuts218nm", "nuts218cd")
poly_df_nuts3 = get_nuts_polygons("NUTS_Level_3_January_2018_Generalised_Clipped_Boundaries_in_the_United_Kingdom/NUTS_Level_3_January_2018_Generalised_Clipped_Boundaries_in_the_United_Kingdom", "nuts318nm", "nuts318cd")
poly_df_nuts3 = update_nuts3_polygons(poly_df_nuts3, poly_df_nuts2) # See note below<jupyter_output><empty_output><jupyter_text>**Note on individuals born in Northern Ireland:**In the UKB, individuals born in Northern Ireland have no reported birth coordinates. Their place of birth is only reported through the `CountryOfBirth_UK` field as "Northern Ireland". We are therefore unable to reverse-geocode their birth location. As we still want to include them in our analysis, we manually report their NUTS-2 and NUTS-3 birth location as "Northern Ireland".Since "Northern Ireland" is a NUTS-2 and not a NUTS-3 area, we modified `poly_df_nuts3` above with `update_nuts3_polygons` to include a polygon for the area of "Northern Ireland". We also remove the unused 11 NUTS-3 areas located in Northern Ireland which we won't use. --- Finally, we reverse-geocode birth locations. We input individuals' birth coordinates and map those to a NUTS area. We do so using as reference both NUTS-2 and NUTS-3 subdivisions.The function `reverse-geocode` takes as input the UKB dataframe `ukb_df` which contains individuals' BNG birth coordinates and the dataframe of polygons `poly_df` created in the previous step. For each birth coordinate tuple (easting and northing), the function outputs the name of the NUTS (2 or 3) area in which the point encoded by the coordinates is located in. The function operates in 3 steps:1. `checkCoordinates` verifies that the coordinates are not `NaN`. Note that in the UKB, place of birth coordinates are only given for individuals born in Great Britain.2. `OSGB36toWGS84` converts the BNG coordinates (OSGB36) to latitude/longitude coordinates (WGS84).3. `LatLonToAddress` uses the `within()` function to map the latitude/longitude coordinates to one of the polygons in `poly_df` and outputs the name of the polygon.Note: if your coordinates are given in latitudes and longitudes, you can modify `reverse-geocode` to skip step (2).<jupyter_code>from bng_to_latlon import OSGB36toWGS84
# Library: https://pypi.org/project/bng-latlon/

def checkCoordinates(easting,northing):
    return (not(np.isnan(easting) and np.isnan(northing)))

def LatLonToAddress(latlon, poly_df):
    p = Point(latlon[1],latlon[0])
    # Find polygon
    for index in poly_df.index:
        if p.within(poly_df.at[index,'geometry']):
            return poly_df.at[index,'name']
            break
    return np.nan

def reverse_geocode(ukb_df, poly_df):
    for index in ukb_df.index:
        easting = ukb_df.at[index,'PlaceOfBirthUK_East']
        northing = ukb_df.at[index,'PlaceOfBirthUK_North']
        if checkCoordinates(easting, northing):
            try:
                latlon = OSGB36toWGS84(easting, northing)
                address = LatLonToAddress(latlon, poly_df)
                ukb_df.loc[index, 'PlaceOfBirth_NUTS'] = address
            except:
                print("Exception at ", index)
    # Manually map all individuals born in Northern Ireland to "Northern Ireland" (see note above)
    ukb_df.loc[ukb_df['CountryOfBirth_UK']=='Northern Ireland', 'PlaceOfBirth_NUTS'] = 'Northern Ireland'
    return ukb_df
# Note: this step will take a few hours to run. A more efficient approach would be to run it on different parts 
# of the dataframe in parallel then merge the outputs.
ukb_nuts2_df = reverse_geocode(ukb_nuts2_df, poly_df_nuts2)
ukb_nuts3_df = reverse_geocode(ukb_nuts3_df, poly_df_nuts3)
# Save
ukb_nuts2_df.to_csv("ukb_reverse_geocoded_nuts2.csv", index=False)
ukb_nuts3_df.to_csv("ukb_reverse_geocoded_nuts3.csv", index=False)<jupyter_output><empty_output><jupyter_text>**1.1 Output files:*** `ukb_reverse_geocoded_nuts2.csv` - UKB dataframe with birth locations reverse-geocoded to NUTS-2 areas.* `ukb_reverse_geocoded_nuts3.csv` - UKB dataframe with birth locations reverse-geocoded to NUTS-3 areas.  1.2 Clean the dataset by removing any inconsistencies in the reverse-geocoding of birth locationsAmong the individuals with reported birth coordinates that were reverse-geocoded in the previous step, we suggest to delete individuals who:1. **Were not mapped to any NUTS area.** That is usually the case for individuals with birth coordinates mapping to bodies of water that are outside NUTS polygons. For example, one individual's birth coordinates mapped to a point in the Atlantic Ocean.    2. **Were mapped to a NUTS area that is not located within their reported country of birth.** For example, an individual's birth coordinates mapped to an area in East Wales while their country of birth was reported as 'England' in the `CountryOfBirth_UK` column. **1.2 Input files:*** `ukb_reverse_geocoded_nuts2.csv` - UKB dataframe with birth locations reverse-geocoded to NUTS-2 areas.* `ukb_reverse_geocoded_nuts2.csv` - UKB dataframe with birth locations reverse-geocoded to NUTS-3 areas.* `UK_NUTS_Level_2_January_2018_Names_and_Countries.csv` - Name and country of NUTS-2 areas.* `UK_NUTS_Level_3_January_2018_Names_and_Countries.csv` - Name and country of NUTS-3 areas.<jupyter_code># Select individuals born in Great Britain with reported birth coordinates.
def get_gb_inds(ukb_df):
    countries = ['England','Scotland','Wales']
    GB_inds = ukb_df[ukb_df['CountryOfBirth_UK'].isin(countries)]
    GB_inds = GB_inds[GB_inds['PlaceOfBirthUK_East'].notnull()]
    return GB_inds

# Read dataframes of nuts countries.
def get_nuts_countries(countries_path):
    nuts_countries = pd.read_csv(countries_path, index_col='name')
    return nuts_countries

# Get list of individuals with 'NaN' or inconsistent birth locations.
def get_inds_to_delete(GB_inds, countries_path):
    deleted_inds = list()
    # Select individuals that were not mapped to a NUTS area
    deleted_inds.extend(list(GB_inds[GB_inds['PlaceOfBirth_NUTS'].isnull()].index.values))
    GB_inds = GB_inds.drop(deleted_inds)
    # Select individuals that were mapped to a NUTS area that is not located in their reported country of birth
    nuts_countries = get_nuts_countries(countries_path)
    for index in GB_inds.index:
        nuts = GB_inds.at[index, 'PlaceOfBirth_NUTS']
        nuts_country = nuts_countries.at[nuts,'country']
        country_of_birth = GB_inds.at[index, 'CountryOfBirth_UK']
        if nuts_country!=country_of_birth:
            deleted_inds.append(index)
    return deleted_inds
 
# Delete individuals with no or inconsistent reverse-geocoding of birth locations from dataframe.
def clean_reverse_geocoding(ukb_df, countries_path):
    GB_inds = get_gb_inds(ukb_df)
    deleted_inds = get_inds_to_delete(GB_inds, countries_path)
    ukb_df = ukb_df.drop(deleted_inds)
    return ukb_df
# Get cleaned dataframes
ukb_nuts2_df = clean_reverse_geocoding(ukb_nuts2_df, "UK_NUTS_Level_2_January_2018_Names_and_Countries.csv")
ukb_nuts3_df = clean_reverse_geocoding(ukb_nuts3_df, "UK_NUTS_Level_3_January_2018_Names_and_Countries.csv")
# Save
ukb_nuts2_df.to_csv("ukb_reverse_geocoded_cleaned_nuts2.csv", index=False)
ukb_nuts3_df.to_csv("ukb_reverse_geocoded_cleaned_nuts3.csv", index=False)<jupyter_output><empty_output><jupyter_text>**1.2 Output files*** `ukb_reverse_geocoded_cleaned_nuts2.csv` - UKB dataframe after reverse-geocoding to NUTS-2 and cleaning.* `ukb_reverse_geocoded_cleaned_nuts3.csv` - UKB dataframe after reverse-geocoding to NUTS-3 and cleaning. ---  2. Compute GNN matrix As mentioned earlier, this part requires as input a tree sequence constructed from the dataset.In this section, we compute the GNN proportions of each haplotype following the approach outlined in the Methods on Genealogical Nearest Neighbours and in Supplementary Figure 19. Here, our reference sets correspond to the NUTS areas to which individuals were mapped earlier. Before computing the GNNs, we extract the list of augmented samples with `run_compute_ukbb_gnn_get_augmented`. Augmented samples are the samples that were sequentially added to the ancestors tree sequence so that path compression could generate additional ancestors through signals of shared recombination breakpoints. To avoid biasing the analysis, we remove these samples when computing the GNNS in `run_compute_ukbb_gnn`. **2. Input files:*** `ukb_reverse_geocoded_cleaned_nuts2.csv` - UKB dataframe after reverse-geocoding to NUTS-2 and cleaning.* `ukb_reverse_geocoded_cleaned_nuts3.csv` - UKB dataframe after reverse-geocoding to NUTS-3 and cleaning.* `ts_path` - Path to the tree sequence constructed from the dataset.<jupyter_code># Functions to compute the GNN matrix

"""
Various utilities for manipulating tree sequences and running tsinfer.
"""
import argparse
import subprocess
import time
import collections
import json
import sys
import io
import csv
import itertools
import os.path

import msprime
import tsinfer
import daiquiri
import numpy as np
import pandas as pd
import tqdm
import humanize
import cyvcf2

# Create dictionary of individuals' ethnicity and NUTS birth location
def get_ukb_dict(ukb_df):
    ukb_dict = {}
    for index in ukb_df.index:
        ukb_dict[ukb_df.at[index,'SampleID']] = {
            'Ethnicity': ukb_df.at[index,'Ethnicity'],
            'PlaceOfBirth_NUTS': ukb_df.at[index,'PlaceOfBirth_NUTS'],
    }
    return ukb_dict

def get_augmented_samples(tables):
    # Shortcut. Iterating over all the IDs is very slow here.
    # Note that we don't necessarily recover all of the samples that were 
    # augmented here because they might have been simplified out.
    nodes = tables.nodes
    ids = np.where(nodes.flags == tsinfer.NODE_IS_SAMPLE_ANCESTOR)[0]
    sample_ids = np.zeros(len(ids), dtype=int)
    for j, node_id in enumerate(tqdm.tqdm(ids)):
        offset = nodes.metadata_offset[node_id: node_id + 2]       
        buff = bytearray(nodes.metadata[offset[0]: offset[1]])
        md = json.loads(buff.decode())       
        sample_ids[j] = md["sample"]
    return sample_ids   

def run_compute_ukbb_gnn_get_augmented(ts_path):
    ts = tskit.load(ts_path)
    tables = ts.tables
    before = time.time()
    augmented_samples = set(get_augmented_samples(tables))
    duration = time.time() - before
    print("Got augmented:", len(augmented_samples), "in ", duration)
    return augmented_samples
 
def run_compute_ukbb_gnn(ts_path, augmented_samples, ukb_dict):
    ts = tskit.load(ts_path)
    tables = ts.tables
    before = time.time()

    reference_sets_map = collections.defaultdict(list)
    ind_metadata = [None for _ in range(ts.num_individuals)]
    all_samples = []

    for ind in ts.individuals():
        md = json.loads(ind.metadata.decode())
        
        try:
            sample = ukb_dict[int(md["SampleID"])]
            if type(sample["PlaceOfBirth_NUTS"])!=float:
                md["PlaceOfBirth"] = ukb_dict[int(md["SampleID"])]["PlaceOfBirth_NUTS"]
            else:
                md["PlaceOfBirth"]=np.nan
                
        except KeyError:
            # Individuals not in dictionary (e.g. removed in step 1.2)
            md["PlaceOfBirth"] = np.nan
            pass

        ind_metadata[ind.id] = md
        # Remove augmented samples to avoid biasing the analysis
        for node in ind.nodes:
            if node not in augmented_samples:
                reference_sets_map[md["PlaceOfBirth"]].append(node)
                all_samples.append(node)

    reference_set_names = list(reference_sets_map.keys())
    reference_sets = [reference_sets_map[key] for key in reference_set_names]
    
    cols = {
        "place_of_birth": [
            ind_metadata[ts.node(u).individual]["PlaceOfBirth"] for u in all_samples],
        "sample_id": [
            ind_metadata[ts.node(u).individual]["SampleID"] for u in all_samples],
        "ethnicity": [
            ind_metadata[ts.node(u).individual]["Ethnicity"] for u in all_samples],
    }

    print("Computing GNNs for ", len(all_samples), "samples")
    before = time.time()
    A = ts.genealogical_nearest_neighbours(all_samples, reference_sets, num_threads=16)
    duration = time.time() - before
    print("Done in {:.2f} mins".format(duration / 60))

    for j, name in enumerate(reference_set_names):
        cols[name] = A[:, j]
    gnn_df = pd.DataFrame(cols)
    return gnn_df

# Main function to compute GNNs for the UKB using NUTS areas as reference sets
def compute_gnn(ts_path, ukb_df, augmented_samples):
    ukb_dict = get_ukb_dict(ukb_df)
    gnn_df = run_compute_ukbb_gnn(ts_path, augmented_samples, ukb_dict)
    return gnn_df
# Get augmented samples. 
ts_path = "" # Add path to tree sequence
augmented_samples = run_compute_ukbb_gnn_get_augmented(ts_path)
# Compute GNNs for NUTS-2 and NUTS-3 mappings
print('Computing GNNs for NUTS-2')
gnn_nuts2_df = compute_gnn(ts_path, ukb_nuts2_df, augmented_samples)
print('Computing GNNs for NUTS-3')
gnn_nuts3_df = compute_gnn(ts_path, ukb_nuts3_df, augmented_samples)
# Save
gnn_nuts2_df.to_csv("gnn_nuts2.csv", index=False)
gnn_nuts3_df.to_csv("gnn_nuts3.csv", index=False)<jupyter_output><empty_output><jupyter_text>**2. Output files:*** `gnn_nuts2.csv` - GNN matrix computed using NUTS-2 areas as reference sets.* `gnn_nuts3.csv` - GNN matrix computed using NUTS-3 areas as reference sets. ---  3. Plot a heatmap and geographical representation of the GNN matrixSee Figure 5 on the tsinfer paper **3. Input files:*** `gnn_nuts2.csv` - GNN matrix computed using NUTS-2 areas as reference sets.* `gnn_nuts3.csv` - GNN matrix computed using NUTS-3 areas as reference sets.* `poly_df_nuts2` - Dataframe of NUTS-2 polygons created in (1.1).* `poly_df_nuts3` - Dataframe of NUTS-3 polygons created in (1.1).* `gadm36_IRL_shp/gadm36_IRL_0` - Republic of Ireland boundary vectors. The shapefile can be downloaded from the Database of Global Administrative Areas [here](https://gadm.org/download_country_v3.html) by inputting "Ireland" as country and selecting the 'Shapefile' link.<jupyter_code>import matplotlib
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import seaborn as sns
from scipy import stats
from scipy.cluster import hierarchy<jupyter_output><empty_output><jupyter_text>3.1 Pre-processing  3.1.1 Reformat GNN dataframes<jupyter_code># Function to average GNN proportions by NUTS area and z-score columns
def group_normalise_gnns(gnn_df):
    # Group rows by NUTS area and take the average of their GNN scores
    gnn_dfg = gnn_df.iloc[:,3:].groupby(gnn_df.place_of_birth).mean()
    # Z-score GNN proportions
    for col in list(gnn_dfg):
        gnn_dfg[col] = stats.zscore(gnn_dfg[col])
    return gnn_dfg

# Remove nan or empty columns and rows that do not correspond to a NUTS area.
# These result from individuals who have not been mapped to any NUTS-area
def clean_gnn(gnn_dfg, poly_df):
    for col in gnn_dfg.columns:
        if str(col) not in list(poly_df['name']):
            gnn_dfg = gnn_dfg.drop(col, axis=1)
    for index in gnn_dfg:
        if str(index) not in list(poly_df['name']):
            gnn_dfg = gnn_dfg.drop(index)
    return gnn_dfg
gnn_nuts2_dfg = clean_gnn(group_normalise_gnns(gnn_nuts2_df), poly_df_nuts2)
gnn_nuts3_dfg = clean_gnn(group_normalise_gnns(gnn_nuts3_df), poly_df_nuts3)<jupyter_output><empty_output><jupyter_text>3.1.2 Import boundary data for Ireland<jupyter_code>area_path_RI = "gadm36_IRL_shp/gadm36_IRL_0"
reader_RI = shpreader.Reader(area_path_RI)
geometries_RI = list(reader_RI.geometries())
names_RI = [record.attributes["NAME_0"] for record in reader_RI.records()]<jupyter_output><empty_output><jupyter_text>3.2 Plot the NUTS-2 GNN matrix heatmap and UK geographical map color-coded by hierarchical clustering See Figure 5A, 5B on the tsinfer paper  3.2.1 Cluster NUTS-2 areas by hierarchical clustering and divide the resulting tree into *n* subtrees The aim of this section is to cluster the reference areas in `gnn_dfg` by hierachical clustering and divide the resulting tree into n subtrees or clusters. In (3.2.3), we create a map of the UK where each area is colour-coded based on the subtree it falls in.<jupyter_code># Function that hierarchically clusters a dataframe and outputs a tree object.
def get_tree(df, method='average'):
    row_linkage = hierarchy.linkage(df, method=method)
    tree = hierarchy.to_tree(row_linkage, rd=True)
    return tree

# Function that takes as input a scipy tree object and a desired number of subtrees and outputs the list of nodes 
# at the root of each subtree. The function traverses the tree nodes starting from the root node and stops once
# it has traversed a number of nodes equal to the number of desired subtrees.
def get_cluster_nodes(tree, n_subtrees):
    all_nodes = tree[1][::-1]  
    nodes = []
    for node in all_nodes:
        if len(nodes)>=n_subtrees:break
        if node in nodes:
            nodes.remove(node)
        nodes.extend([node.get_left(), node.get_right()])
    return nodes

# Function that takes as input a node id and returns the ids of all leaf nodes under that node.
def get_node_leaves(node):
    if node.is_leaf():
        return [node.id]
    leaves = node.right.pre_order(lambda x: x.id)+node.left.pre_order(lambda x: x.id)
    return leaves

# Function to split a tree into n subtrees.
# The function takes as input a scipy hierarchical tree object and a desired number of subtrees.
# It outputs 2d list of node ids corresponding to the ids of the leaf nodes in each subtree.
def get_tree_clusters(tree, n_subtrees):
    nodes = get_cluster_nodes(tree, n_subtrees)
    leaves = []
    for node in nodes:
        leaves.append(get_node_leaves(node))
    return leaves
# Set number of subtrees
n_subtrees = 6
# Get tree object from hierarchical clustering of GNN matrix
tree_nuts2 = get_tree(gnn_nuts2_dfg)
# Divide tree into n subtrees (clusters) and get a list of the leaf nodes in each subtree
tree_clusters_nuts2 = get_tree_clusters(tree_nuts2, n_subtrees)<jupyter_output><empty_output><jupyter_text>3.2.2 Create a colourmap Here, we colour-code each NUTS-2 area based on the subtree it is located in.<jupyter_code># This function uses the tree object and the subtrees defined above and identifies the subtree each NUTS area is located in.
def get_nuts_cluster(tree_clusters, gnn_dfg, poly_df):
    poly_df['clusters'] = 0
    for i in range(len(tree_clusters)):
        for node_id in tree_clusters[i]:
            node = gnn_dfg.index.values[node_id]
#             index = poly_df.loc[poly_df['code']==node].index.values[0]
            index = poly_df[poly_df['name']==node].index.values[0]
            print(index)
            poly_df.loc[index,'clusters'] = i
    return poly_df

# This function creates a colourmap based on the subtrees identified above and assign to each NUTS area a colour
def get_colourmap(n_subtrees, poly_df):
    cmap = plt.cm.get_cmap('Set2', n_subtrees+1)
    colourmap = dict(zip(poly_df['code'], cmap(poly_df['clusters'])))
    return colourmap
poly_df_nuts2 = get_nuts_cluster(tree_clusters_nuts2, gnn_nuts2_dfg, poly_df_nuts2)
colourmap = get_colourmap(n_subtrees, poly_df_nuts2)<jupyter_output><empty_output><jupyter_text>3.2.3 Plot a geographical map of the UK<jupyter_code>fig, ax = plt.subplots(figsize=(10,20))
projection = ccrs.OSGB()
ax = plt.axes(projection=projection)
ax.coastlines(resolution='10m', linewidth=1.6)

# Plot polygons

#Default parameters
facecolor = [0.9375, 0.9375, 0.859375]
edgecolor = 'black'
linewidth=10
# Row-specific parameters
for index in poly_df_nuts2.index:
    name = poly_df_nuts2.at[index,'code']
    geometry = poly_df_nuts2.at[index,'geometry']
    facecolor = colourmap[name]
    ax.add_geometries([geometry], ccrs.PlateCarree(), facecolor=facecolor, edgecolor=edgecolor)    
    
# Add the coastline of Ireland
facecolor = 'white'
edgecolor = 'black'
name = names_RI
ax.add_geometries(geometries_RI, ccrs.PlateCarree(), facecolor=facecolor, edgecolor=edgecolor)

plt.savefig("UK_map.pdf")
plt.show()<jupyter_output><empty_output><jupyter_text>3.2.4 Plot a heatmap representation of the GNN matrix<jupyter_code># Function to rotate poition of branchpoints in the hierarchical tree
def rotate_linkage(linkage, index):
    x, y = linkage[index][0:2]
    linkage[index][0] = y
    linkage[index][1] = x
    
# Function to rename NUTS areas by code instead of name
def rename_gnn_sets(poly_df, gnn_dfg):
    name_code_dict = dict(zip(poly_df['name'], poly_df['code']))
    gnn_dfg.rename(columns=name_code_dict, inplace=True)
    gnn_dfg.rename(index=name_code_dict, inplace=True)
    return gnn_dfg
gnn_nuts2_dfg = rename_gnn_sets(poly_df_nuts2, gnn_nuts2_dfg)

row_linkage = scipy.cluster.hierarchy.linkage(gnn_nuts2_dfg, method="average")
rotate_linkage(row_linkage, -1)
order = scipy.cluster.hierarchy.leaves_list(row_linkage)
x_pop = gnn_nuts2_dfg.index.values[order]
row_colors = pd.Series(gnn_nuts2_dfg.reindex(x_pop).index, index=gnn_nuts2_dfg.reindex(x_pop).index).map(colourmap)

cg = sns.clustermap(gnn_nuts2_dfg[x_pop], row_linkage=row_linkage, col_cluster=False, row_colors=row_colors, figsize=(20, 20))
cg.ax_row_dendrogram.set_visible(False)  # Set to True to display the dendrogram
plt.savefig("gnn_heatmap.pdf")<jupyter_output><empty_output><jupyter_text>**3.2 Output files:*** `UK_map.pdf` - Geographical map of the UK divided into NUTS-2 areas, where each area is colour-coded based on its position in the hierarchical tree.* `gnn_heatmap.pdf` - Heatmap plot of the NUTS-2 GNN matrix.  3.3 Plot the geographical map of a specific row of the NUTS-3 GNN matrixSee Figure 5 C,D,E on the tsinfer paper In this part we use `gnn_nuts3_dfg`to plot a geographical representation of the GNN proportions of a specific NUTS-3 area.<jupyter_code>def transpose_df(df):
    df_t = df.transpose()
    df_t.index.name = df.index.name
    return df_t
# Transpose GNN matrix to plot a column. This facilitates the manipulation of the dataframe.
gnn_nuts3_dfg_t = transpose_df(gnn_nuts3_dfg)
def plot_GNN_row(gnn_dfg_t, poly_df, nuts_area):
    
    # Extract column containing to the GNN proportions of the indicated NUTS area
    area_gnn = poly_df.merge(gnn_dfg_t[nuts_area], left_on='name', right_on='place_of_birth')

    # Define log-normalization scale based on min and max values in the gnn matrix
    vmin = gnn_dfg_t.loc[gnn_dfg_t.stack().idxmin()]
    vmax = gnn_dfg_t.loc[gnn_dfg_t.stack().idxmax()]
    diff = 1-vmin     # translate values by 1
    norm = colors.LogNorm(vmin=vmin+diff, vmax=vmax+diff)
    
    # Create colourmap 
    cmap = plt.get_cmap('pink_r') # more colormaps here: https://matplotlib.org/3.1.0/gallery/color/colormap_reference.html
    colours = cmap(norm((area_gnn[nuts_area].fillna(0).values)+diff))
    colourmap = dict(zip(area_gnn['name'], colours))
    
    # Plot
    fig, ax = plt.subplots(figsize=(10,20))
    projection = ccrs.OSGB()
    ax = plt.axes(projection=projection)
    edgecolor = None
    ax.coastlines(resolution='10m', linewidth=0.8)
    # Plot polygons
    # Default parameters
    facecolor = [0.9375, 0.9375, 0.859375]
    edgecolor = None
    # Row-specific parameters
    for index in area_gnn.index:
        facecolor = colours[index]
        geometry = area_gnn.at[index,'geometry']
        ax.add_geometries([geometry], ccrs.PlateCarree(),facecolor=facecolor, edgecolor=edgecolor)
    
    # Add Ireland coastline
    facecolor = 'white'
    edgecolor = None
    name = names_RI
    ax.add_geometries(geometries_RI, ccrs.PlateCarree(), facecolor=facecolor, edgecolor=edgecolor)
    
    names = list(poly_df['name'])
    geometries = list(poly_df['geometry'])
    area_index = names.index(nuts_area)
    point = geometries[area_index].centroid
    x = point.coords[0][0]
    y = point.coords[0][1]
    ax.text(x, y, nuts_area, fontsize=14, fontweight='bold', transform=ccrs.PlateCarree())
    
    # Add colourbar    
    mapper = matplotlib.cm.ScalarMappable(norm=norm, cmap=cmap)
    mapper.set_array(list(area_gnn[nuts_area]))
    plt.colorbar(mapper, shrink=0.6, ticks=[1,2,4,6,8,10,12,14], format='%.0f')
    
    plt.savefig("gnn_row_map_"+nuts_area+".pdf")
    plt.show()
# Plot row of the GNN matrix
# Input NUTS-3 area name below. A list of all area names can be found in the dataframe poly_df_nuts3 or in the file "UK_NUTS_Level_3_January_2018_Names_and_Countries.csv".
nuts_area = "North Yorkshire CC" 
plot_GNN_row(gnn_nuts3_dfg_t, poly_df_nuts3, nuts_area)<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup