diff --git a/Dockerfile_processor b/Dockerfile_processor index 6960ba6..d46b5ef 100644 --- a/Dockerfile_processor +++ b/Dockerfile_processor @@ -6,4 +6,4 @@ COPY . /app # RUN pip install ./arroyo RUN pip install --no-cache-dir --upgrade . -CMD ["python", "-m", "tr_ap_xps.apps.processor_cli"] +ENTRYPOINT ["arroyo", "run"] diff --git a/block_configs/labview_processor_block.yaml b/block_configs/labview_processor_block.yaml new file mode 100644 index 0000000..a1b9da0 --- /dev/null +++ b/block_configs/labview_processor_block.yaml @@ -0,0 +1,18 @@ +blocks: + - name: XPS LabVIEW Operator + description: Real-time AP-XPS processor for LabVIEW data + + operator: + class: tr_ap_xps.pipeline.xps_operator.build_xps_operator + kwargs: + build_heatmaps: true + + listeners: + - class: tr_ap_xps.labview.xps_labview_listener_factory + kwargs: + zmq_connection_address: "tcp://localhost:5555" + + publishers: + - class: tr_ap_xps.websockets.xps_ws_publisher_factory + kwargs: + ws_url: "ws://0.0.0.0:8001/xps_operator" \ No newline at end of file diff --git a/block_configs/timepix_processor_block.yaml b/block_configs/timepix_processor_block.yaml new file mode 100644 index 0000000..9e38315 --- /dev/null +++ b/block_configs/timepix_processor_block.yaml @@ -0,0 +1,18 @@ +blocks: + - name: XPS Timepix Operator + description: Real-time AP-XPS processor for Timepix data + + operator: + class: tr_ap_xps.pipeline.xps_operator.build_xps_operator + kwargs: + build_heatmaps: false + + listeners: + - class: tr_ap_xps.timepix.xps_timepix_listener_factory + kwargs: + zmq_connection_address: "tcp://host.docker.internal:5657" + + publishers: + - class: tr_ap_xps.websockets.xps_ws_publisher_factory + kwargs: + ws_url: "ws://0.0.0.0:8001/xps_operator" \ No newline at end of file diff --git a/docker-compose.yml b/docker-compose.yml index 3caea88..7eae1da 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -1,39 +1,37 @@ services: tiled: - # see the file ./tiled/deploy/config.yml for detailed configuration of tiled - image: ghcr.io/bluesky/tiled:main - ports: - - "8000:8000" - environment: - - TILED_SINGLE_USER_API_KEY=${TILED_SINGLE_USER_API_KEY} - volumes: - - ./services/tiled/deploy:/deploy:Z - - ./data:/data:Z - restart: unless-stopped - healthcheck: - test: ["CMD", "curl", "-f", "http://tiled:8000/healthz"] - interval: 10s - timeout: 5s - retries: 3 - networks: - aps_net: + # see the file ./tiled/deploy/config.yml for detailed configuration of tiled + image: ghcr.io/bluesky/tiled:main + ports: + - "8000:8000" + environment: + - TILED_SINGLE_USER_API_KEY=${TILED_SINGLE_USER_API_KEY} + volumes: + - ./services/tiled/deploy:/deploy:Z + - ./data:/data:Z + restart: unless-stopped + healthcheck: + test: ["CMD", "curl", "-f", "http://tiled:8000/healthz"] + interval: 10s + timeout: 5s + retries: 3 + networks: + - mle_net + - processor: - command: python -m tr_ap_xps.apps.processor_cli_tpx + timepix_processor: build: context: . dockerfile: Dockerfile_processor + command: "block_configs/timepix_processor_block.yaml" restart: unless-stopped ports: - "8001:8001" volumes: - .:/app:Z - - ./settings_container.yaml:/app/settings.yaml:Z networks: - aps_net: - depends_on: - tiled: - condition: service_healthy + mle_net: + frontend: build: @@ -45,8 +43,9 @@ services: volumes: - .:/app:Z networks: - aps_net: + mle_net: networks: - aps_net: + mle_net: + name: mle_net driver: bridge diff --git a/pyproject.toml b/pyproject.toml index 5a76b81..a57198d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,7 +13,7 @@ classifiers = [ ] dependencies = [ - "arroyopy", + "arroyopy @ git+https://github.com/als-computing/arroyopy.git", "astropy", "dynaconf", "python-dotenv", diff --git a/settings.yaml b/settings.yaml deleted file mode 100644 index a6fccdc..0000000 --- a/settings.yaml +++ /dev/null @@ -1,23 +0,0 @@ -# xps_operator: -# log_level: "INFO" -# tiled_uri: "http://localhost:8000" -# lv_zmq_listener: -# zmq_pub_address: "tcp://localhost" -# zmq_pub_port: 5555 -# tpx_zmq_listener: -# zmq_pub_address: "tcp://localhost" -# zmq_pub_port: 5657 -# websockets_publisher: -# host: "0.0.0.0" -# port: 8001 -xps_operator: - log_level: "DEBUG" - tiled_uri: "https://tiled-staging.computing.als.lbl.gov" - tiled_api_key: afecf30c4fcf142a9c30142dc8b7d5212a55a100715c580f034afcddc317be822174ead4 - lv_zmq_listener: - zmq_pub_address: "tcp://localhost" - zmq_pub_port: 5555 - tpx_zmq_listener: - zmq_pub_address: "tcp://131.243.191.226" - zmq_pub_port: 5657 - websocket_url: "ws://localhost:8766" \ No newline at end of file diff --git a/settings_container.yaml b/settings_container.yaml deleted file mode 100644 index ff5c8ff..0000000 --- a/settings_container.yaml +++ /dev/null @@ -1,10 +0,0 @@ -xps: - log_level: "INFO" - tiled_uri: "http://tiled:8000" - tiled_api_key: "@env TILED_SINGLE_USER_API_KEY" - lv_zmq_listener: - zmq_pub_address: "tcp://simulator" - zmq_pub_port: 5555 - websockets_publisher: - host: "0.0.0.0" - port: 8001 diff --git a/src/tr_ap_xps/apps/__init__.py b/src/tr_ap_xps/apps/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/tr_ap_xps/apps/processor_cli.py b/src/tr_ap_xps/apps/processor_cli.py deleted file mode 100644 index 9d52d11..0000000 --- a/src/tr_ap_xps/apps/processor_cli.py +++ /dev/null @@ -1,77 +0,0 @@ -import asyncio -import logging -import signal - -import typer -from tiled.client import from_uri -from tiled.client.node import Container - -from ..config import settings -from ..labview import XPSLabviewZMQListener, setup_zmq -from ..log_utils import setup_logger -from ..pipeline.xps_operator import XPSOperator -from ..tiled import TiledPublisher -from ..websockets import XPSWSResultPublisher - -app = typer.Typer() -logger = logging.getLogger("tr_ap_xps") -setup_logger(logger) - -app_settings = settings.xps_operator - -def tiled_runs_container() -> Container: - try: - client = from_uri(app_settings.tiled_uri, api_key=app_settings.tiled_api_key) - if client.get("runs") is None: # TODO test case - client.create_container("runs") - return client["runs"] - except Exception as e: - logger.error(f"Error connecting to Tiled: {e}") - - -@app.command() -async def listen() -> None: - try: - logger.setLevel(app_settings.log_level.upper()) - logger.debug("DEBUG LOGGING SET") - logger.info( - f"lv_zmq_pub_address: {app_settings.lv_zmq_listener.zmq_pub_address}" - ) - logger.info(f"lv_zmq_pub_address: {app_settings.lv_zmq_listener.zmq_pub_port}") - logger.info(f"tiled_uri: {app_settings.tiled_uri}") - logger.info( - f"tiled_api_key: {'****' if app_settings.tiled_api_key else 'NOT SET!!!'}" - ) - - received_sigterm = {"received": False} # Define the variable received_sigterm - - # setup websocket server - operator = XPSOperator() - ws_publisher = XPSWSResultPublisher(app_settings.websocket_url) - # tiled_pub = TiledPublisher(tiled_runs_container()) - - operator.add_publisher(ws_publisher) - # operator.add_publisher(tiled_pub) - # connect to labview zmq - - lv_zmq_socket = setup_zmq() - listener = XPSLabviewZMQListener(operator=operator, zmq_socket=lv_zmq_socket) - - # Wait for both tasks to complete - await asyncio.gather(listener.start(), ws_publisher.start()) - - def handle_sigterm(signum, frame): - logger.info("SIGTERM received, stopping...") - received_sigterm["received"] = True - asyncio.create_task(listener.stop()) - asyncio.create_task(ws_publisher.stop()) - - # Register the handler for SIGTERM - signal.signal(signal.SIGTERM, handle_sigterm) - except Exception as e: - logger.error(f"Error setting up XPS processor {e}") - raise e - - -if __name__ == "__main__": - asyncio.run(listen()) diff --git a/src/tr_ap_xps/apps/processor_cli_tpx.py b/src/tr_ap_xps/apps/processor_cli_tpx.py deleted file mode 100644 index d4e484a..0000000 --- a/src/tr_ap_xps/apps/processor_cli_tpx.py +++ /dev/null @@ -1,78 +0,0 @@ -import asyncio -import logging -import signal - -import typer -from tiled.client import from_uri -from tiled.client.node import Container - -from ..config import settings -# from ..labview import XPSLabviewZMQListener, setup_zmq -from ..timepix import XPSTimepixZMQListener, setup_zmq -from ..log_utils import setup_logger -from ..pipeline.xps_operator import XPSOperator -from ..tiled import TiledPublisher -from ..websockets import XPSWSResultPublisher - -app = typer.Typer() -logger = logging.getLogger("tr_ap_xps") -setup_logger(logger) - -app_settings = settings.xps_operator - -def tiled_runs_container() -> Container: - try: - client = from_uri(app_settings.tiled_uri, api_key=app_settings.tiled_api_key) - if client.get("runs") is None: # TODO test case - client.create_container("runs") - return client["runs"] - except Exception as e: - logger.error(f"Error connecting to Tiled: {e}") - - -@app.command() -async def listen() -> None: - try: - logger.setLevel(app_settings.log_level.upper()) - logger.debug("DEBUG LOGGING SET") - logger.info( - f"tpx_zmq_socket_address: {app_settings.tpx_zmq_listener.zmq_pub_address}" - ) - logger.info(f"tpx_zmq_socket_port: {app_settings.tpx_zmq_listener.zmq_pub_port}") - logger.info(f"tiled_uri: {app_settings.tiled_uri}") - logger.info( - f"tiled_api_key: {'****' if app_settings.tiled_api_key else 'NOT SET!!!'}" - ) - - received_sigterm = {"received": False} # Define the variable received_sigterm - - # setup websocket server - operator = XPSOperator() - ws_publisher = XPSWSResultPublisher(app_settings.websocket_url) - # tiled_pub = TiledPublisher(tiled_runs_container()) - - operator.add_publisher(ws_publisher) - # operator.add_publisher(tiled_pub) - # connect to labview zmq - - tpx_zmq_socket = setup_zmq() - listener = XPSTimepixZMQListener(operator=operator, zmq_socket=tpx_zmq_socket) - - # Wait for both tasks to complete - await asyncio.gather(listener.start(), ws_publisher.start()) - - def handle_sigterm(signum, frame): - logger.info("SIGTERM received, stopping...") - received_sigterm["received"] = True - asyncio.create_task(listener.stop()) - asyncio.create_task(ws_publisher.stop()) - - # Register the handler for SIGTERM - signal.signal(signal.SIGTERM, handle_sigterm) - except Exception as e: - logger.error(f"Error setting up XPS processor {e}") - raise e - - -if __name__ == "__main__": - asyncio.run(listen()) diff --git a/src/tr_ap_xps/labview.py b/src/tr_ap_xps/labview.py index 1050c20..0bfc77c 100644 --- a/src/tr_ap_xps/labview.py +++ b/src/tr_ap_xps/labview.py @@ -7,7 +7,6 @@ from arroyopy.zmq import ZMQListener -from .config import settings from .schemas import NumpyArrayModel, XPSImageInfo, XPSRawEvent, XPSStart, XPSStop # Maintain a map of LabView datatypes. LabView sends BigE, @@ -26,25 +25,21 @@ "Double Float": np.dtype(np.double).newbyteorder(">"), } -app_settings = settings.xps_operator - logger = logging.getLogger(__name__) -def setup_zmq(): +# removed app_settings; accepts zmq_connection_address as parameters instead +def setup_zmq( + zmq_connection_address: str = "tcp://localhost:5555", +): ctx = zmq.asyncio.Context() lv_zmq_socket = ctx.socket(zmq.SUB) lv_zmq_socket.setsockopt(zmq.RCVHWM, 100000) - logger.info( - f"binding to: {app_settings.lv_zmq_listener.zmq_pub_address}:{app_settings.lv_zmq_listener.zmq_pub_port}" - ) - lv_zmq_socket.connect( - f"{app_settings.lv_zmq_listener.zmq_pub_address}:{app_settings.lv_zmq_listener.zmq_pub_port}" - ) + logger.info(f"connecting to: {zmq_connection_address}") + lv_zmq_socket.connect(zmq_connection_address) lv_zmq_socket.setsockopt(zmq.SUBSCRIBE, b"") return lv_zmq_socket - class XPSLabviewZMQListener(ZMQListener): stop_signal = False @@ -135,3 +130,11 @@ def _build_stop(message: dict) -> XPSStop: if logger.getEffectiveLevel() == logging.DEBUG: logger.debug(f"stop: {message}") return XPSStop(**message) + + +def xps_labview_listener_factory( + operator, + zmq_connection_address: str = "tcp://localhost:5555", +) -> XPSLabviewZMQListener: + socket = setup_zmq(zmq_connection_address) + return XPSLabviewZMQListener(operator=operator, zmq_socket=socket) \ No newline at end of file diff --git a/src/tr_ap_xps/pipeline/xps_operator.py b/src/tr_ap_xps/pipeline/xps_operator.py index 99b71bf..0d358c4 100644 --- a/src/tr_ap_xps/pipeline/xps_operator.py +++ b/src/tr_ap_xps/pipeline/xps_operator.py @@ -1,25 +1,61 @@ import asyncio import logging +import os +import numpy as np from arroyopy.operator import Operator from arroyopy.schemas import Message -from ..schemas import XPSResult, XPSRawEvent, XPSResultStop, XPSStart, XPSStop +from ..log_utils import setup_logger +from ..schemas import ( + DataFrameModel, + NumpyArrayModel, + XPSRawEvent, + XPSResult, + XPSResultStart, + XPSResultStop, + XPSStart, + XPSStop, +) from ..timing import timer from .xps_processor import XPSProcessor -logger = logging.getLogger(__name__) +setup_logger( + logging.getLogger("tr_ap_xps"), log_level=os.getenv("LOGGING_LEVEL", "DEBUG") +) +logger = logging.getLogger(__name__) # resolves to "tr_ap_xps.pipeline.xps_operator" class XPSOperator(Operator): - """ - XPSOperator is responsible for handling XPS-related messages and processing frames. - - """ - def __init__(self, build_heatmaps: bool = False) -> None: + super().__init__() self.xps_processor = None self.build_heatmaps = build_heatmaps + self.cumulative_sum = None + self.total_cycles = 0 + + def _compute_timepix_arrays(self, message: XPSRawEvent): + raw_array = message.image.array.astype(np.float64) + cycles_in_flush = message.image_info.cycles_in_flush or 1 + logger.debug( + f"_compute_timepix_arrays: raw_array.shape={raw_array.shape}, cycles_in_flush={cycles_in_flush}" + ) + + if self.cumulative_sum is None: + self.cumulative_sum = raw_array.copy() + self.total_cycles = cycles_in_flush + else: + self.cumulative_sum += raw_array + self.total_cycles += cycles_in_flush + + integrated_2d = raw_array if raw_array.ndim == 2 else np.sum(raw_array, axis=1) + average = self.cumulative_sum / self.total_cycles + shot_mean_2d = average if average.ndim == 2 else np.sum(average, axis=1) + logger.debug( + f"_compute_timepix_arrays: integrated_2d.shape={integrated_2d.shape}, shot_mean_2d.shape={shot_mean_2d.shape}" + ) + + return integrated_2d, shot_mean_2d async def process(self, message: Message) -> None: """ @@ -28,20 +64,25 @@ async def process(self, message: Message) -> None: Args: message (Message): The message to be processed. It can be one of the following types: - - XPSStart: Initializes the XPSProcessor and publishes the start message. + - XPSStart: Initializes the XPSProcessor and publishes XPSResultStart. - XPSRawEvent: Processes a frame using the XPSProcessor and publishes the result. - - XPSStop: Finalizes the XPSProcessor and publishes the stop message. + - XPSStop: Finalizes the XPSProcessor and publishes XPSResultStop. Returns: None """ if isinstance(message, XPSStart): + logger.info(f"Start message received: scan_name={message.scan_name}") timer.reset() self.xps_processor = XPSProcessor(message) - await self.publish(message) + self.cumulative_sum = None + self.total_cycles = 0 + await self.publish(XPSResultStart(scan_name=message.scan_name)) elif isinstance(message, XPSRawEvent): - + logger.debug( + f"XPSRawEvent received: frame_number={message.image_info.frame_number}" + ) if self.build_heatmaps: if not self.xps_processor: logger.error( @@ -52,23 +93,53 @@ async def process(self, message: Message) -> None: self.xps_processor.process_frame, message ) else: - result = XPSResult( - shot_num=message.image_info.frame_number, - integrated_frames=message.image, - rolling_mean=None, - rolling_std=None, - frame_number=message.image_info.frame_number, - detected_peaks=None, - vfft=None, - ifft=None, - shot_recent=None, - shot_mean=None, - shot_std=None, - ) + try: + integrated_2d, shot_mean_2d = self._compute_timepix_arrays(message) + result = XPSResult( + shot_num=message.image_info.frame_number, + integrated_frames=NumpyArrayModel(array=integrated_2d), + frame_number=message.image_info.frame_number, + detected_peaks=None, + vfft=None, + ifft=None, + shot_recent=None, + shot_mean=NumpyArrayModel(array=shot_mean_2d), + shot_std=None, + ) + timer.end_frame() + except Exception as e: + logger.error(f"Error computing timepix arrays: {e}", exc_info=True) + return if result: - await self.publish(result) - # elif isinstance(message, XPSStop): - # data_frame_model = DataFrameModel(df=timer.timing_dataframe) - # new_msg = XPSResultStop(function_timings=data_frame_model) - # await self.publish(new_msg) - # self.xps_processor = None + logger.debug( + f"Publishing XPSResult: frame_number={result.frame_number}, shot_num={result.shot_num}" + ) + await self.publish(result) + else: + logger.debug("Result is None, not publishing") + + elif isinstance(message, XPSStop): + logger.info("Stop message received") + self.cumulative_sum = None + self.total_cycles = 0 + self.xps_processor = None + try: + timings_df = timer.timing_dataframe + await self.publish( + XPSResultStop( + function_timings=( + DataFrameModel(df=timings_df) + if not timings_df.empty + else None + ) + ) + ) + except Exception as e: + logger.error(f"Error publishing XPSResultStop: {e}", exc_info=True) + + else: + logger.warning(f"Unknown message type received: {type(message)}") + + +def build_xps_operator(build_heatmaps: bool = False) -> XPSOperator: + return XPSOperator(build_heatmaps=build_heatmaps) diff --git a/src/tr_ap_xps/pipeline/xps_processor.py b/src/tr_ap_xps/pipeline/xps_processor.py index f601569..12957d8 100644 --- a/src/tr_ap_xps/pipeline/xps_processor.py +++ b/src/tr_ap_xps/pipeline/xps_processor.py @@ -7,7 +7,7 @@ from .fft import calculate_fft_items from .peak_fitting import peak_fit -logger = logging.getLogger("tr_ap_xps.processor") +logger = logging.getLogger(__name__) class XPSProcessor: @@ -97,8 +97,9 @@ def process_frame(self, message: XPSRawEvent) -> None: shot_std=NumpyArrayModel(array=self.shot_rolling_std), ) self.shot_cache = None + timer.end_frame() return result except Exception as e: logger.exception(f"Error processing frame: {e}") return None - timer.end_frame() + diff --git a/src/tr_ap_xps/schemas.py b/src/tr_ap_xps/schemas.py index a061f02..34eb1b0 100644 --- a/src/tr_ap_xps/schemas.py +++ b/src/tr_ap_xps/schemas.py @@ -1,6 +1,6 @@ from typing import Literal, Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field from arroyopy.schemas import DataFrameModel, Event, Message, NumpyArrayModel, Start, Stop @@ -35,10 +35,9 @@ class XPSMessage(Message): class XPSStart(Start, XPSMessage): - msg_type: str = Literal["start"] - binding_energy: float = Field(..., alias="Binding Energy") - msg_type: str = Field("start", alias="msg_type") """ + Incoming start message from ZMQ. Supports both LabVIEW and Timepix sources. + LabVIEW Message: Incoming message from LabView at the start of a scan. @@ -71,27 +70,65 @@ class XPSStart(Start, XPSMessage): "File Ver": "1.0.0" } + Timepix Message (msgpack-encoded, from splash_timepix server): + + { + "msg_type": "start", + "scan_name": "acquisition_20250128T143022Z_a1b2c3d4", + "tdc_frequency_hz": 1000.0, + "t_delta_ns": 10.0, + "t_cycle_ns": 1000000.0, + "n_bins": 100, + "detector_size_x": 256, + "detector_size_y": 256, + "flush_interval_s": 1.0, + "cycles_per_flush": 1000, + "tdc_channel": 1, + "tdc_edge": "rising", + "collapse_y": false, + "zmq_port": 5657, + "tcp_port": 9090 + } """ - binding_energy: float = Field(..., alias="Binding Energy") - scan_name: str = Field(..., alias="scan_name") - f_trigger: int = Field(..., alias="F_Trigger") - f_untrigger: int = Field(..., alias="F_Un-Trigger") - f_dead: int = Field(..., alias="F_Dead") - f_reset: int = Field(..., alias="F_Reset") - ccd_nx: int = Field(..., alias="CCD_nx") - ccd_ny: int = Field(..., alias="CCD_ny") - pass_energy: float = Field(..., alias="Pass Energy") - center_energy: float = Field(..., alias="Center Energy") - offset_energy: float = Field(..., alias="Offset Energy") - lens_mode: str = Field(..., alias="Lens Mode") - rectangle: Rectangle = Field(..., alias="Rectangle") - dt: float = Field(..., alias="dt") - photon_energy: float = Field(..., alias="Photon Energy") - binding_energy: float = Field(..., alias="Binding Energy") - file_ver: str = Field(..., alias="File Ver") - data_type: str = Field(..., alias="data_type") + model_config = ConfigDict(populate_by_name=True, extra="allow") + + msg_type: str = Field("start", alias="msg_type") + + # LabVIEW fields — all Optional so Timepix start messages are also accepted + binding_energy: Optional[float] = Field(None, alias="Binding Energy") + scan_name: Optional[str] = Field(None, alias="scan_name") + f_trigger: Optional[int] = Field(None, alias="F_Trigger") + f_untrigger: Optional[int] = Field(None, alias="F_Un-Trigger") + f_dead: Optional[int] = Field(None, alias="F_Dead") + f_reset: Optional[int] = Field(None, alias="F_Reset") + ccd_nx: Optional[int] = Field(None, alias="CCD_nx") + ccd_ny: Optional[int] = Field(None, alias="CCD_ny") + pass_energy: Optional[float] = Field(None, alias="Pass Energy") + center_energy: Optional[float] = Field(None, alias="Center Energy") + offset_energy: Optional[float] = Field(None, alias="Offset Energy") + lens_mode: Optional[str] = Field(None, alias="Lens Mode") + rectangle: Optional[Rectangle] = Field(None, alias="Rectangle") + dt: Optional[float] = Field(None, alias="dt") + photon_energy: Optional[float] = Field(None, alias="Photon Energy") + file_ver: Optional[str] = Field(None, alias="File Ver") + data_type: Optional[str] = Field(None, alias="data_type") + + # Timepix fields from splash_timepix start message + tdc_frequency_hz: Optional[float] = None + t_delta_ns: Optional[float] = None + t_cycle_ns: Optional[float] = None + n_bins: Optional[int] = None + detector_size_x: Optional[int] = None + detector_size_y: Optional[int] = None + flush_interval_s: Optional[float] = None + cycles_per_flush: Optional[int] = None + tdc_channel: Optional[int] = None + tdc_edge: Optional[str] = None + collapse_y: Optional[bool] = None + zmq_port: Optional[int] = None + tcp_port: Optional[int] = None class XPSImageInfo(BaseModel): @@ -99,6 +136,10 @@ class XPSImageInfo(BaseModel): width: int height: int data_type: str + # Timepix event metadata + timestamp: Optional[float] = None + cycles_in_flush: Optional[int] = None + total_cycles: Optional[int] = None class XPSRawEvent(Event, XPSMessage): @@ -118,15 +159,41 @@ class XPSRawEvent(Event, XPSMessage): class XPSStop(Stop, XPSMessage): """ + LabVIEW Message: { "msg_type": "stop", "Num Frames": 1 } + Timepix Message: + { + "msg_type": "stop", + "scan_name": "acquisition_20250128T143022Z_a1b2c3d4", + "total_flushes": 9, + "total_cycles": 99, + "total_packets": 50000, + "acquisition_duration_s": 28.91 + } """ - pass - # num_frames: int = Field(..., alias="Num Frames") + model_config = ConfigDict(extra="allow") + + scan_name: Optional[str] = None + total_flushes: Optional[int] = None + total_cycles: Optional[int] = None + total_packets: Optional[int] = None + acquisition_duration_s: Optional[float] = None + + +# ADDED: operator output types — clean separation from ZMQ input types + +class XPSResultStart(Start, XPSMessage): + """ + Published by XPSOperator when a new scan begins. + Downstream publishers use this to signal that a new acquisition has started. + """ + msg_type: str = Literal["result_start"] + scan_name: Optional[str] = None class XPSResult(Event, XPSMessage): @@ -147,5 +214,8 @@ class XPSResult(Event, XPSMessage): class XPSResultStop(Stop, XPSMessage): + """ + Published by XPSOperator when processing ends. + """ msg_type: str = Literal["result_stop"] - function_timings: DataFrameModel + function_timings: Optional[DataFrameModel] = None diff --git a/src/tr_ap_xps/timepix.py b/src/tr_ap_xps/timepix.py index df9eec2..2c967c5 100644 --- a/src/tr_ap_xps/timepix.py +++ b/src/tr_ap_xps/timepix.py @@ -7,28 +7,26 @@ import zmq.asyncio from arroyopy.zmq import ZMQListener -from .config import settings from .schemas import NumpyArrayModel, XPSImageInfo, XPSRawEvent, XPSStart, XPSStop from arroyopy.operator import Operator -app_settings = settings.xps_operator -logger = logging.getLogger(__name__) +logger = logging.getLogger(__name__) -def setup_zmq(): + +# removed app_settings; accepts zmq_connection_address as parameters instead +def setup_zmq( + zmq_connection_address: str = "tcp://localhost:5657", +): ctx = zmq.asyncio.Context() - tpx_zmq_socket = ctx.socket(zmq.SUB) - tpx_zmq_socket.setsockopt(zmq.RCVHWM, 100000) - logger.info( - f"binding to: {app_settings.tpx_zmq_listener.zmq_pub_address}:{app_settings.tpx_zmq_listener.zmq_pub_port}" - ) - tpx_zmq_socket.connect( - f"{app_settings.tpx_zmq_listener.zmq_pub_address}:{app_settings.tpx_zmq_listener.zmq_pub_port}" - ) - tpx_zmq_socket.setsockopt(zmq.SUBSCRIBE, b"") - return tpx_zmq_socket + lv_zmq_socket = ctx.socket(zmq.SUB) + lv_zmq_socket.setsockopt(zmq.RCVHWM, 100000) + logger.info(f"connecting to: {zmq_connection_address}") + lv_zmq_socket.connect(zmq_connection_address) + lv_zmq_socket.setsockopt(zmq.SUBSCRIBE, b"") + return lv_zmq_socket class XPSTimepixZMQListener(ZMQListener): stop_signal = False @@ -116,51 +114,20 @@ async def start(self): @staticmethod def _build_start(metadata: dict) -> XPSStart: - """Build XPSStart from metadata dict. + """Build XPSStart from Timepix metadata. - Note: This is a basic conversion. splash_timepix start messages have different - fields than XPSStart expects, so this may need adjustment based on your needs. + CHANGED: XPSStart fields are all Optional now, so pass metadata directly. + No need to fabricate fake LabVIEW fields (Rectangle, F_Reset, etc.). """ - # Create Rectangle using a dict with alias keys - detector_x = metadata.get("detector_size_x", 256) - detector_y = metadata.get("detector_size_y", 256) - rectangle_dict = { - "Left": 0, - "Top": 0, - "Right": detector_x, - "Bottom": detector_y, - "Rotation": 0, - } - - # Create XPSStart using a dict with alias keys (as Pydantic expects) - # This matches the LabVIEW JSON format - start_dict = { - "msg_type": "start", - "scan_name": metadata.get("scan_name", "unknown"), - "Binding Energy": 0.0, # Not applicable for TimePix - "F_Trigger": 0, - "F_Un-Trigger": 0, - "F_Dead": 0, - "F_Reset": 0, - "CCD_nx": metadata.get("detector_size_x", 256), - "CCD_ny": metadata.get("detector_size_y", 256), - "Pass Energy": 0.0, - "Center Energy": 0.0, - "Offset Energy": 0.0, - "Lens Mode": "TimePix3", - "Rectangle": rectangle_dict, # Pass dict, not object - "dt": metadata.get("t_delta_ns", 10.0) / 1e9, # Convert ns to seconds - "Photon Energy": 0.0, - "File Ver": "1.0.0", - "data_type": "uint32", - } - return XPSStart(**start_dict) + return XPSStart(**metadata) @staticmethod def _build_stop(metadata: dict) -> XPSStop: - """Build XPSStop from metadata dict.""" - # XPSStop doesn't have required fields, so we can create an empty one - return XPSStop() + """Build XPSStop from Timepix metadata. + + CHANGED: XPSStop accepts extra fields now, so pass metadata directly. + """ + return XPSStop(**metadata) @staticmethod def _build_event( @@ -171,11 +138,17 @@ def _build_event( dtype = metadata["dtype"] image_info = XPSImageInfo( - frame_number=0, width=shape[0], height=shape[1], data_type=dtype + frame_number=metadata.get("flush_number", 0), + width=shape[0], + height=shape[1], + data_type=dtype, + timestamp=metadata.get("timestamp"), + cycles_in_flush=metadata.get("cycles_in_flush"), + total_cycles=metadata.get("total_cycles"), ) array_received = np.frombuffer(image, dtype=dtype).reshape(shape) - image_info.frame_number = metadata.get("flush_number", 0) + return XPSRawEvent( image=NumpyArrayModel(array=array_received), image_info=image_info ) @@ -208,3 +181,11 @@ async def process(self, message): import asyncio asyncio.run(listener.start()) + +# factory function for YAML instantiation +def xps_timepix_listener_factory( + operator, + zmq_connection_address: str = "tcp://localhost:5657", +) -> XPSTimepixZMQListener: + socket = setup_zmq(zmq_connection_address) + return XPSTimepixZMQListener(operator=operator, zmq_socket=socket) \ No newline at end of file diff --git a/src/tr_ap_xps/websockets.py b/src/tr_ap_xps/websockets.py index 8c14673..daba062 100644 --- a/src/tr_ap_xps/websockets.py +++ b/src/tr_ap_xps/websockets.py @@ -11,7 +11,7 @@ from arroyopy.publisher import Publisher -from .schemas import XPSResult, XPSResultStop, XPSStart +from .schemas import XPSResult, XPSResultStart, XPSResultStop logger = logging.getLogger(__name__) @@ -43,9 +43,9 @@ async def start( logger.info(f"Websocket server started at ws://{parsed_url.hostname}:{parsed_url.port}") await server.wait_closed() - async def publish(self, message: XPSResult) -> None: + async def publish(self, message) -> None: if self.connected_clients: # Only send if there are clients connected - asyncio.gather( + await asyncio.gather( *(self.publish_ws(client, message) for client in self.connected_clients) ) @@ -53,30 +53,38 @@ async def publish_ws( self, # client: websockets.client.ClientConnection, client, - message: Union[XPSResult | XPSStart | XPSResultStop], + message: Union[XPSResultStart, XPSResult, XPSResultStop], ) -> None: - if isinstance(message, XPSResultStop): - self.current_start_message = None - return + # CHANGED: now handles all three result types cleanly - if isinstance(message, XPSStart): + if isinstance(message, XPSResultStart): self.current_start_message = message - await client.send(json.dumps(message.model_dump())) + await client.send(json.dumps({ + "msg_type": "start", + "scan_name": message.scan_name, + })) + return + + elif isinstance(message, XPSResultStop): + self.current_start_message = None + await client.send(json.dumps({"msg_type": "stop"})) return - # send basic info - await client.send( - json.dumps( - { - # "result_info": message.result_info, - "frame_number": message.frame_number, - } + elif isinstance(message, XPSResult): + # send basic info + await client.send( + json.dumps( + { + # "result_info": message.result_info, + "msg_type": "event", + "frame_number": message.frame_number, + } + ) ) - ) - # send image data separately to client memory issues - image_bundle = await asyncio.to_thread(pack_images, message) - logger.info(f"Sending image bundle to client of size {len(image_bundle)}") - await client.send(image_bundle) + # send image data separately to avoid client memory issues + image_bundle = await asyncio.to_thread(pack_images, message) + logger.info(f"Sending image bundle to client of size {len(image_bundle)}") + await client.send(image_bundle) async def websocket_handler(self, websocket): logger.info(f"New connection from {websocket.remote_address}") @@ -102,10 +110,10 @@ def convert_to_uint8(image: np.ndarray) -> bytes: # scaled = (image - image.min()) / (image.max() - image.min()) * 255 # return scaled.astype(np.uint8).tobytes() - + if np.allclose(image, 0): return image.astype(np.uint8).tobytes() - + image_normalized = (image - image.min()) / (image.max() - image.min()) # Apply logarithmic stretch @@ -142,8 +150,16 @@ def pack_images(message: XPSResult) -> bytes: "raw": convert_to_uint8(message.integrated_frames.array), "width": message.shot_mean.array.shape[0], "height": message.shot_mean.array.shape[1], - "fitted": json.dumps(peaks_output(message.detected_peaks.df)), + "fitted": json.dumps(peaks_output(message.detected_peaks.df)) if message.detected_peaks else None, "shot_num": message.shot_num, - "shot_recent": convert_to_uint8(message.shot_recent.array), + "shot_recent": convert_to_uint8(message.shot_recent.array) if message.shot_recent else None, + "shot_mean": convert_to_uint8(message.shot_mean.array), } ) + + +# ADDED: factory function for YAML instantiation +def xps_ws_publisher_factory( + ws_url: str = "ws://localhost:8001/xps_operator", +) -> XPSWSResultPublisher: + return XPSWSResultPublisher(ws_url=ws_url) \ No newline at end of file