Skip to content

Commit

Permalink
Add QdAcquisitionConfig
Browse files Browse the repository at this point in the history
This is a combination of the acquisition header and the first frame
header, meaning we can accurately get the detector shape

Still exporting the `QdAcquisitionHeader` to Python, which can be used
to parse from raw bytes
  • Loading branch information
sk1p committed Sep 16, 2024
1 parent bb449bf commit b1c275c
Show file tree
Hide file tree
Showing 3 changed files with 79 additions and 30 deletions.
38 changes: 21 additions & 17 deletions libertem_qd_mpx/src/background_thread.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,21 +10,20 @@ use std::{
use common::{
background_thread::{BackgroundThread, BackgroundThreadSpawnError, ControlMsg, ReceiverMsg},
frame_stack::{FrameMeta, FrameStackForWriting, FrameStackWriteError, WriteGuard},
generic_connection::AcquisitionConfig,
tcp::{self, ReadExactError},
utils::{num_from_byte_slice, NumParseError},
};
use ipc_test::{slab::ShmError, SharedSlabAllocator};
use log::{debug, error, info, trace, warn};

use crate::base_types::{
AcqHeaderParseError, FrameMetaParseError, QdAcquisitionHeader, QdDetectorConnConfig,
QdFrameMeta, RecoveryStrategy, PREFIX_SIZE,
AcqHeaderParseError, FrameMetaParseError, QdAcquisitionConfig, QdAcquisitionHeader,
QdDetectorConnConfig, QdFrameMeta, RecoveryStrategy, PREFIX_SIZE,
};

type QdControlMsg = ControlMsg<()>;

type QdReceiverMsg = ReceiverMsg<QdFrameMeta, QdAcquisitionHeader>;
type QdReceiverMsg = ReceiverMsg<QdFrameMeta, QdAcquisitionConfig>;

#[derive(Debug, thiserror::Error)]
pub enum AcquisitionError {
Expand Down Expand Up @@ -407,7 +406,10 @@ fn acquisition(
let mut last_control_check = Instant::now();

from_thread_s.send(ReceiverMsg::AcquisitionStart {
pending_acquisition: acquisition_header.clone(),
pending_acquisition: QdAcquisitionConfig::new(
acquisition_header.clone(),
first_frame_meta.clone(),
),
})?;

let mut frame_stack = make_frame_stack(shm, config, first_frame_meta, to_thread_r)?;
Expand Down Expand Up @@ -714,7 +716,7 @@ impl QdBackgroundThread {

impl BackgroundThread for QdBackgroundThread {
type FrameMetaImpl = QdFrameMeta;
type AcquisitionConfigImpl = QdAcquisitionHeader;
type AcquisitionConfigImpl = QdAcquisitionConfig;
type ExtraControl = ();

fn channel_to_thread(
Expand Down Expand Up @@ -760,7 +762,9 @@ mod test {
time::timeout,
};

use crate::base_types::{QdAcquisitionHeader, QdDetectorConnConfig, RecoveryStrategy};
use crate::base_types::{
QdAcquisitionConfig, QdAcquisitionHeader, QdDetectorConnConfig, RecoveryStrategy,
};

use super::QdBackgroundThread;

Expand Down Expand Up @@ -934,7 +938,7 @@ End

// we let the `GenericConnection` drive the background thread:
let bg = QdBackgroundThread::spawn(config, &shm).unwrap();
let mut conn: GenericConnection<QdBackgroundThread, QdAcquisitionHeader> =
let mut conn: GenericConnection<QdBackgroundThread, QdAcquisitionConfig> =
GenericConnection::new(bg, &shm).unwrap();
conn.start_passive(
|| Ok::<(), ConnectionError>(()),
Expand Down Expand Up @@ -1006,7 +1010,7 @@ End

// we let the `GenericConnection` drive the background thread:
let bg = QdBackgroundThread::spawn(config, &shm).unwrap();
let mut conn: GenericConnection<QdBackgroundThread, QdAcquisitionHeader> =
let mut conn: GenericConnection<QdBackgroundThread, QdAcquisitionConfig> =
GenericConnection::new(bg, &shm).unwrap();
conn.start_passive(
|| Ok::<(), ConnectionError>(()),
Expand Down Expand Up @@ -1093,7 +1097,7 @@ End

// we let the `GenericConnection` drive the background thread:
let bg = QdBackgroundThread::spawn(config, &shm).unwrap();
let mut conn: GenericConnection<QdBackgroundThread, QdAcquisitionHeader> =
let mut conn: GenericConnection<QdBackgroundThread, QdAcquisitionConfig> =
GenericConnection::new(bg, &shm).unwrap();
conn.start_passive(
|| Ok::<(), ConnectionError>(()),
Expand Down Expand Up @@ -1208,7 +1212,7 @@ End

// we let the `GenericConnection` drive the background thread:
let bg = QdBackgroundThread::spawn(config, &shm).unwrap();
let mut conn: GenericConnection<QdBackgroundThread, QdAcquisitionHeader> =
let mut conn: GenericConnection<QdBackgroundThread, QdAcquisitionConfig> =
GenericConnection::new(bg, &shm).unwrap();
conn.start_passive(
|| Ok::<(), ConnectionError>(()),
Expand Down Expand Up @@ -1313,7 +1317,7 @@ End

// we let the `GenericConnection` drive the background thread:
let bg = QdBackgroundThread::spawn(config, &shm).unwrap();
let mut conn: GenericConnection<QdBackgroundThread, QdAcquisitionHeader> =
let mut conn: GenericConnection<QdBackgroundThread, QdAcquisitionConfig> =
GenericConnection::new(bg, &shm).unwrap();
conn.start_passive(
|| Ok::<(), ConnectionError>(()),
Expand Down Expand Up @@ -1423,7 +1427,7 @@ End

// we let the `GenericConnection` drive the background thread:
let bg = QdBackgroundThread::spawn(config, &shm).unwrap();
let mut conn: GenericConnection<QdBackgroundThread, QdAcquisitionHeader> =
let mut conn: GenericConnection<QdBackgroundThread, QdAcquisitionConfig> =
GenericConnection::new(bg, &shm).unwrap();
conn.start_passive(
|| Ok::<(), ConnectionError>(()),
Expand Down Expand Up @@ -1535,7 +1539,7 @@ End

// we let the `GenericConnection` drive the background thread:
let bg = QdBackgroundThread::spawn(config, &shm).unwrap();
let mut conn: GenericConnection<QdBackgroundThread, QdAcquisitionHeader> =
let mut conn: GenericConnection<QdBackgroundThread, QdAcquisitionConfig> =
GenericConnection::new(bg, &shm).unwrap();
conn.start_passive(
|| Ok::<(), ConnectionError>(()),
Expand Down Expand Up @@ -1639,7 +1643,7 @@ End

// we let the `GenericConnection` drive the background thread:
let bg = QdBackgroundThread::spawn(config, &shm).unwrap();
let mut conn: GenericConnection<QdBackgroundThread, QdAcquisitionHeader> =
let mut conn: GenericConnection<QdBackgroundThread, QdAcquisitionConfig> =
GenericConnection::new(bg, &shm).unwrap();
conn.start_passive(
|| Ok::<(), ConnectionError>(()),
Expand Down Expand Up @@ -1774,7 +1778,7 @@ End

// we let the `GenericConnection` drive the background thread:
let bg = QdBackgroundThread::spawn(config, &shm).unwrap();
let mut conn: GenericConnection<QdBackgroundThread, QdAcquisitionHeader> =
let mut conn: GenericConnection<QdBackgroundThread, QdAcquisitionConfig> =
GenericConnection::new(bg, &shm).unwrap();
conn.start_passive(
|| Ok::<(), ConnectionError>(()),
Expand Down Expand Up @@ -1921,7 +1925,7 @@ End

// we let the `GenericConnection` drive the background thread:
let bg = QdBackgroundThread::spawn(config, &shm).unwrap();
let mut conn: GenericConnection<QdBackgroundThread, QdAcquisitionHeader> =
let mut conn: GenericConnection<QdBackgroundThread, QdAcquisitionConfig> =
GenericConnection::new(bg, &shm).unwrap();
conn.start_passive(
|| Ok::<(), ConnectionError>(()),
Expand Down
60 changes: 52 additions & 8 deletions libertem_qd_mpx/src/base_types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -518,12 +518,6 @@ pub struct QdAcquisitionHeader {
raw_kv: HashMap<String, String>,
}

impl AcquisitionConfig for QdAcquisitionHeader {
fn num_frames(&self) -> usize {
self.frames_in_acquisition
}
}

#[pymethods]
impl QdAcquisitionHeader {
#[classmethod]
Expand All @@ -539,8 +533,14 @@ impl QdAcquisitionHeader {
fn frames_per_trigger(&self) -> usize {
self.frames_per_trigger
}
}

fn nav_shape(&self) -> Option<(usize, usize)> {
impl QdAcquisitionHeader {
pub fn num_frames(&self) -> usize {
self.frames_in_acquisition
}

pub fn nav_shape(&self) -> Option<(usize, usize)> {
if let (Some(scan_x), Some(scan_y)) = (self.scan_x, self.scan_y) {
return Some((scan_y, scan_x));
}
Expand Down Expand Up @@ -742,9 +742,53 @@ impl DetectorConnectionConfig for QdDetectorConnConfig {
}
}

#[pyclass]
#[derive(Debug, Clone)]
pub struct QdAcquisitionConfig {
acq_header: QdAcquisitionHeader,
frame_header: QdFrameMeta,
}

impl QdAcquisitionConfig {
pub fn new(acq_header: QdAcquisitionHeader, frame_header: QdFrameMeta) -> Self {
Self {
acq_header,
frame_header,
}
}
}

impl AcquisitionConfig for QdAcquisitionConfig {
fn num_frames(&self) -> usize {
self.acq_header.frames_in_acquisition
}
}

#[pymethods]
impl QdAcquisitionConfig {
fn frames_in_acquisition(&self) -> usize {
self.acq_header.frames_in_acquisition
}

fn frames_per_trigger(&self) -> usize {
self.acq_header.frames_per_trigger
}

fn nav_shape(&self) -> Option<(usize, usize)> {
self.acq_header.nav_shape()
}

fn detector_shape(&self) -> (u32, u32) {
(
self.frame_header.height_in_pixels,
self.frame_header.width_in_pixels,
)
}
}

#[cfg(test)]
mod test {
use common::{frame_stack::FrameMeta, generic_connection::AcquisitionConfig};
use common::frame_stack::FrameMeta;

use crate::base_types::{DType, Layout};

Expand Down
11 changes: 6 additions & 5 deletions libertem_qd_mpx/src/main_py.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ use pyo3::{

use common::{impl_py_cam_client, impl_py_connection};

use crate::base_types::{QdDetectorConnConfig, QdFrameMeta, RecoveryStrategy};
use crate::base_types::{QdAcquisitionConfig, QdDetectorConnConfig, QdFrameMeta, RecoveryStrategy};
use crate::decoder::QdDecoder;
use crate::{background_thread::QdBackgroundThread, base_types::QdAcquisitionHeader};

Expand All @@ -19,6 +19,7 @@ fn libertem_qd_mpx(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<QdFrameStack>()?;
m.add_class::<CamClient>()?;
m.add_class::<QdAcquisitionHeader>()?;
m.add_class::<QdAcquisitionConfig>()?;

let env = env_logger::Env::default()
.filter_or("LIBERTEM_QD_LOG_LEVEL", "error")
Expand All @@ -35,7 +36,7 @@ impl_py_connection!(
QdFrameStack,
QdFrameMeta,
QdBackgroundThread,
QdAcquisitionHeader,
QdAcquisitionConfig,
libertem_qd_mpx
);

Expand Down Expand Up @@ -90,13 +91,13 @@ impl QdConnection {
);

let shm =
GenericConnection::<QdBackgroundThread, QdAcquisitionHeader>::shm_from_config(&config)
GenericConnection::<QdBackgroundThread, QdAcquisitionConfig>::shm_from_config(&config)
.map_err(|e| PyConnectionError::new_err(e.to_string()))?;

let bg_thread = QdBackgroundThread::spawn(&config, &shm)
.map_err(|e| PyConnectionError::new_err(e.to_string()))?;
let generic_conn =
GenericConnection::<QdBackgroundThread, QdAcquisitionHeader>::new(bg_thread, &shm)
GenericConnection::<QdBackgroundThread, QdAcquisitionConfig>::new(bg_thread, &shm)
.map_err(|e| PyConnectionError::new_err(e.to_string()))?;

let conn = _PyQdConnection::new(shm, generic_conn);
Expand All @@ -108,7 +109,7 @@ impl QdConnection {
&mut self,
timeout: Option<f32>,
py: Python<'_>,
) -> PyResult<Option<QdAcquisitionHeader>> {
) -> PyResult<Option<QdAcquisitionConfig>> {
self.conn.wait_for_arm(timeout, py)
}

Expand Down

0 comments on commit b1c275c

Please sign in to comment.