A Stochastic Expectation Maximization Algorithm for the Estimation of Wastewater Treatment Plant Ammonium Concentration

In this study, we address the intricate challenge of reconciling environmental sustainability with economic viability within wastewater treatment plants (WWTPs). Our primary objective is to minimize fossil energy consumption and reduce nitrogen concentrations. Current controllers struggle to adapt to fluctuating electricity prices and the variable conditions within WWTPs. While Model Predictive Control and Dynamic Programming offer promising control strategies, their effective deployment hinges on the availability of a robust system dynamics model. To address the stochastic and nonlinear nature of WWTP processes, we introduce a stochastic model and estimation method combining a Monte Carlo Sequential smoothing algorithm with a Stochastic Expectation Maximization method. The proposed methodology results in accurate 24-hour confidence interval predictions, outperforming the conventional estimation method, Prediction Error Minimization (PEM), offering a reliable model for control of WWTPs.


I. INTRODUCTION
Wastewater treatment plants (WWTPs) play a pivotal role in addressing the contemporary challenge of harmonizing environmental sustainability with economic viability.The surge in energy costs has accentuated the need for tighter control of energy consumption within industrial processes.Simultaneously, stringent environmental commitments mandate the swift reduction of pollutant emissions to mitigate the adverse effects of climate change.These imperatives align the industrial sector toward a shared objective: minimizing fossil energy consumption.
In the context of WWTPs, the central control objective centers around nitrogen removal.More precisely, the primary goal is to reduce the concentration of ammonium (N H + 4 ) originating from the incoming water, while simultaneously minimizing energy consumption.This reduction can be accomplished through the utilization of Alternating Activated Sludge (AAS) plants, characterized by alternating periods of aeration (with the introduction of oxygen (O 2 ) into the aeration tank) and non-aerated phases, as illustrated in Fig. 1a.Presently, most of the WWTPs are governed by rudimentary controllers based on specific rules, such as clock-based control or min/max level rules [1].However, these controllers struggle to adapt to fluctuating electricity prices and the variable conditions within WWTPs.Conversely, Model Predictive Control (MPC) [2] and Dynamic Programming (DP) [3] hold significant potential as they leverage dynamic process models to facilitate precise equipment control, adhering to predetermined constraints and cost functions [2].Nevertheless, their effective deployment hinges on the availability of a robust system dynamics model, which constitutes the central focal point of our investigation.
Modeling in this domain presents unique challenges.Our study aims to introduce a model and estimation method for precise predictions of N H + 4 concentrations within an AAS plant over a 24-hour horizon.WWTPs typically have a limited number of sensors, making the accurate measurement of all species impossible (partially observed system).Furthermore, sensor data often contain noise, usually originating from a single sensor within a tank.Additionally, WWTP processes are characterized by significant disturbances, nonlinearity, and complex dynamics (see Fig. 1b).They undergo diurnal and seasonal variations influenced by influent flow fluctuations, influent composition, and process dynamics.
Our contribution entails the development of a stochastic model capable of learning the perturbations affecting the system.The model is estimated using nonlinear filters and smoothers [4] based on Sequential Monte Carlo methods and the "Expectation-Maximization" algorithm [5] designed to handle nonlinearity and stochasticity effectively without reliance on Gaussian assumptions.Notably, this model is straightforward and easily integrable into a model-based controller.We rigorously tested it on simulated data derived from the Activated Sludge Model (ASM) [6], constating better performance compared to the conventional estimation method, Prediction Error Minimization (PEM) [7].Section 2 discusses the current state of the field, including existing methodologies and their limitations.In Section 3, we introduce our model and outline our model parameter estimation methodology.Section 4 covers performance evaluation and results, while Section 5 concludes and explores implications for wastewater treatment plant control.

