document cryojax.constants

docs/api/constants/scattering_factor_parameters.md

@@ -0,0 +1,11 @@
+# Scattering factor parameters
+
+Modeling the electron scattering amplitudes of individual atoms is an important component of modeling cryo-EM images, as these are typically used to approximate the electrostatic potential. Typically, the scattering factor for each individual atom is numerically approximated with a fixed functional form but varying parameters for different atoms. These parameters are stored in lookup tables in the literature. This documentation provides these lookup tables so that they may be used to compute electrostatic potentials in cryoJAX.
+
+## Extracting parameters from a lookup table
+
+::: cryojax.constants.get_tabulated_scattering_factor_parameters
+
+## Lookup tables
+
+::: cryojax.constants.read_peng_element_scattering_factor_parameter_table

docs/api/constants/units.md

@@ -0,0 +1,5 @@
+# Working with physical units
+
+Here, convenience methods for working with physical units are described.
+
+::: cryojax.constants.convert_keV_to_angstroms

mkdocs.yml

@@ -100,6 +100,9 @@ nav:
         - cryojax.rotations:
             - 'api/rotations/lie_groups.md'
             - 'api/rotations/conversions.md'
+        - cryojax.constants:
+            - 'api/constants/units.md'
+            - 'api/constants/scattering_factor_parameters.md'
         - cryojax.utils:
             - 'api/utils/filter_specs.md'
             - 'api/utils/transforms.md'

src/cryojax/constants/__init__.py

