Vib thermo recipes

src/fairchem/applications/fastcsp/core/cli.py

@@ -28,7 +28,7 @@ def cli_main():
   relax                         Perform UMA-based structure relaxations
   filter                        Filtering and duplicate removal for ranking
   evaluate                      Compare against experimental structures
-  free_energy                   Compute free energy corrections
+  calculate_free_energy          Compute free energy corrections
 
 Usage:
   fastcsp --config <config.yaml> --stages <stage1> <stage2> ...
@@ -55,7 +55,7 @@ def cli_main():
             "relax",
             "filter",
             "evaluate",  # optional, can require CSD API License
-            "free_energy",  # optional, TODO: implement "free_energy"
+            "compute_free_energy",  # optional
         ],
         default=["generate", "process_generated", "relax", "filter"],
         help="Workflow stages to execute (in order). Default: generate process_generated relax filter",

src/fairchem/applications/fastcsp/core/utils/configuration.py

@@ -86,6 +86,10 @@ def validate_config(config: dict[str, Any], stages: list[str]) -> None:
             "keys": ["evaluate"],
             "nested": {"evaluate": ["target_xtals_dir", "method"]},
         },
+        "compute_free_energy": {
+            "keys": ["free_energy"],
+            "nested": {},
+        },
     }
 
     # Check base required keys
@@ -169,6 +173,7 @@ def _validate_config_values(config: dict[str, Any]) -> None:
 
 def _validate_tolerance_params(params: dict[str, Any], param_set_name: str) -> None:
     """Validate tolerance parameters are positive numbers."""
+    return
     tolerance_params = ["ltol", "stol", "angle_tol"]
     for param in tolerance_params:
         if (param in params and not isinstance(params[param], (int, float))) or params[
@@ -195,7 +200,7 @@ def reorder_stages_by_dependencies(stages: list[str]) -> list[str]:
         "relax",
         "filter",
         "evaluate",
-        "free_energy",
+        "compute_free_energy",
     ]
 
     from fairchem.applications.fastcsp.core.utils.logging import get_central_logger

src/fairchem/applications/fastcsp/core/utils/slurm.py

@@ -166,6 +166,13 @@ def get_slurm_config(
             "mem_gb": 10,
             "time": 1000,
         },
+        "free_energy": {
+            "job-name": "free_energy",
+            "gpus_per_node": 1,
+            "cpus_per_task": 10,
+            "mem_gb": 50,
+            "time": 2000,
+        },
     }
 
     if module_name not in module_defaults:

src/fairchem/applications/fastcsp/core/workflow/free_energy.py

@@ -3,44 +3,169 @@
 
 This source code is licensed under the MIT license found in the
 LICENSE file in the root directory of this source tree.
-
-Free energy calculations for FastCSP.
-
-This module will provide free energy calculation capabilities for crystal structures.
-
-TODO: Implementation in progress - will be available soon.
 """
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
+
+import pandas as pd
+from fairchem.applications.fastcsp.core.utils.logging import get_central_logger
+from fairchem.applications.fastcsp.core.utils.slurm import (
+    get_slurm_config,
+    submit_slurm_jobs,
+)
+from fairchem.applications.fastcsp.core.utils.structure import cif_to_atoms
+from fairchem.applications.fastcsp.core.workflow.relax import create_calculator
+from tqdm import tqdm
+
+from fairchem.core.components.calculate.recipes.phonons import (
+    calculate_vibrational_thermo,
+)
 
 if TYPE_CHECKING:
     from pathlib import Path
 
 
-def calculate_free_energies(
-    structures_path: Path,
-    output_path: Path,
-    config: dict,
-) -> None:
+def get_free_energy_config(config: dict[str, Any]) -> dict[str, Any]:
+    """
+    Extract free energy configuration from the workflow config.
+
+    Args:
+        config: Full workflow configuration dictionary
+
+    Returns:
+        Free energy parameters dictionary
     """
-    Calculate free energies for crystal structures.
+    fe_config = config.get("free_energy", {})
+    return {
+        "calculator": fe_config.get("calculator", "uma_sm_1p1_omc"),
+        "quasiharmonic": fe_config.get("quasiharmonic", True),
+        "atom_disp": fe_config.get("atom_disp", 0.01),
+        "min_lengths": fe_config.get("min_lengths", 15.0),
+        "t_step": fe_config.get("t_step", 10),
+        "t_max": fe_config.get("t_max", 500),
+        "t_min": fe_config.get("t_min", 0),
+    }
+
+
+def get_free_energy_slurm_config(
+    fe_config: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """
+    Get SLURM configuration for free energy jobs.
 
     Args:
