feat(deps): update dependencies and add scripts and config files for generation and relaxation

pyproject.toml

@@ -5,6 +5,7 @@ description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.10, <3.11"
 dependencies = [
+    "torch",
     "ase==3.23.0",
     "networkx==3.2.1",
     "numpy==1.26.4",
@@ -12,9 +13,7 @@ dependencies = [
     "rdkit-pypi>=2022.9.5",
     "scipy==1.14.1",
     "setuptools>=82.0.1",
-    "torch",
     "torch-geometric>=2.7.0",
-    "torch-scatter==2.1.2",
     "tqdm==4.67.1",
     "mace-torch",
     "joblib>=1.5.3",
@@ -23,7 +22,7 @@ dependencies = [
 ]
 
 [tool.uv.sources]
-torch = { path = "../whl/torch-2.5.1+das.opt1.dtk25042-cp310-cp310-manylinux_2_28_x86_64.whl" }
+torch = { path = "/data/wangleping/torch-2.5.1+das.opt1.dtk25041-cp310-cp310-manylinux_2_28_x86_64.whl" }
 mace-torch = { path = "./mace", editable = true } 
 cuequivariance-torch = { path = "./cuequivariance_torch/cuequivariance_torch", editable = true }
 fasteq = { path = "./FastEq", editable = true } 
@@ -41,6 +40,5 @@ build-backend = "hatchling.build"
 packages = ["basic_function", "mace-bench/src/batchopt"]
 
 [tool.uv.extra-build-dependencies]
-torch-scatter = ["torch"]
 cuequivariance-torch = ["hatchling"] 
 

scripts/config.sh

-#!/bin/bash
-# =============================================================================
-# BOMLIP-CSP Configuration File
-# =============================================================================
-# This file contains all configurable parameters for the crystal structure
-# search and generation pipeline. Modify the values below to customize
-# your run.
-# =============================================================================
-
-# -----------------------------------------------------------------------------
-# [Molecular Parameters]
-# SMILES string of the input molecule(s).
-# Use '.' (dot) to separate multiple molecules for co-crystal generation.
-# Example single molecule:  "C1CC2=COC=C12"
-# Example co-crystal:       "C1CC2=COC=C12.CCO"
-SMILES="C1CC2=COC=C12"
-
-# Number of conformers to generate during the conformer search step.
-# Higher values explore more conformational space but take longer.
-# Set to 0 to skip generation and only load existing conformers.
-GENERATE_CONFORMERS=10
-
-# Number of conformers to actually use for crystal structure generation.
-# Must be <= the number of generated conformers.
-# Set to 0 to skip structure generation.
-USE_CONFORMERS=4
-
-# Number of molecules in the unit cell (Z').
-# Use comma-separated values for multiple molecule types (co-crystal),
-# e.g. "1,1" means 1 copy of each molecule in the asymmetric unit.
-# Use space-separated values for multiple packings, e.g. "1 2".
-MOLECULE_NUM_IN_CELL=1
-
-# -----------------------------------------------------------------------------
-# [Crystal Structure Parameters]
-# Space group numbers for structure generation.
-# Use comma-separated values within a packing, and space-separated
-# values for multiple packings.
-# Example: "14,61" means search space groups P21/c and Pbca in one packing.
-# Example: "14 61" means P21/c in packing 1, Pbca in packing 2.
-SPACE_GROUP_LIST="14,61"
-
-# Prefix name added to the output CIF files.
-# Use space-separated values for multiple packings.
-ADD_NAME="XULDUD"
-
-# Number of crystal structures to generate per (space group, conformer) combination.
-# Use space-separated values for multiple packings.
-NUM_GENERATION=100
-
-# -----------------------------------------------------------------------------
-# [Compute Parameters]
-# Maximum number of parallel workers for structure generation.
-# Should not exceed the number of available CPU cores.
-MAX_WORKERS=16
-
-# -----------------------------------------------------------------------------
-# [Run Mode]
-# Execution mode controlling which steps are performed.
-# Available options:
-#   all              - Run conformer search followed by structure generation (default)
-#   conformer_only   - Only perform conformer search
-#   structure_only   - Skip conformer search, use existing conformers to generate structures
-MODE="all"
-
-# -----------------------------------------------------------------------------
-# [Path Parameters]
-# Directory for storing intermediate conformers and output CIF structures.
-# Relative to the project root directory.
-OUTPUT_DIR="csp_results"

scripts/sh_search_gen.sh → scripts/generation.sh

@@ -3,13 +3,13 @@
 # BOMLIP-CSP: Conformer Search & Crystal Structure Generation Pipeline
 # =============================================================================
 # Usage: ./search_gen_proc.sh [config_file]
-#   config_file  Path to the parameter configuration file (default: config.sh)
+#   config_file  Path to the parameter configuration file (default: generation_config.sh)
 # =============================================================================
 
-# --- Determine config file path (default: config.sh in the same directory) ---
+# --- Determine config file path (default: generation_config.sh in the same directory) ---
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 PROJECT_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)"
-CONFIG_FILE="${1:-${SCRIPT_DIR}/config.sh}"
+CONFIG_FILE="${1:-${SCRIPT_DIR}/generation_config.sh}"
 
 if [ ! -f "${CONFIG_FILE}" ]; then
     echo "Error: Configuration file not found: ${CONFIG_FILE}"
@@ -50,7 +50,7 @@ echo " Max workers:          ${MAX_WORKERS}"
 echo "=========================================="
 
 # --- Run the pipeline ---
-python "${SCRIPT_DIR}/search_gen.py" \
+python "${SCRIPT_DIR}/generation.py" \
     --path "${TAR_DIR}" \
     --smiles "${SMILES}" \
     --molecule_num_in_cell ${MOLECULE_NUM_IN_CELL} \

scripts/generation_config.sh

+#!/bin/bash
+# =============================================================================
+# BOMLIP-CSP: Crystal Structure Generation Configuration
+# =============================================================================
+#
+# [Molecular Parameters]
+#   SMILES              SMILES string of the input molecule(s).
+#                       Use '.' (dot) to separate multiple molecules for co-crystal.
+#   GENERATE_CONFORMERS Number of conformers to generate during conformer search.
+#                       Higher values explore more space but take longer.
+#                       Set to 0 to skip generation and only load existing conformers.
+#   USE_CONFORMERS      Number of conformers to use for crystal structure generation.
+#                       Must be <= GENERATE_CONFORMERS. Set to 0 to skip generation.
+#   MOLECULE_NUM_IN_CELL Number of molecules in the unit cell (Z').
+#                       Comma-separated for co-crystal, space-separated for multiple packings.
+#
+# [Crystal Structure Parameters]
+#   SPACE_GROUP_LIST    Space group numbers. Comma-separated within a packing,
+#                       space-separated for multiple packings.
+#   ADD_NAME            Prefix name for output CIF files.
+#   NUM_GENERATION      Number of structures per (space group, conformer) combination.
+#
+# [Compute Parameters]
+#   MAX_WORKERS         Maximum parallel workers for structure generation.
+#
+# [Run Mode]
+#   MODE                Execution mode: all / conformer_only / structure_only
+#
+# [Path Parameters]
+#   OUTPUT_DIR          Directory for intermediate conformers and output CIF files.
+# =============================================================================
+
+SMILES="C1CC2=COC=C12"
+GENERATE_CONFORMERS=10
+USE_CONFORMERS=4
+MOLECULE_NUM_IN_CELL=1
+SPACE_GROUP_LIST="14,61"
+ADD_NAME="XULDUD"
+NUM_GENERATION=100
+MAX_WORKERS=16
+MODE="all"
+OUTPUT_DIR="csp_results"

scripts/relaxation.sh

+#!/bin/bash
+# =============================================================================
+# BOMLIP-CSP: Relaxation
+# =============================================================================
+# Usage: ./relaxation.sh [config_file]
+#   config_file  Path to the parameter configuration file (default: relaxation_config.sh)
+# =============================================================================
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)"
+
+# Load configuration
+CONFIG_FILE="${1:-${SCRIPT_DIR}/relaxation_config.sh}"
+if [ ! -f "${CONFIG_FILE}" ]; then
+    echo "Error: Configuration file not found: ${CONFIG_FILE}" >&2
+    exit 1
+fi
+source "${CONFIG_FILE}"
+
+# Generate timestamped output directory
+if [ -z "${OUTPUT_PATH}" ]; then
+    TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
+    OUTPUT_PATH="${PROJECT_DIR}/csp_results/run_${TIMESTAMP}"
+fi
+mkdir -p "${OUTPUT_PATH}"
+
+echo "Output directory: ${OUTPUT_PATH}"
+
+# Save execution parameters
+cat > "${OUTPUT_PATH}/run_params.txt" <<EOF
+# BOMLIP-CSP Relaxation Run Parameters
+# Generated: $(date +"%Y-%m-%d %H:%M:%S")
+# Config:    ${CONFIG_FILE}
+
+[Input]
+TARGET_FOLDER       = ${TARGET_FOLDER}
+MOLECULE_SINGLE     = ${MOLECULE_SINGLE}
+
+[Compute]
+N_GPUS              = ${N_GPUS}
+GPU_OFFSET          = ${GPU_OFFSET}
+NUM_WORKERS         = ${NUM_WORKERS}
+BATCH_SIZE          = ${BATCH_SIZE}
+NUM_THREADS         = ${NUM_THREADS}
+BIND_CORES          = ${BIND_CORES:-<unset>}
+
+[Optimization]
+MAX_STEPS           = ${MAX_STEPS}
+FMAX                = ${FMAX}
+SCALAR_PRESSURE     = ${SCALAR_PRESSURE}
+SKIP_SECOND_STAGE   = ${SKIP_SECOND_STAGE}
+
+[Optimizer]
+OPTIMIZER1          = ${OPTIMIZER1}
+OPTIMIZER2          = ${OPTIMIZER2}
+FILTER1             = ${FILTER1:-<unset>}
+FILTER2             = ${FILTER2:-<unset>}
+
+[Model]
+MODEL               = ${MODEL}
+CUEQ                = ${CUEQ}
+USE_ORDERED_FILES   = ${USE_ORDERED_FILES}
+COMPILE_MODE        = ${COMPILE_MODE:-<unset>}
+
+[Profiling]
+PROFILE             = ${PROFILE}
+
+[Output]
+OUTPUT_PATH         = ${OUTPUT_PATH}
+EOF
+
+echo "Parameters saved to ${OUTPUT_PATH}/run_params.txt"
+
+# Build optional arguments
+FILTER1_ARG=""
+FILTER2_ARG=""
+[ "${FILTER1}" != "none" ] && [ -n "${FILTER1}" ] && FILTER1_ARG="--filter1 ${FILTER1}"
+[ "${FILTER2}" != "none" ] && [ -n "${FILTER2}" ] && FILTER2_ARG="--filter2 ${FILTER2}"
+
+COMPILE_ARG=""
+[ -n "${COMPILE_MODE}" ] && COMPILE_ARG="--compile_mode ${COMPILE_MODE}"
+
+BIND_CORES_ARG=""
+[ -n "${BIND_CORES}" ] && BIND_CORES_ARG="--bind_cores ${BIND_CORES}"
+
+PROFILE_ARG="--profile ${PROFILE}"
+
+ORDERED_FILES_ARG=""
+[ "${USE_ORDERED_FILES}" = true ] && ORDERED_FILES_ARG="--use_ordered_files true"
+
+CUEQ_ARG=""
+[ "${CUEQ}" = true ] && CUEQ_ARG="--cueq true"
+
+python "${PROJECT_DIR}/mace-bench/scripts/mace_opt_batch.py" \
+    --target_folder "${TARGET_FOLDER}" \
+    --molecule_single ${MOLECULE_SINGLE} \
+    --gpu_offset ${GPU_OFFSET} \
+    --n_gpus ${N_GPUS} \
+    --num_workers ${NUM_WORKERS} \
+    --batch_size ${BATCH_SIZE} \
+    --max_steps ${MAX_STEPS} \
+    --fmax ${FMAX} \
+    --optimizer1 "${OPTIMIZER1}" \
+    --optimizer2 "${OPTIMIZER2}" \
+    --num_threads ${NUM_THREADS} \
+    --model ${MODEL} \
+    --output_path "${OUTPUT_PATH}" \
+    ${FILTER1_ARG} ${FILTER2_ARG} ${COMPILE_ARG} ${BIND_CORES_ARG} ${PROFILE_ARG} ${ORDERED_FILES_ARG} ${CUEQ_ARG} \
+    > "${OUTPUT_PATH}/opt.log" 2>&1
+
+echo "Relaxation completed. Results in ${OUTPUT_PATH}"

