cylindra.widgets.batch

CylindraBatchWidget

Methods are available in the namespace ui.batch.

Source code in cylindra/widgets/batch/main.py

@magicclass(
    widget_type="split",
    layout="horizontal",
    name="Batch Analysis",
    properties={"min_height": 360},
    symbol=Expr("getattr", [Symbol("ui"), "batch"]),
)
class CylindraBatchWidget(MagicTemplate):
    constructor = field(ProjectSequenceEdit)
    sta = field(BatchSubtomogramAveraging)
    loader_infos = BatchLoaderAccessor()

    def __init__(self):
        self._loaders = LoaderList()
        self._loaders.events.inserted.connect(self.reset_choices)
        self._loaders.events.removed.connect(self.reset_choices)
        self._loaders.events.moved.connect(self.reset_choices)

    def __post_init__(self):
        self.native.setSizePolicy(
            QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding
        )
        self.sta.visible = False

    def _get_loader_paths(self, *_) -> list[PathInfo]:
        return [prj._get_loader_paths() for prj in self.constructor.projects]

    def _get_expression(self, *_):
        return self.constructor._get_expression()

    def _get_constructor_scale(self, *_) -> float:
        return self.constructor.scale.value

    @set_design(text=capitalize, location=ProjectSequenceEdit.File)
    @do_not_record
    def new_projects(
        self,
        paths: Path.Multiple[FileFilter.IMAGE],
        save_root: Path.Save[FileFilter.DIRECTORY],
        ref_paths: Path.Multiple[FileFilter.IMAGE] = [],
        scale: Annotated[Optional[float], {"text": "Use image original scale", "options": {"min": 0.01, "step": 0.0001},}] = None,
        tilt_model: Annotated[dict, {"widget_type": TiltModelEdit}] = None,
        bin_size: list[int] = [1],
        invert: bool = False,
        extension: Literal["", ".zip", ".tar"] = "",
        strip_prefix: str = "",
        strip_suffix: str = "",
        overwrite: bool = True,
    ):  # fmt: skip
        """Create new projects from images.

        This method is usually used for batch processing, efficient visual inspection,
        and particle picking.

        Parameters
        ----------
        paths : list of str or Path
            A list of image paths or wildcard patterns, such as "path/to/*.mrc".
        save_root : str or Path
            The root directory to save the output projects.
        ref_paths : list of str or Path, optional
            A list of reference image paths. Reference images are usually a binned or
            denoised version of the original images.
        scale : float, optional
            The scale of the images in nanometers. If None, the original scale of the
            images will be used.
        tilt_model : dict, optional
            A tilt model that describes the tilt angles and axis.
        bin_size : list of int, default [1]
            Initial bin size of image. Binned image will be used for visualization in
            the viewer. You can use both binned and non-binned image for analysis.
        invert : bool, default False
            Whether to invert the image intensities.
        extension : str, default ""
            The file extension (or directory) of the saved project files.
        strip_prefix : str, default ""
            A prefix to strip from the project name.
        strip_suffix : str, default ""
            A suffix to strip from the project name.
        overwrite : bool, default True
            If child project files of the same name already exist under the save root,
            they will be overwritten. This is useful when cylindra batch project is
            imported from file outputs of a long-running job from other softwares.
        """
        _paths = unwrap_wildcard(paths)
        self._new_projects_from_table(
            _paths,
            save_root=save_root,
            ref_paths=unwrap_wildcard(ref_paths) or None,
            scale=[scale] * len(_paths),
            tilt_model=[tilt_model] * len(_paths),
            bin_size=[bin_size] * len(_paths),
            invert=[invert] * len(_paths),
            extension=extension,
            strip_prefix=strip_prefix,
            strip_suffix=strip_suffix,
            overwrite=overwrite,
        )

    def _new_projects_from_table(
        self,
        path: list[Path],
        save_root: Path,
        ref_paths: list[Path] | None = None,
        scale: list[float | None] | None = None,
        tilt_model: list[dict | None] | None = None,
        bin_size: list[list[int]] | None = None,
        invert: list[bool] | None = None,
        splines: list[list["CylSpline"]] | None = None,
        molecules: list[dict[str, Molecules]] | None = None,
        molecule_sources: list[dict[str, int | None]] | None = None,
        extension: Literal["", ".zip", ".tar"] = "",
        strip_prefix: str = "",
        strip_suffix: str = "",
        overwrite: bool = True,
    ):
        projects = list[tuple[CylindraProject, str]]()
        num_projects = len(path)
        for img_path, _scale, _ref, tlt, _bin_size, _inv in zip(
            path,
            _or_default_list(scale, None, num_projects),
            _or_default_list(ref_paths, None, num_projects),
            _or_default_list(tilt_model, None, num_projects),
            _or_default_list(bin_size, [1], num_projects),
            _or_default_list(invert, False, num_projects),
            strict=True,
        ):
            each_project = CylindraProject.new(
                img_path,
                scale=_scale,
                image_reference=_ref,
                multiscales=_bin_size,
                missing_wedge=tlt,
                invert=_inv,
            )
            prj_name = img_path.stem
            projects.append((each_project, prj_name))
        self._add_new_projects(
            projects,
            save_root=save_root,
            splines=splines,
            molecules=molecules,
            molecule_sources=molecule_sources,
            extension=extension,
            strip_prefix=strip_prefix,
            strip_suffix=strip_suffix,
            overwrite=overwrite,
        )

    def _add_new_projects(
        self,
        projects: list[tuple[CylindraProject, str]],
        save_root: Path,
        splines: list[list["CylSpline"]] | None = None,
        molecules: list[dict[str, Molecules]] | None = None,
        molecule_sources: list[dict[str, int | None]] | None = None,
        extension: Literal["", ".zip", ".tar"] = "",
        strip_prefix: str = "",
        strip_suffix: str = "",
        overwrite: bool = True,
    ):
        """Save the new (project, name) pairs under `save_root` and add them."""
        num_projects = len(projects)
        if num_projects == 0:
            raise ValueError("No projects created.")
        save_root = Path(save_root)
        save_root.mkdir(parents=True, exist_ok=True)
        self.constructor.projects.clear()
        for (prj, prj_name), spl, mole, mole_src in zip(
            projects,
            splines or [[]] * num_projects,
            molecules or [{}] * num_projects,
            molecule_sources or [{}] * num_projects,
            strict=True,
        ):
            if strip_prefix and prj_name.startswith(strip_prefix):
                prj_name = prj_name[len(strip_prefix) :]
            if strip_suffix and prj_name.endswith(strip_suffix):
                prj_name = prj_name[: -len(strip_suffix)]
            save_path = save_root / f"{prj_name}{extension}"
            if save_path.exists() and not overwrite:
                prj = CylindraProject.from_file(save_path)
            else:
                prj.save(
                    save_path, splines=spl, molecules=mole, molecule_sources=mole_src
                )
                prj.project_path = save_path
            wdt = self.constructor.projects._add(prj.project_path)
            # same as `add_projects`
            self.constructor.scale.value = wdt.project.scale
        self.save_batch_project(save_path=save_root)
        self.show()

    @set_design(text=capitalize, location=constructor)
    @thread_worker
    def construct_loader(
        self,
        paths: Annotated[Any, {"bind": _get_loader_paths}],
        predicate: Annotated[str | pl.Expr | None, {"bind": _get_expression}] = None,
        name: str = "Loader",
        scale: Annotated[float | None, {"bind": _get_constructor_scale}] = None,
    ):  # fmt: skip
        """Construct a batch loader object from the given paths and predicate.

        Parameters
        ----------
        paths : list of (Path, list[Path]) or list of (Path, list[Path], Path)
            List of tuples of image path, list of molecule paths, and project path. The
            project path is optional.
        predicate : str or polars expression, optional
            Filter predicate of molecules.
        name : str, default "Loader"
            Name of the loader.
        """
        if name == "":
            raise ValueError("Name must be given.")

        yield 0.0, 0.0  # this function yields the progress
        loader = BatchLoader()
        image_paths = dict[int, Path]()
        invert = dict[int, bool]()
        _temp_feat = TempFeatures()
        for img_id, path_info in enumerate(paths):
            path_info = PathInfo(*path_info)
            # subtomograms of a tilt-series project are reconstructed from the tilt
            # series, while the others are loaded from the tomograms.
            ts_loader = path_info.tilt_series_loader()
            img = path_info.lazy_imread() if ts_loader is None else None
            image_paths[img_id] = Path(path_info.image)
            invert[img_id] = path_info.need_invert
            if scale is None:
                if prj := path_info.project_instance():
                    scale = prj.scale
                else:
                    scale = img.scale.x
            if prj := path_info.project_instance():
                tilt = prj.missing_wedge.as_param()
                if tilt is not None:
                    tilt = parse_tilt_model(tilt)
            else:
                tilt = None
            # molecules of a tilt series are in nm of the tilt series, so they are
            # not rescaled
            mole_scale = scale if ts_loader is None else ts_loader.scale
            for molecule_id, mole in enumerate(
                path_info.iter_molecules(_temp_feat, mole_scale)
            ):
                if ts_loader is None:
                    loader.add_tomogram(img.value, mole, img_id, tilt_model=tilt)
                else:
                    loader.add_loader(ts_loader.replace(molecules=mole), img_id)
                yield img_id / len(paths), molecule_id / len(path_info.molecules)
            yield (img_id + 1) / len(paths), 0.0
        yield 1.0, 1.0

        if predicate is not None:
            if isinstance(predicate, str):
                predicate = eval(predicate, POLARS_NAMESPACE, {})
            loader = loader.filter(predicate)
        new = loader.replace(
            molecules=loader.molecules.drop_features(_temp_feat.to_drop),
            scale=scale,
        )

        @thread_worker.callback
        def _on_return():
            self._add_loader(new, name, image_paths, invert)

        return _on_return

    @construct_loader.yielded.connect
    def _on_construct_loader_yielded(self, prog: tuple[float, float]):
        btn = get_button(self.construct_loader, cache=True)
        btn.text = f"Constructing... ({prog[0]:.1%}, {prog[1]:.1%})"

    @construct_loader.finished.connect
    def _on_construct_loader_finished(self):
        btn = get_button(self.construct_loader, cache=True)
        btn.text = "Construct loader"

    @set_design(text=capitalize, location=ProjectSequenceEdit.File)
    def construct_loader_by_list(
        self,
        project_paths: Path.Multiple[FileFilter.PROJECT],
        mole_pattern: str = "*",
        predicate: Annotated[str | pl.Expr | None, {"bind": _get_expression}] = None,
        name: str = "Loader",
    ):
        """Construct a batch loader from a list of project paths and a molecule pattern.

        Parameters
        ----------
        project_paths : list of path-like
            All the project paths to be used for construction. Entries can contain
            glob patterns such as "*" and "?".
        mole_pattern : str, default "*"
            A glob pattern for molecule file names. For example, "*-ALN1.csv" will only
            collect the molecule file names ends with "-ALN1.csv".
        predicate : str or polars expression, optional
            Filter predicate of molecules.
        name : str, default "Loader"
            Name of the loader.
        """
        self.constructor.add_projects(project_paths, clear=True)
        self.constructor.select_molecules_by_pattern(mole_pattern)
        self.construct_loader(self._get_loader_paths(), predicate=predicate, name=name)

    def _add_loader(
        self,
        loader: BatchLoader,
        name: str,
        image_paths: dict[int, Path],
        invert: dict[int, bool],
    ):
        self._loaders.append(LoaderInfo(loader, name, image_paths, invert))
        self.sta.visible = True
        try:
            self.sta["loader_name"].value = self.sta["loader_name"].choices[-1]
        except Exception:
            pass  # Updating the value is not important. Silence just in case.

    @set_design(text=capitalize, location=ProjectSequenceEdit.MacroMenu)
    @do_not_record
    def show_macro(self):
        """Show the macro widget of the batch analyzer."""
        from cylindra import instance

        ui = instance()
        assert ui is not None
        macro_str = self.macro.widget.textedit.value
        ui.OthersMenu.Macro._get_macro_window(macro_str, "Batch")

    @set_design(text=capitalize, location=ProjectSequenceEdit.MacroMenu)
    @do_not_record
    def show_native_macro(self):
        """Show the native macro widget of the batch analyzer."""
        self.macro.widget.show()
        ACTIVE_WIDGETS.add(self.macro.widget)

    @set_design(text="Load batch analysis project", location=ProjectSequenceEdit.File)
    @confirm(
        text="Are you sure to clear all loaders?", condition="len(self._loaders) > 0"
    )
    def load_batch_project(self, path: Path.Read[FileFilter.PROJECT]):
        """Load a batch project from a JSON file.

        Parameters
        ----------
        path : path-like
            Path to the JSON file.
        """
        self._loaders.clear()
        return CylindraBatchProject.from_file(path)._to_gui(self)

    @set_design(
        text="Save as batch analysis project", location=ProjectSequenceEdit.File
    )
    @do_not_record
    def save_batch_project(
        self,
        save_path: Path.Save,
        molecules_ext: Literal[".csv", ".parquet"] = ".csv",
    ):
        """Save the GUI state to a JSON file.

        Parameters
        ----------
        save_path : path-like
            Path to the JSON file.
        molecules_ext : str, default ".csv"
            Extension of the molecule files.
        """
        return CylindraBatchProject.save_gui(self, Path(save_path), molecules_ext)

    @nogui
    def iter_projects(self) -> Iterator[CylindraProject]:
        """Iterate over all projects in the batch project."""
        for prj_widget in self.constructor.projects:
            yield prj_widget.project

construct_loader(paths, predicate=None, name='Loader', scale=None)

Construct a batch loader object from the given paths and predicate.

Parameters:

Name	Type	Description	Default
paths	list of (Path, list[Path]) or list of (Path, list[Path], Path)	List of tuples of image path, list of molecule paths, and project path. The project path is optional.	required
predicate	str or polars expression	Filter predicate of molecules.	None
name	str	Name of the loader.	"Loader"

Source code in cylindra/widgets/batch/main.py

