add fixed_mask device

Author: josephhanrahan

sim/blop_sim/backends/__init__.py

@@ -3,5 +3,6 @@
 from .core import SimBackend
 from .simple import SimpleBackend
 from .xrt import XRTBackend
+from .xrt_xas import XRTXASBackend
 
 __all__ = ["SimBackend", "SimpleBackend", "XRTBackend", "XRTXASBackend"]

sim/blop_sim/backends/models/xrt_xas_model.py

@@ -17,19 +17,41 @@
 
 import numpy as np
 
-import xrt.backends.raycing.sources as rsources
+# import xrt.backends.raycing.sources as rsources
 import xrt.backends.raycing.screens as rscreens
 import xrt.backends.raycing.materials as rmats
-import xrt.backends.raycing.materials.elemental as rmatsel
-import xrt.backends.raycing.materials.compounds as rmatsco
-import xrt.backends.raycing.materials.crystals as rmatscr
+# import xrt.backends.raycing.materials.elemental as rmatsel
+# import xrt.backends.raycing.materials.compounds as rmatsco
+# import xrt.backends.raycing.materials.crystals as rmatscr
 import xrt.backends.raycing.oes as roes
 import xrt.backends.raycing.apertures as rapts
 import xrt.backends.raycing.run as rrun
 import xrt.backends.raycing as raycing
 import xrt.plotter as xrtplot
 import xrt.runner as xrtrun
 
+def build_histRGB(lb, gb, limits=None, isScreen=False, shape=None):
+    if shape is None:
+        shape = [256, 256]
+    good = (lb.state == 1) | (lb.state == 2)
+    if isScreen:
+        x, y, z = lb.x[good], lb.z[good], lb.y[good]
+    else:
+        x, y, z = lb.x[good], lb.y[good], lb.z[good]
+    goodlen = len(lb.x[good])
+    hist2dRGB = np.zeros((shape[1], shape[0], 3), dtype=np.float64)
+    hist2d = np.zeros((shape[1], shape[0]), dtype=np.float64)
+
+    if limits is None and goodlen > 0:
+        limits = np.array([[np.min(x), np.max(x)], [np.min(y), np.max(y)], [np.min(z), np.max(z)]])
+
+    if goodlen > 0:
+        beamLimits = [limits[1], limits[0]] or None
+        flux = gb.Jss[good] + gb.Jpp[good]
+        hist2d, _, _ = np.histogram2d(y, x, bins=[shape[1], shape[0]], range=beamLimits, weights=flux)
+        hist2dRGB = None
+    return hist2d, hist2dRGB, limits
+
 CVD = rmats.Material(
     elements=['C'],
     quantities=[1.0],
@@ -51,8 +73,7 @@
     V=160.17128543981727,
     elements=['Si'],
     quantities=[1.0],
-    name=r"Si220",
-    kind=r"crystal")
+    name=r"Si220")
 
 CVDcoating = rmats.Material(
     elements=['C'],
@@ -91,20 +112,20 @@
     beta=1.5707963267948966,
     gamma=1.5707963267948966,
     atomsFraction=[1, 1, 1, 1, 1, 1, 1, 1],
-    d=1.916079064786341,
-    V=159.1751463370933,
-    elements=['Si', 'Si', 'Si', 'Si', 'Si', 'Si', 'Si', 'Si'],
-    quantities=[1, 1, 1, 1, 1, 1, 1, 1],
-    rho=2.3439368026411915,
-    kind=r"crystal harmonics")
+    # V=159.1751463370933,
+    # elements=['Si', 'Si', 'Si', 'Si', 'Si', 'Si', 'Si', 'Si'],
+    # quantities=[1, 1, 1, 1, 1, 1, 1, 1],
+    # rho=2.3439368026411915,
+    # kind=r"crystal harmonics"
+    )
 
 
 def build_beamline():
     BioXAS_Main = raycing.BeamLine(
-        alignE=8000,
-        name=r"BioXAS_Main")
+        alignE=8000)
 
-    BioXAS_Main.Wiggler = raycing.sources.synchr.Wiggler(
+    # BioXAS_Main.Wiggler = raycing.sources.synchr.Wiggler(
+    BioXAS_Main.Wiggler = raycing.sources_synchr.Wiggler(
         bl=BioXAS_Main,
         name=r"Wiggler",
         center=[0, 0, 0],
@@ -202,9 +223,10 @@ def build_beamline():
         center=[0, 26000, r"auto"],
         x=[1.0, -0.0, 0.0],
         z=[0.0, 0.0, 1.0],
-        limPhysX=[0.0, 0.0],
-        limPhysY=[0.0, 0.0],
-        cLimits=[0.0, 0.0])
+        # limPhysX=[0.0, 0.0],
+        # limPhysY=[0.0, 0.0],
+        # cLimits=[0.0, 0.0]
+        )
 
     BioXAS_Main.Mirror2 = roes.ToroidMirror(
         bl=BioXAS_Main,
@@ -266,9 +288,10 @@ def build_beamline():
         center=[0, 30650, r"auto"],
         x=[1.0, -0.0, 0.0],
         z=[0.0, 0.0, 1.0],
-        limPhysX=[0.0, 0.0],
-        limPhysY=[0.0, 0.0],
-        cLimits=[0.0, 0.0])
+        # limPhysX=[0.0, 0.0],
+        # limPhysY=[0.0, 0.0],
+        # cLimits=[0.0, 0.0])
+    )
 
     return BioXAS_Main
 

