Exploring scientific machine learning pipelines through the SimulAI toolkit
Overview
SciML, short for Scientific Machine Learning, encompasses work that merges quantitative sciences with machine learning. It has gained significant traction over the past decade, driven by the widespread availability of specialized hardware (such as GPUs and TPUs) and datasets. Additionally, it has been propelled by the overarching influence of the machine learning wave, now ingrained in the zeitgeist of our times. In this context, we’d like to introduce SimulAI, an open-source toolkit under the Apache 2.0 license. SimulAI is designed to be user-friendly, providing a high-level Python interface for managing scientific machine learning pipelines. This article aims to showcase its current workflow and utility in constructing scientific experiments. We encourage feedback and potential contributions from the interested community, with plans to delve into more advanced topics in future articles.
SciML
The past decade has seen the emergence of a groundbreaking field known as Physics-Informed Neural Networks (PINNs). These have reignited the interest of the research community in leveraging neural networks to solve ordinary and partial differential equations, commonly encountered in physics. Following PINNs, a new wave of innovation brought forth the concept of neural operators . This includes advancements like Deep Operator Networks (DeepONets) and Fourier Neural Operators (FNO). These developments have enabled the modeling of intricate mathematical operations, expanding the versatility of neural networks to address multi-scenario problems. Subsequent to these milestones, numerous enhancements have been made across various auxiliary algorithms. These include the introduction of new loss functions, penalization methods, and even domain decomposition techniques.
The realm of scientific applications for deep learning continues to witness a surge in both the development of new applications and the creation of deep learning architectures. These are tailored to address scientific problems by incorporating both data and mathematical expressions into the training loop. The spectrum of application areas has expanded significantly, ranging from turbulence modeling and earthquake hypocenter localization to chemical kinetics. This new wave of machine learning has not only fostered innovation in addressing specific challenges but has also spurred the creation of dedicated libraries and packages for constructing SciML workloads. Notable mentions include Modulus by NVIDIA, DeepXDE from Brown University, and more recently, the SimulAI toolkit, sponsored by IBM.
SimulAI aims to be a versatile and user-friendly Python package, designed to empower both research and industry communities by integrating machine learning models into their existing work. Leveraging the robust foundations of PyTorch, NumPy, and SciPy as its mathematical engines, SimulAI provides access to a diverse array of state-of-the-art models and architectures. These include DeepONets, Variational AutoEncoders (VAEs), Operator Inference (OpInf), and now extends its offerings to include transformers and U-Net.
In this context, SimulAI’s current workflow will be showcased by solving a canonical problem within the SciML domain — creating a Physics-Informed Neural Network (PINN) to approximate the solution of a Partial Differential Equation (PDE). This demonstration underscores SimulAI’s capability to handle complex scientific challenges through its accessible and comprehensive set of tools and models.
Physics-Informed Neural Networks in a Nutshell
PINNs can be categorized as supervised or unsupervised learning architectures. They incorporate physics principles, expressed as governing equations, into their loss functions as penalization terms. The inclusion of these equations in the loss functions aims to enforce conservation laws since PDEs are the foundation of mathematical models that represents the physical universe. In order to illustrate it, let us consider a specific case of the canonical Allen-Cahn PDE, which models phase-separation in multi-component systems. Let us consider the following PDE equation defined over a 1D domain of size L and a time interval [0, T]:
The PDE is constrained by the following boundary and initial conditions:
Considering that we can evaluate the expressions above for any sample of spatiotemporal positions within the domain, we can easily assess root mean-squared metrics (RMSE) for them. In this way, we can rewrite the mathematical expressions in the form of a loss function, as seen as following:
The boundary conditions are written as subtrations between the domain extremes since they are periodic for this problem.
The lambda-penalties in the aforementioned expression play a crucial role in balancing various terms, considering potential differences in their orders of magnitude. This strategic use aims to exert control over and enhance the convergence speed during training. Once the loss function is defined, the next step involves minimizing it through an optimization algorithm. Subsequently, the model coefficients are updated accordingly. Despite the apparent simplicity of PINNs, there are intricacies in their convergence dynamics, presenting an ongoing research challenge. However, for the specific case we are addressing, a solution can be achieved without delving into more sophisticated algorithms.
Defining PINN expressions
The Allen-Cahn PDE, along with its corresponding initial and boundary conditions, can be succinctly expressed in Python using the SimulAI syntax.
f = 'D(u, t) - mu*D(D(u, x), x) + alpha*(u**3) + beta*u'
g_u = 'u'
g_ux = 'D(u, x)'
input_labels = ['x', 't']
output_labels = ['u']Utilizing strings to represent the equations for minimization provides a clear and concise approach to express physics. However, for cases requiring increased flexibility in input representation, SimulAI supports alternative methods such as Python functions or SymPy objects within its symbolic framework.
Instantiating a neural network
In our experiment, we aim to train a neural network parameterized by a set of coefficients (weights and biases). This network receives coordinates (representing time and one space dimension) and predicts the corresponding value for the Allen-Cahn state variable . SimulAI offers a versatile range of pre-implemented neural network models designed to function as Physics-Informed Neural Networks (PINNs). Additionally, users can explore combinations of these models. A prime example is the special class of Multilayer Perceptron (MLP) networks, exemplified in the following template function:
def model():
from simulai.models import ImprovedDenseNetwork from simulai.regression
import SLFNN, ConvexDenseNetwork
n_inputs = len(input_labels) n_outputs = len(output_labels)
# Configuration for the fully-connected network
config = { "layers_units": 6*[128],
"activations": "tanh",
"input_size": n_inputs,
"output_size": n_outputs,
"name": "net",
}
# Instantiating the subnetworks
densenet = ConvexDenseNetwork(**config)
encoder_u = SLFNN(input_size=n_inputs, output_size=128, activation="tanh")
encoder_v = SLFNN(input_size=n_inputs, output_size=128, activation="tanh")
net = ImprovedDenseNetwork(network=densenet,
encoder_u=encoder_u,
encoder_v=encoder_v,
devices="gpu")
return net
net = model()
For a general overview of currently available models on the repository, take a look at documentation main page.
Datasets
The datasets used to train the PINN, as shown in the code blocks below (data, data_boudary_xL, data_boundary_x0, data_boundary_t0), consist of coordinate samples taken across the rectangular domain and can even be generated randomly. Users have the flexibility to define the construction of these input datasets. Notably, for unsupervised training, no target datasets are required, as the vanilla PINN relies solely on the governing equations and their respective boundary and initial conditions. Nevertheless, users have the option to include additional complementary datasets (measured or sensed) to create a hybrid loss function in a semi-supervised manner.
Converting symbolic expressions into PyTorch tensor objects
In SimulAI, mathematical expressions undergo processing and conversion into tensor-like objects using the symbolic
engine, resulting in the creation of a ‘PDE residual’ object, as demonstrated in the code snippet below:
residual = SymbolicOperator(expressions=[f],
input_vars=input_labels,
auxiliary_expressions={'periodic_u': g_u,
'periodic_du': g_ux},
constants={'mu':1e-4,
'alpha':5,
'beta':-5},
output_vars=output_labels,
function=net,
engine='torch',
device='gpu')It’s important to highlight that boundary conditions are added as auxiliary expressions, and the constants mu, alpha, and beta are provided within a dictionary. Subsequently, the residual is forwarded to the optimizer instance.
The optimizer
The optimization framework serves as an API for both the PyTorch engine and internal optimization algorithms (currently under construction). It offers a range of configurations, including learning rate schedulers, early-stopping mechanisms, Tensorboard registers, and more.
lr = 1e-3 # Initial learning rate
optimizer_config = {'lr': lr}
optimizer = Optimizer('adam', params=optimizer_config,
lr_decay_scheduler_params={'name': 'ExponentialLR',
'gamma': 0.9,
'decay_frequency': 5_000},
shuffle=False,
summary_writer=True)In addition to the basic optimizer configuration, it is necessary to provide a dictionary with extra parameters:
params = {'residual': residual,
'initial_input': data_boundary_t0,
'initial_state': u_init,
'boundary_input': {'periodic_u': [data_boundary_xL,data_boundary_x0],
'periodic_du': [data_boundary_xL, data_boundary_x0]},
'boundary_penalties': [1, 1],
'initial_penalty': 100}It’s important to note that you can define fixed lambda-penalties for the loss function terms, or you can dynamically adjust them using loss
adjusters. For detailed configuration examples, refer to this
example. Boundary conditions require special handling, especially for periodic conditions in the context of Allen-Cahn. Finally, to initiate training, use the fit method of the optimizer class.
optimizer.fit(op=net,
input_data=data,
n_epochs=100_000,
loss="pirmse", params=params,
device='gpu')For a detailed walkthrough of the process, briefly touched on here, please refer to notebook.
Evaluations using the model
Now that our neural network is trained, we can utilize it for making predictions. To do so, we first define an evaluation spatiotemporal grid with any resolution:
X_DIM_F = 500
T_DIM_F = 500
x_f = np.linspace(*x_interval, X_DIM_F)
t_f = np.linspace(*t_interval, T_DIM_F)
T_f, X_f = np.meshgrid(t_f, x_f, indexing='ij')
data_f = np.hstack([X_f.flatten()[:, None],
T_f.flatten()[:, None]])Next, we use the eval method, common to all NetworkTemplate objects.
approximated_data = net.eval(input_data=torch.from_numpy(data_f.astype("float32")))
U_f = approximated_data.reshape(T_DIM_F, X_DIM_F)Figure 2 illustrates a sample result from the neural network training.
Saving and restoring the model
Once the neural network is trained, it can be saved using the built-in save/restore mechanism:
saver = SPFile(compact=False)
saver.write(save_dir='.', name='allen_cahn_net', model=net, template=model)This operation ensures the model’s storage in a dedicated directory, comprising a file for the model coefficients (with the extension pth), a Python module housing a copy of the model template (the function model mentioned earlier), and potentially a pickle file for storing configuration arguments. The model is easily restored from disk as follows:
loader = SPFile(compact=False)
model = loader.read(model_path=<path to the model>, device=<destiny device>)Conclusion
We trust that this concise overview has effectively covered the fundamental aspects of Physics-Informed Neural Networks (PINNs) and the toolkit’s functionalities. While leveraging PINNs to address Partial Differential Equation (PDE) systems stands as a pivotal undertaking to gauge the potency of neural networks in handling Physics datasets, it is crucial to emphasize that the realm of Scientific Machine Learning (SciML) extends well beyond PINNs. It encompasses more traditional data-driven approaches, wherein extensive datasets, either simulated or collected from sensor devices, are utilized to train algorithms tailored for diverse prediction tasks.
For instance, algorithms like neural operators can undergo training on simulated data, exhibiting prowess in tasks such as time forecasting, dimensionality reduction, and even seamlessly integrating with Physics-informed approaches in a hybrid fashion. SimulAI already accommodates many of these pipelines, and we extend an invitation to the interested community to contribute, evaluate, and offer feedback on their usage. To those eager to contribute directly through code and test cases, rest assured that your contributions are warmly welcomed and highly valued.
Acknowledgments
A heartfelt appreciation to IBM Research for their support and the robust infrastructure that has played a pivotal role in the success of this project. Our gratitude extends to the entire IBM Research community for fostering an environment of innovation and collaboration.
