Fix theta-search solver and polarization handling

CHANGELOG.md

@@ -1,6 +1,11 @@
 # Changelog
 
-## Unreleased
+## 0.4.8 - 2026-09-03
+
+- grax log records no longer leak to the terminal when the host application has not called `grax.setup_logging`. A `NullHandler` is attached to the `grax` logger at import, so Python's `logging.lastResort` handler no longer prints `WARNING`s to stderr -- most visibly from the spawned batch workers, which re-import grax but never configure logging. `setup_logging` now also sets `propagate = False` on the `grax` logger and refuses to attach a second file handler for the same path on a repeat call.
+- The multilayer theta-search `Requested theta half-width ... reaches near/into 0 deg. Reducing to ...` notice is now `INFO`, not `WARNING`, and is emitted once per distinct `(center, requested half-width)` instead of on every rough/fine re-centring attempt. At small grazing angles the clamp is expected behaviour, not a fault; it belongs in the log file, not the console.
+- `multilayer_theta_search_cases` gained `solver` and `solver_options`, matching `run_multilayer_theta_search` and the other case generators. The batch runner already read a per-case `"solver"` key; the generator just could not set one, so cases built directly could only take the solver from the runner default.
+- The `multilayer_theta_search` example takes `--polarization {s,p,TE,TM}` alongside `--solver`, and writes each run under `results/<solver>_<polarization>/` (logs and checkpoints included) so an s run and a p run, or an rcwa and a neviere run, sit side by side.
 
 - Added `grax_opt.optimize_to_joint_measurements` for fitting one parameter set jointly against several measured curves, each with its own measurement file and energy grid, with a configurable `joint_loss_reduction` (`mean`, `sum`, `pooled`, or `weighted`). Each trial evaluates every measurement in a single `BatchSimulationRunner` batch so trial-level `max_workers` parallelizes across measurements as well as energies.
 - Joint fits are not limited to curves that differ by grazing angle. `angle_mode`, `grazing_angle_deg`, `cff`, `diffraction_order` and `polarization` are run-level defaults that any individual `MeasurementSpec` may override, so one fit can span angles, angle modes, diffraction orders and polarizations. Numerical settings (`fourier_orders`, `solver`, `backend`, `max_workers`) stay run-level, since they describe how a curve is computed rather than what was measured.

docs/tutorials/multilayer-theta-search.md

@@ -116,3 +116,8 @@ for all three stages — rough scan, precise scan and final solve — so the sel
 angle and the reported efficiency come from the same polarization. Because the
 scan maximizes the selected order, `s` and `p` generally settle on slightly
 different angles.
+
+The runnable example exposes both choices as flags,
+`--solver {rcwa,neviere}` and `--polarization {s,p,TE,TM}`, and writes each
+combination under its own `results/<solver>_<polarization>/` directory so runs
+do not overwrite one another.

examples/simulation/multilayer_theta_search/multilayer_theta_search.py

@@ -13,13 +13,36 @@
 import grax
 
 
+parser = argparse.ArgumentParser(description="Multilayer theta-search sweep")
+parser.add_argument(
+    "--solver",
+    choices=("rcwa", "neviere"),
+    default="rcwa",
+    help="Electromagnetic solver to run. Both compute every diffraction order; "
+    "they differ only in how each layer is crossed in z.",
+)
+parser.add_argument(
+    "--polarization",
+    choices=("s", "p", "TE", "TM"),
+    default="s",
+    help="Incident polarization used for every stage of the search "
+    "(s/TE or p/TM, canonicalized to s/p).",
+)
+args = parser.parse_args()
+
+# Keep each solver/polarization run in its own directory so an s and a p run,
+# or an rcwa and a neviere run, sit side by side instead of overwriting.
+run_tag = f"{args.solver}_{args.polarization.lower()}"
+
 example_root = Path(__file__).resolve().parent
-output_dir = example_root / "results"
+output_dir = example_root / "results" / run_tag
 optical_constants_dir = example_root / "optical_constants"
 log_dir = output_dir / "logs"
 log_dir.mkdir(parents=True, exist_ok=True)
 
-grax.setup_logging(level="INFO", log_dir=str(log_dir), run_id="multilayer_theta_search")
+grax.setup_logging(
+    level="INFO", log_dir=str(log_dir), run_id=f"multilayer_theta_search_{run_tag}"
+)
 logger = logging.getLogger("grax.examples.multilayer_theta_search")
 
 energies_ev = np.arange(500.0, 6000.1, 10, dtype=float)
@@ -65,23 +88,14 @@
     z_resolution_nm=1.0,
 )
 