scripts/relaxation_config.sh

+#!/bin/bash
+# =============================================================================
+# BOMLIP-CSP: Crystal Structure Relaxation Configuration
+# =============================================================================
+#
+# [Input Parameters]
+#   TARGET_FOLDER       Directory containing crystal structure files to optimize.
+#   MOLECULE_SINGLE     Number of atoms per molecule (default: 64).
+#
+# [Compute Parameters]
+#   N_GPUS              Number of GPUs to use for optimization.
+#   GPU_OFFSET          Starting GPU device index.
+#   NUM_WORKERS         Number of parallel workers distributing the files.
+#   BATCH_SIZE          Number of crystal structures to process in a single batch.
+#   NUM_THREADS         Number of CPU threads per process.
+#   BIND_CORES          Comma-separated list of core ranges for each worker, e.g.
+#                       '0-15,16-31,...'. Must match --num_workers count.
+#
+# [Optimization Parameters]
+#   MAX_STEPS           Maximum optimization steps per structure (default: 100).
+#   FMAX                Force convergence criterion in eV/A (default: 0.01).
+#   SCALAR_PRESSURE     Scalar pressure for cell optimization in GPa (default: 0.0006).
+#   SKIP_SECOND_STAGE   Whether to skip the second optimization stage (true/false).
+#
+# [Optimizer Parameters]
+#   OPTIMIZER1          Optimizer for the first stage. Choices:
+#                       LBFGS, QuasiNewton, BFGS, BFGSLineSearch, BFGSFusedLS.
+#   OPTIMIZER2          Optimizer for the second stage. Same choices as OPTIMIZER1.
+#   FILTER1             Cell filter for the first stage: none / UnitCellFilter.
+#   FILTER2             Cell filter for the second stage: none / UnitCellFilter.
+#
+# [Model Parameters]
+#   MODEL               ML potential model to use. Choices:
+#                       mace, chgnet, sevennet, uma-s-1.
+#   CUEQ                Whether to use cuEquivariance Library (true/false).
+#   USE_ORDERED_FILES   Whether to sort files by atomic number descending (true/false).
+#   COMPILE_MODE        torch.compile mode: none / default / reduce-overhead /
+#                       max-autotune / max-autotune-no-cudagraphs.
+#
+# [Profiling Parameters]
+#   PROFILE             Enable profiling. Set to 'false' to disable, 'true' for
+#                       basic profiling, or a JSON string with options for
+#                       wait, warmup, active, repeat.
+#
+# [Path Parameters]
+#   OUTPUT_PATH         Absolute path for output files. Leave empty "" for
+#                       auto-generated timestamped directory.
+# =============================================================================
+
+TARGET_FOLDER="${PROJECT_DIR}/csp_results/structures"
+MOLECULE_SINGLE=13
+N_GPUS=8
+GPU_OFFSET=0
+NUM_WORKERS=8
+BATCH_SIZE=25
+NUM_THREADS=4
+BIND_CORES=""
+MAX_STEPS=50
+FMAX=0.01
+SCALAR_PRESSURE=0.0006
+SKIP_SECOND_STAGE=false
+OPTIMIZER1="BFGSFusedLS"
+OPTIMIZER2="BFGS"
+FILTER1="UnitCellFilter"
+FILTER2="UnitCellFilter"
+MODEL="mace"
+CUEQ=true
+USE_ORDERED_FILES=false
+COMPILE_MODE=""
+PROFILE="false"
+OUTPUT_PATH=""

