Phase weight rework 2

src/dara/hardware.py

@@ -0,0 +1,96 @@
+"""Portable CPU core-count detection for Ray + BGMN concurrency sizing.
+
+Priority, most authoritative first:
+
+1. An explicit SLURM allocation (``SLURM_CPUS_PER_TASK``, then
+   ``SLURM_CPUS_ON_NODE``) -- trusts a scheduler-assigned budget over
+   anything detected locally.
+2. OS-level CPU affinity (``os.sched_getaffinity``, Linux only), narrowed by
+   a cgroup CPU quota if one is set and implies fewer cores than affinity --
+   this matters on shared workstations, containers, and HPC nodes.
+3. ``os.cpu_count()`` -- universal fallback, also the only option on macOS,
+   where ``sched_getaffinity`` does not exist.
+"""
+
+from __future__ import annotations
+
+import math
+import os
+from pathlib import Path
+
+
+def _slurm_cpu_count() -> int | None:
+    for var in ("SLURM_CPUS_PER_TASK", "SLURM_CPUS_ON_NODE"):
+        value = os.environ.get(var)
+        if not value:
+            continue
+        try:
+            n = int(value)
+        except ValueError:
+            continue
+        if n > 0:
+            return n
+    return None
+
+
+def _cgroup_quota_cpu_count() -> int | None:
+    """Return the core count implied by a cgroup CPU quota, if one is set."""
+    cgroup_v2_path = Path("/sys/fs/cgroup/cpu.max")
+    if cgroup_v2_path.exists():
+        try:
+            quota_str, period_str = cgroup_v2_path.read_text().split()
+        except (OSError, ValueError):
+            quota_str = period_str = None
+        if quota_str is not None and quota_str != "max":
+            try:
+                quota, period = int(quota_str), int(period_str)
+                if quota > 0 and period > 0:
+                    return max(1, math.floor(quota / period))
+            except ValueError:
+                pass
+
+    quota_path = Path("/sys/fs/cgroup/cpu/cpu.cfs_quota_us")
+    period_path = Path("/sys/fs/cgroup/cpu/cpu.cfs_period_us")
+    if quota_path.exists() and period_path.exists():
+        try:
+            quota = int(quota_path.read_text().strip())
+            period = int(period_path.read_text().strip())
+            if quota > 0 and period > 0:
+                return max(1, math.floor(quota / period))
+        except (OSError, ValueError):
+            pass
+
+    return None
+
+
+def _affinity_cpu_count() -> int | None:
+    sched_getaffinity = getattr(os, "sched_getaffinity", None)
+    if sched_getaffinity is None:
+        return None
+    try:
+        affinity = len(sched_getaffinity(0))
+    except OSError:
+        return None
+    return affinity if affinity > 0 else None
+
+
+def detect_available_cores() -> int:
+    """
+    Detect the number of CPU cores actually available to this process.
+
+    Returns
+    -------
+        the detected core count, always >= 1. Never returns None or 0.
+    """
+    slurm = _slurm_cpu_count()
+    if slurm is not None:
+        return slurm
+
+    affinity = _affinity_cpu_count()
+    if affinity is not None:
+        cgroup = _cgroup_quota_cpu_count()
+        if cgroup is not None:
+            return max(1, min(affinity, cgroup))
+        return affinity
+
+    return max(1, os.cpu_count() or 1)

src/dara/search/core.py

@@ -11,6 +11,7 @@
 
 from dara.search.data_model import PeakMatchingStrategy
 from dara.search.tree import BaseSearchTree, SearchTree
+from dara.settings import DaraSettings
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -26,7 +27,21 @@
     "b1": "0_0^0.005",
     "rp": 4,
 }
-DEFAULT_REFINEMENT_PARAMS = {"n_threads": 8, "eps1": 0, "eps2": "0_-0.05^0.05"}
+
+# Portable CPU/thread budget: BGMN threads default to 1 (see
+# DaraSettings.BGMN_N_THREADS) so worker *concurrency*, not per-task
+# threading, scales with available cores (see dara.hardware and the
+# num_cpus=... passed to ray.init() below and to each refinement task's
+# .options(num_cpus=...) in dara.search.tree.batch_refinement). Both are
+# overridable via DaraSettings (~/.dara.yaml or the DARA_RAY_NUM_CPUS /
+# DARA_BGMN_N_THREADS env vars).
+_settings = DaraSettings()
+
+DEFAULT_REFINEMENT_PARAMS = {
+    "n_threads": _settings.BGMN_N_THREADS,
+    "eps1": 0,
+    "eps2": "0_-0.05^0.05",
+}
 DEFAULT_PEAK_MATCHING_STRATEGY = PeakMatchingStrategy.default()
 
 
