relaxengine.py

"""
Copyright (c) 2025 {Chengxi Zhao, Zhaojia Ma, Dingrui Fan}

This source code is licensed under the MIT license found in the
LICENSE file in the root directory of this source tree.
"""

from ase.io import read

# from ase.optimize import ASE_LBFGS
import torch
from torch.multiprocessing import Process, set_start_method
from batchopt.atoms_to_graphs import AtomsToGraphs
from batchopt.utils import data_list_collater
from batchopt.relaxation.optimizers import (
    BFGS,
    BFGSFusedLS,
)
from batchopt.relaxation import OptimizableBatch, OptimizableUnitCellBatch
import logging
import time
import csv
from multiprocessing import Queue
import os
import psutil
import multiprocessing
import json
import subprocess

try:
    from chgnet.model.dynamics import CHGNetCalculator
except ImportError:
    logging.warning("Failed to import CHGNet modules")

try:
    from sevenn.calculator import SevenNetCalculator, SevenNetD3Calculator
except ImportError:
    logging.warning("Failed to import SevenNet modules")

try:
    from fairchem.core import pretrained_mlip, FAIRChemCalculator
except ImportError:
    logging.warning("Failed to import FAIRChem modules")

try:
    from mace.calculators import mace_off
except ImportError:
    logging.warning("Failed to import MACE modules")

import threading
from .utils import count_atoms_cif
from collections import deque