-parser = argparse.ArgumentParser(description="Multilayer theta-search sweep")
-parser.add_argument(
-    "--solver",
-    choices=("rcwa", "neviere"),
-    default="rcwa",
-    help="Electromagnetic solver to run. Both compute every diffraction order; "
-    "they differ only in how each layer is crossed in z.",
-)
-args = parser.parse_args()
-
 # Guard the executable part: the batch runner spawns worker processes on macOS,
 # and a spawned worker re-imports this file by path. Without the guard each
 # worker re-runs the whole example and recursively spawns more, which fails with
 # BrokenProcessPool before any case completes.
 if __name__ == "__main__":
     sweep = grax.run_multilayer_theta_search_sweep(
         solver=args.solver,
+        polarization=args.polarization,
         grating=grating,
         energies_ev=energies_ev,
         output_dir=output_dir,
@@ -108,7 +122,7 @@
         show_progress=True,
         live_plot=True,
         on_error="fail_fast",
-        checkpoint_dir=output_dir / f"checkpoints_{args.solver}",
+        checkpoint_dir=output_dir / "checkpoints",
         resume=True,
         theta_tracking_mode="auto",
         max_tracking_energy_step_ev=None,

src/grax/__init__.py

@@ -4,6 +4,13 @@
 
 import logging
 
+# Attach a no-op handler to the package logger so that when an application has
+# not configured logging, grax log records are not emitted to stderr by
+# ``logging.lastResort``. This matters for the spawned batch workers, which
+# re-import grax but never call ``setup_logging``: without this, every
+# ``logger.warning`` inside a worker leaks to the terminal.
+logging.getLogger("grax").addHandler(logging.NullHandler())
+
 from .afm_grating import AFMGrating
 from .afm_preprocessing import AFMPreprocessing
 from .materials import (
@@ -155,6 +162,23 @@ def setup_logging(
 
     root_logger = logging.getLogger("grax")
     root_logger.setLevel(level)
+
+    # Do not let grax records bubble to the root logger: they are already
+    # written to this file handler, and a StreamHandler attached to the root
+    # by the host application would otherwise re-print every one of them.
+    root_logger.propagate = False
+
+    # Guard against duplicate file handlers if setup_logging is called twice
+    # (e.g. an example re-run in the same interpreter): a second identical
+    # FileHandler would write every record to the log twice.
+    new_target = getattr(handler, "baseFilename", None)
+    for existing in list(root_logger.handlers):
+        if (
+            isinstance(existing, logging.FileHandler)
+            and getattr(existing, "baseFilename", None) == new_target
+        ):
+            root_logger.removeHandler(existing)
+            existing.close()
     root_logger.addHandler(handler)
 
     logging.getLogger("numpy").setLevel(level)

src/grax/simulation/cases.py

@@ -187,6 +187,8 @@ def multilayer_theta_search_cases(
     roughness_sigma_nm: float | None = None,
     precise_peak_selection_mode: PeakSelectionMode = "max",
     polarization: str | None = None,
+    solver: str | None = None,
+    solver_options: object | None = None,
 ) -> Iterator[dict[str, object]]:
     """Yield energy-only cases for the multilayer theta-search workflow.
 
@@ -219,6 +221,10 @@ def multilayer_theta_search_cases(
         roughness_sigma_nm: Optional rms roughness in nanometers.
         polarization: Incident polarization, ``"s"``/``"TE"`` or ``"p"``/``"TM"``.
             When ``None`` the runner's ``polarization`` is used.
+        solver: Electromagnetic solver (``"rcwa"`` or ``"neviere"``) used by the
+            theta search. When ``None`` the runner's ``solver`` is used.
+        solver_options: Integration settings for ``solver="neviere"``. When
+            ``None`` the runner's ``solver_options`` is used.
         precise_peak_selection_mode: Mode used to select the final theta from the
             precise scan. ``"max"`` uses the sampled maximum, ``"gauss"`` fits a
             local Gaussian, and ``"voigt"`` fits a local Voigt profile.
@@ -257,6 +263,8 @@ def multilayer_theta_search_cases(
             "roughness_sigma_nm": roughness_sigma_nm,
             "precise_peak_selection_mode": precise_peak_selection_mode,
             **({} if polarization is None else {"polarization": polarization}),
+            **({} if solver is None else {"solver": solver}),
+            **({} if solver_options is None else {"solver_options": solver_options}),
             "case_id": f"{_THETA_SEARCH_CASE_ID_PREFIX}-{index:08d}",
             "grating": grating,
             "energy_ev": float(energy_ev),

src/grax/simulation/theta_search.py

@@ -21,6 +21,13 @@
 
 logger = logging.getLogger(__name__)
 
+# Half-width clamps that have already been reported, keyed by
+# ``(round(center_deg, 3), round(requested_half_width_deg, 3))``. The rough and
+# fine scans re-enter the clamp on every re-centring attempt for the same
+# energy, so without this the identical notice is logged up to eight times per
+# energy point.
+_reported_half_width_clamps: set[tuple[float, float]] = set()
+
 
 def _worker_identity() -> str:
     """Return a compact worker identity for theta-search logs."""
@@ -120,13 +127,19 @@ def _safe_theta_scan_half_width_deg(
         raise ValueError("Theta scan center must be > 0 deg.")
     max_safe_half_width = 0.95 * center_deg
     if requested_half_width_deg > max_safe_half_width:
-        logger.warning(
-            "Requested theta half-width %.6f deg around center %.6f deg reaches near/into 0 deg. "
-            "Reducing to %.6f deg.",
-            requested_half_width_deg,
-            center_deg,
-            max_safe_half_width,
-        )
+        clamp_key = (round(center_deg, 3), round(requested_half_width_deg, 3))
+        if clamp_key not in _reported_half_width_clamps:
+            _reported_half_width_clamps.add(clamp_key)
+            # INFO, not WARNING: at small grazing angles this is the expected
+            # adjustment, not a fault. It goes to the configured log file and
+            # never to the terminal (see the NullHandler in grax/__init__.py).
+            logger.info(
+                "Requested theta half-width %.6f deg around center %.6f deg reaches near/into 0 deg. "
+                "Reducing to %.6f deg.",
+                requested_half_width_deg,
+                center_deg,
+                max_safe_half_width,
+            )
     return float(min(requested_half_width_deg, max_safe_half_width))
 
 

src/grax/web/app.py

@@ -1901,6 +1901,7 @@ def _cases_for_workflow(
     run_x_resolution_nm = float(form.get("run_x_resolution_nm") or grating.x_resolution_nm)
     run_z_resolution_nm = float(form.get("run_z_resolution_nm") or grating.z_resolution_nm)
     polarization = _normalized_polarization(form.get("polarization", "s"))
+    solver = _normalized_solver(form.get("solver", "rcwa"))
 
     if workflow == "fixed_angle":
         cases = list(
@@ -1942,6 +1943,7 @@ def _cases_for_workflow(
         case["case_id"] = f"{workflow}-{index:08d}"
         case["fourier_orders"] = fourier_orders
         case["polarization"] = polarization
+        case["solver"] = solver
         if workflow != "multilayer_theta_search":
             case["x_resolution_nm"] = run_x_resolution_nm
             case["z_resolution_nm"] = run_z_resolution_nm

tests/smoke/test_simulation_examples.py

@@ -185,6 +185,13 @@ def test_multilayer_theta_search_docs_use_grouped_canonical_arguments() -> None:
     tutorial_call = tutorial_source.split("run_multilayer_theta_search_sweep(", maxsplit=1)[1].split(")\n", maxsplit=1)[0]
     assert "run_multilayer_theta_search(" not in tutorial_source
 
+    # The example exposes both physics choices as flags and keeps each
+    # solver/polarization run in its own directory.
+    assert '"--solver"' in example_source
+    assert '"--polarization"' in example_source
+    assert 'run_tag = f"{args.solver}_{args.polarization.lower()}"' in example_source
+    assert "polarization=args.polarization" in example_source
+
     for call_block in (example_call, tutorial_call):
         assert call_block.index("multilayer_bragg_order") < call_block.index("rough_scan_half_width_deg")
         assert call_block.index("rough_scan_half_width_deg") < call_block.index("rough_fourier_orders")

tests/unit/test_logging_setup.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator
+from pathlib import Path
+
+import pytest
+
+import grax
+
+
+def test_grax_logger_has_null_handler_on_import() -> None:
+    """Without setup_logging, grax records must not fall through to stderr."""
+
+    handlers = logging.getLogger("grax").handlers
+    assert any(isinstance(handler, logging.NullHandler) for handler in handlers)
+
+
+@pytest.fixture
+def restore_grax_logger() -> Iterator[logging.Logger]:
+    grax_logger = logging.getLogger("grax")
+    original_handlers = list(grax_logger.handlers)
+    original_propagate = grax_logger.propagate
+    original_level = grax_logger.level
+    try:
+        yield grax_logger
+    finally:
+        for handler in list(grax_logger.handlers):
+            if handler not in original_handlers:
+                grax_logger.removeHandler(handler)
+                handler.close()
+        grax_logger.propagate = original_propagate
+        grax_logger.setLevel(original_level)
+
+
+def test_setup_logging_disables_propagation_and_dedupes_file_handler(
+    restore_grax_logger: logging.Logger,
+    tmp_path: Path,
+) -> None:
+    grax_logger = restore_grax_logger
+    log_file = tmp_path / "run.log"
+
+    grax.setup_logging(log_file=str(log_file))
+    grax.setup_logging(log_file=str(log_file))
+
+    file_handlers = [
+        handler
+        for handler in grax_logger.handlers
+        if isinstance(handler, logging.FileHandler)
+    ]
+    assert len(file_handlers) == 1
+    assert grax_logger.propagate is False

tests/unit/test_polarization_aliases.py

@@ -321,3 +321,35 @@ def test_theta_search_batch_workflow_honours_polarization() -> None:
     results = list(runner.run_cases(cases))
 
     assert results[0].polarization == "p"
+
+
+@pytest.mark.unit
+def test_theta_search_cases_carry_requested_solver_into_the_batch() -> None:
+    """A case-level solver from the generator must beat the runner default."""
+
+    common = dict(
+        grating=_multilayer_grating(),
+        energies_ev=[500.0],
+        diffraction_order=2,
+        rough_scan_half_width_deg=0.4,
+        rough_scan_points=5,
+        rough_fourier_orders=2,
+        rough_x_resolution_nm=4.0,
+        rough_z_resolution_nm=2.0,
+        fine_scan_half_width_deg=0.1,
+        fine_scan_points=5,
+        fine_fourier_orders=2,
+        fine_x_resolution_nm=4.0,
+        fine_z_resolution_nm=2.0,
+        final_fourier_orders=3,
+        final_x_resolution_nm=4.0,
+        final_z_resolution_nm=2.0,
+    )
+    cases = list(multilayer_theta_search_cases(solver="neviere", **common))
+    assert cases[0]["solver"] == "neviere"
+
+    # Runner default is rcwa; the case must override it.
+    runner = BatchSimulationRunner(fourier_orders=3, solver="rcwa", on_error="fail_fast")
+    results = list(runner.run_cases(cases))
+
+    assert results[0].solver == "neviere"

tests/unit/test_simulation_core.py

@@ -1430,6 +1430,24 @@ def test_lazy_case_helpers_yield_expected_cases() -> None:
     assert first_theta_search["energy_ev"] == 100.0
     assert first_theta_search["workflow"] == "multilayer_theta_search"
     assert "grazing_angle_deg" not in first_theta_search
+    # solver / solver_options are only stamped when explicitly requested.
+    assert "solver" not in first_theta_search
+    assert "solver_options" not in first_theta_search
+
+
+def test_multilayer_theta_search_cases_stamp_requested_solver() -> None:
+    grating = build_test_grating()
+    options = object()
+    case = next(
+        multilayer_theta_search_cases(
+            grating=grating,
+            energies_ev=iter([100.0]),
+            solver="neviere",
+            solver_options=options,
+        )
+    )
+    assert case["solver"] == "neviere"
+    assert case["solver_options"] is options
 
 
 def test_case_helpers_reject_removed_public_override_arguments() -> None:

tests/unit/test_simulation_theta_search.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+import logging
 from collections.abc import Iterator
 from pathlib import Path
 
@@ -1149,3 +1150,44 @@ def fake_run_multilayer_theta_search(**kwargs: object) -> SingleSimulationResult
     assert calls == [1800.0]
 
 
+
+
+def test_safe_theta_scan_half_width_clamp_logs_once_at_info(
+    caplog: pytest.LogCaptureFixture,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from grax.simulation import theta_search as theta_search_module
+
+    # caplog captures through the root logger, so propagation must be on for
+    # this assertion regardless of whether an earlier test called setup_logging.
+    monkeypatch.setattr(logging.getLogger("grax"), "propagate", True)
+    theta_search_module._reported_half_width_clamps.clear()
+
+    with caplog.at_level(logging.INFO, logger="grax.simulation.theta_search"):
+        first = theta_search_module._safe_theta_scan_half_width_deg(
+            center_deg=0.9, requested_half_width_deg=1.0
+        )
+        second = theta_search_module._safe_theta_scan_half_width_deg(
+            center_deg=0.9, requested_half_width_deg=1.0
+        )
+
+    assert first == pytest.approx(0.95 * 0.9)
+    assert second == pytest.approx(0.95 * 0.9)
+    clamp_records = [
+        record
+        for record in caplog.records
+        if record.name == "grax.simulation.theta_search"
+        and "reaches near/into 0 deg" in record.getMessage()
+    ]
+    assert len(clamp_records) == 1
+    assert clamp_records[0].levelno == logging.INFO
+
+
+def test_safe_theta_scan_half_width_returns_request_when_within_bounds() -> None:
+    from grax.simulation import theta_search as theta_search_module
+
+    theta_search_module._reported_half_width_clamps.clear()
+    value = theta_search_module._safe_theta_scan_half_width_deg(
+        center_deg=5.0, requested_half_width_deg=0.5
+    )
+    assert value == pytest.approx(0.5)