From 695d53ed838dfc8e0c8b14e2e44857876334e290 Mon Sep 17 00:00:00 2001 From: Sander Vandenhaute Date: Sun, 28 Jul 2024 05:14:36 -0400 Subject: [PATCH] linting --- psiflow/data/utils.py | 7 ++++--- psiflow/execution.py | 39 +++++++++++++++++++----------------- psiflow/functions.py | 4 ++-- psiflow/metrics.py | 3 +-- psiflow/models/_mace.py | 2 +- psiflow/models/mace_utils.py | 3 ++- psiflow/reference/_cp2k.py | 4 ++-- psiflow/sampling/client.py | 2 +- psiflow/sampling/optimize.py | 2 +- psiflow/sampling/sampling.py | 4 ++-- psiflow/sampling/server.py | 13 +++++++----- tests/conftest.py | 2 +- tests/test_data.py | 2 +- 13 files changed, 47 insertions(+), 40 deletions(-) diff --git a/psiflow/data/utils.py b/psiflow/data/utils.py index 25a1458..80ac9f7 100644 --- a/psiflow/data/utils.py +++ b/psiflow/data/utils.py @@ -26,7 +26,7 @@ def _write_frames( all_states.append(extra_states) with open(outputs[0], "w") as f: for state in all_states: # avoid double newline by using strip! - f.write(state.to_string().strip() + '\n') + f.write(state.to_string().strip() + "\n") write_frames = python_app(_write_frames, executors=["default_threads"]) @@ -605,8 +605,9 @@ def _batch_frames( else: # write and clear pass if len(data) > 0: - with open(outputs[batch_index], "w") as g: - g.write("\n".join(data)) + with open(outputs[batch_index], "w") as f: + f.write("\n".join([d.strip() for d in data if d is not None])) + f.write("\n") batch_index += 1 assert batch_index == len(outputs) diff --git a/psiflow/execution.py b/psiflow/execution.py index 4cdc1c2..d2dca24 100644 --- a/psiflow/execution.py +++ b/psiflow/execution.py @@ -1,9 +1,8 @@ from __future__ import annotations # necessary for type-guarding class methods -import re - import logging import math +import re import shutil import sys from pathlib import Path @@ -97,7 +96,7 @@ def create_executor(self, path: Path, **kwargs) -> ParslExecutor: worker_options.append("--gpus={}".format(self.max_workers)) # ensure proper scale in - if getattr(self.parsl_provider, 'nodes_per_block', 1) > 1: + if getattr(self.parsl_provider, "nodes_per_block", 1) > 1: worker_options.append("--idle-timeout={}".format(20)) else: worker_options.append("--idle-timeout={}".format(int(1e6))) @@ -180,10 +179,10 @@ def __init__( self.timeout = timeout default_env_vars = { - 'OMP_NUM_THREADS': str(self.cores_per_worker), - 'KMP_AFFINITY': 'granularity=fine,compact,1,0', - 'OMP_PROC_BIND': 'false', - 'PYTHONUNBUFFERED': 'TRUE', + "OMP_NUM_THREADS": str(self.cores_per_worker), + "KMP_AFFINITY": "granularity=fine,compact,1,0", + "OMP_PROC_BIND": "false", + "PYTHONUNBUFFERED": "TRUE", } if env_vars is None: env_vars = default_env_vars @@ -254,17 +253,19 @@ def __init__( assert max_training_time * 60 < self.max_runtime self.max_training_time = max_training_time if self.max_workers > 1: - message = ('the max_training_time keyword does not work ' - 'in combination with multi-gpu training. Adjust ' - 'the maximum number of epochs to control the ' - 'duration of training') + message = ( + "the max_training_time keyword does not work " + "in combination with multi-gpu training. Adjust " + "the maximum number of epochs to control the " + "duration of training" + ) assert self.max_training_time is None, message default_env_vars = { - 'OMP_NUM_THREADS': str(self.cores_per_worker), - 'KMP_AFFINITY': 'granularity=fine,compact,1,0', - 'OMP_PROC_BIND': 'false', - 'PYTHONUNBUFFERED': 'TRUE', + "OMP_NUM_THREADS": str(self.cores_per_worker), + "KMP_AFFINITY": "granularity=fine,compact,1,0", + "OMP_PROC_BIND": "false", + "PYTHONUNBUFFERED": "TRUE", } if env_vars is None: env_vars = default_env_vars @@ -293,7 +294,9 @@ def wq_resources(self): resource_specification["gpus"] = nworkers # one per GPU resource_specification["cores"] = self.cores_available - resource_specification["disk"] = 1000 * nworkers # some random nontrivial amount? + resource_specification["disk"] = ( + 1000 * nworkers + ) # some random nontrivial amount? memory = 1000 * self.cores_available # similarly rather random resource_specification["memory"] = int(memory) resource_specification["running_time_min"] = self.max_training_time @@ -360,8 +363,8 @@ def command(self): def parse_size(size): size = size.upper() - if not re.match(r' ', size): - size = re.sub(r'([KMGT]?B)', r' \1', size) + if not re.match(r" ", size): + size = re.sub(r"([KMGT]?B)", r" \1", size) number, unit = [string.strip() for string in size.split()] return int(float(number) * units[unit]) diff --git a/psiflow/functions.py b/psiflow/functions.py index d3d382b..d68d404 100644 --- a/psiflow/functions.py +++ b/psiflow/functions.py @@ -218,8 +218,8 @@ def __post_init__(self): from mace.tools import torch_tools, utils torch_tools.set_default_dtype(self.dtype) - if self.device == 'gpu': # when it's not a specific GPU, use any - self.device = 'cuda' + if self.device == "gpu": # when it's not a specific GPU, use any + self.device = "cuda" self.device = torch_tools.init_device(self.device) torch.set_num_threads(self.ncores) diff --git a/psiflow/metrics.py b/psiflow/metrics.py index 357f947..2f73386 100644 --- a/psiflow/metrics.py +++ b/psiflow/metrics.py @@ -172,7 +172,6 @@ def _to_wandb( os.environ["WANDB_API_KEY"] = wandb_api_key os.environ["WANDB_SILENT"] = "True" import tempfile - import numpy as np from pathlib import Path import colorcet as cc @@ -533,6 +532,6 @@ def to_wandb(self): return to_wandb( self.wandb_id, self.wandb_project, - os.environ['WANDB_API_KEY'], + os.environ["WANDB_API_KEY"], inputs=[self.metrics], ) diff --git a/psiflow/models/_mace.py b/psiflow/models/_mace.py index ea7d770..95e3b95 100644 --- a/psiflow/models/_mace.py +++ b/psiflow/models/_mace.py @@ -184,7 +184,7 @@ def train( command_cd = "cd $mytmpdir;" command_env = "" for key, value in env_vars.items(): - command_env += ' export {}={}; '.format(key, value) + command_env += " export {}={}; ".format(key, value) command_list = [ command_tmp, command_cd, diff --git a/psiflow/models/mace_utils.py b/psiflow/models/mace_utils.py index ec1d664..28411a2 100644 --- a/psiflow/models/mace_utils.py +++ b/psiflow/models/mace_utils.py @@ -676,9 +676,10 @@ def run(rank: int, args: argparse.Namespace, world_size: int) -> None: def main(): + import signal + import torch from mace import tools - import signal signal.signal(signal.SIGTERM, timeout_handler) # main() diff --git a/psiflow/reference/_cp2k.py b/psiflow/reference/_cp2k.py index 095564b..1cf7c1f 100644 --- a/psiflow/reference/_cp2k.py +++ b/psiflow/reference/_cp2k.py @@ -260,8 +260,8 @@ def get_single_atom_references(self, element): "ot": {"minimizer": "CG"} } # necessary for oxygen calculation, at least in 2024.1 - key = 'ignore_convergence_failure' - cp2k_input_dict['force_eval']['dft']['scf'][key] = "TRUE" + key = "ignore_convergence_failure" + cp2k_input_dict["force_eval"]["dft"]["scf"][key] = "TRUE" reference = CP2K( dict_to_str(cp2k_input_dict), diff --git a/psiflow/sampling/client.py b/psiflow/sampling/client.py index 92cc75b..82b4e92 100644 --- a/psiflow/sampling/client.py +++ b/psiflow/sampling/client.py @@ -77,7 +77,7 @@ def main(): try: t0 = time.time() function([template] * 10) # torch warmp-up before simulation - print('time for 10 evaluations: {}'.format(time.time() - t0)) + print("time for 10 evaluations: {}".format(time.time() - t0)) run_driver( unix=True, address=str(Path.cwd() / args.address), diff --git a/psiflow/sampling/optimize.py b/psiflow/sampling/optimize.py index 14c6782..21f4c09 100644 --- a/psiflow/sampling/optimize.py +++ b/psiflow/sampling/optimize.py @@ -11,7 +11,7 @@ import psiflow from psiflow.data import Dataset -from psiflow.data.utils import write_frames, read_frames +from psiflow.data.utils import write_frames from psiflow.geometry import Geometry from psiflow.hamiltonians import Hamiltonian from psiflow.utils.io import save_xml diff --git a/psiflow/sampling/sampling.py b/psiflow/sampling/sampling.py index 640d7b6..040acac 100644 --- a/psiflow/sampling/sampling.py +++ b/psiflow/sampling/sampling.py @@ -358,7 +358,7 @@ def _execute_ipi( for i, plumed_str in enumerate(plumed_list): write_command += 'echo "{}" > metad_input{}.txt; '.format(plumed_str, i) for key, value in env_vars.items(): - write_command += ' export {}={}; '.format(key, value) + write_command += " export {}={}; ".format(key, value) command_start = command_server + " --nwalkers={}".format(nwalkers) command_start += " --input_xml={}".format(inputs[0].filepath) command_start += " --start_xyz={}".format(inputs[1].filepath) @@ -533,7 +533,7 @@ def _sample( command_server = definition.server_command() command_client = definition.client_command() resources = definition.wq_resources(max_nclients) - print('ENV VARS') + print("ENV VARS") print(definition.env_vars) result = execute_ipi( len(walkers), diff --git a/psiflow/sampling/server.py b/psiflow/sampling/server.py index a1fb1eb..2be50e6 100644 --- a/psiflow/sampling/server.py +++ b/psiflow/sampling/server.py @@ -1,7 +1,7 @@ import argparse -import os import ast import glob +import os import signal import xml.etree.ElementTree as ET from copy import deepcopy @@ -24,6 +24,7 @@ def remdsort(inputfile, prefix="SRT_"): from ipi.inputs.simulation import InputSimulation from ipi.utils.io.inputs import io_xml from ipi.utils.messages import verbosity + verbosity.level = "low" # opens & parses the input file ifile = open(inputfile, "r") @@ -133,8 +134,8 @@ def remdsort(inputfile, prefix="SRT_"): ) else: # FIX - if o.format == 'ase': - extension = 'extxyz' + if o.format == "ase": + extension = "extxyz" else: extension = o.format filename = filename + "_" + padb + "." + extension @@ -343,6 +344,7 @@ def anisotropic_barostat_h0(input_xml, data_start): def start(args): from ipi.engine.simulation import Simulation from ipi.utils.softexit import softexit + data_start = read(args.start_xyz, index=":") assert len(data_start) == args.nwalkers for i in range(args.nwalkers): @@ -371,6 +373,7 @@ def start(args): def cleanup(args): from psiflow.data.utils import _write_frames + with open("input.xml", "r") as f: content = f.read() if "vibrations" in content: @@ -458,9 +461,9 @@ def main(): if not args.cleanup: start(args) else: - #try: + # try: cleanup(args) - #except BaseException as e: # noqa: B036 + # except BaseException as e: # noqa: B036 # print(e) # print("i-PI cleanup failed!") # print("files in directory:") diff --git a/tests/conftest.py b/tests/conftest.py index 39bc2eb..97ad51e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -46,7 +46,7 @@ def context(request, tmp_path_factory): psiflow_config = yaml.safe_load(f) psiflow_config["path"] = tmp_path_factory.mktemp("psiflow_internal") psiflow.load(psiflow_config) - context = psiflow.context() + context = psiflow.context() # noqa: F841 yield parsl.dfk().cleanup() diff --git a/tests/test_data.py b/tests/test_data.py index 9b8d567..3657625 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -181,7 +181,7 @@ def test_readwrite_cycle(dataset, tmp_path): geometry = Geometry.from_data(np.ones(2), np.zeros((2, 3)), cell=None) geometry.stress = np.array([np.nan] * 9).reshape(3, 3) - assert 'nan' not in geometry.to_string() + assert "nan" not in geometry.to_string() def test_dataset_empty(tmp_path):