II. WASTEWATER TREATMENT PLANT AND CONTROL
The optimization of WWTPs performance is a significant subject in the existing literature.This focus emerged with the development of the ASM No. 1 (ASM1) in 1987 by the International Association on Water Pollution Research and Control (IAWPRC) [6].The ASM1 model was specifically designed to describe the denitrification process with activated sludge.Subsequently, various models were introduced to address more intricate phenomena [8].Additionally, the Benchmark Simulation Models [9] were developed using ASM models to facilitate comparisons of various control strategies.The objective of our study is to model N H + 4 concentrations within an AAS plant over a 24-hour period.While the ASM1 model aligns with the goals of our research, it presents certain challenges.Firstly, the benchmark was primarily designed for large WWTPs where multiple species are measured, which differs from our specific focus on AAS plants, prevalent in many municipal WWTPs (see Fig. 1a).Furthermore, ASM1's complexity, with its 13 variables and 21 parameters, renders it unsuitable for real-time WWTP control.Estimating these parameters can be demanding, and integrating the model into a closed-loop control system can lead to computational time issues To streamline this complex model, various simplifications have been proposed in the literature [10], [11], [12], [13].Most of these simplifications are based on linear models [10], [11].However, linear models have limitations when it comes to short-term predictions, typically up to 1 to 2 hours, because of the strong nonlinearity and stochasticity of the processes.In contrast, other models employ physically motivated simplifications by eliminating slow dynamics and employing heuristic approaches [12], [14].These models offer viable solutions and, under specific conditions, may even demonstrate observability [14].Nevertheless, they require precise calibration and can still present integration challenges in an online controller.An intermediary approach between these two categories of models focuses solely on the dynamics of interest, particularly ammonium concentration, as exemplified by Lukasse in 1999 [15] and Stentoft in 2018 [16].These models exclusively center on the target species and do not consider interactions with unmeasurable or unregulated variables, making them a viable alternative for modeling and control [16].
However, the majority of these models do not adequately address the crucial aspect of modeling uncertainty, which becomes significant when simplifications are introduced.This has prompted our exploration of a modeling approach that generates confidence bands.Commonly used methods such as the Prediction Error Minimization (PEM) strategy [7] focus on finding optimal model parameters but do not quantify uncertainty.In contrast, models proposed by Lukasse [15] and Stentoft [16] employ the Extended Kalman Filter (EKF) [17] and Kalman Filter (KF) [18], respectively, to capture model dynamics and quantify uncertainties.However, these models rely on Gaussian and local linearity assumptions to estimate mean and variance.Such assumptions can be problematic when dealing with systems exhibiting strong non-linearities, as illustrated in Fig. 1b.In addressing this challenge, we have chosen a modeling approach similar to Lukasse [15] and Stentoft [16], but with a significant difference: the utilization of a non-linear and non-gaussian model using a Particle Filter (PF) [4].The PF offers distinctive advantages over the EKF and KF, particularly when dealing with non-linearity and complex, uncertain systems [19].Its proficiency in modeling non-linear and non-Gaussian systems is invaluable for predicting N H + 4 concentrations over a 24-hour horizon with non-parametric confidence intervals.

III. METHODOLOGY
In this section, we present the methodology adopted to achieve our primary objective: the quantification of ammonium concentration within the aeration tank of an AAS plant.Our methodology encompasses the development of a simplified ASM1 model (grey-box) trained on data coming from ASM1 model as a proxy of real data (see Fig. 1c.), thoughtfully tailored to our specific research objectives.It is imperative to emphasize that the full ASM1 model presents challenges in achieving precise identification due to the constraints associated with online sensors.These sensors are often technologically unavailable, expensive, or impractical for real-world applications.Consequently, we operate within the framework of a partially observed system, where, in practice, many WWTPs are equipped with oxygen (O 2 ) and/or oxidation-reduction potential (ORP) sensors, and occasionally ammonium (N H + 4 ) sensors.For the purposes of our study, we exclusively consider the availability of an ammonium sensor.Our research methodology involves the utilization of simulated data generated using the ASM1 model to train a reduced model, as outlined in the subsequent subsection.This data-driven approach facilitates comprehensive testing of our modeling and estimation methods under controlled conditions.Looking forward, we intend to employ the 24-hour horizon predictions generated by our model within a model-based controller framework, such as MPC or DP, for the regulation of the aeration process in WWTPs.This control mechanism plays a pivotal role in deciding when aeration is necessary, particularly in WWTPs with AAS plants.

A. ASM1 model
The ASM1 model [6], initially developed for the removal of nitrogen, provides a comprehensive description of how nitrogen is processed within a tank reactor with a constant volume (V ) and a continuous inflow rate (Q in ).This type of reactor is commonly known as a Continuously Stirred Tank Reactor (CSTR).The evolution of various species' concentrations within the reactor can be described by the following equation: where S(t) ∈ R 13 + represents respectively the vector of species concentrations, Q in and S in denote the inflow and inlet water concentrations, and u(t) ∈ {0, 1} represents the aeration control for the reactor.
The ASM1 model serves as a mathematical representation of the function (t, S, u) → r(t, S(t), u(t)), encapsulating the reactions that occur within the reactor tank.A detailed description of this function is provided in the following references [6], [8], [20].This model will be utilized to generate the data required for our experiments, combining a variable inflow, as observed in real WWTPs [9], but with fixed inlet concentrations.Unlike most of the studies [14], [16], we have chosen to separate the estimation of reactor dynamics from the estimation of inlet concentrations, recognizing them as distinct challenges.The former involves modeling dynamics, while the latter entails predicting external exogenous signals influenced by various factors such as weather and human behavior.The estimation of the inlet concentrations is deferred to a separate study.This separation is possible because we are working with simulated data, where we precisely know the inlet concentrations, allowing us to address each aspect with greater precision.

