"""
This module provides the GeometryOptimizationTask class for performing geometry optimization tasks.
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.registry import resolve_tasks
from gnrs.core.task import TaskABC
from gnrs.gnrsutil.molecule_bonding import get_vdw_distance_cutoff_matrix
from gnrs.parallel.io import read_parallel
from gnrs.parallel.structs import DistributedStructs
AVAILABLE_METHODS = ["LBFGS", "BFGS", "RIGID_PRESS", "SYMM_RIGID_PRESS"]
AVAILABLE_ENERGY_METHODS = ["DFTB", "AIMS", "MACEOFF", "UMA", "VASP", "AIMNET"]
logger = logging.getLogger("GeoOptTask")
[docs]
class GeometryOptimizationTask(TaskABC):
"""
Task for relaxing geometry of crystal structures.
"""
[docs]
def __init__(
self,
comm: MPI.Comm,
config: dict,
gnrs_info: dict,
optimizer: str,
energy_method: str | None = None,
instance_id: str | None = None,
) -> None:
"""
Initialize the geometry optimization task.
Args:
comm: MPI communicator
config: Config dictionary
gnrs_info: Genarris info dictionary
optimizer: Optimizer
energy_method: Energy calculator(optional)
instance_id: Unique ID for this task instance
"""
super().__init__(comm, config, gnrs_info, instance_id=instance_id)
self.opt_name = optimizer.lower()
self.opt_class = f"{self.opt_name.upper()}Optimizer"
# Set task name and energy method
if energy_method is not None:
self.energy_method = energy_method.lower()
self.energy_class = f"{self.energy_method.upper()}Energy"
self.task_name = f"{self.opt_name}_{self.energy_method}"
else:
self.energy_method = None
self.task_name = self.opt_name
self.energy_set = {}
self.structs = None
self.dsdict = None
self.rank_calc_dir = None
self.opt_calc = None
self.energy_calc = None
[docs]
def initialize(self) -> None:
"""
Initialize the optimization task.
"""
iid = self._instance_id or self.task_name
title = f"Geometry Optimization: {iid}"
super().initialize(self.task_name, title)
logger.info(f"Starting geometry optimization task: {iid}")
# If struct_path specified, use it instead of default
spath = self.config[self.opt_name].get("struct_path")
if spath is not None and self._reads_struct_path():
logger.info(f"Reading from user given file {spath}")
self.structs = read_parallel(spath)
# Log the optimizer being used
if self.energy_method is not None:
gout.emit(
f"Using ASE {self.opt_name} optimizer with {self.energy_method} energy method."
)
else:
gout.emit("Using builtin optimizer.")
self._load_modules()
def _reads_struct_path(self) -> bool:
"""
Whether this task takes its input from ``[<optimizer>] struct_path``.
Only the first task in the workflow that uses this optimizer does;
later ones continue from the previous task's output. Decided from the
task list rather than by consuming the setting, so it holds on a
restart that skips the first task.
Returns:
True if this task reads ``struct_path``.
"""
specs = resolve_tasks(self.config.get("workflow", {}).get("tasks", []))
first = next(s.instance_id for s in specs if self.opt_name in s.sections)
return first == self._active_instance_id
def _load_modules(self) -> None:
"""
Load the required optimizer and energy calculator modules.
Raises:
ImportError: If the requested optimization method or energy calculator is not found.
"""
# Check if optimization method is implemented
self.opt_file = f"gnrs.optimize.{self.opt_name}"
try:
opt_module = importlib.import_module(self.opt_file)
self.opt_calc = getattr(opt_module, self.opt_class)
except (ImportError, AttributeError):
logger.error("Unable to find requested optimization method.")
logger.error(f"Available methods: {AVAILABLE_METHODS}")
raise
# If the builtin optimizer is used instead of ASE
if self.energy_method is None:
return
# Check if energy calculator is implemented within gnrs
self.energy_file = f"gnrs.energy.{self.energy_method}"
try:
energy_module = importlib.import_module(self.energy_file)
self.energy_calc = getattr(energy_module, self.energy_class)
except (ImportError, AttributeError):
logger.error("Unable to find requested energy calculation method.")
logger.error(f"Available methods: {AVAILABLE_ENERGY_METHODS}")
raise
[docs]
def pack_settings(self) -> dict:
"""
Pack settings for the optimization task.
Returns:
dict: Task settings dictionary
"""
iid = self._active_instance_id
overrides = self.config.get(iid, {}) if iid != self.task_name else {}
task_set = {**self.config.get(self.opt_name, {}), **overrides}
task_set.pop("struct_path", None)
if self.opt_name in ["rigid_press", "symm_rigid_press"]:
task_set["z"] = self.config["master"]["z"]
task_set["mol_path"] = self.gnrs_info["molecule_path"][0]
sr = task_set.pop("sr")
cutoff_mult = task_set.pop("natural_cutoff_mult")
cutoff_matrix, hbond = get_vdw_distance_cutoff_matrix(
self.gnrs_info["molecule_path"], task_set["z"], sr, cutoff_mult
)
task_set["cutoff_matrix"] = cutoff_matrix
# Pack settings for energy method separately in self.energy_set
if self.energy_method is not None:
energy_method = task_set.pop("energy_method")
self.energy_set = dict(self.config[energy_method])
return task_set
[docs]
def print_settings(self, task_set: dict) -> None:
"""
Print the task settings.
Args:
task_set: Task settings dictionary
"""
gout.emit("Optimization Settings:")
gout.print_dict_table(task_set, ["Option", "Value"], skip=("cutoff_matrix"))
if self.energy_method is not None:
gout.emit("Energy Settings:")
super().print_settings(self.energy_set)
[docs]
def create_folders(self) -> None:
"""
Create the necessary folders for the optimization task.
"""
super().create_folders()
[docs]
def collect_results(self):
super().collect_results()
[docs]
def analyze(self) -> None:
"""
Analyze the results of the optimization task.
"""
dsdict = DistributedStructs(self.structs)
vol_stat = dsdict.get_statistics("get_volume", ptype="method")
gout.print_sub_section("Unit Cell Volume Statistics")
gout.print_dict_table(vol_stat, header=["Stat", "Volume (A^3)"])
if self.energy_method is not None:
energy_stat = dsdict.get_statistics(f"{self.opt_name}_{self.energy_method}")
gout.print_sub_section("Energy Statistics")
gout.print_dict_table(energy_stat, header=["Stat", "Energy (eV)"])
[docs]
def finalize(self) -> None:
"""
Finalize the optimization task.
"""
logger.info("Completed optimization task")
if self.energy_method is not None:
self.gnrs_info["energy_list"].append(self.energy_method)
super().finalize(self.task_name)