Efficient water desalination with graphene nanopores obtained using artificial intelligence

Single-layer graphene, as an iconic two-dimensional (2D) material, has drawn much scientific attention in recent decades. Because of its ultrathin thickness and outstanding mechanical properties, graphene with artificial pores has been demonstrated to have great potentials in many engineering applications, such as effective hydrogen gas separator, next-generation energy storage or supercapacitor building, and high-resolution DNA sequencing. Given the potential imminent global water scarcity crisis, another important application for nanoporous graphene is energy-efficient water desalination. Equipped with nanoporous 2D material membranes like graphene, the reverse osmosis (RO) water desalination process can expect 2–3 orders improvement in water flux compared with traditional polymeric membranes. In RO, the geometry of nanopores in 2D materials plays a determinant role in water desalination performance. In general, a large pore that allows high water flux is likely to perform poorly in rejecting ions; a small pore that rejects 100% undesired ions, on the other hand, usually has limited water flux. Thus, an optimal nanopore for water desalination is expected to allow as high water flux as possible while maintaining a high ion rejection rate. However, finding the optimal nanopore geometry on graphene can be challenging due to high computational and experimental cost associated with extensive experiments: there are countless possible shapes for a pore on a 4 nm × 4 nm graphene membrane, but evaluating the water flux and ion rejection of a single pore using 10 ns MD simulation takes roughly 36 h on a 56-core CPU cluster. Given this time benchmark, evaluating the water desalination performance of 1000 graphene nanopores can take more than 4 years. Therefore, to discover the optimal graphene nanopore for water desalination, an efficient nanopore screening method with a fast nanopore water desalination performance predictor is needed. Inspired by the recent success of deep learning and reinforcement learning (RL), we create an AI framework consisting of a state-of-the-art deep reinforcement learning (DRL) algorithm coupled with a convolutional neural network (CNN) to solve this challenge. The main idea of RL is to train an agent to find an optimal policy that maximizes the expected return through actively interacting with the environment to achieve a goal. Recently, DRL has proven to be an efficient tool in material-related engineering fields, such as material design and molecule optimization. In this work, we designed and implemented an AI framework consisting of DRL, which is capable of creating a nanopore on a single-layer graphene membrane to reach optimal water desalination performance. By a series of decisions on whether or not to remove carbon atoms and which atom to be removed, the DRL agent can eventually create a pore that allows the highest water flux while maintaining ion rejection rate above an acceptable threshold. Such precisely controlled atom-by-atom removal nanopore synthesis can be conducted by electrochemical reaction. Perforation technologies can also offer the opportunities to control the formation of pores, gaps, and bridges with nanometer dimensions on 2D materials such as graphene experimentally. During training, the DRL agent learns from the feedback based on the water desalination performance (reward for high water flux and penalty for lower ion rejection). However, conventional methods to calculate desalination performance, like MD simulation, are too time-consuming to be implemented directly in our DRL model. To evaluate DRL-designed nanopores fast and accurately, we implemented a CNN-based model that uses the geometry of porous graphene membrane to directly predict the water flux and ion rejection rate under certain external pressure. To this end, a ResNet model is trained on the dataset we collected through MD simulation of water desalination using various graphene nanopores. With the CNN-accelerated desalination performance prediction, the DRL model can rapidly discover the optimal graphene nanopore for water desalination. MD simulations on top-performing DRL-created graphene nanopores prove that they have higher water flux while maintaining a similar ion rejection rate compared to the circular nanopores. Further investigation of molecular trajectories reveals the reason that DRL-created nanopores outperform the conventional circular nanopores and provides insights for energy-efficient water desalination. Lastly, our AI-driven framework can be potentially applied to various application areas of 2D materials besides water desalination, such as gas permeation and separation, battery and super-capacitor applications, and biomolecular translocation.

