bfgs_torch.py

"""
Copyright (c) 2025 Ma Zhaojia

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

from __future__ import annotations

import logging
import torch
from torch_scatter import scatter

from ..optimizable import OptimizableBatch

class BFGS:
    def __init__(
        self,
        optimizable_batch: OptimizableBatch,
        maxstep: float = 0.2,
        alpha: float = 70.0,
        early_stop = False,
    ) -> None:
        """
        Args:
        """
        self.optimizable = optimizable_batch
        self.maxstep = maxstep
        self.alpha = alpha
        # self.H0 = 1.0 / self.alpha
        self.trajectories = None
        self.device=self.optimizable.device

        self.fmax = None
        self.steps = None

        self.initialize()
        self.early_stop = early_stop
    
    
    def initialize(self):
        # initial hessian
        self.H0 = [
            torch.eye(3 * size, device=self.optimizable.device, dtype=torch.float64) * self.alpha 
            for size in self.optimizable.elem_per_group
        ]

        self.H = [None] * self.optimizable.batch_size
        self.pos0 = torch.zeros_like(self.optimizable.get_positions().reshape(-1), device=self.device, dtype=torch.float64)
        self.forces0 = torch.zeros_like(self.pos0, device=self.device, dtype=torch.float64)

    def restart_from_earlystop(self, restart_indices, old_batch_indices):
        H_new = []
        pos0_new = torch.zeros_like(self.optimizable.get_positions().reshape(-1), device=self.device, dtype=torch.float64)
        forces0_new = torch.zeros_like(pos0_new, device=self.device, dtype=torch.float64)

        # collect the preserved historical data by old_batch_indices
        for i, idx in enumerate(restart_indices):
            mask_old = (idx==old_batch_indices.repeat_interleave(3))
            mask = (i==self.optimizable.batch_indices.repeat_interleave(3))
            H_new.append(self.H[idx])
            pos0_new[mask] = self.pos0[mask_old]
            forces0_new[mask] = self.forces0[mask_old]

        # append new info for the new batch
        for i in range(len(H_new), self.optimizable.batch_size):
            H_new.append(None)

        self.H = H_new
        self.pos0 = pos0_new
        self.forces0 = forces0_new
        

    def run(self, fmax, maxstep, is_restart_earlystop=False, restart_indices=None, old_batch_indices=None):
        logging.info("Enter bfgs's main program.")
        self.fmax = fmax
        self.max_iter = maxstep

        if is_restart_earlystop:
            self.restart_from_earlystop(restart_indices, old_batch_indices)

        iteration = 0
        max_forces = self.optimizable.get_max_forces(apply_constraint=True)
        logging.info("Step   Fmax(eV/A)")

        while iteration < self.max_iter and not self.optimizable.converged(
            forces=None, fmax=self.fmax, max_forces=max_forces, f_upper_limit=1e25,
        ):
            if self.early_stop and iteration > 0:
                converge_indices = self.optimizable.converge_indices_list
                if len(converge_indices) > 0:
                    logging.info(f"Early stopping at iteration {iteration}")
                    break

            logging.info(
                f"{iteration} " + " ".join(f"{x:18.15g}" for x in max_forces.tolist())
            )

            self.step()
            max_forces = self.optimizable.get_max_forces(apply_constraint=True)
            iteration += 1

        logging.info(
            f"{iteration} " + " ".join(f"{x:18.15g}" for x in max_forces.tolist())
        )

        # GPU memory usage as per nvidia-smi seems to gradually build up as
        # batches are processed. This releases unoccupied cached memory.
        torch.cuda.empty_cache()

        # set predicted values to batch
        for name, value in self.optimizable.results.items():
            setattr(self.optimizable.batch, name, value)

        self.nsteps = iteration

        if self.early_stop:
            converge_indices_list = self.optimizable.converge_indices_list
            return converge_indices_list
        else:
            return self.optimizable.converged(
                forces=None, fmax=self.fmax, max_forces=max_forces
            )
        

    def step(self): 
        forces = self.optimizable.get_forces(apply_constraint=True).to(
            dtype=torch.float64
        )
        pos = self.optimizable.get_positions().to(dtype=torch.float64)
        dpos, steplengths = self.prepare_step(pos, forces)
        dpos = self.determine_step(dpos, steplengths)
        self.optimizable.set_positions(pos+dpos)


    def prepare_step(self, pos, forces):
        forces = forces.reshape(-1)
        pos = pos.view(-1)
        self.update(pos, forces, self.pos0, self.forces0)

        dpos_list = []
        cur_indices = self.optimizable.batch_indices.repeat_interleave(3)
        # 预初始化结果列表
        dpos_list = [None] * len(self.H)
        
        # 分离计算任务：仅对需要计算的H矩阵创建流
        calc_indices = [i for i, need_update in enumerate(self.optimizable.update_mask) if need_update]
        streams = [torch.cuda.Stream() for _ in calc_indices]
        
        # 并行执行实际计算
        for i, stream in zip(calc_indices, streams):
            with torch.cuda.stream(stream):
                omega, V = torch.linalg.eigh(self.H[i])
                dpos_list[i] = (V @ (forces[cur_indices==i].t() @ V / torch.abs(omega)).t())

        # 同步所有计算流
        torch.cuda.current_stream().synchronize()
        
        # 在主线程处理零张量
        for i in range(len(self.H)):
            if not self.optimizable.update_mask[i]:
                dpos_list[i] = torch.zeros_like(forces[cur_indices==i])

        # 同步所有流
        for stream in streams:
            stream.synchronize()
        
        # dpos = torch.vstack(dpos_list)
        dpos = torch.zeros_like(forces)
        for i in torch.unique(cur_indices):
            mask = (cur_indices == i)
            dpos[mask] = dpos_list[i]
        dpos = dpos.reshape(-1, 3)

        steplengths = (dpos ** 2).sum(dim=-1).sqrt()
        self.pos0 = pos
        self.forces0 = forces

        return dpos, steplengths


    def determine_step(self, dpos, steplengths):
        longest_steps = scatter(
            steplengths, self.optimizable.batch_indices, reduce="max"
        )
        longest_steps = longest_steps[self.optimizable.batch_indices]
        maxstep = longest_steps.new_tensor(self.maxstep)
        scale = (longest_steps).reciprocal() * torch.min(longest_steps, maxstep)
        dpos *= scale.unsqueeze(1)
        return dpos

    def update(self, pos, forces, pos0, forces0):
        if self.H is None:
            self.H = self.H0
            return
        dpos = pos - pos0
        dforces = forces - forces0
        batch_indices_flatten = self.optimizable.batch_indices.repeat_interleave(3)
        dg = torch.zeros_like(dforces)
        all_size = self.optimizable.elem_per_group

        for i in range(self.optimizable.batch_size):
            if self.H[i] is None:
                continue
            mask = (i==batch_indices_flatten)
            if torch.abs(dpos[mask]).max() < 1e-7:
                continue

            dg[mask] = self.H[i] @ dpos[mask]

        a = self._batched_dot_1d(dforces, dpos)
        b = self._batched_dot_1d(dpos, dg)

        for i in range(self.optimizable.batch_size):
            if self.H[i] is None:
                self.H[i] = torch.eye(3*all_size[i], device=self.device, dtype=torch.float64) * self.alpha
                continue
            mask = (i==batch_indices_flatten)
            if not self.optimizable.update_mask[i]:
                continue
            if torch.abs(dpos[mask]).max() < 1e-7:
                continue

            outer_force = torch.outer(dforces[mask], dforces[mask])
            outer_dg = torch.outer(dg[mask], dg[mask])
            self.H[i] -= outer_force / a[i] + outer_dg / b[i]

        

    def update_parallel(self, pos, forces, pos0, forces0):
        if self.H is None:
            self.H = self.H0
            return

        dpos = pos - pos0

        if torch.abs(dpos).max() < 1e-7:
            return

        dforces = forces - forces0
        cur_indices = self.optimizable.batch_indices.repeat_interleave(3)
        a = self._batched_dot_1d(dforces, dpos)
        # DONE: There is a bug using hstack.
        # dg = torch.hstack([self.H[i] @ dpos[cur_indices == i] for i in range(len(self.H))])
        # DONE: parallel this part
        # dg_list = [self.H[i] @ dpos[cur_indices == i] for i in range(len(self.H))]
        dg_list = [None] * len(self.H)
        streams = [torch.cuda.Stream() for _ in dg_list]
        for i, stream in zip(range(len(dg_list)), streams):
            with torch.cuda.stream(stream):
                dg_list[i] = self.H[i] @ dpos[cur_indices == i]

        torch.cuda.current_stream().synchronize()
        for stream in streams: 
            stream.synchronize()

        dg = torch.zeros_like(dforces)
        for i in torch.unique(cur_indices):
            mask = (cur_indices == i)
            dg[mask] = dg_list[i]
        b = self._batched_dot_1d(dpos, dg)

        # DONE: parallel this part
        for i, stream in zip(range(len(self.H)), streams):
            if not self.optimizable.update_mask[i]:
                continue
            with torch.cuda.stream(stream):
                outer_force = torch.outer(dforces[cur_indices==i], dforces[cur_indices==i])
                outer_dg = torch.outer(dg[cur_indices==i], dg[cur_indices==i])
                self.H[i] -= outer_force / a[i] + outer_dg / b[i]
        
        torch.cuda.current_stream().synchronize()
        for stream in streams: 
            stream.synchronize()


    def _batched_dot_2d(self, x: torch.Tensor, y: torch.Tensor):
        return scatter(
            (x * y).sum(dim=-1), self.optimizable.batch_indices, reduce="sum"
        )
    
    def _batched_dot_1d(self, x: torch.Tensor, y: torch.Tensor):
        return scatter(
            (x * y), self.optimizable.batch_indices.repeat_interleave(3), reduce="sum"
        )