Published based on research by IoT++ AI COE ( Center of Excellence)
Pump failure is a critical issue in drilling operations that can lead to significant downtime and increased costs. It refers to a situation where the mud pumps, which are essential for circulating drilling fluid, malfunction or cease to operate effectively.
Pump failures can halt all drilling work until the issue is resolved, making them one of the major risks in industries like drilling industries or other industries that could cause substantial non-productive time (NPT).
Understanding and predicting pump failure events is crucial for industries to improve operational efficiency, minimize downtime, and reduce costs.
Traditional methods for predicting these events often rely on historical data analysis and expert judgment, which can be subjective and may not capture the complex interactions between various drilling parameters and pump performance indicators.
Many techniques have been used to address this problem. Statistical methods and traditional machine learning approaches are some of the well-known methods employed so far.
However, these methods have often proved insufficient to fully solve the problem of pump failure prediction, as they may not adequately capture the intricate relationships between multiple variables and the dynamic nature of drilling operations.
Pump failure can be predicted by analyzing deviations in sensor data from normal behavior. The LSTM Autoencoder, a self-supervised model, processes time-series data to detect these anomalies, using reconstruction error to predict failures.
Key steps include data pre-processing, windowing, model training, and inferencing. The model is converted to ONNX format for real-time deployment across platforms, improving efficiency and operational safety.
Classical Analogy
Imagine a musician is very familiar with classical music and has spent years listening to and playing pieces by composers like Bach, Mozart, and Beethoven.
They are excellent at remembering and recreating melodies within this genre because they understand the typical patterns, scales, and structures used in classical music.
Now, suppose this musician is suddenly asked to listen to and recreate a complex jazz improvisation by a different musician—a style they are not familiar with.
Jazz improvisation is often characterized by complex rhythms, unconventional scales, and spontaneous changes that are very different from the structured patterns of classical music.
- Jazz Melody (Unfamiliar Sequential Data): The jazz piece represents a sequence of data that is very different from what the musician (or LSTM autoencoder) is trained on. The rhythms, scales, and overall structure are foreign to them.
- Listening and Remembering the Jazz Melody (Encoding with LSTM): As the musician listens to the jazz piece, they try to apply the same encoding techniques they use for classical music.
However, because the patterns and structures in jazz are unfamiliar, they struggle to create a meaningful summary. The encoding process fails to capture the essence of the jazz melody accurately because the musician’s mental model is based on classical music, not jazz. - Recreating the Jazz Melody (Decoding with LSTM): When the musician tries to recreate the jazz piece, the result is poor. They might inadvertently apply classical music patterns to the jazz piece, leading to a melody that sounds more classical than jazz.
The reconstructed sequence is inaccurate because the encoded information didn’t properly capture the unique characteristics of the jazz melody.
Breaking it Down:
- Jazz Melody: Represents data that is very different from the data the LSTM autoencoder (or the musician) is trained on. It has patterns and structures that the model is unfamiliar with.
- Inadequate Encoding: The LSTM autoencoder struggles to create a meaningful compressed representation of the jazz melody because it doesn’t recognize the patterns and relationships in the data.
- Poor Reconstruction: The decoded output doesn’t accurately reflect the original jazz melody. Instead, it might resemble something more familiar to the model, such as a classical piece, rather than capturing the essence of jazz.
The Key Idea:
An LSTM autoencoder, like the musician, is trained to recognize and reproduce patterns it has learned from a specific type of data. If it encounters data that is very different from what it knows—like a classical musician trying to recreate a jazz improvisation—the encoding process won’t accurately capture the data’s essential characteristics.
As a result, the reconstructed output will be flawed and may not resemble the original data.
This limitation illustrates how LSTM autoencoders (and machine learning models in general) perform best when the input data is similar to what they were trained on. If the data is too different, the model’s performance can degrade significantly, leading to poor or incorrect reconstructions.
Our Objective