@set_design(text=capitalize, location=constructor)
@thread_worker
def construct_loader(
    self,
    paths: Annotated[Any, {"bind": _get_loader_paths}],
    predicate: Annotated[str | pl.Expr | None, {"bind": _get_expression}] = None,
    name: str = "Loader",
    scale: Annotated[float | None, {"bind": _get_constructor_scale}] = None,
):  # fmt: skip
    """Construct a batch loader object from the given paths and predicate.

    Parameters
    ----------
    paths : list of (Path, list[Path]) or list of (Path, list[Path], Path)
        List of tuples of image path, list of molecule paths, and project path. The
        project path is optional.
    predicate : str or polars expression, optional
        Filter predicate of molecules.
    name : str, default "Loader"
        Name of the loader.
    """
    if name == "":
        raise ValueError("Name must be given.")

    yield 0.0, 0.0  # this function yields the progress
    loader = BatchLoader()
    image_paths = dict[int, Path]()
    invert = dict[int, bool]()
    _temp_feat = TempFeatures()
    for img_id, path_info in enumerate(paths):
        path_info = PathInfo(*path_info)
        # subtomograms of a tilt-series project are reconstructed from the tilt
        # series, while the others are loaded from the tomograms.
        ts_loader = path_info.tilt_series_loader()
        img = path_info.lazy_imread() if ts_loader is None else None
        image_paths[img_id] = Path(path_info.image)
        invert[img_id] = path_info.need_invert
        if scale is None:
            if prj := path_info.project_instance():
                scale = prj.scale
            else:
                scale = img.scale.x
        if prj := path_info.project_instance():
            tilt = prj.missing_wedge.as_param()
            if tilt is not None:
                tilt = parse_tilt_model(tilt)
        else:
            tilt = None
        # molecules of a tilt series are in nm of the tilt series, so they are
        # not rescaled
        mole_scale = scale if ts_loader is None else ts_loader.scale
        for molecule_id, mole in enumerate(
            path_info.iter_molecules(_temp_feat, mole_scale)
        ):
            if ts_loader is None:
                loader.add_tomogram(img.value, mole, img_id, tilt_model=tilt)
            else:
                loader.add_loader(ts_loader.replace(molecules=mole), img_id)
            yield img_id / len(paths), molecule_id / len(path_info.molecules)
        yield (img_id + 1) / len(paths), 0.0
    yield 1.0, 1.0

    if predicate is not None:
        if isinstance(predicate, str):
            predicate = eval(predicate, POLARS_NAMESPACE, {})
        loader = loader.filter(predicate)
    new = loader.replace(
        molecules=loader.molecules.drop_features(_temp_feat.to_drop),
        scale=scale,
    )

    @thread_worker.callback
    def _on_return():
        self._add_loader(new, name, image_paths, invert)

    return _on_return

construct_loader_by_list(project_paths, mole_pattern='*', predicate=None, name='Loader')

Construct a batch loader from a list of project paths and a molecule pattern.

Parameters:

Name	Type	Description	Default
project_paths	list of path-like	All the project paths to be used for construction. Entries can contain glob patterns such as "*" and "?".	required
mole_pattern	str	A glob pattern for molecule file names. For example, "*-ALN1.csv" will only collect the molecule file names ends with "-ALN1.csv".	"*"
predicate	str or polars expression	Filter predicate of molecules.	None
name	str	Name of the loader.	"Loader"

Source code in cylindra/widgets/batch/main.py

@set_design(text=capitalize, location=ProjectSequenceEdit.File)
def construct_loader_by_list(
    self,
    project_paths: Path.Multiple[FileFilter.PROJECT],
    mole_pattern: str = "*",
    predicate: Annotated[str | pl.Expr | None, {"bind": _get_expression}] = None,
    name: str = "Loader",
):
    """Construct a batch loader from a list of project paths and a molecule pattern.

    Parameters
    ----------
    project_paths : list of path-like
        All the project paths to be used for construction. Entries can contain
        glob patterns such as "*" and "?".
    mole_pattern : str, default "*"
        A glob pattern for molecule file names. For example, "*-ALN1.csv" will only
        collect the molecule file names ends with "-ALN1.csv".
    predicate : str or polars expression, optional
        Filter predicate of molecules.
    name : str, default "Loader"
        Name of the loader.
    """
    self.constructor.add_projects(project_paths, clear=True)
    self.constructor.select_molecules_by_pattern(mole_pattern)
    self.construct_loader(self._get_loader_paths(), predicate=predicate, name=name)

iter_projects()

Iterate over all projects in the batch project.

Source code in cylindra/widgets/batch/main.py

@nogui
def iter_projects(self) -> Iterator[CylindraProject]:
    """Iterate over all projects in the batch project."""
    for prj_widget in self.constructor.projects:
        yield prj_widget.project

load_batch_project(path)

Load a batch project from a JSON file.

Parameters:

Name	Type	Description	Default
path	path - like	Path to the JSON file.	required

Source code in cylindra/widgets/batch/main.py

@set_design(text="Load batch analysis project", location=ProjectSequenceEdit.File)
@confirm(
    text="Are you sure to clear all loaders?", condition="len(self._loaders) > 0"
)
def load_batch_project(self, path: Path.Read[FileFilter.PROJECT]):
    """Load a batch project from a JSON file.

    Parameters
    ----------
    path : path-like
        Path to the JSON file.
    """
    self._loaders.clear()
    return CylindraBatchProject.from_file(path)._to_gui(self)

new_projects(paths, save_root, ref_paths=[], scale=None, tilt_model=None, bin_size=[1], invert=False, extension='', strip_prefix='', strip_suffix='', overwrite=True)

Create new projects from images.

This method is usually used for batch processing, efficient visual inspection, and particle picking.

Parameters:

Name	Type	Description	Default
paths	list of str or Path	A list of image paths or wildcard patterns, such as "path/to/*.mrc".	required
save_root	str or Path	The root directory to save the output projects.	required
ref_paths	list of str or Path	A list of reference image paths. Reference images are usually a binned or denoised version of the original images.	[]
scale	float	The scale of the images in nanometers. If None, the original scale of the images will be used.	None
tilt_model	dict	A tilt model that describes the tilt angles and axis.	None
bin_size	list of int	Initial bin size of image. Binned image will be used for visualization in the viewer. You can use both binned and non-binned image for analysis.	[1]
invert	bool	Whether to invert the image intensities.	False
extension	str	The file extension (or directory) of the saved project files.	""
strip_prefix	str	A prefix to strip from the project name.	""
strip_suffix	str	A suffix to strip from the project name.	""
overwrite	bool	If child project files of the same name already exist under the save root, they will be overwritten. This is useful when cylindra batch project is imported from file outputs of a long-running job from other softwares.	True

Source code in cylindra/widgets/batch/main.py

@set_design(text=capitalize, location=ProjectSequenceEdit.File)
@do_not_record
def new_projects(
    self,
    paths: Path.Multiple[FileFilter.IMAGE],
    save_root: Path.Save[FileFilter.DIRECTORY],
    ref_paths: Path.Multiple[FileFilter.IMAGE] = [],
    scale: Annotated[Optional[float], {"text": "Use image original scale", "options": {"min": 0.01, "step": 0.0001},}] = None,
    tilt_model: Annotated[dict, {"widget_type": TiltModelEdit}] = None,
    bin_size: list[int] = [1],
    invert: bool = False,
    extension: Literal["", ".zip", ".tar"] = "",
    strip_prefix: str = "",
    strip_suffix: str = "",
    overwrite: bool = True,
):  # fmt: skip
    """Create new projects from images.

    This method is usually used for batch processing, efficient visual inspection,
    and particle picking.

    Parameters
    ----------
    paths : list of str or Path
        A list of image paths or wildcard patterns, such as "path/to/*.mrc".
    save_root : str or Path
        The root directory to save the output projects.
    ref_paths : list of str or Path, optional
        A list of reference image paths. Reference images are usually a binned or
        denoised version of the original images.
    scale : float, optional
        The scale of the images in nanometers. If None, the original scale of the
        images will be used.
    tilt_model : dict, optional
        A tilt model that describes the tilt angles and axis.
    bin_size : list of int, default [1]
        Initial bin size of image. Binned image will be used for visualization in
        the viewer. You can use both binned and non-binned image for analysis.
    invert : bool, default False
        Whether to invert the image intensities.
    extension : str, default ""
        The file extension (or directory) of the saved project files.
    strip_prefix : str, default ""
        A prefix to strip from the project name.
    strip_suffix : str, default ""
        A suffix to strip from the project name.
    overwrite : bool, default True
        If child project files of the same name already exist under the save root,
        they will be overwritten. This is useful when cylindra batch project is
        imported from file outputs of a long-running job from other softwares.
    """
    _paths = unwrap_wildcard(paths)
    self._new_projects_from_table(
        _paths,
        save_root=save_root,
        ref_paths=unwrap_wildcard(ref_paths) or None,
        scale=[scale] * len(_paths),
        tilt_model=[tilt_model] * len(_paths),
        bin_size=[bin_size] * len(_paths),
        invert=[invert] * len(_paths),
        extension=extension,
        strip_prefix=strip_prefix,
        strip_suffix=strip_suffix,
        overwrite=overwrite,
    )

save_batch_project(save_path, molecules_ext='.csv')

Save the GUI state to a JSON file.

Parameters:

Name	Type	Description	Default
save_path	path - like	Path to the JSON file.	required
molecules_ext	str	Extension of the molecule files.	".csv"

Source code in cylindra/widgets/batch/main.py

@set_design(
    text="Save as batch analysis project", location=ProjectSequenceEdit.File
)
@do_not_record
def save_batch_project(
    self,
    save_path: Path.Save,
    molecules_ext: Literal[".csv", ".parquet"] = ".csv",
):
    """Save the GUI state to a JSON file.

    Parameters
    ----------
    save_path : path-like
        Path to the JSON file.
    molecules_ext : str, default ".csv"
        Extension of the molecule files.
    """
    return CylindraBatchProject.save_gui(self, Path(save_path), molecules_ext)

show_macro()

Show the macro widget of the batch analyzer.

Source code in cylindra/widgets/batch/main.py

@set_design(text=capitalize, location=ProjectSequenceEdit.MacroMenu)
@do_not_record
def show_macro(self):
    """Show the macro widget of the batch analyzer."""
    from cylindra import instance

    ui = instance()
    assert ui is not None
    macro_str = self.macro.widget.textedit.value
    ui.OthersMenu.Macro._get_macro_window(macro_str, "Batch")

show_native_macro()

Show the native macro widget of the batch analyzer.

Source code in cylindra/widgets/batch/main.py

@set_design(text=capitalize, location=ProjectSequenceEdit.MacroMenu)
@do_not_record
def show_native_macro(self):
    """Show the native macro widget of the batch analyzer."""
    self.macro.widget.show()
    ACTIVE_WIDGETS.add(self.macro.widget)

BatchSubtomogramAveraging

Methods are available in the namespace ui.batch.sta.

Source code in cylindra/widgets/batch/sta.py

