Crystal Structure Prediction Using Optuna in Matlantis CSP

Introduction

Crystal structure exploration is a crucial first step in materials development and property prediction. However, the exploration space is vast, evaluation is computationally expensive, and the actual stable structures and compositions vary greatly depending on the elemental system. As a result, the success and efficiency of the exploration tend to depend heavily on "experience and trial and error," and those working on the same problem often feel that "there must be a smarter way to do it."

We recently released Matlantis CSP (MTCSP), a crystal structure discovery service, on Matlantis™, a cloud-based versatile atomistic simulation platform for materials discovery. MTCSP utilizes Optuna™, a black-box optimization framework led by PFN, as an approach to efficiently and effectively advance crystal structure discovery with fewer trials. This article will introduce how Optuna is incorporated into MTCSP.

What is Optuna, the black-box optimization framework?

Optuna is a black-box optimization framework designed to find good solutions with a small number of trials for problems where "evaluation is possible, but it's difficult to know how to optimize them." A typical example is hyperparameter search in machine learning, but essentially, it's not limited to hyperparameter search as long as there is a mechanism that "gives input parameters and returns an output score."

In the context of crystal structure exploration, the goal is to find energetically stable structures given constraints such as a given elemental system and compositional ratio range. The input is the crystal structure, and the output is the formation energy of the structure. Traditionally, time-consuming DFT calculations were used, but Matlantis allows for the rapid evaluation of formation energies using the general-purpose neural network potential PFP.

What is the crystal structure search service MTCSP?

Matlantis CSP (MTCSP) is a crystal structure search service provided on Matlantis, a cloud-based versatile atomistic simulation platform. Based on the composition and conditions provided by the user, it generates candidate crystal structures (atomic arrangements) and searches for more stable (or promising) structures while evaluating them using indicators such as energy.

The difficulty in crystal structure exploration lies not only in the vast search space, but also in the fact that evaluating the structure takes a certain amount of time, and the effectiveness of each search method varies greatly depending on the element system and conditions. In other words, the search is not simply about increasing the number of candidates, but also within a limited computational budget.

• How to create candidates
• How to evaluate the structure
• How to construct a search loop

This requires some ingenuity. MTCSP is designed to make computationally intensive searches more practical by treating the execution of searches and the analysis of results as a service, allowing users to focus on specifying conditions according to their purpose and performing analyses tailored to their specific applications.

Using Optuna in MTCSP

MTCSP offers numerous features for efficiently performing crystal structure searches and easily analyzing results, but Optuna is deeply involved in the search phase. In particular, it is closely related to the structure generation algorithm (how to create candidate crystal structures to evaluate) and the configuration of the search loop, so we will explain those. The part that evaluates the energy of the structure can be done quickly using PFP available on Matlantis, so it is outside the scope of this explanation. A conceptual diagram of MTCSP is shown in Figure 1. MTCSP is a service that processes phases of structure generation and search and phases of structure relaxation and energy evaluation alternately, and finally performs phase diagram evaluation to output a new crystal structure. Of these, Optuna is deeply involved in the "structure generation and search" phase.

Figure 1: Conceptual diagram of MTCSP

The search algorithm using Optuna in MTCSP is not simply the search algorithm implemented in Optuna, but rather one that has been specifically designed and implemented for crystal structure search. For details, please refer to the published technical paper. Historically, crystal structure search has used probabilistic sampling algorithms and genetic algorithms based on them, and we have also designed and implemented a method based on the NSGA-II genetic algorithm implemented in Optuna. In particular, Optuna supports high-speed search through large-scale asynchronous parallel optimization, and in order to benefit from this, the algorithm needs to operate asynchronously and in parallel, so it is implemented in this way in MTCSP as well.

The search loop configuration in MTCSP using Optuna is not simply using Optuna, but rather it is implemented by linking various Optuna functions while adding custom functions as needed. First, the sequence diagram of the initialization phase of the implemented search loop is shown below. The search is controlled by a class called Experiment. This is a wrapper for Optuna's search control class called Study, and provides various convenient functions for MTCSP users. For example, add_pure_atoms adds monatomic crystals corresponding to a specified element system, and create_initial_population adds an initial population for a genetic algorithm specific to MTCSP. These internally generate and persist structures by calling Optuna APIs. When persisting structures, instead of saving them directly to Optuna's storage, a Structures Store is built and used to store structures externally. The Structure Store is a file-based storage. Optuna's storage is typically an RDB such as MySQL, and it is not efficient to directly save relatively large amounts of crystal structures as strings there, so this is used for crystal structures. The initialization phase is performed by a single process and a single thread.

Figure 2: Sequence diagram of the initialization phase in MTCSP

Next, the sequence diagram of the search loop itself is shown below. The search loop is started by the `search` method of the `Experiment` class. This entire loop is controlled by Optuna's `Study` class, which works in conjunction with the MTCSP-specific algorithms mentioned above, a `Relaxer` that evaluates the structures generated by the algorithms, and a `Rejecter` that performs pruning, to efficiently generate and evaluate candidate crystal structures asynchronously and in parallel, and prune any unnecessary ones. To reiterate, this search loop can be executed asynchronously and in parallel, and the same framework supports everything from thread-process parallelism on a single node to multi-node parallel execution on a large-scale computing infrastructure.

Figure 3: Sequence diagram of the search loop in MTCSP

Thus, MTCSP is closely related to Optuna and their relationship is inseparable. There are several reasons for this configuration.

• PFN's large-scale computing cluster is well-suited to Optuna. Optuna's asynchronous processing support makes it easy to evaluate structures with hundreds of thousands of elements.
• The ability to run on both computing clusters and other platforms (i.e., Matlantis) simply by rewriting a few settings.
• MTCSP can also benefit from the continuous development of Optuna, which is open source software. In fact, MTCSP's workload is continuously becoming more efficient thanks to benefits such as faster database speeds.

These are some examples.

Conclusion

This article describes how Matlantis CSP (MTCSP) incorporates the black-box optimization framework Optuna to "smartly run crystal structure searches with fewer trials." Crystal structure searches are challenging due to the large search space, evaluation costs, and the fact that the optimal search strategy changes depending on the element system. However, MTCSP combines algorithms specifically designed for crystal structure searches with Optuna's expertise in asynchronous parallel search loop control, enabling efficient searches even within a realistic computational budget.

We will continue developing both Optuna and MTCSP, striving to improve our service levels by maximizing the synergy between the two. We hope that MTCSP will be useful in your research and development work as a tool to advance the trial and error process of materials discovery.

The post "Crystal structure search using Optuna on Matlantis CSP" first appeared on the Preferred Networks Tech Blog.