In this article we will explain how a more advanced deep learning based “Autoencoder LSTM” model can be more useful for the pump failure prediction ahead of time.
Autoencoder LSTM is a self supervised deep learning model which is used for many purposes such as video, text and time series.
Our data from a pump failure is also similar to a time series because it involves many parameters whose values depend on the time. They show normal behavior during normal activity and deviate from normal behavior during pump failure situations.
This deviation during pump failure situations can be used to predict pump failure situations. Our autoencoder LSTM model can be a best fit for this kind of problem.
An LSTM Autoencoder is a type of neural network designed for sequence data. It combines an autoencoder, which learns a compressed representation of input data, with LSTM layers that handle sequences.
The model uses an Encoder-Decoder LSTM architecture to read, encode, decode, and recreate sequences. The performance is evaluated based on how well it can reconstruct the input sequences.
Let’s discuss the architecture of the model in detail which we have used for this project.
Model architecture:
1. First the input sequence as a window size of 150 i.e. 150 continuous rows is passed or feeded as one window at once. This input sequence goes inside the encoder which has two LSTM layers in it. The encoder compresses the sequence after making its summary with the help of LSTM layers.
2. The bottleneck is the part where the compressed or summarized part comes from the encoder and it sends it to the decoder.
3. After this the bottleneck sends this summarized part to the decoder which again with the help of two LSTM layers reconstructs it again similar to the original sequence and gives it as the output.
4. Now we have the original sequence and reconstructed sequence. As every machine learning model has some error during prediction, in this case also the reconstructed part is not exactly similar to the original sequence.
This difference of error is known as reconstruction error. This reconstruction error is the most significant part as it is the base on which we will decide that the pump will fail or not in the near future. We calculate this reconstruction error for this window using MSE(mean squared error) and store it. Then we do the same process for the next window consisting of 150 rows.
Note – The training is done only on the normal dataset i.e. where the pump is showing normal behavior. So, the model only sees the normal dataset and it is not aware of abnormal data. Therefore, when the abnormal data comes, the model tries to reconstruct it like normal data which shows deviation between the actual data given and predicted data. This deviation is used to predict the pump failure for our case. This deviation is calculated by using the reconstruction error.

Data Science Life Cycle
Setting up the Jupyter Environment in vs code:
First you need to install both jupyter notebook(anaconda installation) and vs code. Then you can follow the following links –
1. To setup the virtual environment – Creating virtual environments in VS Code and Jupyter notebook | by Abhishek Jain | Medium
2. To set up a Jupyter notebook in vs code – How to Install Jupyter Notebook in VSCode | Jupyter Notebook in Visual Studio Code (Easy)
Now let’s start step by step going through the code:
1. Importing Libraries
We used sensory data from pumps for this project in which there is data from 52 different sensors and there is a target column which tells us about the situation of the pump. The different sensors are placed at different places in the pump system which keeps an eye on the health of the pump.
They are installed on the pump to monitor key parameters such as vibration, temperature, pressure, flow rate, and motor current. These sensors continuously collect data, which reflects the pump’s operational status.
(The link for the data – pump_sensor_data)
We started with pre-processing of the data and then trained it using autoencoder LSTM architecture.
Let us now discuss in detail about the pre-processing and training process. We start by importing all the libraries and modules required for this project.

After this we load the csv file data using pandas.

2. Pre-processing
In any machine learning project , pre-processing is a very important step because in this step , we prepare and increase the quality of data by doing mathematical operations on the raw data.
During these operations, we handle nan values and missing values by imputation or by removing them , and deal with outliers by removal or by transformation and scale of the data by normalization or standardization.
These operations make the quality of data better and eventually can increase the overall accuracy of the model.
Let’s discuss them one by one,
If we have unwanted columns ,we can drop them. In our case there is an unwanted column “unnamed: 0”, we drop it and also make the time_ref as the index so that we can easily visualize our results in the end.

Next, we will see the nan and missing values in the dataset. We calculate the total number of nan values in each column

In our data, the sensor 15 has all nan val with 0 numerical value. So we can drop it as it doesn’t affect our results since there is no value to influence.

Next, we introduce a new term called forward filling. Forward filling is a method used in data pre-processing to handle missing data in a dataset.
It is particularly useful in time series data or sequences where the values are expected to change gradually over time. It involves taking the last known (non-missing) value and carrying it forward to fill in subsequent missing values.
We forward fill the nan values with limit 20. This means during forward only 20 continuous nan values can be replaced by previous non nan value.

In our target (machine_status) column, we have 3 categorical values – normal, broken and recovering. We modify these categories and do operations on them such that normal is considered as 0 which means the pump is working properly. Recovering and broken are both considered as 1 which means the pump is not working properly.

