ENH: Provide a pyFAI Geometry Optimization Task

lute/execution/executor.py

@@ -1151,7 +1151,7 @@ def _submit_cmd(self, executable_path: str, params: str) -> str:
         nprocs: int = max(
             int(os.environ.get("SLURM_NPROCS", len(os.sched_getaffinity(0)))) - 1, 1
         )
-        mpi_cmd: str = f"mpirun -np {nprocs}"
+        mpi_cmd: str = f"mpirun -np {nprocs} --map-by core"
         if __debug__:
             py_cmd = f"python -B -u -m mpi4py.run {executable_path} {params}"
         else:

lute/execution/logging.py

@@ -94,10 +94,22 @@ def get_logger(name: str, is_task: bool = True) -> logging.Logger:
             other_logger, logging.PlaceHolder
         ):
             other_logger.disabled = True
+        elif "pyFAI" in other_name and not isinstance(
+            other_logger, logging.PlaceHolder
+        ):
+            other_logger.disabled = True
+        elif "goniometer" in other_name and not isinstance(
+            other_logger, logging.PlaceHolder
+        ):
+            other_logger.disabled = True
         elif "numpy" in other_name and not isinstance(
             other_logger, logging.PlaceHolder
         ):
             other_logger.setLevel(logging.CRITICAL)
+        elif "numexpr.utils" in other_name and not isinstance(
+            other_logger, logging.PlaceHolder
+        ):
+            other_logger.setLevel(logging.WARNING)
     logger: logging.Logger = logging.getLogger(name)
     logger.propagate = False
     handler: logging.Handler

lute/io/models/__init__.py

@@ -11,3 +11,4 @@
 from .tests import *
 from .mpi_tests import *
 from .geometry import *
+from .geom_opt import *

lute/io/models/geom_opt.py

@@ -0,0 +1,160 @@
+"""Models for optimizing detector geometry using PyFAI and Bayesian optimization.
+
+Classes:
+    - OptimizePyFAIGeometryParameters(TaskParameters):
+        Parameters for optimizing detector geometry using PyFAI and Bayesian optimization.
+"""
+
+__all__ = ["OptimizePyFAIGeometryParameters"]
+__author__ = "Louis Conreux"
+
+from typing import Any, Dict, Optional, Union, Tuple
+
+from pydantic import BaseModel, Field, validator
+
+from lute.io.models.base import TaskParameters
+from lute.io.models.validators import validate_smd_path, validate_calib_path
+
+
+class OptimizePyFAIGeometryParameters(TaskParameters):
+    """Parameters for optimizing detector geometry using PyFAI and Bayesian optimization.
+
+    The Bayesian Optimization has default hyperparameters that can be overriden by the user.
+    """
+
+    class Config(TaskParameters.Config):
+        set_result: bool = True
+        """Whether the Executor should mark a specified parameter as a result."""
+
+    class BayesGeomOptParameters(BaseModel):
+        """Bayesian optimization hyperparameters."""
+
+        bounds: Dict[str, Union[float, Tuple[float, float]]] = Field(
+            {
+                "dist": (0.02, 0.6),
+                "poni1": (-0.01, 0.01),
+                "poni2": (-0.01, 0.01),
+            },
+            description="Bounds defining the parameter search space for the Bayesian optimization.",
+        )
+
+        res: float = Field(
+            None,
+            description="Resolution of the grid used to discretize the parameter search space.",
+        )
+
+        max_rings: int = Field(
+            5,
+            description="Maximum number of rings to be used for the Bayesian optimization.",
+        )
+
+        n_samples: Optional[int] = Field(
+            50,
+            description="Number of random starts to initialize the Bayesian optimization.",
+        )
+
+        n_iterations: Optional[int] = Field(
+            50,
+            description="Number of iterations to run the Bayesian optimization.",
+        )
+
+        prior: Optional[bool] = Field(
+            True,
+            description="Whether to use a gaussian prior centered on the search space for the Bayesian optimization or randomly pick samples.",
+        )
+
+        af: Optional[str] = Field(
+            "ucb",
+            description="Acquisition function to be used by the Bayesian optimization.",
+        )
+
+        hyperparams: Optional[Dict[str, float]] = Field(
+            {
+                "beta": 1.96,
+                "epsilon": 0.01,
+            },
+            description="Hyperparameters for the acquisition function.",
+        )
+
+        seed: Optional[int] = Field(
+            None,
+            description="Seed for the random number generator for potential reproducibility.",
+        )
+
+    _find_smd_path = validate_smd_path("powder")
+
+    _find_in_file_path = validate_calib_path("in_file")
+
+    exp: str = Field(
+        "",
+        description="Experiment name.",
+    )
+
+    run: Union[str, int] = Field(
+        None,
+        description="Run number.",
+    )
+
+    det_type: str = Field(
+        "",
+        description="Detector type. Currently supported: 'ePix10k2M', 'ePix10kaQuad', 'Rayonix', 'Jungfrau1M', 'Jungfrau4M'",
+    )
+
+    work_dir: str = Field(
+        "",
+        description="Main working directory for LUTE.",
+    )
+
+    in_file: str = Field(
+        "",
+        description="Path to the input .data file containing the detector geometry info to be calibrated.",
+    )
+
+    powder: str = Field(
+        "",
+        description="Powder diffraction pattern to be used for the calibration.",
+    )
+
+    calibrant: str = Field(
+        "",
+        description="Calibrant used for the calibration supported by pyFAI: https://github.com/silx-kit/pyFAI/tree/main/src/pyFAI/resources/calibration",
+    )
+
+    out_file: str = Field(
+        "",
+        description="Path to the output .data file containing the optimized detector geometry.",
+        is_result=True,
+    )
+
+    bo_params: BayesGeomOptParameters = Field(
+        BayesGeomOptParameters(),
+        description="Bayesian optimization parameters containing bounds and resolution for defining space search and hyperparameters.",
+    )
+
+    @validator("exp", always=True)
+    def validate_exp(cls, exp: str, values: Dict[str, Any]) -> str:
+        if not exp:
+            exp = values["lute_config"].experiment
+        return exp
+
+    @validator("run", always=True)
+    def validate_run(
+        cls, run: Union[str, int], values: Dict[str, Any]
+    ) -> Union[str, int]:
+        if not run:
+            run = values["lute_config"].run
+        return run
+
+    @validator("work_dir", always=True)
+    def validate_work_dir(cls, work_dir: str, values: Dict[str, Any]) -> str:
+        if not work_dir:
+            work_dir = values["lute_config"].work_dir
+        return work_dir
+
+    @validator("out_file", always=True)
+    def validate_out_file(cls, out_file: str, values: Dict[str, Any]) -> str:
+        if not out_file:
+            in_file = values["in_file"]
+            run = values["run"]
+            out_file = in_file.replace("0-end.data", f"{run}-end.data")
+        return out_file

