Make FastCSP compatible w/ Genarris3 updates (#1959)

Author: gvahe

packages/fairchem-core/pyproject.toml

@@ -19,7 +19,7 @@ dependencies = [
     "huggingface_hub>=0.27.1",
     "ase>=3.26.0",
     "ase-db-backends>=0.10.0",
-    "monty>=v2025.1.3",
+    "monty>=2026.2.18",
     "clusterscope==0.0.18",
     "setuptools<81.0.0",
     "requests",
@@ -38,7 +38,7 @@ dev = ["pre-commit", "pytest", "pytest-cov", "coverage", "syrupy", "ruff==0.5.1"
 docs = ["jupyter-book", "jupytext", "sphinx","sphinx-autoapi==3.3.3", "astroid<4", "umap-learn", "vdict", "ipywidgets", "jupyter_book>=2.0", "torch-dftd"]
 adsorbml = ["dscribe", "x3dase", "scikit-image"]
 torchsim = ["torch-sim-atomistic>=0.5.2; python_version >= '3.12'"]
-extras = ["pymatgen", "quacc[phonons]>=0.15.3", "pandas", "nvalchemi-toolkit-ops==0.2.0", "pyarrow"]
+extras = ["pymatgen>=2025.1.9", "quacc[phonons]>=0.15.3", "pandas", "nvalchemi-toolkit-ops==0.2.0", "pyarrow"]
 
 [project.scripts]
 fairchem = "fairchem.core._cli:main"

packages/fairchem-demo-ocpapi/pyproject.toml

@@ -22,7 +22,7 @@ dependencies = [
     "requests",
     "responses == 0.23.2",
     "tenacity == 8.2.3",
-    "tqdm == 4.66.1",
+    "tqdm",
 ]
 
 [project.urls]

src/fairchem/applications/fastcsp/core/configs/example_config.yaml

@@ -7,10 +7,10 @@ molecules: configs/example_systems.csv # File containing list of target molecule
 # Genarris configuration for structure generation
 genarris:
   # Execution environment
-  python_cmd: /path/to/python/with/genarris/installed # Python interpreter with Genarris installed
-  genarris_script: /path/to/genarris_master.py        # Path to Genarris master script
   mpi_launcher: /path/to/mpirun                       # MPI launcher: mpirun, srun, or custom path
-  base_config: configs/gnrs_base.conf                 # Base Genarris configuration template
+  python_cmd: /path/to/python/with/genarris/installed # Python interpreter with Genarris installed
+  genarris_cli: /path/to/genarris_cli.py              # Path to Genarris cli script
+  genarris_base_config: configs/gnrs_base.conf        # Base Genarris configuration template
   # Structure generation parameters
   vars:
     Z: [1, 2, 3, 4, 6, 8] # [1] by default

src/fairchem/applications/fastcsp/core/configs/genarris_base.conf

@@ -6,25 +6,26 @@
 name                        = ???
 molecule_path               = ???
 Z                           = ???
-log_level                   = debug
-downselection_workflow      = robust_flow
-restart                     = True
+log_level                   = info
 
-[generation] 
-num_structures_per_spg      = ??? 
-spg_distribution_type       = ??? 
-specific_radius_proportion  = 0.95 
-max_attempts_per_spg        = 100000000 
-tol                         = 0.01 
-unit_cell_volume_mean       = predict 
-volume_mult                 = 1.5 
-max_attempts_per_volume     = 10000000 
-generation_type             = crystal 
-natural_cutoff_mult         = 1.2 
+[workflow]
+tasks                       = ['generation', 'symm_rigid_press']
 
-[symm_rigid_press] 
-sr                          = 0.85 
-method                      = BFGS 
-natural_cutoff_mult         = 1.2 
+[generation]
+num_structures_per_spg      = ???
+spg_distribution_type       = ???
+specific_radius_proportion  = 0.95
+max_attempts_per_spg        = 100000000
 tol                         = 0.01
-debug_flag                  = False
+unit_cell_volume_mean       = predict
+volume_mult                 = 1.5
+max_attempts_per_volume     = 10000000
+generation_type             = crystal
+natural_cutoff_mult         = 1.2
+
+[symm_rigid_press]
+sr                          = 0.85
+method                      = BFGS
+natural_cutoff_mult         = 1.2
+tol                         = 0.01
+debug_flag                  = False

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

@@ -50,7 +50,9 @@ def validate_config(config: dict[str, Any], stages: list[str]) -> None:
     stage_requirements = {
         "generate": {
             "keys": ["molecules", "genarris"],
-            "nested": {"genarris": ["python_cmd", "genarris_script", "base_config"]},
+            "nested": {
+                "genarris": ["python_cmd", "genarris_cli", "genarris_base_config"]
+            },
         },
         "process_generated": {
             "keys": ["pre_relaxation_filter"],

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

@@ -27,6 +27,7 @@
 from __future__ import annotations
 
 import ast
+import json
 import shutil
 from configparser import ConfigParser
 from pathlib import Path
@@ -64,7 +65,7 @@ def create_gnrs_submit_script(
         gnrs_config: Genarris configuration containing execution parameters:
             - mpi_launcher: MPI command to use (default: "mpirun")
             - python_cmd: Python executable path (default: "python")
-            - genarris_script: Genarris main script name (default: "genarris_master.py")
+            - genarris_cli: Genarris main script name (default: "genarris_cli.py")
         genarris_slurm_config: SLURM resource allocation parameters
         single_gnrs_folder: Directory where the SLURM script will be created
 
@@ -83,7 +84,7 @@ def create_gnrs_submit_script(
 export OMP_NUM_THREADS=1
 
 {gnrs_config.get("mpi_launcher", "mpirun")} -np {genarris_slurm_config.get("nodes", 1) * genarris_slurm_config["ntasks-per-node"]} \\
-    {gnrs_config.get("python_cmd", "python")} {gnrs_config.get("genarris_script", "genarris_master.py")} {single_gnrs_folder}/ui.conf > {single_gnrs_folder}/Genarris.out
+    {gnrs_config.get("python_cmd", "python")} {gnrs_config.get("genarris_cli", "cli.py")} --config {single_gnrs_folder}/ui.conf > {single_gnrs_folder}/Genarris.out
 """
 
     with open(single_gnrs_folder / "slurm.sh", "w") as f:
@@ -96,7 +97,7 @@ def create_gnrs_config(
     mol_name: str,
     geometry_path: str | Path,
     num_structures: int,
-    spg_info: str,
+    spg_distribution_type: str | list[int],
     Z: int,
 ):
     """
@@ -112,7 +113,7 @@ def create_gnrs_config(
         mol_name: Identifier for the molecule being processed
         geometry_path: Path to the molecular geometry file (XYZ, SDF, etc.)
         num_structures: Number of crystal structures to generate
-        spg_info: Space group specification (number or list or "standard")
+        spg_distribution_type: Space group distribution type ("standard" or custom list[int])
         Z: Number of molecules per unit cell (Z-value)
 
     Side Effects:
@@ -123,10 +124,10 @@ def create_gnrs_config(
         config.read_file(config_file)
 
     config["master"]["name"] = mol_name
-    config["master"]["molecule_path"] = str(geometry_path)
+    config["master"]["molecule_path"] = json.dumps([str(geometry_path)])
     config["master"]["Z"] = str(Z)
     config["generation"]["num_structures_per_spg"] = str(num_structures)
-    config["generation"]["spg_distribution_type"] = spg_info
+    config["generation"]["spg_distribution_type"] = spg_distribution_type
 
     with open(output_dir / "ui.conf", "w") as f:
         config.write(f)
@@ -143,9 +144,11 @@ def create_genarris_jobs(
     logger = get_central_logger()
     logger.info(f"Starting Genarris generation for {mol_info['name']}")
 
-    gnrs_base_config = gnrs_config.get("base_config")
+    gnrs_base_config = gnrs_config.get("genarris_base_config")
     if gnrs_base_config is None:
-        raise KeyError("Genarris 'base_config' section is missing in the config file.")
+        raise KeyError(
+            "Genarris 'genarris_base_config' section is missing in the config file."
+        )
     logger.info(f"Using Genarris base configuration: {gnrs_base_config}")
 
     # Parameters for each Genarris run
@@ -159,24 +162,24 @@ def create_genarris_jobs(
 
     z_list = [str(z) for z in gnrs_vars.get("Z", [1])]
     num_structures_per_spg = gnrs_vars.get("num_structures_per_spg", 500)
-    spg_info = gnrs_vars.get("spg_info", "standard")
+    spg_distribution_type = gnrs_vars.get("spg_distribution_type", "standard")
 
     # molecule specific spg and z_list info from csv file if provided
     # mol_info["z"] is a list of z_values
     # mol_info["spg"] is a list of lists of spg values that correspond to each z_value
     if gnrs_vars.get("read_z_from_file", False):
         z_list = [str(z) for z in ast.literal_eval(mol_info["z"])]
     if gnrs_vars.get("read_spg_from_file", False):
-        spg_info = ast.literal_eval(mol_info["spg"])
-    if spg_info == "standard":
-        spg_info = ["standard"] * len(z_list)
-    if isinstance(spg_info[0], int):
-        spg_info = [spg_info]
-    if len(spg_info) == 1 and len(z_list) > 1:
-        spg_info = spg_info * len(z_list)
-    if len(spg_info) != len(z_list):
+        spg_distribution_type = ast.literal_eval(mol_info["spg"])
+    if spg_distribution_type == "standard":
+        spg_distribution_type = ["standard"] * len(z_list)
+    if isinstance(spg_distribution_type[0], int):
+        spg_distribution_type = [spg_distribution_type]
+    if len(spg_distribution_type) == 1 and len(z_list) > 1:
+        spg_distribution_type = spg_distribution_type * len(z_list)
+    if len(spg_distribution_type) != len(z_list):
         raise ValueError(
-            f"Length of spg_info {spg_info} does not match length of z_list {z_list} for molecule {mol_info['name']}."
+            f"Length of spg_distribution_type {spg_distribution_type} does not match length of z_list {z_list} for molecule {mol_info['name']}."
         )
     mol = mol_info["name"]  # System name
 
@@ -222,7 +225,7 @@ def create_genarris_jobs(
                     geometry_path=new_conf_path,
                     Z=z,
                     num_structures=num_structures_per_spg,
-                    spg_info=str(spg_info[i]),
+                    spg_distribution_type=str(spg_distribution_type[i]),
                 )
 
             # Create SLURM submission script if it doesn't exist

src/fairchem/core/common/utils.py

@@ -75,7 +75,13 @@ def decorator(func):
         @wraps(func)
         def cls_method(self, *args, **kwargs):
             f = func
-            if self.regress_forces and not getattr(self, "direct_forces", 0):
+            if hasattr(self, "regress_config"):
+                regress_forces = self.regress_config.forces
+                direct_forces = self.regress_config.direct_forces
+            else:
+                regress_forces = self.regress_forces
+                direct_forces = getattr(self, "direct_forces", 0)
+            if regress_forces and not direct_forces:
                 f = dec(func)
             return f(self, *args, **kwargs)
 

src/fairchem/core/models/base.py

@@ -161,13 +161,6 @@ def tasks(self) -> dict[str, Task]:
         """
         return self._tasks
 
-    @property
-    def direct_forces(self) -> bool:
-        """
-        Whether this model uses direct force prediction.
-        """
-        return getattr(self.backbone, "direct_forces", False)
-
     @property
     def dataset_to_tasks(self) -> dict[str, list]:
         """
@@ -191,8 +184,10 @@ def _validate_task_compatibility(self, task: Task) -> None:
         """
         derivative_properties = ("forces", "stress", "hessian")
 
+        backbone_regress_config = getattr(self.backbone, "regress_config", None)
         if (
-            self.direct_forces
+            backbone_regress_config is not None
+            and backbone_regress_config.direct_forces
             and task.inference_only
             and task.property in derivative_properties
         ):

src/fairchem/core/models/uma/escn_md.py

@@ -512,18 +512,6 @@ def __init__(
         )
         self.register_buffer("coefficient_index", coefficient_index, persistent=False)
 
-    @property  # deprecate this
-    def direct_forces(self) -> bool:
-        return self.regress_config.direct_forces
-
-    @property  # deprecate this
-    def regress_forces(self) -> bool:
-        return self.regress_config.forces
-
-    @property  # deprecate this
-    def regress_stress(self) -> bool:
-        return self.regress_config.stress
-
     def balance_channels(
         self,
         x_message_prime: torch.Tensor,
@@ -903,7 +891,7 @@ def get_default_untrained_tasks(
         stress computation requires energy-conserving force computation.
         """
         # Direct force models can't compute stress via autograd
-        if self.direct_forces:
+        if self.regress_config.direct_forces:
             return []
 
         tasks = []
@@ -1065,14 +1053,6 @@ def __init__(
         backbone.force_block = None
         self.regress_config = backbone.regress_config
 
-    @property
-    def regress_forces(self) -> bool:
-        return self.regress_config.forces
-
-    @property
-    def regress_stress(self) -> bool:
-        return self.regress_config.stress
-
     @conditional_grad(torch.enable_grad())
     def forward(
         self, data: AtomicData, emb: dict[str, torch.Tensor]

src/fairchem/core/models/uma/escn_moe.py

@@ -449,14 +449,6 @@ def __init__(
         self.merged_on_dataset = None
         self.non_merged_dataset_names: list[str] = []
 
-    @property
-    def regress_forces(self) -> bool:
-        return self.regress_config.forces
-
-    @property
-    def regress_stress(self) -> bool:
-        return self.regress_config.stress
-
     @staticmethod
     def _build_expert_mapping(
         dataset_names: list[str] | None,
@@ -608,14 +600,6 @@ def __init__(
         # keep track if this head has been merged or not
         self.merged_on_dataset = None
 
-    @property
-    def regress_forces(self) -> bool:
-        return self.regress_config.forces
-
-    @property
-    def regress_stress(self) -> bool:
-        return self.regress_config.stress
-
     def merge_MOLE_model(self, data):
         self.merged_on_dataset = data.dataset[0]
         self.non_merged_dataset_names = [