Crystal Structure Prediction Using Optuna in Matlantis CSP

Area

Chemoinformatics / Materials Science

Data Science / AutoML

Tag

# Matlantis

# Optuna

# 結晶構造予測(CSP)

Hideaki Imamura

Introduction

Crystal structure prediction represents a critical initial step in materials development and property prediction. However, this search space is vast, with computational evaluation being both time-consuming and resource-intensive. The actual stability of specific compositions and structures varies significantly across different elemental systems. As a result, the success and efficiency of the search process often heavily depend on “experience and trial-and-error,” making it natural for those working on the same problems to frequently encounter situations where they wish they could “make the process more intelligent.”

Recently, we launched the crystal structure prediction service Matlantis CSP (MTCSP) on our cloud-based atomic-level AI simulation platform Matlantis™︎. For MTCSP, we have implemented the black-box optimization framework Optuna™︎—which PFN leads in developing—as an approach to facilitate efficient and effective crystal structure prediction through minimal trial iterations. In this article, we explain how Optuna has been incorporated into MTCSP.

Optuna, The Black-Box Optimization Framework

Optuna is a black-box optimization framework designed to help find optimal solutions efficiently for problems where “evaluation is possible but determining how to optimize remains challenging.” While typical applications include hyperparameter tuning for machine learning, its fundamental principle applies broadly—any system where “input parameters yield corresponding output scores” can serve as the target domain.

In crystal structure prediction, this involves searching for energetically stable structures that satisfy given constraints such as elemental composition ranges. The input is the crystal structure itself, while the output is the structure’s formation energy. Conventional approaches relied on computationally intensive DFT calculations, but on Matlantis, this can be achieved rapidly using PFP, our universal machine learned interatomic potential (uMLIP) for formation energy evaluation.

Matlantis CSP, Crystal Structure Prediction Service

Matlantis CSP(MTCSP) is a crystal structure prediction service offered on the cloud-based atomic-level AI simulation platform Matlantis. Upon receiving user-specified conditions, the service generates candidate crystal structures (atomic configurations), evaluates them based on energy metrics, and systematically explores more stable (or potentially promising) structures.

The challenge of crystal structure prediction stems from two key factors: first, the vast search space involved, and second, the significant computational time required for structure evaluation. The effectiveness of different search approaches varies dramatically depending on the elemental composition and experimental conditions. This means that structure exploration isn’t simply a matter of increasing candidate generation – within limited computational resources, careful consideration must be given to:

how to efficiently generate candidate structures
how to effectively evaluate structures
how to configure the search loop

MTCSP was developed to enable practical execution of extensive searches by treating both the search execution and result analysis as service components. This allows users to focus exclusively on specifying purpose-appropriate conditions and conducting analysis tailored to their specific applications.

How to Use Optuna in MTCSP

While MTCSP provides numerous features for efficiently conducting crystal structure predictions and easily analyzing results, its primary involvement with Optuna lies in the search process. Specifically, it plays a critical role in both structural generation algorithms (how to create candidate crystal structures for evaluation) and the configuration of the search loop. Therefore, we will focus on explaining these aspects. The energy evaluation component for structures can be implemented efficiently using PFP available in Matlantis, so this aspect is outside the scope of discussion. Figure 1 shows the conceptual diagram of MTCSP. MTCSP operates by alternating between phases for structural generation/search and phases for structural relaxation/energy evaluation, ultimately performing phase diagram evaluation and outputting new crystal structures. Within this process, Optuna’s primary contribution is in the “structural generation and search” phase.

Fig1: MTCSP overview

The search algorithm implemented in MTCSP using Optuna is not simply utilizing Optuna’s built-in search algorithms but has been specially designed and implemented specifically for crystal structure prediction. See our technical paper for more details. For context, historical approaches in crystal structure prediction have predominantly employed stochastic sampling algorithms and their derivative genetic algorithms. Following this tradition, we designed and implemented a method based on NSGA-II, a well-known genetic algorithm implemented in Optuna. Notably, Optuna supports high-speed search through large-scale asynchronous parallel optimization. To fully benefit from this capability, the algorithm must operate in an asynchronous parallel manner, which is precisely how it has been implemented in MTCSP.

The search loop implementation in MTCSP combines Optuna with custom functionalities as needed, rather than simply utilizing Optuna’s core features. Below is a sequence diagram illustrating the initialization phase of the search process. The search is controlled by an “Experiment” class, which serves as a wrapper around Optuna’s “Study” class responsible for managing the search process. This class provides various convenient features for MTCSP users, including:

– add_pure_atoms: Adds single-atom crystals corresponding to specified elemental systems

– create_initial_population: Generates initial populations for MTCSP-specific genetic algorithms

These functionalities internally invoke Optuna’s API while handling structure generation and persistence. When persisting structures, instead of storing them directly in Optuna’s storage, we use a separate external “Structures Store” for storage. The Structures Store is a file-based storage system. While Optuna’s storage typically uses relational databases like MySQL, storing relatively large crystal structures as strings directly in this database would be inefficient, so we use this alternative approach for crystal structures. The initialization phase is executed in a single process with a single thread.

Fig2: Sequence diagram for the initialization phase in MTCSP

Below is the sequence diagram for the core search loop mechanism. The search loop is initiated by the Experiment’s search method. This entire loop is managed by Optuna’s Study class, which coordinates with specialized MTCSP algorithms, the structure evaluation Relaxer, and the branch pruning Rejecter to generate, evaluate, and efficiently prune candidate crystal structures asynchronously and in parallel. This search loop can be executed both asynchronously and in parallel – supporting both thread/process parallelism on a single node and multi-node execution across large-scale computational infrastructures, all within the same unified framework.

Fig3: Sequence diagram of the search loop in MTCSP

This integration demonstrates the inseparable relationship between MTCSP and Optuna. There are several key reasons for this configuration:

The excellent compatibility between PFN’s large-scale computing cluster and Optuna. Optuna’s asynchronous processing support enables efficient evaluation of hundreds of thousands of structures.
The ability to run the same code on both the computing cluster and alternative platforms (such as Matlantis) with minimal configuration changes.
The ability to benefit from Optuna’s ongoing development as an open source project. MTCSP has indeed gained advantages like database performance improvements, resulting in continuous efficiency gains in its workload processing.

Conclusion

This article explained how Matlantis CSP (MTCSP) incorporates the black-box optimization framework Optuna to enable “efficient, intelligent crystal structure prediction with minimal computational effort.” While crystal structure prediction presents challenges due to its vast search space, high evaluation costs, and the need to adjust search strategies based on different elemental systems, MTCSP addresses these by combining specialized algorithms designed for crystal structure prediction with Optuna’s expertise in asynchronous parallel search loop management. This approach allows efficient exploration even within realistic computational budgets.