import logging
import time
import warnings
from pathlib import Path
import h5py
import numpy
from ewoksxrpd.tasks.utils.data_utils import create_hdf5_link
from silx.io import h5py_utils
from ..scripts.resources import BACKGROUND_WORKFLOW, CNMF_WORKFLOW, INFERENCE_WORKFLOW
from .slurm import SlurmTask
logger = logging.getLogger(__name__)
NXPROCESS_PATH_FORMAT = "{entry_name}/{nxprocess_name}"
NXDATA_PATH_DIFFMAP_FORMAT = "{entry_name}/{nxprocess_name}/diffmap"
NXDATA_PATH_BACKGROUND_FORMAT = "{entry_name}/{nxprocess_name}/filtered"
NXDATA_URL_FORMAT = "{processed_data_filename}::{nxdata}"
[docs]
class BackgroundRemoval(
SlurmTask,
input_names=[
"nxdata_url",
"wavelength",
],
optional_input_names=[
"output_filename",
"external_output_filename",
"nxprocess_name",
"radial_limits",
"model_path",
"force_training",
"use_neuralnetwork",
"do_background_removal",
"processing_info",
],
output_names=[
"nxdata_url",
"diffmap_nxprocess_path",
],
):
[docs]
def run(self):
nxdata_url = self.inputs.nxdata_url
if not nxdata_url or not self.get_input_value("do_background_removal", True):
self.outputs.nxdata_url = None
self.outputs.diffmap_nxprocess_path = None
logger.warning(
"Background removal is disabled. Skipping BackgroundRemoval task."
)
return
input_filename, nxdata_diffmap_path = nxdata_url.split("::")
h5path_parts = Path(f"/{nxdata_diffmap_path}").parts
if len(h5path_parts) == 4:
entry_name = h5path_parts[1] # e.g. "1.1"
nxprocess_diffmap_name = h5path_parts[2] # e.g. "eiger_diffmap"
else:
raise ValueError(
f"Invalid NXdata path format: {nxdata_diffmap_path}. Expected format: {NXDATA_PATH_DIFFMAP_FORMAT}"
)
# Get output files
external_output_filename: str = self.get_input_value(
"external_output_filename", None
)
output_filename: str = self.get_input_value("output_filename", None)
if not external_output_filename and not output_filename:
external_output_filename = output_filename = input_filename
elif not external_output_filename:
external_output_filename = output_filename
elif not output_filename:
output_filename = external_output_filename
# Get the NXentry group name from the diffmap_nxdata_path
nxprocess_background_name = self.get_input_value(
"nxprocess_name", nxprocess_diffmap_name.replace("diffmap", "filtered")
)
if nxprocess_background_name == nxprocess_diffmap_name:
nxprocess_background_name = nxprocess_diffmap_name + "_filtered"
nxprocess_background_path = NXPROCESS_PATH_FORMAT.format(
entry_name=entry_name,
nxprocess_name=nxprocess_background_name,
)
nxdata_background_path = NXDATA_PATH_BACKGROUND_FORMAT.format(
entry_name=entry_name,
nxprocess_name=nxprocess_background_name,
)
self.outputs.nxdata_url = NXDATA_URL_FORMAT.format(
processed_data_filename=external_output_filename,
nxdata=nxdata_background_path,
)
self.outputs.diffmap_nxprocess_path = nxprocess_background_path
# Submit to Slurm if requested
if self.get_input_value("submit_to_slurm", False):
inputs = [
{"name": k, "value": v, "all": True}
for k, v in self.named_input_values.items()
if k != "submit_to_slurm"
] + [{"name": "submit_to_slurm", "value": False, "all": True}]
_ = self.execute_remote(workflow=BACKGROUND_WORKFLOW, inputs=inputs)
return
mode = "r+" if external_output_filename == input_filename else "r"
with h5py_utils.open_item(
filename=input_filename, name="/", mode=mode
) as file_input:
if nxdata_diffmap_path not in file_input:
raise FileNotFoundError(
f"Diffraction map {nxdata_diffmap_path} not found in the file {input_filename}. Skipping background removal."
)
# Retrieve the diffraction map data
nxdata_diffmap = file_input[nxdata_diffmap_path]
dset_intensity = (
nxdata_diffmap["intensity_norm"]
if "intensity_norm" in nxdata_diffmap
else nxdata_diffmap["intensity"]
)
diffmap_array = dset_intensity[:]
nxdata_axes = nxdata_diffmap.attrs.get("axes", [""])
radial_name = next((i for i in nxdata_axes if i != "."), "2th")
radial_dset = (
nxdata_diffmap[radial_name] if radial_name in nxdata_diffmap else None
)
radial_array = radial_dset[:] if radial_dset is not None else None
if radial_array is None:
raise KeyError(
f"Radial dataset '{radial_name}' not found in {nxdata_diffmap_path}. It's required for background removal."
)
# Open the external file
with h5py_utils.open_item(
filename=external_output_filename, name="/", mode="r+"
) as file_output:
if nxprocess_background_path in file_output:
raise FileExistsError(
f"NXprocess group {nxprocess_background_path} already exists in the file {external_output_filename}. Please change the nxprocess_name parameter to avoid overwriting existing data."
)
# Predict/Calculate the filtered diffractograms
model_path = self.get_input_value("model_path", None)
if model_path is None:
model_path = external_output_filename.replace(
".h5", f"_{nxdata_diffmap_path.split('/')[0]}"
)
Path(model_path).parent.mkdir(exist_ok=True)
filtered_array = self.get_filtered_array(
diffmap_array=diffmap_array,
radial_array=radial_array,
model_path=model_path,
radial_limits=self.get_input_value("radial_limits", None),
force_training=self.get_input_value("force_training", False),
wavelength=self.inputs.wavelength,
use_neuralnetwork=self.get_input_value("use_neuralnetwork", True),
)
# Create the NXprocess group for filtered diffmap data
nxprocess_background_grp = file_output.create_group(
nxprocess_background_path
)
nxprocess_background_grp["name"] = "Background removed (auto)"
nxprocess_background_grp.attrs.update(
{
"NX_class": "NXprocess",
"default": "filtered",
}
)
# Save processing information
nxinfo_grp = nxprocess_background_grp.create_group("processing_info")
nxinfo_grp.attrs["NX_class"] = "NXcollection"
for key, value in self.get_input_value("processing_info", {}).items():
nxinfo_grp.attrs[key] = value
# Create the NXdata group for the filtered diffmap data
nxdata_background = file_output.create_group(nxdata_background_path)
nxdata_background.attrs.update(
{
"NX_class": "NXdata",
"signal": "map",
}
)
# Create the dataset for the filtered intensity
dset_intensity_background = nxdata_background.create_dataset(
name="intensity",
data=filtered_array,
chunks=(1, *diffmap_array.shape[1:]),
dtype="float32",
)
# Copy the radial information, and all the axes
for ax_name in nxdata_axes:
nxdata_diffmap_axis = nxdata_diffmap[ax_name]
virtual_source = h5py.VirtualSource(nxdata_diffmap_axis)
virtual_layout = h5py.VirtualLayout(
shape=nxdata_diffmap_axis.shape,
dtype=nxdata_diffmap_axis.dtype,
)
virtual_layout[:] = virtual_source[:]
nxdata_background.create_virtual_dataset(ax_name, virtual_layout)
# NxData attrs
nxdata_background.attrs.update(
{
"interpretation": "image",
"signal": "intensity",
"axes": nxdata_axes,
}
)
# Map with shifted dimensions
y_nb_points, x_nb_points, nbpt_rad = dset_intensity_background.shape
layout = h5py.VirtualLayout(
shape=(nbpt_rad, y_nb_points, x_nb_points),
dtype=dset_intensity_background.dtype,
)
source = h5py.VirtualSource(dset_intensity_background)
for i in range(y_nb_points):
for j in range(x_nb_points):
layout[:, i, j] = source[i, j]
nxdata_background.create_virtual_dataset(
"map", layout, fillvalue=numpy.nan
).attrs["interpretation"] = "image"
nxdata_background.attrs["signal"] = "map"
# Link the diffmap nexus process to output_filename
if output_filename != external_output_filename:
with h5py_utils.open_item(
filename=output_filename,
name=entry_name,
mode="r+",
) as scan_entry:
create_hdf5_link(
parent=scan_entry,
link_name=nxprocess_background_name,
target=nxprocess_background_grp,
relative=True,
)
nxdata_background.attrs["finished"] = True
[docs]
def get_filtered_array(
self,
diffmap_array,
radial_array,
model_path,
wavelength,
radial_limits=None,
force_training=False,
use_neuralnetwork=True, # If False, skip the neural network filtering/training
):
from spi.math.filter_background import deep_filter_background, filter_background
from spi.utils import reformat
sh = diffmap_array.shape
flat_dataset = diffmap_array.reshape((sh[0] * sh[1], sh[2]))
flat_dataset[flat_dataset < 0] = 0
if radial_limits is not None:
radial_min, radial_max = radial_limits
x_min = numpy.argmin(abs(radial_array - radial_min))
x_max = numpy.argmin(abs(radial_array - radial_max))
else:
x_min = 0
x_max = len(radial_array)
# Filtered the dataset with pyobjcryst
if not use_neuralnetwork:
dataset_filtered = numpy.zeros_like(flat_dataset)
logger.warning("Skipping neural network filtering.")
dataset_filtered[:, x_min:x_max] = numpy.array(
filter_background(
dataset=flat_dataset[:, x_min:x_max],
wavelength=wavelength,
tth=radial_array[x_min:x_max],
)
)
return dataset_filtered.reshape(sh[0], sh[1], sh[2])
flat_dataset[:, 0:x_min] = 0.0
flat_dataset[:, x_max::] = 0.0
model_path = Path(model_path)
weights_file = model_path.with_suffix(".weights.h5")
if not weights_file.is_file():
# Train the model (will be done only once, then it will find the weights file)
logger.warning(
"There is not weight file. Prediction cannot be done without training! Redoing the training step..."
)
model_path = self.train_data(
flat_dataset=flat_dataset,
radial_array=radial_array,
wavelength=wavelength,
model_path=model_path,
)
filtered_map, norms = deep_filter_background(flat_dataset, model_path)
for i in range(norms.shape[0]):
filtered_map[i] = filtered_map[i] * norms[i]
nx, ny, ntth = diffmap_array.shape
return reformat(filtered_map.reshape(ny * nx, ntth), diffmap_array.shape)
[docs]
def train_data(self, flat_dataset, radial_array, model_path, wavelength) -> str:
from torch.cuda import get_device_name
logger.info(f"CUDA devices: {get_device_name()}")
logger.info(
"Training the neural network to remove the background. This may take some minutes..."
)
t0 = time.perf_counter()
from spi.math.filter_background import train_filtering_network
train_filtering_network(
flat_dataset,
radial_array,
wavelength,
Path(model_path),
downsampling_rate=0.1,
)
t1 = time.perf_counter()
logger.info(f"Training time: {t1 - t0} s")
return model_path
[docs]
class ConstrainedNMF(
SlurmTask,
input_names=[
"nxdata_url",
"wavelength",
"references_directory",
],
optional_input_names=[
"nb_components",
"output_filename",
"external_output_filename",
"nxprocess_name",
"radial_limits",
"do_matrix_factorization",
"processing_info",
],
output_names=[
"nxprocess_url",
"cnmf_nxprocess_path",
],
):
[docs]
def run(self):
warnings.simplefilter("ignore")
nxdata_url = self.inputs.nxdata_url
if not nxdata_url or not self.get_input_value("do_matrix_factorization", True):
self.outputs.nxprocess_url = None
self.outputs.cnmf_nxprocess_path = None
logger.warning("Matrix factorization is disabled. Skipping the CNMF task.")
return
input_filename, nxdata_filtered_path = nxdata_url.split("::")
h5path_parts = Path(f"/{nxdata_filtered_path}").parts
if len(h5path_parts) == 4:
entry_name = h5path_parts[1] # e.g. "1.1"
nxprocess_filtered_name = h5path_parts[2] # e.g. "eiger_diffmap"
else:
raise ValueError(
f"Invalid NXdata path format: {nxdata_filtered_path}. Expected format: {NXDATA_PATH_BACKGROUND_FORMAT}"
)
references_directory = self.inputs.references_directory
if not Path(references_directory).is_dir():
raise ValueError(
f"References directory '{references_directory}' does not exist or is not a directory."
)
nb_cif_files = len(list(Path(references_directory).glob("*.cif")))
nb_components = self.get_input_value("nb_components", nb_cif_files)
logger.info(
f"There are {nb_cif_files} .cif files in the target directory. {nb_components=}"
)
# Get output files
external_output_filename: str = self.get_input_value(
"external_output_filename", None
)
output_filename: str = self.get_input_value("output_filename", None)
if not external_output_filename and not output_filename:
external_output_filename = output_filename = input_filename
elif not external_output_filename:
external_output_filename = output_filename
elif not output_filename:
output_filename = external_output_filename
nxprocess_cnmf_name = self.get_input_value(
"nxprocess_name", nxprocess_filtered_name.replace("filtered", "cnmf")
)
if nxprocess_cnmf_name == nxprocess_filtered_name:
nxprocess_cnmf_name = nxprocess_filtered_name + "_cnmf"
nxprocess_cnmf_path = NXPROCESS_PATH_FORMAT.format(
entry_name=entry_name,
nxprocess_name=nxprocess_cnmf_name,
)
self.outputs.cnmf_nxprocess_path = nxprocess_cnmf_path
# Submit to Slurm if requested
if self.get_input_value("submit_to_slurm", False):
inputs = [
{"name": k, "value": v, "all": True}
for k, v in self.named_input_values.items()
if k != "submit_to_slurm"
] + [{"name": "submit_to_slurm", "value": False, "all": True}]
_ = self.execute_remote(workflow=CNMF_WORKFLOW, inputs=inputs)
return
# Open and wait until the file is ready (not recommended)
# timeout = 10
# t0 = time.perf_counter()
# filtered_data_available = False
# while not filtered_data_available:
# with h5py_utils.open_item(
# filename=input_filename, name=nxdata_filtered_path,
# ) as nxdata_filtered:
# if nxdata_filtered and nxdata_filtered.attrs.get("finished") == True:
# filtered_data_available = True
# if not filtered_data_available:
# time.sleep(10)
# if time.perf_counter() - t0 > timeout:
# raise TimeoutError(f"Timeout reached after {timeout} seconds.")
# logger.info("Filtered data is finished and available")
mode = "r+" if external_output_filename == input_filename else "r"
with h5py_utils.open_item(
filename=input_filename, name="/", mode=mode
) as file_input:
# Retrieve the filtered diffraction map data
nxdata_filtered = file_input[nxdata_filtered_path]
if "intensity" not in nxdata_filtered:
raise KeyError(
f"Filtered diffraction map {nxdata_filtered_path} does not contain 'intensity' dataset."
)
filtered_array = nxdata_filtered["intensity"][:]
nxdata_axes = nxdata_filtered.attrs.get("axes", [""])
radial_name = nxdata_axes[0]
radial_dset = (
nxdata_filtered[radial_name] if radial_name in nxdata_filtered else None
)
radial_array = radial_dset[:] if radial_dset is not None else None
if radial_array is None:
raise KeyError(
f"Radial dataset '{radial_name}' not found in {nxdata_filtered_path}. It's required for CNMF."
)
# Calculate the CNMF maps
cnmf_maps, diffractograms, phase_indexes = self.get_cnmf_maps(
diffmap_array=filtered_array,
radial_array=radial_array,
references_directory=Path(references_directory),
nb_components=nb_components,
radial_limits=self.get_input_value("radial_limits", None),
wavelength=self.inputs.wavelength,
)
# Write the result
with h5py_utils.open_item(
filename=external_output_filename,
name="/",
mode="r+",
) as file_output:
if nxprocess_cnmf_path in file_output:
raise FileExistsError(
f"NXprocess group {nxprocess_cnmf_path} already exists in the file {external_output_filename}. Please change the nxprocess_name parameter to avoid overwriting existing data."
)
# Create the NXprocess for CNMF results
nxprocess_cnmf_grp = file_output.create_group(nxprocess_cnmf_path)
nxprocess_cnmf_grp["name"] = "CNMF maps"
nxprocess_cnmf_grp.attrs.update(
{
"NX_class": "NXprocess",
}
)
for index, phase_name in phase_indexes.items():
phase_map = cnmf_maps[:, :, index]
diff = diffractograms[index]
nxdata_cnmf = nxprocess_cnmf_grp.create_group(
f"result_{phase_name}"
)
nxdata_cnmf.attrs.update(
{
"NX_class": "NXdata",
"signal": "intensity",
"title": f"CNMF: {phase_name}",
"axes": nxdata_axes[1:],
}
)
dset_cnmf_intensity = nxdata_cnmf.create_dataset(
name="intensity",
data=phase_map,
dtype="float32",
)
dset_cnmf_intensity.attrs["interpretation"] = "image"
nxdata_cnmf.create_dataset(
name="diffractogram",
data=diff,
dtype="float32",
)
if index == 0:
nxprocess_cnmf_grp.attrs.update(
{"default": f"result_{phase_name}"}
)
nxinfo_grp = nxprocess_cnmf_grp.create_group("processing_info")
nxinfo_grp.attrs["NX_class"] = "NXcollection"
for key, value in self.get_input_value(
"processing_info", {}
).items():
nxinfo_grp.attrs[key] = value
self.outputs.nxprocess_url = NXDATA_URL_FORMAT.format(
processed_data_filename=external_output_filename,
nxdata=f"{nxprocess_cnmf_path}/result_{phase_name}",
)
# Copy the radial information, and all the axes
for ax_name in nxdata_axes[1:]:
nxdata_filtered_axis = nxdata_filtered[ax_name]
virtual_source = h5py.VirtualSource(nxdata_filtered_axis)
virtual_layout = h5py.VirtualLayout(
shape=nxdata_filtered_axis.shape,
dtype=nxdata_filtered_axis.dtype,
)
virtual_layout[:] = virtual_source[:]
nxdata_cnmf.create_virtual_dataset(ax_name, virtual_layout)
# Link the diffmap nexus process to output_filename
if output_filename != external_output_filename:
with h5py_utils.open_item(
filename=output_filename,
name=entry_name,
mode="r+",
) as scan_entry:
create_hdf5_link(
parent=scan_entry,
link_name=nxprocess_cnmf_name,
target=nxprocess_cnmf_grp,
relative=True,
)
[docs]
def get_cnmf_maps(
self,
diffmap_array: numpy.ndarray, # shape=(nb_x_points, nb_y_points, pyFAI_radial_bins)
radial_array: numpy.ndarray, # shape=(pyFAI_radial_bins, )
references_directory: Path,
nb_components: int,
radial_limits=None,
wavelength=0.953734,
):
from spi import cnmf_analysis
from spi.utils import reformat
reference_phases = [
str(_.stem) for _ in Path(references_directory).glob("*.cif")
]
reference_phases.sort()
nb_cif_phases = len(reference_phases)
PHASE_INDEXES = {}
for ind, phase_file in enumerate(reference_phases):
PHASE_INDEXES[ind] = phase_file
nb_additional_phases = nb_components - nb_cif_phases
for ind_phase in range(nb_additional_phases):
PHASE_INDEXES[ind_phase + nb_cif_phases] = f"residual_{ind_phase:02}"
sh = diffmap_array.shape
flat_dataset = diffmap_array.reshape(sh[0] * sh[1], sh[2])
radial_min, radial_max = radial_limits
index_min = numpy.argmin(abs(radial_array - radial_min))
index_max = numpy.argmin(abs(radial_array - radial_max))
# Check on the wavelength
if not (0.1 < wavelength < 10):
raise ValueError(
f"Wavelength {wavelength} is not valid. It should be in Angstroms (0.1 - 10)."
)
logger.warning("Factorizing the matrix. This may take a while...")
t0 = time.perf_counter()
cnmf = cnmf_analysis(
flat_dataset,
nb_components,
Path(references_directory),
radial_array,
wavelength,
index_interval=(index_min, index_max),
instrument=None,
)
t1 = time.perf_counter()
logger.info(f"Time to factorize: {t1 - t0} s")
maps = reformat(cnmf.W.detach().numpy(), (sh[0], sh[1], nb_components))
diffractograms = cnmf.H.detach().numpy()
return maps, diffractograms, PHASE_INDEXES
[docs]
class PhaseInference(
SlurmTask,
input_names=[
"nxdata_url",
"references_directory",
],
optional_input_names=[
"output_filename",
"external_output_filename",
"nxprocess_name",
"radial_limits",
"do_phase_inference",
"do_training",
"inference_weights_filename",
"processing_info",
],
output_names=[
"nxprocess_url",
"inference_nxprocess_path",
],
):
[docs]
def run(self):
warnings.simplefilter("ignore")
nxdata_url = self.inputs.nxdata_url
if not nxdata_url or not self.get_input_value("do_phase_inference", True):
self.outputs.nxprocess_url = None
self.outputs.inference_nxprocess_path = None
logger.warning("Phase inference is disabled. Skipping this task.")
return
input_filename, nxdata_filtered_path = nxdata_url.split("::")
h5path_parts = Path(f"/{nxdata_filtered_path}").parts
if len(h5path_parts) == 4:
entry_name = h5path_parts[1] # e.g. "1.1"
nxprocess_filtered_name = h5path_parts[2] # e.g. "eiger_diffmap"
else:
raise ValueError(
f"Invalid NXdata path format: {nxdata_filtered_path}. Expected format: {NXDATA_PATH_BACKGROUND_FORMAT}"
)
references_directory = self.inputs.references_directory
if not Path(references_directory).is_dir():
raise ValueError(
f"References directory '{references_directory}' does not exist or is not a directory."
)
# Get output files
external_output_filename: str = self.get_input_value(
"external_output_filename", None
)
output_filename: str = self.get_input_value("output_filename", None)
if not external_output_filename and not output_filename:
external_output_filename = output_filename = input_filename
elif not external_output_filename:
external_output_filename = output_filename
elif not output_filename:
output_filename = external_output_filename
nxprocess_inference_name = self.get_input_value(
"nxprocess_name", nxprocess_filtered_name.replace("filtered", "inference")
)
if nxprocess_inference_name == nxprocess_filtered_name:
nxprocess_inference_name = nxprocess_filtered_name + "_inference"
nxprocess_inference_path = NXPROCESS_PATH_FORMAT.format(
entry_name=entry_name,
nxprocess_name=nxprocess_inference_name,
)
self.outputs.inference_nxprocess_path = nxprocess_inference_path
inference_weights_filename = self.get_input_value(
"inference_weights_filename", None
)
if not inference_weights_filename:
logger.error(
"There is no trained model for PhaseInference. Training has to be done before the experiment!"
)
return
elif not Path(inference_weights_filename).is_file():
raise FileNotFoundError(f"{inference_weights_filename} could not be found")
# Submit to Slurm if requested
if self.get_input_value("submit_to_slurm", False):
inputs = [
{"name": k, "value": v, "all": True}
for k, v in self.named_input_values.items()
if k != "submit_to_slurm"
] + [{"name": "submit_to_slurm", "value": False, "all": True}]
_ = self.execute_remote(workflow=INFERENCE_WORKFLOW, inputs=inputs)
return
mode = "r+" if external_output_filename == input_filename else "r"
with h5py_utils.open_item(
filename=input_filename, name="/", mode=mode
) as file_input:
if nxdata_filtered_path not in file_input:
raise FileNotFoundError(
f"Filtered diffraction map {nxdata_filtered_path} not found in the file {input_filename}. Skipping CNMF task."
)
# Retrieve the filtered diffraction map data
nxdata_filtered = file_input[nxdata_filtered_path]
if "intensity" not in nxdata_filtered:
raise KeyError(
f"Filtered diffraction map {nxdata_filtered_path} does not contain 'intensity' dataset."
)
filtered_array = nxdata_filtered["intensity"][:]
nxdata_axes = nxdata_filtered.attrs.get("axes", [""])
radial_name = nxdata_axes[0]
radial_dset = (
nxdata_filtered[radial_name] if radial_name in nxdata_filtered else None
)
radial_array = radial_dset[:] if radial_dset is not None else None
if radial_array is None:
raise KeyError(
f"Radial dataset '{radial_name}' not found in {nxdata_filtered_path}. It's required for CNMF."
)
# Calculate the maps
probabilistic_maps, phase_indexes = self.get_probabilistic_maps(
diffmap_array=filtered_array,
references_directory=references_directory,
weights_filename=inference_weights_filename,
)
# Write the result
with h5py_utils.open_item(
filename=external_output_filename,
name="/",
mode="r+",
) as file_output:
if nxprocess_inference_name in file_output:
raise FileExistsError(
f"NXprocess group {nxprocess_inference_name} already exists in the file {external_output_filename}. Please change the nxprocess_name parameter to avoid overwriting existing data."
)
# Create the NXprocess for Phase Inference results
nxprocess_inference = file_output.create_group(nxprocess_inference_path)
nxprocess_inference.attrs.update(
{
"NX_class": "NXprocess",
}
)
for index, phase_name in phase_indexes.items():
phase_map = probabilistic_maps[:, :, index]
nxdata_inference = nxprocess_inference.create_group(
f"result_{phase_name}"
)
nxdata_inference.attrs.update(
{
"NX_class": "NXdata",
"signal": "intensity",
"title": f"Inference: {phase_name}",
"axes": nxdata_axes[1:],
}
)
nxdata_inference.create_dataset(
name="intensity",
data=phase_map.T,
dtype="float32",
)
if index == 0:
nxprocess_inference.attrs.update(
{"default": f"result_{phase_name}"}
)
nxinfo_grp = nxprocess_inference.create_group("processing_info")
nxinfo_grp.attrs["NX_class"] = "NXcollection"
for key, value in self.get_input_value(
"processing_info", {}
).items():
nxinfo_grp.attrs[key] = value
self.outputs.nxprocess_url = NXDATA_URL_FORMAT.format(
processed_data_filename=external_output_filename,
nxdata=f"{nxprocess_inference_path}/result_{phase_name}",
)
# Copy the radial information, and all the axes
for ax_name in nxdata_axes[1:]:
nxdata_filtered_axis = nxdata_filtered[ax_name]
virtual_source = h5py.VirtualSource(nxdata_filtered_axis)
virtual_layout = h5py.VirtualLayout(
shape=nxdata_filtered_axis.shape,
dtype=nxdata_filtered_axis.dtype,
)
virtual_layout[:] = virtual_source[:]
nxdata_inference.create_virtual_dataset(ax_name, virtual_layout)
# Link the diffmap nexus process to output_filename
if output_filename != external_output_filename:
with h5py_utils.open_item(
filename=output_filename,
name=entry_name,
mode="r+",
) as scan_entry:
create_hdf5_link(
parent=scan_entry,
link_name=nxprocess_inference_name,
target=nxprocess_inference,
relative=True,
)
[docs]
def get_probabilistic_maps(
self,
diffmap_array,
references_directory,
weights_filename,
):
from spi.math import normalize_signal_wise
from spi.models.ProbabilisticMixtureClassifier import (
ProbabilisticMixtureClassifier,
)
from spi.utils import reformat
sh = diffmap_array.shape
data_norm = normalize_signal_wise(diffmap_array.reshape(sh[0] * sh[1], sh[2]))
reference_phases = [
str(_.name) for _ in Path(references_directory).glob("*.cif")
]
reference_phases.sort()
PHASE_INDEXES = {}
phase_index = 0
for phase_file in reference_phases:
PHASE_INDEXES[phase_index] = phase_file
phase_index += 1
N_PHASES = len(PHASE_INDEXES)
model = ProbabilisticMixtureClassifier(n_phases=N_PHASES)
model.build(input_shape=(1, sh[2], 1))
model.summary()
model.load_weights(Path(weights_filename))
N_ITER = 25
result = model.MC_predict(
data_norm.reshape(sh[0] * sh[1], sh[2], 1), n_iter=N_ITER
)
# probabilist_maps.shape = (sh[1], sh[0], N_PHASES)
probabilistic_maps = reformat(result, (sh[1], sh[0], N_PHASES))
return probabilistic_maps, PHASE_INDEXES
[docs]
class PhaseInferenceTrainModel(
SlurmTask,
input_names=[
"nxdata_url",
"references_directory",
"wavelength",
"radial_limits",
],
optional_input_names=[
"inference_weights_filename",
"nb_synthetic_samples",
"nb_samples_mixture",
"max_phases_mixture",
"processing_info",
],
output_names=[
"inference_weights_filename",
],
):
[docs]
def run(self):
references_directory = self.inputs.references_directory
if not Path(references_directory).is_dir():
raise FileNotFoundError(f"{references_directory} could not be found")
# Retrieve the radial array from the file, needed to get the radial limits
filename_synthetic_data = self.generate_synthetic_data()
filename_weights_latest = self.train_with_synthetic_data(
filename_synthetic_data
)
self.outputs.inference_weights_filename = filename_weights_latest
[docs]
def generate_synthetic_data(self):
warnings.simplefilter("ignore")
from spi.data.Dataset import make_mixtures_dataset
from spi.data.utils import store_dataset_classification
nxdata_url = self.inputs.nxdata_url
input_filename, nxdata_path_integrate = nxdata_url.split("::")
radial_array = None
with h5py_utils.open_item(
filename=input_filename, name=nxdata_path_integrate, mode="r"
) as nxdata_integrate:
nxdata_integrate_axes = nxdata_integrate.attrs.get("axes", [""])
radial_name = nxdata_integrate_axes[-1]
radial_array = nxdata_integrate[radial_name][:]
if radial_array is None:
raise KeyError(
f"Radial dataset '{radial_name}' not found in {input_filename}:{nxdata_path_integrate}."
)
t0 = time.perf_counter()
logger.warning("Generating synthetic data. This may take a while...")
# Check on the wavelength
wavelength = self.inputs.wavelength
if not (0.1 < wavelength < 10):
raise ValueError(
f"Wavelength {wavelength} is not valid. It should be in Angstroms (0.1 - 10)."
)
synthetic_patterns, labels = make_mixtures_dataset(
Path(self.inputs.references_directory),
n_samples=self.get_input_value("nb_synthetic_samples", 1000),
n_samples_mixture=self.get_input_value("nb_samples_mixture", 750),
patterns_size=len(radial_array),
max_phases_mixture=self.get_input_value("nb_samples_mixture", 3),
min_angle=radial_array.min(),
max_angle=radial_array.max(),
source_wavelength=wavelength,
min_contribution=0,
)
t1 = time.perf_counter()
logger.warning(
f"Time to calculate {synthetic_patterns.shape[0]} synthetic patterns: {t1-t0} s"
)
# TODO: spi more flexible with the .cif files and the models
model_file = self.get_input_value("inference_weights_filename", None)
if model_file:
subdirectory_models = Path(model_file).parent
else:
subdirectory_models = (
Path(self.inputs.references_directory).parent / "models_inference"
)
subdirectory_models.mkdir(parents=True, exist_ok=True)
filename_synthetic_data = subdirectory_models / "synthetic"
_ = store_dataset_classification(
path=filename_synthetic_data,
patterns=synthetic_patterns,
labels=labels,
)
return filename_synthetic_data
[docs]
def train_with_synthetic_data(self, filename_synthetic_data: str):
import keras
from spi.io import simple_load_training_dataset
from spi.models import ProbabilisticMixtureClassifier
nxdata_url = self.inputs.nxdata_url
input_filename, nxdata_path_integrate = nxdata_url.split("::")
radial_array = None
with h5py_utils.open_item(
filename=input_filename, name=nxdata_path_integrate, mode="r"
) as nxdata_integrate:
nxdata_integrate_axes = nxdata_integrate.attrs.get("axes", [""])
radial_name = nxdata_integrate_axes[-1]
radial_array = nxdata_integrate[radial_name][:]
if radial_array is None:
raise KeyError(
f"Radial dataset '{radial_name}' not found in {input_filename}:{nxdata_path_integrate}."
)
radial_limits = self.inputs.radial_limits
radial_min, radial_max = radial_limits
index_min = numpy.argmin(abs(radial_array - radial_min))
index_max = numpy.argmin(abs(radial_array - radial_max))
dataset_load = simple_load_training_dataset(filename_synthetic_data)
dataset_load["x_train"][:, :index_min] = 0
dataset_load["x_train"][:, index_max:] = 0
dataset_load["x_val"][:, :index_min] = 0
dataset_load["x_val"][:, index_max:] = 0
N_PHASES = dataset_load["y_train"].shape[1]
labels_list_train = []
labels_list_val = []
labels_list_test = []
train_samples = dataset_load["y_train"].shape[0]
val_samples = dataset_load["y_val"].shape[0]
test_samples = dataset_load["y_test"].shape[0]
for p in range(N_PHASES):
labels_array_train = dataset_load["y_train"][:, p].reshape(train_samples, 1)
labels_array_val = dataset_load["y_val"][:, p].reshape(val_samples, 1)
labels_array_test = dataset_load["y_test"][:, p].reshape(test_samples, 1)
labels_list_train.append(labels_array_train)
labels_list_val.append(labels_array_val)
labels_list_test.append(labels_array_test)
loss = keras.losses.BinaryCrossentropy()
metric = keras.metrics.BinaryAccuracy()
learning_rate = 2e-4 # dividing factor applied to the loss gradients : take small step towards the optimum
gradient_optimizer = keras.optimizers.Adam(learning_rate)
BATCH_SIZE = 64
model = ProbabilisticMixtureClassifier(n_phases=N_PHASES)
ntth = dataset_load["x_train"].shape[1]
t0 = time.perf_counter()
model.build(input_shape=(BATCH_SIZE, ntth, 1))
losses = []
metrics = []
for m in range(N_PHASES):
losses.append("binary_crossentropy")
metrics.append("binary_accuracy")
model.compile(loss=loss, optimizer=gradient_optimizer, metrics=[metric])
model.summary()
model_file = self.get_input_value("inference_weights_filename", None)
if model_file:
subdirectory_models = Path(model_file).parent
else:
subdirectory_models = (
Path(self.inputs.references_directory).parent / "models_inference"
)
model_file = subdirectory_models / "inference_model.weights.h5"
subdirectory_models.mkdir(parents=True, exist_ok=True)
save_best = keras.callbacks.ModelCheckpoint(
Path(model_file),
monitor="val_loss",
verbose=0,
save_best_only=True,
save_weights_only=True,
mode="min",
save_freq="epoch",
initial_value_threshold=None,
)
t0 = time.perf_counter()
logger.warning("Training the model. This may take a while...")
_ = model.fit(
x=dataset_load["x_train"].reshape(
dataset_load["x_train"].shape[0], ntth, 1
),
y=dataset_load["y_train"],
batch_size=BATCH_SIZE,
epochs=64,
verbose=2,
validation_split=0.1,
callbacks=[save_best],
)
logger.info(f"Time to train the model: {time.perf_counter()-t0}s")
# filename_weights_latest = str(model_file).replace(
# ".weights.h5", "_latest.weights.h5"
# )
model.save_weights(Path(model_file))
print(f"Successfully saved the model weights to {model_file}")
return str(model_file)