pytme 0.2.2__cp311-cp311-macosx_14_0_arm64.whl → 0.2.3__cp311-cp311-macosx_14_0_arm64.whl

scripts/preprocess.py

@@ -1,93 +1,132 @@
 #!python3
-""" Apply tme.preprocessor.Preprocessor methods to an input file based
-    on a provided yaml configuration obtaiend from preprocessor_gui.py.
+""" Preprocessing routines for template matching.
 
     Copyright (c) 2023 European Molecular Biology Laboratory
 
     Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
-import yaml
+import warnings
 import argparse
-import textwrap
-from tme import Preprocessor, Density
+import numpy as np
+
+from tme import Density, Structure
+from tme.backends import backend as be
+from tme.preprocessing.frequency_filters import BandPassFilter
 
 
 def parse_args():
     parser = argparse.ArgumentParser(
-        description=textwrap.dedent(
-            """
-        Apply preprocessing to an input file based on a provided YAML configuration.
-
-        Expected YAML file format:
-        ```yaml
-        <method_name>:
-            <parameter1>: <value1>
-            <parameter2>: <value2>
-            ...
-        ```
-        """
-        ),
-        formatter_class=argparse.RawDescriptionHelpFormatter,
+        description="Perform template matching preprocessing.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
-    parser.add_argument(
-        "-i",
-        "--input_file",
+
+    io_group = parser.add_argument_group("Input / Output")
+    io_group.add_argument(
+        "-m",
+        "--data",
+        dest="data",
         type=str,
         required=True,
-        help="Path to the input data file in CCP4/MRC format.",
+        help="Path to a file in PDB/MMCIF, CCP4/MRC, EM, H5 or a format supported by "
+        "tme.density.Density.from_file "
+        "https://kosinskilab.github.io/pyTME/reference/api/tme.density.Density.from_file.html",
     )
-    parser.add_argument(
-        "-y",
-        "--yaml_file",
+    io_group.add_argument(
+        "-o",
+        "--output",
+        dest="output",
         type=str,
         required=True,
-        help="Path to the YAML configuration file.",
+        help="Path the output should be written to.",
     )
-    parser.add_argument(
-        "-o",
-        "--output_file",
-        type=str,
+
+    box_group = parser.add_argument_group("Box")
+    box_group.add_argument(
+        "--box_size",
+        dest="box_size",
+        type=int,
         required=True,
-        help="Path to output file in CPP4/MRC format..",
+        help="Box size of the output",
     )
-    parser.add_argument(
-        "--compress", action="store_true", help="Compress the output file using gzip."
+    box_group.add_argument(
+        "--sampling_rate",
+        dest="sampling_rate",
+        type=float,
+        required=True,
+        help="Sampling rate of the output file.",
     )
 
+    modulation_group = parser.add_argument_group("Modulation")
+    modulation_group.add_argument(
+        "--invert_contrast",
+        dest="invert_contrast",
+        action="store_true",
+        required=False,
+        help="Inverts the template contrast.",
+    )
+    modulation_group.add_argument(
+        "--lowpass",
+        dest="lowpass",
+        type=float,
+        required=False,
+        default=None,
+        help="Lowpass filter the template to the given resolution. Nyquist by default. "
+        "A value of 0 disables the filter.",
+    )
+    modulation_group.add_argument(
+        "--no_centering",
+        dest="no_centering",
+        action="store_true",
+        help="Assumes the template is already centered and omits centering.",
+    )
     args = parser.parse_args()
-
     return args
 
 
 def main():
     args = parse_args()
-    with open(args.yaml_file, "r") as f:
-        preprocess_settings = yaml.safe_load(f)
 
-    if len(preprocess_settings) > 1:
-        raise NotImplementedError(
-            "Multiple preprocessing methods specified. "
-            "The script currently supports one method at a time."
+    try:
+        data = Structure.from_file(args.data)
+        data = Density.from_structure(data, sampling_rate=args.sampling_rate)
+    except NotImplementedError:
+        data = Density.from_file(args.data)
+
+    if not args.no_centering:
+        data, _ = data.centered(0)
+
+    recommended_box = be.compute_convolution_shapes([args.box_size], [1])[1][0]
+    if recommended_box != args.box_size:
+        warnings.warn(
+            f"Consider using --box_size {recommended_box} instead of {args.box_size}."
         )
 
-    method_name = list(preprocess_settings.keys())[0]
-    if not hasattr(Preprocessor, method_name):
-        raise ValueError(f"Method {method_name} does not exist in Preprocessor.")
+    data.pad(
+        np.multiply(args.box_size, np.divide(args.sampling_rate, data.sampling_rate)),
+        center=True,
+    )
+
+    bpf_mask = 1
+    lowpass = 2 * args.sampling_rate if args.lowpass is None else args.lowpass
+    if args.lowpass != 0:
+        bpf_mask = BandPassFilter(
+            lowpass=lowpass,
+            highpass=None,
+            use_gaussian=True,
+            return_real_fourier=True,
+            shape_is_real_fourier=False,
+        )(shape=data.shape)["data"]
 
-    density = Density.from_file(args.input_file)
-    output = density.empty
+    data_ft = np.fft.rfftn(data.data, s=data.shape)
+    data_ft = np.multiply(data_ft, bpf_mask, out=data_ft)
+    data.data = np.fft.irfftn(data_ft, s=data.shape).real
 
-    method_params = preprocess_settings[method_name]
-    preprocessor = Preprocessor()
-    method = getattr(preprocessor, method_name, None)
-    if not method:
-        raise ValueError(
-            f"{method} does not exist in dge.preprocessor.Preprocessor class."
-        )
+    data = data.resample(args.sampling_rate, method="spline", order=3)
 
-    output.data = method(template=density.data, **method_params)
-    output.to_file(args.output_file, gzip=args.compress)
+    if args.invert_contrast:
+        data.data = data.data * -1
 
+    data.to_file(args.output)
 
 if __name__ == "__main__":
-    main()
+    main()

scripts/preprocessor_gui.py

@@ -132,14 +132,6 @@ def local_gaussian_filter(
     )
 
 
-def ntree(
-    template: NDArray,
-    sigma_range: Tuple[float, float],
-    **kwargs: dict,
-) -> NDArray:
-    return preprocessor.ntree_filter(template=template, sigma_range=sigma_range)
-
-
 def mean(
     template: NDArray,
     width: int,
@@ -197,6 +189,10 @@ def compute_power_spectrum(template: NDArray) -> NDArray:
     return np.fft.fftshift(np.log(np.abs(np.fft.fftn(template))))
 
 
+def invert_contrast(template: NDArray) -> NDArray:
+    return template * -1
+
+
 def widgets_from_function(function: Callable, exclude_params: List = ["self"]):
     """
     Creates list of magicui widgets by inspecting function typing ann
@@ -252,13 +248,13 @@ WRAPPED_FUNCTIONS = {
     "gaussian_filter": gaussian_filter,
     "bandpass_filter": bandpass_filter,
     "edge_gaussian_filter": edge_gaussian_filter,
-    "ntree_filter": ntree,
     "local_gaussian_filter": local_gaussian_filter,
     "difference_of_gaussian_filter": difference_of_gaussian_filter,
     "mean_filter": mean,
     "wedge_filter": wedge,
     "power_spectrum": compute_power_spectrum,
     "ctf": ctf_filter,
+    "invert_contrast": invert_contrast,
 }
 
 EXCLUDED_FUNCTIONS = [
@@ -634,6 +630,7 @@ class MaskWidget(widgets.Container):
 
         data = active_layer.data.copy()
         cutoff = np.quantile(data, self.percentile_range_edit.value / 100)
+        cutoff = max(cutoff, np.finfo(np.float32).resolution)
         data[data < cutoff] = 0
 
         center_of_mass = Density.center_of_mass(np.abs(data), 0)

tme/__version__.py

@@ -1 +1 @@
-__version__ = "0.2.2"
+__version__ = "0.2.3"

tme/analyzer.py

@@ -459,14 +459,14 @@ class PeakCaller(ABC):
 
         final_order = top_scores[
             filter_points_indices(
-                coordinates=peak_positions[top_scores],
+                coordinates=peak_positions[top_scores, :],
                 min_distance=self.min_distance,
                 batch_dims=self.batch_dims,
             )
         ]
 
-        self.peak_list[0] = peak_positions[final_order,]
-        self.peak_list[1] = rotations[final_order,]
+        self.peak_list[0] = peak_positions[final_order, :]
+        self.peak_list[1] = rotations[final_order, :]
         self.peak_list[2] = peak_scores[final_order]
         self.peak_list[3] = peak_details[final_order]
 
@@ -475,6 +475,7 @@ class PeakCaller(ABC):
         fast_shape: Tuple[int],
         targetshape: Tuple[int],
         templateshape: Tuple[int],
+        convolution_shape: Tuple[int] = None,
         fourier_shift: Tuple[int] = None,
         convolution_mode: str = None,
         shared_memory_handler=None,
@@ -488,6 +489,7 @@ class PeakCaller(ABC):
             return self
 
         # Wrap peaks around score space
+        convolution_shape = be.to_backend_array(convolution_shape)
         fast_shape = be.to_backend_array(fast_shape)
         if fourier_shift is not None:
             fourier_shift = be.to_backend_array(fourier_shift)
@@ -501,10 +503,9 @@ class PeakCaller(ABC):
             )
 
         # Remove padding to fast Fourier (and potential full convolution) shape
+        output_shape = convolution_shape
         targetshape = be.to_backend_array(targetshape)
         templateshape = be.to_backend_array(templateshape)
-        fast_shape = be.minimum(be.add(targetshape, templateshape) - 1, fast_shape)
-        output_shape = fast_shape
         if convolution_mode == "same":
             output_shape = targetshape
         elif convolution_mode == "valid":
@@ -515,7 +516,7 @@ class PeakCaller(ABC):
 
         output_shape = be.to_backend_array(output_shape)
         starts = be.astype(
-            be.divide(be.subtract(fast_shape, output_shape), 2),
+            be.divide(be.subtract(convolution_shape, output_shape), 2),
             be._int_dtype,
         )
         stops = be.add(starts, output_shape)
@@ -1019,6 +1020,7 @@ class MaxScoreOverRotations:
         self,
         targetshape: Tuple[int],
         templateshape: Tuple[int],
+        convolution_shape: Tuple[int],
         fourier_shift: Tuple[int] = None,
         convolution_mode: str = None,
         shared_memory_handler=None,
@@ -1039,6 +1041,7 @@ class MaxScoreOverRotations:
             "s1": targetshape,
             "s2": templateshape,
             "convolution_mode": convolution_mode,
+            "convolution_shape": convolution_shape,
         }
         if convolution_mode is not None:
             scores = apply_convolution_mode(scores, **convargs)

tme/backends/__init__.py

@@ -147,7 +147,7 @@ class BackendManager:
         _dependencies = {
             "numpyfftw": "numpy",
             "cupy": "cupy",
-            "pytorch": "pytorch",
+            "pytorch": "torch",
             "mlx": "mlx",
             "jax": "jax",
         }

tme/backends/_jax_utils.py

@@ -19,9 +19,9 @@ def _correlate(template: BackendArray, ft_target: BackendArray) -> BackendArray:
     """
     Computes :py:meth:`tme.matching_exhaustive.cc_setup`.
     """
-    template_ft = jnp.fft.rfftn(template)
+    template_ft = jnp.fft.rfftn(template, s=template.shape)
     template_ft = template_ft.at[:].multiply(ft_target)
-    correlation = jnp.fft.irfftn(template_ft)
+    correlation = jnp.fft.irfftn(template_ft, s=template.shape)
     return correlation
 
 
@@ -77,14 +77,15 @@ def _reciprocal_target_std(
     --------
     :py:meth:`tme.matching_exhaustive.flc_scoring`.
     """
-    ft_template_mask = jnp.fft.rfftn(template_mask)
+    ft_shape = template_mask.shape
+    ft_template_mask = jnp.fft.rfftn(template_mask, s=ft_shape)
 
     # E(X^2)- E(X)^2
-    exp_sq = jnp.fft.irfftn(ft_target2 * ft_template_mask)
+    exp_sq = jnp.fft.irfftn(ft_target2 * ft_template_mask, s=ft_shape)
     exp_sq = exp_sq.at[:].divide(n_observations)
 
     ft_template_mask = ft_template_mask.at[:].multiply(ft_target)
-    sq_exp = jnp.fft.irfftn(ft_template_mask)
+    sq_exp = jnp.fft.irfftn(ft_template_mask, s=ft_shape)
     sq_exp = sq_exp.at[:].divide(n_observations)
     sq_exp = sq_exp.at[:].power(2)
 
@@ -99,7 +100,7 @@ def _reciprocal_target_std(
 
 
 def _apply_fourier_filter(arr: BackendArray, arr_filter: BackendArray) -> BackendArray:
-    arr_ft = jnp.fft.rfftn(arr)
+    arr_ft = jnp.fft.rfftn(arr, s=arr.shape)
     arr_ft = arr_ft.at[:].multiply(arr_filter)
     return arr.at[:].set(jnp.fft.irfftn(arr_ft, s=arr.shape))
 
@@ -107,6 +108,7 @@ def _apply_fourier_filter(arr: BackendArray, arr_filter: BackendArray) -> Backen
 def _identity(arr: BackendArray, arr_filter: BackendArray) -> BackendArray:
     return arr
 
+
 @partial(
     pmap,
     in_axes=(0,) + (None,) * 6,
@@ -127,8 +129,8 @@ def scan(
     if hasattr(target_filter, "shape"):
         target = _apply_fourier_filter(target, target_filter)
 
-    ft_target = jnp.fft.rfftn(target)
-    ft_target2 = jnp.fft.rfftn(jnp.square(target))
+    ft_target = jnp.fft.rfftn(target, s=fast_shape)
+    ft_target2 = jnp.fft.rfftn(jnp.square(target), s=fast_shape)
     inv_denominator, target, scoring_func = None, None, _flc_scoring
     if not rotate_mask:
         n_observations = jnp.sum(template_mask)

tme/backends/cupy_backend.py

@@ -149,10 +149,10 @@ class CupyBackend(NumpyFFTWBackend):
             cache.clear()
 
         def rfftn(arr: CupyArray, out: CupyArray) -> CupyArray:
-            return cufft.rfftn(arr)
+            return cufft.rfftn(arr, s=fast_shape)
 
         def irfftn(arr: CupyArray, out: CupyArray) -> CupyArray:
-            return cufft.irfftn(arr)
+            return cufft.irfftn(arr, s=fast_shape)
 
         PLAN_CACHE[current_device] = [fast_shape, fast_ft_shape]
 
@@ -167,11 +167,6 @@ class CupyBackend(NumpyFFTWBackend):
         fast_shape = [next_fast_len(x, real=True) for x in convolution_shape]
         fast_ft_shape = list(fast_shape[:-1]) + [fast_shape[-1] // 2 + 1]
 
-        # This almost never happens but avoid cuFFT casting errors
-        is_odd = fast_shape[-1] % 2
-        fast_shape[-1] += is_odd
-        fast_ft_shape[-1] += is_odd
-
         return convolution_shape, fast_shape, fast_ft_shape
 
     def max_filter_coordinates(self, score_space, min_distance: Tuple[int]):

tme/backends/jax_backend.py

@@ -119,19 +119,6 @@ class JaxBackend(NumpyFFTWBackend):
 
         return rfftn, irfftn
 
-    def compute_convolution_shapes(
-        self, arr1_shape: Tuple[int], arr2_shape: Tuple[int]
-    ) -> Tuple[List[int], List[int], List[int]]:
-        conv_shape, fast_shape, fast_ft_shape = super().compute_convolution_shapes(
-            arr1_shape, arr2_shape
-        )
-
-        is_odd = fast_shape[-1] % 2
-        fast_shape[-1] += is_odd
-        fast_ft_shape[-1] += is_odd
-
-        return conv_shape, fast_shape, fast_ft_shape
-
     def rigid_transform(
         self,
         arr: BackendArray,
@@ -144,8 +131,8 @@ class JaxBackend(NumpyFFTWBackend):
         **kwargs,
     ) -> Tuple[BackendArray, BackendArray]:
         rotate_mask = arr_mask is not None
-        center = self.divide(self.to_backend_array(arr.shape), 2)[:, None]
 
+        center = self.divide(self.to_backend_array(arr.shape) - 1, 2)[:, None]
         indices = self._array_backend.indices(arr.shape, dtype=self._float_dtype)
         indices = indices.reshape((arr.ndim, -1))
         indices = indices.at[:].add(-center)
@@ -200,7 +187,7 @@ class JaxBackend(NumpyFFTWBackend):
         target_shape = tuple(
             (x.stop - x.start + p) for x, p in zip(splits[0][0], target_pad)
         )
-        fast_shape, fast_ft_shape, shift = matching_data._fourier_padding(
+        conv_shape, fast_shape, fast_ft_shape, shift = matching_data._fourier_padding(
             target_shape=self.to_numpy_array(target_shape),
             template_shape=self.to_numpy_array(matching_data._template.shape),
             pad_fourier=False,
@@ -210,13 +197,25 @@ class JaxBackend(NumpyFFTWBackend):
             "convolution_mode": convolution_mode,
             "fourier_shift": shift,
             "targetshape": target_shape,
-            "templateshape": matching_data._template.shape,
+            "templateshape": matching_data.template.shape,
+            "convolution_shape": conv_shape,
         }
 
         create_target_filter = matching_data.target_filter is not None
         create_template_filter = matching_data.template_filter is not None
         create_filter = create_target_filter or create_template_filter
 
+        # Applying the filter leads to more FFTs
+        fastt_shape = matching_data._template.shape
+        if create_template_filter:
+            _, fastt_shape, _, tshift = matching_data._fourier_padding(
+                target_shape=self.to_numpy_array(matching_data._template.shape),
+                template_shape=self.to_numpy_array(
+                    [1 for _ in matching_data._template.shape]
+                ),
+                pad_fourier=False,
+            )
+
         ret, template_filter, target_filter = [], 1, 1
         rotation_mapping = {
             self.tobytes(matching_data.rotations[i]): i
@@ -246,12 +245,12 @@ class JaxBackend(NumpyFFTWBackend):
 
             if create_template_filter:
                 template_filter = matching_data.template_filter(
-                    shape=matching_data._template.shape, **filter_args
+                    shape=fastt_shape, **filter_args
                 )["data"]
                 template_filter = template_filter.at[(0,) * template_filter.ndim].set(0)
 
             if create_target_filter:
-                target_filter = matching_data.template_filter(
+                target_filter = matching_data.target_filter(
                     shape=fast_shape, **filter_args
                 )["data"]
                 target_filter = target_filter.at[(0,) * target_filter.ndim].set(0)
@@ -260,8 +259,8 @@ class JaxBackend(NumpyFFTWBackend):
             base, targets = None, self._array_backend.stack(targets)
             scores, rotations = scan_inner(
                 targets,
-                matching_data.template,
-                matching_data.template_mask,
+                self.topleft_pad(matching_data.template, fastt_shape),
+                self.topleft_pad(matching_data.template_mask, fastt_shape),
                 matching_data.rotations,
                 template_filter,
                 target_filter,
@@ -280,3 +279,18 @@ class JaxBackend(NumpyFFTWBackend):
                 ret.append(tuple(temp._postprocess(**analyzer_args)))
 
         return ret
+
+    def get_available_memory(self) -> int:
+        import jax
+
+        _memory = {"cpu": 0, "gpu": 0}
+        for device in jax.devices():
+            if device.platform == "cpu":
+                _memory["cpu"] = super().get_available_memory()
+            else:
+                mem_stats = device.memory_stats()
+                _memory["gpu"] += mem_stats.get("bytes_limit", 0)
+
+        if _memory["gpu"] > 0:
+            return _memory["gpu"]
+        return _memory["cpu"]

tme/backends/npfftw_backend.py

@@ -186,7 +186,7 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
     def to_sharedarr(
         self, arr: NDArray, shared_memory_handler: type = None
     ) -> shm_type:
-        if type(shared_memory_handler) == SharedMemoryManager:
+        if isinstance(shared_memory_handler, SharedMemoryManager):
             shm = shared_memory_handler.SharedMemory(size=arr.nbytes)
         else:
             shm = shared_memory.SharedMemory(create=True, size=arr.nbytes)
@@ -347,7 +347,8 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
         cache: bool = False,
     ) -> Tuple[NDArray, NDArray]:
         translation = self.zeros(arr.ndim) if translation is None else translation
-        center = self.divide(self.to_backend_array(arr.shape), 2)
+
+        center = self.divide(self.to_backend_array(arr.shape) - 1, 2)
         if not use_geometric_center:
             center = self.center_of_mass(arr, cutoff=0)
 

tme/backends/pytorch_backend.py

@@ -81,13 +81,13 @@ class PytorchBackend(NumpyFFTWBackend):
 
     def max(self, *args, **kwargs) -> NDArray:
         ret = self._array_backend.amax(*args, **kwargs)
-        if type(ret) == self._array_backend.Tensor:
+        if isinstance(ret, self._array_backend.Tensor):
             return ret
         return ret[0]
 
     def min(self, *args, **kwargs) -> NDArray:
         ret = self._array_backend.amin(*args, **kwargs)
-        if type(ret) == self._array_backend.Tensor:
+        if isinstance(ret, self._array_backend.Tensor):
             return ret
         return ret[0]
 
@@ -154,7 +154,7 @@ class PytorchBackend(NumpyFFTWBackend):
             1, -1
         )
         if unraveled_coords.size(0) == 1:
-            return tuple(unraveled_coords[0, :].tolist())
+            return (unraveled_coords[0, :],)
 
         else:
             return tuple(unraveled_coords.T)
@@ -206,7 +206,9 @@ class PytorchBackend(NumpyFFTWBackend):
         else:
             raise NotImplementedError("Operation only implemented for 2 and 3D inputs.")
 
-        pool = func(kernel_size=min_distance, return_indices=True)
+        pool = func(
+            kernel_size=min_distance, padding=min_distance // 2, return_indices=True
+        )
         _, indices = pool(score_space.reshape(1, 1, *score_space.shape))
         coordinates = self.unravel_index(indices.reshape(-1), score_space.shape)
         coordinates = self.transpose(self.stack(coordinates))
@@ -217,7 +219,7 @@ class PytorchBackend(NumpyFFTWBackend):
 
     def from_sharedarr(self, args) -> TorchTensor:
         if self.device == "cuda":
-            return args[0]
+            return args
 
         shm, shape, dtype = args
         required_size = int(self._array_backend.prod(self.to_backend_array(shape)))
@@ -235,13 +237,12 @@ class PytorchBackend(NumpyFFTWBackend):
 
         nbytes = arr.numel() * arr.element_size()
 
-        if type(shared_memory_handler) == SharedMemoryManager:
+        if isinstance(shared_memory_handler, SharedMemoryManager):
             shm = shared_memory_handler.SharedMemory(size=nbytes)
         else:
             shm = shared_memory.SharedMemory(create=True, size=nbytes)
 
         shm.buf[:nbytes] = arr.numpy().tobytes()
-
         return shm, arr.shape, arr.dtype
 
     def transpose(self, arr):
@@ -415,6 +416,8 @@ class PytorchBackend(NumpyFFTWBackend):
             yield None
 
     def device_count(self) -> int:
+        if self.device == "cpu":
+            return 1
         return self._array_backend.cuda.device_count()
 
     def reverse(self, arr: TorchTensor) -> TorchTensor:

tme/density.py

@@ -116,8 +116,8 @@ class Density:
         response = "Density object at {}\nOrigin: {}, Sampling Rate: {}, Shape: {}"
         return response.format(
             hex(id(self)),
-            tuple(np.round(self.origin, 3)),
-            tuple(np.round(self.sampling_rate, 3)),
+            tuple(round(float(x), 3) for x in self.origin),
+            tuple(round(float(x), 3) for x in self.sampling_rate),
             self.shape,
         )
 
@@ -306,6 +306,10 @@ class Density:
                 "std": float(mrc.header.rms),
             }
 
+        non_standard_crs = not np.all(crs_index == (0, 1, 2))
+        if non_standard_crs:
+            warnings.warn("Non standard MAPC, MAPR, MAPS, adapting data and origin.")
+
         if is_gzipped(filename):
             if use_memmap:
                 warnings.warn(
@@ -315,6 +319,10 @@ class Density:
             use_memmap = False
 
         if subset is not None:
+            subset = tuple(
+                subset[i] if i < len(subset) else slice(0, data_shape[i])
+                for i in crs_index
+            )
             subset_shape = tuple(x.stop - x.start for x in subset)
             if np.allclose(subset_shape, data_shape):
                 return cls._load_mrc(
@@ -328,18 +336,16 @@ class Density:
                 dtype=data_type,
                 header_size=1024 + extended_header,
             )
-            return data, origin, sampling_rate, metadata
-
-        if not use_memmap:
+        elif subset is None and not use_memmap:
             with mrcfile.open(filename, header_only=False) as mrc:
                 data = mrc.data.astype(np.float32, copy=False)
         else:
             with mrcfile.mrcmemmap.MrcMemmap(filename, header_only=False) as mrc:
                 data = mrc.data
 
-        if not np.all(crs_index == (0, 1, 2)):
+        if non_standard_crs:
             data = np.transpose(data, crs_index)
-            start = np.take(start, crs_index)
+            origin = np.take(origin, crs_index)
 
         return data, origin, sampling_rate, metadata
 
@@ -873,6 +879,7 @@ class Density:
             mrc.header.nzstart, mrc.header.nystart, mrc.header.nxstart = np.rint(
                 np.divide(self.origin, self.sampling_rate)
             )
+            mrc.header.origin = tuple(x for x in self.origin)
             # mrcfile library expects origin to be in xyz format
             mrc.header.mapc, mrc.header.mapr, mrc.header.maps = (1, 2, 3)
             mrc.header["origin"] = tuple(self.origin[::-1])
@@ -1594,7 +1601,7 @@ class Density:
         rotation_matrix: NDArray,
         translation: NDArray = None,
         order: int = 3,
-        use_geometric_center: bool = False,
+        use_geometric_center: bool = True,
     ) -> "Density":
         """
         Performs a rigid transform of the class instance.