scripts/structure_generate.py

@@ -9,9 +9,9 @@ if __name__ == '__main__':
 
     time_start = time.time()
 
-    # --- Load configuration from config.sh ---
+    # --- Load configuration from generation_config.sh ---
     script_dir = os.path.dirname(os.path.abspath(__file__))
-    config_path = sys.argv[1] if len(sys.argv) > 1 else os.path.join(script_dir, "config.sh")
+    config_path = sys.argv[1] if len(sys.argv) > 1 else os.path.join(script_dir, "generation_config.sh")
     cfg = {}
     exec(open(config_path).read(), cfg)
 

uv.lock

@@ -109,7 +109,6 @@ dependencies = [
     { name = "setuptools" },
     { name = "torch" },
     { name = "torch-geometric" },
-    { name = "torch-scatter" },
     { name = "tqdm" },
 ]
 
@@ -126,9 +125,8 @@ requires-dist = [
     { name = "rdkit-pypi", specifier = ">=2022.9.5" },
     { name = "scipy", specifier = "==1.14.1" },
     { name = "setuptools", specifier = ">=82.0.1" },
-    { name = "torch", path = "../whl/torch-2.5.1+das.opt1.dtk25042-cp310-cp310-manylinux_2_28_x86_64.whl" },
+    { name = "torch", path = "/data/wangleping/torch-2.5.1+das.opt1.dtk25041-cp310-cp310-manylinux_2_28_x86_64.whl" },
     { name = "torch-geometric", specifier = ">=2.7.0" },
-    { name = "torch-scatter", specifier = "==2.1.2" },
     { name = "tqdm", specifier = "==4.67.1" },
 ]
 
