From 8c89a783b044615f1bed97898d276e99a2456017 Mon Sep 17 00:00:00 2001 From: lauren-walters Date: Wed, 29 Jul 2026 21:57:11 -0700 Subject: [PATCH 1/5] Bound BGMN threads per refinement and size Ray concurrency to cores Ray's num_cpus=1 default per refinement task let as many concurrent tasks run as detected CPUs, decoupled from the n_threads a task's BGMN subprocess actually used -- at the old default (n_threads=8, 12 concurrent tasks) this oversubscribed the machine ~8-10x in OS threads, which can starve Ray's own GCS/raylet under sustained load. - Add dara.hardware.detect_available_cores(): portable core-count detection (SLURM allocation, then OS affinity/cgroup quota, then os.cpu_count()). - Add DaraSettings.RAY_NUM_CPUS (default: detected cores) and DaraSettings.BGMN_N_THREADS (default: 1), both overridable via DARA_RAY_NUM_CPUS / DARA_BGMN_N_THREADS env vars or ~/.dara.yaml. - search_phases()'s ray.init() now passes num_cpus=RAY_NUM_CPUS explicitly instead of relying on Ray's own auto-detection (which ignores SLURM allocations and cgroup/container limits). - Each refinement Ray task now requests num_cpus=n_threads via .options(...) at submission, instead of a hardcoded @ray.remote(num_cpus=1), so Ray throttles concurrency to floor(cores / n_threads) automatically -- with defaults this is 12 workers x 1 thread on a 12-core machine; overriding n_threads upward safely drops concurrency to match. Threads/concurrency/config only -- no change to search logic, pruning, or node counts. Co-Authored-By: Claude Sonnet 5 --- src/dara/hardware.py | 96 +++++++++++++++++++++++++++++++++++++++++ src/dara/search/core.py | 24 ++++++++++- src/dara/search/tree.py | 18 +++++++- src/dara/settings.py | 25 +++++++++++ 4 files changed, 159 insertions(+), 4 deletions(-) create mode 100644 src/dara/hardware.py diff --git a/src/dara/hardware.py b/src/dara/hardware.py new file mode 100644 index 00000000..cef2f811 --- /dev/null +++ b/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) diff --git a/src/dara/search/core.py b/src/dara/search/core.py index fc2167b7..d1d1362f 100644 --- a/src/dara/search/core.py +++ b/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} diff --git a/src/dara/search/tree.py b/src/dara/search/tree.py index e55e91e7..45ec5798 100644 --- a/src/dara/search/tree.py +++ b/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, diff --git a/src/dara/settings.py b/src/dara/settings.py index 85a58bb4..ecee36f7 100644 --- a/src/dara/settings.py +++ b/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") From 588ec0d2b3c2c2047a841d7c4c3cd12c5b8bd3d8 Mon Sep 17 00:00:00 2001 From: lauren-walters Date: Thu, 30 Jul 2026 15:58:25 -0700 Subject: [PATCH 2/5] Keep out-of-order-intensity branches when they materially improve fit The intensity-order pruning heuristic assumed phases are discovered in decreasing order of abundance, terminating any branch where a newly-added phase had more calculated peak intensity than one added earlier -- even when that branch's overall fit was clearly better, discarding the correct answer. - Add should_prune_low_weight_fraction(intensity_out_of_order, parent_rwp, child_rwp, material_improvement_threshold=LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT), a small standalone helper in dara.search.tree. The intensity-order signal is still computed exactly as before; the helper only decides whether it's fatal. LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT defaults to 0.10: an out-of-order branch is now pruned only if it does NOT improve Rwp by at least 10% relative to its parent. With no parent to compare against (or a degenerate parent_rwp <= 0), the ordering signal alone still decides, same as before. - Wire the helper into BaseSearchTree.expand_node() in place of the old inline boolean check. - Add tests/test_search_tree.py: 9 direct unit tests of the helper plus 3 integration-level tests driving the real expand_node() (score_phases/ refine_phases mocked to avoid needing BGMN; calculate_fom_and_strain patched since it only reads a phase's own CIF file) confirming a materially-better branch is retained, a marginally-better one is still pruned, and in-order additions are never flagged regardless of fit. Decision-only change: no recovery/branch-respawning mechanism, no change to how many nodes/branches get explored beyond retaining branches this rule would previously have discarded. On a deterministic 30-CIF subset of dara-clustering-problem-2 (first 30 alphabetically), before/after node count, statuses, and best Rwp are identical (15 nodes, best Rwp 29.73) -- the out-of-order case that arose there didn't clear the 10% bar either way, consistent with the unit/integration tests being the primary proof this works as designed. Co-Authored-By: Claude Sonnet 5 --- src/dara/search/tree.py | 72 +++++++++- tests/test_search_tree.py | 291 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 361 insertions(+), 2 deletions(-) create mode 100644 tests/test_search_tree.py diff --git a/src/dara/search/tree.py b/src/dara/search/tree.py index 45ec5798..59f919d2 100644 --- a/src/dara/search/tree.py +++ b/src/dara/search/tree.py @@ -338,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.10 means the child +# must fit at least 10% better than its parent, relatively, to survive. +LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT = 0.10 + + +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]: @@ -505,10 +563,20 @@ 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 = ( + # 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 diff --git a/tests/test_search_tree.py b/tests/test_search_tree.py new file mode 100644 index 00000000..b79964fc --- /dev/null +++ b/tests/test_search_tree.py @@ -0,0 +1,291 @@ +import unittest +from pathlib import Path +from unittest.mock import patch + +import pandas as pd +from treelib import Node, Tree + +from dara.refine import RefinementPhase +from dara.result import DiaResult, LstResult, PhaseResult, RefinementResult +from dara.search.data_model import SearchNodeData +from dara.search.tree import ( + LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT, + BaseSearchTree, + should_prune_low_weight_fraction, +) + + +def _phase(name: str) -> RefinementPhase: + return RefinementPhase.make(Path(f"/fake/cifs/{name}.cif")) + + +def _phase_result() -> PhaseResult: + """A structurally-valid but otherwise-meaningless PhaseResult. + + `calculate_fom_and_strain` (the only code that reads these fields) is + mocked out in the integration test below, so content doesn't matter -- + it only needs to satisfy LstResult.phases_results' pydantic validation. + """ + return PhaseResult( + SpacegroupNo=None, + HermannMauguin=None, + XrayDensity=None, + Rphase=None, + UNIT="nm", + GEWICHT=1.0, + GEWICHT_NAME=None, + ) + + +def _make_refinement_result( + rwp: float, + phase_intensities: dict[str, float], +) -> RefinementResult: + """Build a minimal but genuinely valid (pydantic-validated) RefinementResult. + + `phase_intensities` maps phase stem -> total calculated peak intensity + for that phase; both `peak_data` (read by the intensity-ordering check) + and `plot_data.structs` (read by `remove_unnecessary_phases`) are built + from it so removing any one phase changes the calculated pattern enough + for `remove_unnecessary_phases` to consider it necessary. + """ + peak_rows = [ + {"phase": phase, "2theta": 10.0 + i, "intensity": intensity} + for i, (phase, intensity) in enumerate(phase_intensities.items()) + ] + structs = {phase: [intensity] for phase, intensity in phase_intensities.items()} + y_calc_total = sum(phase_intensities.values()) + + return RefinementResult( + lst_data=LstResult( + raw_lst="", + pattern_name="fake", + num_steps=1, + Rp=0.0, + Rpb=rwp, + R=0.0, + Rwp=rwp, + Rexp=1.0, + d=0.0, + **{"1-rho": rwp}, + phases_results={phase: _phase_result() for phase in phase_intensities}, + ), + plot_data=DiaResult( + x=[0.0], + y_obs=[y_calc_total], + y_calc=[y_calc_total], + y_bkg=[0.0], + structs=structs, + ), + peak_data=pd.DataFrame(peak_rows), + ) + + +class TestShouldPruneLowWeightFraction(unittest.TestCase): + """Direct unit tests for the standalone pruning-decision helper.""" + + def test_in_order_never_pruned(self): + """Not out-of-order: the parent/child fit is irrelevant, never prune.""" + self.assertFalse( + should_prune_low_weight_fraction( + intensity_out_of_order=False, parent_rwp=40.0, child_rwp=39.9 + ) + ) + self.assertFalse( + should_prune_low_weight_fraction( + intensity_out_of_order=False, parent_rwp=40.0, child_rwp=100.0 + ) + ) + self.assertFalse( + should_prune_low_weight_fraction( + intensity_out_of_order=False, parent_rwp=None, child_rwp=39.9 + ) + ) + + def test_out_of_order_no_parent_falls_back_to_ordering_only(self): + """With no parent to compare against (root's first child), the + ordering signal alone determines the outcome -- prune.""" + self.assertTrue( + should_prune_low_weight_fraction( + intensity_out_of_order=True, parent_rwp=None, child_rwp=10.0 + ) + ) + + def test_out_of_order_zero_or_negative_parent_rwp_is_pruned(self): + """Degenerate parent_rwp <= 0 can't support a relative-improvement + calculation (division by zero); fall back to the ordering-only rule.""" + self.assertTrue( + should_prune_low_weight_fraction( + intensity_out_of_order=True, parent_rwp=0.0, child_rwp=5.0 + ) + ) + self.assertTrue( + should_prune_low_weight_fraction( + intensity_out_of_order=True, parent_rwp=-1.0, child_rwp=5.0 + ) + ) + + def test_out_of_order_marginal_improvement_still_pruned(self): + """Out-of-order + improvement below the material threshold: pruned.""" + parent_rwp = 40.0 + # 5% relative improvement, below the 10% default threshold + child_rwp = 38.0 + self.assertTrue( + should_prune_low_weight_fraction( + intensity_out_of_order=True, parent_rwp=parent_rwp, child_rwp=child_rwp + ) + ) + + def test_out_of_order_material_improvement_kept(self): + """Out-of-order + material improvement: this is the bug fix -- a + branch that fits clearly better than its parent must survive even + though the newest phase isn't the smallest.""" + parent_rwp = 40.0 + # 50% relative improvement + child_rwp = 20.0 + self.assertFalse( + should_prune_low_weight_fraction( + intensity_out_of_order=True, parent_rwp=parent_rwp, child_rwp=child_rwp + ) + ) + + def test_improvement_exactly_at_threshold_is_kept(self): + """Exactly LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT relative improvement + counts as material (the check is `< threshold`, not `<= threshold`).""" + parent_rwp = 40.0 + child_rwp = parent_rwp * (1 - LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT) + self.assertFalse( + should_prune_low_weight_fraction( + intensity_out_of_order=True, parent_rwp=parent_rwp, child_rwp=child_rwp + ) + ) + + def test_improvement_just_under_threshold_is_pruned(self): + parent_rwp = 40.0 + child_rwp = parent_rwp * (1 - LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT) + 1e-6 + self.assertTrue( + should_prune_low_weight_fraction( + intensity_out_of_order=True, parent_rwp=parent_rwp, child_rwp=child_rwp + ) + ) + + def test_worse_fit_is_pruned(self): + """Out-of-order + the child fits WORSE than the parent (negative + 'improvement'): definitely pruned.""" + self.assertTrue( + should_prune_low_weight_fraction( + intensity_out_of_order=True, parent_rwp=20.0, child_rwp=30.0 + ) + ) + + def test_custom_threshold_is_respected(self): + parent_rwp, child_rwp = 40.0, 30.0 # 25% relative improvement + self.assertTrue( + should_prune_low_weight_fraction( + intensity_out_of_order=True, + parent_rwp=parent_rwp, + child_rwp=child_rwp, + material_improvement_threshold=0.5, + ) + ) + self.assertFalse( + should_prune_low_weight_fraction( + intensity_out_of_order=True, + parent_rwp=parent_rwp, + child_rwp=child_rwp, + material_improvement_threshold=0.2, + ) + ) + + +def _make_test_tree() -> BaseSearchTree: + """A minimal, real (treelib-backed) BaseSearchTree for driving the actual + `expand_node()` code path without needing a real pattern file or BGMN. + """ + tree = BaseSearchTree.__new__(BaseSearchTree) + Tree.__init__(tree) + tree.pinned_phases = [] + tree.all_phases_result = {} + tree.max_phases = 5 + tree.rpb_threshold = 0.5 + tree.record_peak_matcher_scores = False + tree.express_mode = False + tree.maximum_grouping_distance = 0.1 + tree.peak_obs = pd.DataFrame({"2theta": [10.0, 11.0], "intensity": [100.0, 50.0]})[ + ["2theta", "intensity"] + ].values + return tree + + +@patch("dara.search.tree.calculate_fom_and_strain", return_value=(0.5, 0.0)) +class TestExpandNodeLowWeightFractionIntegration(unittest.TestCase): + """Integration-level tests exercising the real `expand_node()` method + (real `remove_unnecessary_phases`, real intensity-ordering computation, + real status-decision branching) with `score_phases`/`refine_phases` + mocked to avoid needing BGMN, and `calculate_fom_and_strain` mocked + since it reads a phase's own CIF file (irrelevant to this decision). + """ + + def _run_expand(self, parent_rwp, child_rwp, phase_a_intensity, phase_b_intensity): + """Set up a root -> phase_a chain, then expand with phase_b added. + + phase_b's calculated intensity is set higher than phase_a's, so the + addition is out-of-order by construction; parent_rwp/child_rwp + control whether it's a material improvement. + """ + tree = _make_test_tree() + phase_a, phase_b = _phase("A"), _phase("B") + + parent_result = _make_refinement_result(parent_rwp, {"A": phase_a_intensity}) + root = Node( + data=SearchNodeData( + current_result=parent_result, + current_phases=[phase_a], + status="pending", + ) + ) + tree.add_node(root) + + child_result = _make_refinement_result( + child_rwp, {"A": phase_a_intensity, "B": phase_b_intensity} + ) + + with ( + patch.object(tree, "score_phases", return_value=([phase_b], {}, 0)), + patch.object(tree, "refine_phases", return_value={phase_b: child_result}), + ): + tree.expand_node(root.identifier) + + children = tree.children(root.identifier) + self.assertEqual(len(children), 1) + return children[0].data.status + + def test_material_improvement_is_retained(self, mock_fom): + """Out-of-order (B's intensity > A's) but a clearly better fit + (50% relative Rwp improvement): must NOT be pruned as low-weight-fraction.""" + status = self._run_expand( + parent_rwp=40.0, child_rwp=20.0, phase_a_intensity=60.0, phase_b_intensity=200.0 + ) + self.assertNotEqual(status, "low_weight_fraction") + self.assertEqual(status, "pending") + + def test_marginal_improvement_is_still_pruned(self, mock_fom): + """Out-of-order (B's intensity > A's) with only a marginal (5%) + Rwp improvement: still pruned as low-weight-fraction, as before.""" + status = self._run_expand( + parent_rwp=40.0, child_rwp=38.0, phase_a_intensity=60.0, phase_b_intensity=200.0 + ) + self.assertEqual(status, "low_weight_fraction") + + def test_in_order_addition_never_flagged_regardless_of_fit(self, mock_fom): + """B's intensity is LOWER than A's (properly ordered): never flagged, + even with a poor Rwp improvement -- the ordering check simply + doesn't fire, independent of the material-improvement carve-out.""" + status = self._run_expand( + parent_rwp=40.0, child_rwp=39.9, phase_a_intensity=200.0, phase_b_intensity=60.0 + ) + self.assertNotEqual(status, "low_weight_fraction") + + +if __name__ == "__main__": + unittest.main() From c74645e81f5323ddc3bc199296706fc509bb2afc Mon Sep 17 00:00:00 2001 From: lauren-walters Date: Thu, 30 Jul 2026 16:11:08 -0700 Subject: [PATCH 3/5] cleanup: lint + trim comments - ruff format: collapse two multi-line assignments in tree.py's expand_node() and the whole of the new tests/test_search_tree.py to match the project's configured formatter (line-length=120). No other files in this branch's diff needed reformatting (hardware.py and settings.py were already compliant); pre-existing formatting drift in core.py/tree.py outside this branch's added lines was left untouched. - ruff check across the whole repo: zero new findings in this branch's files. The 13 existing findings are all pre-existing, in files this branch never touched (notebooks, scripts/filter_cod.py, scripts/filter_icsd.py, tests/test_cif2str.py) -- left as-is per scope. - Reviewed this branch's comments/docstrings against origin/main...HEAD for LLM-flavored padding, step narration, and restated-the-obvious comments. Found none to remove: every comment already documents non-obvious rationale (the hardware-detection priority order, the LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT threshold's meaning, the .options(num_cpus=...) oversubscription-safety note) or is a concise, accurate docstring on a public helper. No comment/docstring text changed. No logic or behavior change. Co-Authored-By: Claude Sonnet 5 --- src/dara/search/tree.py | 10 ++---- tests/test_search_tree.py | 72 +++++++++------------------------------ 2 files changed, 18 insertions(+), 64 deletions(-) diff --git a/src/dara/search/tree.py b/src/dara/search/tree.py index 59f919d2..318c2641 100644 --- a/src/dara/search/tree.py +++ b/src/dara/search/tree.py @@ -564,14 +564,8 @@ def expand_node(self, nid: str) -> list[str]: reverse=True, ) # 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 - ) + 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, diff --git a/tests/test_search_tree.py b/tests/test_search_tree.py index b79964fc..e16fc807 100644 --- a/tests/test_search_tree.py +++ b/tests/test_search_tree.py @@ -87,43 +87,25 @@ class TestShouldPruneLowWeightFraction(unittest.TestCase): def test_in_order_never_pruned(self): """Not out-of-order: the parent/child fit is irrelevant, never prune.""" self.assertFalse( - should_prune_low_weight_fraction( - intensity_out_of_order=False, parent_rwp=40.0, child_rwp=39.9 - ) + should_prune_low_weight_fraction(intensity_out_of_order=False, parent_rwp=40.0, child_rwp=39.9) ) self.assertFalse( - should_prune_low_weight_fraction( - intensity_out_of_order=False, parent_rwp=40.0, child_rwp=100.0 - ) + should_prune_low_weight_fraction(intensity_out_of_order=False, parent_rwp=40.0, child_rwp=100.0) ) self.assertFalse( - should_prune_low_weight_fraction( - intensity_out_of_order=False, parent_rwp=None, child_rwp=39.9 - ) + should_prune_low_weight_fraction(intensity_out_of_order=False, parent_rwp=None, child_rwp=39.9) ) def test_out_of_order_no_parent_falls_back_to_ordering_only(self): """With no parent to compare against (root's first child), the ordering signal alone determines the outcome -- prune.""" - self.assertTrue( - should_prune_low_weight_fraction( - intensity_out_of_order=True, parent_rwp=None, child_rwp=10.0 - ) - ) + self.assertTrue(should_prune_low_weight_fraction(intensity_out_of_order=True, parent_rwp=None, child_rwp=10.0)) def test_out_of_order_zero_or_negative_parent_rwp_is_pruned(self): """Degenerate parent_rwp <= 0 can't support a relative-improvement calculation (division by zero); fall back to the ordering-only rule.""" - self.assertTrue( - should_prune_low_weight_fraction( - intensity_out_of_order=True, parent_rwp=0.0, child_rwp=5.0 - ) - ) - self.assertTrue( - should_prune_low_weight_fraction( - intensity_out_of_order=True, parent_rwp=-1.0, child_rwp=5.0 - ) - ) + self.assertTrue(should_prune_low_weight_fraction(intensity_out_of_order=True, parent_rwp=0.0, child_rwp=5.0)) + self.assertTrue(should_prune_low_weight_fraction(intensity_out_of_order=True, parent_rwp=-1.0, child_rwp=5.0)) def test_out_of_order_marginal_improvement_still_pruned(self): """Out-of-order + improvement below the material threshold: pruned.""" @@ -131,9 +113,7 @@ def test_out_of_order_marginal_improvement_still_pruned(self): # 5% relative improvement, below the 10% default threshold child_rwp = 38.0 self.assertTrue( - should_prune_low_weight_fraction( - intensity_out_of_order=True, parent_rwp=parent_rwp, child_rwp=child_rwp - ) + should_prune_low_weight_fraction(intensity_out_of_order=True, parent_rwp=parent_rwp, child_rwp=child_rwp) ) def test_out_of_order_material_improvement_kept(self): @@ -144,9 +124,7 @@ def test_out_of_order_material_improvement_kept(self): # 50% relative improvement child_rwp = 20.0 self.assertFalse( - should_prune_low_weight_fraction( - intensity_out_of_order=True, parent_rwp=parent_rwp, child_rwp=child_rwp - ) + should_prune_low_weight_fraction(intensity_out_of_order=True, parent_rwp=parent_rwp, child_rwp=child_rwp) ) def test_improvement_exactly_at_threshold_is_kept(self): @@ -155,28 +133,20 @@ def test_improvement_exactly_at_threshold_is_kept(self): parent_rwp = 40.0 child_rwp = parent_rwp * (1 - LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT) self.assertFalse( - should_prune_low_weight_fraction( - intensity_out_of_order=True, parent_rwp=parent_rwp, child_rwp=child_rwp - ) + should_prune_low_weight_fraction(intensity_out_of_order=True, parent_rwp=parent_rwp, child_rwp=child_rwp) ) def test_improvement_just_under_threshold_is_pruned(self): parent_rwp = 40.0 child_rwp = parent_rwp * (1 - LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT) + 1e-6 self.assertTrue( - should_prune_low_weight_fraction( - intensity_out_of_order=True, parent_rwp=parent_rwp, child_rwp=child_rwp - ) + should_prune_low_weight_fraction(intensity_out_of_order=True, parent_rwp=parent_rwp, child_rwp=child_rwp) ) def test_worse_fit_is_pruned(self): """Out-of-order + the child fits WORSE than the parent (negative 'improvement'): definitely pruned.""" - self.assertTrue( - should_prune_low_weight_fraction( - intensity_out_of_order=True, parent_rwp=20.0, child_rwp=30.0 - ) - ) + self.assertTrue(should_prune_low_weight_fraction(intensity_out_of_order=True, parent_rwp=20.0, child_rwp=30.0)) def test_custom_threshold_is_respected(self): parent_rwp, child_rwp = 40.0, 30.0 # 25% relative improvement @@ -211,9 +181,7 @@ def _make_test_tree() -> BaseSearchTree: tree.record_peak_matcher_scores = False tree.express_mode = False tree.maximum_grouping_distance = 0.1 - tree.peak_obs = pd.DataFrame({"2theta": [10.0, 11.0], "intensity": [100.0, 50.0]})[ - ["2theta", "intensity"] - ].values + tree.peak_obs = pd.DataFrame({"2theta": [10.0, 11.0], "intensity": [100.0, 50.0]})[["2theta", "intensity"]].values return tree @@ -246,9 +214,7 @@ def _run_expand(self, parent_rwp, child_rwp, phase_a_intensity, phase_b_intensit ) tree.add_node(root) - child_result = _make_refinement_result( - child_rwp, {"A": phase_a_intensity, "B": phase_b_intensity} - ) + child_result = _make_refinement_result(child_rwp, {"A": phase_a_intensity, "B": phase_b_intensity}) with ( patch.object(tree, "score_phases", return_value=([phase_b], {}, 0)), @@ -263,27 +229,21 @@ def _run_expand(self, parent_rwp, child_rwp, phase_a_intensity, phase_b_intensit def test_material_improvement_is_retained(self, mock_fom): """Out-of-order (B's intensity > A's) but a clearly better fit (50% relative Rwp improvement): must NOT be pruned as low-weight-fraction.""" - status = self._run_expand( - parent_rwp=40.0, child_rwp=20.0, phase_a_intensity=60.0, phase_b_intensity=200.0 - ) + status = self._run_expand(parent_rwp=40.0, child_rwp=20.0, phase_a_intensity=60.0, phase_b_intensity=200.0) self.assertNotEqual(status, "low_weight_fraction") self.assertEqual(status, "pending") def test_marginal_improvement_is_still_pruned(self, mock_fom): """Out-of-order (B's intensity > A's) with only a marginal (5%) Rwp improvement: still pruned as low-weight-fraction, as before.""" - status = self._run_expand( - parent_rwp=40.0, child_rwp=38.0, phase_a_intensity=60.0, phase_b_intensity=200.0 - ) + status = self._run_expand(parent_rwp=40.0, child_rwp=38.0, phase_a_intensity=60.0, phase_b_intensity=200.0) self.assertEqual(status, "low_weight_fraction") def test_in_order_addition_never_flagged_regardless_of_fit(self, mock_fom): """B's intensity is LOWER than A's (properly ordered): never flagged, even with a poor Rwp improvement -- the ordering check simply doesn't fire, independent of the material-improvement carve-out.""" - status = self._run_expand( - parent_rwp=40.0, child_rwp=39.9, phase_a_intensity=200.0, phase_b_intensity=60.0 - ) + status = self._run_expand(parent_rwp=40.0, child_rwp=39.9, phase_a_intensity=200.0, phase_b_intensity=60.0) self.assertNotEqual(status, "low_weight_fraction") From 8e8b33fac7049ecd812be0b878c2b24071c1378e Mon Sep 17 00:00:00 2001 From: lauren-walters Date: Thu, 30 Jul 2026 16:18:50 -0700 Subject: [PATCH 4/5] Change 3: lower material-improvement threshold to 8% LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT: 0.10 -> 0.08. Updated the constant's doc comment to match. tests/test_search_tree.py: the exactly-at-threshold and just-under-threshold tests already derived their boundary values from LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT itself, so they needed no value changes to track the new threshold -- only the exactly-at-threshold case needed a tiny epsilon nudge, since 40.0 * (1 - 0.08) does not round-trip to exactly 0.08 in floating point and was landing on the wrong side of the `<` vs `<=` boundary (0.07999999999999989 < 0.08), an artifact of the boundary construction, not a bug in should_prune_low_weight_fraction. The marginal (5%, still below 8%) and material (50%) cases needed no changes; only a stale "10%" comment was updated to "8%". Co-Authored-By: Claude Sonnet 5 --- src/dara/search/tree.py | 6 +++--- tests/test_search_tree.py | 11 ++++++++--- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/src/dara/search/tree.py b/src/dara/search/tree.py index 318c2641..73f994e2 100644 --- a/src/dara/search/tree.py +++ b/src/dara/search/tree.py @@ -341,9 +341,9 @@ def remove_unnecessary_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.10 means the child -# must fit at least 10% better than its parent, relatively, to survive. -LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT = 0.10 +# 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( diff --git a/tests/test_search_tree.py b/tests/test_search_tree.py index e16fc807..6f9c9cd9 100644 --- a/tests/test_search_tree.py +++ b/tests/test_search_tree.py @@ -110,7 +110,7 @@ def test_out_of_order_zero_or_negative_parent_rwp_is_pruned(self): def test_out_of_order_marginal_improvement_still_pruned(self): """Out-of-order + improvement below the material threshold: pruned.""" parent_rwp = 40.0 - # 5% relative improvement, below the 10% default threshold + # 5% relative improvement, below the 8% default threshold child_rwp = 38.0 self.assertTrue( should_prune_low_weight_fraction(intensity_out_of_order=True, parent_rwp=parent_rwp, child_rwp=child_rwp) @@ -129,9 +129,14 @@ def test_out_of_order_material_improvement_kept(self): def test_improvement_exactly_at_threshold_is_kept(self): """Exactly LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT relative improvement - counts as material (the check is `< threshold`, not `<= threshold`).""" + counts as material (the check is `< threshold`, not `<= threshold`). + + Nudged by a tiny epsilon so the comparison lands unambiguously on + the "kept" side of the boundary regardless of floating-point + rounding in the parent_rwp * (1 - threshold) computation. + """ parent_rwp = 40.0 - child_rwp = parent_rwp * (1 - LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT) + child_rwp = parent_rwp * (1 - LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT) - 1e-9 self.assertFalse( should_prune_low_weight_fraction(intensity_out_of_order=True, parent_rwp=parent_rwp, child_rwp=child_rwp) ) From aae94a368424a19b4562ba2f8d6fdb19fba140f6 Mon Sep 17 00:00:00 2001 From: lauren-walters Date: Thu, 30 Jul 2026 17:47:30 -0700 Subject: [PATCH 5/5] Merge upstream/main; fix test fixture for required refinement_metrics Merges upstream/main (15 commits ahead of where this branch forked, including PR #30's non-reverted RefinementMetrics addition to RefinementResult) into this branch -- no conflicts. RefinementResult.refinement_metrics is now a required field (RefinementMetrics, whose only required field is rwp). The _make_refinement_result test helper in tests/test_search_tree.py built a RefinementResult without it, so pydantic validation failed: "refinement_metrics Field required [type=missing]" -- this is what broke CI (GitHub tests PR branches merged into the current base branch tip, not the branch in isolation, which is why this only showed up there and not in any isolated-branch reproduction). Fix: pass a real RefinementMetrics(rwp=rwp) instance, matching how production code constructs one (see get_result() / refine.py). No model changes, no extra="ignore", nothing stubbed out. Co-Authored-By: Claude Sonnet 5 --- tests/test_search_tree.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/test_search_tree.py b/tests/test_search_tree.py index 6f9c9cd9..3508c0d7 100644 --- a/tests/test_search_tree.py +++ b/tests/test_search_tree.py @@ -6,7 +6,7 @@ from treelib import Node, Tree from dara.refine import RefinementPhase -from dara.result import DiaResult, LstResult, PhaseResult, RefinementResult +from dara.result import DiaResult, LstResult, PhaseResult, RefinementMetrics, RefinementResult from dara.search.data_model import SearchNodeData from dara.search.tree import ( LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT, @@ -48,6 +48,13 @@ def _make_refinement_result( and `plot_data.structs` (read by `remove_unnecessary_phases`) are built from it so removing any one phase changes the calculated pattern enough for `remove_unnecessary_phases` to consider it necessary. + + `refinement_metrics` is a real `RefinementMetrics` instance, matching how + production code builds one (`RefinementMetrics(rwp=lst_data.rwp)` in + `get_result`/`refine.py`) -- its only required field is `rwp`, so this + passes the same `rwp` used for `lst_data`. `missing_peaks`/`extra_peaks`/ + `intensity_mismatch_peaks` are irrelevant to the pruning-decision logic + under test and are left at their `None` defaults. """ peak_rows = [ {"phase": phase, "2theta": 10.0 + i, "intensity": intensity} @@ -78,6 +85,7 @@ def _make_refinement_result( structs=structs, ), peak_data=pd.DataFrame(peak_rows), + refinement_metrics=RefinementMetrics(rwp=rwp), )