Speed up subset of python dictionary

I have optimized my code using numba but one line of my code which basically subset a large dictionary is taking more than 99% of the computation time of the code (among the lines i would like to optimize). The line of code is the following :

all_weight = np.fromiter(map(dic_transition.get, all_trans), dtype=np.float)

Are there any ways to make this operation faster ?

Here is the whole script that can be reproduced:

import networkx as nx
import numpy as np
import itertools
import pandas as pd
import numba
import osmnx as ox
from sklearn.neighbors import BallTree, DistanceMetric

def swap_lon_lat(array):
    array_bis = np.copy(array)
    array_bis[:,[0, 1]] = array_bis[:,[1, 0]]
    return(array_bis)

def get_emission_matrix(MP_traj, df_coord, tree, radius = 0.500/6371, sigma=0.5):
    MP_traj_swaped = np.radians(swap_lon_lat(MP_traj))
    neighbors, dist_array  = tree.query_radius(MP_traj_swaped, radius, return_distance=True)
    emit_p = np.zeros((len(MP_traj), len(df_coord)))
    for i in range(len(MP_traj_swaped)):
        np.put(emit_p[i, :], neighbors[i], dist_array[i] *6371)

    max_nb_neighbors = max([len(i) for i in neighbors])
    new_neighbors = np.zeros((len(MP_traj), max_nb_neighbors))

    for i in range(len(MP_traj_swaped)):
        np.put(new_neighbors[i, :], list(range(len(neighbors[i]))), neighbors[i])
    new_neighbors = new_neighbors.astype(int)
    nb_points = np.shape(new_neighbors)[0]
    emit_p = fill_emission_matrix(emit_p, new_neighbors, nb_points, sigma=sigma)
    return(emit_p, neighbors)

@numba.jit(nopython=True)
def fill_emission_matrix(emit_p, new_neighbors, nb_points, sigma=0.1):
    for i in range(nb_points):
        for j in new_neighbors[i]:
            if j!=0:
                emit_p[i,j] = np.exp(-(emit_p[i,j]/(np.sqrt(2) * sigma))**2)/(np.sqrt(2*np.pi) * sigma)
    return(emit_p)

def get_transition_matrix(neighbors, dic_transition, sub_nodes):
    all_trans = [itertools.product(neighbors[i], neighbors[i+1]) for i in range(len(neighbors)-1)]
    all_trans = itertools.chain.from_iterable(all_trans)
    all_weight = np.fromiter(map(dic_transition.get, all_trans), dtype=np.float)
    all_weight[np.isnan(all_weight)] = 0
    return(all_weight)

radius, sigma = 1.5/6371, 0.5
MP_traj = np.array([[45.79752 ,  4.928819],[45.796373,  4.928486],[45.795357,  4.927367],[45.791835,  4.912801], [45.784338,  4.895901]])
graph = ox.graph_from_bbox(45.808288928626446, 45.77078558386208,  4.945873034649236,  4.876255464119831, network_type='drive')
ListNodes = sorted(graph.nodes())

dic_transition = {}
for id_node, ori_node in enumerate(ListNodes):
    trans_dict = nx.single_source_dijkstra_path_length(graph, ori_node, cutoff=5000, weight='length')
    for dest_node, weigth in trans_dict.items():
        dic_transition[(ori_node, dest_node)] = np.int32(1000 * np.exp(-weigth/5000))

df_coord = pd.DataFrame([[graph.nodes[i]['x'], graph.nodes[i]['y']] for i in ListNodes], columns=['y','x']).values
tree = BallTree(np.radians(swap_lon_lat(df_coord)), metric=DistanceMetric.get_metric('haversine'))

# Part of the script that i would like to optimize
getNodeID = dict(zip(range(len(ListNodes)), ListNodes))
emit_p, neighbors = get_emission_matrix(swap_lon_lat(MP_traj), df_coord, tree, radius = radius, sigma=sigma)
neighbors = np.array([np.array([getNodeID[j] for j in i]) for i in neighbors])
emit_p_bis = np.take(emit_p, np.where(np.sum(emit_p, axis=0) != 0)[0], axis=1)
sub_nodes = np.array([getNodeID[i] for i in  np.where(np.sum(emit_p, axis=0) != 0)[0]])
all_weight = get_transition_matrix(neighbors, dic_transition, sub_nodes)

Thanks in advance.