# Geochemistry and machine learning: methods and benchmarking

N I Prasianakis, E. Laloy, D. Jacques, J C L Meeussen, G D Miron, D A Kulik, A. Idiart, E. Demirer, E. Coene, B. Cochepin, et al.

► To cite this version:

N I Prasianakis, E. Laloy, D. Jacques, J C L Meeussen, G D Miron, et al. Geochemistry and machine learning: methods and benchmarking. Environmental Earth Sciences, 2025, 84 (5), pp.121. 10.1007/s12665-024-12066-3. hal-04988116

HAL Id: hal-04988116

https://brgm.hal.science/hal-04988116v1

Submitted on 12 Mar 2025

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_1.png" /></figure>

Distributed under a Creative Commons Attribution 4.0 International License
# Geochemistry and machine learning: methods and benchmarking

N. I. Prasianakis<fnref n="1" /> · E. Laloy<fnref n="2" /> · D. Jacques<fnref n="2" /> · J. C. L. Meeussen<fnref n="3" /> · G. D. Miron<fnref n="1" /> · D. A. Kulik<fnref n="1" /> · A. Idiart<fnref n="4" /> · E. Demirer<fnref n="4" /> · E. Coene<fnref n="4" /> · B. Cochepin<fnref n="5" /> · M. Leconte<fnref n="5" /> · M. E. Savino<fnref n="5" />,<fnref n="6" /> · J. Samper-Pilar<fnref n="7" /> · M. De Lucia<fnref n="8" /> · S. V. Churakov<fnref n="1" />,<fnref n="9" /> · O. Kolditz<fnref n="10" /> · C. Yang<fnref n="7" /> · J. Samper<fnref n="7" /> · F. Claret<fnref n="11" />

Received: 8 July 2024 / Accepted: 23 December 2024 / Published online: 18 February 2025

© The Author(s) 2025

## Abstract

Thanks to the recent progress in numerical methods and computer technology, the application fields of artificial intelligence (AI) and machine learning methods (ML) are growing at a very fast pace. The field of geochemistry for nuclear waste management has recently started using ML for the acceleration of numerical simulations of reactive transport processes, for the improvement of multiscale and multiphysics couplings efficiency, and for uncertainty quantification and sensitivity analysis. Several case studies indicate that the use of ML based approaches brings an overall acceleration of geochemical and reactive transport simulations between one and four orders of magnitude. This paper presents a benchmarking exercise that aims at providing a set of reference data and models for developing and applying ML techniques for geochemical and reactive transport simulations. Several well-known geochemical speciation codes are used to generate systematically a consistent set of high-quality chemical equilibrium data, to be used as input for the training of several ML methods. Two benchmarks are formulated, each with multiple levels of gradually increasing degree of complexity. The first benchmark focuses on cement chemistry, while the second one considers uranium sorption on a clay mineral. The performance of different ML techniques is then evaluated in terms of their numerical efficiency and accuracy. A speedup of several orders of magnitude is observed. The benefits and the limitations of different ML based techniques are then analysed and highlighted.

**Keywords** Machine learning · Geochemistry · Nuclear waste management · Numerical methods

## Introduction

Geochemical and reactive transport modelling is a necessary tool used to investigate subsurface processes (Kolditz et al. 2012, 2023; Steefel et al. 2015; Huang et al. 2018) including geological disposal of nuclear waste (Claret et al. 2022) across several scales (Águilat et al. 2021; Samper

et al. 2009; Patel et al. 2018). Reactive transport models couple fluid flow, heat transfer and solute transport with geochemical reactions under both equilibrium and kinetically controlled conditions. The results of such models can be used to design and assess the long-term performance of underground facilities (de Windt and Spycher 2019). In a typical reactive transport simulation, a flow and transport

<footnote>N. I. Prasianakis
nikolaos.prasianakis@psi.ch</footnote>

<footnote n="1">¹ Laboratory for Waste Management, Paul Scherrer Institute, Forschungsstrasse 111, Villigen PSI CH-5232, Switzerland</footnote>

<footnote n="2">² Engineered and Geosystems Analysis, Belgian Nuclear Research Centre (SCK CEN), Mol, Belgium</footnote>

<footnote n="3">³ Nuclear Research and Consultancy Group (NRG), Petten, The Netherlands</footnote>

<footnote n="4">⁴ AMPHOS 21 Consulting, S.L., Calle Venezuela, 103, 08019 Barcelona, Spain</footnote>

<footnote n="5">⁵ Andra, 1/7 Rue Jean Monnet, 92290 Chatenay-Malabry, France</footnote>

<footnote n="6">⁶ Universite Paris-Saclay, AgroParisTech, INRAE, UMR MIA Paris-Saclay, 91120 Palaiseau, France</footnote>

<footnote n="7">⁷ Interdisciplinary Center of Chemistry and Biology (CICA), Civil Engineering, Campus de Elviña, University of A Coruña, 15071 A Coruña, Spain</footnote>

<footnote n="8">⁸ GFZ Helmholtz Centre for Geosciences, Telegrafenberg, 14473 Potsdam, Germany</footnote>

<footnote n="9">⁹ Institute of Geological Sciences, University of Bern, Bern, Baltzerstrasse 1+3, CH-3012, Switzerland</footnote>

<footnote n="10">¹⁰ Department of Environmental Informatics, Helmholtz Centre for Environmental Research, UFZ, Leipzig, Germany</footnote>

<footnote n="11">¹¹ BRGM, 3 Avenue Claude Guillemin, 45060 Orléans, France</footnote>
solver is coupled with a geochemical solver using either an operator splitting or a global implicit approach (Steefel et al. 2015). The flow solver provides the numerical solution of the differential equations that describe the advection and diffusion of the chemical species. At the same time, the geochemical solver simulates the progress of chemical reactions between chemical species based on their thermodynamic properties. In an operator splitting type of reactive transport code, the geochemical solver can be easily replaced by any other module, e.g. by a surrogate model. Recently, machine learning (ML) techniques have emerged as a way to enhance computational performance of geochemical speciation solvers and reactive transport codes (Jatnieks et al. 2016; De Lucia et al. 2017; Laloy and Jacques 2019, 2022; Leal et al. 2020; Prasianakis et al. 2020; Demirer et al. 2023; Savino et al. 2022; Sochala et al. 2022; Silva et al. 2024). In Laloy and Jacques (2019), surrogates based on Gaussian processes, polynomial chaos expansion and deep neural networks (DNN) were employed to emulate time series of concentrations at the outflow of a 1D transport problem involving kinetic sorption of uranium. The performance of the surrogate models was evaluated in an application related to global sensitivity, uncertainty and optimization analyses. Laloy and Jacques (2022) applied surrogate models of the geochemical solver to describe one- and two-dimensional leaching of cement matrices for different cement systems using k-Nearest Neighbours (kNN) and DNN. In Prasianakis et al. (2020), neural networks (NN) have been used for upscaling reactive transport simulations from the pore to the field scale, and to accelerate via surrogate modelling the simulation of mineral precipitation in microfluidic lab-on-a-chip devices. Demirer et al. (2023) used DNN to accelerate 3D reactive transport simulations of dolomitization of a carbonate reservoir under non-isothermal conditions. In

Silva et al. (2024) the use of ML models for tabular data specialized methods was explored in reactive transport setup. Recently, emerging benchmarking opportunities have been discussed by Bildstein et al. (2021) in a guest editorial to the subsurface environmental simulation. Amongst others, ML benchmarking was identified as an important aspect of model development. The present study contributes in filling this gap. In this paper, the potential to accelerate the geochemical part of the reactive transport codes using ML which shall result in the subsequent acceleration of the full code, is thoroughly investigated. The main concept is illustrated in Fig. 1.

In conventional reactive transport simulations, geochemical computations are performed at each grid point of the discretized computational domain and at every time integration step. The number of grid points which are needed to resolve a complex three-dimensional domain can range between 1 million ($10^6$) to 1 billion ($10^9$) grid points. The number of time-steps of a simulation highly depends on the problem at hand and the selected time discretization, and typically lies between $10^4$ and $10^6$ discrete time steps. A rough estimate of the number of geochemical calculations per simulation can therefore be in the range of $10^{10}-10^{15}$ for realistic system descriptions, depending on the system complexity. It is not a surprise that due to the higher complexity of finding the thermodynamic equilibrium of a geochemical system, compared to mass transport equations, the chemical part of the coupled solver is typically the computational bottleneck of a reactive transport simulation. For example, in pore-scale lattice Boltzmann simulators, the geochemical calculations account for more than 99% of the total computational effort, with the flow calculation being less than 1% of the computational cost. The exact ratio

<figure><img src="image_2.png" /><figcaption>**Fig. 1** For most reactive transport codes, and for typical couplings, the geochemical solver is the bottleneck in terms of computational cost. ML techniques as surrogate models have the potential to significantly</figcaption></figure>

accelerate the reactive transport simulations. The approach followed in this work is also illustrated

<figure><img src="image_3.png" /></figure>
of geochemical to transport computational costs depends on the solver type and the details of the coupling algo- rithm. It is then evident that a significant acceleration of the overall code is to be expected once the geochemical solver is accelerated. Moreover, with the typical use and performance of the geochemical codes being at 10³–10⁴ calculations per second (more details on this aspect fol- low in Section “Comparison of computational efficiency between the geochemical solvers and the surrogate mod- els”), the more demanding of the aforementioned cases would require 10¹² s to complete or ~ 10⁴ years, if run on a single central processing unit (CPU) core, which is abso- lutely prohibitive. The gains can be significant when the geochemical part is accelerated. For instance, a reduction in calculation times of the reactive transport simulations by a factor of 5 to 10 is reported by Laloy and Jacques (2022) and by 2–3 orders of magnitude in Prasianakis et al. (2020), which opens up routes for addressing sensitivity, uncertainty and optimization applications. We note that a possible acceleration technique is also the use of look- up tables instead of the geochemical solver, which have been found to provide a speedup of the reactive transport simulator by 1–2 orders of magnitude for a very simple cement-based geochemical problem (Huang et al. 2018). However, for higher dimension geochemical problems, the look-up tables grow exponentially to prohibitive sizes. The dimensionality issue seems to be efficiently addressed also by some dedicated ML algorithms such as the On Demand ML, or ODML (Leal et al. 2020), shown to speed up the geochemical part of a reactive transport simulation by a factor ranging from 1 to 3 orders of magnitude (Kyas et al. 2022).

Another possibility to accelerate the calculations is to
parallelize the simulations so that the computational effort
is split among several computing units. Many flow solvers
are designed to run in parallel on multicore CPU's or on
CPU clusters or even on hybrid CPU/GPGPU cluster and
supercomputer architectures (Kolditz et al. 2012; Lichtner
et al. 2015; Safi et al. 2019; Molins et al. 2020; Krause et al.
2021; Soulaine et al. 2021). While several flow solvers
run on multi-GPU clusters, to the best of our knowledge
none of the commonly used geochemical solvers which are
programmed in Fortran, C, or Java, can run directly on the
GPU architecture. An efficient GPU implementation would
require to re-write the full geochemical code for example in
the CUDA programming language by simultaneously taking
into account the GPU hardware architecture, which is not
anticipated by their developers to occur in the near future.
The practical implementation of porting geochemical cal-
culations on the GPU can be based on the use of surrogate
models. The surrogate models explored here, being simpler
mathematical representations than the geochemical codes
they attempt to emulate, can be readily ported on the GPU.

In addition, several open source ML algorithms already
perform the training and predictions (inference) directly
on GPU hardware, thereby benefiting from GPU hardware
acceleration.

The benchmarks presented in this paper provide a point of
reference for selecting the optimal acceleration strategy, test-
ing codes and models and for addressing the challenges to:

* Streamline the production of high-quality consistent training datasets, using the most widespread geochemical solvers (PHREEQC (Parkhurst 1995), ORCHESTRA (Meeussen 2003) and GEMS (Kulik et al. 2013)).

* Set the specifications and provide a framework which allows the produced datasets to be directly used by all available and future ML techniques.

* Compare the accuracy and consistency of the geochemical solvers.

* Use Deep Neural Networks, Gaussian Processes with Active Learning, and tree based methods to develop surrogate models from the generated data.

* Setup appropriate metrics for the critical evaluation of the accuracy of ML models.

* Provide a guide on the problems and circumstances where ML can be beneficial.

For the scope of this paper, two reference benchmarks have been defined. The first system is relevant to cement hydration and degradation under leaching conditions and uses the CEMDATA-18 thermodynamic database (Lothenbach et al. 2019). The simplest chemical system considers Ca–Si–H–O basis and includes simple C–S–H solid solution models. With the increasing complexity, more realistic systems are also addressed by adding Al–Mg–S–C–(Na–K) species. Output parameters of interest are the aqueous speciation and the resulting mineral precipitates. The second system is related to sorption of U in claystone formations (e.g. Callovo-Oxfordian, Opalinus or Boom clay) using the two site non-electrostatic surface complexation and cation exchange model (2SP-NE SC/CE) for U sorption on montmorillonite by Marques Fernandes et al. (2012). Regarding the chemical complexity, a system containing Na–Cl–U–H–O is considered as the base case, and a more complex system is defined with the addition of Ca and carbonate (CO₂) to augment the aqueous speciation of U. Parameters of interest, among others, are the resulting concentrations of U sorbed on edges (surface complexes), of U on ion exchange sites, and the amount of precipitating metaschoepite, together with the resulting K<sub>d</sub>'s.

Since there is a large number of machine learning mod-
els available, and with the domain growing at a fast pace,
it is not feasible to address all the existing methods at an
adequate depth. Therefore, a few of the most common
approaches were selected and operated by the respective
experts in the geochemistry field (Churakov et al. 2024; Claret et al. 2024). The work presented here focuses on the step related to the creation of the surrogate models and the comparison of the accuracy and numerical efficiency of the produced models, and not on coupled reactive transport simulations. The same surrogate geochemical models can be used for several applications which do not include transport, as for example for sensitivity analyses (Sochala et al. 2022; Turunen and Lipping 2023), and the target accuracy can be much lower compared to what is needed for a reactive transport simulation. Creating the surrogate models is the first step for a number of applications including reactive transport. At the same time there are several implementation pathways for the coupling and use of surrogate models in a reactive transport setup with some methodologies being able to operate at relatively low accuracy by parallel enforcement of, e.g., mass and charge balance. Addressing these topics in depth goes beyond the scope of this paper.

The paper is organised as follows. In Section “Geochemistry methods and codes”, the used geochemical codes are presented, the outline of the workflow is described in detail, and the sampling techniques which are used to generate the training datasets for the ML methods are discussed. In Section “Definition of geochemical systems with increasing complexity”, the cementitious and uranium-clay chemical systems are described at different levels of complexity. Section “Machine learning methods” introduces the machine learning methods which are explored in the benchmarks. In Section “Metrics of accuracy”, a comprehensive outline of the metrics of accuracy which are used to scrutinize the efficiency of the models is presented. In Section “Comparison between geochemical solvers”, an analysis of the training sets is performed by comparing the different geochemical solvers. Section “Performance of the machine learning algorithms and discussion of results” presents and discusses the major benchmark results of the ML modelling. Section “Comparison of computational efficiency between the geochemical solvers and the surrogate models” evaluates the performance of the surrogate models in terms of potential gain in numerical efficiency. Section “Remarks on the application and performance of current ML modelling approaches” and “Conclusions” contain respectively, discussion and interpretation of the results, and conclusions and lessons learned. Accompanying material, and in depth details of the chemical systems and extended results are provided in the supplementary information file while the open access codes, datasets and results are provided in the data repository Zenodo as described in the Supplementary Information (SI).