sim/blop_sim/backends/xrt_xas.py

@@ -17,7 +17,6 @@ def __init__(self, noise: bool = False):
         """Initialize XRT backend."""
         super().__init__()
         self._beamline = None
-        # self._limits = [[-0.6, 0.6], [-0.45, 0.45]]
         self._noise = noise
 
     def _ensure_beamline(self):
@@ -33,12 +32,13 @@ async def generate_beam(self) -> np.ndarray:
         """
         self._ensure_beamline()
 
-        # Get KB mirror radii from devices
-        # mirror_radii = await self._get_mirror_radii()
+        # Get fixed mask position (x, y, z)
+        fixed_mask_position = await self._get_fixed_mask_position()
 
+        print(fixed_mask_position)
         # Update XRT beamline mirror parameters
-        # self._beamline.toroidMirror01.R = mirror_radii[0]  # Vertical mirror
-        # self._beamline.toroidMirror02.R = mirror_radii[1]  # Horizontal mirror
+        if fixed_mask_position:
+            self._beamline.FEMask.center = fixed_mask_position
 
         # Run ray tracing
         outDict = run_process(self._beamline)
@@ -53,5 +53,17 @@ async def generate_beam(self) -> np.ndarray:
             image += 1e-3 * np.abs(np.random.standard_normal(size=image.shape))
 
         return image
-
-__all__ = ["XRTBackend"]
+    
+    async def _get_fixed_mask_position(self) -> list[float]:
+        for name, device in self._device_states.items():
+            # print(name, device)
+            print(device["type"])
+            if device["type"] == "fixed_mask_xrt":
+                state = await self._get_device_state(name)
+                print(state)
+                # print("DEBUG STATE: {}".format(state))
+                position = [state["x"], state["y"], state["z"]]
+                return position
+
+
+__all__ = ["XRTXASBackend"]

sim/blop_sim/devices/xrt/__init__.py

@@ -1,5 +1,6 @@
 """XRTBackend-specific devices."""
 
 from .kb_mirror import KBMirror
+from .fixed_mask import FixedMask
 
-__all__ = ["KBMirror"]
+__all__ = ["KBMirror", "FixedMask"]

sim/blop_sim/devices/xrt/fixed_mask.py

@@ -0,0 +1,54 @@
+"""Fixed Mask devices for XRTBackend."""
+
+from ophyd_async.core import StandardReadable, soft_signal_rw
+from ophyd_async.core import StandardReadableFormat as Format
+
+from ...backends import SimBackend
+
+
+class FixedMask(StandardReadable):
+    """Fixed mask position control (for XRTBackend).
+
+    Exposes the coordinates that directly control the XRT fixed mask position.
+    Used with XRTBackend for ray-tracing simulation.
+
+    Args:
+        backend: Simulation backend (should be XRTBackend)
+        position: list of floats that define the [x, y, z] position of the mask
+        name: Device name
+    """
+
+    def __init__(
+        self,
+        backend: SimBackend,
+        initial_position: list[float],
+        name: str = "",
+    ):
+        self._backend = backend
+
+        # Curvature radius signal
+        with self.add_children_as_readables(Format.HINTED_SIGNAL):
+            self.x = soft_signal_rw(float, initial_position[0])
+            self.y = soft_signal_rw(float, initial_position[1])
+            self.z = soft_signal_rw(float, initial_position[2])
+
+        super().__init__(name=name)
+
+        # Register with backend
+        backend.register_device(
+            device_name=name,
+            device_type="fixed_mask_xrt",
+            get_state_callback=self._get_state,
+        )
+
+    async def _get_state(self) -> dict:
+        """Get current mirror state for backend (async)."""
+        # position = self.position.get_value()
+        return {
+            "x": await self.x.get_value(),
+            "y": await self.y.get_value(),
+            "z": await self.z.get_value(),
+        }
+
+
+__all__ = ["FixedMask"]

sim/pyproject.toml

@@ -18,7 +18,7 @@ dependencies = [
   "numpy",
   "scipy",
   "matplotlib",
-  "xrt>=1.6.1,<2",
+  "xrt==1.6.1",
 ]
 
 [tool.hatch.build.targets.wheel]