B. Reduced model
In this study, we focus on a reduced model that specifically targets the variables crucial for regulating nitrogen removal, namely, ammonium (N H + 4 ) and nitrate (N O − 3 ).However, in our pursuit of reducing energy consumption, nitrate (N O − 3 ) poses no significant challenges since its concentration increases with aeration (an energy-consuming process).Consequently, elevated nitrate levels are not a concern for energy optimization.As a result, we have chosen to quantify solely the ammonium concentration, designated as S N H (t), which will serve as the sole constraint during control.
To achieve this simplification, we adopt the reduced model used by Stentoft [16] and subsequently remove the nitrate (N O − 3 ) equation.Furthermore, the dynamics of dissolved oxygen (O 2 ) within the system are characterized by rapid changes, often settling at a setpoint during aeration (in AAS plants).Consequently, we simplify the representation of oxygen dynamics by assuming that oxygen levels are at zero when there is no aeration and remain constant when aeration is in progress.These simplifications yield a greatly simplified model, which can be described by the following equation Here, V , β, and K represent the parameters of the model.These simplifications streamline the model, focusing solely on the critical dynamics for control purposes.It allows the model to remain both simple and precise, effectively capturing the principal non-linearity of the system.This 'grey box' model will serve as the framework for quantifying ammonium concentration throughout the rest of the study.

C. State Space Model
With the model's dynamics defined, our focus shifts to its estimation and integration, which are essential for quantifying its parameters and uncertainties.In the context of model fitting, one could directly use sensor values for model learning, minimizing the Root Mean Squared Error (RMSE), akin to the "Prediction Error Method" [7].However, this approach lacks the ability to quantify uncertainties originating from dynamics not captured by the simplified model and uncertainties arising from observations.To address this issue, we adopt a discrete-time state space formulation.This transition involves two key components: an observation equation and a transition equation.The observation equation links the model's state variables to available sensor measurements y t , primarily consisting of ammonium (N H + 4 ) concentrations.Meanwhile, the transition equation outlines how state variables evolve over time, taking into account control inputs and process noise: Here, σ 2 ϵ and σ 2 m represent the variances of the observation and the process, two new parameters introduced into the framework.The motivation behind employing state-space models lies in their utility for achieving precise 24-hour confidence bands of ammonium concentration.This capability is crucial for implementing effective control measures aimed at strictly preventing non-compliant nitrogen discharge in a robust controller.However, state-space models come with a significant challenge: estimating the distribution of the hidden state (S N H (t)) based on observed data ({y 1 , . . ., y T }).This process involves both filtering (estimating the state when observations are available up to time 't') and smoothing (when observations extend beyond time 't').These filtering and smoothing challenges are closely linked to the estimation problem, as discussed in the following section.

D. Parameter Estimation
Now that we have defined the structure of the model, the next crucial step is the estimation of its parameters.Our primary focus lies in maximizing the likelihood of the model concerning a set of measurements {y 1 , . . ., y T }.These measurements help us to determine the unknown parameters, denoted as θ = (V, β, K, σ 2 ϵ , σ 2 m ).The likelihood can be expressed as: where y 1:t−1 = (y 1 , . . ., y t−1 ).Accurately determining the unknown parameters θ is crucial for precise state reconstruction.The maximization of likelihood can be achieved through two methods: direct maximization via gradient ascent or the well-known "Expectation-Maximization" (EM) algorithm [5].The first method involves filtering, while the second combines filtering and smoothing but offers greater stability.In the case of a linear system, the Kalman Filter and Kalman Smoother [18] provide closed-form solutions for filtering and smoothing.However, in nonlinear systems, the filtering and smoothing distributions become intractable.
One approach is to linearize the system and use the Extended Kalman Filter (EKF) [21], as applied in the work by Stentoft et al. [16].However, EKF may not be suitable for highly nonlinear systems as indicated by Stentoft et al. [16].An alternative strategy is the use of sequential Monte Carlo methods to address diverse distribution forms, such as the Ensemble Kalman Filter and Smoother (EnKF and EnKS) [21], [22] or the Particle Filter and Particle Smoother (PF and PS) [4], [23].Chau et al. [19] compared the use of EKF, EnKF and PF for parameter estimation in nonlinear models.Following their conclusion, we choose to employ the Particle Filter combined with a backward smoother (BS) with the Monte-Carlo Expectation Maximization (MCEM) method [24], a stochastic extension of the EM tailored for nonlinear systems.This strategy is ideal for addressing non gaussian systems that will be the case with unknown inlet concentrations (deferred to a separate study).
Without delving into intricate details, Particle Filtering is a sequential Monte Carlo method that approximates the posterior distribution of the dynamic system's state using a set of weighted samples (particles).The PF algorithm iteratively updates the weights based on the current observations and system dynamics.The BS algorithm operates by propagating PF estimates backward in time, accounting for both the system dynamics and observations.A detailed description of the PF and BS can be found in [4], [23], [19], [25].The Monte-Carlo Expectation Maximization (MCEM) [24] algorithm iteratively estimates model parameters.It alternates between an expectation step and a maximization step, with the expectation step using results from the BS algorithm to approximate the expected value of the log-likelihood function with Monte-Carlo under the current parameter estimates, while the maximization step updates parameter estimates to maximize the estimation from the expectation step.Once the model is fitted, we can efficiently generate particles and utilize them to obtain mean and prediction intervals for 24-hour forecasts into the future.This comprehensive methodology aligns with our primary objective of achieving a robust 24-hour prediction of ammonium concentration, considering both parameter uncertainties and the inherently nonlinear model dynamics.