@magicclass(name="Batch Subtomogram Analysis")
@_shared_doc.update_cls
class BatchSubtomogramAveraging(MagicTemplate):
    def _get_parent(self):
        from .main import CylindraBatchWidget

        return self.find_ancestor(CylindraBatchWidget, cache=True)

    def _get_loader_names(self, _=None) -> list[str]:
        try:
            parent = self._get_parent()
        except Exception:
            return []
        return parent.loader_infos.names()

    # Menus
    BatchSubtomogramAnalysis = field(
        BatchSubtomogramAnalysis, name="Subtomogram Analysis"
    )
    BatchRefinement = field(BatchRefinement, name="Refinement")
    BatchLoaderMenu = field(BatchLoaderMenu, name="Loader")

    @magicclass(layout="horizontal", properties={"margins": (0, 0, 0, 0)})
    class Header(MagicTemplate):
        loader_name = abstractapi()
        show_loader_info = abstractapi()
        remove_loader = abstractapi()

    loader_name = vfield(str, location=Header, record=False).with_choices(
        choices=_get_loader_names
    )

    def _get_current_loader_name(self, _=None) -> str:
        return self.loader_name

    @set_design(text="??", max_width=36, location=Header)
    @do_not_record
    def show_loader_info(self):
        """Show information about this loader"""
        info = self._get_parent().loader_infos[self.loader_name]
        loader = info.loader
        img_info = "\n" + "\n".join(
            f"{img_id}: {img_path}" for img_id, img_path in info.image_paths.items()
        )

        info_text = (
            f"name: {info.name}\nmolecule: n={loader.count()}\nimages:{img_info}"
        )
        view = DataFrameView(value=loader.molecules.to_dataframe())
        txt = ConsoleTextEdit(value=info_text)
        txt.read_only = True
        cnt = Container(widgets=[txt, view], layout="horizontal", labels=False)
        cnt.native.setParent(self.native, cnt.native.windowFlags())
        cnt.show()

    @set_design(text="✕", max_width=36, location=Header)
    def remove_loader(
        self, loader_name: Annotated[str, {"bind": _get_current_loader_name}]
    ):
        """Remove this loader"""
        del self._get_parent().loader_infos[loader_name]

    params = field(StaParameters)

    def _get_selected_loader_choice(self, *_) -> list[str]:
        try:
            loader = self.get_loader(self.loader_name)
            return loader.molecules.features.columns
        except Exception:
            return []

    def _get_template_path(self, _=None):
        return self.params.template_path.value

    def _get_mask_params(self, _=None):
        return self.params._get_mask_params()

    @set_design(text="Split loader", location=BatchLoaderMenu)
    def split_loader(
        self,
        loader_name: Annotated[str, {"bind": _get_current_loader_name}],
        by: Annotated[str, {"choices": _get_selected_loader_choice}],
        delete_old: bool = False,
    ):
        """Split the selected loader by the values of the given column.

        Parameters
        ----------
        loader_name : str
            Name of the input loader
        by : str
            Column name to split the loader
        delete_old : bool, default False
            If true, the original loader will be deleted.
        """
        parent = self._get_parent()
        loaders = parent._loaders
        batch_info = loaders[loader_name]
        batch_loader = batch_info.loader
        n_unique = batch_loader.molecules.features[by].n_unique()
        if n_unique > 48:
            raise ValueError(
                f"Too many groups ({n_unique}). Did you choose a float column?"
            )
        for _key, loader in batch_loader.groupby(by):
            existing_id = set(loader.features[Mole.image])
            image_paths = {
                k: v for k, v in batch_info.image_paths.items() if v in existing_id
            }
            invert = {k: v for k, v in batch_info.invert.items() if v in existing_id}
            parent._add_loader(loader, f"{loader_name}_{_key}", image_paths, invert)

        if delete_old:
            idx = -1
            for i, info in enumerate(loaders):
                if info.loader is batch_loader:
                    idx = i
                    break
            else:
                idx = -1
            if idx < 0:
                raise RuntimeError("Loader not found.")
            del loaders[idx]

    @set_design(text="Filter loader", location=BatchLoaderMenu)
    def filter_loader(
        self,
        loader_name: Annotated[str, {"bind": _get_current_loader_name}],
        expression: PolarsExprStr,
    ):
        """Filter the selected loader and add the filtered one to the list.

        Parameters
        ----------
        loader_name : str
            Name of the input loader
        expression : str
            polars expression that will be used to filter the loader. For example,
            `col("score") > 0.7` will filter out all low-score molecules.
        """
        loaderlist = self._get_parent()._loaders
        info = loaderlist[loader_name]
        loader = info.loader
        new = loader.filter(norm_polars_expr(expression))
        existing_id = set(new.features[Mole.image])
        loaderlist.add_loader(
            new,
            name=f"{info.name}-Filt",
            image_paths={k: v for k, v in info.image_paths.items() if v in existing_id},
            invert={k: v for k, v in info.invert.items() if v in existing_id},
        )
        return None

    @nogui
    def get_loader(self, name: str) -> BatchLoader:
        """Return the acryo.BatchLoader object with the given name"""
        if not isinstance(name, str):
            raise TypeError(f"Name must be a string, got {type(name).__name__}")
        return self._get_parent().loader_infos[name].loader

    @set_design(text="Average all molecules", location=BatchSubtomogramAnalysis)
    @dask_thread_worker.with_progress(desc="Averaging all molecules in projects")
    def average_all(
        self,
        loader_name: Annotated[str, {"bind": _get_current_loader_name}],
        size: _SubVolumeSize = None,
        interpolation: Annotated[int, {"choices": INTERPOLATION_CHOICES}] = 1,
        bin_size: _BINSIZE = 1,
    ):
        """Average all the molecules in the selected loader.

        Parameters
        ----------
        {loader_name}{size}{interpolation}{bin_size}
        """
        t0 = timer()
        loader = (
            self._get_parent()
            .loader_infos[loader_name]
            .loader.replace(order=interpolation)
            .binning(bin_size, compute=False)
        )
        shape = self._get_shape_in_px(size, loader)
        img = ip.asarray(
            loader.replace(output_shape=shape).order_optimize().average(),
            axes="zyx",
        ).set_scale(zyx=loader.scale, unit="nm")
        t0.toc()
        return self._show_rec.with_args(img, f"[AVG]{loader_name}")

    @set_design(text="Average group-wise", location=BatchSubtomogramAnalysis)
    @dask_thread_worker.with_progress(desc="Grouped subtomogram averaging")
    def average_groups(
        self,
        loader_name: Annotated[str, {"bind": _get_current_loader_name}],
        size: _SubVolumeSize = None,
        by: PolarsExprStr = "col('pf-id')",
        interpolation: Annotated[int, {"choices": INTERPOLATION_CHOICES}] = 1,
        bin_size: _BINSIZE = 1,
    ):
        """Groupwise subtomogram averaging using molecules grouped by the given expression.

        This method first group molecules by its features, and then average each group.
        This method is useful for such as get average of each protofilament and
        segmented subtomogram averaging.

        Parameters
        ----------
        {loader_name}{size}
        by : str or polars expression
            Expression to group molecules.
        {interpolation}{bin_size}
        """
        t0 = timer()
        loader = (
            self._get_parent()
            .loader_infos[loader_name]
            .loader.replace(order=interpolation)
            .binning(bin_size, compute=False)
        )
        shape = self._get_shape_in_px(size, loader)
        img = ip.asarray(
            loader.replace(output_shape=shape)
            .order_optimize()
            .groupby(norm_polars_expr(by))
            .average()
            .value_stack(axis=0),
            axes="pzyx",
        ).set_scale(zyx=loader.scale, unit="nm")
        t0.toc()
        return self._show_rec.with_args(img, f"[AVG]{loader_name}", store=False)

    @set_design(text="Align all molecules", location=BatchRefinement)
    @dask_thread_worker.with_progress(desc="Aligning all molecules")
    def align_all(
        self,
        loader_name: Annotated[str, {"bind": _get_current_loader_name}],
        template_path: Annotated[str | Path, {"bind": _get_template_path}],
        mask_params: Annotated[Any, {"bind": _get_mask_params}],
        max_shifts: _MaxShifts = (1.0, 1.0, 1.0),
        rotations: _Rotations = ((0.0, 0.0), (0.0, 0.0), (0.0, 0.0)),
        cutoff: _CutoffFreq = 0.5,
        interpolation: Annotated[int, {"choices": INTERPOLATION_CHOICES}] = 3,
        method: Annotated[str, {"choices": METHOD_CHOICES}] = "zncc",
        bin_size: _BINSIZE = 1,
    ):  # fmt: skip
        """
        Align all the molecules in the selected loader.

        Parameters
        ----------
        {loader_name}{template_path}{mask_params}{max_shifts}{rotations}{cutoff}
        {interpolation}{method}{bin_size}
        """
        t0 = timer()
        loaderlist = self._get_parent()._loaders
        info = loaderlist[loader_name]
        loader = info.loader.replace(order=interpolation).binning(
            bin_size, compute=False
        )
        template, mask = loader.normalize_input(
            template=self.params._norm_template_param(template_path),
            mask=self.params._get_mask(params=mask_params),
        )
        _Logger.print(f"Aligning {loader.molecules.count()} molecules ...")
        aligned = (
            loader.replace(output_shape=template.shape)
            .order_optimize()
            .align(
                template=template,
                mask=mask,
                max_shifts=max_shifts,
                rotations=rotations,
                cutoff=cutoff,
                alignment_model=_get_alignment(method),
            )
            .order_restore()
        )
        loaderlist.add_loader(
            aligned,
            name=_coerce_aligned_name(info.name, loaderlist),
            image_paths=info.image_paths,
            invert=info.invert,
        )
        t0.toc()

    @set_design(text="Align all (template-free)", location=BatchRefinement)
    @dask_thread_worker.with_progress(desc="Align all started")
    def align_all_template_free(
        self,
        loader_name: Annotated[str, {"bind": _get_current_loader_name}],
        mask_params: Annotated[Any, {"bind": _get_mask_params}],
        size: _SubVolumeSize = 12.0,
        max_shifts: _MaxShifts = (0.8, 0.8, 0.8),
        max_rotations: _MaxRotations = (3.0, 3.0, 3.0),
        min_rotation_step: Annotated[float, {"min": 0.1, "step": 0.1}] = 1.0,
        interpolation: Annotated[int, {"choices": INTERPOLATION_CHOICES}] = 3,
        method: Annotated[str, {"choices": METHOD_CHOICES}] = "zncc",
        bin_size: _BINSIZE = 1,
        max_num_iters: Annotated[int, {"min": 3, "max": 100}] = 20,
    ):  # fmt: skip
        """Create initial model by iteratively aligning molecules without template.

        This method iteratively calculate average, evaluate using FSC, and align
        molecules to the current average.

        Parameters
        ----------
        {loader_name}{mask_params}{size}{max_shifts}{max_rotations}{min_rotation_step}
        {interpolation}{method}{bin_size}{max_num_iters}
        """
        t0 = timer()
        rng = np.random.default_rng(1428)
        loaderlist = self._get_parent()._loaders
        info = loaderlist[loader_name]
        mask = self.params._get_mask(params=mask_params)
        loader = info.loader.replace(order=interpolation).binning(
            bin_size, compute=False
        )
        shape = self._get_shape_in_px(size, loader)
        loader = loader.replace(output_shape=shape).order_optimize()
        _Logger.print(f"Aligning {loader.molecules.count()} molecules ...")
        _alignment_state = template_free.AlignmentState(
            rng=rng,
            mask=mask,
            alignment_model=_get_alignment(method),
            min_rotation_step=min_rotation_step,
        )
        # loaders of each tomogram, which have their own tilt models
        sub_loaders = list(loader.loaders.iter_loaders())
        num_images = len(sub_loaders)
        with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdir:
            # Subtomograms are extracted to files only once, as spline coefficients so
            # that they can be resampled without the prefilter. In each iteration, the
            # extracted subtomograms are resampled at the aligned molecules for both
            # alignment and FSC calculation.
            extracted: dict[Hashable, ExtractedSubvolumeLoader] = {}
            for ith, (img_id, sub) in enumerate(sub_loaders):
                yield thread_worker.description(
                    f"Extracting subtomograms ({ith + 1}/{num_images})"
                )
                extracted[img_id] = sub.extract_subtomograms(
                    Path(tmpdir) / f"image-{ith}",
                    chunksize=64,
                    apply_prefilter=True,
                )
            # the image IDs are kept, as the aligned molecules are the output
            joined = join_extracted_loaders(extracted)
            yield thread_worker.description(_pdesc.align_tf_0(_alignment_state))
            result0 = _alignment_state.fsc_step_init(joined, max_shifts, max_rotations)
            _Logger.print(f"Iteration {int(_alignment_state.num_iter)}")
            yield _plot_current_fsc.with_args(result0.fsc, result0.avg).with_desc(_pdesc.align_tf_1(_alignment_state))  # fmt: skip
            while True:
                _Logger.print(_alignment_state.next_params(loader.scale).format())
                _exceeded = max_num_iters <= _alignment_state.num_iter
                if _alignment_state.is_converged() or _exceeded:
                    _Logger.print("Maximum iteration exceeded" if _exceeded else "FSC converged.")  # fmt: skip
                    break
                # the extracted subtomograms are resampled at the aligned molecules
                joined = _alignment_state.align_step(joined)
                yield thread_worker.description(_pdesc.align_tf_0(_alignment_state))
                result = _alignment_state.fsc_step(joined)
                _Logger.print(f"Iteration {int(_alignment_state.num_iter)}")
                yield _plot_current_fsc.with_args(result.fsc, result.avg).with_desc(_pdesc.align_tf_1(_alignment_state))  # fmt: skip

        mole_aligned = _as_alignment_output(loader.molecules, [joined.molecules])
        loaderlist.add_loader(
            loader.replace(molecules=mole_aligned).order_restore(),
            name=_coerce_aligned_name(info.name, loaderlist),
            image_paths=info.image_paths,
            invert=info.invert,
        )
        t0.toc()
        return self._show_rec.with_args(
            _alignment_state.results[-1].avg, f"[AVG]{loader_name}"
        )

    @set_design(text="RMA alignment", location=BatchRefinement)
    @dask_thread_worker.with_progress(desc="Simulated annealing (RMA)")
    def align_all_rma(
        self,
        loader_name: Annotated[str, {"bind": _get_current_loader_name}],
        template_path: Annotated[str | Path, {"bind": _get_template_path}],
        mask_params: Annotated[Any, {"bind": _get_mask_params}],
        max_shifts: _MaxShifts = (0.8, 0.8, 0.8),
        rotations: _Rotations = ((0.0, 0.0), (0.0, 0.0), (0.0, 0.0)),
        cutoff: _CutoffFreq = 0.5,
        interpolation: Annotated[int, {"choices": INTERPOLATION_CHOICES}] = 3,
        method: Annotated[str, {"choices": METHOD_CHOICES}] = "zncc",
        range_long: _DistRangeLon = (4.0, 4.28),
        range_lat: _DistRangeLat = ("d.mean() - 0.1", "d.mean() + 0.1"),
        angle_max: _AngleMaxLon = 5.0,
        bin_size: _BINSIZE = 1,
        temperature_time_const: Annotated[float, {"min": 0.01, "max": 10.0}] = 1.0,
        lj_const: Annotated[float, {"min": 0.00, "max": 1000.0, "step": 0.1, "label": "LJ const"}] = 0.0,
        upsample_factor: Annotated[int, {"min": 1, "max": 20}] = 5,
        num_trials: Annotated[int, {"min": 1, "max": 100}] = 5,
        seed: _SeedType = 0,
    ):  # fmt: skip
        """Run RMA alignment on all molecules in the selected loader.

        Parameters
        ----------
        {loader_name}{template_path}{mask_params}{max_shifts}{rotations}{cutoff}
        {interpolation}{method}{range_long}{range_lat}{angle_max}{bin_size}
        {temperature_time_const}{lj_const}{upsample_factor}{num_trials}{seed}
        """
        t0 = timer()
        batch = self._get_parent()
        loaderlist = batch._loaders
        info = loaderlist[loader_name]
        loader = info.loader.replace(order=interpolation).binning(
            bin_size, compute=False
        )
        template, mask = loader.normalize_input(
            template=self.params._norm_template_param(template_path),
            mask=self.params._get_mask(params=mask_params),
        )
        loader = loader.replace(output_shape=template.shape)
        aln = _get_alignment(method)
        sub_inputs = list(self._group_loader_by_spline(loader_name, loader))
        num_splines = len(sub_inputs)
        _Logger.print(f"{num_splines} splines found for RMA alignment.")

        @thread_worker.callback
        def _plot_annealing_result(results):
            with _Logger.set_plt():
                _annealing.plot_annealing_result(results)

        for ith, inputs in enumerate(sub_inputs):
            _p_ind = f"({ith + 1}/{num_splines})"
            yield thread_worker.description(f"Landscape construction {_p_ind}")
            landscape = Landscape.from_loader(
                loader=inputs.loader,
                template=template,
                mask=mask,
                max_shifts=max_shifts,
                upsample_factor=upsample_factor,
                alignment_model=aln.with_params(
                    rotations=rotations,
                    cutoff=cutoff,
                    tilt=inputs.loader.tilt_model,
                ),
            ).normed()
            yield thread_worker.description(f"RMA alignment {_p_ind}")
            mole, results = landscape.run_annealing_along_spline(
                inputs.spline,
                range_long,
                range_lat,
                angle_max=angle_max,
                temperature_time_const=temperature_time_const,
                lj_nstd=None if lj_const < 1e-3 else 1 / lj_const,
                random_seeds=create_random_seeds(num_trials, seed),
            )
            yield _plot_annealing_result.with_args(results)
            mole = mole.with_features(
                pl.lit(inputs.image_id).alias(Mole.image),
                pl.lit(inputs.molecule_id).alias(Mole.id),
            )
            inputs.loader = inputs.loader.replace(molecules=mole)

        loaderlist.add_loader(
            join_loaders(loader, sub_inputs),
            name=_coerce_aligned_name(info.name, loaderlist),
            image_paths=info.image_paths,
            invert=info.invert,
        )
        t0.toc()

    @set_design(text="RMA alignment (template-free)", location=BatchRefinement)
    @dask_thread_worker.with_progress(desc="RMA alignment started")
    def align_all_rma_template_free(
        self,
        loader_name: Annotated[str, {"bind": _get_current_loader_name}],
        mask_params: Annotated[Any, {"bind": _get_mask_params}],
        size: _SubVolumeSize = 12.0,
        max_shifts: _MaxShifts = (0.8, 0.8, 0.8),
        max_rotations: _MaxRotations = (0.0, 0.0, 0.0),
        interpolation: Annotated[int, {"choices": INTERPOLATION_CHOICES}] = 3,
        method: Annotated[str, {"choices": METHOD_CHOICES}] = "zncc",
        range_long: _DistRangeLon = (4.0, 4.28),
        range_lat: _DistRangeLat = ("d.mean() - 0.1", "d.mean() + 0.1"),
        angle_max: _AngleMaxLon = 5.0,
        bin_size: _BINSIZE = 1,
        temperature_time_const: Annotated[float, {"min": 0.01, "max": 10.0}] = 1.0,
        upsample_factor: Annotated[int, {"min": 1, "max": 20}] = 3,
        num_trials: Annotated[int, {"min": 1, "max": 100}] = 5,
        max_num_iters: Annotated[int, {"min": 3, "max": 100}] = 8,
        seed: _SeedType = 0,
    ):
        """Create initial model by iteratively aligning molecules by RMA without template.

        The annealing schedule of RMA is split into `max_num_iters` iterations. In
        each iteration, annealing is resumed from the last state and temperature of
        the previous iteration using the current average as the template, and then the
        average is updated and evaluated using FSC. If FSC converged, the rest of the
        annealing will be finished in the next iteration.

        Parameters
        ----------
        {loader_name}{mask_params}{size}{max_shifts}{max_rotations}
        {interpolation}{method}{range_long}{range_lat}{angle_max}{bin_size}
        {temperature_time_const}{upsample_factor}{num_trials}{max_num_iters}{seed}
        """
        t0 = timer()
        rng = np.random.default_rng(seed)
        batch = self._get_parent()
        loaderlist = batch._loaders
        info = loaderlist[loader_name]
        mask = self.params._get_mask(params=mask_params)
        loader = info.loader.replace(order=interpolation).binning(
            bin_size, compute=False
        )
        shape = self._get_shape_in_px(size, loader)
        loader = loader.replace(output_shape=shape)

        _alignment_state = template_free.RMAAlignmentState(
            rng=rng,
            mask=mask,
            alignment_model=_get_alignment(method),
            max_num_iters=max_num_iters,
        )
        sub_inputs = list(self._group_loader_by_spline(loader_name, loader))
        num_splines = len(sub_inputs)
        _Logger.print(f"{num_splines} splines found for RMA alignment.")
        annealings: list[template_free.ResumableAnnealing | None] = [None] * num_splines
        with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdir:
            # Subtomograms are extracted to files only once, as spline coefficients so
            # that they can be resampled without the prefilter. Landscapes are always
            # constructed around the initial molecules so that the annealing can be
            # resumed in the next iteration, and the averages for FSC are resampled
            # from the extracted subtomograms at the aligned molecules.
            for ith, inputs in enumerate(sub_inputs):
                yield thread_worker.description(
                    f"Extracting subtomograms ({ith + 1}/{num_splines})"
                )
                inputs.loader = inputs.loader.extract_subtomograms(
                    Path(tmpdir) / f"spline-{ith}", chunksize=64, apply_prefilter=True
                )
            extracted = [inputs.loader for inputs in sub_inputs]
            yield thread_worker.description(_pdesc.align_tf_0(_alignment_state))
            result = _alignment_state.fsc_step_init(
                join_extracted_loaders(extracted), max_shifts, max_rotations,
                upsample_factor, temperature_time_const, num_trials,
            )  # fmt: skip
            _final = False
            while True:
                _Logger.print(f"Iteration {int(_alignment_state.num_iter)}")
                yield _plot_current_fsc.with_args(result.fsc, result.avg).with_desc(_pdesc.align_tf_1(_alignment_state))  # fmt: skip
                if _final:
                    break
                if _final := _alignment_state.is_final_iteration():
                    _converged = _alignment_state.is_converged()
                    _Logger.print("FSC converged." if _converged else "Last iteration.")
                _Logger.print(_alignment_state.next_params(loader.scale).format())
                _ratio = _alignment_state.temperature_ratio(_final)
                _Logger.print(f"Annealing until temperature reaches {_ratio:.2e} of T0")
                _iter = _alignment_state.num_iter + 1
                molecules: list[Molecules] = []
                for ith, inputs in enumerate(sub_inputs):
                    _p_ind = f"({ith + 1}/{num_splines})"
                    yield thread_worker.description(
                        f"Landscape construction {_p_ind} of iteration {_iter}"
                    )
                    landscape = _alignment_state.built_landscape_step(inputs.loader)
                    annealing = annealings[ith]
                    if annealing is None:
                        annealing = _alignment_state.prep_annealing_step(
                            landscape, inputs.spline, range_long, range_lat, angle_max
                        )
                        annealings[ith] = annealing
                    else:
                        annealing.update_landscape(landscape)
                    yield thread_worker.description(
                        f"RMA step {_p_ind} of iteration {_iter}"
                    )
                    mole = _alignment_state.rma_step(annealing, final=_final)
                    molecules.append(
                        mole.with_features(
                            pl.lit(inputs.image_id).alias(Mole.image),
                            pl.lit(inputs.molecule_id).alias(Mole.id),
                        )
                    )
                yield thread_worker.description(_pdesc.align_tf_0(_alignment_state))
                result = _alignment_state.fsc_step(
                    join_extracted_loaders(extracted, molecules)
                )

        loaderlist.add_loader(
            join_loaders(loader, sub_inputs, molecules),
            name=_coerce_aligned_name(info.name, loaderlist),
            image_paths=info.image_paths,
            invert=info.invert,
        )
        t0.toc()
        return self._show_rec.with_args(
            _alignment_state.results[-1].avg, f"[AVG]{loader_name}"
        )

    @set_design(text="RFA alignment (template-free)", location=BatchRefinement)
    @dask_thread_worker.with_progress(desc="RFA alignment started")
    def align_all_rfa_template_free(
        self,
        loader_name: Annotated[str, {"bind": _get_current_loader_name}],
        mask_params: Annotated[Any, {"bind": _get_mask_params}],
        size: _SubVolumeSize = 12.0,
        max_shifts: _MaxShifts = (2, 2, 2),
        max_rotations: _MaxRotations = (0.0, 180.0, 0.0),
        interpolation: Annotated[int, {"choices": INTERPOLATION_CHOICES}] = 3,
        method: Annotated[str, {"choices": METHOD_CHOICES}] = "zncc",
        range_long: _DistRangeLon = (4.0, 10.0),
        angle_max: _AngleMaxLon = 5.0,
        optimize_polarity: bool = True,
        randomize_roll_angle: bool = False,
        bin_size: _BINSIZE = 1,
        temperature_time_const: Annotated[float, {"min": 0.01, "max": 50.0}] = 5.0,
        upsample_factor: Annotated[int, {"min": 1, "max": 20}] = 3,
        num_trials: Annotated[int, {"min": 1, "max": 100}] = 5,
        max_num_iters: Annotated[int, {"min": 3, "max": 100}] = 8,
        seed: _SeedType = 0,
    ):
        """Create initial model of filaments by iterative RFA without template.

        This method is useful to create an initial model from molecules mapped along
        manually initialized splines (such as by `map_along_spline`), where the
        lattice parameters (and the polarity of the filaments) are unknown.

        Parameters
        ----------
        {loader_name}{mask_params}{size}{max_shifts}{max_rotations}
        {interpolation}{method}{range_long}{angle_max}
        optimize_polarity : bool, default True
            If true, the polarity of each filament is optimized by running annealing
            in both orientations. Otherwise, the orientation of the input molecules
            is kept.
        randomize_roll_angle : bool, default False
            If true, molecules are rotated around their y-axis (the filament axis) by
            random angles before alignment, so that the missing wedges of the
            subtomograms are distributed around the axis in the average. Rotation
            around the y-axis should be searched (such as `max_rotations=(0, 180, 0)`)
            to recover the roll angles.
        {bin_size}{temperature_time_const}{upsample_factor}{num_trials}
        {max_num_iters}{seed}
        """
        t0 = timer()
        rng = np.random.default_rng(seed)
        batch = self._get_parent()
        loaderlist = batch._loaders
        info = loaderlist[loader_name]
        mask = self.params._get_mask(params=mask_params)
        loader = info.loader.replace(order=interpolation).binning(
            bin_size, compute=False
        )
        shape = self._get_shape_in_px(size, loader)
        loader = loader.replace(output_shape=shape)

        _alignment_state = template_free.RFAAlignmentState(
            rng=rng,
            mask=mask,
            alignment_model=_get_alignment(method),
            max_num_iters=max_num_iters,
        )
        sub_inputs = list(self._group_loader_by_spline(loader_name, loader))
        num_splines = len(sub_inputs)
        _Logger.print(f"{num_splines} splines found for RFA alignment.")
        # annealings in the forward (and reverse) orientations
        annealings: list[tuple[template_free.ResumableAnnealing, ...] | None] = [
            None
        ] * num_splines
        # molecules around which the first landscapes are built, in the forward (and
        # reverse) orientations
        initial_refs: list[tuple[Molecules, ...]] = []
        with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdir:
            # Same as RMA, subtomograms are extracted only once. Landscapes are always
            # built at the initial positions of molecules, but in the orientations of
            # the best results of the last iteration, so that the annealing is resumed
            # with rotations searched around them. Subvolumes in any orientation,
            # including the reverse one, are resampled from the extracted subtomograms.
            for ith, inputs in enumerate(sub_inputs):
                yield thread_worker.description(
                    f"Extracting subtomograms ({ith + 1}/{num_splines})"
                )
                inputs.loader = inputs.loader.extract_subtomograms(
                    Path(tmpdir) / f"spline-{ith}", chunksize=64, apply_prefilter=True
                )
                mole = inputs.loader.molecules
                if randomize_roll_angle:
                    mole = template_free.randomize_roll(mole, rng)
                if optimize_polarity:
                    mole_rev = mole.rotate_by_rotvec_internal([np.pi, 0, 0])
                    initial_refs.append((mole, mole_rev))
                else:
                    initial_refs.append((mole,))
            extracted = [inputs.loader for inputs in sub_inputs]
            yield thread_worker.description(_pdesc.align_tf_0(_alignment_state))
            result = _alignment_state.fsc_step_init(
                join_extracted_loaders(extracted, [refs[0] for refs in initial_refs]),
                max_shifts, max_rotations, upsample_factor, temperature_time_const,
                num_trials,
            )  # fmt: skip
            # intermediate results for debugging
            history = [
                _history_entry(result, [refs[0] for refs in initial_refs], sub_inputs)
            ]
            _final = False
            while True:
                _Logger.print(f"Iteration {int(_alignment_state.num_iter)}")
                yield _plot_current_fsc.with_args(result.fsc, result.avg).with_desc(_pdesc.align_tf_1(_alignment_state))  # fmt: skip
                if _final:
                    break
                if _final := _alignment_state.is_final_iteration():
                    _converged = _alignment_state.is_converged()
                    _Logger.print("FSC converged." if _converged else "Last iteration.")
                _Logger.print(_alignment_state.next_params(loader.scale).format())
                _ratio = _alignment_state.temperature_ratio(_final)
                _Logger.print(f"Annealing until temperature reaches {_ratio:.2e} of T0")
                _iter = _alignment_state.num_iter + 1
                molecules: list[Molecules] = []
                step_results: list[template_free.RFAStepResult] = []
                # molecules around which the landscapes are built
                references: list[tuple[Molecules, ...]] = []
                for ith, inputs in enumerate(sub_inputs):
                    _p_ind = f"({ith + 1}/{num_splines})"
                    yield thread_worker.description(
                        f"Landscape construction {_p_ind} of iteration {_iter}"
                    )
                    if (pair := annealings[ith]) is None:
                        refs = initial_refs[ith]
                    else:
                        refs = tuple(annealing.next_reference() for annealing in pair)
                    references.append(refs)
                    # landscapes are not normalized to compare the two orientations
                    landscapes = _alignment_state.built_landscapes_step(
                        [
                            template_free.resample_extracted(extracted[ith], ref)
                            for ref in refs
                        ],
                        norm=not optimize_polarity,
                    )
                    if pair is None:
                        pair = tuple(
                            _alignment_state.prep_annealing_step(
                                landscape, inputs.spline, range_long, angle_max
                            )
                            for landscape in landscapes
                        )
                        annealings[ith] = pair
                    else:
                        for annealing, landscape in zip(pair, landscapes, strict=True):
                            annealing.update_landscape(landscape)
                    yield thread_worker.description(
                        f"RFA step {_p_ind} of iteration {_iter}"
                    )
                    step = _alignment_state.rfa_step(*pair, final=_final)
                    step_results.append(step)
                    molecules.append(
                        step.molecules.with_features(
                            pl.lit(inputs.image_id).alias(Mole.image),
                            pl.lit(inputs.molecule_id).alias(Mole.id),
                        )
                    )
                if optimize_polarity:
                    _num_rev = sum(step.is_reversed for step in step_results)
                    _Logger.print(
                        f"Selected orientations: {num_splines - _num_rev} forward, "
                        f"{_num_rev} reverse"
                    )
                yield thread_worker.description(_pdesc.align_tf_0(_alignment_state))
                result = _alignment_state.fsc_step(
                    join_extracted_loaders(extracted, molecules)
                )
                history.append(
                    _history_entry(result, molecules, sub_inputs, references)
                )

        loaderlist.add_loader(
            join_loaders(loader, sub_inputs, molecules),
            name=_coerce_aligned_name(info.name, loaderlist),
            image_paths=info.image_paths,
            invert=info.invert,
        )
        t0.toc()

        @thread_worker.callback
        def _on_return():
            layer = self._show_rec(
                _alignment_state.results[-1].avg, f"[AVG]{loader_name}"
            )
            layer.metadata[TEMPLATE_FREE_HISTORY] = history

        return _on_return

    @set_design(text="Calculate FSC", location=BatchSubtomogramAnalysis)
    @dask_thread_worker.with_progress(desc="Calculating FSC")
    def calculate_fsc(
        self,
        loader_name: Annotated[str, {"bind": _get_current_loader_name}],
        template_path: Annotated[str | Path | None, {"bind": _get_template_path}] = None,
        mask_params: Annotated[Any, {"bind": _get_mask_params}] = None,
        size: _SubVolumeSize = None,
        seed: Annotated[Optional[int], {"text": "Do not use random seed."}] = 0,
        interpolation: Annotated[int, {"choices": INTERPOLATION_CHOICES}] = 1,
        n_pairs: Annotated[int, {"min": 1, "label": "number of image pairs"}] = 1,
        show_average: bool = True,
        dfreq: FSCFreq = None,
    ):  # fmt: skip
        """Calculate Fourier Shell Correlation using the selected loader.

        Parameters
        ----------
        {loader_name}{template_path}{mask_params}{size}
        seed : int, optional
            Random seed used for subtomogram sampling.
        {interpolation}
        n_pairs : int, default 1
            How many sets of image pairs will be generated to average FSC.
        show_average : bool, default True
            If true, subtomogram averaging will be shown after FSC calculation.
        dfreq : float, default 0.02
            Precision of frequency to calculate FSC. "0.02" means that FSC will be
            calculated at frequency 0.01, 0.03, 0.05, ..., 0.45.
        """
        t0 = timer()
        loader = (
            self._get_parent()
            .loader_infos[loader_name]
            .loader.replace(order=interpolation)
        )

        template, mask = loader.normalize_input(
            template=self.params._norm_template_param(template_path, allow_none=True),
            mask=self.params._get_mask(params=mask_params),
        )

        fsc, (img_0, img_1), img_mask = (
            loader.reshape(
                template=template if size is None else None,
                mask=mask,
                shape=None if size is None else self._get_shape_in_px(size, loader),
            )
            .order_optimize()
            .fsc_with_halfmaps(
                mask, seed=seed, n_set=n_pairs, dfreq=dfreq, squeeze=False
            )
        )

        def _as_imgarray(im, axes: str = "zyx") -> ip.ImgArray | None:
            if np.isscalar(im):
                return None
            return ip.asarray(im, axes=axes).set_scale(zyx=loader.scale, unit="nm")

        if show_average:
            avg = (img_0[0] + img_1[0]) / 2
            img_avg = _as_imgarray(avg)
        else:
            img_avg = None

        result = FscResult.from_dataframe(fsc, loader.scale)
        criteria = [0.5, 0.143]
        t0.toc()

        @thread_worker.callback
        def _calculate_fsc_on_return():
            _Logger.print_html(f"<b>Fourier Shell Correlation of {loader_name!r}</b>")
            with _Logger.set_plt():
                result.plot(criteria)
                plt.tight_layout()
                plt.show()
            for _c in criteria:
                _r = result.get_resolution(_c)
                _Logger.print_html(f"Resolution at FSC={_c:.3f} ... <b>{_r:.3f} nm</b>")

            if img_avg is not None:
                _imlayer: Image = self._show_rec(img_avg, name=f"[AVG]{loader_name}")
                _imlayer.metadata["fsc"] = result
                _imlayer.metadata["fsc_halfmaps"] = (
                    _as_imgarray(img_0, axes="izyx"),
                    _as_imgarray(img_1, axes="izyx"),
                )
                _imlayer.metadata["fsc_mask"] = _as_imgarray(img_mask)

        return _calculate_fsc_on_return

    @magictoolbar
    class STATools(MagicTemplate):
        show_template = abstractapi()
        show_template_original = abstractapi()
        show_mask = abstractapi()

    @set_design(icon="ic:baseline-view-in-ar", location=STATools)
    @do_not_record
    def show_template(self):
        """Load and show template image in the scale of the tomogram."""
        template = self._get_template_image()
        self._show_rec(template, name="Template image", store=False)

    @set_design(icon="material-symbols:view-in-ar", location=STATools)
    @do_not_record
    def show_template_original(self):
        """Load and show template image in the original scale."""
        _input = self.params._get_template_input(allow_multiple=True)
        if _input is None:
            raise ValueError("No template path provided.")
        elif isinstance(_input, Path):
            self._show_rec(ip.imread(_input), name="Template image", store=False)
        else:
            for i, fp in enumerate(_input):
                img = ip.imread(fp)
                self._show_rec(img, name=f"Template image [{i}]", store=False)

    @set_design(icon="fluent:shape-organic-20-filled", location=STATools)
    @do_not_record
    def show_mask(self):
        """Load and show mask image in the scale of the tomogram."""
        loader = self.get_loader(self.loader_name)
        _, mask = loader.normalize_input(
            self.params._norm_template_param(
                self.params._get_template_input(allow_multiple=False),
                allow_none=True,
            ),
            self.params._get_mask(),
        )
        if mask is None:
            raise ValueError("No mask to show.")
        mask = ip.asarray(mask, axes="zyx").set_scale(zyx=loader.scale, unit="nm")
        self._show_rec(mask, name="Mask image", store=False, threshold=0.5)

    @thread_worker.callback
    def _show_rec(
        self, img: ip.ImgArray, name: str, store: bool = True, threshold=None
    ):
        return self.params._show_reconstruction(img, name, store, threshold)

    def _get_shape_in_px(
        self, default: "nm | None", loader: BatchLoader
    ) -> tuple[int, ...]:
        if default is None:
            tmp = self._get_template_image(loader.scale)
            return tmp.sizesof("zyx")
        else:
            return (roundint(default / loader.scale),) * 3

    def _get_template_image(self, scale: nm | None = None) -> ip.ImgArray:
        if scale is None:
            scale = self.get_loader(self.loader_name).scale

        template = self.params._norm_template_param(
            self.params._get_template_input(allow_multiple=True),
            allow_none=False,
            allow_multiple=True,
        ).provide(scale)
        if isinstance(template, list):
            template = ip.asarray(np.stack(template, axis=0), axes="zyx")
        else:
            template = ip.asarray(template, axes="zyx")
        return template.set_scale(zyx=scale, unit="nm")

    def _group_loader_by_spline(
        self, loader_name: str, loader: BatchLoader
    ) -> Iterator["LoaderOnSpline"]:
        """Split `loader` (derived from the loader `loader_name`) by source splines.

        `loader` may be binned. Molecules of the yielded loaders are in the same
        coordinates as `loader`, so they can be joined back by `join_loaders`.
        """
        batch = self._get_parent()
        info = batch.loader_infos[loader_name]
        # TODO: batch.iter_projects() does not properly work if no batch project is
        # loaded beforehand.
        img_project_map = {
            path: prj
            for prj in batch.iter_projects()
            if (path := prj._batch_image_path()) is not None
        }
        for (mole_id, img_id), each in loader.group_by(["molecules-id", "image-id"]):
            img_path = info.image_paths[img_id]
            prj = img_project_map[img_path]
            for mole_info in prj.molecules_info:
                if mole_info.name.split(".")[0] == mole_id:
                    if (spl_id := mole_info.source) is None:
                        raise ValueError(
                            f"Molecule {mole_id} does not have source spline."
                        )
                    spl = prj.load_spline(spl_id)
                    break
            else:
                raise ValueError(f"Spline for molecule {mole_id} not found.")
            # the loader of the tomogram, or of the tilt series, of the image
            loader_single = loader.loaders[img_id].replace(
                molecules=each.molecules.with_features(
                    pl.lit(img_id).alias(Mole.image),
                    pl.lit(mole_id).alias(Mole.id),
                ),
            )
            yield LoaderOnSpline(
                loader=loader_single, image_id=img_id, molecule_id=mole_id, spline=spl
            )