## Geochemistry methods and codes

In this section, the used geochemical codes GEMS (Kulik et al. 2013), ORCHESTRA (Meeussen 2003) and PHREEQC (Parkhurst and Appelo 2013) are briefly presented. All details including the exact chemical system files, input files and output files for each code used in this benchmark can be found in the online Zenodo digital repository. All three geochemical codes were used to generate training datasets, for all system definitions, closely following the respective specifications. Small differences in the output of the geochemical solvers can be traced to the different methodologies and numerical methods which are used to obtain the equilibrium. In order to eliminate any uncertainties of machine learning model results, as for example when the training and test datasets are produced from a different geochemical solver, the validation of the ML models is performed using training and test datasets originating from one and the same geochemical solver. The random sampling of the training and testing datasets are produced using a different seed number as described in Section “Benchmark Workflow”. The seed number is provided for reproducibility purposes.

## GEMS geochemical software

The Gibbs Energy Minimization Software (GEMS) is a collection of open-source codes developed and maintained at the Laboratory for Waste Management (LES) at PSI which includes three main code packages and default thermodynamic databases:

* GEMS3K: the standalone solver of geochemical equilibria with C/C++ API (Kulik et al. 2013; https://gems.web.psi.ch/GEMS3K), also in form of xGEMS with modern C++ and Python APIs, all using the TSolMod library of activity models for phases-solutions (Wagner et al. 2012);

* GEM-Selektor v.3 (GEMS3): the interactive code package for geochemical modelling by Gibbs energy minimisation (https://gems.web.psi.ch/GEMS3), which includes the default chemical thermodynamic database GEMS TDB in thermochemical and reaction data formats, derived from the PSI-Nagra chemical thermodynamic database 12/07 (https://www.psi.ch/en/les/database);

* GEMSFITS: the multiple parameter optimisation code coupled with GEMS3K (Miron et al. 2015) (https://gems.web.psi.ch/GEMSFITS)

The GEMS3K code represents an efficient numerical engine of GEMS using the Gibbs energy minimisation (GEM) numerical algorithm for solving chemical equilibria
in complex heterogeneous systems with multiple non-ideal solutions. More than 25 activity models of phases-solutions are provided in the built-in TSolMod library (Wagner et al. 2012). Along with the ability to run various models of sorption and ion exchange, this makes GEMS relevant for a wide range of (geo-)chemical applications, in particular in geological waste disposal and in cement chemistry. The GEMS3K open-source code can be coupled with other codes such as reactive mass transport simulators, parameter optimisers, CemGEMS web app back-end (https://cemgems.app), and JupyterHub Python notebooks. The latter is used via the xGEMS common C++/Python interface, for instance, in this study. The input files (in JSON or key-value text format) to start GEMS3K, including thermodynamic data, can be exported per mouse click from the GEM-Selektor graphical user interface (GUI), or prepared manually using an ASCII text editor.

The GEM-Selektor is a code package for interactive thermodynamic modelling of heterogeneous aqueous and non-aqueous chemical systems, especially those involving metastability and dispersity of mineral phases, solid solution—aqueous solution equilibria, adsorption and ion exchange. It includes an advanced multi widget GUI with a context-sensitive help viewer. Since 2000, GEM-Selektor is intensively used at PSI for various projects in waste disposal (with recently added ClaySor sorption model for clays) and nuclear engineering. Equipped with the Cemdata18 thermodynamic database (Lothenbach et al. 2019) available both in GEMS and in PHREEQC formats, GEM-Selektor is widely used in the cement chemistry community since 2007. To date, the code has a broad user community (11,000+ downloads, 1100+ active users, 1500+ publications citing its use). In this study, GEM-Selektor has been used for initial definition and testing of cement and uranyl-clay chemical systems, and has been used further on for generating GEMS training and validation datasets.

The GEM algorithm for computing equilibrium phase assemblages and chemical speciation is based on a mass balance for the entire chemical system set up by the total amounts of chemical elements and [zero] charge ("Independent Components" IC). All chemical species present in all phases are called "Dependent Components" (DC), as their stoichiometry formulas are built from ICs. Each DC is provided at input with its formula (elemental stoichiometry), molar volume, and a value of the standard Gibbs energy per mole G° taken from the project database after corrections to pressure P and temperature T of interest. Each phase definition includes a DC (pure substance phase) or several DCs (solution phase), with optional properties such as the specific surface area or interaction parameters. Multi-component phases are called solutions, with mixing behavior described by a chosen activity model (ideal or non-ideal mixing) and its parameters.

In GEMS, the activities and concentrations of DCs are treated separately for each phase, taking into account the appropriate standard/reference states and activity coefficients. The equilibrium phase assemblage conforming to the Gibbs phase rule is selected according to the energy minimum criteria from all feasible phases. The equilibrium partitioning in a multiphase system, for example including aqueous solution, gas mixture, one or several solid solutions, many pure solid phases, and, optionally, sorption phases, is computed simultaneously for all phases in iterations of the “Interior Points Method” (IPM-3) non-linear minimization algorithm. In addition to the chemical speciation vector n(x) (mole amounts of DCs—the primal solution), IPM-3 computes a dual solution vector u(b) holding equilibrium chemical potentials of ICs at the chemical equilibrium state of interest.

## ORCHESTRA

ORCHESTRA (Meeussen 2003) is a free software framework for implementing chemical equilibrium models. Although it has a similar area of application as e.g. GEMS or PHREEQC, and can read the same thermodynamic database as PHREEQC, the solver module has a very different internal structure. In contrast with traditional codes, the chemical models handled by ORCHESTRA are not defined in its executable source code but in an external text file that is read as input during run time.

This modular separation between abstract solver code and chemical model type definitions has several advantages:

- It is transparent, as all chemical model type definitions are present in external text file(s).

- It is very flexible as it allows users to add or adapt chemical model types without the need to gain access to or update the executable solver code.

- It makes ORCHESTRA a convenient tool for developing new (adsorption) model types of which a large library with state of the art adsorption models is available.

- The resulting generic solver code is very small and efficient, and well suited for parallel processing.

- The approach also makes it possible to change/update/ modernize the solver executable without affecting the chemical model definitions.

The possibility to change the solver executable without affecting the chemical models has made it possible to implement an open source C++ version of the solver, with exactly the same functionality as the original Java version. The advantage of a C++ version is that it is easier to couple with external transport codes and to be used from e.g. Python. Both versions can read the same input files and thus
automatically have the same advanced chemical modelling
capabilities.

The standard chemical set up of ORCHESTRA follows
the LMA (Law of Mass Action) approach, like PHREEQC,
in which the chemical reaction system is summarized as a set
of non-linear equations. In these equations, the log activities
of the selected components act as unknowns, and their given
mass balances as equation results. This set of equations is
solved numerically by Newton-Raphson iterations. For
PHREEQC these equations are defined in its source code,
while in case of ORCHESTRA in its input files, however,
effectively the same equations are solved.

This, combined with the fact that ORCHESTRA can
directly read PHREEQC thermodynamic database files as
input, leads to closely matching results between these two
geochemical codes.

ORCHESTRA has a graphical user interface supporting
interactive set up a chemical system and can use PHREEQC
format database files as input. In this study the following
chemical models were used in the benchmarks: aqueous
complexation, activity correction according to Davies
model, pure phase mineral precipitation, ideal solid solu-
tions, ion exchange (Gaines-Thomas, Vanselow convention),
and surface complexation.

PHREEQC

PHREEQC-3 (Parkhurst and Appelo 2013) is an open-source
geochemical solver based on the law of mass action approach
to obtain chemical speciation under thermodynamic equi-
librium conditions or for mixed kinetic-equilibrium con-
ditions with kinetic homogeneous and heterogeneous rate
equations implemented via the embedded BASIC interpreter.
It contains state-of-the-art models for aqueous speciation
(Davies, Debye-Hückel, specific ion interaction model and
Pitzer model), exchange reactions (Gaines-Tomas, Gapon,

and Vanselow convention, Rothmund-Kornfeld equations),
surface complexation (diffuse double layer model, constant
capacitance model, CD-MUSIC/Triple layer model, compo-
sition of diffuse double layer), solid solutions (ideal, binary
non-ideal model), exchange with gas phase and mineral
precipitation and dissolution. To enhance its use via other
software or to embed it in coupled reactive transport mod-
els, the interfaces iPhreeqc (Charlton and Parkhurst 2011)
and PhreeqcRm (Parkhurst and Wissmeier 2015) have been
developed.

The source code and compiled programs are available at
https://www.usgs.gov/software/phreeqc-version-3. The code
comes with several thermodynamic databases, but other
databases can be integrated as is done in this study. The
version used in this study is PHREEQC-3.7.3 as embedded
and modified in the HPGeochemistry code which is part of
the HYDRUS software (Šimůnek et al. 2024), available at
https://www.pc-progress.com. This adapted version gives
direct interaction with Python (version 3.8) used here for
generating random samples for training.

Benchmark workflow

A systematic workflow employed for generating the training
dataset files is outlined in Fig. 2. The three building blocks
required for the production of the datasets are indicated in
blue: (i) the definition of the system including the range
of the variables (see Section "Definition of geochemical
systems with increasing complexity"), (ii) the input files
for each system and each geochemical solver (GEMS-G,
ORCHESTRA-O, PHREEQC-P), (iii) a script that gener-
ates the samples of the input vector. The latter is created in
Python with the possibility to adapt the number of samples,
the seed number and the method of sampling (Monte Carlo
or Latin Hypercube Sampling). After selection of the system
and the sampling variables, a file with the sampled input

<figure><img src="image_4.png" /><figcaption>**Fig. 2** The benchmark workflow for creating the training set files (G: GEMS, P: PHREEQC, O: ORCHESTRA)</figcaption></figure>

<figure><img src="image_5.png" /></figure>
vectors is created which, combined with the input file for the geochemical solver, will give the file with the training data. The training data files (seed 54,854) have been the pool out of which the ML methods obtain the samples for the training process. It is noted that some of the methods further split these training data samples to three categories during the training process (e.g. training 80% validation 10%, testing 10% of the total samples for NN training), some other methods in two categories (e.g. RF-GP used 80% for training and 20% for testing) or even using 100% of the points for the training (e.g. DecTree training). Additional validation data sample files (distinct from the aforementioned training data samples) were produced in the same way, however by using a different seed number. Validation files should not be confused here with the validation split of the training data used by the NN methods. The validation files with different seeds (seed 69,696) are datasets used to assess the accuracy of the ML models at ‘blind’ test conditions (Validation sample dataset = blind test sample dataset). The used input and script files are available in the supplementary information and the online repository.

A crucial step in the training process is the generation of the training datasets which provide a representative sampling of the multidimensional space of interest. For most of the ML regression methods, the computational effort needed to achieve a certain accuracy strongly depends on the size of the training dataset. Therefore, it is essential to ensure that at least minimal representative datasets are used for training. The number of necessary samples or, in other words, the resolution of the training set, greatly depends on the dimension of the problem as well as on how sharply the function varies in space, and if some discontinuities exist (e.g. a mineral phase appearing).

One of the following four sampling options are commonly used to create a training set: Grid sampling (1) simply consists of selecting points on a regular grid. This simple strategy is however not efficient in more than two or three dimensions. Plain “random” Monte Carlo (MC) sampling (2) randomly draws points from their given probability distribution. Latin hypercube sampling (LHS) (3) improves the coverage of the input space by using a stratified sampling scheme, e.g. (Loh 1996). In its most basic form, LHS partitions each input dimension into equal-sized segments and then ensures that the produced sample contains just a point in each segment. Lastly, low-discrepancy sequences (4), also called “quasi random” sequences or quasi-Monte Carlo (QMC) sampling, are algorithmically generated numbers that are better equally distributed in a given volume than pseudo-random numbers, under certain conditions. The most known low-discrepancy sequence is perhaps the so-called Sobol sequence, e.g. (Joe and Kuo 2003), In this study, the training sets were generated using LHS. There are several open-source (Python, R or other) packages that implement

various LHS and low-discrepancy sequence variants. Here we resorted to the Python-written scikit-optimize package (https://scikit-optimize.github.io/stable/). In addition, it is stressed that, although not used herein, low-discrepancy sequence designs and, to some extent, certain specific LHS variants (Sheikholeslami and Razavi 2017), can be incremented from an initial design.

## Definition of geochemical systems with increasing complexity

Hereafter, the two geochemical systems and their complexity variants that are used for the ML benchmarks are thoroughly described. Specific implementation details for the three geochemical solvers, and the exact input files to produce the training sets are available as supplementary material and on the online Zenodo repository.

## Benchmark 1—cementitious system

The first series of benchmarks concerns the hydration and evolution of cementitious systems for isothermal conditions (25 °C). Based on the masses of different oxides and of water, the mass of cement hydrates, the aqueous concentrations, and the supplementary output variables (see below) as shown in Table 1 are calculated with the geochemical solvers at equilibrium conditions. Kinetically controlled processes have not been considered in this study. The variations in oxide amounts mimic also the cement aging conditions representative for leaching and carbonation. Six different systems with increasing complexity (addition of more chemical components results in a more complex cement hydrates assemblage) were defined as summarized in Table 1. The thermodynamic database CEMDATA18 (Lothenbach et al. 2019) is used with the Davies equation for the aqueous activity coefficients. The input chemical files for each geochemical solver are provided as supplementary material.

The output variables in the training set consist of (i) amount of the chemical component in the aqueous phase (mole, indicated by X\_aq with X the chemical component), (ii) amount of the chemical component in the solid phase (mole, indicated by X\_s), (iii) mass of water after reaction (kg, indicated by MassWater), (iv) pH, (v) amount of the pure phases (mole, indicated by the mineral name as listed in Table 3), (vi) mass of the solid solution (kg, mCSHQ for C–S–H and mET\_ss for SS-AFT), (vii) chemical components in the solid solution (mole, X\_ss for C–S–H and X\_ET\_ss for SS-AFT), (viii) total volume of solids (cm³, indicated by V\_s), and (ix) volume of gel water in C–S–H (cm³, indicated by Gel\_water) calculated as described in (Kulik 2011):
<caption>**Table 1** Benchmark 1—overview of cementitious systems and ranges of amount of oxides (mole) and mass of water (kg). This represents the input model variables, while the output is the resulting thermodynamic equilibrium speciation including amounts of aqueous species, pH, and amounts of mineral precipitates</caption>

<table><thead><tr><th>System</th><th>Acronym</th><th>Oxides</th><th>Hydrates</th><th>CaO</th><th>SiO<sub>2</sub></th><th>CO<sub>2</sub></th><th>Al<sub>2</sub>O<sub>3</sub></th><th>SO<sub>3</sub></th><th>K<sub>2</sub>O</th><th>H<sub>2</sub>O</th></tr></thead><tbody><tr><td>Primitive</td><td>10_PC</td><td>CaO<br>SiO<sub>2</sub></td><td>Portlandite<br>Amorphous SiO<sub>2</sub><br>C-S-H<sup>1</sup></td><td>0.1–1.8</td><td>0.2–0.7</td><td></td><td></td><td></td><td></td><td>0.05–0.15</td></tr><tr><td>Primitive + C</td><td>15_PCc</td><td>CaO<br>SiO<sub>2</sub><br>CO<sub>2</sub></td><td>10_PC +<br>calcite</td><td>0.1–1.8</td><td>0.2–0.7</td><td>0.001–0.96</td><td></td><td></td><td></td><td>0.05–0.15</td></tr><tr><td>Minimal</td><td>20 ANC</td><td>CaO<br>SiO<sub>2</sub><br>Al<sub>2</sub>O<sub>3</sub></td><td>15_PC +<br>gibsite<br>katoite<sup>2</sup><br>chabazite<br>straetlingite</td><td>0.9–1.4</td><td>0.3–0.6</td><td></td><td>0.03–0.07</td><td></td><td></td><td>0.05–0.15</td></tr><tr><td>Minimal + C</td><td>25 ANC</td><td>CaO<br>SiO<sub>2</sub><br>Al<sub>2</sub>O<sub>3</sub><br>CO<sub>2</sub></td><td>20 ANC +<br>calcite<br>hemicarboaluminate<br>monocarboaluminate</td><td>0.9–1.4</td><td>0.3–0.6</td><td>0.001–0.8</td><td>0.03–0.07</td><td></td><td></td><td>0.05–0.15</td></tr><tr><td>Simple</td><td>30_SC</td><td>CaO<br>SiO<sub>2</sub><br>Al<sub>2</sub>O<sub>3</sub><br>SO<sub>3</sub><br>K<sub>2</sub>O</td><td>25 ANC +<br>monosulfoaluminate<sup>3</sup><br>ettringite<br>gypsum</td><td>0.9–1.4</td><td>0.3–0.6</td><td></td><td>0.03–0.07</td><td>0.02–0.05</td><td>0.006–0.012</td><td>0.05–0.15</td></tr><tr><td>Simple + C</td><td>35_SCc</td><td>CaO<br>SiO<sub>2</sub><br>Al<sub>2</sub>O<sub>3</sub><br>SO<sub>3</sub><br>K<sub>2</sub>O<br>CO<sub>2</sub></td><td>25 ANC +<br>monosulfoaluminate<sup>3</sup><br>gypsum<br>thaumasite<br>SS-AFt<sup>4</sup></td><td>0.9–1.4</td><td>0.3–0.6</td><td>0.001–0.8</td><td>0.03–0.07</td><td>0.02–0.05</td><td>0.006–0.012</td><td>0.05–0.15</td></tr></tbody></table>

¹Ideal solid solution, CSHQ model represented by endmembers CHSQ_JenD, CSHQ_JenH, CSHQ_TobD, and CSHQ_TobH from CEMDATA18 for systems PC, Pc, MC and McC. For SC and ScC, KSiOH endmember is added

²C3AH6 from CEMDATA18

³Monosulphate14 from CEMDATA18

⁴Ideal solid solution of the endmembers ettringite03_ss and tricarboalู03 from CEMDATA18

$$Gel\_water = -3.67474 + 14.75335(C/S) - 19.00002(C/S)^2 \\ + 10.931594(C/S)^3 - 2.266734(C/S)^4$$

where C/S is the Ca to Si mole ratio of C-S-H. Specific details on the output variables for each system are given in the supplementary information file (SI).

## Benchmark 2—uranium sorption on clay

The second benchmark considers the uranium sorption on a claystone which is modeled with the two-site protolysis non-electrostatic surface complexation and cation exchange model 2SPNE SC/CE (Bradbury and Baeyens 1997) with parameters for U(VI) sorption on

<caption>**Table 2** Benchmark 2—overview of three systems with increasing complexity for U sorption and ranges of pH, C (mole), Ca (mole) and U (mole). The pH is controlled by the addition of NaOH or HCl. Surrogate models could be created by considering as input variable the pH and U (2D-input) or the amount of acid, base and U (3D-input)</caption>

<table><thead><tr><th>System</th><th>Acronym</th><th>pH</th><th>U</th><th>C</th><th>Ca</th></tr></thead><tbody><tr><td>Primitive</td><td>10_PMU</td><td>2–12</td><td>1E-9–1E-2</td><td></td><td></td></tr><tr><td>Primitive + C</td><td>15_PMuc</td><td>2–12</td><td>1E-9–1E-2</td><td>1E-6–1E-2</td><td></td></tr><tr><td>Primitive + C + Ca</td><td>25_MMuC</td><td>2–12</td><td>1E-8–1E-3</td><td>1E-6–1E-2</td><td>1E-6–1E-2</td></tr></tbody></table>
montmorillonite (Tables 2, 5 and 6 in (Marques Fernandes et al. 2012), see also SI file). For the cation exchange reactions, both the Gaines-Thomas and Vanselow conventions are implemented. Three levels of complexity are defined (Table 2) in a solution with a constant background concentration of NaCl of 0.1 M (pH conditions were controlled by adding corresponding amounts of NaOH or HCl). Adding C changes the aqueous complexation of U and the speciation of U surface complexation containing U-carbonate complexes. Additional consideration of Ca further augments the aqueous complexation of U and also has an effect on the cation exchange. The training set is made by combining 1 g montmorillonite and 1 L of water.

The output variables in the training set consist of (i) amount of the chemical component in the aqueous phase (mole, indicated by X\_aq with X the chemical component), (ii) amount of the chemical component sorbed (mole, indicated by X\_s) with distribution on the edge surface complex (mole, indicated by X\_sc) and the cation exchange complex (indicated by X\_ex), (iii) amount of the pure phases (mole, indicated by the mineral name), and (iv) K_d values of U and Ca by all processes (surface complexation, cation exchange and precipitated) or by surface complexation or cation exchange only (L/kg, indicated by K_d, K_d\_X\_sc, K_d\_X\_ex). Specific details on the output variables for each system are given in the workflow document which is provided as supplementary material.

## Machine learning methods

This section describes the ML methods which have been considered in this benchmark. ML is an inherent component of artificial intelligence which allows software to learn directly from data. It is used for two main problem categories: (a) classification problems, (b) regression problems (Alpaydin 2009). This paper investigates the regression capabilities of ML models. Many applications and learning processes may be classified in three main classes (Sarker 2021): (a) supervised learning, (b) unsupervised learning, and (c) reinforcement learning. The methods reviewed here belong to the first class, meaning that they are used for regressing geochemical data produced by the geochemical solvers in order to create geochemical surrogate models.

### Physics-based recursive partitioning of parameter space: DecTree

The DecTree approach (De Lucia and Kuhn 2021) uses a non-linear recursive partitioning of the multivariate parameter space relying on the identification of smooth functional relationships between each individual response to be predicted and a minimal number of features. This approach

allows to embed into the algorithm the knowledge about the physical processes depicted by the numerical variables in the training dataset. This knowledge is used for example to split the training data based on occurrence of mineral phases at equilibrium, to reduce the number of responses to be predicted applying mass balance equations, and to define engineered features based on underlying governing equations (e.g., laws of mass action). The partitioning can be represented as a decision tree in which each node represents a condition, and the leaves are regions of the parameter space in which mono- or bivariate functional dependencies between response variables and the minimum number of features were identified. In each leaf then interpolation algorithms are used to predict new values based on the samples available in the training dataset.

The DecTree method has been applied and tested in the uranium benchmark case. An in-depth description of the algorithm and its results are described in De Lucia (2024). In the current implementation, the algorithm starts by sub-setting the training dataset based on the mineral occurrence in the responses; a convex hull of the region where either minerals precipitate or not, is stored as a node of the decision tree, so that new predictions can decide in which subregion they belong to with an inexpensive point-in-polygon check. Note that both convex hull and point-in-polygon algorithms are well defined in multidimensional spaces. However, the convexity of any set of points in the training dataset is not guaranteed. This is dealt within DecTree by computing a convex hull of both the region where a condition is true (e.g., a mineral occurs) and its complement and by computing the confusion table of the resulting in sample classifications. The absence of false positives for one of the hulls is heuristically interpreted as confirmation of the convexity of that region, which is then retained by the model. In case where neither subregions are convex, DecTree reverts to a more general but less precise contouring operation. In each leaf then the algorithm screens each feature-response pair for functional dependence based on the absolute value of the Spearman rank-based correlation coefficient. If the sum of functional relationships and applicable mass balance equations equals the number of responses to predict, the leaf is completely determined by means of bivariate interpolations, carried out using cubic splines. If not, then a multivariate interpolation is applied to predict the still unexplained responses. In this study, a Multilevel B-spline Approximation method is used for interpolation (Lee et al. 1997; Hjelle 2001).

For the Uranium test case, DecTree considers as input the dataset and some additional “physical knowledge” of the chemical system: which variable corresponds to a mineral phase (Metaschoepite) and known mass balance equations, which in this case is the sum of uranium sorbed, exchanged, precipitated as Metaschoepite and in the aqueous solution equals the total amount of U in input. DecTree operates a
first split based on acidic or basic conditions; each of these subspaces is further split by learning the region in the two-dimensional domain {U, pH} where mineral Metaschoepite appears. Within the four resulting subspaces, DecTree then identifies functional dependencies between each output and the inputs; when monovariate functional dependencies are found, the corresponding output variables are marked as “solved” and will be predicted by cubic spline in the prediction phase; the remaining variables can be predicted either by multivariate interpolation or by applying mass balance equation if all other variables in the equation are already explained. Multivariate interpolations conducted by multi-level B-splines approximation require some heuristic parametrization such as the degree of the basis splines and the ratio of the dimensions’ ranges. Note that the actual training data, split in the different subregions, are retained by Dec-Tree as interpolation sets to be used in the prediction phase.

For this benchmark, all input and output variables were log-transformed with exception of pH, left untouched, and the amount of Metaschoepite, which was scaled by a factor of 1000. Except for the use of mass balance to predict an output, which is always applied on back-transformed variables, all interpolations are operated in the transformed parameter space, thus introducing a slight systematic bias in DecTree predictions.

The current implementation of DecTree achieves remarkable results in the first simple uranium benchmark, with virtually no training time, competitive prediction speed, and high accuracy. DecTree is able to produce results for all combinations of available training (3) and validation datasets (4) in seconds using a single CPU core. The misclassification of prediction points in which Metaschoepite occurs is very limited (i.e., 4 points in the 20k validation set when DecTree is trained on the 50k training dataset), and this accentuates the importance of refined sampling near the “physical boundaries” or transition regions of the parameter space. Conversely, within each region, the sampling density could be reduced without large loss of accuracy but enhancement of prediction speed, which depends quadratically on the number of points used by the interpolation. Relying on geometric approaches to learn regions in the parameter space obviously requires enough data points for this to be reliable.

Several aspects still need to be further developed to make the DecTree approach truly general. In particular, at the time of writing DecTree is unable to deal with multiple minerals at equilibrium, especially for disjointed or nested occurrences, or with solid solutions, for which no mechanism to embed “physical knowledge” into the model has yet been established. This prevented its application to other benchmarks, such as those for the cementitious systems.

On the other hand, within the considered benchmark, DecTree outputs a recursive split of the training data in form of a decision tree and correctly identifies latent constraints in

some of the resulting subregions: in those where Metaschoepite occurs, for example, DecTree recognizes from the data the bijectivity between most of the output variables and a single input, as expected by the domain expert interpreting its results, since the Law of Mass Action constraints the total degrees of freedom of the system in these regions. This is further leveraged by DecTree by reducing the number of multivariate interpolations, hence performing a model reduction in classical sense within each subspace.

## Gaussian processes

The two approaches described below aim at estimating a real-valued function $f$ defined on a compact subset $A \subset \mathbb{R}^d$ by using the points from the different databases available at which $f$ is evaluated. Function $f$ is then considered as sample of a zero-mean Gaussian process (GP) having a covariance function $k$. The advantage of GP approaches is that, conditionally on a set of $t$ observations, the posterior distribution is still a GP having a mean $\mu_t$ and a covariance function $k_t$ given by

$$ \mu_t(u) = k_t(u)' K_t^{-1} y_t \quad (1) $$

$$ k_t(u, v) = k(u, v) - k_t(u)' K_t^{-1} k_t(v) \quad (2) $$

where $k_t(u) = [k(x_1, u) \dots k(x_t, u)]'$. Here [...] denotes the matrix transposition, $u$ and $v$ are in $A$ and $K_t = [k(x_i, x_j)]_{1 \le i,j \le t}$, where the $x_i$'s are in $A$. For further details on GP, the reader is referred to (Rasmussen 2004) in which their properties are thoroughly presented. In the systems considered here, $f$ models a physical quantity, e.g. pH or species concentration, that is assumed to be smooth and $k$ is either the squared exponential covariance function or the Matérn covariance function, both of them commonly used in this case.

## Gaussian processes with active learning

A detailed description of Gaussian Processes with active learning (AL) is given in (Savino et al. 2022). If the function $f$ is described as a sample of a Gaussian process, it allows to provide an error estimation to choose sequentially the most relevant input data until a given stopping criterion is fulfilled. The main advantage of this approach is that the number of required evaluations of the function to estimate is very limited and that there is no parameter to tune.

Inspired by (Srinivas et al. 2012) who proposed a sequential approach for maximizing a function by modelling it using a Gaussian process, we propose a strategy which consists in adding the new point to the set of observations at which $f$ needs to be evaluated as follows:
$$x_{t+1} \in \operatorname{Argmax}_{x \in A} \sigma_{t}(x) \quad (3)$$

where:

$$\sigma_t(x)^2 = k_t t(x, x) \quad (4)$$

$k_t$ being defined in Eq. 2 and $\operatorname{Argmax}_{x \in A} \sigma_t(x)$ being the set of $x \in A$ where $\sigma_t(x)$ reaches its maximum.

We also propose using the maximum-likelihood strategy described in (Rasmussen 2004) to estimate the length scales $l = (l_i)_{(1 \le i \le d)}$. The method has been implemented by using the GaussianProcessReessor class of the scikit-learn 1.0.2 module of Python (Buitinck et al. 2011) which only provides the computation of $\mu_t$ and $\sigma_t$. This method aims at building an optimized training dataset to get the best estimation on the validation dataset.

To illustrate the importance of such an optimization, we first use a table of 500 randomly chosen observation points to compute the parameters of a Gaussian process and make a first set of predictions on 50,000 (50k) validation points. Thereafter, the AL method is used to choose sequentially 500 observation points from a random table of 50k points and different from those used for the validation step and make a second set of prediction on the 50k validation points. The results are shown in Table 3. As expected, the indicators are better in the later case and emphasizes the importance to choose smartly the training dataset. The normalised Sup-norm indicator is of particular interest since it is not subject to averaging effect due to many validation points and shows a gain of an order of magnitude with the AL method compared to the non-AL method.

Figure 3 demonstrates an exponential decay of the normalized MAE as the number of observation points increases. This means that the method can automatically optimize the size of the training dataset to match a chosen stopping criterion, for instance the ratio of two consecutive maximum variance, the computation of a mobile average or a maximal variance. This method is then especially relevant when the production of a training dataset costs a lot of time and becomes a limitation to the use of other classical learning methods which may require a very large training dataset.

In this paper, two sets of results are presented. The first one without AL using all the 500, 5k and 20k observation points from, respectively, the 500 points, 5k and 20k points training sets (details in the next sections). The second one with AL using 1k observation points chosen from the pool of the 50k training database. Note that a matrix inversion of a growing size matrix must be performed when adding observation points. Consequently, the AL method was limited to 1k observation points beyond which the problem cannot be

<table><caption>**Table 3** Normalized MAE and normalized supnorm for a few species from a training dataset of 500 observation points without and with active learning (AL)</caption><thead><tr><th rowspan="2"></th><th colspan="2">Normalized MAE</th><th colspan="2">Normalized supnorm</th></tr><tr><th>Without AL</th><th>With AL</th><th>Without AL</th><th>With AL</th></tr></thead><tbody><tr><td>Ca_aq</td><td>2.966E-02</td><td>1.387E-02</td><td>2.385E+00</td><td>2.894E-01</td></tr><tr><td>Ca_s</td><td>7.258E-04</td><td>2.660E-05</td><td>3.090E-03</td><td>3.827E-04</td></tr><tr><td>Si_aq</td><td>3.583E-02</td><td>2.101E-02</td><td>3.214E+00</td><td>4.628E-01</td></tr><tr><td>Si_s</td><td>5.452E-04</td><td>2.800E-05</td><td>4.839E-03</td><td>5.612E-04</td></tr><tr><td>H_aq</td><td>3.743E-02</td><td>1.319E-02</td><td>2.94633285</td><td>2.349E-01</td></tr><tr><td>H_s</td><td>1.566E-03</td><td>6.289E-04</td><td>00.7195149</td><td>5.959E-03</td></tr></tbody></table>

<figure><img src="image_6.png" /><figcaption>**Fig. 3** Evolution of the normalized MAE of aqueous calcium with the number of observation points. An exponential decay is observed with the increase of the number of the training samples</figcaption></figure>

handled in an efficient way with the current implementation at the existing hardware, and without further algorithmic optimization.

## Gaussian processes with random forest (RF-GP)

This approach implements two steps of ML: Step 1 uses the Random Forest (RF) algorithm to classify a geochemical system into sub-geochemical systems based on mineral presence and system acidity, and Step 2 uses the GP regressor to simulate geochemical processes in each of the sub-geochemical systems. The supervised RF algorithm employs an ensemble of decision trees (Breiman 2001). Predictions from all individual trees are pooled to make the final prediction. The RF algorithm implemented in the scikit-learn 1.2.2 of Python package is used in this approach (https://pypi.org/project/scikit-learn/). RF is based on the definition of groups
based on several criteria such as pH ranges and mineral combinations. The supervised RF grouping for the groups was trained with 80% of the data points selected randomly with a shuffle method from the training dataset. Input and output variables are used to train the RF. After the training, the RF is used to estimate the group of each data point in the validation data set. The estimation of the grouping is performed by using only the input variables.

As mentioned in the previous section, a kernel is a crucial ingredient of GPs which determines the shape of prior and posterior of the GP. This approach uses the Matérn kernel which is stationary and a generalization of the radial basis function kernel (Seeger 2004). This kernel is given by

$$k(x_i, x_j) = \frac{1}{\Gamma(\nu)2^{\nu-1}} \left( \frac{\sqrt{2\nu}}{l} d(x_i, x_j) \right)^{\nu} K_{\nu} \left( \frac{\sqrt{2\nu}}{l} d(x_i, x_j) \right) \quad (5)$$

where $d(\cdot, \cdot)$ is the Euclidean distance, $K_ν(\cdot)$ is a modified Bessel function and $\Gamma(\cdot)$ is the gamma function. $l > 0$ is a length-scale parameter and $ν$ is another model parameter which controls the smoothness of the resulting function.

The feature engineering which is carried out before the RF classifier includes log-transformation (if chosen) and scaling by subtracting the mean and dividing by the standard deviation. Working with log-transformation has the advantage of ensuring that predictions are always positive thus maintaining the physical consistency of the predictions. However, the accuracy of the training may get worse by working with logarithmic values of concentrations compared to work with concentrations. Typically, 80% of samples are used for training and 20% for testing.

RF-GP identifies three groups for the primitive system 10_PC. The first group corresponds to the absence of portlandite and amorphous silica. The second group corresponds to the presence of portlandite and the third group corresponds to the presence of amorphous silica. 5000 samples are used with the following distribution: 25.7%, 62.6% and 11.7% for groups 1, 2, and 3 respectively. This percentage distribution among groups is the automatic outcome of the RF classifier. It is observed that Group 2 is the most abundantly represented, containing more than half of the dataset.

RF-GP defines 8 groups for the system 30_SC. The first group corresponds to the presence of portlandite, katoite and monosulfate and includes 71.6% of the samples. The second group corresponds to the presence of portlandite, monosulfate and ettringite and includes 15% of the samples. The size of the rest of the groups is much smaller (< 4%). Therefore, most of the points fall within Groups 1 and 2, with some of the groups getting very low percentages of the data set.

RF-GP considers four groups for the uranium benchmark. The first one corresponds to acid conditions and the presence of metaschoepite. The second is characterized by acid conditions and absence of metaschoepite. The third group corresponds to basic conditions and presence of metaschoepite and the fourth one involves basic conditions, and no metaschoepite. Groups are automatically identified by RF. RF-GP uses a model for each output parameter within each of the 4 groups. Training of the RF-GP was performed with 4000 points to maintain CPU time and memory requirements at an acceptable level. All four groups are equally represented, each taking 25% of the training data set. All features except pH are log transformed. The log1p transformation is used for metaschoepite, total acid and total base. The pH in the second group was trained with input data from all groups to improve modeling performance.

The main advantage of RF-GP is the low number of required parameters which reduces greatly the chances of model overfitting. The main challenges and limitations of RF-GP include the large CPU time and memory requirements for training models with large data sets, especially with $N > 10,000$ points. Future work should be devoted to improve the performance of the relative metrics of the cement case and the very low aqueous and sorbed U concentrations in the U case and make use of the confidence intervals of the output estimates.

# Neural networks

## Fundamentals

Neural networks (NN) define the (possibly complex) relationships existing between input, $x$, and output, $y$, data vectors by using combinations of computational units that are called neurons. A neuron is an operator of the form

$$h(\mathbf{x}) = f(\langle \mathbf{x}, \mathbf{w} \rangle + b) \quad (6)$$

where $h(\cdot)$ is the scalar output of the neuron, $f(\cdot)$ is a nonlinear function that is called the “activation function”, $\langle \cdot, \cdot \rangle$ signifies the scalar product, $\mathbf{w} = [w_1, \dots, w_N]$ is a set of weights of same dimension $N$, as $\mathbf{x}$, and $b$ represents the bias associated with the neuron. For a given task, the values for $\mathbf{w}$ and $b$ associated with each neuron must be optimized or “learned” such that the resulting NN performs with a required accuracy. When $f(\cdot)$ is differentiable, $\mathbf{w}$ and $b$ can be learned for example by a gradient descent algorithm. There are many possibilities for $f(\cdot)$. The most commonly used include the rectified linear unit (ReLU), the sigmoid function and hyperbolic tangent function.

When there is no directed loops or cycles across neurons or combinations thereof, the network is said to be a feedforward network (FFN). In the FFN architecture, the neurons
are organized in layers. In addition, when every neuron of a given layer is connected to every neuron of the preceding layer, the network is referred to as “fully-connected”. In this work, all used NNs are fully-connected FFNs, and the performance of more complex NN architectures has not been investigated.

A standard fully-connected FFN layer is given by

$$h(\mathbf{x}) = f(\mathbf{W}\mathbf{x} + \mathbf{b}) \quad (7)$$

where **W** and **b** are now a matrix of weights and a vector of biases, respectively. The name multilayer perceptron (MLP) designates a fully-connected FFN with more than a layer. A most typical network is the 2-layer MLP, which consists of two layers with the outputs of the first-layer neurons becoming inputs to the second-layer neurons

$$y = g[h(\mathbf{x})] = f_2[\mathbf{W}_2 f_1(\mathbf{W}_1 \mathbf{x} + \mathbf{b}_1) + \mathbf{b}_2] \quad (8)$$

where g(·) and h(·) are referred to as output layer and hidden layer, respectively. In theory, the two layer MLP described above is a universal approximator as it can approximate any underlying process between **x** and **y** (Cybenko 1989; Hornik 1991). However, this only works if the dimension of **h**(·) is (potentially many orders of magnitudes) larger than that of the input **x**, thereby making learning practically unfeasible and the two-layer MLP approximator not adequate in many cases. It has been found that the use of many hidden layers is more efficient than increasing the size of a single hidden layer (e.g., (Goodfellow et al. 2016)). When a FFN/MLP has more than a hidden layer it is considered as deep. Note that deep neural networks (DNN) are not necessarily purely FFN but may mix different aspects of FFN, such as convolutional neural networks (CNN) with recurrent neural networks (RNN) or transformer neural networks (TNN).

## Architecture and training

In this study, several fully-connected FFN with different architectures and different optimization algorithms were trained, but using the same mean squared error (MSE) loss function and a similar strategy to monitor training performance and avoid overfitting by means of an independent validation dataset.

Three of the major machine learning platforms available as open-source software or commercial product have been used, namely: TensorFlow (Abadi et al. 2016), PyTorch (Paszke et al. 2019), and Matlab (Matlab 2021). The parametrization of the ML models implemented with TensorFlow was as follows. The training datasets produced by the geochemical solvers were further split for the training of the NN into training, validation and testing using an 80-10-10% split. It should be mentioned that in principle geochemical solvers can be used to produce many distinct validation and testing datasets to be used during the NN training. However, this would increase the number of total samples needed. The tuning of hyperparameters (e.g. number of layers/neurons, activation function, learning function parametrization) is critical to optimize the performance of the NN. Such an optimization is generally time and computation-intensive as well as case-dependent. To streamline this part of a regular ML pipeline, various optimization and sampling frameworks were considered. Indeed, in the context of NN, the size of the hyperparameter space grows quickly. Among the expert judgement and all the alternatives considered, Optuna (Akiba et al. 2019) is selected in Tensorflow due to its wide selection of algorithms as well as ease of implementation in an already established pipeline.

For Tensorflow implementation, the number of hidden layers, size of hidden layers and hidden layer activation function to be optimized were selected with a tree-structured Parzen estimator (Ozaki et al. 2022) and pruning algorithm. The single-objective optimization is run for 500 trials on each system. As expected, results show that there is no global optimum for any of the chemical systems. However, a global trend emerges. The optimization algorithm converges to a deeper NN for increasing geochemical complexity, as summarized in Table 4. More details are provided in the supplementary information.

The learning rate is chosen using the learning rate range test (Smith 2017) which hovers in the 0.001–0.01 range. In addition, the learning rate is set to decay during training when the validation loss tends to a plateau. Finally, the data is scaled and preprocessed with a min-max scaler and a base 10 logarithm, respectively.

The optimisation strategy with PyTorch used 80% of the available training samples for actual training and the remaining 20% (validation set) to monitor potential overfitting and

<table><caption>**Table 4** Tensorflow NN model architecture as used for the cement benchmark</caption><thead><tr><th>System</th><th>Primitive (10_PC)</th><th>Primitive + C (15_PCc)</th><th>Minimal (20 MC)</th><th>Minimal + C (25_McC)</th><th>Simple (30_SC)</th><th>Simple + C (35_SCc)</th></tr></thead><tbody><tr><td>Number of layers</td><td>5</td><td>10</td><td>9</td><td>10</td><td>7</td><td>10</td></tr><tr><td>Total number of neurons</td><td>100</td><td>~700</td><td>~700</td><td>~500</td><td>~300</td><td>~700</td></tr></tbody></table>
select the best training epoch for each trained NN. Since limited experimentation with the Optuna hyperparameter optimization framework (Akiba et al. 2019) did not reveal any global optimum for any of the considered test cases, the NN architecture and its hyperparameters were chosen based on expert judgment. For all geochemical systems and training set sizes, a 5-layer DNN with 192 neurons per hidden layer and the Mish activation function (Misra 2019) were selected. It is stressed, however, that other activation functions, such as Relu together with, to some extent, other architectures, were found to perform overall similarly. The loss function was minimized by stochastic gradient descent with the Adam algorithm (Kingma and Ba 2015) using an initial learning rate of 0.001. This learning rate was then automatically decreased by 50% each time the validation loss did not decrease for 10 epochs, until a minimum value of the learning rate $1 \times 10^{-6}$ was reached.

The optimisation strategy with Matlab is as follows. During the training phase, the samples contained in the training datasets were split randomly to 80% of the samples for training, 10% of the samples for validation and 10% for testing the model. The training algorithm which showed better performance for small to medium size NN was the “trainlm” which is an implementation of the Levenberg-Marquardt algorithm (Marquardt 1963; Hagan and Menhaj 1994). Similar behavior was experienced with the “trainbr” algorithm which stands for the Bayesian regularization backpropagation algorithm (MacKay 1992; Dan Foresee and Hagan 1997), which focuses in producing networks which can generalize well. The learning function hyperparameters where automatically adjusted during the training. The size of the NN which were trained was different in each case. For the simplest 10_PC cement case, a NN consisting of four hidden layers with 20 neurons at each layer was found adequate, while for the most complex case addressed, four hidden layers with 40 neurons at each layer were selected. The activation function was chosen to be the hyperbolic tangent sigmoid transfer function (“tansig”). With the increase of the number of neurons and data samples, these algorithms demonstrated a much slower convergence rate compared to gradient descent algorithms.

## Metrics of accuracy

A very crucial aspect for characterizing the performance of a surrogate model along with the computational complexity is the resulting model accuracy. Typically, there is a trade-off between obtained accuracy and model complexity. By increasing the complexity of the model and the number of training samples all methods show an improvement in the model predictions. However, the more complex the model is, the more the computational effort is required for training,

which may set practical limitations. To measure the accuracy of the model, the predictions have to be compared with the results of the geochemical solvers. A single accuracy metric seems not to be adequate to give a good sense of the model performance, especially when outputs are different physical quantities and span over several orders of magnitude (e.g. pH, concentrations, precipitate amount). For that reason, several metrics of accuracy have been calculated in an effort to provide a holistic description of the accuracy of the model predictions.

## Accuracy indicators

For measuring the accuracy of the models, test datasets which cover the same range of input parameters but are different from the training datasets (using different seed) have been created using the geochemical solvers. The same independent test datasets have been used by all ML methods for estimating their accuracy. For each specific geochemical system, the datasets are uniquely identified by the size, sampling technique and seed number which are also provided in the supplementary data. The metrics of accuracy which have been used in this work are the following:

The Mean Absolute Error MAE:

$$MAE = \frac{1}{N} \sum_{i=1}^{N} |y_i - \hat{y}_i| \qquad (9)$$

where N is the number of samples, $y_i$ is the ground truth (geochemical code calculation), and $\hat{y}_i$ is the ML model prediction.

The Mean Square Error MSE:

$$MSE = \frac{1}{N} \sum_{i=1}^{N} (y_i - \hat{y}_i)^2 \qquad (10)$$

The Root Mean Square Error RMSE:

$$RMSE = \sqrt{MSE} = \sqrt{\frac{1}{N} \sum_{i=1}^{N} (y_i - \hat{y}_i)^2} \qquad (11)$$

The coefficient of determination R²:

$$R^2 = 1 - \frac{\sum_{i=1}^{N} (y_i - \hat{y}_i)^2}{\sum_{i=1}^{N} (y_i - \bar{y})^2} \qquad (12)$$

where $\bar{y} = \frac{1}{N} \sum_{i=1}^{N} y_i$ is the mean value of the ground truth data. Since the R² for all models is very close to 1, the more indicative metric 1-R² has been used.

Along the typical metrics of accuracy, and due to the nature of the data which cover several length scales, the following normalized metrics are also considered.

The Normalized Mean Absolute error NMAE:
$$NMAE = \frac{1}{N} \sum_{i=1}^{N} \frac{|y_i - \hat{y}_i|}{y_{max} - y_{min}} \qquad (13)$$

where $y_{max}$ and $y_{min}$ are the maximum and minimum values of the ground truth data.

The Normalized Root Mean Square Error NRMSE,

$$NRMSE = \frac{RMSE}{y_{max} - y_{min}} \qquad (14)$$

The Normalized sup norm Normsupnorm

$$Normsupnorm = \max_{1 \le i \le N} \frac{|y_i - \hat{y}_i|}{y_{max} - y_{min}} \qquad (15)$$

## Relative measures of accuracy

While the above introduced standard and normalized measures of accuracy are typically used in error estimates, the denominator is a constant (e.g., the range of true variable or its mean, or standard deviation). In these cases, the measure does underestimate potential relative discrepancies for very small values. In practice the measures of the previous section provide a good insight only when the distribution of the variable of interest is uniform on a small range or symmetric distributed, e.g., Gaussian. Many variables dealt within geochemical models are, however, rather lognormally distributed or uniformly distributed on a logarithmic scale. In those cases, the accuracy indicators defined above are therefore biased towards large values. This problem can be partially solved by operating on logarithms of the variables. Plugging the natural logarithms of the true and predicted variables y and $\hat{y}$ in the MAElog and RMSElog formulas, and by assuming both variables strictly > 0:

The Mean Absolute Error of the logarithms MAElog and RMSElog

$$MAElog = \frac{1}{N} \sum_{i=1}^{N} |\log y_i - \log \hat{y}_i| = \frac{1}{N} \sum_{i=1}^{N} \left| \log \frac{y_i}{\hat{y}_i} \right| \qquad (16)$$

$$RMSElog = \sqrt{\frac{1}{N} \sum_{i=1}^{N} (\log y_i - \log \hat{y}_i)^2} = \sqrt{\frac{1}{N} \sum_{i=1}^{N} \left( \log \frac{y_i}{\hat{y}_i} \right)^2} \qquad (17)$$

It is also usual to define a slightly different variant of RMSElog, called RMSLE (Root Mean Square Logarithmic Error) by adding 1 to both the predicted and the true value to avoid dividing by 0:

$$RMSLE = \sqrt{\frac{1}{N} \sum_{i=1}^{N} (\log(y_i + 1) - \log(\hat{y}_i + 1))^2} = \sqrt{\frac{1}{N} \sum_{i=1}^{N} \left( \log \frac{y_i + 1}{\hat{y}_i + 1} \right)^2} \qquad (18)$$

All these measures yield 0 if y and $\hat{y}$ are identical. Note that the $MAE_{log}$ resembles the geometric mean of the absolute values of the quotients GMAQ (Geometric Mean of Absolute Quotients) of y and $\hat{y}$:

$$GMAQ = \left( \prod_{i=1}^{N} \frac{\hat{y}_i}{y_i} \right)^{\frac{1}{N}} \qquad (19)$$

Similar to the R², the GMAQ is expected to have values very close to 1. So, the absolute GMAQabs is calculated:

$$GMAQabs = abs(1 - GMAQ) \qquad (20)$$

The geometric mean of the quotients would be equal to 1 if the two variables are identical. The connection with the $MAE_{log}$ may be seen in a straightforward way since by definition a geometric mean is the N-th root of the product of the terms:

$$\exp \left[ \frac{1}{N} \sum \log a_i \right] = \left( \prod a_i \right) \qquad (21)$$

So, the $MAE_{log}$ is the logarithm of the actual geometric mean of the (absolute) logarithms of the quotients between the variables.

Instead of the simple quotient of the true and predicted values y and $\hat{y}$ we can define a relative (per observation) error $\alpha_i$ as:

$$a_i = \begin{cases} \frac{y_i - \hat{y}_i}{y_i} & \text{if } y_i, \hat{y}_i \neq 0 \\ 1 & \text{if } y_i = 0 \text{ and } \hat{y}_i \neq 0 \\ 0 & \text{if } y_i = 0 \text{ and } \hat{y}_i = 0 \end{cases} \qquad (22)$$

The same treatment of the case when a variable is 0 of course also applies to the above introduced relative measures. Without using logarithms, we can then define some relative measures such as the Mean Absolute Percentage Error MAPE or the Relative RMSE (RRMSE):

$$MAPE = \frac{100\%}{N} \sum |\alpha_i| \qquad (23)$$

$$RRMSE = \sqrt{\frac{1}{N} \sum (\alpha_i)^2} \qquad (24)$$

MAPE and RRMSE as defined above might not be appropriate for estimating the similarity of variables which have physically meaningful zero values, such as minerals.
However, they capture discrepancies occurring at very low values, which is a common case for geochemical variables distributed over many orders of magnitude, whereas classical metrics such as RMSE are biased towards discrepancies occurring at large values.

These relative measures account for discrepancies across all magnitudes of the y and ŷ variables and preserve the physical meaning of 0.

Finally, it is desired to report the number of negative predictions NNEG, as well as the Range of the predictions.

$$ \text{Range} = \hat{y}_{\text{max}} - \hat{y}_{\text{min}} \qquad (25) $$

To summarize, 15 relevant metrics have been defined, each highlighting a different facet of the accuracy of the ML model. The relevant headers in the metrics files are: [MAE, MSE, RMSE, 1-R2, NMAE, NRMSE, Normsupnorm, MAElog, RMSElog, RMSLE, GMAQabs, MAPE, RRMSE, NNEG, Range]. For the consistent calculation of all metrics, a single routine was used which can be found in the online Zenodo repository.

## Comparison between geochemical solvers

As a starting point for the training of the ML models, training sets of the size of 500, 5000, 50,000 and even 500,000 samples were produced for each system and each geochemical solver using the same seed number across all geochemical solvers. The produced training datasets are stored at the online Zenodo repository to facilitate the reproducibility of the current work as well as to serve as a high quality training dataset basis for future ML algorithmic developments. The different size of the training datasets is effectively providing the description of the space of solutions at different resolution levels. It is noted that there is no general rule which can define a priori the number of samples which are necessary to produce a surrogate model of a specific accuracy, since this is highly dependent on the complexity and degree of nonlinearity of the space that needs to be represented. As will be shown in the next sections, typically a larger number of training samples tends to improve the accuracy of the model but it also increases the computational effort of the training. All output variables were visualized with plots to capture efficiently the differences between the three geochemical solvers. Examples of some variables for the primitive cement system and primitive sorption system are shown in Fig. 4. Since ORCHESTRA and PHREEQC have the same algorithmic approach and provide similar results, only the difference between GEMS and ORCHESTRA will be discussed.

A more detailed comparison between the equilibration data produced by GEMS and ORCHESTRA is shown in

Fig. 5. The output of the two solvers for the amount of amorphous SiO₂ solid phase is in excellent agreement except for a single point in the lower concentrations range, which can only be visualized when plotted with logarithmic axes. The underlying reason for the large deviations of this parameter lies in the fact that for this specific data point the fraction of Si present in amorphous SiO₂ is extremely small (0.2%) as Si is for 99.8% present in CSH solid solution. The allowed inaccuracy in the total Si mass balance in this system (convergence criterion) can then lead to a large relative inaccuracy in the amount of amorphous SiO₂ when this is present in very small amounts. This inaccuracy would certainly not affect transport calculations, as in that case only the ratio of total dissolved vs. total in solid phase is relevant, and that ratio is predicted much more accurately.

The only metrics capable of capturing these few outliers are the relative metrics RRMSE and MAPE and, in a lesser measure, the RMSElog, computed on the logarithms of the variable. This highlights the importance of the metric adopted when the variable values span across several orders of magnitude, four orders in this case. At the same time only an outlier point can significantly change the relative metrics, even if its significance is small as in the example of Fig. 5 where the amorphous SiO₂ outlier concentration is too small to have a practical influence in the interpretation of results using the geochemical solvers. For the presentation of the machine learning model results we selected to present for completeness a normalized (RMSE) and a relative metric (RRMSE) of accuracy, while the computation of the rest of the metrics can be found in the supplemental material.

An accuracy overview over the cementitious systems is provided in the following figures by comparing the results between GEMS and ORCHESTRA. In the comparison, the aqueous species common to all cases, and a few common characteristic output parameters are plotted in Fig. 6 (RMSE) and Fig. 7. (RRMSE). A general trend is observed: with increasing complexity of the chemical system, the difference between the geochemical solvers slightly increases.

The difference between the output of the geochemical solvers provides an indicator which can be used to assess a practical theoretical lower limit of meaningful surrogate model accuracy. This limit may be used as a guide for the training of the surrogate models.

## Performance of the machine learning algorithms and discussion of results

The ML methods have been trained using the well defined training datasets produced by the geochemical solvers. The resulting models have been subsequently tested
<figure><img src="image_7.png" /><figcaption>**Fig. 4** Comparison of results of the three geochemical solvers for 500 samples. Top is for 10_PC system with pH (left) and amount of C–S–H (middle) as a function of bulk CaO, and capillary water (right) as a function of bulk SiO₂. Bottom is for the uranium PU</figcaption></figure>

system (with Vanselow convention) with aqueous (left) and sorbed (sum of surface complex, cation exchange complex and precipitated metaschoepite, middle) as a function of pH and U on the surface complex as a function of total U (right)

against separate test datasets (produced with a different seed), which contain sample points that the models have not been exposed during the training phase. Two metrics of accuracy (RMSE and RRMSE) are plotted in the next sections considering the systems in the order of increasing complexity. The plots for all metrics considered in Section "Metrics of accuracy" and for all systems considered in the benchmark are provided in the supplementary material.

For the legends of the plots the following convention is used `[method]_[Case]_[number of training points]: the legend for the respective method starts with the method which can be GP_AL (GP active learning), RF-GP (Random Forest GP), NN_Matlab, NN_Pytorch, NN_TensorFl (Neural network model trained in Matlab, Pytorch or Tensor flow according to Section "Neural Networks"), DecTree (DecTree method). It is followed with the case acronym e.g. 10_PC, 20麦克, 30_SC, 35_SC and the training dataset information: O for Orchestra, P for PhreeqC, and G for GEMS accompanied with the number of samples (500, 1000, 5000, 50,000).

## Cementitious system

### Primitive cement CaO–SiO₂–H₂O (10_PC)

The basic cementitious system considered in this benchmark has a three dimensional input (CaO, SiO₂ and H₂O) and 18 dimensional output. Output is the aqueous species after equilibration, as well as the amounts and composition of the resulting solid phases and solid solutions, including the amount of gel water within the resulting CSH phases. In Figs. 8 and 9, the metrics of the output (RMSE and RRMSE respectively) are presented for five different models, two based on Gaussian Processes and three based on neural networks.

The RMSE plot is interpretable since it has the same units as the output parameters (Molar concentrations for aqueous species for example). The training has been conducted using a clearly defined set of 5000 sample points, produced by the geochemical solvers, while the metrics are calculated using an independent sample set, containing 50,000 samples. The details of the parametrization, training procedure
and validation for each method can be found in the sup-
plementary information and the accompanying digital files.
The main observations and interpretation of the results is
as follows:

* All methods achieve different accuracy for different output parameter. Here GP methods have focused in creating one model for each output parameter, resulting in as many models as the number of outputs (18 models each one with 3 input/1 output for the 10_PC model). RF-GP requires a GP for each output within a group. The neural network approaches have focused in creating a single model connecting the input-output (1 model 3 input/18 output for the 10_PC model).

* GP methods show variable performance. RF-GP provides good values of RMSE for most outputs. The RRMSE for amorphous silica for RF-GP is having its highest value. GP tends to achieve less accurate results compared to the RF-GP and the NN models. As expected, the results of RF-GP are slightly better than those of GP because RF-GP combines the benefits of a classification algorithm and a GP. Another option for GP models, is the active learning to choose iteratively the most relevant training data to improve the performance.

* NN methods also show variable performance. Among the NN methods, the NN_Matlab seems to achieve overall the best accuracy with relatively small size model in terms of number of neurons (NN_Matlab ~ 160 neurons, TF TensorFlow ~ 1K neurons, PyTorch ~ 1K neurons). This is attributed mainly to the implemented Levenberg-Marquardt algorithm, compared to the ADAM (and its variations) algorithm used by NN_TensorFL and NN_PyTorch. Similar findings were also observed in (Wang et al. 2021) signifying the continuous development and optimization of NN algorithms. Moreover, different hyperparametrization (L-M: tanh, TF/PyTorch: RELU type of activation function) of the NN models was implemented (OPTUNA optimization and Expert knowledge). The Levenberg-Marquardt algorithm performed exceptionally in smaller dimensional systems and relatively small NN sizes. For the more complex cases the training effort using the Levenberg-Marquardt algorithm increases dramatically and the achieved accuracy, while it is maintained at relatively good levels, seems to be not significantly better than the ADAM based accuracy on larger networks. Indeed, the Levenberg-Marquardt (LM) algorithm uses a modified version of the Gauss-Newton algorithm, which incorporates an approximation of the Hessian matrix. The Hessian gives information about

<figure><img src="image_8.png" /><figcaption>**Fig. 5** Regression plot between GEMS and ORCHESTRA output for the primitive cementitious system 10_PC (5k training datasets) shown for AmorfSi, on linear and logarithmic axes respectively. The red line is the line where the results of GEMS and ORCHESTRA are found if they are identical. The blue crosses are the different results for the case 10_PC. Most of the points are very close to the red line. The two plots highlight the effect of a single outlier (the blue cross</figcaption></figure>

far from the red line on the righth log plot) on the normal and relative
metrics as perceived also schematically. A single outlier at the lower
end of the values has no effect on the normal metrics while it greatly
affects the values of the relative metrics. The metrics applied for the
comparison of GEMS and ORCHESTRA are also reported within the
figures
<figure><img src="image_9.png" /><figcaption>**Fig. 6** Output Metrics of cementitious system of different complexity between the geochemical solvers GEMS (_G_) and ORCHESTRA (_O_). Each bar represents a different complexity cementitious system</figcaption></figure>

(10_PC, 15_PCc, 25_MCc, 30_SC and 35_SCc cases). Comparison is done for the characteristic output parameters which are common to all systems. RMSE is shown

<figure><img src="image_10.png" /><figcaption>**Fig. 7** Output Metrics of cementitious system of different complexity between the geochemical solvers GEMS (_G_) and ORCHESTRA (_O_). Each bar represents a different complexity cementitious system</figcaption></figure>

(10_PC, 15_PCc, 25_MCc, 30_SC and 35_SCc cases). Comparison is done for the characteristic output parameters which are common to all systems. RRMSE is shown

second-order derivatives. However, this comes at the cost of computing the Hessian matrix, an O(N²) operation where N² is the size of the matrix, and storing it. Furthermore, LM generally uses the full dataset to compute gradients and update directions which also affects the

performance. Conversely, gradient descent can make full use of batch processing to reduce memory requirements. The findings highlight the potential of the neural network methods in improving their performance based on new training algorithms and implementations yet to come.
* It seems that the solid phases are predicted with better accuracy, compared to the aqueous phases for all methods. This can be partially attributed to two factors: (a) when total masses are used as input and the equilibrium predicts the presence of solid phases, typically the amount of mass which remains in the aqueous phases can be smaller by several orders of magnitude compared to

the mass in the solid phases, (b) a simpler relationship between the input, and the output solids, compared to the output aqueous parameters might exist. We note that for the ML methods the input data are typically normalized to increase the accuracy in the case where the output parameters span a large range of scales. The NN_Matlab on top of the normalization has been trained using the log

<figure><img src="image_11.png" /><figcaption>**Fig. 8** Output metrics for cementitious system 10_PC (CaO–SiO$_2$–H$_2$O). Each bar represents a different method of creating the surrogate model. The difference between GEMS and ORCHESTRA is also plotted. RMSE is shown</figcaption></figure>

<figure><img src="image_12.png" /><figcaption>**Fig. 9** Output metrics for cementitious system 10_PC (see also Table 1). Each bar represents a different method of creating the surrogate model. RMSE is shown</figcaption></figure>

<figure><img src="image_13.png" /></figure>
<figure><img src="image_14.png" /><figcaption>**Fig. 10** Improvement of accuracy for the cementitious system 10_PC (CaO-SiO₂-H₂O) for the neural network implementation Pytorch. Each bar represents a different training dataset: 500, 5000 and 50,000</figcaption></figure>

samples. A clear trend can be observed: The increase of the number of training samples has a positive effect on the resulting accuracy

of the input/output values in an effort to bridge this gap between different parameters. GP log transformation did not improve the results, and a min-max normalization was used. For the NN-PyTorch training, each input and output dimension was standardized (that is, mean and standard deviation set to zero and unity, respectively). It is also important to mention that for a reactive transport simulation, the accurate prediction of the aqueous phases plays typically a more significant role and should be prioritized over the accuracy of the solid phases.

* The achieved accuracies can be compared with the differences between the three geochemical solvers for the evaluation of the same system. While the differences between PHREEQC and ORCHESTRA results are extremely small, the differences with GEMS results are slightly larger due to the fundamental differences between LMA and GEMS approaches and algorithms. Therefore, this difference sets a lower threshold of accuracy beyond which the further training and fidelity of the surrogate models would have no practical use. From Figs. 6, 7 follows, however, that the differences between ORCHESTRA and GEMS are generally smaller than the differences between the surrogate models and the respective geochemical solvers, which for some outputs it can be 1 or even 2 orders of magnitude.

An interesting feature of all ML methods is the effect of number of samples on the accuracy of the resulting model. In principle, as shown in Fig. 3 for GP, Fig. 10 for NN and Figs. 17, 18 for the DecTree method, a larger training dataset will result in an improvement of the accuracy. In Fig. 10 the RMSE metrics of a NN model (PyTorch) for the system 10_PC are plotted. For the training, the exact same network architecture, activation functions and algorithm were used in order to see the effect of the size of the training dataset while keeping the rest of the hyperparameters the same. With the increase of the number of the training samples, the space of solutions is sampled at a higher resolution and the model is exposed to more detailed features. A similar plot for the GP, GP-AL method along with a comparison with the MARS and linear interpolation method can be found in the Supplementary material file, highlighting the increase of accuracy with increasing number of training points (or by enabling the active learning approach), as well as the increased accuracy compared to a linear interpolation approach.
## Minimal system CaO-SiO$_2$-Al$_2$O$_3$-H$_2$O-CO$_2$ (20 MC)

The next system under consideration is the slightly more complex 20\_MC with 5 input and 21 output. The effort in training and the complexity of the system increases mainly due to the number of input parameters rather than the output.

In Figs. 11 and 12 the results of the case 20\_MC are shown for two GP and two NN approaches.

Interpretation of the results leads to the following conclusions:

* For some of the outputs and GP-N/AL models show better accuracy compared to the NN models. This probably

<figure><img src="image_15.png" /><figcaption>**Fig. 11** Output metrics of cementitious system MC_20 (CaO-SiO$_2$-Al$_2$O$_3$-H$_2$O-CO$_2$). Each bar represents a different method of creating the surrogate model. RMSE is shown</figcaption></figure>

<figure><img src="image_16.png" /><figcaption>**Fig. 12** Output metrics of cementitious system MC_20 (CaO-SiO$_2$-Al$_2$O$_3$-H$_2$O-CO$_2$). Each bar represents a different method of creating the surrogate model. RRMSE is shown</figcaption></figure>

<figure><img src="image_17.png" /></figure>
means that these chemical species (solid phase) depend
only on a very few input variables. However, on average
the NN and GP seem to perform equally good in this
case. (Or all models perform relatively well).

* The GP-AL with the implementation of active learning seems to provide similar results to the naive GP but using a significant smaller training sample (1000 samples instead of 5000). This suggests that a careful selection of the training samples (done here by the active-learning algorithm) can have a significant effect on the accuracy for all methods.

* The aqueous species are still predicted by all methods with less accuracy compared to the solid phases when the RRMSE metric is considered but with higher accuracy when the RMSE metric is considered.

* The overall accuracy of the models for this case is close to the accuracy of the PC_10 case.

Complex system CaO-SiO₂-Al₂O₃-SO₃-K₂O-H₂O (SC_30)
and case CaO-SiO₂-Al₂O₃-SO₃-K₂O-H₂O-CO₂ (ScC_35)

The next systems under consideration are the more com-
plex systems: the SC_30 with 6 input and 28 output, and
the SC_35 with 7 input and 44 output. The latter one is
the most complex system explored in this benchmark. The
effort in training and the complexity of the system increases
again mainly due to the increase of the number of input
parameters rather than the output. The results are illustrated
in Figs. 13, 14, 15, 16.

The main observations and interpretation of the results in
the more complex cases is as follows:

* Solids amounts are better predicted because most of mass is in solid phase when precipitation occurs. However, the aqueous species are more important for the transport and their mass conservation is also crucial for reactive transport simulations.

* RF-GP provides good values of most metrics for most outputs. However, the RRMSE are large for some mineral phases.

* The surrogate models required a larger training set to achieve similar accuracy compared to the simpler cases. This is likely due to the increase of both input and output parameters together with more nonlinear relationships between them. For the most complex case (7 input) it seemed that 50,000 samples were adequate, which is an order of magnitude more samples compared to the simplest case considered (3 input) which provided reasonable results with 5000 points.

* The computational effort needed for training the surrogate models increases following the increased dimensionality of the complex case, since in principle larger models with more variables have to be implemented to describe the existing non-linearities between the input and output data. At the same time more samples have to be used for the training of the models increasing for example the time needed for an epoch during the neural network training process.

<figure><img src="image_18.png" /><figcaption>**Fig. 13** Output metrics of cementitious system SC_30 (CaO-SiO₂-Al₂O₃-SO₃-K₂O-H₂O). Each bar represents a different method of creating the surrogate model. RMSE is shown. For comparison, the difference between GEMS-ORCHESTRA is also plotted</figcaption></figure>
<figure><img src="image_19.png" /><figcaption>**Fig. 14** Output metrics of cementitious system SC_30 (CaO–SiO₂–Al₂O₃–SO₃–K₂O–H₂O). Each bar represents a different method of creating the surrogate model. RRMSE is shown. For comparison, the difference between GEMS-ORCHESTRA is also plotted</figcaption></figure>

* For the more complex cases, the GP and GP-AL methods could not handle at the moment larger training datasets and the results with the maximum amount of samples is reported. The GP methods are dependent on the inversion of a matrix which is built from the training samples. With increasing number of samples the size of the matrix reaches the limits of the computer memory and CPU time needed for the calculation. The GP methods have used several random subset of the training datasets to ensure that the order of magnitude of the accuracy of the results is not affected.

* The surrogate model's accuracy remains in principle close to the difference between the two geochemical solvers GEMS-ORCHESTRA and further optimization of the surrogate models or an increase of the training datasets (for the NN models), while possible, was not further considered.

## Uranium system

The base-case for the Uranium system is presented in this section. The original system consists of a three dimension input (HCl/NaOH/U) and a 9 dimension output. Due to the nature of the problem at hand, there have been different modelling approaches when considering the training dataset. In this benchmark the geometric method DecTree was also tested. A dimension reduction was implemented by the Geometric method which considered the pH as the governing variable (2-dimension input pH/U, 8 output) as described in detail in Section "Machine learning methods".

The RF-GP implemented 9 models of 4 dimension input (total uranium U, total acid HCl, total base NaOH, and group integer identifier) and 1-dimension output. NN- Matlab approach considered the data as is, implementing a single model with three dimensional input (3-dimension HCl/NaOH/U) and 9 dimension output), described by 4 hidden layers with 20 neurons each. The DecTree method was trained using three different training datasets with increasing resolution: 5K, 20K and 50K datasets. The training of the RF-GP was performed with 4000 points (selected from the 500K data file with all 4 groups represented equally). The neural network method addressed the problem by using a 5K training sample dataset. DecTree and RF-GP validation was made by using the 20K validation file, while the validation of the neural network model was conducted using the 50K validation file. The accuracy of the surrogate models for the considered 8 output parameters (omitting the 9th parameter pH) appear in Figs. 17 and 18

The main observations and interpretation of the results can be summarized as follows:

* All methods provide consistent results with the best accuracy provided by the DecTree method when the high resolution 50K training dataset was used. The RF-GP method at this stage provide highly accurate results as shown in Fig. 17 at the RMSE plot. However, the RRMSE plot shows a higher error for the RF-GP for aqueous, exchanged and sorbed U compared to the rest of the methods which can be attributed to less accurate

<figure><img src="image_20.png" /></figure>
<figure><img src="image_21.png" /><figcaption>**Fig. 15** Output metrics of cementitious system ScC_35 (CaO–SiO₂–Al₂O₃–SO₃–K₂O–H₂O–CO₂). Each bar represents a different method of creating the surrogate model. RMSE is shown. For comparison, the difference between GEMS-ORCHESTRA is also plotted</figcaption></figure>

<figure><img src="image_22.png" /><figcaption>**Fig. 16** Output metrics of cementitious system ScC_35 (CaO–SiO₂–Al₂O₃–SO₃–K₂O–H₂O–CO₂). Each bar represents a different method of creating the surrogate model. RRMSE is shown. For comparison, the difference between GEMS-ORCHESTRA is also plotted</figcaption></figure>

results for very low aqueous, exchanged and sorbed U concentrations.

* The DecTree results show as expected in most of the outputs a clear trend of improved accuracy by increasing the number of samples. This is mostly driven by the

increasing accuracy of the geometric split by convex hull, while the refinement in the interior regions of the different subspaces plays a less important role.

* The neural network model results, with the exception of the Metaschoepite is in close agreement with
<figure><img src="image_23.png" /><figcaption>**Fig. 17** Benchmark results of uranium case. The accuracy of the DecTree method for training datasets of 5K, 20K and 50K points is shown along with the results of the neural network NN_Matlab (5K</figcaption></figure>

training points) and the RF-GP methods (4K training points). The RMSE metric is plot

the results of DecTree for 5K and 20K samples. The RRMSE accuracy of Metaschoepite is at the same level as the RF-GP method.

* The accuracy of the surrogate models in the uranium case seems to be somewhat lower compared to the simple and complex cementitious cases. For example, the RRMSE of U_aq is at the level of 10⁻²-10⁻¹, while the RRMSE for the aqueous species in the cementitious case is in principle an order of magnitude lower. This signifies that the achieved accuracy of the surrogate models is not dependent only on the dimensionality of the chemical system of interest, but also on the respective multi-dimensional topology of the solution space and the range of variables. The latter can be complex enough even for low dimensional problems. It is noted that a possible challenge might arise because of the Uranium concentrations which cover much larger ranges in log units.

* Another possibility of the difference in accuracy for the cementitious and uranium system is the fact that the cementitious system is dominated by solid phase equilibria which effectively reduce the degrees of freedom (or dimensionality) by making input parameters dependent on each other. The degrees of freedom in a thermodynamic system are given by the Gibbs phase rule: F = 2 + C - P (where F is the number of degrees of freedom, C is the number of [independent] components, and P is the number of phases). In the cementitious sys-

tem, the number of phases P is related to the number of components C, the phases lock specific ratios of components, and dictate the concentration of various ions thus reducing the degrees of freedom F.

## Comparison of computational efficiency between the geochemical solvers and the surrogate models

In this section, the trained surrogate models are benchmarked against the geochemical solvers to evaluate the difference in the computational efficiency of native and surrogate algorithms. Performance on single-core and multi-core CPU as well as on GPU computational units is considered. All the geochemical solvers are currently limited to run on CPU architectures, while most of the surrogate models can be executed both on CPU or GPU. PHREEQC and GEMS do not offer a direct parallelization option for batch calculations while ORCHESTRA can use by default all the available CPU cores. In practice, during a reactive transport simulation the PHREEQC and GEMS solvers can also be used in parallel by using multiple solver instances on the same CPU chip and launching parallel tasks, a process which has to be programmed explicitly.

Figure 19 shows the computational performance of equilibrating 500,000 samples of two traditional geochemical solvers, ORCHESTRA and PHREEQC, and of the surrogate

<figure><img src="image_24.png" /></figure>
<figure><img src="image_25.png" /><figcaption>**Fig. 18** Benchmark results of uranium case. The accuracy of the DecTree method for training datasets of 5K, 20K and 50K points is shown along with the results of the neural network NN_Matlab (5K</figcaption></figure>

training points) and the RF-GP methods (4K training points). The RRMSE metric is plotted

models trained on the cement system 10_PC. The performance of GEMS algorithm was similar to that of PHREEQC and is not plotted in the graphs. All the benchmarks are run on the same machine, equipped with an Intel i9 10900F @ 2.80 GHz (10-core/20-hyperthreads, produced in year 2020) CPU, 64 GB of RAM @ 2933 MHz and a relatively old GPU, the Nvidia Geforce GTX 1070 with 8 Gb of dedicated RAM (produced in year 2016).

Results are reported (Fig. 19) for PHREEQC with a single core and with PHREEQC used in batch mode. Note that a different version of PHREEQC with parallel computation exists called PHREEQCRM. However, this is effectively N instances of PHREEQC running in parallel, therefore the performance should scale linearly with the number of CPU cores. The typical performance of all three geochemical solvers (PHREEQC, GEMS, ORCHESTRA) for a random sample file was in the range of 10²–10⁴ equilibrations/s. However, for the traditional solvers the performance on a large random sample dataset may not represent the efficiency that can be achieved in a reactive transport simulation. Random input would represent a case where the geochemical solver is used for sensitivity analysis or in a case where drastic geochemical changes occur relatively fast. The geochemical solvers calculate a chemical equilibrium system by iteratively solving (or minimizing) a set of non-linear equations, which requires start estimations (initial guess) of the “unknowns” as input. The number of required iterations (which is approximately proportional to calculation

time) to reach a solution depends strongly on the accuracy of the start estimations. For transport systems, generally very good start estimations (so-called warm-start initial guess) are available, resulting from chemical speciation obtained at the previous time step for the same computational grid point. For the random sample dataset, there is no relation between the subsequent datapoints. As a result, solving each independent datapoint requires the “cold-start initial guess” leading to many more iterations. To demonstrate the order of magnitude of this effect, the results of transport calculations with ORCHESTRA for the cement as well as the uranium system are added (Fig. 19).

GEMS and PHREEQC (the PhreeqcRM reaction module) also are reported to have the possibility to use information from previous calculations as a warm-start initial guess for a new one, but for these codes the efficiency and resulting accuracy was not evaluated here.

Surrogate models’ performance is reported on both a 10 physical CPU cores (20 threads), and TensorFlow it is also reported on the GPU. Although many current tools in scientific computing do not yet take full advantage of GPU architectures, they are particularly advantageous when training a NN or inferring from it. At the same time there exist several flow solver implementations (commercial and scientific) which can run on single or parallel GPU which can be coupled to a GPU machine learning model. A NN model expressed as a C function and running on the CPU achieves the best efficiency followed by the NN TensorFlow model
on the GPU. NN implementations are approximately two orders of magnitude faster than the GP and the ORCHES-TRA in reactive transport mode, and five orders of magni-tude faster from PHREEQC in batch mode random input on a single CPU core. The latter holds also for GEMS and ORCHESTRA.

When looking at surrogate methods, it appears that all of them yield a significant speedup in computational effi-ciency compared to the standard “cold start”/random input use of the geochemical solvers. Under transport conditions and when ORCHESTRA is operated with “warm initializa-tions” it achieves similar efficiency to that of GP-AL and TensorFlow on the tested CPU (about 600K calculations/s). At the same time, RF-GP reach an efficiency of about 2000 calculations per second for all the 18 model outputs describ-ing the cementitious case.

Interestingly, inference of NN output from Matlab, auto-matically coded and compiled in C source code, results in a drastic increase in performance compared to the rest of the implementations. Similar performance is observed within the Matlab environment. The C executable is able to per-form about ~40 million calculations/s using the 10-CPU cores, outperforming even the inference of TensorFlow on a GPU (~30 million calculations per second). The very sig-nificant gains in speed is explainable by diverse factors. The observed efficiency can be attributed to the simplicity of the standalone C-code, which is essentially a simple func-tion which can be readily integrated in any CPU reactive transport code. It is also attributed to the NN model itself, which is smaller and has much less variables compared to the native TensorFlow model, for the same system. As the number of parameters increases, the inference time increases as well. The exact functions, models and executables are provided in the online Zenodo repository. Furthermore, although a lot of effort is invested in closing the gap between

a C++ and a python interpreted program (e.g. Cython), an interpreted language is not likely to outperform a statically compiled language in raw speed. The NN, expressed as a simple mathematical C function can benefit directly from the compiler optimization procedures, compared to a model running within a software environment. Finally, the GPU that is used in this benchmark is far from the state-of-the-art standards now available to both the industry and consumers (8 years old), while the CPU is relative newer and high-end (4 years old). The NN C-code function was observed to exhibit a perfect parallelization on the CPU cores using the openmpi protocol. The conjunction of these factors can explain the difference in the relative performance of the CPU and GPU implementations presented here. While it is also possible to export the Matlab function in CUDA code for-mat, which would be native also for running on the GPU and a further speed-up is to be expected, this possibility was not further explored in this benchmark.

Figure 20 summarizes the results of the code perfor-mance on equilibrating 500K samples for uranium system Mc. Identically to the cement system, ORCHESTRA’s performance is on par with the surrogate models on CPU, with about 330K calculations/s in a “warm start” condition. PHREEQC, ORCHESTRA and GEMS performance under “cold start” remains at the same relatively low levels as the cement case. The DecTree approach reaches more than 500K calculations/s (on a single CPU thread), while NN inference with TensorFlow on CPU yields 200K calculations/s. The use of GPU architecture for NN-based codes results in 2 orders of magnitude increase in performance when com-pared to the traditional geochemical solver.

As expected, surrogate models generally outperform traditional geochemical solvers by several orders of mag-nitude when “cold start”/random input conditions are con-sidered highlighting a significant acceleration potential

<figure><img src="image_26.png" /><figcaption>**Fig. 19** Computational effi-ciency of the geochemical solv-ers and of the machine learning models. Neural network model expressed as a C function and running on the CPU achieves the best efficiency followed by the neural network TensorFlow model on the GPU. Neural network implementations are approximately two orders of magnitude faster than the GP and the ORCHESTRA in reac-tive transport mode, and five orders of magnitude faster from PHREEQC in batch mode ran-dom input on a single CPU core</figcaption></figure>

<figure><img src="image_27.png" /></figure>
for applications where many geochemical calculations are needed (e.g. sensitivity analysis). However, the gap between the performance of surrogate models and geochemical solvers may be less extreme in reactive transport simulations, where good initial guesses for the solution are usually available from previous time steps. As general conclusion, it can be strongly recommended to improve and optimize the warm start capabilities of the conventional geochemical codes. This has been demonstrated for the ORCHESTRA code when used in “warm start” mode during a reactive transport simulation. Such an increase in efficiency is still to be proven for more complex cases and with the geochemical codes GEMS and PHREEQC, which shall be pursued in the future.

While these results point out promising speedup in view of potential applications to reactive transport simulations, it remains to be assessed how the resulting approximations would propagate in coupled simulations. The effect of the accuracy of the surrogate model on the reactive transport simulation during a large number of timesteps has been a subject of several works (Jatnieks et al. 2016; Prasianakis et al. 2020; Laloy and Jacques 2022; Silva et al. 2024) where it is highlighted that low accuracy surrogates may result in a simulation drift when for example the mass balance accuracy error is accumulated at the course of the simulation. In general terms, surrogate predictions are associated with an expected error variance. Methods such as Gaussian Processes can automatically deliver such estimations, but this is not the case for other approaches or when the surrogates are coupled with another solver. De Lucia and Kühn (2021) suggest to employ mass and charge balance equations to check plausibility of data-driven surrogate predictions, and to solve numerically the rejected ones. Similar considerations and implementations are tested in Silva et al. (2024). Mass and charge balance could also be enforced by recasting the regression problem in terms of reaction rates, thus effectively departing from a purely data-driven ML approach.

These aspects, however, are outside the scope of the present work and will be the subject of future investigations.

## Remarks on the application and performance of current ML modelling approaches

Based on the conducted benchmarking exercise, it is possible to assess the conditions and the classes of problems for which the application of ML can be highly beneficial. In such considerations, it is crucial to include not only the potential savings in computational time during the simulation, but also the effort needed to create and deploy the surrogate model.

In addition to inference time (calculations/s), it is also necessary to take into account the training time to generate the surrogate models. For most of the systems considered and NN sizes, TensorFlow and Pytorch NNs take from a few minutes (~1K-neuron models) to an hour (~100K-neuron models) to train an instance on a common GPU. In turn, training Matlab NNs on a single CPU core can take between a few hours to a few days for the most complex systems and largest training dataset when the Levenberg-Marquardt optimizer is used. Conversely, training the DecTree model for the Uranium case is almost instantaneous after the expert user has designed the model and integrated the physical constraints. Gaussian Processes with active learning (GP-AL) can be computationally expensive or even prohibitive for large dimension problems but the surrogate model is automatically created with no parameter to tune and with the best and minimal training dataset. A statistical indicator (the variance) for each estimated value is also provided. The estimated CPU time needed for RF-GP training for the 10_PC case is on the order of 11 h, for 30_SC is on the order of one to two weeks and for U case is on the order of 4 h.

<figure><img src="image_28.png" /><figcaption>**Fig. 20** Computational efficiency of the geochemical solvers and of the machine learning models. Neural networks run on the GPU show the best performance being two orders of magnitude faster compared to the CPU implementations of ORHESTRA (in reactive transport conditions using 20 threads), of DecTree (single thread) and of neural networks TensorFlow (20 threads)</figcaption></figure>
When considering implementing ML algorithms in reactive transport simulations or any other application, it is paramount, already at the stage of the design of the computational approach, to estimate the computational cost needed to create the model versus the expected computational benefit (savings in time). In an effort to rationalize this consideration it is necessary to breakdown the overall time costs.

The model development costs *Cost_Dev* in terms of time include:

- (A), the cost for producing an adequate training dataset, e.g. 500–50,000 samples (typically in order of minutes for the geochemical solvers).

- (B), the cost for creating and training the model, e.g. minutes to days, depending on the accuracy and algorithm, including, if any, the dependencies on “expert” intervention to prepare the model/arrange data. At this point, it is stressed that the training effort strongly depends on the desired accuracy of the surrogate model and the available computational resources. In that sense, a criterion of accuracy of the surrogate model needs to be set in advance.

- (C) is the cost to deploy and integrate the model to the transport code or the respective numerical framework. (for automated workflows this is typically negligible, while for first-time implementations and inexperienced users it might require significant time)

In summary, the total costs for the model development is $Cost_Dev = A + B + C$.

An estimate of the simulation computational costs can be similarly calculated when using the geochemical code *Cost_Geochem* or the surrogate ML model *Cost_ML* which include:

- (D) the number of total geochemical calculations per simulation (could range $10^6–10^{15}$), which is the number of discretization points where geochemical calculations take place, multiplied with the number of time steps.

- (E) the number of simulations that need to be performed (e.g. a single one or many as in sensitivity analyses, Inverse modeling, optimization problems, etc.).

The expected time costs for the simulations would be:

- $Cost_Geochem = (D)*(E)/(Geochemical solver speed in calcs/s)$

- $Cost_ML = (D)*(E)/(ML solver speed in calcs/s)$

A simple comparison of the costs for a desired accuracy of the model provides a reference of when it makes sense to switch to a ML enhanced model or to stick on the traditional approach. That is, if $Cost_ML + Cost_Dev < Cost_Geochem$, then it is recommended to proceed with the ML assisted algorithms. As a final step, when considering coupled

processes, the computational cost for the geochemical part has to be compared with the computational cost of the rest of the algorithms (e.g. transport, mechanics), which then allows to calculate a more precise overall gain of computational cost. It is recalled here that in Multiphysics codes the geochemistry is typically one of the most expensive code modules.

It is evident that for small problems which are required to be calculated only once and the number of grid points/time steps is small, the classical direct approach will be preferred. For larger problems and for cases where the same chemical system is repeatedly used, which could be several thousands of total simulations for a sensitivity or uncertainty quantification analysis, the surrogate model implementations will be the only realistic way to proceed. For the latter class of problems the use of ML approaches seems to open new horizons since more complex high resolution problems can be considered. Along with the argument relevant to the feasibility it is important to consider that more efficient simulations not only decrease the times to solution but can also be translated directly to energy savings and smaller environmental footprint. The energy savings are not only relevant when laptops or workstations are used but also when supercomputing units are considered. The top 10 systems in the list of the top500 supercomputers (Top500 2024) have a power consumption of 5–30 MW each. Depending on the energy source which is used to power them up it is theoretically possible to calculate an equivalent amount of possible CO₂ emission reduction, per simulation, when simulation acceleration techniques are used.

It should be noted that once a surrogate model for a specific chemical system and concentration range is created, it is possible to be used for different simulation setups which however consider the same chemical system without the need for retraining or further development. On the contrary, any slight change in the geochemical system, or in the limits where the surrogate model is trained requires the creation of a new surrogate model. For small extensions of the chemical systems, it is possible that a fine tuning of the existing surrogate models could be sufficient. For example, transfer learning techniques could be used for the NN models. It is also anticipated that the more the surrogate models are used for simulations, the larger the (open) libraries of such models will grow to be used across the scientific community.

## Conclusions

In this work, a benchmark for testing geochemical codes and machine learning (ML) algorithms has been established with the purpose of understanding the potential of the latter as surrogate models at a significantly reduced computational cost. Moreover, the benchmark has served as a vehicle
to gain confidence on ML methods applied for geochemical calculations. While the three geochemical solvers used have a distinct conceptual and algorithmic approach, an excellent agreement has been obtained between them, setting up the basis for benchmarking other existing geochemical solvers or for testing novel algorithmic approaches. All necessary files for reproducing the chemical systems and performing the comparisons are provided as supplementary material and information as well as at the online Zenodo repository.

As a next step, a workflow has been established with the aim of producing standardized and high-quality training datasets for the ML algorithms. These datasets can be used as a training and testing base for the development of any kind of ML method even if not related to geochemistry. Several ML approaches have been implemented while a few of them were still further developed during the progress of this benchmark (RF-GP, DecTree, GP-AL). The comparison shows a high consistency among the considered ML methods which reach high levels of accuracy. Finally, the efficiency of the produced models was benchmarked in different computer hardware architectures.

The results highlighted that under circumstances the gains in efficiency by using ML models can be of two to several orders of magnitude, compared to the standard use of the geochemical solvers. A “warm start” initial guess operation mode proved to also provide significant efficiency gains for one of the considered geochemical solvers, without however surpassing the efficiency of the top-efficient surrogate models. The exact performance gains during realistic reactive transport simulations will strongly depend on the coupling and the efficiency of the flow solver and the system under consideration. A framework for assessing the applicability of ML methods by considering the overall effort and expected computational gains has also been provided. It is clear that geochemistry is the bottleneck in current solver implementations and a speedup of geochemical calculations will result in a speedup of the coupled algorithm. The efficiency in reactive transport simulations will be carefully considered in a future dedicated benchmark effort.

Finally, an important factor to be considered for the future success and acceptance of the ML applications in the geochemical modelling community will be the ability of code developers to provide an automated workflow for all stages of model setup. Namely, the data generation, surrogate model design, model training, validation and testing. The broad scientific community is primary interested in geochemical applications rather than the details of the code development. Therefore, the availability of robust workflows and user friendly interfaces to application software guiding the user through the most important stages of the model development process would be an important target for future method development.

**Supplementary Information** The online version contains supplementary material available at https://doi.org/10.1007/s12665-024-12066-3 and at the Zenodo online repository https://doi.org/10.5281/zenodo.11274789.

**Acknowledgements** C. Tournassat is acknowledged for providing the input for the geochemical modelling of the Uranium system.

**Author contributions** The geochemical model and production of training datasets were performed for ORCHESTRA by J.M., for PHREEQC by D.J. and for GEMS by D.K. The machine learning model implementation and results were performed for DecTree by M.D.L. for RF-GP J.S. and J.S.II, for GP/GP-AL by M.S., M.L., F.B., for Neural Network Matlab and C by N.I.P., for Neural network Tensorflow by E.D., E.C., A.I. and for PyTorch by E.L. All authors have contributed to the manuscript, methodology development, interpretation of results. All authors have reviewed the manuscript and approved its publication.

**Funding** Open Access funding provided by Lib4RI – Library for the Research Institutes within the ETH Domain: Eawag, Empa, PSI & WSL. Support through the DONUT EURAD WP is acknowledged, funded from the European Union’s Horizon 2020 research and innovation program under grant agreement No. 847593. M.D.L. gratefully acknowledges the Helmholtz Association for financial support within the T^6 project (grant ZT-1-PF-5-084). UDC contribution received funding from ENRESA (Ref 79-CO-IS-2019-0001), Spanish Ministry of Science and Innovation (PID2023-153202OB-I00 & TED2021-130315B-I00), the Galician Regional Government (Grant ED431C2021/54) and the University of A Coruña. Support through the PREDIS (Pre-disposal management of radioactive waste, Euratom research and training programme, grant agreement No 945098) is also acknowledged.

**Data availability** Codes, geochemical models, training sets, validation sets and some of the ML models are provided at the Zenodo repository: https://doi.org/10.5281/zenodo.11274789

## Declarations

**Conflict of interest** The authors declare no competing interests.

**Open Access** This article is licensed under a Creative Commons Attribution 4.0 International License, which permits use, sharing, adaptation, distribution and reproduction in any medium or format, as long as you give appropriate credit to the original author(s) and the source, provide a link to the Creative Commons licence, and indicate if changes were made. The images or other third party material in this article are included in the article's Creative Commons licence, unless indicated otherwise in a credit line to the material. If material is not included in the article's Creative Commons licence and your intended use is not permitted by statutory regulation or exceeds the permitted use, you will need to obtain permission directly from the copyright holder. To view a copy of this licence, visit http://creativecommons.org/licenses/by/4.0/.

## References

Abadi M, Agarwal A, Barham P et al (2016) TensorFlow: large-scale machine learning on heterogeneous distributed systems. https://arxiv.org/abs/1603.04467v2. Accessed 10 May 2024

Águila JF, Montoya V, Samper J et al (2021) Modeling cesium migration through Opalinus clay: a benchmark for single- and multi-species sorption-diffusion models. Comput Geosci 25:1405–1436. https://doi.org/10.1007/s10596-021-10050-5
Akiba T, Sano S, Yanase T et al (2019) Optuna: a next-generation hyperparameter optimization framework. Proc ACM SIGKDD Int Conf Knowl Discov Data Min. https://doi.org/10.1145/3292500.3330701

Alpaydin E (2009) Introduction to machine learning (adaptive computation and machine learning series). MIT Press, Cambridge

Bildstein O, Claret F, Lagneau V (2021) Guest editorial to the special issue: subsurface environmental simulation benchmarks. Comput Geosci 25(4):1281–1283. https://doi.org/10.1007/s10596-021-10069-8

Bradbury MH, Baeyens B (1997) A mechanistic description of Ni and Zn sorption on Na-montmorillonite Part II: modelling. J Contam Hydrol 27:223–248. https://doi.org/10.1016/S0169-7722(97)00007-7

Breiman L (2001) Random forests. Mach Learn 45:5–32. https://doi.org/10.1023/A:1010933404324/METRICS

Buitinck L, Louppe G, Blondel M et al (2011) Scikit-learn: machine learning in python. J Mach Learn Res 12:2825–2830

Charlton SR, Parkhurst DL (2011) Modules based on the geochemical model PHREEQC for use in scripting and programming languages. Comput Geosci 37:1653–1663. https://doi.org/10.1016/j.CAGEO.2011.02.005

Churakov SV, Claret F, Idiart A, Jacques D, Govaerts J, Kolditz O, Samper J (2024) Position paper on high fidelity simulations for coupled processes, multi-physics and chemistry in geological disposal of nuclear waste. Environ Earth Sci 83(17):521. https://doi.org/10.1007/s12665-024-11832-7

Claret F, Dauzeres A, Jacques D et al (2022) Modelling of the long-term evolution and performance of engineered barrier system. EPJ Nucl Sci Technol. https://doi.org/10.1051/epjn/2022038

Claret F, Prasianakis NI, Baksay A, Lukin D, Pepin G, Ahusborde E, Kolditz O (2024) EURAD state-of-the-art report: development and improvement of numerical methods and tools for modeling coupled processes in the field of nuclear waste disposal. Front Nuclear Eng 3:1437714. https://doi.org/10.3389/fnuen.2024.1437714

Cybenko G (1989) Approximation by superpositions of a sigmoidal function. Math Control Signals Syst 2:303–314. https://doi.org/10.1007/BF02551274/METRICS

Dan Foresee F, Hagan MT (1997) Gauss-Newton approximation to Bayesian learning. Int Conf Neural Netw 3:1930–1935. https://doi.org/10.1109/ICNN.1997.614194

De Lucia M (2024) DecTree: a physics-based geochemical surrogate for surface complexation of uranium on clay. Adv Geosc. https://doi.org/10.5194/adoeo-65-45-2024

De Lucia M, Kuhn M (2021) DecTree v1.0 - chemistry speedup in reactive transport simulations: purely data-driven and physics-based surrogates. Geosci Model Dev 14:4713–4730. https://doi.org/10.5194/GMD-14-4713-2021

De Lucia M, Kempka T, Jatnieks J, Kuhn M (2017) Integrating surrogate models into subsurface simulation framework allows computation of complex reactive transport scenarios. Eur Geosci Union Gen Assem 2017 Egu Div Energy, Resour Environ 125:580–587. https://doi.org/10.1016/j.egypro.2017.08.200

de Windt L, Spycher NF (2019) Reactive transport modeling: a key performance assessment tool for the geologic disposal of nuclear waste. Elements 15:99–102. https://doi.org/10.2138/GSELEMENTS.15.2.99

Demirer E, Coene E, Iraola A, et al (2023) Improving the performance of reactive transport simulations using artificial neural networks. Transp Porous Media 149:271–297. https://doi.org/10.1007/s11242-022-01856-7

Goodfellow I, Bengio Y, Courville A (2016) Deep learning. The MIT Press, Cambridge, Massachusetts

Hagan MT, Menhaj M (1994) Training feed forward networks with the marquardt algorithm. IEEE Trans Neural Networks 5:989–993

Hjelle O (2001) Approximation of scattered data with multilevel B-splines. http://www.sintef.no/upload/IKT/9011/geometri/MBA/MBA.pdf. Accessed 10 May 2024

Hornik K (1991) Approximation capabilities of multilayer feedforward networks. Neural Netw 4:251–257. https://doi.org/10.1016/0893-6080(91)90009-T

Huang Y, Shao H, Wieland E et al (2018) A new approach to coupled two-phase reactive transport simulation for long-term degradation of concrete. Constr Build Mater 190:805–829. https://doi.org/10.1016/j.conbuildmat.2018.09.114

Jatnieks J, De Lucia M, Dransch D, Sips M (2016) Data-driven surrogate model approach for improving the performance of reactive transport simulations. Eur Geosci Union Gen Assem 97:447–453. https://doi.org/10.1016/j.egypro.2016.10.047

Joe S, Kuo FY (2003) Remark on algorithm 659. ACM Trans Math Softw 29:49–57. https://doi.org/10.1145/641876.641879

Kingma DP, Ba JL (2015) Adam: a method for stochastic optimization. In: 3rd international conference on learning representations, ICLR 2015 - conference track proceedings. International conference on learning representations, ICLR

Kolditz O, Bauer S, Bilke L et al (2012) OpenGeoSys: an open-source initiative for numerical simulation of thermo-hydro-mechanical/chemical (THM/C) processes in porous media. Environ Earth Sci 67:589–599

Kolditz O, Jacques D, Claret F et al (2023) Digitalisation for nuclear waste management: predisposal and disposal. Environ Earth Sci 82:42. https://doi.org/10.1007/s12665-022-10675-4

Krause MJ, Kummerländer A, Avis SJ et al (2021) OpenLB—Open source lattice Boltzmann code. Comput Math with Appl 81:258–288. https://doi.org/10.1016/j.CAMWA.2020.04.033

Kulik DA (2011) Improving the structural consistency of C-S-H solid solution thermodynamic models. Cem Concr Res 41:477–495. https://doi.org/10.1016/j.CEMCONRES.2011.01.012

Kulik DA, Wagner T, Dmytrieva SV et al (2013) GEM-Selektor geochemical modeling package: revised algorithm and GEMS3K numerical kernel for coupled simulation codes. Comput Geosci 17:1–24. https://doi.org/10.1007/s10596-012-9310-6

Kyas S, Volpatto D, Saar MO et al (2022) Accelerated reactive transport simulations in heterogeneous porous media using Reaktoro and Firedrake. Comput Geosci 26:295–327. https://doi.org/10.1007/s10596-021-10126-2

Laloy E, Jacques D (2019) Emulation of CPU-demanding reactive transport models: a comparison of Gaussian processes, polynomial chaos expansion, and deep neural networks. Comput Geosci 23:1193–1215. https://doi.org/10.1007/s10596-019-09875-y

Laloy E, Jacques D (2022) Speeding up reactive transport simulations in cement systems by surrogate geochemical modeling: deep neural networks and k-nearest neighbors. Transp Porous Media 143:433–462. https://doi.org/10.1007/s11242-022-01779-3

Leal AMM, Kyas S, Kulik DA, Saar MO (2020) Accelerating reactive transport modeling: on-demand machine learning algorithm for chemical equilibrium calculations. Transp Porous Media. https://doi.org/10.1007/s11242-020-01412-1

Lee S, Wolberg G, Shin SY (1997) Scattered data interpolation with multilevel b-splines. IEEE Trans vis Comput Graph 3:228–244. https://doi.org/10.1109/2945.620490

Lichtner PC, Hammond GE, Lu C et al (2015) PFLOTRAN user manual: a massively parallel reactive flow and transport model for describing surface and subsurface processes. https://doi.org/10.2172/1168703

Loh WL (1996) On Latin hypercube sampling. Ann Stat 24:2058–2080. https://doi.org/10.1214/aos/1069362310

Lothenbach B, Kulik DA, Matschei T et al (2019) Cemdata18: a chemical thermodynamic database for hydrated Portland cements and alkali-activated materials. Cem Concr Res 115:472–506. https://doi.org/10.1016/J.CEMCONRES.2018.04.018
MacKay DJC (1992) Bayesian interpolation. Neural Comput 4:415-447. https://doi.org/10.1162/NECO.1992.4.3.415

Marquardt DW (1963) An algorithm for least-squares estimation of nonlinear parameters. J Soc Ind Appl Math 11:431-441. https://doi.org/10.1137/0111030

Marques Fernandes M, Baeyens B, Dähn R et al (2012) U(VI) sorption on montmorillonite in the absence and presence of carbonate: a macroscopic and microscopic study. Geochim Cosmochim Acta 93:262-277. https://doi.org/10.1016/J.GCA.2012.04.017

Matlab (2021) MATLAB version R2021b. In: The Mathworks Inc., Natick, Massachusetts, United States. https://www.mathworks.com. Accessed 1 July 2024

Meeussen JCL (2003) ORCHESTRA: an object-oriented framework for implementing chemical equilibrium models. Environ Sci Technol 37:1175-1182. https://doi.org/10.1021/es025597s

Miron GD, Kulik DA, Dmytrieva SV, Wagner T (2015) GEMSFITS: code package for optimization of geochemical model parameters and inverse modeling. Appl Geochem 55:28-45. https://doi.org/10.1016/J.APGEOCHEM.2014.10.013

Misra D (2019) Mish: a self regularized non-monotonic activation function. 31st Br Mach Vis Conf BMVC 2020

Molins S, Soulaine C, Prasianakis NI et al (2020) Simulation of mineral dissolution at the pore scale with evolving fluid-solid interfaces: review of approaches and benchmark problem set. Comput Geosci 254(25):1285-1318. https://doi.org/10.1007/S10596-019-09903-X

Ozaki Y, Tanigaki Y, Watanabe S et al (2022) Multiobjective tree-structured Parzen estimator. J Artif Intell Res 73:1209-1250. https://doi.org/10.1613/JAIR.1.13188

Parkhurst DL (1995) User's guide to PHREEQC: a computer program for speciation, reaction-path, advective-transport, and inverse geochemical calculations. US Department of the Interior, US Geological Survey

Parkhurst DL, Appelo CAJ (2013) Description of input and examples for PHREEQC version 3—a computer program for speciation, batch-reaction, one-dimensional transport, and inverse geochemical calculations. US Geol Surv Tech Methods 6:497

Parkhurst DL, Wissmeier L (2015) PhreeqcRM: a reaction module for transport simulators based on the geochemical model PHREEQC. Adv Water Resour 83:176-189. https://doi.org/10.1016/J.ADVWA TRES.2015.06.001

Paszke A, Gross S, Massa F, et al (2019) PyTorch: an imperative style, high-performance deep learning library. Adv Neural Inf Process Syst 32:1-12

Patel RA, Perko J, Jacques D et al (2018) A three-dimensional lattice Boltzmann method based reactive transport model to simulate changes in cement paste microstructure due to calcium leaching. Constr Mater 166:158-170

Prasianakis NI, Haller R, Mahrous M et al (2020) Neural network based process coupling and parameter upscaling in reactive transport simulations. Geochim Cosmochim Acta. https://doi.org/10.1016/j.gca.2020.07.019

Rasmussen CE (2004) Gaussian processes in machine learning. Lect Notes Comput Sci (including Subser Lect Notes Artif Intell Lect Notes Bioinform) 3176:63-71. https://doi.org/10.1007/978-3-540-28650-9_4

Safi MA, Mantzaras J, Prasianakis NI et al (2019) A pore-level direct numerical investigation of water evaporation characteristics under air and hydrogen in the gas diffusion layers of polymer electrolyte fuel cells. Int J Heat Mass Transf 129:1250-1262. https://doi.org/10.1016/J.IJHEATMASTRANSFER.2018.10.042

Samper J, Xu T, Yang C (2009) A sequential partly iterative approach for multicomponent reactive transport with CORE2D. Comput Geosci 13:301-316. https://doi.org/10.1007/s10596-008-9119-5

Sarker IH (2021) Machine learning: algorithms, real-world applications and research directions. SN Comput Sci 2:1-21. https://doi.org/10.1007/S42979-021-00592-X/FIGURES/11

Savino M, Lévy-Leduc C, Leconte M, Cochepin B (2022) An active learning approach for improving the performance of equilibrium based chemical simulations. Comput Geosci 26:365-380

Seeger M (2004) Gaussian processes for machine learning. Int J Neural Syst 14:69-106

Sheikholeslami R, Razavi S (2017) Progressive Latin hypercube sampling: an efficient approach for robust sampling-based analysis of environmental models. Environ Model Softw 93:109-126. https://doi.org/10.1016/j.envsoft.2017.03.010

Silva VL, Regnier G, Salinas P, Heaney CE, Jackson MD, Pain CC (2024) Rapid modelling of reactive transport in porous media using machine learning: limitations and solutions. arXiv preprint arXiv:2405.14548. Accessed 8 Nov 2024

Šimůnek J, Brunetti G, Jacques D et al (2024) Developments and applications of the HYDRUS computer software packages since 2016. Vadose Zo J. https://doi.org/10.1002/VZJ2.20310

Smith LN (2017) Cyclical learning rates for training neural networks. Proc - 2017 IEEE Winter Conf Appl Comput Vision, WACV 2017 464-472. https://doi.org/10.1109/WACV.2017.58

Sochala P, Chiaberge C, Claret F, Tournassat C (2022) Uncertainty propagation in pore water chemical composition calculation using surrogate models. Sci Rep 12:15077. https://doi.org/10.1038/s41598-022-18411-5

Soulaine C, Pavuluri S, Claret F, Tournassat C (2021) porousMedia-4Foam: multi-scale open-source platform for hydro-geochemical simulations with OpenFOAM®. Environ Model Softw 145:105199. https://doi.org/10.1016/J.ENVSOFT.2021.105199

Srinivas N, Krause A, Kakade SM, Seeger MW (2012) Information-theoretic regret bounds for Gaussian process optimization in the bandit setting. IEEE Trans Inf Theory 58:3250-3265. https://doi.org/10.1109/TIT.2011.2182033

Steefel CI, Appelo CAJ, Arora B et al (2015) Reactive transport codes for subsurface environmental simulation. Comput Geosci 19:445-478. https://doi.org/10.1007/s10596-014-9443-x

Top500 June 2024 | TOP500 most powerful commercially available computer systems. https://top500.org/lists/top500/2024/06/. Accessed 21 June 2024

Turunen J, Lipping T (2023) Feasibility of neural network meta-models for emulation and sensitivity analysis of radionuclide transport models. Sci Rep 13:6985. https://doi.org/10.1038/s41598-023-34089-9

Wagner T, Kulik DA, Hingerl FF, Dmytrievava SV (2012) Gem-Selektor geochemical modeling package: TSolMod library and data interface for multicomponent phase models. Can Miner 50:1173-1195. https://doi.org/10.3749/CANMIN.50.5.1173

Wang W, Taylor J, Bala B (2021) Exploiting the power of Levenberg-Marquardt optimizer with anomaly detection in time series. https://arxiv.org/abs/211106060v1. Accessed 1 July 2024

**Publisher's Note** Springer Nature remains neutral with regard to jurisdictional claims in published maps and institutional affiliations.