@@ -1,6 +1,6 @@
 from ._scattering_factor_parameters import (
     get_tabulated_scattering_factor_parameters as get_tabulated_scattering_factor_parameters,  # noqa: E501
-    peng_element_scattering_factor_parameter_table as peng_element_scattering_factor_parameter_table,  # noqa: E501
+    read_peng_element_scattering_factor_parameter_table as read_peng_element_scattering_factor_parameter_table,  # noqa: E501
 )
 from ._unit_conversions import (
     convert_keV_to_angstroms as convert_keV_to_angstroms,

src/cryojax/constants/_scattering_factor_parameters.py

@@ -5,39 +5,22 @@
 
 import importlib.resources as pkg_resources
 import os
+from typing import Optional
 
 import jax.numpy as jnp
 import numpy as np
 from jaxtyping import Array, Float, Int
 
 
-def _load_peng_element_scattering_factor_parameter_table():
-    """Internal function to load the atomic scattering factor parameter
-    table."""
-    with pkg_resources.as_file(
-        pkg_resources.files("cryojax").joinpath("constants")
-    ) as path:
-        atom_scattering_factor_params = jnp.load(
-            os.path.join(path, "peng1996_element_params.npy")
-        )
-
-    return jnp.asarray(atom_scattering_factor_params)
-
-
-peng_element_scattering_factor_parameter_table = (
-    _load_peng_element_scattering_factor_parameter_table()
-)
-
-
 def get_tabulated_scattering_factor_parameters(
     atom_identities: Int[Array, " n_atoms"] | Int[np.ndarray, " n_atoms"],
-    scattering_factor_parameter_table: (
+    scattering_factor_parameter_table: Optional[
         Float[Array, "n_params n_elements n_scattering_factors"]
         | Float[np.ndarray, "n_params n_elements n_scattering_factors"]
-    ) = peng_element_scattering_factor_parameter_table,
+    ] = None,
 ) -> Float[Array, "n_params n_atoms n_scattering_factors"]:
     """Gets the parameters for the scattering factor for each atom in
-    `atom_names`.
+    `atom_identities`.
 
     **Arguments:**
 
@@ -47,11 +30,38 @@ def get_tabulated_scattering_factor_parameters(
     - `scattering_factor_parameter_table`:
         Array containing the table of scattering factors. By default, this
         is the tabulation from "Robust Parameterization of Elastic and
-        Absorptive Electron Atomic Scattering Factors" by Peng et al. (1996).
+        Absorptive Electron Atomic Scattering Factors" by Peng et al. (1996),
+        given by `load_peng_element_scattering_factor_parameter_table`.
 
     **Returns:**
 
     The particular scattering factor parameters stored in
     `scattering_factor_parameter_table` for `atom_identities`.
     """
+    if scattering_factor_parameter_table is None:
+        scattering_factor_parameter_table = (
+            read_peng_element_scattering_factor_parameter_table()
+        )
     return jnp.asarray(scattering_factor_parameter_table)[:, jnp.asarray(atom_identities)]
+
+
+def read_peng_element_scattering_factor_parameter_table() -> (
+    Float[np.ndarray, "2 n_elements 5"]
+):
+    r"""Function to load the atomic scattering factor parameter
+    table from *"Robust Parameterization of Elastic and Absorptive
+    Electron Atomic Scattering Factors" by Peng et al. (1996).*.
+
+    **Returns:**
+
+    The parameter table for parameters $\{a_i\}_{i = 1}^5$ and $\{b_i\}_{i = 1}^5$
+    for each atom, described in the above reference.
+    """
+    with pkg_resources.as_file(
+        pkg_resources.files("cryojax").joinpath("constants")
+    ) as path:
+        atom_scattering_factor_params = jnp.load(
+            os.path.join(path, "peng1996_element_params.npy")
+        )
+
+    return np.asarray(atom_scattering_factor_params)

src/cryojax/constants/_unit_conversions.py

@@ -7,6 +7,19 @@
 def convert_keV_to_angstroms(
     energy_in_keV: Float[Array, ""] | float,
 ) -> Float[Array, ""]:
-    """Get the relativistic electron wavelength at a given accelerating voltage."""
+    """Get the relativistic electron wavelength at a given accelerating voltage. For
+    reference, see Equation 2.5 in Section 2.1 from *Spence, John CH. High-resolution
+    electron microscopy. OUP Oxford, 2013.*.
+
+    **Arguments:**
+
+    - `energy_in_keV`:
+        The energy in kiloelectron volts.
+
+    **Returns:**
+
+    The relativistically corrected electron wavelength in Angstroms corresponding to the
+    energy `energy_in_keV`.
+    """
     energy_in_eV = 1000.0 * energy_in_keV  # keV to eV
-    return jnp.asarray(12.2643 / (energy_in_eV + 0.97845e-6 * energy_in_eV**2) ** 0.5)
+    return jnp.asarray(12.2639 / (energy_in_eV + 0.97845e-6 * energy_in_eV**2) ** 0.5)

src/cryojax/simulator/_potential_representation/atom_potential.py

@@ -16,7 +16,7 @@
 
 from ...constants import (
     get_tabulated_scattering_factor_parameters,
-    peng_element_scattering_factor_parameter_table,
+    read_peng_element_scattering_factor_parameter_table,
 )
 from ...coordinates import make_1d_coordinate_grid
 from ...internal import error_if_negative, error_if_not_positive
@@ -248,18 +248,33 @@ def __init__(
         b_factors: Optional[
             Float[Array, " n_atoms"] | Float[np.ndarray, " n_atoms"]
         ] = None,
+        *,
+        scattering_factor_parameter_table: Optional[
+            Float[Array, "2 n_elements 5"] | Float[np.ndarray, "2 n_elements 5"]
+        ] = None,
     ):
         """**Arguments:**
 
-        - `atom_positions`: The coordinates of the atoms in units of angstroms.
-        - `atom_identities`: Array containing the index of the one-hot encoded atom names.
-                             Hydrogen is "1", Carbon is "6", Nitrogen is "7", etc.
-        - `b_factors`: The B-factors applied to each atom.
+        - `atom_positions`:
+            The coordinates of the atoms in units of angstroms.
+        - `atom_identities`:
+            Array containing the index of the one-hot encoded atom names.
+            Hydrogen is "1", Carbon is "6", Nitrogen is "7", etc.
+        - `b_factors`:
+            The B-factors applied to each atom.
+        - `scattering_factor_parameter_table`:
+            The scattering factor parameter table from Peng et al. (1996). If
+            not provided, load from `cryojax.constants`.
+
         """
+        if scattering_factor_parameter_table is None:
+            scattering_factor_parameter_table = (
+                read_peng_element_scattering_factor_parameter_table()
+            )
         self.atom_positions = jnp.asarray(atom_positions)
         scattering_factor_a, scattering_factor_b = (
             get_tabulated_scattering_factor_parameters(
-                atom_identities, peng_element_scattering_factor_parameter_table
+                atom_identities, scattering_factor_parameter_table
             )
         )
         self.scattering_factor_a = jnp.asarray(scattering_factor_a)