class Scheduler:
    """
    Scheduler distributes relaxation tasks to workers.
    """

    def __init__(
        self,
        files,
        num_workers,
        devices,
        batch_size,
        max_steps,
        filter1,
        filter2,
        optimizer1,
        optimizer2,
        skip_second_stage,
        scalar_pressure,
        compile_mode,
        profile,
        num_threads,
        bind_cores,
        cueq,
        molecule_single,
        output_path,
        model,
    ):

        self.files = files
        self.num_workers = num_workers
        self.devices = devices
        self.batch_size = batch_size
        self.max_steps = max_steps
        self.filter1 = filter1
        self.filter2 = filter2
        self.optimizer1 = optimizer1
        self.optimizer2 = optimizer2
        self.skip_second_stage = skip_second_stage
        self.scalar_pressure = scalar_pressure
        self.compile_mode = compile_mode
        self.profile = profile
        self.num_threads = num_threads
        self.cueq = cueq
        self.molecule_single = molecule_single
        self.output_path = (
            output_path
            if os.path.isabs(output_path)
            else os.path.abspath(output_path)
        )
        self.model = model

        try:
            set_start_method("spawn")
        except RuntimeError:
            logging.warning(
                "set_start_method('spawn') failed, trying 'forkserver' instead."
            )

        if bind_cores is not None:
            self.cpu_mask = self._parse_bind_cores(bind_cores)
        else:
            self.cpu_mask = None

    def _parse_bind_cores(self, bind_cores):
        # Expect custom_bind_str to be like "0-15,16-31,..."
        ranges = bind_cores.split(",")
        if len(ranges) != self.num_workers:
            return None
        binding = []
        for r in ranges:
            try:
                start_str, end_str = r.split("-")
                start = int(start_str)
                end = int(end_str)
            except ValueError:
                logging.error("Custom binding format should be 'start-end'.")
                return None

            binding.append(set(range(start, end + 1)))
        return binding

    def _get_physical_logical_core_mapping(self):
        """Get the mapping between logical cores and their physical core IDs."""
        try:
            # This information is available in Linux systems
            mapping = {}
            logical_cores = psutil.cpu_count(logical=True)

            for i in range(logical_cores):
                try:
                    # Read core_id from /sys/devices/system/cpu/cpu{i}/topology/core_id
                    with open(
                        f"/sys/devices/system/cpu/cpu{i}/topology/core_id"
                    ) as f:
                        core_id = int(f.read().strip())
                    # Read physical_package_id (socket) for more complete information
                    with open(
                        f"/sys/devices/system/cpu/cpu{i}/topology/physical_package_id"
                    ) as f:
                        package_id = int(f.read().strip())
                    mapping[i] = (package_id, core_id)
                except (FileNotFoundError, ValueError, IOError):
                    mapping[i] = None
            return mapping
        except Exception as e:
            logging.error(f"Failed to get core mapping: {e}")
            return {}

    def _get_physical_core_mask(self):
        # Get the number of physical and logical cores
        physical_cores = psutil.cpu_count(logical=False)
        logical_cores = psutil.cpu_count(logical=True)

        if physical_cores is None or physical_cores < 1:
            # Fallback to multiprocessing if psutil fails
            logical_cores = multiprocessing.cpu_count()
            physical_cores = logical_cores // 2
            if physical_cores < 1:
                physical_cores = 1
            print(f"Using estimated physical cores: {physical_cores}")

        # Get the mapping between logical and physical cores
        core_mapping = self._get_physical_logical_core_mapping()

        # Create a CPU mask that includes all physical cores (first core of each physical core)
        physical_core_mask = set()
        if core_mapping:
            # Group by physical core ID
            cores_by_physical = {}
            for logical_id, physical_info in core_mapping.items():
                if physical_info is not None:
                    package_id, core_id = physical_info
                    key = (package_id, core_id)
                    if key not in cores_by_physical:
                        cores_by_physical[key] = []
                    cores_by_physical[key].append(logical_id)

            # Select one logical core from each physical core
            for physical_cores_list in cores_by_physical.values():
                physical_core_mask.add(
                    physical_cores_list[0]
                )  # First logical core of each physical core
        else:
            # If mapping fails, use a simple assumption (may not be accurate on all systems)
            threads_per_core = logical_cores // physical_cores
            physical_core_mask = set(range(0, logical_cores, threads_per_core))

        return physical_core_mask

    def worker_task(
        self, files, device, batch_size, result_queue, physical_cores
    ):
        if physical_cores is not None:
            try:
                # Bind the current process to physical cores
                pid = os.getpid()
                os.sched_setaffinity(pid, physical_cores)
                logging.info(f"bind to physical_core_ids: {physical_cores}")

                # Verify the affinity was set correctly
                current_affinity = os.sched_getaffinity(pid)
                logging.info(
                    f"Process bound to {len(current_affinity)} cores: {sorted(current_affinity)}"
                )

            except AttributeError:
                logging.error(
                    "sched_setaffinity not supported on this platform"
                )
            except Exception as e:
                logging.error(f"Failed to bind to physical cores: {e}")

        # pass the number of processes on each worker
        nproc = self.num_workers // len(self.devices)

        worker = Worker(
            files,
            device,
            batch_size,
            self.max_steps,
            self.filter1,
            self.filter2,
            self.optimizer1,
            self.optimizer2,
            self.skip_second_stage,
            self.scalar_pressure,
            self.compile_mode,
            self.profile,
            self.cueq,
            self.molecule_single,
            self.output_path,
            self.model,
            nproc,
        )
        # results = worker.run()
        results = worker.continuous_run()
        result_queue.put(results)

    def _terminate_processes(self, processes):
        """Helper method to terminate all processes."""
        for i, p in processes:
            if p.is_alive():
                logging.info(f"Terminating process {p.pid}")
                p.terminate()
                p.join(timeout=3)  # Wait for up to 3 seconds
                if p.is_alive():
                    logging.warning(
                        f"Process {p.pid} did not terminate, killing it"
                    )
                    p.kill()
                    p.join()

    # create a thread to conduct "nvidia-smi"
    @staticmethod
    def _monitor_memory(interval=2, gpu_index=1):
        try:
            while True:
                result = subprocess.check_output(
                    [
                        "nvidia-smi",
                        "--query-gpu=memory.used,memory.total",
                        "--format=csv,nounits,noheader",
                    ]
                ).decode("utf-8")

                lines = result.strip().split("\n")
                used, total = map(int, lines[gpu_index].split(","))
                logging.info(
                    f"[nvidia-smi] Memory-Usage on GPU {gpu_index}: {used}MiB / {total}MiB"
                )

                time.sleep(interval)
        except KeyboardInterrupt:
            logging.info("Monitor interrupted.")

        except Exception as e:
            logging.error(f"Unexpected error when monitor memory: {str(e)}")

    def run(self):
        logging.info(f"Starting Scheduler with {self.num_workers} workers.")
        processes = []
        result_queue = Queue()
        start_time = time.perf_counter()

        if self.cpu_mask is not None:
            physical_cores_per_worker = self.cpu_mask
            logging.info(
                f"Use customed cores binding. Physical cores per worker: {physical_cores_per_worker}"
            )
        else:
            # all_physical_cores = self._get_physical_core_mask()
            # num_per_worker = len(all_physical_cores) // self.num_workers
            # physical_cores_per_worker = [
            #     list(all_physical_cores)[i:i + num_per_worker] for i in range(0, len(all_physical_cores), num_per_worker)
            # ]
            # logging.info(f"Physical cores per worker: {physical_cores_per_worker}")
            physical_cores_per_worker = [None] * self.num_workers

        try:
            # Start all worker processes
            for i in range(self.num_workers):
                files_for_worker = self.files[i :: self.num_workers]
                device = self.devices[i % len(self.devices)]
                logging.info(
                    f"Starting worker {i} with {len(files_for_worker)} files on device {device}."
                )
                p = Process(
                    target=self.worker_task,
                    args=(
                        files_for_worker,
                        device,
                        self.batch_size,
                        result_queue,
                        physical_cores_per_worker[i],
                    ),
                )
                p.start()
                processes.append((i, p))

            # monitor gpu memory usage to figure out what makes the differences of footprint among batches
            # in each iteration.
            use_memory_monitor = False
            if use_memory_monitor:
                monitor_proc = Process(
                    target=Scheduler._monitor_memory, args=()
                )
                monitor_proc.start()

            # Monitor processes and collect results
            csv_paths = []
            completed_processes = 0
            while completed_processes < self.num_workers:
                for i, p in processes:
                    if not p.is_alive() and p.exitcode != 0:
                        if p.exitcode == -11 or p.exitcode == 1:
                            # Restart the process if exit code is -11 or -1
                            logging.warning(
                                f"Worker process {p.pid} exited with code {p.exitcode}. Restarting worker {i}."
                            )
                            files_for_worker = self.files[i :: self.num_workers]
                            device = self.devices[i % len(self.devices)]
                            new_process = Process(
                                target=self.worker_task,
                                args=(
                                    files_for_worker,
                                    device,
                                    self.batch_size,
                                    result_queue,
                                    physical_cores_per_worker[i],
                                ),
                            )
                            new_process.start()
                            processes[i] = (
                                i,
                                new_process,
                            )  # Replace the old process with the new one
                        else:
                            # Raise an error for other exit codes
                            raise RuntimeError(
                                f"Worker process {p.pid} failed with exit code {p.exitcode}"
                            )

                # Try to get result from queue with timeout
                try:
                    result = result_queue.get(timeout=10)
                    csv_paths.append(result)
                    completed_processes += 1
                except Exception as e:
                    continue

            # terminate monitor
            if use_memory_monitor:
                monitor_proc.terminate()
                monitor_proc.join()

            # Process results and create final CSV
            merged_results = []
            for csv_path in csv_paths:
                try:
                    with open(csv_path, mode="r") as f:
                        reader = csv.DictReader(f)
                        merged_results.extend(list(reader))
                except Exception as e:
                    logging.error(f"Error processing {csv_path}: {str(e)}")

        except Exception as e:
            # Log the error and elapsed time
            end_time = time.perf_counter()
            elapsed_time = end_time - start_time
            logging.error(
                f"Error occurred after running for {elapsed_time:.2f} seconds: {str(e)}"
            )

            # Create error log file
            error_log = f"scheduler_error_{int(time.time())}.log"
            with open(error_log, "w") as f:
                f.write(f"Error occurred after {elapsed_time:.2f} seconds\n")
                f.write(f"Error message: {str(e)}\n")
                f.write(f"Number of workers: {self.num_workers}\n")
                f.write(f"Batch size: {self.batch_size}\n")

            # Terminate all processes
            self._terminate_processes(processes)
            raise  # Re-raise the exception after cleanup

        finally:
            end_time = time.perf_counter()
            elapsed_time = end_time - start_time

            # Write final results if we have any
            if "merged_results" in locals() and merged_results:
                csv_file = os.path.join(
                    self.output_path, "results_scheduler.csv"
                )
                with open(csv_file, mode="w", newline="") as file:
                    writer = csv.DictWriter(
                        file,
                        fieldnames=[
                            "file",
                            "stage1_steps",
                            "stage1_time",
                            "stage1_energy",
                            "stage1_density",
                            "stage2_steps",
                            "stage2_time",
                            "stage2_energy",
                            "stage2_density",
                            "total_steps",
                            "total_time",
                        ],
                    )
                    writer.writeheader()
                    for row in merged_results:
                        try:
                            processed_row = {
                                "file": row["file"],
                                "stage1_steps": int(row["stage1_steps"]),
                                "stage1_time": float(row["stage1_time"]),
                                "stage1_energy": float(row["stage1_energy"]),
                                "stage1_density": float(row["stage1_density"]),
                                "stage2_steps": int(row["stage2_steps"]),
                                "stage2_time": float(row["stage2_time"]),
                                "stage2_energy": float(row["stage2_energy"]),
                                "stage2_density": float(row["stage2_density"]),
                                "total_steps": int(row["total_steps"]),
                                "total_time": float(row["total_time"]),
                            }
                            writer.writerow(processed_row)
                        except (KeyError, ValueError) as e:
                            logging.error(
                                f"Invalid data format in row {row}: {str(e)}"
                            )

            # Write summary
            summary_csv_file = os.path.join(
                self.output_path, "summary_scheduler.csv"
            )
            with open(summary_csv_file, mode="w", newline="") as file:
                writer = csv.DictWriter(
                    file,
                    fieldnames=["elapsed_time", "num_workers", "batch_size"],
                )
                writer.writeheader()
                writer.writerow(
                    {
                        "elapsed_time": elapsed_time,
                        "num_workers": self.num_workers,
                        "batch_size": self.batch_size,
                    }
                )

            logging.info(f"Scheduler completed in {elapsed_time:.2f} seconds.")

    def run_debug(self):
        logging.info("Starting Scheduler in debug mode (sequential execution).")

        def worker_task(files, device, batch_size):
            worker = Worker(
                files, device, batch_size, self.max_steps, self.filter1
            )
            worker.run()

        for i in range(self.num_workers):
            files_for_worker = self.files[i :: self.num_workers]
            device = self.devices[i % len(self.devices)]
            logging.info(
                f"Running worker {i} with {len(files_for_worker)} files on device {device}."
            )
            worker_task(files_for_worker, device, self.batch_size)

        logging.info("All workers have completed their tasks in debug mode.")


