Source code for ewoksmx.tasks.characterisation.aggregate_diffraction_data

import tempfile
from typing import List
from typing import Optional

from esrf_pathlib import ESRFPath
from ewokscore.model import BaseInputModel
from ewokscore.model import BaseOutputModel
from ewokscore.task import Task
from pydantic import Field

from ewoksmx.models.characterisation_models import MetadataModel
from ewoksmx.utils.cbf_utils import generate_cbf_sum_images
from ewoksmx.utils.hdf5_utils import aggregate_hdf5_masters_with_vds
from ewoksmx.utils.hdf5_utils import merge_hdf5_wedges




class InputModel(BaseInputModel):
    data_paths: List[ESRFPath] = Field(
        ...,
        min_length=1,
        description="List of paths to CBF images or HDF5 master files",
    )
    metadata: MetadataModel





class OutputModel(BaseOutputModel):
    xds_master_file: Optional[ESRFPath] = Field(
        None,
        description="Optional HDF5 master with XDS image numbering convention, required for fine-sliced data",
    )
    labelit_cbf_files: List[ESRFPath] = Field(
        ..., description="One summed CBF per wedge for Labelit"
    )
    mosflm_master_file: ESRFPath = Field(
        ..., description="HDF5 master with summed images for MOSFLM"
    )
    metadata: MetadataModel





class AggregateDiffractionData(Task, input_model=InputModel, output_model=OutputModel):
    """Aggregates the input diffraction data into HDF5 and CBF files."""



    def run(self):
        # Access inputs via self.inputs (which is an instance of InputModel)
        data_paths = self.inputs.data_paths

        working_dir = ESRFPath(tempfile.mkdtemp(prefix="ewoksmx_agg_diff_data_"))
        aggregated_output_path = working_dir / "ref-fineslice_aggregated_1_master.h5"
        sum_output_path = working_dir / "ref-aggregated_sum_1_master.h5"

        # Merge hdf5 masters with virtual data sets
        aggregate_hdf5_masters_with_vds(data_paths, aggregated_output_path)

        # Create HDF5 master for MOSFLM by summing images
        merge_hdf5_wedges(aggregated_output_path, sum_output_path)

        # Generate summed CBF images for Labelit
        list_cbf_path, metadata = generate_cbf_sum_images(
            aggregated_output_path, working_dir
        )

        self.outputs.xds_master_file = aggregated_output_path
        self.outputs.mosflm_master_file = sum_output_path
        self.outputs.labelit_cbf_files = list_cbf_path
        self.outputs.metadata = self.inputs.metadata