{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ae02282-e75f-40aa-a12b-93e2d7ed81f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "Zircon Geochemical Classification Script\n",
    "#\n",
    "# This script performs a machine learning-based classification of zircon samples\n",
    "# using geochemical data. It trains and evaluates six different classification\n",
    "# models to predict the rock type (Igneous, Metamorphic, or Sedimentary)\n",
    "# of zircon samples. The script also includes a workflow for predicting the\n",
    "# rock type of new, unknown samples from the Roberts Database and exports\n",
    "# the results to a CSV file.\n",
    "#\n",
    "# =============================================================================\n",
    "\n",
    "# -----------------------------------------------------------------------------\n",
    "# 1. IMPORT NECESSARY LIBRARIES\n",
    "# -----------------------------------------------------------------------------\n",
    "# pandas is used for data manipulation and analysis.\n",
    "import pandas as pd\n",
    "# numpy is used for numerical operations, especially for arrays.\n",
    "import numpy as np\n",
    "# matplotlib and seaborn are used for data visualization.\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "# The following classes and functions from scikit-learn are used for\n",
    "# machine learning tasks such as data splitting, scaling, and model training/evaluation.\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.ensemble import GradientBoostingClassifier, RandomForestClassifier\n",
    "from sklearn.tree import DecisionTreeClassifier\n",
    "from sklearn.svm import SVC\n",
    "from sklearn.neighbors import KNeighborsClassifier\n",
    "from sklearn.naive_bayes import GaussianNB\n",
    "from sklearn.metrics import accuracy_score, classification_report\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "\n",
    "# -----------------------------------------------------------------------------\n",
    "# 2. DATA LOADING AND PREPARATION\n",
    "# -----------------------------------------------------------------------------\n",
    "\n",
    "# Load the training data from the specified CSV file.\n",
    "# The file path is hardcoded as provided.\n",
    "file_path = '/Users/Ian/Documents/USGS/Mojave:Mountain Pass/Zircon ML/ZirconMLCode/ZrnLog_3.csv'\n",
    "data = pd.read_csv(file_path)\n",
    "\n",
    "# Define the geochemical features (X) and the target variable (y).\n",
    "# The features are log-transformed elemental ratios and concentrations.\n",
    "features = ['log10_LuHf', 'log10_ThU', 'log10_GdYb', 'log10_UNb', 'log10_UYb',\n",
    "            'log10_ThNb', 'log10_NbYb', 'log10_SmYb', 'log10_Eu',\n",
    "            'log10_ThYb', 'log10_SumHREE', 'log10_CeU']\n",
    "# The target variable is 'RockType', representing the rock of origin.\n",
    "target = 'RockType'\n",
    "\n",
    "# Prepare training data by separating features and target.\n",
    "X = data[features]\n",
    "y = data[target]\n",
    "\n",
    "# -----------------------------------------------------------------------------\n",
    "# 3. DATA SPLITTING AND SCALING\n",
    "# -----------------------------------------------------------------------------\n",
    "\n",
    "# Split the data into training (80%) and testing (20%) sets.\n",
    "# stratify=y ensures that the proportion of each rock type is the same\n",
    "# in both the training and testing sets, which is crucial for imbalanced datasets.\n",
    "# random_state=42 ensures reproducibility of the split.\n",
    "X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)\n",
    "\n",
    "# Initialize the StandardScaler to standardize the numerical features.\n",
    "# StandardScaler removes the mean and scales to unit variance.\n",
    "scaler = StandardScaler()\n",
    "# Fit the scaler on the training data and transform it.\n",
    "X_train_scaled = scaler.fit_transform(X_train)\n",
    "# Transform the test data using the *same* scaler to prevent data leakage from the test set.\n",
    "X_test_scaled = scaler.transform(X_test)\n",
    "\n",
    "# -----------------------------------------------------------------------------\n",
    "# 4. MODEL INITIALIZATION, TRAINING, AND EVALUATION\n",
    "# -----------------------------------------------------------------------------\n",
    "\n",
    "# Initialize a dictionary of the six classification models to be evaluated.\n",
    "# Note: SVC(probability=True) is used to enable the predict_proba method for confidence scores.\n",
    "models = {\n",
    "    'Naive Bayes': GaussianNB(),\n",
    "    'Decision Tree': DecisionTreeClassifier(),\n",
    "    'Random Forest': RandomForestClassifier(),\n",
    "    'Support Vector Machine': SVC(probability=True),\n",
    "    'k-Nearest Neighbors': KNeighborsClassifier(),\n",
    "    'Gradient Boosting': GradientBoostingClassifier(),\n",
    "}\n",
    "\n",
    "# Initialize dictionaries to store results, predictions, and confidence scores.\n",
    "results = {}\n",
    "predictions = {}\n",
    "confidence_scores = {}\n",
    "\n",
    "# Loop through each model to train and evaluate it.\n",
    "for model_name, model in models.items():\n",
    "    # Train the model. SVC uses scaled data, while tree-based models and Naive Bayes\n",
    "    # generally perform well with unscaled data.\n",
    "    if model_name == 'Support Vector Machine' or model_name == 'k-Nearest Neighbors':\n",
    "        model.fit(X_train_scaled, y_train)\n",
    "        y_pred = model.predict(X_test_scaled)\n",
    "    else:\n",
    "        model.fit(X_train, y_train)\n",
    "        y_pred = model.predict(X_test)\n",
    "    \n",
    "    # Calculate and print the model's accuracy and a detailed classification report.\n",
    "    accuracy = accuracy_score(y_test, y_pred)\n",
    "    results[model_name] = accuracy\n",
    "    print(f'{model_name} Accuracy: {accuracy:.2f}')\n",
    "    print(classification_report(y_test, y_pred))\n",
    "    \n",
    "    # Store predictions and the corresponding confidence scores (max probability).\n",
    "    if hasattr(model, 'predict_proba'):\n",
    "        probabilities = model.predict_proba(X_test_scaled if model_name in ['Support Vector Machine', 'k-Nearest Neighbors'] else X_test)\n",
    "        predictions[model_name] = y_pred\n",
    "        confidence_scores[model_name] = np.max(probabilities, axis=1)\n",
    "\n",
    "# -----------------------------------------------------------------------------\n",
    "# 5. UNKNOWN SAMPLE CLASSIFICATION\n",
    "# -----------------------------------------------------------------------------\n",
    "\n",
    "# Load the unknown samples from the specified CSV file.\n",
    "unknown_file_path = r'/Users/Ian/Documents/USGS/Mojave:Mountain Pass/Zircon ML/ZirconMLCode/RobertsDatabaseClassifier.csv'\n",
    "unknown_data = pd.read_csv(unknown_file_path)\n",
    "\n",
    "# Screen and clean the unknown data.\n",
    "# This ensures that only rows with valid numerical feature values are used for prediction.\n",
    "unknown_data_filtered = unknown_data.copy()\n",
    "unknown_data_filtered = unknown_data_filtered[features]  # Select only the feature columns.\n",
    "# Coerce non-numeric values to NaN, and then drop rows with NaNs.\n",
    "unknown_data_filtered = unknown_data_filtered.apply(pd.to_numeric, errors='coerce')\n",
    "unknown_data_filtered.dropna(inplace=True)\n",
    "# Keep track of the original row indices to link back to other data.\n",
    "valid_indices = unknown_data_filtered.index\n",
    "unknown_data_cleaned = unknown_data.loc[valid_indices]\n",
    "\n",
    "# Prepare the cleaned unknown data for prediction by applying the same scaling.\n",
    "X_unknown = unknown_data_filtered\n",
    "X_unknown_scaled = scaler.transform(X_unknown)\n",
    "\n",
    "# Create a DataFrame to store the classification results for the unknown samples.\n",
    "unknown_results = pd.DataFrame(index=unknown_data_cleaned.index)\n",
    "# Include key metadata from the original unknown data.\n",
    "unknown_results['Best Age'] = unknown_data_cleaned['Best Age']\n",
    "unknown_results['Reference'] = unknown_data_cleaned['Reference']\n",
    "\n",
    "# Predict the rock type for the unknown samples using each trained model.\n",
    "for model_name, model in models.items():\n",
    "    # Use scaled data for SVC and KNN, and unscaled for others.\n",
    "    if model_name == 'Support Vector Machine' or model_name == 'k-Nearest Neighbors':\n",
    "        y_unknown_pred = model.predict(X_unknown_scaled)\n",
    "        probabilities = model.predict_proba(X_unknown_scaled)\n",
    "    else:\n",
    "        y_unknown_pred = model.predict(X_unknown)\n",
    "        probabilities = model.predict_proba(X_unknown)\n",
    "    \n",
    "    # Store the predicted rock type and the corresponding confidence score.\n",
    "    unknown_results[model_name] = y_unknown_pred\n",
    "    unknown_results[model_name + ' Confidence'] = np.max(probabilities, axis=1)\n",
    "\n",
    "# Display the classification results and save them to a new CSV file.\n",
    "print(\"\\nUnknown Samples Classification Results:\")\n",
    "print(unknown_results)\n",
    "unknown_results.to_csv('/Users/Ian/Documents/USGS/Mojave:Mountain Pass/Zircon ML/ZirconMLCode/ZirconResults_RockType_Roberts.csv', index=True)\n",
    "\n",
    "# -----------------------------------------------------------------------------\n",
    "# 6. VISUALIZATION\n",
    "# -----------------------------------------------------------------------------\n",
    "\n",
    "# Plot feature importances for tree-based models.\n",
    "feature_importances = {name: m.feature_importances_ for name, m in models.items() if hasattr(m, 'feature_importances_')}\n",
    "plt.figure(figsize=(12, 6))\n",
    "for model_name, importances in feature_importances.items():\n",
    "    plt.barh(features, importances, alpha=0.5, label=model_name)\n",
    "plt.title('Feature Importances from Tree-Based Models')\n",
    "plt.xlabel('Importance')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "# Plot a bar chart comparing the accuracy of all models.\n",
    "plt.figure(figsize=(12, 6))\n",
    "sns.barplot(x=list(results.keys()), y=list(results.values()))\n",
    "plt.title('Model Accuracy Comparison')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.xticks(rotation=45)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:anaconda3]",
   "language": "python",
   "name": "conda-env-anaconda3-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}