Repository: mjo22/cryojax (language: python)

Issue:
Should serialization with dataclasses-json be dropped?
equinox supports model serialization / deserialization in a very Pytree way: https://docs.kidger.site/equinox/api/serialisation/.

The upside to the dataclasses-json approach is that it is very simple to use and saves both static and traced data in the Pytree. `classmethod`s can be used for deserialization, and therefore user doesn’t need to know anything about the metadata that was stored in the saved model.

The downside is that it restricts the way classes are allowed to be written. Because when classes are deserialized they run `__init__`, one needs to make sure all initialization variables are stored as dataclass fields. `InitVars` are not allowed. This is not true in `equinox` because it just updates the values in some existing pytree. 

The other unfortunate part of writing cryojax classes is that we always have to use the `cryojax.core.field` when writing dataclass fields so that fields always have their default serialization metadata.

In the spirit of `cryojax` `Module`s not diverging from `equinox` `Module`s, it is tempting to just accept their serialization scheme.


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 986a250b..622a8111 100644
--- a/README.md
+++ b/README.md
@@ -29,7 +29,7 @@ cd cryojax
 python -m pip install .
 ```
 
-This will install the remaining dependencies, such as [equinox](https://github.com/patrick-kidger/equinox/) for jax-friendly dataclasses, [jaxlie](https://github.com/brentyi/jaxlie) for coordinate rotations and translations, [mrcfile](https://github.com/ccpem/mrcfile) for I/O, and [dataclasses-json](https://github.com/lidatong/dataclasses-json) for serialization.
+This will install the remaining dependencies, such as [equinox](https://github.com/patrick-kidger/equinox/) for jax-friendly dataclasses, [jaxlie](https://github.com/brentyi/jaxlie) for coordinate rotations and translations, and [mrcfile](https://github.com/ccpem/mrcfile) for I/O.
 
 The [jax-finufft](https://github.com/dfm/jax-finufft) package is an optional dependency used for non-uniform fast fourier transforms. These are included as an option for computing image projections. In this case, we recommend first following the `jax_finufft` installation instructions and then installing `cryojax`.
 
@@ -37,21 +37,22 @@ The [jax-finufft](https://github.com/dfm/jax-finufft) package is an optional dep
 
 Please note that this library is currently experimental and the API is subject to change! The following is a basic workflow to generate an image with a gaussian white noise model.
 
-First, instantiate the scattering model ("scattering") and its respective representation
-of an electron density ("density").
+First, instantiate the `ScatteringModel` and its respective representation
+of an `ElectronDensity`.
 
 ```python
 import jax
 import jax.numpy as jnp
 import cryojax.simulator as cs
 
-template = "example.mrc"
-density = cs.VoxelGrid.from_file(template)
+filename = "example.mrc"
+density = cs.VoxelGrid.from_file(filename)
+pixel_size = density.voxel_size
 manager = cs.ImageManager(shape=(320, 320))
-scattering = cs.FourierSliceExtract(manager, pixel_size=density.voxel_size)
+scattering = cs.FourierSliceExtract(manager, pixel_size=pixel_size)
 ```
 
-Here, `template` is a 3D electron density map in MRC format. This could be taken from the [EMDB](https://www.ebi.ac.uk/emdb/), or rasterized from a [PDB](https://www.rcsb.org/). [cisTEM](https://github.com/timothygrant80/cisTEM) provides an excellent rasterization tool in its image simulation program. In the above example, a voxel electron density in fourier space is loaded and the fourier-slice projection theorem is initialized. Note that we must explicitly set the pixel size of the projection image. Here, it is the same as the voxel size of the electron density. We can now instantiate the `Ensemble` of biological specimen.
+Here, `filename` is a 3D electron density map in MRC format. This could be taken from the [EMDB](https://www.ebi.ac.uk/emdb/), or rasterized from a [PDB](https://www.rcsb.org/). [cisTEM](https://github.com/timothygrant80/cisTEM) provides an excellent rasterization tool in its image simulation program. In the above example, a voxel electron density in fourier space is loaded and the fourier-slice projection theorem is initialized. Note that we must explicitly set the pixel size of the projection image. Here, it is the same as the voxel size of the electron density. We can now instantiate the `Ensemble` of biological specimen.
 
 ```python
 # Translations in Angstroms, angles in degrees
@@ -72,49 +73,58 @@ The stack of electron densities is stored in a single `ElectronDensity`, whose p
 Next, the model for the electron microscope. `Optics` and `Detector` models and their respective parameters are initialized. These are stored in the `Instrument` container.
 
 ```python
-new_pixel_size = scattering.pixel_size+0.02
 optics = cs.CTFOptics(defocus_u=10000.0, defocus_v=9800.0, defocus_angle=10.0)
-detector = cs.GaussianDetector(pixel_size=new_pixel_size, variance=cs.Constant(1.0))
+detector = cs.GaussianDetector(variance=cs.Constant(1.0))
 instrument = cs.Instrument(optics=optics, detector=detector)
 ```
 
-Here, the `Detector` has an optional pixel size that interpolates the rasterized pixel size to a new one. This is meant to correct for small mismeasurements in experimentally reported pixel size due to microscope alignment imperfections. The `CTFOptics` has all parameters used in CTFFIND4, which take their default values if not
+Here, the `Detector` is simply modeled by gaussian white noise. The `CTFOptics` has all parameters used in CTFFIND4, which take their default values if not
 explicitly configured here. Finally, we can instantiate the `ImagePipeline`.
 
 ```python
 key = jax.random.PRNGKey(seed=0)
-model = cs.ImagePipeline(scattering=scattering, ensemble=ensemble, instrument=instrument)
-image = model.sample(key)
+pipeline = cs.ImagePipeline(scattering=scattering, ensemble=ensemble, instrument=instrument)
+image = pipeline.sample(key)
 ```
 
 This computes an image using the noise model of the detector. One can also compute an image without the stochastic part of the model.
 
 ```python
-image = model.render()
+image = pipeline.render()
 ```
 
 Imaging models also accept a series of `Filter`s and `Mask`s. For example, one could add a `LowpassFilter`, `WhiteningFilter`, and a `CircularMask`.
 
 ```python
 micrograph = ...  # A micrograph used for whitening
-filter = cs.LowpassFilter(manager, cutoff=1.0)  # Cutoff modes above Nyquist frequency
-          * cs.WhiteningFilter(manager, micrograph=micrograph)
-mask = cs.CircularMask(manager, radius=1.0)     # Cutoff pixels above radius equal to (half) image size
-model = cs.ImagePipeline(
+freqs = manager.padded_frequency_grid  # Get the upsampled frequency grid
+coords = manager.coordinate_grid  # Get the coordinate grid
+filter = cs.LowpassFilter(freqs, cutoff=1.0)  # Cutoff modes above Nyquist frequency
+         * cs.WhiteningFilter(freqs, micrograph=micrograph)
+mask = cs.CircularMask(coords, radius=1.0)    # Cutoff pixels above radius equal to (half) image size
+pipeline = cs.ImagePipeline(
     scattering=scattering, ensemble=ensemble, instrument=instrument, filter=filter, mask=mask
     )
-image = model.sample(key)
+image = pipeline.sample(key)
 ```
 
-`cryojax` also defines a library of `Distribution`s, which inherit from the `ImagePipeline`. If a `GaussianImage` is instantiated, it is equipped with a the log likelihood function.
+`cryojax` also defines a library of `Distribution`s, which take an `ImagePipeline` as input. For example, instantiate an `IndependentFourierGaussian` distribution to call its log likelihood function.
 
 ```python
+from cryojax.utils import rfftn
+
+# Read observed data in real space
 observed = ...
-model = cs.GaussianImage(scattering=scattering, ensemble=ensemble, instrument=instrument)
+# Normalize to mean zero and standard deviation 1
+observed = manager.normalize_image(observed, is_real=True)
+# Upsample observed data in fourier space
+observed = rfftn(manager.pad_to_padded_shape(observed))
+# Instantiate distribution and compute
+model = cs.IndependentFourierGaussian(pipeline)
 log_likelihood = model.log_probability(observed)
 ```
 
-Note that the user may need to do preprocessing of `observed`, such as applying the relevant `Filter`s and `Mask`s.
+Note that in this example, the user must make sure `observed` is the expected shape and is in fourier space.
 
 Additional components can be plugged into the image formation model. For example, modeling the solvent is supported through the `ImagePipeline`'s `Ice` model. Models for exposure to the electron beam are supported through the `Instrument`'s `Exposure` model.
 
@@ -122,7 +132,7 @@ For these more advanced examples, see the tutorials section of the repository. I
 
 ## Creating a loss function
 
-In `jax`, we ultimately want to build a loss function and apply functional transformations to it. Assuming we have already globally configured our model components at our desired initial state, the below creates a loss function at an updated set of parameters. First, we must update the model.
+In `jax`, we may want to build a loss function and apply functional transformations to it. Assuming we have already globally configured our model components at our desired initial state, the below creates a loss function at an updated set of parameters. First, we must update the model.
 
 ```python
 
@@ -131,8 +141,14 @@ def update_model(model, params):
     """
     Update the model with equinox.tree_at (https://docs.kidger.site/equinox/api/manipulation/#equinox.tree_at).
     """
-    where = lambda model: (model.ensemble.pose.view_phi, model.instrument.optics.defocus_u, model.instrument.detector.pixel_size)
-    updated_model = eqx.tree_at(where, model, (params["view_phi"], params["defocus_u"], params["pixel_size"]))
+    where = lambda model: (
+        model.pipeline.ensemble.pose.view_phi,
+        model.pipeline.instrument.optics.defocus_u,
+        model.pipeline.scattering.pixel_size
+    )
+    updated_model = eqx.tree_at(
+        where, model, (params["view_phi"], params["defocus_u"], params["pixel_size"])
+    )
     return updated_model
 ```
 
@@ -149,19 +165,22 @@ def loss(params, model, observed):
 Finally, we can evaluate an updated set of parameters.
 
 ```python
-params = dict(view_phi=jnp.asarray(jnp.pi), defocus_u=jnp.asarray(9000.0), pixel_size=jnp.asarray(1.30))
+params = dict(
+    view_phi=jnp.asarray(jnp.pi),
+    defocus_u=jnp.asarray(9000.0),
+    pixel_size=jnp.asarray(density.voxel_size+0.02),
+)
 log_likelihood, grad = loss(params, model, observed)
 ```
 
-To summarize, this example creates a loss function at an updated set of `Pose`, `Optics`, and `Detector` parameters. In general, any `cryojax` `Module` may contain model parameters. The exception to this is in the `ImageManager`, `Filter`, and `Mask`. These classes do computation upon initialization, so they should not be explicitly instantiated in the loss function evaluation. Another gotcha is that if the `model` is not passed as an argument to the loss, there may be long compilation times because the electron density will be treated as static. However, this may result in slight speedups.
+To summarize, this example creates a loss function at an updated set of `Pose`, `Optics`, and `ScatteringModel` parameters. In general, any `cryojax` `Module` may contain model parameters. The exception to this is in the `ImageManager`, `Filter`, and `Mask`. These classes do computation upon initialization, so they should not be explicitly instantiated in the loss function evaluation. Another gotcha is that if the `model` is not passed as an argument to the loss, there may be long compilation times because the electron density will be treated as static. However, this may result in slight speedups.
 
 In general, there are many ways to write loss functions. See the [equinox](https://github.com/patrick-kidger/equinox/) documentation for more use cases.
 
 ## Features
 
 - Imaging models in `cryojax` support `jax` functional transformations, such as automatic differentiation with `grad`, paralellization with `vmap` and `pmap`, and just-in-time compilation with `jit`. Models also support GPU/TPU acceleration.
-- `cryojax.Module`s, including `ImagePipeline` models, are JSON serializable thanks to the package `dataclasses-json`. The method `Module.dumps` serializes the object as a JSON string, and `Module.loads` instantiates it from the string. For example, write a model to disk with `model.dump("model.json")` and instantiate it with `cs.ImagePipeline.load("model.json")`.
-- A `cryojax.Module` is just an `equinox.Module` with added serialization functionality. Therefore, the `equinox` ecosystem is available for usage!
+- A `cryojax.Module` is just an `equinox.Module`. Therefore, the `equinox` ecosystem is available for usage!
 
 ## Similar libraries
 
diff --git a/setup.py b/setup.py
index ef0b96ba..6e19ba4b 100644
--- a/setup.py
+++ b/setup.py
@@ -24,10 +24,9 @@
     "jaxtyping",
     "jaxlie",
     "mrcfile",
-    "dataclasses-json",
 ]
 EXTRA_REQUIRE = {
-    "test": ["pytest>=3.6", "pycistem"],
+    "test": ["pytest>=3.6", "pycistem", "gemmi"],
     "docs": [],
 }
 EXTRA_REQUIRE["coverage"] = EXTRA_REQUIRE["test"] + ["pytest-cov"]
diff --git a/src/cryojax/core.py b/src/cryojax/core.py
index 3a28f7ed..f794b36d 100644
--- a/src/cryojax/core.py
+++ b/src/cryojax/core.py
@@ -1,132 +1,77 @@
 """
-Core functionality in cryojax, i.e. base classes and serialization.
+Core functionality in cryojax, i.e. base classes and metadata.
 """
 
 from __future__ import annotations
 
-__all__ = ["field", "Module", "Buffer"]
+__all__ = ["field", "Module", "BufferModule"]
 
-from types import FunctionType
-from typing import Any, Union
+from typing import Any
 from jaxtyping import Array, ArrayLike
 
 import jax
 import jax.numpy as jnp
-import numpy as np
 
 import equinox as eqx
 
-import marshal
-import base64
-from dataclasses_json import DataClassJsonMixin, config
-from dataclasses_json.mm import JsonData
-
 
 def field(
     *,
-    encode: Any = Array,
+    stack: bool = True,
     **kwargs: Any,
 ) -> Any:
     """
-    Add metadata to usual dataclass fields.
-
-    Parameters
-    ----------
-    encode : `Any`
-        Type hint for the field's json encoding. If
-        `False`, do not encode the field.
+    Add default metadata to usual dataclass fields through dataclasses
+    and equinox.
+
+    Arguments
+    ---------
+    stack : `bool`
+        Metadata that indicates if a field is to be stacked
+        when giving Module fields a batch dimension. This metadata
+        is not currently supported in every cryojax Module, and is
+        only for internal cryojax batch dimension functionality.
+        In general, a user can add batch dimensions however they want.
+
+        As a rule of thumb, this should only be set to ``False`` if the
+        field in question is cumbersome and unecessary to store. See
+        ``cryojax.simulator.density.VoxelGrid`` for an example. This
+        argument is not used if ``static = True``.
     """
-    # Dataclass kwargs
-    metadata = kwargs.pop("metadata", {})
-    init = kwargs.pop("init", True)
-    # Equinox kwargs
+    # Equinox metadata
     static = kwargs.pop("static", False)
     if static:
+        # ... null converter
         _converter = lambda x: x
     else:
-        # This converter is necessary when a parameter is typed as,
+        # ... this converter is necessary when a parameter is typed as,
         # for example, Optional[Real_].
         _converter = (
             lambda x: jnp.asarray(x) if isinstance(x, ArrayLike) else x
         )
+    # ... set the converter to the passed converter or the default convertere
     converter = kwargs.pop("converter", _converter)
-    # Get serialization metadata
-    if init:
-        if encode is False:
-            encoder = config(decoder=_dummy_decoder, encoder=_dummy_encoder)
-        elif encode == Array:
-            encoder = config(encoder=_np_encoder, decoder=_jax_decoder)
-        elif encode == np.ndarray:
-            encoder = config(encoder=_np_encoder, decoder=_np_decoder)
-        else:
-            encoder = config(encoder=_object_encoder, decoder=_object_decoder)
-    else:
-        encoder = config(decoder=_dummy_decoder, encoder=_dummy_encoder)
-    # Update metadata for serialization
-    metadata["encode"] = encode
-    metadata.update(encoder)
+    # Cryojax metadata
+    metadata = kwargs.pop("metadata", {})
+    # ... add the stack keyword, if the field is traced
+    if not static:
+        metadata["stack"] = stack
 
     return eqx.field(
         converter=converter,
         static=static,
         metadata=metadata,
-        init=init,
         **kwargs,
     )
 
 
-class _Serializable(DataClassJsonMixin):
-    """
-    Base class for serializable ``cryojax`` dataclasses.
-
-    This class implements serialization functionality for cryojax
-    objects. This subclasses DataClassJsonMixin from dataclasses-json
-    and provides custom encoding/decoding for Arrays and cryojax
-    objects.
-
-    See dataclasses-json for more information. The ``load``, ``loads``,
-    ``dump``, and ``dumps`` methods are just for convenience and mimic
-    their ``from_json`` and ``to_json`` methods. The dataclasses-json API
-    is not stable, so hopefully these methods can add some stability.
-    """
-
-    @classmethod
-    def load(cls, filename: str, **kwargs: Any) -> _Serializable:
-        """
-        Load a ``cryojax`` object from a file.
-        """
-        with open(filename, "r", encoding="utf-8") as f:
-            s = f.read()
-        return cls.from_json(s, **kwargs)
-
-    def dump(self, filename: str, **kwargs: Any):
-        """
-        Dump a ``cryojax`` object to a file.
-        """
-        with open(filename, "w", encoding="utf-8") as f:
-            f.write(self.to_json(**kwargs))
-
-    @classmethod
-    def loads(cls, s: JsonData, **kwargs: Any) -> _Serializable:
-        """
-        Load a ``cryojax`` object from a json string.
-        """
-        return cls.from_json(s, **kwargs)
-
-    def dumps(self, **kwargs: Any) -> JsonData:
-        """
-        Dump a ``cryojax`` object to a json string.
-        """
-        return self.to_json(**kwargs)
-
-
-class Module(eqx.Module, _Serializable):
+class Module(eqx.Module):
     """
     Base class for ``cryojax`` objects.
     """
 
 
-class Buffer(Module):
+class BufferModule(Module):
     """
     A Module composed of buffers (do not take gradients).
     """
@@ -137,116 +82,3 @@ def __getattribute__(self, __name: str) -> Any:
             return jax.lax.stop_gradient(value)
         else:
             return value
-
-
-#
-# Encoders
-#
-def _dummy_encoder(x: Any) -> str:
-    """Encode nothing"""
-    pass
-
-
-def _dummy_decoder(x: Any) -> Any:
-    """Decode nothing"""
-    pass
-
-
-def _np_encoder(x: Any) -> Any:
-    """Numpy array encoder"""
-    if isinstance(x, (np.ndarray, Array)):
-        x = np.asarray(x)
-        data_b64 = base64.b64encode(np.ascontiguousarray(x).data)
-        return dict(
-            __ndarray__=data_b64.decode(),
-            dtype=str(x.dtype),
-            shape=x.shape,
-        )
-    else:
-        return _object_encoder(x)
-
-
-def _np_decoder(x: Any) -> Any:
-    """Numpy array decoder"""
-    if isinstance(x, dict) and "__ndarray__" in x:
-        data = base64.b64decode(x["__ndarray__"])
-        return np.frombuffer(data, x["dtype"]).reshape(x["shape"])
-    else:
-        return _object_decoder(x)
-
-
-def _jax_decoder(x: Any) -> Any:
-    """Jax array decoder"""
-    a = _np_decoder(x)
-    if isinstance(a, np.ndarray):
-        return jnp.asarray(a)
-    else:
-        return a
-
-
-def _object_encoder(x: Any) -> Any:
-    """Encoder for python objects, or collections of them."""
-    if isinstance(x, Module):
-        return _cryojax_encoder(x)
-    elif isinstance(x, FunctionType):
-        return _function_encoder(x)
-    elif isinstance(x, complex):
-        return _complex_encoder(x)
-    elif isinstance(x, list):
-        return [_object_encoder(xi) for xi in x]
-    else:
-        return x
-
-
-def _object_decoder(x: Any) -> Any:
-    """Decoder for python objects, or collections of them."""
-    if isinstance(x, dict) and "__class__" in x:
-        return _cryojax_decoder(x)
-    if isinstance(x, dict) and "__code__" in x:
-        return _function_decoder(x)
-    elif isinstance(x, dict) and "__complex__" in x:
-        return _complex_decoder(x)
-    elif isinstance(x, list):
-        return [_object_decoder(xi) for xi in x]
-    else:
-        return x
-
-
-def _cryojax_encoder(x: Module) -> dict[str, Union[str, dict]]:
-    """Encode a Module"""
-    return dict(__class__=x.__class__.__name__, __dict__=x.to_dict())
-
-
-def _cryojax_decoder(x: dict[str, Union[str, dict]]) -> Module:
-    """Decode a Module"""
-    from . import simulator
-
-    cls = getattr(simulator, x["__class__"])
-    return cls.from_dict(x["__dict__"])
-
-
-def _function_encoder(x: FunctionType) -> dict[str, str]:
-    """Encode a FunctionType"""
-    return dict(
-        __name__=x.__name__,
-        __code__=marshal.dumps(x.__code__).decode("raw_unicode_escape"),
-    )
-
-
-def _function_decoder(x: dict[str, str]) -> FunctionType:
-    """Decode a FunctionType"""
-    return FunctionType(
-        marshal.loads(bytes(x["__code__"], "raw_unicode_escape")),
-        globals(),
-        x["__name__"],
-    )
-
-
-def _complex_encoder(x: complex) -> dict[str, Union[str, float]]:
-    """Encode a complex number"""
-    return dict(__complex__=True, real=x.real, imag=x.imag)
-
-
-def _complex_decoder(x: dict[str, Union[str, float]]) -> complex:
-    """Decode a complex number"""
-    return complex(x["real"], x["imag"])
diff --git a/src/cryojax/io/__init__.py b/src/cryojax/io/__init__.py
index a9e68b04..cfb94b10 100644
--- a/src/cryojax/io/__init__.py
+++ b/src/cryojax/io/__init__.py
@@ -1,6 +1,8 @@
-from . import load_voxels, load_atoms
+from . import load_atoms, _mrc, _pdb
 
-from .load_voxels import *
 from .load_atoms import *
+from ._mrc import *
+from ._pdb import *
+from ._cif import *
 
-__all__ = load_voxels.__all__ + load_atoms.__all__
+__all__ = load_atoms.__all__ + _mrc.__all__ + _pdb.__all__ + _cif.__all__
diff --git a/src/cryojax/io/_cif.py b/src/cryojax/io/_cif.py
new file mode 100644
index 00000000..bcc641e0
--- /dev/null
+++ b/src/cryojax/io/_cif.py
@@ -0,0 +1,42 @@
+from .load_atoms import clean_gemmi_structure
+
+
+__all__ = [
+    "read_atomic_model_from_cif",
+]
+
+
+def read_atomic_model_from_cif(path, i_model=0, clean=True, assemble=True):
+    """Read Gemmi Model from CIF file.
+
+    Parameters
+    ----------
+    path : string
+        Path to mmCIF file.
+    i_model : integer
+        Optional, default: 0
+        Index of the returned model in the Gemmi Structure.
+    clean : bool
+        Optional, default: True
+        If True, use Gemmi remove_* methods to clean up structure.
+    assemble: bool
+        Optional, default: True
+        If True, use Gemmi make_assembly to build biological object.
+
+    Returns
+    -------
+    model : Gemmi Class
+        Gemmi model
+    """
+    import gemmi
+
+    cif_block = gemmi.cif.read(path)[0]
+    structure = gemmi.make_structure_from_block(cif_block)
+    if clean:
+        structure = clean_gemmi_structure(structure)
+    model = structure[i_model]
+    if assemble:
+        assembly = structure.assemblies[i_model]
+        chain_naming = gemmi.HowToNameCopiedChain.AddNumber
+        model = gemmi.make_assembly(assembly, model, chain_naming)
+    return model
diff --git a/src/cryojax/io/_mrc.py b/src/cryojax/io/_mrc.py
new file mode 100644
index 00000000..8139d5f1
--- /dev/null
+++ b/src/cryojax/io/_mrc.py
@@ -0,0 +1,52 @@
+"""
+Routines for working with MRC files.
+"""
+
+__all__ = ["load_mrc"]
+
+import numpy as np
+import mrcfile
+
+
+def load_mrc(filename: str) -> tuple[np.ndarray, float]:
+    """
+    Read MRC data to ``numpy`` array.
+
+    Parameters
+    ----------
+    filename : `str`
+        Path to data.
+
+    Returns
+    -------
+    data : `ArrayLike`, shape `(N1, N2, N3)` or `(N1, N2)`
+        Model in cartesian coordinates.
+    voxel_size : `ArrayLike`, shape `(3,)` or `(2,)`
+        The voxel_size in each dimension, stored
+        in the MRC file.
+    """
+    # Read MRC
+    with mrcfile.open(filename) as mrc:
+        data = np.asarray(mrc.data, dtype=float)
+        if data.ndim == 2:
+            voxel_size = np.asarray(
+                [mrc.voxel_size.x, mrc.voxel_size.y], dtype=float
+            )
+        elif data.ndim == 3:
+            voxel_size = np.asarray(
+                [mrc.voxel_size.x, mrc.voxel_size.y, mrc.voxel_size.z],
+                dtype=float,
+            )
+        else:
+            raise NotImplementedError(
+                "MRC files with 2D and 3D data are supported."
+            )
+
+    assert all(
+        voxel_size != np.zeros(data.ndim)
+    ), "MRC file must set the voxel size."
+    assert all(
+        voxel_size == voxel_size[0]
+    ), "Voxel size must be same in all dimensions."
+
+    return data, voxel_size[0]
diff --git a/src/cryojax/io/_pdb.py b/src/cryojax/io/_pdb.py
new file mode 100644
index 00000000..d6e5f664
--- /dev/null
+++ b/src/cryojax/io/_pdb.py
@@ -0,0 +1,50 @@
+"""
+Read and write atomic models in various formats.
+Large amounts of the code are adapted from the ioSPI package
+"""
+
+__all__ = [
+    "read_atomic_model_from_pdb",
+]
+
+from .load_atoms import clean_gemmi_structure
+
+
+def read_atomic_model_from_pdb(path, i_model=0, clean=True, assemble=True):
+    """Read Gemmi Model from PDB file.
+
+    Parameters
+    ----------
+    path : string
+        Path to PDB file.
+    i_model : integer
+        Optional, default: 0
+        Index of the returned model in the Gemmi Structure.
+    clean : bool
+        Optional, default: True
+        If True, use Gemmi remove_* methods to clean up structure.
+    assemble: bool
+        Optional, default: True
+        If True, use Gemmi make_assembly to build biological object.
+
+    Returns
+    -------
+    model : Gemmi Class
+        Gemmi model
+
+    Notes
+    -----
+    Currently Hydrogen atoms are not read in!
+    We should look into adding hydrogens, potentially.
+    """
+    import gemmi
+
+    structure = gemmi.read_structure(path)
+    if clean:
+        structure = clean_gemmi_structure(structure)
+    model = structure[i_model]
+    if assemble:
+        assembly = structure.assemblies[i_model]
+        chain_naming = gemmi.HowToNameCopiedChain.AddNumber
+        model = gemmi.make_assembly(assembly, model, chain_naming)
+    return model
diff --git a/src/cryojax/io/load_atoms.py b/src/cryojax/io/load_atoms.py
index 58f2455d..89b75277 100644
--- a/src/cryojax/io/load_atoms.py
+++ b/src/cryojax/io/load_atoms.py
@@ -1,5 +1,123 @@
 """
 Routines for loading atomic structures.
+Large amounts of the code are adapted from the ioSPI package
 """
 
