diff --git a/psiflow/data/dataset.py b/psiflow/data/dataset.py index 65b869a..d0fa168 100644 --- a/psiflow/data/dataset.py +++ b/psiflow/data/dataset.py @@ -43,6 +43,7 @@ class Dataset: This class provides methods for manipulating and analyzing collections of atomic structures. """ + extxyz: psiflow._DataFuture def __init__( @@ -415,6 +416,7 @@ def _concatenate_multiple(*args: list[np.ndarray]) -> list[np.ndarray]: Note: This function is wrapped as a Parsl app and executed using the default_threads executor. """ + def pad_arrays( arrays: list[np.ndarray], pad_dimension: int = 1, @@ -621,6 +623,7 @@ class Computable: outputs (ClassVar[tuple[str, ...]]): Names of output quantities. batch_size (ClassVar[Optional[int]]): Default batch size for computation. """ + outputs: ClassVar[tuple[str, ...]] = () batch_size: ClassVar[Optional[int]] = None diff --git a/psiflow/execution.py b/psiflow/execution.py index 73b6c5b..1b9a529 100644 --- a/psiflow/execution.py +++ b/psiflow/execution.py @@ -131,8 +131,8 @@ def from_config( provider_cls = SlurmProvider provider_kwargs = kwargs.pop("slurm") # do not allow empty dict provider_kwargs["init_blocks"] = 0 - if not 'exclusive' in provider_kwargs: - provider_kwargs['exclusive'] = False + if "exclusive" not in provider_kwargs: + provider_kwargs["exclusive"] = False else: provider_cls = LocalProvider # noqa: F405 provider_kwargs = kwargs.pop("local", {}) @@ -452,7 +452,7 @@ def from_config( max_idletime: float = 20, internal_tasks_max_threads: int = 10, default_threads: int = 4, - htex_address: str = '127.0.0.1', + htex_address: str = "127.0.0.1", zip_staging: Optional[bool] = None, container_uri: Optional[str] = None, container_engine: str = "apptainer", @@ -552,11 +552,14 @@ def from_config( context = ExecutionContext(config, definitions, path / "context_dir") if make_symlinks: - src, dest = Path.cwd() / f'psiflow_log', path / 'parsl.log' + src, dest = Path.cwd() / "psiflow_log", path / "parsl.log" _create_symlink(src, dest) - src, dest = Path.cwd() / f'psiflow_submit_scripts', path / '000' / 'submit_scripts' + src, dest = ( + Path.cwd() / "psiflow_submit_scripts", + path / "000" / "submit_scripts", + ) _create_symlink(src, dest, is_dir=True) - src, dest = Path.cwd() / f'psiflow_task_logs', path / '000' / 'task_logs' + src, dest = Path.cwd() / "psiflow_task_logs", path / "000" / "task_logs" _create_symlink(src, dest, is_dir=True) return context @@ -684,5 +687,3 @@ def _create_symlink(src: Path, dest: Path, is_dir: bool = False) -> None: else: dest.touch(exist_ok=True) src.symlink_to(dest, target_is_directory=is_dir) - - diff --git a/psiflow/geometry.py b/psiflow/geometry.py index e930dfc..791bea6 100644 --- a/psiflow/geometry.py +++ b/psiflow/geometry.py @@ -274,26 +274,30 @@ def from_string(cls, s: str, natoms: Optional[int] = None) -> Optional[Geometry] # read and format per_atom data column_indices = {} - if 'Properties' in comment_dict: - properties = comment_dict['Properties'].split(':') + if "Properties" in comment_dict: + properties = comment_dict["Properties"].split(":") count = 0 for i in range(len(properties) // 3): name = properties[3 * i] ncolumns = int(properties[3 * i + 2]) column_indices[name] = count count += ncolumns - assert 'pos' in column_indices # positions need to be there + assert "pos" in column_indices # positions need to be there per_atom = np.recarray(natoms, dtype=per_atom_dtype) per_atom.forces[:] = np.nan - POS_INDEX = column_indices.get('pos', 1) - FORCES_INDEX = column_indices.get('forces', None) + POS_INDEX = column_indices.get("pos", 1) + FORCES_INDEX = column_indices.get("forces", None) for i in range(natoms): values = lines[i + 1].split() per_atom.numbers[i] = chemical_symbols.index(values[0]) - per_atom.positions[i, :] = [float(_) for _ in values[POS_INDEX:POS_INDEX + 3]] + per_atom.positions[i, :] = [ + float(_) for _ in values[POS_INDEX : POS_INDEX + 3] + ] if FORCES_INDEX is not None: - per_atom.forces[i, :] = [float(_) for _ in values[FORCES_INDEX:FORCES_INDEX + 3]] + per_atom.forces[i, :] = [ + float(_) for _ in values[FORCES_INDEX : FORCES_INDEX + 3] + ] order = {} for key, value in comment_dict.items(): diff --git a/psiflow/learning.py b/psiflow/learning.py index a9e867d..f1e72e3 100644 --- a/psiflow/learning.py +++ b/psiflow/learning.py @@ -17,7 +17,7 @@ from psiflow.models import Model from psiflow.reference import Reference, evaluate from psiflow.sampling import SimulationOutput, Walker, sample -from psiflow.utils.apps import boolean_or, setup_logger, unpack_i, isnan +from psiflow.utils.apps import boolean_or, isnan, setup_logger, unpack_i logger = setup_logger(__name__) diff --git a/psiflow/reference/_cp2k.py b/psiflow/reference/_cp2k.py index cd84130..8b3e8cf 100644 --- a/psiflow/reference/_cp2k.py +++ b/psiflow/reference/_cp2k.py @@ -68,9 +68,9 @@ def set_global_section(cp2k_input_dict: dict, properties: tuple): global_dict = cp2k_input_dict["global"] # override low/silent print levels - level = global_dict.pop('print_level', 'MEDIUM') - if level in ['SILENT', 'LOW']: - global_dict['print_level'] = 'MEDIUM' + level = global_dict.pop("print_level", "MEDIUM") + if level in ["SILENT", "LOW"]: + global_dict["print_level"] = "MEDIUM" if properties == ("energy",): global_dict["run_type"] = "ENERGY" @@ -156,11 +156,11 @@ def _prepare_input( if "forces" in properties: cp2k_input_dict["force_eval"]["print"] = {"FORCES": {}} cp2k_input_str = dict_to_str(cp2k_input_dict) - with open(outputs[0], 'w') as f: + with open(outputs[0], "w") as f: f.write(cp2k_input_str) -prepare_input = python_app(_prepare_input, executors=['default_threads']) +prepare_input = python_app(_prepare_input, executors=["default_threads"]) # typeguarding for some reason incompatible with WQ @@ -175,14 +175,9 @@ def cp2k_singlepoint_pre( cd_command = "cd $mytmpdir" cp_command = "cp {} cp2k.inp".format(inputs[0].filepath) - command_list = [ - tmp_command, - cd_command, - cp_command, - cp2k_command - ] + command_list = [tmp_command, cd_command, cp_command, cp2k_command] - return ' && '.join(command_list) + return " && ".join(command_list) @typeguard.typechecked @@ -242,7 +237,7 @@ def wrapped_app_pre(geometry, stdout: str, stderr: str): geometry, cp2k_input_dict=self.cp2k_input_dict, properties=tuple(self.outputs), - outputs=[psiflow.context().new_file('cp2k_', '.inp')], + outputs=[psiflow.context().new_file("cp2k_", ".inp")], ) return app_pre( cp2k_command=cp2k_command, diff --git a/psiflow/reference/reference.py b/psiflow/reference/reference.py index 8776faf..8c9d83c 100644 --- a/psiflow/reference/reference.py +++ b/psiflow/reference/reference.py @@ -55,7 +55,7 @@ def _nan_if_unsuccessful( return result -nan_if_unsuccessful = python_app(_nan_if_unsuccessful, executors=['default_threads']) +nan_if_unsuccessful = python_app(_nan_if_unsuccessful, executors=["default_threads"]) @join_app diff --git a/psiflow/utils/apps.py b/psiflow/utils/apps.py index 7bf807d..74c9c84 100644 --- a/psiflow/utils/apps.py +++ b/psiflow/utils/apps.py @@ -141,4 +141,4 @@ def _isnan(a: Union[float, np.ndarray]) -> bool: return bool(np.any(np.isnan(a))) -isnan = python_app(_isnan, executors=['default_threads']) +isnan = python_app(_isnan, executors=["default_threads"])