AI framework: The system couples a deep reinforcement learning (DRL) agent with a convolutional neural network (CNN)-based performance predictor to design graphene nanopores for reverse osmosis desalination. At each timestep, the agent may remove at most one carbon atom from a graphene sheet. The CNN predicts water flux and ion rejection for instantaneous feedback. State representation includes a Morgan fingerprint of the graphene membrane, Cartesian coordinates of atoms, and CNN-extracted geometrical features. Candidate actions are restricted to a subset of M atoms near the pore edge to keep the action space tractable. Dataset generation and augmentation: A graphene desalination MD setup includes a graphene piston applying constant pressure, a saline section (KCl, ~2.28 M), a single-layer graphene membrane with a pore, and a freshwater section. Simulation cell ~4 nm × 4 nm × 13 nm with periodic boundary conditions. Water flux is obtained as the slope of filtered water count vs. time; ion rejection is 1 minus the fraction of ions entering the freshwater region. A total of 185 distinct porous graphene structures were simulated. To enlarge the dataset for CNN training, pores were augmented by flips along x or y axes and translations (−4 to 4 Å) in x and y (validated to preserve performance). Augmentation multiplicity: DRL-generated pores ×32; zero-flux pores ×6; others ×24. Final dataset: 3937 samples. A reverse sigmoid captures the general trade-off between flux and ion rejection across samples. Performance prediction models: CNNs (VGG16, ResNet18, ResNet50) and an MLP head predict standardized flux and ion rejection; XGBoost with 1024-dim one-hot Morgan fingerprints (cutoff 5 Å) serves as a baseline. CNN input representation: atoms colored on a 380×380 pixel map resized to 224×224; CNN outputs a 1000-dim feature vector then passed to a two-layer MLP (256 and 64 neurons) with residual blocks and ReLU to predict flux or ion rejection (two separate CNN+MLP models). Models implemented in PyTorch, CNNs pre-trained on ImageNet; Adam optimizer with learning rates 1e-4 (CNN) and 1e-3 (MLP). Train/test split 4:1; 600 epochs; best test MSE model selected. DRL formulation: Finite MDP with discount γ=1. State s_t includes Morgan fingerprint, atom coordinates (with removed atoms set to origin to maintain fixed input size), and CNN-extracted features. Predicted flux f_t and ion rejection i_t are used for reward r_t = α f_t + σ(i_t) − σ(1), with α=0.01 and a generalized logistic function σ(x) parameterized by A=−15, K=0, B=13, Q=100, β=0.01, C=1. An extra +0.05 reward is given when removing an atom to encourage early growth. Candidate action set c_t contains up to M atoms nearest the pore center; if insufficient edge atoms, nearby non-edge atoms are included. Optimization via Deep Q-Learning with experience replay: Q-network and target network (identical architecture), Adam optimizer lr=0.001, replay buffer size 10,000, batch 128, target update every 10 steps. Ten random seeds were trained for 2000 episodes each. MD simulation details: Simulations performed in LAMMPS. Water model: SPC/E with SHAKE constraints. Nonbonded: Lennard-Jones (cutoff 12 Å; Lorentz–Berthelot mixing; parameters in Supplementary Table 1) plus long-range Coulomb via PPPM Ewald (RMS error 0.005). Graphene membrane and piston treated as rigid entities (internal interactions not computed) to reduce cost. Protocol: energy minimization (1000 iterations), NPT at 300 K for 5 ps (velocities Gaussian-initialized), then NVT at 300 K for 10 ns with Nosé–Hoover thermostat (τ=0.5 ps). External pressure of 100 MPa applied by the piston along z to mimic RO; flux is approximately linear with pressure, enabling extrapolation to lower pressures. Trajectories saved every 5 ps. Data augmentation used ASE; pore area and perimeter computed via computer vision routines. Feature visualization: t-SNE on 1000-dim CNN features maps DRL-created (7999) and dataset pores (3937) into 2D, clustering by similar flux or ion rejection to validate learned geometric descriptors.