-__all__ = []
+__all__ = [
+    "clean_gemmi_structure",
+    "extract_gemmi_atoms",
+    "extract_atomic_parameter",
+    "get_scattering_info_from_gemmi_model",
+]
+
+import itertools
+import numpy as np
+
+
+def get_scattering_info_from_gemmi_model(model):
+    atoms = extract_gemmi_atoms(model)
+    coords = extract_atomic_parameter(atoms, "cartesian_coordinates")
+    a_vals = extract_atomic_parameter(atoms, "electron_form_factor_a")
+    b_vals = extract_atomic_parameter(atoms, "electron_form_factor_b")
+    return coords, a_vals, b_vals
+
+
+def clean_gemmi_structure(structure=None):
+    """Clean Gemmi Structure.
+
+    Parameters
+    ----------
+    structure : Gemmi Class
+        Gemmi Structure object
+
+    Returns
+    -------
+    structure : Gemmi Class
+        Same object, cleaned up of unnecessary atoms.
+
+    """
+    if structure is not None:
+        structure.remove_alternative_conformations()
+        structure.remove_hydrogens()
+        structure.remove_waters()
+        structure.remove_ligands_and_waters()
+        structure.remove_empty_chains()
+
+    return structure
+
+
+def extract_gemmi_atoms(model, chains=None, split_chains=False):
+    """
+    Extract Gemmi atoms from the input Gemmi model, separated by chain.
+
+    Parameters
+    ----------
+    model : Gemmi Class
+        Gemmi model
+    chains : list of strings
+        chains to select, optional.
+        If not provided, retrieve atoms from all chains.
+    split_chains : bool
+        Optional, default: False
+        if True, keep the atoms from different chains in separate lists
+
+    Returns
+    -------
+    atoms : list (or list of list(s)) of Gemmi atoms
+        Gemmi atom objects, either concatenated or separated by chain
+    """
+    if chains is None:
+        chains = [ch.name for ch in model]
+
+    atoms = []
+    for ch in model:
+        if ch.name in chains:
+            atoms.append([at for res in ch for at in res])
+
+    if not split_chains:
+        atoms = list(itertools.chain.from_iterable(atoms))
+
+    return atoms
+
+
+def extract_atomic_parameter(atoms, parameter_type, split_chains=False):
+    """
+    Interpret Gemmi atoms and extract a single parameter type.
+
+    Parameters
+    ----------
+    atoms : list (of list(s)) of Gemmi atoms
+        Gemmi atom objects associated with each chain
+    parameter_type : string
+        'cartesian_coordinates', 'form_factor_a', or 'form_factor_b'
+    split_chains : bool
+        Optional, default: False
+        if True, keep the atoms from different chains in separate lists
+
+    Returns
+    -------
+    atomic_parameter : list of floats, or list of lists of floats
+        atomic parameter associated with each atom, optionally split by chain
+    """
+    # if list of Gemmi atoms, convert into a list of list
+    if type(atoms[0]) != list:
+        atoms = [atoms]
+
+    if parameter_type == "cartesian_coordinates":
+        atomic_parameter = [at.pos.tolist() for ch in atoms for at in ch]
+    elif parameter_type == "electron_form_factor_a":
+        atomic_parameter = [at.element.c4322.a for ch in atoms for at in ch]
+    elif parameter_type == "electron_form_factor_b":
+        atomic_parameter = [at.element.c4322.b for ch in atoms for at in ch]
+    else:
+        raise ValueError("Atomic parameter type not recognized.")
+
+    # optionally preserve the list of lists (separated by chain) structure
+    if split_chains:
+        reshape = [0] + [len(ch) for ch in atoms]
+        atomic_parameter = [
+            atomic_parameter[reshape[i] : reshape[i] + reshape[i + 1]]
+            for i in range(len(reshape) - 1)
+        ]
+
+    return np.array(atomic_parameter)
diff --git a/src/cryojax/io/load_voxels.py b/src/cryojax/io/load_voxels.py
deleted file mode 100644
index ed26eb8a..00000000
--- a/src/cryojax/io/load_voxels.py
+++ /dev/null
@@ -1,138 +0,0 @@
-"""
-Routines for loading voxel-based electron densities.
-"""
-
-from __future__ import annotations
-
-__all__ = ["load_mrc", "load_voxel_cloud", "load_fourier_grid"]
-
-import mrcfile, os
-import numpy as np
-import jax.numpy as jnp
-from typing import Any
-from ..utils import fftn, pad, make_frequencies, flatten_and_coordinatize
-
-
-def load_voxel_cloud(filename: str, **kwargs: Any) -> dict[str, Any]:
-    """
-    Read a 3D template on a cartesian grid
-    to a point cloud.
-
-    This is used to instantiate ``cryojax.simulator.VoxelCloud``.
-
-    Parameters
-    ----------
-    filename : `str`
-        Path to template.
-    kwargs :
-        Keyword arguments passed to
-        ``cryojax.utils.coordinates.flatten_and_coordinatize``.
-
-    Returns
-    -------
-    cloud : `dict`
-        Electron density in a point cloud representation,
-        generated from a 3D voxel template. By default,
-        voxels with zero density are masked.
-    """
-    # Load template
-    filename = os.path.abspath(filename)
-    template, voxel_size = load_mrc(filename)
-    # Change how template sits in the box to match cisTEM
-    # Ideally we would have this read in from the MRC and be the
-    # same for all I/O methods. However, the algorithms used all
-    # have their own xyz conventions. This is made to match
-    # jax-finufft.
-    template = jnp.transpose(template, axes=[1, 2, 0])
-    # Load flattened density and coordinates
-    density, coordinates = flatten_and_coordinatize(
-        template, voxel_size, **kwargs
-    )
-    # Gather fields to instantiate an ElectronCloud
-    cloud = dict(
-        weights=density, coordinates=coordinates, voxel_size=voxel_size
-    )
-
-    return cloud
-
-
-def load_fourier_grid(filename: str, pad_scale: float = 1.0) -> dict[str, Any]:
-    """
-    Read a 3D template in Fourier space on a cartesian grid.
-
-    This is used to instantiate ``cryojax.simulator.VoxelGrid``.
-
-    Parameters
-    ----------
-    filename : `str`
-        Path to template.
-
-    Returns
-    -------
-    voxels : `dict`
-        3D electron density in a 3D voxel grid representation.
-        Instantiates a ``cryojax.simulator.ElectronGrid``
-    """
-    # Load template
-    filename = os.path.abspath(filename)
-    template, voxel_size = load_mrc(filename)
-    # Change how template sits in box to match cisTEM
-    template = jnp.transpose(template, axes=[2, 1, 0])
-    # Pad template
-    padded_shape = tuple([int(s * pad_scale) for s in template.shape])
-    template = pad(template, padded_shape)
-    # Load density and coordinates
-    density = fftn(template)
-    coordinates = make_frequencies(template.shape, voxel_size)
-    # Get central z slice
-    coordinates = jnp.expand_dims(coordinates[:, :, 0, :], axis=2)
-    # Gather fields to instantiate an ElectronGrid
-    voxels = dict(
-        weights=density, coordinates=coordinates, voxel_size=voxel_size
-    )
-
-    return voxels
-
-
-def load_mrc(filename: str) -> tuple[np.ndarray, float]:
-    """
-    Read MRC data to ``numpy`` array.
-
-    Parameters
-    ----------
-    filename : `str`
-        Path to data.
-
-    Returns
-    -------
-    data : `ArrayLike`, shape `(N1, N2, N3)` or `(N1, N2)`
-        Model in cartesian coordinates.
-    voxel_size : `ArrayLike`, shape `(3,)` or `(2,)`
-        The voxel_size in each dimension, stored
-        in the MRC file.
-    """
-    # Read MRC
-    with mrcfile.open(filename) as mrc:
-        data = np.asarray(mrc.data, dtype=float)
-        if data.ndim == 2:
-            voxel_size = np.asarray(
-                [mrc.voxel_size.x, mrc.voxel_size.y], dtype=float
-            )
-        elif data.ndim == 3:
-            voxel_size = np.asarray(
-                [mrc.voxel_size.x, mrc.voxel_size.y, mrc.voxel_size.z],
-                dtype=float,
-            )
-        else:
-            raise NotImplementedError(
-                "MRC files with 2D and 3D data are supported."
-            )
-
-    assert all(
-        voxel_size != np.zeros(data.ndim)
-    ), "MRC file must set the voxel size."
-    assert all(
-        voxel_size == voxel_size[0]
-    ), "Voxel size must be same in all dimensions."
-
-    return data, float(voxel_size[0])
diff --git a/src/cryojax/simulator/__init__.py b/src/cryojax/simulator/__init__.py
index fefb2bd1..f6f3409d 100644
--- a/src/cryojax/simulator/__init__.py
+++ b/src/cryojax/simulator/__init__.py
@@ -8,12 +8,13 @@
     scattering,
     filter,
     mask,
+    noise,
     ice,
     optics,
     exposure,
     detector,
     instrument,
-    image,
+    pipeline,
     distribution,
 )
 
@@ -26,12 +27,13 @@
 from .scattering import *
 from .filter import *
 from .mask import *
+from .noise import *
 from .ice import *
 from .optics import *
 from .exposure import *
 from .detector import *
 from .instrument import *
-from .image import *
+from .pipeline import *
 from .distribution import *
 
 
@@ -45,12 +47,13 @@
     + scattering.__all__
     + filter.__all__
     + mask.__all__
+    + noise.__all__
     + ice.__all__
     + optics.__all__
     + exposure.__all__
     + detector.__all__
     + instrument.__all__
-    + image.__all__
+    + pipeline.__all__
     + distribution.__all__
 )
 
@@ -65,12 +68,13 @@
         scattering,
         filter,
         mask,
+        noise,
         ice,
         optics,
         exposure,
         detector,
         instrument,
-        image,
+        pipeline,
         distribution,
     ]
 )
diff --git a/src/cryojax/simulator/assembly/_assembly.py b/src/cryojax/simulator/assembly/_assembly.py
index 73cf2af2..867c9d16 100644
--- a/src/cryojax/simulator/assembly/_assembly.py
+++ b/src/cryojax/simulator/assembly/_assembly.py
@@ -1,6 +1,6 @@
 """
 Abstraction of a biological assembly. This assembles a structure
-by adding together a collection of subunits, parameterized by
+by computing an Ensemble of subunits, parameterized by
 some geometry.
 """
 
