Grid density

deepmd/utils/data.py

@@ -6,6 +6,7 @@
 from typing import (
     List,
     Optional,
+    Tuple,
 )
 
 import numpy as np
@@ -15,6 +16,10 @@
     GLOBAL_NP_FLOAT_PRECISION,
 )
 from deepmd.utils import random as dp_random
+from deepmd.utils.density import (
+    calculate_density,
+    generate_grid,
+)
 from deepmd.utils.path import (
     DPPath,
 )
@@ -58,6 +63,8 @@ def __init__(
         modifier=None,
         trn_all_set: bool = False,
         sort_atoms: bool = True,
+        density_grid_size: Tuple[int, int, int] = (100, 100, 100),
+        density_origin: np.ndarray = np.zeros(3),
     ):
         """Constructor."""
         root = DPPath(sys_path)
@@ -126,6 +133,8 @@ def __init__(
         self.nframes = np.sum(frames_list)
         # The prefix sum stores the range of indices contained in each directory, which is needed by get_item method
         self.prefix_sum = np.cumsum(frames_list).tolist()
+        self.density_grid_size = density_grid_size
+        self.density_origin = density_origin
 
     def add(
         self,
@@ -607,7 +616,35 @@ def _load_data(
         else:
             dtype = GLOBAL_NP_FLOAT_PRECISION
         path = set_name / (key + ".npy")
-        if path.is_file():
+        if key == "grid":
+            box_path = set_name / "box.npy"
+            boxes = box_path.load_numpy()
+            data = []
+            for box in boxes:
+                box = box.reshape(3,3)
+                grid = generate_grid(box, self.density_grid_size, self.density_origin)  # [ngrids, 3]
+                data.append(grid)
+            data = np.stack(data)  # [nframes, ngrids, 3]
+            return np.float32(1.0), data
+        elif key == "density" and path.is_file():
+            path_list = path.load_numpy()
+            box_path = set_name / "box.npy"
+            boxes = box_path.load_numpy()
+            data = []
+            for idx, path in enumerate(path_list):
+                filename = set_name / path[0]
+                densities = []
+                box = boxes[idx]
+                box = box.reshape(3,3)
+                for _, batch_densities in calculate_density(
+                    str(filename), box, self.density_grid_size, self.density_origin
+                ):
+                    densities.append(batch_densities)
+                densities = np.concatenate(densities)  # [ngrids]
+                data.append(densities)
+            data = np.stack(data)  # [nframes, ngrids]
+            return np.float32(1.0), data
+        elif path.is_file():
             data = path.load_numpy().astype(dtype)
             try:  # YWolfeee: deal with data shape error
                 if atomic:

deepmd/utils/density.py

@@ -0,0 +1,135 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import numpy as np
+from typing import Tuple
+class DensityCalculator:
+    def __init__(self, binary_file: str, lattice_constant: float):
+        self.lattice_constant = lattice_constant  # in Bohr
+        self.read_binary_file(binary_file)
+
+    def read_binary_file(self, filename: str):
+        # Read the binary file for g_vectors & rhog
+        with open(filename, 'rb') as f:
+            # Read header
+            size = np.fromfile(f, dtype=np.int32, count=1)[0]
+            assert size == 3, f"Unexpected header value: {size}"
+
+            self.gammaonly, self.ngm_g, self.nspin = np.fromfile(f, dtype=np.int32, count=3)
+            _ = np.fromfile(f, dtype=np.int32, count=1)[0]  # Skip the last '3'
+
+            # Read reciprocal lattice vectors
+            size = np.fromfile(f, dtype=np.int32, count=1)[0]
+            assert size == 9, f"Unexpected header value: {size}"
+            self.bmat = np.fromfile(f, dtype=np.float64, count=9).reshape(3, 3)
+            _ = np.fromfile(f, dtype=np.int32, count=1)[0]  # Skip the last '9'
+
+            # Read Miller indices
+            size = np.fromfile(f, dtype=np.int32, count=1)[0]
+            assert size == self.ngm_g * 3, f"Unexpected Miller indices size: {size}"
+            self.miller_indices = np.fromfile(f, dtype=np.int32, count=self.ngm_g*3).reshape(self.ngm_g, 3)
+            _ = np.fromfile(f, dtype=np.int32, count=1)[0]  # Skip the size at the end
+
+            # Read rhog
+            size = np.fromfile(f, dtype=np.int32, count=1)[0]
+            assert size == self.ngm_g, f"Unexpected rhog size: {size}"
+            self.rhog = np.fromfile(f, dtype=np.complex128, count=self.ngm_g)
+            _ = np.fromfile(f, dtype=np.int32, count=1)[0]  # Skip the size at the end
+
+            # If nspin == 2, read second spin component
+            if self.nspin == 2:
+                size = np.fromfile(f, dtype=np.int32, count=1)[0]
+                assert size == self.ngm_g, f"Unexpected rhog_spin2 size: {size}"
+                self.rhog_spin2 = np.fromfile(f, dtype=np.complex128, count=self.ngm_g)
+                _ = np.fromfile(f, dtype=np.int32, count=1)[0]  # Skip the size at the end
+
+        # Calculate real space lattice vectors (dimensionless)
+        self.lat_vec = np.linalg.inv(self.bmat).T
+
+        # Calculate cell volume
+        self.cell_volume = np.abs(np.linalg.det(self.lattice_constant * self.lat_vec))
+
+        # Calculate G vectors (in units of 2π/lattice_constant)
+        self.bmat = 2 * np.pi * self.bmat / self.lattice_constant
+        # self.g_vectors = 2 * np.pi * np.dot(self.miller_indices, self.bmat.T)
+        self.g_vectors = np.dot(self.miller_indices, self.bmat.T)
+    def calculate_density_batch(self, points: np.ndarray) -> np.ndarray:
+        """
+        Calculate the exact density using Fourier transform.
+
+        Args:
+            points (np.ndarray): Array of points in real space, shape (n, 3), in units of lattice_constant
+
+        Returns:
+            np.ndarray: Exact density values
+        """
+        phases = np.exp(1j * np.dot(points, self.g_vectors.T))
+        densities = np.real(np.dot(phases, self.rhog))
+        return densities
+
+    def get_lattice_vectors(self) -> np.ndarray:
+        """
+        Get the real space lattice vectors.
+
+        Returns:
+            np.ndarray: Real space lattice vectors, shape (3, 3), in Bohr units
+        """
+        return self.lattice_constant * self.lat_vec
+
+
+# def generate_grid(
+#     lattice_vectors: np.ndarray,
+#     grid_size: Tuple[int, int, int],
+#     origin: np.ndarray = np.zeros(3),
+# ) -> np.ndarray:
+#     """
+#     生成给定晶格和网格大小的实空间网格点。
+
+#     参数:
+#     lattice_vectors (np.ndarray): 形状为 (3, 3) 的晶格向量
+#     grid_size (Tuple[int, int, int]): 三个方向上的网格点数
+#     origin (np.ndarray): 网格原点坐标，默认为 (0, 0, 0)
+
+#     返回:
+#     np.ndarray: 形状为 (N, 3) 的网格点坐标数组
+#     """
+#     nx, ny, nz = grid_size
+#     x = np.linspace(0, 1, nx, endpoint=False)
+#     y = np.linspace(0, 1, ny, endpoint=False)
+#     z = np.linspace(0, 1, nz, endpoint=False)
+
+#     grid = np.meshgrid(x, y, z, indexing="ij")
+#     points = np.stack(grid, axis=-1).reshape(-1, 3)
+
+#     # 将分数坐标转换为实空间坐标
+#     real_points = np.dot(points, lattice_vectors) + origin
+
+#     return real_points
+
+
+# def calculate_density(
+#     filename: str,
+#     lattice_vectors: np.ndarray,
+#     grid_size: Tuple[int, int, int],
+#     origin: np.ndarray = np.zeros(3),
+#     batch_size: int = 1000,
+# ) -> Iterator[Tuple[np.ndarray, np.ndarray]]:
+#     """
+#     计算给定网格的电子密度，以迭代器形式返回结果。
+
+#     参数:
+#     filename (str): 包含 n(G) 数据的二进制文件路径
+#     grid_size (Tuple[int, int, int]): 三个方向上的网格点数
+#     origin (np.ndarray): 网格原点坐标，默认为 (0, 0, 0)
+#     lattice_vectors (np.ndarray): 形状为 (3, 3) 的晶格向量
+#     batch_size (int): 每批处理的点数
+
+#     返回:
+#     Iterator[Tuple[np.ndarray, np.ndarray]]: yielding (坐标, 密度值) 对
+#     """
+#     calculator = DensityCalculator(filename)
+
+#     points = generate_grid(lattice_vectors, grid_size, origin)
+
+#     for i in range(0, len(points), batch_size):
+#         batch_points = points[i : i + batch_size]
+#         batch_densities = calculator.calculate_density_batch(batch_points)
+#         yield batch_points, batch_densities

source/tests/dentest/dataset/Al13Cu3/type.raw

@@ -0,0 +1,16 @@
+0
+0
+0
+0
+0
+0
+0
+0
+0
+0
+0
+0
+0
+1
+1
+1

source/tests/dentest/dataset/Al13Cu3/type_map.raw

@@ -0,0 +1,2 @@
+Al
+Cu

source/tests/dentest/dataset/Mg1Al20Cu11/.completed

@@ -0,0 +1 @@
+done

source/tests/dentest/dataset/Mg1Al20Cu11/.lbg-18431-14407614_base.sh

@@ -0,0 +1 @@
+ulimit -s unlimited;ulimit -m unlimited;cd /home/input_lbg-18431-14407614;touch /home/input_lbg-18431-14407614/STDOUTERR;/bin/bash /home/input_lbg-18431-14407614/lbg-18431-14407614.sh >/home/input_lbg-18431-14407614/STDOUTERR 2>&1

source/tests/dentest/dataset/Mg1Al20Cu11/.lbg-compute-completed

@@ -0,0 +1 @@
+done

source/tests/dentest/dataset/Mg1Al20Cu11/INPUT

@@ -0,0 +1,13 @@
+INPUT_PARAMETERS
+calculation scf
+ecutwfc 100.000000
+scf_thr 1.000000e-07
+scf_nmax 100
+basis_type pw
+dft_functional pbe
+gamma_only 0
+kspacing 0.1
+symmetry 0
+mixing_beta 0.1
+out_chg 1 7
+ks_solver dav_subspace

source/tests/dentest/dataset/Mg1Al20Cu11/KPT

@@ -0,0 +1,4 @@
+K_POINTS
+0
+Gamma
+5 4 3 0 0 0