- CNN performance predictor accuracy (standardized labels): • XGBoost: Flux MSE 0.011, R² 0.988; Ion rejection MSE 0.008, R² 0.992. • VGG16: Flux MSE 0.0448, R² 0.957; Ion rejection MSE 0.0156, R² 0.985. • ResNet18: Flux MSE 0.0024, R² 0.998; Ion rejection MSE 0.0039, R² 0.996. • ResNet50: Flux MSE 0.0022, R² 0.998; Ion rejection MSE 0.0038, R² 0.996. ResNet50 chosen for DRL reward estimation. - DRL training (10 seeds, 2000 episodes) learns to balance flux increase with ion-rejection penalties. The agent ceases pore growth when further enlargement would excessively reduce ion rejection. Predicted performance of DRL-created pores averages ~40 ns⁻¹ water flux with ~96% ion rejection. - Through DRL+CNN, 7999 candidate nanopores were generated and evaluated in under a week; equivalent brute-force MD evaluation would take ~33 years on a 56-core CPU node (~36 h per sample). - DRL extrapolates beyond the training dataset to discover nanopores with higher water flux at fixed ion rejection thresholds (e.g., >90%). - MD validation shows DRL-created pores achieve higher ion rejection at similar flux compared to circular pores: • For equal area (113 Ų): DRL pore maintains >90% ion rejection, whereas a circular pore rejects ~65% of ions under 100 MPa. • For matched flux (~125 ns⁻¹): DRL pore (113 Ų) rejects ~7% more ions than a circular pore (88 Ų). - Geometric insight: High-performance DRL pores exhibit semi-oval shapes with rough edges and corners that create ion-free zones due to steric exclusion of hydrated ions. DRL pores generally have higher perimeter-to-area ratios than circular pores, correlating with improved ion rejection at comparable flux. - t-SNE of CNN features clusters pores by desalination performance and geometry, indicating the CNN captures relevant geometric descriptors.

The study addresses the challenge of discovering optimal graphene nanopore geometries for RO desalination by integrating DRL with a fast, accurate CNN performance predictor. This closed-loop AI framework enables rapid, informed exploration of pore designs, learning policies that maximize flux while preserving high ion rejection. Compared with circular pores, DRL-designed pores emphasize rough edges and corner features that create steric barriers to hydrated ions, thereby enhancing ion rejection without sacrificing water transport. The ability to evaluate thousands of designs in days (rather than decades via MD alone) demonstrates the framework’s practical value for materials discovery. The findings validate that geometric factors—especially perimeter/area ratio and the presence of small interior corners—are key determinants of desalination performance, and that DRL can reliably exploit these to optimize pore topology. The approach is generalizable to other 2D materials and separation or transport applications where geometry critically governs performance.

This work introduces an AI-driven framework that combines a ResNet50-based CNN performance predictor with a deep Q-learning DRL agent to design graphene nanopores for efficient water desalination. The method accelerates pore screening by orders of magnitude, creating and assessing 7999 pores in under a week. DRL-designed pores achieve higher ion rejection at similar flux compared to circular pores by leveraging rough edges and corner-induced steric hindrance, as quantified by increased perimeter/area ratios. The framework demonstrates strong extrapolative capability, discovering pores with improved flux for a given ion rejection threshold. With minor modifications, the approach can be extended to other nanomaterial design problems, including gas separation, energy storage membranes, and biomolecular translocation, provided a reliable ML-based property predictor is available.

- Simulation conditions differ from typical seawater RO: saline molarity is ~2.28 M (higher than seawater) and pressure is 100 MPa, chosen for computational efficiency and later extrapolated assuming linearity with pressure. - Base MD dataset contains 185 unique pores, expanded to 3937 via data augmentation (flips/translations and small Gaussian noise on labels). While validated for invariance, augmentation may limit diversity compared to fully independent simulations. - DRL relies on CNN predictions for rewards; inaccuracies, though small, can influence policy learning. Only selected high-performance designs were validated by full MD. - Action space reduction to candidate edge atoms improves tractability but may omit some global edits; pore stabilization (e.g., edge passivation) is discussed but not explicitly simulated across all designs. - Experimental realization and long-term stability of irregular, rough-edged pores were not tested; results are based on atomistic simulations.