Updated Bayesian Optimization Workflow

src/rxn_ca/cli.py

@@ -293,15 +293,43 @@ def suggest_precursors():
         action='store_true',
         help='Output as JSON instead of human-readable format',
     )
+    parser.add_argument(
+        '--anions',
+        type=str,
+        default=None,
+        help='Comma-separated anion formulas for precursors (default: O,CO3,OH,NO3)',
+    )
+    parser.add_argument(
+        '--metathesis',
+        type=str,
+        default=None,
+        help='Comma-separated metathesis anion formulas (e.g., Cl or Cl,Br)',
+    )
+    parser.add_argument(
+        '--counter-cations',
+        type=str,
+        default=None,
+        help='Comma-separated counter-cations for metathesis (e.g., Na,K)',
+    )
 
     args = parser.parse_args()
 
     target = args.target
     print(f"Finding precursor combinations for: {target}", file=sys.stderr)
 
+    # Parse anion/cation arguments
+    anions = args.anions.split(',') if args.anions else None
+    metathesis_anions = args.metathesis.split(',') if args.metathesis else None
+    counter_cations = args.counter_cations.split(',') if args.counter_cations else None
+
     # Expand elements to include precursor anions (C, N, H, etc.)
     print("Expanding element set for precursor phases...", file=sys.stderr)
-    elements = get_expanded_elements(target)
+    elements = get_expanded_elements(
+        target,
+        anions=anions,
+        metathesis_anions=metathesis_anions,
+        counter_cations=counter_cations,
+    )
     print(f"  Elements: {', '.join(sorted(elements))}", file=sys.stderr)
 
     # Fetch entries from Materials Project

src/rxn_ca/optimization/objective.py

@@ -49,7 +49,7 @@ class ObjectiveConfig:
     num_realizations: int = 3
     cache_results: bool = True
     live_compress: bool = True
-    compress_freq: int = 50
+    compress_freq: int = 500
 
 
 class ObjectiveFunction:

src/rxn_ca/optimization/precursor_selection.py

@@ -65,11 +65,11 @@ def __post_init__(self):
     AnionType("phosphate", "PO4", -3, frozenset({"P", "O"})),
 ]
 
-# Default anion types for typical solid-state synthesis
-DEFAULT_PRECURSOR_ANIONS: List[str] = ["oxide", "carbonate", "hydroxide", "nitrate"]
+# Default anion types for typical solid-state synthesis (by formula)
+DEFAULT_PRECURSOR_ANIONS: List[str] = ["O", "CO3", "OH", "NO3"]
 
-# Anion types useful for metathesis reactions
-METATHESIS_ANIONS: List[str] = ["chloride", "bromide", "nitrate", "sulfate", "acetate"]
+# Anion types useful for metathesis reactions (by formula)
+METATHESIS_ANIONS: List[str] = ["Cl", "Br", "NO3", "SO4", "C2H3O2"]
 
 # Counter-cations for metathesis reactions (provide leaving groups)
 # Maps cation symbol to its oxidation state
@@ -81,6 +81,9 @@ def __post_init__(self):
     "Cs": 1,
 }
 
+# Build formula -> AnionType lookup (populated after COMMON_ANION_TYPES is defined)
+_ANION_BY_FORMULA: Dict[str, "AnionType"] = {}
+
 
 # =============================================================================
 # Precursor formula generation
