make slurm_submit() time and account optional

WBM add 'find large structures that changed symmetry during relaxation'
update all slurm_submit(account="LEE-SL...->matgen")
ruff unignore PT013 S301 and apply fixes

Author: janosh

.pre-commit-config.yaml

@@ -8,7 +8,7 @@ default_install_hook_types: [pre-commit, commit-msg]
 
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.4.6
+    rev: v0.4.8
     hooks:
       - id: ruff
         args: [--fix]
@@ -57,7 +57,7 @@ repos:
         exclude: ^(site/src/figs/.+\.svelte|data/wbm/20.+\..+|site/src/(routes|figs).+\.(yaml|json)|changelog.md)$
 
   - repo: https://github.com/pre-commit/mirrors-eslint
-    rev: v9.3.0
+    rev: v9.4.0
     hooks:
       - id: eslint
         types: [file]
@@ -71,15 +71,15 @@ repos:
           - typescript-eslint
 
   - repo: https://github.com/python-jsonschema/check-jsonschema
-    rev: 0.28.4
+    rev: 0.28.5
     hooks:
       - id: check-jsonschema
         files: ^models/(.+)/\1.*\.yml$
         args: [--schemafile, tests/model-schema.yml]
       - id: check-github-actions
 
   - repo: https://github.com/RobertCraigie/pyright-python
-    rev: v1.1.365
+    rev: v1.1.366
     hooks:
       - id: pyright
         args: [--level, error]

data/wbm/compile_wbm_test_set.py

@@ -541,7 +541,7 @@ def fix_bad_struct_index_mismatch(material_id: str) -> str:
 
 # %%
 with gzip.open(DATA_FILES.mp_patched_phase_diagram, "rb") as zip_file:
-    ppd_mp: PatchedPhaseDiagram = pickle.load(zip_file)
+    ppd_mp: PatchedPhaseDiagram = pickle.load(zip_file)  # noqa: S301
 
 
 # %% calculate e_above_hull for each material

data/wbm/eda_wbm.py

@@ -9,7 +9,7 @@
 import pandas as pd
 import plotly.express as px
 from matplotlib.colors import SymLogNorm
-from pymatgen.core import Composition
+from pymatgen.core import Composition, Structure
 from pymatviz import (
     count_elements,
     ptable_heatmap,
@@ -18,6 +18,7 @@
     spacegroup_sunburst,
 )
 from pymatviz.io import save_fig
+from pymatviz.structure_viz import plot_structure_2d
 from pymatviz.utils import si_fmt, si_fmt_int
 
 from matbench_discovery import PDF_FIGS, ROOT, SITE_FIGS, STABILITY_THRESHOLD
@@ -371,3 +372,39 @@
 img_name = "mp-vs-wbm-arity-hist"
 save_fig(fig, f"{SITE_FIGS}/{img_name}.svelte")
 save_fig(fig, f"{PDF_FIGS}/{img_name}.pdf", width=450, height=280)
+
+
+# %% find large structures that changed symmetry during relaxation
+df_sym_change = (
+    df_wbm.query(f"{Key.init_wyckoff} != {Key.wyckoff}")
+    .filter(regex="wyckoff|sites")
+    .nlargest(10, Key.n_sites)
+)
+
+
+# %%
+df_wbm_structs = pd.read_json(DATA_FILES.wbm_cses_plus_init_structs).set_index(
+    Key.mat_id
+)
+
+
+# %%
+for wbm_id in df_sym_change.index:
+    init_struct = Structure.from_dict(df_wbm_structs.loc[wbm_id][Key.init_struct])
+    final_struct = Structure.from_dict(df_wbm_structs.loc[wbm_id][Key.cse]["structure"])
+    init_struct.properties[Key.mat_id] = f"{wbm_id}-init"
+    final_struct.properties[Key.mat_id] = f"{wbm_id}-final"
+
+    plot_structure_2d([init_struct, final_struct])
+
+
+# %% export initial and final structures with symmetry change to CIF
+wbm_id = df_sym_change.index[0]
+
+struct = Structure.from_dict(df_wbm_structs.loc[wbm_id][Key.cse]["structure"])
+struct.to(f"{module_dir}/{wbm_id}.cif")
+struct.to(f"{module_dir}/{wbm_id}.json")
+
+struct = Structure.from_dict(df_wbm_structs.loc[wbm_id][Key.init_struct])
+struct.to(f"{module_dir}/{wbm_id}-init.cif")
+struct.to(f"{module_dir}/{wbm_id}-init.json")

matbench_discovery/data.py

