Skip to content

Commit

Permalink
Update flows (#18)
Browse files Browse the repository at this point in the history 
* update grid scan schemas

* updates redis connection

* update screening schema

* fix redis connections

* more logging

* bump version

* pre-commit
  • Loading branch information
@fhernandezvivanco
fhernandezvivanco authored Oct 28, 2024
1 parent 9b4cfad commit f065f02
Show file tree
Hide file tree
Showing 5 changed files with 83 additions and 46 deletions.
20 changes: 16 additions & 4 deletions mxcubecore/HardwareObjects/ANSTO/PrefectWorkflow.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -104,8 +104,13 @@ def __init__(self, name: str) -> None:
self.command_failed = False
self.gevent_event = None
self.workflow_name = None
self.redis_port = int(os.environ.get("DATA_PROCESSING_REDIS_PORT", "6379"))
self.redis_host = os.environ.get("DATA_PROCESSING_REDIS_HOST", "mx_redis")

self.REDIS_HOST = os.environ.get("MXCUBE_REDIS_HOST", "mx_redis")
self.REDIS_PORT = int(os.environ.get("MXCUBE_REDIS_PORT", "6379"))
self.REDIS_USERNAME = os.environ.get("MXCUBE_REDIS_USERNAME", None)
self.REDIS_PASSWORD = os.environ.get("MXCUBE_REDIS_PASSWORD", None)
self.REDIS_DB = int(os.environ.get("MXCUBE_REDIS_DB", "0"))

self.mxcubecore_workflow_aborted = False

def _init(self) -> None:
Expand Down Expand Up @@ -133,7 +138,13 @@ def init(self) -> None:

# self.beamline = HWR.beamline

self.redis_connection = redis.StrictRedis(self.redis_host, self.redis_port)
self.redis_connection = redis.StrictRedis(
host=self.REDIS_HOST,
port=self.REDIS_PORT,
username=self.REDIS_USERNAME,
password=self.REDIS_PASSWORD,
db=self.REDIS_DB,
)

self.raster_flow = None

Expand Down Expand Up @@ -269,7 +280,7 @@ def set_values_map(self, params):

def get_available_workflows(self) -> list[dict]:
"""
Gets the available workflows specified in the edna_params.xml file
Gets the available workflows specified in the prefect_flow.xml file
Returns
-------
Expand All @@ -290,6 +301,7 @@ def get_available_workflows(self) -> list[dict]:
dict_workflow["requires"] = []
dict_workflow["doc"] = ""
workflow_list.append(dict_workflow)
logging.getLogger("HWR").info(f"Available workflows: {workflow_list}")
return workflow_list

def abort(self) -> None:
Expand Down
67 changes: 45 additions & 22 deletions mxcubecore/HardwareObjects/ANSTO/prefect_flows/grid_scan_flow.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,11 @@
GRID_SCAN_DEPLOYMENT_NAME = environ.get(
"GRID_SCAN_DEPLOYMENT_NAME", "mxcube-grid-scan/plans"
)
_number_of_processes = environ.get("GRID_SCAN_NUMBER_OF_PROCESSES", None)
if _number_of_processes is not None:
GRID_SCAN_NUMBER_OF_PROCESSES = int(_number_of_processes)
else:
GRID_SCAN_NUMBER_OF_PROCESSES = None


class GridScanFlow(AbstractPrefectWorkflow):
Expand Down Expand Up @@ -66,8 +71,9 @@ def run(self, dialog_box_parameters: dict) -> None:
width = round(grid.width)
height = round(grid.height)

dialog_box_model = GridScanDialogBox.parse_obj(dialog_box_parameters)
dialog_box_model = GridScanDialogBox.model_validate(dialog_box_parameters)

# TODO: sample_id should be obtained from the database!
redis_grid_scan_id = self.redis_connection.get(
f"mxcube_grid_scan_id:{dialog_box_model.sample_id}"
)
Expand All @@ -85,11 +91,12 @@ def run(self, dialog_box_parameters: dict) -> None:
beam_position=beam_position,
number_of_columns=num_cols,
number_of_rows=num_rows,
exposure_time=dialog_box_model.exposure_time,
omega_range=dialog_box_model.omega_range,
hardware_trigger=dialog_box_model.hardware_trigger,
detector_distance=dialog_box_model.detector_distance,
photon_energy=dialog_box_model.photon_energy,
omega_range=dialog_box_model.omega_range,
md3_alignment_y_speed=dialog_box_model.md3_alignment_y_speed,
hardware_trigger=dialog_box_model.hardware_trigger,
number_of_processes=GRID_SCAN_NUMBER_OF_PROCESSES,
)

self.redis_connection.set(
Expand All @@ -99,7 +106,8 @@ def run(self, dialog_box_parameters: dict) -> None:
f"Parameters sent to prefect flow: {prefect_parameters}"
)
grid_scan_flow = MX3PrefectClient(
name=GRID_SCAN_DEPLOYMENT_NAME, parameters=prefect_parameters.dict()
name=GRID_SCAN_DEPLOYMENT_NAME,
parameters=prefect_parameters.model_dump(exclude_none=True),
)

try:
Expand Down Expand Up @@ -129,20 +137,26 @@ def run(self, dialog_box_parameters: dict) -> None:
grid_size = num_cols * num_rows
logging.getLogger("user_level_log").warning("Processing data...")
number_of_spots_array = np.zeros((num_rows, num_cols))
resolution_array = np.zeros((num_rows, num_cols))
for _ in range(grid_size):
data, last_id = self.read_message_from_redis_streams(
topic=f"number_of_spots_{prefect_parameters.grid_scan_id}:{prefect_parameters.sample_id}",
id=last_id,
)
number_of_spots = int(data[b"number_of_spots"])
number_of_spots = float(data[b"number_of_spots"])
resolution = float(data[b"resolution"])
heatmap_coordinate = pickle.loads(data[b"heatmap_coordinate"])
logging.getLogger("HWR").info(
f"heatmap coordinate: {heatmap_coordinate}, "
f"resolution: {resolution}, "
f"number of spots {number_of_spots}"
)
number_of_spots_array[
heatmap_coordinate[1], heatmap_coordinate[0]
] = number_of_spots
resolution_array[heatmap_coordinate[1], heatmap_coordinate[0]] = (
resolution
)

logging.getLogger("user_level_log").warning("Data processing finished")

Expand All @@ -155,14 +169,21 @@ def run(self, dialog_box_parameters: dict) -> None:
num_rows=num_rows,
number_of_spots_array=number_of_spots_array,
)
crystalmap_array = self.create_heatmap(
num_cols=num_cols,
num_rows=num_rows,
number_of_spots_array=resolution_array,
)

heatmap = {}
crystalmap = {}

if grid:
for i in range(1, num_rows * num_cols + 1):
heatmap[i] = [i, list(heatmap_array[i - 1])]
crystalmap[i] = [i, list(crystalmap_array[i - 1])]

heat_and_crystal_map = {"heatmap": heatmap, "crystalmap": heatmap}
heat_and_crystal_map = {"heatmap": heatmap, "crystalmap": crystalmap}
self.sample_view.set_grid_data(
sid, heat_and_crystal_map, data_file_path="this_is_not_used"
)
Expand Down Expand Up @@ -268,40 +289,42 @@ def dialog_box(self) -> dict:
"""
dialog = {
"properties": {
"exposure_time": {
"title": "exposure time",
"md3_alignment_y_speed": {
"title": "Alignment Y Speed [mm/s]",
"type": "number",
"minimum": 0,
"default": 1,
"maximum": 14.8,
"default": 10,
"widget": "textarea",
},
"omega_range": {
"title": "omega range",
"title": "Omega Range [degrees]",
"type": "number",
"minimum": 0,
"exclusiveMaximum": 361,
"maximum": 360,
"default": 0,
"widget": "textarea",
},
"detector_distance": {
"title": "detector distance",
"title": "Detector Distance [m]",
"type": "number",
"default": -0.298,
"minimum": 0,
"maximum": 1,
"default": 0.396,
"widget": "textarea",
},
"photon_energy": {
"title": "photon energy",
"title": "Photon Energy [keV]",
"type": "number",
"minimum": 0,
"default": 12700,
"minimum": 5,
"maximum": 25,
"default": 13,
"widget": "textarea",
},
"hardware_trigger": {
"title": "Hardware trigger (dev only)",
"title": "Hardware Trigger (dev only)",
"type": "boolean",
"minimum": 0,
"exclusiveMaximum": 361,
"default": False,
"default": True,
"widget": "textarea",
},
"sample_id": {
Expand All @@ -311,7 +334,7 @@ def dialog_box(self) -> dict:
"widget": "textarea",
},
},
"required": ["exposure_time", "omega_range"],
"required": ["detector_distance", "photon_energy"],
"dialogName": "Grid scan parameters",
}

Expand Down
18 changes: 11 additions & 7 deletions mxcubecore/HardwareObjects/ANSTO/prefect_flows/schemas/grid_scan.py
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
from pydantic.v1 import BaseModel
from typing import Optional, Union

from pydantic import BaseModel


class GridScanDialogBox(BaseModel):
exposure_time: float
md3_alignment_y_speed: float
omega_range: float
detector_distance: float
photon_energy: float
Expand All @@ -13,15 +14,18 @@ class GridScanDialogBox(BaseModel):

class GridScanParams(BaseModel):
sample_id: str
grid_scan_id: int = 0
grid_scan_id: Union[str, int]
grid_top_left_coordinate: Union[tuple[int, int], list[int]]
grid_height: int
grid_width: int
beam_position: Optional[Union[tuple[int, int], list[int]]] = [612, 512]
number_of_columns: int
number_of_rows: int
exposure_time: float
omega_range: float
hardware_trigger: bool = False
detector_distance: float
photon_energy: float
omega_range: float = 0
md3_alignment_y_speed: float = 10
beam_position: Union[tuple[int, int], list[int]] = (612, 512)
count_time: Optional[float] = None
hardware_trigger: bool = True
crystal_finder_threshold: int = 1
number_of_processes: Optional[float] = None
22 changes: 10 additions & 12 deletions mxcubecore/HardwareObjects/ANSTO/prefect_flows/screening_flow.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,48 +81,48 @@ def dialog_box(self) -> dict:
dialog = {
"properties": {
"exposure_time": {
"title": "exposure time",
"title": "Total Exposure Time [s]",
"type": "number",
"minimum": 0,
"default": 1,
"widget": "textarea",
},
"omega_range": {
"title": "omega range",
"title": "Omega Range [degrees]",
"type": "number",
"minimum": 0,
"exclusiveMaximum": 361,
"default": 10,
"widget": "textarea",
},
"number_of_frames": {
"title": "number of frames",
"title": "Number of Frames",
"type": "number",
"minimum": 1,
"default": 100,
"widget": "textarea",
},
"detector_distance": {
"title": "detector distance",
"title": "Detector Distance [m]",
"type": "number",
"default": -0.298,
"default": 0.396,
"widget": "textarea",
},
"photon_energy": {
"title": "photon energy",
"title": "Photon Energy [keV]",
"type": "number",
"minimum": 0,
"default": 12700,
"default": 13,
"widget": "textarea",
},
"processing_pipeline": {
"title": "Data processing pipeline",
"title": "Data Processing Pipeline",
"type": "string",
"enum": ["dials", "fast_dp", "dials_and_fast_dp"],
"default": "dials",
},
"crystal_counter": {
"title": "Crystal counter",
"title": "Crystal Counter",
"type": "number",
"minimum": 0,
"default": 0,
Expand All @@ -131,9 +131,7 @@ def dialog_box(self) -> dict:
"hardware_trigger": {
"title": "Hardware trigger (dev only)",
"type": "boolean",
"minimum": 0,
"exclusiveMaximum": 361,
"default": False,
"default": True,
"widget": "textarea",
},
"sample_id": {
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[tool.poetry]
name = "mxcubecore"
version = "1.79.0.2"
version = "1.79.0.3"
license = "LGPL-3.0-or-later"
description = "Core libraries for the MXCuBE application"
authors = ["The MXCuBE collaboration <mxcube@esrf.fr>"]
Expand Down

0 comments on commit f065f02

Please sign in to comment.