Source code for gnrs.energy.energy_task

"""
This module computes the energy.

This source code is licensed under the BSD-3-Clause license found in the
LICENSE file in the root directory of this source tree.
"""

from __future__ import annotations

__author__ = ["Yi Yang", "Rithwik Tom"]
__email__ = "yiy5@andrew.cmu.edu"
__group__ = "https://www.noamarom.com/"

import importlib
import json
import logging
import os
from functools import partial

from mpi4py import MPI

import gnrs.output as gout
from gnrs.core.task import TaskABC
from gnrs.parallel.structs import DistributedStructs

AVAILABLE_CALCULATORS = ["DFTB", "AIMS", "MACEOFF", "UMA", "VASP", "AIMNET"]
logger = logging.getLogger("EnergyCalcTask")




class EnergyCalculationTask(TaskABC):
    """
    Task for computing energy using DFT, semi-empirical, or MLIP methods.
    Uses ASE calculators for energy evaluation.
    """



    def __init__(
        self,
        comm: MPI.Comm,
        config: dict,
        gnrs_info: dict,
        energy_method: str,
        instance_id: str | None = None,
    ) -> None:
        """
        Initialize the energy calculation task.

        Args:
            comm: MPI communicator
            config: Config dictionary
            gnrs_info: Genarris info dictionary
            energy_method: Energy calculation method
            instance_id: Unique ID for this task instance
        """
        super().__init__(comm, config, gnrs_info, instance_id=instance_id)
        self.energy_name = energy_method.lower()
        self.energy_class = self.energy_name.upper() + "Energy"
        self.energy_file = f"gnrs.energy.{self.energy_name}"

        try:
            energy_module = importlib.import_module(self.energy_file)
            self.energy_calc = getattr(energy_module, self.energy_class)
        except (ImportError, AttributeError):
            logger.warn("Unable to find requested energy calculation method.")
            logger.warn(f"Available calculators: {AVAILABLE_CALCULATORS}")
            raise




    def initialize(self) -> None:
        """
        Initialize the energy calculation task.
        """
        iid = self._instance_id or self.energy_name
        title = f"Energy Calculation: {iid}"
        super().initialize(self.energy_name, title)
        logger.info(f"Starting energy calculation task: {iid}")




    def pack_settings(self) -> dict:
        """
        Pack the settings for the energy calculation task.
        """
        task_set = {**self._merge_config(self.energy_name, self._active_instance_id)}
        return task_set




    def print_settings(self, task_settings: dict) -> None:
        """
        Print the settings for the energy calculation task.
        """
        super().print_settings(task_settings)




    def create_folders(self) -> None:
        """
        Create the folders for the energy calculation task.
        """
        super().create_folders()




    def perform_task(self, task_settings: dict) -> None:
        """
        Perform the energy calculation task.
        """
        if "energy_settings_path" in task_settings:
            set_file = task_settings["energy_settings_path"]
            with open(set_file) as jfile:
                task_settings["energy_settings"] = json.load(jfile)

        # Create and change dirs
        os.chdir(self.calc_dir)
        dir_name = "rank_" + str(self.rank)
        if not os.path.exists(dir_name):
            os.mkdir(dir_name)
        os.chdir(dir_name)

        self.rank_calc_dir = os.path.join(self.calc_dir, dir_name)
        self._load_checkpoints()

        # Calculate energy
        calc = self.energy_calc(self.comm, task_settings, self.energy_name)
        calc.run_batch(
            self.structs,
            on_structure_done=partial(self.dsdict.checkpoint_save, self.rank_calc_dir),
            done=self.dsdict.done,
        )




    def collect_results(self) -> None:
        """
        Collect the results from the energy calculation task.
        """
        super().collect_results()




    def analyze(self) -> None:
        """
        Analyze the results from the energy calculation task.
        """
        dsdict = DistributedStructs(self.structs)
        stat_dict = dsdict.get_statistics(self.energy_name)
        gout.print_sub_section("Energy Statistics")
        gout.print_dict_table(stat_dict, header=["Stat", "eV"])




    def finalize(self) -> None:
        """
        Finalize the energy calculation task.
        """
        logger.info("Completed energy calculation")
        super().finalize(self.energy_name)