@@ -117,7 +117,7 @@ def load(
     print(f"Loading {key!r} from cached file at {cache_path!r}")
     if ".pkl" in file_path:  # handle key='mp_patched_phase_diagram' separately
         with gzip.open(cache_path, "rb") as zip_file:
-            return pickle.load(zip_file)
+            return pickle.load(zip_file)  # noqa: S301
     if ".pth" in file_path:  # handle model checkpoints (e.g. key='alignn_checkpoint')
         return cache_path
 

matbench_discovery/slurm.py

@@ -29,8 +29,9 @@ def _get_calling_file_path(frame: int = 1) -> str:
 def slurm_submit(
     job_name: str,
     out_dir: str,
-    time: str,
-    account: str,
+    *,
+    time: str | None = None,
+    account: str | None = None,
     partition: str | None = None,
     py_file_path: str | None = None,
     slurm_flags: str | Sequence[str] = (),
@@ -72,30 +73,34 @@ def slurm_submit(
 
     os.makedirs(out_dir, exist_ok=True)  # slurm fails if out_dir is missing
 
+    # ensure pre_cmd ends with a semicolon
+    if pre_cmd and not pre_cmd.strip().endswith(";"):
+        pre_cmd += ";"
+
     cmd = [
-        *f"sbatch --{account=} --{time=}".replace("'", "").split(),
-        *("--job-name", job_name),
+        *("sbatch", "--job-name", job_name),
         *("--output", f"{out_dir}/slurm-%A{'-%a' if array else ''}.log"),
         *(slurm_flags.split() if isinstance(slurm_flags, str) else slurm_flags),
-        *("--wrap", f"{pre_cmd} python {py_file_path}".strip()),
+        *("--wrap", f"{pre_cmd or ''} python {py_file_path}".strip()),
     ]
-    if partition:
-        cmd += ["--partition", partition]
-    if array:
-        cmd += ["--array", array]
+    for flag in (f"{time=}", f"{account=}", f"{partition=}", f"{array=}"):
+        key, val = flag.split("=")
+        if val != "None":
+            cmd += (f"--{key}", val)
 
     is_log_file = not sys.stdout.isatty()
     is_slurm_job = "SLURM_JOB_ID" in os.environ
 
     slurm_vars = {
-        f"slurm_{key}": val
+        f"slurm_{key}": os.environ[f"SLURM_{key}".upper()]
         for key in SLURM_KEYS
-        if (val := os.getenv(f"SLURM_{key}".upper()))
+        if f"SLURM_{key}".upper() in os.environ
     }
-    slurm_vars["slurm_timelimit"] = time
-    if slurm_flags:
+    if time is not None:
+        slurm_vars["slurm_timelimit"] = time
+    if slurm_flags != ():
         slurm_vars["slurm_flags"] = str(slurm_flags)
-    if pre_cmd:
+    if pre_cmd not in ("", None):
         slurm_vars["pre_cmd"] = pre_cmd
 
     # print sbatch command into slurm log file and at job submission time

models/alignn/test_alignn.py

@@ -58,15 +58,13 @@
 
 slurm_vars = slurm_submit(
     job_name=job_name,
-    partition="ampere",
-    account="LEE-SL3-GPU",
+    account="matgen",
     time="11:55:0",
     out_dir=out_dir,
     slurm_flags="--nodes 1 --gpus-per-node 1",
     # pre_cmd is platform specific, remove when running on other systems
     # just left here for reference
-    pre_cmd=". /etc/profile.d/modules.sh; module load rhel8/default-amp;"
-    "module load cuda/11.8",
+    pre_cmd="module load cuda/11.8",
 )
 
 

models/bowsr/test_bowsr.py

@@ -46,8 +46,7 @@
 slurm_vars = slurm_submit(
     job_name=job_name,
     out_dir=out_dir,
-    partition="skylake",
-    account="LEE-SL3-CPU",
+    account="matgen",
     time="11:55:0",
     # --time=2:0:0 is probably enough but best be safe.
     array=f"1-{slurm_array_task_count}%{slurm_max_parallel}",

models/cgcnn/test_cgcnn.py

@@ -35,8 +35,7 @@
 
 slurm_vars = slurm_submit(
     job_name=job_name,
-    partition="ampere",
-    account="LEE-SL3-GPU",
+    account="matgen",
     time="2:0:0",
     out_dir=out_dir,
     slurm_flags="--nodes 1 --gpus-per-node 1",

models/cgcnn/train_cgcnn.py

@@ -41,8 +41,7 @@
 
 slurm_vars = slurm_submit(
     job_name=job_name,
-    partition="ampere",
-    account="LEE-SL3-GPU",
+    account="matgen",
     time="11:55:0",
     array=f"1-{ensemble_size}",
     out_dir=out_dir,

models/megnet/test_megnet.py

@@ -39,10 +39,9 @@
 slurm_vars = slurm_submit(
     job_name=job_name,
     out_dir=module_dir,
-    partition="icelake-himem",
-    account="LEE-SL3-CPU",
+    account="matgen",
     time="11:55:0",
-    slurm_flags=("--mem", "30G"),
+    slurm_flags="--mem 30G",
     array=f"1-{slurm_array_task_count}",
     # TF_CPP_MIN_LOG_LEVEL=2 means INFO and WARNING logs are not printed
     # https://stackoverflow.com/a/40982782

models/voronoi_rf/train_test_voronoi_rf.py

@@ -41,8 +41,7 @@
 slurm_vars = slurm_submit(
     job_name=job_name,
     out_dir=out_dir,
-    partition="icelake-himem",
-    account="LEE-SL3-CPU",
+    account="matgen",
     time="6:0:0",
 )
 

models/voronoi_rf/voronoi_featurize_dataset.py

@@ -43,8 +43,7 @@
 
 slurm_vars = slurm_submit(
     job_name=job_name,
-    partition="icelake-himem",
-    account="LEE-SL3-CPU",
+    account="matgen",
     time="11:55:0",
     array=f"1-{slurm_array_task_count}",
     slurm_flags=("--mem", "15G") if data_name == "mp" else (),

models/wrenformer/test_wrenformer.py

@@ -32,8 +32,7 @@
 
 slurm_vars = slurm_submit(
     job_name=job_name,
-    partition="ampere",
-    account="LEE-SL3-GPU",
+    account="matgen",
     time="2:0:0",
     out_dir=out_dir,
     slurm_flags="--nodes 1 --gpus-per-node 1",

models/wrenformer/train_wrenformer.py

@@ -32,8 +32,7 @@
 
 slurm_vars = slurm_submit(
     job_name=job_name,
-    partition="ampere",
-    account="LEE-SL3-GPU",
+    account="matgen",
     time="8:0:0",
     array=f"1-{ensemble_size}",
     out_dir=out_dir,

pyproject.toml

@@ -104,10 +104,8 @@ ignore = [
   "PLR",     # pylint refactor
   "PLW2901", # redefined-loop-name
   "PT006",   # pytest-parametrize-names-wrong-type
-  "PT013",   # pytest-incorrect-pytest-import
   "PTH",
   "S108",
-  "S301",
   "S310",
   "S311",
   "S603",

scripts/compute_struct_fingerprints.py

@@ -39,8 +39,7 @@
 slurm_vars = slurm_submit(
     job_name=f"{data_name}-struct-fingerprints",
     out_dir=out_dir,
-    partition="icelake-himem",
-    account="LEE-SL3-CPU",
+    account="matgen",
     time="6:0:0",
     array=f"1-{slurm_array_task_count}",
     slurm_flags=("--mem", "30G"),

scripts/project_compositions.py

@@ -30,8 +30,7 @@
 slurm_vars = slurm_submit(
     job_name=f"{data_name}-{projection_type}-{out_dim}d",
     out_dir=out_dir,
-    partition="icelake-himem",
-    account="LEE-SL3-CPU",
+    account="matgen",
     time="6:0:0",
 )
 

tests/test_data.py

@@ -8,7 +8,6 @@
 import pandas as pd
 import pytest
 from pymatgen.core import Lattice, Structure
-from pytest import CaptureFixture
 
 from matbench_discovery import FIGSHARE_DIR, ROOT
 from matbench_discovery.data import (
@@ -45,7 +44,7 @@ def test_load(
     df_float: pd.DataFrame,
     # df with Structures and ComputedStructureEntries as dicts
     df_with_pmg_objects: pd.DataFrame,
-    capsys: CaptureFixture[str],
+    capsys: pytest.CaptureFixture[str],
     tmp_path: Path,
     key: str,
     hydrate: bool,
@@ -148,7 +147,7 @@ def test_load_no_mock(
     version: str,
     expected_shape: tuple[int, int],
     expected_cols: set[str],
-    capsys: CaptureFixture[str],
+    capsys: pytest.CaptureFixture[str],
     tmp_path: Path,
 ) -> None:
     assert os.listdir(tmp_path) == [], "cache_dir should be empty"

tests/test_energy.py

@@ -5,7 +5,6 @@
 from pymatgen.analysis.phase_diagram import PDEntry
 from pymatgen.core import Lattice, Structure
 from pymatgen.entries.computed_entries import ComputedEntry, Entry
-from pytest import approx
 
 from matbench_discovery.energy import (
     get_e_form_per_atom,
@@ -59,5 +58,5 @@ def test_mp_ref_energies() -> None:
     """Test MP elemental reference energies are in sync with PDEntries saved to disk."""
     for key, val in mp_elemental_ref_energies.items():
         actual = mp_elem_reference_entries[key].energy_per_atom
-        assert actual == approx(val, abs=1e-3), f"{key=}"
-        assert actual == approx(val, abs=1e-3), f"{key=}"
+        assert actual == pytest.approx(val, abs=1e-3), f"{key=}"
+        assert actual == pytest.approx(val, abs=1e-3), f"{key=}"

tests/test_metrics.py

@@ -3,7 +3,6 @@
 import numpy as np
 import pandas as pd
 import pytest
-from pytest import approx
 
 from matbench_discovery.enums import Key
 from matbench_discovery.metrics import classify_stable, stable_metrics
@@ -48,7 +47,7 @@ def test_stable_metrics() -> None:
         RMSE=1.157,
         R2=-3.030,
     ).items():
-        assert metrics[key] == approx(val, abs=1e-3), f"{key=}"
+        assert metrics[key] == pytest.approx(val, abs=1e-3), f"{key=}"
 
     assert math.isnan(metrics["F1"])