@@ -932,8 +930,8 @@ wheels = [
 
 [[package]]
 name = "torch"
-version = "2.5.1+das.opt1.dtk25042"
-source = { path = "../whl/torch-2.5.1+das.opt1.dtk25042-cp310-cp310-manylinux_2_28_x86_64.whl" }
+version = "2.5.1+das.opt1.dtk25041"
+source = { path = "/data/wangleping/torch-2.5.1+das.opt1.dtk25041-cp310-cp310-manylinux_2_28_x86_64.whl" }
 dependencies = [
     { name = "filelock" },
     { name = "fsspec" },
@@ -943,7 +941,7 @@ dependencies = [
     { name = "typing-extensions" },
 ]
 wheels = [
-    { filename = "torch-2.5.1+das.opt1.dtk25042-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:c5064c2643f6587ae226afeca9359b04c06da5621a8cdf428bd1ed6b6caf3e89" },
+    { filename = "torch-2.5.1+das.opt1.dtk25041-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:8d8586325865a527617c8e8767344c9a806eab20f9701e54209e81873101fef1" },
 ]
 
 [package.metadata]
@@ -993,12 +991,6 @@ wheels = [
     { url = "https://mirrors.bfsu.edu.cn/pypi/web/packages/1e/d3/4dffd7300500465e0b4a2ae917dcb2ce771de0b9a772670365799a27c024/torch_geometric-2.7.0-py3-none-any.whl", hash = "sha256:6e0cd3ad824d484651ef5d308fc66c687bfcf5ba040d56d1e0fe0f81f365e292", size = 1275346, upload-time = "2025-10-15T20:48:01.949Z" },
 ]
 
-[[package]]
-name = "torch-scatter"
-version = "2.1.2"
-source = { registry = "https://mirrors.bfsu.edu.cn/pypi/web/simple" }
-sdist = { url = "https://mirrors.bfsu.edu.cn/pypi/web/packages/f5/ab/2a44ecac0f891dd0d765fc59ac8d277c6283a31907626560e72685df2ed6/torch_scatter-2.1.2.tar.gz", hash = "sha256:69b3aa435f2424ac6a1bfb6ff702da6eb73b33ca0db38fb26989c74159258e47" }
-
 [[package]]
 name = "torchmetrics"
 version = "1.9.0"