TDM 40100: Project 13 - XGBoost

import pandas as pd
worldbank_data = pd.read_csv("/anvil/projects/tdm/data/worldbank/worldbank_data.csv")

# TODO: Display first five rows

# Sort by country and year
worldbank_data = worldbank_data.sort_values(by=["country", "year"])

# Define numeric columns (excluding 'year')
numeric_cols = worldbank_data.select_dtypes(include="number").columns.difference(["year"])

# Interpolate only values still missing
for col in numeric_cols:
    mask = worldbank_data[col].isna()
    interpolated = (
        worldbank_data
        .groupby("country")[col]
        .transform(lambda group: group.interpolate(method=____, limit_direction=____).ffill().bfill()) # For YOU to fill in
    )
    worldbank_data.loc[mask, col] = interpolated[mask]

# Drop non-numeric column
numeric_data = worldbank_data.drop(columns=["country"])

# Compute correlation matrix
correlations = numeric_data.corr()

# Extract correlations with fertility_rate
fert_corr = correlations["fertility_rate"].sort_values(ascending=False)

# TODO: Report the 5 features most positively and most negatively correlated with fertility rate.

latest_data = worldbank_data.loc[worldbank_data.groupby("country")["year"].idxmax()]

fertility_by_country = latest_data[["____", "____", "____"]] # For YOU to fill in with correct column names

import geopandas as gpd, matplotlib.pyplot as plt, pandas as pd

# Load, merge, and plot
import geopandas as gpd, matplotlib.pyplot as plt, pandas as pd

# Load, merge, and plot
(gpd.read_file("https://raw.githubusercontent.com/nvkelso/natural-earth-vector/master/geojson/ne_110m_admin_0_countries.geojson")
 .assign(ADMIN_norm=lambda d: d["ADMIN"].str.lower())
 .merge(fertility_by_country[fertility_by_country["year"] == 2024]
        .assign(fertility_rate=lambda d: pd.to_numeric(d["fertility_rate"], errors="coerce"),
                name_norm=lambda d: d["country"].str.lower().replace({
                    "united states": "united states of america",
                    "north macedonia": "republic of north macedonia",
                    "kyrgyz republic": "kyrgyzstan"}))
        .rename(columns={"country": "name"})
        .query("~name.str.contains('income|world|OECD|IDA|IBRD|region|fragile', case=False)"),
        left_on="ADMIN_norm", right_on="name_norm", how="left")
 [lambda d: ~d["ADMIN"].isin(["Antarctica", "Falkland Islands", "French Southern and Antarctic Lands"])]
 .plot(column="fertility_rate", cmap="YlOrRd", legend=True, figsize=(15, 8), missing_kwds={"color": "lightgrey"}))

plt.title("Fertility Rate by Country (2024)")
plt.axis("off")
plt.show()

from sklearn.model_selection import train_test_split

# correlations with fertility_rate
correlations = worldbank_data.corr(numeric_only=True)["fertility_rate"]

# Select the top 5 positively correlated features (excluding fertility_rate itself)
top5_features = correlations.drop("fertility_rate").sort_values(ascending=False).head(5).index.tolist()

# Define X and y using those features
X_small = worldbank_data[top5_features]
y = worldbank_data["fertility_rate"]

# Split into 80% train/ 20% test sets
X_train, X_test, y_train, y_test = train_test_split(X_small, y, test_size=______, random_state=____) # For YOU to fill in

from xgboost import XGBRegressor

# Only fill in the missing piece use the correct parameter

small_model = XGBRegressor(random_state=__) # For YOU to fill in

# Fit the model on training data
small_model.fit(______________, ______________) # For YOU to fill in

from sklearn.metrics import mean_squared_error, r2_score

# Predict on the test set
y_pred = small_model.predict(X_test)

# Calculate RMSE and R^2
rmse = mean_squared_error(y_test, y_pred) ** 0.5
r2 = r2_score(y_test, y_pred)

# TODO: Print the RMSE and R^2 values

from sklearn.model_selection import train_test_split

# Define features and target
X_full = worldbank_data.drop(columns=["__________", "__________"]) # For YOU to fill in
y = worldbank_data["__________"] # For YOU to fill in

# Split into training and test sets
X_train_full, X_test_full, y_train_full, y_test_full = train_test_split(
    X_full, y, test_size=___, random_state=__) # For YOU to fill in

from xgboost import XGBRegressor, plot_importance
from sklearn.metrics import mean_squared_error, r2_score
import matplotlib.pyplot as plt

# Initialize and train model with pretuned parameters
tuned_model = XGBRegressor(
    random_state=____, # For YOU to fill in
    max_depth=____, # For YOU to fill in
    learning_rate=____, # For YOU to fill in
    n_estimators=______ # For YOU to fill in
)
tuned_model.fit(X_train_full, y_train_full)

y_pred = tuned_model.predict(X_test_full)
rmse = mean_squared_error(y_test_full, y_pred)**0.5
r2 = r2_score(y_test_full, y_pred)

# TODO: Print the RMSE and R^2 values

from xgboost import plot_importance
import matplotlib.pyplot as plt

plt.figure(figsize=(10, 6))
plot_importance(tuned_model, max_num_features=__) # For YOU to fill in
plt.title("________") # For YOU to fill in
plt.tight_layout()
plt.show()

import matplotlib.pyplot as plt

countries_to_plot = ["_________", "______"] # For YOU to fill in

plt.figure(figsize=(10, 6))
for country in countries_to_plot:
    data = worldbank_data[worldbank_data["country"] == country]
    plt.plot(data["year"], data["fertility_rate"], label=country)

plt.title("_____") # For YOU to fill in
plt.xlabel("_____") # For YOU to fill in
plt.ylabel("_______") # For YOU to fill in
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()