@@ -164,22 +167,49 @@ def generate_precursor_formula(
     return cation_str + anion_str
 
 
-def get_anion_by_name(name: str) -> AnionType:
-    """Get an AnionType by its name.
+def _build_anion_lookup() -> None:
+    """Build the formula -> AnionType lookup dict."""
+    for anion in COMMON_ANION_TYPES:
+        _ANION_BY_FORMULA[anion.formula] = anion
+
+
+def get_anion(identifier: str) -> AnionType:
+    """Get an AnionType by its formula or name.
 
     Args:
-        name: Anion name (e.g., "oxide", "carbonate")
+        identifier: Anion formula (e.g., "CO3", "Cl") or name (e.g., "carbonate")
 
     Returns:
         Matching AnionType
 
     Raises:
-        ValueError: If anion name not found
+        ValueError: If anion not found
+
+    Examples:
+        >>> get_anion("CO3")
+        AnionType(name='carbonate', formula='CO3', ...)
+        >>> get_anion("Cl")
+        AnionType(name='chloride', formula='Cl', ...)
+        >>> get_anion("carbonate")  # name also works for backwards compat
+        AnionType(name='carbonate', formula='CO3', ...)
     """
+    # Ensure lookup is populated
+    if not _ANION_BY_FORMULA:
+        _build_anion_lookup()
+
+    # Try formula first
+    if identifier in _ANION_BY_FORMULA:
+        return _ANION_BY_FORMULA[identifier]
+
+    # Fall back to name lookup for backwards compatibility
     for anion in COMMON_ANION_TYPES:
-        if anion.name == name:
+        if anion.name == identifier:
             return anion
-    raise ValueError(f"Unknown anion type: {name}")
+
+    raise ValueError(
+        f"Unknown anion: '{identifier}'. "
+        f"Valid formulas: {list(_ANION_BY_FORMULA.keys())}"
+    )
 
 
 def generate_practical_precursors(
@@ -214,7 +244,7 @@ def generate_practical_precursors(
     if not oxidation_states:
         return []
 
-    anions = [get_anion_by_name(name) for name in anion_types]
+    anions = [get_anion(a) for a in anion_types]
 
     precursors = []
     seen = set()
@@ -253,7 +283,7 @@ def generate_metathesis_sources(
     if counter_cations is None:
         counter_cations = ["Na", "K"]
 
-    anion = get_anion_by_name(target_anion)
+    anion = get_anion(target_anion)
 
     sources = []
     for cation in counter_cations:
@@ -266,73 +296,90 @@ def generate_metathesis_sources(
     return sources
 
 
-def get_elements_from_anion_types(anion_types: Optional[List[str]] = None) -> Set[str]:
-    """Get all elements introduced by a set of anion types.
+def get_elements_from_anions(anions: Optional[List[str]] = None) -> Set[str]:
+    """Get all elements introduced by a set of anions.
 
     Args:
-        anion_types: List of anion type names. Defaults to DEFAULT_PRECURSOR_ANIONS.
+        anions: List of anion formulas (e.g., ["CO3", "Cl"]).
+            Defaults to DEFAULT_PRECURSOR_ANIONS.
 
     Returns:
         Set of element symbols
     """
-    if anion_types is None:
-        anion_types = DEFAULT_PRECURSOR_ANIONS
+    if anions is None:
+        anions = DEFAULT_PRECURSOR_ANIONS
 
     elements: Set[str] = set()
-    for name in anion_types:
-        anion = get_anion_by_name(name)
+    for identifier in anions:
+        anion = get_anion(identifier)
         elements.update(anion.elements)
 
     return elements
 
 
 def get_expanded_elements(
     target_phase: str,
-    anion_types: Optional[List[str]] = None,
-    include_metathesis: bool = True,
+    anions: Optional[List[str]] = None,
+    metathesis_anions: Optional[List[str]] = None,
+    counter_cations: Optional[List[str]] = None,
 ) -> Set[str]:
     """Get the full set of elements needed for precursor selection.
 
     This expands beyond the target phase elements to include elements from
-    common precursor anions (e.g., C from carbonates, N from nitrates).
+    precursor anions (e.g., C from CO3, N from NO3) and optionally metathesis
+    reagents.
 
     Use this to determine what elements to pass to get_entries().
 
     Args:
         target_phase: Target product formula (e.g., "BaTiO3")
-        anion_types: Anion types to consider. Defaults to DEFAULT_PRECURSOR_ANIONS.
-        include_metathesis: If True, also include elements from metathesis anions.
+        anions: Base anion formulas for precursors. Defaults to ["O", "CO3", "OH", "NO3"].
+        metathesis_anions: Additional anions for metathesis (e.g., ["Cl"]).
+            Defaults to None (no metathesis anions).
+        counter_cations: Counter-cations for metathesis (e.g., ["Na", "K"]).
+            Defaults to None (no counter-cations).
 
     Returns:
         Set of element symbols needed for get_entries()
 
     Examples:
         >>> get_expanded_elements("BaTiO3")
         {'Ba', 'Ti', 'O', 'C', 'N', 'H'}
-        >>> get_expanded_elements("BaTiO3", anion_types=["oxide"])
+        >>> get_expanded_elements("BaTiO3", anions=["O"])
         {'Ba', 'Ti', 'O'}
+        >>> get_expanded_elements("BaTiO3", metathesis_anions=["Cl"], counter_cations=["Na"])
+        {'Ba', 'Ti', 'O', 'C', 'N', 'H', 'Cl', 'Na'}
     """
     # Start with target phase elements
     target_comp = Composition(target_phase)
     elements = {str(el) for el in target_comp.elements}
 
-    # Add elements from precursor anions
-    if anion_types is None:
-        anion_types = list(DEFAULT_PRECURSOR_ANIONS)
-
-    if include_metathesis:
-        # Add metathesis anion elements too
-        anion_types = list(set(anion_types + METATHESIS_ANIONS))
-
-    elements.update(get_elements_from_anion_types(anion_types))
-
-    # Add counter-cation elements if including metathesis
-    if include_metathesis:
-        elements.update(METATHESIS_COUNTER_CATIONS.keys())
-        # Remove NH4 as it's not a real element, add N and H instead
-        elements.discard("NH4")
-        elements.add("N")
-        elements.add("H")
+    # Add elements from base precursor anions
+    if anions is None:
+        anions = list(DEFAULT_PRECURSOR_ANIONS)
+
+    all_anions = list(anions)
+
+    # Add metathesis anions if specified
+    if metathesis_anions:
+        all_anions = list(set(all_anions + metathesis_anions))
+
+    elements.update(get_elements_from_anions(all_anions))
+
+    # Add counter-cation elements if specified
+    if counter_cations:
+        for cation in counter_cations:
+            if cation == "NH4":
+                # NH4 is not a real element, add N and H instead
+                elements.add("N")
+                elements.add("H")
+            elif cation not in METATHESIS_COUNTER_CATIONS:
+                raise ValueError(
+                    f"Unknown counter-cation: '{cation}'. "
+                    f"Valid options: {list(METATHESIS_COUNTER_CATIONS.keys())}"
+                )
+            else:
+                elements.add(cation)
 
     return elements
 

src/rxn_ca/optimization/search_space.py

@@ -124,19 +124,26 @@ def add_precursor_ratio(
         name: str,
         low: float,
         high: float,
+        step: float = 0.05,
     ) -> "SearchSpace":
-        """Add a precursor ratio parameter.
+        """Add a precursor ratio parameter (discrete with given step size).
+
+        Discrete rather than continuous: BayBE enumerates all candidate values
+        and selects via Thompson Sampling, avoiding the L-BFGS-B boundary-hang
+        that occurs when a continuous parameter's optimum sits at its lower bound
+        (e.g. the stoichiometric ratio is also the minimum of the search range).
 
         Args:
             name: Parameter name (should match a precursor slot name + "_ratio")
-            low: Minimum ratio (typically 0-1)
-            high: Maximum ratio (typically 0-1)
+            low: Minimum ratio value
+            high: Maximum ratio value
+            step: Grid spacing between ratio values (default 0.05)
 
         Returns:
             self for method chaining
         """
         ratio_name = f"{name}_ratio" if not name.endswith("_ratio") else name
-        param = ContinuousParameter(name=ratio_name, low=low, high=high)
+        param = DiscreteParameter(name=ratio_name, low=low, high=high, step=step)
         return self._add_parameter(param)
 
     def add_continuous(
@@ -295,6 +302,41 @@ def parameter_names(self) -> List[str]:
         """Get list of all parameter names."""
         return [p.name for p in self.parameters]
 
+    def as_dict(self) -> dict:
+        """Serialize to a JSON-safe dict for passing between jobflow jobs."""
+        params = []
+        for p in self.parameters:
+            d: Dict[str, Any] = {"type": p.param_type.value, "name": p.name}
+            if isinstance(p, DiscreteParameter):
+                d["low"] = p.low
+                d["high"] = p.high
+                d["step"] = p.step
+            elif isinstance(p, ContinuousParameter):
+                d["low"] = p.low
+                d["high"] = p.high
+            elif isinstance(p, PrecursorSlotParameter):
+                d["candidates"] = p.candidates
+            elif isinstance(p, CategoricalParameter):
+                d["choices"] = p.choices
+            params.append(d)
+        return {"parameters": params}
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "SearchSpace":
+        """Reconstruct a SearchSpace from a serialized dict."""
+        space = cls()
+        for p in d["parameters"]:
+            ptype = p["type"]
+            if ptype == "discrete":
+                space.add_discrete(p["name"], p["low"], p["high"], p["step"])
+            elif ptype == "continuous":
+                space.add_continuous(p["name"], p["low"], p["high"])
+            elif ptype == "precursor_slot":
+                space.add_precursor_slot(p["name"], p["candidates"])
+            elif ptype == "categorical":
+                space.add_categorical(p["name"], p["choices"])
+        return space
+
     def __repr__(self) -> str:
         param_strs = []
         for p in self.parameters:

src/rxn_ca/utilities/get_scored_rxns.py

@@ -1,3 +1,7 @@
+import os
+import multiprocessing as mp
+from typing import List
+
 from rxn_network.reactions.reaction_set import ReactionSet
 
 from ..core import HeatingSchedule
@@ -6,11 +10,14 @@
 from ..reactions import ReactionLibrary, ScoredReaction, ScoredReactionSet, score_rxns
 from ..reactions.scorers import BasicScore, TammanScore
 
-from typing import List
+_scoring_globals = {}
 
-import multiprocessing as mp
 
-_scoring_globals = {}
+def _pool_initializer(data: dict):
+    """Populate worker-process globals (called by forkserver/spawn pool workers)."""
+    global _scoring_globals
+    _scoring_globals.update(data)
+
 
 def fn(temp):
     score_class = _scoring_globals.get('score_class')

src/rxn_ca/workflow/__init__.py

@@ -8,13 +8,16 @@
 """
 
 from .schemas import SimulationOutput, ReactionLibraryData
-from .jobs import setup_reaction_library, run_simulation
-from .flows import create_simulation_flow
+from .jobs import setup_reaction_library, run_simulation, init_bo_campaign, bo_trial_step
+from .flows import create_simulation_flow, BOFlowMaker
 
 __all__ = [
     "SimulationOutput",
     "ReactionLibraryData",
     "setup_reaction_library",
     "run_simulation",
+    "init_bo_campaign",
+    "bo_trial_step",
     "create_simulation_flow",
+    "BOFlowMaker",
 ]