Source code for ewoksxrdct.tasks.create_nxxrdct

from typing import Optional
from typing import Union

from ewokscore import Task
from ewokscore.model import BaseInputModel
from ewokscore.model import BaseOutputModel
from nxxrdct import NXxrdct
from pydantic import Field




class CreateNxxrdctInputModel(BaseInputModel, validate_assignment=True):
    nx_path: str = Field(
        ...,
        description="Full path to the NX file.",
    )
    bliss_path: Optional[str] = Field(
        default=None, description="Path to the HDF5 master file."
    )
    title: Optional[str] = Field(default=None, description="Title of experiment.")
    start_time: Optional[str] = Field(
        default=None, description="Start time of the scan."
    )
    end_time: Optional[str] = Field(default=None, description="End time of the scan.")
    beam_incident_energy: Optional[float] = Field(
        default=None,
        description="Incident energy of the beam.",
    )
    detector_data: Optional[str] = Field(
        default=None, description="Data url to the AI data."
    )
    detector_polar_angle: Optional[str] = Field(
        default=None, description="Polar angle outputted from AI."
    )
    detector_count_time: Optional[float] = Field(
        default=None, description="Time to record a single frame."
    )
    detector_distance: Optional[float] = Field(
        default=None, description="Distance between the detector and sample."
    )
    detector_x_pixel_size: Optional[float] = Field(
        default=None, description="The detector pixel size in the x direction."
    )
    detector_y_pixel_size: Optional[float] = Field(
        default=None, description="The detector pixel size in the y direction."
    )
    detector_diffraction_channel: Optional[str] = Field(
        default=None, description="The diffraction channel of the detector."
    )
    source_name: Optional[str] = Field(default=None, description="Name of the source.")
    source_type: Optional[str] = Field(default=None, description="Type of the source.")
    source_probe: Optional[str] = Field(
        default="XRAYS", description="Probe of the source."
    )
    monochromator_wavelength: Optional[float] = Field(
        default=None,
        description="Wavelength.",
    )
    sample_name: Optional[str] = Field(default=None, description="Name of the sample.")
    sample_rotation_angles: Union[str, list, None] = Field(
        default=None, description="Rotation angles (one per frame)."
    )
    sample_translation_values: Union[str, list, None] = Field(
        default=None, description="Translation values (one per frame)."
    )
    sample_x_translation: Union[str, list, None] = Field(
        default=None, description="Translation values of x."
    )
    sample_y_translation: Union[str, list, None] = Field(
        default=None, description="Translation values of y."
    )
    sample_z_translation: Union[str, list, None] = Field(
        default=None, description="Translation values of z."
    )
    monitor_data: Optional[list] = Field(
        default=None, description="Data from the fpico meter."
    )
    monitor_mode: Optional[str] = Field(
        default=None, description="Mode of the monitor."
    )
    monitor_preset: Optional[str] = Field(
        default=None, description="Preset of the monitor."
    )
    monitor_integral: Optional[float] = Field(default=None, description="Integral.")





class CreateNxxrdctOutputModel(BaseOutputModel):
    nx_path: str = Field(
        ...,
        description="Full path to the NX file.",
    )





class CreateNxxrdct(
    Task,
    input_model=CreateNxxrdctInputModel,
    output_model=CreateNxxrdctOutputModel,
):


    def build_nx_object(self) -> NXxrdct:
        """Create the NX object."""
        nx = NXxrdct()
        nx.title = self.inputs.title
        nx.start_time = self.inputs.start_time
        nx.end_time = self.inputs.end_time
        nx.beam.incident_energy = self.inputs.beam_incident_energy
        nx.instrument.detector.data = self.inputs.detector_data
        nx.instrument.detector.polar_angle = self.inputs.detector_polar_angle
        nx.instrument.detector.count_time = self.inputs.detector_count_time
        nx.instrument.detector.distance = self.inputs.detector_distance
        nx.instrument.detector.x_pixel_size = self.inputs.detector_x_pixel_size
        nx.instrument.detector.y_pixel_size = self.inputs.detector_y_pixel_size
        nx.instrument.detector.diffraction_channel = (
            self.inputs.detector_diffraction_channel
        )
        nx.instrument.source.name = self.inputs.source_name
        nx.instrument.source.type = self.inputs.source_type
        nx.instrument.source.probe = self.inputs.source_probe
        nx.instrument.monochromator.wavelength = self.inputs.monochromator_wavelength
        nx.sample.name = self.inputs.sample_name
        nx.sample.rotation_angle = self.inputs.sample_rotation_angles
        nx.sample.translation_values = self.inputs.sample_translation_values
        nx.sample.x_translation = self.inputs.sample_x_translation
        nx.sample.y_translation = self.inputs.sample_y_translation
        nx.sample.z_translation = self.inputs.sample_z_translation
        nx.control.data = self.inputs.monitor_data
        nx.control.mode = self.inputs.monitor_mode
        nx.control.preset = self.inputs.monitor_preset
        nx.control.integral = self.inputs.monitor_integral

        return nx




    def save_nx_file(self, nx: NXxrdct, save_path: str) -> None:
        """Save the NX object at the specified file path."""
        nx.save(file_path=save_path, data_path="entry")




    def run(self) -> None:
        """Execute the task to create an NX file from the given inputs."""
        nx = self.build_nx_object()
        output_path = self.inputs.nx_path

        self.save_nx_file(nx, output_path)
        self.outputs.nx_path = output_path