@@ -100,7 +115,12 @@ def search_phases(
         refinement_params = {}
 
     if not ray.is_initialized():
-        ray.init(runtime_env={"working_dir": None})
+        # num_cpus is set explicitly to the detected/allocated core count
+        # (DaraSettings.RAY_NUM_CPUS) rather than left to Ray's own
+        # auto-detection, which just calls os.cpu_count() and would ignore
+        # an explicit SLURM allocation or cgroup/container CPU affinity
+        # restriction (see dara.hardware.detect_available_cores).
+        ray.init(runtime_env={"working_dir": None}, num_cpus=_settings.RAY_NUM_CPUS)
 
     phase_params = {**DEFAULT_PHASE_PARAMS, **phase_params}
     refinement_params = {**DEFAULT_REFINEMENT_PARAMS, **refinement_params}

src/dara/search/tree.py

@@ -20,6 +20,7 @@
 from dara.refine import RefinementPhase
 from dara.search.data_model import PeakMatchingStrategy, SearchNodeData, SearchResult
 from dara.search.peak_matcher import PeakMatcher
+from dara.settings import DaraSettings
 from dara.utils import (
     estimate_rpb_threshold,
     find_optimal_intensity_threshold,
@@ -39,8 +40,20 @@
 
 logger = get_logger(__name__, level="INFO")
 
+# Defensive fallback for callers that build a SearchTree directly (bypassing
+# search_phases()'s merge of DEFAULT_REFINEMENT_PARAMS, which always sets
+# n_threads) without an explicit n_threads in refinement_params. See
+# dara.hardware and dara.search.core for the portable core/thread sizing
+# this is part of.
+_DEFAULT_BGMN_N_THREADS = DaraSettings().BGMN_N_THREADS
 
-@ray.remote(num_cpus=1)
+
+# num_cpus is intentionally not fixed here -- it is set per-submission via
+# .options(num_cpus=...) in batch_refinement() to match the actual n_threads
+# a given task's BGMN subprocess will use, so Ray's scheduler throttles
+# concurrency to roughly floor(cores / n_threads) instead of assuming a
+# hardcoded num_cpus=1 regardless of how many OS threads the task spawns.
+@ray.remote
 def remote_do_refinement_no_saving(
     pattern_path: Path,
     cif_paths: list[Path],
@@ -129,8 +142,9 @@ def batch_refinement(
     phase_params: dict[str, ...] | None = None,
     refinement_params: dict[str, float] | None = None,
 ) -> list[RefinementResult]:
+    n_threads = (refinement_params or {}).get("n_threads", _DEFAULT_BGMN_N_THREADS)
     handles = [
-        remote_do_refinement_no_saving.remote(
+        remote_do_refinement_no_saving.options(num_cpus=n_threads).remote(
             pattern_path,
             cif_paths,
             wavelength=wavelength,
@@ -324,6 +338,64 @@ def remove_unnecessary_phases(
     return new_phases
 
 
+# Minimum relative Rwp improvement over the parent node required to keep an
+# out-of-order-intensity branch instead of pruning it (see
+# `should_prune_low_weight_fraction`). Lower Rwp is better, so relative
+# improvement is (parent_rwp - child_rwp) / parent_rwp; 0.08 means the child
+# must fit at least 8% better than its parent, relatively, to survive.
+LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT = 0.08
+
+
+def should_prune_low_weight_fraction(
+    intensity_out_of_order: bool,
+    parent_rwp: float | None,
+    child_rwp: float,
+    material_improvement_threshold: float = LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT,
+) -> bool:
+    """
+    Decide whether an out-of-order-intensity branch should be pruned.
+
+    Phases are expected to be discovered in roughly decreasing order of
+    abundance (the naive peak-match search tends to find the biggest
+    contributor to the pattern first), so a newly-added phase that turns out
+    to have *more* calculated peak intensity than a phase added earlier in
+    the same branch (`intensity_out_of_order`) is treated as suspicious by
+    default -- it can mean the refinement is using this phase to absorb
+    residual intensity that doesn't belong to it (overfitting), rather than
+    genuinely explaining new pattern features.
+
+    But that's only actually suspicious if the phase isn't earning its keep.
+    If adding it materially improved the fit (Rwp) over the parent node, the
+    improvement is real signal, not an ordering artifact, regardless of what
+    order the search happened to try phases in -- so the branch is kept.
+
+    Args:
+        intensity_out_of_order: whether the newly-added phase has more
+            calculated peak intensity than an earlier phase in the same
+            branch (the raw ordering signal; computed by the caller).
+        parent_rwp: the Rwp of the parent node, or None if the parent is the
+            (unrefined) root -- with no parent to compare against, the
+            ordering signal alone determines the outcome.
+        child_rwp: the Rwp of the new (child) node being evaluated.
+        material_improvement_threshold: the minimum relative Rwp
+            improvement required to treat an out-of-order branch as
+            materially improved rather than pruning it.
+
+    Returns
+    -------
+        True if the branch should be pruned as a low-weight-fraction
+        (suspicious, non-material) out-of-order addition; False if it
+        should be kept.
+    """
+    if not intensity_out_of_order:
+        return False
+    if parent_rwp is None or parent_rwp <= 0:
+        return True
+
+    relative_improvement = (parent_rwp - child_rwp) / parent_rwp
+    return relative_improvement < material_improvement_threshold
+
+
 def get_natural_break_results(
     results: list[SearchResult], sorting: bool = True
 ) -> list[SearchResult]:
@@ -491,9 +563,13 @@ def expand_node(self, nid: str) -> list[str]:
                         ]["intensity"].sum(),
                         reverse=True,
                     )
-                    # make sure the newly added phase has the lowest peak intensity
-                    is_low_weight_fraction = (
-                        sorted_searched_phases[-1] != searched_phases[-1]
+                    # the newly added phase does not have the lowest peak intensity
+                    intensity_out_of_order = sorted_searched_phases[-1] != searched_phases[-1]
+                    parent_rwp = node.data.current_result.lst_data.rwp if node.data.current_result is not None else None
+                    is_low_weight_fraction = should_prune_low_weight_fraction(
+                        intensity_out_of_order,
+                        parent_rwp,
+                        new_result.lst_data.rwp,
                     )
                 else:
                     is_low_weight_fraction = False

src/dara/settings.py

@@ -6,6 +6,8 @@
 from pydantic import Field, model_validator
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
+from dara.hardware import detect_available_cores
+
 _DEFAULT_CONFIG_FILE_PATH = "~/.dara.yaml"
 
 
@@ -23,6 +25,29 @@ class DaraSettings(BaseSettings):
     PATH_TO_ICSD: Path = Field(Path("~/ICSD_2024/ICSD_2024_experimental_inorganic/experimental_inorganic").expanduser())
     PATH_TO_COD: Path = Field(Path("~/COD_2024").expanduser())
 
+    RAY_NUM_CPUS: int = Field(
+        default_factory=detect_available_cores,
+        description=(
+            "Number of CPUs to give Ray's own resource accounting (passed as "
+            "num_cpus= to ray.init()). Auto-detected from an explicit SLURM "
+            "allocation, then OS-level CPU affinity/cgroup quota, then the "
+            "total CPU count -- see dara.hardware.detect_available_cores. "
+            "Override via the DARA_RAY_NUM_CPUS env var or ~/.dara.yaml."
+        ),
+    )
+    BGMN_N_THREADS: int = Field(
+        1,
+        description=(
+            "Threads per BGMN refinement subprocess (the n_threads refinement "
+            "param). Kept low (default 1) so worker concurrency -- not "
+            "per-task threading -- scales with available cores: refinement "
+            "Ray tasks are submitted with num_cpus=BGMN_N_THREADS, so "
+            "concurrency stays near floor(cores / BGMN_N_THREADS) "
+            "automatically even if this is overridden upward. Override via "
+            "the DARA_BGMN_N_THREADS env var or ~/.dara.yaml."
+        ),
+    )
+
     model_config = SettingsConfigDict(env_prefix="dara_")  # prepend dara_ to env vars
 
     @model_validator(mode="before")