Actually, we don’t use target column during the training of the autoencoder LSTM model as it is an unsupervised kind of model. But we are still doing some operations on the target column because we have to select the most important features for our training purpose.
This is achieved by using correlation of all columns against the target column (machine_status) and then only considering values which have correlation above 0.7 .


Now, after this step we are left with 7 columns only. So let’s calculate the mean and standard deviation for the remaining columns.

We can save this dictionary (mean_std_train) in a JSON file to use it for testing if required or to use it in the future. The reason for saving it is to use these stats for scaling the data.
The scaling is done for the following reasons:
Equalizes Feature Ranges: Ensures all features contribute equally by bringing them to a similar scale.
Improves Model Performance: Prevents features with larger values from dominating, leading to better accuracy.
Faster Convergence: Helps gradient-based algorithms converge more quickly and reliably.
Avoids Numerical Instabilities: Keeps values manageable to prevent overflow/underflow errors.

Note – We use the stats from the json to scale the data instead of a general Sklearn scaler. The reason for doing so is that we have lots of nan values still in the data which may create a problem if we use a scaler of Sklearn. So we scale manually by calculating numpy nan mean and numpy nan standard deviation. The remaining nan values are handled during the window making where we ignore a whole window if it contains nan values. We only feed valid windows which don’t have any nan values.
Now, we don’t require the target column as well as during training we are not going to feed it into the model. So, let’s drop it.

3. Splitting the data
The pre-processed data is split into train, val and test data in the ratio of 0.6:0.2:0.2 respectively. We can use the Sklearn library for this purpose.

4. Sliding window
We can make sliding windows for the training and validation sets which are feeded into the model during training.
In our case, we have taken 150 as the window_size which means only 150 values are feeded into the model and it will predict a reconstruction error corresponding to it.
We have taken strides of size 25 during sliding windows. We also have established a condition that if any of the windows have nan values in it, we neglect those windows and in the end we will have only valid windows for scaling and feeding.
To make the data consistent at the end and to visualize it, we store the last indices of the both valid and ignored windows. Scaling is done for each window separately and applied to only valid windows.


Let’s make the sliding window and data loader for our dataset to feed it into the model.

5. Modeling and Training
For this project, we use an Autoencoder LSTM model which takes the input parameters and analyzes them and recreates it through its trained parameters. The loss or accuracy is given in the form of reconstruction error which is calculated by the difference between actual and predicted values. We have used huber loss for this purpose.
The parameters used in the architecture during training process –
Number of layers of LSTM = 2
Hidden dimension = 128
Latent dimension = 64
Number of epochs =30
Loss function = huber loss
Optimizer = adam
We have also used the decaying learning rate for training purpose so that overfitting can be reduced if there is any.


In our architecture we have used 2 layers of Lstm in each encoder and decoder. For the output purpose we have also used ‘relu’ function instead of ‘tanh’ which can be useful if there is a gradient vanishing problem during training.
In short the model summary is:

Now we can proceed with the training of the model with the training set and validate it using the validation dataset.


We can also see the loss curve for this training:


6. Reconstruction Error Calculation
Our model has been trained now.
Now, we define a function to calculate the reconstruction error which is the most important part as our model’s prediction is based on this reconstruction error value. So let’s define it in a function called calculate_reconstruction_error:

7. Inferencing
During training, a machine learning model learns patterns and relationships from a dataset by adjusting its parameters. Once training is complete, the model is evaluated and fine-tuned.
The inferencing stage utilizes this trained model to analyze new input data and generate predictions. The data fed into a model during inferencing is new and was not used in the training process.
Now we test the model by using the test set which we defined during the splitting of the data. To start with, we first make sliding windows and then call the data loader function to these valid test windows.

Now the test loader is ready to feed into the model to calculate the reconstruction error.
Let’s calculate the reconstruction error for the test loader.

Note – During the calculation of reconstruction error for this test set, only valid windows will be fed which are filtered during the sliding window generation. Although, we keep tracking the indexes of the valid and ignored windows so that we can use them in the end to map the reconstruction error back to the original data points
Now , we define a function which will map the reconstruction error to the original data points in the test set. This is achieved by mapping the reconstruction error of a window to the last index of that window.
For example – in our case , we have taken 150 as window size so for the first valid window the reconstruction error is mapped to the data point which is at 149 index(considering 0 indexing).
For the second valid window taking stride as 1, the starting and the ending indices are 1 and 150 respectively. So, for the second valid window, the reconstruction error is mapped to 150 index point. This pattern is followed for all the valid windows.
Now, if the window is ignored, then we fill 0 to the last index of that window. For the first 149 values , there will be no values as there are only in the window first. So we also fill 0 in the reconstruction error at their place also.