align_all(loader_name, template_path, mask_params, max_shifts=(1.0, 1.0, 1.0), rotations=((0.0, 0.0), (0.0, 0.0), (0.0, 0.0)), cutoff=0.5, interpolation=3, method='zncc', bin_size=1)

Align all the molecules in the selected loader.

Parameters:

Name	Type	Description	Default
loader_name	str	Name of the batch subtomogram loader to be used.	required
template_path	Path or str	Path to template image.	required
mask_params	str or (float, float)	Mask image path or dilation/Gaussian blur parameters. If a path is given, image must in the same shape as the template.	required
max_shifts	int or tuple of int	Maximum shift between subtomograms and template in nm. ZYX order.	(1.0, 1.0, 1.0)
rotations	((float, float), (float, float), (float, float))	Rotation in external degree around each axis.	((0.0, 0.0), (0.0, 0.0), (0.0, 0.0))
cutoff	float	Cutoff frequency of low-pass filter applied in each subtomogram.	0.5
interpolation	int	Interpolation order.	3
method	str	Correlation metrics for alignment.	'zncc'
bin_size	int	Bin size of multiscale image to be used. Set to >1 to boost performance.	1

Source code in cylindra/widgets/batch/sta.py

@set_design(text="Align all molecules", location=BatchRefinement)
@dask_thread_worker.with_progress(desc="Aligning all molecules")
def align_all(
    self,
    loader_name: Annotated[str, {"bind": _get_current_loader_name}],
    template_path: Annotated[str | Path, {"bind": _get_template_path}],
    mask_params: Annotated[Any, {"bind": _get_mask_params}],
    max_shifts: _MaxShifts = (1.0, 1.0, 1.0),
    rotations: _Rotations = ((0.0, 0.0), (0.0, 0.0), (0.0, 0.0)),
    cutoff: _CutoffFreq = 0.5,
    interpolation: Annotated[int, {"choices": INTERPOLATION_CHOICES}] = 3,
    method: Annotated[str, {"choices": METHOD_CHOICES}] = "zncc",
    bin_size: _BINSIZE = 1,
):  # fmt: skip
    """
    Align all the molecules in the selected loader.

    Parameters
    ----------
    {loader_name}{template_path}{mask_params}{max_shifts}{rotations}{cutoff}
    {interpolation}{method}{bin_size}
    """
    t0 = timer()
    loaderlist = self._get_parent()._loaders
    info = loaderlist[loader_name]
    loader = info.loader.replace(order=interpolation).binning(
        bin_size, compute=False
    )
    template, mask = loader.normalize_input(
        template=self.params._norm_template_param(template_path),
        mask=self.params._get_mask(params=mask_params),
    )
    _Logger.print(f"Aligning {loader.molecules.count()} molecules ...")
    aligned = (
        loader.replace(output_shape=template.shape)
        .order_optimize()
        .align(
            template=template,
            mask=mask,
            max_shifts=max_shifts,
            rotations=rotations,
            cutoff=cutoff,
            alignment_model=_get_alignment(method),
        )
        .order_restore()
    )
    loaderlist.add_loader(
        aligned,
        name=_coerce_aligned_name(info.name, loaderlist),
        image_paths=info.image_paths,
        invert=info.invert,
    )
    t0.toc()