lute/io/models/validators.py

@@ -5,8 +5,8 @@
         parameters for validation.
 """
 
-__all__ = ["template_parameter_validator", "validate_smd_path"]
-__author__ = "Gabriel Dorlhiac"
+__all__ = ["template_parameter_validator", "validate_smd_path", "validate_calib_path"]
+__author__ = ["Gabriel Dorlhiac", "Louis Conreux"]
 
 import os
 from typing import Dict, Any, Optional
@@ -15,6 +15,9 @@
 
 from lute.io.db import read_latest_db_entry
 
+import psana  # type: ignore
+from PSCalib.CalibFileFinder import find_calib_file  # type: ignore
+
 
 def template_parameter_validator(template_params_name: str):
     """Populates a TaskParameters model with a set of validated TemplateParameters.
@@ -68,3 +71,31 @@ def _validate_smd_path(cls, smd_path: str, values: Dict[str, Any]) -> str:
         return smd_path
 
     return validator(smd_path_name, always=True, allow_reuse=True)(_validate_smd_path)
+
+
+def validate_calib_path(calib_path_name: str):
+    """Finds the path to a valid calibration file or raises an error."""
+
+    def _validate_calib_path(cls, calib_path: str, values: Dict[str, Any]) -> str:
+        if calib_path == "":
+            exp: str = values["lute_config"].experiment
+            run: int = int(values["lute_config"].run)
+            try:
+                det_type: str = values["det_type"]
+            except KeyError:
+                det_type = values["detname"]
+            cdir = f"/sdf/data/lcls/ds/{exp[:3]}/{exp}/calib"
+            ds_args = f"exp={exp}:run={run}:idx"
+            ds = psana.DataSource(ds_args)
+            det = psana.Detector(det_type, ds.env())
+            src = str(det.name)
+            type = "geometry"
+            calib_path = find_calib_file(cdir, src, type, run, pbits=1)
+            if calib_path is None:
+                raise ValueError(f"No calibration file found for {det_type} in {cdir}")
+
+        return calib_path
+
+    return validator(calib_path_name, always=True, allow_reuse=True)(
+        _validate_calib_path
+    )

lute/managed_tasks.py

@@ -130,3 +130,8 @@
 
 PeakFinderPsocake: Executor = Executor("FindPeaksPsocake")
 """Performs Bragg peak finding using psocake - *DEPRECATED*."""
+
+# PyFAI
+#######
+PyFAIGeometryOptimizer: MPIExecutor = MPIExecutor("OptimizePyFAIGeometry")
+"""Optimize detector geometry using PyFAI coupled with Bayesian Optimization."""

lute/tasks/__init__.py

@@ -86,6 +86,11 @@ def import_task(task_name: str) -> Type[Task]:
 
         return TestMultiNodeCommunication
 
+    if task_name == "OptimizePyFAIGeometry":
+        from .geom_opt import OptimizePyFAIGeometry
+
+        return OptimizePyFAIGeometry
+
     if task_name == "OptimizeAgBhGeometryExhaustive":
         from .geometry import OptimizeAgBhGeometryExhaustive