Now we have mapped a reconstruction error value to every point in the test dataset. Let’s define a function to plot the curve – reconstruction error vs time.


In the above plot we can see that we have marked the data points as the anomalies above the threshold line which means there is a possibility of pump failure.
This threshold line is based on the value of the threshold. There is no direct way to predict the threshold value. We have to calculate it manually by hit and trial methods.
One way to calculate it is through the training or validation error itself. We can calculate the reconstruction error for training and validation data sets and can check how we can adjust them for our testing dataset.
8. Saving the model and trained weights
Till now we have tested the model only on the previous data set(our test set is from previous data not current or real time data). So, to predict in real time we need to save our model and its weights so that we can use them during inference of the real time data.

9. Conversion of Pytorch model into ONNX model
Pytorch models have limitations when it comes to using trained models for deployment in another platform or environment because they are specific to the environment in which they are trained.
Advantages of ONNX model:
1. Cross-Platform Compatibility
Interoperability: ONNX is an open standard format that allows models to be used across different frameworks and platforms. By converting a PyTorch model to ONNX, it can be run in a variety of environments that support ONNX, such as TensorFlow, Caffe2, Microsoft’s ONNX Runtime, and many more. This makes it easier to deploy models in production environments that might not directly support PyTorch.
2. Optimized Performance
Inference Optimization: ONNX Runtime, and other ONNX-compatible engines, provide optimizations specifically for inferencing. These optimizations can significantly reduce the latency and increase the throughput of predictions, making ONNX models well-suited for real-time applications.
Hardware Acceleration: ONNX models can leverage hardware-specific optimizations, such as those provided by NVIDIA TensorRT, Intel OpenVINO, and other libraries that optimize models for specific hardware accelerators (like GPUs, TPUs, and FPGAs). This can drastically improve inference speed and efficiency on the target hardware.
3. Deployment Flexibility
Cloud and Edge Deployment: ONNX models are versatile and can be easily deployed in various environments, including cloud services, on-premises servers, or edge devices. This flexibility is crucial for real-time applications, where deployment constraints and requirements can vary widely.
Compatibility with IoT Devices: For edge computing and IoT applications, where resources are often limited, using a lightweight and optimized model format like ONNX allows for efficient real-time inferencing on less powerful hardware.
So, let’s convert the pytorch model into onnx model:

Note- For conversion of pytorch model to ONNX model , we have to give a sample dataset similar to the original dataset in dimension and size. We have to give the data with the same batch size, window size and input size (number of parameters).
10. Real time prediction
We can deploy our ONNX model on the platform on which we want to give real time predictions. We can take raw data, then pre-process it and then predict reconstruction error by feeding it to the ONNX model.
Then we can analyze this reconstruction error by using a predetermined or dynamic threshold value to generate alarms for abnormal situations. We can use tools like grafana and kafka for this purpose.
Apache Kafka handles real-time data ingestion and pre-processing. It streams raw data to the ONNX model, which then generates reconstruction errors for anomaly detection.
Grafana visualizes these prediction results through dashboards and charts, providing real-time insights and alerting for abnormal situations based on reconstruction errors and predefined thresholds.
Conclusion
In this article, we have discussed how to predict pump failure using the LSTM autoencoder model. The pump failure is a very serious problem which can lead to loss of money and time.
Nowadays we have sensors in the pumps and on the outputs of these sensors we can predict the pump failure in advance. We have taken data from sensors and then pre-processed it and then split it into the training ,validation and testing datasets.
We have used training and validation datasets for training of the model and testing dataset for inference the model. Then we decide a temporary threshold value (0.2 in our case) to plot the reconstruction error for the test dataset.
To make the model more accurate, we can fine tune the threshold value using hit and trial methods or any statistical method. After this for real time prediction we save the model and then convert it into the ONNX model so that we can deploy it under different environments and on different platforms without worrying about libraries and packages.

