Merge pull request #556 from michael-0brien/padded-shape

Deprecate `BasicImageConfig(..., pad_options=...)` -> `BasicImageConfig(..., padded_shape=...)`

Author: michael-0brien

benchmarks/benchmark_jit.py

@@ -67,13 +67,11 @@ def make_particle_parameters(key: PRNGKeyArray, config: cxs.BasicImageConfig):
         }
 
     # Generate particle parameters. First, the image config
-    pad_options = dict(shape=(200, 200))
-
     config = cxs.BasicImageConfig(
         shape=(150, 150),
         pixel_size=2.0,
         voltage_in_kilovolts=300.0,
-        pad_options=pad_options,
+        padded_shape=(200, 200),
     )
     # ... RNG keys
     keys = jax.random.split(jax.random.key(0), num_images)

benchmarks/benchmark_projection_method_tradeoff.py

@@ -66,12 +66,12 @@ def setup_volumes_and_configs(n_iterations, n_atoms, box_size, pixel_size=2.0):
     )
 
     # Image config
-    pad_options = dict(shape=(int(box_size * 1.5), int(box_size * 1.5)))
+    padded_shape = (int(box_size * 1.5), int(box_size * 1.5))
     config = cxs.BasicImageConfig(
         shape=(box_size, box_size),
         pixel_size=pixel_size,
         voltage_in_kilovolts=300.0,
-        pad_options=pad_options,
+        padded_shape=padded_shape,
     )
 
     return volume_fourier_grid, avg_time, volume_gmm, atom_volume, config

cryojax/dataset/_particle_data/relion.py

@@ -1148,9 +1148,10 @@ def _make_config(
     voltage_in_kilovolts,
     pad_options,
 ):
+    padded_shape = None if "shape" not in pad_options else pad_options["shape"]
     return eqx.tree_at(
         lambda x: (x.pixel_size, x.voltage_in_kilovolts),
-        BasicImageConfig(image_shape, 1.0, 1.0, pad_options=pad_options),
+        BasicImageConfig(image_shape, 1.0, 1.0, padded_shape=padded_shape),
         (pixel_size, voltage_in_kilovolts),
     )
 

cryojax/io/_pdb.py

@@ -11,6 +11,7 @@
 from typing import Any, Literal, TypedDict, overload
 from xml.etree import ElementTree
 
+import equinox.internal as eqxi
 import jax
 import mdtraj
 import mmdf
@@ -23,10 +24,10 @@
 
 
 if hasattr(typing, "GENERATING_DOCUMENTATION"):
-    AtomProperties = dict  # pyright: ignore[reportAssignmentType]
+    _AtomProperties = dict[str, Any]  # pyright: ignore[reportAssignmentType]
 else:
 
-    class AtomProperties(TypedDict):
+    class _AtomProperties(TypedDict):
         masses: Float[np.ndarray, "... n_atoms"]
         b_factors: Float[np.ndarray, "... n_atoms"]
         charges: Float[np.ndarray, "... n_atoms"]
@@ -62,7 +63,7 @@ def read_atoms_from_pdb(  # type: ignore
 ) -> tuple[
     Float[np.ndarray, "... n_atoms 3"],
     Int[np.ndarray, " n_atoms"],
-    AtomProperties,
+    dict[str, Any],
 ]: ...
 
 
@@ -81,6 +82,7 @@ def read_atoms_from_pdb(
 ) -> tuple[Float[np.ndarray, "... n_atoms 3"], Int[np.ndarray, "... n_atoms"]]: ...
 
 
