1""" 

2@file 

3@brief Function to test others functionalities 

4""" 

5import os 

6import pandas 

7from pyquickhelper.loghelper import fLOG 

8from ..faq.faq_matplotlib import graph_cities 

9from ..special import tsp_kruskal_algorithm, distance_haversine 

10 

11 

12def american_cities(df_or_filename, nb_cities=-1, img=None, fLOG=fLOG): 

13 """ 

14 Computes the :epkg:`TSP` for american cities. 

15 

16 @param df_or_filename dataframe 

17 @param nb_cities number of cities to keep 

18 @param img image to produce 

19 @param fLOG logging function 

20 @return dataframe (results) 

21 """ 

22 def haversine(p1, p2): 

23 return distance_haversine(p1[0], p1[1], p2[0], p2[1]) 

24 

25 if isinstance(df_or_filename, str): 

26 df = pandas.read_csv(df_or_filename) 

27 else: 

28 df = df_or_filename 

29 df["Longitude"] = -df["Longitude"] 

30 df = df[df.Latitude < 52] 

31 df = df[df.Longitude > -130].copy() 

32 fLOG(df.columns) 

33 df = df.dropna() 

34 

35 if nb_cities > 0: 

36 df = df[:nb_cities].copy() 

37 

38 fLOG(df.shape) 

39 points = [(row[1], row[2], row[3]) 

40 for row in df.itertuples(index=False)] 

41 fLOG("number of cities:", len(points)) 

42 trip = tsp_kruskal_algorithm( 

43 points, distance=haversine, fLOG=fLOG, max_iter=10) 

44 

45 # trip 

46 dftrip = pandas.DataFrame( 

47 trip, columns=["Latitude", "Longitude", "City"]) 

48 

49 # graph 

50 for i in range(0, dftrip.shape[0]): 

51 if i % 10 != 0: 

52 dftrip.loc[i, "City"] = "" 

53 

54 if img is not None: 

55 import matplotlib.pyplot as plt 

56 fig, ax = graph_cities(dftrip, markersize=3, linked=True, fLOG=fLOG, 

57 fontcolor="red", fontsize='16', loop=True, figsize=(32, 32)) 

58 assert ax is not None 

59 fig.savefig(img) 

60 assert os.path.exists(img) 

61 plt.close('all') 

62 fLOG("end") 

63 return dftrip