diff --git a/emmet-builders/emmet/builders/vasp/materials.py b/emmet-builders/emmet/builders/vasp/materials.py index 7a2d08301e..b156ae3c5f 100644 --- a/emmet-builders/emmet/builders/vasp/materials.py +++ b/emmet-builders/emmet/builders/vasp/materials.py @@ -13,6 +13,7 @@ # from emmet.core import SETTINGS from emmet.builders import SETTINGS +from emmet.builders.settings import EmmetBuildSettings from emmet.core.utils import group_structures, jsanitize from emmet.core.vasp.calc_types import TaskType from emmet.core.vasp.material import MaterialsDoc @@ -46,12 +47,7 @@ def __init__( materials: Store, task_validation: Optional[Store] = None, query: Optional[Dict] = None, - allowed_task_types: Optional[List[str]] = None, - symprec: float = SETTINGS.SYMPREC, - ltol: float = SETTINGS.LTOL, - stol: float = SETTINGS.STOL, - angle_tol: float = SETTINGS.ANGLE_TOL, - build_tags: List[str] = SETTINGS.BUILD_TAGS, + settings: Optional[EmmetBuildSettings] = None, **kwargs, ): """ @@ -60,30 +56,13 @@ def __init__( materials: Store of materials documents to generate task_validation: Store for storing task validation results query: dictionary to limit tasks to be analyzed - allowed_task_types: list of task_types that can be processed - symprec: tolerance for SPGLib spacegroup finding - ltol: StructureMatcher tuning parameter for matching tasks to materials - stol: StructureMatcher tuning parameter for matching tasks to materials - angle_tol: StructureMatcher tuning parameter for matching tasks to materials """ self.tasks = tasks self.materials = materials self.task_validation = task_validation - self.allowed_task_types = ( - [t.value for t in TaskType] - if allowed_task_types is None - else allowed_task_types - ) - - self._allowed_task_types = {TaskType(t) for t in self.allowed_task_types} - - self.build_tags = build_tags if build_tags else [] self.query = query if query else {} - self.symprec = symprec - self.ltol = ltol - self.stol = stol - self.angle_tol = angle_tol + self.settings = settings or SETTINGS self.kwargs = kwargs sources = [tasks] @@ -124,7 +103,7 @@ def get_items(self) -> Iterator[List[Dict]]: self.logger.info("Materials builder started") self.logger.info( - f"Allowed task types: {[task_type.value for task_type in self._allowed_task_types]}" + f"Allowed task types: {[task_type.value for task_type in self.settings.VASP_ALLOWED_VASP_TYPES]}" ) self.logger.info("Setting indexes") @@ -136,8 +115,13 @@ def get_items(self) -> Iterator[List[Dict]]: # Get all processed tasks: temp_query = dict(self.query) temp_query["state"] = "successful" - if self.build_tags: - temp_query["tags"] = {"$in": self.build_tags} + if len(self.settings.BUILD_TAGS) > 0 and len(self.settings.EXCLUDED_TAGS) > 0: + temp_query["$and"] = [ + {"tags": {"$in": self.settings.BUILD_TAGS}}, + {"tags": {"$nin": self.settings.EXCLUDED_TAGS}}, + ] + elif len(self.settings.BUILD_TAGS) > 0: + temp_query["tags"] = {"$in": self.settings.BUILD_TAGS} self.logger.info("Finding tasks to process") all_tasks = { @@ -267,7 +251,7 @@ def filter_and_group_tasks(self, tasks: List[TaskDocument]) -> Iterator[List[Dic for task in tasks if any( allowed_type is task.task_type - for allowed_type in self._allowed_task_types + for allowed_type in self.settings.VASP_ALLOWED_VASP_TYPES ) ] @@ -280,10 +264,10 @@ def filter_and_group_tasks(self, tasks: List[TaskDocument]) -> Iterator[List[Dic grouped_structures = group_structures( structures, - ltol=self.ltol, - stol=self.stol, - angle_tol=self.angle_tol, - symprec=self.symprec, + ltol=self.settings.LTOL, + stol=self.settings.STOL, + angle_tol=self.settings.ANGLE_TOL, + symprec=self.settings.SYMPREC, ) for group in grouped_structures: grouped_tasks = [filtered_tasks[struc.index] for struc in group] diff --git a/emmet-builders/emmet/builders/vasp/task_validator.py b/emmet-builders/emmet/builders/vasp/task_validator.py index 66792200d4..a2ab69b6c7 100644 --- a/emmet-builders/emmet/builders/vasp/task_validator.py +++ b/emmet-builders/emmet/builders/vasp/task_validator.py @@ -1,9 +1,9 @@ -from typing import Dict, List, Union +from typing import Dict, List, Union, Optional import numpy as np from maggma.builders import MapBuilder from maggma.core import Store -from pymatgen import Structure +from pymatgen.core import Structure from emmet.builders import SETTINGS from emmet.core.vasp.calc_types import run_type, task_type @@ -20,6 +20,7 @@ def __init__( self, tasks: Store, task_validation: Store, + settings: Optional[EmmetBuildSettings] = None, **kwargs, ): """ @@ -31,8 +32,7 @@ def __init__( """ self.tasks = tasks self.task_validation = task_validation - self.settings = SETTINGS - + self.settings = settings or SETTINGS self.kwargs = kwargs super().__init__(