diff --git a/janus_core/calculations/base.py b/janus_core/calculations/base.py index bc3fccdfb..5c5f2d037 100644 --- a/janus_core/calculations/base.py +++ b/janus_core/calculations/base.py @@ -30,17 +30,17 @@ class BaseCalculation(FileNameMixin): struct_path : Optional[PathLike] Path of structure to simulate. Required if `struct` is None. Default is None. - arch : Architectures + arch : MaybeSequence[Architectures] MLIP architecture to use for calculations. Default is "mace_mp". - device : Devices + device : MaybeSequence[Devices] Device to run model on. Default is "cpu". - model_path : Optional[PathLike] + model_path : Optional[MaybeSequence[PathLike]] Path to MLIP model. Default is `None`. read_kwargs : ASEReadArgs Keyword arguments to pass to ase.io.read. Default is {}. sequence_allowed : bool Whether a sequence of Atoms objects is allowed. Default is True. - calc_kwargs : Optional[dict[str, Any]] + calc_kwargs : Optional[MaybeSequence[dict[str, Any]]] Keyword arguments to pass to the selected calculator. Default is {}. set_calc : Optional[bool] Whether to set (new) calculators for structures. Default is None. @@ -73,12 +73,12 @@ def __init__( calc_name: str = "base", struct: Optional[MaybeSequence[Atoms]] = None, struct_path: Optional[PathLike] = None, - arch: Architectures = "mace_mp", - device: Devices = "cpu", - model_path: Optional[PathLike] = None, + arch: MaybeSequence[Architectures] = "mace_mp", + device: MaybeSequence[Devices] = "cpu", + model_path: Optional[MaybeSequence[PathLike]] = None, read_kwargs: Optional[ASEReadArgs] = None, sequence_allowed: bool = True, - calc_kwargs: Optional[dict[str, Any]] = None, + calc_kwargs: Optional[MaybeSequence[dict[str, Any]]] = None, set_calc: Optional[bool] = None, attach_logger: bool = False, log_kwargs: Optional[dict[str, Any]] = None, @@ -101,17 +101,17 @@ def __init__( struct_path : Optional[PathLike] Path of structure to simulate. Required if `struct` is None. Default is None. - arch : Architectures + arch : MaybeSequence[Architectures] MLIP architecture to use for calculations. Default is "mace_mp". - device : Devices + device : MaybeSequence[Devices] Device to run MLIP model on. Default is "cpu". - model_path : Optional[PathLike] + model_path : Optional[MaybeSequence[PathLike]] Path to MLIP model. Default is `None`. read_kwargs : Optional[ASEReadArgs] Keyword arguments to pass to ase.io.read. Default is {}. sequence_allowed : bool Whether a sequence of Atoms objects is allowed. Default is True. - calc_kwargs : Optional[dict[str, Any]] + calc_kwargs : Optional[MaybeSequence[dict[str, Any]]] Keyword arguments to pass to the selected calculator. Default is {}. set_calc : Optional[bool] Whether to set (new) calculators for structures. Default is None. diff --git a/janus_core/calculations/single_point.py b/janus_core/calculations/single_point.py index a4171f547..1079aadf2 100644 --- a/janus_core/calculations/single_point.py +++ b/janus_core/calculations/single_point.py @@ -35,12 +35,12 @@ class SinglePoint(BaseCalculation): struct_path : Optional[PathLike] Path of structure to simulate. Required if `struct` is None. Default is None. - arch : Architectures + arch : MaybeSequence[Architectures] MLIP architecture to use for single point calculations. Default is "mace_mp". - device : Devices + device : MaybeSequence[Devices] Device to run model on. Default is "cpu". - model_path : Optional[PathLike] + model_path : Optional[MaybeSequence[PathLike]] Path to MLIP model. Default is `None`. read_kwargs : ASEReadArgs Keyword arguments to pass to ase.io.read. By default, @@ -82,9 +82,9 @@ def __init__( *, struct: Optional[MaybeSequence[Atoms]] = None, struct_path: Optional[PathLike] = None, - arch: Architectures = "mace_mp", - device: Devices = "cpu", - model_path: Optional[PathLike] = None, + arch: MaybeSequence[Architectures] = "mace_mp", + device: MaybeSequence[Devices] = "cpu", + model_path: Optional[MaybeSequence[PathLike]] = None, read_kwargs: Optional[ASEReadArgs] = None, calc_kwargs: Optional[dict[str, Any]] = None, set_calc: Optional[bool] = None, @@ -107,12 +107,12 @@ def __init__( struct_path : Optional[PathLike] Path of structure to simulate. Required if `struct` is None. Default is None. - arch : Architectures + arch : MaybeSequence[Architectures] MLIP architecture to use for single point calculations. Default is "mace_mp". - device : Devices + device : MaybeSequence[Devices] Device to run MLIP model on. Default is "cpu". - model_path : Optional[PathLike] + model_path : Optional[MaybeSequence[PathLike]] Path to MLIP model. Default is `None`. read_kwargs : Optional[ASEReadArgs] Keyword arguments to pass to ase.io.read. By default, diff --git a/janus_core/cli/singlepoint.py b/janus_core/cli/singlepoint.py index ead77ed76..527e32ff8 100644 --- a/janus_core/cli/singlepoint.py +++ b/janus_core/cli/singlepoint.py @@ -7,11 +7,11 @@ from typer_config import use_config from janus_core.cli.types import ( - Architecture, - CalcKwargs, - Device, + ArchitectureList, + CalcKwargsList, + DeviceList, LogPath, - ModelPath, + ModelPathList, ReadKwargsAll, StructPath, Summary, @@ -28,9 +28,9 @@ def singlepoint( # numpydoc ignore=PR02 ctx: Context, struct: StructPath, - arch: Architecture = "mace_mp", - device: Device = "cpu", - model_path: ModelPath = None, + arch: ArchitectureList = ("mace_mp",), + device: DeviceList = ("cpu",), + model_path: ModelPathList = None, properties: Annotated[ Optional[list[str]], Option( @@ -50,7 +50,7 @@ def singlepoint( ), ] = None, read_kwargs: ReadKwargsAll = None, - calc_kwargs: CalcKwargs = None, + calc_kwargs: CalcKwargsList = None, write_kwargs: WriteKwargs = None, log: LogPath = None, tracker: Annotated[ diff --git a/janus_core/cli/types.py b/janus_core/cli/types.py index 31143a941..3cecb216b 100644 --- a/janus_core/cli/types.py +++ b/janus_core/cli/types.py @@ -72,6 +72,14 @@ def __str__(self) -> str: Device = Annotated[Optional[str], Option(help="Device to run calculations on.")] ModelPath = Annotated[Optional[str], Option(help="Path to MLIP model.")] +ArchitectureList = Annotated[ + Optional[list[str]], Option(help="MLIP architecture to use for calculations.") +] +DeviceList = Annotated[ + Optional[list[str]], Option(help="Device to run calculations on.") +] +ModelPathList = Annotated[Optional[list[str]], Option(help="Path to MLIP model.")] + ReadKwargsAll = Annotated[ Optional[TyperDict], Option( @@ -117,6 +125,21 @@ def __str__(self) -> str: ), ] +CalcKwargsList = Annotated[ + Optional[list[TyperDict]], + Option( + parser=parse_dict_class, + help=( + """ + Keyword arguments to pass to selected calculator. Must be passed as a + dictionary wrapped in quotes, e.g. "{'key' : value}". For the default + architecture ('mace_mp'), "{'model':'small'}" is set unless overwritten. + """ + ), + metavar="DICT", + ), +] + WriteKwargs = Annotated[ Optional[TyperDict], Option( diff --git a/janus_core/helpers/mlip_calculators.py b/janus_core/helpers/mlip_calculators.py index e8b48020a..ded239b1a 100644 --- a/janus_core/helpers/mlip_calculators.py +++ b/janus_core/helpers/mlip_calculators.py @@ -68,6 +68,47 @@ def _set_model_path( return model_path +def _set_torch(calculate: callable, dtype: torch.dtype): + """ + Wrap calculate function to set torch default dtype before calculations. + + Parameters + ---------- + calculate : callable + Function to wrap. + dtype : torch.dtype + Default dtype to set. + + Returns + ------- + callable + Wrapped function. + """ + + def wrapper(*args, **kwargs) -> callable: + """ + Wrap function to set torch default dtype. + + Parameters + ---------- + *args + Arguments passed to calculate. + **kwargs + Additional keyword arguments passed to calculate. + + Returns + ------- + callable + Wrapped function. + """ + import torch + + torch.set_default_dtype(dtype) + return calculate(*args, **kwargs) + + return wrapper + + def choose_calculator( arch: Architectures = "mace", device: Devices = "cpu", @@ -123,12 +164,14 @@ def choose_calculator( elif arch == "mace_mp": from mace import __version__ from mace.calculators import mace_mp + import torch # Default to "small" model and float64 precision model = model_path if model_path else "small" kwargs.setdefault("default_dtype", "float64") calculator = mace_mp(model=model, device=device, **kwargs) + calculator.calculate = _set_torch(calculator.calculate, torch.float64) elif arch == "mace_off": from mace import __version__ @@ -164,6 +207,7 @@ def choose_calculator( potential = load_model("M3GNet-MP-2021.2.8-DIRECT-PES") calculator = M3GNetCalculator(potential=potential, **kwargs) + calculator.calculate = _set_torch(calculator.calculate, torch.float32) elif arch == "chgnet": from chgnet import __version__ @@ -186,6 +230,7 @@ def choose_calculator( model = None calculator = CHGNetCalculator(model=model, use_device=device, **kwargs) + calculator.calculate = _set_torch(calculator.calculate, torch.float32) elif arch == "alignn": from alignn import __version__ diff --git a/janus_core/helpers/multi_calc.py b/janus_core/helpers/multi_calc.py new file mode 100644 index 000000000..ccd2fa3a4 --- /dev/null +++ b/janus_core/helpers/multi_calc.py @@ -0,0 +1,113 @@ +"""Define MultiCalc ASE Calculator.""" + +from collections.abc import Sequence +from typing import Any + +from ase.calculators.calculator import ( + BaseCalculator, + Calculator, + CalculatorSetupError, + PropertyNotImplementedError, +) + + +class MultiCalc(BaseCalculator): + """ + ASE MultiCalc class. + + Parameters + ---------- + calcs : Sequence[Calculator] + Calculators to use. + """ + + def __init__(self, calcs: Sequence[Calculator]): + """ + Initialise class. + + Parameters + ---------- + calcs : Sequence[Calculator] + Calculators to use. + """ + super().__init__() + + if len(calcs) == 0: + raise CalculatorSetupError("Please provide a list of Calculators") + + common_properties = set.intersection( + *(set(calc.implemented_properties) for calc in calcs) + ) + + self.implemented_properties = list(common_properties) + if not self.implemented_properties: + raise PropertyNotImplementedError( + "The provided Calculators have" " no properties in common!" + ) + + self.calcs = calcs + + def __str__(self) -> str: + """ + Return string representation of the calculator. + + Returns + ------- + str + String representation. + """ + calcs = ", ".join(calc.__class__.__name__ for calc in self.calcs) + return f"{self.__class__.__name__}({calcs})" + + def get_properties(self, properties, atoms) -> dict[str, Any]: + """ + Get properties from each listed calculator. + + Parameters + ---------- + properties : list[str] + List of properties to be calculated. + atoms : Atoms + Atoms object to calculate properties for. + + Returns + ------- + dict + Dictionary of results. + """ + results = {} + + def get_property(prop: str) -> None: + """ + Get property from each listed calculator. + + Parameters + ---------- + prop : str + Property to get. + """ + contribs = [calc.get_property(prop, atoms) for calc in self.calcs] + results[prop] = contribs + + for prop in properties: # get requested properties + get_property(prop) + for prop in self.implemented_properties: # cache all available props + if all(prop in calc.results for calc in self.calcs): + get_property(prop) + return results + + def calculate(self, atoms, properties, system_changes) -> None: + """ + Calculate properties for each calculator and return values as list. + + Parameters + ---------- + atoms : Atoms + Atoms object to calculate properties for. + properties : list[str] + List of properties to be calculated. + system_changes : list[str] + List of what has changed since last calculation. + """ + self.atoms = atoms.copy() # for caching of results + self.results = self.get_properties(properties, atoms) diff --git a/janus_core/helpers/struct_io.py b/janus_core/helpers/struct_io.py index 87e4049f2..5b541b976 100644 --- a/janus_core/helpers/struct_io.py +++ b/janus_core/helpers/struct_io.py @@ -20,6 +20,8 @@ Properties, ) from janus_core.helpers.mlip_calculators import choose_calculator +from janus_core.helpers.multi_calc import MultiCalc +from janus_core.helpers.utils import param_to_sequence def results_to_info( @@ -65,10 +67,10 @@ def results_to_info( def attach_calculator( struct: MaybeSequence[Atoms], *, - arch: Architectures = "mace_mp", - device: Devices = "cpu", - model_path: Optional[PathLike] = None, - calc_kwargs: Optional[dict[str, Any]] = None, + arch: MaybeSequence[Architectures] = "mace_mp", + device: MaybeSequence[Devices] = "cpu", + model_path: Optional[MaybeSequence[PathLike]] = None, + calc_kwargs: Optional[MaybeSequence[dict[str, Any]]] = None, ) -> None: """ Configure calculator and attach to structure(s). @@ -77,29 +79,56 @@ def attach_calculator( ---------- struct : Optional[MaybeSequence[Atoms]] ASE Atoms structure(s) to attach calculators to. - arch : Architectures + arch : MaybeSequence[Architectures] MLIP architecture to use for calculations. Default is "mace_mp". - device : Devices + device : MaybeSequence[Devices] Device to run model on. Default is "cpu". - model_path : Optional[PathLike] + model_path : Optional[MaybeSequence[PathLike]] Path to MLIP model. Default is `None`. - calc_kwargs : Optional[dict[str, Any]] + calc_kwargs : Optional[MaybeSequence[dict[str, Any]]] Keyword arguments to pass to the selected calculator. Default is {}. """ calc_kwargs = calc_kwargs if calc_kwargs else {} - calculator = choose_calculator( - arch=arch, - device=device, - model_path=model_path, - **calc_kwargs, - ) + # Convert parameters to sequence + arch = param_to_sequence(arch) + device = param_to_sequence(device) + model_path = param_to_sequence(model_path) + calc_kwargs = param_to_sequence(calc_kwargs) + + calculators = [] + for i, _arch in enumerate(arch): + if i < len(device): + _device = device[i] + else: + _device = "cpu" + + if i < len(model_path): + _model_path = model_path[i] + else: + _model_path = None + if i < len(calc_kwargs): + _calc_kwargs = calc_kwargs[i] + else: + _calc_kwargs = {} + calculators.append( + choose_calculator( + arch=_arch, + device=_device, + model_path=_model_path, + **_calc_kwargs, + ) + ) + if len(calculators) == 1: + calc = calculators[0] + else: + calc = MultiCalc(calculators) if isinstance(struct, Sequence): for image in struct: - image.calc = copy(calculator) + image.calc = copy(calc) else: - struct.calc = calculator + struct.calc = calc def input_structs( @@ -108,10 +137,10 @@ def input_structs( struct_path: Optional[PathLike] = None, read_kwargs: Optional[ASEReadArgs] = None, sequence_allowed: bool = True, - arch: Architectures = "mace_mp", - device: Devices = "cpu", - model_path: Optional[PathLike] = None, - calc_kwargs: Optional[dict[str, Any]] = None, + arch: MaybeSequence[Architectures] = "mace_mp", + device: MaybeSequence[Devices] = "cpu", + model_path: Optional[MaybeSequence[PathLike]] = None, + calc_kwargs: Optional[MaybeSequence[dict[str, Any]]] = None, set_calc: Optional[bool] = None, logger: Optional[logging.Logger] = None, ) -> MaybeSequence[Atoms]: @@ -129,13 +158,13 @@ def input_structs( Keyword arguments to pass to ase.io.read. Default is {}. sequence_allowed : bool Whether a sequence of Atoms objects is allowed. Default is True. - arch : Architectures + arch : MaybeSequence[Architectures] MLIP architecture to use for calculations. Default is "mace_mp". - device : Devices + device : MaybeSequence[Devices] Device to run model on. Default is "cpu". - model_path : Optional[PathLike] + model_path : Optional[MaybeSequence[PathLike]] Path to MLIP model. Default is `None`. - calc_kwargs : Optional[dict[str, Any]] + calc_kwargs : Optional[MaybeSequence[dict[str, Any]]] Keyword arguments to pass to the selected calculator. Default is {}. set_calc : Optional[bool] Whether to set (new) calculators for structures. Default is True if diff --git a/janus_core/helpers/utils.py b/janus_core/helpers/utils.py index 085572579..fd60d413d 100644 --- a/janus_core/helpers/utils.py +++ b/janus_core/helpers/utils.py @@ -166,6 +166,25 @@ def none_to_dict(dictionaries: Sequence[Optional[dict]]) -> Generator[dict, None yield from (dictionary if dictionary else {} for dictionary in dictionaries) +def param_to_sequence(param: Any) -> Sequence[Any]: + """ + Convert parameter to sequence. + + Parameters + ---------- + param : Any + Parameter to be converted. + + Returns + ------- + Sequence[Any] + Parameter as sequence. + """ + if param is None or param == {} or isinstance(param, (str, Path)): + param = (param,) + return param + + def write_table( fmt: Literal["ascii", "csv"], file: Optional[TextIO] = None,