From 756e12eac20712533eb3a63810d913d06f937013 Mon Sep 17 00:00:00 2001 From: Phil Blowey Date: Thu, 22 Oct 2026 17:05:45 +0100 Subject: [PATCH 01/14] Initial cli tool commit --- setup.py | 1 + src/dlstbx/cli/dc_sim.py | 441 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 442 insertions(+) create mode 100644 src/dlstbx/cli/dc_sim.py diff --git a/setup.py b/setup.py index 9a512d7f7..9ccfec8a8 100644 --- a/setup.py +++ b/setup.py @@ -13,6 +13,7 @@ "dials.swirly_eyes=dlstbx.cli.swirly_eyes:run", "dlstbx.align_crystal=dlstbx.cli.align_crystal:run", "dlstbx.archive_ancient_visits=dlstbx.cli.archive_ancient_visits:run", + "dlstbx.dc_sim=dlstbx.cli.dc_sim:run", "dlstbx.dc_sim_verify=dlstbx.cli.dc_sim_verify:run", "dlstbx.ep_predict_phase=dlstbx.cli.ep_predict_phase:run", "dlstbx.ep_predict_results=dlstbx.cli.ep_predict_results:runmain", diff --git a/src/dlstbx/cli/dc_sim.py b/src/dlstbx/cli/dc_sim.py new file mode 100644 index 000000000..f69c35d97 --- /dev/null +++ b/src/dlstbx/cli/dc_sim.py @@ -0,0 +1,441 @@ +from __future__ import annotations + +import argparse +import logging +import os +import re +import shutil +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path + +import ispyb.sqlalchemy +import sqlalchemy +import sqlalchemy.orm +from ispyb.sqlalchemy import ( + BLSample, + BLSession, + DataCollection, + DataCollectionGroup, + GridInfo, + Proposal, +) +from sqlalchemy.orm import Load + +from dlstbx.util import INDUSTRIAL_CODES + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class DataCollectionRecords: + """Related ISPyB records needed to simulate one data collection.""" + + data_collection: DataCollection + data_collection_group: DataCollectionGroup + bl_session: BLSession + proposal: Proposal + grid_info: GridInfo | None + bl_sample: BLSample | None + + +def big_query(session: sqlalchemy.orm.Session, dcid: int) -> DataCollectionRecords: + """Fetch a data collection and its related records in one database query. + + GridInfo is optional and may be associated directly with the data collection + or, for legacy records, with its data collection group. If both forms exist, + the direct association is preferred. BLSample is optional; the + DataCollection value is preferred, falling back to the group for legacy + records. + """ + + result = ( + session.query( + DataCollection, + DataCollectionGroup, + BLSession, + Proposal, + GridInfo, + BLSample, + ) + .join( + DataCollectionGroup, + DataCollection.dataCollectionGroupId + == DataCollectionGroup.dataCollectionGroupId, + ) + .join(BLSession, DataCollectionGroup.sessionId == BLSession.sessionId) + .join(Proposal, BLSession.proposalId == Proposal.proposalId) + .outerjoin( + GridInfo, GridInfo.dataCollectionId == DataCollection.dataCollectionId + ) + .outerjoin( + BLSample, + BLSample.blSampleId + == sqlalchemy.func.coalesce( + DataCollection.BLSAMPLEID, DataCollectionGroup.blSampleId + ), + ) + .filter(DataCollection.dataCollectionId == dcid) + .one_or_none() + ) + if result is None: + logger.error( + f"No data collection or related ISPyB records found for dcid={dcid}" + ) + raise ValueError( + f"Could not find a data collection and its related records for dcid {dcid}" + ) + + return DataCollectionRecords(*result) + + +def retrieve_sessionid(db_session: sqlalchemy.orm.Session, visit: str) -> int: + query = ( + db_session.query(BLSession, Proposal) + .options( + Load(BLSession).load_only( + BLSession.sessionId, BLSession.visit_number, BLSession.proposalId + ), + Load(Proposal).load_only( + Proposal.proposalId, Proposal.proposalCode, Proposal.proposalNumber + ), + ) + .join( + Proposal, + Proposal.proposalId == BLSession.proposalId, + ) + .filter( + sqlalchemy.func.concat( + Proposal.proposalCode, + Proposal.proposalNumber, + "-", + BLSession.visit_number, + ) + == visit + ) + ) + query_results = query.first() + if query_results is None: + raise ValueError(f"Query to obtain sessionid failed for {visit}") + if query_results[0].sessionId is None: + raise ValueError(f"Could not find sessionid for visit {visit}") + return query_results[0].sessionId + + +def get_data_collection_group_id( + db_session: sqlalchemy.orm.Session, dcgid: int +) -> sqlalchemy.orm.Query: + return ( + db_session.query(DataCollectionGroup) + .filter(DataCollectionGroup.dataCollectionGroupId == dcgid) + .first() + ) + + +def copy_data_collection_group( + source_dcg: DataCollectionGroup, + new_session_id: int, + new_start_time: datetime, + new_end_time: datetime, +) -> DataCollectionGroup: + """ + Create a copy of a DataCollectionGroup row with different sessionId and times. + + Args: + source_dcg: The DataCollectionGroup to copy + new_session_id: The new sessionId for the copy + new_start_time: The new startTime for the copy + new_end_time: The new endTime for the copy + + Returns: + A new DataCollectionGroup object with all attributes copied except + the ones specified (which are replaced with the new values) + """ + # Get all columns from the source object + new_dcg = DataCollectionGroup() + + # Copy all attributes from source to new, except the primary key and fields to update + for column in DataCollectionGroup.__table__.columns: + attr_name = column.name + + # Skip the primary key + if attr_name == "dataCollectionGroupId": + continue + + # Use new values for these attributes + if attr_name == "sessionId": + setattr(new_dcg, attr_name, new_session_id) + elif attr_name == "startTime": + setattr(new_dcg, attr_name, new_start_time) + elif attr_name == "endTime": + setattr(new_dcg, attr_name, new_end_time) + else: + # Copy from source + setattr(new_dcg, attr_name, getattr(source_dcg, attr_name, None)) + + return new_dcg + + +def copy_data_collection( + source_dc: DataCollection, + new_session_id: int, + new_data_collection_group_id: int, + **kwargs, +) -> DataCollection: + """ + Create a copy of a DataCollection row with different sessionId and/or dataCollectionGroupId. + + Args: + source_dc: The DataCollection to copy + new_session_id: The new SESSIONID for the copy + new_data_collection_group_id: The new dataCollectionGroupId for the copy + **kwargs: Additional attributes to override (e.g., startTime, endTime, comments) + + Returns: + A new DataCollection object with all attributes copied except + the ones specified (which are replaced with the new values) + """ + # Create a new DataCollection object + new_dc = DataCollection() + + # Copy all attributes from source to new, except the primary key and fields to update + for column in DataCollection.__table__.columns: + attr_name = column.name + + # Skip the primary key + if attr_name == "dataCollectionId": + continue + + # Use new values for these attributes + if attr_name == "SESSIONID": + setattr(new_dc, attr_name, new_session_id) + elif attr_name == "dataCollectionGroupId": + setattr(new_dc, attr_name, new_data_collection_group_id) + elif attr_name in kwargs: + # Use provided override values + setattr(new_dc, attr_name, kwargs[attr_name]) + else: + # Copy from source + setattr(new_dc, attr_name, getattr(source_dc, attr_name, None)) + + return new_dc + + +def simulate_data_collection(dcid: int) -> None: + """Look up and copy the records needed to simulate data collection.""" + start_time = datetime.now() + dflt_proposal = "nt44198" + logger.info(f"Starting data collection simulation for dcid={dcid}") + + sessionmaker = sqlalchemy.orm.sessionmaker( + bind=sqlalchemy.create_engine( + ispyb.sqlalchemy.url(), connect_args={"use_pure": True} + ) + ) + with sessionmaker() as session: + logger.info("Loading source data collection and related ISPyB records") + dc_records = big_query(session, dcid) + + if dc_records.proposal.proposalCode in INDUSTRIAL_CODES: + logger.error( + f"Source data collection {dcid} is industrial; simulation is not allowed" + ) + raise SystemExit(1) + + beamline = dc_records.bl_session.beamLineName + dest_visit_dir, dest_visit = get_beamline_visit_and_dir_for_proposal( + dflt_proposal, beamline + ) + + dc = dc_records.data_collection + dcg = dc_records.data_collection_group + logger.info( + f"Source records loaded: dcid={dcid}, beamline={beamline}, " + f"destination visit={dest_visit}" + ) + + now = datetime.now() + source_image_dir = Path(dc.imageDirectory) + random_uuid = os.urandom(4).hex() # Generate a random UUID for uniqueness + dest_dir = ( + dest_visit_dir + / "tmp" + / f"{now:%Y-%m-%d}" + / f"{now:%H}-{now:%M}-{now:%S}-{random_uuid}" + ) + dest_dir.mkdir(parents=True, exist_ok=True) + logger.info(f"Created image destination directory {dest_dir}") + # Get the session id for the destination visit + logger.info(f"Finding ISPyB session for destination visit {dest_visit}") + new_bl_session_id = retrieve_sessionid(session, dest_visit) + # Create a copy with new sessionId and times + new_dcg = copy_data_collection_group( + dcg, + new_session_id=new_bl_session_id, + new_start_time=start_time, + new_end_time=datetime.now(), + ) + + # Add and commit the new DataCollectionGroup + session.add(new_dcg) + session.commit() + logger.info( + f"Created destination DataCollectionGroup {new_dcg.dataCollectionGroupId}" + ) + + # Copy the image files to the new destination directory + image_suffix = dc.imageSuffix + if image_suffix == "cbf": + file_template = dc.fileTemplate + match = re.search(r"(#+)(?=\.cbf$)", file_template) + if not match: + raise ValueError( + f"File template {file_template} does not contain a valid '#' pattern for image numbering." + ) + n_digits = len(match.group(1)) + start_image_number = dc.startImageNumber + number_of_images = dc.numberOfImages + copied_images = 0 + missing_images = 0 + logger.info( + f"Copying {number_of_images} CBF images from {source_image_dir} " + f"to {dest_dir}" + ) + for image_number in range( + start_image_number, start_image_number + number_of_images + ): + image_filename = ( + file_template[: match.start()] + + str(image_number).zfill(n_digits) + + file_template[match.end() :] + ) + src_image_path = source_image_dir / image_filename + dest_image_path = dest_dir / image_filename + if src_image_path.exists(): + shutil.copy2(src_image_path, dest_image_path) + copied_images += 1 + else: + logger.warning( + f"Source image {src_image_path} does not exist; skipping copy" + ) + missing_images += 1 + logger.info( + f"Finished copying CBF images: copied={copied_images}, " + f"missing={missing_images}" + ) + + elif image_suffix in ["h5", "nxs"]: + file_prefix = f"{dc.imagePrefix}_{dc.dataCollectionNumber}" + files = { + path + for pattern in (f"{file_prefix}_*.h5", f"{file_prefix}.nxs") + for path in source_image_dir.glob(pattern) + } + logger.info( + f"Copying {len(files)} HDF5/Nexus file(s) from {source_image_dir} " + f"to {dest_dir}" + ) + for file in files: + shutil.copy2(file, dest_dir) + logger.info("Finished copying image data files") + else: + logger.error( + f"Unsupported image suffix {image_suffix!r} for source dcid={dcid}" + ) + raise ValueError( + f"Invalid image suffix {image_suffix} for source data collection" + ) + + # Create new DataCollection + new_dc = copy_data_collection( + dc, + new_session_id=new_bl_session_id, + new_data_collection_group_id=new_dcg.dataCollectionGroupId, + startTime=start_time, + endTime=datetime.now(), + comments=f"Simulated data collection copied from dcid {dcid}", + imageDirectory=str(dest_dir) + "/", + ) + + # Create destination directory for JPEG snapshots + jpgs_dest_dir = ( + dest_visit_dir + / "jpegs" + / "tmp" + / f"{now:%Y-%m-%d}" + / f"{now:%H}-{now:%M}-{now:%S}-{random_uuid}" + ) + jpgs_dest_dir.mkdir(parents=True, exist_ok=True) + logger.info(f"Created snapshot destination directory {jpgs_dest_dir}") + + # Mapping of snapshot attributes to their names + snapshot_attrs = [ + "xtalSnapshotFullPath1", + "xtalSnapshotFullPath2", + "xtalSnapshotFullPath3", + "xtalSnapshotFullPath4", + ] + + # Copy and update snapshot paths + copied_snapshots = 0 + missing_snapshots = 0 + for attr_name in snapshot_attrs: + src_snapshot_path = getattr(dc, attr_name, None) + if src_snapshot_path is not None: + src_snapshot = Path(src_snapshot_path) + if src_snapshot.exists(): + # Create a unique filename for the snapshot + dest_snapshot_path = jpgs_dest_dir / src_snapshot.name + shutil.copy2(src_snapshot, dest_snapshot_path) + # Update the new DataCollection with the new snapshot path + setattr(new_dc, attr_name, str(dest_snapshot_path)) + copied_snapshots += 1 + else: + logger.warning( + f"Crystal snapshot {src_snapshot} does not exist; skipping copy" + ) + missing_snapshots += 1 + + session.add(new_dc) + session.commit() + logger.info( + f"Created destination DataCollection {new_dc.dataCollectionId}; " + f"copied snapshots={copied_snapshots}, missing snapshots={missing_snapshots}" + ) + logger.info( + f"Data collection simulation for dcid={dcid} completed in " + f"{(datetime.now() - start_time).total_seconds():.1f} seconds" + ) + + +def get_beamline_visit_and_dir_for_proposal( + proposal: str, beamline: str +) -> tuple[Path, str]: + data_root = Path("/dls", beamline, "data", str(datetime.now().year)) + logger.info(f"Searching for proposal {proposal} under {data_root}") + for cm_dir in data_root.iterdir(): + if cm_dir.name.startswith(proposal): + dest_visit_dir = cm_dir + dest_visit = cm_dir.name + break + else: + raise ValueError( + f"Could not determine destination directory for proposal {proposal}" + ) + logger.info(f"Using destination visit directory {dest_visit_dir}") + return dest_visit_dir, dest_visit + + +def run(args: list[str] | None = None) -> None: + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + parser = argparse.ArgumentParser( + description="Simulate data collection. Takes a dcid as input and copies data and database records to simulate a data collection." + ) + parser.add_argument("dcid", type=int, help="Data collection ID") + parsed_args = parser.parse_args(args) + + simulate_data_collection(parsed_args.dcid) + + +if __name__ == "__main__": + run() From 2aafa197248a5bde71c35cfe97e675c761df3e12 Mon Sep 17 00:00:00 2001 From: Phil Blowey Date: Fri, 23 Oct 2026 12:47:56 +0100 Subject: [PATCH 02/14] Call mimas --- src/dlstbx/cli/dc_sim.py | 51 ++++++++++++++++++++++++++++++++++++---- 1 file changed, 47 insertions(+), 4 deletions(-) diff --git a/src/dlstbx/cli/dc_sim.py b/src/dlstbx/cli/dc_sim.py index f69c35d97..61da0e0fd 100644 --- a/src/dlstbx/cli/dc_sim.py +++ b/src/dlstbx/cli/dc_sim.py @@ -1,10 +1,12 @@ from __future__ import annotations import argparse +import getpass import logging import os import re import shutil +import socket from dataclasses import dataclass from datetime import datetime from pathlib import Path @@ -12,6 +14,9 @@ import ispyb.sqlalchemy import sqlalchemy import sqlalchemy.orm +import workflows.recipe +import workflows.transport +import zocalo.configuration from ispyb.sqlalchemy import ( BLSample, BLSession, @@ -221,7 +226,7 @@ def copy_data_collection( return new_dc -def simulate_data_collection(dcid: int) -> None: +def simulate_data_collection(dcid: int) -> int: """Look up and copy the records needed to simulate data collection.""" start_time = datetime.now() dflt_proposal = "nt44198" @@ -406,6 +411,7 @@ def simulate_data_collection(dcid: int) -> None: f"Data collection simulation for dcid={dcid} completed in " f"{(datetime.now() - start_time).total_seconds():.1f} seconds" ) + return new_dc.dataCollectionId def get_beamline_visit_and_dir_for_proposal( @@ -428,13 +434,50 @@ def get_beamline_visit_and_dir_for_proposal( def run(args: list[str] | None = None) -> None: logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + zc = zocalo.configuration.from_file() + zc.activate() parser = argparse.ArgumentParser( - description="Simulate data collection. Takes a dcid as input and copies data and database records to simulate a data collection." + description="Simulate data collection. Takes a dcid as input and copies data and database records to simulate a data collection, then triggers mimas.", + usage="dlstbx.dc_sim [options] dcid", ) parser.add_argument("dcid", type=int, help="Data collection ID") - parsed_args = parser.parse_args(args) + zc.add_command_line_options(parser) + workflows.transport.add_command_line_options(parser, transport_argument=True) - simulate_data_collection(parsed_args.dcid) + parsed_args = parser.parse_args(args) + new_dcid = simulate_data_collection(parsed_args.dcid) + + logger.info("Triggering Mimas") + headers = { + "zocalo.go.user": getpass.getuser(), + "zocalo.go.host": socket.gethostname(), + } + + transport = workflows.transport.lookup(parsed_args.transport)() + transport.connect() + transport.send( + "processing_recipe", + { + "recipes": ["mimas"], + "parameters": { + "ispyb_dcid": new_dcid, + "event": "start", + }, + }, + headers=headers, + ) + transport.send( + "processing_recipe", + { + "recipes": ["mimas"], + "parameters": { + "ispyb_dcid": new_dcid, + "event": "end", + }, + }, + headers=headers, + ) + transport.disconnect() if __name__ == "__main__": From 6d00b5980389d71a282e28e3f0b00ff2621dc2b1 Mon Sep 17 00:00:00 2001 From: Phil Blowey Date: Fri, 23 Oct 2026 16:15:26 +0100 Subject: [PATCH 03/14] Add grid_scan functionality --- src/dlstbx/cli/dc_sim.py | 63 +++++++++++++++++++++++++++++++--------- 1 file changed, 50 insertions(+), 13 deletions(-) diff --git a/src/dlstbx/cli/dc_sim.py b/src/dlstbx/cli/dc_sim.py index 61da0e0fd..c11e3c29c 100644 --- a/src/dlstbx/cli/dc_sim.py +++ b/src/dlstbx/cli/dc_sim.py @@ -127,16 +127,6 @@ def retrieve_sessionid(db_session: sqlalchemy.orm.Session, visit: str) -> int: return query_results[0].sessionId -def get_data_collection_group_id( - db_session: sqlalchemy.orm.Session, dcgid: int -) -> sqlalchemy.orm.Query: - return ( - db_session.query(DataCollectionGroup) - .filter(DataCollectionGroup.dataCollectionGroupId == dcgid) - .first() - ) - - def copy_data_collection_group( source_dcg: DataCollectionGroup, new_session_id: int, @@ -226,6 +216,47 @@ def copy_data_collection( return new_dc +def copy_grid_info( + source_gi: GridInfo, + new_dcid: int, + **kwargs, +) -> DataCollection: + """ + Create a copy of a DataCollection row with different sessionId and/or dataCollectionGroupId. + + Args: + source_dc: The DataCollection to copy + new_session_id: The new SESSIONID for the copy + new_data_collection_group_id: The new dataCollectionGroupId for the copy + **kwargs: Additional attributes to override (e.g., startTime, endTime, comments) + + Returns: + A new DataCollection object with all attributes copied except + the ones specified (which are replaced with the new values) + """ + # Create a new GridInfo object + new_gi = GridInfo() + + # Copy all attributes from source to new, except the primary key and fields to update + for column in GridInfo.__table__.columns: + attr_name = column.name + + # Skip the primary key + if attr_name == "gridInfoId": + continue + + # Use new values for these attributes + if attr_name == "dataCollectionId": + setattr(new_gi, attr_name, new_dcid) + elif attr_name in kwargs: + # Use provided override values + setattr(new_gi, attr_name, kwargs[attr_name]) + else: + # Copy from source + setattr(new_gi, attr_name, getattr(source_gi, attr_name, None)) + return new_gi + + def simulate_data_collection(dcid: int) -> int: """Look up and copy the records needed to simulate data collection.""" start_time = datetime.now() @@ -361,6 +392,14 @@ def simulate_data_collection(dcid: int) -> int: comments=f"Simulated data collection copied from dcid {dcid}", imageDirectory=str(dest_dir) + "/", ) + session.add(new_dc) + session.commit() + new_dcid = new_dc.dataCollectionId + + if dc_records.grid_info is not None: + new_gi = copy_grid_info(source_gi=dc_records.grid_info, new_dcid=new_dcid) + session.add(new_gi) + session.commit() # Create destination directory for JPEG snapshots jpgs_dest_dir = ( @@ -401,8 +440,6 @@ def simulate_data_collection(dcid: int) -> int: ) missing_snapshots += 1 - session.add(new_dc) - session.commit() logger.info( f"Created destination DataCollection {new_dc.dataCollectionId}; " f"copied snapshots={copied_snapshots}, missing snapshots={missing_snapshots}" @@ -411,7 +448,7 @@ def simulate_data_collection(dcid: int) -> int: f"Data collection simulation for dcid={dcid} completed in " f"{(datetime.now() - start_time).total_seconds():.1f} seconds" ) - return new_dc.dataCollectionId + return new_dcid def get_beamline_visit_and_dir_for_proposal( From 64592297a49d5ffade3ef115c950a982cfdf6eec Mon Sep 17 00:00:00 2001 From: Phil Blowey Date: Tue, 27 Oct 2026 13:50:43 +0000 Subject: [PATCH 04/14] Tidy up logging --- src/dlstbx/cli/dc_sim.py | 82 ++++++++++++++++++++++++---------------- 1 file changed, 50 insertions(+), 32 deletions(-) diff --git a/src/dlstbx/cli/dc_sim.py b/src/dlstbx/cli/dc_sim.py index c11e3c29c..64be5368b 100644 --- a/src/dlstbx/cli/dc_sim.py +++ b/src/dlstbx/cli/dc_sim.py @@ -14,7 +14,6 @@ import ispyb.sqlalchemy import sqlalchemy import sqlalchemy.orm -import workflows.recipe import workflows.transport import zocalo.configuration from ispyb.sqlalchemy import ( @@ -27,7 +26,8 @@ ) from sqlalchemy.orm import Load -from dlstbx.util import INDUSTRIAL_CODES +from dlstbx.util import INDUSTRIAL_CODES, DowngradeErrorsFilter +from dlstbx.util.colorstreamhandler import ColorStreamHandler logger = logging.getLogger(__name__) @@ -470,7 +470,23 @@ def get_beamline_visit_and_dir_for_proposal( def run(args: list[str] | None = None) -> None: - logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + # Enable logging to console + console = ColorStreamHandler() + console.setLevel(logging.INFO) + logging.getLogger("dlstbx").setLevel(logging.INFO) + for pika_logger in ( + "base_connection", + "blocking_connection", + "utils.connection_workflow", + "utils.io_services_utils", + ): + logging.getLogger(f"pika.adapters.{pika_logger}").addFilter( + DowngradeErrorsFilter() + ) + logging.getLogger("workflows").setLevel(logging.INFO) + logging.getLogger("zocalo").setLevel(logging.INFO) + logging.getLogger().setLevel(logging.WARNING) + logging.getLogger().addHandler(console) zc = zocalo.configuration.from_file() zc.activate() parser = argparse.ArgumentParser( @@ -478,43 +494,45 @@ def run(args: list[str] | None = None) -> None: usage="dlstbx.dc_sim [options] dcid", ) parser.add_argument("dcid", type=int, help="Data collection ID") + parser.add_argument("--suppress-mimas", action="store_true") zc.add_command_line_options(parser) workflows.transport.add_command_line_options(parser, transport_argument=True) parsed_args = parser.parse_args(args) new_dcid = simulate_data_collection(parsed_args.dcid) - logger.info("Triggering Mimas") - headers = { - "zocalo.go.user": getpass.getuser(), - "zocalo.go.host": socket.gethostname(), - } - - transport = workflows.transport.lookup(parsed_args.transport)() - transport.connect() - transport.send( - "processing_recipe", - { - "recipes": ["mimas"], - "parameters": { - "ispyb_dcid": new_dcid, - "event": "start", + if not parsed_args.suppress_mimas: + logger.info("Triggering Mimas") + headers = { + "zocalo.go.user": getpass.getuser(), + "zocalo.go.host": socket.gethostname(), + } + + transport = workflows.transport.lookup(parsed_args.transport)() + transport.connect() + transport.send( + "processing_recipe", + { + "recipes": ["mimas"], + "parameters": { + "ispyb_dcid": new_dcid, + "event": "start", + }, }, - }, - headers=headers, - ) - transport.send( - "processing_recipe", - { - "recipes": ["mimas"], - "parameters": { - "ispyb_dcid": new_dcid, - "event": "end", + headers=headers, + ) + transport.send( + "processing_recipe", + { + "recipes": ["mimas"], + "parameters": { + "ispyb_dcid": new_dcid, + "event": "end", + }, }, - }, - headers=headers, - ) - transport.disconnect() + headers=headers, + ) + transport.disconnect() if __name__ == "__main__": From 2395c18f019bc0d71d17230db34e2b56621644fc Mon Sep 17 00:00:00 2001 From: Phil Blowey Date: Tue, 27 Oct 2026 16:35:11 +0000 Subject: [PATCH 05/14] Add dcg functionality --- src/dlstbx/cli/dc_sim.py | 459 ++++++++++++++++++++++----------------- 1 file changed, 256 insertions(+), 203 deletions(-) diff --git a/src/dlstbx/cli/dc_sim.py b/src/dlstbx/cli/dc_sim.py index 64be5368b..dc26ee358 100644 --- a/src/dlstbx/cli/dc_sim.py +++ b/src/dlstbx/cli/dc_sim.py @@ -7,7 +7,6 @@ import re import shutil import socket -from dataclasses import dataclass from datetime import datetime from pathlib import Path @@ -17,7 +16,6 @@ import workflows.transport import zocalo.configuration from ispyb.sqlalchemy import ( - BLSample, BLSession, DataCollection, DataCollectionGroup, @@ -26,25 +24,16 @@ ) from sqlalchemy.orm import Load +from dlstbx.crud import get_dcids_for_data_collection_group from dlstbx.util import INDUSTRIAL_CODES, DowngradeErrorsFilter from dlstbx.util.colorstreamhandler import ColorStreamHandler logger = logging.getLogger(__name__) -@dataclass(frozen=True) -class DataCollectionRecords: - """Related ISPyB records needed to simulate one data collection.""" - - data_collection: DataCollection - data_collection_group: DataCollectionGroup - bl_session: BLSession - proposal: Proposal - grid_info: GridInfo | None - bl_sample: BLSample | None - - -def big_query(session: sqlalchemy.orm.Session, dcid: int) -> DataCollectionRecords: +def get_dc_and_grid_info_from_dcid( + session: sqlalchemy.orm.Session, dcid: int +) -> tuple[DataCollection, GridInfo | None]: """Fetch a data collection and its related records in one database query. GridInfo is optional and may be associated directly with the data collection @@ -57,41 +46,76 @@ def big_query(session: sqlalchemy.orm.Session, dcid: int) -> DataCollectionRecor result = ( session.query( DataCollection, - DataCollectionGroup, - BLSession, - Proposal, GridInfo, - BLSample, - ) - .join( - DataCollectionGroup, - DataCollection.dataCollectionGroupId - == DataCollectionGroup.dataCollectionGroupId, ) - .join(BLSession, DataCollectionGroup.sessionId == BLSession.sessionId) - .join(Proposal, BLSession.proposalId == Proposal.proposalId) .outerjoin( GridInfo, GridInfo.dataCollectionId == DataCollection.dataCollectionId ) - .outerjoin( - BLSample, - BLSample.blSampleId - == sqlalchemy.func.coalesce( - DataCollection.BLSAMPLEID, DataCollectionGroup.blSampleId - ), - ) .filter(DataCollection.dataCollectionId == dcid) .one_or_none() ) if result is None: + logger.error(f"No data collection found for dcid={dcid}") + raise ValueError(f"Could not find data collection for dcid={dcid}") + + return result + + +def get_data_collection_group_from_dcgid( + db_session: sqlalchemy.orm.Session, dcgid: int +) -> DataCollectionGroup: + dcg = ( + db_session.query(DataCollectionGroup).filter( + DataCollectionGroup.dataCollectionGroupId == dcgid + ) + ).one_or_none() + if dcg is None: + logger.error( + f"No data collection or related ISPyB records found for dcgid={dcgid}" + ) + raise ValueError( + f"Could not find a data collection and its related records for dcgid {dcgid}" + ) + return dcg + + +def get_data_collection_group_from_dcid( + db_session: sqlalchemy.orm.Session, dcid: int +) -> DataCollectionGroup: + dcg = ( + db_session.query(DataCollectionGroup) + .join( + DataCollection, + DataCollection.dataCollectionGroupId + == DataCollectionGroup.dataCollectionGroupId, + ) + .filter(DataCollection.dataCollectionId == dcid) + ).one_or_none() + if dcg is None: logger.error( f"No data collection or related ISPyB records found for dcid={dcid}" ) raise ValueError( - f"Could not find a data collection and its related records for dcid {dcid}" + f"Could not find a data collection and its related records for dcgid={dcid}" ) + return dcg + - return DataCollectionRecords(*result) +def get_bl_session_and_proposal_code_from_dcgid( + db_session: sqlalchemy.orm.Session, dcgid: int +) -> tuple[str, BLSession]: + results = ( + db_session.query(Proposal.proposalCode, BLSession) + .join(BLSession, BLSession.proposalId == Proposal.proposalId) + .join(DataCollectionGroup, DataCollectionGroup.sessionId == BLSession.sessionId) + .filter(DataCollectionGroup.dataCollectionGroupId == dcgid) + .one_or_none() + ) + + if results is None: + logger.error(f"No proposal code or BLSession found for dcgid={dcgid}") + raise ValueError(f"Could not find proposal code or BLSession for dcgid={dcgid}") + return results def retrieve_sessionid(db_session: sqlalchemy.orm.Session, visit: str) -> int: @@ -257,11 +281,9 @@ def copy_grid_info( return new_gi -def simulate_data_collection(dcid: int) -> int: +def simulate(dcid_or_dcgid: int, is_dcg: bool = False) -> list[int]: """Look up and copy the records needed to simulate data collection.""" - start_time = datetime.now() dflt_proposal = "nt44198" - logger.info(f"Starting data collection simulation for dcid={dcid}") sessionmaker = sqlalchemy.orm.sessionmaker( bind=sqlalchemy.create_engine( @@ -269,29 +291,33 @@ def simulate_data_collection(dcid: int) -> int: ) ) with sessionmaker() as session: - logger.info("Loading source data collection and related ISPyB records") - dc_records = big_query(session, dcid) - - if dc_records.proposal.proposalCode in INDUSTRIAL_CODES: - logger.error( - f"Source data collection {dcid} is industrial; simulation is not allowed" + if is_dcg: + logger.info( + f"Starting data collection simulation for dcgid={dcid_or_dcgid}" ) + dcg = get_data_collection_group_from_dcgid(session, dcid_or_dcgid) + dcgid = dcid_or_dcgid + dcids = get_dcids_for_data_collection_group(dcgid, session) + else: + logger.info(f"Starting data collection simulation for dcid={dcid_or_dcgid}") + dcg = get_data_collection_group_from_dcid(session, dcid_or_dcgid) + dcgid = dcg.dataCollectionGroupId + dcids = [dcid_or_dcgid] + + proposal_code, bl_session = get_bl_session_and_proposal_code_from_dcgid( + session, dcgid + ) + + if proposal_code in INDUSTRIAL_CODES: + logger.error("Source data is industrial; simulation is not allowed") raise SystemExit(1) - beamline = dc_records.bl_session.beamLineName + beamline = bl_session.beamLineName dest_visit_dir, dest_visit = get_beamline_visit_and_dir_for_proposal( dflt_proposal, beamline ) - dc = dc_records.data_collection - dcg = dc_records.data_collection_group - logger.info( - f"Source records loaded: dcid={dcid}, beamline={beamline}, " - f"destination visit={dest_visit}" - ) - now = datetime.now() - source_image_dir = Path(dc.imageDirectory) random_uuid = os.urandom(4).hex() # Generate a random UUID for uniqueness dest_dir = ( dest_visit_dir @@ -301,153 +327,174 @@ def simulate_data_collection(dcid: int) -> int: ) dest_dir.mkdir(parents=True, exist_ok=True) logger.info(f"Created image destination directory {dest_dir}") + # Create destination directory for JPEG snapshots + jpgs_dest_dir = ( + dest_visit_dir + / "jpegs" + / "tmp" + / f"{now:%Y-%m-%d}" + / f"{now:%H}-{now:%M}-{now:%S}-{random_uuid}" + ) + jpgs_dest_dir.mkdir(parents=True, exist_ok=True) + logger.info(f"Created snapshot destination directory {jpgs_dest_dir}") # Get the session id for the destination visit logger.info(f"Finding ISPyB session for destination visit {dest_visit}") new_bl_session_id = retrieve_sessionid(session, dest_visit) # Create a copy with new sessionId and times + logger.info("Copying data collection group") new_dcg = copy_data_collection_group( dcg, new_session_id=new_bl_session_id, - new_start_time=start_time, + new_start_time=datetime.now(), new_end_time=datetime.now(), ) # Add and commit the new DataCollectionGroup session.add(new_dcg) session.commit() - logger.info( - f"Created destination DataCollectionGroup {new_dcg.dataCollectionGroupId}" - ) + new_dcgid = new_dcg.dataCollectionGroupId + logger.info(f"Created destination DataCollectionGroup dcgid={new_dcgid}") - # Copy the image files to the new destination directory - image_suffix = dc.imageSuffix - if image_suffix == "cbf": - file_template = dc.fileTemplate - match = re.search(r"(#+)(?=\.cbf$)", file_template) - if not match: - raise ValueError( - f"File template {file_template} does not contain a valid '#' pattern for image numbering." - ) - n_digits = len(match.group(1)) - start_image_number = dc.startImageNumber - number_of_images = dc.numberOfImages - copied_images = 0 - missing_images = 0 - logger.info( - f"Copying {number_of_images} CBF images from {source_image_dir} " - f"to {dest_dir}" - ) - for image_number in range( - start_image_number, start_image_number + number_of_images - ): - image_filename = ( - file_template[: match.start()] - + str(image_number).zfill(n_digits) - + file_template[match.end() :] - ) - src_image_path = source_image_dir / image_filename - dest_image_path = dest_dir / image_filename - if src_image_path.exists(): - shutil.copy2(src_image_path, dest_image_path) - copied_images += 1 - else: - logger.warning( - f"Source image {src_image_path} does not exist; skipping copy" - ) - missing_images += 1 - logger.info( - f"Finished copying CBF images: copied={copied_images}, " - f"missing={missing_images}" + logger.info("Loading source data collection and related ISPyB records") + new_dcids = [] + for dcid in dcids: + new_dcid = simulate_datacollection( + session, dcid, dest_dir, jpgs_dest_dir, new_bl_session_id, new_dcgid ) + new_dcids.append(new_dcid) + return new_dcids + + +def simulate_datacollection( + db_session: sqlalchemy.orm.Session, + dcid: int, + dest_dir: Path, + jpgs_dest_dir: Path, + new_bl_session_id: int, + new_dcgid: int, +) -> int: + dc, gi = get_dc_and_grid_info_from_dcid(db_session, dcid) + source_image_dir = Path(dc.imageDirectory) + start_time = datetime.now() - elif image_suffix in ["h5", "nxs"]: - file_prefix = f"{dc.imagePrefix}_{dc.dataCollectionNumber}" - files = { - path - for pattern in (f"{file_prefix}_*.h5", f"{file_prefix}.nxs") - for path in source_image_dir.glob(pattern) - } - logger.info( - f"Copying {len(files)} HDF5/Nexus file(s) from {source_image_dir} " - f"to {dest_dir}" - ) - for file in files: - shutil.copy2(file, dest_dir) - logger.info("Finished copying image data files") - else: - logger.error( - f"Unsupported image suffix {image_suffix!r} for source dcid={dcid}" - ) + # Copy the image files to the new destination directory + image_suffix = dc.imageSuffix + if image_suffix == "cbf": + file_template = dc.fileTemplate + match = re.search(r"(#+)(?=\.cbf$)", file_template) + if not match: raise ValueError( - f"Invalid image suffix {image_suffix} for source data collection" + f"File template {file_template} does not contain a valid '#' pattern for image numbering." + ) + n_digits = len(match.group(1)) + start_image_number = dc.startImageNumber + number_of_images = dc.numberOfImages + copied_images = 0 + missing_images = 0 + logger.info( + f"Copying {number_of_images} CBF images from {source_image_dir} " + f"to {dest_dir}" + ) + for image_number in range( + start_image_number, start_image_number + number_of_images + ): + image_filename = ( + file_template[: match.start()] + + str(image_number).zfill(n_digits) + + file_template[match.end() :] ) + src_image_path = source_image_dir / image_filename + dest_image_path = dest_dir / image_filename + if src_image_path.exists(): + shutil.copy2(src_image_path, dest_image_path) + copied_images += 1 + else: + logger.warning( + f"Source image {src_image_path} does not exist; skipping copy" + ) + missing_images += 1 + logger.info( + f"Finished copying CBF images: copied={copied_images}, " + f"missing={missing_images}" + ) - # Create new DataCollection - new_dc = copy_data_collection( - dc, - new_session_id=new_bl_session_id, - new_data_collection_group_id=new_dcg.dataCollectionGroupId, - startTime=start_time, - endTime=datetime.now(), - comments=f"Simulated data collection copied from dcid {dcid}", - imageDirectory=str(dest_dir) + "/", + elif image_suffix in ["h5", "nxs"]: + file_prefix = f"{dc.imagePrefix}_{dc.dataCollectionNumber}" + files = { + path + for pattern in (f"{file_prefix}_*.h5", f"{file_prefix}.nxs") + for path in source_image_dir.glob(pattern) + } + logger.info( + f"Copying {len(files)} HDF5/Nexus file(s) from {source_image_dir} " + f"to {dest_dir}" + ) + for file in files: + shutil.copy2(file, dest_dir) + logger.info("Finished copying image data files") + else: + logger.error( + f"Unsupported image suffix {image_suffix!r} for source dcid={dcid}" + ) + raise ValueError( + f"Invalid image suffix {image_suffix} for source data collection" ) - session.add(new_dc) - session.commit() - new_dcid = new_dc.dataCollectionId - if dc_records.grid_info is not None: - new_gi = copy_grid_info(source_gi=dc_records.grid_info, new_dcid=new_dcid) - session.add(new_gi) - session.commit() + # Create new DataCollection + new_dc = copy_data_collection( + dc, + new_session_id=new_bl_session_id, + new_data_collection_group_id=new_dcgid, + startTime=start_time, + endTime=datetime.now(), + comments=f"Simulated data collection copied from dcid {dcid}", + imageDirectory=str(dest_dir) + "/", + ) + db_session.add(new_dc) + db_session.commit() + new_dcid = new_dc.dataCollectionId - # Create destination directory for JPEG snapshots - jpgs_dest_dir = ( - dest_visit_dir - / "jpegs" - / "tmp" - / f"{now:%Y-%m-%d}" - / f"{now:%H}-{now:%M}-{now:%S}-{random_uuid}" - ) - jpgs_dest_dir.mkdir(parents=True, exist_ok=True) - logger.info(f"Created snapshot destination directory {jpgs_dest_dir}") + if gi is not None: + new_gi = copy_grid_info(source_gi=gi, new_dcid=new_dcid) + db_session.add(new_gi) + db_session.commit() # Mapping of snapshot attributes to their names - snapshot_attrs = [ - "xtalSnapshotFullPath1", - "xtalSnapshotFullPath2", - "xtalSnapshotFullPath3", - "xtalSnapshotFullPath4", - ] - - # Copy and update snapshot paths - copied_snapshots = 0 - missing_snapshots = 0 - for attr_name in snapshot_attrs: - src_snapshot_path = getattr(dc, attr_name, None) - if src_snapshot_path is not None: - src_snapshot = Path(src_snapshot_path) - if src_snapshot.exists(): - # Create a unique filename for the snapshot - dest_snapshot_path = jpgs_dest_dir / src_snapshot.name - shutil.copy2(src_snapshot, dest_snapshot_path) - # Update the new DataCollection with the new snapshot path - setattr(new_dc, attr_name, str(dest_snapshot_path)) - copied_snapshots += 1 - else: - logger.warning( - f"Crystal snapshot {src_snapshot} does not exist; skipping copy" - ) - missing_snapshots += 1 + snapshot_attrs = [ + "xtalSnapshotFullPath1", + "xtalSnapshotFullPath2", + "xtalSnapshotFullPath3", + "xtalSnapshotFullPath4", + ] + + # Copy and update snapshot paths + copied_snapshots = 0 + missing_snapshots = 0 + for attr_name in snapshot_attrs: + src_snapshot_path = getattr(dc, attr_name, None) + if src_snapshot_path is not None: + src_snapshot = Path(src_snapshot_path) + if src_snapshot.exists(): + # Create a unique filename for the snapshot + dest_snapshot_path = jpgs_dest_dir / src_snapshot.name + shutil.copy2(src_snapshot, dest_snapshot_path) + # Update the new DataCollection with the new snapshot path + setattr(new_dc, attr_name, str(dest_snapshot_path)) + copied_snapshots += 1 + else: + logger.warning( + f"Crystal snapshot {src_snapshot} does not exist; skipping copy" + ) + missing_snapshots += 1 - logger.info( - f"Created destination DataCollection {new_dc.dataCollectionId}; " - f"copied snapshots={copied_snapshots}, missing snapshots={missing_snapshots}" - ) - logger.info( - f"Data collection simulation for dcid={dcid} completed in " - f"{(datetime.now() - start_time).total_seconds():.1f} seconds" - ) + logger.info( + f"Created destination DataCollection {new_dc.dataCollectionId}; " + f"copied snapshots={copied_snapshots}, missing snapshots={missing_snapshots}" + ) + logger.info( + f"Data collection simulation for dcid={dcid} completed in " + f"{(datetime.now() - start_time).total_seconds():.1f} seconds" + ) return new_dcid @@ -493,45 +540,51 @@ def run(args: list[str] | None = None) -> None: description="Simulate data collection. Takes a dcid as input and copies data and database records to simulate a data collection, then triggers mimas.", usage="dlstbx.dc_sim [options] dcid", ) - parser.add_argument("dcid", type=int, help="Data collection ID") + parser.add_argument( + "dcid_or_dcgid", + type=int, + help="Data collection ID or Data collection group ID - To use Data collection group id, set --dcgid option", + ) + parser.add_argument("--is-dcgid", action="store_true") parser.add_argument("--suppress-mimas", action="store_true") zc.add_command_line_options(parser) workflows.transport.add_command_line_options(parser, transport_argument=True) parsed_args = parser.parse_args(args) - new_dcid = simulate_data_collection(parsed_args.dcid) + new_dcids = simulate(parsed_args.dcid_or_dcgid, parsed_args.is_dcgid) if not parsed_args.suppress_mimas: - logger.info("Triggering Mimas") - headers = { - "zocalo.go.user": getpass.getuser(), - "zocalo.go.host": socket.gethostname(), - } - transport = workflows.transport.lookup(parsed_args.transport)() transport.connect() - transport.send( - "processing_recipe", - { - "recipes": ["mimas"], - "parameters": { - "ispyb_dcid": new_dcid, - "event": "start", + for new_dcid in new_dcids: + logger.info(f"Triggering Mimas for dcid={new_dcid}") + headers = { + "zocalo.go.user": getpass.getuser(), + "zocalo.go.host": socket.gethostname(), + } + + transport.send( + "processing_recipe", + { + "recipes": ["mimas"], + "parameters": { + "ispyb_dcid": new_dcid, + "event": "start", + }, }, - }, - headers=headers, - ) - transport.send( - "processing_recipe", - { - "recipes": ["mimas"], - "parameters": { - "ispyb_dcid": new_dcid, - "event": "end", + headers=headers, + ) + transport.send( + "processing_recipe", + { + "recipes": ["mimas"], + "parameters": { + "ispyb_dcid": new_dcid, + "event": "end", + }, }, - }, - headers=headers, - ) + headers=headers, + ) transport.disconnect() From fe1642cd6917f9afd3877ba0d76f6e8c14f2ffab Mon Sep 17 00:00:00 2001 From: Phil Blowey Date: Wed, 28 Oct 2026 16:15:22 +0000 Subject: [PATCH 06/14] Make self contained wrapper --- src/dlstbx/wrapper/dc_sim.py | 516 ++++++++++++++++++++++++++++++----- 1 file changed, 449 insertions(+), 67 deletions(-) diff --git a/src/dlstbx/wrapper/dc_sim.py b/src/dlstbx/wrapper/dc_sim.py index 56426e91a..9cc6c39f9 100644 --- a/src/dlstbx/wrapper/dc_sim.py +++ b/src/dlstbx/wrapper/dc_sim.py @@ -1,12 +1,349 @@ from __future__ import annotations -import dlstbx.dc_sim.definitions +import os +import re +import shutil +from datetime import datetime +from pathlib import Path + +import ispyb.sqlalchemy +import sqlalchemy +import sqlalchemy.orm +from ispyb.sqlalchemy import ( + BLSession, + DataCollection, + DataCollectionGroup, + GridInfo, + Proposal, +) +from sqlalchemy.orm import Load + +from dlstbx.util import INDUSTRIAL_CODES from dlstbx.wrapper import Wrapper class DCSimWrapper(Wrapper): _logger_name = "dlstbx.wrap.dc_sim" + def retrieve_sessionid(self, db_session: sqlalchemy.orm.Session, visit: str) -> int: + query = ( + db_session.query(BLSession, Proposal) + .options( + Load(BLSession).load_only( + BLSession.sessionId, BLSession.visit_number, BLSession.proposalId + ), + Load(Proposal).load_only( + Proposal.proposalId, Proposal.proposalCode, Proposal.proposalNumber + ), + ) + .join( + Proposal, + Proposal.proposalId == BLSession.proposalId, + ) + .filter( + sqlalchemy.func.concat( + Proposal.proposalCode, + Proposal.proposalNumber, + "-", + BLSession.visit_number, + ) + == visit + ) + ) + query_results = query.first() + if query_results is None: + raise ValueError(f"Query to obtain sessionid failed for {visit}") + if query_results[0].sessionId is None: + raise ValueError(f"Could not find sessionid for visit {visit}") + return query_results[0].sessionId + + def copy_data_collection_group( + self, db_session: sqlalchemy.orm.Session, dcgid: int, new_session_id: int + ) -> int: + """Copy a row from DataCollectionGroup table, replace session id, insert to the database and return new dcgid""" + now = datetime.now() + src_dcg = ( + db_session.query(DataCollectionGroup).filter( + DataCollectionGroup.dataCollectionGroupId == dcgid + ) + ).one_or_none() + if src_dcg is None: + self.log.error(f"No data collection group found for dcgid={dcgid}") + raise ValueError(f"Could not find data collection group dcgid={dcgid}") + + new_dcg = DataCollectionGroup() + # Copy all attributes from source to new, except the primary key and fields to update + for column in DataCollectionGroup.__table__.columns: + attr_name = column.name + + # Skip the primary key + if attr_name == "dataCollectionGroupId": + continue + + # Use new values for these attributes + if attr_name == "sessionId": + setattr(new_dcg, attr_name, new_session_id) + elif attr_name == "startTime": + setattr(new_dcg, attr_name, now) + elif attr_name == "endTime": + setattr(new_dcg, attr_name, now) + else: + # Copy from source + setattr(new_dcg, attr_name, getattr(src_dcg, attr_name, None)) + + # Add and commit new dcg + db_session.add(new_dcg) + db_session.commit() + return new_dcg.dataCollectionGroupId + + def get_dc_and_grid_info_from_dcid( + self, db_session: sqlalchemy.orm.Session, dcid: int + ) -> tuple[DataCollection, GridInfo | None]: + """Fetch a data collection and its related records in one database query. + + GridInfo is optional and may be associated directly with the data collection + or, for legacy records, with its data collection group. If both forms exist, + the direct association is preferred. BLSample is optional; the + DataCollection value is preferred, falling back to the group for legacy + records. + """ + + result = ( + db_session.query( + DataCollection, + GridInfo, + ) + .outerjoin( + GridInfo, GridInfo.dataCollectionId == DataCollection.dataCollectionId + ) + .filter(DataCollection.dataCollectionId == dcid) + .one_or_none() + ) + if result is None: + self.log.error(f"No data collection found for dcid={dcid}") + raise ValueError(f"Could not find data collection for dcid={dcid}") + + return result + + def copy_data_collection( + self, + db_session: sqlalchemy.orm.Session, + source_dc: DataCollection, + new_session_id: int, + new_data_collection_group_id: int, + **kwargs, + ) -> DataCollection: + """ + Create a copy of a DataCollection row with different sessionId and dataCollectionGroupId. + and insert into the database + + Args: + db_session: Database session + source_dc: The DataCollection to copy + new_session_id: The new SESSIONID for the copy + new_data_collection_group_id: The new dataCollectionGroupId for the copy + **kwargs: Additional attributes to override (e.g., startTime, endTime, comments) + + Returns: + The new data collection id + """ + # Create a new DataCollection object + new_dc = DataCollection() + + # Copy all attributes from source to new, except the primary key and fields to update + for column in DataCollection.__table__.columns: + attr_name = column.name + + # Skip the primary key + if attr_name == "dataCollectionId": + continue + + # Use new values for these attributes + if attr_name == "SESSIONID": + setattr(new_dc, attr_name, new_session_id) + elif attr_name == "dataCollectionGroupId": + setattr(new_dc, attr_name, new_data_collection_group_id) + elif attr_name in kwargs: + # Use provided override values + setattr(new_dc, attr_name, kwargs[attr_name]) + else: + # Copy from source + setattr(new_dc, attr_name, getattr(source_dc, attr_name, None)) + + # Add and commit new data collection + db_session.add(new_dc) + db_session.commit() + return new_dc.dataCollectionId + + def copy_grid_info( + self, + db_session: sqlalchemy.orm.Session, + source_gi: GridInfo, + new_dcid: int, + **kwargs, + ) -> DataCollection: + """ + Create a copy of a DataCollection row with different sessionId and/or dataCollectionGroupId. + + Args: + db_session: Database session + source_gc: The GridInfo to copy + new_dcid: The new dcid to link the copy to + + Returns: + None + """ + # Create a new GridInfo object + new_gi = GridInfo() + + # Copy all attributes from source to new, except the primary key and fields to update + for column in GridInfo.__table__.columns: + attr_name = column.name + + # Skip the primary key + if attr_name == "gridInfoId": + continue + + # Use new values for these attributes + if attr_name == "dataCollectionId": + setattr(new_gi, attr_name, new_dcid) + elif attr_name in kwargs: + # Use provided override values + setattr(new_gi, attr_name, kwargs[attr_name]) + else: + # Copy from source + setattr(new_gi, attr_name, getattr(source_gi, attr_name, None)) + + # Add and commit new GridInfo + db_session.add(new_gi) + db_session.commit() + + def simulate_datacollection( + self, + db_session: sqlalchemy.orm.Session, + dcid: int, + dest_dir: Path, + jpgs_dest_dir: Path, + new_bl_session_id: int, + new_dcgid: int, + ) -> int: + dc, gi = self.get_dc_and_grid_info_from_dcid(db_session, dcid) + source_image_dir = Path(dc.imageDirectory) + start_time = datetime.now() + + # Copy the image files to the new destination directory + image_suffix = dc.imageSuffix + if image_suffix == "cbf": + file_template = dc.fileTemplate + match = re.search(r"(#+)(?=\.cbf$)", file_template) + if not match: + raise ValueError( + f"File template {file_template} does not contain a valid '#' pattern for image numbering." + ) + n_digits = len(match.group(1)) + start_image_number = dc.startImageNumber + number_of_images = dc.numberOfImages + copied_images = 0 + missing_images = 0 + self.log.info( + f"Copying {number_of_images} CBF images from {source_image_dir} " + f"to {dest_dir}" + ) + for image_number in range( + start_image_number, start_image_number + number_of_images + ): + image_filename = ( + file_template[: match.start()] + + str(image_number).zfill(n_digits) + + file_template[match.end() :] + ) + src_image_path = source_image_dir / image_filename + dest_image_path = dest_dir / image_filename + if src_image_path.exists(): + shutil.copy2(src_image_path, dest_image_path) + copied_images += 1 + else: + self.log.warning( + f"Source image {src_image_path} does not exist; skipping copy" + ) + missing_images += 1 + self.log.info( + f"Finished copying CBF images: copied={copied_images}, " + f"missing={missing_images}" + ) + + elif image_suffix in ["h5", "nxs"]: + file_prefix = f"{dc.imagePrefix}_{dc.dataCollectionNumber}" + files = { + path + for pattern in (f"{file_prefix}_*.h5", f"{file_prefix}.nxs") + for path in source_image_dir.glob(pattern) + } + self.log.info( + f"Copying {len(files)} HDF5/Nexus file(s) from {source_image_dir} " + f"to {dest_dir}" + ) + for file in files: + shutil.copy2(file, dest_dir) + self.log.info("Finished copying image data files") + else: + self.log.error( + f"Unsupported image suffix {image_suffix!r} for source dcid={dcid}" + ) + raise ValueError( + f"Invalid image suffix {image_suffix} for source data collection" + ) + + # Extra params to set in new dc + dc_extra_params = { + "startTime": start_time, + "endTime": datetime.now(), + "comments": f"Simulated data collection copied from dcid {dcid}", + "imageDirectory": str(dest_dir) + "/", + } + + # Handle crystal snapshots + # Mapping of snapshot attributes to their names + snapshot_attrs = [ + "xtalSnapshotFullPath1", + "xtalSnapshotFullPath2", + "xtalSnapshotFullPath3", + "xtalSnapshotFullPath4", + ] + for attr_name in snapshot_attrs: + src_snapshot_path = getattr(dc, attr_name, None) + if src_snapshot_path is not None: + src_snapshot = Path(src_snapshot_path) + if src_snapshot.exists(): + # Create a unique filename for the snapshot + dest_snapshot_path = jpgs_dest_dir / src_snapshot.name + shutil.copy2(src_snapshot, dest_snapshot_path) + # Update the new DataCollection with the new snapshot path + dc_extra_params[attr_name] = str(dest_snapshot_path) + else: + self.log.warning( + f"Crystal snapshot {src_snapshot} does not exist; skipping copy" + ) + + # Create new DataCollection + new_dcid = self.copy_data_collection( + db_session=db_session, + source_dc=dc, + new_session_id=new_bl_session_id, + new_data_collection_group_id=new_dcgid, + **dc_extra_params, + ) + + if gi is not None: + self.copy_grid_info(db_session=db_session, source_gi=gi, new_dcid=new_dcid) + + self.log.info(f"Created DataCollection {new_dcid}; ") + self.log.info( + f"Data collection simulation for dcid={dcid} completed in " + f"{(datetime.now() - start_time).total_seconds():.1f} seconds" + ) + return new_dcid + def run(self): assert hasattr(self, "recwrap"), "No recipewrapper object found" # Get parameters from the recipe file (as a copy so that changes can be made without affecting anything) @@ -14,14 +351,8 @@ def run(self): # A list of placeholder values to check for in params placeholders = [ - "{src_dir}", - "{src_prefix}", - "{src_run_num}", - "{sample_id}", - "{visit}", - "{src_dcid}", - "{src_dcg}", - "{is_dcg}", + "{dest_visit}", + "{copy_dcg}", ] # Replace any remaining placeholder values in params with None. @@ -29,67 +360,118 @@ def run(self): if value in placeholders: params[key] = None - self.log.info( - "Running simulated data collection '%s' on beamline '%s'", - params["scenario"], - params["beamline"], + beamline = params["beamline"] + src_visit = params["src_visit"] + dest_visit = params.get("dest_visit") + dcid = int(params["dcid"]) + now = datetime.now() + random_uuid = os.urandom(4).hex() # Generate a random UUID for uniqueness + + if src_visit.startswith(tuple(INDUSTRIAL_CODES)) or not src_visit.startswith( + tuple(params["src_allowed_visits"]) + ): + raise ValueError(f"{src_visit} is not an allowed source visit") + + if dest_visit is not None: + if not params["dest_visit"].startswith( + tuple(params["dest_allowed_visits"]) + ): + raise ValueError(f"{dest_visit} is not an allowed destination visit") + proposal = dest_visit.split("-")[0] + else: + proposal = params["dflt_proposal"] + if beamline.startswith("i02"): + if dest_visit is None: + try: + dest_visit = f"{proposal}-{params['dflt_i02_visits'][beamline]}" + except KeyError: + self.log.error("Could not determine which destination visit to use") + raise + dest_visit_dir = Path("/dls/mx/data", proposal, dest_visit) + else: + if dest_visit is None: + for visit_dir in Path( + "/dls", beamline, "data", str(now.year) + ).iterdir(): + if visit_dir.name.startswith(proposal): + dest_visit = visit_dir.name + break + dest_visit_dir = Path("/dls", beamline, "data", str(now.year), dest_visit) + assert dest_visit_dir.is_dir(), ( + f"Could not find {dest_visit_dir} directory for the specified visit and beamline." ) - # Convert command line input of certain parameters into a list - for key in ["src_prefix", "src_run_num"]: - try: - value = eval(params[key]) - if isinstance(value, list) or isinstance(value, tuple): - params[key] = value - else: - params[key] = [ - value, - ] - except (SyntaxError, NameError, TypeError): - # Case for dealing with non-evaluatable input (e.g. string) - if params[key] is not None: - params[key] = [ - params[key], - ] - # Check that supplied destination visit is allowed - if params["visit"] is not None: - if not params["visit"].startswith(tuple(params["dest_allowed_visits"])): - raise ValueError(f"{params['visit']} is not an allowed visit") - - # Set whether to output as dcg - is_dcg = params["src_dcg"] is not None - if params["is_dcg"] is not None: - # Check that is_dcg is an acceptable boolean value - if params["is_dcg"].lower() in ["true", "1", "y", "yes"]: - is_dcg = True - elif params["is_dcg"].lower() in ["false", "0", "n", "no"]: - is_dcg = False - else: - raise ValueError( - f"{params['is_dcg']} is not a valid value for is_dcg. Acceptable values are 'true', '1', 'y', 'yes' for True and 'false', '0', 'n', 'no' for False - case insensitive" - ) + dcgid = params["dc_info"]["dataCollectionGroupId"] + dcids = [dcid] + # Set whether to copy entire dcg + copy_dcg = params["copy_dcg"] + if copy_dcg: + dcids = ( + params["dcg_dcids"] + if dcid in params["dcg_dcids"] + else [dcid] + params["dcg_dcids"] + ) + dcids.sort() - # Simulate the data collection - result = dlstbx.dc_sim.call_sim( - test_name=params["scenario"], - beamline=params["beamline"], - src_dir=params["src_dir"], - src_prefix=params["src_prefix"], - src_run_num=params["src_run_num"], - sample_id=params["sample_id"], - dest_visit=params["visit"], - dflt_proposals=params["dflt_proposals"], - dflt_i02_visits=params["dflt_i02_visits"], - src_dcid=params["src_dcid"], - src_allowed_visits=params["src_allowed_visits"], - is_dcg=is_dcg, - src_dcg=params["src_dcg"], + dest_dir = ( + dest_visit_dir + / "tmp" + / f"{now:%Y-%m-%d}" + / f"{now:%H}-{now:%M}-{now:%S}-{random_uuid}" ) + dest_dir.mkdir(parents=True, exist_ok=True) + self.log.info(f"Created image destination directory {dest_dir}") + # Create destination directory for JPEG snapshots + jpgs_dest_dir = ( + dest_visit_dir + / "jpegs" + / "tmp" + / f"{now:%Y-%m-%d}" + / f"{now:%H}-{now:%M}-{now:%S}-{random_uuid}" + ) + jpgs_dest_dir.mkdir(parents=True, exist_ok=True) + self.log.info(f"Created snapshot destination directory {jpgs_dest_dir}") - if result: - self.log.info(f"Simulated data collection completed with {result!r}") - else: - self.log.error("Simulated data collection failed") + sessionmaker = sqlalchemy.orm.sessionmaker( + bind=sqlalchemy.create_engine( + ispyb.sqlalchemy.url(), connect_args={"use_pure": True} + ) + ) + + with sessionmaker() as session: + dest_bl_session_id = self.retrieve_sessionid(session, dest_visit) + + new_dcgid = self.copy_data_collection_group( + db_session=session, + dcgid=dcgid, + new_session_id=dest_bl_session_id, + ) + + for dcid in dcids: + new_dcid = self.simulate_datacollection( + db_session=session, + dcid=dcid, + dest_dir=dest_dir, + jpgs_dest_dir=jpgs_dest_dir, + new_bl_session_id=dest_bl_session_id, + new_dcgid=new_dcgid, + ) + + self.log.info(f"Triggering mimas start event for dcid={new_dcid}") + self.recwrap.send_to( + "mimas", + { + "recipes": ["mimas"], + "parameters": {"ispyb_dcid": new_dcid, "event": "start"}, + }, + ) + self.log.info(f"Triggering mimas end event for dcid={new_dcid}") + self.recwrap.send_to( + "mimas", + { + "recipes": ["mimas"], + "parameters": {"ispyb_dcid": new_dcid, "event": "end"}, + }, + ) - self.recwrap.send_to("dc_sim", result._asdict()) - return bool(result) + return True From 116355fa852c022724cb217f448f1093b0f2b414 Mon Sep 17 00:00:00 2001 From: Phil Blowey Date: Thu, 29 Oct 2026 10:05:27 +0000 Subject: [PATCH 07/14] Add input validation --- src/dlstbx/wrapper/dc_sim.py | 107 ++++++++++++++++++++++------------- 1 file changed, 69 insertions(+), 38 deletions(-) diff --git a/src/dlstbx/wrapper/dc_sim.py b/src/dlstbx/wrapper/dc_sim.py index 9cc6c39f9..e71798cbe 100644 --- a/src/dlstbx/wrapper/dc_sim.py +++ b/src/dlstbx/wrapper/dc_sim.py @@ -7,6 +7,7 @@ from pathlib import Path import ispyb.sqlalchemy +import pydantic import sqlalchemy import sqlalchemy.orm from ispyb.sqlalchemy import ( @@ -22,6 +23,26 @@ from dlstbx.wrapper import Wrapper +@pydantic.dataclasses.dataclass(config=pydantic.ConfigDict(extra="allow")) +class DataCollectionInfo: + dataCollectionGroupId: pydantic.PositiveInt + + +@pydantic.dataclasses.dataclass +class DCSimParameters: + beamline: str + src_visit: str + src_allowed_visits: list[str] + dest_allowed_visits: list[str] + dflt_proposal: str + dflt_i02_visits: dict[str, int] + dc_info: DataCollectionInfo + dcg_dcids: list[pydantic.PositiveInt] + dcid: pydantic.PositiveInt + dest_visit: str | None = None + copy_dcg: bool | None = False + + class DCSimWrapper(Wrapper): _logger_name = "dlstbx.wrap.dc_sim" @@ -226,8 +247,9 @@ def simulate_datacollection( jpgs_dest_dir: Path, new_bl_session_id: int, new_dcgid: int, - ) -> int: + ) -> tuple[int, bool]: dc, gi = self.get_dc_and_grid_info_from_dcid(db_session, dcid) + is_grid_scan = gi is not None source_image_dir = Path(dc.imageDirectory) start_time = datetime.now() @@ -334,7 +356,7 @@ def simulate_datacollection( **dc_extra_params, ) - if gi is not None: + if is_grid_scan: self.copy_grid_info(db_session=db_session, source_gi=gi, new_dcid=new_dcid) self.log.info(f"Created DataCollection {new_dcid}; ") @@ -342,7 +364,7 @@ def simulate_datacollection( f"Data collection simulation for dcid={dcid} completed in " f"{(datetime.now() - start_time).total_seconds():.1f} seconds" ) - return new_dcid + return new_dcid, is_grid_scan def run(self): assert hasattr(self, "recwrap"), "No recipewrapper object found" @@ -360,30 +382,30 @@ def run(self): if value in placeholders: params[key] = None - beamline = params["beamline"] - src_visit = params["src_visit"] - dest_visit = params.get("dest_visit") - dcid = int(params["dcid"]) + params = DCSimParameters(**params) + + beamline = params.beamline + src_visit = params.src_visit + dest_visit = params.dest_visit + dcid = params.dcid now = datetime.now() random_uuid = os.urandom(4).hex() # Generate a random UUID for uniqueness if src_visit.startswith(tuple(INDUSTRIAL_CODES)) or not src_visit.startswith( - tuple(params["src_allowed_visits"]) + tuple(params.src_allowed_visits) ): raise ValueError(f"{src_visit} is not an allowed source visit") if dest_visit is not None: - if not params["dest_visit"].startswith( - tuple(params["dest_allowed_visits"]) - ): + if not dest_visit.startswith(tuple(params.dest_allowed_visits)): raise ValueError(f"{dest_visit} is not an allowed destination visit") proposal = dest_visit.split("-")[0] else: - proposal = params["dflt_proposal"] + proposal = params.dflt_proposal if beamline.startswith("i02"): if dest_visit is None: try: - dest_visit = f"{proposal}-{params['dflt_i02_visits'][beamline]}" + dest_visit = f"{proposal}-{params.dflt_i02_visits[beamline]}" except KeyError: self.log.error("Could not determine which destination visit to use") raise @@ -401,15 +423,15 @@ def run(self): f"Could not find {dest_visit_dir} directory for the specified visit and beamline." ) - dcgid = params["dc_info"]["dataCollectionGroupId"] + dcgid = params.dc_info.dataCollectionGroupId dcids = [dcid] # Set whether to copy entire dcg - copy_dcg = params["copy_dcg"] + copy_dcg = params.copy_dcg if copy_dcg: dcids = ( - params["dcg_dcids"] - if dcid in params["dcg_dcids"] - else [dcid] + params["dcg_dcids"] + params.dcg_dcids + if dcid in params.dcg_dcids + else [dcid] + params.dcg_dcids ) dcids.sort() @@ -446,9 +468,9 @@ def run(self): dcgid=dcgid, new_session_id=dest_bl_session_id, ) - + grid_scan_dcids: list[int] = [] for dcid in dcids: - new_dcid = self.simulate_datacollection( + new_dcid, is_grid_scan = self.simulate_datacollection( db_session=session, dcid=dcid, dest_dir=dest_dir, @@ -456,22 +478,31 @@ def run(self): new_bl_session_id=dest_bl_session_id, new_dcgid=new_dcgid, ) - - self.log.info(f"Triggering mimas start event for dcid={new_dcid}") - self.recwrap.send_to( - "mimas", - { - "recipes": ["mimas"], - "parameters": {"ispyb_dcid": new_dcid, "event": "start"}, - }, - ) - self.log.info(f"Triggering mimas end event for dcid={new_dcid}") - self.recwrap.send_to( - "mimas", - { - "recipes": ["mimas"], - "parameters": {"ispyb_dcid": new_dcid, "event": "end"}, - }, - ) - + # 3D grid scans expect all dcids in dcg to be registered before processing + if is_grid_scan: + grid_scan_dcids.append(new_dcid) + # Trigger processing for non grid scans as soon as they are ready. + else: + self.trigger_mimas_for_dcid(new_dcid) + # Process grid scans once all new dcids are in the database + for grid_scan_dcid in grid_scan_dcids: + self.trigger_mimas_for_dcid(grid_scan_dcid) return True + + def trigger_mimas_for_dcid(self, new_dcid): + self.log.info(f"Triggering mimas start event for dcid={new_dcid}") + self.recwrap.send_to( + "mimas", + { + "recipes": ["mimas"], + "parameters": {"ispyb_dcid": new_dcid, "event": "start"}, + }, + ) + self.log.info(f"Triggering mimas end event for dcid={new_dcid}") + self.recwrap.send_to( + "mimas", + { + "recipes": ["mimas"], + "parameters": {"ispyb_dcid": new_dcid, "event": "end"}, + }, + ) From fcc4581e705a7e82656676bf59893362534277ba Mon Sep 17 00:00:00 2001 From: Phil Blowey Date: Thu, 29 Oct 2026 10:05:52 +0000 Subject: [PATCH 08/14] Add unit tests --- tests/wrapper/test_dc_sim.py | 321 +++++++++++++++++++++++++++++++++++ 1 file changed, 321 insertions(+) create mode 100644 tests/wrapper/test_dc_sim.py diff --git a/tests/wrapper/test_dc_sim.py b/tests/wrapper/test_dc_sim.py new file mode 100644 index 000000000..65b3dee6c --- /dev/null +++ b/tests/wrapper/test_dc_sim.py @@ -0,0 +1,321 @@ +from __future__ import annotations + +from datetime import datetime +from types import SimpleNamespace +from unittest.mock import ANY, Mock + +import pydantic +import pytest +from ispyb.sqlalchemy import DataCollection, DataCollectionGroup, GridInfo + +from dlstbx.wrapper.dc_sim import DCSimParameters, DCSimWrapper + + +class FakeSession: + def __init__(self, query_result=None): + self.query_result = query_result + self.added = [] + self.commits = 0 + + def query(self, *_args): + return self + + def options(self, *_args): + return self + + def join(self, *_args): + return self + + def outerjoin(self, *_args): + return self + + def filter(self, *_args): + return self + + def first(self): + return self.query_result + + def one_or_none(self): + return self.query_result + + def add(self, instance): + self.added.append(instance) + + def commit(self): + self.commits += 1 + for instance in self.added: + if isinstance(instance, DataCollectionGroup): + instance.dataCollectionGroupId = 900 + elif isinstance(instance, DataCollection): + instance.dataCollectionId = 901 + + +def values_for_model(model, prefix="source"): + return { + column.name: f"{prefix}_{column.name}" + for column in model.__table__.columns + if column.name + not in {"dataCollectionGroupId", "dataCollectionId", "gridInfoId"} + } + + +def test_dc_sim_parameters_validate_and_expose_attributes(): + params = DCSimParameters( + beamline="i02-1", + src_visit="nt44198-1", + src_allowed_visits=["nt44198"], + dest_allowed_visits=["nt44198"], + dflt_proposal="nt44198", + dflt_i02_visits={"i02-1": 7}, + dc_info={"dataCollectionGroupId": "42", "runStatus": "DataCollection"}, + dcg_dcids=["100", 101], + dcid="100", + dest_visit=None, + copy_dcg=None, + ) + + assert params.beamline == "i02-1" + assert params.dcid == 100 + assert params.dcg_dcids == [100, 101] + assert params.dc_info.dataCollectionGroupId == 42 + assert params.copy_dcg is None + + +def test_dc_sim_parameters_reject_invalid_values(): + with pytest.raises(pydantic.ValidationError): + DCSimParameters( + beamline="i02-1", + src_visit="nt44198-1", + src_allowed_visits=["nt44198"], + dest_allowed_visits=["nt44198"], + dflt_proposal="nt44198", + dflt_i02_visits={"i02-1": 7}, + dc_info={"dataCollectionGroupId": 42}, + dcg_dcids=[], + dcid=0, + ) + + +def test_retrieve_sessionid_returns_session_id(): + wrapper = DCSimWrapper() + session = FakeSession((SimpleNamespace(sessionId=123), SimpleNamespace())) + + assert wrapper.retrieve_sessionid(session, "nt44198-1") == 123 + + +@pytest.mark.parametrize( + ("result", "message"), + [ + (None, "Query to obtain sessionid failed"), + ( + (SimpleNamespace(sessionId=None), SimpleNamespace()), + "Could not find sessionid", + ), + ], +) +def test_retrieve_sessionid_raises_when_visit_is_not_found(result, message): + wrapper = DCSimWrapper() + + with pytest.raises(ValueError, match=message): + wrapper.retrieve_sessionid(FakeSession(result), "nt44198-1") + + +def test_get_dc_and_grid_info_returns_records(): + wrapper = DCSimWrapper() + dc = SimpleNamespace(dataCollectionId=42) + grid_info = SimpleNamespace(gridInfoId=7) + + assert wrapper.get_dc_and_grid_info_from_dcid(FakeSession((dc, grid_info)), 42) == ( + dc, + grid_info, + ) + + +def test_get_dc_and_grid_info_raises_for_missing_record(): + wrapper = DCSimWrapper() + + with pytest.raises(ValueError, match="Could not find data collection for dcid=42"): + wrapper.get_dc_and_grid_info_from_dcid(FakeSession(None), 42) + + +def test_copy_data_collection_overrides_ids_and_fields(): + wrapper = DCSimWrapper() + source_dc = SimpleNamespace(**values_for_model(DataCollection)) + session = FakeSession() + + new_dcid = wrapper.copy_data_collection( + session, + source_dc, + new_session_id=101, + new_data_collection_group_id=202, + comments="simulation", + ) + + copied = session.added[0] + assert new_dcid == 901 + assert copied.SESSIONID == 101 + assert copied.dataCollectionGroupId == 202 + assert copied.comments == "simulation" + assert copied.imageSuffix == source_dc.imageSuffix + assert session.commits == 1 + + +def test_copy_data_collection_group_updates_session_and_timestamps(): + wrapper = DCSimWrapper() + source = SimpleNamespace(**values_for_model(DataCollectionGroup)) + session = FakeSession(source) + + new_dcgid = wrapper.copy_data_collection_group(session, 12, 34) + + copied = session.added[0] + assert new_dcgid == 900 + assert copied.sessionId == 34 + assert isinstance(copied.startTime, datetime) + assert copied.endTime == copied.startTime + assert copied.comments == source.comments + assert session.commits == 1 + + +def test_copy_grid_info_relinks_record_to_new_data_collection(): + wrapper = DCSimWrapper() + source = SimpleNamespace( + **{ + column.name: f"source_{column.name}" + for column in GridInfo.__table__.columns + if column.name not in {"gridInfoId", "dataCollectionId"} + } + ) + session = FakeSession() + + wrapper.copy_grid_info(session, source, new_dcid=45) + + copied = session.added[0] + assert copied.dataCollectionId == 45 + assert copied.gridInfoId is None + for column in GridInfo.__table__.columns: + if column.name not in {"gridInfoId", "dataCollectionId"}: + assert getattr(copied, column.name) == getattr(source, column.name) + assert session.commits == 1 + + +def test_simulate_datacollection_copies_cbf_images_and_snapshots(tmp_path, monkeypatch): + wrapper = DCSimWrapper() + source_dir = tmp_path / "source" + dest_dir = tmp_path / "destination" + jpgs_dest_dir = tmp_path / "snapshots" + source_dir.mkdir() + dest_dir.mkdir() + jpgs_dest_dir.mkdir() + (source_dir / "image_0001.cbf").write_text("one") + (source_dir / "image_0003.cbf").write_text("three") + snapshot = tmp_path / "crystal.jpg" + snapshot.write_text("snapshot") + dc = SimpleNamespace( + imageDirectory=str(source_dir), + imageSuffix="cbf", + fileTemplate="image_####.cbf", + startImageNumber=1, + numberOfImages=3, + xtalSnapshotFullPath1=str(snapshot), + xtalSnapshotFullPath2=None, + xtalSnapshotFullPath3=None, + xtalSnapshotFullPath4=None, + ) + grid_info = SimpleNamespace(gridInfoId=4) + monkeypatch.setattr( + wrapper, "get_dc_and_grid_info_from_dcid", lambda *_: (dc, grid_info) + ) + copy_dc = Mock(return_value=501) + copy_grid = Mock() + monkeypatch.setattr(wrapper, "copy_data_collection", copy_dc) + monkeypatch.setattr(wrapper, "copy_grid_info", copy_grid) + + new_dcid = wrapper.simulate_datacollection( + FakeSession(), 25, dest_dir, jpgs_dest_dir, 30, 40 + ) + + assert new_dcid == 501 + assert sorted(path.name for path in dest_dir.iterdir()) == [ + "image_0001.cbf", + "image_0003.cbf", + ] + assert (jpgs_dest_dir / "crystal.jpg").read_text() == "snapshot" + assert copy_dc.call_args.kwargs["xtalSnapshotFullPath1"] == str( + jpgs_dest_dir / "crystal.jpg" + ) + assert copy_dc.call_args.kwargs["new_session_id"] == 30 + assert copy_dc.call_args.kwargs["new_data_collection_group_id"] == 40 + copy_grid.assert_called_once_with(db_session=ANY, source_gi=grid_info, new_dcid=501) + + +def test_simulate_datacollection_copies_hdf5_and_nexus_files(tmp_path, monkeypatch): + wrapper = DCSimWrapper() + source_dir = tmp_path / "source" + dest_dir = tmp_path / "destination" + jpgs_dest_dir = tmp_path / "snapshots" + source_dir.mkdir() + dest_dir.mkdir() + jpgs_dest_dir.mkdir() + for name in ( + "sample_2_000001.h5", + "sample_2_000002.h5", + "sample_2.nxs", + "other.h5", + ): + (source_dir / name).write_text(name) + dc = SimpleNamespace( + imageDirectory=str(source_dir), + imageSuffix="h5", + imagePrefix="sample", + dataCollectionNumber=2, + xtalSnapshotFullPath1=None, + xtalSnapshotFullPath2=None, + xtalSnapshotFullPath3=None, + xtalSnapshotFullPath4=None, + ) + monkeypatch.setattr( + wrapper, "get_dc_and_grid_info_from_dcid", lambda *_: (dc, None) + ) + copy_dc = Mock(return_value=502) + copy_grid = Mock() + monkeypatch.setattr(wrapper, "copy_data_collection", copy_dc) + monkeypatch.setattr(wrapper, "copy_grid_info", copy_grid) + + assert ( + wrapper.simulate_datacollection( + FakeSession(), 25, dest_dir, jpgs_dest_dir, 30, 40 + ) + == 502 + ) + assert {path.name for path in dest_dir.iterdir()} == { + "sample_2_000001.h5", + "sample_2_000002.h5", + "sample_2.nxs", + } + copy_grid.assert_not_called() + + +@pytest.mark.parametrize( + ("suffix", "template", "message"), + [ + ("cbf", "image.cbf", "does not contain a valid '#' pattern"), + ("tif", "image_####.tif", "Invalid image suffix tif"), + ], +) +def test_simulate_datacollection_rejects_bad_image_format( + tmp_path, monkeypatch, suffix, template, message +): + wrapper = DCSimWrapper() + dc = SimpleNamespace( + imageDirectory=str(tmp_path), + imageSuffix=suffix, + fileTemplate=template, + startImageNumber=1, + numberOfImages=1, + ) + monkeypatch.setattr( + wrapper, "get_dc_and_grid_info_from_dcid", lambda *_: (dc, None) + ) + + with pytest.raises(ValueError, match=message): + wrapper.simulate_datacollection(FakeSession(), 25, tmp_path, tmp_path, 30, 40) From ada1a9f9984755f05b8aaec4fb1e72d757f16615 Mon Sep 17 00:00:00 2001 From: Phil Blowey Date: Thu, 29 Oct 2026 10:14:02 +0000 Subject: [PATCH 09/14] Add unit tests --- tests/wrapper/test_dc_sim.py | 22 ++++++++++------------ 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/tests/wrapper/test_dc_sim.py b/tests/wrapper/test_dc_sim.py index 65b3dee6c..f4bd08411 100644 --- a/tests/wrapper/test_dc_sim.py +++ b/tests/wrapper/test_dc_sim.py @@ -221,20 +221,20 @@ def test_simulate_datacollection_copies_cbf_images_and_snapshots(tmp_path, monke xtalSnapshotFullPath3=None, xtalSnapshotFullPath4=None, ) - grid_info = SimpleNamespace(gridInfoId=4) monkeypatch.setattr( - wrapper, "get_dc_and_grid_info_from_dcid", lambda *_: (dc, grid_info) + wrapper, "get_dc_and_grid_info_from_dcid", lambda *_: (dc, None) ) copy_dc = Mock(return_value=501) copy_grid = Mock() monkeypatch.setattr(wrapper, "copy_data_collection", copy_dc) monkeypatch.setattr(wrapper, "copy_grid_info", copy_grid) - new_dcid = wrapper.simulate_datacollection( + new_dcid, is_grid_scan = wrapper.simulate_datacollection( FakeSession(), 25, dest_dir, jpgs_dest_dir, 30, 40 ) assert new_dcid == 501 + assert not is_grid_scan assert sorted(path.name for path in dest_dir.iterdir()) == [ "image_0001.cbf", "image_0003.cbf", @@ -245,7 +245,7 @@ def test_simulate_datacollection_copies_cbf_images_and_snapshots(tmp_path, monke ) assert copy_dc.call_args.kwargs["new_session_id"] == 30 assert copy_dc.call_args.kwargs["new_data_collection_group_id"] == 40 - copy_grid.assert_called_once_with(db_session=ANY, source_gi=grid_info, new_dcid=501) + copy_grid.assert_not_called() def test_simulate_datacollection_copies_hdf5_and_nexus_files(tmp_path, monkeypatch): @@ -273,26 +273,24 @@ def test_simulate_datacollection_copies_hdf5_and_nexus_files(tmp_path, monkeypat xtalSnapshotFullPath3=None, xtalSnapshotFullPath4=None, ) + grid_info = SimpleNamespace(gridInfoId=4) monkeypatch.setattr( - wrapper, "get_dc_and_grid_info_from_dcid", lambda *_: (dc, None) + wrapper, "get_dc_and_grid_info_from_dcid", lambda *_: (dc, grid_info) ) copy_dc = Mock(return_value=502) copy_grid = Mock() monkeypatch.setattr(wrapper, "copy_data_collection", copy_dc) monkeypatch.setattr(wrapper, "copy_grid_info", copy_grid) - assert ( - wrapper.simulate_datacollection( - FakeSession(), 25, dest_dir, jpgs_dest_dir, 30, 40 - ) - == 502 - ) + assert wrapper.simulate_datacollection( + FakeSession(), 25, dest_dir, jpgs_dest_dir, 30, 40 + ) == (502, True) assert {path.name for path in dest_dir.iterdir()} == { "sample_2_000001.h5", "sample_2_000002.h5", "sample_2.nxs", } - copy_grid.assert_not_called() + copy_grid.assert_called_once_with(db_session=ANY, source_gi=grid_info, new_dcid=502) @pytest.mark.parametrize( From 61203f71c409738f259cf72f544e27bb3c2d8eda Mon Sep 17 00:00:00 2001 From: Phil Blowey Date: Thu, 29 Oct 2026 10:16:05 +0000 Subject: [PATCH 10/14] Remove old dc_sim code --- src/dlstbx/cli/dc_sim.py | 592 ---------------------- src/dlstbx/dc_sim/__init__.py | 755 ---------------------------- src/dlstbx/dc_sim/check.py | 517 ------------------- src/dlstbx/dc_sim/dbserverclient.py | 413 --------------- src/dlstbx/dc_sim/definitions.py | 412 --------------- src/dlstbx/dc_sim/mydb.py | 173 ------- tests/dc_sim/test_dc_sim.py | 279 ---------- tests/dc_sim/test_em_sim.py | 280 ----------- 8 files changed, 3421 deletions(-) delete mode 100644 src/dlstbx/cli/dc_sim.py delete mode 100644 src/dlstbx/dc_sim/__init__.py delete mode 100644 src/dlstbx/dc_sim/check.py delete mode 100644 src/dlstbx/dc_sim/dbserverclient.py delete mode 100644 src/dlstbx/dc_sim/definitions.py delete mode 100644 src/dlstbx/dc_sim/mydb.py delete mode 100644 tests/dc_sim/test_dc_sim.py delete mode 100644 tests/dc_sim/test_em_sim.py diff --git a/src/dlstbx/cli/dc_sim.py b/src/dlstbx/cli/dc_sim.py deleted file mode 100644 index dc26ee358..000000000 --- a/src/dlstbx/cli/dc_sim.py +++ /dev/null @@ -1,592 +0,0 @@ -from __future__ import annotations - -import argparse -import getpass -import logging -import os -import re -import shutil -import socket -from datetime import datetime -from pathlib import Path - -import ispyb.sqlalchemy -import sqlalchemy -import sqlalchemy.orm -import workflows.transport -import zocalo.configuration -from ispyb.sqlalchemy import ( - BLSession, - DataCollection, - DataCollectionGroup, - GridInfo, - Proposal, -) -from sqlalchemy.orm import Load - -from dlstbx.crud import get_dcids_for_data_collection_group -from dlstbx.util import INDUSTRIAL_CODES, DowngradeErrorsFilter -from dlstbx.util.colorstreamhandler import ColorStreamHandler - -logger = logging.getLogger(__name__) - - -def get_dc_and_grid_info_from_dcid( - session: sqlalchemy.orm.Session, dcid: int -) -> tuple[DataCollection, GridInfo | None]: - """Fetch a data collection and its related records in one database query. - - GridInfo is optional and may be associated directly with the data collection - or, for legacy records, with its data collection group. If both forms exist, - the direct association is preferred. BLSample is optional; the - DataCollection value is preferred, falling back to the group for legacy - records. - """ - - result = ( - session.query( - DataCollection, - GridInfo, - ) - .outerjoin( - GridInfo, GridInfo.dataCollectionId == DataCollection.dataCollectionId - ) - .filter(DataCollection.dataCollectionId == dcid) - .one_or_none() - ) - if result is None: - logger.error(f"No data collection found for dcid={dcid}") - raise ValueError(f"Could not find data collection for dcid={dcid}") - - return result - - -def get_data_collection_group_from_dcgid( - db_session: sqlalchemy.orm.Session, dcgid: int -) -> DataCollectionGroup: - dcg = ( - db_session.query(DataCollectionGroup).filter( - DataCollectionGroup.dataCollectionGroupId == dcgid - ) - ).one_or_none() - if dcg is None: - logger.error( - f"No data collection or related ISPyB records found for dcgid={dcgid}" - ) - raise ValueError( - f"Could not find a data collection and its related records for dcgid {dcgid}" - ) - return dcg - - -def get_data_collection_group_from_dcid( - db_session: sqlalchemy.orm.Session, dcid: int -) -> DataCollectionGroup: - dcg = ( - db_session.query(DataCollectionGroup) - .join( - DataCollection, - DataCollection.dataCollectionGroupId - == DataCollectionGroup.dataCollectionGroupId, - ) - .filter(DataCollection.dataCollectionId == dcid) - ).one_or_none() - if dcg is None: - logger.error( - f"No data collection or related ISPyB records found for dcid={dcid}" - ) - raise ValueError( - f"Could not find a data collection and its related records for dcgid={dcid}" - ) - return dcg - - -def get_bl_session_and_proposal_code_from_dcgid( - db_session: sqlalchemy.orm.Session, dcgid: int -) -> tuple[str, BLSession]: - results = ( - db_session.query(Proposal.proposalCode, BLSession) - .join(BLSession, BLSession.proposalId == Proposal.proposalId) - .join(DataCollectionGroup, DataCollectionGroup.sessionId == BLSession.sessionId) - .filter(DataCollectionGroup.dataCollectionGroupId == dcgid) - .one_or_none() - ) - - if results is None: - logger.error(f"No proposal code or BLSession found for dcgid={dcgid}") - raise ValueError(f"Could not find proposal code or BLSession for dcgid={dcgid}") - return results - - -def retrieve_sessionid(db_session: sqlalchemy.orm.Session, visit: str) -> int: - query = ( - db_session.query(BLSession, Proposal) - .options( - Load(BLSession).load_only( - BLSession.sessionId, BLSession.visit_number, BLSession.proposalId - ), - Load(Proposal).load_only( - Proposal.proposalId, Proposal.proposalCode, Proposal.proposalNumber - ), - ) - .join( - Proposal, - Proposal.proposalId == BLSession.proposalId, - ) - .filter( - sqlalchemy.func.concat( - Proposal.proposalCode, - Proposal.proposalNumber, - "-", - BLSession.visit_number, - ) - == visit - ) - ) - query_results = query.first() - if query_results is None: - raise ValueError(f"Query to obtain sessionid failed for {visit}") - if query_results[0].sessionId is None: - raise ValueError(f"Could not find sessionid for visit {visit}") - return query_results[0].sessionId - - -def copy_data_collection_group( - source_dcg: DataCollectionGroup, - new_session_id: int, - new_start_time: datetime, - new_end_time: datetime, -) -> DataCollectionGroup: - """ - Create a copy of a DataCollectionGroup row with different sessionId and times. - - Args: - source_dcg: The DataCollectionGroup to copy - new_session_id: The new sessionId for the copy - new_start_time: The new startTime for the copy - new_end_time: The new endTime for the copy - - Returns: - A new DataCollectionGroup object with all attributes copied except - the ones specified (which are replaced with the new values) - """ - # Get all columns from the source object - new_dcg = DataCollectionGroup() - - # Copy all attributes from source to new, except the primary key and fields to update - for column in DataCollectionGroup.__table__.columns: - attr_name = column.name - - # Skip the primary key - if attr_name == "dataCollectionGroupId": - continue - - # Use new values for these attributes - if attr_name == "sessionId": - setattr(new_dcg, attr_name, new_session_id) - elif attr_name == "startTime": - setattr(new_dcg, attr_name, new_start_time) - elif attr_name == "endTime": - setattr(new_dcg, attr_name, new_end_time) - else: - # Copy from source - setattr(new_dcg, attr_name, getattr(source_dcg, attr_name, None)) - - return new_dcg - - -def copy_data_collection( - source_dc: DataCollection, - new_session_id: int, - new_data_collection_group_id: int, - **kwargs, -) -> DataCollection: - """ - Create a copy of a DataCollection row with different sessionId and/or dataCollectionGroupId. - - Args: - source_dc: The DataCollection to copy - new_session_id: The new SESSIONID for the copy - new_data_collection_group_id: The new dataCollectionGroupId for the copy - **kwargs: Additional attributes to override (e.g., startTime, endTime, comments) - - Returns: - A new DataCollection object with all attributes copied except - the ones specified (which are replaced with the new values) - """ - # Create a new DataCollection object - new_dc = DataCollection() - - # Copy all attributes from source to new, except the primary key and fields to update - for column in DataCollection.__table__.columns: - attr_name = column.name - - # Skip the primary key - if attr_name == "dataCollectionId": - continue - - # Use new values for these attributes - if attr_name == "SESSIONID": - setattr(new_dc, attr_name, new_session_id) - elif attr_name == "dataCollectionGroupId": - setattr(new_dc, attr_name, new_data_collection_group_id) - elif attr_name in kwargs: - # Use provided override values - setattr(new_dc, attr_name, kwargs[attr_name]) - else: - # Copy from source - setattr(new_dc, attr_name, getattr(source_dc, attr_name, None)) - - return new_dc - - -def copy_grid_info( - source_gi: GridInfo, - new_dcid: int, - **kwargs, -) -> DataCollection: - """ - Create a copy of a DataCollection row with different sessionId and/or dataCollectionGroupId. - - Args: - source_dc: The DataCollection to copy - new_session_id: The new SESSIONID for the copy - new_data_collection_group_id: The new dataCollectionGroupId for the copy - **kwargs: Additional attributes to override (e.g., startTime, endTime, comments) - - Returns: - A new DataCollection object with all attributes copied except - the ones specified (which are replaced with the new values) - """ - # Create a new GridInfo object - new_gi = GridInfo() - - # Copy all attributes from source to new, except the primary key and fields to update - for column in GridInfo.__table__.columns: - attr_name = column.name - - # Skip the primary key - if attr_name == "gridInfoId": - continue - - # Use new values for these attributes - if attr_name == "dataCollectionId": - setattr(new_gi, attr_name, new_dcid) - elif attr_name in kwargs: - # Use provided override values - setattr(new_gi, attr_name, kwargs[attr_name]) - else: - # Copy from source - setattr(new_gi, attr_name, getattr(source_gi, attr_name, None)) - return new_gi - - -def simulate(dcid_or_dcgid: int, is_dcg: bool = False) -> list[int]: - """Look up and copy the records needed to simulate data collection.""" - dflt_proposal = "nt44198" - - sessionmaker = sqlalchemy.orm.sessionmaker( - bind=sqlalchemy.create_engine( - ispyb.sqlalchemy.url(), connect_args={"use_pure": True} - ) - ) - with sessionmaker() as session: - if is_dcg: - logger.info( - f"Starting data collection simulation for dcgid={dcid_or_dcgid}" - ) - dcg = get_data_collection_group_from_dcgid(session, dcid_or_dcgid) - dcgid = dcid_or_dcgid - dcids = get_dcids_for_data_collection_group(dcgid, session) - else: - logger.info(f"Starting data collection simulation for dcid={dcid_or_dcgid}") - dcg = get_data_collection_group_from_dcid(session, dcid_or_dcgid) - dcgid = dcg.dataCollectionGroupId - dcids = [dcid_or_dcgid] - - proposal_code, bl_session = get_bl_session_and_proposal_code_from_dcgid( - session, dcgid - ) - - if proposal_code in INDUSTRIAL_CODES: - logger.error("Source data is industrial; simulation is not allowed") - raise SystemExit(1) - - beamline = bl_session.beamLineName - dest_visit_dir, dest_visit = get_beamline_visit_and_dir_for_proposal( - dflt_proposal, beamline - ) - - now = datetime.now() - random_uuid = os.urandom(4).hex() # Generate a random UUID for uniqueness - dest_dir = ( - dest_visit_dir - / "tmp" - / f"{now:%Y-%m-%d}" - / f"{now:%H}-{now:%M}-{now:%S}-{random_uuid}" - ) - dest_dir.mkdir(parents=True, exist_ok=True) - logger.info(f"Created image destination directory {dest_dir}") - # Create destination directory for JPEG snapshots - jpgs_dest_dir = ( - dest_visit_dir - / "jpegs" - / "tmp" - / f"{now:%Y-%m-%d}" - / f"{now:%H}-{now:%M}-{now:%S}-{random_uuid}" - ) - jpgs_dest_dir.mkdir(parents=True, exist_ok=True) - logger.info(f"Created snapshot destination directory {jpgs_dest_dir}") - # Get the session id for the destination visit - logger.info(f"Finding ISPyB session for destination visit {dest_visit}") - new_bl_session_id = retrieve_sessionid(session, dest_visit) - # Create a copy with new sessionId and times - logger.info("Copying data collection group") - new_dcg = copy_data_collection_group( - dcg, - new_session_id=new_bl_session_id, - new_start_time=datetime.now(), - new_end_time=datetime.now(), - ) - - # Add and commit the new DataCollectionGroup - session.add(new_dcg) - session.commit() - new_dcgid = new_dcg.dataCollectionGroupId - logger.info(f"Created destination DataCollectionGroup dcgid={new_dcgid}") - - logger.info("Loading source data collection and related ISPyB records") - new_dcids = [] - for dcid in dcids: - new_dcid = simulate_datacollection( - session, dcid, dest_dir, jpgs_dest_dir, new_bl_session_id, new_dcgid - ) - new_dcids.append(new_dcid) - return new_dcids - - -def simulate_datacollection( - db_session: sqlalchemy.orm.Session, - dcid: int, - dest_dir: Path, - jpgs_dest_dir: Path, - new_bl_session_id: int, - new_dcgid: int, -) -> int: - dc, gi = get_dc_and_grid_info_from_dcid(db_session, dcid) - source_image_dir = Path(dc.imageDirectory) - start_time = datetime.now() - - # Copy the image files to the new destination directory - image_suffix = dc.imageSuffix - if image_suffix == "cbf": - file_template = dc.fileTemplate - match = re.search(r"(#+)(?=\.cbf$)", file_template) - if not match: - raise ValueError( - f"File template {file_template} does not contain a valid '#' pattern for image numbering." - ) - n_digits = len(match.group(1)) - start_image_number = dc.startImageNumber - number_of_images = dc.numberOfImages - copied_images = 0 - missing_images = 0 - logger.info( - f"Copying {number_of_images} CBF images from {source_image_dir} " - f"to {dest_dir}" - ) - for image_number in range( - start_image_number, start_image_number + number_of_images - ): - image_filename = ( - file_template[: match.start()] - + str(image_number).zfill(n_digits) - + file_template[match.end() :] - ) - src_image_path = source_image_dir / image_filename - dest_image_path = dest_dir / image_filename - if src_image_path.exists(): - shutil.copy2(src_image_path, dest_image_path) - copied_images += 1 - else: - logger.warning( - f"Source image {src_image_path} does not exist; skipping copy" - ) - missing_images += 1 - logger.info( - f"Finished copying CBF images: copied={copied_images}, " - f"missing={missing_images}" - ) - - elif image_suffix in ["h5", "nxs"]: - file_prefix = f"{dc.imagePrefix}_{dc.dataCollectionNumber}" - files = { - path - for pattern in (f"{file_prefix}_*.h5", f"{file_prefix}.nxs") - for path in source_image_dir.glob(pattern) - } - logger.info( - f"Copying {len(files)} HDF5/Nexus file(s) from {source_image_dir} " - f"to {dest_dir}" - ) - for file in files: - shutil.copy2(file, dest_dir) - logger.info("Finished copying image data files") - else: - logger.error( - f"Unsupported image suffix {image_suffix!r} for source dcid={dcid}" - ) - raise ValueError( - f"Invalid image suffix {image_suffix} for source data collection" - ) - - # Create new DataCollection - new_dc = copy_data_collection( - dc, - new_session_id=new_bl_session_id, - new_data_collection_group_id=new_dcgid, - startTime=start_time, - endTime=datetime.now(), - comments=f"Simulated data collection copied from dcid {dcid}", - imageDirectory=str(dest_dir) + "/", - ) - db_session.add(new_dc) - db_session.commit() - new_dcid = new_dc.dataCollectionId - - if gi is not None: - new_gi = copy_grid_info(source_gi=gi, new_dcid=new_dcid) - db_session.add(new_gi) - db_session.commit() - - # Mapping of snapshot attributes to their names - snapshot_attrs = [ - "xtalSnapshotFullPath1", - "xtalSnapshotFullPath2", - "xtalSnapshotFullPath3", - "xtalSnapshotFullPath4", - ] - - # Copy and update snapshot paths - copied_snapshots = 0 - missing_snapshots = 0 - for attr_name in snapshot_attrs: - src_snapshot_path = getattr(dc, attr_name, None) - if src_snapshot_path is not None: - src_snapshot = Path(src_snapshot_path) - if src_snapshot.exists(): - # Create a unique filename for the snapshot - dest_snapshot_path = jpgs_dest_dir / src_snapshot.name - shutil.copy2(src_snapshot, dest_snapshot_path) - # Update the new DataCollection with the new snapshot path - setattr(new_dc, attr_name, str(dest_snapshot_path)) - copied_snapshots += 1 - else: - logger.warning( - f"Crystal snapshot {src_snapshot} does not exist; skipping copy" - ) - missing_snapshots += 1 - - logger.info( - f"Created destination DataCollection {new_dc.dataCollectionId}; " - f"copied snapshots={copied_snapshots}, missing snapshots={missing_snapshots}" - ) - logger.info( - f"Data collection simulation for dcid={dcid} completed in " - f"{(datetime.now() - start_time).total_seconds():.1f} seconds" - ) - return new_dcid - - -def get_beamline_visit_and_dir_for_proposal( - proposal: str, beamline: str -) -> tuple[Path, str]: - data_root = Path("/dls", beamline, "data", str(datetime.now().year)) - logger.info(f"Searching for proposal {proposal} under {data_root}") - for cm_dir in data_root.iterdir(): - if cm_dir.name.startswith(proposal): - dest_visit_dir = cm_dir - dest_visit = cm_dir.name - break - else: - raise ValueError( - f"Could not determine destination directory for proposal {proposal}" - ) - logger.info(f"Using destination visit directory {dest_visit_dir}") - return dest_visit_dir, dest_visit - - -def run(args: list[str] | None = None) -> None: - # Enable logging to console - console = ColorStreamHandler() - console.setLevel(logging.INFO) - logging.getLogger("dlstbx").setLevel(logging.INFO) - for pika_logger in ( - "base_connection", - "blocking_connection", - "utils.connection_workflow", - "utils.io_services_utils", - ): - logging.getLogger(f"pika.adapters.{pika_logger}").addFilter( - DowngradeErrorsFilter() - ) - logging.getLogger("workflows").setLevel(logging.INFO) - logging.getLogger("zocalo").setLevel(logging.INFO) - logging.getLogger().setLevel(logging.WARNING) - logging.getLogger().addHandler(console) - zc = zocalo.configuration.from_file() - zc.activate() - parser = argparse.ArgumentParser( - description="Simulate data collection. Takes a dcid as input and copies data and database records to simulate a data collection, then triggers mimas.", - usage="dlstbx.dc_sim [options] dcid", - ) - parser.add_argument( - "dcid_or_dcgid", - type=int, - help="Data collection ID or Data collection group ID - To use Data collection group id, set --dcgid option", - ) - parser.add_argument("--is-dcgid", action="store_true") - parser.add_argument("--suppress-mimas", action="store_true") - zc.add_command_line_options(parser) - workflows.transport.add_command_line_options(parser, transport_argument=True) - - parsed_args = parser.parse_args(args) - new_dcids = simulate(parsed_args.dcid_or_dcgid, parsed_args.is_dcgid) - - if not parsed_args.suppress_mimas: - transport = workflows.transport.lookup(parsed_args.transport)() - transport.connect() - for new_dcid in new_dcids: - logger.info(f"Triggering Mimas for dcid={new_dcid}") - headers = { - "zocalo.go.user": getpass.getuser(), - "zocalo.go.host": socket.gethostname(), - } - - transport.send( - "processing_recipe", - { - "recipes": ["mimas"], - "parameters": { - "ispyb_dcid": new_dcid, - "event": "start", - }, - }, - headers=headers, - ) - transport.send( - "processing_recipe", - { - "recipes": ["mimas"], - "parameters": { - "ispyb_dcid": new_dcid, - "event": "end", - }, - }, - headers=headers, - ) - transport.disconnect() - - -if __name__ == "__main__": - run() diff --git a/src/dlstbx/dc_sim/__init__.py b/src/dlstbx/dc_sim/__init__.py deleted file mode 100644 index 479ae7740..000000000 --- a/src/dlstbx/dc_sim/__init__.py +++ /dev/null @@ -1,755 +0,0 @@ -# Functions to simulate a data collection -# -# This -# * inserts new entries into the datacollection table using the DbserverClient.py script -# * copies images from the source data collection -# * runs the scripts RunAtStartOfDataCollection.sh and RunAtEndOfDataCollection.sh -# at appropriate times. - -from __future__ import annotations - -import collections -import errno -import glob -import logging -import os -import pathlib -import re -import shutil -import sys -import time -import uuid -from datetime import datetime -from pathlib import Path - -import ispyb.sqlalchemy -import procrunner -import sqlalchemy -from workflows.transport.stomp_transport import StompTransport - -import dlstbx.dc_sim.dbserverclient -import dlstbx.dc_sim.definitions -import dlstbx.dc_sim.mydb as db - -log = logging.getLogger("dlstbx.dc_sim") - -# Constants -MX_SCRIPTS_BINDIR = "/dls_sw/apps/mx-scripts/bin" - -SimulationResult = collections.namedtuple( - "SimulationResult", - [ - "DCIDs", - "JobIDs", - "URLs", - "beamline", - "scenario", - "time_end", - "time_start", - "type", - ], -) - - -def mkdir_p(path): - try: - os.makedirs(path) - except OSError as exc: # Python >2.5 - if exc.errno == errno.EEXIST and os.path.isdir(path): - pass - else: - raise - - -def copy_via_temp_file(source, destination): - dest_dir, dest_file = os.path.split(destination) - temp_dest_file = ".tmp." + dest_file - temp_destination = os.path.join(dest_dir, temp_dest_file) - shutil.copyfile(source, temp_destination) - os.rename(temp_destination, destination) - - -def retrieve_grid_info_values(_db, _src_dcid): - _db.cursor.execute( - "SELECT dx_mm, dy_mm, steps_x, steps_y, micronsperpixelx, micronsperpixely, " - "snapshot_offsetxpixel, snapshot_offsetypixel, patchesx, patchesy, orientation, snaked " - "FROM GridInfo " - "WHERE datacollectionid=%d" % _src_dcid - ) - - desc = [d[0] for d in _db.cursor.description] - result = [dict(zip(desc, line)) for line in _db.cursor] - - if len(result) == 0: - return None - return result[0] - - -def retrieve_blsample_values(_db, _src_blsampleid): - _db.cursor.execute( - "SELECT blsampleid, name, code, location, holderlength, looplength, looptype, wirewidth, comments, " - "blsamplestatus, lastknowncenteringposition " - "FROM BLSample " - "WHERE blsampleid=%d " % _src_blsampleid - ) - - desc = [d[0] for d in _db.cursor.description] - result = [dict(zip(desc, line)) for line in _db.cursor] - - if not result[0].get("blsampleid"): - sys.exit(f"Could not find the blsampleid for {_src_blsampleid}") - - return result[0] - - -def retrieve_max_dcnumber(_db, _sessionid, _dest_dir, _dest_prefix): - rows = _db.doQuery( - "SELECT max(datacollectionnumber) " - "FROM DataCollection " - "WHERE sessionid=%d " - "AND imagedirectory='%s' " - "AND imageprefix='%s'" % (_sessionid, _dest_dir + "/", _dest_prefix) - ) - return rows[0][0] - - -def _simulate( - _dest_visit, - _beamline, - _src_dir, - _src_visit, - _src_prefix, - _src_run_number, - _dest_prefix, - _dest_visit_dir, - _dest_dir, - _sample_id, - proc_params, - db_session, - data_collection_group_id, - scenario_name, - scenario_type, -): - _db = db.DB() - dbsc = dlstbx.dc_sim.dbserverclient.DbserverClient() - ispyb.sqlalchemy.enable_debug_logging() - - log.debug("Getting the source SessionID") - src_sessionid = db.retrieve_sessionid(db_session, _src_visit) - log.debug(f"Source SessionID is {src_sessionid}") - - try: - run_number = int(_src_run_number) - except (TypeError, ValueError): - run_number = None - - row = db.retrieve_datacollection( - db_session, src_sessionid, _src_dir, _src_prefix, run_number - ) - _src_dcid = row.dataCollectionId - _src_dcgid = row.dataCollectionGroupId - start_img_number = row.startImageNumber - filetemplate = row.fileTemplate - src_xtal_snapshot_path = [ - row.xtalSnapshotFullPath1, - row.xtalSnapshotFullPath2, - row.xtalSnapshotFullPath3, - row.xtalSnapshotFullPath4, - ] - log.debug( - f"Source dataset from DCID {_src_dcid}, DCGID {_src_dcgid}, file template {filetemplate}" - ) - - if scenario_type == "em-spa": - # start copying over data files - log.info( - f"Copying first 5 files from {_src_dir} to {pathlib.Path(_dest_dir) / 'raw'}" - ) - - data_dirs = [f for f in pathlib.Path(_src_dir).glob("**/*") if f.is_dir()] - data_files = [f for f in pathlib.Path(_src_dir).glob("**/*") if f.is_file()] - for dd in data_dirs: - mkdir_p(pathlib.Path(_dest_dir) / "raw" / dd.relative_to(_src_dir)) - for df in data_files[:5]: - copy_via_temp_file( - df, pathlib.Path(_dest_dir) / "raw" / df.relative_to(_src_dir) - ) - - i = ispyb.open() - - if data_collection_group_id is None: - # Get the sessionid for the dest_visit - log.debug("(SQL) Getting the destination sessionid") - sessionid = db.retrieve_sessionid(db_session, _dest_visit) - - dcgparams = i.mx_acquisition.get_data_collection_group_params() - dcgparams["parentid"] = sessionid - dcgparams["experimenttype"] = "EM" - dcgparams["comments"] = "Created for simulated data collection" - datacollectiongroupid = i.mx_acquisition.upsert_data_collection_group( - list(dcgparams.values()) - ) - dcparams = i.mx_acquisition.get_data_collection_params() - key_maps = { - "runStatus": "run_status", - "imageSuffix": "imgsuffix", - "fileTemplate": "file_template", - "comments": "comments", - } - for attr, key in key_maps.items(): - dcparams[key] = getattr(row, attr) - dcparams["parentid"] = datacollectiongroupid - dcparams["imgdir"] = str(pathlib.Path(_dest_dir) / "raw") - dcparams["visitid"] = sessionid - datacollectionid = i.mx_acquisition.upsert_data_collection( - list(dcparams.values()) - ) - else: - datacollectiongroupid = data_collection_group_id - datacollectionid = _src_dcid - - log.debug( - "Source dataset from DCID %r, DCGID %r", - _src_dcid, - _src_dcgid, - ) - - # nowstr = datetime.now().strftime("%Y-%m-%d %H:%M:%S") - - proc_job_values = i.mx_processing.get_job_params() - proc_job_values["datacollectionid"] = datacollectionid - proc_job_values["display_name"] = "RELION" - proc_job_values["comments"] = "Submitted as part of simulated data collection" - proc_job_values["recipe"] = "relion4" - proc_job_values["automatic"] = 0 - procjobid = i.mx_processing.upsert_job(list(proc_job_values.values())) - - for k, v in proc_params.items(): - if k != "import_images": - job_param_values = (None, procjobid, k, v) - else: - job_param_values = ( - None, - procjobid, - k, - _dest_dir + "/raw/Frames/*.tiff", - ) - i.mx_processing.upsert_job_parameter(job_param_values) - - default_configuration = "/dls_sw/apps/zocalo/secrets/credentials-live.cfg" - stomp = StompTransport() - stomp.load_configuration_file(default_configuration) - stomp.connect() - - dispatcher_message = {"parameters": {"ispyb_process": procjobid}} - - stomp.send("processing_recipe", dispatcher_message) - - num_data_file_blocks = len(data_files) // 5 - delay = 5 * 60 - for i in range(1, num_data_file_blocks): - log.info( - f"Waiting {delay} seconds (until {datetime.fromtimestamp(time.time() + delay):%H:%M:%S})" - ) - time.sleep(delay) - log.info( - f"Copying further 5 files from {_src_dir} to {pathlib.Path(_dest_dir) / 'raw'}" - ) - for df in data_files[i * 5 : (i + 1) * 5]: - log.debug(df) - copy_via_temp_file( - df, pathlib.Path(_dest_dir) / "raw" / df.relative_to(_src_dir) - ) - log.info( - f"Copying remaining files from {_src_dir} to {pathlib.Path(_dest_dir) / 'raw'}" - ) - for df in data_files[num_data_file_blocks * 5 :]: - log.debug(df) - copy_via_temp_file( - df, pathlib.Path(_dest_dir) / "raw" / df.relative_to(_src_dir) - ) - - return datacollectionid, datacollectiongroupid, procjobid - - if scenario_type == "mx": - if start_img_number is None: - sys.exit("Could not find the first image number for data collection") - no_images = row.numberOfImages - if not no_images: - sys.exit("Could not find the number of images for data collection") - log.debug(f"Source dataset has {no_images} images") - - # Get the sessionid for the dest_visit - log.debug("(SQL) Getting the destination sessionid") - sessionid = db.retrieve_sessionid(db_session, _dest_visit) - - # Get the highest run number for the datacollections of this dest_visit with the particular img.dir and prefix - log.debug( - "(SQL) Getting the currently highest run number for this img. directory + prefix" - ) - # Can't change the run number for .h5 files otherwise the link from the master.h5 to data_*.h5 will be incorrect - if not filetemplate.endswith(".h5"): - run_number = retrieve_max_dcnumber(_db, sessionid, _dest_dir, _dest_prefix) - try: - run_number = int(run_number) + 1 - except (TypeError, ValueError): - run_number = 1 - - log.debug("(SQL) Getting values from the source datacollectiongroup record") - src_blsampleid = row.DataCollectionGroup.blSampleId - - log.debug( - "(filesystem) Copy the xtal snapshot(s) (if any) from source to target directories" - ) - dest_xtal_snapshot_path = ["", "", "", ""] - for x in range(0, 4): - if src_xtal_snapshot_path[x] is not None: - if os.path.exists(src_xtal_snapshot_path[x]): - png = re.sub( - "^.*/(.*)$", _dest_dir + r"/\1", src_xtal_snapshot_path[x] - ) - dest_xtal_snapshot_path[x] = re.sub( - "^" + _dest_visit_dir, _dest_visit_dir + "/jpegs", png - ) - path = os.path.dirname(dest_xtal_snapshot_path[x]) - log.debug("(filesystem) ... 'mkdir -p' %s" % path) - os.makedirs(path, exist_ok=True) - log.debug( - "(filesystem) ... copying %s to %s" - % (src_xtal_snapshot_path[x], dest_xtal_snapshot_path[x]) - ) - copy_via_temp_file( - src_xtal_snapshot_path[x], dest_xtal_snapshot_path[x] - ) - - log.info(f"539 src_blsampleid: {src_blsampleid}") - log.info("539 _sample_id: %s", _sample_id) - # Get a blsampleId either from a copy of the blsample used by the src dc or use the blsampleId provided on the command-line - blsample_id = None - if src_blsampleid is not None: - if _sample_id is None: - log.debug("(SQL) Getting values from the source blsample record") - bls_row = retrieve_blsample_values(_db, src_blsampleid) - - blsample_xml = ( - dlstbx.dc_sim.dbserverclient.populate_blsample_xml_template(bls_row) - ) - print(blsample_xml) - - # Ingest the blsample data using the DbserverClient - log.debug("(dbserver) Ingest the blsample XML") - blsample_id = dbsc.storeBLSample(blsample_xml) - else: - blsample_id = _sample_id - elif _sample_id is not None: - blsample_id = _sample_id - - if data_collection_group_id is None: - # Produce a DataCollectionGroup xml blob from the template - dcg_xml = dlstbx.dc_sim.dbserverclient.populate_dcg_xml_template( - row, sessionid, blsample_id - ) - - # Ingest the DataCollectionGroup xml data using the DbserverClient - log.debug("(dbserver) Ingest the datacollectiongroup XML") - datacollectiongroupid = dbsc.storeDataCollectionGroup(dcg_xml) - else: - datacollectiongroupid = data_collection_group_id - - # Produce a DataCollection xml blob from the template and use the new run number - row_as_dictionary = { - name.lower(): getattr(row, name) - for name in dir(row) - if not name.startswith("_") - } - dc_xml = dlstbx.dc_sim.dbserverclient.populate_dc_xml_template( - row_as_dictionary, - sessionid, - datacollectiongroupid, - no_images, - _dest_dir + "/", - _dest_prefix, - run_number, - dest_xtal_snapshot_path, - blsample_id, - scenario_name=scenario_name, - ) - - # Ingest the DataCollection xml blob data using the DbserverClient - log.debug("(dbserver) Ingest the datacollection XML") - datacollectionid = dbsc.storeDataCollection(dc_xml) - - # Get the grid info values associated with the source dcid - gi_row = retrieve_grid_info_values(_db, _src_dcid) - - # Prouce a GridInfo xml blob from the template if the source DataCollectionGroup has one: - if gi_row is not None: - gridinfo_xml = dlstbx.dc_sim.dbserverclient.populate_grid_info_xml_template( - gi_row, datacollectionid - ) - - # Ingest the GridInfo.xml file data using the DbserverClient - log.debug("(dbserver) Ingest the gridinfo XML") - dbsc.storeGridInfo(gridinfo_xml) - - run_at_params = [ - "automaticProcessing_Yes", - str(datacollectionid), - _dest_visit_dir, - filetemplate, - _dest_dir + "/", - f"{_dest_prefix}_{run_number}_", - os.path.splitext(filetemplate)[-1], - ] - - command = [f"{MX_SCRIPTS_BINDIR}/RunAtStartOfCollect-{_beamline}.sh"] - command.extend(run_at_params) - log.info("command: %s", " ".join(command)) - result = procrunner.run(command, timeout=180) - log.info("runtime: %s", result["runtime"]) - if result["exitcode"] or result["timeout"]: - log.info("timeout: %s", result["timeout"]) - log.debug(result["stdout"]) - log.debug(result["stderr"]) - log.error( - "RunAtStartOfCollect failed with exit code %d", result["exitcode"] - ) - - if filetemplate.endswith(".cbf"): - # Also copy images one by one from source to destination directory. - for x in range(start_img_number, start_img_number + no_images): - num_digits = len(filetemplate.split(".")[0].split("_")[-1]) - _src_img_number = f"%0{num_digits}d" % x - img_number = "%04d" % x - src_prefix = "" - if _src_prefix is not None: - src_prefix = _src_prefix - src_fname = "%s_%s_%s.cbf" % ( - src_prefix, - str(_src_run_number), - str(_src_img_number), - ) - dest_fname = "%s_%d_%s.cbf" % ( - _dest_prefix, - run_number, - str(img_number), - ) - src = os.path.join(_src_dir, src_fname) - target = os.path.join(_dest_dir, dest_fname) - log.info(f"(filesystem) Copy file {src} to {target}") - copy_via_temp_file(src, target) - elif filetemplate.endswith(".h5"): - files = [] - src_prefix = "" - if _src_prefix is not None: - src_prefix = _src_prefix - for ext in ("_*.h5", ".nxs", "_meta.hdf5"): - files.extend( - glob.glob( - os.path.join( - _src_dir, filetemplate.split("_master.h5")[0] + ext - ) - ) - ) - for src in files: - dest_fname = os.path.basename(src).replace( - "%s_%s" % (src_prefix, str(_src_run_number)), - "%s_%d" % (_dest_prefix, run_number), - ) - target = os.path.join(_dest_dir, dest_fname) - log.info(f"(filesystem) Copy file {src} to {target}") - copy_via_temp_file(src, target) - else: - raise RuntimeError("Unsupported file extension for %s" % filetemplate) - - # Populate a datacollection XML blob - nowstr = datetime.now().strftime("%Y-%m-%d %H:%M:%S") - - dc_xml = dlstbx.dc_sim.dbserverclient.dc_endtime_temp_xml % ( - datacollectionid, - nowstr, - ) - print(dc_xml) - log.debug( - "(dbserver) Ingest the datacollection XML to update with the d.c. end time" - ) - dbsc.updateDbObject(dc_xml) - - # Populate a datacollectiongroup XML blob - nowstr = datetime.now().strftime("%Y-%m-%d %H:%M:%S") - - dcg_xml = dlstbx.dc_sim.dbserverclient.dcg_endtime_temp_xml % ( - datacollectiongroupid, - nowstr, - ) - print(dcg_xml) - - # Ingest the DataCollectionGroup xml blob using the DbserverClient - log.debug( - "(dbserver) Ingest the datacollectiongroup XML to update with the d.c.g. end time" - ) - dbsc.updateDbObject(dcg_xml) - - command = [f"{MX_SCRIPTS_BINDIR}/RunAtEndOfCollect-{_beamline}.sh"] - command.extend(run_at_params) - log.info("command: %s", " ".join(command)) - result = procrunner.run(command, timeout=180) - log.info("runtime: %s", result["runtime"]) - if result["exitcode"] or result["timeout"]: - log.info("timeout: %s", result["timeout"]) - log.debug(result["stdout"]) - log.debug(result["stderr"]) - log.error("RunAtEndOfCollect failed with exit code %d", result["exitcode"]) - - return datacollectionid, datacollectiongroupid, None - - raise ValueError("Unknown scenario type %s" % scenario_type) - - -def call_sim( - test_name, - beamline, - src_dir=None, - src_prefix=None, - src_run_num=None, - sample_id=None, - dest_visit=None, - dflt_proposals=None, - dflt_i02_visits=None, - src_dcid=None, - src_allowed_visits=None, - is_dcg=None, - src_dcg=None, -): - scenario = dlstbx.dc_sim.definitions.tests.get(test_name) - if scenario is None: - log.info( - f"{test_name} is not a defined test scenario - attempting to use custom data" - ) - scenario = {"type": "mx", "is_dcg": is_dcg} - - # Create database session - url = ispyb.sqlalchemy.url() - engine = sqlalchemy.create_engine(url, connect_args={"use_pure": True}) - db_session = sqlalchemy.orm.sessionmaker(bind=engine)() - - # Check for values specified twice - if scenario.get("src_dir") and (src_dcid or src_dcg): - for item, inp_value in [ - ("src_dcid", src_dcid), - ("src_dcg", src_dcg), - ]: - if inp_value: - log.warning( - f"{item}: {inp_value} provided alongside non-custom scenario {test_name} - scenario data will be used" - ) - else: - for ref_key, inp_value in [ - ("src_dir", src_dir), - ("src_prefix", src_prefix), - ("src_run_num", src_run_num), - ("sample_id", sample_id), - ]: - if scenario.get(ref_key) and inp_value: - log.warning( - f"{ref_key} read from scenario but also specified in command line - using scenario value" - ) - elif src_dcg and inp_value: - log.warning( - f"{ref_key} Specified at command line but dcg: {src_dcg} also provided - using dcg data" - ) - elif src_dcid and inp_value: - log.warning( - f"{ref_key} Specified at command line but dcid: {src_dcid} also provided - using dcid data" - ) - # Get parameters from datacollection group ID if supplied - if src_dcg is not None: - if src_dcid is not None: - log.warning( - f"dcid: {src_dcid} supplied alongside dcg: {src_dcg} - using dcg only" - ) - log.info(f"Getting source data from dcg: {src_dcg}") - src_dir, src_prefix, src_run_num, sample_id = db.retrieve_dcs_from_dcg( - db_session, src_dcg - ) - # Get parameters from datacollection ID if supplied - elif src_dcid is not None: - log.info(f"Getting source data from dcid: {src_dcid}") - # Lookup database entry for dcid - row = db.retrieve_dc_from_dcid(db_session, src_dcid) - # Set parameters from database entry - src_dir = row.imageDirectory - src_prefix = [row.imagePrefix] - src_run_num = [row.dataCollectionNumber] - sample_id = row.BLSAMPLEID - log.info( - f"Source file path = {src_dir}, prefix = {src_prefix[0]}, run number = {src_run_num[0]}, sample id = {sample_id}" - ) - - # Read in values from the scenario if present, otherwise use command line values - try: - src_dir = Path(scenario.get("src_dir", src_dir)) - except TypeError: - raise ValueError("src_dir source data path not specified") - if not (src_prefix := scenario.get("src_prefix", src_prefix)): - log.warning("src_prefix not specified") - if not (src_run_num := scenario.get("src_run_num", src_run_num)): - log.warning("src_run_num not specified") - try: - sample_id = int(scenario.get("sample_id", sample_id)) - except TypeError: - log.warning("sample_id value not specified") - sample_id = None - proc_params = scenario.get("proc_params") - time_start = time.time() - now = datetime.now() - - # Calculate the destination directory from specified visit number - if dest_visit is not None: - # Initial check to ensure that specified visit is either in-house or commissioning - proposal = dest_visit.split("-")[0] - if beamline.startswith("i02"): - dest_visit_dir = Path("/dls/mx/data", proposal, dest_visit) - else: - dest_visit_dir = Path("/dls", beamline, "data", str(now.year), dest_visit) - assert dest_visit_dir.is_dir(), ( - f"Could not find {dest_visit_dir} directory for the specified visit and beamline." - ) - # Else, calculate the destination directory for default proposal numbers - else: - # Get default proposals if a visit is not specified - if beamline.startswith(("e", "m")): - proposal = dflt_proposals["em"] - else: - proposal = dflt_proposals["mx"] - if beamline.startswith("i02"): - try: - dest_visit = f"{proposal}-{dflt_i02_visits[beamline]}" - except KeyError: - log.error("Could not determine which destination visit to use") - raise - dest_visit_dir = Path("/dls/mx/data", proposal, dest_visit) - elif scenario.get("visit_num"): - dest_visit = f"{proposal}-{scenario['visit_num']}" - dest_visit_dir = Path("/dls", beamline, "data", str(now.year), dest_visit) - else: - for cm_dir in Path("/dls", beamline, "data", str(now.year)).iterdir(): - if cm_dir.name.startswith(proposal): - dest_visit = cm_dir.name - break - else: - raise ValueError( - f"Could not determine destination directory for proposal {proposal}" - ) - - # Set mandatory parameters - dest_visit_dir = Path("/dls", beamline, "data", str(now.year), dest_visit) - - random_uuid = str(uuid.uuid4())[:8] - dest_dir = ( - dest_visit_dir - / "tmp" - / f"{now:%Y-%m-%d}" - / f"{now:%H}-{now:%M}-{now:%S}-{random_uuid}" - ) - - if not src_dir.is_dir(): - sys.exit( - "ERROR: The src_dir parameter does not appear to contain a valid directory." - ) - - # Extract necessary info from the source directory path - m1 = re.search(r"(/dls/(\S+?)/data/\d+/)(\S+)", str(src_dir)) - if m1: - subdir = m1.groups()[2] - m2 = re.search(r"^(\S+?)/", subdir) - if m2: - src_visit = m2.groups()[0] - elif subdir: - src_visit = subdir - else: - m1 = re.search(r"(/dls/mx/data/)(\S+)", str(src_dir)) - if m1: - subdir = m1.groups()[1] - src_visit = subdir.split(os.sep)[1] - - if src_visit is None: - sys.exit( - "ERROR: The src_dir parameter does not appear to contain a valid visit directory." - ) - - # Compare to src_allowed_visits if src_dir not from scenario - if "src_dir" not in scenario and not src_visit.startswith( - tuple(src_allowed_visits) - ): - raise ValueError(f"Supplied src_dir from forbidden visit: {src_visit}") - - if scenario["type"] == "mx": - start_script = f"{MX_SCRIPTS_BINDIR}/RunAtStartOfCollect-{beamline}.sh" - if not os.path.exists(start_script): - log.error(f"The file {start_script} was not found.") - sys.exit(1) - end_script = f"{MX_SCRIPTS_BINDIR}/RunAtEndOfCollect-{beamline}.sh" - if not os.path.exists(end_script): - log.error(f"The file {end_script} was not found.") - sys.exit(1) - - # Create destination directory - log.debug(f"Creating directory {dest_dir}") - try: - dest_dir.mkdir(parents=True, exist_ok=True) - log.info(f"Directory {dest_dir} created successfully.") - except Exception: - sys.exit(f"ERROR: Creating directory {dest_dir} failed.") - - # Call _simulate - dcid_list = [] - dcg_list = [] - jobid_list = [] - for src_run_number in src_run_num: - for src_prefix_item in src_prefix: - dest_prefix = src_prefix_item - if scenario.get("is_dcg") and len(dcg_list): - dcg = dcg_list[0] - else: - dcg = None - dcid, dcg, jobid = _simulate( - dest_visit, - beamline, - str(src_dir), - src_visit, - src_prefix_item, - src_run_number, - dest_prefix, - str(dest_visit_dir), - str(dest_dir), - sample_id, - proc_params, - db_session, - data_collection_group_id=dcg, - scenario_name=test_name, - scenario_type=scenario["type"], - ) - jobid_list.append(jobid) - dcid_list.append(dcid) - dcg_list.append(dcg) - if scenario.get("delay"): - log.info(f"Sleeping for {scenario['delay']} seconds") - time.sleep(scenario["delay"]) - - if not dcid_list: - return None - - return SimulationResult( - DCIDs=dcid_list, - JobIDs=jobid_list, - time_start=time_start, - time_end=time.time(), - type=scenario["type"], - beamline=beamline, - scenario=test_name, - URLs=[ - f"https://ispyb.diamond.ac.uk/dc/visit/{dest_visit}/id/{dcid}" - for dcid in dcid_list - ], - ) diff --git a/src/dlstbx/dc_sim/check.py b/src/dlstbx/dc_sim/check.py deleted file mode 100644 index 9f2b78d12..000000000 --- a/src/dlstbx/dc_sim/check.py +++ /dev/null @@ -1,517 +0,0 @@ -from __future__ import annotations - -import logging -from collections import defaultdict -from operator import attrgetter - -import ispyb -import ispyb.sqlalchemy -import sqlalchemy -from ispyb.sqlalchemy import ( - CTF, - AutoProcIntegration, - AutoProcProgram, - DataCollection, - ImageQualityIndicators, - MotionCorrection, - ParticleClassification, - ParticleClassificationGroup, - ParticlePicker, - ProcessingJob, - RelativeIceThickness, - Screening, - ScreeningOutput, - ScreeningOutputLattice, -) -from sqlalchemy.orm import Load - -import dlstbx.dc_sim.definitions as df - -logger = logging.getLogger("dlxtbx.dc_sim.check") - - -def check_test_outcome(test, db_session=None): - expected_outcome = df.tests.get(test["scenario"], {}).get("results", {}) - if expected_outcome == {}: - print(f"Scenario {test['scenario']} is happy with any outcome.") - test["success"] = True - return - - check_functions = { - "mx": _check_mx_outcome, - "em-spa": _check_relion_outcome, - } - if ( - scenario_type := df.tests.get(test["scenario"], {}).get("type", "mx") - ) not in check_functions: - print(f"Skipping unknown test scenario type {scenario_type}") - return - - check_function = check_functions[scenario_type] - - try: - overall, failed_tests = check_function(test, expected_outcome, db_session) - except Exception as e: - print(f"Test validation failed with {e}, leaving test alone") - logger.error( - f"dc_sim validation for {scenario_type} failed with {e}", exc_info=True - ) - return test - - if failed_tests: - test["success"] = False - test["reason"] = "\n".join(failed_tests) - - if overall and all(overall.values()): - test["success"] = True - - print(test) - return test - - -def _retrieve_mx_processing(dcid, session): - image_count = ( - session.query(sqlalchemy.func.count(1)) - .filter( - ImageQualityIndicators.dataCollectionId == DataCollection.dataCollectionId - ) - .group_by(ImageQualityIndicators.dataCollectionId) - .label("image_count") - ) - - query = ( - session.query( - DataCollection, - AutoProcProgram, - AutoProcIntegration, - ScreeningOutput, - ScreeningOutputLattice, - image_count, - ) - .outerjoin( - ProcessingJob, - ProcessingJob.dataCollectionId == DataCollection.dataCollectionId, - ) - .outerjoin( - AutoProcProgram, - AutoProcProgram.processingJobId == ProcessingJob.processingJobId, - ) - .outerjoin( - AutoProcIntegration, - AutoProcIntegration.autoProcProgramId == AutoProcProgram.autoProcProgramId, - ) - .outerjoin( - Screening, DataCollection.dataCollectionId == Screening.dataCollectionId - ) - .outerjoin( - ScreeningOutput, ScreeningOutput.screeningId == Screening.screeningId - ) - .outerjoin( - ScreeningOutputLattice, - ScreeningOutputLattice.screeningOutputId - == ScreeningOutput.screeningOutputId, - ) - .filter(DataCollection.dataCollectionId == dcid) - .options( - Load(DataCollection).load_only( - "dataCollectionId", - "numberOfImages", - ), - Load(AutoProcProgram).load_only("processingPrograms"), - Load(AutoProcIntegration).load_only( - "autoProcIntegrationId", - "cell_a", - "cell_b", - "cell_c", - "cell_alpha", - "cell_beta", - "cell_gamma", - ), - Load(ScreeningOutput).load_only("screeningId", "program"), - Load(ScreeningOutputLattice).load_only( - "spaceGroup", - "unitCell_a", - "unitCell_b", - "unitCell_c", - "unitCell_alpha", - "unitCell_beta", - "unitCell_gamma", - ), - ) - ) - query_results = query.all() - - return list(query_results) - - -def _check_mx_outcome(test, expected_outcome, session): - failed_tests = [] - outcomes = {} - overall = {} - for dcid in test.get("DCIDs", []): - rows = _retrieve_mx_processing(dcid, session) - for ( - data_collection, - autoprocprogram, - integration, - screening_output, - screening_output_lattice, - iqi_count, - ) in rows: - if getattr(screening_output, "screeningId", None): - check_screening_outcomes( - dcid, - screening_output.program, - screening_output_lattice, - outcomes, - expected_outcome, - ) - elif getattr(integration, "autoProcIntegrationId", None): - check_integration_outcomes( - dcid, - autoprocprogram.processingPrograms, - integration, - outcomes, - expected_outcome, - ) - - if getattr(data_collection, "numberOfImages", None): - outcomes.update( - check_pia_outcomes( - dcid, - data_collection.numberOfImages, - iqi_count, - outcomes, - expected_outcome, - ) - ) - - for program in expected_outcome.get("required", []): - if program not in outcomes or outcomes[program]["success"] is None: - outcomes[program] = { - "success": False, - "reason": [ - "Expected result not present for program {program}, DCID:{dcid}".format( - program=program, dcid=dcid - ) - ], - } - - for program in outcomes: - overall.setdefault(program, True) - if outcomes[program]["success"] is False: - overall[program] = False - failed_tests.extend(outcomes[program]["reason"]) - elif outcomes[program]["success"] is None and overall[program] is not False: - overall[program] = None - return overall, failed_tests - - -def _check_relion_outcome(test, expected_outcome, db): - failed_tests = [] - overall = {} - for jobid in test.get("JobIDs", []): - program = "relion" - motioncorr_data, autoprocpid = _retrieve_motioncorr(db, jobid) - job_results = { - "motion_correction": motioncorr_data, - "ctf": _retrieve_ctf(db, autoprocpid), - "relative_ice_thickness": _retrieve_relative_ice_thickness(db, autoprocpid), - "particle_picker": _retrieve_particle_picker(db, autoprocpid), - "particle_classification": _retrieve_particle_classification( - db, autoprocpid - ), - } - # print(job_results) - if ( - len(job_results["motion_correction"]) == 0 - or len(job_results["ctf"]) == 0 - or len(job_results["relative_ice_thickness"]) == 0 - or len(job_results["particle_picker"]) == 0 - or len(job_results["particle_classification"]) != 50 - ): - overall[program] = None - continue - outcomes = check_relion_outcomes(job_results, expected_outcome, jobid) - - overall.setdefault(program, True) - if outcomes[program]["success"] is False: - overall[program] = False - failed_tests.extend(outcomes[program]["reason"]) - elif outcomes[program]["success"] is None and overall[program] is not False: - overall[program] = None - return overall, failed_tests - - -def check_screening_outcomes( - dcid, program, screening_output_lattice, outcomes, expected_outcome -): - all_screening_programs = [ - "mosflm", - "Stepped transmission 1", - "XOalign", - "dials.align_crystal", - "EDNA MXv1", - ] - if program not in all_screening_programs: - return - outcomes.update( - { - program: {"success": None} - for program in all_screening_programs - if program not in outcomes - } - ) - error_explanation = ( - "{variable}: {value} outside range {expected}, program: {program}, DCID:{dcid}" - ) - if not outcomes[program]["success"]: - failure_reasons = [] - results = {"spacegroup": screening_output_lattice.spaceGroup} - results.update( - { - v: getattr(screening_output_lattice, "unitCell_" + v) - for v in ("a", "b", "c", "alpha", "beta", "gamma") - } - ) - for variable, outcome in results.items(): - if outcome is None or expected_outcome[variable] != outcome: - failure_reasons.append( - error_explanation.format( - variable=variable, - value=outcome, - expected=expected_outcome[variable], - program=program, - dcid=dcid, - ) - ) - if failure_reasons: - outcomes[program]["success"] = False - outcomes[program]["reason"] = failure_reasons - else: - outcomes[program]["success"] = True - outcomes[program]["reason"] = [] - - -def check_integration_outcomes(dcid, program, integration, outcomes, expected_outcome): - all_integration_programs = { - "fast_dp", - "xia2 3dii", - "xia2 dials", - "autoPROC", - "autoPROC+STARANISO", - } - if program not in all_integration_programs: - return - outcomes.update( - { - program: {"success": None} - for program in all_integration_programs - if program not in outcomes - } - ) - error_explanation = ( - "{variable}: {value} outside range {expected}, program: {program}, DCID:{dcid}" - ) - - if not outcomes[program]["success"]: - failure_reasons = [] - # No result registered, may still be running - if integration.cell_a and expected_outcome: - for variable in ("a", "b", "c", "alpha", "beta", "gamma"): - outcome = getattr(integration, "cell_" + variable) - if outcome is None or expected_outcome[variable] != outcome: - failure_reasons.append( - error_explanation.format( - variable=variable, - value=outcome, - expected=expected_outcome[variable], - program=program, - dcid=dcid, - ) - ) - outcomes[program]["success"] = not failure_reasons - outcomes[program]["reason"] = failure_reasons - - -def check_pia_outcomes(dcid, image_count, iqi_count, outcomes, expected_outcome): - error_explanation = "Expected PIA result count for {dcid}: {expected}, got {actual}" - expected_pia_count = expected_outcome.get("pia", min(image_count, 250)) - if iqi_count == expected_pia_count: - outcomes["pia"] = {"success": True} - else: - outcomes["pia"] = { - "success": None, - "reason": [ - error_explanation.format( - dcid=dcid, - expected=expected_pia_count, - actual=iqi_count, - ) - ], - } - return outcomes - - -def _retrieve_motioncorr(db_session, jobid): - autoproc_query = ( - db_session.query(AutoProcProgram) - .options( - Load(AutoProcProgram).load_only("autoProcProgramId", "processingJobId"), - ) - .filter(AutoProcProgram.processingJobId == jobid) - ) - - autoproc_query_result = autoproc_query.order_by( - AutoProcProgram.autoProcProgramId.desc() - ).first() - autoprocpid = autoproc_query_result.autoProcProgramId - - query = db_session.query(MotionCorrection).filter( - MotionCorrection.autoProcProgramId == autoprocpid - ) - - query_results = query.all() - - return list(query_results), autoprocpid - - -def _retrieve_ctf(db_session, autoprocid): - query = ( - db_session.query(CTF, MotionCorrection) - .join( - MotionCorrection, - MotionCorrection.motionCorrectionId == CTF.motionCorrectionId, - ) - .filter(MotionCorrection.autoProcProgramId == autoprocid) - ) - query_results = query.all() - - return [q[0] for q in query_results] - - -def _retrieve_relative_ice_thickness(db_session, autoprocid): - query = ( - db_session.query(RelativeIceThickness, MotionCorrection) - .join( - MotionCorrection, - MotionCorrection.motionCorrectionId - == RelativeIceThickness.motionCorrectionId, - ) - .filter(MotionCorrection.autoProcProgramId == autoprocid) - ) - query_results = query.all() - - return [q[0] for q in query_results] - - -def _retrieve_particle_picker(db_session, autoprocid): - query = ( - db_session.query(ParticlePicker, MotionCorrection) - .join( - MotionCorrection, - MotionCorrection.motionCorrectionId - == ParticlePicker.firstMotionCorrectionId, - ) - .filter(MotionCorrection.autoProcProgramId == autoprocid) - ) - query_results = query.all() - - return [q[0] for q in query_results] - - -def _retrieve_particle_classification(db_session, autoprocid): - query = ( - db_session.query(ParticleClassification, ParticleClassificationGroup) - .join( - ParticleClassificationGroup, - ParticleClassificationGroup.particleClassificationGroupId - == ParticleClassification.particleClassificationGroupId, - ) - .filter(ParticleClassificationGroup.programId == autoprocid) - ) - - query_results = query.all() - - return [q[0] for q in query_results] - - -def check_relion_outcomes(job_results, expected_outcome, jobid): - all_programs = [ - "relion", - ] - outcomes = {program: {"success": None} for program in all_programs} - - failure_reasons = [] - count_based_failure_reasons = [] - for record_type, readable_name, get_micrograph_path in ( - ("motion_correction", "motion correction", attrgetter("micrographFullPath")), - ("ctf", "CTF", attrgetter("MotionCorrection.micrographFullPath")), - ( - "relative_ice_thickness", - "ice thickness", - attrgetter("MotionCorrection.micrographFullPath"), - ), - ( - "particle_picker", - "Particle Picker", - attrgetter("MotionCorrection.micrographFullPath"), - ), - ): - seen_records = defaultdict(int) - for result in job_results[record_type]: - micrograph = get_micrograph_path(result) - if not micrograph: - failure_reasons.append(f"Unexpected {readable_name} result: {result!r}") - continue - expected_record = expected_outcome[record_type].get(micrograph) - if not expected_record: - failure_reasons.append( - f"Unexpected {readable_name} result for micrograph {micrograph}" - ) - continue - seen_records[micrograph] += 1 - for variable in expected_record: - outcome = getattr(result, variable, None) - if outcome is None or expected_record[variable] != outcome: - failure_reasons.append( - f"{readable_name} for {micrograph} {variable}: {outcome} " - f"outside range {expected_record[variable]} in JobID:{jobid}" - ) - if len(seen_records) != len(expected_outcome[record_type]): - # if the lengths don't match the test is inconclusive so will need to set success to None - count_based_failure_reasons.append( - f"Out of {len(expected_outcome[record_type])} expected micrographs " - f"only {len(seen_records)} were seen" - ) - if any(count > 1 for count in seen_records.values()): - failure_reasons.append( - f"{readable_name} micrographs were seen more than once: %r" - % { - micrograph - for micrograph, count in seen_records.items() - if count > 1 - } - ) - - outcomes["relion"]["reason"] = failure_reasons + count_based_failure_reasons - outcomes["relion"]["success"] = not failure_reasons - if outcomes["relion"]["success"] and count_based_failure_reasons: - outcomes["relion"]["success"] = None - return outcomes - - -if __name__ == "__main__": - db = ispyb.open("/dls_sw/apps/zocalo/secrets/credentials-ispyb-sp.cfg") - - check_test_outcome( - { - "time_start": 1539264122.187212, - "scenario": "native", - "success": None, - "time_end": 1539264176.104456, - "DCIDs": [3029691], - "beamline": "i03", - }, - db, - ) diff --git a/src/dlstbx/dc_sim/dbserverclient.py b/src/dlstbx/dc_sim/dbserverclient.py deleted file mode 100644 index 18289a63d..000000000 --- a/src/dlstbx/dc_sim/dbserverclient.py +++ /dev/null @@ -1,413 +0,0 @@ -from __future__ import annotations - -import datetime -import http.client -import re -import sys -import xml.etree.ElementTree as ET - -DBSERVER_HOST = "sci-serv3" -DBSERVER_PORT = "2611" - - -def flatten_xml(xml, tag): - return "".join("\n".join(i for t in xml.iter(tag) for i in t.itertext())) - - -def _f(_v): - if _v is None: - return float("nan") - else: - return float(_v) - - -def _i(_v): - if _v is None: - return -1 - else: - return int(_v) - - -def _s(_v): - if _v is None: - return "null" - else: - return str(_v) - - -def _clean_nan_null_minusone(s): - return re.sub(r"\<[^<>]*\>(null|nan|-1)\]*\>", "", s) - - -def populate_blsample_xml_template(_row): - temp = _blsample_xml % ( - _s(_row["name"]), - _s(_row["code"]), - _s(_row["location"]), - _f(_row["holderlength"]), - _f(_row["looplength"]), - _s(_row["looptype"]), - _f(_row["wirewidth"]), - _s(_row["comments"]), - _s(_row["blsamplestatus"]), - _s(_row["lastknowncenteringposition"]), - ) - - # remove lines with null, nan and -1 values: - temp = _clean_nan_null_minusone(temp) - return temp - - -def populate_dcg_xml_template(datacollection, sessionid, blsample_id): - nowstr = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") - if blsample_id is None: - blsample_id_elem = "" - else: - blsample_id_elem = f"{blsample_id}\n" - - temp = _dcg_temp_xml_format.format( - sessionid=sessionid, - blsample_xml=blsample_id_elem, - comments="Simulated datacollection.", - experimenttype=_s(datacollection.DataCollectionGroup.experimentType), - starttime=nowstr, - crystalclass=_s(datacollection.DataCollectionGroup.crystalClass), - detectormode=_s(datacollection.DataCollectionGroup.detectorMode), - ) - - # remove lines with null, nan and -1 values: - temp = _clean_nan_null_minusone(temp) - return temp - - -def populate_grid_info_xml_template(_row, _dcid): - temp = _grid_info_temp_xml % ( - _f(_row["dx_mm"]), - _f(_row["dy_mm"]), - _i(_row["steps_x"]), - _i(_row["steps_y"]), - _f(_row["snapshot_offsetxpixel"]), - _f(_row["snapshot_offsetypixel"]), - _s(_row["orientation"]), - bool(_row["snaked"]), - _dcid, - _i(_row["patchesx"]), - _i(_row["patchesy"]), - _f(_row["micronsperpixelx"]), - _f(_row["micronsperpixely"]), - ) - - # remove lines with null, nan and -1 values: - temp = _clean_nan_null_minusone(temp) - return temp - - -def populate_dc_xml_template( - _row, - _sessionid, - _dcg_id, - _no_images, - _dir, - _prefix, - _run_number, - _xtal_snapshot_path, - _blsample_id, - scenario_name=None, -): - nowstr = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") - suffix = _row["imagesuffix"] - if suffix == "h5": - file_template = "%s_%d_master.%s" % (_prefix, _run_number, suffix) - else: - file_template = "%s_%d_####.%s" % (_prefix, _run_number, suffix) - if _blsample_id is None: - blsample_id_elem = "" - else: - blsample_id_elem = "%d\n" % _blsample_id - - temp = _dc_temp_xml % ( - _sessionid, - _dcg_id, - blsample_id_elem, - _run_number, - nowstr, - _s(_row["runstatus"]), - _f(_row["axisstart"]), - _f(_row["axisend"]), - _f(_row["axisrange"]), - _f(_row["overlap"]), - _no_images, - _i(_row["startimagenumber"]), - _i(_row["numberofpasses"]), - _f(_row["exposuretime"]), - _dir, - _prefix, - suffix, - file_template, - _f(_row["wavelength"]), - _f(_row["resolution"]), - _f(_row["detectordistance"]), - _f(_row["xbeam"]), - _f(_row["ybeam"]), - _i(_row["printableforreport"]), - _f(_row["slitgapvertical"]), - _f(_row["slitgaphorizontal"]), - _f(_row["transmission"]), - _s(_row["synchrotronmode"]), - _xtal_snapshot_path[0], - _xtal_snapshot_path[1], - _xtal_snapshot_path[2], - _xtal_snapshot_path[3], - _s(_row["rotationaxis"] or None), - _f(_row["phistart"]), - _f(_row["chistart"]), - _f(_row["kappastart"]), - _f(_row["omegastart"]), - _f(_row["undulatorgap1"]), - _f(_row["beamsizeatsamplex"]), - _f(_row["beamsizeatsampley"]), - _f(_row["flux"]), - _i(_row["focalspotsizeatsamplex"]), - _i(_row["focalspotsizeatsampley"]), - ) - temp = temp.format( - comments=( - f"Simulated datacollection ({scenario_name})." - if scenario_name - else "Simulated datacollection." - ) - ) - - # remove lines with null, nan and -1 values: - temp = _clean_nan_null_minusone(temp) - return temp - - -_blsample_xml = ( - '' - "" - "%s" - "%s" - "%s" - "%.6f" - "%.6f" - "%s" - "%.6f" - "%s" - "%s" - "False" - "%s" - "" -) - -_dcg_temp_xml_format = ( - '' - "" - "{sessionid}" - "{blsample_xml}" - "{experimenttype}" - "{starttime}" - "{crystalclass}" - "{detectormode}" - "{comments}" - "" -) - -_grid_info_temp_xml = ( - '' - "" - "%.2f" - "%.2f" - "%d" - "%d" - "%.4f" - "%.4f" - "%s" - "%s" - "%d" - "%d" - "%d" - "%.6f" - "%.6f" - "" -) - -_dc_temp_xml = ( - '' - "" - "%d" - "%d" - "%s%d" - "%s" - "%s" - "%.2f" - "%.2f" - "%.2f" - "%.2f" - "%d" - "%d" - "%d" - "%.3f" - "%s" - "%s" - "%s" - "%s" - "%.6f" - "%.2f" - "%.6f" - "%.6f" - "%.6f" - "{comments}" - "%d" - "%.6f" - "%.6f" - "%.6f" - "%s" - "%s" - "%s" - "%s" - "%s" - "%s" - "%.1f" - "%.1f" - "%.1f" - "%.1f" - "%.6f" - "%.2f" - "%.2f" - "%.6f" - "%d" - "%d" - "" -) - -dc_endtime_temp_xml = ( - '' - "" - "%d" - "%s" - "" -) - -dcg_endtime_temp_xml = ( - '' - "" - "%d" - "%s" - "" -) - - -class DbserverClient: - def __init__(self, _host=DBSERVER_HOST, _port=DBSERVER_PORT): - self.DB_host = _host - self.DB_port = _port - print(f"Connection parameters are: {self.DB_host}:{self.DB_port}") - - def _send(self, _path, _xml): - conn = http.client.HTTPConnection(f"{self.DB_host}:{self.DB_port}") - _xml = _xml.encode("latin-1") - try: - conn.putrequest("POST", str(_path)) - conn.putheader("Host", self.DB_host) - conn.putheader("Content-type", "text/xml") - conn.putheader("Content-length", len(_xml)) - conn.putheader("Accept", "text/xml") - conn.endheaders() - conn.send(_xml) - except OSError: - conn.close() - sys.exit("socket.error - is the server available?") - try: - response = conn.getresponse() - except http.client.BadStatusLine: - conn.close() - sys.exit("http.client.BadStatusLine: Unknown status code.") - except OSError: - conn.close() - sys.exit( - "socket.error - is the client authorised to connect to the server?" - ) - - xml = None - if response.status < 400 or response.status >= 600: - lengthstr = response.getheader("Content-length") - if lengthstr: - length = int(lengthstr) - xml = response.read(length) # get the raw XML - xml = xml.decode("latin-1") - print(xml) - else: - conn.close() - sys.exit("No Content-length in received header.") - else: - conn.close() - sys.exit(str(response.status) + ": " + response.reason) - - conn.close() - return xml - - def storeBLSample(self, xml): - xml_dbstatus_returned = self._send("/store_blsample_request", xml) - - e = ET.fromstring(xml_dbstatus_returned) - flatcode = flatten_xml(e, "code") - if flatcode != "ok": - sys.exit(flatten_xml(e, "message")) - - blsampleid = flatten_xml(e, "blSampleId") - if not blsampleid.isdigit(): - print(xml_dbstatus_returned) - sys.exit("No BLSampleID found in output") - return int(blsampleid) - - def storeDataCollectionGroup(self, xml): - xml_dbstatus_returned = self._send("/store_object_request", xml) - - e = ET.fromstring(xml_dbstatus_returned) - flatcode = flatten_xml(e, "code") - if flatcode != "ok": - sys.exit(flatten_xml(e, "message")) - - dcgid = flatten_xml(e, "dataCollectionGroupId") - if not dcgid.isdigit(): - print(xml_dbstatus_returned) - sys.exit("No DCGID found in output") - return int(dcgid) - - def storeGridInfo(self, xml): - xml_dbstatus_returned = self._send("/store_object_request", xml) - - e = ET.fromstring(xml_dbstatus_returned) - flatcode = flatten_xml(e, "code") - if flatcode != "ok": - sys.exit(flatten_xml(e, "message")) - - gridid = flatten_xml(e, "gridInfoId") - if not gridid.isdigit(): - print(xml_dbstatus_returned) - sys.exit("No GridInfoID found in output") - return int(gridid) - - def storeDataCollection(self, xml): - xml_dbstatus_returned = self._send("/store_object_request", xml) - - e = ET.fromstring(xml_dbstatus_returned) - flatcode = flatten_xml(e, "code") - if flatcode != "ok": - sys.exit(flatten_xml(e, "message")) - - dcid = flatten_xml(e, "dataCollectionId") - if not dcid.isdigit(): - print(xml_dbstatus_returned) - sys.exit("No DCID found in output") - return int(dcid) - - def updateDbObject(self, xml): - xml_dbstatus_returned = self._send("/update_object_request", xml) - - e = ET.fromstring(xml_dbstatus_returned) - flatcode = flatten_xml(e, "code") - if flatcode != "ok": - print(xml_dbstatus_returned) - sys.exit(flatten_xml(e, "message")) diff --git a/src/dlstbx/dc_sim/definitions.py b/src/dlstbx/dc_sim/definitions.py deleted file mode 100644 index 6cb9294c0..000000000 --- a/src/dlstbx/dc_sim/definitions.py +++ /dev/null @@ -1,412 +0,0 @@ -from __future__ import annotations - -from pytest import approx - -tests = { - "i04-eiger-screening": { - "type": "mx", - "src_dir": "/dls/i04/data/2021/cm28182-5/20211118/screening/TestInsulin/ins1", - "src_run_num": (2,), - "src_prefix": ("ins1",), - "sample_id": 3635367, - "results": { - "a": approx(77.8, abs=1), - "b": approx(77.8, abs=1), - "c": approx(77.8, abs=1), - "alpha": approx(90, abs=0), - "beta": approx(90, abs=0), - "gamma": approx(90, abs=0), - "spacegroup": "I23", - }, - }, - "i04-eiger-gridscan": { - "type": "mx", - "src_dir": "/dls/i04/data/2021/cm28182-2/xraycentring/manual", - "src_run_num": (29,), - "src_prefix": ("xrc",), - "results": {}, - }, - "i04-eiger-high-mult": { - "type": "mx", - "src_dir": "/dls/i04/data/2021/cm28182-1/20210226/gw/swmr/TestThaumatin/Se-Thaumatin6", - "src_run_num": (1,), - "src_prefix": ("Se-Thaumatin6_400hz_24images",), - "results": { - "a": approx(100, abs=100), - "b": approx(100, abs=100), - "c": approx(100, abs=100), - "alpha": approx(180, abs=180), - "beta": approx(180, abs=180), - "gamma": approx(180, abs=180), - }, - }, - "i04-eiger-native": { - "type": "mx", - "src_dir": "/dls/i04/data/2023/mx23694-95/auto/TestProteinaseK/protk5", - "src_run_num": (2,), - "src_prefix": ("protk5",), - "sample_id": 4509789, - "results": { - "a": approx(100, abs=100), - "b": approx(100, abs=100), - "c": approx(100, abs=100), - "alpha": approx(180, abs=180), - "beta": approx(180, abs=180), - "gamma": approx(180, abs=180), - }, - }, - "i04-1-eiger-native": { - "type": "mx", - "src_dir": "/dls/i04-1/data/2023/cm33904-3/TestLysozyme/Lys-Br16", - "src_run_num": (1,), - "src_prefix": ("Lys-Br16",), - "sample_id": 4684404, - "results": { - "a": approx(100, abs=100), - "b": approx(100, abs=100), - "c": approx(100, abs=100), - "alpha": approx(180, abs=180), - "beta": approx(180, abs=180), - "gamma": approx(180, abs=180), - }, - }, - "i23-lysozyme": { - "type": "mx", - "src_dir": "/dls/i23/data/2023/cm33851-2/TestLysozyme/lys_1_km70", - "src_run_num": (1,), - "src_prefix": ("data",), - "sample_id": 4493445, - "results": { - "a": approx(100, abs=100), - "b": approx(100, abs=100), - "c": approx(100, abs=100), - "alpha": approx(180, abs=180), - "beta": approx(180, abs=180), - "gamma": approx(180, abs=180), - }, - }, - "i04-insulin": { - "type": "mx", - "src_dir": "/dls/i04/data/2023/cm33903-3/20230808", - "src_run_num": (1,), - "src_prefix": ("ins12",), - "results": { - "a": approx(100, abs=100), - "b": approx(100, abs=100), - "c": approx(100, abs=100), - "alpha": approx(180, abs=180), - "beta": approx(180, abs=180), - "gamma": approx(180, abs=180), - }, - }, - "i04-mad-multi-dcg": { - "type": "mx", - "src_dir": "/dls/i04/data/2023/cm33903-4/TestThaumatin/SeThau7", - "src_run_num": (1,), - "src_prefix": ("Sethau7_2_pk", "Sethau7_2_if", "Sethau7_2_hrm"), - "sample_id": 4873871, - "is_dcg": True, - "results": { - "a": approx(100, abs=100), - "b": approx(100, abs=100), - "c": approx(100, abs=100), - "alpha": approx(180, abs=180), - "beta": approx(180, abs=180), - "gamma": approx(180, abs=180), - }, - }, - "i04-sad-multi-dcg": { - "type": "mx", - "src_dir": "/dls/i04/data/2023/cm33903-4/TestThaumatin/SeThau2", - "src_run_num": ( - 3, - 4, - 5, - ), - "src_prefix": ("SeThau2",), - "sample_id": 4873856, - "is_dcg": True, - "results": { - "a": approx(100, abs=100), - "b": approx(100, abs=100), - "c": approx(100, abs=100), - "alpha": approx(180, abs=180), - "beta": approx(180, abs=180), - "gamma": approx(180, abs=180), - }, - }, - "i03-native-multi-nondcg": { - "type": "mx", - "src_dir": "/dls/i03/data/2021/cm28170-1/Blah/xtal_1", - "src_run_num": ( - 1, - 2, - ), - "src_prefix": ("xtal_1",), - "sample_id": 3221630, - "results": { - "a": approx(100, abs=100), - "b": approx(100, abs=100), - "c": approx(100, abs=100), - "alpha": approx(180, abs=180), - "beta": approx(180, abs=180), - "gamma": approx(180, abs=180), - }, - }, - "i03-gridscan": { - # /dls/i03/data/2022/cm31105-1/xraycentring/TestProteinaseK/protk_16/protk_16_14_master.h5 - # https://ispyb.diamond.ac.uk/dc/visit/cm31105-1/id/7697211 - "type": "mx", - "src_dir": "/dls/i03/data/2022/cm31105-1/xraycentring/TestProteinaseK/protk_16/", - "src_run_num": (14,), - "src_prefix": ("protk_16",), - "sample_id": 3699723, - }, - "i03-native": { - "type": "mx", - "src_dir": "/dls/i03/data/2023/cm33866-2/TestInsulin/ins_1/", - "src_run_num": (2,), - "src_prefix": ("200",), - "sample_id": 4491570, - "results": { - "a": approx(100, abs=100), - "b": approx(100, abs=100), - "c": approx(100, abs=100), - "alpha": approx(180, abs=180), - "beta": approx(180, abs=180), - "gamma": approx(180, abs=180), - }, - }, - "vmxi-rotation": { - "type": "mx", - "src_dir": "/dls/mx/data/nt30330/nt30330-89/VMXi-AB2090/well_110/images", - "src_run_num": (107993,), - "src_prefix": ("image",), - "sample_id": 4315478, - "results": { - "a": approx(100, abs=100), - "b": approx(100, abs=100), - "c": approx(100, abs=100), - "alpha": approx(180, abs=180), - "beta": approx(180, abs=180), - "gamma": approx(180, abs=180), - }, - }, - "protk-au-insitu": { - "type": "mx", - "src_dir": "/dls/i24/data/2017/nr16818-47/Josh/InSitu/CrystalQuickX/ProtK/Au", - "src_run_num": (1,), - "src_prefix": ( - "ProtK_InSitu_Au_47_1", - "ProtK_InSitu_Au_47_2", - "ProtK_InSitu_Au_47_3", - "ProtK_InSitu_Au_47_4", - "ProtK_InSitu_Au_47_5", - "ProtK_InSitu_Au_47_6", - "ProtK_InSitu_Au_47_7", - "ProtK_InSitu_Au_47_8", - "ProtK_InSitu_Au_47_9", - "ProtK_InSitu_Au_47_10", - ), - "results": { - "a": approx(100, abs=100), - "b": approx(100, abs=100), - "c": approx(100, abs=100), - "alpha": approx(180, abs=180), - "beta": approx(180, abs=180), - "gamma": approx(180, abs=180), - }, - "delay": 20, - }, - "smargon-dcg": { - "type": "mx", - "src_dir": "/dls/i03/data/2018/cm19644-4/20180912/chigroup/protk", - "src_run_num": (8, 9, 10), - "src_prefix": ("protk",), - "is_dcg": True, - "results": { - "a": approx(100, abs=100), - "b": approx(100, abs=100), - "c": approx(100, abs=100), - "alpha": approx(180, abs=180), - "beta": approx(180, abs=180), - "gamma": approx(180, abs=180), - }, - }, - "inverse-beam": { - "type": "mx", - "src_dir": "/dls/i23/data/2018/nr19737-17/20180510/lys6", - "src_run_num": (1,), - "src_prefix": ("data_2200eV_A_1_", "data_2200eV_B_1_"), - "is_dcg": True, - "results": { - "a": approx(100, abs=100), - "b": approx(100, abs=100), - "c": approx(100, abs=100), - "alpha": approx(180, abs=180), - "beta": approx(180, abs=180), - "gamma": approx(180, abs=180), - }, - }, - "i04-eiger-sad-single": { - "type": "mx", - "src_dir": "/dls/i04/data/2023/cm33903-3/20230622/TestLysozyme/lysBr5", - "src_run_num": (1,), - "src_prefix": ("lysBr5",), - "sample_id": 4687158, - "results": { - "a": approx(100, abs=100), - "b": approx(100, abs=100), - "c": approx(100, abs=100), - "alpha": approx(180, abs=180), - "beta": approx(180, abs=180), - "gamma": approx(180, abs=180), - }, - }, - "i04-eiger-sad-wedged": { - "type": "mx", - "src_dir": "/dls/i04/data/2021/cm28182-2/TestInsulin/insulin2", - "src_run_num": (4,), - "src_prefix": ("insulin2_wedged",), - "sample_id": 3343403, - "results": { - "a": approx(100, abs=100), - "b": approx(100, abs=100), - "c": approx(100, abs=100), - "alpha": approx(180, abs=180), - "beta": approx(180, abs=180), - "gamma": approx(180, abs=180), - }, - }, - "screening": { - "type": "mx", - "src_dir": "/dls/i04/data/2017/cm16781-1/20170223/group1/Thaum/Th_4", - "src_run_num": (1,), - "src_prefix": ("Th_4",), - "sample_id": 790048, - "results": { - "a": approx(57.92, abs=2), - "b": approx(57.92, abs=2), - "c": approx(150.51, abs=2), - "alpha": approx(90, abs=0), - "beta": approx(90, abs=0), - "gamma": approx(90, abs=0), - "spacegroup": "P4", - }, - }, - "screening-i24": { - "type": "mx", - "src_dir": "/dls/i24/data/2017/cm16788-3/screening/hewl/hewl_1", - "src_run_num": (1,), - "src_prefix": ("hewl_1",), - "sample_id": 1018393, - "results": { - "a": approx(78.30, abs=2), - "b": approx(78.39, abs=2), - "c": approx(36.50, abs=2), - "alpha": approx(90, abs=0), - "beta": approx(90, abs=0), - "gamma": approx(90, abs=0), - "spacegroup": "P4", - "required": ("EDNA MXv1", "mosflm"), - }, - }, - "screening-smargon": { - "type": "mx", - "src_dir": "/dls/i04/data/2017/cm16781-1/20170316/Thaum/Th_4", - "src_run_num": (2,), - "src_prefix": ("Th_4",), - "sample_id": 790048, - "results": {}, - }, - "trp-multi": { - "type": "mx", - "src_dir": "/dls/i04/data/2013/nt5966-4/20131007/TRP/P1_X6", - "src_run_num": (1, 2, 3, 4), - "src_prefix": ("TRP_M1S6",), - "results": { - "a": approx(100, abs=100), - "b": approx(100, abs=100), - "c": approx(100, abs=100), - "alpha": approx(180, abs=180), - "beta": approx(180, abs=180), - "gamma": approx(180, abs=180), - }, - }, - "i19-1-smallmolecule": { - "type": "mx", - "src_dir": "/dls/i19-1/data/2023/cm33867-2/TBB_2_1-col", - "src_run_num": ("01",), - "src_prefix": ("TBB2",), - "sample_id": 4290218, - "results": { - "a": approx(100, abs=100), - "b": approx(100, abs=100), - "c": approx(100, abs=100), - "alpha": approx(180, abs=180), - "beta": approx(180, abs=180), - "gamma": approx(180, abs=180), - }, - }, - "relion": { - "type": "em-spa", - "frames": list(set(range(21, 50)) - {32, 33, 34, 38, 41}), - "src_dir": "/dls/m12/data/2022/cm31111-2/raw", - "src_run_num": (None,), - "src_prefix": (None,), - "visit_num": 2, - "proc_params": { - "acquisition_software": "SerialEM", - "import_images": "/dls/m12/data/2022/cm31111-2/raw/Frames/*.tiff", - "motioncor_gainreference": "/dls/m12/data/2022/cm31111-2/processing/gain.mrc", - "voltage": "200", - "Cs": "2.7", - "ctffind_do_phaseshift": "false", - "angpix": "0.885", - "motioncor_binning": "1", - "motioncor_doseperframe": "1.277", - "stop_after_ctf_estimation": "false", - }, - "results": { - "motion_correction": { - f"MotionCorr/job002/Movies/Frames/20170629_000{image_number}_frameImage.mrc": { - "totalMotion": approx(15, 0.75), - # "early_motion": approx(2.5, 0.5), - # "late_motion": approx(15, 2), - "averageMotionPerFrame": approx(0.5, 1.25), - } - for image_number in (set(range(21, 50)) - {32, 33, 34, 38, 41}) - }, - "ctf": { - f"MotionCorr/job002/Movies/Frames/20170629_000{image_number}_frameImage.mrc": { - "astigmatism": ( - approx(0, abs=0.75) if image_number == 42 else approx(247, 0.75) - ), - "astigmatismAngle": ( - approx(0, abs=360) if image_number == 42 else approx(83, 0.75) - ), - "estimatedResolution": approx(5, 0.75), - "estimatedDefocus": approx(10800, 0.75), - "ccValue": approx(0.15, 0.75), - } - for image_number in (set(range(21, 50)) - {32, 33, 34, 38, 41}) - }, - "relative_ice_thickness": { - f"MotionCorr/job002/Movies/Frames/20170629_000{image_number}_frameImage.mrc": { - "median": approx(250e3, 0.2), - "minimum": approx(250e3, 0.2), - "maximum": approx(250e3, 0.2), - "q1": approx(250e3, 0.2), - "q3": approx(250e3, 0.2), - } - for image_number in (set(range(21, 50)) - {32, 33, 34, 38, 41}) - }, - "particle_picker": { - f"MotionCorr/job002/Movies/Frames/20170629_000{image_number}_frameImage.mrc": { - "numberOfParticles": approx(250, 0.75), - } - for image_number in (set(range(21, 50)) - {32, 33, 34, 38, 41}) - }, - }, - }, -} diff --git a/src/dlstbx/dc_sim/mydb.py b/src/dlstbx/dc_sim/mydb.py deleted file mode 100644 index 90dd76624..000000000 --- a/src/dlstbx/dc_sim/mydb.py +++ /dev/null @@ -1,173 +0,0 @@ -from __future__ import annotations - -import configparser -import logging - -import ispyb.connector.mysqlsp.main -import sqlalchemy -from ispyb.sqlalchemy import BLSession, DataCollection, Proposal -from sqlalchemy.orm import Load - -log = logging.getLogger("dlstbx.dc_sim") - - -class DB: - def __init__(self): - config = configparser.RawConfigParser(allow_no_value=True) - assert config.read("/dls_sw/dasc/mariadb/credentials/ispyb_scripts.cfg") - credentials = dict(config.items("prod")) - self.i = ispyb.connector.mysqlsp.main.ISPyBMySQLSPConnector(**credentials) - self.cursor = self.i.create_cursor() - - def doQuery(self, querystr): - cursor = self.i.create_cursor() - log.debug("DB: %s", querystr) - try: - cursor.execute(querystr) - return cursor.fetchall() - finally: - cursor.close() - - -def retrieve_datacollection(db_session, sessionid, path, prefix, run_number): - records_to_collect = ( - DataCollection.BLSAMPLEID, - DataCollection.FOCALSPOTSIZEATSAMPLEX, - DataCollection.FOCALSPOTSIZEATSAMPLEY, - DataCollection.axisEnd, - DataCollection.axisRange, - DataCollection.axisStart, - DataCollection.beamSizeAtSampleX, - DataCollection.beamSizeAtSampleY, - DataCollection.chiStart, - DataCollection.comments, - DataCollection.dataCollectionGroupId, - DataCollection.dataCollectionId, - DataCollection.detectorDistance, - DataCollection.exposureTime, - DataCollection.fileTemplate, - DataCollection.flux, - DataCollection.imageSuffix, - DataCollection.kappaStart, - DataCollection.numberOfImages, - DataCollection.numberOfPasses, - DataCollection.omegaStart, - DataCollection.overlap, - DataCollection.phiStart, - DataCollection.printableForReport, - DataCollection.resolution, - DataCollection.rotationAxis, - DataCollection.runStatus, - DataCollection.slitGapHorizontal, - DataCollection.slitGapVertical, - DataCollection.startImageNumber, - DataCollection.synchrotronMode, - DataCollection.transmission, - DataCollection.undulatorGap1, - DataCollection.wavelength, - DataCollection.xBeam, - DataCollection.xtalSnapshotFullPath1, - DataCollection.xtalSnapshotFullPath2, - DataCollection.xtalSnapshotFullPath3, - DataCollection.xtalSnapshotFullPath4, - DataCollection.yBeam, - ) - - query = ( - db_session.query(DataCollection) - .options(Load(DataCollection).load_only(*records_to_collect)) - .filter(DataCollection.SESSIONID == sessionid) - .filter(DataCollection.imageDirectory == path + "/") - ) - if run_number is not None: - query = query.filter(DataCollection.dataCollectionNumber == run_number) - if prefix is not None: - query = query.filter(DataCollection.imagePrefix == prefix) - result = query.first() - if not result: - raise ValueError("No matching data collection found") - return result - - -def retrieve_sessionid(db_session, visit): - query = ( - db_session.query(BLSession, Proposal) - .options( - Load(BLSession).load_only( - BLSession.sessionId, BLSession.visit_number, BLSession.proposalId - ), - Load(Proposal).load_only( - Proposal.proposalId, Proposal.proposalCode, Proposal.proposalNumber - ), - ) - .join( - Proposal, - Proposal.proposalId == BLSession.proposalId, - ) - .filter( - sqlalchemy.func.concat( - Proposal.proposalCode, - Proposal.proposalNumber, - "-", - BLSession.visit_number, - ) - == visit - ) - ) - - query_results = query.first() - if query_results is None: - raise ValueError(f"Query to obtain sessionid failed for {visit}") - if query_results[0].sessionId is None: - raise ValueError(f"Could not find sessionid for visit {visit}") - return query_results[0].sessionId - - -def retrieve_dc_from_dcid(db_session, dcid): - records_to_collect = ( - DataCollection.BLSAMPLEID, - DataCollection.dataCollectionGroupId, - DataCollection.dataCollectionId, - DataCollection.dataCollectionNumber, - DataCollection.imageDirectory, - DataCollection.imagePrefix, - DataCollection.imageSuffix, - ) - - query = ( - db_session.query(DataCollection) - .options(Load(DataCollection).load_only(*records_to_collect)) - .filter(DataCollection.dataCollectionId == dcid) - ) - result = query.first() - if not result: - raise ValueError("No matching data collection found") - return result - - -def retrieve_dcs_from_dcg(db_session, dcg): - records_to_collect = ( - DataCollection.BLSAMPLEID, - DataCollection.dataCollectionGroupId, - DataCollection.dataCollectionId, - DataCollection.dataCollectionNumber, - DataCollection.imageDirectory, - DataCollection.imagePrefix, - DataCollection.imageSuffix, - ) - - query = ( - db_session.query(DataCollection) - .options(Load(DataCollection).load_only(*records_to_collect)) - .filter(DataCollection.dataCollectionGroupId == dcg) - ) - rows = query.all() - if not rows: - raise ValueError("No matching data collections found") - - img_dir = rows[0].imageDirectory - sample_id = rows[0].BLSAMPLEID - # Get values for prefixes and run_nums, removing duplicate values - img_prefixes = list({r.imagePrefix for r in rows}) - img_run_num = list({r.dataCollectionNumber for r in rows}) - return img_dir, img_prefixes, img_run_num, sample_id diff --git a/tests/dc_sim/test_dc_sim.py b/tests/dc_sim/test_dc_sim.py deleted file mode 100644 index 96b3a7c76..000000000 --- a/tests/dc_sim/test_dc_sim.py +++ /dev/null @@ -1,279 +0,0 @@ -from __future__ import annotations - -import random -import time -from dataclasses import dataclass -from unittest import mock - -import pytest - -import dlstbx.dc_sim.check -import dlstbx.dc_sim.definitions - -all_programs = ["fast_dp", "xia2 3dii", "xia2 dials", "autoPROC", "autoPROC+STARANISO"] - - -@dataclass(frozen=True) -class DataCollectionResult: - dataCollectionId: int = 12345 - numberOfImages: int = 900 - - -@dataclass(frozen=True) -class AutoProcProgramResult: - processingPrograms: str - - -@dataclass(frozen=True) -class AutoProcIntegrationResult: - cell_a: float - cell_b: float - cell_c: float - cell_alpha: float - cell_beta: float - cell_gamma: float - autoProcIntegrationId: int = 12345 - - -@dataclass(frozen=True) -class ScreeningOutputResult: - program: str - screeningId: int = 12345 - - -@dataclass(frozen=True) -class ScreeningOutputLatticeResult: - SpaceGroup: str - unitCell_a: float = 100.0 - unitCell_b: float = 100.0 - unitCell_c: float = 100.0 - unitCell_alpha: float = 100.0 - unitCell_beta: float = 100.0 - unitCell_gamma: float = 100.0 - - -def make_dummy_db_and_test_dictionary(beamline, scenario, results): - """ - Helper function to create a dummy database interface and a dictionary to pass - to the check function. - - :param beamline: lower case string - :param scenario: scenario name as given in definitions - :param results: dictionary of DCID to list of results. Results are - (string, boolean) tuples of program names - and whether or not this test should be successful. - """ - - test_dictionary = { - "beamline": beamline, - "scenario": scenario, - "DCIDs": [int(i) for i in results], - "success": None, - "time_start": time.time() - 900, - "time_end": time.time() - 1, - } - - result_mocks = [] - for dcid, dcid_results in results.items(): - for program, success in dcid_results: - data_collection = DataCollectionResult(dataCollectionId=dcid) - autoprocprogram = AutoProcProgramResult(processingPrograms=program) - if success: - integration = AutoProcIntegrationResult( - cell_a=mock.ANY, - cell_b=mock.ANY, - cell_c=mock.ANY, - cell_alpha=mock.ANY, - cell_beta=mock.ANY, - cell_gamma=mock.ANY, - ) - else: - integration = AutoProcIntegrationResult( - cell_a=mock.ANY, - cell_b=mock.ANY, - cell_c=mock.ANY, - cell_alpha=mock.ANY, - cell_beta=-2, - cell_gamma=mock.ANY, - ) - screening_output = None - screening_output_lattice = None - iqi_count = 250 - result_mocks.append( - ( - data_collection, - autoprocprogram, - integration, - screening_output, - screening_output_lattice, - iqi_count, - ) - ) - return result_mocks, test_dictionary - - -@mock.patch("dlstbx.dc_sim.check._retrieve_mx_processing") -def test_check_should_return_NONE_if_no_processing_results_have_arrived_yet( - mock_retrieve_mx_processing, -): - result_mocks, result = make_dummy_db_and_test_dictionary( - "i03", "i03-native", {2960726: []} - ) - mock_retrieve_mx_processing.return_value = result_mocks - - dlstbx.dc_sim.check.check_test_outcome(result, mock.Mock()) - - assert result["success"] is None - - -@mock.patch("dlstbx.dc_sim.check._retrieve_mx_processing") -def test_check_should_return_NONE_if_some_successful_processing_results_are_available( - mock_retrieve_mx_processing, -): - result_mocks, result = make_dummy_db_and_test_dictionary( - "i03", "i03-native", {2960726: [(p, True) for p in all_programs[:2]]} - ) - mock_retrieve_mx_processing.return_value = result_mocks - - dlstbx.dc_sim.check.check_test_outcome(result, mock.Mock()) - - assert result["success"] is None - - -@mock.patch("dlstbx.dc_sim.check._retrieve_mx_processing") -def test_check_should_return_PASS_if_all_processing_results_are_good( - mock_retrieve_mx_processing, -): - result_mocks, result = make_dummy_db_and_test_dictionary( - "i03", "i03-native", {2960726: [(p, True) for p in all_programs]} - ) - mock_retrieve_mx_processing.return_value = result_mocks - - dlstbx.dc_sim.check.check_test_outcome(result, mock.Mock()) - - assert result["success"] is True - - -@mock.patch("dlstbx.dc_sim.check._retrieve_mx_processing") -def test_check_should_return_PASS_if_all_processing_results_are_good_no_matter_what_other_programs_say( - mock_retrieve_mx_processing, -): - result_mocks, result = make_dummy_db_and_test_dictionary( - "i03", - "i03-native", - {2960726: [(p, True) for p in all_programs] + [("random program", False)]}, - ) - mock_retrieve_mx_processing.return_value = result_mocks - - dlstbx.dc_sim.check.check_test_outcome(result, mock.Mock()) - - assert result["success"] is True - - -# @pytest.mark.parametrize('program', all_programs) -@mock.patch("dlstbx.dc_sim.check._retrieve_mx_processing") -def test_check_should_ignore_invalid_results_if_another_valid_result_exists_from_the_same_program( - mock_retrieve_mx_processing, -): - program = all_programs[3] - - # case where only a single program returned results: outcome should be undecided - result_mocks, result = make_dummy_db_and_test_dictionary( - "i03", "i03-native", {2960726: [(program, True), (program, False)]} - ) - mock_retrieve_mx_processing.return_value = result_mocks - - dlstbx.dc_sim.check.check_test_outcome(result, mock.Mock()) - - assert result["success"] is None - - # case where all programs passed: outcome should be PASS - result_mocks, result = make_dummy_db_and_test_dictionary( - "i03", - "i03-native", - {2960726: [(program, False)] + [(p, True) for p in all_programs]}, - ) - mock_retrieve_mx_processing.return_value = result_mocks - - dlstbx.dc_sim.check.check_test_outcome(result, mock.Mock()) - - assert result["success"] is True - - -@mock.patch("dlstbx.dc_sim.check._retrieve_mx_processing") -def test_check_should_return_FAIL_if_all_processing_results_are_bad( - mock_retrieve_mx_processing, -): - result_mocks, result = make_dummy_db_and_test_dictionary( - "i03", "i03-native", {2960726: [(p, False) for p in all_programs]} - ) - mock_retrieve_mx_processing.return_value = result_mocks - - dlstbx.dc_sim.check.check_test_outcome(result, mock.Mock()) - - assert result["success"] is False - # and a reason should be given - assert result["reason"] - - -@mock.patch("dlstbx.dc_sim.check._retrieve_mx_processing") -@pytest.mark.parametrize("broken_program", all_programs) -def test_check_should_return_FAIL_if_one_processing_result_is_bad( - mock_retrieve_mx_processing, broken_program -): - broken_results = [(p, p != broken_program) for p in all_programs] - random.shuffle(broken_results) - - result_mocks, result = make_dummy_db_and_test_dictionary( - "i03", "i03-native", {2960726: broken_results} - ) - mock_retrieve_mx_processing.return_value = result_mocks - - dlstbx.dc_sim.check.check_test_outcome(result, mock.Mock()) - - assert result["success"] is False - # and a reason should be given - assert result["reason"] - # and the broken program should be mentioned - assert broken_program in result["reason"] - - -@mock.patch("dlstbx.dc_sim.check._retrieve_mx_processing") -@pytest.mark.parametrize("broken_program", all_programs) -def test_check_should_return_FAIL_if_a_single_processing_result_is_bad( - mock_retrieve_mx_processing, broken_program -): - result_mocks, result = make_dummy_db_and_test_dictionary( - "i03", "i03-native", {2960726: [(broken_program, False)]} - ) - mock_retrieve_mx_processing.return_value = result_mocks - - dlstbx.dc_sim.check.check_test_outcome(result, mock.Mock()) - - assert result["success"] is False - # and a reason should be given - assert result["reason"] - # and the broken program should be mentioned - assert broken_program in result["reason"] - - -@mock.patch("dlstbx.dc_sim.check._retrieve_mx_processing") -@pytest.mark.parametrize("missing_program", all_programs) -@pytest.mark.parametrize("added_program", all_programs + ["random program"]) -def test_check_should_not_be_confused_by_other_programs_appearing_instead_of_known_programs( - mock_retrieve_mx_processing, missing_program, added_program -): - if missing_program == added_program: - return - - program_list = set(all_programs) - program_list.remove(missing_program) - program_list.add(added_program) - result_mocks, result = make_dummy_db_and_test_dictionary( - "i03", "i03-native", {2960726: [(p, True) for p in program_list]} - ) - mock_retrieve_mx_processing.return_value = result_mocks - - dlstbx.dc_sim.check.check_test_outcome(result, mock.Mock()) - - assert result["success"] is None diff --git a/tests/dc_sim/test_em_sim.py b/tests/dc_sim/test_em_sim.py deleted file mode 100644 index fdd58868b..000000000 --- a/tests/dc_sim/test_em_sim.py +++ /dev/null @@ -1,280 +0,0 @@ -from __future__ import annotations - -import time -from dataclasses import dataclass -from unittest import mock - -import pytest - -import dlstbx.dc_sim.check -import dlstbx.dc_sim.definitions - - -@dataclass(frozen=True) -class MotioncorrectionResult: - micrographFullPath: str - totalMotion: float = 15 - averageMotionPerFrame: float = 0.5 - - -@dataclass(frozen=True) -class IceThicknessResult: - MotionCorrection: MotioncorrectionResult - median: float = 250e3 - minimum: float = 250e3 - maximum: float = 250e3 - q1: float = 250e3 - q3: float = 250e3 - - -@dataclass(frozen=True) -class CTFResult: - MotionCorrection: MotioncorrectionResult - astigmatism: float = 247 - astigmatismAngle: float = 83 - estimatedResolution = 5 - estimatedDefocus = 10800 - ccValue = 0.15 - - -@dataclass(frozen=True) -class ParticlePickerResult: - MotionCorrection: MotioncorrectionResult - numberOfParticles: int = 200 - - -def test_check_relion_outcomes_pass_checks(): - frame_numbers = dlstbx.dc_sim.definitions.tests["relion"]["frames"] - - dc_results = { - "motion_correction": [ - MotioncorrectionResult( - micrographFullPath=f"MotionCorr/job002/Movies/Frames/20170629_000{frame}_frameImage.mrc" - ) - for frame in frame_numbers - ], - "relative_ice_thickness": [ - IceThicknessResult( - MotionCorrection=MotioncorrectionResult( - micrographFullPath=f"MotionCorr/job002/Movies/Frames/20170629_000{frame}_frameImage.mrc" - ), - ) - for frame in frame_numbers - ], - "ctf": [ - CTFResult( - MotionCorrection=MotioncorrectionResult( - micrographFullPath=f"MotionCorr/job002/Movies/Frames/20170629_000{frame}_frameImage.mrc" - ), - astigmatism=(0 if frame == 42 else 247), - astigmatismAngle=(-35 if frame == 42 else 83), - ) - for frame in frame_numbers - ], - "particle_picker": [ - ParticlePickerResult( - MotionCorrection=MotioncorrectionResult( - micrographFullPath=f"MotionCorr/job002/Movies/Frames/20170629_000{frame}_frameImage.mrc" - ), - numberOfParticles=200, - ) - for frame in frame_numbers - ], - "particle_classification": range(50), - } - - expected_outcome = dlstbx.dc_sim.definitions.tests.get("relion", {}).get("results") - - check_result = dlstbx.dc_sim.check.check_relion_outcomes( - dc_results, expected_outcome, 1 - ) - assert check_result["relion"]["success"] - - -def test_check_relion_outcomes_fail_checks(): - frame_numbers = dlstbx.dc_sim.definitions.tests["relion"]["frames"] - - def db_motion_corr_f(frame): - if frame == 30: - return MotioncorrectionResult( - micrographFullPath=f"MotionCorr/job002/Movies/Frames/20170629_000{frame}_frameImage.mrc", - averageMotionPerFrame=-16, - ) - else: - return MotioncorrectionResult( - micrographFullPath=f"MotionCorr/job002/Movies/Frames/20170629_000{frame}_frameImage.mrc" - ) - - dc_results_f = { - "motion_correction": [db_motion_corr_f(frame) for frame in frame_numbers], - "relative_ice_thickness": [ - IceThicknessResult( - MotionCorrection=MotioncorrectionResult( - micrographFullPath=f"MotionCorr/job002/Movies/Frames/20170629_000{frame}_frameImage.mrc" - ), - ) - for frame in frame_numbers - ], - "ctf": [ - CTFResult( - MotionCorrection=db_motion_corr_f(frame), - astigmatism=(0 if frame == 42 else 247), - astigmatismAngle=(-35 if frame == 42 else 83), - ) - for frame in frame_numbers - ], - "particle_picker": [ - ParticlePickerResult( - MotionCorrection=MotioncorrectionResult( - micrographFullPath=f"MotionCorr/job002/Movies/Frames/20170629_000{frame}_frameImage.mrc" - ), - numberOfParticles=200, - ) - for frame in frame_numbers - ], - "particle_classification": range(50), - } - - expected_outcome_f = dlstbx.dc_sim.definitions.tests.get("relion", {}).get( - "results" - ) - - check_result = dlstbx.dc_sim.check.check_relion_outcomes( - dc_results_f, expected_outcome_f, 1 - ) - assert not check_result["relion"]["success"] - assert check_result["relion"]["reason"] == [ - f"motion correction for MotionCorr/job002/Movies/Frames/20170629_00030_frameImage.mrc averageMotionPerFrame: -16 outside range {pytest.approx(0.5, 1.25)} in JobID:1" - ] - - -@mock.patch("dlstbx.dc_sim.check._retrieve_motioncorr") -@mock.patch("dlstbx.dc_sim.check._retrieve_relative_ice_thickness") -@mock.patch("dlstbx.dc_sim.check._retrieve_ctf") -@mock.patch("dlstbx.dc_sim.check._retrieve_particle_picker") -@mock.patch("dlstbx.dc_sim.check._retrieve_particle_classification") -@mock.patch("sqlalchemy.create_engine") -@mock.patch("sqlalchemy.orm.Session") -def test_check_test_outcome_success( - mock_sess, - mock_eng, - mock_classification, - mock_parpick, - mock_ctf, - mock_ice_thickness, - mock_mcorr, -): - frame_numbers = dlstbx.dc_sim.definitions.tests["relion"]["frames"] - - def db_motion_corr(i): - motion_corr = mock.Mock() - motion_corr.micrographFullPath = ( - f"MotionCorr/job002/Movies/Frames/20170629_000{i}_frameImage.mrc" - ) - motion_corr.totalMotion = 15 - motion_corr.averageMotionPerFrame = 0.5 - return motion_corr - - mock_mcorr.return_value = [db_motion_corr(_) for _ in frame_numbers], 1 - mock_ice_thickness.return_value = [ - IceThicknessResult(MotionCorrection=db_motion_corr(frame)) - for frame in frame_numbers - ] - mock_ctf.return_value = [ - CTFResult( - MotionCorrection=db_motion_corr(frame), - astigmatism=(0 if frame == 42 else 247), - astigmatismAngle=(-35 if frame == 42 else 83), - ) - for frame in frame_numbers - ] - mock_parpick.return_value = [ - ParticlePickerResult( - MotionCorrection=MotioncorrectionResult( - micrographFullPath=f"MotionCorr/job002/Movies/Frames/20170629_000{frame}_frameImage.mrc" - ), - numberOfParticles=200, - ) - for frame in frame_numbers - ] - mock_classification.return_value = [None for _ in range(50)] - - test = { - "beamline": "m12", - "scenario": "relion", - "DCIDs": [1], - "JobIDs": [1], - "time_start": time.time() - 900, - "time_end": time.time() - 1, - } - test = dlstbx.dc_sim.check.check_test_outcome(test, mock.Mock()) - assert test["success"] - - -@mock.patch("dlstbx.dc_sim.check._retrieve_motioncorr") -@mock.patch("dlstbx.dc_sim.check._retrieve_relative_ice_thickness") -@mock.patch("dlstbx.dc_sim.check._retrieve_ctf") -@mock.patch("dlstbx.dc_sim.check._retrieve_particle_picker") -@mock.patch("dlstbx.dc_sim.check._retrieve_particle_classification") -@mock.patch("sqlalchemy.create_engine") -@mock.patch("sqlalchemy.orm.Session") -def test_check_test_outcome_failure( - mock_sess, - mock_eng, - mock_classification, - mock_parpick, - mock_ctf, - mock_ice_thickness, - mock_mcorr, -): - frame_numbers = dlstbx.dc_sim.definitions.tests["relion"]["frames"] - - def db_motion_corr(i): - motion_corr = mock.Mock() - motion_corr.micrographFullPath = ( - f"MotionCorr/job002/Movies/Frames/20170629_000{i}_frameImage.mrc" - ) - motion_corr.totalMotion = 15 - if i == 30: - motion_corr.averageMotionPerFrame = -0.2 - else: - motion_corr.averageMotionPerFrame = 0.5 - return motion_corr - - mock_mcorr.return_value = [db_motion_corr(frame) for frame in frame_numbers], 1 - mock_ice_thickness.return_value = [ - IceThicknessResult(MotionCorrection=db_motion_corr(frame)) - for frame in frame_numbers - ] - mock_ctf.return_value = [ - CTFResult( - MotionCorrection=db_motion_corr(frame), - astigmatism=(0 if frame == 42 else 247), - astigmatismAngle=(-35 if frame == 42 else 83), - ) - for frame in frame_numbers - ] - mock_parpick.return_value = [ - ParticlePickerResult( - MotionCorrection=MotioncorrectionResult( - micrographFullPath=f"MotionCorr/job002/Movies/Frames/20170629_000{frame}_frameImage.mrc" - ), - numberOfParticles=200, - ) - for frame in frame_numbers - ] - mock_classification.return_value = [None for _ in range(50)] - - test = { - "beamline": "m12", - "scenario": "relion", - "DCIDs": [1], - "JobIDs": [1], - "time_start": time.time() - 900, - "time_end": time.time() - 1, - } - test_checked = dlstbx.dc_sim.check.check_test_outcome(test, mock.Mock()) - assert not test_checked["success"] - assert test_checked["reason"].endswith( - f"averageMotionPerFrame: -0.2 outside range {pytest.approx(0.5, 1.25)} in JobID:1" - ) From 1cf908fdf3ca993e93482bb190f3735cb7b99f06 Mon Sep 17 00:00:00 2001 From: Phil Blowey Date: Thu, 29 Oct 2026 10:27:57 +0000 Subject: [PATCH 11/14] Remove dc_sim cli from setup --- setup.py | 1 - 1 file changed, 1 deletion(-) diff --git a/setup.py b/setup.py index 9ccfec8a8..9a512d7f7 100644 --- a/setup.py +++ b/setup.py @@ -13,7 +13,6 @@ "dials.swirly_eyes=dlstbx.cli.swirly_eyes:run", "dlstbx.align_crystal=dlstbx.cli.align_crystal:run", "dlstbx.archive_ancient_visits=dlstbx.cli.archive_ancient_visits:run", - "dlstbx.dc_sim=dlstbx.cli.dc_sim:run", "dlstbx.dc_sim_verify=dlstbx.cli.dc_sim_verify:run", "dlstbx.ep_predict_phase=dlstbx.cli.ep_predict_phase:run", "dlstbx.ep_predict_results=dlstbx.cli.ep_predict_results:runmain", From 70b64b3e7eaa62ee0312da7a5971d5c187a7038a Mon Sep 17 00:00:00 2001 From: Phil Blowey Date: Thu, 29 Oct 2026 10:28:34 +0000 Subject: [PATCH 12/14] Remove dc_sim verify --- setup.py | 1 - src/dlstbx/cli/dc_sim_verify.py | 233 -------------------------------- 2 files changed, 234 deletions(-) delete mode 100644 src/dlstbx/cli/dc_sim_verify.py diff --git a/setup.py b/setup.py index 9a512d7f7..ae922f3b0 100644 --- a/setup.py +++ b/setup.py @@ -13,7 +13,6 @@ "dials.swirly_eyes=dlstbx.cli.swirly_eyes:run", "dlstbx.align_crystal=dlstbx.cli.align_crystal:run", "dlstbx.archive_ancient_visits=dlstbx.cli.archive_ancient_visits:run", - "dlstbx.dc_sim_verify=dlstbx.cli.dc_sim_verify:run", "dlstbx.ep_predict_phase=dlstbx.cli.ep_predict_phase:run", "dlstbx.ep_predict_results=dlstbx.cli.ep_predict_results:runmain", "dlstbx.find_funny_eiger_frames=dlstbx.cli.find_funny_eiger_frames:run", diff --git a/src/dlstbx/cli/dc_sim_verify.py b/src/dlstbx/cli/dc_sim_verify.py deleted file mode 100644 index 17ba46cc3..000000000 --- a/src/dlstbx/cli/dc_sim_verify.py +++ /dev/null @@ -1,233 +0,0 @@ -# -# dlstbx.dc_sim_verify -# Verify all outstanding simulated data collections against expected -# results. Create a report as junit.xml file. -# -from __future__ import annotations - -import collections -import datetime -import queue -import sys -import time -from optparse import SUPPRESS_HELP, OptionParser -from pprint import pprint - -import ispyb -import junit_xml -import sqlalchemy -import workflows.recipe -from workflows.transport.stomp_transport import StompTransport - -import dlstbx.dc_sim.check -import dlstbx.util.result - -processqueue: queue.Queue = queue.Queue() - -results_queue = "reduce.dc_sim" -test_results: dict[str, list] = {} -# Fail scenarios that have not succeeded after 2 hours. -# Accommodate for possibility of 1 hour time zone difference between workstations. -test_timeout = 2 * 3600 -forget_test_after = 2 * 24 * 3600 # forget test after 2 days - -_ispyb_sessionmaker = sqlalchemy.orm.sessionmaker( - bind=sqlalchemy.create_engine( - ispyb.sqlalchemy.url(), connect_args={"use_pure": True} - ) -) - - -def process_result(rw, header, message): - processqueue.put((header, message)) - - -def run(): - global test_results - - parser = OptionParser(usage="dlstbx.dc_sim_verify [options]") - - parser.add_option("-?", action="help", help=SUPPRESS_HELP) - parser.add_option( - "--test", - action="store_true", - dest="test", - help="Run in ActiveMQ testing (zocdev) namespace", - ) - parser.add_option( - "--xml", action="store_true", dest="xml", help="Generate XML report files" - ) - - default_configuration = "/dls_sw/apps/zocalo/secrets/credentials-live.cfg" - if "--test" in sys.argv: - default_configuration = "/dls_sw/apps/zocalo/secrets/credentials-testing.cfg" - # override default stomp host - StompTransport.load_configuration_file(default_configuration) - - StompTransport.add_command_line_options(parser) - (options, args) = parser.parse_args(sys.argv[1:]) - stomp = StompTransport() - stomp.connect() - txn = stomp.transaction_begin() - - sid = workflows.recipe.wrap_subscribe( - stomp, - results_queue, - process_result, - acknowledgement=True, - exclusive=True, - allow_non_recipe_messages=True, - ) - - try: - while True: - header, message = processqueue.get(True, 5) - # Acknowledge all received messages within transaction - stomp.ack(header, transaction=txn) - - if message.get("summary"): - # aggregated test information - for beamline_scenario_string in message["summary"]: - test_results.setdefault(beamline_scenario_string, []).extend( - message["summary"][beamline_scenario_string] - ) - continue - elif "beamline" in message and "scenario" in message: - test_results.setdefault( - f"{message['beamline']}-{message['scenario']}", [] - ).append(message) - - except queue.Empty: - pass # No more messages coming in - stomp.unsubscribe(sid) - # Further messages are not processed and will be redelivered to the next instance - - def filter_testruns(runlist): - """Filter out all historic test runs older then forget_test_after - and all test runs that are older than the most recent run that - ended in a result.""" - # Forget everything older than forget_test_after - forget_all_before = time.time() - forget_test_after - runlist = [t for t in runlist if t["time_end"] > forget_all_before] - - unfinished = [t for t in runlist if t.get("success") is None] - finished = [t for t in runlist if t.get("success") is not None] - - # Keep the single most recent test run that has a result... - if finished: - latest_outcome = max(t["time_end"] for t in finished) - finished = [t for t in finished if t["time_end"] == latest_outcome] - if finished: - finished = [finished[0]] - - # ...and all test runs that have not yet finished - return finished + unfinished - - test_results = { - setting: filter_testruns(testruns) for setting, testruns in test_results.items() - } - - # Forget tests where all test runs have expired - for key in list(test_results): - if not test_results[key]: - del test_results[key] - - # Check all test runs that do not yet have a definite outcome - for testruns in test_results.values(): - for testrun in testruns: - if testrun.get("success") is None: - with _ispyb_sessionmaker() as db_session: - print("Verifying", testrun) - dlstbx.dc_sim.check.check_test_outcome( - testrun, db_session=db_session - ) - # 3 possible outcomes: - # The test can be successful (testrun['success'] = True) - # it can fail (testrun['success'] = False; testrun['reason'] set) - # or it can be inconclusive (eg. because results are missing) - # in which case no changes are made - if ( - testrun.get("success") is None - and testrun["time_end"] < time.time() - test_timeout - ): - print("Rejecting with timeout:", testrun) - testrun["success"] = False - existing_reason = testrun.get("reason") - testrun["reason"] = "No valid results appeared within timeout" - if existing_reason: - testrun["reason"] += f" ({existing_reason})" - - # Show all known test results - pprint(test_results) - - # If there are results then put summary back on results queue - if test_results: - stomp.send(results_queue, {"summary": test_results}, transaction=txn) - stomp.transaction_commit(txn) - - if options.xml: - generate_xml_files(test_results) - - time.sleep(0.3) # to ensure all stomp messages have been sent - if not test_results: - exit("There are no recent simulation runs to verify") - - -def generate_xml_files(test_results): - # Create JUnit result records - junit_results = collections.defaultdict(list) - for testruns in test_results.values(): - r = dlstbx.util.result.Result() - r.set_name(testruns[0]["scenario"]) - r.set_classname("{test[beamline]}.{test[scenario]}".format(test=testruns[0])) - for test in testruns: - if test.get("success") in (False, True): - r.log_message( - "Started at {start:%Y-%m-%d %H:%M:%S}, finished at {end:%Y-%m-%d %H:%M:%S}, took {elapsed:.1f} seconds.".format( - start=datetime.datetime.fromtimestamp(test["time_start"]), - end=datetime.datetime.fromtimestamp(test["time_end"]), - elapsed=test["time_end"] - test["time_start"], - ) - ) - if test["success"]: - r.log_message("Test successful") - else: - r.log_error(test.get("reason", "Test failed")) - if len(testruns) > 1: - r.log_message( - "%d further run(s) of this test ongoing" % (len(testruns) - 1) - ) - for url in test.get("URLs", []): - r.log_message(url) - r.set_time(test["time_end"] - test["time_start"]) - break - else: - r.log_message( - "Started at {start:%Y-%m-%d %H:%M:%S}, finished at {end:%Y-%m-%d %H:%M:%S}, took {elapsed:.1f} seconds.".format( - start=datetime.datetime.fromtimestamp(testruns[0]["time_start"]), - end=datetime.datetime.fromtimestamp(testruns[0]["time_end"]), - elapsed=testruns[0]["time_end"] - testruns[0]["time_start"], - ) - ) - r.log_skip( - "Waiting on results, %d instance(s) of this test ongoing" - % len(testruns) - ) - for url in test.get("URLs", []): - r.log_message(url) - r.set_time(testruns[0]["time_end"] - testruns[0]["time_start"]) - - junit_results[test.get("type")].append(r) - - # Export grouped results - for group in junit_results: - if group: - filename = f"output-{group}.xml" - else: - filename = "output.xml" - ts = junit_xml.TestSuite( - "Simulated data collections", - junit_results[group], - ) - with open(filename, "w") as f: - junit_xml.TestSuite.to_file(f, [ts], prettyprint=True) From a0558179e0b2af11e44a258ac0fbc8d629426517 Mon Sep 17 00:00:00 2001 From: Phil Blowey Date: Thu, 29 Oct 2026 11:24:05 +0000 Subject: [PATCH 13/14] Add checks for missing data --- src/dlstbx/wrapper/dc_sim.py | 18 ++++++++++++------ tests/wrapper/test_dc_sim.py | 3 +++ 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/src/dlstbx/wrapper/dc_sim.py b/src/dlstbx/wrapper/dc_sim.py index e71798cbe..32d6f9751 100644 --- a/src/dlstbx/wrapper/dc_sim.py +++ b/src/dlstbx/wrapper/dc_sim.py @@ -266,7 +266,7 @@ def simulate_datacollection( start_image_number = dc.startImageNumber number_of_images = dc.numberOfImages copied_images = 0 - missing_images = 0 + missing_images = [] self.log.info( f"Copying {number_of_images} CBF images from {source_image_dir} " f"to {dest_dir}" @@ -285,16 +285,22 @@ def simulate_datacollection( shutil.copy2(src_image_path, dest_image_path) copied_images += 1 else: - self.log.warning( - f"Source image {src_image_path} does not exist; skipping copy" - ) - missing_images += 1 + missing_images.append(image_number) + if not copied_images: + raise FileNotFoundError( + f"Image files not found for dcid={dcid}, data may have been archived" + ) self.log.info( f"Finished copying CBF images: copied={copied_images}, " - f"missing={missing_images}" + f"missing_images={missing_images}" ) elif image_suffix in ["h5", "nxs"]: + main_file_path = source_image_dir / dc.fileTemplate + if not main_file_path.exists(): + raise FileNotFoundError( + f"Image files not found for dcid={dcid}, data may have been archived" + ) file_prefix = f"{dc.imagePrefix}_{dc.dataCollectionNumber}" files = { path diff --git a/tests/wrapper/test_dc_sim.py b/tests/wrapper/test_dc_sim.py index f4bd08411..7642b0e9c 100644 --- a/tests/wrapper/test_dc_sim.py +++ b/tests/wrapper/test_dc_sim.py @@ -257,6 +257,7 @@ def test_simulate_datacollection_copies_hdf5_and_nexus_files(tmp_path, monkeypat dest_dir.mkdir() jpgs_dest_dir.mkdir() for name in ( + "sample_2_master.h5", "sample_2_000001.h5", "sample_2_000002.h5", "sample_2.nxs", @@ -267,6 +268,7 @@ def test_simulate_datacollection_copies_hdf5_and_nexus_files(tmp_path, monkeypat imageDirectory=str(source_dir), imageSuffix="h5", imagePrefix="sample", + fileTemplate="sample_2_master.h5", dataCollectionNumber=2, xtalSnapshotFullPath1=None, xtalSnapshotFullPath2=None, @@ -286,6 +288,7 @@ def test_simulate_datacollection_copies_hdf5_and_nexus_files(tmp_path, monkeypat FakeSession(), 25, dest_dir, jpgs_dest_dir, 30, 40 ) == (502, True) assert {path.name for path in dest_dir.iterdir()} == { + "sample_2_master.h5", "sample_2_000001.h5", "sample_2_000002.h5", "sample_2.nxs", From e21322679fa6d2ee5e2b6a34f1c52360bc0b97d8 Mon Sep 17 00:00:00 2001 From: Irakli Sikharulidze Date: Wed, 4 Nov 2026 15:45:30 +0000 Subject: [PATCH 14/14] Minor refactoring * Remove excessive error handling * Simplify sessionid db query * Fix type definitions --- src/dlstbx/wrapper/dc_sim.py | 35 ++++++++++------------------------- tests/wrapper/test_dc_sim.py | 28 ++-------------------------- 2 files changed, 12 insertions(+), 51 deletions(-) diff --git a/src/dlstbx/wrapper/dc_sim.py b/src/dlstbx/wrapper/dc_sim.py index 32d6f9751..d02afc52d 100644 --- a/src/dlstbx/wrapper/dc_sim.py +++ b/src/dlstbx/wrapper/dc_sim.py @@ -48,14 +48,11 @@ class DCSimWrapper(Wrapper): def retrieve_sessionid(self, db_session: sqlalchemy.orm.Session, visit: str) -> int: query = ( - db_session.query(BLSession, Proposal) + db_session.query(BLSession) .options( Load(BLSession).load_only( BLSession.sessionId, BLSession.visit_number, BLSession.proposalId ), - Load(Proposal).load_only( - Proposal.proposalId, Proposal.proposalCode, Proposal.proposalNumber - ), ) .join( Proposal, @@ -71,12 +68,8 @@ def retrieve_sessionid(self, db_session: sqlalchemy.orm.Session, visit: str) -> == visit ) ) - query_results = query.first() - if query_results is None: - raise ValueError(f"Query to obtain sessionid failed for {visit}") - if query_results[0].sessionId is None: - raise ValueError(f"Could not find sessionid for visit {visit}") - return query_results[0].sessionId + query_results = query.one() + return query_results.sessionId def copy_data_collection_group( self, db_session: sqlalchemy.orm.Session, dcgid: int, new_session_id: int @@ -87,10 +80,7 @@ def copy_data_collection_group( db_session.query(DataCollectionGroup).filter( DataCollectionGroup.dataCollectionGroupId == dcgid ) - ).one_or_none() - if src_dcg is None: - self.log.error(f"No data collection group found for dcgid={dcgid}") - raise ValueError(f"Could not find data collection group dcgid={dcgid}") + ).one() new_dcg = DataCollectionGroup() # Copy all attributes from source to new, except the primary key and fields to update @@ -119,7 +109,7 @@ def copy_data_collection_group( def get_dc_and_grid_info_from_dcid( self, db_session: sqlalchemy.orm.Session, dcid: int - ) -> tuple[DataCollection, GridInfo | None]: + ) -> sqlalchemy.Row[tuple[DataCollection, GridInfo | None]]: """Fetch a data collection and its related records in one database query. GridInfo is optional and may be associated directly with the data collection @@ -138,12 +128,8 @@ def get_dc_and_grid_info_from_dcid( GridInfo, GridInfo.dataCollectionId == DataCollection.dataCollectionId ) .filter(DataCollection.dataCollectionId == dcid) - .one_or_none() + .one() ) - if result is None: - self.log.error(f"No data collection found for dcid={dcid}") - raise ValueError(f"Could not find data collection for dcid={dcid}") - return result def copy_data_collection( @@ -153,7 +139,7 @@ def copy_data_collection( new_session_id: int, new_data_collection_group_id: int, **kwargs, - ) -> DataCollection: + ) -> int: """ Create a copy of a DataCollection row with different sessionId and dataCollectionGroupId. and insert into the database @@ -202,7 +188,7 @@ def copy_grid_info( source_gi: GridInfo, new_dcid: int, **kwargs, - ) -> DataCollection: + ) -> None: """ Create a copy of a DataCollection row with different sessionId and/or dataCollectionGroupId. @@ -315,9 +301,6 @@ def simulate_datacollection( shutil.copy2(file, dest_dir) self.log.info("Finished copying image data files") else: - self.log.error( - f"Unsupported image suffix {image_suffix!r} for source dcid={dcid}" - ) raise ValueError( f"Invalid image suffix {image_suffix} for source data collection" ) @@ -424,6 +407,8 @@ def run(self): if visit_dir.name.startswith(proposal): dest_visit = visit_dir.name break + if not dest_visit: + raise KeyError("Could not determine which destination visit to use") dest_visit_dir = Path("/dls", beamline, "data", str(now.year), dest_visit) assert dest_visit_dir.is_dir(), ( f"Could not find {dest_visit_dir} directory for the specified visit and beamline." diff --git a/tests/wrapper/test_dc_sim.py b/tests/wrapper/test_dc_sim.py index 7642b0e9c..30b912773 100644 --- a/tests/wrapper/test_dc_sim.py +++ b/tests/wrapper/test_dc_sim.py @@ -35,7 +35,7 @@ def filter(self, *_args): def first(self): return self.query_result - def one_or_none(self): + def one(self): return self.query_result def add(self, instance): @@ -98,28 +98,11 @@ def test_dc_sim_parameters_reject_invalid_values(): def test_retrieve_sessionid_returns_session_id(): wrapper = DCSimWrapper() - session = FakeSession((SimpleNamespace(sessionId=123), SimpleNamespace())) + session = FakeSession(SimpleNamespace(sessionId=123)) assert wrapper.retrieve_sessionid(session, "nt44198-1") == 123 -@pytest.mark.parametrize( - ("result", "message"), - [ - (None, "Query to obtain sessionid failed"), - ( - (SimpleNamespace(sessionId=None), SimpleNamespace()), - "Could not find sessionid", - ), - ], -) -def test_retrieve_sessionid_raises_when_visit_is_not_found(result, message): - wrapper = DCSimWrapper() - - with pytest.raises(ValueError, match=message): - wrapper.retrieve_sessionid(FakeSession(result), "nt44198-1") - - def test_get_dc_and_grid_info_returns_records(): wrapper = DCSimWrapper() dc = SimpleNamespace(dataCollectionId=42) @@ -131,13 +114,6 @@ def test_get_dc_and_grid_info_returns_records(): ) -def test_get_dc_and_grid_info_raises_for_missing_record(): - wrapper = DCSimWrapper() - - with pytest.raises(ValueError, match="Could not find data collection for dcid=42"): - wrapper.get_dc_and_grid_info_from_dcid(FakeSession(None), 42) - - def test_copy_data_collection_overrides_ids_and_fields(): wrapper = DCSimWrapper() source_dc = SimpleNamespace(**values_for_model(DataCollection))