@@ -72,11 +72,6 @@ def __init__(
         self.pose = pose or EulerPose()
         self.conformation = conformation
 
-    @property
-    def resolution(self) -> Real_:
-        """Hack to make this class act like a Specimen."""
-        return self.subunit.resolution
-
     @cached_property
     @abstractmethod
     def n_subunits(self) -> int:
diff --git a/src/cryojax/simulator/density/_atom_density.py b/src/cryojax/simulator/density/_atom_density.py
index b493adef..88a0a40d 100644
--- a/src/cryojax/simulator/density/_atom_density.py
+++ b/src/cryojax/simulator/density/_atom_density.py
@@ -6,7 +6,7 @@
 
 __all__ = ["AtomCloud"]
 
-from typing import Type, Any
+from typing import Type, Any, ClassVar
 
 import equinox as eqx
 from jaxtyping import Array
@@ -22,27 +22,23 @@ class AtomCloud(ElectronDensity):
     """
 
     weights: Array = field()
-    coordinates: Array = field()
+    coordinate_list: Array = field()
     variances: Array = field()
     identity: Array = field()
 
-    is_real: bool = field(default=True, static=True)
-
-    def __check_init__(self):
-        if self.is_real is False:
-            raise NotImplementedError(
-                "Fourier atomic densities are not supported."
-            )
+    is_real: ClassVar[bool] = True
 
     def rotate_to(self, pose: Pose) -> AtomCloud:
-        coordinates = pose.rotate(self.coordinates, is_real=self.is_real)
-        return eqx.tree_at(lambda d: d.coordinates, self, coordinates)
+        return eqx.tree_at(
+            lambda d: d.coordinate_list,
+            self,
+            pose.rotate(self.coordinate_list, is_real=self.is_real),
+        )
 
     @classmethod
     def from_file(
         cls: Type[AtomCloud],
         filename: str,
-        config: dict = {},
         **kwargs: Any,
     ) -> AtomCloud:
         """
@@ -50,21 +46,7 @@ def from_file(
 
         TODO: What is the file format appropriate here? Q. for Michael...
         """
-        raise NotImplementedError
-        # return cls.from_mrc(filename, config=config, **kwargs)
+        import gemmi
 
-    @classmethod
-    def from_stack(
-        cls: Type["AtomCloud"], stack: list["AtomCloud"]
-    ) -> "AtomCloud":
-        """
-        Stack a list of electron densities along the leading
-        axis of a single electron density.
-        """
-        raise NotImplementedError
-
-    def __getitem__(self, key: int) -> "AtomCloud":
-        raise NotImplementedError
-
-    def __len__(self) -> int:
         raise NotImplementedError
+        # return cls.from_mrc(filename, config=config, **kwargs)
diff --git a/src/cryojax/simulator/density/_electron_density.py b/src/cryojax/simulator/density/_electron_density.py
index 3f03954e..0390a2b1 100644
--- a/src/cryojax/simulator/density/_electron_density.py
+++ b/src/cryojax/simulator/density/_electron_density.py
@@ -2,17 +2,23 @@
 Base electron density representation.
 """
 
-__all__ = ["ElectronDensity"]
+__all__ = ["ElectronDensity", "ElectronDensityType"]
 
+import dataclasses
 from abc import abstractmethod
-from typing import Optional, Type
+from typing import Type, Any, TypeVar
+from typing_extensions import Self
+from equinox import AbstractClassVar
 
-from equinox import AbstractVar
+import jax.numpy as jnp
 
 from ..pose import Pose
 from ...core import Module, field
 
 
+ElectronDensityType = TypeVar("ElectronDensityType", bound="ElectronDensity")
+
+
 class ElectronDensity(Module):
     """
     Abstraction of an electron density map.
@@ -24,11 +30,11 @@ class ElectronDensity(Module):
         real or fourier space.
     """
 
-    is_real: AbstractVar[bool]
+    is_real: AbstractClassVar[bool]
     _is_stacked: bool = field(static=True, default=False, kw_only=True)
 
     @abstractmethod
-    def rotate_to(self, pose: Pose) -> "ElectronDensity":
+    def rotate_to_pose(self, pose: Pose) -> Self:
         """
         View the electron density at a given pose.
 
@@ -42,32 +48,81 @@ def rotate_to(self, pose: Pose) -> "ElectronDensity":
     @classmethod
     @abstractmethod
     def from_file(
-        cls: Type["ElectronDensity"],
+        cls: Type[ElectronDensityType],
         filename: str,
-        config: Optional[dict] = None,
-    ) -> "ElectronDensity":
+        **kwargs: Any,
+    ) -> ElectronDensityType:
         """
         Load an ElectronDensity from a file.
         """
         raise NotImplementedError
 
     @classmethod
-    @abstractmethod
     def from_stack(
-        cls: Type["ElectronDensity"], stack: list["ElectronDensity"]
-    ) -> "ElectronDensity":
+        cls: Type[ElectronDensityType], stack: list[ElectronDensityType]
+    ) -> ElectronDensityType:
         """
         Stack a list of electron densities along the leading
         axis of a single electron density.
         """
-        raise NotImplementedError
+        if not all([cls == type(density) for density in stack]):
+            raise TypeError(
+                "Electron density stack should all be of the same type."
+            )
+        # Gather static and traced fields separately
+        other, stacked = {}, {}
+        for field in dataclasses.fields(stack[0]):
+            name = field.name
+            if name == "_is_stacked":
+                pass
+            elif ("static" in field.metadata and field.metadata["static"]) or (
+                "stack" in field.metadata and not field.metadata["stack"]
+            ):
+                # Static or unstacked fields should all match, so take the first.
+                other[name] = getattr(stack[0], name)
+            else:
+                # Traced fields, unless specified in metadata, get stacked.
+                stacked[name] = jnp.stack(
+                    [getattr(density, name) for density in stack], axis=0
+                )
+        return cls(**stacked, **other, _is_stacked=True)
 
-    @abstractmethod
-    def __getitem__(self, idx: int) -> "ElectronDensity":
-        """Get a particular electron density in the stack."""
-        raise NotImplementedError
+    def __getitem__(self, idx: int) -> Self:
+        if self._is_stacked:
+            # Gather static and traced fields separately
+            indexed = {}
+            for field in dataclasses.fields(self):
+                name = field.name
+                if name == "_is_stacked":
+                    pass
+                elif (
+                    "static" in field.metadata and field.metadata["static"]
+                ) or (
+                    "stack" in field.metadata and not field.metadata["stack"]
+                ):
+                    pass
+                else:
+                    # Get stacked fields at particular index
+                    indexed[name] = getattr(self, name)[idx]
+            return dataclasses.replace(self, **indexed, _is_stacked=False)
+        else:
+            raise IndexError("Cannot index an non-stacked ElectronDensity.")
 
-    @abstractmethod
     def __len__(self) -> int:
-        """Get the number of electron densities in the stack."""
-        raise NotImplementedError
+        if self._is_stacked:
+            for field in dataclasses.fields(self):
+                if not (
+                    ("static" in field.metadata and field.metadata["static"])
+                    or (
+                        "stack" in field.metadata
+                        and not field.metadata["stack"]
+                    )
+                ):
+                    return getattr(self, field.name).shape[0]
+            raise AttributeError(
+                "Could not get the length of the ElectronDensity stack."
+            )
+        else:
+            raise TypeError(
+                "Cannot get length of non-stacked ElectronDensity."
+            )
diff --git a/src/cryojax/simulator/density/_voxel_density.py b/src/cryojax/simulator/density/_voxel_density.py
index 897e007f..bba6fefa 100644
--- a/src/cryojax/simulator/density/_voxel_density.py
+++ b/src/cryojax/simulator/density/_voxel_density.py
@@ -4,23 +4,46 @@
 
 __all__ = ["Voxels", "VoxelCloud", "VoxelGrid"]
 
-import equinox as eqx
-
+import pathlib
 from abc import abstractmethod
-from typing import Any, Type
-from jaxtyping import Float, Array
+from typing import Any, Tuple, Type, ClassVar, TypeVar, Optional, overload
+from typing_extensions import Self
+from jaxtyping import Complex, Float, Array
 from equinox import AbstractVar
+from functools import cached_property
 
+import equinox as eqx
+import jax
 import jax.numpy as jnp
 
 from ._electron_density import ElectronDensity
 from ..pose import Pose
-from ...io import load_voxel_cloud, load_fourier_grid
+from ...io import (
+    load_mrc,
+    read_atomic_model_from_pdb,
+    read_atomic_model_from_cif,
+    get_scattering_info_from_gemmi_model,
+)
 from ...core import field
-from ...typing import Real_, RealCloud, CloudCoords3D
-
-_CubicVolume = Float[Array, "N N N"]
-_VolumeSliceCoords = Float[Array, "N N 1 3"]
+from cryojax.utils import (
+    make_frequencies,
+    make_coordinates,
+    pad,
+    fftn,
+)
+from cryojax.typing import (
+    RealCloud,
+    RealVolume,
+    VolumeCoords,
+    CloudCoords3D,
+    Real_,
+)
+
+_RealCubicVolume = Float[Array, "N N N"]
+_ComplexCubicVolume = Complex[Array, "N N N"]
+_VolumeSliceCoords = Float[Array, "N N//2+1 1 3"]
+
+VoxelType = TypeVar("VoxelType", bound="Voxels")
 
 
 class Voxels(ElectronDensity):
@@ -31,136 +54,212 @@ class Voxels(ElectronDensity):
     ----------
     weights :
         The electron density.
-    coordinates :
-        The coordinate system.
     voxel_size
         The voxel size of the electron density.
     """
 
     weights: AbstractVar[Array]
-    coordinates: AbstractVar[Array]
-    voxel_size: Real_ = field()
+    voxel_size: Real_ = field(stack=False)
 
+    @overload
     @classmethod
-    def from_file(
-        cls: Type["Voxels"],
-        filename: str,
-        config: dict = {},
+    @abstractmethod
+    def from_density_grid(
+        cls: Type[VoxelType],
+        density_grid: RealVolume,
+        voxel_size: float,
+        coordinate_grid: None,
         **kwargs: Any,
-    ) -> "Voxels":
-        """Load a ElectronDensity."""
-        return cls.from_mrc(
-            filename, config=config, _is_stacked=False, **kwargs
-        )
+    ) -> VoxelType:
+        ...
 
+    @overload
     @classmethod
     @abstractmethod
-    def from_mrc(
-        cls: Type["Voxels"],
-        filename: str,
-        config: dict = {},
+    def from_density_grid(
+        cls: Type[VoxelType],
+        density_grid: RealVolume,
+        voxel_size: float,
+        coordinate_grid: VolumeCoords,
         **kwargs: Any,
-    ) -> "Voxels":
-        """Load a ElectronDensity from MRC file format."""
+    ) -> VoxelType:
+        ...
+
+    @classmethod
+    @abstractmethod
+    def from_density_grid(
+        cls: Type[VoxelType],
+        density_grid: RealVolume,
+        voxel_size: float = 1.0,
+        coordinate_grid: Optional[VolumeCoords] = None,
+        **kwargs: Any,
+    ) -> VoxelType:
+        """
+        Load a Voxels object from real-valued 3D electron
+        density map.
+        """
         raise NotImplementedError
 
     @classmethod
-    def from_stack(cls: Type["Voxels"], stack: list["Voxels"]) -> "Voxels":
+    def from_gemmi(
+        cls: Type[VoxelType],
+        model,
+        n_voxels_per_side: Tuple[int, int, int],
+        voxel_size: float = 1.0,
+        **kwargs: Any,
+    ) -> VoxelType:
         """
-        Stack a list of electron densities along the leading
-        axis of a single electron density.
+        Loads a PDB file as a Voxels subclass.  Uses the Gemmi library.
+        Heavily based on a code from Frederic Poitevin, located at
+
+        https://github.com/compSPI/ioSPI/blob/master/ioSPI/atomic_models.py
         """
-        if not all([cls == type(density) for density in stack]):
-            raise TypeError(
-                "Electron density stack should all be of the same type."
-            )
-        if not all([stack[0].is_real == density.is_real for density in stack]):
-            raise TypeError(
-                "Electron density stack should all be in real or fourier space."
-            )
-        weights = jnp.stack([density.weights for density in stack], axis=0)
-        coordinates = jnp.stack(
-            [density.coordinates for density in stack], axis=0
+        coords, a_vals, b_vals = get_scattering_info_from_gemmi_model(model)
+
+        coordinate_grid_in_angstroms = make_coordinates(
+            n_voxels_per_side, voxel_size
         )
-        voxel_size = jnp.stack(
-            [density.voxel_size for density in stack], axis=0
+        density = _build_real_space_voxels_from_atoms(
+            coords, a_vals, b_vals, coordinate_grid_in_angstroms
         )
-        return cls(
-            weights=weights,
-            coordinates=coordinates,
-            voxel_size=voxel_size,
-            is_real=stack[0].is_real,
-            _is_stacked=True,
+
+        return cls.from_density_grid(
+            density,
+            voxel_size,
+            coordinate_grid_in_angstroms / voxel_size,
+            **kwargs,
         )
 
-    def __getitem__(self, idx: int) -> "Voxels":
-        if self._is_stacked:
-            cls = type(self)
-            return cls(
-                weights=self.weights[idx],
-                coordinates=self.coordinates[idx],
-                voxel_size=self.voxel_size[idx],
-                is_real=self.is_real,
-                _is_stacked=False,
-            )
+    @classmethod
+    def from_file(
+        cls: Type[VoxelType],
+        filename: str,
+        *args: Any,
+        **kwargs: Any,
+    ) -> VoxelType:
+        """Load a voxel-based electron density."""
+        path = pathlib.Path(filename)
+        if path.suffix == ".mrc":
+            return cls.from_mrc(filename, *args, **kwargs)
+        elif path.suffix == ".pdb":
+            return cls.from_pdb(filename, *args, **kwargs)
+        elif path.suffix == ".cif":
+            return cls.from_cif(filename, *args, **kwargs)
         else:
-            return self
+            raise NotImplementedError(
+                f"File format {path.suffix} not supported."
+            )
 
-    def __len__(self) -> int:
-        if self._is_stacked:
-            return self.weights.shape[0]
-        else:
-            return 1
+    @classmethod
+    def from_mrc(
+        cls: Type[VoxelType],
+        filename: str,
+        **kwargs: Any,
+    ) -> VoxelType:
+        """Load Voxels from MRC file format."""
+        density_grid, voxel_size = load_mrc(filename)
+        return cls.from_density_grid(density_grid, voxel_size, **kwargs)
+
+    @classmethod
+    def from_pdb(
+        cls: Type[VoxelType],
+        filename: str,
+        n_voxels_per_side: Tuple[int, int, int],
+        voxel_size: float = 1.0,
+        **kwargs: Any,
+    ) -> VoxelType:
+        """Load Voxels from PDB file format."""
+        model = read_atomic_model_from_pdb(filename)
+        return cls.from_gemmi(model, n_voxels_per_side, voxel_size, **kwargs)
+
+    @classmethod
+    def from_cif(
+        cls: Type[VoxelType],
+        filename: str,
+        n_voxels_per_side: Tuple[int, int, int],
+        voxel_size: float = 1.0,
+        **kwargs: Any,
+    ) -> VoxelType:
+        """Load Voxels from CIF file format."""
+        model = read_atomic_model_from_cif(filename)
+        return cls.from_gemmi(model, n_voxels_per_side, voxel_size, **kwargs)
 
 
 class VoxelGrid(Voxels):
     """
     Abstraction of a 3D electron density voxel grid.
 
-    The voxel grid should be given in Fourier space.
+    The voxel grid should be given in fourier space.
 
     Attributes
     ----------
     weights :
-        3D electron density grid in Fourier space.
-    coordinates :
-        Central slice of cartesian coordinate system.
+        3D electron density grid in fourier space.
+    frequency_slice :
+        Central slice of cartesian coordinate system
+        in fourier space.
     """
 
-    weights: _CubicVolume = field()
-    coordinates: _VolumeSliceCoords = field()
+    weights: _ComplexCubicVolume = field()
+    frequency_slice: _VolumeSliceCoords = field(stack=False)
 
-    is_real: bool = field(default=False, static=True, kw_only=True)
+    is_real: ClassVar[bool] = False
 
-    def __check_init__(self):
-        if self.is_real is True:
-            raise NotImplementedError(
-                "Real voxel grid densities are not supported."
-            )
+    @cached_property
+    def frequency_slice_in_angstroms(self) -> _VolumeSliceCoords:
+        return self.frequency_slice / self.voxel_size
 
-    def rotate_to(self, pose: Pose) -> "VoxelGrid":
+    def rotate_to_pose(self, pose: Pose) -> Self:
         """
         Compute rotations of a central slice in fourier space
         by an imaging pose.
 
         This rotation is the inverse rotation as in real space.
         """
-        coordinates = pose.rotate(self.coordinates, is_real=self.is_real)
-
-        return eqx.tree_at(lambda d: d.coordinates, self, coordinates)
+        return eqx.tree_at(
+            lambda d: d.frequency_slice,
+            self,
+            pose.rotate(self.frequency_slice, is_real=self.is_real),
+        )
 
     @classmethod
-    def from_mrc(
+    def from_density_grid(
         cls: Type["VoxelGrid"],
-        filename: str,
-        config: dict = {},
+        density_grid: RealVolume,
+        voxel_size: float = 1.0,
+        coordinate_grid: Optional[VolumeCoords] = None,
+        pad_scale: float = 1.0,
         **kwargs: Any,
     ) -> "VoxelGrid":
-        """
-        See ``cryojax.io.voxel.load_fourier_grid`` for
-        documentation.
-        """
-        return cls(**load_fourier_grid(filename, **config), **kwargs)
+        # Change how template sits in box to match cisTEM
+        density_grid = jnp.transpose(density_grid, axes=[2, 1, 0])
+        # Pad template
+        padded_shape = tuple([int(s * pad_scale) for s in density_grid.shape])
+        padded_density_grid = pad(density_grid, padded_shape)
+        # Load density and coordinates. For now, do not store the
+        # fourier density only on the half space. Fourier slice extraction
+        # does not currently work if rfftn is used
+        fourier_density_grid = fftn(padded_density_grid)
+        # ... create in-plane frequency slice on the half space
+        frequency_slice = make_frequencies(
+            padded_density_grid.shape[:-1], half_space=True
+        )
+        # ... zero pad to make the slice 3-dimensional
+        frequency_slice = jnp.expand_dims(
+            jnp.pad(
+                frequency_slice,
+                ((0, 0), (0, 0), (0, 1)),
+                mode="constant",
+                constant_values=0.0,
+            ),
+            axis=2,
+        )
+
+        return cls(
+            weights=fourier_density_grid,
+            frequency_slice=frequency_slice,
+            voxel_size=voxel_size,
+        )
 
 
 class VoxelCloud(Voxels):
@@ -172,42 +271,178 @@ class VoxelCloud(Voxels):
     Attributes
     ----------
     weights :
-        3D electron density cloud.
-    coordinates :
-        Cartesian coordinate system for density cloud.
+        Flattened 3D electron density voxel grid into a
+        point cloud.
+    coordinate_list :
+        List of coordinates for the point cloud.
     """
 
     weights: RealCloud = field()
-    coordinates: CloudCoords3D = field()
+    coordinate_list: CloudCoords3D = field(stack=False)
 
-    is_real: bool = field(default=True, static=True, kw_only=True)
+    is_real: ClassVar[bool] = True
 
-    def __check_init__(self):
-        if self.is_real is False:
-            raise NotImplementedError(
-                "Fourier voxel cloud densities are not supported."
-            )
+    @cached_property
+    def coordinate_list_in_angstroms(self) -> CloudCoords3D:
+        return self.voxel_size * self.coordinate_list
 
-    def rotate_to(self, pose: Pose) -> "VoxelCloud":
+    def rotate_to_pose(self, pose: Pose) -> Self:
         """
         Compute rotations of a point cloud by an imaging pose.
 
         This transformation will return a new density cloud
         with rotated coordinates.
         """
-        coordinates = pose.rotate(self.coordinates, is_real=self.is_real)
-
-        return eqx.tree_at(lambda d: d.coordinates, self, coordinates)
+        return eqx.tree_at(
+            lambda d: d.coordinate_list,
+            self,
+            pose.rotate(self.coordinate_list, is_real=self.is_real),
+        )
 
     @classmethod
-    def from_mrc(
+    def from_density_grid(
         cls: Type["VoxelCloud"],
-        filename: str,
-        config: dict = {},
+        density_grid: RealVolume,
+        voxel_size: float = 1.0,
+        coordinate_grid: Optional[VolumeCoords] = None,
+        mask_zeros: bool = True,
         **kwargs: Any,
     ) -> "VoxelCloud":
-        """
-        See ``cryojax.io.voxel.load_grid_as_cloud`` for
-        documentation.
-        """
-        return cls(**load_voxel_cloud(filename, **config), **kwargs)
+        # Change how template sits in the box.
+        # Ideally we would change this in the same way for all
+        # I/O methods. However, the algorithms used all
+        # have their own xyz conventions. The choice here is to
+        # make jax-finufft output match cisTEM.
+        density_grid = jnp.transpose(density_grid, axes=[1, 2, 0])
+        # Make coordinates if not given
+        if coordinate_grid is None:
+            coordinate_grid = make_coordinates(density_grid.shape)
+        # Load flattened density and coordinates
+        if not mask_zeros:
+            flat_density = density_grid.ravel()
+            coordinate_list = coordinate_grid.ravel()
+        else:
+            # ... mask zeros if desired to store smaller arrays. This
+            # option is not jittable.
+            nonzero = jnp.where(~jnp.isclose(density_grid, 0.0, **kwargs))
+            flat_density = density_grid[nonzero]
+            coordinate_list = coordinate_grid[nonzero]
+
+        return cls(
+            weights=flat_density,
+            coordinate_list=coordinate_list,
+            voxel_size=voxel_size,
+        )
+
+
+def _eval_3d_real_space_gaussian(
+    coordinate_system: Float[Array, "N1 N2 N3 3"],
+    atom_position: Float[Array, "3"],
+    a: float,
+    b: float,
+) -> Float[Array, "N1 N2 N3"]:
+    """
+    Evaluate a gaussian on a 3D grid.
+    The naming convention for parameters follows ``Robust
+    Parameterization of Elastic and Absorptive Electron Atomic Scattering
+    Factors'' by Peng et al.
+
+    Parameters
+    ----------
+    coordinate_system : `Array`, shape `(N1, N2, N3, 3)`
+        The coordinate_system of the grid.
+    pos : `Array`, shape `(3,)`
+        The center of the gaussian.
+    a : `float`
+        A scale factor.
+    b : `float`
+        The scale of the gaussian.
+
+    Returns
+    -------
+    density : `Array`, shape `(N1, N2, N3)`
+        The density of the gaussian on the grid.
+    """
+    b_inverse = 4.0 * jnp.pi / b
+    sq_distances = jnp.sum(
+        b_inverse * (coordinate_system - atom_position) ** 2, axis=-1
+    )
+    density = jnp.exp(-jnp.pi * sq_distances) * a * b_inverse ** (3.0 / 2.0)
+    return density
+
+
+def _eval_3d_atom_potential(
+    coordinate_system: Float[Array, "N1 N2 N3 3"],
+    atom_position: Float[Array, "3"],
+    atomic_as: Float[Array, "5"],
+    atomic_bs: Float[Array, "5"],
+) -> Float[Array, "N1 N2 N3"]:
+    """
+    Evaluates the electron potential of a single atom on a 3D grid.
+
+    Parameters
+    ----------
+    coordinate_system : `Array`, shape `(N1, N2, N3, 3)`
+        The coordinate_system of the grid.
+    atom_position : `Array`, shape `(3,)`
+        The location of the atom.
+    atomic_as : `Array`, shape `(5,)`
+        The intensity values for each gaussian in the atom.
+    atomic_bs : `Array`, shape `(5,)`
+        The inverse scale factors for each gaussian in the atom.
+
+    Returns
+    -------
+    potential : `Array`, shape `(N1, N2, N3)`
+        The potential of the atom evaluate on the grid.
+    """
+    eval_fxn = jax.vmap(
+        _eval_3d_real_space_gaussian, in_axes=(None, None, 0, 0)
+    )
+    return jnp.sum(
+        eval_fxn(coordinate_system, atom_position, atomic_as, atomic_bs),
+        axis=0,
+    )
+
+
+@jax.jit
+def _build_real_space_voxels_from_atoms(
+    atom_positions: Float[Array, "N 3"],
+    ff_a: Float[Array, "N 5"],
+    ff_b: Float[Array, "N 5"],
+    coordinate_system: Float[Array, "N1 N2 N3 3"],
+) -> Tuple[_RealCubicVolume, _VolumeSliceCoords]:
+    """
+    Build a voxel representation of an atomic model.
+
+    Parameters
+    ----------
+    atom_coords : `Array`, shape `(N, 3)`
+        The coordinates of the atoms.
+    ff_a : `Array`, shape `(N, 5)` or `(N, 5, 3)`
+        Intensity values for each Gaussian in the atom
+    ff_b : `Array`, shape `(N, 5)` or `(N, 5, 3)`
+        The inverse scale factors for each Gaussian in the atom
+    coordinate_system : `Array`, shape `(N1, N2, N3, 3)`
+        The coordinates of each voxel in the grid.
+
+    Returns
+    -------
+    density :  `Array`, shape `(N1, N2, N3)`
+        The voxel representation of the atomic model.
+    z_plane_coordinates : `Array`, shape `(N1, N2, 3)`
+        The coordinates of each voxel in the z=0 plane.
+    """
+    density = jnp.zeros(coordinate_system.shape[:-1])
+
+    def add_gaussian_to_density(i, density):
+        density += _eval_3d_atom_potential(
+            coordinate_system, atom_positions[i], ff_a[i], ff_b[i]
+        )
+        return density
+
+    density = jax.lax.fori_loop(
+        0, atom_positions.shape[0], add_gaussian_to_density, density
+    )
+
+    return density
diff --git a/src/cryojax/simulator/detector.py b/src/cryojax/simulator/detector.py
index 89c5e965..d517c90b 100644
--- a/src/cryojax/simulator/detector.py
+++ b/src/cryojax/simulator/detector.py
@@ -5,72 +5,28 @@
 __all__ = [
     "Detector",
     "NullDetector",
-    "NoiselessDetector",
     "GaussianDetector",
-    "rescale_pixel_size",
 ]
 
 from abc import abstractmethod
-from typing import Optional, Any
 from typing_extensions import override
-from functools import partial
 
-import jax
 import jax.numpy as jnp
 from jaxtyping import PRNGKeyArray
 
 from .noise import GaussianNoise
-from .kernel import Kernel, Constant
-from ..utils import scale, ifftn
+from .kernel import KernelType, Constant
 from ..core import field, Module
-from ..typing import Real_, RealImage, ImageCoords
+from ..typing import ComplexImage, ImageCoords
 
 
 class Detector(Module):
     """
     Base class for an electron detector.
-
-    Attributes
-    ----------
-    pixel_size :
-        The pixel size measured by the detector.
-        This is in dimensions of physical length.
-        If this is given, the pixel size of the
-        image will be interpolated to this new pixel
-        size.
-    interpolation_method :
-        The interpolation method used for measuring
-        the image at the new ``pixel_size``.
     """
 
-    pixel_size: Optional[Real_] = field(default=None)
-    interpolation_method: str = field(static=True, default="bicubic")
-
-    def measure_at_pixel_size(
-        self, image: RealImage, current_pixel_size: Real_
-    ) -> RealImage:
-        """
-        Measure an image at a given pixel size to
-        the detector pixel size.
-        """
-        if self.pixel_size is None:
-            return image
-        else:
-            return rescale_pixel_size(
-                image,
-                current_pixel_size,
-                new_pixel_size=self.pixel_size,
-                method=self.interpolation_method,
-                antialias=False,
-            )
-
     @abstractmethod
-    def sample(
-        self,
-        key: PRNGKeyArray,
-        freqs: ImageCoords,
-        image: Optional[RealImage] = None,
-    ) -> RealImage:
+    def sample(self, key: PRNGKeyArray, freqs: ImageCoords) -> ComplexImage:
         """Sample a realization from the detector noise model."""
         raise NotImplementedError
 
@@ -81,28 +37,8 @@ class NullDetector(Detector):
     """
 
     @override
-    def sample(
-        self,
-        key: PRNGKeyArray,
-        freqs: ImageCoords,
-        image: Optional[RealImage] = None,
-    ) -> RealImage:
-        return jnp.zeros(jnp.asarray(freqs).shape[0:-1])
-
-
-class NoiselessDetector(Detector):
-    """
-    A detector with no noise model.
-    """
-
-    @override
-    def sample(
-        self,
-        key: PRNGKeyArray,
-        freqs: ImageCoords,
-        image: Optional[RealImage] = None,
-    ) -> RealImage:
-        return jnp.zeros(jnp.asarray(freqs).shape[0:-1])
+    def sample(self, key: PRNGKeyArray, freqs: ImageCoords) -> ComplexImage:
+        return jnp.zeros(jnp.asarray(freqs).shape[0:-1], dtype=complex)
 
 
 class GaussianDetector(GaussianNoise, Detector):
@@ -118,39 +54,8 @@ class GaussianDetector(GaussianNoise, Detector):
         ``Constant()``.
     """
 
-    variance: Kernel = field(default_factory=Constant)
+    variance: KernelType = field(default_factory=Constant)  # type: ignore
 
     @override
-    def sample(
-        self,
-        key: PRNGKeyArray,
-        freqs: ImageCoords,
-        image: Optional[RealImage] = None,
-    ) -> RealImage:
-        return ifftn(super().sample(key, freqs)).real
-
-
-@partial(jax.jit, static_argnames=["method", "antialias"])
-def rescale_pixel_size(
-    image: RealImage,
-    current_pixel_size: Real_,
-    new_pixel_size: Real_,
-    **kwargs: Any,
-) -> RealImage:
-    """
-    Measure an image at a given pixel size using interpolation.
-
-    For more detail, see ``cryojax.utils.interpolation.scale``.
-
-    Parameters
-    ----------
-    image :
-        The image to be magnified.
-    current_pixel_size :
-        The pixel size of the input image.
-    new_pixel_size :
-        The new pixel size after interpolation.
-    """
-    scale_factor = current_pixel_size / new_pixel_size
-    s = jnp.array([scale_factor, scale_factor])
-    return scale(image, image.shape, s, **kwargs)
+    def sample(self, key: PRNGKeyArray, freqs: ImageCoords) -> ComplexImage:
+        return super().sample(key, freqs)
diff --git a/src/cryojax/simulator/distribution.py b/src/cryojax/simulator/distribution.py
index b3ef090d..8432f5c4 100644
--- a/src/cryojax/simulator/distribution.py
+++ b/src/cryojax/simulator/distribution.py
@@ -4,110 +4,157 @@
 
 from __future__ import annotations
 
-__all__ = ["Distribution", "GaussianImage"]
+__all__ = ["Distribution", "IndependentFourierGaussian"]
 
 from abc import abstractmethod
-from typing import Union, Optional
+from typing import Union, Optional, Any
 from typing_extensions import override
 from functools import cached_property
 
 import jax.numpy as jnp
 from jaxtyping import PRNGKeyArray
 
-from .ice import NullIce, GaussianIce
-from .detector import NullDetector, GaussianDetector
-from .image import ImagePipeline, _PRNGKeyArrayLike
-from ..utils import fftn
-from ..typing import Real_, Image, RealImage
+from .noise import GaussianNoise
+from .ice import GaussianIce
+from .detector import GaussianDetector
+from .pipeline import ImagePipeline
+from ..typing import Real_, RealImage, ComplexImage, Image
+from ..core import Module, field
 
 
-class Distribution(ImagePipeline):
+class Distribution(Module):
     """
     An imaging pipeline equipped with a probabilistic model.
     """
 
+    pipeline: ImagePipeline = field()
+
     @abstractmethod
-    def log_probability(self, observed: RealImage) -> Real_:
+    def log_probability(self, observed: Image) -> Real_:
         """
         Evaluate the log-probability.
 
-        Attributes
+        Parameters
         ----------
         observed :
-            The observed data in real space. This must be the same
-            shape as ``scattering.shape``. Note that the user
-            should preprocess the observed data before passing it
-            to the image, such as applying the ``filters`` and
-            ``masks``.
+            The observed data in real or fourier space.
         """
         raise NotImplementedError
 
-    @override
-    def __call__(
-        self,
-        *,
-        key: Optional[Union[PRNGKeyArray, _PRNGKeyArrayLike]] = None,
-        observed: Optional[RealImage] = None,
-        view: bool = True,
-    ) -> Union[Image, Real_]:
+    def sample(self, key: PRNGKeyArray, **kwargs: Any) -> RealImage:
         """
-        If ``observed = None``, sample an image from
-        a noise model. Otherwise, compute the log likelihood.
+        Sample from the distribution.
+
+        Parameters
+        ----------
+        key :
+            The RNG key or key(s). See ``ImagePipeline.sample`` for
+            documentation.
         """
-        if key is not None:
-            return self.sample(key, view=view)
-        elif observed is not None:
-            return self.log_probability(observed)
-        else:
-            return self.render(view=view)
+        return self.pipeline.sample(key, **kwargs)
 
 
-class GaussianImage(Distribution):
-    """
-    Sample an image from a gaussian noise model, or compute
-    the log-likelihood.
+class IndependentFourierGaussian(Distribution):
+    r"""
+    A gaussian noise model, where each fourier mode is independent.
 
-    Note that this computes the likelihood in Fourier space,
+    This computes the likelihood in Fourier space,
     which allows one to model an arbitrary noise power spectrum.
+
+    If no `GaussianNoise` model is explicitly passed, the variance is computed as
+
+    .. math::
+        Var[D(q)] + CTF(q)^2 Var[I(q)]
+
+    where :math:`D(q)` and :math:`I(q)` are independent gaussian random variables in fourier
+    space for the detector and ice, respectively, for a given fourier mode :math:`q`.
+
+    Attributes
+    ----------
+    noise :
+        The gaussian noise model. If not given, use the detector and ice noise
+        models.
     """
 
+    noise: Optional[GaussianNoise] = field(default=None)
+
     def __check_init__(self):
-        if not isinstance(self.solvent, GaussianIce) and not isinstance(
-            self.instrument.detector, GaussianDetector
-        ):
+        if (
+            not isinstance(self.pipeline.solvent, GaussianIce)
+            and not isinstance(
+                self.pipeline.instrument.detector, GaussianDetector
+            )
+        ) and self.noise is None:
             raise ValueError(
-                "Either GaussianIce or GaussianDetector are required."
+                "Either a GaussianIce, GaussianDetector, or GaussianNoise model are required."
             )
 
     @override
-    def log_probability(self, observed: RealImage) -> Real_:
-        """Evaluate the log-likelihood of the data given a parameter set."""
-        # Get variance
-        variance = self.variance
+    def sample(self, key: PRNGKeyArray, **kwargs: Any) -> RealImage:
+        """Sample from the Gaussian noise model."""
+        if self.noise is None:
+            return super().sample(key, **kwargs)
+        else:
+            freqs = self.pipeline.scattering.padded_frequency_grid_in_angstroms
+            noise = self.noise.sample(key, freqs)
+            image = self.pipeline.render(view=False, get_real=False)
+            return self.pipeline._postprocess_image(image + noise, **kwargs)
+
+    @override
+    def log_probability(self, observed: ComplexImage) -> Real_:
+        """
+        Evaluate the log-likelihood of the gaussian noise model.
+
+        This evaluates the log probability in the super sampled
+        coordinate system. Therefore, ``observed.shape`` must
+        match ``ImageManager.padded_shape``.
+
+        Parameters
+        ----------
+        observed :
+           The observed data in fourier space. This must match
+           the ImageManager.padded_shape shape.
+        """
+        pipeline = self.pipeline
+        if (
+            observed.shape
+            != pipeline.scattering.manager.padded_frequency_grid.shape[:-1]
+        ):
+            raise ValueError(
+                "Shape of observed must match ImageManager.padded_shape"
+            )
         # Get residuals
-        residuals = fftn(self.render() - observed)
-        # Crop redundant frequencies
-        _, N2 = self.manager.shape
-        z = N2 // 2 + 1
-        residuals = residuals[:, :z]
-        if not isinstance(variance, Real_):
-            variance = variance[:, :z]
-        loss = jnp.sum((residuals * jnp.conjugate(residuals)) / (2 * variance))
+        residuals = pipeline.render(view=False, get_real=False) - observed
+        # Apply filters, crop, and mask
+        residuals = pipeline._postprocess_image(
+            residuals, view=True, get_real=False
+        )
+        # Compute loss
+        loss = jnp.sum(
+            (residuals * jnp.conjugate(residuals)) / (2 * self.variance)
+        )
+        # Divide by number of modes (parseval's theorem)
         loss = loss.real / residuals.size
 
         return loss
 
     @cached_property
     def variance(self) -> Union[Real_, RealImage]:
-        # Gather image configuration
-        freqs = self.manager.freqs / self.pixel_size
-        # Variance from detector
-        if not isinstance(self.instrument.detector, NullDetector):
-            variance = self.instrument.detector.variance(freqs)
+        pipeline = self.pipeline
+        # Gather frequency coordinates
+        freqs = pipeline.scattering.frequency_grid_in_angstroms
+        if self.noise is None:
+            # Variance from detector
+            if isinstance(pipeline.instrument.detector, GaussianDetector):
+                variance = pipeline.instrument.detector.variance(freqs)
+            else:
+                variance = jnp.asarray(0.0)
+            # Variance from ice
+            if isinstance(pipeline.solvent, GaussianIce):
+                ctf = pipeline.instrument.optics(
+                    freqs, pose=pipeline.ensemble.pose
+                )
+                variance += ctf**2 * pipeline.solvent.variance(freqs)
+            return variance
         else:
-            variance = 0.0
-        # Variance from ice
-        if not isinstance(self.solvent, NullIce):
-            ctf = self.instrument.optics(freqs, pose=self.ensemble.pose)
-            variance += ctf**2 * self.solvent.variance(freqs)
-        return variance
+            return self.noise.variance(freqs)
diff --git a/src/cryojax/simulator/ensemble.py b/src/cryojax/simulator/ensemble.py
index 0e3e2f87..1e5149c0 100644
--- a/src/cryojax/simulator/ensemble.py
+++ b/src/cryojax/simulator/ensemble.py
@@ -49,4 +49,4 @@ def realization(self) -> ElectronDensity:
                 lambda i=i: self.density[i] for i in range(len(self.density))
             ]
             density = jax.lax.switch(self.conformation, funcs)
-        return density.rotate_to(self.pose)
+        return density.rotate_to_pose(self.pose)
diff --git a/src/cryojax/simulator/exposure.py b/src/cryojax/simulator/exposure.py
index cc143dda..5b4ff763 100644
--- a/src/cryojax/simulator/exposure.py
+++ b/src/cryojax/simulator/exposure.py
@@ -75,7 +75,7 @@ class UniformExposure(Exposure):
     mu: Intensity offset.
     """
 
-    N: Real_ = field(default=1e5)
+    N: Real_ = field(default=1.0)
     mu: Real_ = field(default=0.0)
 
     def scaling(self, freqs: ImageCoords, **kwargs: Any) -> Real_:
diff --git a/src/cryojax/simulator/filter.py b/src/cryojax/simulator/filter.py
index 67349fe9..e1bb0dfd 100644
--- a/src/cryojax/simulator/filter.py
+++ b/src/cryojax/simulator/filter.py
@@ -5,29 +5,28 @@
 from __future__ import annotations
 
 __all__ = [
-    "compute_lowpass_filter",
-    "compute_whitening_filter",
     "Filter",
-    "ProductFilter",
+    "FilterType",
     "LowpassFilter",
     "WhiteningFilter",
 ]
 
 from abc import abstractmethod
-from typing import Any
+from typing import Any, TypeVar
 from typing_extensions import override
 
-import numpy as np
+import jax
 import jax.numpy as jnp
 
-from .manager import ImageManager
+from ..utils import powerspectrum, make_frequencies, rfftn
+from ..core import field, BufferModule
+from ..typing import Image, ImageCoords, RealImage
 
-from ..utils import powerspectrum, make_frequencies
-from ..core import field, Buffer
-from ..typing import Image, ImageCoords, Real_, RealImage, ComplexImage
+FilterType = TypeVar("FilterType", bound="Filter")
+"""TypeVar for the Filter base class."""
 
 
-class Filter(Buffer):
+class Filter(BufferModule):
     """
     Base class for computing and applying an image filter.
 
@@ -40,11 +39,8 @@ class Filter(Buffer):
 
     filter: Image = field(init=False)
 
-    def __post_init__(self, *args: Any, **kwargs: Any):
-        self.filter = self.evaluate(*args, **kwargs)
-
     @abstractmethod
-    def evaluate(self, *args: Any, **kwargs: Any) -> Image:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         """Compute the filter."""
         raise NotImplementedError
 
@@ -52,21 +48,27 @@ def __call__(self, image: Image) -> Image:
         """Apply the filter to an image."""
         return self.filter * image
 
-    def __mul__(self, other: Filter) -> Filter:
-        return ProductFilter(self, other)
+    def __mul__(self: FilterType, other: FilterType) -> _ProductFilter:
+        return _ProductFilter(filter1=self, filter2=other)
 
-    def __rmul__(self, other: Filter) -> Filter:
-        return ProductFilter(other, self)
+    def __rmul__(self: FilterType, other: FilterType) -> _ProductFilter:
+        return _ProductFilter(filter1=other, filter2=self)
 
 
-class ProductFilter(Filter):
+class _ProductFilter(Filter):
     """A helper to represent the product of two filters."""
 
-    filter1: Filter = field()
-    filter2: Filter = field()
+    filter1: FilterType  # type: ignore
+    filter2: FilterType  # type: ignore
+
+    @override
+    def __init__(self, filter1: FilterType, filter2: FilterType) -> None:
+        self.filter1 = filter1
+        self.filter2 = filter2
+        self.filter = filter1.filter * filter2.filter
 
-    def evaluate(self) -> Image:
-        return self.filter1.filter * self.filter2.filter
+    def __repr__(self):
+        return f"{repr(self.filter1)} * {repr(self.filter2)}"
 
 
 class LowpassFilter(Filter):
@@ -86,43 +88,39 @@ class LowpassFilter(Filter):
         By default, ``0.05``.
     """
 
-    manager: ImageManager = field()
+    cutoff: float = field(static=True)
+    rolloff: float = field(static=True)
 
-    cutoff: float = field(static=True, default=0.95)
-    rolloff: float = field(static=True, default=0.05)
-
-    @override
-    def evaluate(self, **kwargs) -> RealImage:
-        return compute_lowpass_filter(
-            self.manager.padded_freqs, self.cutoff, self.rolloff, **kwargs
-        )
+    def __init__(
+        self,
+        freqs: ImageCoords,
+        cutoff: float = 0.95,
+        rolloff: float = 0.05,
+    ) -> None:
+        self.cutoff = cutoff
+        self.rolloff = rolloff
+        self.filter = _compute_lowpass_filter(freqs, self.cutoff, self.rolloff)
 
 
 class WhiteningFilter(Filter):
     """
-    Apply an whitening filter to an image.
-
-    See documentation for
-    ``cryojax.simulator.compute_whitening_filter``
-    for more information.
+    Apply a whitening filter to an image.
     """
 
-    manager: ImageManager = field()
-
-    micrograph: ComplexImage = field()
-
-    @override
-    def evaluate(self, **kwargs: Any) -> RealImage:
-        return compute_whitening_filter(
-            self.manager.padded_freqs, self.micrograph, **kwargs
+    def __init__(
+        self,
+        frequency_grid: ImageCoords,
+        micrograph: RealImage,
+        *,
+        grid_spacing: float = 1.0,
+    ):
+        self.filter = _compute_whitening_filter(
+            frequency_grid, micrograph, grid_spacing
         )
 
 
-def compute_lowpass_filter(
-    freqs: ImageCoords,
-    cutoff: float = 0.667,
-    rolloff: float = 0.05,
-    **kwargs: Any,
+def _compute_lowpass_filter(
+    freqs: ImageCoords, cutoff: float = 0.667, rolloff: float = 0.05
 ) -> RealImage:
     """
     Create a low-pass filter.
@@ -137,13 +135,11 @@ def compute_lowpass_filter(
     rolloff :
         The rolloff width as a fraction of the Nyquist frequency.
         By default, ``0.05``.
-    kwargs :
-        Keyword arguments passed to ``cryojax.utils.make_coordinates``.
 
     Returns
     -------
     mask : `Array`, shape `shape`
-        An array representing the anti-aliasing filter.
+        An array representing the low pass filter.
     """
 
     k_max = 1.0 / 2.0
@@ -167,32 +163,54 @@ def compute_lowpass_filter(
     return mask
 
 
-def compute_whitening_filter(
-    freqs: ImageCoords, micrograph: ComplexImage, **kwargs: Any
+def _compute_whitening_filter(
+    freqs: ImageCoords,
+    micrograph: RealImage,
+    grid_spacing: float = 1.0,
 ) -> RealImage:
     """
     Compute a whitening filter from a micrograph. This is taken
     to be the inverse square root of the 2D radially averaged
     power spectrum.
 
+    This implementation follows the cisTEM whitening filter
+    algorithm.
+
     Parameters
     ----------
     freqs :
         The image coordinates.
     micrograph :
-        The micrograph in fourier space.
+        The micrograph in real space.
 
     Returns
     -------
-    spectrum :
-        The power spectrum isotropically averaged onto a coordinate
-        system whose shape is set by ``shape``.
+    filter :
+        The whitening filter.
     """
-    micrograph = jnp.asarray(micrograph)
     # Make coordinates
-    micrograph_freqs = make_frequencies(micrograph.shape, *kwargs)
+    micrograph_frequency_grid = make_frequencies(
+        micrograph.shape, grid_spacing
+    )
+    # Transform to fourier space
+    fourier_micrograph = rfftn(micrograph)
+    # Compute norms
+    radial_frequency_grid = jnp.linalg.norm(micrograph_frequency_grid, axis=-1)
+    interpolating_radial_frequency_grid = jnp.linalg.norm(freqs, axis=-1)
     # Compute power spectrum
-    micrograph /= jnp.sqrt(np.prod(micrograph.shape))
-    spectrum, _ = powerspectrum(micrograph, micrograph_freqs, grid=freqs)
-
-    return 1 / jnp.sqrt(spectrum)
+    spectrum, _ = powerspectrum(
+        fourier_micrograph,
+        radial_frequency_grid,
+        k_max=jnp.sqrt(2.0) / 2.0,
+        interpolating_radial_frequency_grid=interpolating_radial_frequency_grid,
+    )
+    # Compute inverse square root
+    filter = jax.lax.rsqrt(spectrum)
+    # Divide filter by maximum, excluding zero mode
+    maximum = jnp.max(jnp.delete(filter, jnp.asarray((0, 0), dtype=int)))
+    filter /= maximum
+    # Set zero mode manually to 1 (this diverges from the cisTEM
+    # algorithm).
+    filter = filter.at[0, 0].set(1.0)
+
+    return filter
diff --git a/src/cryojax/simulator/ice.py b/src/cryojax/simulator/ice.py
index a41f978b..f9ad951a 100644
--- a/src/cryojax/simulator/ice.py
+++ b/src/cryojax/simulator/ice.py
@@ -10,7 +10,7 @@
 import jax.numpy as jnp
 from jaxtyping import PRNGKeyArray
 
-from .kernel import Kernel, Exp
+from .kernel import KernelType, Exp
 from .noise import GaussianNoise
 from ..core import field, Module
 from ..typing import RealImage, ComplexImage, Image, ImageCoords
@@ -58,7 +58,7 @@ class GaussianIce(GaussianNoise, Ice):
         ``Exp()``.
     """
 
-    variance: Kernel = field(default_factory=Exp)
+    variance: KernelType = field(default_factory=Exp)  # type: ignore
 
     def sample(
         self,
diff --git a/src/cryojax/simulator/image.py b/src/cryojax/simulator/image.py
deleted file mode 100644
index d796fe3b..00000000
--- a/src/cryojax/simulator/image.py
+++ /dev/null
@@ -1,288 +0,0 @@
-"""
-Image formation models.
-"""
-
-from __future__ import annotations
-
-__all__ = ["ImagePipeline"]
-
-from typing import Union, get_args, Optional
-
-import equinox as eqx
-import jax
-import jax.numpy as jnp
-import jax.tree_util as jtu
-from jaxtyping import PRNGKeyArray, Shaped
-
-from .filter import Filter
-from .mask import Mask
-from .ensemble import Ensemble
-from .assembly import Assembly
-from .scattering import ScatteringModel
-from .manager import ImageManager
-from .instrument import Instrument
-from .exposure import NullExposure
-from .detector import NullDetector
-from .optics import NullOptics
-from .ice import Ice, NullIce
-from ..utils import fftn, ifftn
-from ..core import field, Module
-from ..typing import RealImage, ComplexImage, Image, Real_
-
-
-_PRNGKeyArrayLike = Shaped[PRNGKeyArray, "M"]
-
-
-class ImagePipeline(Module):
-    """
-    Base class for an imaging model.
-
-    Call an ``ImagePipeline`` or its ``render``, ``sample``,
-    or ``log_probability`` routines.
-
-    Attributes
-    ----------
-    ensemble :
-        The ensemble from which to render images.
-    scattering :
-        The scattering model.
-    instrument :
-        The abstraction of the electron microscope.
-    solvent :
-        The solvent around the specimen.
-    filter :
-        A filter to apply to the image.
-    mask :
-        A mask to apply to the image.
-
-    Properties
-    ----------
-    manager :
-        Exposes the API of the scattering model's image
-        manager.
-    pixel_size :
-        The pixel size of the image.
-    """
-
-    ensemble: Union[Ensemble, Assembly] = field()
-    scattering: ScatteringModel = field()
-    instrument: Instrument = field(default_factory=Instrument)
-    solvent: Ice = field(default_factory=NullIce)
-
-    filter: Optional[Filter] = field(default=None)
-    mask: Optional[Mask] = field(default=None)
-
-    @property
-    def manager(self) -> ImageManager:
-        """The ImageManager"""
-        return self.scattering.manager
-
-    @property
-    def pixel_size(self) -> Real_:
-        """The image pixel size."""
-        if self.instrument.detector.pixel_size is not None:
-            return self.instrument.detector.pixel_size
-        else:
-            return self.scattering.pixel_size
-
-    def render(self, view: bool = True) -> RealImage:
-        """
-        Render an image of a Specimen.
-
-        Parameters
-        ----------
-        view : `bool`, optional
-            If ``True``, view the cropped, masked,
-            and filtered image.
-        """
-        if isinstance(self.ensemble, Ensemble):
-            image = self._render_ensemble()
-        elif isinstance(self.ensemble, Assembly):
-            image = self._render_assembly()
-        else:
-            raise ValueError(
-                "The ensemble must an instance of an Ensemble or an Assembly."
-            )
-
-        if view:
-            image = self._view(image)
-
-        return image
-
-    def sample(
-        self, key: Union[PRNGKeyArray, _PRNGKeyArrayLike], view: bool = True
-    ) -> RealImage:
-        """
-        Sample the an image from a realization of the noise.
-
-        Parameters
-        ----------
-        key :
-            The random number generator key(s). If
-            the ``ImagePipeline`` is configured with
-            more than one stochastic model (e.g. a detector
-            and solvent model), then this parameter could be
-            ``jax.random.split(jax.random.PRNGKey(seed), 2)``.
-        view : `bool`, optional
-            If ``True``, view the cropped, masked,
-            and filtered image.
-        """
-        # Check PRNGKey
-        idx = 0  # Keep track of number of stochastic models
-        if isinstance(key, get_args(PRNGKeyArray)):
-            key = jnp.expand_dims(key, axis=0)
-        # Frequencies
-        freqs = self.manager.padded_freqs / self.pixel_size
-        # The image of the specimen drawn from the ensemble
-        image = self.render(view=False)
-        if not isinstance(self.solvent, NullIce):
-            # The image of the solvent
-            ice_image = ifftn(
-                self.instrument.optics(freqs)
-                * self.solvent.sample(key[idx], freqs)
-            ).real
-            image = image + ice_image
-            idx += 1
-        if not isinstance(self.instrument.detector, NullDetector):
-            # Measure the detector readout
-            noise = self.instrument.detector.sample(
-                key[idx], freqs, image=image
-            )
-            image = image + noise
-            idx += 1
-        if view:
-            image = self._view(image)
-
-        return image
-
-    def __call__(
-        self,
-        key: Optional[Union[PRNGKeyArray, _PRNGKeyArrayLike]] = None,
-        view: bool = True,
-    ) -> Image:
-        """
-        Sample or render an image.
-        """
-        if key is None:
-            return self.render(view=view)
-        else:
-            return self.sample(key, view=view)
-
-    def _view(self, image: Image, is_real: bool = True) -> RealImage:
-        """
-        View the image. This function applies
-        filters, crops the image, then applies masks.
-        """
-        # Apply filters to the image
-        if self.filter is not None:
-            if is_real:
-                image = fftn(image)
-            image = ifftn(self.filter(image)).real
-        # Crop the image
-        image = self.manager.crop(image)
-        # Mask the image
-        if self.mask is not None:
-            image = self.mask(image)
-        return image
-
-    def _render_ensemble(self, get_real: bool = True) -> Image:
-        """Render an image of a Specimen."""
-        pixel_size = self.scattering.pixel_size
-        freqs = self.manager.padded_freqs / pixel_size
-        # Draw the electron density at a particular conformation and pose
-        density = self.ensemble.realization
-        # Compute the scattering image in fourier space
-        image = self.scattering.scatter(density)
-        # Normalize to cisTEM conventions
-        image = self.manager.normalize_to_cistem(image, is_real=False)
-        # Apply translation
-        image *= self.ensemble.pose.shifts(freqs)
-        # Measure the image with the instrument
-        image = self._measure_with_instrument(image, get_real=get_real)
-
-        return image
-
-    def _render_assembly(self, get_real: bool = True) -> Image:
-        """Render an image of an Assembly from its subunits."""
-        # Get the subunits
-        subunits = self.ensemble.subunits
-        # Split up computation with two different instruments
-        optics_instrument = eqx.tree_at(
-            lambda ins: (ins.exposure, ins.detector),
-            self.instrument,
-            (NullExposure(), NullDetector()),
-        )
-        no_optics_instrument = eqx.tree_at(
-            lambda ins: ins.optics, self.instrument, NullOptics()
-        )
-        # Create an ImagePipeline for the subunits
-        subunit_pipeline = eqx.tree_at(
-            lambda m: (m.ensemble, m.instrument),
-            self,
-            (subunits, optics_instrument),
-        )
-        # Setup vmap over poses and conformations
-        if self.ensemble.conformation is None:
-            where_vmap = lambda m: (m.ensemble.pose,)
-            n_vmap = 1
-        else:
-            where_vmap = lambda m: (m.ensemble.pose, m.ensemble.conformation)
-            n_vmap = 2
-        to_vmap = eqx.tree_at(
-            where_vmap,
-            jtu.tree_map(lambda _: False, subunit_pipeline),
-            tuple(n_vmap * [True]),
-        )
-        vmap, novmap = eqx.partition(subunit_pipeline, to_vmap)
-        # Compute all subunit images and sum
-        compute_stack = jax.vmap(
-            lambda vmap, novmap: eqx.combine(vmap, novmap)._render_ensemble(
-                get_real=False
-            ),
-            in_axes=(0, None),
-        )
-        compute_stack_and_sum = jax.jit(
-            lambda vmap, novmap: jnp.sum(
-                compute_stack(vmap, novmap),
-                axis=0,
-            )
-        )
-        image = compute_stack_and_sum(vmap, novmap)
-        # Finally, measure the image without applying the CTF
-        no_optics_pipeline = eqx.tree_at(
-            lambda m: m.instrument, self, no_optics_instrument
-        )
-        image = no_optics_pipeline._measure_with_instrument(
-            image, get_real=get_real
-        )
-
-        return image
-
-    def _measure_with_instrument(
-        self, image: ComplexImage, get_real: bool = True
-    ) -> Image:
-        """Measure an image with the instrument"""
-        instrument = self.instrument
-        current_pixel_size = self.scattering.pixel_size
-        freqs = self.manager.padded_freqs / current_pixel_size
-        # Compute and apply CTF
-        ctf = instrument.optics(freqs, pose=self.ensemble.pose)
-        image = ctf * image
-        # Apply the electron exposure model
-        scaling, offset = instrument.exposure.scaling(
-            freqs
-        ), instrument.exposure.offset(freqs)
-        image = scaling * image + offset
-        # Add some detector logic to avoid unecessary FFTs
-        if instrument.detector.pixel_size is not None:
-            # Measure at the detector pixel size
-            image = instrument.detector.measure_at_pixel_size(
-                ifftn(image).real, current_pixel_size
-            )
-            if not get_real:
-                image = fftn(image)
-        else:
-            if get_real:
-                image = ifftn(image).real
-
-        return image
diff --git a/src/cryojax/simulator/kernel.py b/src/cryojax/simulator/kernel.py
index 24c0f2b2..c5604bb2 100644
--- a/src/cryojax/simulator/kernel.py
+++ b/src/cryojax/simulator/kernel.py
@@ -9,8 +9,7 @@
 
 __all__ = [
     "Kernel",
-    "ProductKernel",
-    "SumKernel",
+    "KernelType",
     "Constant",
     "Exp",
     "Gaussian",
@@ -19,7 +18,8 @@
 ]
 
 from abc import abstractmethod
-from typing import Any, Union, Callable, Concatenate, ParamSpec, Optional
+from typing import overload, Any, Union, Callable, Concatenate, TypeVar
+from typing_extensions import override
 from jaxtyping import Array
 
 import jax.numpy as jnp
@@ -27,7 +27,8 @@
 from ..core import field, Module
 from ..typing import Real_, ImageCoords, RealImage, Image
 
-P = ParamSpec("P")
+KernelType = TypeVar("KernelType", bound="Kernel")
+"""TypeVar for the Kernel base class"""
 
 
 class Kernel(Module):
@@ -44,8 +45,18 @@ class Kernel(Module):
         2) Overrwrite :func:`Kernel.evaluate`.
     """
 
+    @overload
     @abstractmethod
-    def evaluate(self, coords: ImageCoords, **kwargs: Any) -> Array:
+    def evaluate(self, freqs: ImageCoords, **kwargs: Any) -> Array:
+        ...
+
+    @overload
+    @abstractmethod
+    def evaluate(self, freqs: None, **kwargs: Any) -> Array:
+        ...
+
+    @abstractmethod
+    def evaluate(self, freqs: ImageCoords | None, **kwargs: Any) -> Array:
         """
         Evaluate the kernel at a set of coordinates.
 
@@ -57,53 +68,61 @@ def evaluate(self, coords: ImageCoords, **kwargs: Any) -> Array:
         pass
 
     def __call__(
-        self, coords: ImageCoords, *args: Any, **kwargs: Any
+        self, freqs: ImageCoords | None = None, **kwargs: Any
     ) -> Array:
-        coords = jnp.asarray(coords)
-        return self.evaluate(coords, *args, **kwargs)
+        return self.evaluate(freqs, **kwargs)
 
-    def __add__(self, other: Union[Kernel, float]) -> Kernel:
+    def __add__(self, other: Union[KernelType, Real_]) -> _SumKernel:
         if isinstance(other, Kernel):
-            return SumKernel(self, other)
-        return SumKernel(self, Constant(other))
+            return _SumKernel(self, other)
+        return _SumKernel(self, Constant(other))
 
-    def __radd__(self, other: Any) -> Kernel:
-        # We'll hit this first branch when using the `sum` function
-        if other == 0:
-            return self
+    def __radd__(self, other: Any) -> _SumKernel:
         if isinstance(other, Kernel):
-            return SumKernel(other, self)
-        return SumKernel(Constant(other), self)
+            return _SumKernel(other, self)
+        return _SumKernel(Constant(other), self)
 
-    def __mul__(self, other: Union[Kernel, float]) -> Kernel:
+    def __mul__(self, other: Union[Kernel, Real_]) -> _ProductKernel:
         if isinstance(other, Kernel):
-            return ProductKernel(self, other)
-        return ProductKernel(self, Constant(other))
+            return _ProductKernel(self, other)
+        return _ProductKernel(self, Constant(other))
 
-    def __rmul__(self, other: Any) -> Kernel:
+    def __rmul__(self, other: Any) -> _ProductKernel:
         if isinstance(other, Kernel):
-            return ProductKernel(other, self)
-        return ProductKernel(Constant(other), self)
+            return _ProductKernel(other, self)
+        return _ProductKernel(Constant(other), self)
 
 
-class SumKernel(Kernel):
+class _SumKernel(Kernel):
     """A helper to represent the sum of two kernels"""
 
-    kernel1: Kernel = field()
-    kernel2: Kernel = field()
+    kernel1: KernelType = field()  # type: ignore
+    kernel2: KernelType = field()  # type: ignore
+
+    @override
+    def evaluate(
+        self, freqs: ImageCoords | None = None, **kwargs: Any
+    ) -> Array:
+        return self.kernel1(freqs) + self.kernel2(freqs)
 
-    def evaluate(self, coords: ImageCoords) -> Array:
-        return self.kernel1(coords) + self.kernel2(coords)
+    def __repr__(self):
+        return f"{repr(self.kernel1)} + {repr(self.kernel2)}"
 
 
-class ProductKernel(Kernel):
+class _ProductKernel(Kernel):
     """A helper to represent the product of two kernels"""
 
-    kernel1: Kernel = field()
-    kernel2: Kernel = field()
+    kernel1: KernelType = field()  # type: ignore
+    kernel2: KernelType = field()  # type: ignore
+
+    @override
+    def evaluate(
+        self, freqs: ImageCoords | None = None, **kwargs: Any
+    ) -> Array:
+        return self.kernel1(freqs) * self.kernel2(freqs)
 
-    def evaluate(self, coords: ImageCoords) -> Array:
-        return self.kernel1(coords) * self.kernel2(coords)
+    def __repr__(self):
+        return f"{repr(self.kernel1)} * {repr(self.kernel2)}"
 
 
 class Constant(Kernel):
@@ -118,7 +137,10 @@ class Constant(Kernel):
 
     value: Real_ = field(default=1.0)
 
-    def evaluate(self, coords: Optional[ImageCoords] = None) -> Real_:
+    @override
+    def evaluate(
+        self, freqs: ImageCoords | None = None, **kwargs: Any
+    ) -> Real_:
         if jnp.ndim(self.value) != 0:
             raise ValueError("The value of a constant kernel must be a scalar")
         return self.value
@@ -136,9 +158,15 @@ class ZeroMode(Kernel):
 
     value: Real_ = field(default=1.0)
 
-    def evaluate(self, freqs: ImageCoords) -> RealImage:
-        N1, N2 = freqs.shape[0:-1]
-        return jnp.zeros((N1, N2)).at[0, 0].set(N1 * N2 * self.value)
+    @override
+    def evaluate(self, freqs: ImageCoords | None, **kwargs: Any) -> RealImage:
+        if freqs is None:
+            raise ValueError(
+                "The frequency grid must be given as an argument to the Kernel call."
+            )
+        else:
+            N1, N2 = freqs.shape[0:-1]
+            return jnp.zeros((N1, N2)).at[0, 0].set(N1 * N2 * self.value)
 
 
 class Exp(Kernel):
@@ -177,11 +205,17 @@ class Exp(Kernel):
     scale: Real_ = field(default=1.0)
     offset: Real_ = field(default=0.0)
 
-    def evaluate(self, freqs: ImageCoords) -> RealImage:
-        k_sqr = jnp.sum(freqs**2, axis=-1)
-        scaling = 1.0 / (k_sqr + jnp.divide(1, (self.scale) ** 2)) ** 1.5
-        scaling *= jnp.divide(self.amplitude, 2 * jnp.pi * self.scale**3)
-        return scaling + self.offset
+    @override
+    def evaluate(self, freqs: ImageCoords | None, **kwargs: Any) -> RealImage:
+        if freqs is None:
+            raise ValueError(
+                "The frequency grid must be given as an argument to the Kernel call."
+            )
+        else:
+            k_sqr = jnp.sum(freqs**2, axis=-1)
+            scaling = 1.0 / (k_sqr + jnp.divide(1, (self.scale) ** 2)) ** 1.5
+            scaling *= jnp.divide(self.amplitude, 2 * jnp.pi * self.scale**3)
+            return scaling + self.offset
 
 
 class Gaussian(Kernel):
@@ -218,10 +252,16 @@ class Gaussian(Kernel):
     b_factor: Real_ = field(default=1.0)
     offset: Real_ = field(default=0.0)
 
-    def evaluate(self, freqs: ImageCoords) -> RealImage:
-        k_sqr = jnp.sum(freqs**2, axis=-1)
-        scaling = self.amplitude * jnp.exp(-0.5 * self.b_factor * k_sqr)
-        return scaling + self.offset
+    @override
+    def evaluate(self, freqs: ImageCoords | None, **kwargs: Any) -> RealImage:
+        if freqs is None:
+            raise ValueError(
+                "The frequency grid must be given as an argument to the Kernel call."
+            )
+        else:
+            k_sqr = jnp.sum(freqs**2, axis=-1)
+            scaling = self.amplitude * jnp.exp(-0.5 * self.b_factor * k_sqr)
+            return scaling + self.offset
 
 
 class Empirical(Kernel):
@@ -244,7 +284,10 @@ class Empirical(Kernel):
     amplitude: Real_ = field(default=1.0)
     offset: Real_ = field(default=0.0)
 
-    def evaluate(self, coords: Optional[ImageCoords] = None) -> Image:
+    @override
+    def evaluate(
+        self, freqs: ImageCoords | None = None, **kwargs: Any
+    ) -> Image:
         """Return the scaled and offset measurement."""
         return self.amplitude * self.measurement + self.offset
 
@@ -260,9 +303,15 @@ class Custom(Kernel):
         :func:`Kernel.evaluate`.
     """
 
-    function: Callable[Concatenate[ImageCoords, P], Array] = field(static=True)
+    function: Callable[Concatenate[ImageCoords, ...], Array] = field(
+        static=True
+    )
 
+    @override
     def evaluate(
-        self, coords: ImageCoords, *args: P.args, **kwargs: P.kwargs
+        self, freqs: ImageCoords | None = None, *args: Any, **kwargs: Any
     ) -> Array:
-        return self.function(coords, *args, **kwargs)
+        if freqs is None:
+            return self.function(*args, **kwargs)
+        else:
+            return self.function(freqs, *args, **kwargs)
diff --git a/src/cryojax/simulator/manager.py b/src/cryojax/simulator/manager.py
index a88fe217..2a9e0f9e 100644
--- a/src/cryojax/simulator/manager.py
+++ b/src/cryojax/simulator/manager.py
@@ -6,11 +6,9 @@
 
 __all__ = ["ImageManager"]
 
-from typing import Any, Union, Callable
+from typing import Any, Union, Callable, Optional
 
-import jax.numpy as jnp
-
-from ..core import field, Buffer
+from ..core import field, BufferModule
 from ..typing import (
     Image,
     ImageCoords,
@@ -21,11 +19,11 @@
     crop,
     pad,
     crop_or_pad,
-    resize,
+    normalize_image,
 )
 
 
-class ImageManager(Buffer):
+class ImageManager(BufferModule):
     """
     Configuration and utilities for an electron microscopy image.
 
@@ -43,14 +41,14 @@ class ImageManager(Buffer):
     pad_mode :
         The method of image padding. By default, ``"edge"``.
         For all options, see ``jax.numpy.pad``.
-    freqs :
+    frequency_grid :
         The fourier wavevectors in the imaging plane.
-    padded_freqs :
+    padded_frequency_grid :
         The fourier wavevectors in the imaging plane
         in the padded coordinate system.
-    coords :
+    coordinate_grid :
         The coordinates in the imaging plane.
-    padded_coords :
+    padded_coordinate_grid :
         The coordinates in the imaging plane
         in the padded coordinate system.
     """
@@ -61,58 +59,48 @@ class ImageManager(Buffer):
 
     padded_shape: tuple[int, int] = field(static=True, init=False)
 
-    freqs: ImageCoords = field(init=False)
-    padded_freqs: ImageCoords = field(init=False)
-    coords: ImageCoords = field(init=False)
-    padded_coords: ImageCoords = field(init=False)
+    frequency_grid: ImageCoords = field(init=False)
+    padded_frequency_grid: ImageCoords = field(init=False)
+    coordinate_grid: ImageCoords = field(init=False)
+    padded_coordinate_grid: ImageCoords = field(init=False)
 
     def __post_init__(self):
         # Set shape after padding
         padded_shape = tuple([int(s * self.pad_scale) for s in self.shape])
         self.padded_shape = padded_shape
         # Set coordinates
-        self.freqs = make_frequencies(self.shape)
-        self.padded_freqs = make_frequencies(self.padded_shape)
-        self.coords = make_coordinates(self.shape)
-        self.padded_coords = make_coordinates(self.padded_shape)
+        self.frequency_grid = make_frequencies(self.shape)
+        self.padded_frequency_grid = make_frequencies(self.padded_shape)
+        self.coordinate_grid = make_coordinates(self.shape)
+        self.padded_coordinate_grid = make_coordinates(self.padded_shape)
 
-    def crop(self, image: Image) -> Image:
+    def crop_to_shape(self, image: Image) -> Image:
         """Crop an image."""
         return crop(image, self.shape)
 
-    def pad(self, image: Image, **kwargs: Any) -> Image:
+    def pad_to_padded_shape(self, image: Image, **kwargs: Any) -> Image:
         """Pad an image."""
         return pad(image, self.padded_shape, mode=self.pad_mode, **kwargs)
 
-    def crop_or_pad(self, image: Image, **kwargs: Any) -> Image:
+    def crop_or_pad_to_padded_shape(
+        self, image: Image, **kwargs: Any
+    ) -> Image:
         """Reshape an image using cropping or padding."""
         return crop_or_pad(
             image, self.padded_shape, mode=self.pad_mode, **kwargs
         )
 
-    def downsample(
-        self, image: Image, method="lanczos5", **kwargs: Any
+    def normalize_image(
+        self,
+        image: Image,
+        is_real: bool = True,
+        half_space: bool = True,
+        shape_in_real_space: Optional[tuple[int, int]] = None,
     ) -> Image:
-        """Downsample an image using interpolation."""
-        return resize(
-            image, self.shape, antialias=False, method=method, **kwargs
+        """Normalize an image so that it is mean 0 and standard deviation 1 in real space."""
+        return normalize_image(
+            image,
+            is_real=is_real,
+            half_space=half_space,
+            shape_in_real_space=shape_in_real_space,
         )
-
-    def upsample(self, image: Image, method="bicubic", **kwargs: Any) -> Image:
-        """Upsample an image using interpolation."""
-        return resize(image, self.padded_shape, method=method, **kwargs)
-
-    def normalize_to_cistem(
-        self, image: Image, is_real: bool = False
-    ) -> Image:
-        """Normalize images on the exit plane according to cisTEM conventions."""
-        M1, M2 = image.shape
-        if is_real:
-            raise NotImplementedError(
-                "Normalization to cisTEM conventions not supported for real input."
-            )
-        else:
-            # Set zero frequency component to zero
-            image = image.at[0, 0].set(0.0 + 0.0j)
-            # cisTEM normalization convention for projections
-            return image / jnp.sqrt(M1 * M2)
diff --git a/src/cryojax/simulator/mask.py b/src/cryojax/simulator/mask.py
index 78affc60..235eb128 100644
--- a/src/cryojax/simulator/mask.py
+++ b/src/cryojax/simulator/mask.py
@@ -5,25 +5,25 @@
 from __future__ import annotations
 
 __all__ = [
-    "compute_circular_mask",
     "Mask",
-    "ProductMask",
+    "MaskType",
     "CircularMask",
 ]
 
 from abc import abstractmethod
-from typing import Any
-from typing_extensions import override
+from typing import Any, TypeVar
 
 import jax.numpy as jnp
 
-from .manager import ImageManager
-
-from ..core import field, Buffer
+from ..core import field, BufferModule
 from ..typing import RealImage, ImageCoords
 
 
-class Mask(Buffer):
+MaskType = TypeVar("MaskType", bound="Mask")
+"""TypeVar for the Mask base class."""
+
+
+class Mask(BufferModule):
     """
     Base class for computing and applying an image mask.
 
@@ -36,33 +36,35 @@ class Mask(Buffer):
 
     mask: RealImage = field(init=False)
 
-    def __post_init__(self, *args: Any, **kwargs: Any):
-        self.mask = self.evaluate(*args, **kwargs)
-
     @abstractmethod
-    def evaluate(self, *args: Any, **kwargs: Any) -> RealImage:
-        """Compute the mask."""
+    def __init__(self, *args: Any, **kwargs: Any):
+        """Compute the mask"""
         raise NotImplementedError
 
     def __call__(self, image: RealImage) -> RealImage:
         """Apply the mask to an image."""
         return self.mask * image
 
-    def __mul__(self, other: Mask) -> Mask:
-        return ProductMask(self, other)
+    def __mul__(self: MaskType, other: MaskType) -> _ProductMask:
+        return _ProductMask(mask1=self, mask2=other)
 
-    def __rmul__(self, other: Mask) -> Mask:
-        return ProductMask(other, self)
+    def __rmul__(self: MaskType, other: MaskType) -> _ProductMask:
+        return _ProductMask(mask1=other, mask2=self)
 
 
-class ProductMask(Mask):
-    """A helper to represent the product of two masks."""
+class _ProductMask(Mask):
+    """A helper to represent the product of two filters."""
 
-    mask1: Mask = field()
-    mask2: Mask = field()
+    mask1: MaskType  # type: ignore
+    mask2: MaskType  # type: ignore
 
-    def evaluate(self) -> RealImage:
-        return self.mask1.mask * self.mask2.mask
+    def __init__(self, mask1: MaskType, mask2: MaskType):
+        self.mask1 = mask1
+        self.mask2 = mask2
+        self.mask = mask1.mask * mask2.mask
+
+    def __repr__(self):
+        return f"{repr(self.mask1)} * {repr(self.mask2)}"
 
 
 class CircularMask(Mask):
@@ -81,23 +83,22 @@ class CircularMask(Mask):
         By default, ``0.05``.
     """
 
-    manager: ImageManager = field()
-
-    radius: float = field(static=True, default=0.95)
-    rolloff: float = field(static=True, default=0.05)
+    radius: float = field(static=True)
+    rolloff: float = field(static=True)
 
-    @override
-    def evaluate(self, **kwargs: Any) -> RealImage:
-        return compute_circular_mask(
-            self.manager.coords, self.radius, self.rolloff, **kwargs
-        )
+    def __init__(
+        self,
+        freqs: ImageCoords,
+        radius: float = 0.95,
+        rolloff: float = 0.05,
+    ) -> None:
+        self.radius = radius
+        self.rolloff = rolloff
+        self.mask = _compute_circular_mask(freqs, self.radius, self.rolloff)
 
 
-def compute_circular_mask(
-    coords: ImageCoords,
-    cutoff: float = 0.95,
-    rolloff: float = 0.05,
-    **kwargs: Any,
+def _compute_circular_mask(
+    coords: ImageCoords, cutoff: float = 0.95, rolloff: float = 0.05
 ) -> RealImage:
     """
     Create a circular mask.
diff --git a/src/cryojax/simulator/noise.py b/src/cryojax/simulator/noise.py
index fc7cc02d..23a8bd6d 100644
--- a/src/cryojax/simulator/noise.py
+++ b/src/cryojax/simulator/noise.py
@@ -11,7 +11,7 @@
 from jaxtyping import PRNGKeyArray
 
 from .kernel import Kernel, Constant
-from ..utils import fftn
+from ..utils import rfftn
 from ..core import field, Module
 from ..typing import ImageCoords, ComplexImage
 
@@ -50,5 +50,5 @@ class GaussianNoise(Noise):
     @override
     def sample(self, key: PRNGKeyArray, freqs: ImageCoords) -> ComplexImage:
         spectrum = self.variance(freqs)
-        white_noise = fftn(jr.normal(key, shape=freqs.shape[0:-1]))
+        white_noise = jr.normal(key, shape=freqs.shape[0:-1])
         return spectrum * white_noise
diff --git a/src/cryojax/simulator/pipeline.py b/src/cryojax/simulator/pipeline.py
new file mode 100644
index 00000000..9a5dfa83
--- /dev/null
+++ b/src/cryojax/simulator/pipeline.py
@@ -0,0 +1,295 @@
+"""
+Image formation models.
+"""
+
+from __future__ import annotations
+
+__all__ = ["ImagePipeline", "SuperpositionPipeline"]
+
+from typing import Union, Optional
+from typing_extensions import override
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import jax.tree_util as jtu
+from jaxtyping import PRNGKeyArray
+
+from .filter import Filter
+from .mask import Mask
+from .ensemble import Ensemble
+from .scattering import ScatteringModel
+from .instrument import Instrument
+from .detector import NullDetector
+from .ice import Ice, NullIce
+from ..utils import rfftn, irfftn
+from ..core import field, Module
+from ..typing import ComplexImage, RealImage, Image
+
+
+class ImagePipeline(Module):
+    """
+    Base class for an imaging model.
+
+    Call an ``ImagePipeline``'s ``render`` and ``sample``,
+    routines.
+
+    Attributes
+    ----------
+    ensemble :
+        The ensemble from which to render images.
+    scattering :
+        The scattering model.
+    instrument :
+        The abstraction of the electron microscope.
+    solvent :
+        The solvent around the specimen.
+    filter :
+        A filter to apply to the image.
+    mask :
+        A mask to apply to the image.
+    """
+
+    ensemble: Ensemble = field()
+    scattering: ScatteringModel = field()
+    instrument: Instrument = field(default_factory=Instrument)
+    solvent: Ice = field(default_factory=NullIce)
+
+    filter: Optional[Filter] = field(default=None)
+    mask: Optional[Mask] = field(default=None)
+
+    def render(
+        self,
+        *,
+        view: bool = True,
+        get_real: bool = True,
+        normalize: bool = False,
+    ) -> Image:
+        """
+        Render an image of a Specimen.
+
+        Parameters
+        ----------
+        view : `bool`, optional
+            If ``True``, view the cropped, masked,
+            and filtered image.
+        get_real : `bool`, optional
+            If ``True``, return the image in real space.
+        normalize : `bool`, optional
+            If ``True``, normalize the image to mean zero
+            and standard deviation 1.
+        """
+        freqs = self.scattering.padded_frequency_grid_in_angstroms
+        # Draw the electron density at a particular conformation and pose
+        density = self.ensemble.realization
+        # Compute the scattering image in fourier space
+        image = self.scattering(density)
+        # Apply translation
+        image *= self.ensemble.pose.shifts(freqs)
+        # Measure the image with the instrument
+        # ... first compute and apply CTF
+        ctf = self.instrument.optics(freqs, pose=self.ensemble.pose)
+        image = ctf * image
+        # ... then apply the electron exposure model
+        scaling, offset = self.instrument.exposure.scaling(
+            freqs
+        ), self.instrument.exposure.offset(freqs)
+        image = scaling * image + offset
+
+        return self._postprocess_image(
+            image, view=view, get_real=get_real, normalize=normalize
+        )
+
+    def sample(
+        self,
+        key: PRNGKeyArray,
+        *,
+        view: bool = True,
+        get_real: bool = True,
+        normalize: bool = False,
+    ) -> RealImage:
+        """
+        Sample the an image from a realization of the noise.
+
+        Parameters
+        ----------
+        key :
+            The random number generator key.
+        view : `bool`, optional
+            If ``True``, view the cropped, masked,
+            and filtered image.
+        get_real : `bool`, optional
+            If ``True``, return the image in real space.
+        normalize : `bool`, optional
+            If ``True``, normalize the image to mean zero
+            and standard deviation 1.
+        """
+        idx = 0  # Keep track of number of stochastic models
+        if not isinstance(self.solvent, NullIce) and not isinstance(
+            self.instrument.detector, NullDetector
+        ):
+            keys = jax.random.split(key)
+        else:
+            keys = jnp.expand_dims(key, axis=0)
+        # Frequencies
+        freqs = self.scattering.padded_frequency_grid_in_angstroms
+        # The image of the specimen drawn from the ensemble
+        image = self.render(view=False, get_real=False)
+        if not isinstance(self.solvent, NullIce):
+            # The image of the solvent
+            ice_image = self.instrument.optics(freqs) * self.solvent.sample(
+                keys[idx], freqs
+            )
+            image = image + ice_image
+            idx += 1
+        if not isinstance(self.instrument.detector, NullDetector):
+            # Measure the detector readout
+            noise = self.instrument.detector.sample(keys[idx], freqs)
+            image = image + noise
+            idx += 1
+
+        return self._postprocess_image(
+            image, view=view, get_real=get_real, normalize=normalize
+        )
+
+    def __call__(
+        self,
+        key: Optional[Union[PRNGKeyArray, _PRNGKeyArrayLike]] = None,
+        *,
+        view: bool = True,
+        get_real: bool = True,
+        normalize: bool = False,
+    ) -> Image:
+        """
+        Sample or render an image.
+        """
+        if key is None:
+            return self.render(
+                view=view, get_real=get_real, normalize=normalize
+            )
+        else:
+            return self.sample(
+                key, view=view, get_real=get_real, normalize=normalize
+            )
+
+    def _postprocess_image(
+        self,
+        image: ComplexImage,
+        *,
+        view: bool = True,
+        get_real: bool = True,
+        normalize: bool = False,
+    ) -> Image:
+        """
+        Return an image postprocessed with filters, cropping, and masking
+        in either real or fourier space.
+        """
+        manager = self.scattering.manager
+        if view:
+            return self._filter_crop_mask(
+                image, get_real=get_real, normalize=normalize
+            )
+        else:
+            if normalize:
+                image = manager.normalize_image(
+                    image,
+                    is_real=False,
+                    shape_in_real_space=manager.padded_shape,
+                )
+            return irfftn(image, s=manager.padded_shape) if get_real else image
+
+    def _filter_crop_mask(
+        self,
+        image: ComplexImage,
+        *,
+        get_real: bool = True,
+        normalize: bool = False,
+    ) -> Image:
+        """
+        View the image. This function applies
+        filters, crops the image, then applies masks.
+        """
+        manager = self.scattering.manager
+        # Apply filter
+        if self.filter is not None:
+            image = self.filter(image)
+        # Crop and apply mask
+        if self.mask is None and manager.padded_shape == manager.shape:
+            # ... if there are no masks and we don't need to crop,
+            # minimize moving back and forth between real and fourier space
+            if normalize:
+                image = manager.normalize_image(
+                    image, is_real=False, shape_in_real_space=manager.shape
+                )
+            return irfftn(image, s=manager.shape) if get_real else image
+        else:
+            # ... otherwise, inverse transform, crop, mask, and normalize
+            image = manager.crop_to_shape(
+                irfftn(image, s=manager.padded_shape)
+            )
+            if self.mask is not None:
+                image = self.mask(image)
+            if normalize:
+                image = manager.normalize_image(image, is_real=True)
+            return image if get_real else rfftn(image)
+
+
+class SuperpositionPipeline(ImagePipeline):
+    """
+    Compute an image from a superposition of states in
+    ``Ensemble``. This assumes that either ``Ensemble.Pose``
+    and/or ``Ensemble.conformation`` has a batch dimension.
+
+    This class can be used to compute a micrograph, where there
+    are many specimen in the field of view. Or it can be used to
+    compute an image from ``Assembly.subunits``.
+    """
+
+    @override
+    def render(
+        self,
+        *,
+        view: bool = True,
+        get_real: bool = True,
+        normalize: bool = False,
+    ) -> Image:
+        """Render the superposition of states in the Ensemble."""
+        # Setup vmap over poses and conformations
+        false_pytree = jtu.tree_map(lambda _: False, self)
+        if self.ensemble.conformation is None:
+            # ... only vmap over poses
+            to_vmap_args = (
+                lambda m: (m.ensemble.pose,),
+                false_pytree,
+                (True,),
+            )
+        else:
+            # ... vmap over poses and conformations
+            to_vmap_args = (
+                lambda m: (m.ensemble.pose, m.ensemble.conformation),
+                false_pytree,
+                (True, True),
+            )
+        vmap, novmap = eqx.partition(self, eqx.tree_at(*to_vmap_args))
+        # Compute all images and sum
+        compute_image = lambda model: super(type(model), model).render(
+            view=False, get_real=False
+        )
+        # ... vmap to compute a stack of images to superimpose
+        compute_stack = jax.vmap(
+            lambda vmap, novmap: compute_image(eqx.combine(vmap, novmap)),
+            in_axes=(0, None),
+        )
+        # ... sum over the stack of images and jit
+        compute_stack_and_sum = jax.jit(
+            lambda vmap, novmap: jnp.sum(
+                compute_stack(vmap, novmap),
+                axis=0,
+            )
+        )
+        # ... compute the superposition
+        image = compute_stack_and_sum(vmap, novmap)
+
+        return self._postprocess_image(
+            image, view=view, get_real=get_real, normalize=normalize
+        )
diff --git a/src/cryojax/simulator/pose.py b/src/cryojax/simulator/pose.py
index 06fac697..746ee23c 100644
--- a/src/cryojax/simulator/pose.py
+++ b/src/cryojax/simulator/pose.py
@@ -14,7 +14,7 @@
 ]
 
 from abc import abstractmethod
-from typing import Union
+from typing import Union, Optional, Any
 from typing_extensions import override
 from jaxtyping import Float, Array
 from functools import cached_property
@@ -172,10 +172,16 @@ class QuaternionPose(Pose):
     @override
     def rotation(self) -> SO3:
         """Generate rotation from a unit quaternion."""
-        wxyz = jnp.array(
+        q = jnp.asarray(
             [self.view_qw, self.view_qx, self.view_qy, self.view_qz]
         )
-        R = SO3(wxyz=wxyz)
+        q_norm = jnp.sqrt(
+            self.view_qx**2
+            + self.view_qx**2
+            + self.view_qy**2
+            + self.view_qz**2
+        )
+        R = SO3(wxyz=q / q_norm)
         return R.inverse() if self.inverse else R
 
 
diff --git a/src/cryojax/simulator/scattering/_fourier_slice_extract.py b/src/cryojax/simulator/scattering/_fourier_slice_extract.py
index e3a1cdeb..fbe287cf 100644
--- a/src/cryojax/simulator/scattering/_fourier_slice_extract.py
+++ b/src/cryojax/simulator/scattering/_fourier_slice_extract.py
@@ -9,62 +9,50 @@
 from typing import Any
 
 import jax.numpy as jnp
+from jax.scipy.ndimage import map_coordinates
 
 from ._scattering_model import ScatteringModel
 from ..density import VoxelGrid
 from ...core import field
 from ...typing import (
-    Real_,
     ComplexImage,
     ComplexVolume,
     VolumeCoords,
 )
-from ...utils import (
-    ifftn,
-    fftn,
-    map_coordinates,
-)
 
 
 class FourierSliceExtract(ScatteringModel):
     """
     Scatter points to the image plane using the
     Fourier-projection slice theorem.
+
+    Attributes ``order``, ``mode``, and ``cval``
+    are passed to ``jax.scipy.map_coordinates``.
     """
 
     order: int = field(static=True, default=1)
     mode: str = field(static=True, default="wrap")
     cval: complex = field(static=True, default=0.0 + 0.0j)
 
-    def scatter(
-        self,
-        density: VoxelGrid,
-    ) -> ComplexImage:
+    def scatter(self, density: VoxelGrid) -> ComplexImage:
         """
         Compute an image by sampling a slice in the
         rotated fourier transform and interpolating onto
         a uniform grid in the object plane.
         """
-        fourier_projection = extract_slice(
+        return extract_slice(
             density.weights,
-            density.coordinates,
-            self.pixel_size,
+            density.frequency_slice,
             order=self.order,
             mode=self.mode,
             cval=self.cval,
         )
-        if self.manager.padded_shape != fourier_projection.shape:
-            fourier_projection = fftn(
-                self.manager.crop_or_pad(ifftn(fourier_projection).real)
-            )
-
-        return fourier_projection
 
 
 def extract_slice(
     weights: ComplexVolume,
-    coordinates: VolumeCoords,
-    pixel_size: Real_,
+    frequency_slice: VolumeCoords,
+    order: int = 1,
     **kwargs: Any,
 ) -> ComplexImage:
     """
@@ -73,36 +61,33 @@ def extract_slice(
 
     Arguments
     ---------
-    weights : shape `(N1, N2, N3)`
+    weights : shape `(N, N, N)`
         Density grid in fourier space.
-    coordinates : shape `(N1, N2, 1, 3)`
+    frequency_slice : shape `(N, N//2+1, 1, 3)`
         Frequency central slice coordinate system.
-    pixel_size :
-        The pixel_size of ``coordinates``.
-    kwargs:
-        Passed to ``cryojax.utils.interpolate.map_coordinates``.
+    order : int
+        Spline order of interpolation. By default, ``1``.
+    kwargs
+        Keyword arguments passed to ``jax.scipy.ndimage.map_coordinates``.
 
     Returns
     -------
-    projection :
+    projection : shape `(N, N//2+1)`
         The output image in fourier space.
     """
-    weights, coordinates = jnp.asarray(weights), jnp.asarray(coordinates)
     N1, N2, N3 = weights.shape
-    if not all([Ni == N1 for Ni in [N1, N2, N3]]):
-        raise ValueError("Only cubic boxes are supported for fourier slice.")
-    dx = pixel_size
-    box_size = jnp.array([N1 * dx, N2 * dx, N3 * dx])
-    # Need to convert to "array index coordinates".
-    # Make coordinates dimensionless
-    coordinates *= box_size
-    # Interpolate on the upper half plane get the slice
-    # z = N2 // 2 + 1
-    # fourier_projection = map_coordinates(
-    #    weights, coordinates[:, :z], **kwargs
-    # )[..., 0]
-    # Transform back to real space
-    # fourier_projection = irfftn(fourier_projection, s=(N1, N2))
-    #
+    N = N1
+    if (N1, N2, N3) != (N, N, N):
+        raise ValueError(
+            "Only cubic boxes are supported for fourier slice theorem."
+        )
+    # Need to convert to logical coordinates, so make coordinates dimensionless
+    grid_shape = jnp.asarray([N, N, N], dtype=float)
+    frequency_slice *= grid_shape
+    # Convert arguments to map_coordinates convention and compute
+    k_x, k_y, k_z = jnp.transpose(frequency_slice, axes=[3, 0, 1, 2])
+    projection = map_coordinates(weights, (k_x, k_y, k_z), order, **kwargs)[
+        :, :, 0
+    ]
 
-    return map_coordinates(weights, coordinates, **kwargs)[..., 0]
+    return projection
diff --git a/src/cryojax/simulator/scattering/_gaussian_mixture.py b/src/cryojax/simulator/scattering/_gaussian_mixture.py
index c216db79..e039257f 100644
--- a/src/cryojax/simulator/scattering/_gaussian_mixture.py
+++ b/src/cryojax/simulator/scattering/_gaussian_mixture.py
@@ -101,6 +101,8 @@ def _eval_Gaussian_kernel(sq_distances, atom_variances) -> ImageCoords:
 #          but identical coordinate systems. To make sure no messiness happens, I use
 #          the same functions from cryojax.utils.coordinates.
 # 2) For 3D coordinates I also load from cryojax.utils.coordinates.
+# Erik: Happy to change this, I just didn't know abouc cryojax.utils.coordinates could do this:
+# The call to fftfreq mislead me into thinking it only did Fourier-space.
 def _build_pixel_grid(
     npixels_per_side: int, pixel_size: Real_
 ) -> tuple[jnp.ndarray, jnp.ndarray]:
diff --git a/src/cryojax/simulator/scattering/_nufft_project.py b/src/cryojax/simulator/scattering/_nufft_project.py
index 6b887f19..06b27072 100644
--- a/src/cryojax/simulator/scattering/_nufft_project.py
+++ b/src/cryojax/simulator/scattering/_nufft_project.py
@@ -19,13 +19,11 @@
 from ._scattering_model import ScatteringModel
 from ...core import field
 from ...typing import (
-    Real_,
     ComplexImage,
     RealCloud,
     CloudCoords2D,
     CloudCoords3D,
 )
-from ...utils import nufft
 
 
 class NufftProject(ScatteringModel):
@@ -47,18 +45,16 @@ def scatter(self, density: Union[VoxelCloud, AtomCloud]) -> ComplexImage:
         if isinstance(density, VoxelCloud):
             fourier_projection = project_with_nufft(
                 density.weights,
-                density.coordinates,
-                self.pixel_size,
+                density.coordinate_list,
                 self.manager.padded_shape,
                 eps=self.eps,
             )
         elif isinstance(density, AtomCloud):
             fourier_projection = project_atoms_with_nufft(
                 density.weights,
-                density.coordinates,
+                density.coordinate_list,
                 density.variances,
                 density.identity,
-                self.pixel_size,
                 self.manager.padded_shape,
                 eps=self.eps,
             )
@@ -74,7 +70,6 @@ def project_atoms_with_nufft(
     coordinates,
     variances,
     identity,
-    pixel_size: Real_,
     shape: tuple[int, int],
     **kwargs: Any,
 ) -> ComplexImage:
@@ -87,18 +82,17 @@ def project_atoms_with_nufft(
         # kernel_i = atom_density_kernel[atom_type_i]
 
         # Build an
-        atom_i_image = project_with_nufft(
-            weights_i, coords_i, pixel_size, shape, **kwargs
-        )
+        atom_i_image = project_with_nufft(weights_i, coords_i, shape, **kwargs)
 
         # img += atom_i_image * kernel_i
         img += atom_i_image
 
+    return img
+
 
 def project_with_nufft(
     weights: RealCloud,
     coordinates: Union[CloudCoords2D, CloudCoords3D],
-    pixel_size: Real_,
     shape: tuple[int, int],
     **kwargs: Any,
 ) -> ComplexImage:
@@ -112,31 +106,36 @@ def project_with_nufft(
     ---------
     weights : shape `(N,)`
         Density point cloud.
-    coordinates : shape `(N, 3)`
+    coordinates : shape `(N, 2)` or shape `(N, 3)`
         Coordinate system of point cloud.
-    pixel_size :
-        The rasterization pixel_size.
     shape :
         Shape of the imaging plane in pixels.
         ``width, height = shape[0], shape[1]``
         is the size of the desired imaging plane.
-    kwargs:
-        Passed to ``cryojax.utils.integration.nufft``.
 
     Returns
     -------
     projection :
         The output image in fourier space.
     """
-    weights, coordinates = jnp.asarray(weights), jnp.asarray(coordinates)
+    from jax_finufft import nufft1
+
+    weights, coordinates = jnp.asarray(weights).astype(complex), jnp.asarray(
+        coordinates
+    )
+    # Flip and negate x and y to convert to cryojax conventions
+    coordinates = -jnp.flip(coordinates[:, :2], axis=-1)
+    # Normalize coordinates betweeen -pi and pi
     M1, M2 = shape
-    image_size = jnp.array(np.array([M1, M2]) * pixel_size)
-    coordinates = jnp.flip(coordinates[:, :2], axis=-1)
-    fourier_projection = nufft(
-        weights, coordinates, image_size, shape, **kwargs
+    image_size = jnp.asarray((M1, M2), dtype=float)
+    periodic_coords = 2 * jnp.pi * coordinates / image_size
+    # Compute and shift origin to cryojax conventions
+    x, y = periodic_coords.T
+    fourier_projection = jnp.fft.ifftshift(
+        nufft1(shape, weights, x, y, **kwargs)
     )
 
-    return fourier_projection
+    return fourier_projection[:, : M2 // 2 + 1]
 
 
 """
diff --git a/src/cryojax/simulator/scattering/_scattering_model.py b/src/cryojax/simulator/scattering/_scattering_model.py
index 3c392800..eb3e8af8 100644
--- a/src/cryojax/simulator/scattering/_scattering_model.py
+++ b/src/cryojax/simulator/scattering/_scattering_model.py
@@ -5,15 +5,22 @@
 
 from __future__ import annotations
 
-__all__ = ["ScatteringModel"]
+__all__ = ["ScatteringModel", "rescale_pixel_size"]
 
 from abc import abstractmethod
+from functools import partial, cached_property
+from typing import Any
 
-from ..density import ElectronDensity
+import jax
+import jax.numpy as jnp
+from jax.image import scale_and_translate
+
+from ..density import ElectronDensity, Voxels, VoxelGrid
 from ..manager import ImageManager
 
+from ...utils import rfftn, irfftn
 from ...core import field, Module
-from ...typing import Real_, ComplexImage
+from ...typing import Real_, RealImage, ComplexImage
 
 
 class ScatteringModel(Module):
@@ -29,8 +36,14 @@ class ScatteringModel(Module):
         Handles image configuration and
         utility routines.
     pixel_size :
-        Rasterization pixel size. This is in
-        dimensions of length.
+        The pixel size of the image in Angstroms.
+        For voxel-based ``ElectronDensity`` representations,
+        if the pixel size is different than the voxel size,
+        images will be interpolated in real space.
+    method :
+        The interpolation method used for measuring
+        the image at the new ``pixel_size``. Passed to
+        ``jax.image.scale_and_translate``.
 
     Methods
     -------
@@ -41,6 +54,8 @@ class ScatteringModel(Module):
     manager: ImageManager = field()
     pixel_size: Real_ = field()
 
+    method: str = field(static=True, default="bicubic")
+
     @abstractmethod
     def scatter(self, density: ElectronDensity) -> ComplexImage:
         """
@@ -53,3 +68,113 @@ def scatter(self, density: ElectronDensity) -> ComplexImage:
             The electron density representation.
         """
         raise NotImplementedError
+
+    def __call__(
+        self, density: ElectronDensity, **kwargs: Any
+    ) -> ComplexImage:
+        """
+        Compute an image at the exit plane, measured at the ScatteringModel
+        pixel size and post-processed with the ImageManager utilities.
+        """
+        image = self.scatter(density, **kwargs)
+        if isinstance(density, VoxelGrid):
+            # Resize the image to match the ImageManager.padded_shape
+            image = self.manager.crop_or_pad_to_padded_shape(
+                irfftn(image, s=density.weights.shape[0:2])
+            )
+        else:
+            # ... otherwise, assume the image is already at the padded_shape
+            image = irfftn(image, s=self.manager.padded_shape)
+        # Rescale the pixel size if different from the voxel size
+        if isinstance(density, Voxels):
+            current_pixel_size = density.voxel_size
+            new_pixel_size = self.pixel_size
+            rescale_fn = lambda image: rescale_pixel_size(
+                image,
+                current_pixel_size,
+                new_pixel_size,
+                method=self.method,
+            )
+            null_fn = lambda image: image
+            image = jax.lax.cond(
+                jnp.isclose(current_pixel_size, new_pixel_size),
+                null_fn,
+                rescale_fn,
+                image,
+            )
+        # Transform back to fourier space and give the image zero mean
+        image = rfftn(image).at[0, 0].set(0.0 + 0.0j)
+        return image
+
+    @cached_property
+    def coordinate_grid_in_angstroms(self):
+        return self.pixel_size * self.manager.coordinate_grid
+
+    @cached_property
+    def frequency_grid_in_angstroms(self):
+        return self.manager.frequency_grid / self.pixel_size
+
+    @cached_property
+    def padded_coordinate_grid_in_angstroms(self):
+        return self.pixel_size * self.manager.padded_coordinate_grid
+
+    @cached_property
+    def padded_frequency_grid_in_angstroms(self):
+        return self.manager.padded_frequency_grid / self.pixel_size
+
+
+@partial(jax.jit, static_argnames=["method", "antialias"])
+def rescale_pixel_size(
+    image: RealImage,
+    current_pixel_size: Real_,
+    new_pixel_size: Real_,
+    method: str = "bicubic",
+    antialias: bool = False,
+    **kwargs: Any,
+) -> RealImage:
+    """
+    Measure an image at a given pixel size using interpolation.
+
+    For more detail, see ``cryojax.utils.interpolation.scale``.
+
+    Parameters
+    ----------
+    image :
+        The image to be magnified.
+    current_pixel_size :
+        The pixel size of the input image.
+    new_pixel_size :
+        The new pixel size after interpolation.
+    method :
+        Interpolation method. See ``jax.image.scale_and_translate``
+        for documentation.
+    kwargs :
+        Keyword arguments passed to ``jax.image.scale_and_translate``.
+
+    Returns
+    -------
+    rescaled_image :
+        An image with pixels whose size are rescaled by
+        ``current_pixel_size / new_pixel_size``.
+    """
+    # Compute scale factor for pixel size rescaling
+    scale_factor = current_pixel_size / new_pixel_size
+    # Scaling in both dimensions is the same
+    scaling = jnp.asarray([scale_factor, scale_factor])
+    # Compute the translation in the jax.image convention that leaves
+    # cryojax images untranslated
+    N1, N2 = image.shape
+    translation = (1 - scaling) * jnp.array([N1 // 2, N2 // 2], dtype=float)
+    # Rescale pixel sizes
+    rescaled_image = scale_and_translate(
+        image,
+        image.shape,
+        (0, 1),
+        scaling,
+        translation,
+        method,
+        antialias=antialias,
+        **kwargs,
+    )
+
+    return rescaled_image
diff --git a/src/cryojax/utils/__init__.py b/src/cryojax/utils/__init__.py
index 2777c6ab..f1a6d16d 100644
--- a/src/cryojax/utils/__init__.py
+++ b/src/cryojax/utils/__init__.py
@@ -3,8 +3,6 @@
     coordinates,
     average,
     spectrum,
-    integrate,
-    interpolate,
     normalize,
     edges,
 )
@@ -13,8 +11,6 @@
 from .coordinates import *
 from .average import *
 from .spectrum import *
-from .integrate import *
-from .interpolate import *
 from .normalize import *
 from .edges import *
 
@@ -23,8 +19,6 @@
     + coordinates.__all__
     + average.__all__
     + spectrum.__all__
-    + integrate.__all__
-    + interpolate.__all__
     + normalize.__all__
     + edges.__all__
 )
@@ -35,8 +29,6 @@
         coordinates,
         average,
         spectrum,
-        integrate,
-        interpolate,
         normalize,
         edges,
     ]
diff --git a/src/cryojax/utils/average.py b/src/cryojax/utils/average.py
index dc6cafc3..7c9243dd 100644
--- a/src/cryojax/utils/average.py
+++ b/src/cryojax/utils/average.py
@@ -4,20 +4,21 @@
 
 __all__ = ["radial_average"]
 
-from typing import Optional, Union
+from typing import Optional, Union, Any
 
 import jax
 import jax.numpy as jnp
 
-from ..typing import RealVector, Vector, Image, ImageCoords
+from ..typing import RealVector, Vector, Image, RealImage
 
 
 @jax.jit
 def radial_average(
     image: Image,
-    norm: Image,
+    radial_grid: RealImage,
     bins: RealVector,
-    grid: Optional[ImageCoords] = None,
+    interpolating_radial_grid: Optional[RealImage] = None,
+    **kwargs: Any,
 ) -> Union[Image, Vector]:
     """
     Radially average vectors r with a given magnitude
@@ -27,15 +28,15 @@ def radial_average(
     ---------
     image :
         Two-dimensional image.
-    norm :
+    radial_grid :
         Radial coordinate system of image.
     bins :
         Radial bins for averaging. These
         must be evenly spaced.
-    grid :
+    interpolating_radial_grid :
         If ``None``, evalulate the spectrum as a 1D
         profile. Otherwise, evaluate the spectrum on this
-        2D grid of frequencies.
+        2D grid of frequencies using linear interpolation.
 
     Returns
     -------
@@ -47,13 +48,14 @@ def radial_average(
     dr = bins[1] - bins[0]
 
     def average(carry, r_i):
-        mask = jnp.logical_and(norm >= r_i, norm < r_i + dr)
+        mask = jnp.logical_and(radial_grid >= r_i, radial_grid < r_i + dr)
         return carry, jnp.sum(jnp.where(mask, image, 0.0)) / jnp.sum(mask)
 
     _, profile = jax.lax.scan(average, None, bins)
 
-    if grid is not None:
-        idxs = jnp.digitize(grid, bins).ravel()
-        return jax.vmap(lambda idx: profile[idx])(idxs).reshape(grid.shape)
+    if interpolating_radial_grid is not None:
+        return jnp.interp(
+            interpolating_radial_grid.ravel(), bins, profile, **kwargs
+        ).reshape(interpolating_radial_grid.shape)
     else:
         return profile
diff --git a/src/cryojax/utils/coordinates.py b/src/cryojax/utils/coordinates.py
index f491c1e5..109ed87e 100644
--- a/src/cryojax/utils/coordinates.py
+++ b/src/cryojax/utils/coordinates.py
@@ -6,92 +6,82 @@
     "make_coordinates",
     "make_frequencies",
     "cartesian_to_polar",
-    "flatten_and_coordinatize",
 ]
 
-from typing import Union, Any
+from typing import Union, Any, Optional
 from jaxtyping import Array, Float
 
 import jax.numpy as jnp
-import numpy as np
 
 from ..typing import RealImage, RealVolume, Image, ImageCoords
 
 
-def make_coordinates(*args: Any, **kwargs: Any) -> Float[Array, "... D"]:
+def make_coordinates(
+    shape: tuple[int, ...], grid_spacing: float = 1.0, indexing: str = "xy"
+) -> Float[Array, "*shape len(shape)"]:
     """
-    Wraps ``_fftfreqs`` for ``real = True``.
-    """
-    return _fftfreqs(*args, real=True, **kwargs)
+    Create a real-space cartesian coordinate system on a grid.
 
+    Arguments
+    ---------
+    shape :
+        Shape of the voxel grid, with
+        ``ndim = len(shape)``.
+    grid_spacing :
+        The grid spacing, in units of length.
+    indexing :
+        Either ``"xy"`` or ``"ij"``, passed to
+        ``jax.numpy.meshgrid``.
 
-def make_frequencies(*args: Any, **kwargs: Any) -> Float[Array, "... D"]:
-    """
-    Wraps ``_fftfreqs`` for ``real = False``.
+    Returns
+    -------
+    coordinate_grid :
+        Cartesian coordinate system in real space.
     """
-    return _fftfreqs(*args, real=False, **kwargs)
+    coordinate_grid = _make_coordinates_or_frequencies(
+        shape, grid_spacing=grid_spacing, real_space=True, indexing=indexing
+    )
+    return coordinate_grid
 
 
-def flatten_and_coordinatize(
-    template: Union[RealVolume, RealImage],
-    resolution: float,
-    mask: bool = True,
-    indexing="xy",
-    **kwargs: Any,
-) -> tuple[Float[Array, "N"], Float[Array, "N ndim"]]:
+def make_frequencies(
+    shape: tuple[int, ...],
+    grid_spacing: float = 1.0,
+    half_space: bool = True,
+    indexing: str = "xy",
+) -> Float[Array, "*shape len(shape)"]:
     """
-    Returns 3D volume or 2D image and its coordinate system.
-    By default, coordinates are shape ``(N, ndim)``, where
-    ``ndim = template.ndim`` and ``N = N1*N2*N3 - M`` or
-    ``N = N2*N3 - M``, where ``M`` is a number of points
-    close to zero that are masked out. The coordinate system
-    is set with dimensions of length with zero in the center.
-
-    Parameters
-    ----------
-    template : `np.ndarray`, shape `(N1, N2, N3)` or `(N1, N2)`
-        3D volume or 2D image on a cartesian grid.
-    resolution : float
-        Resolution of the template.
-    mask : `bool`
-        If ``True``, run template through ``numpy.isclose``
-        to remove coordinates with zero electron density.
-    kwargs
-        Keyword arguments passed to ``numpy.isclose``.
-        Disabled for ``mask = False``.
+    Create a fourier-space cartesian coordinate system on a grid.
+    The zero-frequency component is in the beginning.
+
+    Arguments
+    ---------
+    shape :
+        Shape of the voxel grid, with
+        ``ndim = len(shape)``.
+    grid_spacing :
+        The grid spacing, in units of length.
+    half_space :
+        Return a frequency grid on the half space.
+        ``shape[-1]`` is the axis on which the negative
+        frequencies are omitted.
+    indexing :
+        Either ``"xy"`` or ``"ij"``, passed to
+        ``jax.numpy.meshgrid``.
 
     Returns
     -------
-    density : `Array`, shape `(N, ndim)`
-        Volume or image.
-    coords : `Array`, shape `(N, ndim)`
-        Cartesian coordinate system.
+    frequency_grid :
+        Cartesian coordinate system in frequency space.
     """
-    template = jnp.asarray(template)
-    ndim, shape = template.ndim, template.shape
-    # Mask out points where the electron density is close
-    # to zero.
-    flat = template.ravel()
-    if mask:
-        nonzero = np.where(~np.isclose(flat, 0.0, **kwargs))
-        density = flat[nonzero]
-    else:
-        nonzero = True
-        density = flat
-
-    # Create coordinate buffer
-    N = density.size
-    coords = np.zeros((N, ndim))
-
-    # Generate rectangular grid and fill coordinate array.
-    R = _fftfreqs(shape, resolution, real=True, indexing=indexing)
-    for i in range(ndim):
-        if mask:
-            coords[..., i] = R[..., i].ravel()[nonzero]
-        else:
-            coords[..., i] = R[..., i].ravel()
-
-    return jnp.array(density), jnp.array(coords)
+    frequency_grid = _make_coordinates_or_frequencies(
+        shape,
+        grid_spacing=grid_spacing,
+        real_space=False,
+        half_space=half_space,
+        indexing=indexing,
+    )
+    return frequency_grid
 
 
 def cartesian_to_polar(
@@ -118,59 +108,57 @@ def cartesian_to_polar(
         return kr, theta
 
 
-def _fftfreqs(
+def _make_coordinates_or_frequencies(
     shape: tuple[int, ...],
-    pixel_size: Union[float, np.ndarray] = 1.0,
-    real: bool = False,
+    grid_spacing: float = 1.0,
+    real_space: bool = False,
+    half_space: bool = True,
     indexing: str = "xy",
-) -> Float[Array, "... D"]:
-    """
-    Create a cartesian coordinate system on a grid.
-    This can be used for real and fourier space.
-    If used for fourier space, the
-    zero-frequency component is in the beginning.
-
-    Note that the current implementations create coordinates in
-    ``numpy``, then transfer to ``jax``. Therefore, these should
-    not be run inside jit compilation.
-
-    Arguments
-    ---------
-    shape :
-        Shape of the voxel grid, with
-        ``ndim = len(shape)``.
-    pixel_size :
-        Image pixel size.
-    real :
-        Choose whether to create coordinate system
-        in real or fourier space.
-
-
-    Returns
-    -------
-    coords :
-        2D or 3D cartesian coordinate system.
-    """
+) -> Float[Array, "*shape len(shape)"]:
     ndim = len(shape)
     shape = (*shape[:2][::-1], *shape[2:]) if indexing == "xy" else shape
-    if isinstance(pixel_size, (np.floating, float)):
-        pixel_size = ndim * [pixel_size]
+    coords1D = []
+    for idx in range(ndim):
+        if real_space:
+            c1D = _make_coordinates_or_frequencies_1d(
+                shape[idx], grid_spacing, real_space
+            )
+        else:
+            if not half_space:
+                rfftfreq = False
+            else:
+                if indexing == "xy" and ndim == 2:
+                    rfftfreq = True if idx == 0 else False
+                else:
+                    rfftfreq = False if idx < ndim - 1 else True
+            c1D = _make_coordinates_or_frequencies_1d(
+                shape[idx], grid_spacing, real_space, rfftfreq
+            )
+        coords1D.append(c1D)
+    coords = jnp.stack(jnp.meshgrid(*coords1D, indexing=indexing), axis=-1)
+
+    return coords
+
+
+def _make_coordinates_or_frequencies_1d(
+    size: int,
+    grid_spacing: float,
+    real_space: bool = False,
+    rfftfreq: Optional[bool] = None,
+) -> Float[Array, "size"]:
+    """One-dimensional coordinates in real or fourier space"""
+    if real_space:
+        make_1d = (
+            lambda size, dx: jnp.fft.fftshift(jnp.fft.fftfreq(size, 1 / dx))
+            * size
+        )
     else:
-        pixel_size = list(pixel_size)
-    assert len(pixel_size) == ndim
-    coords1D = [
-        _fftfreqs1d(shape[idx], pixel_size[idx], real) for idx in range(ndim)
-    ]
-    coords = np.stack(np.meshgrid(*coords1D, indexing=indexing), axis=-1)
-
-    return jnp.asarray(coords)
-
+        if rfftfreq is None:
+            raise ValueError(
+                "Argument rfftfreq cannot be None if real_space=False."
+            )
+        else:
+            fn = jnp.fft.rfftfreq if rfftfreq else jnp.fft.fftfreq
+            make_1d = lambda size, dx: fn(size, grid_spacing)
 
-def _fftfreqs1d(s: int, pixel_size: float, real: bool = False) -> np.ndarray:
-    """One-dimensional coordinates in real or fourier space"""
-    fftfreq = (
-        lambda s: np.fft.fftshift(np.fft.fftfreq(s, 1 / pixel_size)) * s
-        if real
-        else np.fft.fftfreq(s, pixel_size)
-    )
-    return fftfreq(s)
+    return make_1d(size, grid_spacing)
diff --git a/src/cryojax/utils/edges.py b/src/cryojax/utils/edges.py
index 5b24e134..b588db9d 100644
--- a/src/cryojax/utils/edges.py
+++ b/src/cryojax/utils/edges.py
@@ -2,7 +2,7 @@
 Routines for dealing with image edges.
 """
 
-__all__ = ["bound", "crop", "pad", "crop_or_pad"]
+__all__ = ["crop", "pad", "crop_or_pad"]
 
 from typing import Union
 from jaxtyping import Array, Float
@@ -13,26 +13,6 @@
 from ..typing import Cloud, CloudCoords2D, Image, Volume
 
 
-@jax.jit
-def bound(
-    density: Cloud, coords: CloudCoords2D, box_size: Float[Array, "2"]
-) -> Cloud:
-    """
-    Use a boolean mask to set density values out of
-    bounds to zero. The mask is ``True`` for
-    all points outside of the box_size, and False otherwise.
-    """
-    coords, box_size = jnp.asarray(coords), jnp.asarray(box_size)
-    x, y = coords.T
-    Lx, Ly = box_size[0], box_size[1]
-    x_mask = jnp.logical_or(x < -Lx / 2, x >= Lx / 2)
-    y_mask = jnp.logical_or(y < -Ly / 2, y >= Ly / 2)
-    mask = jnp.logical_or(x_mask, y_mask)
-    masked_density = jnp.where(mask, complex(0.0), density)
-
-    return masked_density
-
-
 def crop(image: Image, shape: tuple[int, int]) -> Image:
     """
     Crop an image to a new shape.
diff --git a/src/cryojax/utils/fft.py b/src/cryojax/utils/fft.py
index 819d5cc0..f97ccff5 100644
--- a/src/cryojax/utils/fft.py
+++ b/src/cryojax/utils/fft.py
@@ -1,5 +1,5 @@
 """
-Routines to compute FFTs.
+Routines to compute FFTs, in cryojax conventions.
 """
 
 __all__ = ["ifftn", "irfftn", "fftn", "rfftn"]
@@ -8,17 +8,14 @@
 
 import jax.numpy as jnp
 
-from ..typing import RealImage, RealVolume, Image, Volume
-
-
-def irfftn(
-    ft: Union[Image, Volume], **kwargs: Any
-) -> Union[RealImage, RealVolume]:
-    """
-    Convenience wrapper for ``cryojax.utils.fft.ifftn``
-    for ``real = True``.
-    """
-    return ifftn(ft, real=True, **kwargs)
+from ..typing import (
+    RealImage,
+    RealVolume,
+    ComplexImage,
+    ComplexVolume,
+    Image,
+    Volume,
+)
 
 
 def ifftn(
@@ -49,7 +46,9 @@ def ifftn(
         return ift
 
 
-def fftn(ift: Union[Image, Volume], **kwargs: Any) -> Union[Image, Volume]:
+def fftn(
+    ift: Union[Image, Volume], **kwargs: Any
+) -> Union[ComplexImage, ComplexVolume]:
     """
     Helper routine to match the fourier transform of an array
     with the output of a type 1 non-uniform FFT with jax-finufft.
@@ -70,10 +69,10 @@ def fftn(ift: Union[Image, Volume], **kwargs: Any) -> Union[Image, Volume]:
 
 
 def irfftn(
-    ft: Union[Image, Volume],
+    ft: Union[ComplexImage, ComplexVolume],
     s: Optional[tuple[int, ...]] = None,
     **kwargs: Any
-) -> Union[Image, Volume]:
+) -> Union[RealImage, RealVolume]:
     """
     Helper routine to compute the inverse fourier transform
     from real input.
@@ -88,14 +87,14 @@ def irfftn(
     ift :
         Inverse fourier transform.
     """
-    ift = jnp.fft.fftshift(
-        jnp.fft.irfftn(ft, s=s, **kwargs), axes=tuple(range(ift.dim - 1))
-    )
+    ift = jnp.fft.fftshift(jnp.fft.irfftn(ft, s=s, **kwargs))
 
     return ift
 
 
-def rfftn(ift: Union[Image, Volume], **kwargs: Any) -> Union[Image, Volume]:
+def rfftn(
+    ift: Union[RealImage, RealVolume], **kwargs: Any
+) -> Union[ComplexImage, ComplexVolume]:
     """
     Helper routine to compute the fourier transform of an array
     to match the output of a type 1 non-uniform FFT.
@@ -110,8 +109,6 @@ def rfftn(ift: Union[Image, Volume], **kwargs: Any) -> Union[Image, Volume]:
     ft :
         Fourier transform of array.
     """
-    ft = jnp.fft.rfftn(
-        jnp.fft.ifftshift(ift, axes=tuple(range(ift.dim - 1))), **kwargs
-    )
+    ft = jnp.fft.rfftn(jnp.fft.ifftshift(ift), **kwargs)
 
     return ft
diff --git a/src/cryojax/utils/integrate.py b/src/cryojax/utils/integrate.py
deleted file mode 100644
index 9782e878..00000000
--- a/src/cryojax/utils/integrate.py
+++ /dev/null
@@ -1,158 +0,0 @@
-"""
-Integration routines.
-"""
-
-__all__ = ["nufft", "integrate_gaussians"]
-
-from typing import Any, Union
-from jaxtyping import Array, Float
-
-import jax
-import jax.numpy as jnp
-from jax.scipy import special
-
-from .coordinates import _fftfreqs1d
-
-from ..typing import (
-    RealVector,
-    RealCloud,
-    CloudCoords2D,
-    ComplexImage,
-    RealImage,
-)
-
-
-def nufft(
-    density: RealCloud,
-    coords: CloudCoords2D,
-    box_size: Union[Float[Array, ""], Float[Array, "2"]],
-    shape: tuple[int, int],
-    eps: float = 1e-6,
-    **kwargs: Any
-) -> ComplexImage:
-    r"""
-    Helper routine to compute a non-uniform FFT on the
-    imaging plane for a 3D point cloud.
-
-    .. warning::
-        If any values in ``2 * np.pi * coords / box_size``
-        lies out of bounds of :math:`$(-3\pi, 3\pi]$`,
-        this method will crash.
-
-    Arguments
-    ---------
-    density :
-        Density point cloud over which to compute
-        the fourier transform.
-    coords :
-        Coordinate system for density cloud.
-    box_size :
-        2D imaging plane that ``coords`` lies in.
-    shape :
-        Desired output shape of the transform.
-    eps :
-        Precision of the non-uniform FFT. See
-        `finufft <https://finufft.readthedocs.io/en/latest/>`_
-        for more detail.
-    kwargs :
-        Keyword arguments passed to ``jax_finufft.nufft1``.
-    Return
-    ------
-    ft : Fourier transform.
-    """
-    from jax_finufft import nufft1
-
-    coords = jnp.asarray(coords)
-    complex_density = jnp.asarray(density.astype(complex))
-    periodic_coords = 2 * jnp.pi * coords / box_size
-
-    x, y = periodic_coords.T
-    ft = nufft1(shape, complex_density, -x, -y, eps=eps, **kwargs)
-
-    return jnp.fft.ifftshift(ft)
-
-
-def integrate_gaussians(
-    weights: RealCloud,
-    centers: CloudCoords2D,
-    variances: RealCloud,
-    shape: tuple[int, int],
-    pixel_size: float,
-) -> RealImage:
-    """
-    Integrate a sum of Gaussians over a grid given by the given shape.
-
-    Computes the integral of a function which is represented as a weighted sum of Gaussians
-    with the given centers and scales. The coordinate system corresponds to the grid with
-    unit spacing and with the origin at the center of the grid.
-
-    Parameters
-    ----------
-    weights :
-        Weights of Gaussian densities.
-    centers :
-        Centers of Gaussian densities.
-    variances :
-        Variances of Gaussian densities.
-    shape :
-        Shape of grid to integrate over.
-        The number of dimensions is inferred from the length of this sequence.
-    pixel_size :
-        Pixel size.
-
-    Returns
-    -------
-    image : `Array`, shape `shape`
-        Integrals of Gaussian densities over grid.
-    """
-    x, y = [_fftfreqs1d(s + 1, pixel_size, real=True) for s in shape]
-    image = _integrate_gaussians(x, y, weights, centers, variances)
-
-    return image
-
-
-@jax.jit
-def _integrate_gaussians(
-    x: RealVector,
-    y: RealVector,
-    weights: RealCloud,
-    centers: RealCloud,
-    scales: RealCloud,
-) -> RealImage:
-    """
-    Integrate a Gaussian density over a set of intervals given by boundaries.
-
-    Parameters
-    ----------
-    x : shape `(M1 + 1,)`
-        x boundary of intervals to integrate over.
-    y : shape `(M2 + 1,)`
-        y boundary of intervals to integrate over.
-    weights :
-        Gaussian weights
-    centers :
-        Centers of Gaussian densities.
-    scales :
-        Scales of Gaussian densities.
-
-    Returns
-    -------
-    I : shape `(M1, M2)`
-        Integrals of Gaussian densities over intervals.
-    """
-    centers = jnp.expand_dims(centers, axis=-1)
-    scales = jnp.expand_dims(scales, axis=-1)
-
-    # Marginals
-    I_x = 0.5 * jnp.diff(
-        special.erf((x - centers[:, 0]) / (scales * jnp.sqrt(2.0))), axis=-1
-    )
-    I_y = 0.5 * jnp.diff(
-        special.erf((y - centers[:, 1]) / (scales * jnp.sqrt(2.0))), axis=-1
-    )
-    # Contract by summing over batch dimension and outer product over spatial dimensions
-    # In summation notation, this would correspond to the einsum
-    # jnp.einsum('mi,mj,m->ij', *marginals, weights)
-    I = jnp.matmul((I_x * jnp.expand_dims(weights, axis=-1)).T, I_y)
-
-    return I
diff --git a/src/cryojax/utils/interpolate.py b/src/cryojax/utils/interpolate.py
deleted file mode 100644
index eeb89070..00000000
--- a/src/cryojax/utils/interpolate.py
+++ /dev/null
@@ -1,123 +0,0 @@
-"""
-Interpolation routines.
-"""
-
-__all__ = ["resize", "scale", "scale_and_translate", "map_coordinates"]
-
-from typing import Union, Any
-from jaxtyping import Array, Float
-
-import jax
-import jax.numpy as jnp
-from jax.scipy.ndimage import map_coordinates as _map_coordinates
-from jax.image import resize as _resize
-from jax.image import scale_and_translate as _scale_and_translate
-
-
-def resize(
-    image: Array,
-    shape: tuple[int, int],
-    method="lanczos5",
-    align_corners: bool = True,
-    **kwargs: Any
-) -> Array:
-    """
-    Resize an image with interpolation.
-
-    Wraps ``jax.image.resize``.
-    """
-    if align_corners:
-        return _resize_with_aligned_corners(image, shape, method, **kwargs)
-    else:
-        return _resize(image, shape, method, **kwargs)
-
-
-def scale_and_translate(
-    image: Array,
-    shape: tuple[int, int],
-    scale: Float[Array, "2"],
-    translation: Float[Array, "2"],
-    method="lanczos5",
-    **kwargs
-) -> Array:
-    """
-    Resize, scale, and translate an image with interpolation.
-
-    Wraps ``jax.image.scale_and_translate``.
-    """
-    image = jnp.asarray(image)
-    spatial_dims = (0, 1)
-    N1, N2 = image.shape
-    translation += (1 - scale) * jnp.array([N1 // 2, N2 // 2], dtype=float)
-    return _scale_and_translate(
-        image, shape, spatial_dims, scale, translation, method, **kwargs
-    )
-
-
-def scale(
-    image: Array,
-    shape: tuple[int, int],
-    scale: Float[Array, "2"],
-    method="lanczos5",
-    **kwargs
-) -> Array:
-    """
-    Resize and scale an image with interpolation.
-
-    Wraps ``jax.image.scale_and_translate``.
-    """
-    translation = jnp.array([0.0, 0.0])
-    return scale_and_translate(
-        image, shape, scale, translation, method=method, **kwargs
-    )
-
-
-def map_coordinates(
-    input: Array, coordinates: Array, order=1, mode="wrap", cval=0.0
-) -> Array:
-    """
-    Interpolate a set of points in fourier space on a grid
-    with a given coordinate system onto a new coordinate system.
-    """
-    input, coordinates = jnp.asarray(input), jnp.asarray(coordinates)
-    N1, N2, N3 = input.shape
-    box_shape = jnp.array([N1, N2, N3], dtype=float)[:, None, None, None]
-    coordinates = jnp.transpose(coordinates, axes=[3, 0, 1, 2])
-    # Flip negative valued frequencies to get the logical coordinates.
-    coordinates = jnp.where(
-        coordinates < 0, box_shape + coordinates, coordinates
-    )
-    return _map_coordinates(input, coordinates, order, mode=mode, cval=cval)
-
-
-def _resize_with_aligned_corners(
-    image: Array,
-    shape: tuple[int, ...],
-    method: Union[str, jax.image.ResizeMethod],
-    antialias: bool = False,
-) -> Array:
-    """
-    Alternative to jax.image.resize(), which emulates
-    align_corners=True in PyTorch's interpolation functions.
-
-    Adapted from https://github.com/google/jax/issues/11206.
-    ."""
-    image = jnp.asarray(image)
-    spatial_dims = tuple(
-        i
-        for i in range(len(shape))
-        if not jax.core.symbolic_equal_dim(image.shape[i], shape[i])
-    )
-    scale = jnp.array(
-        [(shape[i] - 1.0) / (image.shape[i] - 1.0) for i in spatial_dims]
-    )
-    translation = -(scale / 2.0 - 0.5)
-    return _scale_and_translate(
-        image,
-        shape,
-        method=method,
-        scale=scale,
-        spatial_dims=spatial_dims,
-        translation=translation,
-        antialias=antialias,
-    )
diff --git a/src/cryojax/utils/normalize.py b/src/cryojax/utils/normalize.py
index 8dd1803a..6836c58d 100644
--- a/src/cryojax/utils/normalize.py
+++ b/src/cryojax/utils/normalize.py
@@ -4,23 +4,42 @@
 
 __all__ = ["rescale_image", "normalize_image"]
 
-from typing import Any
+from typing import Any, Optional
 
 import jax.numpy as jnp
 
 from ..typing import Image
 
 
-def normalize_image(image: Image, **kwargs: Any) -> Image:
+def normalize_image(
+    image: Image,
+    *,
+    is_real: bool = True,
+    half_space: bool = True,
+    shape_in_real_space: Optional[tuple[int, int]] = None,
+) -> Image:
     """
     Normalize so that the image is mean 0
     and standard deviation 1 in real space.
     """
-    return rescale_image(image, 1.0, 0.0, **kwargs)
+    return rescale_image(
+        image,
+        1.0,
+        0.0,
+        is_real=is_real,
+        half_space=half_space,
+        shape_in_real_space=shape_in_real_space,
+    )
 
 
 def rescale_image(
-    image: Image, N: float, mu: float, *, real: bool = True
+    image: Image,
+    std: float,
+    mean: float,
+    *,
+    is_real: bool = True,
+    half_space: bool = True,
+    shape_in_real_space: Optional[tuple[int, int]] = None,
 ) -> Image:
     """
     Normalize so that the image is mean mu
@@ -28,32 +47,49 @@ def rescale_image(
 
     Parameters
     ----------
-    image : `jax.Array`, shape `(N1, N2)`
+    image : `Array`, shape `(N1, N2)`
         The image in either real or Fourier space.
         If in Fourier space, the zero frequency
-        component should be in the center of the image.
-    N : `float`
-        Intensity scale factor.
-    mu : `float`
+        component should be in the corner of the image.
+        It should also be on the half space.
+    std : `float`
+        Intensity standard deviation.
+    mean : `float`
         Intensity offset.
-    real : `bool`
+    is_real : `bool`
         If ``True``, the given ``image`` is in real
         space. If ``False``, it is in Fourier space.
 
     Returns
     -------
-    rescaled_image : `jax.Array`, shape `(N1, N2)`
+    rescaled_image : `Array`, shape `(N1, N2)`
         Image rescaled to have mean ``mu`` and standard
         deviation ``N``.
     """
     image = jnp.asarray(image)
     N1, N2 = image.shape
     # First normalize image to zero mean and unit standard deviation
-    if real:
+    if is_real:
         normalized_image = (image - image.mean()) / image.std()
-        rescaled_image = N * normalized_image + mu
+        rescaled_image = std * normalized_image + mean
     else:
+        if half_space:
+            N_modes = (
+                N1 * (2 * N2 - 1)
+                if shape_in_real_space is None
+                else shape_in_real_space[0] * shape_in_real_space[1]
+            )
+        else:
+            N_modes = N1 * N2
         normalized_image = image.at[0, 0].set(0.0)
-        normalized_image /= jnp.linalg.norm(normalized_image) / (N1 * N2)
-        rescaled_image = (normalized_image * N).at[0, 0].set(mu * N1 * N2)
+        normalizing_factor = (
+            jnp.sqrt(
+                jnp.sum(jnp.abs(normalized_image[:, 0]) ** 2)
+                + 2 * jnp.sum(jnp.abs(normalized_image[:, 1:]) ** 2)
+            )
+            if half_space
+            else jnp.linalg.norm(normalize_image)
+        )
+        normalized_image /= normalizing_factor / N_modes
+        rescaled_image = (normalized_image * std).at[0, 0].set(mean * N_modes)
     return rescaled_image
diff --git a/src/cryojax/utils/spectrum.py b/src/cryojax/utils/spectrum.py
index 008be95f..e8f70d9b 100644
--- a/src/cryojax/utils/spectrum.py
+++ b/src/cryojax/utils/spectrum.py
@@ -4,59 +4,62 @@
 
 __all__ = ["powerspectrum"]
 
-from typing import Optional, Union
+from typing import Optional, Union, overload
 
 import jax.numpy as jnp
 
 from .average import radial_average
-from ..typing import RealVector, RealImage, ComplexImage, ImageCoords
+from ..typing import Real_, RealVector, RealImage, ComplexImage
 
 
 def powerspectrum(
-    image: ComplexImage,
-    freqs: ImageCoords,
-    pixel_size: float = 1.0,
-    k_min: Optional[float] = None,
-    k_max: Optional[float] = None,
-    grid: Optional[ImageCoords] = None,
+    fourier_image: ComplexImage,
+    radial_frequency_grid: RealImage,
+    pixel_size: Real_ = 1.0,
+    *,
+    k_min: Optional[Real_] = None,
+    k_max: Optional[Real_] = None,
+    interpolating_radial_frequency_grid: Optional[RealImage] = None,
 ) -> tuple[Union[RealImage, RealVector], RealVector]:
     """
     Compute the power spectrum of an image averaged on a set
     of radial bins. This does not compute the zero mode of
     the spectrum.
 
-    Arguments
-    ---------
-    image :
+    Parameters
+    ----------
+    fourier_image :
         An image in Fourier space.
-    freqs :
+    radial_frequency_grid :
         The frequency range of the desired wavevectors.
     pixel_size :
-        The pixel size of the frequency grid.
-    grid :
+        The pixel size of the radial frequency grid.
+    k_min :
+        Minimum wavenumber bin. By default, ``0.0``.
+    k_max :
+        Maximum wavenumber bin. By default, ``1 / (2 * pixel_size)``.
+    interpolating_radial_frequency_grid :
         If ``None``, evalulate the spectrum as a 1D
         profile. Otherwise, evaluate the spectrum on this
         2D grid of frequencies.
     Returns
     -------
-    spectrum : shape `(max(L1, L2) // 2,)` or `(L1, L2)`
-        Power spectrum up to the Nyquist frequency. ``(L1, L2)`` can
-        be ``(M1, M2)`` or ``(N1, N2)``.
-    k_bins : shape `(max(L1, L2) // 2,)`
-        Radial bins for averaging. The minimum wavenumber
-        and wavenumber spacing is computed as ``1 / max(L1, L2)``.
-        ``(L1, L2)`` can be ``(M1, M2)`` or ``(N1, N2)``.
+    spectrum :
+        Power spectrum up to the Nyquist frequency.
+    bins :
+        Radial wavenumber bins for averaging.
     """
-    power = (image * jnp.conjugate(image)).real
-    k_norm = jnp.linalg.norm(freqs, axis=-1)
-    k_max = 1.0 / (2.0 * pixel_size)
-    if grid is None:
-        q_norm = None
-        k_min = 1.0 / (pixel_size * max(*k_norm.shape))
-    else:
-        q_norm = jnp.linalg.norm(grid, axis=-1)
-        k_min = 1.0 / (pixel_size * max(*q_norm.shape))
-    k_bins = jnp.arange(k_min, k_max, k_min)  # Left edges of bins
-    spectrum = radial_average(power, k_norm, k_bins, grid=q_norm)
-
-    return spectrum, k_bins
+    # Compute power
+    power = (fourier_image * jnp.conjugate(fourier_image)).real
+    # Compute bins
+    k_min = 0.0 if k_min is None else k_min
+    k_max = 1.0 / (pixel_size * 2.0) if k_max is None else k_max
+    k_step = 1.0 / (pixel_size * max(*power.shape))
+    bins = jnp.arange(k_min, k_max, k_step)  # Left edges of bins
+    # Compute radially averaged power spectrum as a 1D profile or
+    # interpolated onto a 2D grid
+    spectrum = radial_average(
+        power, radial_frequency_grid, bins, interpolating_radial_frequency_grid
+    )
+
+    return spectrum, bins
diff --git a/tutorials/test-helix.ipynb b/tutorials/simulate-helix.ipynb
similarity index 100%
rename from tutorials/test-helix.ipynb
rename to tutorials/simulate-helix.ipynb
diff --git a/tutorials/test-pipeline.ipynb b/tutorials/simulate-image.ipynb
similarity index 56%
rename from tutorials/test-pipeline.ipynb
rename to tutorials/simulate-image.ipynb
index 7b327fc6..b269d63a 100644
--- a/tutorials/test-pipeline.ipynb
+++ b/tutorials/simulate-image.ipynb
@@ -58,7 +58,7 @@
    "source": [
     "# Image simulator imports\n",
     "import cryojax.simulator as cs\n",
-    "from cryojax.utils import fftn, ifftn, make_frequencies"
+    "from cryojax.utils import rfftn, irfftn, make_frequencies"
    ]
   },
   {
@@ -75,7 +75,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "First we must read in our template into a given volume representation of the ``ensemble``. We also must choose its ``Pose``."
+    "First we must read in our template into a given volume representation of the ``Ensemble``. We also must choose its ``Pose``."
    ]
   },
   {
@@ -85,9 +85,9 @@
    "outputs": [],
    "source": [
     "# Read template into either an ElectronCloud or ElectronGrid\n",
-    "#density = cs.VoxelCloud.from_file(filename, config=dict(atol=1e-8))\n",
-    "density = cs.VoxelGrid.from_file(filename, config=dict(pad_scale=1.5))\n",
-    "pose = cs.EulerPose(offset_x=-50.0, offset_y=20.0, view_phi=-20, view_theta=100, view_psi=20)\n",
+    "#density = cs.VoxelCloud.from_file(filename)\n",
+    "density = cs.VoxelGrid.from_file(filename, pad_scale=1.5)\n",
+    "pose = cs.EulerPose(offset_x=-50.0, offset_y=20.0, view_phi=-20., view_theta=100., view_psi=20.)\n",
     "ensemble = cs.Ensemble(density=density, pose=pose)"
    ]
   },
@@ -107,9 +107,9 @@
     "# Configure the image formation process\n",
     "shape = (81, 82)\n",
     "pixel_size = 5.28  # Angstroms\n",
-    "pad_scale = 1.5\n",
-    "#scattering = cs.NufftScattering(shape=shape, pad_scale=pad_scale, eps=1e-5)\n",
-    "scattering = cs.FourierSliceExtract(cs.ImageManager(shape=shape, pad_scale=pad_scale), pixel_size=pixel_size)"
+    "manager = cs.ImageManager(shape=shape, pad_scale=1.5)\n",
+    "#scattering = cs.NufftProject(manager, pixel_size=pixel_size, eps=1e-5)\n",
+    "scattering = cs.FourierSliceExtract(manager, pixel_size=pixel_size)"
    ]
   },
   {
@@ -127,8 +127,8 @@
    "source": [
     "# Initialize the instrument\n",
     "optics = cs.CTFOptics(defocus_u=10000, defocus_v=10000, amplitude_contrast=.07)\n",
-    "exposure = cs.UniformExposure(N=1e5, mu=0.0)\n",
-    "detector = cs.GaussianDetector(pixel_size=pixel_size, variance=cs.Constant(1.0))\n",
+    "exposure = cs.UniformExposure(N=1000, mu=0.0)\n",
+    "detector = cs.GaussianDetector(variance=cs.Constant(1.0))\n",
     "instrument_s = cs.Instrument(exposure=exposure)\n",
     "instrument_o = cs.Instrument(exposure=exposure, optics=optics)\n",
     "instrument_d = cs.Instrument(exposure=exposure, optics=optics, detector=detector)"
@@ -148,7 +148,7 @@
    "outputs": [],
    "source": [
     "# Then, choose a model for the solvent\n",
-    "solvent = cs.GaussianIce(variance=cs.Exp(amplitude=1.0, scale=1.0))"
+    "solvent = cs.GaussianIce(variance=cs.Exp(amplitude=5.0, scale=7.5))"
    ]
   },
   {
@@ -177,12 +177,12 @@
    "outputs": [],
    "source": [
     "# Plot models\n",
-    "key = jax.random.split(jax.random.PRNGKey(0))\n",
+    "key = jax.random.PRNGKey(0)\n",
     "fig, axes = plt.subplots(ncols=3, figsize=(12, 6))\n",
     "ax1, ax2, ax3 = axes\n",
     "im1 = plot_image(scattering_model(key), fig, ax1, label=\"Exit plane\")\n",
     "im2 = plot_image(optics_model(key), fig, ax2, label=\"Detector plane\")\n",
-    "im3 = plot_image(detector_model(jax.random.PRNGKey(0)), fig, ax3, label=\"Detector readout\")\n",
+    "im3 = plot_image(detector_model(key), fig, ax3, label=\"Detector readout\")\n",
     "plt.tight_layout()"
    ]
   },
@@ -195,6 +195,13 @@
     "To compute a ``WhiteningFilter``, we first have to generate some synthetic data of structured noise from which we may compute its power spectrum."
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -206,10 +213,12 @@
     "AAAABBB\n",
     "\"\"\", figsize=(6.5, 3.5))\n",
     "key1, key2 = jax.random.split(jax.random.PRNGKey(1234))\n",
-    "freqs = make_frequencies((800, 600), pixel_size=pixel_size)\n",
-    "micrograph = solvent.sample(key1, freqs) * optics(freqs) + fftn(detector.sample(key2, freqs))\n",
-    "whiten = cs.WhiteningFilter(scattering.manager, micrograph=micrograph)\n",
-    "plot_image(ifftn(micrograph.T).real, fig, axes[\"A\"], label=\"Synthetic micrograph\")\n",
+    "micrograph_shape = (500, 400)\n",
+    "freqs = make_frequencies(micrograph_shape, pixel_size)\n",
+    "fourier_micrograph = solvent.sample(key1, freqs) * optics(freqs) + detector.sample(key2, freqs)\n",
+    "micrograph = irfftn(fourier_micrograph, s=micrograph_shape)\n",
+    "whiten = cs.WhiteningFilter(scattering.manager.padded_frequency_grid, micrograph)\n",
+    "plot_image(micrograph, fig, axes[\"A\"], label=\"Synthetic micrograph\")\n",
     "plot_image(whiten.filter, fig, axes[\"B\"], label=\"Whitening filter\")\n",
     "plt.tight_layout()"
    ]
@@ -218,7 +227,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Now generate an image with the toy icy structured noise model with a given whitening filter. We will also apply a low pass filter for antialiasing."
+    "Now generate an image with the toy icy structured noise model with a given whitening filter. We will also apply a low pass filter and a circular mask."
    ]
   },
   {
@@ -229,118 +238,10 @@
    "source": [
     "# Show forward model for a whitened image\n",
     "fig, ax = plt.subplots(figsize=(3.5, 3.5))\n",
-    "filter = cs.LowpassFilter(scattering.manager, cutoff=0.95) * cs.WhiteningFilter(scattering.manager, micrograph=micrograph)\n",
-    "filtered_model = cs.ImagePipeline(scattering=scattering, ensemble=ensemble, instrument=instrument_d, filter=filter)\n",
-    "plot_image(filtered_model.render(), fig, ax)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "If we want to solve the inverse problem and find the model parameters that most likely generated the data, we want to define a function that can be arbitrarily transformed by JAX and evaulated at subsets of the parameters."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def update_model(pixel_size: jax.Array, model: cs.GaussianImage) -> cs.GaussianImage:\n",
-    "    return eqx.tree_at(lambda m: m.instrument.detector.pixel_size, model, pixel_size)\n",
-    "\n",
-    "@jax.jit\n",
-    "def compute_image(pixel_size: jax.Array, model: cs.GaussianImage) -> jax.Array:\n",
-    "    model = update_model(pixel_size, model)\n",
-    "    return model()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Plot jitted model with updated parameters.\n",
-    "fig, axes = plt.subplots(ncols=2, figsize=(8, 6))\n",
-    "ax1, ax2 = axes\n",
-    "plot_image(compute_image(detector.pixel_size, detector_model), fig, ax1)\n",
-    "plot_image(compute_image(1.1*detector.pixel_size, detector_model), fig, ax2)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Now, we can also use the model to compute a likelihood. For this, we can use the ``GaussianImage`` model. This is a subclass of the ``ImagePipeline``, adding the ability to evaulate a likelihood. Let's evaulate the likelihood at the simulated data, and visualize the residuals."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Initialize the model and plot residuals\n",
-    "fig, axes = plt.subplots(ncols=3, figsize=(12, 6))\n",
-    "ax1, ax2, ax3 = axes\n",
-    "mask = cs.CircularMask(scattering.manager)\n",
-    "observation_model = cs.GaussianImage(scattering=scattering, ensemble=ensemble, instrument=instrument_d)\n",
-    "observed = mask(observation_model(key=jax.random.PRNGKey(1244)))\n",
-    "model = cs.GaussianImage(scattering=scattering, ensemble=ensemble, instrument=instrument_d, mask=mask)\n",
-    "simulated, observed = model.render(), observed\n",
-    "plot_image(simulated, fig, ax1)\n",
-    "plot_image(observed, fig, ax2)\n",
-    "plot_image(simulated - observed, fig, ax3)\n",
-    "plt.tight_layout()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Evaluating the ``GaussianImage`` with the observed data set will automatically evaulate the ``GausianImage.log_likelihood`` function."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def update_model(pixel_size: jax.Array, model: cs.GaussianImage) -> cs.GaussianImage:\n",
-    "    return eqx.tree_at(lambda m: m.instrument.detector.pixel_size, model, pixel_size)\n",
-    "\n",
-    "@jax.jit\n",
-    "def compute_loss(pixel_size: jax.Array, model: cs.GaussianImage, observed: jax.Array) -> jax.Array:\n",
-    "    model = update_model(pixel_size, model)\n",
-    "    return model(observed=observed)\n",
-    "\n",
-    "@jax.jit\n",
-    "@jax.grad\n",
-    "def compute_grad(pixel_size: jax.Array, model: cs.GaussianImage, observed: jax.Array) -> jax.Array:\n",
-    "    return compute_loss(pixel_size, model, observed)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Benchmark loss pipeline\n",
-    "%timeit likelihood = compute_loss(detector.pixel_size+.01, model, observed)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Benchmark gradient pipeline\n",
-    "%timeit grad = compute_grad(detector.pixel_size+.01, model, observed)"
+    "filter = cs.LowpassFilter(scattering.manager.padded_frequency_grid, cutoff=0.95) * whiten\n",
+    "mask = cs.CircularMask(scattering.manager.coordinate_grid)\n",
+    "whitened_model = cs.ImagePipeline(scattering=scattering, ensemble=ensemble, instrument=instrument_d, filter=filter, mask=mask)\n",
+    "plot_image(whitened_model.render(), fig, ax)"
    ]
   },
   {


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 250b84d1..c09edf76 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,18 +1,19 @@
 import os
 import pytest
 
+import equinox as eqx
 import jax.random as jr
 from jax import config
 
 import cryojax.simulator as cs
-from cryojax.io import load_fourier_grid
+from cryojax.utils import rfftn
 
 config.update("jax_enable_x64", True)
 
 
-@pytest.fixture()
+@pytest.fixture
 def manager():
-    return cs.ImageManager(shape=(81, 82))
+    return cs.ImageManager(shape=(81, 82), pad_scale=1.1)
 
 
 @pytest.fixture
@@ -34,30 +35,27 @@ def density():
 
 
 @pytest.fixture
-def weights_and_coordinates():
-    filename = os.path.join(
-        os.path.dirname(__file__), "data", "3jar_monomer_bfm1_ps5_28.mrc"
-    )
-    return load_fourier_grid(filename)
+def sample_pdb_path():
+    return os.path.join(os.path.dirname(__file__), "data", "1uao.pdb")
 
 
 @pytest.fixture
 def filters(manager):
-    return cs.LowpassFilter(manager)
+    return cs.LowpassFilter(manager.padded_frequency_grid)
     # return None
 
 
 @pytest.fixture
 def masks(manager):
-    return cs.CircularMask(manager)
+    return cs.CircularMask(manager.coordinate_grid)
 
 
 @pytest.fixture
-def instrument(pixel_size):
+def instrument():
     return cs.Instrument(
         optics=cs.CTFOptics(),
-        exposure=cs.UniformExposure(N=1e5, mu=1.0),
-        detector=cs.GaussianDetector(pixel_size=pixel_size),
+        exposure=cs.UniformExposure(N=1000.0, mu=0.0),
+        detector=cs.GaussianDetector(cs.Constant(1.0)),
     )
 
 
@@ -65,7 +63,13 @@ def instrument(pixel_size):
 def ensemble(density):
     return cs.Ensemble(
         density=density,
-        pose=cs.EulerPose(degrees=False),
+        pose=cs.EulerPose(
+            view_phi=30.0,
+            view_theta=100.0,
+            view_psi=-10.0,
+            offset_x=10.0,
+            offset_y=-5.0,
+        ),
     )
 
 
@@ -75,19 +79,27 @@ def solvent():
 
 
 @pytest.fixture
-def noisy_model(scattering, ensemble, instrument, solvent, filters, masks):
+def noiseless_model(scattering, ensemble, instrument):
+    instrument = eqx.tree_at(
+        lambda ins: ins.detector, instrument, cs.NullDetector()
+    )
+    return cs.ImagePipeline(
+        scattering=scattering, ensemble=ensemble, instrument=instrument
+    )
+
+
+@pytest.fixture
+def noisy_model(scattering, ensemble, instrument, solvent):
     return cs.ImagePipeline(
         scattering=scattering,
         ensemble=ensemble,
         instrument=instrument,
         solvent=solvent,
-        filter=filters,
-        mask=masks,
     )
 
 
 @pytest.fixture
-def maskless_model(scattering, ensemble, instrument, solvent, filters):
+def filtered_model(scattering, ensemble, instrument, solvent, filters):
     return cs.ImagePipeline(
         scattering=scattering,
         ensemble=ensemble,
@@ -98,15 +110,10 @@ def maskless_model(scattering, ensemble, instrument, solvent, filters):
 
 
 @pytest.fixture
-def test_image(noisy_model):
-    return noisy_model.sample(jr.split(jr.PRNGKey(1234), num=2))
-
-
-@pytest.fixture
-def likelihood_model(
+def filtered_and_masked_model(
     scattering, ensemble, instrument, solvent, filters, masks
 ):
-    return cs.GaussianImage(
+    return cs.ImagePipeline(
         scattering=scattering,
         ensemble=ensemble,
         instrument=instrument,
@@ -114,3 +121,21 @@ def likelihood_model(
         filter=filters,
         mask=masks,
     )
+
+
+@pytest.fixture
+def test_image(noisy_model):
+    image = noisy_model.sample(jr.PRNGKey(1234), view=False)
+    return rfftn(image)
+
+
+@pytest.fixture
+def likelihood_model(noisy_model):
+    return cs.IndependentFourierGaussian(noisy_model)
+
+
+@pytest.fixture
+def likelihood_model_with_custom_variance(noiseless_model):
+    return cs.IndependentFourierGaussian(
+        noiseless_model, noise=cs.GaussianNoise(variance=cs.Constant(1.0))
+    )
diff --git a/tests/data/1uao.pdb b/tests/data/1uao.pdb
new file mode 100644
index 00000000..6e1325a8
--- /dev/null
+++ b/tests/data/1uao.pdb
@@ -0,0 +1,282 @@
+HEADER    DE NOVO PROTEIN                         13-MAR-03   1UAO
+TITLE     NMR STRUCTURE OF DESIGNED PROTEIN, CHIGNOLIN, CONSISTING OF ONLY TEN
+TITLE    2 AMINO ACIDS (ENSEMBLES)
+COMPND    MOL_ID: 1;
+COMPND   2 MOLECULE: CHIGNOLIN;
+COMPND   3 CHAIN: A;
+COMPND   4 ENGINEERED: YES
+SOURCE    MOL_ID: 1;
+SOURCE   2 SYNTHETIC: YES;
+SOURCE   3 OTHER_DETAILS: THE SEQUENCE WAS DESIGNED ON THE BASIS OF STATISTICS
+SOURCE   4 DERIVED FROM NUMEROUS PROTEIN SEGMENTS.
+KEYWDS    DE NOVO PROTEIN, BETA-HAIRPIN, MINI-PROTEIN, G-PEPTIDE, AUTONOMOUS
+KEYWDS   2 ELEMENT
+EXPDTA    SOLUTION NMR
+NUMMDL    18
+AUTHOR    S.HONDA,K.YAMASAKI
+REVDAT   4   26-AUG-15 1UAO    1       JRNL   VERSN
+REVDAT   3   24-FEB-09 1UAO    1       VERSN
+REVDAT   2   17-AUG-04 1UAO    1       JRNL
+REVDAT   1   13-APR-04 1UAO    0
+JRNL        AUTH   S.HONDA,K.YAMASAKI,Y.SAWADA,H.MORII
+JRNL        TITL   10 RESIDUE FOLDED PEPTIDE DESIGNED BY SEGMENT STATISTICS
+JRNL        REF    STRUCTURE                     V.  12  1507 2004
+JRNL        REFN                   ISSN 0969-2126
+JRNL        PMID   15296744
+JRNL        DOI    10.1016/J.STR.2004.05.022
+REMARK   2
+REMARK   2 RESOLUTION. NOT APPLICABLE.
+REMARK   3
+REMARK   3 REFINEMENT.
+REMARK   3   PROGRAM     : X-PLOR 3.1
+REMARK   3   AUTHORS     : BRUNGER
+REMARK   3
+REMARK   3  OTHER REFINEMENT REMARKS: THE STRUCTURES ARE BASED ON A TOTAL OF
+REMARK   3  185 RESTRAINTS, 172 ARE NOE-DERIVED DISTANCE CONSTRAINTS, 12
+REMARK   3  DIHEDRAL ANGLE RESTRAINTS,1 DISTANCE RESTRAINTS FROM HYDROGEN
+REMARK   3  BONDS.
+REMARK   4
+REMARK   4 1UAO COMPLIES WITH FORMAT V. 3.15, 01-DEC-08
+REMARK 100
+REMARK 100 THIS ENTRY HAS BEEN PROCESSED BY PDBJ ON 17-MAR-03.
+REMARK 100 THE RCSB ID CODE IS RCSB005623.
+REMARK 210
+REMARK 210 EXPERIMENTAL DETAILS
+REMARK 210  EXPERIMENT TYPE                : NMR
+REMARK 210  TEMPERATURE           (KELVIN) : 277
+REMARK 210  PH                             : 5.5
+REMARK 210  IONIC STRENGTH                 : 20MM SODIUM PHOSOHATE
+REMARK 210  PRESSURE                       : AMBIENT
+REMARK 210  SAMPLE CONTENTS                : 2MM CHIGNOLIN
+REMARK 210
+REMARK 210  NMR EXPERIMENTS CONDUCTED      : 2D NOESY; 2D ROESY; DQF-COSY; 2D
+REMARK 210                                   TOCSY
+REMARK 210  SPECTROMETER FIELD STRENGTH    : 500 MHZ
+REMARK 210  SPECTROMETER MODEL             : AMX
+REMARK 210  SPECTROMETER MANUFACTURER      : BRUKER
+REMARK 210
+REMARK 210  STRUCTURE DETERMINATION.
+REMARK 210   SOFTWARE USED                 : XWINNMR 2.6, FELIX 2000, X-PLOR
+REMARK 210                                   3.1
+REMARK 210   METHOD USED                   : SIMULATED ANNEALING
+REMARK 210
+REMARK 210 CONFORMERS, NUMBER CALCULATED   : 50
+REMARK 210 CONFORMERS, NUMBER SUBMITTED    : 18
+REMARK 210 CONFORMERS, SELECTION CRITERIA  : STRUCTURES WITH THE LOWEST ENERGY
+REMARK 210
+REMARK 210 BEST REPRESENTATIVE CONFORMER IN THIS ENSEMBLE : 1
+REMARK 210
+REMARK 210 REMARK: THIS STRUCTURE WAS DETERMINED USING STANDARD 2D
+REMARK 210  HOMONUCLEAR TECHNIQUES.
+REMARK 215
+REMARK 215 NMR STUDY
+REMARK 215 THE COORDINATES IN THIS ENTRY WERE GENERATED FROM SOLUTION
+REMARK 215 NMR DATA.  PROTEIN DATA BANK CONVENTIONS REQUIRE THAT
+REMARK 215 CRYST1 AND SCALE RECORDS BE INCLUDED, BUT THE VALUES ON
+REMARK 215 THESE RECORDS ARE MEANINGLESS.
+REMARK 300
+REMARK 300 BIOMOLECULE: 1
+REMARK 300 SEE REMARK 350 FOR THE AUTHOR PROVIDED AND/OR PROGRAM
+REMARK 300 GENERATED ASSEMBLY INFORMATION FOR THE STRUCTURE IN
+REMARK 300 THIS ENTRY. THE REMARK MAY ALSO PROVIDE INFORMATION ON
+REMARK 300 BURIED SURFACE AREA.
+REMARK 350
+REMARK 350 COORDINATES FOR A COMPLETE MULTIMER REPRESENTING THE KNOWN
+REMARK 350 BIOLOGICALLY SIGNIFICANT OLIGOMERIZATION STATE OF THE
+REMARK 350 MOLECULE CAN BE GENERATED BY APPLYING BIOMT TRANSFORMATIONS
+REMARK 350 GIVEN BELOW.  BOTH NON-CRYSTALLOGRAPHIC AND
+REMARK 350 CRYSTALLOGRAPHIC OPERATIONS ARE GIVEN.
+REMARK 350
+REMARK 350 BIOMOLECULE: 1
+REMARK 350 AUTHOR DETERMINED BIOLOGICAL UNIT: MONOMERIC
+REMARK 350 APPLY THE FOLLOWING TO CHAINS: A
+REMARK 350   BIOMT1   1  1.000000  0.000000  0.000000        0.00000
+REMARK 350   BIOMT2   1  0.000000  1.000000  0.000000        0.00000
+REMARK 350   BIOMT3   1  0.000000  0.000000  1.000000        0.00000
+REMARK 500
+REMARK 500 GEOMETRY AND STEREOCHEMISTRY
+REMARK 500 SUBTOPIC: TORSION ANGLES
+REMARK 500
+REMARK 500 TORSION ANGLES OUTSIDE THE EXPECTED RAMACHANDRAN REGIONS:
+REMARK 500 (M=MODEL NUMBER; RES=RESIDUE NAME; C=CHAIN IDENTIFIER;
+REMARK 500 SSEQ=SEQUENCE NUMBER; I=INSERTION CODE).
+REMARK 500
+REMARK 500 STANDARD TABLE:
+REMARK 500 FORMAT:(10X,I3,1X,A3,1X,A1,I4,A1,4X,F7.2,3X,F7.2)
+REMARK 500
+REMARK 500 EXPECTED VALUES: GJ KLEYWEGT AND TA JONES (1996). PHI/PSI-
+REMARK 500 CHOLOGY: RAMACHANDRAN REVISITED. STRUCTURE 4, 1395 - 1400
+REMARK 500
+REMARK 500  M RES CSSEQI        PSI       PHI
+REMARK 500  4 TRP A   9     -167.59    -69.03
+REMARK 500  7 THR A   6      -53.23   -121.01
+REMARK 500  8 TRP A   9       38.27    -87.59
+REMARK 500 11 TRP A   9      -72.89    -84.72
+REMARK 500 12 TRP A   9       57.70    -93.27
+REMARK 500 13 TRP A   9       42.80    -90.88
+REMARK 500 18 TRP A   9       79.66   -101.87
+REMARK 500
+REMARK 500 REMARK: NULL
+REMARK 900
+REMARK 900 RELATED ENTRIES
+REMARK 900 RELATED ID: 5694   RELATED DB: BMRB
+REMARK 900 NMR STRUCTURE OF DESIGNED PROTEIN, CHIGNOLIN, ONSISTING OF
+REMARK 900 ONLY TEN AMINO ACIDS
+REMARK 900 RELATED ID: 2RVD   RELATED DB: PDB
+REMARK 900 NMR STRUCTURE OF A MUTANT OF CHIGNOLIN, CLN025
+REMARK 900 RELATED ID: 5AWL   RELATED DB: PDB
+REMARK 900 CRYSTAL STRUCTURE OF A MUTANT OF CHIGNOLIN, CLN025
+DBREF  1UAO A    1    10  PDB    1UAO     1UAO             1     10
+SEQRES   1 A   10  GLY TYR ASP PRO GLU THR GLY THR TRP GLY
+SHEET    1   A 2 TYR A   2  ASP A   3  0
+SHEET    2   A 2 THR A   8  TRP A   9 -1  O  THR A   8   N  ASP A   3
+CRYST1    1.000    1.000    1.000  90.00  90.00  90.00 P 1           1
+ORIGX1      1.000000  0.000000  0.000000        0.00000
+ORIGX2      0.000000  1.000000  0.000000        0.00000
+ORIGX3      0.000000  0.000000  1.000000        0.00000
+SCALE1      1.000000  0.000000  0.000000        0.00000
+SCALE2      0.000000  1.000000  0.000000        0.00000
+SCALE3      0.000000  0.000000  1.000000        0.00000
+MODEL        1
+ATOM      1  N   GLY A   1      -6.778  -1.424   4.200  1.00  0.00           N
+ATOM      2  CA  GLY A   1      -6.878  -0.708   2.896  1.00  0.00           C
+ATOM      3  C   GLY A   1      -5.557  -0.840   2.138  1.00  0.00           C
+ATOM      4  O   GLY A   1      -4.640  -1.504   2.579  1.00  0.00           O
+ATOM      5  H1  GLY A   1      -5.778  -1.527   4.462  1.00  0.00           H
+ATOM      6  H2  GLY A   1      -7.214  -2.365   4.112  1.00  0.00           H
+ATOM      7  H3  GLY A   1      -7.273  -0.879   4.933  1.00  0.00           H
+ATOM      8  HA2 GLY A   1      -7.677  -1.140   2.309  1.00  0.00           H
+ATOM      9  HA3 GLY A   1      -7.085   0.336   3.073  1.00  0.00           H
+ATOM     10  N   TYR A   2      -5.452  -0.212   0.999  1.00  0.00           N
+ATOM     11  CA  TYR A   2      -4.189  -0.302   0.213  1.00  0.00           C
+ATOM     12  C   TYR A   2      -3.197   0.744   0.717  1.00  0.00           C
+ATOM     13  O   TYR A   2      -3.252   1.898   0.342  1.00  0.00           O
+ATOM     14  CB  TYR A   2      -4.490  -0.048  -1.264  1.00  0.00           C
+ATOM     15  CG  TYR A   2      -3.216  -0.141  -2.068  1.00  0.00           C
+ATOM     16  CD1 TYR A   2      -2.780  -1.383  -2.544  1.00  0.00           C
+ATOM     17  CD2 TYR A   2      -2.474   1.015  -2.342  1.00  0.00           C
+ATOM     18  CE1 TYR A   2      -1.601  -1.470  -3.293  1.00  0.00           C
+ATOM     19  CE2 TYR A   2      -1.295   0.927  -3.091  1.00  0.00           C
+ATOM     20  CZ  TYR A   2      -0.858  -0.315  -3.568  1.00  0.00           C
+ATOM     21  OH  TYR A   2       0.304  -0.402  -4.307  1.00  0.00           O
+ATOM     22  H   TYR A   2      -6.203   0.317   0.660  1.00  0.00           H
+ATOM     23  HA  TYR A   2      -3.763  -1.288   0.328  1.00  0.00           H
+ATOM     24  HB2 TYR A   2      -5.189  -0.788  -1.616  1.00  0.00           H
+ATOM     25  HB3 TYR A   2      -4.918   0.937  -1.380  1.00  0.00           H
+ATOM     26  HD1 TYR A   2      -3.352  -2.273  -2.332  1.00  0.00           H
+ATOM     27  HD2 TYR A   2      -2.811   1.973  -1.975  1.00  0.00           H
+ATOM     28  HE1 TYR A   2      -1.265  -2.429  -3.660  1.00  0.00           H
+ATOM     29  HE2 TYR A   2      -0.722   1.818  -3.303  1.00  0.00           H
+ATOM     30  HH  TYR A   2       0.765   0.437  -4.234  1.00  0.00           H
+ATOM     31  N   ASP A   3      -2.287   0.349   1.560  1.00  0.00           N
+ATOM     32  CA  ASP A   3      -1.287   1.320   2.084  1.00  0.00           C
+ATOM     33  C   ASP A   3      -0.301   1.677   0.964  1.00  0.00           C
+ATOM     34  O   ASP A   3       0.433   0.826   0.502  1.00  0.00           O
+ATOM     35  CB  ASP A   3      -0.523   0.687   3.249  1.00  0.00           C
+ATOM     36  CG  ASP A   3       0.418   1.724   3.866  1.00  0.00           C
+ATOM     37  OD1 ASP A   3      -0.069   2.762   4.283  1.00  0.00           O
+ATOM     38  OD2 ASP A   3       1.609   1.463   3.911  1.00  0.00           O
+ATOM     39  H   ASP A   3      -2.259  -0.588   1.846  1.00  0.00           H
+ATOM     40  HA  ASP A   3      -1.795   2.207   2.427  1.00  0.00           H
+ATOM     41  HB2 ASP A   3      -1.224   0.346   3.997  1.00  0.00           H
+ATOM     42  HB3 ASP A   3       0.055  -0.151   2.889  1.00  0.00           H
+ATOM     43  N   PRO A   4      -0.311   2.922   0.557  1.00  0.00           N
+ATOM     44  CA  PRO A   4       0.579   3.407  -0.513  1.00  0.00           C
+ATOM     45  C   PRO A   4       2.002   3.580   0.021  1.00  0.00           C
+ATOM     46  O   PRO A   4       2.943   3.738  -0.731  1.00  0.00           O
+ATOM     47  CB  PRO A   4      -0.033   4.752  -0.911  1.00  0.00           C
+ATOM     48  CG  PRO A   4      -0.876   5.221   0.298  1.00  0.00           C
+ATOM     49  CD  PRO A   4      -1.199   3.960   1.121  1.00  0.00           C
+ATOM     50  HA  PRO A   4       0.567   2.732  -1.353  1.00  0.00           H
+ATOM     51  HB2 PRO A   4       0.750   5.467  -1.123  1.00  0.00           H
+ATOM     52  HB3 PRO A   4      -0.670   4.630  -1.773  1.00  0.00           H
+ATOM     53  HG2 PRO A   4      -0.308   5.921   0.895  1.00  0.00           H
+ATOM     54  HG3 PRO A   4      -1.792   5.679  -0.042  1.00  0.00           H
+ATOM     55  HD2 PRO A   4      -0.977   4.125   2.167  1.00  0.00           H
+ATOM     56  HD3 PRO A   4      -2.232   3.677   0.991  1.00  0.00           H
+ATOM     57  N   GLU A   5       2.168   3.542   1.315  1.00  0.00           N
+ATOM     58  CA  GLU A   5       3.531   3.694   1.893  1.00  0.00           C
+ATOM     59  C   GLU A   5       4.420   2.567   1.368  1.00  0.00           C
+ATOM     60  O   GLU A   5       5.577   2.766   1.055  1.00  0.00           O
+ATOM     61  CB  GLU A   5       3.450   3.614   3.418  1.00  0.00           C
+ATOM     62  CG  GLU A   5       4.652   4.334   4.030  1.00  0.00           C
+ATOM     63  CD  GLU A   5       5.629   3.303   4.597  1.00  0.00           C
+ATOM     64  OE1 GLU A   5       5.649   2.193   4.090  1.00  0.00           O
+ATOM     65  OE2 GLU A   5       6.341   3.640   5.529  1.00  0.00           O
+ATOM     66  H   GLU A   5       1.397   3.406   1.904  1.00  0.00           H
+ATOM     67  HA  GLU A   5       3.944   4.648   1.600  1.00  0.00           H
+ATOM     68  HB2 GLU A   5       2.537   4.084   3.756  1.00  0.00           H
+ATOM     69  HB3 GLU A   5       3.457   2.580   3.727  1.00  0.00           H
+ATOM     70  HG2 GLU A   5       5.147   4.918   3.267  1.00  0.00           H
+ATOM     71  HG3 GLU A   5       4.318   4.985   4.823  1.00  0.00           H
+ATOM     72  N   THR A   6       3.881   1.383   1.261  1.00  0.00           N
+ATOM     73  CA  THR A   6       4.684   0.239   0.748  1.00  0.00           C
+ATOM     74  C   THR A   6       4.218  -0.101  -0.669  1.00  0.00           C
+ATOM     75  O   THR A   6       4.941  -0.687  -1.449  1.00  0.00           O
+ATOM     76  CB  THR A   6       4.486  -0.974   1.661  1.00  0.00           C
+ATOM     77  OG1 THR A   6       3.204  -1.540   1.423  1.00  0.00           O
+ATOM     78  CG2 THR A   6       4.589  -0.535   3.123  1.00  0.00           C
+ATOM     79  H   THR A   6       2.944   1.248   1.514  1.00  0.00           H
+ATOM     80  HA  THR A   6       5.729   0.510   0.729  1.00  0.00           H
+ATOM     81  HB  THR A   6       5.248  -1.709   1.456  1.00  0.00           H
+ATOM     82  HG1 THR A   6       3.211  -2.440   1.757  1.00  0.00           H
+ATOM     83 HG21 THR A   6       5.309   0.265   3.208  1.00  0.00           H
+ATOM     84 HG22 THR A   6       4.906  -1.371   3.728  1.00  0.00           H
+ATOM     85 HG23 THR A   6       3.625  -0.188   3.464  1.00  0.00           H
+ATOM     86  N   GLY A   7       3.014   0.272  -1.008  1.00  0.00           N
+ATOM     87  CA  GLY A   7       2.498  -0.018  -2.375  1.00  0.00           C
+ATOM     88  C   GLY A   7       1.780  -1.368  -2.385  1.00  0.00           C
+ATOM     89  O   GLY A   7       1.815  -2.091  -3.362  1.00  0.00           O
+ATOM     90  H   GLY A   7       2.450   0.749  -0.363  1.00  0.00           H
+ATOM     91  HA2 GLY A   7       1.808   0.760  -2.670  1.00  0.00           H
+ATOM     92  HA3 GLY A   7       3.323  -0.049  -3.071  1.00  0.00           H
+ATOM     93  N   THR A   8       1.127  -1.719  -1.310  1.00  0.00           N
+ATOM     94  CA  THR A   8       0.411  -3.025  -1.274  1.00  0.00           C
+ATOM     95  C   THR A   8      -0.641  -3.013  -0.162  1.00  0.00           C
+ATOM     96  O   THR A   8      -0.633  -2.164   0.707  1.00  0.00           O
+ATOM     97  CB  THR A   8       1.418  -4.147  -1.011  1.00  0.00           C
+ATOM     98  OG1 THR A   8       2.633  -3.587  -0.532  1.00  0.00           O
+ATOM     99  CG2 THR A   8       1.682  -4.912  -2.308  1.00  0.00           C
+ATOM    100  H   THR A   8       1.108  -1.125  -0.531  1.00  0.00           H
+ATOM    101  HA  THR A   8      -0.073  -3.195  -2.224  1.00  0.00           H
+ATOM    102  HB  THR A   8       1.018  -4.825  -0.273  1.00  0.00           H
+ATOM    103  HG1 THR A   8       3.275  -4.296  -0.450  1.00  0.00           H
+ATOM    104 HG21 THR A   8       0.813  -5.501  -2.562  1.00  0.00           H
+ATOM    105 HG22 THR A   8       2.533  -5.563  -2.176  1.00  0.00           H
+ATOM    106 HG23 THR A   8       1.885  -4.210  -3.104  1.00  0.00           H
+ATOM    107  N   TRP A   9      -1.545  -3.955  -0.184  1.00  0.00           N
+ATOM    108  CA  TRP A   9      -2.598  -4.011   0.868  1.00  0.00           C
+ATOM    109  C   TRP A   9      -1.960  -4.395   2.205  1.00  0.00           C
+ATOM    110  O   TRP A   9      -1.576  -5.528   2.417  1.00  0.00           O
+ATOM    111  CB  TRP A   9      -3.640  -5.067   0.487  1.00  0.00           C
+ATOM    112  CG  TRP A   9      -3.848  -5.068  -0.996  1.00  0.00           C
+ATOM    113  CD1 TRP A   9      -3.216  -5.882  -1.873  1.00  0.00           C
+ATOM    114  CD2 TRP A   9      -4.743  -4.236  -1.786  1.00  0.00           C
+ATOM    115  NE1 TRP A   9      -3.664  -5.596  -3.150  1.00  0.00           N
+ATOM    116  CE2 TRP A   9      -4.608  -4.591  -3.148  1.00  0.00           C
+ATOM    117  CE3 TRP A   9      -5.648  -3.218  -1.454  1.00  0.00           C
+ATOM    118  CZ2 TRP A   9      -5.349  -3.958  -4.146  1.00  0.00           C
+ATOM    119  CZ3 TRP A   9      -6.398  -2.577  -2.455  1.00  0.00           C
+ATOM    120  CH2 TRP A   9      -6.248  -2.947  -3.799  1.00  0.00           C
+ATOM    121  H   TRP A   9      -1.532  -4.627  -0.892  1.00  0.00           H
+ATOM    122  HA  TRP A   9      -3.075  -3.047   0.957  1.00  0.00           H
+ATOM    123  HB2 TRP A   9      -3.295  -6.041   0.800  1.00  0.00           H
+ATOM    124  HB3 TRP A   9      -4.575  -4.843   0.980  1.00  0.00           H
+ATOM    125  HD1 TRP A   9      -2.478  -6.628  -1.618  1.00  0.00           H
+ATOM    126  HE1 TRP A   9      -3.364  -6.042  -3.968  1.00  0.00           H
+ATOM    127  HE3 TRP A   9      -5.766  -2.927  -0.422  1.00  0.00           H
+ATOM    128  HZ2 TRP A   9      -5.228  -4.248  -5.180  1.00  0.00           H
+ATOM    129  HZ3 TRP A   9      -7.093  -1.796  -2.187  1.00  0.00           H
+ATOM    130  HH2 TRP A   9      -6.827  -2.451  -4.564  1.00  0.00           H
+ATOM    131  N   GLY A  10      -1.844  -3.461   3.109  1.00  0.00           N
+ATOM    132  CA  GLY A  10      -1.229  -3.774   4.431  1.00  0.00           C
+ATOM    133  C   GLY A  10      -1.877  -5.031   5.014  1.00  0.00           C
+ATOM    134  O   GLY A  10      -1.289  -6.092   4.882  1.00  0.00           O
+ATOM    135  OXT GLY A  10      -2.949  -4.912   5.584  1.00  0.00           O
+ATOM    136  H   GLY A  10      -2.159  -2.553   2.918  1.00  0.00           H
+ATOM    137  HA2 GLY A  10      -0.169  -3.939   4.305  1.00  0.00           H
+ATOM    138  HA3 GLY A  10      -1.387  -2.946   5.106  1.00  0.00           H
+TER     139      GLY A  10
+ENDMDL
+MASTER      102    0    0    0    2    0    0    6   77    1    0    1
+END
diff --git a/tests/test_ctf.py b/tests/test_ctf.py
index bb84bd73..f7ebf7e3 100644
--- a/tests/test_ctf.py
+++ b/tests/test_ctf.py
@@ -1,4 +1,5 @@
 import pytest
+import jax.numpy as jnp
 import numpy as np
 from pycistem.core import CTF
 
@@ -27,7 +28,7 @@ def test_ctf_with_cistem(
 
     Modified from https://github.com/jojoelfe/contrasttransferfunction"""
     shape = (512, 512)
-    freqs = make_frequencies(shape, pixel_size=pixel_size)
+    freqs = make_frequencies(shape, pixel_size)
     k_sqr, theta = cartesian_to_polar(freqs, square=True)
     # Compute cryojax CTF
     optics = CTFOptics(
@@ -37,9 +38,8 @@ def test_ctf_with_cistem(
         voltage=kV,
         spherical_aberration=cs,
         amplitude_contrast=ac,
-        envelope=None,
     )
-    ctf = np.array(optics(freqs, normalize=False))
+    ctf = np.array(optics(freqs))
     # Compute cisTEM CTF
     cisTEM_optics = CTF(
         kV=kV,
@@ -52,13 +52,12 @@ def test_ctf_with_cistem(
     )
     cisTEM_ctf = np.vectorize(
         lambda k_sqr, theta: cisTEM_optics.Evaluate(k_sqr, theta)
-    )(k_sqr.ravel() * pixel_size**2, theta.ravel()).reshape(shape)
+    )(k_sqr.ravel() * pixel_size**2, theta.ravel()).reshape(freqs.shape[0:2])
 
     # Compute cryojax and cisTEM power spectrum
-    spectrum1D, _ = powerspectrum(ctf, freqs, pixel_size=pixel_size)
-    cisTEM_spectrum1D, _ = powerspectrum(
-        cisTEM_ctf, freqs, pixel_size=pixel_size
-    )
+    radial_freqs = jnp.linalg.norm(freqs, axis=-1)
+    spectrum1D, _ = powerspectrum(ctf, radial_freqs, pixel_size)
+    cisTEM_spectrum1D, _ = powerspectrum(cisTEM_ctf, radial_freqs, pixel_size)
 
     np.testing.assert_allclose(ctf, cisTEM_ctf, atol=5e-2)
     np.testing.assert_allclose(spectrum1D, cisTEM_spectrum1D, atol=5e-3)
diff --git a/tests/test_distribution.py b/tests/test_distribution.py
new file mode 100644
index 00000000..30614411
--- /dev/null
+++ b/tests/test_distribution.py
@@ -0,0 +1,29 @@
+import pytest
+
+import cryojax.simulator as cs
+import equinox as eqx
+import numpy as np
+import jax.numpy as jnp
+import jax.random as jr
+
+
+def test_custom_variance(
+    likelihood_model,
+    likelihood_model_with_custom_variance,
+    test_image,
+):
+    likelihood_model = eqx.tree_at(
+        lambda m: m.pipeline.solvent, likelihood_model, cs.NullIce()
+    )
+    np.testing.assert_allclose(
+        likelihood_model.variance,
+        likelihood_model_with_custom_variance.variance,
+    )
+    np.testing.assert_allclose(
+        likelihood_model.log_probability(test_image),
+        likelihood_model_with_custom_variance.log_probability(test_image),
+    )
+    np.testing.assert_allclose(
+        likelihood_model.sample(jr.PRNGKey(seed=0)),
+        likelihood_model_with_custom_variance.sample(jr.PRNGKey(seed=0)),
+    )
diff --git a/tests/test_rescale.py b/tests/test_exposure.py
similarity index 87%
rename from tests/test_rescale.py
rename to tests/test_exposure.py
index c82e6217..403ee5d4 100644
--- a/tests/test_rescale.py
+++ b/tests/test_exposure.py
@@ -8,9 +8,7 @@
 config.update("jax_enable_x64", True)
 
 
-@pytest.mark.parametrize(
-    "rescaled_model", ["noisy_model", "maskless_model", "likelihood_model"]
-)
+@pytest.mark.parametrize("rescaled_model", ["noisy_model", "noiseless_model"])
 def test_rescale(rescaled_model, request):
     rescaled_model = request.getfixturevalue(rescaled_model)
     exposure = rescaled_model.instrument.exposure
diff --git a/tests/test_fft.py b/tests/test_fft.py
index bddb3ca8..07b0ee4d 100644
--- a/tests/test_fft.py
+++ b/tests/test_fft.py
@@ -3,39 +3,36 @@
 import jax.numpy as jnp
 import numpy as np
 from cryojax.utils import fftn, ifftn
+from jax import config
 
+config.update("jax_enable_x64", True)
 
-@pytest.mark.parametrize(
-    "model", ["noisy_model", "maskless_model", "likelihood_model"]
-)
+
+@pytest.mark.parametrize("model", ["noisy_model", "noiseless_model"])
 def test_fft(model, request):
     model = request.getfixturevalue(model)
-    image = fftn(model.render())
+    image = model.render()
     random = jnp.asarray(np.random.randn(*image.shape))
-    # Set tolerance based on tests with jnp.fft + random data
-    rkwargs = dict(zip(["atol", "rtol"], [1e-6, 5e-4]))
-    fkwargs = dict(zip(["atol", "rtol"], [5e-5, 1e-6]))
     # Run tests with jnp.fft and random data
     np.testing.assert_allclose(
-        random, jnp.fft.ifftn(jnp.fft.fftn(random)).real, **rkwargs
+        random, jnp.fft.ifftn(jnp.fft.fftn(random)).real
     )
     np.testing.assert_allclose(
         jnp.fft.fftn(random),
         jnp.fft.fftn(jnp.fft.ifftn(jnp.fft.fftn(random)).real),
-        **fkwargs
     )
     # Run tests with cryojax.utils and random data
     np.testing.assert_allclose(
-        ifftn(fftn(random)).real,
-        jnp.fft.ifftn(jnp.fft.fftn(random)).real,
-        **rkwargs
+        ifftn(fftn(random)).real, jnp.fft.ifftn(jnp.fft.fftn(random)).real
     )
-    np.testing.assert_allclose(random, ifftn(fftn(random)).real, **rkwargs)
+    np.testing.assert_allclose(random, ifftn(fftn(random)).real)
+    np.testing.assert_allclose(fftn(random), fftn(ifftn(fftn(random))))
+    # Run tests with an image
+    np.testing.assert_allclose(image, ifftn(fftn(image)).real)
+    # ... test zero mode separately
     np.testing.assert_allclose(
-        fftn(random), fftn(ifftn(fftn(random))), **fkwargs
+        fftn(image)[1:, 1:], fftn(ifftn(fftn(image)).real)[1:, 1:]
     )
-    # Run tests with an image
     np.testing.assert_allclose(
-        ifftn(image).real, ifftn(fftn(ifftn(image))).real, **rkwargs
+        fftn(image)[0, 0], fftn(ifftn(fftn(image)).real)[0, 0], atol=1e-12
     )
-    np.testing.assert_allclose(image, fftn(ifftn(image).real), **fkwargs)
diff --git a/tests/test_filters_and_masks.py b/tests/test_filters_and_masks.py
new file mode 100644
index 00000000..31ebcd63
--- /dev/null
+++ b/tests/test_filters_and_masks.py
@@ -0,0 +1,42 @@
+import pytest
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+
+@pytest.mark.parametrize("model", ["noisy_model"])
+def test_compute_with_filters_and_masks(
+    model, filtered_and_masked_model, request, filters, masks
+):
+    """Make sure that adding null filters and masks does not change output"""
+    model = request.getfixturevalue(model)
+    # Add null filters and masks
+    null_mask = eqx.tree_at(lambda m: m.mask, masks, jnp.asarray(1.0))
+    null_filter = eqx.tree_at(lambda f: f.filter, filters, jnp.asarray(1.0))
+    where = lambda m: (m.filter, m.mask)
+    model_with_null_mask = eqx.tree_at(
+        where, filtered_and_masked_model, (None, null_mask)
+    )
+    model_with_null_filter = eqx.tree_at(
+        where, filtered_and_masked_model, (null_filter, None)
+    )
+    model_with_null_filter_and_mask = eqx.tree_at(
+        where,
+        filtered_and_masked_model,
+        (null_filter, null_mask),
+    )
+    # Compute images
+    key = jax.random.PRNGKey(0)
+    image = model.render()
+    noisy_image = model.sample(key)
+    # Check render
+    np.testing.assert_allclose(model_with_null_mask.render(), image)
+    np.testing.assert_allclose(model_with_null_filter.render(), image)
+    np.testing.assert_allclose(model_with_null_filter_and_mask.render(), image)
+    # Check sample
+    np.testing.assert_allclose(model_with_null_mask.sample(key), noisy_image)
+    np.testing.assert_allclose(model_with_null_filter.sample(key), noisy_image)
+    np.testing.assert_allclose(
+        model_with_null_filter_and_mask.sample(key), noisy_image
+    )
diff --git a/tests/test_jit.py b/tests/test_jit.py
index 8bac61b4..7ef84d19 100644
--- a/tests/test_jit.py
+++ b/tests/test_jit.py
@@ -13,8 +13,10 @@
 config.update("jax_enable_x64", True)
 
 
-def test_jit(likelihood_model, test_image):
-    key = jr.split(jr.PRNGKey(0))
+@pytest.mark.parametrize("model", ["likelihood_model"])
+def test_jit(model, test_image, request):
+    likelihood_model = request.getfixturevalue(model)
+    key = jr.PRNGKey(0)
 
     @jax.jit
     def compute_image(model, key):
@@ -35,7 +37,7 @@ def compute_loss(model, test_image):
 
 def test_value_and_grad(likelihood_model, test_image):
     def build_model(model, params):
-        where = lambda m: m.ensemble.pose.offset_z
+        where = lambda m: m.pipeline.ensemble.pose.offset_z
         return eqx.tree_at(where, model, params["offset_z"])
 
     @jax.jit
diff --git a/tests/test_normalize.py b/tests/test_normalize.py
new file mode 100644
index 00000000..f0cfa8b5
--- /dev/null
+++ b/tests/test_normalize.py
@@ -0,0 +1,38 @@
+import pytest
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from cryojax.utils import irfftn
+from jax import config
+
+config.update("jax_enable_x64", True)
+
+
+@pytest.mark.parametrize(
+    "model",
+    [
+        "noisy_model",
+        "noiseless_model",
+        "filtered_model",
+        "filtered_and_masked_model",
+    ],
+)
+def test_compute_with_filters_and_masks(model, request):
+    model = request.getfixturevalue(model)
+    key = jax.random.PRNGKey(1234)
+    im1 = model.render(get_real=True, normalize=True)
+    im2 = model.render(view=False, get_real=True, normalize=True)
+    im3 = model.sample(key, get_real=True, normalize=True)
+    im4 = model.sample(key, view=False, get_real=True, normalize=True)
+    im5 = irfftn(
+        model.render(get_real=False, normalize=True),
+        s=model.scattering.manager.shape,
+    )
+    im6 = irfftn(
+        model.render(get_real=False, normalize=True, view=False),
+        s=model.scattering.manager.padded_shape,
+    )
+    for im in [im1, im2, im3, im4, im5, im6]:
+        np.testing.assert_allclose(jnp.std(im), jnp.asarray(1.0), rtol=1e-3)
+        np.testing.assert_allclose(jnp.mean(im), jnp.asarray(0.0), atol=1e-16)
diff --git a/tests/test_serialize.py b/tests/test_serialize.py
deleted file mode 100644
index 34232771..00000000
--- a/tests/test_serialize.py
+++ /dev/null
@@ -1,68 +0,0 @@
-import pytest
-
-import numpy as np
-from jax import config
-
-import cryojax.simulator as cs
-
-config.update("jax_enable_x64", True)
-
-
-def test_deserialize_instrument(instrument):
-    """Test Instrument deserialization"""
-    assert (
-        cs.Instrument.from_json(instrument.to_json()).to_json()
-        == instrument.to_json()
-    )
-
-
-def test_deserialize_specimen(ensemble):
-    """Test Specimen deserialization"""
-    assert (
-        cs.Ensemble.from_json(ensemble.to_json()).to_json()
-        == ensemble.to_json()
-    )
-
-
-def test_deserialize_density(density):
-    """Test specimen deserialization."""
-    assert (
-        cs.VoxelGrid.from_json(density.to_json()).to_json()
-        == density.to_json()
-    )
-
-
-@pytest.mark.parametrize("filters_or_masks", ["filters", "masks"])
-def test_deserialize_filters_and_masks(filters_or_masks, request):
-    """Test filter and mask deserialization."""
-    filter_or_mask = request.getfixturevalue(filters_or_masks)
-    type = getattr(cs, filter_or_mask.__class__.__name__)
-    test = type.from_json(filter_or_mask.to_json())
-    assert test.to_json() == filter_or_mask.to_json()
-    np.testing.assert_allclose(
-        test.evaluate(), filter_or_mask.evaluate(), rtol=1e-6
-    )
-
-
-@pytest.mark.parametrize(
-    "model", ["noisy_model", "maskless_model", "likelihood_model"]
-)
-def test_deserialize_model(model, request):
-    """Test model deserialization."""
-    model = request.getfixturevalue(model)
-    cls = getattr(cs, model.__class__.__name__)
-    test_model = cls.from_json(model.to_json())
-    assert model.to_json() == test_model.to_json()
-    np.testing.assert_allclose(test_model.render(), model.render(), rtol=1e-6)
-
-
-def test_serialize_function():
-    """Test function serialization as a dataclass field."""
-
-    def f(x):
-        return 2 * x
-
-    x = 10
-    kernel = cs.Custom(function=f)
-    deserialized_kernel = cs.Custom.from_json(kernel.to_json())
-    assert kernel(x) == deserialized_kernel(x)
diff --git a/tests/test_shape.py b/tests/test_shape.py
index 5b6890d2..cdb653b2 100644
--- a/tests/test_shape.py
+++ b/tests/test_shape.py
@@ -1,11 +1,24 @@
 import pytest
 
 
-@pytest.mark.parametrize("model", ["noisy_model", "maskless_model"])
-def test_shape(model, request):
-    """Make sure shapes are as expected"""
+@pytest.mark.parametrize("model", ["noisy_model", "filtered_and_masked_model"])
+def test_real_shape(model, request):
+    """Make sure shapes are as expected in real space."""
     model = request.getfixturevalue(model)
     image = model()
     padded_image = model(view=False)
-    assert image.shape == model.manager.shape
-    assert padded_image.shape == model.manager.padded_shape
+    assert image.shape == model.scattering.manager.shape
+    assert padded_image.shape == model.scattering.manager.padded_shape
+
+
+@pytest.mark.parametrize("model", ["noisy_model", "filtered_and_masked_model"])
+def test_fourier_shape(model, request):
+    """Make sure shapes are as expected in fourier space."""
+    model = request.getfixturevalue(model)
+    image = model(get_real=False)
+    padded_image = model(view=False, get_real=False)
+    assert image.shape == model.scattering.manager.frequency_grid.shape[0:2]
+    assert (
+        padded_image.shape
+        == model.scattering.manager.padded_frequency_grid.shape[0:2]
+    )
diff --git a/tests/test_voxels_from_atoms.py b/tests/test_voxels_from_atoms.py
new file mode 100644
index 00000000..018cee37
--- /dev/null
+++ b/tests/test_voxels_from_atoms.py
@@ -0,0 +1,134 @@
+import jax.numpy as jnp
+import numpy as np
+import pytest
+from cryojax.simulator.density import VoxelGrid, VoxelCloud
+from cryojax.simulator.density._voxel_density import (
+    _build_real_space_voxels_from_atoms,
+)
+from cryojax.utils import ifftn, make_coordinates
+from jax import config
+
+config.update("jax_enable_x64", True)
+
+
+def test_VoxelGrid_VoxelCloud_agreement(sample_pdb_path):
+    """
+    Integration test ensuring that the VoxelGrid and VoxelCloud classes
+    produce comparable electron densities when loaded from PDB.
+    """
+    n_voxels_per_side = (128, 128, 128)
+    voxel_size = 0.5
+
+    # Load the PDB file into a VoxelGrid
+    vg = VoxelGrid.from_pdb(
+        sample_pdb_path,
+        n_voxels_per_side=n_voxels_per_side,
+        voxel_size=voxel_size,
+    )
+    # Since Voxelgrid is in Frequency space by default, we have to first
+    # transform back into real space.
+    vg_density = ifftn(vg.weights).real
+    # The constructors each transpose in a unique way in order for
+    # jax-finufft and the fourier slice theorem to match each other
+    # and cisTEM. This operation undos the difference between both transposes
+    vg_density = jnp.transpose(vg_density, axes=[1, 0, 2])
+    # Ravel the grid
+    vg_density = vg_density.ravel()
+
+    vc = VoxelCloud.from_pdb(
+        sample_pdb_path,
+        n_voxels_per_side=n_voxels_per_side,
+        voxel_size=voxel_size,
+        mask_zeros=False,
+    )
+
+    np.testing.assert_allclose(vg_density, vc.weights, atol=1e-12)
+
+
+class TestBuildRealSpaceVoxelsFromAtoms:
+    @pytest.mark.parametrize("largest_atom", range(0, 3))
+    def test_maxima_are_in_write_positions(self, largest_atom):
+        """
+        Test that the maxima of the density are in the correct positions.
+        """
+        atom_positions = jnp.array(
+            [
+                [0.0, 0.0, 0.0],
+                [0.0, 0.0, 1.0],
+                [1.0, 0.0, 0.0],
+                [0.0, 1.0, 0.0],
+            ]
+        )
+        num_atoms = atom_positions.shape[0]
+        ff_a = np.array(
+            num_atoms
+            * [
+                [1.0, 0.5],
+            ]
+        )
+        ff_a[
+            largest_atom
+        ] += 0.5  # Give one atom more weight for testing purposes
+        ff_a = jnp.array(ff_a)
+        ff_b = jnp.array(
+            num_atoms
+            * [
+                [0.3, 0.2],
+            ]
+        )
+
+        # Build the coordinate system
+        n_voxels_per_side = (128, 128, 128)
+        voxel_size = 0.1
+        coordinate_system = make_coordinates(n_voxels_per_side, voxel_size)
+
+        # Build the density
+        density = _build_real_space_voxels_from_atoms(
+            atom_positions, ff_a, ff_b, coordinate_system
+        )
+
+        # Find the maximum
+        maximum_index = jnp.argmax(density)
+        maximum_position = coordinate_system.reshape(-1, 3)[maximum_index]
+
+        # Check that the maximum is in the correct position
+        assert jnp.allclose(maximum_position, atom_positions[largest_atom])
+
+    def test_integral_is_correct(self):
+        """
+        Test that the maxima of the density are in the correct positions.
+        """
+        atom_positions = jnp.array(
+            [
+                [0.0, 0.0, 0.0],
+                [0.0, 0.0, 1.0],
+                [1.0, 0.0, 0.0],
+                [0.0, 1.0, 0.0],
+            ]
+        )
+        num_atoms = atom_positions.shape[0]
+        ff_a = np.array(
+            num_atoms
+            * [
+                [1.0, 0.5],
+            ]
+        )
+
+        ff_b = jnp.array(
+            num_atoms
+            * [
+                [0.3, 0.2],
+            ]
+        )
+
+        n_voxels_per_side = (128, 128, 128)
+        voxel_size = 0.05
+        coordinate_system = make_coordinates(n_voxels_per_side, voxel_size)
+
+        # Build the density
+        density = _build_real_space_voxels_from_atoms(
+            atom_positions, ff_a, ff_b, coordinate_system
+        )
+
+        integral = jnp.sum(density) * voxel_size**3
+        assert jnp.isclose(integral, jnp.sum(ff_a))
diff --git a/tutorials/test-ctf.ipynb b/tutorials/test-ctf.ipynb
deleted file mode 100644
index 1e2b0552..00000000
--- a/tutorials/test-ctf.ipynb
+++ /dev/null
@@ -1,204 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Jax imports\n",
-    "import jax.numpy as jnp\n",
-    "import numpy as np\n",
-    "from jax import config\n",
-    "\n",
-    "config.update(\"jax_enable_x64\", False)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Plotting imports\n",
-    "from matplotlib import pyplot as plt\n",
-    "from mpl_toolkits.axes_grid1 import make_axes_locatable"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def plot_image(image, fig, ax, cmap=\"gray\", **kwargs):\n",
-    "    im = ax.imshow(image, cmap=cmap, origin=\"lower\", **kwargs)\n",
-    "    divider = make_axes_locatable(ax)\n",
-    "    cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
-    "    fig.colorbar(im, cax=cax)\n",
-    "    return fig, ax\n",
-    "\n",
-    "def plot_images(images, cmap=\"gray\", labels=None, **kwargs):\n",
-    "    ncols = len(images)\n",
-    "    fig, axes = plt.subplots(ncols=ncols, figsize=(ncols*5, 5))\n",
-    "    for idx in range(ncols):\n",
-    "        ax = axes[idx]\n",
-    "        image = images[idx]\n",
-    "        im = ax.imshow(image, cmap=cmap, origin=\"lower\", **kwargs)\n",
-    "        label = None if labels is None else labels[idx]\n",
-    "        ax.set(title=label)\n",
-    "        divider = make_axes_locatable(ax)\n",
-    "        cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
-    "        fig.colorbar(im, cax=cax)\n",
-    "    return fig, axes\n",
-    "\n",
-    "def plot_profiles(bins, profiles, labels=None, **kwargs):\n",
-    "    nprofiles = len(profiles)\n",
-    "    fig, ax = plt.subplots()\n",
-    "    ax.set(**kwargs)\n",
-    "    for idx in range(nprofiles):\n",
-    "        profile = profiles[idx]\n",
-    "        label = None if labels is None else labels[idx]\n",
-    "        ax.plot(bins, profile, label=label)\n",
-    "    ax.legend(fontsize=12)\n",
-    "    return fig, ax"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# cisTEM imports\n",
-    "from pycistem.core import CTF"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# cryojax imports\n",
-    "from cryojax.simulator import CTFOptics\n",
-    "from cryojax.utils import make_frequencies, cartesian_to_polar, powerspectrum"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Parameters for the CTF\n",
-    "defocus1, defocus2, asti_angle, kV, cs, ac = 24000, 12000, 30.0, 300.0, 2.7, 0.07"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Frequency coordinates\n",
-    "shape = (512, 512)\n",
-    "pixel_size = 0.9\n",
-    "freqs = make_frequencies(shape, pixel_size=pixel_size)\n",
-    "k_sqr, theta = cartesian_to_polar(freqs, square=True)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# cryojax CTF and power spectrum\n",
-    "optics = CTFOptics(defocus_u=defocus1, defocus_v=defocus2,\n",
-    "                   defocus_angle=asti_angle, voltage=kV,\n",
-    "                   spherical_aberration=cs, amplitude_contrast=ac, envelope=None)\n",
-    "ctf = np.array(optics(freqs))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# cisTEM CTF\n",
-    "cisTEM_optics = CTF(kV=kV, cs=cs, ac=ac, defocus1=defocus1, defocus2=defocus2, astig_angle=asti_angle, pixel_size=pixel_size)\n",
-    "cisTEM_ctf = np.vectorize(lambda k_sqr, theta: cisTEM_optics.Evaluate(k_sqr, theta))(k_sqr.ravel()*pixel_size**2, theta.ravel()).reshape(shape)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Plot CTFs\n",
-    "ctfs = [ctf, cisTEM_ctf]\n",
-    "labels = [\"cryojax\", \"cisTEM\"]\n",
-    "fig, axes = plot_images(ctfs, labels=labels)\n",
-    "plt.tight_layout()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Compute cryojax and cisTEM power spectrum\n",
-    "spectrum1D, k_bins = powerspectrum(ctf, freqs, pixel_size=pixel_size)\n",
-    "cisTEM_spectrum1D, k_bins = powerspectrum(cisTEM_ctf, freqs, pixel_size=pixel_size)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Plot profiles\n",
-    "profiles = [spectrum1D, cisTEM_spectrum1D]\n",
-    "fig, axes = plot_profiles(k_bins, profiles, labels=labels)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Assert CTFs are the same\n",
-    "assert np.allclose(ctf, cisTEM_ctf, atol=5e-2)\n",
-    "assert np.allclose(spectrum1D, cisTEM_spectrum1D, atol=5e-3)"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "jax-2dtm-cpu",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.10.12"
-  },
-  "orig_nbformat": 4
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}