+@eqxi.doc_remove_args("loads_b_factors")
 def read_atoms_from_pdb(
     filename: str | pathlib.Path,
     *,
@@ -203,7 +205,7 @@ def mmdf_to_atoms(  # type: ignore
 ) -> tuple[
     Float[np.ndarray, "... n_atoms 3"],
     Int[np.ndarray, "... n_atoms"],
-    AtomProperties,
+    dict[str, Any],
 ]: ...
 
 
@@ -408,7 +410,7 @@ def read_topology_from_pdb(
 class _AtomicModelInfo(TypedDict):
     positions: Float[np.ndarray, "M N 3"]
     numbers: Int[np.ndarray, "M N 3"]
-    properties: AtomProperties
+    properties: _AtomProperties
 
 
 def _load_atom_info(df: pd.DataFrame, model_index: int | None, stack_models: bool):
@@ -450,7 +452,7 @@ def _load_atom_info(df: pd.DataFrame, model_index: int | None, stack_models: boo
             )
 
     # Gather atom info and return
-    properties = AtomProperties(
+    properties = _AtomProperties(
         charges=np.asarray(charges, dtype=int),
         b_factors=np.asarray(b_factors, dtype=float),
         masses=np.asarray(atom_masses, dtype=float),

cryojax/simulator/_image_config.py

@@ -3,9 +3,10 @@
 import math
 import warnings
 from functools import cached_property
-from typing import Any, Literal, TypedDict
+from typing import Literal, TypedDict
 
 import equinox as eqx
+import equinox.internal as eqxi
 import jax
 import jax.numpy as jnp
 from jaxtyping import Array, Float
@@ -29,8 +30,8 @@
 
 # Not currently public API
 class PrecomputedGrids(eqx.Module, strict=True):
-    only_rfft: bool
-    only_fourier: bool
+    only_rfft: bool = eqx.field(static=True)
+    only_fourier: bool = eqx.field(static=True)
 
     _frequency_grid: Float[Array, "_ _ 2"]
     _coordinate_grid: Float[Array, "_ _ 2"] | None
@@ -138,8 +139,7 @@ class AbstractImageConfig(eqx.Module, strict=True):
     pixel_size: eqx.AbstractVar[Float[Array, ""]]
     voltage_in_kilovolts: eqx.AbstractVar[Float[Array, ""]]
 
-    pad_options: eqx.AbstractVar[dict[str, Any]]
-
+    padded_shape: eqx.AbstractVar[tuple[int, int]]
     precompute_mode: eqx.AbstractVar[
         Literal["none", "rfft", "fft", "all", "compile_time_eval"]
     ]
@@ -160,12 +160,6 @@ def __check_init__(self):
                 f"Found that `{cls}.padded_shape` is less than `{cls}.shape` in one or "
                 " more dimensions."
             )
-        if set(self.pad_options.keys()) != {"shape"}:
-            raise AttributeError(
-                f"Found that `{cls}.pad_options` was not a dictionary with "
-                "key 'shape'. "
-                f"Instead, had keys {set(self.pad_options.keys())}."
-            )
 
     @property
     def wavelength_in_angstroms(self) -> Float[Array, ""]:
@@ -278,10 +272,6 @@ def _get_grid_impl(_self):
 
         return frequency_grid
 
-    @property
-    def padded_shape(self):
-        return self.pad_options["shape"]
-
     @property
     def n_pixels(self) -> int:
         """Convenience property for `math.prod(shape)`"""
@@ -544,21 +534,24 @@ class BasicImageConfig(AbstractImageConfig, strict=True):
     pixel_size: Float[Array, ""]
     voltage_in_kilovolts: Float[Array, ""]
 
-    pad_options: PadOptions
-
-    precompute_mode: Literal["none", "rfft", "fft", "all", "compile_time_eval"]
+    padded_shape: tuple[int, int]
+    precompute_mode: Literal["none", "rfft", "fft", "all", "compile_time_eval"] = (
+        eqx.field(static=True)
+    )
     precomputed_grids: PrecomputedGrids | None
 
+    @eqxi.doc_remove_args("pad_options")
     def __init__(
         self,
         shape: tuple[int, int],
         pixel_size: FloatLike,
         voltage_in_kilovolts: FloatLike,
         *,
+        padded_shape: tuple[int, int] | None = None,
         precompute_mode: Literal[
             "none", "rfft", "fft", "all", "compile_time_eval"
         ] = "none",
-        pad_options: dict[str, Any] = {},
+        pad_options: dict = {},
     ):
         """**Arguments:**
 
@@ -568,11 +561,9 @@ def __init__(
             The pixel size of the image in angstroms.
         - `voltage_in_kilovolts`:
             The incident energy of the electron beam.
-        - `pad_options`:
-            Options that control image padding.
-            - 'shape':
-                The shape of the image after padding. By default, equal
-                to `shape`.
+        - `padded_shape`:
+            The shape of the image after padding. By default, equal
+            to `shape`.
         - `precompute_mode`:
             How to pre-compute coordinate and frequency grids stored in
             the `image_config`. Options are
@@ -597,10 +588,17 @@ def __init__(
         # Set parameters
         self.pixel_size = jnp.asarray(pixel_size, dtype=float)
         self.voltage_in_kilovolts = jnp.asarray(voltage_in_kilovolts, dtype=float)
-        # Set shape
+        # Set shape and padded shape
+        if "shape" in pad_options:
+            warnings.warn(
+                "`BasicImageConfig(..., pad_options=...)` is deprecated and will "
+                "be removed in cryoJAX 0.6.0. Use `padded_shape` instead.",
+                category=FutureWarning,
+                stacklevel=2,
+            )
+            padded_shape = pad_options["shape"]
         self.shape = shape
-        # Set pad options
-        self.pad_options = _dict_to_pad_options(pad_options, shape)
+        self.padded_shape = shape if padded_shape is None else padded_shape
         # Finally, grid precompute
         if precompute_mode == "rfft":
             self.precomputed_grids = PrecomputedGrids(
@@ -628,22 +626,25 @@ class DoseImageConfig(AbstractImageConfig, strict=True):
     voltage_in_kilovolts: Float[Array, ""]
     electron_dose: Float[Array, ""]
 
-    pad_options: PadOptions
-
-    precompute_mode: Literal["none", "rfft", "fft", "all", "compile_time_eval"]
+    padded_shape: tuple[int, int]
+    precompute_mode: Literal["none", "rfft", "fft", "all", "compile_time_eval"] = (
+        eqx.field(static=True)
+    )
     precomputed_grids: PrecomputedGrids | None
 
+    @eqxi.doc_remove_args("pad_options")
     def __init__(
         self,
         shape: tuple[int, int],
         pixel_size: FloatLike,
         voltage_in_kilovolts: FloatLike,
         electron_dose: FloatLike,
         *,
+        padded_shape: tuple[int, int] | None = None,
         precompute_mode: Literal[
             "none", "rfft", "fft", "all", "compile_time_eval"
         ] = "none",
-        pad_options: dict[str, Any] = {},
+        pad_options: dict = {},
     ):
         """**Arguments:**
 
@@ -656,12 +657,9 @@ def __init__(
         - `electron_dose`:
             The integrated dose rate of the electron beam in
             $e^-/A^2$
-        - `pad_options`:
-            Options that control image padding. This is a dictionary
-            with keys:
-            - 'shape':
-                The shape of the image after padding. By default, equal
-                to `shape`.
+        - `padded_shape`:
+            The shape of the image after padding. By default, equal
+            to `shape`.
         - `precompute_mode`:
             How to pre-compute coordinate and frequency grids stored in
             the `image_config`. Options are
@@ -687,10 +685,17 @@ def __init__(
         self.pixel_size = jnp.asarray(pixel_size, dtype=float)
         self.voltage_in_kilovolts = jnp.asarray(voltage_in_kilovolts, dtype=float)
         self.electron_dose = jnp.asarray(electron_dose, dtype=float)
-        # Set shape
+        # Set shape and padded shape
+        if "shape" in pad_options:
+            warnings.warn(
+                "`BasicImageConfig(..., pad_options=...)` is deprecated and will "
+                "be removed in cryoJAX 0.6.0. Use `padded_shape` instead.",
+                category=FutureWarning,
+                stacklevel=2,
+            )
+            padded_shape = pad_options["shape"]
         self.shape = shape
-        # Set pad options
-        self.pad_options = _dict_to_pad_options(pad_options, shape)
+        self.padded_shape = shape if padded_shape is None else padded_shape
         # Finally, grid precompute
         if precompute_mode == "rfft":
             self.precomputed_grids = PrecomputedGrids(
@@ -728,16 +733,3 @@ def _safe_multiply_by_constant(
     grid = grid.at[:, 1:, 0].multiply(constant)
     grid = grid.at[1:, :, 1].multiply(constant)
     return grid
-
-
-def _dict_to_pad_options(d: dict[str, Any], default_shape: tuple[int, int]) -> PadOptions:
-    if not set(d.keys()).issubset({"shape"}):
-        raise ValueError(
-            "Expected that dictionary `pad_options` passed to "
-            "`BasicImageConfig(..., pad_options=...)` "
-            f"had a subset of keys {{'shape'}}, but found that it had keys "
-            f"{set(d.keys())}."
-        )
-    shape = d["shape"] if "shape" in d else default_shape
-
-    return PadOptions(shape=shape)

docs/api/simulator/config.md

@@ -8,17 +8,14 @@ The `AbstractImageConfig` is an object at the core of simulating images in `cryo
                 - shape
                 - pixel_size
                 - voltage_in_kilovolts
-                - pad_options
+                - padded_shape
 
 ---
 
 ::: cryojax.simulator.BasicImageConfig
         options:
             members:
                 - __init__
-                - shape
-                - pixel_size
-                - voltage_in_kilovolts
                 - wavelength_in_angstroms
                 - lorentz_factor
                 - interaction_constant
@@ -27,7 +24,6 @@ The `AbstractImageConfig` is an object at the core of simulating images in `cryo
                 - n_pixels
                 - y_dim
                 - x_dim
-                - padded_shape
                 - padded_n_pixels
                 - padded_y_dim
                 - padded_x_dim
@@ -38,10 +34,6 @@ The `AbstractImageConfig` is an object at the core of simulating images in `cryo
         options:
             members:
                 - __init__
-                - shape
-                - pixel_size
-                - voltage_in_kilovolts
-                - electron_dose
                 - wavelength_in_angstroms
                 - lorentz_factor
                 - interaction_constant
@@ -50,7 +42,6 @@ The `AbstractImageConfig` is an object at the core of simulating images in `cryo
                 - n_pixels
                 - y_dim
                 - x_dim
-                - padded_shape
                 - padded_n_pixels
                 - padded_y_dim
                 - padded_x_dim

docs/examples/extending-cryojax.ipynb

@@ -86,12 +86,11 @@
     "    psi_angle=-5.0,\n",
     ")\n",
     "# ... the configuration. Add padding with respect to the final image shape.\n",
-    "pad_options = dict(shape=volume.shape[0:2])\n",
     "config = cxs.BasicImageConfig(\n",
     "    shape=(80, 80),\n",
     "    pixel_size=voxel_size,\n",
     "    voltage_in_kilovolts=300.0,\n",
-    "    pad_options=pad_options,\n",
+    "    padded_shape=volume.shape[0:2],\n",
     ")\n",
     "\n",
     "\n",

docs/examples/simulate-image.ipynb

@@ -80,7 +80,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -100,12 +100,11 @@
     ")\n",
     "transfer_theory = cxs.ContrastTransferTheory(ctf, amplitude_contrast_ratio=0.1)\n",
     "# Then the configuration. Add padding to avoid periodic artifacts due to CTF rings\n",
-    "pad_options = dict(shape=volume.shape[0:2])\n",
     "image_config = cxs.BasicImageConfig(\n",
     "    shape=(80, 80),\n",
     "    pixel_size=voxel_size,\n",
     "    voltage_in_kilovolts=300.0,\n",
-    "    pad_options=pad_options,\n",
+    "    padded_shape=volume.shape[0:2],\n",
     ")"
    ]
   },