align_all_rfa_template_free(loader_name, mask_params, size=12.0, max_shifts=(2, 2, 2), max_rotations=(0.0, 180.0, 0.0), interpolation=3, method='zncc', range_long=(4.0, 10.0), angle_max=5.0, optimize_polarity=True, randomize_roll_angle=False, bin_size=1, temperature_time_const=5.0, upsample_factor=3, num_trials=5, max_num_iters=8, seed=0)

Create initial model of filaments by iterative RFA without template.

This method is useful to create an initial model from molecules mapped along manually initialized splines (such as by map_along_spline), where the lattice parameters (and the polarity of the filaments) are unknown.

Parameters:

Name	Type	Description	Default
loader_name	str	Name of the batch subtomogram loader to be used.	required
mask_params	str or (float, float)	Mask image path or dilation/Gaussian blur parameters. If a path is given, image must in the same shape as the template.	required
size	nm	Size of the template in nm. Use the size of template image by default.	12.0
max_shifts	int or tuple of int	Maximum shift between subtomograms and template in nm. ZYX order.	(2, 2, 2)
max_rotations	(float, float, float)	Maximum rotation in degree around Z, Y, X axis. Rotations are searched with steps of at most 30 degrees.	(0.0, 180.0, 0.0)
interpolation	int	Interpolation order.	3
method	str	Correlation metrics for alignment.	'zncc'
range_long	(float, float)	Minimum and maximum allowed distances between longitudinally consecutive monomers	(4.0, 10.0)
angle_max	float	Maximum allowed angle between longitudinally consecutive monomers and the Y axis.	5.0
optimize_polarity	bool	If true, the polarity of each filament is optimized by running annealing in both orientations. Otherwise, the orientation of the input molecules is kept.	True
randomize_roll_angle	bool	If true, molecules are rotated around their y-axis (the filament axis) by random angles before alignment, so that the missing wedges of the subtomograms are distributed around the axis in the average. Rotation around the y-axis should be searched (such as max_rotations=(0, 180, 0)) to recover the roll angles.	False
bin_size	int	Bin size of multiscale image to be used. Set to >1 to boost performance.	1
temperature_time_const	float	Time constant of the temperature decay during annealing. Larger value results in slower annealing. 1.0 is a moderate value.	5.0
upsample_factor	int	Upsampling factor of ZNCC landscape. Be careful not to set this parameter too large. Calculation will take much longer for larger upsample_factor.	3
num_trials	int	Number of independent trials to perform. The best result among all trials will be returned.	5
max_num_iters	int	Maximum number of iterations to perform.	8
seed	int	Random seed for reproducibility.	0

Source code in cylindra/widgets/batch/sta.py

