pytme 0.1.8__cp311-cp311-macosx_14_0_arm64.whl → 0.2.0b0__cp311-cp311-macosx_14_0_arm64.whl

scripts/preprocessor_gui.py

@@ -23,7 +23,7 @@ from qtpy.QtWidgets import QFileDialog
 from numpy.typing import NDArray
 
 from tme import Preprocessor, Density
-from tme.matching_utils import create_mask
+from tme.matching_utils import create_mask, load_pickle
 
 preprocessor = Preprocessor()
 SLIDER_MIN, SLIDER_MAX = 0, 25
@@ -416,7 +416,12 @@ class FilterWidget(widgets.Container):
 
 
 def sphere_mask(
-    template: NDArray, center_x: float, center_y: float, center_z: float, radius: float
+    template: NDArray,
+    center_x: float,
+    center_y: float,
+    center_z: float,
+    radius: float,
+    **kwargs,
 ) -> NDArray:
     return create_mask(
         mask_type="ellipse",
@@ -434,6 +439,7 @@ def ellipsod_mask(
     radius_x: float,
     radius_y: float,
     radius_z: float,
+    **kwargs,
 ) -> NDArray:
     return create_mask(
         mask_type="ellipse",
@@ -451,6 +457,7 @@ def box_mask(
     height_x: int,
     height_y: int,
     height_z: int,
+    **kwargs,
 ) -> NDArray:
     return create_mask(
         mask_type="box",
@@ -469,6 +476,7 @@ def tube_mask(
     inner_radius: float,
     outer_radius: float,
     height: int,
+    **kwargs,
 ) -> NDArray:
     return create_mask(
         mask_type="tube",
@@ -492,6 +500,7 @@ def wedge_mask(
     omit_negative_frequencies: bool = False,
     extrude_plane: bool = True,
     infinite_plane: bool = True,
+    **kwargs,
 ) -> NDArray:
     if tilt_step <= 0:
         wedge_mask = preprocessor.continuous_wedge_mask(
@@ -524,7 +533,7 @@ def wedge_mask(
 
 
 def threshold_mask(
-    template: NDArray, standard_deviation: float = 5.0, invert: bool = False
+    template: NDArray, standard_deviation: float = 5.0, invert: bool = False, **kwargs
 ) -> NDArray:
     template_mean = template.mean()
     template_deviation = standard_deviation * template.std()
@@ -537,7 +546,7 @@ def threshold_mask(
     return mask
 
 
-def lowpass_mask(template: NDArray, sigma: float = 1.0):
+def lowpass_mask(template: NDArray, sigma: float = 1.0, **kwargs):
     template = template / template.max()
     template = (template > np.exp(-2)) * 128.0
     template = preprocessor.gaussian_filter(template=template, sigma=sigma)
@@ -546,6 +555,20 @@ def lowpass_mask(template: NDArray, sigma: float = 1.0):
     return mask
 
 
+def shape_mask(template, shapes_layer, expansion_dim):
+    ret = np.zeros_like(template)
+    mask_shape = tuple(x for i, x in enumerate(template.shape) if i != expansion_dim)
+    masks = shapes_layer.to_masks(mask_shape=mask_shape)
+    for index, shape_type in enumerate(shapes_layer.shape_type):
+        mask = np.expand_dims(masks[index], axis=expansion_dim)
+        mask = np.repeat(
+            mask, repeats=template.shape[expansion_dim], axis=expansion_dim
+        )
+        np.logical_or(ret, mask, out=ret)
+
+    return ret
+
+
 class MaskWidget(widgets.Container):
     def __init__(self, viewer):
         super().__init__(layout="vertical")
@@ -564,6 +587,7 @@ class MaskWidget(widgets.Container):
             "Wedge": wedge_mask,
             "Threshold": threshold_mask,
             "Lowpass": lowpass_mask,
+            "Shape": shape_mask,
         }
 
         self.method_dropdown = widgets.ComboBox(
@@ -581,16 +605,19 @@ class MaskWidget(widgets.Container):
             self._update_action_button_state
         )
 
-        self.align_button = widgets.PushButton(text="Align to axis", enabled=False)
-        self.align_button.changed.connect(self._align_with_axis)
         self.density_field = widgets.Label()
         # self.density_field.value = f"Positive Density in Mask: {0:.2f}%"
 
+        self.shapes_layer_dropdown = widgets.ComboBox(
+            name="shapes_layer", choices=self._get_shape_layers()
+        )
+        self.viewer.layers.events.inserted.connect(self._update_shape_layer_choices)
+        self.viewer.layers.events.removed.connect(self._update_shape_layer_choices)
+
         self.append(self.method_dropdown)
         self.append(self.adapt_button)
         self.append(self.percentile_range_edit)
 
-        self.append(self.align_button)
         self.append(self.action_button)
         self.append(self.density_field)
 
@@ -598,43 +625,9 @@ class MaskWidget(widgets.Container):
         self._on_method_changed(None)
 
     def _update_action_button_state(self, event):
-        self.align_button.enabled = bool(self.viewer.layers.selection.active)
         self.action_button.enabled = bool(self.viewer.layers.selection.active)
         self.adapt_button.enabled = bool(self.viewer.layers.selection.active)
 
-    def _align_with_axis(self):
-        active_layer = self.viewer.layers.selection.active
-
-        if active_layer.metadata.get("is_aligned", False):
-            return
-
-        coords = np.array(np.where(active_layer.data > 0)).T
-        centered_coords = coords - np.mean(coords, axis=0)
-        cov_matrix = np.cov(centered_coords, rowvar=False)
-
-        eigenvalues, eigenvectors = np.linalg.eigh(cov_matrix)
-        principal_eigenvector = eigenvectors[:, np.argmax(eigenvalues)]
-
-        rotation_axis = np.cross(principal_eigenvector, [1, 0, 0])
-        rotation_angle = np.arccos(np.dot(principal_eigenvector, [1, 0, 0]))
-        k = rotation_axis / np.linalg.norm(rotation_axis)
-        K = np.array([[0, -k[2], k[1]], [k[2], 0, -k[0]], [-k[1], k[0], 0]])
-        rotation_matrix = np.eye(3)
-        rotation_matrix += np.sin(rotation_angle) * K
-        rotation_matrix += (1 - np.cos(rotation_angle)) * np.dot(K, K)
-
-        rotated_data = Density.rotate_array(
-            arr=active_layer.data,
-            rotation_matrix=rotation_matrix,
-            use_geometric_center=False,
-            order=1,
-        )
-        eps = np.finfo(rotated_data.dtype).eps
-        rotated_data[rotated_data < eps] = 0
-
-        active_layer.metadata["is_aligned"] = True
-        active_layer.data = rotated_data
-
     def _update_initial_values(self, event=None):
         active_layer = self.viewer.layers.selection.active
 
@@ -673,17 +666,43 @@ class MaskWidget(widgets.Container):
         self.action_widgets.clear()
 
         function = self.methods.get(self.method_dropdown.value)
-        widgets = widgets_from_function(function)
-        for widget in widgets:
+        function_widgets = widgets_from_function(function)
+        for widget in function_widgets:
             self.action_widgets.append(widget)
             self.insert(1, widget)
 
+        for name, param in inspect.signature(function).parameters.items():
+            if name == "shapes_layer":
+                self.action_widgets.append(self.shapes_layer_dropdown)
+                self.insert(1, self.shapes_layer_dropdown)
+
+    def _get_shape_layers(self):
+        layers = [
+            layer.name
+            for layer in self.viewer.layers
+            if isinstance(layer, napari.layers.Shapes)
+        ]
+        return layers
+
+    def _update_shape_layer_choices(self, event):
+        """Update the choices in the shapes layer dropdown."""
+        self.shapes_layer_dropdown.choices = self._get_shape_layers()
+
     def _action(self):
         function = self.methods.get(self.method_dropdown.value)
 
         selected_layer = self.viewer.layers.selection.active
         kwargs = {widget.name: widget.value for widget in self.action_widgets}
+
+        if "shapes_layer" in kwargs:
+            layer_name = kwargs["shapes_layer"]
+            if layer_name not in self.viewer.layers:
+                return None
+            kwargs["shapes_layer"] = self.viewer.layers[layer_name]
+            kwargs["expansion_dim"] = self.viewer.dims.order[0]
+
         processed_data = function(template=selected_layer.data, **kwargs)
+
         new_layer_name = f"{selected_layer.name} ({self.method_dropdown.value})"
 
         if new_layer_name in self.viewer.layers:
@@ -705,14 +724,114 @@ class MaskWidget(widgets.Container):
             metadata["origin_layer"] = selected_layer.name
             new_layer.metadata = metadata
 
-        origin_layer = metadata["origin_layer"]
-        if origin_layer in self.viewer.layers:
-            origin_layer = self.viewer.layers[origin_layer]
-            if np.allclose(origin_layer.data.shape, processed_data.shape):
-                in_mask = np.sum(np.fmax(origin_layer.data * processed_data, 0))
-                in_mask /= np.sum(np.fmax(origin_layer.data, 0))
-                in_mask *= 100
-                self.density_field.value = f"Positive Density in Mask: {in_mask:.2f}%"
+            if self.method_dropdown.value == "Shape":
+                new_layer.metadata = {}
+
+        # origin_layer = metadata["origin_layer"]
+        # if origin_layer in self.viewer.layers:
+        #     origin_layer = self.viewer.layers[origin_layer]
+        #     if np.allclose(origin_layer.data.shape, processed_data.shape):
+        #         in_mask = np.sum(np.fmax(origin_layer.data * processed_data, 0))
+        #         in_mask /= np.sum(np.fmax(origin_layer.data, 0))
+        #         in_mask *= 100
+        #         self.density_field.value = f"Positive Density in Mask: {in_mask:.2f}%"
+
+
+class AlignmentWidget(widgets.Container):
+    def __init__(self, viewer):
+        super().__init__(layout="vertical")
+
+        self.viewer = viewer
+
+        align_button = widgets.PushButton(text="Align to axis", enabled=True)
+        self.align_axis = widgets.ComboBox(
+            value=None, nullable=True, choices=self._get_active_layer_dims
+        )
+        self.viewer.layers.selection.events.changed.connect(self._update_align_axis)
+
+        align_button.changed.connect(self._align_with_axis)
+        container = widgets.Container(
+            widgets=[align_button, self.align_axis], layout="horizontal"
+        )
+        self.append(container)
+
+        rot90 = widgets.PushButton(text="Rotate 90", enabled=True)
+        rotneg90 = widgets.PushButton(text="Rotate -90", enabled=True)
+
+        rot90.changed.connect(self._rot90)
+        rotneg90.changed.connect(self._rotneg90)
+
+        container = widgets.Container(widgets=[rot90, rotneg90], layout="horizontal")
+        self.append(container)
+
+    def _rot90(self, swap_axes: bool = False):
+        active_layer = self.viewer.layers.selection.active
+        if active_layer is None:
+            return None
+        elif self.viewer.dims.ndisplay != 2:
+            return None
+
+        align_axis = self.align_axis.value
+        if self.align_axis.value is None:
+            align_axis = self.viewer.dims.order[0]
+
+        axes = [
+            align_axis,
+            *[i for i in range(len(self.viewer.dims.order)) if i != align_axis],
+        ][:2]
+        axes = axes[::-1] if swap_axes else axes
+        active_layer.data = np.rot90(active_layer.data, k=1, axes=axes)
+
+    def _rotneg90(self):
+        return self._rot90(swap_axes=True)
+
+    def _get_active_layer_dims(self, *args):
+        active_layer = self.viewer.layers.selection.active
+        if active_layer is None:
+            return ()
+        return [i for i in range(active_layer.data.ndim)]
+
+    def _update_align_axis(self, *args):
+        self.align_axis.choices = self._get_active_layer_dims()
+
+    def _align_with_axis(self):
+        active_layer = self.viewer.layers.selection.active
+
+        if self.align_axis.value is None:
+            return None
+
+        if active_layer.metadata.get("is_aligned", None) == self.align_axis.value:
+            return None
+
+        alignment_axis = np.zeros(active_layer.data.ndim)
+        alignment_axis[int(self.align_axis.value)] = 1
+
+        coords = np.array(np.where(active_layer.data > 0)).T
+        centered_coords = coords - np.mean(coords, axis=0)
+        cov_matrix = np.cov(centered_coords, rowvar=False)
+
+        eigenvalues, eigenvectors = np.linalg.eigh(cov_matrix)
+        principal_eigenvector = eigenvectors[:, np.argmax(eigenvalues)]
+
+        rotation_axis = np.cross(principal_eigenvector, alignment_axis)
+        rotation_angle = np.arccos(np.dot(principal_eigenvector, alignment_axis))
+        k = rotation_axis / np.linalg.norm(rotation_axis)
+        K = np.array([[0, -k[2], k[1]], [k[2], 0, -k[0]], [-k[1], k[0], 0]])
+        rotation_matrix = np.eye(3)
+        rotation_matrix += np.sin(rotation_angle) * K
+        rotation_matrix += (1 - np.cos(rotation_angle)) * np.dot(K, K)
+
+        rotated_data = Density.rotate_array(
+            arr=active_layer.data,
+            rotation_matrix=rotation_matrix,
+            use_geometric_center=False,
+            order=1,
+        )
+        eps = np.finfo(rotated_data.dtype).eps
+        rotated_data[rotated_data < eps] = 0
+
+        active_layer.metadata["is_aligned"] = int(self.align_axis.value)
+        active_layer.data = rotated_data
 
 
 class ExportWidget(widgets.Container):
@@ -803,16 +922,36 @@ class PointCloudWidget(widgets.Container):
             options=options,
         )
 
-        if filename:
-            layer = self.viewer.layers.selection.active
-            if layer and isinstance(layer, napari.layers.Points):
-                original_dataframe = self.dataframes.get(layer.name, pd.DataFrame())
+        if not filename:
+            return None
 
-                export_data = pd.DataFrame(layer.data, columns=["z", "y", "x"])
-                merged_data = pd.merge(
-                    export_data, original_dataframe, on=["z", "y", "x"], how="left"
-                )
-                merged_data.to_csv(filename, sep="\t", index=False)
+        layer = self.viewer.layers.selection.active
+        if layer and isinstance(layer, napari.layers.Points):
+            original_dataframe = self.dataframes.get(
+                layer.name, pd.DataFrame(columns=["z", "y", "x"])
+            )
+
+            export_data = pd.DataFrame(layer.data, columns=["z", "y", "x"])
+            merged_data = pd.merge(
+                export_data, original_dataframe, on=["z", "y", "x"], how="left"
+            )
+
+            merged_data["z"] = merged_data["z"].astype(int)
+            merged_data["y"] = merged_data["y"].astype(int)
+            merged_data["x"] = merged_data["x"].astype(int)
+
+            euler_columns = ["euler_z", "euler_y", "euler_x"]
+            for col in euler_columns:
+                if col not in merged_data.columns:
+                    continue
+                merged_data[col] = merged_data[col].fillna(0)
+
+            if "score" in merged_data.columns:
+                merged_data["score"] = merged_data["score"].fillna(1)
+            if "detail" in merged_data.columns:
+                merged_data["detail"] = merged_data["detail"].fillna(2)
+
+            merged_data.to_csv(filename, sep="\t", index=False)
 
     def _get_load_path(self, event):
         options = QFileDialog.Options()
@@ -830,6 +969,13 @@ class PointCloudWidget(widgets.Container):
         dataframe = pd.read_csv(filename, sep="\t")
         points = dataframe[["z", "y", "x"]].values
         layer_name = filename.split("/")[-1]
+
+        if "score" not in dataframe.columns:
+            dataframe["score"] = 1
+
+        if "detail" not in dataframe.columns:
+            dataframe["detail"] = -2
+
         point_properties = {
             "score": np.array(dataframe["score"].values),
             "detail": np.array(dataframe["detail"].values),
@@ -849,6 +995,38 @@ class PointCloudWidget(widgets.Container):
         self.dataframes[layer_name] = dataframe
 
 
+class MatchingWidget(widgets.Container):
+    def __init__(self, viewer):
+        super().__init__(layout="vertical")
+
+        self.viewer = viewer
+        self.dataframes = {}
+
+        self.import_button = widgets.PushButton(name="Import", text="Import Pickle")
+        self.import_button.clicked.connect(self._get_load_path)
+
+        self.append(self.import_button)
+
+    def _get_load_path(self, event):
+        options = QFileDialog.Options()
+        filename, _ = QFileDialog.getOpenFileName(
+            self.native,
+            "Open Pickle File...",
+            "",
+            "Pickle Files (*.pickle);;All Files (*)",
+            options=options,
+        )
+        if filename:
+            self._load_data(filename)
+
+    def _load_data(self, filename):
+        data = load_pickle(filename)
+
+        _ = self.viewer.add_image(data=data[2], name="Rotations", colormap="orange")
+
+        _ = self.viewer.add_image(data=data[0], name="Scores", colormap="turbo")
+
+
 def main():
     viewer = napari.Viewer()
 
@@ -856,10 +1034,16 @@ def main():
     mask_widget = MaskWidget(viewer)
     export_widget = ExportWidget(viewer)
     point_cloud = PointCloudWidget(viewer)
+    matching_widget = MatchingWidget(viewer)
+    alignment_widget = AlignmentWidget(viewer)
 
     viewer.window.add_dock_widget(widget=filter_widget, name="Preprocess", area="right")
+    viewer.window.add_dock_widget(
+        widget=alignment_widget, name="Alignment", area="right"
+    )
     viewer.window.add_dock_widget(widget=mask_widget, name="Mask", area="right")
     viewer.window.add_dock_widget(widget=point_cloud, name="PointCloud", area="left")
+    viewer.window.add_dock_widget(widget=matching_widget, name="Matching", area="left")
 
     viewer.window.add_dock_widget(widget=export_widget, name="Export", area="right")
 

scripts/refine_matches.py

@@ -0,0 +1,218 @@
+#!python3
+""" CLI to refine template matching candidates.
+
+    Copyright (c) 2023-2024 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+import argparse
+from time import time
+
+import numpy as np
+from numpy.typing import NDArray
+
+from tme.backends import backend
+from tme import Density, Structure
+from tme.matching_data import MatchingData
+from tme.analyzer import MaxScoreOverRotations, MaxScoreOverTranslations
+from tme.matching_utils import (
+    load_pickle,
+    get_rotation_matrices,
+    compute_parallelization_schedule,
+)
+from tme.matching_exhaustive import scan_subsets, MATCHING_EXHAUSTIVE_REGISTER
+
+from postprocess import Orientations
+from match_template import load_and_validate_mask
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Refine Template Matching Orientations."
+    )
+    parser.add_argument(
+        "--input_file",
+        required=True,
+        type=str,
+        help="Path to the output of match_template.py.",
+    )
+    parser.add_argument(
+        "--orientations",
+        required=True,
+        type=str,
+        help="Path to orientations from postprocess.py.",
+    )
+    parser.add_argument(
+        "--output_file",
+        required=True,
+        help="Path to the refined output orientations.",
+    )
+    parser.add_argument(
+        "-n",
+        dest="cores",
+        required=False,
+        type=int,
+        default=4,
+        help="Number of cores used for template matching.",
+    )
+
+    args = parser.parse_args()
+
+    return args
+
+
+def load_template(filepath: str, sampling_rate: NDArray) -> "Density":
+    try:
+        template = Density.from_file(filepath)
+    except ValueError:
+        template = Structure.from_file(filepath)
+        template = Density.from_structure(template, sampling_rate=sampling_rate)
+
+    return template
+
+
+def main():
+    args = parse_args()
+    meta = load_pickle(args.input_file)[-1]
+    target_origin, _, sampling_rate, cli_args = meta
+
+    orientations = Orientations.from_file(
+        filename=args.orientations, file_format="text"
+    )
+
+    template = load_template(cli_args.template, sampling_rate)
+    template_mask = load_and_validate_mask(
+        mask_target=template, mask_path=cli_args.template_mask
+    )
+
+    if not cli_args.no_centering:
+        template, translation = template.centered(0)
+
+    if template_mask is None:
+        template_mask = template.empty
+        if not cli_args.no_centering:
+            enclosing_box = template.minimum_enclosing_box(
+                0, use_geometric_center=False
+            )
+            template_mask.adjust_box(enclosing_box)
+
+        template_mask.data[:] = 1
+        translation = np.zeros_like(translation)
+
+    template_mask.pad(template.shape, center=False)
+    origin_translation = np.divide(
+        np.subtract(template.origin, template_mask.origin), template.sampling_rate
+    )
+    translation = np.add(translation, origin_translation)
+
+    template_mask = template_mask.rigid_transform(
+        rotation_matrix=np.eye(template_mask.data.ndim),
+        translation=-translation,
+        order=1,
+    )
+    template_mask.origin = template.origin.copy()
+
+    target = Density.from_file(cli_args.target)
+    peaks = orientations.translations.astype(int)
+    half_shape = np.divide(template.shape, 2).astype(int)
+    observation_starts = np.subtract(peaks, half_shape)
+    observation_stops = np.add(peaks, half_shape) + np.mod(template.shape, 2).astype(
+        int
+    )
+
+    pruned_starts = np.maximum(observation_starts, 0)
+    pruned_stops = np.minimum(observation_stops, target.shape)
+
+    keep_peaks = (
+        np.sum(
+            np.multiply(
+                observation_starts == pruned_starts, observation_stops == pruned_stops
+            ),
+            axis=1,
+        )
+        == observation_starts.shape[1]
+    )
+    observation_starts = observation_starts[keep_peaks]
+    observation_stops = observation_stops[keep_peaks]
+
+    observation_slices = [
+        tuple(slice(s, e) for s, e in zip(start_row, stop_row))
+        for start_row, stop_row in zip(observation_starts, observation_stops)
+    ]
+
+    matching_data = MatchingData(target=target, template=template)
+    matching_data.rotations = np.eye(template.data.ndim).reshape(1, 3, 3)
+
+    target_pad = matching_data.target_padding(pad_target=True)
+    out_shape = np.add(target_pad, template.shape).astype(int)
+
+    observations = np.zeros((len(observation_slices), *out_shape))
+
+
+    for idx, obs_slice in enumerate(observation_slices):
+        subset = matching_data.subset_by_slice(
+            target_slice=obs_slice,
+            target_pad=target_pad,
+            invert_target=cli_args.invert_target_contrast,
+        )
+        xd = template.copy()
+        xd.pad(subset.target.shape, center = True)
+        # observations[idx] = subset.target
+        observations[idx] = xd.data
+
+
+    matching_data = MatchingData(target=observations, template=template)
+    matching_data._set_batch_dimension(target_dims=0, template_dims=None)
+    matching_data.rotations = get_rotation_matrices(15)
+    if template_mask is not None:
+        matching_data.template_mask = template_mask.data
+
+    template_box = np.ones(len(matching_data._output_template_shape), dtype=int)
+    target_padding = np.zeros_like(matching_data._output_target_shape)
+
+    scoring_method = "FLC"
+    callback_class = MaxScoreOverRotations
+    splits, schedule = compute_parallelization_schedule(
+        shape1=matching_data._output_target_shape,
+        shape2=template_box,
+        shape1_padding=target_padding,
+        max_cores=args.cores,
+        max_ram=backend.get_available_memory(),
+        split_only_outer=False,
+        matching_method=scoring_method,
+        analyzer_method=callback_class.__name__,
+        backend=backend._backend_name,
+        float_nbytes=backend.datatype_bytes(backend._default_dtype),
+        complex_nbytes=backend.datatype_bytes(backend._complex_dtype),
+        integer_nbytes=backend.datatype_bytes(backend._default_dtype_int),
+    )
+
+    matching_setup, matching_score = MATCHING_EXHAUSTIVE_REGISTER[scoring_method]
+
+    start = time()
+    candidates = scan_subsets(
+        matching_data=matching_data,
+        matching_score=matching_score,
+        matching_setup=matching_setup,
+        callback_class=callback_class,
+        callback_class_args={
+            # "score_space_shape" : (
+                # matching_data.rotations.shape[0], observations.shape[0]
+            # ),
+            # "score_space_dtype" : backend._default_dtype,
+            # "template_shape" : (matching_data.rotations.shape[0], *matching_data._template.shape)
+        },
+        target_splits=splits,
+        job_schedule=schedule,
+        pad_target_edges=False,
+        pad_fourier=False,
+        interpolation_order=cli_args.interpolation_order,
+    )
+    print(candidates[0].max())
+    Density(candidates[0][0]).to_file("scores.mrc")
+
+    runtime = time() - start
+    print(f"\nRuntime real: {runtime:.3f}s user: {(runtime * args.cores):.3f}s.")
+
+
+if __name__ == "__main__":
+    main()

tme/__init__.py

@@ -1,6 +1,7 @@
+from . import extensions
 from .__version__ import __version__
 from .density import Density
 from .preprocessor import Preprocessor
 from .structure import Structure
+from .orientations import Orientations
 from .matching_optimization import FitRefinement
-from . import extensions

tme/__version__.py

@@ -1 +1 @@
-__version__ = "0.1.8"
+__version__ = "0.2.0b"