import logging
from pathlib import Path
from typing import Dict, List, Union
from ewokstools.submit import save_and_execute, wait_to_finish_queue
from ewoksutils.task_utils import task_inputs
from ._utils import (
CNMF_BASH_PRE_SCRIPT,
SLURM_JOB_PARAMETERS_CNMF,
EwoksWorkflow,
_confirm_submission,
_decide_background_removal,
_get_destination_filename,
_get_nb_components,
_normalize_slurm_job_parameters,
_print_checking_template_format,
_print_job_progress,
_print_number_of_jobs_to_submit,
_print_template_invalid,
_print_template_valid,
_transfer_cif_files,
_validate_yaml_template_model,
_warning_dry_run_mode,
_warning_if_too_many_jobs,
get_list_scan_info_from_bliss_filenames_id13,
_find_wavelength,
)
from .resources import BACKGROUND_CNMF_WORKFLOW
from .resources.models import CnmfModel
logger = logging.getLogger(__name__)
[docs]
def main_cnmf(args):
for file in args.FILES:
if file.endswith((".yaml", ".yml")):
_main_cnmf_from_template(filename=file)
elif file.endswith(".h5"):
kwargs = vars(args)
for key in ("FILES", "command"):
kwargs.pop(key, None)
_main_cnmf_from_cli(filename=file, **kwargs)
else:
logger.warning(
f"File {file} has an unsupported extension. Skipping it. Supported extensions are .yaml, .yml and .h5."
)
def _main_cnmf_from_template(filename: Union[str, Path]):
_print_checking_template_format()
is_valid, result = _validate_yaml_template_model(filename, CnmfModel)
if not is_valid:
_print_template_invalid(filename, result)
return
_print_template_valid()
_main_cnmf(**result)
def _main_cnmf_from_cli(filename: Union[str, Path], **kwargs):
integration_params = {
"processed_filenames": [filename],
"nxprocess_path_diffmap": kwargs.pop("nxprocess_path_diffmap"),
"nxprocess_path_filtered": kwargs.pop("nxprocess_path_filtered"),
"nxprocess_path_cnmf": kwargs.pop("nxprocess_path_cnmf"),
}
submit_parameters = {
"submit": not kwargs.pop("dry_run"),
"slurm_job_parameters": kwargs.pop("slurm_job_parameters") or {},
"queue": kwargs.pop("queue", None),
"local_execution": kwargs.pop("local_execution", False),
}
_main_cnmf(
submit_parameters=submit_parameters,
**integration_params,
**kwargs,
)
def _main_cnmf(
processed_filenames: Union[str, List[str]],
nxprocess_path_filtered: str,
nxprocess_path_cnmf: str,
nxprocess_path_diffmap: str = None,
submit_parameters: Dict = None,
**kwargs,
):
workflow_cnmf = EwoksWorkflow(BACKGROUND_CNMF_WORKFLOW)
submit_parameters = submit_parameters or {}
dry_run = not submit_parameters.get("submit")
list_scans_info = get_list_scan_info_from_bliss_filenames_id13(
bliss_filenames=processed_filenames,
**kwargs,
)
nb_total_jobs = len(list_scans_info)
_print_number_of_jobs_to_submit(nb_total_jobs)
_confirm_submission()
if dry_run:
_warning_dry_run_mode()
else:
_warning_if_too_many_jobs(nb_total_jobs)
submitted_jobs = []
for index_job, scan_info in enumerate(list_scans_info):
output_filename = scan_info.filename_raw_dataset
external_output_filename = output_filename
do_background = _decide_background_removal(
nxdata_url_to_filtered=f"{output_filename}::{nxprocess_path_filtered}/filtered"
)
cnmf_parameters = kwargs.get("cnmf", {}).copy()
wavelength_A = cnmf_parameters.pop("wavelength", None)
if wavelength_A is None:
wavelength_pyfai = _find_wavelength(
processed_filename=output_filename,
nxprocess_diffmap=nxprocess_path_diffmap,
)
wavelength_A = wavelength_pyfai * 1e10
radial_limits = cnmf_parameters.pop("radial_limits")
references_directory = cnmf_parameters.pop("references_directory")
inputs_ewoks = []
if do_background:
inputs_ewoks += task_inputs(
id=workflow_cnmf.background.id,
task_identifier=workflow_cnmf.background.task_identifier,
inputs={
"do_background_removal": True,
"nxdata_url": f"{output_filename}::{nxprocess_path_diffmap}/diffmap",
"wavelength": wavelength_A,
"nxprocess_name": nxprocess_path_filtered.split("/")[-1],
"radial_limits": radial_limits,
"use_neuralnetwork": False,
"worker_module": "scattering",
"output_filename": output_filename,
"external_output_filename": external_output_filename,
"submit_to_slurm": False, # It's already SLURM
"slurm_client": False, # Offline processing
"force_training": False,
**kwargs.get("background", {}),
},
)
else:
inputs_ewoks += task_inputs(
id=workflow_cnmf.background.id,
task_identifier=workflow_cnmf.background.task_identifier,
inputs={
"do_background_removal": False,
"nxdata_url": "",
"wavelength": wavelength_A,
**kwargs.get("background", {}),
},
)
nxprocess_name_cnmf = nxprocess_path_cnmf.split("/")[-1]
target_cif_directory = _transfer_cif_files(
references_directory=str(Path(references_directory).resolve()),
target_directory=Path(external_output_filename).parent
/ f".{nxprocess_name_cnmf}_cif",
)
nb_components = _get_nb_components(
target_cif_directory=target_cif_directory,
nb_components=cnmf_parameters.pop("nb_components", None),
nb_residues=cnmf_parameters.pop("nb_residues", 1),
)
inputs_ewoks += task_inputs(
id=workflow_cnmf.cnmf.id,
task_identifier=workflow_cnmf.cnmf.task_identifier,
inputs={
"nxdata_url": f"{output_filename}::{nxprocess_path_filtered}/filtered",
"wavelength": wavelength_A,
"nxprocess_name": nxprocess_name_cnmf,
"radial_limits": radial_limits,
"worker_module": "scattering",
"output_filename": output_filename,
"external_output_filename": external_output_filename,
"submit_to_slurm": False,
"references_directory": target_cif_directory,
"nb_components": nb_components,
"do_matrix_factorization": True,
**cnmf_parameters,
},
)
submit_parameters_ = submit_parameters.copy()
submitted = save_and_execute(
python_package="ewoksid13",
dry_run=dry_run,
workflow=BACKGROUND_CNMF_WORKFLOW,
inputs=inputs_ewoks,
bash_pre_script=CNMF_BASH_PRE_SCRIPT,
destination_filename=_get_destination_filename(
output_filename, scan_info.scan_nb, nxprocess_name_cnmf
),
slurm_job_parameters=_normalize_slurm_job_parameters(
SLURM_JOB_PARAMETERS_CNMF,
submit_parameters_.pop("slurm_job_parameters", {}),
),
**submit_parameters_,
)
submitted_jobs.append(submitted)
_print_job_progress(index_job + 1, nb_total_jobs, dry_run)
wait_to_finish_queue(submitted_jobs)