@set_design(text="RFA alignment (template-free)", location=BatchRefinement)
@dask_thread_worker.with_progress(desc="RFA alignment started")
def align_all_rfa_template_free(
    self,
    loader_name: Annotated[str, {"bind": _get_current_loader_name}],
    mask_params: Annotated[Any, {"bind": _get_mask_params}],
    size: _SubVolumeSize = 12.0,
    max_shifts: _MaxShifts = (2, 2, 2),
    max_rotations: _MaxRotations = (0.0, 180.0, 0.0),
    interpolation: Annotated[int, {"choices": INTERPOLATION_CHOICES}] = 3,
    method: Annotated[str, {"choices": METHOD_CHOICES}] = "zncc",
    range_long: _DistRangeLon = (4.0, 10.0),
    angle_max: _AngleMaxLon = 5.0,
    optimize_polarity: bool = True,
    randomize_roll_angle: bool = False,
    bin_size: _BINSIZE = 1,
    temperature_time_const: Annotated[float, {"min": 0.01, "max": 50.0}] = 5.0,
    upsample_factor: Annotated[int, {"min": 1, "max": 20}] = 3,
    num_trials: Annotated[int, {"min": 1, "max": 100}] = 5,
    max_num_iters: Annotated[int, {"min": 3, "max": 100}] = 8,
    seed: _SeedType = 0,
):
    """Create initial model of filaments by iterative RFA without template.

    This method is useful to create an initial model from molecules mapped along
    manually initialized splines (such as by `map_along_spline`), where the
    lattice parameters (and the polarity of the filaments) are unknown.

    Parameters
    ----------
    {loader_name}{mask_params}{size}{max_shifts}{max_rotations}
    {interpolation}{method}{range_long}{angle_max}
    optimize_polarity : bool, default True
        If true, the polarity of each filament is optimized by running annealing
        in both orientations. Otherwise, the orientation of the input molecules
        is kept.
    randomize_roll_angle : bool, default False
        If true, molecules are rotated around their y-axis (the filament axis) by
        random angles before alignment, so that the missing wedges of the
        subtomograms are distributed around the axis in the average. Rotation
        around the y-axis should be searched (such as `max_rotations=(0, 180, 0)`)
        to recover the roll angles.
    {bin_size}{temperature_time_const}{upsample_factor}{num_trials}
    {max_num_iters}{seed}
    """
    t0 = timer()
    rng = np.random.default_rng(seed)
    batch = self._get_parent()
    loaderlist = batch._loaders
    info = loaderlist[loader_name]
    mask = self.params._get_mask(params=mask_params)
    loader = info.loader.replace(order=interpolation).binning(
        bin_size, compute=False
    )
    shape = self._get_shape_in_px(size, loader)
    loader = loader.replace(output_shape=shape)

    _alignment_state = template_free.RFAAlignmentState(
        rng=rng,
        mask=mask,
        alignment_model=_get_alignment(method),
        max_num_iters=max_num_iters,
    )
    sub_inputs = list(self._group_loader_by_spline(loader_name, loader))
    num_splines = len(sub_inputs)
    _Logger.print(f"{num_splines} splines found for RFA alignment.")
    # annealings in the forward (and reverse) orientations
    annealings: list[tuple[template_free.ResumableAnnealing, ...] | None] = [
        None
    ] * num_splines
    # molecules around which the first landscapes are built, in the forward (and
    # reverse) orientations
    initial_refs: list[tuple[Molecules, ...]] = []
    with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdir:
        # Same as RMA, subtomograms are extracted only once. Landscapes are always
        # built at the initial positions of molecules, but in the orientations of
        # the best results of the last iteration, so that the annealing is resumed
        # with rotations searched around them. Subvolumes in any orientation,
        # including the reverse one, are resampled from the extracted subtomograms.
        for ith, inputs in enumerate(sub_inputs):
            yield thread_worker.description(
                f"Extracting subtomograms ({ith + 1}/{num_splines})"
            )
            inputs.loader = inputs.loader.extract_subtomograms(
                Path(tmpdir) / f"spline-{ith}", chunksize=64, apply_prefilter=True
            )
            mole = inputs.loader.molecules
            if randomize_roll_angle:
                mole = template_free.randomize_roll(mole, rng)
            if optimize_polarity:
                mole_rev = mole.rotate_by_rotvec_internal([np.pi, 0, 0])
                initial_refs.append((mole, mole_rev))
            else:
                initial_refs.append((mole,))
        extracted = [inputs.loader for inputs in sub_inputs]
        yield thread_worker.description(_pdesc.align_tf_0(_alignment_state))
        result = _alignment_state.fsc_step_init(
            join_extracted_loaders(extracted, [refs[0] for refs in initial_refs]),
            max_shifts, max_rotations, upsample_factor, temperature_time_const,
            num_trials,
        )  # fmt: skip
        # intermediate results for debugging
        history = [
            _history_entry(result, [refs[0] for refs in initial_refs], sub_inputs)
        ]
        _final = False
        while True:
            _Logger.print(f"Iteration {int(_alignment_state.num_iter)}")
            yield _plot_current_fsc.with_args(result.fsc, result.avg).with_desc(_pdesc.align_tf_1(_alignment_state))  # fmt: skip
            if _final:
                break
            if _final := _alignment_state.is_final_iteration():
                _converged = _alignment_state.is_converged()
                _Logger.print("FSC converged." if _converged else "Last iteration.")
            _Logger.print(_alignment_state.next_params(loader.scale).format())
            _ratio = _alignment_state.temperature_ratio(_final)
            _Logger.print(f"Annealing until temperature reaches {_ratio:.2e} of T0")
            _iter = _alignment_state.num_iter + 1
            molecules: list[Molecules] = []
            step_results: list[template_free.RFAStepResult] = []
            # molecules around which the landscapes are built
            references: list[tuple[Molecules, ...]] = []
            for ith, inputs in enumerate(sub_inputs):
                _p_ind = f"({ith + 1}/{num_splines})"
                yield thread_worker.description(
                    f"Landscape construction {_p_ind} of iteration {_iter}"
                )
                if (pair := annealings[ith]) is None:
                    refs = initial_refs[ith]
                else:
                    refs = tuple(annealing.next_reference() for annealing in pair)
                references.append(refs)
                # landscapes are not normalized to compare the two orientations
                landscapes = _alignment_state.built_landscapes_step(
                    [
                        template_free.resample_extracted(extracted[ith], ref)
                        for ref in refs
                    ],
                    norm=not optimize_polarity,
                )
                if pair is None:
                    pair = tuple(
                        _alignment_state.prep_annealing_step(
                            landscape, inputs.spline, range_long, angle_max
                        )
                        for landscape in landscapes
                    )
                    annealings[ith] = pair
                else:
                    for annealing, landscape in zip(pair, landscapes, strict=True):
                        annealing.update_landscape(landscape)
                yield thread_worker.description(
                    f"RFA step {_p_ind} of iteration {_iter}"
                )
                step = _alignment_state.rfa_step(*pair, final=_final)
                step_results.append(step)
                molecules.append(
                    step.molecules.with_features(
                        pl.lit(inputs.image_id).alias(Mole.image),
                        pl.lit(inputs.molecule_id).alias(Mole.id),
                    )
                )
            if optimize_polarity:
                _num_rev = sum(step.is_reversed for step in step_results)
                _Logger.print(
                    f"Selected orientations: {num_splines - _num_rev} forward, "
                    f"{_num_rev} reverse"
                )
            yield thread_worker.description(_pdesc.align_tf_0(_alignment_state))
            result = _alignment_state.fsc_step(
                join_extracted_loaders(extracted, molecules)
            )
            history.append(
                _history_entry(result, molecules, sub_inputs, references)
            )

    loaderlist.add_loader(
        join_loaders(loader, sub_inputs, molecules),
        name=_coerce_aligned_name(info.name, loaderlist),
        image_paths=info.image_paths,
        invert=info.invert,
    )
    t0.toc()

    @thread_worker.callback
    def _on_return():
        layer = self._show_rec(
            _alignment_state.results[-1].avg, f"[AVG]{loader_name}"
        )
        layer.metadata[TEMPLATE_FREE_HISTORY] = history

    return _on_return

align_all_rma(loader_name, template_path, mask_params, max_shifts=(0.8, 0.8, 0.8), rotations=((0.0, 0.0), (0.0, 0.0), (0.0, 0.0)), cutoff=0.5, interpolation=3, method='zncc', range_long=(4.0, 4.28), range_lat=('d.mean() - 0.1', 'd.mean() + 0.1'), angle_max=5.0, bin_size=1, temperature_time_const=1.0, lj_const=0.0, upsample_factor=5, num_trials=5, seed=0)

Run RMA alignment on all molecules in the selected loader.

Parameters:

Name	Type	Description	Default
loader_name	str	Name of the batch subtomogram loader to be used.	required
template_path	Path or str	Path to template image.	required
mask_params	str or (float, float)	Mask image path or dilation/Gaussian blur parameters. If a path is given, image must in the same shape as the template.	required
max_shifts	int or tuple of int	Maximum shift between subtomograms and template in nm. ZYX order.	(0.8, 0.8, 0.8)
rotations	((float, float), (float, float), (float, float))	Rotation in external degree around each axis.	((0.0, 0.0), (0.0, 0.0), (0.0, 0.0))
cutoff	float	Cutoff frequency of low-pass filter applied in each subtomogram.	0.5
interpolation	int	Interpolation order.	3
method	str	Correlation metrics for alignment.	'zncc'
range_long	(float, float)	Minimum and maximum allowed distances between longitudinally consecutive monomers	(4.0, 4.28)
range_lat	(float, float)	Minimum and maximum allowed distances between laterally consecutive monomers	('d.mean() - 0.1', 'd.mean() + 0.1')
angle_max	float	Maximum allowed angle between longitudinally consecutive monomers and the Y axis.	5.0
bin_size	int	Bin size of multiscale image to be used. Set to >1 to boost performance.	1
temperature_time_const	float	Time constant of the temperature decay during annealing. Larger value results in slower annealing. 1.0 is a moderate value.	1.0
lj_const	float	If greater than 0, the annealing will use a Lennard-Jones-like potential to allow molecules to separate further apart than the cutoff distance. Larger value will prefer longer distances between molecules. 0.25 - 0.5 is usually a good choice.	0.0
upsample_factor	int	Upsampling factor of ZNCC landscape. Be careful not to set this parameter too large. Calculation will take much longer for larger upsample_factor.	5
num_trials	int	Number of independent trials to perform. The best result among all trials will be returned.	5
seed	int	Random seed for reproducibility.	0

Source code in cylindra/widgets/batch/sta.py

@set_design(text="RMA alignment", location=BatchRefinement)
@dask_thread_worker.with_progress(desc="Simulated annealing (RMA)")
def align_all_rma(
    self,
    loader_name: Annotated[str, {"bind": _get_current_loader_name}],
    template_path: Annotated[str | Path, {"bind": _get_template_path}],
    mask_params: Annotated[Any, {"bind": _get_mask_params}],
    max_shifts: _MaxShifts = (0.8, 0.8, 0.8),
    rotations: _Rotations = ((0.0, 0.0), (0.0, 0.0), (0.0, 0.0)),
    cutoff: _CutoffFreq = 0.5,
    interpolation: Annotated[int, {"choices": INTERPOLATION_CHOICES}] = 3,
    method: Annotated[str, {"choices": METHOD_CHOICES}] = "zncc",
    range_long: _DistRangeLon = (4.0, 4.28),
    range_lat: _DistRangeLat = ("d.mean() - 0.1", "d.mean() + 0.1"),
    angle_max: _AngleMaxLon = 5.0,
    bin_size: _BINSIZE = 1,
    temperature_time_const: Annotated[float, {"min": 0.01, "max": 10.0}] = 1.0,
    lj_const: Annotated[float, {"min": 0.00, "max": 1000.0, "step": 0.1, "label": "LJ const"}] = 0.0,
    upsample_factor: Annotated[int, {"min": 1, "max": 20}] = 5,
    num_trials: Annotated[int, {"min": 1, "max": 100}] = 5,
    seed: _SeedType = 0,
):  # fmt: skip
    """Run RMA alignment on all molecules in the selected loader.

    Parameters
    ----------
    {loader_name}{template_path}{mask_params}{max_shifts}{rotations}{cutoff}
    {interpolation}{method}{range_long}{range_lat}{angle_max}{bin_size}
    {temperature_time_const}{lj_const}{upsample_factor}{num_trials}{seed}
    """
    t0 = timer()
    batch = self._get_parent()
    loaderlist = batch._loaders
    info = loaderlist[loader_name]
    loader = info.loader.replace(order=interpolation).binning(
        bin_size, compute=False
    )
    template, mask = loader.normalize_input(
        template=self.params._norm_template_param(template_path),
        mask=self.params._get_mask(params=mask_params),
    )
    loader = loader.replace(output_shape=template.shape)
    aln = _get_alignment(method)
    sub_inputs = list(self._group_loader_by_spline(loader_name, loader))
    num_splines = len(sub_inputs)
    _Logger.print(f"{num_splines} splines found for RMA alignment.")

    @thread_worker.callback
    def _plot_annealing_result(results):
        with _Logger.set_plt():
            _annealing.plot_annealing_result(results)

    for ith, inputs in enumerate(sub_inputs):
        _p_ind = f"({ith + 1}/{num_splines})"
        yield thread_worker.description(f"Landscape construction {_p_ind}")
        landscape = Landscape.from_loader(
            loader=inputs.loader,
            template=template,
            mask=mask,
            max_shifts=max_shifts,
            upsample_factor=upsample_factor,
            alignment_model=aln.with_params(
                rotations=rotations,
                cutoff=cutoff,
                tilt=inputs.loader.tilt_model,
            ),
        ).normed()
        yield thread_worker.description(f"RMA alignment {_p_ind}")
        mole, results = landscape.run_annealing_along_spline(
            inputs.spline,
            range_long,
            range_lat,
            angle_max=angle_max,
            temperature_time_const=temperature_time_const,
            lj_nstd=None if lj_const < 1e-3 else 1 / lj_const,
            random_seeds=create_random_seeds(num_trials, seed),
        )
        yield _plot_annealing_result.with_args(results)
        mole = mole.with_features(
            pl.lit(inputs.image_id).alias(Mole.image),
            pl.lit(inputs.molecule_id).alias(Mole.id),
        )
        inputs.loader = inputs.loader.replace(molecules=mole)

    loaderlist.add_loader(
        join_loaders(loader, sub_inputs),
        name=_coerce_aligned_name(info.name, loaderlist),
        image_paths=info.image_paths,
        invert=info.invert,
    )
    t0.toc()

align_all_rma_template_free(loader_name, mask_params, size=12.0, max_shifts=(0.8, 0.8, 0.8), max_rotations=(0.0, 0.0, 0.0), interpolation=3, method='zncc', range_long=(4.0, 4.28), range_lat=('d.mean() - 0.1', 'd.mean() + 0.1'), angle_max=5.0, bin_size=1, temperature_time_const=1.0, upsample_factor=3, num_trials=5, max_num_iters=8, seed=0)