IV. EXPERIMENTAL VALIDATION
To validate our proposed methodology for estimating model parameters, we conducted an empirical analysis, comparing it with two established approaches applied to the same reduced model.

A. Data Generation
We generated a synthetic dataset using the ASM1 model (described in Section III-A).The ASM1 model's parameters and the influent flow rate (Q in (t)) were initialized with standard values from the literature [6], [9].The influent ammonium concentration (S in N H (t)) was assumed constant and known, deferring its estimation to a separate study (as explained in Section III-A).A control strategy resembling ORP control [1] was implemented, activating or deactivating aeration based on predefined thresholds for N H + 4 or N O − 3 concentrations.The system was simulated for 20 days to reach a steady state, followed by a 5-day period dedicated to model parameter estimation.Observations were collected at 5-minute intervals and corrupted with Gaussian white noise with a standard deviation of σ ϵ = 0.2.

B. Experimental Setup
We employed a Stochastic Expectation Maximization (SEM) algorithm in conjunction with Particle Filtering (PF) and Bootstrap Sampling to estimate the model parameters (SEM-PF).The SEM algorithm's M-Step was optimized using the BFGS algorithm [26].Utilizing 300 particles, the algorithm converged after approximately 30 iterations.The analysis was conducted in the Julia programming language.The code for this experiment is available online 1 .We compared the performance of our methodology with two other approaches using the same reduced model: • Standard Prediction Error Minimization (PEM) This approach minimizes one-step ahead prediction errors of the model using available noisy observations [7].• Extended Kalman Filter (EKF) and Extended Kalman Smoother (EKS) with Expectation Maximization (EM -EKF): This approach estimates the parameters as described in Dreano et al. [27].To assess the robustness and performance of our method under various conditions, we conducted a comprehensive comparative analysis.We systematically varied the influent concentration and used diverse sets of observational data.Specifically, we tested four distinct influent concentration values and repeated the estimation process ten times for each condition.In each case, we estimated the "grey-box" model Fig. 2: Filtering concentration in the past for EKF and PF and 24-hour prediction in the future for all the approaches.using the 5-day training phase with all three approaches.Subsequently, each approach was evaluated over the next 24 hours following the training periods to simulate control conditions.Visually, the results suggest that all methodologies show promise in learning the model parameters.State-space approaches (SEM-PF or EM-EKF) appear to outperform the traditional PEM method.Notably, the predicted model distribution in the past closely matches the observed data for the state-space approaches, and their 24-hour predictions yield informative confidence intervals.However, Fig. 2 also highlights a limitation of the proposed model: when the model is unaerated for an extended period and the ammonium concentration is high, it tends to underestimate the concentration.This limitation is inherent to using a simpler model.However, the prediction interval takes this into account and widens at higher ammonium concentrations, demonstrating the value of using a stochastic model.Fig. 3 illustrates the evolution of the RMSE according to the prediction horizon for all three approaches.The results confirm that methodologies based on state-space formulations (SEM-PF or EM-EKF) outperform the classical PEM criterion.State-space approaches exhibit a notably lower initial RMSE, averaging around 0.14 compared to 0.23 for PEM.This suggests that they consistently produce more accurate short-term predictions.As the prediction horizon extends, the RMSE of all approaches stabilizes, indicating that our model's predictive performance remains reliable over longer forecasting periods, such as a 48-hour prediction.Importantly, the long-term prediction RMSE is lower with SEM-PF or EM-EKF (0.21) than with PEM (0.27).It is worth noting the dip in performance around 5 hours in Fig. 3.This can be explained by variations in system performance depending on the applied control strategy (as explained in the second paragraph of Section IV-C).