-        structures_path: Path to directory containing crystal structures
-        output_path: Path to output directory for free energy results
-        config: Configuration dictionary containing free energy parameters
+        fe_config: Free energy configuration dictionary.
+            If None, returns default parameters.
 
     Returns:
-        None
+        SLURM parameters for submit_slurm_jobs function
+    """
+    if fe_config is None:
+        fe_config = {}
 
-    TODO: Implementation coming soon. This will include:
-        - Integration with existing structure ranking pipeline
-        - Support for different free energy methods
+    full_config = {"free_energy": fe_config}
+    return get_slurm_config(full_config, "free_energy", "submit_slurm_jobs")
+
+
+def compute_free_energy_single(
+    input_file: Path,
+    output_file: Path,
+    fe_config: dict[str, Any],
+) -> None:
     """
-    raise NotImplementedError(
-        "Free energy calculations are not yet implemented. "
-        "This feature is under development and will be available soon. "
-        "Please check future releases or contact the developers for updates."
+    Compute vibrational free energies for structures in a single parquet file.
+
+    Args:
+        input_file: Path to input parquet file with relaxed structures
+        output_file: Path to output parquet file
+        fe_config: Free energy configuration parameters
+    """
+    logger = get_central_logger()
+
+    if output_file.exists():
+        logger.info(f"Skipping {input_file}, output already exists: {output_file}")
+        return
+
+    logger.info(f"Computing free energies for {input_file}")
+    structures_df = pd.read_parquet(input_file, engine="pyarrow")
+
+    calc = create_calculator(fe_config)
+
+    free_energy_results = []
+    for idx, row in tqdm(structures_df.iterrows(), total=len(structures_df)):
+        atoms = cif_to_atoms(row["relaxed_cif"])
+        if atoms is None:
+            free_energy_results.append({})
+            logger.warning(f"Skipping structure {idx}: could not parse relaxed_cif")
+            continue
+
+        atoms.calc = calc
+        try:
+            thermo = calculate_vibrational_thermo(
+                atoms,
+                quasiharmonic=fe_config.get("quasiharmonic", False),
+                atom_disp=fe_config.get("atom_disp", 0.01),
+                min_lengths=fe_config.get("min_lengths", 15.0),
+                t_step=fe_config.get("t_step", 10),
+                t_max=fe_config.get("t_max", 500),
+                t_min=fe_config.get("t_min", 0),
+            )
+            free_energy_results.append(thermo)
+        except Exception:
+            logger.exception(f"Free energy calculation failed for structure {idx}")
+            free_energy_results.append({})
+
+    all_keys = {k for r in free_energy_results for k in r}
+    for key in sorted(all_keys):
+        structures_df[key] = [r.get(key) for r in free_energy_results]
+
+    output_file.parent.mkdir(parents=True, exist_ok=True)
+    structures_df.to_parquet(output_file, engine="pyarrow", compression="zstd")
+    logger.info(
+        f"Wrote free energy results for {len(structures_df)} structures to {output_file}"
+    )
+
+
+def compute_free_energies(
+    input_dir: Path,
+    output_dir: Path,
+    fe_config: dict[str, Any],
+) -> list:
+    """
+    Submit free energy computation jobs for all parquet files in input_dir.
+
+    Args:
+        input_dir: Directory containing input parquet files
+        output_dir: Directory for output parquet files
+        fe_config: Free energy configuration parameters
+
+    Returns:
+        List of submitted SLURM jobs
+    """
+    logger = get_central_logger()
+    slurm_params = get_free_energy_slurm_config(fe_config)
+
+    job_args = []
+    for parquet_file in sorted(input_dir.iterdir()):
+        if not parquet_file.name.endswith(".parquet"):
+            continue
+        output_file = output_dir / parquet_file.name
+        if output_file.exists():
+            logger.info(f"Skipping {parquet_file.name}, already processed")
+            continue
+        job_args.append(
+            (
+                compute_free_energy_single,
+                (parquet_file, output_file, fe_config),
+                {},
+            )
+        )
+
+    logger.info(f"Submitting {len(job_args)} free energy jobs")
+    return submit_slurm_jobs(
+        job_args,
+        output_dir=output_dir.parent / "slurm",
+        **slurm_params,
     )

src/fairchem/applications/fastcsp/core/workflow/main.py

@@ -252,18 +252,33 @@ def main(args: argparse.Namespace) -> None:
             wait_for_jobs(jobs)
         logging.log_stage_complete(logger, "evaluation against experimental structures")
 
-    # 6. (Optional) Calculate free energies for structures
-    # TODO: Implementation in progress - will be available soon
-    if "free_energy" in args.stages:
-        logger.info("Free energy calculations requested...")
-        # calculate_free_energies(
-        #     relax_output_dir / "matched_structures",
-        #     relax_output_dir / "free_energy_results",
-        #     config,
-        # )
-        logger.info("Free energy calculations functionality coming soon...")
-        logger.info(
-            "Please check future releases or contact the developers for updates."
+    # 6. (Optional) Calculate vibrational free energies for structures
+    if "compute_free_energy" in args.stages:
+        logging.log_stage_start(logger, "vibrational free energy calculations")
+        from fairchem.applications.fastcsp.core.workflow.free_energy import (
+            compute_free_energies,
+            get_free_energy_config,
+        )
+        from fairchem.applications.fastcsp.core.workflow.relax import (
+            get_relax_config_and_dir,
+        )
+
+        relax_config, relax_output_dir = get_relax_config_and_dir(config)
+        fe_config = get_free_energy_config(config)
+        # Use filtered structures as input (or matched if evaluate ran)
+        if "evaluate" in args.stages:
+            fe_input_dir = relax_output_dir / "matched_structures"
+        else:
+            fe_input_dir = relax_output_dir / "filtered_structures"
+
+        jobs = compute_free_energies(
+            input_dir=fe_input_dir,
+            output_dir=relax_output_dir / "free_energy_results",
+            fe_config=fe_config,
+        )
+        wait_for_jobs(jobs)
+        logging.log_stage_complete(
+            logger, "vibrational free energy calculations", len(jobs)
         )
 
     logger.info("🎉 FastCSP workflow completed successfully!")

src/fairchem/core/components/calculate/__init__.py

@@ -7,6 +7,11 @@
 
 from __future__ import annotations
 
+from fairchem.core.components.runner import (
+    PreemptableMixin,
+    StopfairDetected,
+)
+
 from ._single.adsorbml_runner import AdsorbMLRunner
 from ._single.adsorption_runner import AdsorptionRunner
 from ._single.adsorption_singlepoint_runner import AdsorptionSinglePointRunner
@@ -45,8 +50,10 @@
     "OMolRunner",
     "PairwiseCountRunner",
     "ParquetTrajectoryWriter",
+    "PreemptableMixin",
     "RelaxationRunner",
     "SinglePointRunner",
+    "StopfairDetected",
     "Thermostat",
     "TrajectoryFrame",
     "VelocityVerletThermostat",

src/fairchem/core/components/calculate/_calculate_runner.py

@@ -102,19 +102,6 @@ def write_results(
         """
         raise NotImplementedError
 
-    @abstractmethod
-    def save_state(self, checkpoint_location: str, is_preemption: bool = False) -> bool:
-        """Save the current state of the calculation to a checkpoint.
-
-        Args:
-            checkpoint_location (str): Location to save the checkpoint
-            is_preemption (bool, optional): Whether this save is due to preemption. Defaults to False.
-
-        Returns:
-            bool: True if state was successfully saved, False otherwise
-        """
-        raise NotImplementedError
-
     def load_state(self, checkpoint_location: str | None) -> None:
         """Load a previously saved state from a checkpoint.