diff --git a/docs/source/conf.py b/docs/source/conf.py index 3452cda4..0eee9d0f 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -93,6 +93,12 @@ if "doctest" in sys.argv or os.environ.get("BLOP_DOCS_NO_EXEC"): nb_execution_mode = "off" +# The TES tutorial needs the trained emulator weights, which are derived from +# measured beamline data and not yet publicly downloadable. Skip executing it +# (the page still renders) until a weights URL is configured in +# blop_sim.backends.models.tes_model.WEIGHTS_URL — then remove this exclusion. +nb_execution_excludepatterns = ["tutorials/tes-kb-jacks*"] + # Enable ELLIPSIS option for doctest to match any substring including empty doctest_default_flags = doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE diff --git a/docs/source/tutorials.rst b/docs/source/tutorials.rst index cdc14995..a6d0702d 100644 --- a/docs/source/tutorials.rst +++ b/docs/source/tutorials.rst @@ -7,3 +7,4 @@ Tutorials tutorials/simple-experiment.md tutorials/queueserver.md tutorials/xrt-kb-mirrors.md + tutorials/tes-kb-jacks.md diff --git a/docs/source/tutorials/tes-kb-jacks.md b/docs/source/tutorials/tes-kb-jacks.md new file mode 100644 index 00000000..7de9010b --- /dev/null +++ b/docs/source/tutorials/tes-kb-jacks.md @@ -0,0 +1,299 @@ +--- +jupytext: + text_representation: + extension: .md + format_name: myst + format_version: 0.13 + jupytext_version: 1.17.3 +kernelspec: + display_name: dev + language: python + name: python3 +--- + +# Aligning a Measured KB Mirror System + +In this tutorial, you will align a Kirkpatrick-Baez (KB) mirror system whose response comes from **measured beamline data** rather than an analytic or ray-traced model. The workflow is the same as in the [XRT KB mirror tutorial](./xrt-kb-mirrors.md) — DOFs, objectives, an evaluation function, and an `Agent` — with two differences: + +- The beam is produced by `TESBackend`, an emulator trained on 28,561 detector frames recorded on a complete 13⁴ grid of the four KB jack motors of the NSLS-II TES beamline (3 keV). Astigmatism, non-Gaussian streaks, and flux variation are those of the real machine. +- The degrees of freedom are the **four jack positions** (upstream/downstream per mirror) instead of two mirror radii, so the optimizer works in the same parameter space operators use. + +```{note} +The trained weights (16 MB, derived from measured data) are not part of this repository. Download `emulator_weights.npz` from the [tes-emulator](https://github.com/FLlorente/tes-emulator) package and point the `BLOP_SIM_TES_WEIGHTS` environment variable at it before running this tutorial. +``` + +## Setting Up the Environment + +As in the other tutorials, Blop uses [Bluesky](https://blueskyproject.io/) to run experiments and [Tiled](https://blueskyproject.io/tiled/) to store and retrieve data. + +```{code-cell} ipython3 +import logging +import warnings +from pathlib import PurePath + +import bluesky.plan_stubs as bps +import bluesky.plans as bp +import matplotlib.pyplot as plt +import numpy as np +from ax.api.protocols import IMetric +from bluesky.run_engine import RunEngine +from bluesky_tiled_plugins import TiledWriter +from ophyd_async.core import StaticPathProvider, UUIDFilenameProvider +from tiled.client import from_uri # type: ignore[import-untyped] +from tiled.client.container import Container +from tiled.server import SimpleTiledServer + +from blop.ax import Agent, Objective, RangeDOF +from blop.ax.objective import OutcomeConstraint +from blop.protocols import EvaluationFunction + +# Import simulation devices (requires: pip install -e sim/) +from blop_sim.backends.tes import TESBackend +from blop_sim.devices import DetectorDevice +from blop_sim.devices.simple import KBMirror + +# Suppress noisy logs from httpx and dependency deprecation warnings +logging.getLogger("httpx").setLevel(logging.WARNING) +warnings.filterwarnings("ignore", category=FutureWarning) + +# Enable interactive plotting +plt.ion() + +DETECTOR_STORAGE = "/tmp/blop/sim" +``` + +```{code-cell} ipython3 +tiled_server = SimpleTiledServer(readable_storage=[DETECTOR_STORAGE]) +tiled_client = from_uri(tiled_server.uri) +tiled_writer = TiledWriter(tiled_client) + +RE = RunEngine({}) +RE.subscribe(tiled_writer) +``` + +## Defining Degrees of Freedom + +The emulator is valid inside the motor box it was trained on, and `TESBackend.motor_box()` exposes those bounds directly — using them as DOF bounds means the optimizer can never ask for a configuration the model has not seen: + +```{code-cell} ipython3 +# Create the TES emulator backend (reads BLOP_SIM_TES_WEIGHTS) +backend = TESBackend() +box = backend.motor_box() +box +``` + +The backend reads the same jack-style KB mirror devices as `SimpleBackend`: one horizontal and one vertical mirror, each with an upstream and a downstream jack. We start every jack at the center of its range and let the optimizer explore from there: + +```{code-cell} ipython3 +kbh = KBMirror(backend, orientation="horizontal", name="kbh") +kbv = KBMirror(backend, orientation="vertical", name="kbv") + +# Create detector device +det = DetectorDevice(backend, StaticPathProvider(UUIDFilenameProvider(), PurePath(DETECTOR_STORAGE)), name="det") + +# Move the jacks to the center of the recorded motor box +RE( + bps.mv( + kbh.downstream, np.mean(box["kbh_downstream"]), + kbh.upstream, np.mean(box["kbh_upstream"]), + kbv.downstream, np.mean(box["kbv_downstream"]), + kbv.upstream, np.mean(box["kbv_upstream"]), + ) +) + +# Four DOFs: one per jack, bounded by the recorded motor box +dofs = [ + RangeDOF(actuator=kbh.downstream, bounds=box["kbh_downstream"], parameter_type="float"), + RangeDOF(actuator=kbh.upstream, bounds=box["kbh_upstream"], parameter_type="float"), + RangeDOF(actuator=kbv.downstream, bounds=box["kbv_downstream"], parameter_type="float"), + RangeDOF(actuator=kbv.upstream, bounds=box["kbv_upstream"], parameter_type="float"), +] +``` + +## Defining the Objective and Constraints + +Exactly as in the XRT tutorial: minimize the beam FWHM, and use an intensity `OutcomeConstraint` as a safety net against configurations where the beam is lost: + +```{code-cell} ipython3 +objectives = [ + Objective(name="fwhm", minimize=True), +] + +# Track intensity without optimizing it +intensity_metric = IMetric(name="intensity") + +# Soft constraint: reject configurations with too little flux on the detector +outcome_constraints = [ + OutcomeConstraint(constraint="i >= 10000", i=intensity_metric), +] +``` + +## Writing an Evaluation Function + +The evaluation function is identical to the XRT tutorial's — it reads beam images back from Tiled and computes the FWHM of the marginal profiles, so it does not care which backend produced the image: + +```{code-cell} ipython3 +class DetectorEvaluation(EvaluationFunction): + def __init__(self, tiled_client: Container): + self.tiled_client = tiled_client + + def _fwhm_from_profile(self, profile: np.ndarray) -> float: + """Compute FWHM from a 1D marginal profile. + + Finds the half-maximum crossing points with sub-pixel interpolation. + Returns a large value if the beam is too dim or fills the entire detector. + """ + peak = profile.max() + if peak == 0: + return float(len(profile)) # No signal — return detector width as penalty + + half_max = peak / 2.0 + above = profile >= half_max + if not above.any(): + return float(len(profile)) + + indices = np.where(above)[0] + left_idx = indices[0] + right_idx = indices[-1] + + # Sub-pixel interpolation at left crossing + if left_idx > 0: + left = left_idx - 1 + (half_max - profile[left_idx - 1]) / (profile[left_idx] - profile[left_idx - 1]) + else: + left = 0.0 + + # Sub-pixel interpolation at right crossing + if right_idx < len(profile) - 1: + right = right_idx + (half_max - profile[right_idx]) / (profile[right_idx + 1] - profile[right_idx]) + else: + right = float(len(profile) - 1) + + return right - left + + def _compute_stats(self, image: np.ndarray) -> tuple[float, float]: + """Compute FWHM and integrated intensity from a beam image. + + Returns + ------- + fwhm : float + Geometric mean of the horizontal and vertical FWHM (in pixels). + intensity : float + Total integrated intensity (sum of all pixel values). + """ + gray = image.squeeze().astype(np.float64) + if gray.ndim == 3: + gray = gray.mean(axis=-1) + + # Integrated intensity (total flux on detector) + intensity = gray.sum() + + if intensity == 0: + return float(max(gray.shape)), 0.0 # No beam — return max FWHM penalty + + # Marginal profiles: project onto each axis + x_profile = gray.sum(axis=0) # sum along Y rows -> X profile + y_profile = gray.sum(axis=1) # sum along X cols -> Y profile + + fwhm_x = self._fwhm_from_profile(x_profile) + fwhm_y = self._fwhm_from_profile(y_profile) + + # Geometric mean FWHM — targets a small, round spot + fwhm = np.sqrt(fwhm_x * fwhm_y) + + return float(fwhm), float(intensity) + + def __call__(self, uid: str, suggestions: list[dict]) -> list[dict]: + outcomes = [] + run = self.tiled_client[uid] + + # Read beam images from detector + images = run["primary/det_image"].read() + + # Suggestion IDs stored in start document metadata + suggestion_ids = [suggestion["_id"] for suggestion in run.metadata["start"]["blop_suggestions"]] + + # Compute statistics from each image + for idx, sid in enumerate(suggestion_ids): + image = images[idx] + fwhm, intensity = self._compute_stats(image) + + outcome = { + "_id": sid, + "fwhm": fwhm, + "intensity": intensity, + } + outcomes.append(outcome) + return outcomes +``` + +## Creating and Running the Agent + +```{code-cell} ipython3 +agent = Agent( + sensors=[det], + dofs=dofs, + objectives=objectives, + evaluation_function=DetectorEvaluation(tiled_client), + outcome_constraints=outcome_constraints, + name="tes-blop-demo", + description="Blop agent aligning the measured TES KB system", + experiment_type="demo", +) + +# Register intensity as a tracking metric (monitored but not optimized) +agent.ax_client.configure_metrics([intensity_metric]) +``` + +With four DOFs the parameter space is larger than in the two-radius tutorial, so we use a somewhat larger initial exploration batch before switching to model-driven suggestions: + +```{code-cell} ipython3 +# Initial exploration: one batch of 16 quasi-random points +RE(agent.optimize(1, n_points=16)) +``` + +```{code-cell} ipython3 +# Model-driven optimization +RE(agent.optimize(16)) +``` + +## Understanding the Results + +```{code-cell} ipython3 +_ = agent.ax_client.compute_analyses() +``` + +```{code-cell} ipython3 +agent.ax_client.summarize() +``` + +The physics of a KB mirror ties focus to the **pitch** of each mirror — the difference between its two jacks — so a useful way to look at the surrogate is one jack against its partner: + +```{code-cell} ipython3 +_ = agent.plot_objective(x_dof_name="kbh-downstream", y_dof_name="kbh-upstream", objective_name="fwhm") +``` + +## Applying the Optimal Configuration + +```{code-cell} ipython3 +optimal_parameters, metrics, _, _ = agent.ax_client.get_best_parameterization(use_model_predictions=False) +optimal_parameters +``` + +```{code-cell} ipython3 +# Move the jacks to the optimum, record one frame, and read it back from Tiled +moves = [] +for dof in dofs: + moves += [dof.actuator, optimal_parameters[dof.actuator.name.replace("_", "-")]] +RE(bps.mv(*moves)) +(uid,) = RE(bp.count([det])) + +image = tiled_client[uid]["primary/det_image"].read().squeeze() + +plt.figure(figsize=(6, 4)) +plt.imshow(np.log1p(image), origin="lower") +plt.title("Beam at the optimized jack positions (log scale)") +plt.colorbar(label="log(1 + counts)") +plt.show() +``` + +Because the emulator was trained on real frames, the optimum found here corresponds to the actual focal pitches of the TES KB system — the same configuration a human operator would align to. diff --git a/sim/README.md b/sim/README.md index 84bffa45..9c9e4f20 100644 --- a/sim/README.md +++ b/sim/README.md @@ -18,6 +18,13 @@ The package uses a component-based architecture with individual devices: - `SimpleBackend`: Mathematical Gaussian beam simulation - `XRTBackend`: Full ray-tracing simulation using XRT + - `TESBackend`: Replay of a data-trained emulator of the NSLS-II TES beamline + (28,561 measured KB-jack/detector frames at 3 keV). Reuses the + `blop_sim.devices.simple.KBMirror` jack devices. The trained weights (16 MB, + derived from measured data) are not in this repository: point the + `BLOP_SIM_TES_WEIGHTS` environment variable at the `emulator_weights.npz` + shipped with the [tes-emulator](https://github.com/FLlorente/tes-emulator) + package, or pass `TESBackend(weights_path=...)`. - **Devices**: Individual ophyd-async devices diff --git a/sim/blop_sim/__init__.py b/sim/blop_sim/__init__.py index 6d66a330..4c37a032 100644 --- a/sim/blop_sim/__init__.py +++ b/sim/blop_sim/__init__.py @@ -1,10 +1,24 @@ """blop_sim: Simulation devices for BLOP documentation and tutorials.""" -# Backend exports +from typing import TYPE_CHECKING + from .backends.simple import SimpleBackend -from .backends.xrt import XRTBackend + +if TYPE_CHECKING: + from .backends.tes import TESBackend + from .backends.xrt import XRTBackend __all__ = [ "SimpleBackend", + "TESBackend", "XRTBackend", ] + + +def __getattr__(name: str): + # Lazy imports — see blop_sim.backends.__getattr__. + if name in ("XRTBackend", "TESBackend"): + from . import backends + + return getattr(backends, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/sim/blop_sim/backends/__init__.py b/sim/blop_sim/backends/__init__.py index 78f79df8..7b3fe5ca 100644 --- a/sim/blop_sim/backends/__init__.py +++ b/sim/blop_sim/backends/__init__.py @@ -1,7 +1,27 @@ """Backend simulation infrastructure for blop_sim.""" +from typing import TYPE_CHECKING + from .core import SimBackend from .simple import SimpleBackend -from .xrt import XRTBackend -__all__ = ["SimBackend", "SimpleBackend", "XRTBackend"] +if TYPE_CHECKING: + from .tes import TESBackend + from .xrt import XRTBackend + +__all__ = ["SimBackend", "SimpleBackend", "TESBackend", "XRTBackend"] + + +def __getattr__(name: str): + # Lazy imports: xrt pulls in pyopencl (which can fail to import on hosts + # without OpenCL devices) and tes pulls in torch — neither should break + # `import blop_sim` for users of the other backends. + if name == "XRTBackend": + from .xrt import XRTBackend + + return XRTBackend + if name == "TESBackend": + from .tes import TESBackend + + return TESBackend + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/sim/blop_sim/backends/core.py b/sim/blop_sim/backends/core.py index 5849e357..9ddcff05 100644 --- a/sim/blop_sim/backends/core.py +++ b/sim/blop_sim/backends/core.py @@ -14,8 +14,12 @@ class SimBackend(ABC): _instances: dict[type, "SimBackend"] = {} - def __new__(cls): - """Singleton pattern: return existing instance or create new.""" + def __new__(cls, *args, **kwargs): + """Singleton pattern: return existing instance or create new. + + Constructor arguments are accepted (and consumed by ``__init__``) so + backends can take options, e.g. ``SimpleBackend(noise=True)``. + """ if cls not in cls._instances: instance = super().__new__(cls) cls._instances[cls] = instance diff --git a/sim/blop_sim/backends/models/tes_model.py b/sim/blop_sim/backends/models/tes_model.py new file mode 100644 index 00000000..b0ecb004 --- /dev/null +++ b/sim/blop_sim/backends/models/tes_model.py @@ -0,0 +1,218 @@ +"""Data-driven TES beamline model: motors -> beam image + flux. + +A trimmed inference-only port of the ``tes-emulator`` package: a differentiable +motor->beam emulator of the NSLS-II TES beamline KB system, trained on 28,561 real +detector frames recorded on a complete 13^4 grid of the four KB jack motors (3 keV). + +Motor vector columns (KB jack positions, in motor units): + 0 = KBH downstream jack, 1 = KBH upstream, 2 = KBV downstream, 3 = KBV upstream + +The trained weights are NOT part of this repository (16 MB, derived from measured +beamline data). ``resolve_weights_path`` finds them from, in order: an explicit path, +the ``BLOP_SIM_TES_WEIGHTS`` environment variable, or a cached download from +``TES_WEIGHTS_URL`` when a release URL is configured. +""" + +import os +import urllib.request +from pathlib import Path +from typing import Any + +import numpy as np + +# Release URL for the trained weights (filled in once the weights are published; +# can be overridden with the TES_WEIGHTS_URL environment variable). +WEIGHTS_URL: str | None = None + +_CACHE_DIR = Path("~/.cache/blop-sim").expanduser() + + +def resolve_weights_path(weights_path: str | os.PathLike[str] | None = None) -> Path: + """Locate the trained TES emulator weights (.npz). + + Resolution order: explicit ``weights_path`` argument, the ``BLOP_SIM_TES_WEIGHTS`` + environment variable, then a cached download from ``TES_WEIGHTS_URL`` / + ``WEIGHTS_URL``. Raises with instructions if none is available. + """ + if weights_path is not None: + path = Path(weights_path).expanduser() + if not path.exists(): + raise FileNotFoundError(f"TES emulator weights not found at {path}") + return path + + env_path = os.environ.get("BLOP_SIM_TES_WEIGHTS") + if env_path: + path = Path(env_path).expanduser() + if not path.exists(): + raise FileNotFoundError(f"BLOP_SIM_TES_WEIGHTS points to a missing file: {path}") + return path + + url = os.environ.get("TES_WEIGHTS_URL", WEIGHTS_URL) + if url: + cached = _CACHE_DIR / "tes_emulator_weights.npz" + if not cached.exists(): + _CACHE_DIR.mkdir(parents=True, exist_ok=True) + tmp = cached.with_suffix(".part") + urllib.request.urlretrieve(url, tmp) # noqa: S310 + tmp.rename(cached) + return cached + + raise RuntimeError( + "TES emulator weights not found. Either pass weights_path=..., set the " + "BLOP_SIM_TES_WEIGHTS environment variable to the .npz file, or set " + "TES_WEIGHTS_URL to a download location. The weights ship with the " + "tes-emulator package (https://github.com/FLlorente/tes-emulator)." + ) + + +class TESModel: + """Inference wrapper around the trained TES emulator weights. + + ``predict(m)`` maps (n, 4) KB jack positions to L1-normalized beam images, + detector-count flux, and an in-domain flag. Heads: ``mlp`` (smooth, default), + ``interp`` (exact at the recorded grid points), ``poly`` (physics ridge). + """ + + def __init__(self, weights_path: str | os.PathLike[str] | None = None, head: str | None = None): + import torch # deferred: keep blop_sim importable without torch + + self._torch = torch + w: dict[str, Any] = dict(np.load(resolve_weights_path(weights_path), allow_pickle=False)) + self.K = int(w["pca_K"]) + self.patch = int(w["patch"]) + self.win = int(w["win"]) + self.frame_shape = (int(w["frame_h"]), int(w["frame_w"])) + self.head = head or str(w.get("default_head", "mlp")) + t = lambda k: torch.tensor(w[k], dtype=torch.float64) # noqa: E731 + self._w = {k: t(k) for k in w if k.startswith(("mlp_", "poly_", "cenres_", "interp_", "cen_coef", "pca_"))} + self.box_lo = np.asarray(w["box_lo"], dtype=float) + self.box_hi = np.asarray(w["box_hi"], dtype=float) + self._noise = None + if "noise_ptr" in w: + self._noise = {k: w[f"noise_{k}"] for k in ("rows", "cols", "vals", "ptr")} + + # -- heads --------------------------------------------------------------- + + @staticmethod + def _kb_basis(m): + """Degree-2 KB physics basis with intercept: (n, 4) -> (n, 15).""" + import torch + + p_h, h_h = m[:, 0] - m[:, 1], m[:, 0] + m[:, 1] + p_v, h_v = m[:, 2] - m[:, 3], m[:, 2] + m[:, 3] + f = [p_h, h_h, p_v, h_v] + feats = list(f) + for i in range(4): + for j in range(i, 4): + feats.append(f[i] * f[j]) + return torch.stack([torch.ones_like(p_h), *feats], dim=1) + + def _interp(self, m): + """4-D multilinear interpolation on the full-grid table (clamped to the box).""" + torch = self._torch + table, axes = self._w["interp_table"], self._w["interp_axes"] + n = m.shape[0] + idxs, fracs = [], [] + for c in range(4): + ax = axes[c] + mc = torch.clamp(m[:, c], ax[0], ax[-1]) + i = torch.clamp(torch.searchsorted(ax.contiguous(), mc.contiguous(), right=True) - 1, 0, len(ax) - 2) + fracs.append((mc - ax[i]) / (ax[i + 1] - ax[i])) + idxs.append(i) + out = torch.zeros(n, table.shape[-1], dtype=m.dtype) + for corner in range(16): + wgt = torch.ones(n, dtype=m.dtype) + ix = [] + for c in range(4): + d = (corner >> c) & 1 + wgt = wgt * (fracs[c] if d else 1.0 - fracs[c]) + ix.append(idxs[c] + d) + out = out + wgt[:, None] * table[ix[0], ix[1], ix[2], ix[3]] + return out + + def _anchor(self, m): + """Centroid anchor (cx, cy): interp table, or physics ridge + MLP residual.""" + torch, w = self._torch, self._w + if self.head == "interp": + return self._interp(m)[:, :2] + a = self._kb_basis(m) @ w["cen_coef"] + if "cenres_W0" in w: + z = (m - w["cenres_xmean"]) / w["cenres_xstd"] + z = torch.nn.functional.silu(z @ w["cenres_W0"].T + w["cenres_b0"]) + z = torch.nn.functional.silu(z @ w["cenres_W1"].T + w["cenres_b1"]) + z = z @ w["cenres_W2"].T + w["cenres_b2"] + a = a + z * w["cenres_ystd"] + w["cenres_ymean"] + return a + + def _codes_flux(self, m): + """Head evaluation -> (PCA codes (n, K), flux (n,)).""" + torch, w = self._torch, self._w + if self.head == "interp": + y = self._interp(m) + return y[:, 2 : 2 + self.K], y[:, 2 + self.K] + if self.head == "mlp": + z = (m - w["mlp_xmean"]) / w["mlp_xstd"] + z = torch.nn.functional.silu(z @ w["mlp_W0"].T + w["mlp_b0"]) + z = torch.nn.functional.silu(z @ w["mlp_W1"].T + w["mlp_b1"]) + y = (z @ w["mlp_W2"].T + w["mlp_b2"]) * w["mlp_ystd"] + w["mlp_ymean"] + elif self.head == "poly": + x = self._kb_basis(m) + xs = torch.cat([torch.ones_like(m[:, :1]), (x[:, 1:] - w["poly_xmean"]) / w["poly_xstd"]], dim=1) + y = xs @ w["poly_coef"] * w["poly_ystd"] + w["poly_ymean"] + else: + raise ValueError(f"unknown head {self.head!r}") + return y[:, : self.K], y[:, self.K] + + def _decode(self, codes, anchor): + """PCA codes + anchor -> L1-normalized full frames via grid_sample placement.""" + torch = self._torch + n = codes.shape[0] + p, win = self.patch, self.win + h, w_ = self.frame_shape + small = (codes @ self._w["pca_components"] + self._w["pca_mean"]).reshape(n, 1, p, p) + small = torch.clamp(small, min=0.0) + ys = torch.arange(h, dtype=codes.dtype) + xs = torch.arange(w_, dtype=codes.dtype) + wx = xs[None, None, :] - (anchor[:, 0, None, None] - win // 2) + wy = ys[None, :, None] - (anchor[:, 1, None, None] - win // 2) + scale = (p - 1) / (win - 1) + gx = 2.0 * (wx * scale) / (p - 1) - 1.0 + gy = 2.0 * (wy * scale) / (p - 1) - 1.0 + grid = torch.stack([gx.expand(n, h, w_), gy.expand(n, h, w_)], dim=-1) + canvas = torch.nn.functional.grid_sample(small, grid, mode="bilinear", padding_mode="zeros", align_corners=True)[ + :, 0 + ] + return canvas / torch.clamp(canvas.sum(dim=(1, 2), keepdim=True), min=1e-12) + + # -- public API ---------------------------------------------------------- + + def in_domain(self, m: np.ndarray) -> np.ndarray: + """True where the motor vector lies inside the recorded motor box.""" + m = np.atleast_2d(np.asarray(m, dtype=float)) + return ((m >= self.box_lo) & (m <= self.box_hi)).all(axis=1) + + def predict(self, m: np.ndarray, sample: bool = False, rng: np.random.Generator | None = None) -> dict[str, np.ndarray]: + """(n, 4) motors -> dict(image (n, H, W) L1-normalized, flux (n,), in_domain (n,)). + + ``sample=True`` adds real background-speck patterns from the training data. + """ + torch = self._torch + mt = torch.as_tensor(np.atleast_2d(np.asarray(m, dtype=float)), dtype=torch.float64) + with torch.no_grad(): + anchor = self._anchor(mt) + codes, flux = self._codes_flux(mt) + image = self._decode(codes, anchor).numpy() + if sample: + if self._noise is None: + raise RuntimeError("weights file has no noise bank; use sample=False") + rng = rng or np.random.default_rng() + nb = self._noise + n_bank = len(nb["ptr"]) - 1 + for i in range(image.shape[0]): + k = int(rng.integers(0, n_bank)) + s, e = nb["ptr"][k], nb["ptr"][k + 1] + image[i, nb["rows"][s:e], nb["cols"][s:e]] += nb["vals"][s:e] + return {"image": image, "flux": flux.numpy(), "in_domain": self.in_domain(m)} + + +__all__ = ["TESModel", "resolve_weights_path", "WEIGHTS_URL"] diff --git a/sim/blop_sim/backends/tes.py b/sim/blop_sim/backends/tes.py new file mode 100644 index 00000000..93648883 --- /dev/null +++ b/sim/blop_sim/backends/tes.py @@ -0,0 +1,79 @@ +"""TES beamline simulation backend, driven by a data-trained emulator.""" + +import numpy as np + +from .core import SimBackend +from .models.tes_model import TESModel + + +class TESBackend(SimBackend): + """Simulation backend backed by the NSLS-II TES beamline emulator. + + Unlike :class:`SimpleBackend` (analytic Gaussian) and :class:`XRTBackend` + (ray-tracing), this backend replays a model trained on 28,561 *measured* + detector frames, so the beam response to the four KB jack motors — including + astigmatism, non-Gaussian streaks, and flux variation — is that of the real + beamline at 3 keV. + + It reads the same KB mirror devices as ``SimpleBackend`` + (:class:`blop_sim.devices.simple.KBMirror`): the *horizontal* mirror's + (downstream, upstream) jacks map to motor columns (0, 1) and the *vertical* + mirror's to columns (2, 3). Valid jack positions live inside the recorded + motor box (see :meth:`motor_box`); outside it the model extrapolates and + should not be trusted. + + Args: + weights_path: Path to the trained weights (.npz). Defaults to the + ``BLOP_SIM_TES_WEIGHTS`` environment variable (see + ``models.tes_model.resolve_weights_path``). + noise: If True, add real background-speck patterns to each image. + seed: Seed for the noise sampling. + """ + + def __init__(self, weights_path: str | None = None, noise: bool = False, seed: int | None = None): + super().__init__() + self._model = TESModel(weights_path) + self._image_shape = self._model.frame_shape + self._noise = noise + self._rng = np.random.default_rng(seed) + + def motor_box(self) -> dict[str, tuple[float, float]]: + """Per-motor (lo, hi) validity bounds, keyed by column role. + + Use these as ``RangeDOF`` bounds so the optimizer never leaves the box + the model was trained on. + """ + lo, hi = self._model.box_lo, self._model.box_hi + names = ("kbh_downstream", "kbh_upstream", "kbv_downstream", "kbv_upstream") + return {name: (float(lo[i]), float(hi[i])) for i, name in enumerate(names)} + + async def generate_beam(self) -> np.ndarray: + """Generate a detector-count beam image at the current jack positions. + + Returns: + 2D numpy array with shape ``self._image_shape`` (counts: the model's + L1-normalized beam scaled by its predicted flux, so intensity-based + outcome constraints behave like on a real detector). + """ + m = await self._get_jack_positions() + out = self._model.predict(m[None], sample=self._noise, rng=self._rng) + return (out["flux"][0] * out["image"][0]).astype(np.float64) + + async def _get_jack_positions(self) -> np.ndarray: + """Assemble the 4-motor vector from registered KB mirror devices. + + Motors default to the center of the recorded box, so a beamline with no + mirrors registered (or only one) still produces a valid beam. + """ + m = (self._model.box_lo + self._model.box_hi) / 2 + for name, device in self._device_states.items(): + if device["type"] == "kb_mirror_simple": + state = await self._get_device_state(name) + if state["orientation"] == "horizontal": + m[0], m[1] = state["downstream"], state["upstream"] + elif state["orientation"] == "vertical": + m[2], m[3] = state["downstream"], state["upstream"] + return m + + +__all__ = ["TESBackend"] diff --git a/sim/pyproject.toml b/sim/pyproject.toml index 2e5a09c2..a232a670 100644 --- a/sim/pyproject.toml +++ b/sim/pyproject.toml @@ -18,6 +18,7 @@ dependencies = [ "numpy", "scipy", "matplotlib", + "torch", "xrt>=2.0.0b1", ]