C. Performance Evaluation
Furthermore, Fig. 3 displays the confidence interval of the RMSE for different models.The prediction interval for the state-space formulations hardly overlaps with PEM for 24hour predictions, illustrating their consistent superiority with a high probability.However, we observe minimal difference between the proposed methodology (SEM-PF) and EM-EKF in terms of RMSE (mean values are identical and close confidence intervals).This can be attributed to the small time interval between observations, resulting in a discretized model which is almost linear [27].Additionally, both SEM-PF and EM-EKF allow for the definition of prediction intervals, with a mean coverage probability averaging around 0.86 for both, aligning with a 95% confidence level.The mean width of the prediction intervals is 0.75 and 0.77 for SEM-PF and EM-EKF, respectively.This slight difference could be due to the ability of PF to handle non-Gaussian distributions, potentially leading to narrower confidence intervals in some cases.To further solidify our findings, Table I provides an overview of the mean estimated parameter values and their corresponding standard deviations.The estimates in Table I reveal accurate estimation of observation noise when it is estimated.While some model parameters lack direct physical interpretations, values like V and K appear to be within expected ranges.The proposed methodology and EM-EKF approach show very similar results.Furthermore, the most significant difference between state-space approaches (SEM-PF and EM-EKF) and PEM is observed in the parameters K and β.We hypothesize that the PEM criterion struggles to differentiate between model imperfections and observation noise, potentially leading to bias in these parameters, contributing to its lower performance.
In conclusion, the methodologies employing SEM-PF or EM-EKF demonstrate strong performance and reliability in learning model parameters.They outperform the classical identification PEM method, achieving this despite all approaches utilizing the same underlying models, as evidenced by the RMSE comparisons and the reliability of prediction intervals.In this study, there is no significant difference between using PF or EKF to construct the filtering distribution.However, PF offers the advantage of working with non-Gaussian distributions with no degradation in performance when the distribution is close to a Gaussian, as in this study.In future work, incorporating the estimation of the inlet concentration, which is assumed constant here, could lead to a far less Gaussian distribution, highlighting the benefit of PF's adaptability.

V. CONCLUSION
In conclusion, our application of the SEM algorithm in conjunction with PF and BS has yielded a precise model for predicting ammonium concentration over a 24-hour horizon.The proposed methodology surpasses traditional identification method, such as PEM, and highlights its potential as a valuable tool for optimizing wastewater treatment plant control.While our study found no significant disparity between the use of PF and EKF in this context, the flexibility offered by PF for further investigation without compromising results makes it an attractive choice.Before implementing this approach in practical settings, we must address the challenge of accurately estimating the inlet ammonium concentration.A promising avenue for future research lies in the development of an ammonium prediction model that integrates data from oxygen and ORP sensors, potentially eliminating the need for a dedicated ammonium sensor.Furthermore, extending our investigation to compare various model types, including physical and data-based models, using the methodology proposed in this paper, would provide valuable insights into their effectiveness.These efforts aim to improve the effectiveness of our approach, leading to more efficient wastewater treatment processes.

Fig. 1 :
Fig. 1: Schematic representation of the research methodology.a) Secondary treatment process in an AAS plant.b) Concentrations S(t) generated by ASM1 physical model as a proxy of real data observations.c) Grey box model for 24-hour S N H (t) prediction using ASM1 data.

Fig. 2
Fig. 2 presents an example of the predicted model distributions and 24-hour predictions obtained from the different approaches.The left portion of the figure illustrates the predicted distribution (p(S N H (t)|y 1 , . . ., y t−1 )) during the filtering process, while the right part shows a 24-hour prediction of the model.Visually, the results suggest that all methodologies show promise in learning the model parameters.State-space approaches (SEM-PF or EM-EKF) appear to outperform the traditional PEM method.Notably, the predicted model distribution in the past closely matches the observed data for the state-space approaches, and their 24-hour predictions yield informative confidence intervals.However, Fig.2also highlights a limitation of the proposed model: when the model is unaerated for an extended period and the ammonium concentration is high, it tends to underestimate the concentration.This limitation is inherent to using a simpler model.However, the prediction interval takes this into account and widens at higher ammonium concentrations, demonstrating the value of using a stochastic model.

TABLE I :
Parameter Estimation over different intervals and different inlet concentrations (Mean ± Standard deviation).