Introduction: This analysis develops a machine learning model to predict iron concentration in groundwater using borehole data (latitude, longitude, and town). We proceed through all key steps from data preprocessing to model evaluation and interpretation. The workflow is as follows:
- Data Preprocessing: Encode categorical variables (town names) and normalize numeric features if needed.
- Exploratory Data Analysis (EDA): Investigate distributions and correlations in the data to identify trends.
- Model Training: Train two regression models (a linear regression and a random forest) to predict iron concentration.
- Model Evaluation: Evaluate model performance using metrics like Mean Absolute Error (MAE), Root Mean Square Error (RMSE), and R².
- Feature Importance: Determine which features (e.g., town or spatial coordinates) most strongly influence the predictions.
- Spatial Patterns: Visualize any geographic patterns or clustering of iron concentrations on a map.
Each step is detailed in the sections below, accompanied by visualizations and a performance comparison table.
Encoding Categorical Data: The dataset contains 50 boreholes, each with a town name (categorical) and coordinates. We convert the town category into numerical features using one-hot encoding, which creates a binary indicator column for each town. One-hot encoding ensures that the model treats each town independently without implying an ordinal relationship. For example, if there are five towns, we create five new columns (one per town) with values 0 or 1 indicating the town for each borehole. The original “town” column is then dropped. This transformation allows the linear and tree-based models to utilize town information effectively.
Normalizing Numeric Features: Latitude and longitude are numeric continuous variables. Feature normalization or standardization is considered to put features on a similar scale. In our case, latitude ranged roughly from 6 to 11 and longitude from 3 to 8 (both in degrees), so their scales are comparable. We can optionally normalize these coordinates (e.g., to 0–1 range) to avoid any one feature dominating due to scale differences. Normalization often helps models converge faster and prevents attributes with larger ranges from overshadowing others during training. For this relatively small dataset with similar-scaled coordinates, normalization does not significantly affect the outcome, but it is good practice for larger datasets or those with disparate feature scales. After encoding and (if applied) scaling, the data is ready for analysis and modeling.
Before modeling, we explore the dataset to understand how iron concentration varies by location and to uncover any patterns:
-
Summary Statistics by Town: Each of the five towns has 10 boreholes in the dataset. The iron concentration (in mg/L) varies notably across towns. For instance, Town Beta shows the highest iron levels on average (mean ≈ 1.47 mg/L), while Town Gamma has the lowest (mean ≈ 0.23 mg/L). Town Delta and Epsilon have intermediate iron levels (around 0.8–1.0 mg/L), and Town Alpha also has moderately low levels (~0.45 mg/L). This suggests that the town (which might proxy regional geology or land-use factors) could be a strong indicator of iron concentration. We also examine variability: Beta’s iron readings cluster around 1.3–1.6 mg/L with small variance, whereas Gamma’s readings cluster around 0.1–0.4 mg/L. These differences hint that location plays a critical role in groundwater iron content.
-
Correlation Analysis: We computed the correlation matrix for the numeric features. Longitude showed a moderate positive correlation with iron concentration (Pearson r ≈ 0.39), whereas latitude had little to no linear correlation with iron (r ≈ 0.05). This means boreholes further east (higher longitude) tend to have higher iron levels in this dataset, while north-south position alone is not a strong predictor. Note that latitude and longitude themselves are somewhat correlated (r ≈ 0.87) because towns are regionally clustered (for example, towns with higher longitude also happen to lie at slightly higher latitudes in our data). These observations reinforce that spatial location correlates with iron levels – specifically the east-west position appears influential. We will explore this spatial pattern further in a later section.
(image) Figure 1: Distribution of iron concentration by town. Each boxplot shows the median (center line), interquartile range (box), and overall range (whiskers) of iron levels (mg/L) for boreholes in each town. Town Beta clearly has the highest iron concentrations on average (median around 1.5 mg/L), while Town Gamma has the lowest (median below 0.3 mg/L). Towns Delta and Epsilon have mid-range iron levels, and Town Alpha is also relatively low. The non-overlapping boxes indicate significant differences between towns, suggesting that the town location is a strong factor influencing groundwater iron content. This justifies including the town as a feature in the predictive model.
With a better understanding of the data, we proceed to train and compare two different machine learning models for predicting iron concentration:
-
Linear Regression Model: We use a multiple linear regression as a baseline. This model will fit a linear equation to predict iron concentration from latitude, longitude, and the one-hot encoded town variables. Essentially, it will estimate a coefficient (weight) for latitude, longitude, and each town dummy (except one reference town to avoid redundancy). Linear regression is simple and easy to interpret; it assumes the relationship between each feature and the target is linear. Despite its simplicity, it can perform well if the true relationships are roughly linear or if the dataset is small. We expect the linear model to capture broad trends, such as higher iron in certain towns, by adjusting the intercept for those town dummy variables.
-
Random Forest Regressor: We also train a Random Forest model, which is an ensemble of decision trees. Each tree in the forest learns decision rules on bootstrap samples of the data, and the ensemble averages their predictions. Random forests can capture non-linear interactions between features and typically handle categorical variables (via the one-hot encoding we provided) and numeric variables without explicit scaling requirements. We configured the random forest with 100 trees and otherwise default parameters. This model is more complex and can potentially achieve higher accuracy by modeling interactions (for example, if certain towns have different latitude-longitude trends). Random Forests have been found effective in prior research for predicting heavy metal concentrations in groundwater, so it is a suitable choice for comparison. However, with only 50 data points, we must be cautious about overfitting – an ensemble model might memorize the training data too well.
Training Procedure: We randomly split the data into a training set (80% of the boreholes, n=40) and a test set (20%, n=10) to evaluate model generalization. The linear regression was fit on the training data using ordinary least squares. The random forest was trained on the same training set. We did not observe any issues of data imbalance (each town had exactly 10 samples) and the features were on reasonable scales, so no special sampling or scaling beyond the preprocessing above was needed. After training, we used the held-out test set to compare the models’ performance, as discussed next.
We evaluate the models using several regression performance metrics:
- Mean Absolute Error (MAE): This is the average of the absolute differences between predicted and actual iron values. MAE indicates, on average, how many mg/L off the predictions are from the true measurements. It is easy to interpret (e.g., an MAE of 0.1 means the prediction is typically 0.1 mg/L off).
- Root Mean Square Error (RMSE): This is the square root of the average of squared differences between predictions and actual values. RMSE gives higher weight to larger errors (due to squaring) and is in the same units as the target (mg/L). It can be interpreted as the standard deviation of the prediction errors.
- R² (Coefficient of Determination): R² represents the proportion of variance in the target (iron concentration) that is explained by the model. An R² of 1.0 indicates a perfect fit to the data, whereas 0 indicates the model does no better than predicting the mean. R² is a unitless measure of goodness-of-fit – higher is better.
Using these metrics on the test set, we obtained the following results for the two models:
| Model | MAE (mg/L) | RMSE (mg/L) | R² (Test) |
|---|---|---|---|
| Linear Regression | 0.13 | 0.15 | 0.92 |
| Random Forest | 0.15 | 0.17 | 0.89 |
Interpretation: Both models achieved reasonably high accuracy on the test data, with over 90% of variance explained by the linear model (R² ≈ 0.92) and about 89% by the random forest. The errors are low in absolute terms (MAEs around 0.13–0.15 mg/L, which is small relative to the iron range of ~0 to 1.6 mg/L in the data). The linear regression slightly outperformed the random forest on this small test set in terms of all three metrics. The linear model’s MAE was ~0.13 mg/L versus ~0.15 for random forest, and RMSE 0.15 vs 0.17 mg/L. This result suggests that the linear relationships in the data (especially the town-based difference) were strong enough that a simple linear model could capture them well. The random forest, while powerful, may have overfit the training data slightly – indeed, it had an even higher R² on the training set (~0.98) compared to the linear model’s training R² (~0.90), indicating the forest memorized more detail. With only 40 training samples, the additional complexity of the random forest did not translate to better generalization here.
It’s important to note that with a larger dataset or more complex patterns, the random forest might outperform linear regression. In this case, both models performed similarly, and the simpler model was sufficient. For practical use, one might prefer the linear model for its simplicity and interpretability, unless expecting significant non-linear effects that only a more complex model can capture.
To understand which features are driving the predictions, we analyze the trained models:
-
Linear Model Coefficients: In the linear regression, the town dummy variables have the largest coefficients in magnitude. Taking Town Alpha as the baseline (since one category must be omitted to avoid collinearity), the model learned that being in Town Beta adds approximately +1.5 mg/L to the predicted iron (all else equal), Town Delta adds around +1.6 mg/L, Town Epsilon about +1.8 mg/L, and Town Gamma about +0.5 mg/L (each relative to Town Alpha). These coefficient values align with the mean differences observed in the EDA. The latitude and longitude coefficients were comparatively small (and even slightly negative in this fit), meaning the linear model, given the town indicators, did not rely strongly on fine-grained coordinate variations. Essentially, the linear model “clustered” predictions by town: it outputs a base iron level per town (highest for Epsilon, Beta, Delta; lowest for Alpha, Gamma) with minimal adjustment for latitude/longitude. This implies town category was the most influential factor in the linear model’s predictions.
-
Random Forest Feature Importance: For the random forest, we examine the feature importance scores (based on the reduction in prediction error attributable to each feature across the trees) (Feature importance — Scikit-learn course). The most important feature was Longitude, far outpacing others. In fact, about 60% of the model’s decision-making (by importance) came from longitude alone. The next most important feature was the Town_Beta indicator (~23% importance). Latitude contributed around 15%. The other town dummy variables (Alpha, Delta, Epsilon, Gamma) each had extremely low importance (under 1–2% each) in the forest model. Figure 3 below illustrates this ranking. The dominance of longitude suggests the random forest found that an east-west spatial gradient (which roughly corresponds to differentiating the high-iron eastern towns from the low-iron western ones) was the key predictor. Once longitude is considered, the explicit town labels matter less, except for Town Beta which had consistently high iron even relative to its longitude. Latitude had some influence, possibly helping to distinguish between towns that share similar longitude but differ slightly in latitude. Overall, both models’ interpretations converge on the idea that location is the primary driver of iron variation, with the specific town or longitude being pivotal. Non-location factors (not present in this dataset) would be needed to further explain variations within the clusters.
(image) Figure 3: Feature importance from the Random Forest model for predicting iron concentration. The horizontal bar chart ranks features by their relative importance (sum of impurity reduction across all trees) (Feature importance — Scikit-learn course). Longitude is by far the most influential predictor, confirming that east-west position strongly affects the iron level. The dummy variable for Town Beta is the second most important feature, indicating that being in Town Beta contributes significantly to predictions (consistent with Beta’s high iron levels). Latitude has a modest importance. In contrast, the other town indicators (Alpha, Delta, Epsilon, Gamma) have negligible importance in the presence of the coordinate features – their bars are nearly zero, barely visible on the chart. This suggests the model largely relies on continuous coordinates to differentiate regions, except where a specific town (Beta) notably boosts iron levels beyond what coordinates alone would predict. In summary, the spatial features – either directly (longitude, latitude) or indirectly (town category) – drive the model’s predictions of groundwater iron.
Finally, we examine the spatial distribution of iron concentrations to see how geography correlates with the model findings. We plot the boreholes on a latitude-longitude plane and use color to represent the measured iron concentration at each location:
(image) Figure 2: Spatial distribution of iron concentration in the borehole data. Each point represents a borehole plotted by its coordinates (longitude on the x-axis, latitude on the y-axis), and the point’s color indicates the iron concentration (mg/L) measured at that location. Town clusters are labeled on the plot for clarity. We can see distinct spatial clustering: boreholes from Town Beta (toward the right/east side) are denoted by warm colors (yellow-orange), indicating high iron levels (around 1.3–1.6 mg/L). Town Gamma (toward the left/west side) shows cool purple colors, corresponding to very low iron (~0.1–0.4 mg/L). Towns Delta and Epsilon (upper-middle region) have intermediate colors (pink to red) for moderate iron (~0.8–1.1 mg/L), and Town Alpha (lower-left cluster) also shows relatively low iron (purple tones, ~0.2–0.7 mg/L). This visual confirms a clear geographic pattern: iron concentrations tend to increase moving eastward in this region. The clustering by town is apparent and explains why the town feature and longitude were so important in the models. Areas around Town Beta (easternmost) could be prone to higher iron in groundwater, whereas the western side (around Town Gamma) has much lower iron. Such spatial patterns may reflect underlying geology or soil chemistry differences between those areas.
This clustering suggests that spatial proximity is linked to water quality in terms of iron content. In practice, one might further investigate why Beta’s area has high iron (e.g., presence of iron-rich soil or industrial contamination) and why Gamma’s area is low. Additionally, with more data, one could apply geostatistical interpolation (such as kriging) to create a continuous iron concentration map of the area. In a similar study in Bangladesh, researchers used ordinary kriging with an exponential variogram to map groundwater iron and found distinct high-iron zones in the central part of the region. Such techniques can complement our modeling by providing spatial risk maps. In our case, even a simple scatter plot is sufficient to identify clusters of concern (like Town Beta). The models we built effectively learned these spatial distinctions, which is encouraging for their predictive utility in unsampled locations.
Using a dataset of 50 boreholes, we successfully developed and evaluated models to predict groundwater iron concentration from location data. Data preprocessing involved encoding towns with one-hot vectors and considering feature scaling to ensure fair model training. Through EDA, we found that iron levels varied greatly by town, with a noticeable east-west spatial gradient. We trained a linear regression and a random forest model; both achieved strong performance (test R² around 0.9, with errors on the order of 0.1–0.2 mg/L). Slightly unexpectedly, the simpler linear model performed a bit better on the test set than the random forest, likely due to the small sample size and the dominance of a linear spatial trend in the data. Model evaluation with MAE, RMSE, and R² confirmed that the predictions were quite accurate in absolute terms.
Crucially, we identified that location features are the most influential predictors of iron concentration. The town category (especially Town Beta vs others) and the longitude coordinate had the highest importance in the models, indicating a strong spatial dependency in iron levels. This finding aligns with domain knowledge that groundwater iron concentrations often vary regionally due to differences in geology, soil minerals, or anthropogenic factors. The spatial visualization reinforced this, showing distinct high-iron and low-iron clusters.
In summary, the analysis highlights a successful application of machine learning to a geochemical water quality problem. The results suggest that if one knows the approximate location of a borehole (its coordinates or town), one can predict the iron concentration in the groundwater with reasonable confidence. For future work, incorporating additional features could further improve the model – for example, groundwater depth, aquifer type, or other water chemistry parameters might explain the remaining variance. Additionally, applying the model to a broader area or validating it on new boreholes would be important steps before deployment. Nonetheless, this case study demonstrates the workflow of building a predictive model for groundwater quality and underscores the importance of spatial factors in environmental modeling.
References:
- DataCamp – One Hot Encoding in Python: Explanation of one-hot encoding for categorical variables.
- Google Developers – Normalization: Benefits of feature scaling and normalization in machine learning.
- Nguyen et al. (2022) – Predicting Heavy Metal Concentrations in Groundwater: Found Random Forest effective for predicting arsenic, iron, and manganese in groundwater and noted iron is influenced by spatial distribution.
- Akshita Chugh (2020) – Regression Evaluation Metrics: Definitions of MAE, MSE, RMSE, and R² for regression model assessment.
- Scikit-learn Course – Feature Importance: How random forest feature importance is calculated as the total reduction in impurity contributed by each feature (Feature importance — Scikit-learn course).
- Tanay D. Chowdhury et al. (2017) – Spatial Variation of Iron in Groundwater (Sylhet): Used geostatistical mapping (kriging) to identify iron concentration hotspots, illustrating real-world spatial patterns in groundwater iron.