class Worker:
    """
    Worker is single process that runs a batch of optimization tasks.
    """

    def __init__(
        self,
        files,
        device,
        batch_size,
        max_steps,
        filter1,
        filter2,
        optimizer1,
        optimizer2,
        skip_second_stage,
        scalar_pressure,
        compile_mode,
        profile,
        cueq,
        molecule_single,
        output_path,
        model,
        nproc,
    ):
        self.files = files
        self.device = device
        self.batch_size = batch_size
        self.max_steps = max_steps
        self.filter1 = filter1
        self.filter2 = filter2
        self.optimizer1 = optimizer1
        self.optimizer2 = optimizer2
        self.skip_second_stage = skip_second_stage  # Store skip_second_stage
        self.scalar_pressure = scalar_pressure
        self.compile_mode = compile_mode
        self.profile = profile
        self.cueq = cueq
        self.molecule_single = molecule_single
        self.output_path = (
            output_path
            if os.path.isabs(output_path)
            else os.path.abspath(output_path)
        )
        self.model = model
        self.nproc = nproc

        # Parse profiler options if provided
        self.use_profiler = False
        self.profiler_schedule_config = {
            "wait": 48,
            "warmup": 1,
            "active": 1,
            "repeat": 1,
        }
        self.profiler_log_dir = None

        if self.profile and self.profile != "False":
            self.use_profiler = True
            # Create directory for profiler output
            self.profiler_log_dir = os.path.join(self.output_path, "log")
            os.makedirs(self.profiler_log_dir, exist_ok=True)
            if self.profile != "True":
                try:
                    # Try to parse profile as a JSON string with schedule config
                    profile_config = json.loads(self.profile)
                    if isinstance(profile_config, dict):
                        for key in ["wait", "warmup", "active", "repeat"]:
                            if key in profile_config and isinstance(
                                profile_config[key], int
                            ):
                                self.profiler_schedule_config[key] = (
                                    profile_config[key]
                                )
                except json.JSONDecodeError:
                    logging.warning(
                        f"Could not parse profile config: {self.profile}, using defaults"
                    )

        # For monitor thread
        self.stop_event = threading.Event()

    def run(self):
        logging.info(
            f"Worker started on device {self.device} with {len(self.files)} files."
        )
        a2g = AtomsToGraphs(r_edges=False, r_pbc=True)
        # model = torch.load("/home/mazhaojia/.cache/mace/MACE-OFF23_small.model", map_location=self.device)
        # z_table = utils.AtomicNumberTable([int(z) for z in model.atomic_numbers])
        calculator = mace_off(model="small", device=self.device)

        results = []

        for batch_files in self._batch_files(self.files, self.batch_size):
            logging.info(f"Processing batch with {len(batch_files)} files.")
            start_time = time.perf_counter()

            atoms_list = []
            for file in batch_files:
                atoms = read(file)
                atoms_list.append(atoms)
            gbatch = data_list_collater(
                [a2g.convert(atoms) for atoms in atoms_list]
            )

            gbatch = gbatch.to(self.device)
            if self.filter1 == "UnitCellFilter":
                from batchopt.relaxation import OptimizableUnitCellBatch

                obatch = OptimizableUnitCellBatch(
                    gbatch,
                    trainer=calculator,
                    numpy=False,
                    scalar_pressure=self.scalar_pressure,
                )
            else:
                obatch = OptimizableBatch(
                    gbatch, trainer=calculator, numpy=False
                )

            # First optimization stage
            if self.optimizer1 == "LBFGS":
                batch_optimizer1 = LBFGS(
                    obatch, damping=1.0, alpha=70.0, maxstep=0.2
                )
            elif self.optimizer1 == "BFGS":
                batch_optimizer1 = BFGS(obatch, alpha=70.0, maxstep=0.2)
            elif self.optimizer1 == "BFGSLineSearch":
                batch_optimizer1 = BFGSLineSearch(obatch, device=self.device)
            elif self.optimizer1 == "BFGSFusedLS":
                batch_optimizer1 = BFGSFusedLS(obatch, device=self.device)
            else:
                raise ValueError(f"Unknown optimizer: {self.optimizer1}")

            start_time1 = time.perf_counter()
            batch_optimizer1.run(0.01, self.max_steps)
            end_time1 = time.perf_counter()
            elapsed_time1 = end_time1 - start_time1

            # Save intermediate results
            atoms_list = obatch.get_atoms_list()
            for atoms, file_path in zip(atoms_list, batch_files):
                file_name = file_path.split("/")[-1]
                output_file = os.path.join(
                    self.output_path,
                    "cif_result_press",
                    file_name.replace(".cif", "_press.cif"),
                )
                atoms.write(output_file)

            # Capture maximum force after first optimization stage
            max_force1 = obatch.get_max_forces(apply_constraint=True)

            steps1 = batch_optimizer1.nsteps

            if self.skip_second_stage:
                # If skipping second stage, set metrics to zero
                for file, force in zip(batch_files, max_force1):
                    results.append(
                        {
                            "file": file,
                            "stage1_time": elapsed_time1,
                            "stage1_steps": steps1,
                            "stage2_time": 0.0,
                            "stage2_steps": 0,
                            "total_time": elapsed_time1,
                            "total_steps": steps1,
                            "force1": force.item(),
                            "force2": 0.0,
                        }
                    )
                continue

            # Only proceed with second stage if not skipping
            # Reload intermediate structures for second stage
            atoms_list = []
            for file_path in batch_files:
                file_name = file_path.split("/")[-1]
                press_file = os.path.join(
                    self.output_path,
                    "cif_result_press",
                    file_name.replace(".cif", "_press.cif"),
                )
                atoms = read(press_file)
                atoms_list.append(atoms)

            # Rebuild batch from optimized structures
            gbatch = data_list_collater(
                [a2g.convert(atoms) for atoms in atoms_list]
            )
            gbatch = gbatch.to(self.device)

            # Second optimization stage
            if self.filter2 == "UnitCellFilter":
                obatch2 = OptimizableUnitCellBatch(
                    gbatch, trainer=calculator, numpy=False, scalar_pressure=0.0
                )
            else:
                obatch2 = OptimizableBatch(
                    gbatch, trainer=calculator, numpy=False
                )

            if self.optimizer2 == "LBFGS":
                batch_optimizer2 = LBFGS(
                    obatch2, damping=1.0, alpha=70.0, maxstep=0.2
                )
            elif self.optimizer2 == "BFGS":
                batch_optimizer2 = BFGS(obatch2, alpha=70.0, maxstep=0.2)
            elif self.optimizer2 == "BFGSLineSearch":
                batch_optimizer2 = BFGSLineSearch(obatch2, device=self.device)
            elif self.optimizer2 == "BFGSFusedLS":
                batch_optimizer2 = BFGSFusedLS(obatch2, device=self.device)
            else:
                raise ValueError(f"Unknown optimizer: {self.optimizer2}")
            start_time2 = time.perf_counter()
            batch_optimizer2.run(0.01, self.max_steps)
            end_time2 = time.perf_counter()
            elapsed_time2 = end_time2 - start_time2

            # Save final results
            atoms_list = obatch2.get_atoms_list()
            for atoms, file_path in zip(atoms_list, batch_files):
                file_name = file_path.split("/")[-1]
                output_file = os.path.join(
                    self.output_path,
                    "cif_result_final",
                    file_name.replace(".cif", "_opt.cif"),
                )
                atoms.write(output_file)

            # Capture maximum force after second optimization stage
            max_force2 = obatch2.get_max_forces(apply_constraint=True)

            steps2 = batch_optimizer2.nsteps

            for file, f1, f2 in zip(batch_files, max_force1, max_force2):
                results.append(
                    {
                        "file": file,
                        "stage1_time": elapsed_time1,
                        "stage1_steps": steps1,
                        "stage2_time": elapsed_time2,
                        "stage2_steps": steps2,
                        "total_time": elapsed_time1 + elapsed_time2,
                        "total_steps": steps1 + steps2,
                        "force1": f1.item(),
                        "force2": f2.item(),
                    }
                )

        return results

    def _batch_files(self, files, batch_size):
        for i in range(0, len(files), batch_size):
            yield files[i : i + batch_size]

    @staticmethod
    def _torch_memory_monitor(interval=2, device=None, stop_event=None):
        try:
            # explicitly CUDA initialization
            torch.cuda._lazy_init()
            while not stop_event.is_set():
                allocated = torch.cuda.memory_allocated(device=device)
                reserved = torch.cuda.memory_reserved(device=device)
                logging.info(
                    f"[torch] Allocated Memory: {allocated / 1024**2:.2f} MiB"
                )
                logging.info(
                    f"[torch] Reserved Memory: {reserved / 1024**2:.2f} MiB"
                )
                time.sleep(interval)
        except Exception as e:
            logging.error(f"Unexpected error when monitor memory: {str(e)}")

    def continuous_run(self):
        """
        Execute a continuous run of the batching optimization process.
        """
        logging.info("Starting continuous_run with two rounds of optimization.")

        # torch memory monitor api
        use_torch_memory_monitor = False
        if use_torch_memory_monitor:
            memory_monitor = threading.Thread(
                target=Worker._torch_memory_monitor,
                args=(2, self.device, self.stop_event),
            )
            memory_monitor.start()

        # First round of optimization
        try:
            logging.info("Starting first round of optimization.")
            results_round1, new_atoms_files = self.continuous_batching(
                atoms_path=self.files,
                result_path_prefix=os.path.join(
                    self.output_path, "cif_result_press/"
                ),
                fmax=0.01,
                maxstep=self.max_steps,
                use_filter=self.filter1,
                optimizer=self.optimizer1,
                scalar_pressure=self.scalar_pressure,
                dtype=torch.float64,
            )
            logging.info(
                f"Completed first round of optimization. Results: {len(results_round1)}"
            )
        except KeyboardInterrupt as e:
            if use_torch_memory_monitor:
                self.stop_event.set()
                memory_monitor.join()
            logging.error(f"Error during first round of optimization: {e}")
            raise
        except Exception as e:
            logging.error(f"Error during first round of optimization: {e}")
            raise

        if self.skip_second_stage:
            logging.info("Skipping second round of optimization.")
            return results_round1

        # Second round of optimization without pressure
        try:
            logging.info("Starting second round of optimization.")
            results_round2, _ = self.continuous_batching(
                atoms_path=new_atoms_files,
                result_path_prefix=os.path.join(
                    self.output_path, "cif_result_final/"
                ),
                fmax=0.01,
                maxstep=self.max_steps,
                # maxstep=3000,
                use_filter=self.filter2,
                optimizer=self.optimizer2,
                scalar_pressure=0.0,
                dtype=torch.float64,
            )
            logging.info(
                f"Completed second round of optimization. Results: {len(results_round2)}"
            )
        except KeyboardInterrupt as e:
            if use_torch_memory_monitor:
                self.stop_event.set()
                memory_monitor.join()
            logging.error(f"Error during second round of optimization: {e}")
            raise
        except Exception as e:
            logging.error(f"Error during second round of optimization: {e}")
            raise

        if use_torch_memory_monitor:
            self.stop_event.set()
            memory_monitor.join()

        return self._save_results_to_csv(results_round1, results_round2)

    def _save_results_to_csv(self, results_round1, results_round2):
        """Helper method to save results to CSV file and return the path."""
        combined_results = []
        results_map = {}

        # Process first round results
        for result in results_round1:
            file_name = result["file"]
            results_map[file_name] = {
                "file": file_name,
                "stage1_steps": result["steps"],
                "stage1_time": result["runtime"],
                "stage1_energy": result["energy"],
                "stage1_density": result["density"],
                "stage2_steps": 0,
                "stage2_time": 0.0,
                "stage2_energy": 0.0,
                "stage2_density": 0,
                "total_steps": result["steps"],
                "total_time": result["runtime"],
            }

        # Process second round results
        for result in results_round2:
            file_name = result["file"]
            if file_name in results_map:
                results_map[file_name].update(
                    {
                        "stage2_steps": result["steps"],
                        "stage2_time": result["runtime"],
                        "stage2_energy": result["energy"],
                        "stage2_density": result["density"],
                        "total_steps": results_map[file_name]["stage1_steps"]
                        + result["steps"],
                        "total_time": results_map[file_name]["stage1_time"]
                        + result["runtime"],
                    }
                )
            else:
                results_map[file_name] = {
                    "file": file_name,
                    "stage1_steps": 0,
                    "stage1_time": 0.0,
                    "stage1_energy": 0.0,
                    "stage1_density": 0,
                    "stage2_steps": result["steps"],
                    "stage2_time": result["runtime"],
                    "stage2_energy": result["energy"],
                    "stage2_density": result["density"],
                    "total_steps": result["steps"],
                    "total_time": result["runtime"],
                }

        # Convert map to list
        combined_results = list(results_map.values())

        logging.info(
            f"Combined results from both rounds. Total results: {len(combined_results)}"
        )

        worker_id = os.getpid()
        timestamp = int(time.time())
        csv_filename = f"worker_{worker_id}_{timestamp}.csv"
        csv_path = os.path.join(
            self.output_path, "worker_results", csv_filename
        )
        os.makedirs(os.path.dirname(csv_path), exist_ok=True)

        with open(csv_path, mode="w", newline="") as csvfile:
            fieldnames = [
                "file",
                "stage1_steps",
                "stage1_time",
                "stage1_energy",
                "stage1_density",
                "stage2_steps",
                "stage2_time",
                "stage2_energy",
                "stage2_density",
                "total_steps",
                "total_time",
            ]
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
            writer.writeheader()
            for result in combined_results:
                writer.writerow(result)

        return csv_path

    def _get_density(self, crystal):
        # 计算总质量，ASE 中的 get_masses 方法返回一个数组，包含了所有原子的质量
        total_mass = sum(crystal.get_masses())  # 转换为克

        # 获取体积，ASE 的 get_volume 方法返回晶胞的体积，单位是 Å^3
        # 1 Å^3 = 1e-24 cm^3
        volume = crystal.get_volume()  # 转换为立方厘米

        # 计算密度，质量除以体积
        density = (
            total_mass / (volume * 10**-24) / (6.022140857 * 10**23)
        )  # 单位是 g/cm^3
        return density

    @staticmethod
    def select_factor(history: deque):
        # TODO: when history is mix of different size, the smaller `values` should be selected.
        boundaries = [0, 50, 100, 200, 400, 800]
        values = [0.4, 0.8, 0.9, 0.6, 0.5, 0.4]
        factor_result = []
        for graph_size in history:
            for i in range(len(boundaries) - 1):
                if boundaries[i] <= graph_size < boundaries[i + 1]:
                    factor_result.append(values[i])
                    break
        if len(factor_result) == 0:
            return 0.4
        else:
            return min(factor_result)

    def continuous_batching(
        self,
        atoms_path,
        result_path_prefix,
        fmax,
        maxstep,
        use_filter,
        optimizer,
        scalar_pressure,
        dtype=torch.float64,
    ):
        """
        Performs continuous batched optimization of atomic structures.

        This method implements a continuous batching strategy for optimizing multiple atomic structures,
        where converged structures are replaced with new ones to maintain batch efficiency.

        Parameters
        ----------
        atoms_path : list
            List of file paths to atomic structure files to be optimized
        result_path_prefix : str
            Prefix for output file paths where optimized structures will be saved
        fmax : float, optional
            Maximum force criterion for convergence, by default 0.01
        maxstep : int, optional
            Maximum number of optimization steps per batch, by default 3000
        use_filter : str, optional
            Filter to be used for optimization, by default "UnitCellFilter"
        optimizer : str, optional
            Optimizer to be used for optimization, by default "LBFGS"
        scalar_pressure : float, optional
            Scalar pressure to be applied, by default 0.0

        Returns
        -------
        None
            The optimized structures are saved to disk

        Notes
        -----
        The method:
        - Processes structures in batches of predefined size
        - Uses MACE neural network potential for energy/force calculations
        - Employs LBFGS optimization with unit cell relaxation
        - Dynamically replaces converged structures with new ones in the batch
        - Tracks convergence and optimization steps for each structure
        """
        # Load saved structures
        result = []
        optimized_atoms_paths = []

        json_dir = result_path_prefix.replace("cif", "json")

        remove_list = []
        # TODO: Why we read all CIF here?
        for pre_cif in atoms_path:
            cif_path = os.path.join(result_path_prefix, pre_cif.split("/")[-1])
            json_path = os.path.join(
                json_dir, pre_cif.split("/")[-1].replace(".cif", ".json")
            )
            if (
                os.path.exists(cif_path)
                and os.path.exists(json_path)
                and os.path.getsize(cif_path) > 0
                and os.path.getsize(json_path) > 0
            ):
                with open(json_path, "r") as f:
                    result_data = json.load(f)
                    result.append(result_data)
                    optimized_atoms_paths.append(cif_path)
                    remove_list.append(pre_cif)
                    logging.info(f"File {cif_path} already exists, loaded.")
            # else:
            #     try:
            #         read(pre_cif)
            #     except Exception as e:
            #         logging.info(f"Failed to read {pre_cif}: {e}")
            #         remove_list.append(pre_cif)
        for i in remove_list:
            atoms_path.remove(i)

        if self.batch_size > 0:
            # Initialize variables
            room_in_batch = self.batch_size
            indices_to_process = 0
            cur_batch_path = atoms_path[
                indices_to_process : indices_to_process + room_in_batch
            ]
            if len(cur_batch_path) == 0:
                logging.info("No structures to process.")
                return result, optimized_atoms_paths
            room_in_batch -= len(cur_batch_path)
            indices_to_process += len(cur_batch_path)
            cur_atoms_list = [read(path) for path in cur_batch_path]
            a2g = AtomsToGraphs(r_edges=False, r_pbc=True)
            gbatch = data_list_collater(
                [a2g.convert(read(path)) for path in cur_batch_path]
            )
        else:
            # Set Maximum Number of atoms per batch
            history = deque(maxlen=10)
            history.append(1000)
            max_bnatoms = 24080
            safe_factor = self.select_factor(history)

            indices_to_process = 0
            bnatoms = 0
            cur_batch_path = []
            graphs_list = []
            a2g = AtomsToGraphs(r_edges=False, r_pbc=True)

            while indices_to_process < len(atoms_path):
                graph_natoms = count_atoms_cif(atoms_path[indices_to_process])
                if (
                    bnatoms + graph_natoms
                    > max_bnatoms * safe_factor // self.nproc
                ):
                    break
                graph = a2g.convert(read(atoms_path[indices_to_process]))
                bnatoms += graph_natoms
                cur_batch_path.append(atoms_path[indices_to_process])
                graphs_list.append(graph)
                indices_to_process += 1
                history.append(graph_natoms)
            safe_factor = self.select_factor(history)
            if len(graphs_list) == 0:
                logging.info("No structures to process.")
                return result, optimized_atoms_paths
            gbatch = data_list_collater(graphs_list)
            logging.info(f"current batch size: {len(cur_batch_path)}")

            total_natoms = sum([graph.natoms for graph in graphs_list])
            logging.info(f"total_natoms: {total_natoms}")

        gbatch = gbatch.to(self.device)
        batch_optimizer = None

        # Initial calculator
        if self.model == "mace":
            if dtype == torch.float32:
                calculator = mace_off(
                    model="small",
                    device=self.device,
                    enable_cueq=self.cueq,
                    default_dtype="float32",
                )
            else:
                calculator = mace_off(
                    model="small", device=self.device, enable_cueq=self.cueq
                )
        elif self.model == "chgnet":
            calculator = CHGNetCalculator(
                use_device=self.device, enable_cueq=self.cueq
            )
        elif self.model == "sevennet":
            # calculator = SevenNetCalculator(device=self.device, enable_cueq=self.cueq)
            calculator = SevenNetD3Calculator(
                device=self.device,
                enable_cueq=self.cueq,
                batch_size=self.batch_size,
            )
            # calculator = SevenNetCalculator('7net-mf-ompa', modal='mpa', device=self.device)
        # calculator = MACECalculator(model_paths="/home/mazhaojia/.cache/mace/MACE-OFF23_small.model", device=self.device, compile_mode=self.compile_mode)
        if use_filter == "UnitCellFilter":
            obatch = OptimizableUnitCellBatch(
                gbatch,
                trainer=calculator,
                numpy=False,
                scalar_pressure=scalar_pressure,
            )
        else:
            obatch = OptimizableBatch(gbatch, trainer=calculator, numpy=False)

        orig_cells = obatch.orig_cells.clone()

        converged_atoms_count = 0
        converge_indices = []
        all_indices = []
        cur_batch_steps = [0] * len(cur_batch_path)
        cur_batch_times = [time.perf_counter()] * len(
            cur_batch_path
        )  # Track start times

        while converged_atoms_count < len(atoms_path):
            # Update batch
            if len(all_indices) > 0:
                if self.batch_size > 0:
                    room_in_batch += len(all_indices)
                    new_batch_path = atoms_path[
                        indices_to_process : indices_to_process + room_in_batch
                    ]
                    logging.info(f"new_batch_path: {new_batch_path}")
                    room_in_batch -= len(new_batch_path)
                    indices_to_process += len(new_batch_path)

                    optimized_atoms_new = []
                    cur_batch_path_new = []
                    cur_batch_steps_new = []
                    cur_batch_times_new = []
                    orig_cells_new = torch.zeros(
                        [self.batch_size - room_in_batch, 3, 3],
                        device=self.device,
                    )
                    cell_offset = 0

                    restart_indices = []
                    old_batch_indices = obatch.batch_indices
                    for i in range(len(optimized_atoms)):
                        if i in all_indices:
                            continue
                        else:
                            restart_indices.append(i)
                        optimized_atoms_new.append(optimized_atoms[i])
                        cur_batch_path_new.append(cur_batch_path[i])
                        cur_batch_steps_new.append(cur_batch_steps[i])
                        cur_batch_times_new.append(cur_batch_times[i])

                        orig_cells_new[cell_offset] = orig_cells[i]
                        cell_offset += 1

                    for new_path in new_batch_path:
                        optimized_atoms_new.append(read(new_path))
                        cur_batch_path_new.append(new_path)
                        cur_batch_steps_new.append(0)
                        cur_batch_times_new.append(time.perf_counter())

                    # Update the batch with new structures
                    optimized_atoms = optimized_atoms_new
                    cur_batch_path = cur_batch_path_new
                    cur_batch_steps = cur_batch_steps_new
                    cur_batch_times = cur_batch_times_new
                else:
                    bnatoms = 0
                    optimized_atoms_new = []
                    cur_batch_path_new = []
                    cur_batch_steps_new = []
                    cur_batch_times_new = []

                    restart_indices = []
                    old_batch_indices = obatch.batch_indices
                    for i in range(len(optimized_atoms)):
                        if i in all_indices:
                            continue
                        restart_indices.append(i)
                        optimized_atoms_new.append(optimized_atoms[i])
                        cur_batch_path_new.append(cur_batch_path[i])
                        cur_batch_steps_new.append(cur_batch_steps[i])
                        cur_batch_times_new.append(cur_batch_times[i])
                        bnatoms += a2g.convert(read(cur_batch_path[i])).natoms

                    while indices_to_process < len(atoms_path):
                        new_path = atoms_path[indices_to_process]
                        graph_natoms = count_atoms_cif(new_path)
                        if (
                            bnatoms + graph_natoms
                            > max_bnatoms * safe_factor // self.nproc
                        ):
                            break
                        bnatoms += graph_natoms
                        optimized_atoms_new.append(read(new_path))
                        cur_batch_path_new.append(new_path)
                        cur_batch_steps_new.append(0)
                        cur_batch_times_new.append(time.perf_counter())
                        indices_to_process += 1
                        history.append(graph_natoms)
                    safe_factor = self.select_factor(history)

                    orig_cells_new = torch.zeros(
                        [len(optimized_atoms_new), 3, 3], device=self.device
                    )
                    cell_offset = 0
                    for i in range(len(optimized_atoms)):
                        if i in all_indices:
                            continue
                        orig_cells_new[cell_offset] = orig_cells[i]
                        cell_offset += 1

                    # Update the batch with new structures
                    optimized_atoms = optimized_atoms_new
                    cur_batch_path = cur_batch_path_new
                    cur_batch_steps = cur_batch_steps_new
                    cur_batch_times = cur_batch_times_new

                    logging.info(f"current batch size: {len(optimized_atoms)}")

                graphs_list = [a2g.convert(atoms) for atoms in optimized_atoms]
                total_natoms = sum([graph.natoms for graph in graphs_list])
                logging.info(f"total_natoms: {total_natoms}")
                logging.info(f"cur_batch_path to processing: {cur_batch_path}")
                gbatch = data_list_collater(graphs_list)
                gbatch = gbatch.to(self.device)
                if self.model == "sevennet":
                    # calculator = SevenNetCalculator('7net-mf-ompa', modal='mpa', device=self.device)
                    calculator = SevenNetD3Calculator(
                        device=self.device,
                        enable_cueq=self.cueq,
                        batch_size=self.batch_size,
                    )
                if use_filter == "UnitCellFilter":
                    obatch = OptimizableUnitCellBatch(
                        gbatch,
                        trainer=calculator,
                        numpy=False,
                        scalar_pressure=scalar_pressure,
                    )
                else:
                    obatch = OptimizableBatch(
                        gbatch, trainer=calculator, numpy=False
                    )
                for i in range(cell_offset):
                    obatch.orig_cells[i] = orig_cells_new[i]
                orig_cells = obatch.orig_cells.clone()

            # Optimize the current batch
            if optimizer == "LBFGS":
                batch_optimizer = LBFGS(
                    obatch,
                    damping=1.0,
                    alpha=70.0,
                    maxstep=0.2,
                    early_stop=True,
                )
            elif optimizer == "BFGS":
                if len(all_indices) > 0:
                    logging.info(f"Restarting with indices: {restart_indices}")
                    batch_optimizer.optimizable = obatch
                else:
                    batch_optimizer = BFGS(
                        obatch, alpha=70.0, maxstep=0.2, early_stop=True
                    )
            elif optimizer == "BFGSLineSearch":
                batch_optimizer = BFGSLineSearch(
                    obatch,
                    device=self.device,
                    early_stop=True,
                    use_profiler=self.use_profiler,
                    profiler_log_dir=self.profiler_log_dir,
                    profiler_schedule_config=self.profiler_schedule_config,
                )
            elif optimizer == "BFGSFusedLS":
                if len(all_indices) > 0:
                    logging.info(f"Restarting with indices: {restart_indices}")
                    batch_optimizer.optimizable = obatch
                else:
                    batch_optimizer = BFGSFusedLS(
                        obatch,
                        device=self.device,
                        early_stop=True,
                        use_profiler=self.use_profiler,
                        profiler_log_dir=self.profiler_log_dir,
                        profiler_schedule_config=self.profiler_schedule_config,
                    )
            else:
                raise ValueError(f"Unknown optimizer: {optimizer}")

            # 动态计算剩余可用步数（基于当前批次最大已执行步数）
            current_max_steps = max(cur_batch_steps) if cur_batch_steps else 0
            remaining_steps = max(
                maxstep - current_max_steps, 1
            )  # 保证至少运行1步

            # 执行优化并获取收敛的索引
            if (optimizer == "BFGSFusedLS" or optimizer == "BFGS") and len(
                all_indices
            ) > 0:
                converge_indices = batch_optimizer.run(
                    fmax,
                    remaining_steps,
                    is_restart_earlystop=True,
                    restart_indices=restart_indices,
                    old_batch_indices=old_batch_indices,
                )
            else:
                converge_indices = batch_optimizer.run(fmax, remaining_steps)

            # Print energies of all structures
            # logging.info(f"Final energies of all structures: {batch_optimizer.energies}")
            energies_list = (
                batch_optimizer.optimizable.get_potential_energies().tolist()
            )
            logging.info(f"Final energies of all structures: {energies_list}")

            # 更新所有结构的累计步数
            cur_batch_steps = [
                steps + batch_optimizer.nsteps for steps in cur_batch_steps
            ]

            # 找出超过最大步数的结构索引
            over_maxstep_indices = [
                i
                for i, steps in enumerate(cur_batch_steps)
                if steps >= maxstep - 1
            ]

            # 合并收敛和超限的索引（去重）
            all_indices = list(set(converge_indices + over_maxstep_indices))

            # Get optimized atoms
            optimized_atoms = obatch.get_atoms_list()
            converged_atoms_count += len(all_indices)

            end_time = time.perf_counter()
            # 处理所有需要退出的结构（包括收敛和超限）
            for idx in all_indices:
                runtime = end_time - cur_batch_times[idx]

                energy_per_mol = (
                    energies_list[idx]
                    / (
                        len(optimized_atoms[idx].get_atomic_numbers())
                        / self.molecule_single
                    )
                    * 96.485
                )
                density = self._get_density(optimized_atoms[idx])

                # Save results
                result_data = {
                    "file": cur_batch_path[idx].split("/")[-1].split(".")[0],
                    "steps": cur_batch_steps[idx],
                    "runtime": runtime,
                    "energy": energy_per_mol,
                    "density": density,
                }
                result.append(result_data)

                # Save optimized structure
                # converged_atoms_path = os.path.join(result_path_prefix, cur_batch_path[idx].split('/')[-1].replace('.cif', '.traj'))
                converged_atoms_path = os.path.join(
                    result_path_prefix, cur_batch_path[idx].split("/")[-1]
                )
                optimized_atoms[idx].write(converged_atoms_path)
                optimized_atoms_paths.append(converged_atoms_path)

                # write a json file to store reslt_data
                os.makedirs(json_dir, exist_ok=True)
                # json_path = os.path.join(json_dir, cur_batch_path[idx].split('/')[-1]+'.json')
                json_path = os.path.join(
                    json_dir,
                    cur_batch_path[idx].split("/")[-1].replace(".cif", ".json"),
                )
                with open(json_path, "w") as f:
                    json.dump(result_data, f)

            logging.info(f"cur_batch_path: {cur_batch_path}")
            logging.info(f"cur_batch_steps: {cur_batch_steps}")
            logging.info(f"all_indices: {all_indices}")
            logging.info(f"length of optimized_atoms: {len(optimized_atoms)}")

        return result, optimized_atoms_paths