Create initial model by iteratively aligning molecules by RMA without template.

The annealing schedule of RMA is split into max_num_iters iterations. In each iteration, annealing is resumed from the last state and temperature of the previous iteration using the current average as the template, and then the average is updated and evaluated using FSC. If FSC converged, the rest of the annealing will be finished in the next iteration.

Parameters:

Name	Type	Description	Default
loader_name	str	Name of the batch subtomogram loader to be used.	required
mask_params	str or (float, float)	Mask image path or dilation/Gaussian blur parameters. If a path is given, image must in the same shape as the template.	required
size	nm	Size of the template in nm. Use the size of template image by default.	12.0
max_shifts	int or tuple of int	Maximum shift between subtomograms and template in nm. ZYX order.	(0.8, 0.8, 0.8)
max_rotations	(float, float, float)	Maximum rotation in degree around Z, Y, X axis. Rotations are searched with steps of at most 30 degrees.	(0.0, 0.0, 0.0)
interpolation	int	Interpolation order.	3
method	str	Correlation metrics for alignment.	'zncc'
range_long	(float, float)	Minimum and maximum allowed distances between longitudinally consecutive monomers	(4.0, 4.28)
range_lat	(float, float)	Minimum and maximum allowed distances between laterally consecutive monomers	('d.mean() - 0.1', 'd.mean() + 0.1')
angle_max	float	Maximum allowed angle between longitudinally consecutive monomers and the Y axis.	5.0
bin_size	int	Bin size of multiscale image to be used. Set to >1 to boost performance.	1
temperature_time_const	float	Time constant of the temperature decay during annealing. Larger value results in slower annealing. 1.0 is a moderate value.	1.0
upsample_factor	int	Upsampling factor of ZNCC landscape. Be careful not to set this parameter too large. Calculation will take much longer for larger upsample_factor.	3
num_trials	int	Number of independent trials to perform. The best result among all trials will be returned.	5
max_num_iters	int	Maximum number of iterations to perform.	8
seed	int	Random seed for reproducibility.	0

Source code in cylindra/widgets/batch/sta.py

@set_design(text="RMA alignment (template-free)", location=BatchRefinement)
@dask_thread_worker.with_progress(desc="RMA alignment started")
def align_all_rma_template_free(
    self,
    loader_name: Annotated[str, {"bind": _get_current_loader_name}],
    mask_params: Annotated[Any, {"bind": _get_mask_params}],
    size: _SubVolumeSize = 12.0,
    max_shifts: _MaxShifts = (0.8, 0.8, 0.8),
    max_rotations: _MaxRotations = (0.0, 0.0, 0.0),
    interpolation: Annotated[int, {"choices": INTERPOLATION_CHOICES}] = 3,
    method: Annotated[str, {"choices": METHOD_CHOICES}] = "zncc",
    range_long: _DistRangeLon = (4.0, 4.28),
    range_lat: _DistRangeLat = ("d.mean() - 0.1", "d.mean() + 0.1"),
    angle_max: _AngleMaxLon = 5.0,
    bin_size: _BINSIZE = 1,
    temperature_time_const: Annotated[float, {"min": 0.01, "max": 10.0}] = 1.0,
    upsample_factor: Annotated[int, {"min": 1, "max": 20}] = 3,
    num_trials: Annotated[int, {"min": 1, "max": 100}] = 5,
    max_num_iters: Annotated[int, {"min": 3, "max": 100}] = 8,
    seed: _SeedType = 0,
):
    """Create initial model by iteratively aligning molecules by RMA without template.

    The annealing schedule of RMA is split into `max_num_iters` iterations. In
    each iteration, annealing is resumed from the last state and temperature of
    the previous iteration using the current average as the template, and then the
    average is updated and evaluated using FSC. If FSC converged, the rest of the
    annealing will be finished in the next iteration.

    Parameters
    ----------
    {loader_name}{mask_params}{size}{max_shifts}{max_rotations}
    {interpolation}{method}{range_long}{range_lat}{angle_max}{bin_size}
    {temperature_time_const}{upsample_factor}{num_trials}{max_num_iters}{seed}
    """
    t0 = timer()
    rng = np.random.default_rng(seed)
    batch = self._get_parent()
    loaderlist = batch._loaders
    info = loaderlist[loader_name]
    mask = self.params._get_mask(params=mask_params)
    loader = info.loader.replace(order=interpolation).binning(
        bin_size, compute=False
    )
    shape = self._get_shape_in_px(size, loader)
    loader = loader.replace(output_shape=shape)

    _alignment_state = template_free.RMAAlignmentState(
        rng=rng,
        mask=mask,
        alignment_model=_get_alignment(method),
        max_num_iters=max_num_iters,
    )
    sub_inputs = list(self._group_loader_by_spline(loader_name, loader))
    num_splines = len(sub_inputs)
    _Logger.print(f"{num_splines} splines found for RMA alignment.")
    annealings: list[template_free.ResumableAnnealing | None] = [None] * num_splines
    with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdir:
        # Subtomograms are extracted to files only once, as spline coefficients so
        # that they can be resampled without the prefilter. Landscapes are always
        # constructed around the initial molecules so that the annealing can be
        # resumed in the next iteration, and the averages for FSC are resampled
        # from the extracted subtomograms at the aligned molecules.
        for ith, inputs in enumerate(sub_inputs):
            yield thread_worker.description(
                f"Extracting subtomograms ({ith + 1}/{num_splines})"
            )
            inputs.loader = inputs.loader.extract_subtomograms(
                Path(tmpdir) / f"spline-{ith}", chunksize=64, apply_prefilter=True
            )
        extracted = [inputs.loader for inputs in sub_inputs]
        yield thread_worker.description(_pdesc.align_tf_0(_alignment_state))
        result = _alignment_state.fsc_step_init(
            join_extracted_loaders(extracted), max_shifts, max_rotations,
            upsample_factor, temperature_time_const, num_trials,
        )  # fmt: skip
        _final = False
        while True:
            _Logger.print(f"Iteration {int(_alignment_state.num_iter)}")
            yield _plot_current_fsc.with_args(result.fsc, result.avg).with_desc(_pdesc.align_tf_1(_alignment_state))  # fmt: skip
            if _final:
                break
            if _final := _alignment_state.is_final_iteration():
                _converged = _alignment_state.is_converged()
                _Logger.print("FSC converged." if _converged else "Last iteration.")
            _Logger.print(_alignment_state.next_params(loader.scale).format())
            _ratio = _alignment_state.temperature_ratio(_final)
            _Logger.print(f"Annealing until temperature reaches {_ratio:.2e} of T0")
            _iter = _alignment_state.num_iter + 1
            molecules: list[Molecules] = []
            for ith, inputs in enumerate(sub_inputs):
                _p_ind = f"({ith + 1}/{num_splines})"
                yield thread_worker.description(
                    f"Landscape construction {_p_ind} of iteration {_iter}"
                )
                landscape = _alignment_state.built_landscape_step(inputs.loader)
                annealing = annealings[ith]
                if annealing is None:
                    annealing = _alignment_state.prep_annealing_step(
                        landscape, inputs.spline, range_long, range_lat, angle_max
                    )
                    annealings[ith] = annealing
                else:
                    annealing.update_landscape(landscape)
                yield thread_worker.description(
                    f"RMA step {_p_ind} of iteration {_iter}"
                )
                mole = _alignment_state.rma_step(annealing, final=_final)
                molecules.append(
                    mole.with_features(
                        pl.lit(inputs.image_id).alias(Mole.image),
                        pl.lit(inputs.molecule_id).alias(Mole.id),
                    )
                )
            yield thread_worker.description(_pdesc.align_tf_0(_alignment_state))
            result = _alignment_state.fsc_step(
                join_extracted_loaders(extracted, molecules)
            )

    loaderlist.add_loader(
        join_loaders(loader, sub_inputs, molecules),
        name=_coerce_aligned_name(info.name, loaderlist),
        image_paths=info.image_paths,
        invert=info.invert,
    )
    t0.toc()
    return self._show_rec.with_args(
        _alignment_state.results[-1].avg, f"[AVG]{loader_name}"
    )

align_all_template_free(loader_name, mask_params, size=12.0, max_shifts=(0.8, 0.8, 0.8), max_rotations=(3.0, 3.0, 3.0), min_rotation_step=1.0, interpolation=3, method='zncc', bin_size=1, max_num_iters=20)

Create initial model by iteratively aligning molecules without template.

This method iteratively calculate average, evaluate using FSC, and align molecules to the current average.

Parameters:

Name	Type	Description	Default
loader_name	str	Name of the batch subtomogram loader to be used.	required
mask_params	str or (float, float)	Mask image path or dilation/Gaussian blur parameters. If a path is given, image must in the same shape as the template.	required
size	nm	Size of the template in nm. Use the size of template image by default.	12.0
max_shifts	int or tuple of int	Maximum shift between subtomograms and template in nm. ZYX order.	(0.8, 0.8, 0.8)
max_rotations	(float, float, float)	Maximum rotation in degree around Z, Y, X axis. Rotations are searched with steps of at most 30 degrees.	(3.0, 3.0, 3.0)
min_rotation_step	float	Minimum rotation search precision in degree.	1.0
interpolation	int	Interpolation order.	3
method	str	Correlation metrics for alignment.	'zncc'
bin_size	int	Bin size of multiscale image to be used. Set to >1 to boost performance.	1
max_num_iters	int	Maximum number of iterations to perform.	20

Source code in cylindra/widgets/batch/sta.py

@set_design(text="Align all (template-free)", location=BatchRefinement)
@dask_thread_worker.with_progress(desc="Align all started")
def align_all_template_free(
    self,
    loader_name: Annotated[str, {"bind": _get_current_loader_name}],
    mask_params: Annotated[Any, {"bind": _get_mask_params}],
    size: _SubVolumeSize = 12.0,
    max_shifts: _MaxShifts = (0.8, 0.8, 0.8),
    max_rotations: _MaxRotations = (3.0, 3.0, 3.0),
    min_rotation_step: Annotated[float, {"min": 0.1, "step": 0.1}] = 1.0,
    interpolation: Annotated[int, {"choices": INTERPOLATION_CHOICES}] = 3,
    method: Annotated[str, {"choices": METHOD_CHOICES}] = "zncc",
    bin_size: _BINSIZE = 1,
    max_num_iters: Annotated[int, {"min": 3, "max": 100}] = 20,
):  # fmt: skip
    """Create initial model by iteratively aligning molecules without template.

    This method iteratively calculate average, evaluate using FSC, and align
    molecules to the current average.

    Parameters
    ----------
    {loader_name}{mask_params}{size}{max_shifts}{max_rotations}{min_rotation_step}
    {interpolation}{method}{bin_size}{max_num_iters}
    """
    t0 = timer()
    rng = np.random.default_rng(1428)
    loaderlist = self._get_parent()._loaders
    info = loaderlist[loader_name]
    mask = self.params._get_mask(params=mask_params)
    loader = info.loader.replace(order=interpolation).binning(
        bin_size, compute=False
    )
    shape = self._get_shape_in_px(size, loader)
    loader = loader.replace(output_shape=shape).order_optimize()
    _Logger.print(f"Aligning {loader.molecules.count()} molecules ...")
    _alignment_state = template_free.AlignmentState(
        rng=rng,
        mask=mask,
        alignment_model=_get_alignment(method),
        min_rotation_step=min_rotation_step,
    )
    # loaders of each tomogram, which have their own tilt models
    sub_loaders = list(loader.loaders.iter_loaders())
    num_images = len(sub_loaders)
    with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdir:
        # Subtomograms are extracted to files only once, as spline coefficients so
        # that they can be resampled without the prefilter. In each iteration, the
        # extracted subtomograms are resampled at the aligned molecules for both
        # alignment and FSC calculation.
        extracted: dict[Hashable, ExtractedSubvolumeLoader] = {}
        for ith, (img_id, sub) in enumerate(sub_loaders):
            yield thread_worker.description(
                f"Extracting subtomograms ({ith + 1}/{num_images})"
            )
            extracted[img_id] = sub.extract_subtomograms(
                Path(tmpdir) / f"image-{ith}",
                chunksize=64,
                apply_prefilter=True,
            )
        # the image IDs are kept, as the aligned molecules are the output
        joined = join_extracted_loaders(extracted)
        yield thread_worker.description(_pdesc.align_tf_0(_alignment_state))
        result0 = _alignment_state.fsc_step_init(joined, max_shifts, max_rotations)
        _Logger.print(f"Iteration {int(_alignment_state.num_iter)}")
        yield _plot_current_fsc.with_args(result0.fsc, result0.avg).with_desc(_pdesc.align_tf_1(_alignment_state))  # fmt: skip
        while True:
            _Logger.print(_alignment_state.next_params(loader.scale).format())
            _exceeded = max_num_iters <= _alignment_state.num_iter
            if _alignment_state.is_converged() or _exceeded:
                _Logger.print("Maximum iteration exceeded" if _exceeded else "FSC converged.")  # fmt: skip
                break
            # the extracted subtomograms are resampled at the aligned molecules
            joined = _alignment_state.align_step(joined)
            yield thread_worker.description(_pdesc.align_tf_0(_alignment_state))
            result = _alignment_state.fsc_step(joined)
            _Logger.print(f"Iteration {int(_alignment_state.num_iter)}")
            yield _plot_current_fsc.with_args(result.fsc, result.avg).with_desc(_pdesc.align_tf_1(_alignment_state))  # fmt: skip

    mole_aligned = _as_alignment_output(loader.molecules, [joined.molecules])
    loaderlist.add_loader(
        loader.replace(molecules=mole_aligned).order_restore(),
        name=_coerce_aligned_name(info.name, loaderlist),
        image_paths=info.image_paths,
        invert=info.invert,
    )
    t0.toc()
    return self._show_rec.with_args(
        _alignment_state.results[-1].avg, f"[AVG]{loader_name}"
    )

average_all(loader_name, size=None, interpolation=1, bin_size=1)

Average all the molecules in the selected loader.

Parameters:

Name	Type	Description	Default
loader_name	str	Name of the batch subtomogram loader to be used.	required
size	nm	Size of the template in nm. Use the size of template image by default.	None
interpolation	int	Interpolation order.	1
bin_size	int	Bin size of multiscale image to be used. Set to >1 to boost performance.	1

Source code in cylindra/widgets/batch/sta.py

@set_design(text="Average all molecules", location=BatchSubtomogramAnalysis)
@dask_thread_worker.with_progress(desc="Averaging all molecules in projects")
def average_all(
    self,
    loader_name: Annotated[str, {"bind": _get_current_loader_name}],
    size: _SubVolumeSize = None,
    interpolation: Annotated[int, {"choices": INTERPOLATION_CHOICES}] = 1,
    bin_size: _BINSIZE = 1,
):
    """Average all the molecules in the selected loader.

    Parameters
    ----------
    {loader_name}{size}{interpolation}{bin_size}
    """
    t0 = timer()
    loader = (
        self._get_parent()
        .loader_infos[loader_name]
        .loader.replace(order=interpolation)
        .binning(bin_size, compute=False)
    )
    shape = self._get_shape_in_px(size, loader)
    img = ip.asarray(
        loader.replace(output_shape=shape).order_optimize().average(),
        axes="zyx",
    ).set_scale(zyx=loader.scale, unit="nm")
    t0.toc()
    return self._show_rec.with_args(img, f"[AVG]{loader_name}")

average_groups(loader_name, size=None, by="col('pf-id')", interpolation=1, bin_size=1)

Groupwise subtomogram averaging using molecules grouped by the given expression.

This method first group molecules by its features, and then average each group. This method is useful for such as get average of each protofilament and segmented subtomogram averaging.

Parameters:

Name	Type	Description	Default
loader_name	str	Name of the batch subtomogram loader to be used.	required
size	nm	Size of the template in nm. Use the size of template image by default.	None
by	str or polars expression	Expression to group molecules.	"col('pf-id')"
interpolation	int	Interpolation order.	1
bin_size	int	Bin size of multiscale image to be used. Set to >1 to boost performance.	1

Source code in cylindra/widgets/batch/sta.py

@set_design(text="Average group-wise", location=BatchSubtomogramAnalysis)
@dask_thread_worker.with_progress(desc="Grouped subtomogram averaging")
def average_groups(
    self,
    loader_name: Annotated[str, {"bind": _get_current_loader_name}],
    size: _SubVolumeSize = None,
    by: PolarsExprStr = "col('pf-id')",
    interpolation: Annotated[int, {"choices": INTERPOLATION_CHOICES}] = 1,
    bin_size: _BINSIZE = 1,
):
    """Groupwise subtomogram averaging using molecules grouped by the given expression.

    This method first group molecules by its features, and then average each group.
    This method is useful for such as get average of each protofilament and
    segmented subtomogram averaging.

    Parameters
    ----------
    {loader_name}{size}
    by : str or polars expression
        Expression to group molecules.
    {interpolation}{bin_size}
    """
    t0 = timer()
    loader = (
        self._get_parent()
        .loader_infos[loader_name]
        .loader.replace(order=interpolation)
        .binning(bin_size, compute=False)
    )
    shape = self._get_shape_in_px(size, loader)
    img = ip.asarray(
        loader.replace(output_shape=shape)
        .order_optimize()
        .groupby(norm_polars_expr(by))
        .average()
        .value_stack(axis=0),
        axes="pzyx",
    ).set_scale(zyx=loader.scale, unit="nm")
    t0.toc()
    return self._show_rec.with_args(img, f"[AVG]{loader_name}", store=False)

calculate_fsc(loader_name, template_path=None, mask_params=None, size=None, seed=0, interpolation=1, n_pairs=1, show_average=True, dfreq=None)

Calculate Fourier Shell Correlation using the selected loader.

Parameters:

Name	Type	Description	Default
loader_name	str	Name of the batch subtomogram loader to be used.	required
template_path	Path or str	Path to template image.	None
mask_params	str or (float, float)	Mask image path or dilation/Gaussian blur parameters. If a path is given, image must in the same shape as the template.	None
size	nm	Size of the template in nm. Use the size of template image by default.	None
seed	int	Random seed used for subtomogram sampling.	0
interpolation	int	Interpolation order.	1
n_pairs	int	How many sets of image pairs will be generated to average FSC.	1
show_average	bool	If true, subtomogram averaging will be shown after FSC calculation.	True
dfreq	float	Precision of frequency to calculate FSC. "0.02" means that FSC will be calculated at frequency 0.01, 0.03, 0.05, ..., 0.45.	0.02

Source code in cylindra/widgets/batch/sta.py

@set_design(text="Calculate FSC", location=BatchSubtomogramAnalysis)
@dask_thread_worker.with_progress(desc="Calculating FSC")
def calculate_fsc(
    self,
    loader_name: Annotated[str, {"bind": _get_current_loader_name}],
    template_path: Annotated[str | Path | None, {"bind": _get_template_path}] = None,
    mask_params: Annotated[Any, {"bind": _get_mask_params}] = None,
    size: _SubVolumeSize = None,
    seed: Annotated[Optional[int], {"text": "Do not use random seed."}] = 0,
    interpolation: Annotated[int, {"choices": INTERPOLATION_CHOICES}] = 1,
    n_pairs: Annotated[int, {"min": 1, "label": "number of image pairs"}] = 1,
    show_average: bool = True,
    dfreq: FSCFreq = None,
):  # fmt: skip
    """Calculate Fourier Shell Correlation using the selected loader.

    Parameters
    ----------
    {loader_name}{template_path}{mask_params}{size}
    seed : int, optional
        Random seed used for subtomogram sampling.
    {interpolation}
    n_pairs : int, default 1
        How many sets of image pairs will be generated to average FSC.
    show_average : bool, default True
        If true, subtomogram averaging will be shown after FSC calculation.
    dfreq : float, default 0.02
        Precision of frequency to calculate FSC. "0.02" means that FSC will be
        calculated at frequency 0.01, 0.03, 0.05, ..., 0.45.
    """
    t0 = timer()
    loader = (
        self._get_parent()
        .loader_infos[loader_name]
        .loader.replace(order=interpolation)
    )

    template, mask = loader.normalize_input(
        template=self.params._norm_template_param(template_path, allow_none=True),
        mask=self.params._get_mask(params=mask_params),
    )

    fsc, (img_0, img_1), img_mask = (
        loader.reshape(
            template=template if size is None else None,
            mask=mask,
            shape=None if size is None else self._get_shape_in_px(size, loader),
        )
        .order_optimize()
        .fsc_with_halfmaps(
            mask, seed=seed, n_set=n_pairs, dfreq=dfreq, squeeze=False
        )
    )

    def _as_imgarray(im, axes: str = "zyx") -> ip.ImgArray | None:
        if np.isscalar(im):
            return None
        return ip.asarray(im, axes=axes).set_scale(zyx=loader.scale, unit="nm")

    if show_average:
        avg = (img_0[0] + img_1[0]) / 2
        img_avg = _as_imgarray(avg)
    else:
        img_avg = None

    result = FscResult.from_dataframe(fsc, loader.scale)
    criteria = [0.5, 0.143]
    t0.toc()

    @thread_worker.callback
    def _calculate_fsc_on_return():
        _Logger.print_html(f"<b>Fourier Shell Correlation of {loader_name!r}</b>")
        with _Logger.set_plt():
            result.plot(criteria)
            plt.tight_layout()
            plt.show()
        for _c in criteria:
            _r = result.get_resolution(_c)
            _Logger.print_html(f"Resolution at FSC={_c:.3f} ... <b>{_r:.3f} nm</b>")

        if img_avg is not None:
            _imlayer: Image = self._show_rec(img_avg, name=f"[AVG]{loader_name}")
            _imlayer.metadata["fsc"] = result
            _imlayer.metadata["fsc_halfmaps"] = (
                _as_imgarray(img_0, axes="izyx"),
                _as_imgarray(img_1, axes="izyx"),
            )
            _imlayer.metadata["fsc_mask"] = _as_imgarray(img_mask)

    return _calculate_fsc_on_return

filter_loader(loader_name, expression)

Filter the selected loader and add the filtered one to the list.

Parameters:

Name	Type	Description	Default
loader_name	str	Name of the input loader	required
expression	str	polars expression that will be used to filter the loader. For example, col("score") > 0.7 will filter out all low-score molecules.	required

Source code in cylindra/widgets/batch/sta.py

@set_design(text="Filter loader", location=BatchLoaderMenu)
def filter_loader(
    self,
    loader_name: Annotated[str, {"bind": _get_current_loader_name}],
    expression: PolarsExprStr,
):
    """Filter the selected loader and add the filtered one to the list.

    Parameters
    ----------
    loader_name : str
        Name of the input loader
    expression : str
        polars expression that will be used to filter the loader. For example,
        `col("score") > 0.7` will filter out all low-score molecules.
    """
    loaderlist = self._get_parent()._loaders
    info = loaderlist[loader_name]
    loader = info.loader
    new = loader.filter(norm_polars_expr(expression))
    existing_id = set(new.features[Mole.image])
    loaderlist.add_loader(
        new,
        name=f"{info.name}-Filt",
        image_paths={k: v for k, v in info.image_paths.items() if v in existing_id},
        invert={k: v for k, v in info.invert.items() if v in existing_id},
    )
    return None

get_loader(name)

Return the acryo.BatchLoader object with the given name

Source code in cylindra/widgets/batch/sta.py

@nogui
def get_loader(self, name: str) -> BatchLoader:
    """Return the acryo.BatchLoader object with the given name"""
    if not isinstance(name, str):
        raise TypeError(f"Name must be a string, got {type(name).__name__}")
    return self._get_parent().loader_infos[name].loader

remove_loader(loader_name)

Remove this loader

Source code in cylindra/widgets/batch/sta.py

@set_design(text="✕", max_width=36, location=Header)
def remove_loader(
    self, loader_name: Annotated[str, {"bind": _get_current_loader_name}]
):
    """Remove this loader"""
    del self._get_parent().loader_infos[loader_name]

show_loader_info()

Show information about this loader

Source code in cylindra/widgets/batch/sta.py

@set_design(text="??", max_width=36, location=Header)
@do_not_record
def show_loader_info(self):
    """Show information about this loader"""
    info = self._get_parent().loader_infos[self.loader_name]
    loader = info.loader
    img_info = "\n" + "\n".join(
        f"{img_id}: {img_path}" for img_id, img_path in info.image_paths.items()
    )

    info_text = (
        f"name: {info.name}\nmolecule: n={loader.count()}\nimages:{img_info}"
    )
    view = DataFrameView(value=loader.molecules.to_dataframe())
    txt = ConsoleTextEdit(value=info_text)
    txt.read_only = True
    cnt = Container(widgets=[txt, view], layout="horizontal", labels=False)
    cnt.native.setParent(self.native, cnt.native.windowFlags())
    cnt.show()

show_mask()

Load and show mask image in the scale of the tomogram.

Source code in cylindra/widgets/batch/sta.py

@set_design(icon="fluent:shape-organic-20-filled", location=STATools)
@do_not_record
def show_mask(self):
    """Load and show mask image in the scale of the tomogram."""
    loader = self.get_loader(self.loader_name)
    _, mask = loader.normalize_input(
        self.params._norm_template_param(
            self.params._get_template_input(allow_multiple=False),
            allow_none=True,
        ),
        self.params._get_mask(),
    )
    if mask is None:
        raise ValueError("No mask to show.")
    mask = ip.asarray(mask, axes="zyx").set_scale(zyx=loader.scale, unit="nm")
    self._show_rec(mask, name="Mask image", store=False, threshold=0.5)

show_template()

Load and show template image in the scale of the tomogram.

Source code in cylindra/widgets/batch/sta.py

@set_design(icon="ic:baseline-view-in-ar", location=STATools)
@do_not_record
def show_template(self):
    """Load and show template image in the scale of the tomogram."""
    template = self._get_template_image()
    self._show_rec(template, name="Template image", store=False)

show_template_original()

Load and show template image in the original scale.

Source code in cylindra/widgets/batch/sta.py

@set_design(icon="material-symbols:view-in-ar", location=STATools)
@do_not_record
def show_template_original(self):
    """Load and show template image in the original scale."""
    _input = self.params._get_template_input(allow_multiple=True)
    if _input is None:
        raise ValueError("No template path provided.")
    elif isinstance(_input, Path):
        self._show_rec(ip.imread(_input), name="Template image", store=False)
    else:
        for i, fp in enumerate(_input):
            img = ip.imread(fp)
            self._show_rec(img, name=f"Template image [{i}]", store=False)

split_loader(loader_name, by, delete_old=False)

Split the selected loader by the values of the given column.

Parameters:

Name	Type	Description	Default
loader_name	str	Name of the input loader	required
by	str	Column name to split the loader	required
delete_old	bool	If true, the original loader will be deleted.	False

Source code in cylindra/widgets/batch/sta.py

@set_design(text="Split loader", location=BatchLoaderMenu)
def split_loader(
    self,
    loader_name: Annotated[str, {"bind": _get_current_loader_name}],
    by: Annotated[str, {"choices": _get_selected_loader_choice}],
    delete_old: bool = False,
):
    """Split the selected loader by the values of the given column.

    Parameters
    ----------
    loader_name : str
        Name of the input loader
    by : str
        Column name to split the loader
    delete_old : bool, default False
        If true, the original loader will be deleted.
    """
    parent = self._get_parent()
    loaders = parent._loaders
    batch_info = loaders[loader_name]
    batch_loader = batch_info.loader
    n_unique = batch_loader.molecules.features[by].n_unique()
    if n_unique > 48:
        raise ValueError(
            f"Too many groups ({n_unique}). Did you choose a float column?"
        )
    for _key, loader in batch_loader.groupby(by):
        existing_id = set(loader.features[Mole.image])
        image_paths = {
            k: v for k, v in batch_info.image_paths.items() if v in existing_id
        }
        invert = {k: v for k, v in batch_info.invert.items() if v in existing_id}
        parent._add_loader(loader, f"{loader_name}_{_key}", image_paths, invert)

    if delete_old:
        idx = -1
        for i, info in enumerate(loaders):
            if info.loader is batch_loader:
                idx = i
                break
        else:
            idx = -1
        if idx < 0:
            raise RuntimeError("Loader not found.")
        del loaders[idx]