Plot F1 and K_SRME metrics over time on new `/tasks/(discovery|phonon…

…s)` pages (#200) * extract TableColumnToggleMenu.svelte component from landing page - refactor +page.svelte to use new component and remove inline column toggle menu styles - .gitignore exclude model prediction JSON files - use svelte-zoo click_outside action for column toggle menu closing behavior * record phonon Figshare article ID in figshare.ARTICLE_IDS - add missing eqV2-s-dens-mp.yml figshare download URL for kappa_103 prediction file * calc_kappa_mace_ray_parallelized.py remove main() function - clean up remote parameter passing to calc_kappa_for_structure - calc_kappa_for_structure doc str add missing displacement_distance keyword * add MetricScatter.svelte component for showing metric improvement over time based on model submission dates - add new pages /tasks/discovery and /tasks/phonons using MetricScatter for metric visualization - update elementari package to v0.2.5 * faff to trigger CI
janosh · Feb 5, 2025 · 991f41f · 991f41f
1 parent a7c48dc
commit 991f41f
Show file tree

Hide file tree

Showing 17 changed files with 749 additions and 233 deletions.
diff --git a/.gitignore b/.gitignore
@@ -27,11 +27,12 @@ data/wbm/*symprec=*.csv*
 wandb/
 models/**/*.log
 
-# model checkpoints
+# model checkpoints and prediction files
 models/**/*.model
 models/**/*.pt
 models/**/*.tgz
 models/**/*.csv*
+models/**/*.json*
 
 # WIP model submissions
 models/_*

diff --git a/.pre-commit-config.yaml b/.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.9.3
+    rev: v0.9.4
     hooks:
       - id: ruff
         args: [--fix]
@@ -31,13 +31,13 @@ repos:
       - id: trailing-whitespace
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.14.1
+    rev: v1.15.0
     hooks:
       - id: mypy
         additional_dependencies: [types-pyyaml, types-requests]
 
   - repo: https://github.com/codespell-project/codespell
-    rev: v2.4.0
+    rev: v2.4.1
     hooks:
       - id: codespell
         stages: [pre-commit, commit-msg]
@@ -73,7 +73,7 @@ repos:
           - "@stylistic/eslint-plugin"
 
   - repo: https://github.com/python-jsonschema/check-jsonschema
-    rev: 0.31.0
+    rev: 0.31.1
     hooks:
       - id: check-jsonschema
         files: ^models/.+/.+\.yml$
@@ -84,7 +84,7 @@ repos:
       - id: check-github-actions
 
   - repo: https://github.com/RobertCraigie/pyright-python
-    rev: v1.1.392.post0
+    rev: v1.1.393
     hooks:
       - id: pyright
         args: [--level, error]
diff --git a/matbench_discovery/figshare.py b/matbench_discovery/figshare.py
@@ -22,7 +22,7 @@
 ARTICLE_IDS: Final[dict[str, int | None]] = {
     "model_preds_discovery": 28187990,
     "model_preds_geo_opt": 28187999,
-    "model_preds_phonons": None,
+    "model_preds_phonons": 28347251,
     "data_files": 22715158,
 }
 

diff --git a/matbench_discovery/phonons/thermal_conductivity.py b/matbench_discovery/phonons/thermal_conductivity.py
@@ -228,8 +228,8 @@ def calculate_conductivity(
         **kwargs (Any): Passed to Phono3py.run_thermal_conductivity().
 
     Returns:
-        tuple[Phono3py, dict[str, np.ndarray], Any]: Tuple of (Phono3py object,
-            conductivity dict, conductivity object)
+        tuple[Phono3py, dict[str, np.ndarray], Any]: (Phono3py object, conductivity
+            dict, conductivity object)
     """
     ph3.init_phph_interaction(symmetrize_fc3q=False)
 
@@ -290,7 +290,8 @@ def calc_mode_kappa_tot(
             (T, q-points, bands)
 
     Returns:
-        np.ndarray: Total mode kappa
+        np.ndarray: Total (particle-like + wave-like) thermal conductivity per phonon
+            mode with shape (T, q-points, bands)
     """
     # Temporarily silence divide warnings since we handle NaN values below
     with np.errstate(divide="ignore", invalid="ignore"):

diff --git a/models/eqV2/eqV2-m-omat-mp-salex.yml b/models/eqV2/eqV2-m-omat-mp-salex.yml
@@ -76,7 +76,6 @@ hyperparams:
   stochastic_depth: 0.1
   batch_size: 256
   epochs: 16
-  # train_config_url: #TODO add URL
 
 notes:
   Description: |

diff --git a/models/eqV2/eqV2-s-dens-mp.yml b/models/eqV2/eqV2-s-dens-mp.yml
@@ -77,7 +77,6 @@ hyperparams:
   stochastic_depth: 0.1
   batch_size: 512
   epochs: 150
-  # train_config_url: #TODO add URL
 
 notes:
   Description: |
@@ -90,8 +89,8 @@ metrics:
   phonons:
     kappa_103:
       κ_SRME: 1.665 # eqV2 S without denoising (no DeNS) achieves slightly worse κ_SRME=1.772
-      pred_file: TODO find this file or regenerate it
-      pred_file_url: https://figshare.com/TODO add this URL
+      pred_file: models/eqV2/eqV2-s-dens-mp/2024-11-08-kappa-103-FIRE-dist=0.01-fmax=1e-4-symprec=1e-5.json.gz
+      pred_file_url: https://figshare.com/ndownloader/files/52151558
   geo_opt:
     pred_file: models/eqV2/eqV2-s-dens-mp/2024-10-18-wbm-geo-opt.json.gz
     pred_file_url: https://figshare.com/ndownloader/files/52062392

diff --git a/models/mace/calc_kappa_mace_ray_parallelized.py b/models/mace/calc_kappa_mace_ray_parallelized.py
@@ -64,7 +64,8 @@ def calc_kappa_for_structure(
     Args:
         atoms (Atoms): ASE Atoms object with fc2_supercell, fc3_supercell,
             q_point_mesh keys in its info dict.
-        calc (Calculator): ASE calculator
+        displacement_distance (float): Displacement distance for phono3py
+        checkpoint (str): File path or download URL to model checkpoint
         temperatures (list[float]): Which temperatures to calculate kappa at in Kelvin
         ase_optimizer (str): ASE optimizer to use
         ase_filter (str): ASE filter to use
@@ -216,149 +217,128 @@ def calc_kappa_for_structure(
         return mat_id, info_dict | relax_dict | freqs_dict, force_results
 
 
-def main() -> None:
-    """Main function to run kappa predictions."""
-    # Model parameters
-
-    # Relaxation parameters
-    ase_optimizer: Literal["FIRE", "LBFGS", "BFGS"] = "FIRE"
-    ase_filter: Literal["frechet", "exp"] = "frechet"
-    max_steps = 300
-    fmax = 1e-4  # Run until the forces are smaller than this in eV/A
-
-    # Symmetry parameters
-    symprec = 1e-5  # symmetry precision for enforcing relaxation and conductivity calcs
-    enforce_relax_symm = True  # Enforce symmetry with during relaxation if broken
-    # Conductivity to be calculated if symmetry group changed during relaxation
-    conductivity_broken_symm = False
-    save_forces = True  # Save force sets to file
-    temperatures: list[float] = [300]
-
-    # Initialize Ray (point this at the head node of the cluster)
-    ray_ip: Final[str] = {
-        "lambda-staging-with-ray-2.40": "100.82.154.22",
-    }.get(os.getenv("RAY_IP_KEY", "")) or "100.82.154.22"
-    ray_address = os.getenv("RAY_ADDRESS", f"ray://{ray_ip}:10001")
-    print(f"{ray_address=}")
-
-    # Ray initialization
-    if ray_address:
-        # Connect to existing Ray cluster
-        ray.init(
-            address=ray_address,
-            runtime_env={
-                "py_modules": [matbench_discovery, pymatviz],
-                # "working_dir": PKG_DIR,  # add matbench-discovery root to PYTHONPATH
-                "uv": [
-                    # "mace-torch",
-                    "git+https://github.com/janosh/mace",
-                    # "cuequivariance-ops-torch-cu12",
-                    # "cuequivariance-torch",
-                    # "cuequivariance",
-                    "phono3py",
-                    "ase",
-                    "moyopy",
-                    # "pymatviz",
-                ],
-            },
-        )
-    else:
-        # Start Ray locally with optimized settings for M3 Max
-        ray.init(
-            num_cpus=8,  # Use 8/14 cores (leaving some for system + efficiency cores)
-            num_gpus=1,  # M3 Max GPU will be treated as 1 GPU
-        )
-
-    print(f"\nConnected to Ray cluster: {ray.cluster_resources()}")
-    ray_resources = ray.available_resources()
-    ray_mem = ray_resources.get("memory", 0) / 1e9
-    print(f"Available memory: {ray_mem:.1f} GB")
-    obj_store_mem = ray_resources.get("object_store_memory", 0)
-    print(f"Object store memory: {obj_store_mem / 1e9:.1f} GB")
-
-    model_name = "mace-omat-0-medium"
-    checkpoint = f"https://github.com/ACEsuit/mace-mp/releases/download/mace_omat_0/{model_name}.model"
-
-    task_type = "LTC"  # lattice thermal conductivity
-    displacement_distance = 0.01
-    job_name = (
-        f"phononDB-{task_type}-{ase_optimizer}-dist={displacement_distance}-"
-        f"{fmax=}-{symprec=}"
+# Relaxation parameters
+ase_optimizer: Literal["FIRE", "LBFGS", "BFGS"] = "FIRE"
+ase_filter: Literal["frechet", "exp"] = "frechet"
+max_steps = 300
+fmax = 1e-4  # Run until the forces are smaller than this in eV/A
+
+# Symmetry parameters
+symprec = 1e-5  # symmetry precision for enforcing relaxation and conductivity calcs
+enforce_relax_symm = True  # Enforce symmetry with during relaxation if broken
+# Conductivity to be calculated if symmetry group changed during relaxation
+conductivity_broken_symm = False
+save_forces = True  # Save force sets to file
+temperatures: list[float] = [300]
+
+# Initialize Ray (point this at the head node of the cluster)
+ray_ip: Final[str] = {
+    "lambda-staging-with-ray-2.40": "100.82.154.22",
+}.get(os.getenv("RAY_IP_KEY", "")) or "100.82.154.22"
+ray_address = os.getenv("RAY_ADDRESS", f"ray://{ray_ip}:10001")
+print(f"{ray_address=}")
+
+# Ray initialization
+if ray_address:
+    # Connect to existing Ray cluster
+    ray.init(
+        address=ray_address,
+        runtime_env={
+            "py_modules": [matbench_discovery, pymatviz],
+            # "working_dir": PKG_DIR,  # add matbench-discovery root to PYTHONPATH
+            "uv": [
+                # "mace-torch",
+                "git+https://github.com/janosh/mace",
+                # "cuequivariance-ops-torch-cu12",
+                # "cuequivariance-torch",
+                # "cuequivariance",
+                "phono3py",
+                "ase",
+                "moyopy",
+                # "pymatviz",
+            ],
+        },
     )
-    out_dir = os.getenv(
-        "SBATCH_OUTPUT", f"{module_dir}/{model_name}/{today}-{job_name}"
+else:
+    # Start Ray locally with optimized settings for M3 Max
+    ray.init(
+        num_cpus=8,  # Use 8/14 cores (leaving some for system + efficiency cores)
+        num_gpus=1,  # M3 Max GPU will be treated as 1 GPU
     )
-    os.makedirs(out_dir, exist_ok=True)
-
-    timestamp = f"{datetime.now().astimezone():%Y-%m-%d@%H-%M-%S}"
-    print(f"\nJob {job_name} with {model_name} started {timestamp}")
-
-    atoms_list = ase.io.read(DataFiles.phonondb_pbe_103_structures.path, index=":")
-    # Save run parameters
-    run_params = dict(
-        model_name=model_name,
-        checkpoint=checkpoint,
-        ase_optimizer=ase_optimizer,
-        ase_filter=ase_filter,
-        max_steps=max_steps,
-        force_max=fmax,
-        symprec=symprec,
-        enforce_relax_symm=enforce_relax_symm,
-        conductivity_broken_symm=conductivity_broken_symm,
-        task_type=task_type,
-        struct_data_path=DataFiles.phonondb_pbe_103_structures.path,
-        n_structures=len(atoms_list),
-        versions={dep: version(dep) for dep in ("numpy", "torch", "ray")},
-        temperatures=temperatures,
-    )
-
-    with open(f"{out_dir}/run_params.json", mode="w") as file:
-        json.dump(run_params, file, indent=4)
-
-    # Process structures in parallel
-    futures = [
-        calc_kappa_for_structure.remote(
-            atoms=atoms,
-            displacement_distance=displacement_distance,
-            checkpoint=checkpoint,
-            temperatures=temperatures,
-            ase_optimizer=ase_optimizer,
-            ase_filter=ase_filter,
-            max_steps=max_steps,
-            force_max=fmax,
-            symprec=symprec,
-            enforce_relax_symm=enforce_relax_symm,
-            conductivity_broken_symm=conductivity_broken_symm,
-            save_forces=save_forces,
-            out_dir=out_dir,
-            task_id=idx,
-        )
-        for idx, atoms in enumerate(atoms_list[:])
-    ]
-
-    # Process results as they complete
-    kappa_results: dict[str, dict[str, Any]] = {}
-    force_results: dict[str, dict[str, Any]] = {}
-
-    for future in tqdm(futures, desc=f"Predicting kappa with {model_name}"):
-        mat_id, result_dict, force_dict = ray.get(future)
-        kappa_results[mat_id] = result_dict
-        if force_dict is not None:
-            force_results[mat_id] = force_dict
-
-        # Save intermediate results
-        df_kappa = pd.DataFrame(kappa_results).T
-        df_kappa.index.name = Key.mat_id
-        df_kappa.reset_index().to_json(f"{out_dir}/kappa.json.gz")
-
-        if save_forces:
-            df_force = pd.DataFrame(force_results).T
-            df_force = pd.concat([df_kappa, df_force], axis=1)
-            df_force.index.name = Key.mat_id
-            df_force.reset_index().to_json(f"{out_dir}/force-sets.json.gz")
-
-    print(f"\nResults saved to {out_dir!r}")
 
+print(f"\nConnected to Ray cluster: {ray.cluster_resources()}")
+ray_resources = ray.available_resources()
+ray_mem = ray_resources.get("memory", 0) / 1e9
+print(f"Available memory: {ray_mem:.1f} GB")
+obj_store_mem = ray_resources.get("object_store_memory", 0)
+print(f"Object store memory: {obj_store_mem / 1e9:.1f} GB")
+
+model_name = "mace-omat-0-medium"
+checkpoint = f"https://github.com/ACEsuit/mace-mp/releases/download/mace_omat_0/{model_name}.model"
+
+task_type = "LTC"  # lattice thermal conductivity
+displacement_distance = 0.01
+job_name = (
+    f"phononDB-{task_type}-{ase_optimizer}-dist={displacement_distance}-"
+    f"{fmax=}-{symprec=}"
+)
+out_dir = os.getenv("SBATCH_OUTPUT", f"{module_dir}/{model_name}/{today}-{job_name}")
+os.makedirs(out_dir, exist_ok=True)
+
+timestamp = f"{datetime.now().astimezone():%Y-%m-%d@%H-%M-%S}"
+print(f"\nJob {job_name} with {model_name} started {timestamp}")
+
+atoms_list = ase.io.read(DataFiles.phonondb_pbe_103_structures.path, index=":")
+# Save run parameters
+remote_params = dict(
+    model_name=model_name,
+    checkpoint=checkpoint,
+    ase_optimizer=ase_optimizer,
+    ase_filter=ase_filter,
+    max_steps=max_steps,
+    force_max=fmax,
+    symprec=symprec,
+    enforce_relax_symm=enforce_relax_symm,
+    conductivity_broken_symm=conductivity_broken_symm,
+    task_type=task_type,
+    temperatures=temperatures,
+    out_dir=out_dir,
+)
+run_params = dict(
+    **remote_params,
+    n_structures=len(atoms_list),
+    struct_data_path=DataFiles.phonondb_pbe_103_structures.path,
+    versions={dep: version(dep) for dep in ("numpy", "torch", "ray")},
+)
 
-if __name__ == "__main__":
-    main()
+with open(f"{out_dir}/run_params.json", mode="w") as file:
+    json.dump(run_params, file, indent=4)
+
+# Process structures in parallel
+futures = [
+    calc_kappa_for_structure.remote(atoms=atoms, **remote_params, task_id=idx)
+    for idx, atoms in enumerate(atoms_list[:])
+]
+
+# Process results as they complete
+kappa_results: dict[str, dict[str, Any]] = {}
+force_results: dict[str, dict[str, Any]] = {}
+
+for future in tqdm(futures, desc=f"Predicting kappa with {model_name}"):
+    mat_id, result_dict, force_dict = ray.get(future)
+    kappa_results[mat_id] = result_dict
+    if force_dict is not None:
+        force_results[mat_id] = force_dict
+
+    # Save intermediate results
+    df_kappa = pd.DataFrame(kappa_results).T
+    df_kappa.index.name = Key.mat_id
+    df_kappa.reset_index().to_json(f"{out_dir}/kappa.json.gz")
+
+    if save_forces:
+        df_force = pd.DataFrame(force_results).T
+        df_force = pd.concat([df_kappa, df_force], axis=1)
+        df_force.index.name = Key.mat_id
+        df_force.reset_index().to_json(f"{out_dir}/force-sets.json.gz")
+
+print(f"\nResults saved to {out_dir!r}")