Repository: swagnercarena/paltax (language: python)

Issue:
Fix for pixel PSF and source model tests
Tiny fixes:
  * The Pixel PSF test case was missed since it was indented inside the Gaussian PSF test case. With minor tweaks it passes now.
  * Fix a remaining self-import (remove `third_party` etc.)


Discussion / hints:

Fix patch:
diff --git a/jaxstronomy/__init__.py b/jaxstronomy/__init__.py
index b51b867..32f7c36 100644
--- a/jaxstronomy/__init__.py
+++ b/jaxstronomy/__init__.py
@@ -1,6 +1,7 @@
 from . import cosmology_utils
 from . import image_simulation
 from . import lens_models
+from . import los
 from . import power_spectrum
 from . import psf_models
 from . import source_models
diff --git a/jaxstronomy/cosmology_utils.py b/jaxstronomy/cosmology_utils.py
index e2849cd..6560cc8 100644
--- a/jaxstronomy/cosmology_utils.py
+++ b/jaxstronomy/cosmology_utils.py
@@ -1,5 +1,4 @@
 # coding=utf-8
-# Copyright 2022 The Google Research Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -7,20 +6,6 @@
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Copyright 2022 Google LLC
-
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-
-#     https://www.apache.org/licenses/LICENSE-2.0
-
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -43,6 +28,9 @@
 from jaxstronomy import power_spectrum
 
 SPEED_OF_LIGHT = 2.99792458e8
+GRAVITATIONAL_CONSTANT = 6.67384e-11
+MEGAPARSEC_METERS = 3.0856775800000002e+22
+MASS_SUN = 1.9891000000000002e+30
 # Linear density threshold for collapse.
 DELTA_COLLAPSE_UNCORRECTED = 1.68647
 # Critical density of the universe at redshift 0 in units of
@@ -52,510 +40,574 @@
 GROWTH_Z_MAX = 3e3
 
 
-def _e_z(cosmology_params,
-         z_values):
-  """Calculate the hubble parameter.
-
-  Args:
-    cosmology_params: Cosmological parameters that define the universe's
-      expansion, potentially excluding lookup table.
-    z_values: Redshifts at which to calculate the hubble parameter.
-
-  Returns:
-    The hubble parameter values in units of H_0 (km/s/Mpc).
-
-  Notes:
-    This assumes Lambda dark energy (i.e. no redshift dependence in the dark
-    energy equation of state).
-  """
-  t = cosmology_params['omega_m_zero'] * (1.0 + z_values)**3.0
-  t += cosmology_params['omega_de_zero']
-  t += cosmology_params['omega_rad_zero'] * (1.0 + z_values)**4.0
-
-  return jnp.sqrt(t)
-
-
-def _e_z_rad_to_dark(cosmology_params,
-                     z_values):
-  """Calculate the hubble parameter with radiation treated as dark energy.
-
-  Args:
-    cosmology_params: Cosmological parameters that define the universe's
-      expansion, potentially excluding lookup table.
-    z_values: Redshifts at which to calculate the hubble parameter.
-
-  Returns:
-    The hubble parameter values in units of H_0 (km/s/Mpc).
-
-  Notes:
-    Used for growth factor calculations.
-  """
-  t = cosmology_params['omega_m_zero'] * (1.0 + z_values)**3.0
-  t += cosmology_params['omega_de_zero']
-  t += cosmology_params['omega_rad_zero']
-  return jnp.sqrt(t)
-
-
-def _comoving_distance_numerical(init_cosmology_params,
-                                 z_min, z_max):
-  """Calculates the comoving distance at a given redshift numerically.
-
-  Args:
-    init_cosmology_params: Cosmological parameters that define the universe's
-      expansion.
-    z_min: Lower redshift of distance interval.
-    z_max: Upper redhsift of distance interval.
-
-  Returns:
-    Comoving distance in units of Mpc.
-  """
-  # We probably want a better / less expensive integration scheme than this.
-  z_samples = jnp.linspace(z_min, z_max, 1000)
-  e_z_samples = _e_z(init_cosmology_params, z_samples)
-
-  one_over_ez_int = jnp.trapz(1 / e_z_samples, z_samples)
-
-  # Factor in the speed of light and the Hubble constant. The factor of 1e-3
-  # comes from accounting for integral output being in Mpc * s / km and speed of
-  # light being in m/s as well as h being H0/100.
-  one_over_ez_int *= SPEED_OF_LIGHT * 1e-3
-  one_over_ez_int /= init_cosmology_params['hubble_constant']
-
-  return one_over_ez_int
-
-
-def _sigma_k_integrand(init_cosmology_params,
-                       log_k, radius):
-  """Return the integrand for the variance of the linear density field.
-
-  Args:
-    init_cosmology_params: Cosmological parameters that define the universe's
-      expansion, excluding the lookup table for the comoving_distance
-      calculation.
-    log_k: Log of the wavenumber values at which to evaluate integrand in units
-      of h/Mpc.
-    radius: Scale at which to calculate the rms variance in units of Mpc/h.
-
-  Returns:
-    Integrand at provided wavenumbers.
-  """
-  k = jnp.exp(log_k)
-  # Fourier transform of the tophat filter function.
-  k_times_r = k * radius
-  # Use masking to default to 1 when below fidelity limit of float32.
-  is_below_cut = k_times_r < 3e-2
-  ft_filer = jnp.zeros_like(k)
-  ft_filer += is_below_cut
-  ft_filer += (
-      jnp.logical_not(is_below_cut) * 3.0 / k_times_r**3 *
-      (jnp.sin(k_times_r) - k_times_r * jnp.cos(k_times_r)))
-
-  ps = power_spectrum.matter_power_spectrum(
-      k, init_cosmology_params['omega_m_zero'],
-      init_cosmology_params['omega_b_zero'],
-      init_cosmology_params['temp_cmb_zero'],
-      init_cosmology_params['hubble_constant'], init_cosmology_params['n_s'])
-
-  return ps * ft_filer**2 * k**3
-
-
-def _sigma_numerical(init_cosmology_params,
-                     radius):
-  """Calculate the variance of the linear density field numercially.
-
-  Args:
-    init_cosmology_params: Cosmological parameters that define the universe's
-      expansion, excluding the lookup table for the comoving_distance
-      calculation.
-    radius: Scale at which to calculate the rms variance in units of Mpc/h.
-
-  Returns:
-    RMS variance of the linear density field.
-  """
-  # Integartion bins are in log momentum space (k).
-  log_k_bins = jnp.linspace(
-      jnp.log10(1e-6 / radius), jnp.log10(1e6 / radius), 1000)
-  sigma_k_integrand = jnp.nan_to_num(
-      _sigma_k_integrand(init_cosmology_params, log_k_bins, radius))
-  sigma_squared = jnp.trapz(sigma_k_integrand, log_k_bins)
-
-  return jnp.sqrt(sigma_squared / (2.0 * jnp.pi**2))
-
-
-def _sigma_norm(init_cosmology_params):
-  """Calculate the normalization of the variance of the linear density field.
-
-  Args:
-    init_cosmology_params: Cosmological parameters that define the universe's
-      expansion, excluding the lookup table for the comoving_distance
-      calculation.
-
-  Returns:
-    Normalized RMS variance of the linear density field.
-  """
-  # Normalize the raw output to sigma_eight.
-  sigma_eight_raw = _sigma_numerical(init_cosmology_params, 8.0)
-
-  return init_cosmology_params['sigma_eight'] / sigma_eight_raw
-
-
-def _growth_factor_exact_unormalized(init_cosmology_params,
-                                     z):
-  """Calculate the linear growth factor, unormalized.
-
-  Args:
-    init_cosmology_params: Cosmological parameters that define the universe's
-      expansion, excluding the lookup table for the comoving_distance
-      calculation.
-    z: Redshift at which to calculate the growth factor.
-
-  Returns:
-    Growth factor at input redshift.
-
-  Notes:
-    Approximations only hold to redshift 10.
-  """
-  z_samples = jnp.logspace(
-      jnp.log10(jax.lax.max(z, 1e-6)), jnp.log10(GROWTH_Z_MAX), 10000)
-  integral = jnp.trapz(
-      (1.0 + z_samples) / _e_z_rad_to_dark(init_cosmology_params, z_samples)**3,
-      z_samples)
-
-  return (5.0 / 2.0 * init_cosmology_params['omega_m_zero'] *
-          _e_z_rad_to_dark(init_cosmology_params, z) * integral)
+def _e_z(cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z_values: Union[float, jnp.ndarray]) -> Union[float, jnp.ndarray]:
+    """Calculate the hubble parameter.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion, potentially excluding lookup table.
+        z_values: Redshifts at which to calculate the hubble parameter.
+
+    Returns:
+        The hubble parameter values in units of H_0 (km/s/Mpc).
+
+    Notes:
+        This assumes Lambda dark energy (i.e. no redshift dependence in the dark
+        energy equation of state).
+    """
+    t = cosmology_params['omega_m_zero'] * (1.0 + z_values)**3.0
+    t += cosmology_params['omega_de_zero']
+    t += cosmology_params['omega_rad_zero'] * (1.0 + z_values)**4.0
+
+    return jnp.sqrt(t)
+
+
+def _e_z_rad_to_dark(
+    cosmology_params: Mapping[str, Union[float, int]],
+        z_values: Union[float, jnp.ndarray]) -> Union[float, jnp.ndarray]:
+    """Calculate the hubble parameter with radiation treated as dark energy.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion, potentially excluding lookup table. Does not require
+            lookup tables.
+        z_values: Redshifts at which to calculate the hubble parameter.
+
+    Returns:
+        The hubble parameter values in units of H_0 (km/s/Mpc).
+
+    Notes:
+        Used for growth factor calculations.
+    """
+    t = cosmology_params['omega_m_zero'] * (1.0 + z_values)**3.0
+    t += cosmology_params['omega_de_zero']
+    t += cosmology_params['omega_rad_zero']
+    return jnp.sqrt(t)
+
+
+def _comoving_distance_numerical(
+    cosmology_params: Mapping[str, Union[float, int]],
+        z_min: float, z_max: float) -> float:
+    """Calculates the comoving distance at a given redshift numerically.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion. Does not require lookup tables.
+        z_min: Lower redshift of distance interval.
+        z_max: Upper redhsift of distance interval.
+
+    Returns:
+        Comoving distance in units of Mpc.
+    """
+    # We probably want a better / less expensive integration scheme than this.
+    z_samples = jnp.linspace(z_min, z_max, 1000)
+    e_z_samples = _e_z(cosmology_params, z_samples)
+
+    one_over_ez_int = jnp.trapz(1 / e_z_samples, z_samples)
+
+    # Factor in the speed of light and the Hubble constant. The factor of 1e-3
+    # comes from accounting for integral output being in Mpc * s / km and speed of
+    # light being in m/s as well as h being H0/100.
+    one_over_ez_int *= SPEED_OF_LIGHT * 1e-3
+    one_over_ez_int /= cosmology_params['hubble_constant']
+
+    return one_over_ez_int
+
+
+def _sigma_k_integrand(
+    cosmology_params: Mapping[str, Union[float, int]],
+    log_k: jnp.ndarray, radius: jnp.ndarray) -> jnp.ndarray:
+    """Return the integrand for the variance of the linear density field.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion. Does not require lookup tables.
+        log_k: Log of the wavenumber values at which to evaluate integrand in units
+            of h/Mpc.
+        radius: Scale at which to calculate the rms variance in units of Mpc/h.
+
+    Returns:
+        Integrand at provided wavenumbers.
+    """
+    k = jnp.exp(log_k)
+    # Fourier transform of the tophat filter function.
+    k_times_r = k * radius
+    # Use masking to default to 1 when below fidelity limit of float32.
+    is_below_cut = k_times_r < 3e-2
+    ft_filer = jnp.zeros_like(k)
+    ft_filer += is_below_cut
+    ft_filer += (
+            jnp.logical_not(is_below_cut) * 3.0 / k_times_r**3 *
+            (jnp.sin(k_times_r) - k_times_r * jnp.cos(k_times_r)))
+
+    ps = power_spectrum.matter_power_spectrum(
+            k, cosmology_params['omega_m_zero'],
+            cosmology_params['omega_b_zero'],
+            cosmology_params['temp_cmb_zero'],
+            cosmology_params['hubble_constant'], cosmology_params['n_s'])
+
+    return ps * ft_filer**2 * k**3
+
+
+def _sigma_numerical(cosmology_params: Mapping[str, Union[float, int]],
+    radius: float) -> float:
+    """Calculate the variance of the linear density field numercially.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion. Does not require lookup tables.
+        radius: Scale at which to calculate the rms variance in units of Mpc/h.
+
+    Returns:
+        RMS variance of the linear density field.
+    """
+    # Integartion bins are in log momentum space (k).
+    log_k_bins = jnp.linspace(
+            jnp.log10(1e-6 / radius), jnp.log10(1e6 / radius), 1000)
+    sigma_k_integrand = jnp.nan_to_num(
+            _sigma_k_integrand(cosmology_params, log_k_bins, radius))
+    sigma_squared = jnp.trapz(sigma_k_integrand, log_k_bins)
+
+    return jnp.sqrt(sigma_squared / (2.0 * jnp.pi**2))
+
+
+def _sigma_norm(cosmology_params: Mapping[str, Union[float, int]]) -> float:
+    """Calculate the normalization of the variance of the linear density field.
+
+    Args:
+        init_cosmology_params: Cosmological parameters that define the universe's
+            expansion. Does not require lookup tables.
+
+    Returns:
+        Normalized RMS variance of the linear density field.
+    """
+    # Normalize the raw output to sigma_eight.
+    sigma_eight_raw = _sigma_numerical(cosmology_params, 8.0)
+
+    return cosmology_params['sigma_eight'] / sigma_eight_raw
+
+
+def _growth_factor_exact_unormalized(
+    cosmology_params: Mapping[str, Union[float, int]], z: float) -> float:
+    """Calculate the linear growth factor, unormalized.
+
+    Args:
+        init_cosmology_params: Cosmological parameters that define the universe's
+            expansion, excluding the lookup table for the comoving_distance
+            calculation.
+        z: Redshift at which to calculate the growth factor.
+
+    Returns:
+        Growth factor at input redshift.
+
+    Notes:
+        Approximations only hold to redshift 10.
+    """
+    z_samples = jnp.logspace(
+            jnp.log10(jax.lax.max(z, 1e-6)), jnp.log10(GROWTH_Z_MAX), 10000)
+    integral = jnp.trapz(
+            (1.0 + z_samples) / _e_z_rad_to_dark(cosmology_params, z_samples)**3,
+            z_samples)
+
+    return (5.0 / 2.0 * cosmology_params['omega_m_zero'] *
+                    _e_z_rad_to_dark(cosmology_params, z) * integral)
 
 
 def add_lookup_tables_to_cosmology_params(
-    init_cosmology_params,
-    z_lookup_max,
-    dz,
-    r_min,
-    r_max,
-    n_r_bins = 1000,
-):
-  """Add lookup tables to cosmology params.
-
-  Args:
-    init_cosmology_params: Cosmological parameters that define the universe's
-      expansion, excluding the lookup tables.
-    z_lookup_max: Maximum z to calculate in the lookup tables.
-    dz: Step size in redshift space at which to calculate lookup tables.
-    r_min: Logarithm of minimum lagrangian radius to include in lookup tables.
-      In units of Mpc/h.
-    r_max: Logarithm of maximum lagrangian radius to include in lookup tables.
-      In units of Mpc/h.
-    n_r_bins: Number of radial bins to include in sigma lookup tables.
-
-  Returns:
-    Cosmological parameters with lookup table.
-  """
-  # Jitted function for faster computation.
-  comoving_distance_numerical = jax.jit(_comoving_distance_numerical)
-  sigma_exact = jax.jit(_sigma_numerical)
-  sigma_norm = jax.jit(_sigma_norm)
-  growth_factor = jax.jit(_growth_factor_exact_unormalized)
-
-  z_range = jnp.arange(0, z_lookup_max + dz, dz)
-  radius_range = jnp.linspace(r_min, r_max, n_r_bins)
-
-  # Normalize growth factor at redshift 0 to 1.
-  growth_norm = 1.0 / growth_factor(init_cosmology_params, 0.0)
-  sigma_norm = sigma_norm(init_cosmology_params)
-
-  comoving_lookup_table = jax.vmap(
-      jax.vmap(comoving_distance_numerical, in_axes=[None, None, 0]),
-      in_axes=[None, 0, None])(init_cosmology_params, z_range, z_range)
-  sigma_lookup_table = jax.vmap(
-      sigma_exact, in_axes=[None, 0])(init_cosmology_params, radius_range)
-  growth_lookup_table = (
-      jax.vmap(growth_factor, in_axes=[None, 0])(init_cosmology_params,
-                                                 z_range))
-
-  cosmology_params = {
-      'dz': dz,
-      'z_lookup_max': z_lookup_max,
-      'r_min': r_min,
-      'r_max': r_max,
-      'dr': (r_max - r_min) / (n_r_bins - 1),
-      'comoving_lookup_table': comoving_lookup_table,
-      'sigma_lookup_table': sigma_lookup_table * sigma_norm,
-      'growth_lookup_table': growth_lookup_table * growth_norm
-  }
-  cosmology_params.update(init_cosmology_params)
-
-  return cosmology_params
-
-
-def comoving_distance(cosmology_params,
-                      z_min, z_max):
-  """Calculates the comoving distance at a given redshift using lookup.
-
-  Args:
-    cosmology_params: Cosmological parameters that define the universe's
-      expansion.
-    z_min: Lower redshift of distance interval.
-    z_max: Upper redhsift of distance interval.
-
-  Returns:
-    Comoving distance in units of Mpc.
-
-  Notes:
-    If `z_min` or `z_max` are larger than `z_lookup_max` in `cosmology_params`
-    the return will be equivalent to replacing the offending parameter(s) with
-    `z_lookup_max`.
-  """
-  # Interpolate between the four nearest binds to the query.
-  unrounded_i = z_min / cosmology_params['dz']
-  unrounded_j = z_max / cosmology_params['dz']
-
-  lookup_i_upper = jax.lax.min(
-      jnp.ceil(unrounded_i).astype(int),
-      len(cosmology_params['comoving_lookup_table'] - 1))
-  lookup_i_lower = jax.lax.max(jnp.floor(unrounded_i).astype(int), 0)
-  lookup_j_upper = jax.lax.min(
-      jnp.ceil(unrounded_j).astype(int),
-      len(cosmology_params['comoving_lookup_table'] - 1))
-  lookup_j_lower = jax.lax.max(jnp.floor(unrounded_j).astype(int), 0)
-
-  # Conduct a bilinear interpolation.
-  frac_i = unrounded_i % 1
-  frac_j = unrounded_j % 1
-
-  # Replacing these lines with a matrix multiplication would be cleaner, but
-  # leads to slightly slower code due to fraction matrix initialization.
-  interpolated = (
-      (1 - frac_i) * (1 - frac_j) *
-      cosmology_params['comoving_lookup_table'][lookup_i_lower, lookup_j_lower])
-  interpolated += (
-      (frac_i) * (1 - frac_j) *
-      cosmology_params['comoving_lookup_table'][lookup_i_upper, lookup_j_lower])
-  interpolated += (
-      (1 - frac_i) * (frac_j) *
-      cosmology_params['comoving_lookup_table'][lookup_i_lower, lookup_j_upper])
-  interpolated += (
-      (frac_i) * (frac_j) *
-      cosmology_params['comoving_lookup_table'][lookup_i_upper, lookup_j_upper])
-
-  return interpolated
-
-
-def angular_diameter_distance(cosmology_params,
-                              z):
-  """Calculate the angular diameter distance at a given redshift.
-
-  Args:
-    cosmology_params: Cosmological parameters that define the universe's
-      expansion.
-    z: Redshift (i.e. distance) at which to conduct the calculation.
-
-  Returns:
-    Angular diameter distance in units of Mpc.
-  """
-  return comoving_distance(cosmology_params, 0.0, z) / (1.0 + z)
+        cosmology_params: Mapping[str, Union[float, int]],
+        z_lookup_max: float, dz: float, r_min: float, r_max: float,
+        n_r_bins: Optional[float] = 1000,
+) -> Mapping[str, Union[float, int, jnp.ndarray]]:
+    """Add lookup tables to cosmology params.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion. Does not include lookup tables.
+        z_lookup_max: Maximum z to calculate in the lookup tables.
+        dz: Step size in redshift space at which to calculate lookup tables.
+        r_min: Logarithm of minimum lagrangian radius to include in lookup tables.
+            In units of Mpc/h.
+        r_max: Logarithm of maximum lagrangian radius to include in lookup tables.
+            In units of Mpc/h.
+        n_r_bins: Number of radial bins to include in sigma lookup tables.
+
+    Returns:
+        Cosmological parameters with lookup table.
+    """
+    # Jitted function for faster computation.
+    comoving_distance_numerical = jax.jit(_comoving_distance_numerical)
+    sigma_exact = jax.jit(_sigma_numerical)
+    sigma_norm = jax.jit(_sigma_norm)
+    growth_factor = jax.jit(_growth_factor_exact_unormalized)
+
+    z_range = jnp.arange(0, z_lookup_max + dz, dz)
+    radius_range = jnp.linspace(r_min, r_max, n_r_bins)
+
+    # Normalize growth factor at redshift 0 to 1.
+    growth_norm = 1.0 / growth_factor(cosmology_params, 0.0)
+    sigma_norm = sigma_norm(cosmology_params)
+
+    comoving_lookup_table = jax.vmap(
+            jax.vmap(comoving_distance_numerical, in_axes=[None, None, 0]),
+            in_axes=[None, 0, None])(cosmology_params, z_range, z_range)
+    sigma_lookup_table = jax.vmap(
+            sigma_exact, in_axes=[None, 0])(cosmology_params, radius_range)
+    growth_lookup_table = (
+            jax.vmap(growth_factor, in_axes=[None, 0])(cosmology_params,
+                                                                                                 z_range))
+
+    cosmology_params_lookup = {
+            'dz': dz,
+            'z_lookup_max': z_lookup_max,
+            'r_min': r_min,
+            'r_max': r_max,
+            'dr': (r_max - r_min) / (n_r_bins - 1),
+            'comoving_lookup_table': comoving_lookup_table,
+            'sigma_lookup_table': sigma_lookup_table * sigma_norm,
+            'growth_lookup_table': growth_lookup_table * growth_norm
+    }
+    cosmology_params_lookup.update(cosmology_params)
+
+    return cosmology_params_lookup
+
+
+def comoving_distance(
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z_min: float, z_max: float) -> float:
+    """Calculates the comoving distance at a given redshift using lookup.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z_min: Lower redshift of distance interval.
+        z_max: Upper redhsift of distance interval.
+
+    Returns:
+        Comoving distance in units of Mpc.
+
+    Notes:
+        If `z_min` or `z_max` are larger than `z_lookup_max` in `cosmology_params`
+        the return will be equivalent to replacing the offending parameter(s) with
+        `z_lookup_max`.
+    """
+    # Interpolate between the four nearest binds to the query.
+    unrounded_i = z_min / cosmology_params['dz']
+    unrounded_j = z_max / cosmology_params['dz']
+
+    lookup_i_upper = jax.lax.min(
+            jnp.ceil(unrounded_i).astype(int),
+            len(cosmology_params['comoving_lookup_table'] - 1))
+    lookup_i_lower = jax.lax.max(jnp.floor(unrounded_i).astype(int), 0)
+    lookup_j_upper = jax.lax.min(
+            jnp.ceil(unrounded_j).astype(int),
+            len(cosmology_params['comoving_lookup_table'] - 1))
+    lookup_j_lower = jax.lax.max(jnp.floor(unrounded_j).astype(int), 0)
+
+    # Conduct a bilinear interpolation.
+    frac_i = unrounded_i % 1
+    frac_j = unrounded_j % 1
+
+    # Replacing these lines with a matrix multiplication would be cleaner, but
+    # leads to slightly slower code due to fraction matrix initialization.
+    interpolated = (
+            (1 - frac_i) * (1 - frac_j) *
+            cosmology_params['comoving_lookup_table'][lookup_i_lower, lookup_j_lower])
+    interpolated += (
+            (frac_i) * (1 - frac_j) *
+            cosmology_params['comoving_lookup_table'][lookup_i_upper, lookup_j_lower])
+    interpolated += (
+            (1 - frac_i) * (frac_j) *
+            cosmology_params['comoving_lookup_table'][lookup_i_lower, lookup_j_upper])
+    interpolated += (
+            (frac_i) * (frac_j) *
+            cosmology_params['comoving_lookup_table'][lookup_i_upper, lookup_j_upper])
+
+    return interpolated
+
+
+def angular_diameter_distance(
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z: float) -> float:
+    """Calculate the angular diameter distance at a given redshift.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z: Redshift (i.e. distance) at which to conduct the calculation.
+
+    Returns:
+        Angular diameter distance in units of Mpc.
+    """
+    return comoving_distance(cosmology_params, 0.0, z) / (1.0 + z)
 
 
 def angular_diameter_distance_between(
-    cosmology_params,
-    z_min,
-    z_max,
-):
-  """Calculate the angular diameter distance between two redshifts.
-
-  Args:
-    cosmology_params: Cosmological parameters that define the universe's
-      expansion.
-    z_min: Redshift (i.e. distance) at which to start the calculation.
-    z_max: Redhisft (i.e. distance) at which to stop the calculation.
-
-  Returns:
-    Angular diameter distance in units of Mpc.
-  """
-  return comoving_distance(cosmology_params, z_min, z_max) / (1.0 + z_max)
-
-
-def kpc_per_arcsecond(cosmology_params,
-                      z):
-  """Calculate the physical kpc per arcsecond at a given redshift.
-
-  Args:
-    cosmology_params: Cosmological parameters that define the universe's
-      expansion.
-    z: Redshift (i.e. distance) at which to conduct the calculation.
-
-  Returns:
-    Kpc per arcseconds.
-  """
-  radians_per_arcsecond = jnp.pi / 180 / 3600
-  kpc_per_megapc = 1e3
-  return (angular_diameter_distance(cosmology_params, z) *
-          radians_per_arcsecond * kpc_per_megapc)
-
-
-def comoving_to_angle(comv_x, comv_y,
-                      cosmology_params,
-                      z_lens):
-  """Convert from comoving coordinates to angular units.
-
-  Args:
-    comv_x: Comoving x-coordinate.
-    comv_y: Comoving y-coordinate.
-    cosmology_params: Cosmological parameters that define the universe's
-      expansion.
-    z_lens: Redshift of the lens (i.e. current redshift).
-
-  Returns:
-    X- and y-coordinate in angular units (radians).
-  """
-  arcsecond_per_megaparsec = 1 / comoving_distance(cosmology_params, 0, z_lens)
-  return comv_x * arcsecond_per_megaparsec, comv_y * arcsecond_per_megaparsec
+        cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+        z_min: float, z_max: float) -> float:
+    """Calculate the angular diameter distance between two redshifts.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z_min: Redshift (i.e. distance) at which to start the calculation.
+        z_max: Redhisft (i.e. distance) at which to stop the calculation.
+
+    Returns:
+        Angular diameter distance in units of Mpc.
+    """
+    return comoving_distance(cosmology_params, z_min, z_max) / (1.0 + z_max)
+
+
+def kpc_per_arcsecond(
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z: float) -> float:
+    """Calculate the physical kpc per arcsecond at a given redshift.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z: Redshift (i.e. distance) at which to conduct the calculation.
+
+    Returns:
+        Kpc per arcseconds.
+    """
+    radians_per_arcsecond = jnp.pi / 180 / 3600
+    kpc_per_megapc = 1e3
+    return (angular_diameter_distance(cosmology_params, z) *
+                    radians_per_arcsecond * kpc_per_megapc)
+
+
+def comoving_to_angle(comv_x: Union[float, jnp.ndarray],
+    comv_y: Union[float, jnp.ndarray],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z_lens: float) -> Union[float, jnp.ndarray]:
+    """Convert from comoving coordinates to angular units.
+
+    Args:
+        comv_x: Comoving x-coordinate.
+        comv_y: Comoving y-coordinate.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z_lens: Redshift of the lens (i.e. current redshift).
+
+    Returns:
+        X- and y-coordinate in angular units (radians).
+    """
+    arcsecond_per_megaparsec = 1 / comoving_distance(cosmology_params, 0, z_lens)
+    return comv_x * arcsecond_per_megaparsec, comv_y * arcsecond_per_megaparsec
 
 
 def reduced_to_physical(
-    reduced,
-    cosmology_params,
-    z_lens,
-    z_source,
-):
-  """Transform from reduced deflection angle to physical coordinates.
-
-  Args:
-    reduced: Reduced deflection angle.
-    cosmology_params: Cosmological parameters that define the universe's
-      expansion.
-    z_lens: Redshift of the lens (i.e. current redshift).
-    z_source: Redshift of the source.
-
-  Returns:
-    Physical coordinates.
-  """
-  return reduced * (
-      angular_diameter_distance_between(cosmology_params, 0.0, z_source) /
-      angular_diameter_distance_between(cosmology_params, z_lens, z_source))
-
-
-def rho_matter(cosmology_params,
-               z):
-  """Return the matter density at redshift z.
-
-  Args:
-    cosmology_params: Cosmological parameters that define the universe's
-      expansion.
-    z: Redshift at which to calculate the matter density.
-
-  Returns:
-    Matter density at redshift z in units M_sun * h ** 2 / Mpc ** 3.
-  """
-  return RHO_CRIT_0 * cosmology_params['omega_m_zero'] * (1.0 + z)**3
+        reduced: Union[float, jnp.ndarray],
+        cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+        z_lens: float, z_source: float) -> Union[float, jnp.ndarray]:
+    """Transform from reduced deflection angle to physical coordinates.
+
+    Args:
+        reduced: Reduced deflection angle.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z_lens: Redshift of the lens (i.e. current redshift).
+        z_source: Redshift of the source.
+
+    Returns:
+        Physical coordinates.
+    """
+    return reduced * (
+            angular_diameter_distance_between(cosmology_params, 0.0, z_source) /
+            angular_diameter_distance_between(cosmology_params, z_lens, z_source))
+
+
+def rho_matter(cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z: float) -> float:
+    """Return the matter density at redshift z.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z: Redshift at which to calculate the matter density.
+
+    Returns:
+        Matter density at redshift z in units M_sun * h ** 2 / Mpc ** 3.
+    """
+    return RHO_CRIT_0 * cosmology_params['omega_m_zero'] * (1.0 + z)**3
 
 
 def collapse_overdensity(
-    cosmology_params,
-    z,
-):
-  """Return the collapse overdensity including first-order corrections.
-
-  Args:
-    cosmology_params: Cosmological parameters that define the universe's
-      expansion.
-    z: Redshift at which to calculate the collapse overdensity.
-
-  Returns:
-    Collapse overdensity.
-  """
-  return (DELTA_COLLAPSE_UNCORRECTED *
-          (cosmology_params['omega_m_zero'] * (1.0 + z)**3 /
-           (_e_z(cosmology_params, z))**2)**0.0055)
-
-
-def lagrangian_radius(cosmology_params,
-                      mass):
-  """Return the lagrangian radius of halo of mass M.
-
-  Args:
-    cosmology_params: Cosmological parameters that define the universe's
-      expansion.
-    mass: Masses at which to calculate the lagrangian radius in units M_sun / h.
-
-  Returns:
-    Lagrangian radius in units of comoving Mpc / h.
-  """
-  # rho_matter returns in h units, so convert here for radius to be returned in
-  # Mpc.
-  return ((3.0 * mass / 4.0 / jnp.pi /
-           rho_matter(cosmology_params, 0.0))**(1.0 / 3.0))
-
-
-def sigma_tophat(cosmology_params,
-                 lagrangian_r, z):
-  """The RMS variance of the linear density field, including the growth factor.
-
-  Args:
-    cosmology_params: Cosmological parameters that define the universe's
-      expansion.
-    lagrangian_r: Lagrangian radius at which to evaluate the linear density
-      field in units of comoving Mpc / h.
-    z: Redshift at which to evaluate the linear density field.
-
-  Returns:
-    RMS variance of the linear density field.
-  """
-  # For both calculations, use linear interpolation between bins.
-  lookup_z_unrounded = z / cosmology_params['dz']
-  frac_z = lookup_z_unrounded % 1
-
-  lookup_z_upper = jax.lax.min(
-      jnp.ceil(lookup_z_unrounded).astype(int),
-      len(cosmology_params['growth_lookup_table']) - 1)
-  lookup_z_lower = jax.lax.max(jnp.floor(lookup_z_unrounded).astype(int), 0)
-
-  lookup_sigma_unrounded = ((lagrangian_r - cosmology_params['r_min']) /
-                            cosmology_params['dr'])
-  frac_sigma = lookup_sigma_unrounded % 1
-
-  lookup_sigma_upper = jax.lax.min(
-      jnp.ceil(lookup_sigma_unrounded).astype(int),
-      len(cosmology_params['sigma_lookup_table']) - 1)
-  lookup_sigma_lower = jax.lax.max(
-      jnp.floor(lookup_sigma_unrounded).astype(int), 0)
-
-  sigma_no_growth = (
-      frac_sigma * cosmology_params['sigma_lookup_table'][lookup_sigma_upper] +
-      (1 - frac_sigma) *
-      cosmology_params['sigma_lookup_table'][lookup_sigma_lower])
-  growth = (
-      frac_z * cosmology_params['growth_lookup_table'][lookup_z_upper] +
-      (1 - frac_z) * cosmology_params['growth_lookup_table'][lookup_z_lower])
-
-  return sigma_no_growth * growth
+        cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+        z: float) -> float:
+    """Return the collapse overdensity including first-order corrections.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z: Redshift at which to calculate the collapse overdensity.
+
+    Returns:
+        Collapse overdensity.
+    """
+    return (DELTA_COLLAPSE_UNCORRECTED *
+                    (cosmology_params['omega_m_zero'] * (1.0 + z)**3 /
+                     (_e_z(cosmology_params, z))**2)**0.0055)
+
+
+def lagrangian_radius(
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    mass: Union[float, jnp.ndarray]) -> Union[float, jnp.ndarray]:
+    """Return the lagrangian radius of halo of mass M.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        mass: Mass at which to calculate the lagrangian radius in units M_sun / h.
+
+    Returns:
+        Lagrangian radius in units of comoving Mpc / h.
+    """
+    # rho_matter returns in h units, so convert here for radius to be returned in
+    # Mpc.
+    return ((3.0 * mass / 4.0 / jnp.pi /
+                     rho_matter(cosmology_params, 0.0))**(1.0 / 3.0))
+
+
+def sigma_tophat(cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    lagrangian_r: float, z: float) -> float:
+    """The RMS variance of the linear density field, including the growth factor.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        lagrangian_r: Lagrangian radius at which to evaluate the linear density
+            field in units of comoving Mpc / h.
+        z: Redshift at which to evaluate the linear density field.
+
+    Returns:
+        RMS variance of the linear density field.
+    """
+    # For both calculations, use linear interpolation between bins.
+    lookup_z_unrounded = z / cosmology_params['dz']
+    frac_z = lookup_z_unrounded % 1
+
+    lookup_z_upper = jax.lax.min(
+            jnp.ceil(lookup_z_unrounded).astype(int),
+            len(cosmology_params['growth_lookup_table']) - 1)
+    lookup_z_lower = jax.lax.max(jnp.floor(lookup_z_unrounded).astype(int), 0)
+
+    lookup_sigma_unrounded = ((lagrangian_r - cosmology_params['r_min']) /
+                                                        cosmology_params['dr'])
+    frac_sigma = lookup_sigma_unrounded % 1
+
+    lookup_sigma_upper = jax.lax.min(
+            jnp.ceil(lookup_sigma_unrounded).astype(int),
+            len(cosmology_params['sigma_lookup_table']) - 1)
+    lookup_sigma_lower = jax.lax.max(
+            jnp.floor(lookup_sigma_unrounded).astype(int), 0)
+
+    sigma_no_growth = (
+            frac_sigma * cosmology_params['sigma_lookup_table'][lookup_sigma_upper] +
+            (1 - frac_sigma) *
+            cosmology_params['sigma_lookup_table'][lookup_sigma_lower])
+    growth = (
+            frac_z * cosmology_params['growth_lookup_table'][lookup_z_upper] +
+            (1 - frac_z) * cosmology_params['growth_lookup_table'][lookup_z_lower])
+
+    return sigma_no_growth * growth
 
 
 def derivative_log_sigma_log_r(
-    cosmology_params,
-    lagrangian_r,
-    z,
-):
-  """Derivative of the log RMS variance with respect to log radius.
-
-  Args:
-    cosmology_params: Cosmological parameters that define the universe's
-      expansion.
-    lagrangian_r: Lagrangian radius at which to evaluate the derivative in units
-      of comoving Mpc / h.
-    z: Redshift at which to evaluate the derivative.
-
-  Returns:
-    Derivative of the log RMS variance.
-  """
-  # Calculate the derivative of sigma with respect to radius, then use the chain
-  # rule to get the log of sigma with respect to the log of the radius.
-  dr = cosmology_params['dr']
-
-  # Calculate the derivative using the five-point stencil.
-  derivative = -sigma_tophat(cosmology_params, lagrangian_r + 2 * dr, z)
-  derivative += 8 * sigma_tophat(cosmology_params, lagrangian_r + dr, z)
-  derivative += -8 * sigma_tophat(cosmology_params, lagrangian_r - dr, z)
-  derivative += sigma_tophat(cosmology_params, lagrangian_r - 2 * dr, z)
-  derivative /= 12 * dr
-
-  sigma = sigma_tophat(cosmology_params, lagrangian_r, z)
-  return derivative * lagrangian_r / sigma
+        cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+        lagrangian_r: float, z: float) -> float:
+    """Derivative of the log RMS variance with respect to log radius.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        lagrangian_r: Lagrangian radius at which to evaluate the derivative in units
+            of comoving Mpc / h.
+        z: Redshift at which to evaluate the derivative.
+
+    Returns:
+        Derivative of the log RMS variance.
+    """
+    # Calculate the derivative of sigma with respect to radius, then use the chain
+    # rule to get the log of sigma with respect to the log of the radius.
+    dr = cosmology_params['dr']
+
+    # Calculate the derivative using the five-point stencil.
+    derivative = -sigma_tophat(cosmology_params, lagrangian_r + 2 * dr, z)
+    derivative += 8 * sigma_tophat(cosmology_params, lagrangian_r + dr, z)
+    derivative += -8 * sigma_tophat(cosmology_params, lagrangian_r - dr, z)
+    derivative += sigma_tophat(cosmology_params, lagrangian_r - 2 * dr, z)
+    derivative /= 12 * dr
+
+    sigma = sigma_tophat(cosmology_params, lagrangian_r, z)
+    return derivative * lagrangian_r / sigma
+
+
+def peak_height(cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    mass: float, z: float) -> float:
+    """Peak height for a specific mass and redshift
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        mass: Mass at which to calculate the lagrangian radius in units M_sun / h.
+        z: Redshift of the halo.
+
+    Returns:
+        Peak height.
+    """
+    lagrangian_r = lagrangian_radius(cosmology_params, mass)
+    sigma = sigma_tophat(cosmology_params, lagrangian_r, z)
+    return collapse_overdensity(cosmology_params, z) / sigma
+
+
+def rho_crit(cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z: float) -> float:
+    """Critical density of the universe at redshift z.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z: Redshift at which to calculate critical density.
+
+    Returns:
+        Critical density of the universe at redshift z in units M_sun * h ** 2 /
+            kpc ** 3.
+    """
+    kpc_to_mpc = 1e-3
+    return RHO_CRIT_0 * _e_z(cosmology_params, z) ** 2 * kpc_to_mpc ** 3
+
+
+def calculate_sigma_crit(
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z: float, z_source: float) -> float:
+    """Calculate the critical surface density for the lensing configuration.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z: Redshift at which to calculate critical surface density.
+        z_source: Redshift of the source.
+
+    Returns:
+        The critical surface density in units of M_sun/kpc^2
+    """
+    # Normalization factor for sigma_crit
+    norm = (SPEED_OF_LIGHT ** 2 / (4 * jnp.pi * GRAVITATIONAL_CONSTANT) *
+        MEGAPARSEC_METERS / MASS_SUN)
+    # Get our three angular diameter distances
+    mpc_to_kpc = 1e3
+    dd = angular_diameter_distance(cosmology_params, z)
+    ds = angular_diameter_distance(cosmology_params, z_source)
+    dds = angular_diameter_distance_between(cosmology_params, z, z_source)
+    return ds / (dd * dds) * norm / mpc_to_kpc ** 2
+
+
+def correlation_function(cosmology_params, comoving_r, z_lens):
+    return 0.0
+
+
+def halo_bias(cosmology_params, mass, z_lens):
+    return 0.0
diff --git a/jaxstronomy/image_simulation.py b/jaxstronomy/image_simulation.py
index 723ac6e..cbe2dd5 100644
--- a/jaxstronomy/image_simulation.py
+++ b/jaxstronomy/image_simulation.py
@@ -1,19 +1,4 @@
 # coding=utf-8
-# Copyright 2022 The Google Research Authors.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Copyright 2022 Google LLC
 
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -34,526 +19,690 @@
 """
 
 import functools
-import itertools
-from typing import Mapping, Sequence, Tuple, Union
+from typing import Any, Mapping, Sequence, Tuple, Union
 
 import jax
 import jax.numpy as jnp
 
 from jaxstronomy import cosmology_utils
-from jaxstronomy import lens_models
-from jaxstronomy import psf_models
-from jaxstronomy import source_models
 from jaxstronomy import utils
 
-LENS_MODELS = tuple(
-    [lens_models.__getattribute__(model) for model in lens_models.__all__])
-LENS_MODEL_PARAMETERS = tuple([model.parameters for model in LENS_MODELS])
-ALL_LENS_MODEL_PARAMETERS = tuple(
-    sorted(set(itertools.chain(*LENS_MODEL_PARAMETERS))))
-SOURCE_MODELS = tuple(
-    [source_models.__getattribute__(model) for model in source_models.__all__])
-SOURCE_MODEL_PARAMETERS = tuple([model.parameters for model in SOURCE_MODELS])
-ALL_SOURCE_MODEL_PARAMETERS = tuple(
-    sorted(set(itertools.chain(*SOURCE_MODEL_PARAMETERS))))
-PSF_MODELS = tuple(
-    [psf_models.__getattribute__(model) for model in psf_models.__all__])
-PSF_MODEL_PARAMETERS = tuple([model.parameters for model in PSF_MODELS])
-ALL_PSF_MODEL_PARAMETERS = tuple(
-    sorted(set(itertools.chain(*PSF_MODEL_PARAMETERS))))
-
 
 def generate_image(
-    grid_x,
-    grid_y,
-    kwargs_lens_all,
-    kwargs_source_slice,
-    kwargs_lens_light_slice,
-    kwargs_psf,
-    cosmology_params,
-    z_lens_array,
-    z_source,
-    kwargs_detector,
-):
-  """Generate an image given the source, lens light, and mass profiles.
-
-  Args:
-    grid_x: X-coordinates of grid on which image will be generated.
-    grid_y: Y-coordinates of grid on which image will be generated.
-    kwargs_lens_all: Keyword arguments for each lens model. The kwargs for
-      should contain keyword arguments to pass to derivative functions for each
-      lens model. Keys should include parameters for all lens models (due to use
-      of `jax.lax.switch`) and `model_index` which defines the model to pass the
-      parameters to. Parameter values not relevant for the specified model are
-      discarded. Values should be of type jnp.ndarray and have length equal to
-      the total number of lens models. For more detailed discussion see
-      documentation of `jax.lax.scan`, which is used to iterate over the models.
-    kwargs_source_slice: Keyword arguments to pass to brightness functions for
-      each light model. Keys should include parameters for all source models
-      (due to use of `jax.lax.switch`) and `model_index` which defines the model
-      to pass the parameters to.
-    kwargs_lens_light_slice: Keyword arguments to pass to brightness functions
-      for each light model. Keys should include parameters for all source models
-      (due to use of `jax.lax.switch`) and `model_index` which defines the model
-      to pass the parameters to.
-    kwargs_psf: Keyword arguments defining the point spread function. The psf is
-      applied in the supersampled space, so the size of pixels should be defined
-      with respect to the supersampled space.
-    cosmology_params: Cosmological parameters that define the universe's
-      expansion.
-    z_lens_array: Redshifts for each lens model.
-    z_source: Redshift of the source.
-    kwargs_detector: Keyword arguments defining the detector configuration.
-
-  Returns:
-    Image after gravitational lensing at supersampling resolution. For
-    consistency with lenstronomy, units are per pixel area of the detector, not
-    the supersampling resolution.
-
-  Notes:
-    kwargs_detector must be passed in through functools.partial or equivalent if
-    you want to jit compile this function.
-  """
-  image_array = source_surface_brightness(grid_x, grid_y, kwargs_lens_all,
-                                          kwargs_source_slice, kwargs_detector,
-                                          cosmology_params, z_lens_array,
-                                          z_source)
-  image_array += lens_light_surface_brightness(grid_x, grid_y,
-                                               kwargs_lens_light_slice,
-                                               kwargs_detector)
-  image = jnp.reshape(
-      image_array,
-      (kwargs_detector['n_x'] * kwargs_detector['supersampling_factor'],
-       kwargs_detector['n_y'] * kwargs_detector['supersampling_factor']))
-
-  return psf_convolution(image, kwargs_psf)
+    grid_x: jnp.ndarray,
+    grid_y: jnp.ndarray,
+    kwargs_lens_all: Mapping[str, Union[jnp.ndarray, Mapping[str, jnp.ndarray]]],
+    kwargs_source_slice: Mapping[str, jnp.ndarray],
+    kwargs_lens_light_slice: Mapping[str, jnp.ndarray],
+    kwargs_psf: Mapping[str, Union[float, int, jnp.ndarray]],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z_source: float,
+    kwargs_detector: Mapping[str, Union[float, int]],
+    all_models: Mapping[str, Sequence[Any]]
+) -> jnp.ndarray:
+    """Generate an image given the source, lens light, and mass profiles.
+
+    Args:
+        grid_x: X-coordinates of grid on which image will be generated.
+        grid_y: Y-coordinates of grid on which image will be generated.
+        kwargs_lens_all: Keyword arguments and redshifts for each of the
+            lensing components. This should include los_before, los_after,
+            subhalos, and main_deflector.
+        kwargs_source_slice: Keyword arguments to pass to brightness functions
+            for each light model. Keys should include parameters for all source
+            models (due to use of `jax.lax.switch`) and `model_index` which
+            defines the model to pass the parameters to.
+        kwargs_lens_light_slice: Keyword arguments to pass to brightness
+            functions for each light model. Keys should include parameters for
+            all source models (due to use of `jax.lax.switch`) and `model_index`
+            which defines the model to pass the parameters to.
+        kwargs_psf: Keyword arguments defining the point spread function. The
+            psf is applied in the supersampled space, so the size of pixels
+            should be defined with respect to the supersampled space.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z_source: Redshift of the source.
+        kwargs_detector: Keyword arguments defining the detector configuration.
+        all_models: Tuple of model classes to consider for each component.
+
+    Returns:
+        Image after gravitational lensing at supersampling resolution. For
+        consistency with lenstronomy, units are per pixel area of the detector,
+        not the supersampling resolution.
+
+    Notes:
+        The parameters kwargs_detector and all_models must be made static if to
+        jit compile this function.
+    """
+    image_array = source_surface_brightness(
+        grid_x, grid_y, kwargs_lens_all, kwargs_source_slice, kwargs_detector,
+        cosmology_params, z_source, all_models
+    )
+    image_array += lens_light_surface_brightness(
+        grid_x, grid_y, kwargs_lens_light_slice, kwargs_detector,
+        all_models["all_source_models"]
+    )
+    image = jnp.reshape(
+        image_array,
+        (
+            kwargs_detector["n_x"] * kwargs_detector["supersampling_factor"],
+            kwargs_detector["n_y"] * kwargs_detector["supersampling_factor"],
+        ),
+    )
+
+    return psf_convolution(image, kwargs_psf, all_models["all_psf_models"])
 
 
 def psf_convolution(
-    image,
-    kwargs_psf,
-):
-  """Convolve image with the point spread function.
-
-  Args:
-    image: Image to convolve
-    kwargs_psf: Keyword arguments defining the point spread function.
-
-  Returns:
-    Convolved image.
-  """
-  # Psf not accounting for supersampling by default is prone to user error.
-  # Consider changing.
-  psf_functions = [model.convolve for model in PSF_MODELS]
-  return jax.lax.switch(kwargs_psf['model_index'], psf_functions, image,
-                        kwargs_psf)
+    image: jnp.ndarray,
+    kwargs_psf: Mapping[str, Union[float, int, jnp.ndarray]],
+    all_psf_models: Sequence[Any]
+) -> jnp.ndarray:
+    """Convolve image with the point spread function.
+
+    Args:
+        image: Image to convolve
+        kwargs_psf: Keyword arguments defining the point spread function.
+        all_psf_models: PSF models to use for model_index lookup.
+
+    Returns:
+        Convolved image.
+    """
+    # Psf not accounting for supersampling by default is prone to user error.
+    # Consider changing.
+    psf_functions = [model.convolve for model in all_psf_models]
+    return jax.lax.switch(kwargs_psf["model_index"], psf_functions, image, kwargs_psf)
 
 
 def noise_realization(
-    image, rng,
-    kwargs_detector):
-  """Return noise realization for input image.
-
-  Args:
-    image: Image from which to draw poisson noise.
-    rng: jax PRNG key used for noise realization.
-    kwargs_detector: Keyword arguments defining the detector configuration.
-
-  Returns:
-    Noise realization for image.
-  """
-  rng_normal, rng_poisson = jax.random.split(rng)
-
-  # Calculate expected background noise from detector kwargs.
-  exposure_time_total = (
-      kwargs_detector['exposure_time'] * kwargs_detector['num_exposures'])
-  read_noise = (
-      kwargs_detector['read_noise']**2 * kwargs_detector['num_exposures'])
-  sky_brightness_cps = utils.magnitude_to_cps(
-      kwargs_detector['sky_brightness'],
-      kwargs_detector['magnitude_zero_point'])
-  sky_brightness_tot = (
-      exposure_time_total * sky_brightness_cps *
-      kwargs_detector['pixel_width']**2)
-  background_noise = (
-      jnp.sqrt(read_noise + sky_brightness_tot) / exposure_time_total)
-
-  # By default all simulations are done in units of counts per second, but you
-  # want to calculate poisson statistics in units of counts.
-  flux_noise = jnp.sqrt(jax.nn.relu(image) / exposure_time_total)
-
-  noise = jax.random.normal(rng_normal, image.shape) * background_noise
-  noise += jax.random.normal(rng_poisson, image.shape) * flux_noise
-  return noise
+    image: jnp.ndarray,
+    rng: Sequence[int],
+    kwargs_detector: Mapping[str, Union[float, int]]
+) -> jnp.ndarray:
+    """Return noise realization for input image.
+
+    Args:
+        image: Image from which to draw poisson noise.
+        rng: jax PRNG key used for noise realization.
+        kwargs_detector: Keyword arguments defining the detector configuration.
+
+    Returns:
+        Noise realization for image.
+    """
+    rng_normal, rng_poisson = jax.random.split(rng)
+
+    # Calculate expected background noise from detector kwargs.
+    exposure_time_total = (
+        kwargs_detector["exposure_time"] * kwargs_detector["num_exposures"]
+    )
+    read_noise = kwargs_detector["read_noise"] ** 2 * kwargs_detector["num_exposures"]
+    sky_brightness_cps = utils.magnitude_to_cps(
+        kwargs_detector["sky_brightness"], kwargs_detector["magnitude_zero_point"]
+    )
+    sky_brightness_tot = (
+        exposure_time_total * sky_brightness_cps * kwargs_detector["pixel_width"] ** 2
+    )
+    background_noise = jnp.sqrt(read_noise + sky_brightness_tot) / exposure_time_total
+
+    # By default all simulations are done in units of counts per second, but you
+    # want to calculate poisson statistics in units of counts.
+    flux_noise = jnp.sqrt(jax.nn.relu(image) / exposure_time_total)
+
+    noise = jax.random.normal(rng_normal, image.shape) * background_noise
+    noise += jax.random.normal(rng_poisson, image.shape) * flux_noise
+    return noise
 
 
 def lens_light_surface_brightness(
-    theta_x, theta_y,
-    kwargs_lens_light_slice,
-    kwargs_detector):
-  """Return the lens light surface brightness.
-
-  Args:
-    theta_x: X-coordinates in angular units.
-    theta_y: Y-coordinates in angular units.
-    kwargs_lens_light_slice: Keyword arguments to pass to brightness functions
-      for each light model. Keys should include parameters for all source models
-      (due to use of `jax.lax.switch`) and `model_index` which defines the model
-      to pass the parameters to.
-    kwargs_detector: Keyword arguments defining the detector configuration.
-
-  Returns:
-    Surface brightness of lens light as 1D array.
-  """
-  lens_light_flux = _surface_brightness(theta_x, theta_y,
-                                        kwargs_lens_light_slice)
-  return lens_light_flux * kwargs_detector['pixel_width']**2
+    theta_x: jnp.ndarray,
+    theta_y: jnp.ndarray,
+    kwargs_lens_light_slice: Mapping[str, jnp.ndarray],
+    kwargs_detector: Mapping[str, Union[float, int]],
+    all_source_models: Sequence[Any]
+) -> jnp.ndarray:
+    """Return the lens light surface brightness.
+
+    Args:
+        theta_x: X-coordinates in angular units.
+        theta_y: Y-coordinates in angular units.
+        kwargs_lens_light_slice: Keyword arguments to pass to brightness
+            functions for each light model. Keys should include parameters for
+            all source models (due to use of `jax.lax.switch`) and `model_index`
+            which defines the model to pass the parameters to.
+        kwargs_detector: Keyword arguments defining the detector configuration.
+        all_source_models: Source models to use for model_index lookup.
+
+    Returns:
+        Surface brightness of lens light as 1D array.
+    """
+    lens_light_flux = _surface_brightness(
+        theta_x, theta_y, kwargs_lens_light_slice, all_source_models
+    )
+    return lens_light_flux * kwargs_detector["pixel_width"] ** 2
 
 
 def source_surface_brightness(
-    alpha_x,
-    alpha_y,
-    kwargs_lens_all,
-    kwargs_source_slice,
-    kwargs_detector,
-    cosmology_params,
-    z_lens_array,
-    z_source,
-):
-  """Return the lensed source surface brightness.
-
-  Args:
-    alpha_x: Initial x-component of deflection at each position.
-    alpha_y: Initial y-component of deflection at each position.
-    kwargs_lens_all: Keyword arguments for each lens model. The kwargs for
-      should contain keyword arguments to pass to derivative functions for each
-      lens model. Keys should include parameters for all lens models (due to use
-      of `jax.lax.switch`) and `model_index` which defines the model to pass the
-      parameters to. Parameter values not relevant for the specified model are
-      discarded. Values should be of type jnp.ndarray and have length equal to
-      the total number of lens models. For more detailed discussion see
-      documentation of `jax.lax.scan`, which is used to iterate over the models.
-    kwargs_source_slice: Keyword arguments to pass to brightness functions for
-      each light model. Keys should include parameters for all source models
-      (due to use of `jax.lax.switch`) and `model_index` which defines the model
-      to pass the parameters to.
-    kwargs_detector: Keyword arguments defining the detector configuration. This
-      includes potential supersampling in the lensing calculation.
-    cosmology_params: Cosmological parameters that define the universe's
-      expansion.
-    z_lens_array: Redshifts for each lens model.
-    z_source: Redshift of the source.
-
-  Returns:
-    Lensed source surface brightness as 1D array.
-  """
-  image_flux_array = _image_flux(alpha_x, alpha_y, kwargs_lens_all,
-                                 kwargs_source_slice, cosmology_params,
-                                 z_lens_array, z_source)
-  # Scale by pixel area to go from flux to surface brightness.
-  return image_flux_array * kwargs_detector['pixel_width']**2
-
-
-def _image_flux(alpha_x, alpha_y,
-                kwargs_lens_all,
-                kwargs_source_slice,
-                cosmology_params,
-                z_lens_array, z_source):
-  """Calculate image flux after ray tracing onto the source.
-
-  Args:
-    alpha_x: Initial x-component of deflection at each position.
-    alpha_y: Initial y-component of deflectoin at each position.
-    kwargs_lens_all: Keyword arguments for each lens model. The kwargs for
-      should contain keyword arguments to pass to derivative functions for each
-      lens model. Keys should include parameters for all lens models (due to use
-      of `jax.lax.switch`) and `model_index` which defines the model to pass the
-      parameters to. Parameter values not relevant for the specified model are
-      discarded. Values should be of type jnp.ndarray and have length equal to
-      the total number of lens models. For more detailed discussion see
-      documentation of `jax.lax.scan`, which is used to iterate over the models.
-    kwargs_source_slice: Keyword arguments to pass to brightness functions for
-      each light model. Keys should include parameters for all source models
-      (due to use of `jax.lax.switch`) and `model_index` which defines the model
-      to pass the parameters to. Parameter values not relevant for the specified
-      model are discarded. Values should be of type jnp.ndarray and have length
-      equal to the total number of light models in the slice. For more detailed
-      discussion see documentation of `jax.lax.scan`, which is used to iterate
-      over the models.
-    cosmology_params: Cosmological parameters that define the universe's
-      expansion.
-    z_lens_array: Redshifts for each lens model.
-    z_source: Redshift of the source.
-
-  Returns:
-    Image flux.
-  """
-  x_source_comv, y_source_comv = _ray_shooting(alpha_x, alpha_y,
-                                               kwargs_lens_all,
-                                               cosmology_params, z_lens_array,
-                                               z_source)
-  x_source, y_source = cosmology_utils.comoving_to_angle(
-      x_source_comv, y_source_comv, cosmology_params, z_source)
-  return _surface_brightness(x_source, y_source, kwargs_source_slice)
+    alpha_x: jnp.ndarray,
+    alpha_y: jnp.ndarray,
+    kwargs_lens_all: Mapping[str, Union[jnp.ndarray, Mapping[str, jnp.ndarray]]],
+    kwargs_source_slice: Mapping[str, jnp.ndarray],
+    kwargs_detector: Mapping[str, jnp.ndarray],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z_source: float,
+    all_models: Mapping[str, Sequence[Any]],
+) -> jnp.ndarray:
+    """Return the lensed source surface brightness.
+
+    Args:
+        alpha_x: Initial x-component of deflection at each position.
+        alpha_y: Initial y-component of deflection at each position.
+        kwargs_lens_all: Keyword arguments and redshifts for each of the
+            lensing components. This should include los_before, los_after,
+            subhalos, and main_deflector.
+        kwargs_detector: Keyword arguments defining the detector configuration. This
+            includes potential supersampling in the lensing calculation.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z_source: Redshift of the source.
+        all_models: Tuple of model classes to consider for each component.
+
+    Returns:
+        Lensed source surface brightness as 1D array.
+    """
+    image_flux_array = _image_flux(
+        alpha_x,
+        alpha_y,
+        kwargs_lens_all,
+        kwargs_source_slice,
+        cosmology_params,
+        z_source,
+        all_models
+    )
+    # Scale by pixel area to go from flux to surface brightness.
+    return image_flux_array * kwargs_detector["pixel_width"] ** 2
+
+
+def _image_flux(
+    alpha_x: jnp.ndarray,
+    alpha_y: jnp.ndarray,
+    kwargs_lens_all: Mapping[str, Union[jnp.ndarray, Mapping[str, jnp.ndarray]]],
+    kwargs_source_slice: Mapping[str, jnp.ndarray],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z_source: float,
+    all_models: Mapping[str, Sequence[Any]],
+) -> jnp.ndarray:
+    """Calculate image flux after ray tracing onto the source.
+
+    Args:
+        alpha_x: Initial x-component of deflection at each position.
+        alpha_y: Initial y-component of deflectoin at each position.
+        kwargs_lens_all: Keyword arguments and redshifts for each of the
+            lensing components. This should include los_before, los_after,
+            subhalos, and main_deflector.
+        kwargs_source_slice: Keyword arguments to pass to brightness functions for
+            each light model. Keys should include parameters for all source models
+            (due to use of `jax.lax.switch`) and `model_index` which defines the model
+            to pass the parameters to. Parameter values not relevant for the specified
+            model are discarded. Values should be of type jnp.ndarray and have length
+            equal to the total number of light models in the slice. For more detailed
+            discussion see documentation of `jax.lax.scan`, which is used to iterate
+            over the models.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z_source: Redshift of the source.
+        all_models: Tuple of model classes to consider for each component.
+
+    Returns:
+        Image flux.
+    """
+    x_source_comv, y_source_comv = _ray_shooting(
+        alpha_x,
+        alpha_y,
+        kwargs_lens_all,
+        cosmology_params,
+        z_source,
+        all_models,
+    )
+    x_source, y_source = cosmology_utils.comoving_to_angle(
+        x_source_comv, y_source_comv, cosmology_params, z_source
+    )
+    return _surface_brightness(
+        x_source, y_source, kwargs_source_slice, all_models["all_source_models"]
+    )
 
 
 def _ray_shooting(
-    alpha_x,
-    alpha_y,
-    kwargs_lens_all,
-    cosmology_params,
-    z_lens_array,
-    z_source,
-):
-  """Ray shoot over all of the redshift slices between observer and source.
-
-  Args:
-    alpha_x: Initial x-component of deflection at each position.
-    alpha_y: Initial y-component of deflectoin at each position.
-    kwargs_lens_all: Keyword arguments for each lens model. The kwargs for
-      should contain keyword arguments to pass to derivative functions for each
-      lens model. Keys should include parameters for all lens models (due to use
-      of `jax.lax.switch`) and `model_index` which defines the model to pass the
-      parameters to. Parameter values not relevant for the specified model are
-      discarded. Values should be of type jnp.ndarray and have length equal to
-      the total number of lens models. For more detailed discussion see
-      documentation of `jax.lax.scan`, which is used to iterate over the models.
-    cosmology_params: Cosmological parameters that define the universe's
-      expansion.
-    z_lens_array: Redshifts for each lens model.
-    z_source: Redshift of the source.
-
-  Returns:
-    Comoving x- and y-coordinate after ray shooting.
-  """
-  # Initially all our light rays are localized at the observer.
-  comv_x = jnp.zeros_like(alpha_x)
-  comv_y = jnp.zeros_like(alpha_y)
-
-  z_lens_last = 0.0
-
-  # The state to pass to scan.
-  state = (comv_x, comv_y, alpha_x, alpha_y, z_lens_last)
-
-  ray_shooting_step = functools.partial(
-      _ray_shooting_step, cosmology_params=cosmology_params, z_source=z_source)
-
-  # Scan over all of the lens models in our system to calculate deflection and
-  # ray shoot between lens models.
-  final_state, _ = jax.lax.scan(ray_shooting_step, state, {
-      'z_lens': z_lens_array,
-      'kwargs_lens': kwargs_lens_all
-  })
-
-  comv_x, comv_y, alpha_x, alpha_y, z_lens_last = final_state
-
-  # Continue the ray tracing until the source.
-  delta_t = cosmology_utils.comoving_distance(cosmology_params, z_lens_last,
-                                              z_source)
-  comv_x, comv_y = _ray_step_add(comv_x, comv_y, alpha_x, alpha_y, delta_t)
-
-  return comv_x, comv_y
-
+    alpha_x: jnp.ndarray,
+    alpha_y: jnp.ndarray,
+    kwargs_lens_all:  Mapping[str, Union[jnp.ndarray, Mapping[str, jnp.ndarray]]],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z_source: float,
+    all_models: Mapping[str, Sequence[Any]]
+) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    """Ray shoot over all of the redshift slices between observer and source.
+
+    Args:
+        alpha_x: Initial x-component of deflection at each position.
+        alpha_y: Initial y-component of deflectoin at each position.
+        kwargs_lens_all: Keyword arguments and redshifts for each of the
+            lensing components. This should include los_before, los_after,
+            subhalos, and main_deflector.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z_source: Redshift of the source.
+        all_models: Tuple of model classes to consider for each component.
+
+    Returns:
+        Comoving x- and y-coordinate after ray shooting.
+    """
+    # Initially all our light rays are localized at the observer.
+    comv_x = jnp.zeros_like(alpha_x)
+    comv_y = jnp.zeros_like(alpha_y)
+
+    z_lens_last = 0.0
+
+    # The state to pass to scan.
+    state = (comv_x, comv_y, alpha_x, alpha_y, z_lens_last)
+
+    # Break the ray tracing into four parts: the first batch of line-of-sight
+    # halos, the subhalos, the main deflector, and then the second batch
+    # of line-of-sight halos. This circumvents issues with jax.lax.switch
+    # evaluating the expensive lens models for each deflector.
+    ray_shooting_step_los = functools.partial(
+        _ray_shooting_step,
+        cosmology_params=cosmology_params,
+        z_source=z_source,
+        all_lens_models=all_models["all_los_models"]
+    )
+    ray_shooting_step_main_deflector = functools.partial(
+        _ray_shooting_step,
+        cosmology_params=cosmology_params,
+        z_source=z_source,
+        all_lens_models=all_models["all_main_deflector_models"]
+    )
+
+    # Scan over all of the lens models in our system to calculate deflection and
+    # ray shoot between lens models.
+    state, _ = jax.lax.scan(
+        ray_shooting_step_los,
+        state,
+        {
+            "z_lens": kwargs_lens_all["z_array_los_before"],
+            "kwargs_lens": kwargs_lens_all["kwargs_los_before"]
+        },
+    )
+    # We can do all the subhalos at once, which is a lot faster for large
+    # number of subhalos.
+    state, _ = _ray_shooting_group(
+        state, kwargs_lens_all["kwargs_subhalos"], cosmology_params, z_source,
+        jnp.mean(kwargs_lens_all["z_array_subhalos"]),
+        all_models["all_subhalo_models"])
+    state, _ = _ray_shooting_group(
+        state, kwargs_lens_all["kwargs_main_deflector"], cosmology_params,
+        z_source, jnp.mean(kwargs_lens_all["z_array_main_deflector"]),
+        all_models["all_main_deflector_models"])
+    state, _ = jax.lax.scan(
+        ray_shooting_step_los,
+        state,
+        {
+            "z_lens": kwargs_lens_all["z_array_los_after"],
+            "kwargs_lens": kwargs_lens_all["kwargs_los_after"]
+        },
+    )
+
+    comv_x, comv_y, alpha_x, alpha_y, z_lens_last = state
+
+    # Continue the ray tracing until the source.
+    delta_t = cosmology_utils.comoving_distance(cosmology_params, z_lens_last, z_source)
+    comv_x, comv_y = _ray_step_add(comv_x, comv_y, alpha_x, alpha_y, delta_t)
+
+    return comv_x, comv_y
+
+
+def _ray_shooting_group(
+    state: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, float],
+    kwargs_lens_slice: Mapping[str, jnp.ndarray],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z_source: float,
+    z_lens: float,
+    all_lens_models: Sequence[Any]
+) -> Tuple[Tuple, Tuple]:
+    """Conduct ray shooting for a group of coplanar lens models.
+
+    Args:
+        state: The current comoving positions, deflections, and the previous lens
+            model redshift.
+        kwargs_lens_slice: Keyword arguments specifying the model (through
+            `model_index` value) and the lens model parameters for each lens
+            model in the slice. Due to the nature of `jax.lax.switch` this must
+            have a key-value pair for all lens models included in
+            `all_lens_models`, even if those lens models will not be used.
+            `model_index` of -1 indicates that the model should be ignored.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z_source: Redshift of the source.
+        z_lens: Redshift of the coplanar group of lens models.
+        all_lens_models: Lens models to use for model_index lookup.
+
+    Returns:
+        Two copies of the new state, which is a tuple of new comoving positions,
+        new deflection, and the redshift of the current lens.
+    """
+    comv_x, comv_y, alpha_x, alpha_y, z_lens_last = state
+
+    # The displacement from moving along the deflection direction.
+    delta_t = cosmology_utils.comoving_distance(
+        cosmology_params, z_lens_last, z_lens,
+    )
+    comv_x, comv_y = _ray_step_add(comv_x, comv_y, alpha_x, alpha_y, delta_t)
+    alpha_x, alpha_y = _add_deflection_group(
+        comv_x, comv_y, alpha_x, alpha_y, kwargs_lens_slice,
+        cosmology_params, z_lens, z_source, all_lens_models)
+
+    new_state = (comv_x, comv_y, alpha_x, alpha_y, z_lens)
+
+    # Second return is required by scan, but will be ignored by the compiler.
+    return new_state, new_state
 
 def _ray_shooting_step(
-    state,
-    kwargs_z_lens,
-    cosmology_params,
-    z_source,
-):
-  """Conduct ray shooting between two lens models.
-
-  Args:
-    state: The current comoving positions, deflections, and the previous lens
-      model redshift.
-    kwargs_z_lens: Dict with keys `z_lens`, the redshift of the next lens model,
-      and 'kwargs_lens`, the keyword arguments specifying the next lens model
-      (through `model_index` value) and the lens model parameters. Due to the
-      requirements of `jax.lax.switch`, `kwargs_lens` must have a key-value pair
-      for all lens models included in `lens_models.py`, even if those lens
-      models will not be used.
-    cosmology_params: Cosmological parameters that define the universe's
-      expansion.
-    z_source: Redshift of the source.
-
-  Returns:
-    Two copies of the new state, which is a tuple of new comoving positions, new
-    deflection, and
-    the redshift of the current lens.
-
-  Notes:
-    For use with `jax.lax.scan`.
-  """
-  comv_x, comv_y, alpha_x, alpha_y, z_lens_last = state
-
-  # The displacement from moving along the deflection direction.
-  delta_t = cosmology_utils.comoving_distance(cosmology_params, z_lens_last,
-                                              kwargs_z_lens['z_lens'])
-  comv_x, comv_y = _ray_step_add(comv_x, comv_y, alpha_x, alpha_y, delta_t)
-  alpha_x, alpha_y = _add_deflection(comv_x, comv_y, alpha_x, alpha_y,
-                                     kwargs_z_lens['kwargs_lens'],
-                                     cosmology_params, kwargs_z_lens['z_lens'],
-                                     z_source)
-
-  new_state = (comv_x, comv_y, alpha_x, alpha_y, kwargs_z_lens['z_lens'])
-
-  # Second return is required by scan, but will be ignored by the compiler.
-  return new_state, new_state
-
-
-def _ray_step_add(comv_x, comv_y,
-                  alpha_x, alpha_y,
-                  delta_t):
-  """Conduct ray step by adding deflections to comoving positions.
-
-  Args:
-    comv_x: Comoving x-coordinate.
-    comv_y: Comoving y-coordinate.
-    alpha_x: Current physical x-component of deflection at each position.
-    alpha_y: Current physical y-component of deflection at each position.
-    delta_t: Comoving distance between current slice and next slice.
-
-  Returns:
-    Updated x- and y-coordinate in comoving Mpc.
-  """
-  return comv_x + alpha_x * delta_t, comv_y + alpha_y * delta_t
-
-
-def _add_deflection(comv_x, comv_y,
-                    alpha_x, alpha_y,
-                    kwargs_lens,
-                    cosmology_params,
-                    z_lens,
-                    z_source):
-  """Calculate the deflection for a specific lens model.
-
-  Args:
-    comv_x: Comoving x-coordinate.
-    comv_y: Comoving y-coordinate.
-    alpha_x: Current physical x-component of deflection at each position.
-    alpha_y: Current physical y-component of deflection at each position.
-    kwargs_lens: Keyword arguments specifying the model (through `model_index`
-      value) and the lens model parameters. Due to the nature of
-      `jax.lax.switch` this must have a key-value pair for all lens models
-      included in `lens_models.py`, even if those lens models will not be used.
-      `model_index` of -1 indicates that the previous total should be returned.
-    cosmology_params: Cosmological parameters that define the universe's
-      expansion.
-    z_lens: Redshift of the slice (i.e. current redshift).
-    z_source: Redshift of the source.
-
-  Returns:
-    New x- and y-component of the deflection at each specified position.
-  """
-  theta_x, theta_y = cosmology_utils.comoving_to_angle(comv_x, comv_y,
-                                                       cosmology_params, z_lens)
-
-  # All of our derivatives are defined in reduced coordinates.
-  alpha_x_reduced, alpha_y_reduced = _calculate_derivatives(
-      kwargs_lens, theta_x, theta_y)
-
-  alpha_x_update = cosmology_utils.reduced_to_physical(alpha_x_reduced,
-                                                       cosmology_params, z_lens,
-                                                       z_source)
-  alpha_y_update = cosmology_utils.reduced_to_physical(alpha_y_reduced,
-                                                       cosmology_params, z_lens,
-                                                       z_source)
-
-  return alpha_x - alpha_x_update, alpha_y - alpha_y_update
+    state: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, float],
+    kwargs_z_lens: Mapping[str, Union[jnp.ndarray, Mapping[str, jnp.ndarray]]],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z_source: float,
+    all_lens_models: Sequence[Any]
+) -> Tuple[Tuple, Tuple]:
+    """Conduct ray shooting between two lens models.
+
+    Args:
+        state: The current comoving positions, deflections, and the previous lens
+            model redshift.
+        kwargs_z_lens: Dict with keys `z_lens`, the redshift of the next lens model,
+            and 'kwargs_lens`, the keyword arguments specifying the next lens model
+            (through `model_index` value) and the lens model parameters. Due to the
+            requirements of `jax.lax.switch`, `kwargs_lens` must have a key-value pair
+            for all lens models included in `all_lens_models`, even if those lens
+            models will not be used.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z_source: Redshift of the source.
+        all_lens_models: Lens models to use for model_index lookup.
+
+    Returns:
+        Two copies of the new state, which is a tuple of new comoving positions,
+        new deflection, and the redshift of the current lens.
+
+    Notes:
+        For use with `jax.lax.scan`.
+    """
+    comv_x, comv_y, alpha_x, alpha_y, z_lens_last = state
+
+    # The displacement from moving along the deflection direction.
+    delta_t = cosmology_utils.comoving_distance(
+        cosmology_params, z_lens_last, kwargs_z_lens["z_lens"]
+    )
+    comv_x, comv_y = _ray_step_add(comv_x, comv_y, alpha_x, alpha_y, delta_t)
+    alpha_x, alpha_y = _add_deflection(
+        comv_x,
+        comv_y,
+        alpha_x,
+        alpha_y,
+        kwargs_z_lens["kwargs_lens"],
+        cosmology_params,
+        kwargs_z_lens["z_lens"],
+        z_source,
+        all_lens_models,
+    )
+
+    new_state = (comv_x, comv_y, alpha_x, alpha_y, kwargs_z_lens["z_lens"])
+
+    # Second return is required by scan, but will be ignored by the compiler.
+    return new_state, new_state
+
+
+def _ray_step_add(
+    comv_x: jnp.ndarray,
+    comv_y: jnp.ndarray,
+    alpha_x: jnp.ndarray,
+    alpha_y: jnp.ndarray,
+    delta_t: float
+) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    """Conduct ray step by adding deflections to comoving positions.
+
+    Args:
+        comv_x: Comoving x-coordinate.
+        comv_y: Comoving y-coordinate.
+        alpha_x: Current physical x-component of deflection at each position.
+        alpha_y: Current physical y-component of deflection at each position.
+        delta_t: Comoving distance between current slice and next slice.
+
+    Returns:
+        Updated x- and y-coordinate in comoving Mpc.
+    """
+    return comv_x + alpha_x * delta_t, comv_y + alpha_y * delta_t
+
+
+def _add_deflection_group(
+    comv_x: jnp.ndarray,
+    comv_y: jnp.ndarray,
+    alpha_x: jnp.ndarray,
+    alpha_y: jnp.ndarray,
+    kwargs_lens_slice: Mapping[str, jnp.ndarray],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z_lens: float,
+    z_source: float,
+    all_lens_models: Sequence[Any],
+) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    """Calculate the deflection angle for a group of co-planar lens models.
+
+    Args:
+        comv_x: Comoving x-coordinate.
+        comv_y: Comoving y-coordinate.
+        alpha_x: Current physical x-component of deflection at each position.
+        alpha_y: Current physical y-component of deflection at each position.
+        kwargs_lens_slice: Keyword arguments specifying the model (through
+            `model_index` value) and the lens model parameters for each lens
+            model in the slice. Due to the nature of `jax.lax.switch` this must
+            have a key-value pair for all lens models included in
+            `all_lens_models`, even if those lens models will not be used.
+            `model_index` of -1 indicates that the model should be ignored.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z_lens: Redshift of the slice (i.e. current redshift).
+        z_source: Redshift of the source.
+        all_lens_models: Lens models to use for model_index lookup.
+
+    Returns:
+        New x- and y-component of the deflection at each specified position.
+    """
+    theta_x, theta_y = cosmology_utils.comoving_to_angle(
+        comv_x, comv_y, cosmology_params, z_lens
+    )
+
+    # All of our derivatives are defined in reduced coordinates. We want to
+    # scan over all the calculation in reduced coordinates to improve
+    # parallelization performance.
+    calculate_derivatives = functools.partial(
+        _calculate_derivatives, theta_x=theta_x, theta_y=theta_y,
+        all_lens_models=all_lens_models)
+    alpha_x_reduced_array, alpha_y_reduced_array = jax.vmap(
+        calculate_derivatives, in_axes=0)(kwargs_lens_slice)
+    alpha_x_reduced = jnp.sum(alpha_x_reduced_array, axis=0)
+    alpha_y_reduced = jnp.sum(alpha_y_reduced_array, axis=0)
+
+    alpha_x_update = cosmology_utils.reduced_to_physical(
+        alpha_x_reduced, cosmology_params, z_lens, z_source
+    )
+    alpha_y_update = cosmology_utils.reduced_to_physical(
+        alpha_y_reduced, cosmology_params, z_lens, z_source
+    )
+
+    return alpha_x - alpha_x_update, alpha_y - alpha_y_update
+
+
+def _add_deflection(
+    comv_x: jnp.ndarray,
+    comv_y: jnp.ndarray,
+    alpha_x: jnp.ndarray,
+    alpha_y: jnp.ndarray,
+    kwargs_lens: Mapping[str, Union[int, float, jnp.ndarray]],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z_lens: float,
+    z_source: float,
+    all_lens_models: Sequence[Any],
+) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    """Calculate the deflection for a specific lens model.
+
+    Args:
+        comv_x: Comoving x-coordinate.
+        comv_y: Comoving y-coordinate.
+        alpha_x: Current physical x-component of deflection at each position.
+        alpha_y: Current physical y-component of deflection at each position.
+        kwargs_lens: Keyword arguments specifying the model (through `model_index`
+            value) and the lens model parameters. Due to the nature of
+            `jax.lax.switch` this must have a key-value pair for all lens models
+            included in `all_lens_models`, even if those lens models will not be used.
+            `model_index` of -1 indicates that the previous total should be returned.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z_lens: Redshift of the slice (i.e. current redshift).
+        z_source: Redshift of the source.
+        all_lens_models: Lens models to use for model_index lookup.
+
+    Returns:
+        New x- and y-component of the deflection at each specified position.
+    """
+    theta_x, theta_y = cosmology_utils.comoving_to_angle(
+        comv_x, comv_y, cosmology_params, z_lens
+    )
+
+    # All of our derivatives are defined in reduced coordinates.
+    alpha_x_reduced, alpha_y_reduced = _calculate_derivatives(
+        kwargs_lens, theta_x, theta_y, all_lens_models
+    )
+
+    alpha_x_update = cosmology_utils.reduced_to_physical(
+        alpha_x_reduced, cosmology_params, z_lens, z_source
+    )
+    alpha_y_update = cosmology_utils.reduced_to_physical(
+        alpha_y_reduced, cosmology_params, z_lens, z_source
+    )
+
+    return alpha_x - alpha_x_update, alpha_y - alpha_y_update
 
 
 def _calculate_derivatives(
-    kwargs_lens,
-    theta_x,
-    theta_y,
-):
-  """Calculate the derivatives for the specified lens model.
-
-  Args:
-    kwargs_lens: Keyword arguments specifying the model (through `model_index`
-      value) and the lens model parameters. Due to the nature of
-      `jax.lax.switch` this must have a key-value pair for all lens models
-      included in `lens_models.py`, even if those lens models will not be used.
-      `model_index` of -1 indicates that the previous total should be returned.
-    theta_x: X-coordinate at which to evaluate the derivative.
-    theta_y: Y-coordinate at which to evaluate the derivative.
-
-  Returns:
-    Change in x- and y-component of derivative caused by lens model. X- and
-    y-component stacked along axis=0.
-  """
-  # The jax.lax.switch compilation requires that all of the functions take the
-  # same inputs. We accomplish this using our wrapper and picking out only the
-  # parameters required for that model.
-  derivative_functions = [
-      utils.unpack_parameters_xy(LENS_MODELS[model_index].derivatives,
-                                 LENS_MODEL_PARAMETERS[model_index])
-      for model_index in range(len(LENS_MODELS))
-  ]
-
-  # Condition on model_index to allow for padding.
-  def calculate_derivative():
-    alpha_x, alpha_y = jax.lax.switch(kwargs_lens['model_index'],
-                                      derivative_functions, theta_x, theta_y,
-                                      kwargs_lens)
-    return [alpha_x, alpha_y]
-
-  def identity():
-    return [jnp.zeros_like(theta_x), jnp.zeros_like(theta_y)]
-
-  return jax.lax.cond(kwargs_lens['model_index'] == -1, identity,
-                      calculate_derivative)
+    kwargs_lens: Mapping[str, Union[int, float, jnp.ndarray]],
+    theta_x: jnp.ndarray,
+    theta_y: jnp.ndarray,
+    all_lens_models: Sequence[Any],
+) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    """Calculate the derivatives for the specified lens model.
+
+    Args:
+        kwargs_lens: Keyword arguments specifying the model (through
+            `model_index` value) and the lens model parameters. Due to the
+            nature of `jax.lax.switch` this must have a key-value pair for all
+            lens models included in `all_lens_models`, even if those lens models
+            will not be used. `model_index` of -1 indicates that the previous
+            total should be returned.
+        theta_x: X-coordinate at which to evaluate the derivative.
+        theta_y: Y-coordinate at which to evaluate the derivative.
+        all_lens_models: Lens models to use for model_index lookup.
+
+    Returns:
+        Change in x- and y-component of derivative caused by lens model.
+    """
+    # The jax.lax.switch compilation requires that all of the functions take the
+    # same inputs. We accomplish this using our wrapper and picking out only the
+    # parameters required for that model.
+    derivative_functions = [
+        utils.unpack_parameters_xy(model.derivatives, model.parameters)
+        for model in all_lens_models
+    ]
+
+    # Condition on model_index to allow for padding.
+    def calculate_derivative():
+        alpha_x, alpha_y = jax.lax.switch(
+            kwargs_lens["model_index"],
+            derivative_functions,
+            theta_x,
+            theta_y,
+            kwargs_lens,
+        )
+        return (alpha_x, alpha_y)
+
+    def identity():
+        return (jnp.zeros_like(theta_x), jnp.zeros_like(theta_y))
+
+    return jax.lax.cond(
+        kwargs_lens["model_index"] == -1, identity, calculate_derivative
+    )
 
 
 def _surface_brightness(
-    theta_x, theta_y,
-    kwargs_source_slice):
-  """Return the surface brightness for a slice of light models.
-
-  Args:
-    theta_x: X-coordinate at which to evaluate the surface brightness.
-    theta_y: Y-coordinate at which to evaluate the surface brightness.
-    kwargs_source_slice: Keyword arguments to pass to brightness functions for
-      each light model. Keys should include parameters for all source models
-      (due to use of `jax.lax.switch`) and `model_index` which defines the model
-      to pass the parameters to. Parameter values not relevant for the specified
-      model are discarded. Values should be of type jnp.ndarray and have length
-      equal to the total number of light models in the slice. For more detailed
-      discussion see documentation of `jax.lax.scan`, which is used to iterate
-      over the models.
-
-  Returns:
-    Surface brightness summed over all sources.
-  """
-  add_surface_brightness = functools.partial(
-      _add_surface_brightness, theta_x=theta_x, theta_y=theta_y)
-  brightness_total = jnp.zeros_like(theta_x)
-  brightness_total, _ = jax.lax.scan(add_surface_brightness, brightness_total,
-                                     kwargs_source_slice)
-  return brightness_total
-
-
-def _add_surface_brightness(prev_brightness,
-                            kwargs_source,
-                            theta_x,
-                            theta_y):
-  """Return the surface brightness for a single light model.
-
-  Args:
-    prev_brightness: Previous brightness.
-    kwargs_source: Kwargs to evaluate the source surface brightness.
-    theta_x: X-coordinate at which to evaluate the surface brightness.
-    theta_y: Y-coordinate at which to evaluate the surface brightness.
-
-  Returns:
-    Surface brightness of the source at given coordinates.
-  """
-  source_functions = [
-      utils.unpack_parameters_xy(SOURCE_MODELS[model_index].function,
-                                 SOURCE_MODEL_PARAMETERS[model_index])
-      for model_index in range(len(SOURCE_MODELS))
-  ]
-  brightness = jax.lax.switch(kwargs_source['model_index'], source_functions,
-                              theta_x, theta_y, kwargs_source)
-
-  return prev_brightness + brightness, brightness
+    theta_x: jnp.ndarray,
+    theta_y: jnp.ndarray,
+    kwargs_source_slice: Mapping[str, jnp.ndarray],
+    all_source_models: Sequence[Any]
+) -> jnp.ndarray:
+    """Return the surface brightness for a slice of light models.
+
+    Args:
+        theta_x: X-coordinate at which to evaluate the surface brightness.
+        theta_y: Y-coordinate at which to evaluate the surface brightness.
+        kwargs_source_slice: Keyword arguments to pass to brightness functions for
+            each light model. Keys should include parameters for all source models
+            included in `all_source_models`.
+            (due to use of `jax.lax.switch`) and `model_index` which defines the model
+            to pass the parameters to. Parameter values not relevant for the specified
+            model are discarded. Values should be of type jnp.ndarray and have length
+            equal to the total number of light models in the slice. For more detailed
+            discussion see documentation of `jax.lax.scan`, which is used to iterate
+            over the models.
+        all_source_models: Source models to use for model_index lookup.
+
+    Returns:
+        Surface brightness summed over all sources.
+    """
+    add_surface_brightness = functools.partial(
+        _add_surface_brightness,
+        theta_x=theta_x,
+        theta_y=theta_y,
+        all_source_models=all_source_models,
+    )
+    brightness_total = jnp.zeros_like(theta_x)
+    brightness_total, _ = jax.lax.scan(
+        add_surface_brightness, brightness_total, kwargs_source_slice
+    )
+    return brightness_total
+
+
+def _add_surface_brightness(
+    prev_brightness: jnp.ndarray,
+    kwargs_source: Mapping[str, Union[int, float, jnp.ndarray]],
+    theta_x: jnp.ndarray,
+    theta_y: jnp.ndarray,
+    all_source_models: Sequence[Any]
+) -> jnp.ndarray:
+    """Return the surface brightness for a single light model.
+
+    Args:
+        prev_brightness: Previous brightness.
+        kwargs_source: Kwargs to evaluate the source surface brightness.
+        theta_x: X-coordinate at which to evaluate the surface brightness.
+        theta_y: Y-coordinate at which to evaluate the surface brightness.
+        all_source_models: Source models to use for model_index lookup.
+
+    Returns:
+        Surface brightness of the source at given coordinates.
+    """
+    source_functions = [
+        utils.unpack_parameters_xy(model.function, model.parameters)
+        for model in all_source_models
+    ]
+    brightness = jax.lax.switch(
+        kwargs_source["model_index"], source_functions, theta_x, theta_y,
+        kwargs_source
+    )
+
+    return prev_brightness + brightness, brightness
diff --git a/jaxstronomy/input_config.py b/jaxstronomy/input_config.py
new file mode 100644
index 0000000..002f3a6
--- /dev/null
+++ b/jaxstronomy/input_config.py
@@ -0,0 +1,136 @@
+# coding=utf-8
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Configuration file for generating jaxstronomy image outputs.
+"""
+
+import jax
+import jax.numpy as jnp
+
+from jaxstronomy.input_pipeline import encode_normal, encode_uniform
+from jaxstronomy.input_pipeline import encode_constant
+from jaxstronomy import lens_models
+from jaxstronomy import psf_models
+from jaxstronomy import source_models
+
+
+def get_config():
+    """Return the config for input generation.
+    """
+    config = {}
+
+    # Specify the configuration for each one of the parameters.
+    config['lensing_config'] = {
+        'los_params':{
+            'delta_los': encode_constant(0.0),
+            'r_min': encode_constant(0.5),
+            'r_max': encode_constant(10.0),
+            'm_min': encode_constant(1e7),
+            'm_max': encode_constant(1e10),
+            'dz': encode_constant(0.01),
+            'cone_angle': encode_constant(8.0),
+            'angle_buffer': encode_constant(0.8),
+            'c_zero': encode_constant(18),
+            'conc_zeta': encode_constant(-0.2),
+            'conc_beta': encode_constant(0.8),
+            'conc_m_ref': encode_constant(1e8),
+            'conc_dex_scatter': encode_constant(0.0)
+        },
+        'main_deflector_params': {
+            'mass': encode_constant(1e13),
+            'z_lens': encode_constant(0.5),
+            'theta_e': encode_normal(mean=1.1, std=0.15),
+            'slope': encode_normal(mean=2.0, std=0.1),
+            'center_x': encode_normal(mean=0.0, std=0.16),
+            'center_y': encode_normal(mean=0.0, std=0.16),
+            'axis_ratio': encode_normal(mean=1.0, std=0.1),
+            'angle': encode_uniform(minimum=0.0, maximum=2 * jnp.pi),
+            'gamma_ext': encode_normal(mean=0.0, std=0.1)
+        },
+        'subhalo_params':{
+            'sigma_sub': encode_normal(mean=2.0e-3, std=1.1e-3),
+            'shmf_plaw_index': encode_uniform(minimum=-2.02, maximum=-1.92),
+            'm_pivot': encode_constant(1e10),
+            'm_min': encode_constant(7e7),
+            'm_max': encode_constant(1e10),
+            'k_one': encode_constant(0.0),
+            'k_two': encode_constant(0.0),
+            'c_zero': encode_constant(18),
+            'conc_zeta': encode_constant(-0.2),
+            'conc_beta': encode_constant(0.8),
+            'conc_m_ref': encode_constant(1e8),
+            'conc_dex_scatter': encode_constant(0.0)
+        },
+        'source_params':{
+            'z_source': encode_constant(1.5),
+            'amp': encode_uniform(minimum=1.0, maximum=10.0),
+            'sersic_radius': encode_uniform(minimum=1.0, maximum=3.0),
+            'n_sersic': encode_uniform(minimum=1.0, maximum=4.0),
+            'axis_ratio': encode_normal(mean=1.0, std=0.05),
+            'angle': encode_uniform(minimum=0.0, maximum=2 * jnp.pi),
+            'center_x': encode_normal(mean=0.0, std=0.16),
+            'center_y': encode_normal(mean=0.0, std=0.16)
+        },
+        'lens_light_params':{
+            'amp': encode_constant(0.0),
+            'sersic_radius': encode_uniform(minimum=1.0, maximum=3.0),
+            'n_sersic': encode_uniform(minimum=1.0, maximum=4.0),
+            'axis_ratio': encode_normal(mean=1.0, std=0.05),
+            'angle': encode_uniform(minimum=0.0, maximum=2 * jnp.pi),
+            'center_x': encode_normal(mean=0.0, std=0.16),
+            'center_y': encode_normal(mean=0.0, std=0.16)
+        }
+    }
+
+    # The remaining parameters should not be drawn from random distributions.
+    config['kwargs_detector'] = {
+        'n_x': 128, 'n_y': 128, 'pixel_width': 0.04, 'supersampling_factor': 2,
+        'exposure_time': 1024, 'num_exposures': 2.0, 'sky_brightness': 22,
+        'magnitude_zero_point': 25, 'read_noise': 3.0
+    }
+    config['all_models'] = {
+        'all_los_models': (lens_models.NFW,),
+        'all_subhalo_models': (lens_models.TNFW,),
+        'all_main_deflector_models': (lens_models.EPL, lens_models.Shear),
+        'all_source_models': (source_models.SersicElliptic,),
+        'all_psf_models': (psf_models.Gaussian,)
+    }
+    config['cosmology_params'] = {
+        'omega_m_zero': encode_constant(0.3089),
+        'omega_b_zero': encode_constant(0.0486),
+        'omega_de_zero': encode_constant(0.6910088292453472),
+        'omega_rad_zero': encode_constant(9.117075466e-5),
+        'temp_cmb_zero': encode_constant(2.7255),
+        'hubble_constant': encode_constant(67.74),
+        'n_s': encode_constant(0.9667),
+        'sigma_eight': encode_constant(0.815)
+    }
+    config['kwargs_simulation'] = {
+        'num_z_bins': 1000,
+        'los_pad_length': 10,
+        'subhalos_pad_length': 750,
+        'sampling_pad_length': 10000,
+    }
+
+    config['rng'] = jax.random.PRNGKey(0)
+
+    config['kwargs_psf'] = {'model_index': 0, 'fwhm': 0.04, 'pixel_width': 0.02}
+
+    config['principal_md_index'] = 0
+    config['principal_source_index'] = 0
+    config['truth_parameters'] = (
+        ['main_deflector_params', 'main_deflector_params',
+         'main_deflector_params', 'main_deflector_params', 'subhalo_params'],
+        ['theta_e', 'slope', 'center_x', 'center_y', 'sigma_sub'])
+
+    return config
diff --git a/jaxstronomy/input_pipeline.py b/jaxstronomy/input_pipeline.py
new file mode 100644
index 0000000..fc47a48
--- /dev/null
+++ b/jaxstronomy/input_pipeline.py
@@ -0,0 +1,402 @@
+# coding=utf-8
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Code for drawing batches of example images used to train a neural
+network.
+"""
+
+import functools
+from typing import Any, Mapping, Sequence, Tuple, Union
+
+import jax.numpy as jnp
+import jax
+
+from jaxstronomy import cosmology_utils
+from jaxstronomy import los
+from jaxstronomy import subhalos
+from jaxstronomy import image_simulation
+from jaxstronomy import utils
+
+
+def encode_normal(mean: float, std: float) -> jnp.ndarray:
+    """Generate the jax array that encodes a normal distribution.
+
+    Args:
+        mean: Mean of the normal distribution.
+        std: Standard deviation of the normal distribution.
+
+    Returns:
+        Encoding that represents a normal with given mean and standard
+            deviation.
+    """
+
+    # Encoding is currently [uniform_bool, unfirom_min, uniform_max,
+    # normal_bool, normal_mean, normal_std, constant].
+    return jnp.array([0.0, 0.0, 0.0, 1.0, mean, std, 0.0])
+
+
+def encode_uniform(minimum: float, maximum: float) -> jnp.ndarray:
+    """Generate the jax array that encodes a uniform distribution.
+
+    Args:
+        min: Minimum value of the uniform distribution.
+        max: Maximum value of the uniform distribution.
+
+    Returns:
+        Encoding that represents a uniform with given min and max.
+    """
+
+    # Encoding is currently [uniform_bool, unfirom_min, uniform_max,
+    # normal_bool, normal_mean, normal_std, constant].
+    return jnp.array([1.0, minimum, maximum, 0.0, 0.0, 0.0, 0.0])
+
+
+def encode_constant(constant: float) -> jnp.ndarray:
+    """Generate the jax array that encodes a constant value.
+
+    Args:
+        constant: Constant value.
+
+    Returns:
+        Encoding that represents a constant value.
+    """
+    # Encoding is currently [uniform_bool, unfirom_min, uniform_max,
+    # normal_bool, normal_mean, normal_std, constant].
+    return jnp.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, constant])
+
+
+def decode_maximum(encoding: jnp.ndarray) -> float:
+    """Decode the maximum value of the distribution defined by encoding.
+
+    Args:
+        encoding: Encoded distribution
+
+    Returns:
+        Maximum value of encoded distribution. May be an approximation for
+        distributions without a maximum.
+    """
+    # Encoding is currently [uniform_bool, unfirom_min, uniform_max,
+    # normal_bool, normal_mean, normal_std, constant].
+    # If not uniform will give 0.
+    maximum = encoding[0] * encoding[2]
+    # If normal will give mean + five sigma.
+    maximum += encoding[3] * (encoding[4] + encoding[5]*5)
+    # If constant return constant
+    maximum += encoding[6]
+    return maximum
+
+
+def decode_minimum(encoding: jnp.ndarray) -> float:
+    """Decode the minimum value of the dsitribution defined by the encoding.
+
+    Args:
+        encoding: Encoded distribution
+
+    Returns:
+        Minimum value of encoded distribution. May be an approximation for
+        distributions without a minimum.
+    """
+    # Encoding is currently [uniform_bool, unfirom_min, uniform_max,
+    # normal_bool, normal_mean, normal_std, constant].
+    # If not uniform will give 0.
+    minimum = encoding[0] * encoding[1]
+    # If normal will give mean - five sigma.
+    minimum += encoding[3] * (encoding[4] - encoding[5]*5)
+    # If constant return constant
+    minimum += encoding[6]
+    return minimum
+
+
+def draw_from_encoding(encoding: jnp.ndarray, rng: Sequence[int]) -> float:
+    """Draw from encoded distribution.
+
+    Args:
+        encoding: Encoding defining the distribution.
+        rng: jax PRNG key.
+
+    Returns:
+        A draw from the encoded distribution.
+    """
+    # Encoding is currently [uniform_bool, unfirom_min, uniform_max,
+    # normal_bool, normal_mean, normal_std, constant].
+    # Start with the uniform component.
+    draw = encoding[0] * (
+        jax.random.uniform(rng) * (encoding[2] - encoding[1]) + encoding[1])
+    # Now the normal component.
+    draw += encoding[3] * (jax.random.normal(rng) * encoding[5] + encoding[4])
+    # And finally the constant
+    draw += encoding[6]
+    return draw
+
+
+def normalize_param(parameter: float, encoding: jnp.ndarray) -> float:
+    """Return parameter normalized by encoded distribution.
+
+    Args:
+        parameter: Parameter value to normalize
+        encoding: Encoding defining the distribution.
+
+    Returns:
+        Normalized parameter.
+    """
+    # Encoding is currently [uniform_bool, unfirom_min, uniform_max,
+    # normal_bool, normal_mean, normal_std, constant].
+    # Normalize uniform distribution to be between 0 and 1.
+    normalized_param = jax.lax.select(
+        encoding[0] > 0.0,
+        (parameter - encoding[1]) / (encoding[2] - encoding[1]), 0.0)
+    # Normalize normal distribution to mean 0 and standard deviation 1.0
+    normalized_param += jax.lax.select(
+        encoding[3] > 0.0,(parameter - encoding[4]) / encoding[5], 0.0)
+    # Constant will be normalzied to 0.0 by default.
+
+    return normalized_param
+
+def generate_grids(
+        config: Mapping[str, Mapping[str, jnp.ndarray]]
+) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    """Generate the x- and y-grid on which to generate the lensing image.
+
+    Args:
+        config: Configuration dictionary from which the detector kwargs will be
+            drawn.
+
+    Returns:
+        x- and y-grid in units of arcseconds as a tuple.
+    """
+    kwargs_detector = config['kwargs_detector']
+    grid_x, grid_y = utils.coordinates_evaluate(kwargs_detector['n_x'],
+        kwargs_detector['n_y'], kwargs_detector['pixel_width'],
+        kwargs_detector['supersampling_factor'])
+    return grid_x, grid_y
+
+
+def intialize_cosmology_params(
+        config: Mapping[str, Mapping[str, jnp.ndarray]], rng: Sequence[int],
+) -> Mapping[str, Union[float, int, jnp.ndarray]]:
+    """Initialize the cosmology parameters as needed by the config.
+
+    Args:
+        config: Configuration dictionary for input generation.
+        rng: jax PRNG key.
+
+    Returns:
+        Cosmological parameters with appropriate lookup table.
+    """
+    max_source_z = decode_maximum(
+        config['lensing_config']['source_params']['z_source'])
+    dz = decode_minimum(config['lensing_config']['los_params']['dz'])
+    m_max = max(
+        decode_maximum(config['lensing_config']['subhalo_params']['m_max']),
+        decode_maximum(config['lensing_config']['los_params']['m_max']))
+    m_min = min(
+        decode_minimum(config['lensing_config']['subhalo_params']['m_min']),
+        decode_minimum(config['lensing_config']['los_params']['m_min']))
+
+    cosmology_params_init = draw_sample(config['cosmology_params'], rng)
+
+    # Initial bounds on lagrangian radius are just placeholders.
+    cosmology_params = cosmology_utils.add_lookup_tables_to_cosmology_params(
+        cosmology_params_init, max_source_z, dz / 2, 1e-4, 1e3, 2)
+    r_min = cosmology_utils.lagrangian_radius(cosmology_params,
+                                              m_min / 10)
+    r_max = cosmology_utils.lagrangian_radius(cosmology_params,
+                                              m_max * 10)
+    cosmology_params = cosmology_utils.add_lookup_tables_to_cosmology_params(
+        cosmology_params_init, 1.5, dz / 2, r_min, r_max, 10000)
+    extrenal_los_params = {'m_min': m_min, 'm_max': m_max, 'dz': dz}
+    return los.add_los_lookup_tables_to_cosmology_params(
+        extrenal_los_params, cosmology_params, max_source_z)
+
+
+def draw_sample(
+        encoded_configuration: Mapping[str, Mapping[str, float]],
+        rng: Sequence[int]
+) -> Mapping[str, Mapping[str, float]]:
+    """Map an econded configuration into a configuration of randomly draws.
+
+    Args:
+        encoded_configuration: Configuration with encoded distribution for all
+            leaves of the PyTree structure.
+        rng: jax PRNG key.
+
+    Returns:
+        Configuration with encoded distribution replaced by draws from those
+            distributions.
+    """
+    # Generate the rng keys we will need for each leaf.
+    treedef = jax.tree_util.tree_structure(encoded_configuration)
+    rng_keys = jax.random.split(rng, treedef.num_leaves)
+    rng_tree = jax.tree_util.tree_unflatten(treedef, rng_keys)
+    return jax.tree_util.tree_map(draw_from_encoding, encoded_configuration,
+                                  rng_tree)
+
+
+def extract_multiple_models(
+        encoded_configuration: Mapping[str, Mapping[str, float]],
+        rng: Sequence[int], n_models: int
+) -> Mapping[str, jnp.ndarray]:
+    """Extract multiple models from a single configuration.
+
+    Args:
+        encoded_configuration: Encodings for each of the parameters of the
+            model(s).
+        rng: jax PRNG key.
+        n_models: Number of models to draw parameters for.
+
+    Returns:
+        Draws for the parameters for all the models. For each parameters, first
+        dimension will be the number of models.
+    """
+    draw_sample_vmap = jax.vmap(draw_sample, in_axes=[None, 0])
+    rng_list = jax.random.split(rng, n_models)
+    draws = draw_sample_vmap(encoded_configuration, rng_list)
+    # Add the model index list.
+    draws['model_index'] = jnp.arange(n_models)
+    return draws
+
+
+def extract_truth_values(
+        all_params: Mapping[str, Mapping[str, jnp.ndarray]],
+        lensing_config: Mapping[str, Mapping[str, jnp.ndarray]],
+        truth_parameters: Tuple[Sequence[str], Sequence[str]]) -> jnp.ndarray:
+    """Extract the truth parameters and normalize them according to the config.
+
+    Args:
+        all_params: All of the parameters grouped by object.
+        lensing_config: Distribution encodings for each of the parameters.
+        truth_parameters: List of the lensing objects and corresponding
+            parameters to extract.
+
+    Returns:
+        Truth values for each of the requested parameters.
+    """
+    extract_objects, extract_keys = truth_parameters
+    return jnp.array(jax.tree_util.tree_map(
+        lambda x, y: normalize_param(all_params[x][y],
+                                     lensing_config[x][y]),
+        extract_objects, extract_keys))
+
+
+def draw_image_and_truth(
+        lensing_config: Mapping[str, Mapping[str, jnp.ndarray]],
+        cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+        grid_x: jnp.ndarray, grid_y: jnp.ndarray, rng: Sequence[int],
+        all_models: Mapping[str, Sequence[Any]],
+        principal_md_index: int, principal_source_index: int,
+        kwargs_simulation: Mapping[str, int],
+        kwargs_detector:  Mapping[str, Union[int, float]],
+        kwargs_psf: Mapping[str, Union[float, int, jnp.ndarray]],
+        truth_parameters: Tuple[Sequence[str], Sequence[str]]
+) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    """Draw image and truth values for a realization of the lensing config.
+
+    Args:
+        lensing_config: Distribution encodings for each of the objects in the
+            lensing system.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        grid_x: x-grid in units of arcseconds.
+        grid_y: y-grid in units of arcseconds.
+        rng: jax PRNG key.
+        all_models: Tuple of model classes to consider for each component.
+        principal_md_index: Index of the main deflector model to consider when
+            determining the position of the source and substructure.
+        principal_source_index: Index of the source model to consider when
+            determining the position of the substructure.
+        kwargs_simulation: Keyword arguments for the draws of the substructure.
+        kwargs_detector: Keyword arguments defining the detector configuration.
+        kwargs_psf: Keyword arguments defining the point spread function. The
+            psf is applied in the supersampled space, so the size of pixels
+            should be defined with respect to the supersampled space.
+        truth_parameters: List of the lensing objects and corresponding
+            parameters to extract.
+
+    Returns:
+        Image and corresponding truth values.
+
+    Notes:
+        To jit compile, every parameter after rng must be fixed.
+    """
+
+    num_z_bins = kwargs_simulation['num_z_bins']
+    los_pad_length = kwargs_simulation['los_pad_length']
+    subhalos_pad_length = kwargs_simulation['subhalos_pad_length']
+    sampling_pad_length = kwargs_simulation['sampling_pad_length']
+
+    rng_md, rng_source, rng_ll, rng_los, rng_sub, rng = jax.random.split(rng, 6)
+    main_deflector_params = extract_multiple_models(
+        lensing_config['main_deflector_params'], rng_md,
+        len(all_models['all_main_deflector_models'])
+    )
+    source_params = extract_multiple_models(
+        lensing_config['source_params'], rng_source,
+        len(all_models['all_source_models'])
+    )
+    lens_light_params = extract_multiple_models(
+        lensing_config['lens_light_params'], rng_ll,
+        len(all_models['all_source_models'])
+    )
+    los_params = draw_sample(lensing_config['los_params'], rng_los)
+    subhalo_params = draw_sample(lensing_config['subhalo_params'], rng_sub)
+
+    # Extract the principle model for redshifts and substructure draws.
+    main_deflector_params_sub = jax.tree_util.tree_map(
+        lambda x: x[principal_md_index], main_deflector_params
+    )
+    source_params_sub = jax.tree_util.tree_map(
+        lambda x: x[principal_source_index], source_params
+    )
+    lens_light_params_sub = jax.tree_util.tree_map(
+        lambda x: x[principal_source_index], lens_light_params
+    )
+
+    # Repackage the parameters.
+    all_params = {
+        'source_params': source_params_sub,
+        'lens_light_params': lens_light_params_sub,
+        'los_params': los_params, 'subhalo_params': subhalo_params,
+        'main_deflector_params': main_deflector_params_sub
+    }
+
+    rng_los, rng_sub = jax.random.split(rng)
+    los_before_tuple, los_after_tuple = los.draw_los(
+        main_deflector_params_sub, source_params_sub, los_params,
+        cosmology_params, rng_los, num_z_bins, los_pad_length)
+    subhalos_z, subhalos_kwargs = subhalos.draw_subhalos(
+        main_deflector_params_sub, source_params_sub, subhalo_params,
+        cosmology_params, rng_sub, subhalos_pad_length, sampling_pad_length)
+
+    kwargs_lens_all = {
+        'z_array_los_before': los_before_tuple[0],
+        'kwargs_los_before': los_before_tuple[1],
+        'z_array_los_after': los_after_tuple[0],
+        'kwargs_los_after': los_after_tuple[1],
+        'kwargs_main_deflector': main_deflector_params,
+        'z_array_main_deflector': main_deflector_params['z_lens'],
+        'z_array_subhalos': subhalos_z, 'kwargs_subhalos': subhalos_kwargs}
+    z_source = source_params_sub['z_source']
+
+    image_supersampled = image_simulation.generate_image(
+        grid_x, grid_y, kwargs_lens_all, source_params,
+        lens_light_params, kwargs_psf, cosmology_params, z_source,
+        kwargs_detector, all_models)
+    image = utils.downsample(image_supersampled,
+                             kwargs_detector['supersampling_factor'])
+    # Normalize and the image to have standard deviation 1.
+    image /= jnp.std(image)
+
+    # Extract the truth values and normalize them.
+    truth = extract_truth_values(all_params, lensing_config, truth_parameters)
+
+    return image, truth
diff --git a/jaxstronomy/input_pipeline_paltas.py b/jaxstronomy/input_pipeline_paltas.py
new file mode 100644
index 0000000..3e2f2f4
--- /dev/null
+++ b/jaxstronomy/input_pipeline_paltas.py
@@ -0,0 +1,560 @@
+# -*- coding: utf-8 -*-
+"""
+Manipulate strong lensing images to be used with tensorflow models.
+
+This module contains the functions that allow the numpy file and csv outputs
+of generate to be transformed into a format that can be consumed by
+tensorflow.
+"""
+
+import numpy as np
+import tensorflow as tf
+tf.config.experimental.set_visible_devices([], 'GPU')
+import pandas as pd
+import glob, os
+from tqdm import tqdm
+from lenstronomy.SimulationAPI.observation_api import SingleBand
+import warnings
+from scipy.ndimage import rotate
+
+# Global filters on the python warnings. Using this since filter
+# behaviour is a bit weird.
+DEFAULTVALUEWARNING = True
+
+
+def normalize_outputs(metadata,learning_params,input_norm_path,
+    log_learning_params=None):
+    """Normalize the outputs of the network
+
+    Args:
+        metadata(pd.DataFrame): A pandas object containing the metadata
+        learning_params ([str,...]): A list of strings containing the
+            parameters that the network is expected to learn.
+        input_norm_path (str): The path to a csv that contains the
+            normalization to be applied to the output parameters. If the
+            file already exists it will be read, if no file exists it will
+            be written.
+        log_learning_params ([str,...]): A list of strings containing the
+            parameters that the network is expected to learn the log of. Can
+            be None.
+
+    Returns:
+        (pd.DataFrame): A pandas dataframe with the the mean and standard
+        deviation for each parameter.
+    """
+    # If normalization file is provided use / write it
+    if os.path.isfile(input_norm_path):
+        print('Using input normalization found at %s'%(input_norm_path))
+        norm_dict = pd.read_csv(input_norm_path,index_col='parameter')
+        if not all(elem in norm_dict.index for elem in learning_params):
+            raise ValueError('Not all of the learning parameters are ' +
+                'present in the input normalization dictionary.')
+        if log_learning_params is not None:
+            if not all(elem in norm_dict.index for elem in log_learning_params):
+                raise ValueError('Not all of the log learning parameters are ' +
+                    'present in the input normalization dictionary.')
+    else:
+        print('Writing input normalization to %s'%(input_norm_path))
+        norm_dict = pd.DataFrame(columns=['parameter','mean','std'])
+        norm_dict['parameter'] = learning_params
+
+        # Calculate the normalization for each parameter
+        data = metadata[learning_params].to_numpy()
+        norm_dict['mean'] = np.mean(data,axis=0)
+        norm_dict['std'] = np.std(data,axis=0)
+
+        # Append a dictionary of the log parameters if provided.
+        if log_learning_params is not None:
+            log_norm_dict = pd.DataFrame(columns=['parameter','mean','std'])
+            log_norm_dict['parameter'] = log_learning_params
+            # Calculate the normalization for each parameter
+            log_data = metadata[log_learning_params].to_numpy()
+            log_norm_dict['mean'] = np.mean(np.log(log_data),axis=0)
+            log_norm_dict['std'] = np.std(np.log(log_data),axis=0)
+            norm_dict = norm_dict.append(log_norm_dict)
+
+        # Set parameter to the index
+        norm_dict = norm_dict.set_index('parameter')
+        norm_dict.to_csv(input_norm_path)
+
+    return norm_dict
+
+
+def unnormalize_outputs(input_norm_path,learning_params,mean,standard_dev=None,
+	cov_mat=None):
+	"""Given NN outputs, undo the normalization step and return the parameters
+	in the original space
+
+	Args:
+		input_norm_path (str): The path to a csv that contains the
+			normalization to be undone.
+		learning_params ([str,...]): A list of strings containing the
+			parameters that the network is expected to learn. Length n_params.
+		mean (np.array): A numpy array with dimensions (batch_size,n_params)
+			containing the mean estimate for each parameter
+		standard_dev (np.array): A numpy array with dimensions (batch_size,
+			n_params) containing the standard deviation estimate for each
+			parameter.
+		cov_mat (np.array): A numpy array with dimensions (batch_size,n_params,
+			n_params) containing the covariance matrix estiamtes for each
+			image.
+
+	Notes:
+		All values will be modified in place.
+	"""
+	# Read our normalization dictionary
+	norm_dict = pd.read_csv(input_norm_path,index_col='parameter')
+	# Iterate over our parameters
+	for lpi, param in enumerate(learning_params):
+		param_mean = norm_dict['mean'][param]
+		param_std = norm_dict['std'][param]
+
+		# We always want to correct the mean
+		mean[:,lpi] *= param_std
+		mean[:,lpi] += param_mean
+
+		# If provided we want to correct the standard deviation
+		if standard_dev is not None:
+			standard_dev[:,lpi] *= param_std
+
+		# If provided we want to correct the covariance matrix
+		if cov_mat is not None:
+			cov_mat[:,lpi,:] *= param_std
+			cov_mat[:,:,lpi] *= param_std
+
+
+def kwargs_detector_to_tf_noise(image,kwargs_detector):
+	"""Add noise to the tf tensor provided in agreement with kwargs_detector
+
+	Args:
+		image (tf.Tensor): A tensorflow tensor containing the image
+		kwargs_detector (dict): A dictionary containing the detector kwargs
+			used to generate the noise on the fly.
+
+	Returns:
+		(tf.Tensor): A tensorflow tensor containing the noise.
+	"""
+	# Use lenstronomy to do our noise calculations (lenstronomy is well
+	# tested and is very flexible in its noise inputs).
+	single_band = SingleBand(**kwargs_detector)
+
+	# Add the background noise.
+	noise = tf.random.normal(tf.shape(image)) * single_band.background_noise
+
+	# Add the poisson noise. We have to redo the functions to be tf friendly
+	variance = tf.maximum(single_band.flux_iid(image),tf.zeros_like(image))
+	variance = tf.sqrt(variance) / single_band.exposure_time
+	if single_band._data_count_unit == 'ADU':
+		variance /= single_band.ccd_gain
+	noise += tf.random.normal(tf.shape(image)) * variance
+
+	return noise
+
+
+def generate_tf_record(npy_folder,learning_params,metadata_path,
+	tf_record_path):
+	"""Generate a TFRecord file from a directory of numpy files.
+
+	Args:
+		root_path (str): The path to the folder containing the numpy files.
+		learning_params ([str,...]): A list of strings containing the
+			parameters that the network is expected to learn.
+		metadata_path (str):  The path to the csv file containing the
+			image metadata.
+		tf_record_path (str): The path to which the TFRecord will be saved
+	"""
+	# Pull the list of numpy filepaths from the directory
+	npy_file_list = glob.glob(os.path.join(npy_folder,'image_*.npy'))
+	npy_file_list = list(sorted(npy_file_list))
+	# Open label csv
+	metadata = pd.read_csv(metadata_path, index_col=None)
+
+	# Warn the user a default value of 0 will be used for parameters not
+	# present in the metadata file.
+	global DEFAULTVALUEWARNING
+	for param in learning_params:
+		if param not in metadata and DEFAULTVALUEWARNING:
+			warnings.warn('One or more parameters in learning_params is not '+
+				' present in the metadata. A default value of 0 will be used.',
+				category=RuntimeWarning)
+			DEFAULTVALUEWARNING = False
+
+	# Initialize the writer object and write the lens data
+	with tf.io.TFRecordWriter(tf_record_path) as writer:
+		for npy_file in tqdm(npy_file_list):
+			# Pull the index from the filename
+			index = int(npy_file[-11:-4])
+			image_shape = np.load(npy_file).shape
+			# The image must be converted to a tf string feature
+			image_feature = tf.train.Feature(bytes_list=tf.train.BytesList(
+				value=[np.load(npy_file).astype(np.float32).tostring()]))
+			# Initialize a feature dictionary with the image, the height,
+			# and the width
+			feature = {
+				'image': image_feature,
+				'height': tf.train.Feature(
+					int64_list=tf.train.Int64List(value=[image_shape[0]])),
+				'width': tf.train.Feature(
+					int64_list=tf.train.Int64List(value=[image_shape[1]])),
+				'index': tf.train.Feature(
+					int64_list=tf.train.Int64List(value=[index]))
+			}
+			# Add all of the lens parameters to the feature dictionary
+			for param in learning_params:
+				if param in metadata:
+					value = metadata[param][index]
+				else:
+					value = 0
+				# Write the feature
+				feature[param] = tf.train.Feature(
+					float_list=tf.train.FloatList(value=[value]))
+			# Create the tf example object
+			example = tf.train.Example(features=tf.train.Features(
+				feature=feature))
+			# Write out the example to the TFRecord file
+			writer.write(example.SerializeToString())
+
+
+def generate_tf_dataset(tf_record_path,learning_params,batch_size,
+	n_epochs,norm_images=False,input_norm_path=None,kwargs_detector=None,
+	log_learning_params=None,shuffle=True):
+	"""Generate a TFDataset that a model can be trained with.
+
+	Args:
+		tf_record_paths (str, or [str,...]) A string specifying the paths to
+			the TFRecords that will be used in the dataset. Can also be a list
+			of strings for specifying multiple tf_record_paths.
+		learning_params ([str,...]): A list of strings containing the
+			parameters that the network is expected to learn.
+		batch_size (int): The batch size that will be used for training
+		n_epochs (int): The number of training epochs. The dataset object will
+			deal with iterating over the data for repeated epochs.
+		norm_images (bool): If True, images will be normalized to have std 1.
+		input_norm_path (str): The path to a csv that contains the
+			normalization to be applied to the output parameters. If None
+			no normalization will be applied.
+		kwargs_detector (dict): A dictionary containing the detector kwargs
+			used to generate the noise on the fly. If None no additional
+			noise will be added.
+		log_learning_params ([str,...]): A list of strings containing the
+			parameters that the network is expected to learn the log of. Can
+			be None.
+		shuffle (bool): if True (default is True), randomly shuffles dataset.
+
+	Returns:
+		(tf.Dataset): A tf.Dataset object that returns the input image and the
+		output labels.
+
+	Notes:
+		Do not use kwargs_detector if noise was already added during dataset
+		generation.
+	"""
+	# Read the TFRecords
+	raw_dataset = tf.data.TFRecordDataset(tf_record_path)
+
+	# If normalization file is provided use it
+	if input_norm_path is not None:
+		norm_dict = pd.read_csv(input_norm_path,index_col='parameter')
+	else:
+		norm_dict = None
+
+	# Load a noise model if requested.
+	if kwargs_detector is not None:
+		noise_function = kwargs_detector_to_tf_noise
+	else:
+		warnings.warn('No noise will be added')
+		noise_function = None
+
+	# Create the feature decoder that will be used
+	def parse_image_features(example):  # pragma: no cover
+		data_features = {
+			'image': tf.io.FixedLenFeature([],tf.string),
+			'height': tf.io.FixedLenFeature([],tf.int64),
+			'width': tf.io.FixedLenFeature([],tf.int64),
+			'index': tf.io.FixedLenFeature([],tf.int64),
+		}
+		# Set the log learning params to an empy list if no value is provided.
+		if log_learning_params is None:
+			log_learning_params_list = []
+		else:
+			log_learning_params_list = log_learning_params
+
+		for param in learning_params+log_learning_params_list:
+			data_features[param] = tf.io.FixedLenFeature([],tf.float32)
+		parsed_dataset = tf.io.parse_single_example(example,data_features)
+		image = tf.io.decode_raw(parsed_dataset['image'],out_type=float)
+		image = tf.reshape(image,(parsed_dataset['height'],
+			parsed_dataset['width'],1))
+
+		# Add the noise using the baobab noise function (which is a tf graph)
+		if noise_function is not None:
+			image += noise_function(image,kwargs_detector)
+
+		# If the images must be normed divide by the std
+		if norm_images:
+			image = image / tf.math.reduce_std(image)
+
+		# Log the parameter if needed
+		for param in log_learning_params_list:
+			parsed_dataset[param] = tf.math.log(parsed_dataset[param])
+
+		# Normalize if requested
+		if norm_dict is not None:
+			for param in learning_params+log_learning_params_list:
+				parsed_dataset[param] -= norm_dict['mean'][param]
+				parsed_dataset[param] /= norm_dict['std'][param]
+
+		lens_param_values = tf.stack([parsed_dataset[param] for param in
+			learning_params+log_learning_params_list])
+		return image,lens_param_values
+
+	# Select the buffer size to be slightly larger than the batch
+	buffer_size = int(batch_size*1.2)
+
+	# Set the feature decoder as the mapping function. Drop the remainder
+	# in the case that batch_size does not divide the number of training
+	# points exactly
+	dataset = raw_dataset.map(parse_image_features,
+		num_parallel_calls=tf.data.experimental.AUTOTUNE).repeat(
+		n_epochs)
+	if shuffle:
+		dataset = dataset.shuffle(buffer_size=buffer_size)
+	dataset = dataset.batch(batch_size).prefetch(tf.data.experimental.AUTOTUNE)
+	return dataset
+
+
+def rotate_params_batch(learning_params,output,rot_angle):
+	"""Rotate a batch of lensing parameters according to a specified rotation
+	angle.
+
+	Args:
+		learning_params ([str,...]): A list of strings containing the
+			parameters that the network is expected to learn.
+		output (np.array): A numpy array of dimension (batch_size,n_outputs)
+			containing the true parameter values for each image in the batch.
+			Note that n_outputs should be the same as len(learning_params).
+		rot_angle (float): The angle to rotate the image by in radians.
+
+	Notes:
+		output is modified in place.
+	"""
+
+	# Alter the parameters. Hardcoded for now.
+	def rotate_param(x,y,theta):
+		return (x*np.cos(theta)-y*np.sin(theta),
+			x*np.sin(theta)+y*np.cos(theta))
+
+	if 'main_deflector_parameters_center_x' in learning_params:
+		xi = learning_params.index('main_deflector_parameters_center_x')
+		yi = learning_params.index('main_deflector_parameters_center_y')
+		x,y = rotate_param(output[:,xi],output[:,yi],rot_angle)
+		output[:,xi] = x
+		output[:,yi] = y
+	if 'main_deflector_parameters_e1' in learning_params:
+		xi = learning_params.index('main_deflector_parameters_e1')
+		yi = learning_params.index('main_deflector_parameters_e2')
+		x,y = rotate_param(output[:,xi],output[:,yi],2*rot_angle)
+		output[:,xi] = x
+		output[:,yi] = y
+	if 'main_deflector_parameters_gamma1' in learning_params:
+		xi = learning_params.index('main_deflector_parameters_gamma1')
+		yi = learning_params.index('main_deflector_parameters_gamma2')
+		x,y = rotate_param(output[:,xi],output[:,yi],2*rot_angle)
+		output[:,xi] = x
+		output[:,yi] = y
+
+
+def rotate_covariance_batch(learning_params,coavriance_batch,rot_angle):
+	"""Rotate a batch of covariance matrices according to a specified rotation
+	angle.
+
+	Args:
+		learning_params ([str,...]): A list of strings containing the
+			parameters that the network is expected to learn.
+		coavriance_batch (np.array): A numpy array of dimension
+			(batch_size,n_outputs,n_outputs) containing the covariance values
+			for each image in the batch.
+		rot_angle (float): The angle to rotate the image by in radians.
+
+	Notes:
+		coavriance_batch is modified in place.
+	"""
+	# Calculate rotation matrix. Hardcoded for now.
+	def rotation_matrix(theta):
+		return np.array([[np.cos(theta),-np.sin(theta)],
+			[np.sin(theta),np.cos(theta)]])
+
+	# For each possible pair of parameters, conduct the rotation.
+	if 'main_deflector_parameters_center_x' in learning_params:
+		xi = learning_params.index('main_deflector_parameters_center_x')
+		yi = learning_params.index('main_deflector_parameters_center_y')
+		# Ensure the parameters are next to each other. Otherwise this
+		# indexing will fail.
+		assert xi == yi - 1
+		yi+=1
+		rot_mat = rotation_matrix(rot_angle)
+		coavriance_batch[:,xi:yi,xi:yi] = np.dot(rot_mat,
+			np.dot(coavriance_batch[:,xi:yi,xi:yi],rot_mat.T).T).T
+
+	if 'main_deflector_parameters_e1' in learning_params:
+		xi = learning_params.index('main_deflector_parameters_e1')
+		yi = learning_params.index('main_deflector_parameters_e2')
+		# Ensure the parameters are next to each other.
+		assert xi == yi - 1
+		yi+=1
+		rot_mat = rotation_matrix(2*rot_angle)
+		coavriance_batch[:,xi:yi,xi:yi] = np.dot(rot_mat,
+			np.dot(coavriance_batch[:,xi:yi,xi:yi],rot_mat.T).T).T
+
+	if 'main_deflector_parameters_gamma1' in learning_params:
+		xi = learning_params.index('main_deflector_parameters_gamma1')
+		yi = learning_params.index('main_deflector_parameters_gamma2')
+		assert xi == yi - 1
+		yi+=1
+		rot_mat = rotation_matrix(2*rot_angle)
+		coavriance_batch[:,xi:yi,xi:yi] = np.dot(rot_mat,
+			np.dot(coavriance_batch[:,xi:yi,xi:yi],rot_mat.T).T).T
+
+
+def rotate_image_batch(image_batch,learning_params,output,rot_angle):
+	"""Rotate a batch of strong lensing images and the corresponding lensing
+	parameters
+
+	Args:
+		image_batch (np.array): A numpy image array of shape (batch_size,
+			height,width,n_channels) that will be rotated.
+		learning_params ([str,...]): A list of strings containing the
+			parameters that the network is expected to learn.
+		output (np.array): A numpy array of dimension (batch_size,n_outputs)
+			containing the true parameter values for each image in the batch.
+			Note that n_outputs should be the same as len(learning_params).
+		rot_angle (float): The angle to rotate the image by in radians.
+
+	Returns:
+		(np.array): A numpy array containing the rotated images.
+
+	Notes:
+		output is changed in place.
+	"""
+
+	# Rotate the image
+	image_batch = rotate(image_batch,-rot_angle*180/np.pi,reshape=False,
+		axes=(2,1))
+
+	rotate_params_batch(learning_params,output,rot_angle)
+
+	return image_batch
+
+
+def generate_rotations_dataset(tf_record_path,learning_params,batch_size,
+	n_epochs,norm_images=False,input_norm_path=None,kwargs_detector=None,
+	log_learning_params=None,shuffle=True):
+	"""Returns a generator that builds off of a TFDataset by adding random
+	rotations to the images and parameters.
+
+	Args:
+		tf_record_paths (str, or [str,...]) A string specifying the paths to
+			the TFRecords that will be used in the dataset. Can also be a list
+			of strings for specifying multiple tf_record_paths.
+		learning_params ([str,...]): A list of strings containing the
+			parameters that the network is expected to learn.
+		batch_size (int): The batch size that will be used for training
+		n_epochs (int): The number of training epochs. The dataset object will
+			deal with iterating over the data for repeated epochs.
+		norm_images (bool): If True, images will be normalized to have std 1.
+		input_norm_path (str): The path to a csv that contains the
+			normalization to be applied to the output parameters. If None
+			no normalization will be applied.
+		kwargs_detector (dict): A dictionary containing the detector kwargs
+			used to generate the noise on the fly. If None no additional
+			noise will be added.
+		log_learning_params ([str,...]): A list of strings containing the
+			parameters that the network is expected to learn the log of. Can
+			be None.
+		shuffle (bool): if True (default is True), randomly shuffles dataset.
+
+	Returns:
+		(generator): A generator that returns a tuple with the rotated image
+		and the parameter values.
+	"""
+	# Create our base tf dataset without normalization
+	base_dataset = generate_tf_dataset(tf_record_path,learning_params,
+		batch_size,n_epochs,norm_images=norm_images,
+		kwargs_detector=kwargs_detector,log_learning_params=log_learning_params,
+		shuffle=shuffle)
+
+	# If normalization file is provided use it
+	if input_norm_path is not None:
+		norm_dict = pd.read_csv(input_norm_path,index_col='parameter')
+	else:
+		norm_dict = None
+
+	# If there are no log_learning_params, set it to an empty list
+	if log_learning_params is None:
+		log_learning_params = []
+
+	def rotation_generator(dataset):
+		# Iterate through the images and parameters in the dataset
+		for image_batch, lens_param_batch in dataset:
+			image_batch = image_batch.numpy()
+			lens_param_batch = lens_param_batch.numpy()
+			# Pick a rotation angle
+			rot_angle = np.random.uniform()*2*np.pi
+			# Conduct the rotation
+			image_batch = rotate_image_batch(image_batch,
+				learning_params+log_learning_params,lens_param_batch,rot_angle)
+			if norm_dict is not None:
+				for lpi, param in enumerate(learning_params+log_learning_params):
+					lens_param_batch[:,lpi] -= norm_dict['mean'][param]
+					lens_param_batch[:,lpi] /= norm_dict['std'][param]
+			# Yield the rotated image and parameters
+			yield image_batch, lens_param_batch
+
+	# Return a rotation generator on our base dataset.
+	return rotation_generator(base_dataset)
+
+
+def generate_params_as_input_dataset(base_dataset,params_as_inputs,
+	all_params):
+	"""Generate a dataset where some of the lens parameters are treated as
+	inputs to the model.
+
+	Args:
+		base_dataset (generator): A generator that yields a tuple containing
+			the images and the parameters values.
+		params_as_inputs ([str,...]): A list of strings containing the
+			parameters that will be turned into inputs for the model.
+		all_params ([str,...]): A list of strings containing the parameters
+			that will be outputted by the generator.
+
+	Returns:
+		(generator): A generator that returns a tuple with the inputs (which
+		are the image and the scalar input values), and the outputs parameter
+		values.
+	"""
+
+	# Set a boolean mask for the parameters we want to use as inputs and the
+	# parameters we want to use as outputs.
+	param_is_input = np.array([param in params_as_inputs for param in all_params])
+	param_is_output = ~param_is_input
+
+	def param_extractor(dataset):
+		# Iterate through the images and parameters in the dataset and
+		# reorganize the parameters as needed.
+		for image_batch,all_param_batch in dataset:
+			# If the images or parameters are tensors, do the
+			# appropriate conversion.
+			if isinstance(image_batch,tf.Tensor):
+				image_batch = image_batch.numpy()
+			if isinstance(all_param_batch,tf.Tensor):
+				all_param_batch = all_param_batch.numpy()
+
+			scalar_inputs = all_param_batch[:,param_is_input]
+			outputs_batch = all_param_batch[:,param_is_output]
+
+			# Yield the inputs, which are the image and the scalar inputs, along
+			# with the parameters we want to learn.
+			yield [image_batch, scalar_inputs], outputs_batch
+
+	# Return the generator on the base dataset
+	return param_extractor(base_dataset)
diff --git a/jaxstronomy/lens_models.py b/jaxstronomy/lens_models.py
index fa7aba8..ed12e9b 100644
--- a/jaxstronomy/lens_models.py
+++ b/jaxstronomy/lens_models.py
@@ -40,364 +40,370 @@
 
 
 class EPL():
-  """Elliptical Power Law mass profile.
-
-  Elliptical Power Law mass profile functions, with calculation following those
-  described in Tessore & Metcalf (2015) and implementation closely
-  following the EPL_numba class in Lenstronomy.
-  """
-
-  parameters = (
-      'theta_e', 'slope', 'axis_ratio', 'angle', 'center_x', 'center_y'
-  )
-
-  @staticmethod
-  def derivatives(x, y, theta_e, slope,
-                  axis_ratio, angle, center_x,
-                  center_y):
-    """Calculate the derivative of the potential for the EPL mass profile.
-
-    Args:
-      x: X-coordinates at which to evaluate the derivative.
-      y: Y-coordinates at which to evaluate the derivative.
-      theta_e: Einstein radius of the EPL profile.
-      slope: Power-law slope of the EPL profile.
-      axis_ratio: Axis ratio of the major and minor axis of ellipticity.
-      angle: Clockwise angle of orientation of major axis.
-      center_x: X-coordinate center of the EPL profile.
-      center_y: Y-coordinate cetner of the EPL profile.
-
-    Returns:
-      X- and y-component of the derivatives.
-    """
-    # Transform parameters to Tessore & Metcalf (2015) definition
-    z = jnp.exp(-1j * angle) * ((x - center_x) + (y - center_y) * 1j)
-    scale_length = theta_e * jnp.sqrt(axis_ratio)
-    complex_derivative = EPL._complex_derivative(
-        z.real, z.imag, scale_length, axis_ratio, slope) * jnp.exp(1j * angle)
-    return complex_derivative.real, complex_derivative.imag
-
-  @staticmethod
-  def _complex_derivative(x, y, scale_length,
-                          axis_ratio, slope):
-    """Calculate the complex derivative for the EPL mass profile.
-
-    Args:
-      x: X-coordinates at which to evaluate the derivative.
-      y: Y-coordinates at which to evaluate the derivative.
-      scale_length: Scale length of the EPL mass profile (related to Einstein
-        Radius by axis ratio).
-      axis_ratio: Axis ratio of the major and minor axis of ellipticity.
-      slope: Power-law slope of the EPL profile.
-
-    Returns:
-      Complex derivative at each (x,y) coordinate pair.
-    """
-    ellip_vector = x * axis_ratio + 1j * y
-    ellip_radius = jnp.abs(ellip_vector)
-    ellip_angle = jnp.angle(ellip_vector)
-    omega = EPL._hypergeometric_series(ellip_angle, slope, axis_ratio)
-    return (2 * scale_length) / (1 + axis_ratio) * jnp.nan_to_num(
-        (scale_length / ellip_radius)**(slope - 2), copy=False) * omega
-
-  @staticmethod
-  def _hypergeometric_series(ellip_angle, slope,
-                             axis_ratio):
-    """Calculate the hypergeometric series required for the derivative.
-
-    Args:
-      ellip_angle: The elliptical angles at which to evaluate the series.
-      slope: Power-law slope of the EPL profile.
-      axis_ratio: Axis ratio the major and minor axis of ellipticity.
-
-    Returns:
-      The hypergeometric series sum for each angle.
-    """
-    flattening = (1 - axis_ratio) / (1 + axis_ratio)
-    omegas = jnp.zeros_like(ellip_angle)
-
-    # Calculate the zeroeth term in the Fourier-type series expansion.
-    four_n = 1 * jnp.exp(1j * ellip_angle)
-    omegas += four_n
-    four_factor = -flattening * jnp.exp(2j * ellip_angle)
-    # Use the ratio between terms to calculate the rest of the series.
-    # Consider making the number of terms an input to the calculation.
-    for n in range(1, 200):
-      four_n *= (2 * n - (3 - slope)) / (2 * n + (3 - slope)) * four_factor
-      omegas += four_n
+    """Elliptical Power Law mass profile.
 
-    return omegas
-
-
-class NFW():
-  """Navarro Frenk White (NFW) mass profile.
-
-  NFW mass profile functions, with implementation closely following the NFW
-  class in Lenstronomy.
-  """
-
-  parameters = ('scale_radius', 'alpha_rs', 'center_x', 'center_y')
-
-  @staticmethod
-  def derivatives(x, y, scale_radius,
-                  alpha_rs, center_x,
-                  center_y):
-    """Return the NFW profile derivatives.
-
-    Args:
-      x: X-coordinates at which to evaluate the derivative
-      y: Y-coordinates at which to evaluate the derivative
-      scale_radius: Scale radius of NFW profile.
-      alpha_rs: Derivative at the scale radius
-      center_x: X-coordinate center of the NFW profile.
-      center_y: Y-coordinate cetner of the NFW profile.
-
-    Returns:
-      X- and y-component of the derivative.
-    """
-    rho_input = NFW._alpha_to_rho(alpha_rs, scale_radius)
-    x_centered = x - center_x
-    y_centered = y - center_y
-    radius = jnp.sqrt(x_centered**2 + y_centered**2)
-    return NFW._nfw_derivatives(radius, scale_radius, rho_input, x_centered,
-                                y_centered)
-
-  @staticmethod
-  def _alpha_to_rho(alpha_rs, scale_radius):
-    """Return the NFW profile normalization.
-
-    Args:
-      alpha_rs: Derivative at the scale radius
-      scale_radius: Scale radius of NFW profile.
-
-    Returns:
-      NFW profile normalization.
-
-    """
-    return alpha_rs / (4. * scale_radius**2 * (1. + jnp.log(0.5)))
-
-  @staticmethod
-  def _nfw_derivatives(
-      radius, scale_radius, rho_input,
-      x_centered,
-      y_centered):
-    """Return the NFW profile derivatives.
-
-    Args:
-      radius: Radii at which to evaluate the derivatives.
-      scale_radius: Scale radius of NFW profile.
-      rho_input: Normalization of the NFW profile.
-      x_centered: X-coordinate offset from NFW center.
-      y_centered: Y-coordinate offset from NFW center.
-
-    Returns:
-      X- and y-component of the derivative.
+    Elliptical Power Law mass profile functions, with calculation following
+    those described in Tessore & Metcalf (2015) and implementation closely
+    following the EPL_numba class in Lenstronomy.
     """
-    reduced_radius = radius / scale_radius
-    nfw_integral = NFW._nfw_integral(reduced_radius)
-    derivative_norm = 4 * rho_input * scale_radius * nfw_integral
-    derivative_norm /= reduced_radius**2
 
-    return derivative_norm * x_centered, derivative_norm * y_centered
+    parameters = (
+            'theta_e', 'slope', 'axis_ratio', 'angle', 'center_x', 'center_y'
+    )
+
+    @staticmethod
+    def derivatives(x, y, theta_e, slope,
+                                    axis_ratio, angle, center_x,
+                                    center_y):
+        """Calculate the derivative of the potential for the EPL mass profile.
+
+        Args:
+            x: X-coordinates at which to evaluate the derivative.
+            y: Y-coordinates at which to evaluate the derivative.
+            theta_e: Einstein radius of the EPL profile.
+            slope: Power-law slope of the EPL profile.
+            axis_ratio: Axis ratio of the major and minor axis of ellipticity.
+            angle: Clockwise angle of orientation of major axis.
+            center_x: X-coordinate center of the EPL profile.
+            center_y: Y-coordinate cetner of the EPL profile.
+
+        Returns:
+            X- and y-component of the derivatives.
+        """
+        # Transform parameters to Tessore & Metcalf (2015) definition
+        z = jnp.exp(-1j * angle) * ((x - center_x) + (y - center_y) * 1j)
+        scale_length = theta_e * jnp.sqrt(axis_ratio)
+        complex_derivative = EPL._complex_derivative(
+                z.real, z.imag, scale_length, axis_ratio, slope)
+        complex_derivative *= jnp.exp(1j * angle)
+        return complex_derivative.real, complex_derivative.imag
+
+    @staticmethod
+    def _complex_derivative(x, y, scale_length,
+                                                    axis_ratio, slope):
+        """Calculate the complex derivative for the EPL mass profile.
+
+        Args:
+            x: X-coordinates at which to evaluate the derivative.
+            y: Y-coordinates at which to evaluate the derivative.
+            scale_length: Scale length of the EPL mass profile (related to
+                Einstein Radius by axis ratio).
+            axis_ratio: Axis ratio of the major and minor axis of ellipticity.
+            slope: Power-law slope of the EPL profile.
+
+        Returns:
+            Complex derivative at each (x,y) coordinate pair.
+        """
+        ellip_vector = x * axis_ratio + 1j * y
+        ellip_radius = jnp.abs(ellip_vector)
+        ellip_angle = jnp.angle(ellip_vector)
+        omega = EPL._hypergeometric_series(ellip_angle, slope, axis_ratio)
+        return (2 * scale_length) / (1 + axis_ratio) * jnp.nan_to_num(
+                (scale_length / ellip_radius)**(slope - 2), copy=False) * omega
+
+    @staticmethod
+    def _hypergeometric_series(ellip_angle, slope,
+                                                         axis_ratio):
+        """Calculate the hypergeometric series required for the derivative.
+
+        Args:
+            ellip_angle: The elliptical angles at which to evaluate the series.
+            slope: Power-law slope of the EPL profile.
+            axis_ratio: Axis ratio the major and minor axis of ellipticity.
+
+        Returns:
+            The hypergeometric series sum for each angle.
+        """
+        flattening = (1 - axis_ratio) / (1 + axis_ratio)
+        omegas = jnp.zeros_like(ellip_angle)
+
+        # Calculate the zeroeth term in the Fourier-type series expansion.
+        four_n = 1 * jnp.exp(1j * ellip_angle)
+        omegas += four_n
+        four_factor = -flattening * jnp.exp(2j * ellip_angle)
+        # Use the ratio between terms to calculate the rest of the series.
+        # Consider making the number of terms an input to the calculation.
+        for n in range(1, 200):
+            four_n *= (2 * n - (3 - slope)) / (2 * n + (3 - slope))
+            four_n *= four_factor
+            omegas += four_n
+
+        return omegas
 
-  @staticmethod
-  def _nfw_integral(reduced_radius):
-    """Return analytic solution to integral of NFW profile.
 
-    Args:
-      reduced_radius: Upper limits for integrals in reduced units.
+class NFW():
+    """Navarro Frenk White (NFW) mass profile.
 
-    Returns:
-      Solution to integral for each provided upper limit.
+    NFW mass profile functions, with implementation closely following the NFW
+    class in Lenstronomy.
     """
-    solution = jnp.zeros_like(reduced_radius)
-
-    # There are three regimes for the analytic solution, which we will calculate
-    # with masks to avoid indexing.
-    is_below_one = reduced_radius < 1
-    is_one = reduced_radius == 1
-    is_above_one = reduced_radius > 1
-    solution += jnp.nan_to_num(
-        jnp.log(reduced_radius / 2.) + 1 / jnp.sqrt(1 - reduced_radius**2) *
-        jnp.arccosh(1. / reduced_radius)) * is_below_one
-    solution += (1 + jnp.log(1. / 2.)) * is_one
-    solution += jnp.nan_to_num(
-        jnp.log(reduced_radius / 2) + 1 / jnp.sqrt(reduced_radius**2 - 1) *
-        jnp.arccos(1. / reduced_radius)) * is_above_one
 
-    return solution
+    parameters = ('scale_radius', 'alpha_rs', 'center_x', 'center_y')
+
+    @staticmethod
+    def derivatives(x, y, scale_radius,
+                                    alpha_rs, center_x,
+                                    center_y):
+        """Return the NFW profile derivatives.
+
+        Args:
+            x: X-coordinates at which to evaluate the derivative
+            y: Y-coordinates at which to evaluate the derivative
+            scale_radius: Scale radius of NFW profile.
+            alpha_rs: Derivative at the scale radius
+            center_x: X-coordinate center of the NFW profile.
+            center_y: Y-coordinate cetner of the NFW profile.
+
+        Returns:
+            X- and y-component of the derivative.
+        """
+        rho_input = NFW._alpha_to_rho(alpha_rs, scale_radius)
+        x_centered = x - center_x
+        y_centered = y - center_y
+        radius = jnp.sqrt(x_centered**2 + y_centered**2)
+        return NFW._nfw_derivatives(radius, scale_radius, rho_input, x_centered,
+                                                                y_centered)
+
+    @staticmethod
+    def _alpha_to_rho(alpha_rs, scale_radius):
+        """Return the NFW profile normalization.
+
+        Args:
+            alpha_rs: Derivative at the scale radius
+            scale_radius: Scale radius of NFW profile.
+
+        Returns:
+            NFW profile normalization.
+
+        """
+        return alpha_rs / (4. * scale_radius**2 * (1. + jnp.log(0.5)))
+
+    @staticmethod
+    def _nfw_derivatives(
+            radius, scale_radius, rho_input,
+            x_centered,
+            y_centered):
+        """Return the NFW profile derivatives.
+
+        Args:
+            radius: Radii at which to evaluate the derivatives.
+            scale_radius: Scale radius of NFW profile.
+            rho_input: Normalization of the NFW profile.
+            x_centered: X-coordinate offset from NFW center.
+            y_centered: Y-coordinate offset from NFW center.
+
+        Returns:
+            X- and y-component of the derivative.
+        """
+        reduced_radius = radius / scale_radius
+        nfw_integral = NFW._nfw_integral(reduced_radius)
+        derivative_norm = 4 * rho_input * scale_radius * nfw_integral
+        derivative_norm /= reduced_radius**2
+
+        return derivative_norm * x_centered, derivative_norm * y_centered
+
+    @staticmethod
+    def _nfw_integral(reduced_radius):
+        """Return analytic solution to integral of NFW profile.
+
+        Args:
+            reduced_radius: Upper limits for integrals in reduced units.
+
+        Returns:
+            Solution to integral for each provided upper limit.
+        """
+        solution = jnp.zeros_like(reduced_radius)
+
+        # There are three regimes for the analytic solution, which we will
+        # calculate with masks to avoid indexing.
+        is_below_one = reduced_radius < 1
+        is_one = reduced_radius == 1
+        is_above_one = reduced_radius > 1
+        solution += jnp.nan_to_num(
+                jnp.log(reduced_radius / 2.) + 1 /
+                jnp.sqrt(1 - reduced_radius**2) *
+                jnp.arccosh(1. / reduced_radius)) * is_below_one
+        solution += (1 + jnp.log(1. / 2.)) * is_one
+        solution += jnp.nan_to_num(
+                jnp.log(reduced_radius / 2) + 1 /
+                jnp.sqrt(reduced_radius**2 - 1) *
+                jnp.arccos(1. / reduced_radius)) * is_above_one
+
+        return solution
 
 
 class Shear():
-  """Shear mass profile.
-
-  Shear mass profile functions, with implementation closely following the
-  ShearGammaPsi class in Lenstronomy.
-  """
-
-  parameters = ('gamma_ext', 'angle', 'center_x', 'center_y')
-
-  @staticmethod
-  def derivatives(x, y, gamma_ext,
-                  angle, center_x,
-                  center_y):
-    """Return the shear profile derivatives.
-
-    Args:
-      x: X-coordinates at which to evaluate the derivative
-      y: Y-coordinates at which to evaluate the derivative
-      gamma_ext: Strength of the shear profile.
-      angle: Clockwise angle of the shear vector with respect to the simulation
-        grid.
-      center_x: X-coordinate where shear is 0.
-      center_y: Y-coordinate where shear is 0.
-
-    Returns:
-      X- and y-component of the derivative.
-    """
-    gamma_one, gamma_two = Shear._polar_to_cartesian(gamma_ext, angle)
-    x_centered = x - center_x
-    y_centered = y - center_y
-    return (gamma_one * x_centered + gamma_two * y_centered,
-            gamma_two * x_centered - gamma_one * y_centered)
-
-  @staticmethod
-  def _polar_to_cartesian(gamma_ext,
-                          angle):
-    """Return the two shear components used for the shear derivative.
-
-    Args:
-      gamma_ext: Strength of the shear profile.
-      angle: Clockwise angle of the shear vector with respect to the simulation
-        grid.
-
-    Returns:
-      Two shear components.
-    """
-    return gamma_ext * jnp.cos(2 * angle), gamma_ext * jnp.sin(2 * angle)
-
+    """Shear mass profile.
 
-class TNFW(NFW):
-  """Truncated Navarro Frenk White (TNFW) mass profile.
-
-  TNFW mass profile functions, with implementation closely following the TNFW
-  class in Lenstronomy.
-  """
-
-  parameters = (
-      'scale_radius', 'alpha_rs', 'trunc_radius', 'center_x', 'center_y'
-  )
-
-  @staticmethod
-  def derivatives(x, y, scale_radius,
-                  alpha_rs, trunc_radius, center_x,
-                  center_y):
-    """Return the TNFW profile derivatives.
-
-    Args:
-      x: X-coordinates at which to evaluate the derivative
-      y: Y-coordinates at which to evaluate the derivative
-      scale_radius: Scale radius of TNFW profile.
-      alpha_rs: Derivative at the scale radius
-      trunc_radius: Truncation radius for TNFW profile.
-      center_x: X-coordinate center of the TNFW profile.
-      center_y: Y-coordinate cetner of the TNFW profile.
-
-    Returns:
-      X- and y-component of the derivative.
+    Shear mass profile functions, with implementation closely following the
+    ShearGammaPsi class in Lenstronomy.
     """
-    rho_input = TNFW._alpha_to_rho(alpha_rs, scale_radius)
-    x_centered = x - center_x
-    y_centered = y - center_y
-    radius = jnp.sqrt(x_centered**2 + y_centered**2)
-    return TNFW._tnfw_derivatives(radius, scale_radius, rho_input, trunc_radius,
-                                  x_centered, y_centered)
-
-  @staticmethod
-  def _tnfw_derivatives(
-      radius, scale_radius, rho_input,
-      trunc_radius, x_centered,
-      y_centered):
-    """Return the TNFW profile derivatives.
-
-    Args:
-      radius: Radii at which to evaluate the derivatives.
-      scale_radius: Scale radius of TNFW profile.
-      rho_input: Normalization of the TNFW profile.
-      trunc_radius: Truncation radius for TNFW profile.
-      x_centered: X-coordinate offset from TNFW center.
-      y_centered: Y-coordinate offset from TNFW center.
-
-    Returns:
-      X- and y-component of the derivative.
-    """
-    reduced_radius = radius / scale_radius
-    reduced_trunc_radius = trunc_radius / scale_radius
-    tnfw_integral = TNFW._tnfw_integral(reduced_radius, reduced_trunc_radius)
-    derivative_norm = 4. * rho_input * scale_radius * tnfw_integral
-    derivative_norm /= reduced_radius**2
 
-    return derivative_norm * x_centered, derivative_norm * y_centered
+    parameters = ('gamma_ext', 'angle', 'center_x', 'center_y')
+
+    @staticmethod
+    def derivatives(x, y, gamma_ext,
+                                    angle, center_x,
+                                    center_y):
+        """Return the shear profile derivatives.
+
+        Args:
+            x: X-coordinates at which to evaluate the derivative
+            y: Y-coordinates at which to evaluate the derivative
+            gamma_ext: Strength of the shear profile.
+            angle: Clockwise angle of the shear vector with respect to the
+                simulation grid.
+            center_x: X-coordinate where shear is 0.
+            center_y: Y-coordinate where shear is 0.
+
+        Returns:
+            X- and y-component of the derivative.
+        """
+        gamma_one, gamma_two = Shear._polar_to_cartesian(gamma_ext, angle)
+        x_centered = x - center_x
+        y_centered = y - center_y
+        return (gamma_one * x_centered + gamma_two * y_centered,
+                        gamma_two * x_centered - gamma_one * y_centered)
+
+    @staticmethod
+    def _polar_to_cartesian(gamma_ext,
+                                                    angle):
+        """Return the two shear components used for the shear derivative.
+
+        Args:
+            gamma_ext: Strength of the shear profile.
+            angle: Clockwise angle of the shear vector with respect to the
+                simulation grid.
+
+        Returns:
+            Two shear components.
+        """
+        return gamma_ext * jnp.cos(2 * angle), gamma_ext * jnp.sin(2 * angle)
 
-  @staticmethod
-  def _tnfw_integral(reduced_radius,
-                     reduced_trunc_radius):
-    """Return analytic solution to integral of TNFW profile.
 
-    Args:
-      reduced_radius: Upper limits for integrals in reduced units.
-      reduced_trunc_radius: Reduced truncation radius for TNFW profile.
+class TNFW(NFW):
+    """Truncated Navarro Frenk White (TNFW) mass profile.
 
-    Returns:
-      Solution to integral for each provided upper limit.
-    """
-    solution = (reduced_trunc_radius**2 + 1 + 2 *
-                (reduced_radius**2 - 1)) * TNFW._nfw_function(reduced_radius)
-    solution += reduced_trunc_radius * jnp.pi + (
-        reduced_trunc_radius**2 - 1) * jnp.log(reduced_trunc_radius)
-    solution += jnp.sqrt(reduced_trunc_radius**2 + reduced_radius**2) * (
-        -jnp.pi + TNFW._tnfw_log(reduced_radius, reduced_trunc_radius) *
-        (reduced_trunc_radius**2 - 1) * reduced_trunc_radius**-1)
-    solution *= reduced_trunc_radius**2 * (reduced_trunc_radius**2 + 1)**-2
-    return solution
-
-  @staticmethod
-  def _tnfw_log(reduced_radius,
-                reduced_trunc_radius):
-    """Evaluate log expression that appears repeatedly in the TNFW calculations.
-
-    Args:
-      reduced_radius: Upper limits for integrals in reduced units.
-      reduced_trunc_radius: Reduced truncation radius for TNFW profile.
-
-    Returns:
-      Log calculation output.
+    TNFW mass profile functions, with implementation closely following the TNFW
+    class in Lenstronomy.
     """
-    return jnp.log(reduced_radius *
-                   (reduced_trunc_radius +
-                    jnp.sqrt(reduced_radius**2 + reduced_trunc_radius**2))**-1)
 
-  @staticmethod
-  def _nfw_function(reduced_radius):
-    """Evaluate NFW function.
-
-    Args:
-      reduced_radius: Upper limits for integrals in reduced units.
-
-    Returns:
-      NFW function output.
-    """
-    nfw_function = jnp.zeros_like(reduced_radius)
-
-    # There are three regimes for the NFW function, which we will calculate with
-    # masks to avoid indexing.
-    is_below_one = reduced_radius < 1
-    is_one = reduced_radius == 1
-    is_above_one = reduced_radius > 1
-    nfw_function += jnp.nan_to_num((1 - reduced_radius**2)**-.5 * jnp.arctanh(
-        (1 - reduced_radius**2)**.5)) * is_below_one
-    nfw_function += is_one
-    nfw_function += jnp.nan_to_num((reduced_radius**2 - 1)**-.5 * jnp.arctan(
-        (reduced_radius**2 - 1)**.5)) * is_above_one
-
-    return nfw_function
+    parameters = (
+            'scale_radius', 'alpha_rs', 'trunc_radius', 'center_x', 'center_y'
+    )
+
+    @staticmethod
+    def derivatives(x, y, scale_radius, alpha_rs, trunc_radius, center_x,
+                    center_y):
+        """Return the TNFW profile derivatives.
+
+        Args:
+            x: X-coordinates at which to evaluate the derivative
+            y: Y-coordinates at which to evaluate the derivative
+            scale_radius: Scale radius of TNFW profile.
+            alpha_rs: Derivative at the scale radius
+            trunc_radius: Truncation radius for TNFW profile.
+            center_x: X-coordinate center of the TNFW profile.
+            center_y: Y-coordinate cetner of the TNFW profile.
+
+        Returns:
+            X- and y-component of the derivative.
+        """
+        rho_input = TNFW._alpha_to_rho(alpha_rs, scale_radius)
+        x_centered = x - center_x
+        y_centered = y - center_y
+        radius = jnp.sqrt(x_centered**2 + y_centered**2)
+        return TNFW._tnfw_derivatives(radius, scale_radius, rho_input,
+                                      trunc_radius, x_centered, y_centered)
+
+    @staticmethod
+    def _tnfw_derivatives(
+            radius, scale_radius, rho_input,
+            trunc_radius, x_centered,
+            y_centered):
+        """Return the TNFW profile derivatives.
+
+        Args:
+            radius: Radii at which to evaluate the derivatives.
+            scale_radius: Scale radius of TNFW profile.
+            rho_input: Normalization of the TNFW profile.
+            trunc_radius: Truncation radius for TNFW profile.
+            x_centered: X-coordinate offset from TNFW center.
+            y_centered: Y-coordinate offset from TNFW center.
+
+        Returns:
+            X- and y-component of the derivative.
+        """
+        reduced_radius = radius / scale_radius
+        reduced_trunc_radius = trunc_radius / scale_radius
+        tnfw_integral = TNFW._tnfw_integral(reduced_radius,
+                                            reduced_trunc_radius)
+        derivative_norm = 4. * rho_input * scale_radius * tnfw_integral
+        derivative_norm /= reduced_radius**2
+
+        return derivative_norm * x_centered, derivative_norm * y_centered
+
+    @staticmethod
+    def _tnfw_integral(reduced_radius, reduced_trunc_radius):
+        """Return analytic solution to integral of TNFW profile.
+
+        Args:
+            reduced_radius: Upper limits for integrals in reduced units.
+            reduced_trunc_radius: Reduced truncation radius for TNFW profile.
+
+        Returns:
+            Solution to integral for each provided upper limit.
+        """
+        solution = (reduced_trunc_radius**2 + 1 + 2 *
+                    (reduced_radius**2 - 1))
+        solution *= TNFW._nfw_function(reduced_radius)
+        solution += reduced_trunc_radius * jnp.pi + (
+                reduced_trunc_radius**2 - 1) * jnp.log(reduced_trunc_radius)
+        solution += jnp.sqrt(reduced_trunc_radius**2 + reduced_radius**2) * (
+                -jnp.pi + TNFW._tnfw_log(reduced_radius, reduced_trunc_radius) *
+                (reduced_trunc_radius**2 - 1) * reduced_trunc_radius**-1)
+        solution *= reduced_trunc_radius**2 * (reduced_trunc_radius**2 + 1)**-2
+        return solution
+
+    @staticmethod
+    def _tnfw_log(reduced_radius,
+                                reduced_trunc_radius):
+        """Evaluate log expression that appears in the TNFW calculations.
+
+        Args:
+            reduced_radius: Upper limits for integrals in reduced units.
+            reduced_trunc_radius: Reduced truncation radius for TNFW profile.
+
+        Returns:
+            Log calculation output.
+        """
+        return jnp.log(reduced_radius * (reduced_trunc_radius +
+                                         jnp.sqrt(reduced_radius**2 +
+                                                  reduced_trunc_radius**2))**-1)
+
+    @staticmethod
+    def _nfw_function(reduced_radius):
+        """Evaluate NFW function.
+
+        Args:
+            reduced_radius: Upper limits for integrals in reduced units.
+
+        Returns:
+            NFW function output.
+        """
+        nfw_function = jnp.zeros_like(reduced_radius)
+
+        # There are three regimes for the NFW function, which we will calculate
+        # with masks to avoid indexing.
+        is_below_one = reduced_radius < 1
+        is_one = reduced_radius == 1
+        is_above_one = reduced_radius > 1
+        nfw_function += jnp.nan_to_num(
+            (1 - reduced_radius**2)**-.5 *
+            jnp.arctanh((1 - reduced_radius**2)**.5)) * is_below_one
+        nfw_function += is_one
+        nfw_function += jnp.nan_to_num(
+            (reduced_radius**2 - 1)**-.5 *
+            jnp.arctan((reduced_radius**2 - 1)**.5)) * is_above_one
+
+        return nfw_function
diff --git a/jaxstronomy/los.py b/jaxstronomy/los.py
new file mode 100644
index 0000000..e0ad402
--- /dev/null
+++ b/jaxstronomy/los.py
@@ -0,0 +1,611 @@
+# coding=utf-8
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     https://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Draw masses and cocentrations for NFW line-of-sight halos.
+
+This module contains the functions needed to turn the parameters of the
+los halo distribution into masses, concentrations, and positions as defined in
+https://arxiv.org/pdf/1909.02573.pdf.
+"""
+
+from typing import Mapping, Sequence, Tuple, Union
+
+import jax
+import jax.numpy as jnp
+
+from jaxstronomy import cosmology_utils
+from jaxstronomy import nfw_functions
+from jaxstronomy import power_law
+
+# Parameters fit from simulations and fixed so that the integral for the nu
+# function over all nu gives 1.
+NU_FUNC_A = 0.32218
+NU_FUNC_Q = 0.3
+NU_FUNC_LIT_A = 0.707
+
+
+def nu_function(nu: jnp.ndarray):
+    """Return nu f(nu) for the Sheth Tormen 2001 model.
+
+    Args:
+        nu: Values for nu.
+
+    Returns:
+        Value of nu times f(nu) for Sheth Tormen 2001 model.
+    """
+    # Fundamental unit for the equation.
+    nu_squared = NU_FUNC_LIT_A * jnp.square(nu)
+
+    return (2 * NU_FUNC_A * (1 + nu_squared **(-NU_FUNC_Q)) *
+        jnp.sqrt(nu_squared / (2 * jnp.pi)) * jnp.exp(-nu_squared / 2))
+
+
+def mass_function_exact(
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    masses: jnp.ndarray, z: float) -> jnp.ndarray:
+    """Return the exact Sheth Tormen 2001 mass function.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        masses: Masses at which to evaluate mass function. In units of M_sun.
+        z: Redshift at which to evaluate mass function.
+
+    Returns:
+        The value of the mass function for each provided mass. In units of the
+        physical number density 1/(M_sun * kpc^3).
+    """
+    # Convert to units of nu used by nu_function.
+    h = cosmology_params['hubble_constant'] / 100
+    delta_collapse = cosmology_utils.collapse_overdensity(cosmology_params, z)
+    lagrangian_r = cosmology_utils.lagrangian_radius(cosmology_params, masses)
+    sigma = jax.vmap(cosmology_utils.sigma_tophat, in_axes=[None, 0, None])(
+        cosmology_params, lagrangian_r, z)
+    nu = delta_collapse / sigma
+
+    # Calculate the mass function from its components
+    nu_function_eval = nu_function(nu)
+    derivative_sigma = cosmology_utils.derivative_log_sigma_log_r(
+        cosmology_params, lagrangian_r, z)
+    # Matter density is returned in units of M_sun * h ^ 2 / Mpc ^ 3
+    rho_matter = cosmology_utils.rho_matter(cosmology_params, z) * h ** 2
+    rho_matter *= 1e-9  # Convert to / kpc ^ 3
+
+    return (-1 / 3 * nu_function_eval * rho_matter / masses ** 2 *
+        derivative_sigma)
+
+def _mass_function_power_law_numerical(
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]], z: float,
+    m_min: float, m_max: float) -> Tuple[float, float]:
+    """Return the best fit power law parameters for the mass function.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z: Redhist at which to estimate mass function.
+        m_min: Lower bound of masses to consider. In units of M_sun.
+        m_max: Upper bound of masses to consider. In units of M_sun.
+
+    Returns:
+        Power law slope and norm that best approximates the exact mass function.
+    """
+    # TODO: For now the number of samples are hardcoded to avoid compilation
+    # issues. May want to change this in the future.
+    num_masses = 100
+    masses = jnp.logspace(jnp.log10(m_min), jnp.log10(m_max), num_masses)
+    log_masses = jnp.log(masses)
+    log_mass_function = jnp.log(
+        mass_function_exact(cosmology_params, masses, z))
+
+    # Use the MLE estimate for the slope assuming Gaussian scatter on the log
+    # quantity.
+    slope_estimate = 1 / num_masses * jnp.sum(log_mass_function) * jnp.sum(
+        log_masses)
+    slope_estimate -= jnp.sum(log_mass_function * log_masses)
+    slope_estimate /= (-jnp.sum(log_masses ** 2) +
+        1 / num_masses * jnp.sum(log_masses) ** 2)
+
+    # Similar MLE estimate for the norm.
+    norm_estimate = jnp.exp(1 / num_masses *
+        jnp.sum(log_mass_function - slope_estimate * log_masses))
+
+    return slope_estimate, norm_estimate
+
+
+def add_los_lookup_tables_to_cosmology_params(
+    los_params: Mapping[str, float],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z_lookup_max: float
+) -> Mapping[str, Union[float, int, jnp.ndarray]]:
+    """Add lookup tables for los mass function calculation to cosmology params.
+
+    Args:
+        los_parms: Parameters of the los distribution.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z_lookup_max: The maximum redshift to consider in the lookup table.
+
+    Returns:
+        Cosmological parameters with los lookup table.
+    """
+    # Jitted function for faster computation.
+    mass_function_power_law_numerical = jax.jit(
+        _mass_function_power_law_numerical)
+    m_min = los_params['m_min']
+    m_max = los_params['m_max']
+    dz = los_params['dz']
+
+    z_range = jnp.arange(0, z_lookup_max + dz / 2, dz / 2)
+
+    lookup_tables = jax.vmap(mass_function_power_law_numerical,
+        in_axes=[None, 0, None, None])(cosmology_params, z_range, m_min, m_max)
+
+    cosmology_params_lookup = {
+        'los_z_lookup_max': z_lookup_max,
+        'mass_function_slope_lookup_table': lookup_tables[0],
+        'mass_function_norm_lookup_table': lookup_tables[1]
+    }
+    cosmology_params_lookup.update(cosmology_params)
+
+    return cosmology_params_lookup
+
+
+def mass_function_power_law(
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z: float) -> Tuple[float, float]:
+    """Return the best fit power law parameters from lookup table.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion. Must inlcude los lookup table.
+        z: Redhist at which to estimate mass function.
+
+    Returns:
+        Power law slope and norm that best approximates the exact mass function.
+    """
+    # Interpolate between the two nearest binds to the query.
+    lookup_unrounded = z / cosmology_params['dz'] * 2
+    frac = lookup_unrounded % 1
+
+    lookup_upper = jax.lax.min(
+        jnp.ceil(lookup_unrounded).astype(int),
+        len(cosmology_params['mass_function_slope_lookup_table']) - 1
+    )
+    lookup_lower = jax.lax.max(jnp.floor(lookup_unrounded).astype(int), 0)
+
+    slope_interpolated = (frac *
+        cosmology_params['mass_function_slope_lookup_table'][lookup_upper] +
+        (1 - frac) *
+            cosmology_params['mass_function_slope_lookup_table'][lookup_lower])
+    norm_interpolated = (frac *
+        cosmology_params['mass_function_norm_lookup_table'][lookup_upper] +
+        (1 - frac) *
+            cosmology_params['mass_function_norm_lookup_table'][lookup_lower])
+
+    return slope_interpolated, norm_interpolated
+
+
+def two_halo_boost(main_deflector_params: Mapping[str, float],
+    los_params: Mapping[str, float],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z: float) -> float:
+    """Calculate the boost from the two halo term caused by the host halo.
+
+    Args:
+        main_deflector_params: Parameters of the main deflector.
+        los_parms: Parameters of the los distribution.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z: Redshift at which to calculate the boost.
+
+    Returns:
+        Boost at the given redshift.
+    """
+    # Extract the los and main deflector parameters we will need.
+    dz = los_params['dz']
+    host_mass = main_deflector_params['mass']
+    z_lens = main_deflector_params['z_lens']
+    r_min = los_params['r_min']
+    r_max = los_params['r_max']
+    h = cosmology_params['hubble_constant'] / 100
+
+    # TODO: For now use a fixed range of redshift values to calculate the boost.
+    z_range = jnp.linspace(z, z+dz, 100)
+
+    # Only consider the two-point statistics within the radial limits.
+    comoving_r = jnp.abs(jax.vmap(cosmology_utils.comoving_distance,
+        in_axes = [None, 0, None])(cosmology_params, z_range, z_lens))
+
+    # The two halo term consists of the correlation function and the halo bias.
+    two_halo = cosmology_utils.correlation_function(cosmology_params,
+        comoving_r, z_lens)
+    two_halo *= cosmology_utils.halo_bias(cosmology_params, host_mass * h,
+        z_lens)
+
+    return 1 + jnp.mean(two_halo)
+
+
+def cone_angle_to_radius(main_deflector_params: Mapping[str, float],
+    source_params: Mapping[str, float], los_params: Mapping[str, float],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z: float) -> float:
+    """Return the radius in kpc for the lightcone.
+
+    Args:
+        main_deflector_params: Parameters of the main deflector.
+        source_params: Parameters for the source.
+        los_parms: Parameters of the los distribution.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z: Redshift at which to calculate the radius.
+
+    Returns:
+        Radius of the lightcone. In units of kpc.
+    """
+    # Get the source, lens, and los parameters we will need.
+    z_lens = main_deflector_params['z_lens']
+    z_source = source_params['z_source']
+    cone_angle = los_params['cone_angle']
+
+    # What fraction of the contraction to include. 1.0 will have the cone
+    # contract to a point at the source redshift. 0.0 will have the cone not
+    # contract at all after passing the main deflector.
+    angle_buffer = los_params['angle_buffer']
+
+    # If we are in front of the main deflector, the opening radius is a simple
+    # conversion.
+    kpa = cosmology_utils.kpc_per_arcsecond(cosmology_params, z)
+    los_radius = kpa * cone_angle * 0.5
+
+    # If we're behind the main deflector, then the light cone shrinks back to a
+    # point at the source, with an additional buffer.
+    scale_factor = angle_buffer
+    scale_factor *= cosmology_utils.comoving_distance(cosmology_params, 0.0,
+        z_source)
+    scale_factor /= cosmology_utils.comoving_distance(cosmology_params, z_lens,
+        z_source)
+    scale_factor *= cosmology_utils.comoving_distance(cosmology_params, z_lens,
+        z)
+    scale_factor /= cosmology_utils.comoving_distance(cosmology_params, 0.0, z)
+
+    # Define functions for use with jax.lax.cond.
+    def before():
+        return los_radius
+
+    def after():
+        return los_radius * (1-scale_factor)
+
+    return jax.lax.cond(z < z_lens, before, after)
+
+
+def volume_element(main_deflector_params: Mapping[str, float],
+    source_params: Mapping[str, float], los_params: Mapping[str, float],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z: float) -> float:
+    """Return the volume element for the lightcone.
+
+    Args:
+        main_deflector_params: Parameters of the main deflector.
+        source_params: Parameters for the source.
+        los_parms: Parameters of the los distribution.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z: Redshift at which to calculate the volume element.
+
+    Returns:
+        Volume element of the lightcone. In units of kpc ** 3.
+    """
+    # Get the los parameters we will need.
+    dz = los_params['dz']
+
+    los_radius = cone_angle_to_radius(main_deflector_params, source_params,
+        los_params, cosmology_params, z + dz / 2)
+
+    # Get the thickness of the slice in physical units
+    dz_in_kpc = cosmology_utils.comoving_distance(cosmology_params, z, z + dz)
+    # Comoving to physical and Mpc to kpc
+    dz_in_kpc *= 1000 / (1 + z + dz / 2)
+
+    return dz_in_kpc * jnp.pi * los_radius ** 2
+
+
+def expected_num_halos(main_deflector_params: Mapping[str, float],
+    source_params: Mapping[str, float], los_params: Mapping[str, float],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z) -> float:
+    """Return the expected number of halos in a redshift slice.
+
+    Args:
+        main_deflector_params: Parameters of the main deflector.
+        source_params: Parameters for the source.
+        los_parms: Parameters of the los distribution.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z: Redhisft of the slice.
+
+    Returns:
+        Expected number of los halos in the slice.
+    """
+    # Extract the main deflector, source, and los parameters we need.
+    m_min = los_params['m_min']
+    m_max = los_params['m_max']
+    delta_los = los_params['delta_los']
+
+    slope_pl, norm_pl = mass_function_power_law(cosmology_params, z)
+    norm_pl *= volume_element(main_deflector_params, source_params,
+        los_params, cosmology_params, z)
+    norm_pl *= two_halo_boost(main_deflector_params, los_params,
+        cosmology_params, z)
+    norm_pl *= delta_los
+
+    return norm_pl * power_law.power_law_integrate(m_min , m_max, slope_pl)
+
+
+def draw_redshifts(main_deflector_params: Mapping[str, float],
+    source_params: Mapping[str, float], los_params: Mapping[str, float],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z_min: float, z_max: float, rng: Sequence[int],
+    num_z_bins: int, pad_length: int) -> jnp.ndarray:
+    """Draw the redshifts for the los halos within an interval.
+
+    Args:
+        main_deflector_params: Parameters of the main deflector.
+        source_params: Parameters for the source.
+        los_parms: Parameters of the los distribution.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z_min: Minimum redshift of interval. Lower bound of lowest redshift
+            interval considred [z_min, z_min + dz].
+        z_max: Maximum redshift of interval. Upper bound of highest redshift
+            interval considred [z_max - dz, z_max].
+        rng: jax PRNG key.
+        num_z_bins: Number of redshift bins to consider.
+        pad_length: Length to pad redshift draws to.
+
+    Returns:
+        Redshifts of los halos. A redshift of -1.0 indicates a padded halo.
+
+    Notes:
+        z_min is the lower bou
+    """
+    rng_num, rng_cdf = jax.random.split(rng)
+    # Extract the main deflector, source, and los parameters we will need.
+    dz = los_params['dz']
+
+    # Get the redshift bins which we'll use for our cdf interpolation.
+    z_samples = jnp.linspace(z_min, z_max - 1e-6, num_z_bins)
+
+    # Round down the samples so that they are only in dz increments. This means
+    # many samples will be rounded to the same value, so we will also need the
+    # average number of samples being rounded into each bin.
+    z_samples = jnp.floor(z_samples / dz) * dz
+    samples_per_bin = num_z_bins / ((z_max - z_min) / dz)
+
+    # Calculate the number of expected halos in each bin to build our cdf.
+    num_expected_samps = jax.vmap(expected_num_halos,
+        in_axes=[None, None, None, None, 0])(main_deflector_params,
+        source_params, los_params, cosmology_params, z_samples)
+    num_expected_samps /= samples_per_bin
+
+    # Get the number of los halos we will keep
+    total_expected = jnp.sum(num_expected_samps)
+    num_los = jax.random.poisson(rng_num, total_expected)
+
+    # Use the expected samples to draw from the inverse cdf.
+    cdf = jnp.cumsum(num_expected_samps) / total_expected
+    cdf_draws = jax.random.uniform(rng_cdf, shape=(pad_length,))
+    z_draws = jnp.interp(cdf_draws, cdf, z_samples)
+    z_draws = jnp.round(z_draws / dz) * dz
+
+    # More accurate to give the halos a redshift between the two slices.
+    z_draws += dz / 2
+
+    # Calculate which z draws need to be padded out.
+    indices = jnp.arange(pad_length)
+    is_pad = indices >= num_los
+
+    return z_draws * ~is_pad - is_pad
+
+
+def draw_masses(los_params: Mapping[str, float],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z_values: jnp.ndarray, rng: Sequence[int]) -> jnp.ndarray:
+    """Draw the masses for the los halos within an interval.
+
+    Args:
+        los_parms: Parameters of the los distribution.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z_values: Redshifts of los halos to draw masses for.
+        rng: jax PRNG key.
+
+    Returns:
+        Masses of los halos. In units of M_sun.
+    """
+    # Extract the los parameters we will need.
+    m_min = los_params['m_min']
+    m_max = los_params['m_max']
+
+    # Get the position of each los halo in its respective cdf
+    cdf_draws = jax.random.uniform(rng, shape=z_values.shape)
+
+    # Function for mapping the indiviudal cdf draws to a mass given the slope.
+    def draw_from_pl_cdf(m_min, m_max, slope_pl, cdf_draw):
+        s_one = slope_pl + 1
+        return (cdf_draw * (m_max ** s_one - m_min ** s_one) +
+            m_min ** s_one) ** (1 / s_one)
+
+    slopes_pl, _ = jax.vmap(mass_function_power_law, in_axes=[None, 0])(
+        cosmology_params, z_values)
+    return jax.vmap(draw_from_pl_cdf, in_axes=[None, None, 0, 0])(m_min, m_max,
+        slopes_pl, cdf_draws)
+
+
+def draw_positions(main_deflector_params: Mapping[str, float],
+    source_params: Mapping[str, float], los_params: Mapping[str, float],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    z_values: jnp.ndarray, rng: Sequence[int]) -> jnp.ndarray:
+    """Return position draws for the los halos.
+
+    Args:
+        main_deflector_params: Parameters of the main deflector.
+        source_params: Parameters for the source.
+        los_parms: Parameters of the los distribution.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        z_values: Redshifts of the slices for which positions are being drawn.
+        rng: jax PRNG key.
+
+    Returns:
+        x- and y-positions stacked over the last axis. In units of kpc.
+    """
+    rng_radius, rng_theta = jax.random.split(rng)
+
+    # Draw the los positions uniformly within the slice.
+    los_radius = jax.vmap(cone_angle_to_radius,
+        in_axes=[None, None, None, None, 0])(main_deflector_params,
+        source_params, los_params, cosmology_params, z_values)
+    radius_draws = los_radius * jnp.sqrt(
+        jax.random.uniform(rng_radius, shape=z_values.shape))
+    theta_draws = 2 * jnp.pi * jax.random.uniform(rng_theta,
+        shape=z_values.shape)
+    coords = [radius_draws * jnp.cos(theta_draws),
+        radius_draws * jnp.sin(theta_draws)]
+
+    return jnp.stack(coords, axis=-1)
+
+
+def convert_to_lensing(main_deflector_params: Mapping[str, float],
+    source_params: Mapping[str, float], los_params: Mapping[str, float],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    los_masses: jnp.ndarray, los_z: jnp.ndarray, los_cart_pos: jnp.ndarray,
+    rng: Sequence[int]) -> Tuple[jnp.ndarray, Mapping[str, jnp.ndarray]]:
+    """Convert los masses and positions into lensing quantities.
+
+    Args:
+        main_deflector_params: Parameters of the main deflector.
+        source_params: Parameters for the source.
+        los_parms: Parameters of the los distribution.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        los_mass: Masses of the los halos.
+        los_z: Redshift of the los halos
+        los_cart_pos: X- and y-coordinates of the los halos.
+        rng: jax PRNG key.
+
+    Returns:
+        Redshifts and Dictionary with model index and lensing profile properties
+        for all of the los halos.
+    """
+    # Extract the redshifts we need
+    z_lens = main_deflector_params['z_lens']
+    z_source = source_params['z_source']
+
+    # Calculate the concentration and radial position of the subhalos
+    los_c = nfw_functions.mass_concentration(los_params, cosmology_params,
+        los_masses, z_lens, rng)
+
+    # Concert from masses and concentrations to nfw parameters
+    los_r_two_hund = nfw_functions.r_two_hund_from_m(cosmology_params,
+        los_masses, z_lens)
+    los_r_scale = los_r_two_hund / los_c
+    los_rho_nfw = nfw_functions.rho_nfw_from_c(cosmology_params, los_c, z_lens)
+
+    # Convert to lensing units
+    los_r_scale_ang, los_alpha_rs = nfw_functions.convert_to_lensing_nfw(
+        cosmology_params, los_r_scale, los_z, los_rho_nfw, z_source)
+    kpa = cosmology_utils.kpc_per_arcsecond(cosmology_params, z_lens)
+    subhalos_cart_pos_ang = los_cart_pos / kpa
+
+    # There is only one model, the NFW. Subhalos with redshift -1.0 are treated
+    # as padding models.
+    subhalos_model_index = (jnp.full(los_masses.shape, -1) *
+        jnp.int32(los_z == -1.0))
+
+    subhalos_kwargs = {'model_index': subhalos_model_index,
+        'scale_radius': los_r_scale_ang, 'alpha_rs': los_alpha_rs,
+        'center_x': subhalos_cart_pos_ang[:,0],
+        'center_y': subhalos_cart_pos_ang[:,1]}
+
+    return los_z, subhalos_kwargs
+
+
+def draw_los(main_deflector_params: Mapping[str, float],
+    source_params: Mapping[str, float], los_params: Mapping[str, float],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    rng: Sequence[int], num_z_bins: int, los_pad_length: int,
+) -> Tuple[Tuple[jnp.ndarray, Mapping[str, jnp.ndarray]],
+    Tuple[jnp.ndarray, Mapping[str, jnp.ndarray]]]:
+    """Draw line-of-sight halos with redshift and lensing quantities.
+
+    Args:
+        main_deflector_params: Parameters of the main deflector.
+        source_params: Parameters for the source.
+        los_parms: Parameters of the los distribution.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        rng: jax PRNG key.
+        num_z_bins: Number of redshift bins to use in los draws. These bins
+            will be rounded to the nearest unit of dz and should be a large
+            number.
+        los_pad_length: Number of los halos to pad to.
+
+    Returns:
+        Redshifts and Dictionary with model index and lensing profile properties
+        for all of the los halos. First tuple is for the los halos before the
+        main deflector and second tuple is for los halos after the main
+        deflector.
+    """
+    rng_before, rng_after = jax.random.split(rng)
+    rng_before_z, rng_before_mass, rng_before_pos, rng_before_convert = (
+        jax.random.split(rng_before, 4))
+    rng_after_z, rng_after_mass, rng_after_pos, rng_after_convert = (
+        jax.random.split(rng_after, 4))
+
+    # Draw the redshifts, masses and positions for our los halos up to the pad
+    # before the main deflector.
+    z_lens = main_deflector_params['z_lens']
+    los_before_z = draw_redshifts(main_deflector_params, source_params,
+        los_params, cosmology_params, 0.0, z_lens, rng_before_z,
+        num_z_bins, los_pad_length)
+    los_before_z = jnp.sort(los_before_z)
+    los_before_masses = draw_masses(los_params, cosmology_params, los_before_z,
+        rng_before_mass)
+    los_before_cart_pos = draw_positions(main_deflector_params, source_params,
+        los_params, cosmology_params, los_before_z, rng_before_pos)
+    los_before_tuple = convert_to_lensing(main_deflector_params, source_params,
+        los_params, cosmology_params, los_before_masses, los_before_z,
+        los_before_cart_pos, rng_before_convert)
+
+    # Draw the redshifts, masses and positions for our los halos up to the pad
+    # after the main deflector.
+    z_source = source_params['z_source']
+    los_after_z = draw_redshifts(main_deflector_params, source_params,
+        los_params, cosmology_params, z_lens, z_source, rng_after_z,
+        num_z_bins, los_pad_length)
+    los_after_z = jnp.sort(los_after_z)
+    los_after_masses = draw_masses(los_params, cosmology_params, los_after_z,
+        rng_after_mass)
+    los_after_cart_pos = draw_positions(main_deflector_params, source_params,
+        los_params, cosmology_params, los_after_z, rng_after_pos)
+    los_after_tuple = convert_to_lensing(main_deflector_params, source_params,
+        los_params, cosmology_params, los_after_masses, los_after_z,
+        los_after_cart_pos, rng_after_convert)
+
+    # Set the padding to the minimum redshift
+    is_pad = los_before_z == -1.0
+    los_before_z = los_before_z * ~is_pad + is_pad * jnp.min(los_before_z)
+    is_pad = los_after_z == -1.0
+    los_after_z = los_after_z * ~is_pad + is_pad * jnp.min(los_after_z)
+
+    return los_before_tuple, los_after_tuple
diff --git a/jaxstronomy/models.py b/jaxstronomy/models.py
new file mode 100644
index 0000000..01bef00
--- /dev/null
+++ b/jaxstronomy/models.py
@@ -0,0 +1,136 @@
+# Copyright 2022 The Flax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Flax implementation of ResNet V1."""
+
+# See issue #620.
+# pytype: disable=wrong-arg-count
+
+from functools import partial
+from typing import Any, Callable, Sequence, Tuple
+
+from flax import linen as nn
+import jax.numpy as jnp
+
+ModuleDef = Any
+
+
+class ResNetBlock(nn.Module):
+	"""ResNet block."""
+	filters: int
+	conv: ModuleDef
+	norm: ModuleDef
+	act: Callable
+	strides: Tuple[int, int] = (1, 1)
+
+	@nn.compact
+	def __call__(self, x,):
+		residual = x
+		y = self.conv(self.filters, (3, 3), self.strides)(x)
+		y = self.norm()(y)
+		y = self.act(y)
+		y = self.conv(self.filters, (3, 3))(y)
+		y = self.norm(scale_init=nn.initializers.zeros)(y)
+
+		if residual.shape != y.shape:
+			residual = self.conv(self.filters, (1, 1),
+													 self.strides, name='conv_proj')(residual)
+			residual = self.norm(name='norm_proj')(residual)
+
+		return self.act(residual + y)
+
+
+class BottleneckResNetBlock(nn.Module):
+	"""Bottleneck ResNet block."""
+	filters: int
+	conv: ModuleDef
+	norm: ModuleDef
+	act: Callable
+	strides: Tuple[int, int] = (1, 1)
+
+	@nn.compact
+	def __call__(self, x):
+		residual = x
+		y = self.conv(self.filters, (1, 1))(x)
+		y = self.norm()(y)
+		y = self.act(y)
+		y = self.conv(self.filters, (3, 3), self.strides)(y)
+		y = self.norm()(y)
+		y = self.act(y)
+		y = self.conv(self.filters * 4, (1, 1))(y)
+		y = self.norm(scale_init=nn.initializers.zeros)(y)
+
+		if residual.shape != y.shape:
+			residual = self.conv(self.filters * 4, (1, 1),
+													 self.strides, name='conv_proj')(residual)
+			residual = self.norm(name='norm_proj')(residual)
+
+		return self.act(residual + y)
+
+
+class ResNet(nn.Module):
+	"""ResNetV1."""
+	stage_sizes: Sequence[int]
+	block_cls: ModuleDef
+	num_outputs: int
+	num_filters: int = 64
+	dtype: Any = jnp.float32
+	act: Callable = nn.relu
+	conv: ModuleDef = nn.Conv
+
+	@nn.compact
+	def __call__(self, x, train: bool = True):
+		conv = partial(self.conv, use_bias=False, dtype=self.dtype)
+		norm = partial(nn.BatchNorm,
+									 use_running_average=not train,
+									 momentum=0.9,
+									 epsilon=1e-5,
+									 dtype=self.dtype)
+
+		x = conv(self.num_filters, (7, 7), (2, 2),
+						 padding=[(3, 3), (3, 3)],
+						 name='conv_init')(x)
+		x = norm(name='bn_init')(x)
+		x = nn.relu(x)
+		x = nn.max_pool(x, (3, 3), strides=(2, 2), padding='SAME')
+		for i, block_size in enumerate(self.stage_sizes):
+			for j in range(block_size):
+				strides = (2, 2) if i > 0 and j == 0 else (1, 1)
+				x = self.block_cls(self.num_filters * 2 ** i,
+													 strides=strides,
+													 conv=conv,
+													 norm=norm,
+													 act=self.act)(x)
+		x = jnp.mean(x, axis=(1, 2))
+		x = nn.Dense(self.num_outputs, dtype=self.dtype)(x)
+		x = jnp.asarray(x, self.dtype)
+		return x
+
+
+ResNet18 = partial(ResNet, stage_sizes=[2, 2, 2, 2],
+									 block_cls=ResNetBlock)
+ResNet34 = partial(ResNet, stage_sizes=[3, 4, 6, 3],
+									 block_cls=ResNetBlock)
+ResNet50 = partial(ResNet, stage_sizes=[3, 4, 6, 3],
+									 block_cls=BottleneckResNetBlock)
+ResNet101 = partial(ResNet, stage_sizes=[3, 4, 23, 3],
+										block_cls=BottleneckResNetBlock)
+ResNet152 = partial(ResNet, stage_sizes=[3, 8, 36, 3],
+										block_cls=BottleneckResNetBlock)
+ResNet200 = partial(ResNet, stage_sizes=[3, 24, 36, 3],
+										block_cls=BottleneckResNetBlock)
+
+
+ResNet18Local = partial(ResNet, stage_sizes=[2, 2, 2, 2],
+												block_cls=ResNetBlock, conv=nn.ConvLocal)
diff --git a/jaxstronomy/nfw_functions.py b/jaxstronomy/nfw_functions.py
new file mode 100644
index 0000000..1803d27
--- /dev/null
+++ b/jaxstronomy/nfw_functions.py
@@ -0,0 +1,229 @@
+# coding=utf-8
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     https://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Define functions useful for conversion between different NFW conventions.
+
+This module contains the functions used to move between NFW conventions and
+to transform NFW parameters into lensing inputs.
+"""
+
+from typing import Mapping, Sequence, Tuple, Union
+
+import jax
+import jax.numpy as jnp
+
+from jaxstronomy import cosmology_utils
+
+
+def r_two_hund_from_m(
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    masses: Union[jnp.ndarray, float], z: float) -> Union[float, jnp.ndarray]:
+    """Calculate the two-hundred radial overdensity from the mass.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        masses: Mass values to calculate the overdensity radius for.
+        z: Redshift at which to conduct the calculation.
+
+    Returns:
+        The overdensity radius corresponding to each provided mass.
+    """
+    # Calculate the radius using the critical density at the given redshift.
+    h = cosmology_params['hubble_constant'] / 100.0
+    rho_crit = cosmology_utils.rho_crit(cosmology_params, z) * h ** 2
+    return (3 * masses / (4 * jnp.pi * rho_crit * 200)) ** (1.0 / 3.0)
+
+
+def rho_nfw_from_c(
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    c: Union[jnp.ndarray, float], z: float):
+    """Calculate the amplitude of the NFW profile.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        c: Concentration of the NFW.
+        z: Redshift at which to define the amplitude
+
+    Returns:
+        Amplitude of the NFW
+    """
+    # Calculate the radius using the critical density at the given redshift.
+    h = cosmology_params['hubble_constant'] / 100.0
+    rho_crit = cosmology_utils.rho_crit(cosmology_params, z) * h ** 2
+    return (c ** 3 * rho_crit * 200) / (3 * (jnp.log(1 + c) - c / (1+c)))
+
+
+def _cored_nfw_integral(r_tidal: float, rho_nfw: float, r_scale: float,
+    r_upper: jnp.ndarray) -> jnp.ndarray:
+    """Integrate the cored NFW profile.
+
+    Args:
+        r_tidal: Tidal radius of the NFW profile. Units of kpc.
+        rho_nfw: Normalization of the NFW profile.
+        r_scale: Scale radius of NFW profile. Units of kpc.
+        r_upper: Upper bounds to evaluate in units of kpc.
+
+    Returns:
+        Value of the integral for each r_upper.
+    """
+    # Convert to natural units for NFW
+    x_tidal = r_tidal / r_scale
+    x_upper = r_upper / r_scale
+
+    # Get the NFW value in the core.
+    linear_scaling = rho_nfw / (x_tidal * (1 + x_tidal) ** 2)
+
+    # Calculate the cored component of the integral
+    integral_values = (1.0 / 3.0 * linear_scaling *
+        jnp.minimum(r_tidal,r_upper)**3)
+
+    # Add the nfw component for x_upper > x_tidal
+    lower_bound = 1 / (x_tidal + 1) + jnp.log(x_tidal + 1)
+    upper_bound = 1 / (x_upper + 1) + jnp.log(x_upper + 1)
+    nfw_integral = upper_bound - lower_bound
+    add_nfw = r_upper > r_tidal
+    integral_values += nfw_integral * rho_nfw * r_scale**3 * add_nfw
+
+    return integral_values * 4 * jnp.pi
+
+
+def cored_nfw_draws(r_tidal: float, rho_nfw: float, r_scale: float,
+    r_max: float, rng: Sequence[int], n_draws: float) -> jnp.ndarray:
+    """Return radial samples from a cored nfw profile.
+
+    Args:
+        r_tidal: Tidal radius of the NFW
+        rho_nfw: Amplitude of the NFW
+        r_scale: Scale radius of the NFW
+        r_max: Maximum radius to draw halos within
+        rng: Jax PRNG key.
+        n_draws: Number of positions to draw.
+
+    Returns:
+        Radial position draws.
+    """
+    # Default to 1000 values for the inverse cdf interpolation.
+    r_values = jnp.linspace(0, r_max, 1000)
+    cdf_values = _cored_nfw_integral(r_tidal, rho_nfw, r_scale, r_values)
+
+    # Normalize
+    cdf_values /= jnp.max(cdf_values)
+
+    # Interpolate the inverse cdf function and use that to draw radial
+    # samples
+    cdf_draws = jax.random.uniform(rng, shape=(n_draws,))
+    return jnp.interp(cdf_draws, cdf_values, r_values)
+
+
+def convert_to_lensing_nfw(
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    r_scale: float, z: float, rho_nfw: float, z_source: float
+) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    """Convert physical NFW parameters to lensing parameters.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        r_scale: Scale radius of the NFW.
+        z: Redshift of the NFW.
+        rho_nfw: Amplitude of the NFW.
+        z_source: Redshift of the source.
+
+    Returns:
+        The lensing scale radius and the deflection strength at the scale
+            radius.
+    """
+    kpa = cosmology_utils.kpc_per_arcsecond(cosmology_params, z)
+
+    # Two of the parameters just need to be transformed to arcseconds.
+    r_scale_ang = r_scale / kpa
+
+    # Lensing alpha is defined with respect to sigma_crit.
+    sigma_crit = cosmology_utils.calculate_sigma_crit(cosmology_params, z,
+        z_source)
+    alpha_rs = rho_nfw * (4 * r_scale ** 2 * (1 + jnp.log(0.5)))
+    alpha_rs /= kpa * sigma_crit
+
+    return r_scale_ang, alpha_rs
+
+
+def convert_to_lensing_tnfw(
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    r_scale: float, z: float, rho_nfw: float, r_trunc: float, z_source: float
+) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
+    """Convert physical tNFW parameters to lensing parameters.
+
+    Args:
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        r_scale: Scale radius of the tNFW.
+        z: Redshift of the tNFW.
+        rho_nfw: Amplitude of the tNFW.
+        r_trunc: Truncation radius of the tNFW.
+        z_source: Redshift of the source.
+
+    Returns:
+        The lensing scale radius, the deflection strength at the scale
+            radius, and the truncation radius.
+    """
+    kpa = cosmology_utils.kpc_per_arcsecond(cosmology_params, z)
+
+    # First two parameters can be transformed as for an NFW.
+    r_scale_ang, alpha_rs = convert_to_lensing_nfw(cosmology_params, r_scale, z,
+        rho_nfw, z_source)
+
+    # Truncation radius just needs to be converted to radians.
+    r_trunc_ang = r_trunc / kpa
+
+    return r_scale_ang, alpha_rs, r_trunc_ang
+
+
+def mass_concentration(substructure_params: Mapping[str, float],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    masses: Union[float, jnp.ndarray], z: float,
+    rng: Sequence[int]) -> jnp.ndarray:
+    """Return the concentration of halos at a given mass and redshift.
+
+    Args:
+        substructure_params: Parameters of the substructure distribution.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        masses: Masses to pull concentrations for. In units of M_sun.
+        z: Redshift of the halos.
+        rng: Jax PRNG key.
+
+    Returns:
+        Concentration for each halo.
+    """
+    # Pull out the mass-concentration parameters we need.
+    c_zero = substructure_params['c_zero']
+    zeta = substructure_params['conc_zeta']
+    beta = substructure_params['conc_beta']
+    m_ref = substructure_params['conc_m_ref']
+    dex_scatter = substructure_params['conc_dex_scatter']
+
+    # Use peak heights to calculate concentrations
+    h = cosmology_params['hubble_constant'] / 100
+    peak_heights = jax.vmap(cosmology_utils.peak_height,
+        in_axes=[None, 0, None])(cosmology_params, masses * h, z)
+    peak_height_ref = cosmology_utils.peak_height(cosmology_params, m_ref * h,
+        0.0)
+    concentrations = (c_zero * (1+z) ** zeta *
+        (peak_heights / peak_height_ref) ** (-beta))
+
+    # Add scatter and return concentrations
+    conc_scatter = jax.random.normal(rng, shape=concentrations.shape)
+    conc_scatter *= dex_scatter
+    return 10**(jnp.log10(concentrations) + conc_scatter)
diff --git a/jaxstronomy/power_law.py b/jaxstronomy/power_law.py
new file mode 100644
index 0000000..c6c6476
--- /dev/null
+++ b/jaxstronomy/power_law.py
@@ -0,0 +1,79 @@
+# coding=utf-8
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     https://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Power law integration and draws in jax.
+
+This module includes functions to integrate a power law and draw from a power
+law distribution.
+"""
+
+from typing import Sequence
+
+import jax
+import jax.numpy as jnp
+
+
+def power_law_integrate(p_min: float, p_max: float, slope: float) -> float:
+    """Integrate a power law
+
+    Args:
+        p_min: Lower bound of the power law
+        p_max: Upper bound of the power law
+        slope: Slope of the power law
+
+    Returns:
+        Integral of the power law x^slope from p_min to p_max
+    """
+    upper_bound = 1/(slope+1)*p_max**(slope+1)
+    lower_bound = 1/(slope+1)*p_min**(slope+1)
+    return upper_bound-lower_bound
+
+
+def power_law_draw(p_min: float, p_max: float, slope: float, norm:float,
+    rng: Sequence[int], pad_length: int) -> jnp.ndarray:
+    """Sample from a power law
+
+    Args:
+        p_min: Lower bound of the power law
+        p_max: Upper bound of the power law
+        slope: Slope of the power law
+        norm: Normalization of the power law
+        rng: jax PRNG key used for noise realization.
+        max_subhalos: The length to pad the draw to. Must be a static argument
+            for jax.jit compilation. If more subhalos are drawn than the
+            pad_length, they will be dropped.
+
+    Returns:
+        Values drawn from the power law. Padded values are set to 0.
+    """
+    # Get the expected number of draws
+    n_expected = norm * power_law_integrate(p_min,p_max,slope)
+
+    # Split our random key for the number of draws and the cdf draws
+    rng_draws, rng_cdf = jax.random.split(rng)
+
+    # Draw the number of objects as a poisson process
+    n_draws = jax.random.poisson(rng_draws, n_expected)
+
+    # Get the positions in the cdf we want to draw from a uniform and then
+    # convert that to values in the pdf
+    cdf = jax.random.uniform(rng_cdf, shape=(pad_length,))
+    s_one = slope + 1
+    draws = (cdf * (p_max ** s_one - p_min ** s_one) +
+        p_min ** s_one) ** (1 / s_one)
+
+    # Pad the draws
+    indices = jnp.arange(pad_length)
+    is_not_pad = indices < n_draws
+
+    return draws * is_not_pad
diff --git a/jaxstronomy/subhalos.py b/jaxstronomy/subhalos.py
new file mode 100644
index 0000000..c99ef84
--- /dev/null
+++ b/jaxstronomy/subhalos.py
@@ -0,0 +1,315 @@
+# coding=utf-8
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     https://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Draw subhalo masses and cocentrations for NFW subhalos.
+
+This module containts the functions needed to draw subhalos for an underlying
+distribution as defined in https://arxiv.org/pdf/1909.02573.pdf.
+"""
+
+from typing import Mapping, Optional, Sequence, Tuple, Union
+
+import jax
+import jax.numpy as jnp
+
+from jaxstronomy import cosmology_utils
+from jaxstronomy import nfw_functions
+from jaxstronomy import power_law
+
+
+def host_scaling_function(host_mass: float, z_lens: float, k_one: float,
+    k_two: float) -> float:
+    """Return scaling for the subhalos mass function due to host mass.
+
+    Args:
+        host_mass: Mass of the host halo in units of M_sun.
+        z_lens: The redshift of the host halo.
+        k1: Amplitude of halo mass dependence.
+        k2: Amplitude of the redshift scaling.
+
+    Notes:
+        Derived from galacticus in https://arxiv.org/pdf/1909.02573.pdf.
+    """
+    log_f = (k_one * jnp.log10(host_mass / 1e13) +
+        k_two * jnp.log10(z_lens + 0.5))
+    return 10 ** log_f
+
+
+def draw_nfw_masses(main_deflector_params: Mapping[str, float],
+    subhalo_params: Mapping[str, float],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    rng: Sequence[int], pad_length: int) -> jnp.ndarray:
+    """Draw from the subhalo mass function and return the masses.
+
+    Args:
+        main_deflector_params: Parameters of the main deflector.
+        subhalo_parms: Parameters of the subhalo distribution.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        rng: Jax PRNG key.
+        pad_length: Length to pad mass draw to.
+
+    Returns:
+        Masses of the drawn halos in units of M_sun.
+    """
+    # Cap sigma_sub to 0 for draws.
+    sigma_sub = jax.lax.max(0.0, subhalo_params['sigma_sub'])
+
+    # Extract the remaining parameters from the dict for readability
+    shmf_plaw_index = subhalo_params['shmf_plaw_index']
+    m_pivot = subhalo_params['m_pivot']
+    m_min = subhalo_params['m_min']
+    m_max = subhalo_params['m_max']
+    k_one = subhalo_params['k_one']
+    k_two = subhalo_params['k_two']
+    host_mass = main_deflector_params['mass']
+    z_lens = main_deflector_params['z_lens']
+
+    # Calculate the overall norm of the power law, including the host scaling
+    f_host = host_scaling_function(host_mass, z_lens, k_one, k_two)
+    kpa = cosmology_utils.kpc_per_arcsecond(cosmology_params, z_lens)
+    radius_e = kpa * main_deflector_params['theta_e']
+    area_elem = jnp.pi * (3 * radius_e) ** 2
+
+    # Fold the pivot mass into the norm directly
+    norm = f_host * area_elem * sigma_sub * m_pivot ** (-shmf_plaw_index - 1)
+
+    # Draw our masses
+    return power_law.power_law_draw(m_min, m_max, shmf_plaw_index, norm, rng,
+        pad_length)
+
+
+def rejection_sampling(radial_coord: jnp.ndarray, r_two_hund: float,
+    r_bound: float, rng: Sequence[int]) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    """Return three-dimensional sampling of the halo positions with rejection.
+
+    Args:
+        radial_cord: Radial displacement from center for each halo.
+        r_two_hund: Two-hundred times overdensity radius of the host halo,
+            which will be used to set the bounds on the z-coordinate.
+        r_bound: Bound for x- and y-coordinate in units of kpc.
+        rng: jax PRNG key used for noise realization.
+
+    Returns:
+        Acceptance and three dimensional coordinates of each subhalo.
+    """
+    # Split the rng key for each draw.
+    rng_theta, rng_phi = jax.random.split(rng)
+    theta = jax.random.uniform(rng_theta, shape=radial_coord.shape)
+    theta *= 2 * jnp.pi
+    phi = jnp.arccos(1 - 2 * jax.random.uniform(rng_phi,
+        shape=radial_coord.shape))
+
+    # Get each of the cartesian coordinates
+    cart_x = radial_coord * jnp.sin(phi) * jnp.cos(theta)
+    cart_y = radial_coord * jnp.sin(phi) * jnp.sin(theta)
+    cart_z = radial_coord * jnp.cos(phi)
+
+    # Test if the coordinates fall within the boudns
+    r_xy_inside = jnp.sqrt(cart_x ** 2 + cart_y ** 2) < r_bound
+    z_inside = jnp.abs(cart_z) < r_two_hund
+    is_inside = jnp.logical_and(r_xy_inside, z_inside)
+
+    return (is_inside, jnp.stack([cart_x, cart_y, cart_z], axis=-1))
+
+
+def sample_cored_nfw(main_deflector_params: Mapping[str, float],
+    subhalo_params: Mapping[str, float],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    rng: Sequence[int], n_subs: int, sampling_pad_length: int) -> jnp.ndarray:
+    """ Sample positions for NFW subhalos.
+
+    Args:
+        main_deflector_params: Parameters of the main deflector.
+        subhalo_params: Parameters of the subhalo population.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        rng: jax PRNG key used for noise realization.
+        n_subs: Number of subhalos to draw positions for.
+        sampling_pad_length: Number of positions to sample before conducting
+            rejection sampling. For small numbers, the odds of all of the
+            returned positions being within the host R200 and three times
+            the Einstein radius goes down.
+
+    Returns:
+        x-y-z coordinates of the subhalos in units of kpc.
+
+    Notes:
+        Subhalo positions are drawn within the scale radius of the host in the
+        z-direction and three times the Einstein radius in the xy-plane.
+    """
+    # Seperate rng keys for each of our random draws.
+    rng_host, rng_radii, rng_pos = jax.random.split(rng, 3)
+
+    # Extract the relevant host properties
+    host_mass = main_deflector_params['mass']
+    z_lens = main_deflector_params['z_lens']
+    host_c = nfw_functions.mass_concentration(subhalo_params, cosmology_params,
+        jnp.array([host_mass]), z_lens, rng_host)
+    host_r_two_hund = nfw_functions.r_two_hund_from_m(cosmology_params,
+        host_mass, z_lens)
+    host_r_scale = host_r_two_hund / host_c
+    host_rho_nfw = nfw_functions.rho_nfw_from_c(cosmology_params, host_c,
+        z_lens)
+
+    kpa = cosmology_utils.kpc_per_arcsecond(cosmology_params, z_lens)
+    radius_e_three = kpa * main_deflector_params['theta_e'] * 3
+
+    # The maximum radius is the corner for the cylinder in which we're
+    # rendering.
+    r_max = jnp.sqrt((radius_e_three) ** 2 + host_r_two_hund ** 2)
+
+    # Draw coordinates up to the sampling pad length.
+    radial_coord = nfw_functions.cored_nfw_draws(host_r_two_hund / 2,
+        host_rho_nfw, host_r_scale, r_max, rng_radii, sampling_pad_length)
+    is_inside, cart_pos = rejection_sampling(radial_coord, host_r_two_hund,
+        radius_e_three, rng_pos)
+
+    # We'll sort the cartesian coordinates by whether or not they're inside
+    # the bounds and only return n_sub coordinates. This does not guarantee
+    # that all the samples returned will be within the boudns, but for
+    # large enough sampling_pad_length, nearly all of them will be.
+    return cart_pos[jnp.argsort(~is_inside)[:n_subs]]
+
+
+def get_truncation_radius(subhalos_mass: jnp.ndarray,
+    subhalos_radii: jnp.ndarray, m_pivot: Optional[float] = 1e7,
+    r_pivot: Optional[float] = 50) -> jnp.ndarray:
+    """Return truncation radius for the subhalos.
+
+    Args:
+        subhalo_mass: Mass of the subhalo in units of M_sun.
+        subhalos_radii: Radius of the subhalos in units of kpc.
+        m_pivot: Pivot mass for the scaling in units of M_sun.
+        r_pivot: Pivot radius for the relation in units of kpc.
+
+    Returns:
+        Truncation radius in units of kpc.
+    """
+    # Set a minimum truncation radius of 1e-7 kpc to avoid poorly defined
+    # models
+    r_trunc_min = 1e-7
+    return jax.lax.max((1.4 * (subhalos_mass / m_pivot) ** (1 / 3) *
+        (subhalos_radii / r_pivot) ** (2 / 3)), r_trunc_min)
+
+
+def convert_to_lensing(main_deflector_params: Mapping[str, float],
+    source_params: Mapping[str, Union[int, float]],
+    subhalo_params: Mapping[str, float],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    subhalos_masses: jnp.ndarray, subhalos_cart_pos: jnp.ndarray,
+    rng: Sequence[int]) -> Tuple[jnp.ndarray, Mapping[str, jnp.ndarray]]:
+    """Convert subhalo masses and positions into lensing quantities.
+
+    Args:
+        main_deflector_params: Parameters of the main deflector.
+        source_params: Parameters of the source.
+        subhalo_params: Parameters of the subhalo population.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        subhalo_masses: Masses of the subhalos in units of M_sun.
+        subhalos_cart_pos: x-y-z-position of the subhalos in units of kpc.
+        rng: Jax PRNG key.
+
+    Returns:
+        Redshifts for each subhalo and the dictionary of lensing quantities.
+
+    Notes:
+        Model index set to -1 for subhalos with zero mass.
+    """
+    # Extract the redshifts we need
+    z_lens = main_deflector_params['z_lens']
+    z_source = source_params['z_source']
+
+    # All the subhalos are at the main deflector redshift
+    subhalos_z = jnp.full(subhalos_masses.shape, z_lens)
+
+    # Calculate the concentration and radial position of the subhalos
+    subhalos_c = nfw_functions.mass_concentration(subhalo_params, cosmology_params,
+        subhalos_masses, z_lens, rng)
+    subhalos_radii = jnp.sqrt(jnp.sum(subhalos_cart_pos ** 2, axis=-1))
+
+    # Convert from masses and concentrations to nfw parameters
+    subhalos_r_two_hund = nfw_functions.r_two_hund_from_m(cosmology_params,
+        subhalos_masses, z_lens)
+    subhalos_r_scale = subhalos_r_two_hund / subhalos_c
+    subhalos_rho_nfw = nfw_functions.rho_nfw_from_c(cosmology_params,
+        subhalos_c, z_lens)
+    subhalos_r_trunc = get_truncation_radius(subhalos_masses, subhalos_radii)
+
+    # Convert to lensing units
+    lq_tuple = nfw_functions.convert_to_lensing_tnfw(cosmology_params,
+        subhalos_r_scale, subhalos_z, subhalos_rho_nfw, subhalos_r_trunc,
+        z_source)
+    subhalos_r_scale_ang, subhlos_alpha_rs, subhalos_r_trunc_ang = lq_tuple
+    kpa = cosmology_utils.kpc_per_arcsecond(cosmology_params, z_lens)
+    subhalos_cart_pos_ang = subhalos_cart_pos / kpa
+
+    # There is only one model, the tNFW. Subhalos with mass 0 are treated
+    # as padding models.
+    subhalos_model_index = (jnp.full(subhalos_masses.shape, -1) *
+        jnp.int32(subhalos_masses == 0))
+
+    subhalos_kwargs = {'model_index': subhalos_model_index,
+        'scale_radius': subhalos_r_scale_ang, 'alpha_rs': subhlos_alpha_rs,
+        'trunc_radius': subhalos_r_trunc_ang,
+        # Assume the first main deflector model sets the center.
+        'center_x': (
+            subhalos_cart_pos_ang[:,0] + main_deflector_params['center_x']),
+        'center_y': (
+            subhalos_cart_pos_ang[:,1] + main_deflector_params['center_y'])}
+
+    return subhalos_z, subhalos_kwargs
+
+
+def draw_subhalos(main_deflector_params: Mapping[str, float],
+    source_params: Mapping[str, Union[int, float]],
+    subhalo_params: Mapping[str, float],
+    cosmology_params: Mapping[str, Union[float, int, jnp.ndarray]],
+    rng: Sequence[int], subhalos_pad_length: int, sampling_pad_length: int,
+) -> Tuple[jnp.ndarray, Mapping[str, jnp.ndarray]]:
+    """Draw subhalos with redshift and lensing quantities.
+
+    Args:
+        main_deflector_params: Parameters of the main deflector.
+        source_params: Parameters of the source.
+        subhalo_params: Parameters of the subhalo population.
+        cosmology_params: Cosmological parameters that define the universe's
+            expansion.
+        rng: Jax PRNG key.
+        subhalos_pad_length: Number of subhalos to draw, including padding. If
+            this number is too small, you may undercount the number of subhalos
+            in your system.
+        sampling_pad_length: Number of positional samples to draw before
+            rejection sampling the subhalo positions within the volume
+            boundaries. Should be one to two order of magntidue larger than the
+            subhalo pad length.
+
+    Returns:
+        Redshifts for each subhalo and the dictionary of lensing quantities.
+
+    Notes:
+        Model index set to -1 for subhalos with zero mass.
+    """
+    rng_masses, rng_pos, rng_convert = jax.random.split(rng, 3)
+
+    # Draw the masses and positions for our subhalos up to the pad.
+    subhalos_mass = draw_nfw_masses(main_deflector_params, subhalo_params,
+        cosmology_params, rng_masses, subhalos_pad_length)
+    subhalos_cart_pos = sample_cored_nfw(main_deflector_params, subhalo_params,
+        cosmology_params, rng_pos, subhalos_pad_length, sampling_pad_length)
+
+    # Return our lensing keyword dictionary and array of redshifts.
+    return convert_to_lensing(main_deflector_params, source_params,
+        subhalo_params, cosmology_params, subhalos_mass, subhalos_cart_pos,
+        rng_convert)
diff --git a/jaxstronomy/train.py b/jaxstronomy/train.py
new file mode 100644
index 0000000..36495be
--- /dev/null
+++ b/jaxstronomy/train.py
@@ -0,0 +1,345 @@
+# coding=utf-8
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Training script for dark matter substructure inference."""
+
+import functools
+import itertools
+import time
+from typing import Any, Iterator, Mapping, Optional, Sequence, Tuple, Union
+
+from absl import app
+from absl import flags
+
+from clu import metric_writers
+from clu import periodic_actions
+from flax import jax_utils
+from flax.training import checkpoints
+from flax.training import common_utils
+from flax.training import train_state
+import jax
+from jax import lax
+import jax.numpy as jnp
+import ml_collections
+import optax
+
+from jaxstronomy import input_pipeline
+from jaxstronomy import models
+
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string('workdir', None, 'working directory')
+flags.DEFINE_float('learning_rate', 0.001, 'learning rate')
+flags.DEFINE_integer('num_unique_batches', 0,
+    'number of unique batches of data to draw. If 0 produces infinite data.')
+
+
+def initialized(
+        rng: Sequence[int], image_size: int, model: Any
+    ) -> Tuple[Any, Any]:
+    """Initialize the model parameters
+
+    Args:
+        rng: jax PRNG key.
+        image_size: Size of the input image.
+        model: Model class to initialize.
+
+    Returns:
+        Initialized model parameters and batch stats.
+    """
+    input_shape = (1, image_size, image_size, 1)
+    @jax.jit
+    def init(*args):
+        return model.init(*args)
+    variables = init({'params': rng}, jnp.ones(input_shape, model.dtype))
+    return variables['params'], variables['batch_stats']
+
+
+def gaussian_loss(outputs: jnp.ndarray, truth: jnp.ndarray) -> jnp.ndarray:
+    """Gaussian loss calculated on the outputs and truth values.
+
+    Args:
+        outputs: Values outputted by the model.
+        truth: True value of the parameters.
+
+    Returns:
+        Gaussian loss.
+
+    Notes:
+        Loss does not inlcude constant factor of 1 / (2 * pi) ^ (d/2)
+    """
+    mean, log_var = jnp.split(outputs, 2, axis=-1)
+    loss = 0.5 * jnp.sum(
+        jnp.multiply(jnp.square(mean-truth), jnp.exp(-log_var)), axis=-1)
+    loss += 0.5*jnp.sum(log_var, axis=-1)
+    return jnp.mean(loss)
+
+
+def compute_metrics(
+        outputs: jnp.ndarray, truth: jnp.ndarray) -> Mapping[str, jnp.ndarray]:
+    """Compute the performance metrics of the output.
+
+    Args:
+        outputs: Values outputted by the model.
+        truth: True value of the parameters.
+
+    Returns:
+        Value of each of the metrics.
+    """
+    loss = gaussian_loss(outputs, truth)
+    mean, _ = jnp.split(outputs, 2, axis=-1)
+    rmse = jnp.sqrt(jnp.mean(jnp.square(mean - truth)))
+    metrics = {
+        'loss': loss,
+        'rmse': rmse,
+    }
+    metrics = lax.pmean(metrics, axis_name='batch')
+    return metrics
+
+
+def get_learning_rate_schedule(
+        config: ml_collections.ConfigDict,
+        base_learning_rate: float) -> Any:
+    """Return the learning rate schedule function.
+
+    Args:
+        config: Training configuration.
+        base_learning_rate: Base learning rate for the schedule.
+
+    Returns:
+        Mapping from step to learning rate according to the schedule.
+    """
+    warmup_fn = optax.linear_schedule(init_value=0.0,
+        end_value=base_learning_rate,
+        transition_steps=config.warmup_steps)
+    cosine_steps = max(config.num_train_steps - config.warmup_steps, 1)
+    cosine_fn = optax.cosine_decay_schedule(init_value=base_learning_rate,
+        decay_steps=cosine_steps)
+    schedule_fn = optax.join_schedules(schedules=[warmup_fn, cosine_fn],
+        boundaries=[config.warmup_steps])
+    return schedule_fn
+
+
+def train_step(state, batch, learning_rate_schedule):
+    """Perform a single training step."""
+
+    def loss_fn(params):
+        """loss function used for training."""
+        outputs, new_model_state = state.apply_fn(
+            {'params': params, 'batch_stats': state.batch_stats},
+            batch['image'],
+            mutable=['batch_stats'])
+        loss = gaussian_loss(outputs, batch['truth'])
+        # weight_penalty_params = jax.tree_util.tree_leaves(params)
+        # weight_decay = 0.0001
+        # weight_l2 = sum(jnp.sum(x ** 2)
+        #                 for x in weight_penalty_params
+        #                 if x.ndim > 1)
+        # weight_penalty = weight_decay * 0.5 * weight_l2
+        # loss = loss + weight_penalty
+        return loss, (new_model_state, outputs)
+
+    step = state.step
+    lr = learning_rate_schedule(step)
+
+    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
+    aux, grads = grad_fn(state.params)
+    # Re-use same axis_name as in the call to `pmap(...train_step...)` below.
+    grads = lax.pmean(grads, axis_name='batch')
+    new_model_state, outputs = aux[1]
+    metrics = compute_metrics(outputs, batch['truth'])
+    metrics['learning_rate'] = lr
+
+    new_state = state.apply_gradients(
+        grads=grads, batch_stats=new_model_state['batch_stats'])
+
+    return new_state, metrics
+
+
+class TrainState(train_state.TrainState):
+    batch_stats: Any
+
+
+def restore_checkpoint(state, workdir):
+    return checkpoints.restore_checkpoint(workdir, state)
+
+
+def save_checkpoint(state, workdir, keep_every_n_steps=None):
+    if jax.process_index() == 0:
+        # get train state from the first replica
+        state = jax.device_get(jax.tree_util.tree_map(lambda x: x[0], state))
+        step = int(state.step)
+        checkpoints.save_checkpoint(workdir, state, step, keep=3,
+            keep_every_n_steps=keep_every_n_steps)
+
+
+# pmean only works inside pmap because it needs an axis name.
+# This function will average the inputs across all devices.
+cross_replica_mean = jax.pmap(lambda x: lax.pmean(x, 'x'), 'x')
+
+
+def sync_batch_stats(state):
+    """Sync the batch statistics across replicas."""
+    # Each device has its own version of the running average batch statistics and
+    # we sync them before evaluation.
+    return state.replace(batch_stats=cross_replica_mean(state.batch_stats))
+
+
+def create_train_state(rng, config: ml_collections.ConfigDict,
+                       model, image_size, learning_rate_schedule):
+    """Create initial training state."""
+    params, batch_stats = initialized(rng, image_size, model)
+    tx = optax.adam(
+        learning_rate=learning_rate_schedule
+    )
+    state = TrainState.create(
+        apply_fn=model.apply,
+        params=params,
+        tx=tx,
+        batch_stats=batch_stats)
+    return state
+
+
+def train_and_evaluate(config: ml_collections.ConfigDict,
+                       input_config: dict,
+                       workdir: str,
+                       rng: Union[Iterator[Sequence[int]], Sequence[int]],
+                       image_size: int, learning_rate: float) -> TrainState:
+    """
+    """
+
+    writer = metric_writers.create_default_writer(
+        logdir=workdir, just_logging=jax.process_index() != 0)
+
+    if config.batch_size % jax.device_count() > 0:
+        raise ValueError('Batch size must be divisible by the number of devices')
+
+    steps_per_epoch = config.steps_per_epoch
+    num_steps = config.num_train_steps
+
+    base_learning_rate = learning_rate * config.batch_size / 256.
+
+    model_cls = getattr(models, config.model)
+    num_outputs = len(input_config['truth_parameters'][0]) * 2
+    model = model_cls(num_outputs=num_outputs, dtype=jnp.float32)
+
+    learning_rate_schedule = get_learning_rate_schedule(config,
+        base_learning_rate)
+
+    if isinstance(rng, Iterator):
+        rng_state = next(rng)
+    else:
+        rng, rng_state = jax.random.split(rng)
+    state = create_train_state(rng_state, config, model, image_size,
+        learning_rate_schedule)
+    state = restore_checkpoint(state, workdir)
+
+    # step_offset > 0 if restarting from checkpoint
+    step_offset = int(state.step)
+    state = jax_utils.replicate(state)
+
+    p_train_step = jax.pmap(functools.partial(train_step,
+        learning_rate_schedule=learning_rate_schedule),
+        axis_name='batch')
+    draw_image_and_truth_pmap = jax.pmap(jax.jit(jax.vmap(
+        functools.partial(
+            input_pipeline.draw_image_and_truth,
+            all_models=input_config['all_models'],
+            principal_md_index=input_config['principal_md_index'],
+            principal_source_index=input_config['principal_source_index'],
+            kwargs_simulation=input_config['kwargs_simulation'],
+            kwargs_detector=input_config['kwargs_detector'],
+            kwargs_psf=input_config['kwargs_psf'],
+            truth_parameters=input_config['truth_parameters']),
+        in_axes=(None, None, None, None, 0))),
+        in_axes=(None, None, None, None, 0)
+    )
+    if isinstance(rng, Iterator):
+        rng_cosmo = next(rng)
+    else:
+        rng, rng_cosmo = jax.random.split(rng)
+    cosmology_params = input_pipeline.intialize_cosmology_params(input_config,
+                                                                 rng_cosmo)
+    grid_x, grid_y = input_pipeline.generate_grids(input_config)
+
+    train_metrics = []
+    hooks = []
+    if jax.process_index() == 0:
+        hooks += [periodic_actions.Profile(num_profile_steps=5, logdir=workdir)]
+    train_metrics_last_t = time.time()
+    print('Initial compilation, this might take some minutes...')
+    for step in range(step_offset, num_steps):
+
+        # Generate truths and images
+        if isinstance(rng, Iterator):
+            rng_images = next(rng)
+        else:
+            rng, rng_images = jax.random.split(rng)
+        rng_images = jax.random.split(
+            rng_images, num=jax.device_count() * config.batch_size).reshape(
+                (jax.device_count(), config.batch_size, -1)
+            )
+
+        image, truth = draw_image_and_truth_pmap(input_config['lensing_config'],
+                                                 cosmology_params, grid_x,
+                                                 grid_y, rng_images)
+        image = jnp.expand_dims(image, axis=-1)
+
+        batch = {'image': image, 'truth': truth}
+        state, metrics = p_train_step(state, batch)
+        for h in hooks:
+            h(step)
+        if step == step_offset:
+            print('Initial compilation completed.')
+
+        train_metrics.append(metrics)
+        if (step + 1) % steps_per_epoch == 0:
+            train_metrics = common_utils.get_metrics(train_metrics)
+            summary = {
+                f'train_{k}': v
+                for k, v in jax.tree_util.tree_map(lambda x: x.mean(), train_metrics).items()
+            }
+            summary['steps_per_second'] = steps_per_epoch / (
+                time.time() - train_metrics_last_t)
+            writer.write_scalars(step + 1, summary)
+            print(summary)
+            train_metrics = []
+            train_metrics_last_t = time.time()
+
+        if (step + 1) % steps_per_epoch == 0 or step + 1 == num_steps:
+            state = sync_batch_stats(state)
+            save_checkpoint(state, workdir, config.keep_every_n_steps)
+
+    # Wait until computations are done before exiting
+    jax.random.normal(jax.random.PRNGKey(0), ()).block_until_ready()
+
+    return state
+
+
+def main(_):
+    from jaxstronomy import train_config
+    from jaxstronomy import input_config
+
+    config = train_config.get_config()
+    ic = input_config.get_config()
+    image_size = ic['kwargs_detector']['n_x']
+    rng = jax.random.PRNGKey(0)
+    if FLAGS.num_unique_batches > 0:
+        rng_list = jax.random.split(rng, FLAGS.num_unique_batches)
+        rng = itertools.cycle(rng_list)
+    train_and_evaluate(config, ic, FLAGS.workdir, rng, image_size,
+                       FLAGS.learning_rate)
+
+
+if __name__ == '__main__':
+    app.run(main)
diff --git a/jaxstronomy/train_config.py b/jaxstronomy/train_config.py
new file mode 100644
index 0000000..3355c91
--- /dev/null
+++ b/jaxstronomy/train_config.py
@@ -0,0 +1,23 @@
+import ml_collections
+
+
+def get_config():
+    """Get the default hyperparameter configuration."""
+    config = ml_collections.ConfigDict()
+
+    # As defined in the `models` module.
+    config.model = 'ResNet50'
+
+    config.momentum = 0.9
+    config.batch_size = 64
+
+    config.cache = False
+    config.half_precision = False
+
+    # If num_train_steps==-1 then the number of training steps is calculated from
+    # num_epochs using the entire dataset. Similarly for steps_per_eval.
+    config.steps_per_epoch = 7800
+    config.warmup_steps = 10 * config.steps_per_epoch
+    config.num_train_steps = 500 * config.steps_per_epoch
+    config.keep_every_n_steps = 500
+    return config
diff --git a/jaxstronomy/utils.py b/jaxstronomy/utils.py
index 0ab2822..66c0ad2 100644
--- a/jaxstronomy/utils.py
+++ b/jaxstronomy/utils.py
@@ -1,5 +1,4 @@
 # coding=utf-8
-# Copyright 2022 The Google Research Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -7,20 +6,6 @@
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Copyright 2022 Google LLC
-
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-
-#     https://www.apache.org/licenses/LICENSE-2.0
-
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -35,84 +20,84 @@
 
 
 def coordinates_evaluate(
-    n_x, n_y, pixel_width,
-    supersampling_factor):
-  """Returns the coordinate grid in the observer plane in angular units.
-
-  Args:
-    n_x: Number of pixels in x direction
-    n_y: Number of pixels in y direction
-    pixel_width: Size of a pixel in angular units. Pixels are assumed to be
-      square.
-    supersampling_factor: Factor by which to supersample light rays. Number of
-      light rays will scale like supersampling_factor**2.
-
-  Returns:
-    X- and y-coordinates at which to raytrace.
-  """
-  # Center the x- and y-coordinates at 0.
-  x_span = (
-      jnp.arange(n_x * supersampling_factor) / supersampling_factor -
-      (n_x - 1 / supersampling_factor) / 2.)
-  y_span = (
-      jnp.arange(n_y * supersampling_factor) / supersampling_factor -
-      (n_y - 1 / supersampling_factor) / 2.)
-  pix_xy = jnp.stack(
-      [coord.flatten() for coord in jnp.meshgrid(x_span, y_span)], axis=0)
-  radec_xy = pix_xy * pixel_width
-  return radec_xy[0], radec_xy[1]
+        n_x, n_y, pixel_width,
+        supersampling_factor):
+    """Returns the coordinate grid in the observer plane in angular units.
+
+    Args:
+        n_x: Number of pixels in x direction
+        n_y: Number of pixels in y direction
+        pixel_width: Size of a pixel in angular units. Pixels are assumed to be
+            square.
+        supersampling_factor: Factor by which to supersample light rays. Number of
+            light rays will scale like supersampling_factor**2.
+
+    Returns:
+        X- and y-coordinates at which to raytrace.
+    """
+    # Center the x- and y-coordinates at 0.
+    x_span = (
+            jnp.arange(n_x * supersampling_factor) / supersampling_factor -
+            (n_x - 1 / supersampling_factor) / 2.)
+    y_span = (
+            jnp.arange(n_y * supersampling_factor) / supersampling_factor -
+            (n_y - 1 / supersampling_factor) / 2.)
+    pix_xy = jnp.stack(
+            [coord.flatten() for coord in jnp.meshgrid(x_span, y_span)], axis=0)
+    radec_xy = pix_xy * pixel_width
+    return radec_xy[0], radec_xy[1]
 
 
 def unpack_parameters_xy(
-    func,  # Use Concatenate once available internally.
-    parameters,
+        func,  # Use Concatenate once available internally.
+        parameters,
 ):
-  """Returns function wrapper that unpacks parameters for grid functions.
+    """Returns function wrapper that unpacks parameters for grid functions.
 
-  Returns function wrapper that unpacks required parameters for functions whose
-  first two parameters are the x- and y-coordinates at which they should be
-  evaluated.
+    Returns function wrapper that unpacks required parameters for functions whose
+    first two parameters are the x- and y-coordinates at which they should be
+    evaluated.
 
-  Args:
-    func: Function that takes x- and y-coordinates as well as additional args.
-    parameters: Parameters to unpack.
+    Args:
+        func: Function that takes x- and y-coordinates as well as additional args.
+        parameters: Parameters to unpack.
 
-  Returns:
-    Wrapper for func that unpacks keyword parameters, passes them to func, and
-    returns the output.
-  """
+    Returns:
+        Wrapper for func that unpacks keyword parameters, passes them to func, and
+        returns the output.
+    """
 
-  def derivative_wrapper(x, y, kwargs, parameters):
-    return func(x, y, *[kwargs[param] for param in parameters])
+    def derivative_wrapper(x, y, kwargs, parameters):
+        return func(x, y, *[kwargs[param] for param in parameters])
 
-  return functools.partial(derivative_wrapper, parameters=parameters)
+    return functools.partial(derivative_wrapper, parameters=parameters)
 
 
 def downsample(image, supersampling_factor):
-  """Downsamples image to correct for supersampling factor.
+    """Downsamples image to correct for supersampling factor.
 
-  Args:
-    image: Image to downsample.
-    supersampling_factor: Factor by which light rays were supersampled.
+    Args:
+        image: Image to downsample.
+        supersampling_factor: Factor by which light rays were supersampled.
 
-  Returns:
-    Downsampled image.
-  """
-  n_x, n_y = image.shape
-  image = jnp.reshape(image,
-                      (n_x // supersampling_factor, supersampling_factor,
-                       n_y // supersampling_factor, supersampling_factor))
-  return jnp.mean(jnp.mean(image, axis=3), axis=1)
+    Returns:
+        Downsampled image.
+    """
+    n_x, n_y = image.shape
+    image = jnp.reshape(image,
+                                            (n_x // supersampling_factor, supersampling_factor,
+                                             n_y // supersampling_factor, supersampling_factor))
+    return jnp.mean(jnp.mean(image, axis=3), axis=1)
 
 
 def magnitude_to_cps(magnitude, magnitude_zero_point):
-  """Converts magnitude to counts per second.
+    """Converts magnitude to counts per second.
 
-  Args:
-    magnitude: Input magnitude
-    magnitude_zero_point: Zero point magnitude of the detector.
+    Args:
+        magnitude: Input magnitude
+        magnitude_zero_point: Zero point magnitude of the detector.
 
-  Returns:
-    Counts per second corresponding to input magnitude.
-  """
-  return 10**(-(magnitude - magnitude_zero_point) / 2.5)
+    Returns:
+        Counts per second corresponding to input magnitude.
+    """
+    return 10**(-(magnitude - magnitude_zero_point) / 2.5)
diff --git a/notebooks/Comparison.ipynb b/notebooks/Comparison.ipynb
new file mode 100644
index 0000000..517f801
--- /dev/null
+++ b/notebooks/Comparison.ipynb
@@ -0,0 +1,422 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1bcce9f8-5b40-4ca0-8f1e-b84bcdde80cb",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "import paltas\n",
+    "import jaxstronomy\n",
+    "import jax\n",
+    "import jax.numpy as jnp\n",
+    "import numpy as np\n",
+    "import functools\n",
+    "import matplotlib.pyplot as plt\n",
+    "import copy"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "23104dfa-ee4f-4070-b200-b4688adba4ec",
+   "metadata": {},
+   "source": [
+    "The paltas and jaxstronomy codebase are performing the same operations, so \n",
+    "\n",
+    "There are three components to an update step for our model:\n",
+    "\n",
+    "1. Drawing a population of dark matter substructure for the lensing calculations.\n",
+    "2. Generating a lensing image given that dark matter substructure population.\n",
+    "3. Updating the model given a batch of images.\n",
+    "\n",
+    "Below we'll create code that should generate the same outputs (up to random seed issues for the first bullepoint) that we can\n",
+    "compare."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3ee63f8e-47ac-42af-b8c9-a6370b93ff1e",
+   "metadata": {},
+   "source": [
+    "## Drawing a population of dark matter substructure for the lensing calculations."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "54cdcc6e-684e-4938-a8a9-cf3786ab1623",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Setup paltas and jaxstronomy configs\n",
+    "config_handler = paltas.Configs.config_handler.ConfigHandler('comparison_files/input_config_paltas.py')\n",
+    "\n",
+    "from comparison_files import input_config_jaxstronomy\n",
+    "rng = jax.random.PRNGKey(0)\n",
+    "input_config_jax = input_config_jaxstronomy.get_config()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d6799902-8c81-4eaf-82dd-198aabfe4652",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# A bit of setup for the jaxstronomy code. This is drawing the specific distribution parameters for the subhalos, main deflector halos,\n",
+    "# and the source. These are all encoded as constants, so we're just drawing the constant values.\n",
+    "subhalo_params = jaxstronomy.input_pipeline.draw_sample(input_config_jax['lensing_config']['subhalo_params'], rng)\n",
+    "main_deflector_params = jaxstronomy.input_pipeline.draw_sample(input_config_jax['lensing_config']['main_deflector_params'], rng)\n",
+    "source_params = jaxstronomy.input_pipeline.draw_sample(input_config_jax['lensing_config']['source_params'], rng)\n",
+    "\n",
+    "# Initialize the cosmology parameters we need for our cosmology calculations and create a jitted function for fast draws.\n",
+    "cosmology_params = jaxstronomy.input_pipeline.intialize_cosmology_params(input_config_jax, rng)\n",
+    "subhalos_pad_length = 1000\n",
+    "sampling_pad_length = 100000\n",
+    "draw_subhalos_jit = jax.jit(functools.partial(jaxstronomy.subhalos.draw_subhalos, subhalos_pad_length=subhalos_pad_length, \n",
+    "                                              sampling_pad_length=sampling_pad_length))\n",
+    "_ = draw_subhalos_jit(main_deflector_params, source_params, subhalo_params, cosmology_params, rng)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "420c28a1-8ecc-453f-87b1-209b91e6f0c6",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Here I'm drawing from paltas and from jaxstronomy and storing the results of both draws.\n",
+    "n_draws = 1000\n",
+    "paltas_kwargs_list = []\n",
+    "paltas_n_subhalos = np.zeros(n_draws)\n",
+    "for i in range(n_draws):\n",
+    "    paltas_models, paltas_kwargs, paltas_z = config_handler.subhalo_class.draw_subhalos()\n",
+    "    paltas_n_subhalos[i] = len(paltas_models)\n",
+    "    paltas_kwargs_list.append(paltas_kwargs)\n",
+    "\n",
+    "rng_draw, _ = jax.random.split(rng)\n",
+    "jaxstronomy_kwargs_list = []\n",
+    "jaxstronomy_n_subhalos = np.zeros(n_draws)\n",
+    "for i in range(n_draws):\n",
+    "    rng_draw, _ = jax.random.split(rng_draw)\n",
+    "    jaxstronomy_z_list, jaxstronomy_kwargs = draw_subhalos_jit(main_deflector_params, source_params, subhalo_params, cosmology_params, rng_draw)\n",
+    "    jaxstronomy_n_subhalos[i] = jnp.sum(jaxstronomy_kwargs['alpha_rs'] > 0.0)\n",
+    "    jaxstronomy_kwargs_list.append(jaxstronomy_kwargs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c84a4e86-8cc0-4431-a029-e3f3ae1519aa",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Start by comparing the number of subhalos being drawn and make sure they look similar.\n",
+    "_, bins, _ = plt.hist(paltas_n_subhalos, bins=30, histtype='step', lw=3)\n",
+    "plt.hist(jaxstronomy_n_subhalos, bins=bins, histtype='step', lw=3)\n",
+    "plt.xlabel('Number of Subhalos')\n",
+    "plt.ylabel('Count')\n",
+    "plt.legend(['paltas', 'jaxstronomy'])\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8413d3d2-895c-475d-827f-10b5ac860b41",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Now compare the distribution of the parameters.\n",
+    "def extract_np_array_paltas(param, list_of_draws):\n",
+    "    concat_values = []\n",
+    "    for draw in list_of_draws:\n",
+    "        for subhalos in draw:\n",
+    "            concat_values.append(subhalos[param])\n",
+    "    return np.array(concat_values)\n",
+    "\n",
+    "def extract_np_array_jaxstronomy(param, list_of_draws):\n",
+    "    concat_values = []\n",
+    "    for draw in list_of_draws:\n",
+    "        concat_values.append(draw[param])\n",
+    "    return np.concatenate(concat_values)\n",
+    "\n",
+    "def extract_without_zeros_jaxstronomy(param, list_of_draws, key):\n",
+    "    array = extract_np_array_jaxstronomy(param, list_of_draws)\n",
+    "    return array[key > 0.0]\n",
+    "\n",
+    "alpha_rs_jax = extract_np_array_jaxstronomy('alpha_rs', jaxstronomy_kwargs_list)\n",
+    "_, bins, _ = plt.hist(np.log(extract_np_array_paltas('alpha_Rs', paltas_kwargs_list)), bins=100, histtype='step', log=True, lw=3)\n",
+    "plt.hist(np.log(extract_without_zeros_jaxstronomy('alpha_rs', jaxstronomy_kwargs_list, alpha_rs_jax)), bins=bins, histtype='step', log=True, lw=3)\n",
+    "plt.xlabel(r'Max Proxy $(\\alpha_{Rs})$')\n",
+    "plt.ylabel('Count')\n",
+    "plt.legend(['paltas', 'jaxstronomy'])\n",
+    "plt.show()\n",
+    "\n",
+    "_, bins, _ = plt.hist(np.log(extract_np_array_paltas('Rs', paltas_kwargs_list)), bins=100, histtype='step', log=True, lw=3)\n",
+    "plt.hist(np.log(extract_without_zeros_jaxstronomy('scale_radius', jaxstronomy_kwargs_list, alpha_rs_jax)), bins=bins, histtype='step', log=True, lw=3)\n",
+    "plt.xlabel(r'Size Proxy $(Rs)$')\n",
+    "plt.ylabel('Count')\n",
+    "plt.legend(['paltas', 'jaxstronomy'])\n",
+    "plt.show()\n",
+    "\n",
+    "_, bins, _ = plt.hist(np.log(extract_np_array_paltas('r_trunc', paltas_kwargs_list)), bins=100, histtype='step', log=True, lw=3)\n",
+    "plt.hist(np.log(extract_without_zeros_jaxstronomy('trunc_radius', jaxstronomy_kwargs_list, alpha_rs_jax)), bins=bins, histtype='step', log=True, lw=3)\n",
+    "plt.xlabel(r'Truncation Radius')\n",
+    "plt.ylabel('Count')\n",
+    "plt.legend(['paltas', 'jaxstronomy'])\n",
+    "plt.show()\n",
+    "\n",
+    "_, bins, _ = plt.hist(extract_np_array_paltas('center_x', paltas_kwargs_list), bins=100, histtype='step', log=True, lw=3)\n",
+    "plt.hist(extract_without_zeros_jaxstronomy('center_x', jaxstronomy_kwargs_list, alpha_rs_jax), bins=bins, histtype='step', log=True, lw=3)\n",
+    "plt.xlabel(r'Center x')\n",
+    "plt.ylabel('Count')\n",
+    "plt.legend(['paltas', 'jaxstronomy'])\n",
+    "plt.show()\n",
+    "\n",
+    "_, bins, _ = plt.hist(extract_np_array_paltas('center_y', paltas_kwargs_list), bins=100, histtype='step', log=True, lw=3)\n",
+    "plt.hist(extract_without_zeros_jaxstronomy('center_y', jaxstronomy_kwargs_list, alpha_rs_jax), bins=bins, histtype='step', log=True, lw=3)\n",
+    "plt.xlabel(r'Center y')\n",
+    "plt.ylabel('Count')\n",
+    "plt.legend(['paltas', 'jaxstronomy'])\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ffef12b4-9dc4-4eea-87db-ee03c77d2e0d",
+   "metadata": {},
+   "source": [
+    "### Modify two of the substructure parameters, slope and normalization, and make sure that everything still agrees"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a050cbea-e334-4c17-98f9-8da00b30dae1",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "input_config_jax['lensing_config']['subhalo_params']['sigma_sub'] = jaxstronomy.input_pipeline.encode_constant(1.0e-3)\n",
+    "input_config_jax['lensing_config']['subhalo_params']['shmf_plaw_index'] = jaxstronomy.input_pipeline.encode_constant(-1.92)\n",
+    "subhalo_params = jaxstronomy.input_pipeline.draw_sample(input_config_jax['lensing_config']['subhalo_params'], rng)\n",
+    "config_handler = paltas.Configs.config_handler.ConfigHandler('comparison_files/input_config_paltas_two.py')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d5f39bbb-f9a0-42da-8725-9d3fcad4c045",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Again, draw from paltas and from jaxstronomy and storing the results of both draws.\n",
+    "n_draws = 1000\n",
+    "paltas_kwargs_list = []\n",
+    "paltas_n_subhalos = np.zeros(n_draws)\n",
+    "for i in range(n_draws):\n",
+    "    paltas_models, paltas_kwargs, paltas_z = config_handler.subhalo_class.draw_subhalos()\n",
+    "    paltas_n_subhalos[i] = len(paltas_models)\n",
+    "    paltas_kwargs_list.append(paltas_kwargs)\n",
+    "\n",
+    "rng_draw, _ = jax.random.split(rng)\n",
+    "jaxstronomy_kwargs_list = []\n",
+    "jaxstronomy_n_subhalos = np.zeros(n_draws)\n",
+    "for i in range(n_draws):\n",
+    "    rng_draw, _ = jax.random.split(rng_draw)\n",
+    "    jaxstronomy_z_list, jaxstronomy_kwargs = draw_subhalos_jit(main_deflector_params, source_params, subhalo_params, cosmology_params, rng_draw)\n",
+    "    jaxstronomy_n_subhalos[i] = jnp.sum(jaxstronomy_kwargs['alpha_rs'] > 0.0)\n",
+    "    jaxstronomy_kwargs_list.append(jaxstronomy_kwargs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b7647972-454c-4638-865b-0264aefbdb2e",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Start by comparing the number of subhalos being drawn and make sure they look similar.\n",
+    "_, bins, _ = plt.hist(paltas_n_subhalos, bins=30, histtype='step', lw=3)\n",
+    "plt.hist(jaxstronomy_n_subhalos, bins=bins, histtype='step', lw=3)\n",
+    "plt.xlabel('Number of Subhalos')\n",
+    "plt.ylabel('Count')\n",
+    "plt.legend(['paltas', 'jaxstronomy'])\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "76025e56-2cf7-4e81-a327-20e9f36d701e",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "alpha_rs_jax = extract_np_array_jaxstronomy('alpha_rs', jaxstronomy_kwargs_list)\n",
+    "_, bins, _ = plt.hist(np.log(extract_np_array_paltas('alpha_Rs', paltas_kwargs_list)), bins=100, histtype='step', log=True, lw=3)\n",
+    "plt.hist(np.log(extract_without_zeros_jaxstronomy('alpha_rs', jaxstronomy_kwargs_list, alpha_rs_jax)), bins=bins, histtype='step', log=True, lw=3)\n",
+    "plt.xlabel(r'Max Proxy $(\\alpha_{Rs})$')\n",
+    "plt.ylabel('Count')\n",
+    "plt.legend(['paltas', 'jaxstronomy'])\n",
+    "plt.show()\n",
+    "\n",
+    "_, bins, _ = plt.hist(np.log(extract_np_array_paltas('Rs', paltas_kwargs_list)), bins=100, histtype='step', log=True, lw=3)\n",
+    "plt.hist(np.log(extract_without_zeros_jaxstronomy('scale_radius', jaxstronomy_kwargs_list, alpha_rs_jax)), bins=bins, histtype='step', log=True, lw=3)\n",
+    "plt.xlabel(r'Size Proxy $(Rs)$')\n",
+    "plt.ylabel('Count')\n",
+    "plt.legend(['paltas', 'jaxstronomy'])\n",
+    "plt.show()\n",
+    "\n",
+    "_, bins, _ = plt.hist(np.log(extract_np_array_paltas('r_trunc', paltas_kwargs_list)), bins=100, histtype='step', log=True, lw=3)\n",
+    "plt.hist(np.log(extract_without_zeros_jaxstronomy('trunc_radius', jaxstronomy_kwargs_list, alpha_rs_jax)), bins=bins, histtype='step', log=True, lw=3)\n",
+    "plt.xlabel(r'Truncation Radius')\n",
+    "plt.ylabel('Count')\n",
+    "plt.legend(['paltas', 'jaxstronomy'])\n",
+    "plt.show()\n",
+    "\n",
+    "_, bins, _ = plt.hist(extract_np_array_paltas('center_x', paltas_kwargs_list), bins=100, histtype='step', log=True, lw=3)\n",
+    "plt.hist(extract_without_zeros_jaxstronomy('center_x', jaxstronomy_kwargs_list, alpha_rs_jax), bins=bins, histtype='step', log=True, lw=3)\n",
+    "plt.xlabel(r'Center x')\n",
+    "plt.ylabel('Count')\n",
+    "plt.legend(['paltas', 'jaxstronomy'])\n",
+    "plt.show()\n",
+    "\n",
+    "_, bins, _ = plt.hist(extract_np_array_paltas('center_y', paltas_kwargs_list), bins=100, histtype='step', log=True, lw=3)\n",
+    "plt.hist(extract_without_zeros_jaxstronomy('center_y', jaxstronomy_kwargs_list, alpha_rs_jax), bins=bins, histtype='step', log=True, lw=3)\n",
+    "plt.xlabel(r'Center y')\n",
+    "plt.ylabel('Count')\n",
+    "plt.legend(['paltas', 'jaxstronomy'])\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "718365fd-b607-465a-8d82-ec51aeaafc67",
+   "metadata": {},
+   "source": [
+    "## Generating a lensing image and comparing it to nno substructure."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b4fd1ae6-8d42-4829-9722-437f1dfd591e",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Setup paltas and jaxstronomy configs\n",
+    "config_handler_zero = paltas.Configs.config_handler.ConfigHandler('comparison_files/input_config_paltas_zero.py')\n",
+    "config_handler = paltas.Configs.config_handler.ConfigHandler('comparison_files/input_config_paltas.py')\n",
+    "# No noise in paltas and jaxstronomy to make things easier\n",
+    "config_handler.add_noise=False\n",
+    "config_handler_zero.add_noise = False\n",
+    "\n",
+    "from comparison_files import input_config_jaxstronomy\n",
+    "rng = jax.random.PRNGKey(0)\n",
+    "input_config_jax = input_config_jaxstronomy.get_config()\n",
+    "draw_image_and_truth_jit = jax.jit(functools.partial(jaxstronomy.input_pipeline.draw_image_and_truth, all_models=input_config_jax['all_models'],\n",
+    "                                                    principal_md_index=input_config_jax['principal_md_index'], principal_source_index=input_config_jax['principal_source_index'],\n",
+    "                                                    kwargs_simulation=input_config_jax['kwargs_simulation'], kwargs_detector=input_config_jax['kwargs_detector'],\n",
+    "                                                    kwargs_psf=input_config_jax['kwargs_psf'], truth_parameters=input_config_jax['truth_parameters']))\n",
+    "lensing_config_zero = copy.deepcopy(input_config_jax['lensing_config'])\n",
+    "lensing_config_zero['subhalo_params']['sigma_sub'] = jaxstronomy.input_pipeline.encode_constant(0.0)\n",
+    "cosmology_params = jaxstronomy.input_pipeline.intialize_cosmology_params(input_config_jax, rng)\n",
+    "grid_x, grid_y = jaxstronomy.input_pipeline.generate_grids(input_config_jax)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c6c8104a-d170-4584-9427-bdc04a81dadc",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "f, ax = plt.subplots(4, 4, figsize=(21.5, 22), sharex=False, sharey=False,gridspec_kw={'hspace': 0.02,'wspace':0.02},dpi=100)\n",
+    "image_zero, _ = config_handler_zero.draw_image()\n",
+    "for i in range(16):\n",
+    "    image, metadata = config_handler.draw_image()\n",
+    "    ax[i//4,i%4].imshow((image-image_zero)/image_zero,cmap='plasma')\n",
+    "    ax[i//4,i%4].get_xaxis().set_visible(False)\n",
+    "    ax[i//4,i%4].get_yaxis().set_visible(False)\n",
+    "    \n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d7899769-093e-43cc-b78e-cd437de9febc",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "f, ax = plt.subplots(4, 4, figsize=(21.5, 22), sharex=False, sharey=False,gridspec_kw={'hspace': 0.02,'wspace':0.02},dpi=100)\n",
+    "rng_draw, _ = jax.random.split(rng)\n",
+    "image_jax_zero, _ = draw_image_and_truth_jit(lensing_config_zero, cosmology_params, grid_x, grid_y, rng_draw)\n",
+    "for i in range(16):\n",
+    "    rng_draw, _ = jax.random.split(rng_draw)\n",
+    "    image_jax, truth_jax = draw_image_and_truth_jit(input_config_jax['lensing_config'], cosmology_params, grid_x, grid_y, rng_draw)\n",
+    "    ax[i//4,i%4].imshow((image_jax-image_jax_zero)/image_jax_zero,cmap='plasma')\n",
+    "    ax[i//4,i%4].get_xaxis().set_visible(False)\n",
+    "    ax[i//4,i%4].get_yaxis().set_visible(False)\n",
+    "    \n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "88f847bb-487f-4a00-84b9-cc5668161cbb",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "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.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/notebooks/comparison_files/input_config_jaxstronomy.py b/notebooks/comparison_files/input_config_jaxstronomy.py
new file mode 100644
index 0000000..5608d48
--- /dev/null
+++ b/notebooks/comparison_files/input_config_jaxstronomy.py
@@ -0,0 +1,136 @@
+# coding=utf-8
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Configuration file for generating jaxstronomy image outputs.
+"""
+
+import jax
+import jax.numpy as jnp
+
+from jaxstronomy.input_pipeline import encode_normal, encode_uniform
+from jaxstronomy.input_pipeline import encode_constant
+from jaxstronomy import lens_models
+from jaxstronomy import psf_models
+from jaxstronomy import source_models
+
+
+def get_config():
+    """Return the config for input generation.
+    """
+    config = {}
+
+    # Specify the configuration for each one of the parameters.
+    config['lensing_config'] = {
+        'los_params':{
+            'delta_los': encode_constant(0.0),
+            'r_min': encode_constant(0.5),
+            'r_max': encode_constant(10.0),
+            'm_min': encode_constant(1e7),
+            'm_max': encode_constant(1e10),
+            'dz': encode_constant(0.01),
+            'cone_angle': encode_constant(8.0),
+            'angle_buffer': encode_constant(0.8),
+            'c_zero': encode_constant(18),
+            'conc_zeta': encode_constant(-0.2),
+            'conc_beta': encode_constant(0.8),
+            'conc_m_ref': encode_constant(1e8),
+            'conc_dex_scatter': encode_constant(1.0)
+        },
+        'main_deflector_params': {
+            'mass': encode_constant(1e13),
+            'z_lens': encode_constant(0.5),
+            'theta_e': encode_constant(1.1),
+            'slope': encode_constant(2.0),
+            'center_x': encode_constant(0.08),
+            'center_y': encode_constant(-0.16),
+            'axis_ratio': encode_constant(0.9),
+            'angle': encode_constant(0.0),
+            'gamma_ext': encode_constant(0.0)
+        },
+        'subhalo_params':{
+            'sigma_sub': encode_constant(2.0e-3),
+            'shmf_plaw_index': encode_constant(-1.98),
+            'm_pivot': encode_constant(1e10),
+            'm_min': encode_constant(7e7),
+            'm_max': encode_constant(1e10),
+            'k_one': encode_constant(0.0),
+            'k_two': encode_constant(0.0),
+            'c_zero': encode_constant(18),
+            'conc_zeta': encode_constant(-0.2),
+            'conc_beta': encode_constant(0.8),
+            'conc_m_ref': encode_constant(1e8),
+            'conc_dex_scatter': encode_constant(1.0)
+        },
+        'source_params':{
+            'z_source': encode_constant(1.5),
+            'amp': encode_constant(10.0),
+            'sersic_radius': encode_constant(1.5),
+            'n_sersic': encode_constant(3.0),
+            'axis_ratio': encode_constant(0.9),
+            'angle': encode_constant(0.0),
+            'center_x': encode_constant(0.16),
+            'center_y': encode_constant(-0.08)
+        },
+        'lens_light_params':{
+            'amp': encode_constant(0.0),
+            'sersic_radius': encode_uniform(minimum=1.0, maximum=3.0),
+            'n_sersic': encode_uniform(minimum=1.0, maximum=4.0),
+            'axis_ratio': encode_normal(mean=1.0, std=0.05),
+            'angle': encode_uniform(minimum=0.0, maximum=2 * jnp.pi),
+            'center_x': encode_normal(mean=0.0, std=0.16),
+            'center_y': encode_normal(mean=0.0, std=0.16)
+        }
+    }
+
+    # The remaining parameters should not be drawn from random distributions.
+    config['kwargs_detector'] = {
+        'n_x': 128, 'n_y': 128, 'pixel_width': 0.04, 'supersampling_factor': 2,
+        'exposure_time': 1024, 'num_exposures': 2.0, 'sky_brightness': 22,
+        'magnitude_zero_point': 25, 'read_noise': 3.0
+    }
+    config['all_models'] = {
+        'all_los_models': (lens_models.NFW,),
+        'all_subhalo_models': (lens_models.TNFW,),
+        'all_main_deflector_models': (lens_models.EPL, lens_models.Shear),
+        'all_source_models': (source_models.SersicElliptic,),
+        'all_psf_models': (psf_models.Gaussian,)
+    }
+    config['cosmology_params'] = {
+        'omega_m_zero': encode_constant(0.3089),
+        'omega_b_zero': encode_constant(0.0486),
+        'omega_de_zero': encode_constant(0.6910088292453472),
+        'omega_rad_zero': encode_constant(9.117075466e-5),
+        'temp_cmb_zero': encode_constant(2.7255),
+        'hubble_constant': encode_constant(67.74),
+        'n_s': encode_constant(0.9667),
+        'sigma_eight': encode_constant(0.815)
+    }
+    config['kwargs_simulation'] = {
+        'num_z_bins': 1000,
+        'los_pad_length': 10,
+        'subhalos_pad_length': 750,
+        'sampling_pad_length': 10000,
+    }
+
+    config['rng'] = jax.random.PRNGKey(0)
+
+    config['kwargs_psf'] = {'model_index': 0, 'fwhm': 0.04, 'pixel_width': 0.02}
+
+    config['principal_md_index'] = 0
+    config['principal_source_index'] = 0
+    config['truth_parameters'] = (
+        ['main_deflector_params', 'main_deflector_params',
+         'main_deflector_params', 'main_deflector_params', 'subhalo_params'],
+        ['theta_e', 'slope', 'center_x', 'center_y', 'sigma_sub'])
+
+    return config
diff --git a/notebooks/comparison_files/input_config_paltas.py b/notebooks/comparison_files/input_config_paltas.py
new file mode 100644
index 0000000..f77bf14
--- /dev/null
+++ b/notebooks/comparison_files/input_config_paltas.py
@@ -0,0 +1,103 @@
+# Configuration for a tight spread on main deflector with PEMD + SHEAR, sources
+# drawn from COSMOS, and only varying the d_los and sigma_sub of the DG_19
+# subhalo and los classes
+
+import numpy as np
+from scipy.stats import norm, truncnorm, uniform
+from paltas.Substructure.los_dg19 import LOSDG19
+from paltas.Substructure.subhalos_dg19 import SubhalosDG19
+from paltas.MainDeflector.simple_deflectors import PEMDShear
+from paltas.Sources.sersic import SingleSersicSource
+from lenstronomy.Util.kernel_util import degrade_kernel
+from paltas.Sampling import distributions
+from astropy.io import fits
+import pandas as pd
+import paltas
+import os
+
+# Define the numerics kwargs.
+kwargs_numerics = {'supersampling_factor':2,'supersampling_convolution':True}
+# We do not use point_source_supersampling_factor but it must be passed in to
+# surpress a warning.
+kwargs_numerics['point_source_supersampling_factor'] = (
+    kwargs_numerics['supersampling_factor'])
+# This is always the number of pixels for the CCD. If drizzle is used, the
+# final image will be larger.
+numpix = 128
+
+# Define some general image kwargs for the dataset
+mask_radius = 0.0
+mag_cut = 2.0
+
+# Define arguments that will be used multiple times
+output_ab_zeropoint = 25.127
+
+config_dict = {
+    'subhalo':{
+        'class': SubhalosDG19,
+        'parameters':{
+            'sigma_sub':2.0e-3,
+            'shmf_plaw_index':-1.98,
+            'm_pivot': 1e10,
+            'm_min': 7e7,
+            'm_max': 1e10,
+            'c_0':18,
+            'conc_zeta':-0.2,
+            'conc_beta':0.8,
+            'conc_m_ref': 1e8,
+            'k1':0.0,
+            'k2':0.0,
+            'dex_scatter': 1.0
+        }
+    },
+    'main_deflector':{
+        'class': PEMDShear,
+        'parameters':{
+            'M200': 1e13,
+            'z_lens': 0.5,
+            'gamma': 2.0,
+            'theta_E': 1.1,
+            'e1': (1.0-0.9)/(1.0+0.9),
+            'e2': 0.0,
+            'center_x': 0.08,
+            'center_y': -0.16,
+            'gamma1': 0.0,
+            'gamma2': 0.0,
+            'ra_0':0.0, 
+            'dec_0':0.0
+        }
+    },
+    'source':{
+        'class': SingleSersicSource,
+        'parameters':{
+            'magnitude': -29.027384789, # Equivalent to an amplitude of 10.
+            'output_ab_zeropoint': output_ab_zeropoint,
+            'R_sersic': 1.5,
+            'n_sersic':3.0 ,
+            'e1': (1.0-0.9)/(1.0+0.9),
+            'e2': 0.0,
+            'center_x': 0.16,
+            'center_y': -0.08,
+            'z_source':1.5}
+    },
+    'cosmology':{
+        'parameters':{
+            'cosmology_name': 'planck18'
+        }
+    },
+    'psf':{
+        'parameters':{
+            'psf_type':'GAUSSIAN',
+            'fwhm': 0.04,
+            'point_source_supersampling_factor':2
+        }
+    },
+    'detector':{
+        'parameters':{
+            'pixel_scale':0.040,'ccd_gain':1.58,'read_noise':3.0,
+            'magnitude_zero_point':output_ab_zeropoint,
+            'exposure_time':1380,'sky_brightness':21.83,
+            'num_exposures':1,'background_noise':None
+        }
+    }
+}
diff --git a/notebooks/comparison_files/input_config_paltas_two.py b/notebooks/comparison_files/input_config_paltas_two.py
new file mode 100644
index 0000000..eaf8e4a
--- /dev/null
+++ b/notebooks/comparison_files/input_config_paltas_two.py
@@ -0,0 +1,133 @@
+# Configuration for a tight spread on main deflector with PEMD + SHEAR, sources
+# drawn from COSMOS, and only varying the d_los and sigma_sub of the DG_19
+# subhalo and los classes
+
+import numpy as np
+from scipy.stats import norm, truncnorm, uniform
+from paltas.Substructure.los_dg19 import LOSDG19
+from paltas.Substructure.subhalos_dg19 import SubhalosDG19
+from paltas.MainDeflector.simple_deflectors import PEMDShear
+from paltas.Sources.sersic import SingleSersicSource
+from lenstronomy.Util.kernel_util import degrade_kernel
+from paltas.Sampling import distributions
+from astropy.io import fits
+import pandas as pd
+import paltas
+import os
+
+# Define the numerics kwargs.
+kwargs_numerics = {'supersampling_factor':2,'supersampling_convolution':True}
+# We do not use point_source_supersampling_factor but it must be passed in to
+# surpress a warning.
+kwargs_numerics['point_source_supersampling_factor'] = (
+    kwargs_numerics['supersampling_factor'])
+# This is always the number of pixels for the CCD. If drizzle is used, the
+# final image will be larger.
+numpix = 128
+
+# Define some general image kwargs for the dataset
+mask_radius = 0.5
+mag_cut = 2.0
+
+# Define arguments that will be used multiple times
+output_ab_zeropoint = 25.127
+
+# Define the cosmos path
+root_path = paltas.__path__[0][:-7]
+cosmos_folder = root_path + r'/datasets/cosmos/COSMOS_23.5_training_sample/'
+
+# Load the empirical psf. Grab a psf from the middle of the first chip.
+# Degrade to account for the 4x supersample
+hdul = fits.open(os.path.join(root_path,
+    'datasets/hst_psf/emp_psf_f814w.fits'))
+# Don't leave any 0 values in the psf.
+psf_pix_map = degrade_kernel(hdul[0].data[17]-np.min(hdul[0].data[17]),2)
+
+config_dict = {
+    'subhalo':{
+        'class': SubhalosDG19,
+        'parameters':{
+            'sigma_sub':1.0e-3,
+            'shmf_plaw_index':-1.92,
+            'm_pivot': 1e10,
+            'm_min': 7e7,
+            'm_max': 1e10,
+            'c_0':18,
+            'conc_zeta':-0.2,
+            'conc_beta':0.8,
+            'conc_m_ref': 1e8,
+            'k1':0.0,
+            'k2':0.0,
+            'dex_scatter': 1.0
+        }
+    },
+    'los':{
+        'class': LOSDG19,
+        'parameters':{
+            'delta_los':0.0,
+            'm_min':1e7,
+            'm_max':1e10,
+            'z_min':0.01,
+            'dz':0.01,
+            'cone_angle':8.0,
+            'r_min':0.5,
+            'r_max':10.0,
+            'c_0':18,
+            'conc_zeta':-0.2,
+            'conc_beta':0.8,
+            'conc_m_ref':1e8,
+            'alpha_dz_factor':5.0,
+            'dex_scatter': 1.0
+        }
+    },
+    'main_deflector':{
+        'class': PEMDShear,
+        'parameters':{
+            'M200': 1e13,
+            'z_lens': 0.5,
+            'gamma': 2.0,
+            'theta_E': 1.1,
+            'e1': (1.0-0.9)/(1.0+0.9),
+            'e2': 0.0,
+            'center_x': 0.08,
+            'center_y': -0.16,
+            'gamma1': 0.0,
+            'gamma2': 0.0,
+            'ra_0':0.0, 
+            'dec_0':0.0
+        }
+    },
+    'source':{
+        'class': SingleSersicSource,
+        'parameters':{
+            'magnitude': 10.0,
+            'output_ab_zeropoint': output_ab_zeropoint,
+            'R_sersic': 1.5,
+            'n_sersic':3.0 ,
+            'e1': (1.0-0.9)/(1.0+0.9),
+            'e2': 0.0,
+            'center_x': 0.16,
+            'center_y': -0.08,
+            'z_source':1.5}
+    },
+    'cosmology':{
+        'parameters':{
+            'cosmology_name': 'planck18'
+        }
+    },
+    'psf':{
+        'parameters':{
+            'psf_type':'PIXEL',
+            'kernel_point_source': psf_pix_map,
+            'point_source_supersampling_factor':2
+        }
+    },
+    'detector':{
+        'parameters':{
+            'pixel_scale':0.040,'ccd_gain':1.58,'read_noise':3.0,
+            'magnitude_zero_point':output_ab_zeropoint,
+            'exposure_time':1380,'sky_brightness':21.83,
+            'num_exposures':1,'background_noise':None
+        }
+    }
+}
diff --git a/notebooks/comparison_files/input_config_paltas_zero.py b/notebooks/comparison_files/input_config_paltas_zero.py
new file mode 100644
index 0000000..fd62454
--- /dev/null
+++ b/notebooks/comparison_files/input_config_paltas_zero.py
@@ -0,0 +1,103 @@
+# Configuration for a tight spread on main deflector with PEMD + SHEAR, sources
+# drawn from COSMOS, and only varying the d_los and sigma_sub of the DG_19
+# subhalo and los classes
+
+import numpy as np
+from scipy.stats import norm, truncnorm, uniform
+from paltas.Substructure.los_dg19 import LOSDG19
+from paltas.Substructure.subhalos_dg19 import SubhalosDG19
+from paltas.MainDeflector.simple_deflectors import PEMDShear
+from paltas.Sources.sersic import SingleSersicSource
+from lenstronomy.Util.kernel_util import degrade_kernel
+from paltas.Sampling import distributions
+from astropy.io import fits
+import pandas as pd
+import paltas
+import os
+
+# Define the numerics kwargs.
+kwargs_numerics = {'supersampling_factor':2,'supersampling_convolution':True}
+# We do not use point_source_supersampling_factor but it must be passed in to
+# surpress a warning.
+kwargs_numerics['point_source_supersampling_factor'] = (
+    kwargs_numerics['supersampling_factor'])
+# This is always the number of pixels for the CCD. If drizzle is used, the
+# final image will be larger.
+numpix = 128
+
+# Define some general image kwargs for the dataset
+mask_radius = 0.0
+mag_cut = 2.0
+
+# Define arguments that will be used multiple times
+output_ab_zeropoint = 25.127
+
+config_dict = {
+    'subhalo':{
+        'class': SubhalosDG19,
+        'parameters':{
+            'sigma_sub':0.0e-3,
+            'shmf_plaw_index':-1.98,
+            'm_pivot': 1e10,
+            'm_min': 7e7,
+            'm_max': 1e10,
+            'c_0':18,
+            'conc_zeta':-0.2,
+            'conc_beta':0.8,
+            'conc_m_ref': 1e8,
+            'k1':0.0,
+            'k2':0.0,
+            'dex_scatter': 1.0
+        }
+    },
+    'main_deflector':{
+        'class': PEMDShear,
+        'parameters':{
+            'M200': 1e13,
+            'z_lens': 0.5,
+            'gamma': 2.0,
+            'theta_E': 1.1,
+            'e1': (1.0-0.9)/(1.0+0.9),
+            'e2': 0.0,
+            'center_x': 0.08,
+            'center_y': -0.16,
+            'gamma1': 0.0,
+            'gamma2': 0.0,
+            'ra_0':0.0, 
+            'dec_0':0.0
+        }
+    },
+    'source':{
+        'class': SingleSersicSource,
+        'parameters':{
+            'magnitude': -29.027384789, # Equivalent to an amplitude of 10.
+            'output_ab_zeropoint': output_ab_zeropoint,
+            'R_sersic': 1.5,
+            'n_sersic':3.0 ,
+            'e1': (1.0-0.9)/(1.0+0.9),
+            'e2': 0.0,
+            'center_x': 0.16,
+            'center_y': -0.08,
+            'z_source':1.5}
+    },
+    'cosmology':{
+        'parameters':{
+            'cosmology_name': 'planck18'
+        }
+    },
+    'psf':{
+        'parameters':{
+            'psf_type':'GAUSSIAN',
+            'fwhm': 0.04,
+            'point_source_supersampling_factor':2
+        }
+    },
+    'detector':{
+        'parameters':{
+            'pixel_scale':0.040,'ccd_gain':1.58,'read_noise':3.0,
+            'magnitude_zero_point':output_ab_zeropoint,
+            'exposure_time':1380,'sky_brightness':21.83,
+            'num_exposures':1,'background_noise':None
+        }
+    }
+}
diff --git a/requirements.txt b/requirements.txt
index 13ea088..bffdddb 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,3 +3,9 @@ chex
 dm-pix
 immutabledict
 jax
+tensorflow
+immutabledict
+clu
+flax==0.6.2
+ml_collections
+optax


Test patch:
diff --git a/jaxstronomy/cosmology_utils_test.py b/jaxstronomy/cosmology_utils_test.py
index 57a63ce..f7d6d3c 100644
--- a/jaxstronomy/cosmology_utils_test.py
+++ b/jaxstronomy/cosmology_utils_test.py
@@ -1,19 +1,4 @@
 # coding=utf-8
-# Copyright 2022 The Google Research Authors.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Copyright 2022 Google LLC
 
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -44,465 +29,516 @@
 # Using immutabledict is good practice, although we will have to cast to dict
 # everywhere because jax does not accept immutabledict.
 COSMOLOGY_PARAMS_INIT = immutabledict({
-    'omega_m_zero': 0.3089,
-    'omega_b_zero': 0.0486,
-    'omega_de_zero': 0.6910088292453472,
-    'omega_rad_zero': 9.117075466e-5,
-    'temp_cmb_zero': 2.7255,
-    'hubble_constant': 67.74,
-    'n_s': 0.9667,
-    'sigma_eight': 0.8159,
+        'omega_m_zero': 0.3089,
+        'omega_b_zero': 0.0486,
+        'omega_de_zero': 0.6910088292453472,
+        'omega_rad_zero': 9.117075466e-5,
+        'temp_cmb_zero': 2.7255,
+        'hubble_constant': 67.74,
+        'n_s': 0.9667,
+        'sigma_eight': 0.8159,
 })
 COSMOLOGY_PARAMS_LENSTRONOMY = immutabledict({
-    'omega_m_zero': 0.30966,
-    'omega_b_zero': 0.0486,
-    'omega_de_zero': 0.688846306,
-    'omega_rad_zero': 0.0014937,
-    'temp_cmb_zero': 2.7255,
-    'hubble_constant': 67.66,
-    'n_s': 0.9667,
-    'sigma_eight': 0.8159,
+        'omega_m_zero': 0.30966,
+        'omega_b_zero': 0.0486,
+        'omega_de_zero': 0.688846306,
+        'omega_rad_zero': 0.0014937,
+        'temp_cmb_zero': 2.7255,
+        'hubble_constant': 67.66,
+        'n_s': 0.9667,
+        'sigma_eight': 0.8159,
 })
 
 
 def _prepare_cosmology_params(
-    cosmology_params_init,
-    z_lookup_max,
-    dz,
-    r_min=1e-4,
-    r_max=1e3,
-    n_r_bins=2,
+        cosmology_params_init,
+        z_lookup_max,
+        dz,
+        r_min=1e-4,
+        r_max=1e3,
+        n_r_bins=2,
 ):
-  # Only generate a lookup table for values we need.
-  # When 0,0 is specified for the two z values, need to select a small non-zero
-  # values to generate a non-empty table.
-  z_lookup_max = max(z_lookup_max, 1e-7)
-  dz = max(dz, 1e-7)
-  return cosmology_utils.add_lookup_tables_to_cosmology_params(
-      dict(cosmology_params_init), z_lookup_max, dz, r_min, r_max, n_r_bins)
+    # Only generate a lookup table for values we need.
+    # When 0,0 is specified for the two z values, need to select a small non-zero
+    # values to generate a non-empty table.
+    z_lookup_max = max(z_lookup_max, 1e-7)
+    dz = max(dz, 1e-7)
+    return cosmology_utils.add_lookup_tables_to_cosmology_params(
+            dict(cosmology_params_init), z_lookup_max, dz, r_min, r_max, n_r_bins)
 
 
 def _prepare_x_y():
-  rng = jax.random.PRNGKey(3)
-  rng_x, rng_y = jax.random.split(rng)
-  x = jax.random.normal(rng_x, shape=(3,)) * 1e3
-  y = jax.random.normal(rng_y, shape=(3,)) * 1e3
-  return x, y
+    rng = jax.random.PRNGKey(3)
+    rng_x, rng_y = jax.random.split(rng)
+    x = jax.random.normal(rng_x, shape=(3,)) * 1e3
+    y = jax.random.normal(rng_y, shape=(3,)) * 1e3
+    return x, y
 
 
 def _prepare_alpha_x_alpha_y():
-  rng = jax.random.PRNGKey(2)
-  rng_x, rng_y = jax.random.split(rng)
-  alpha_x = jax.random.normal(rng_x, shape=(3,)) * 2
-  alpha_y = jax.random.normal(rng_y, shape=(3,)) * 2
-  return alpha_x, alpha_y
+    rng = jax.random.PRNGKey(2)
+    rng_x, rng_y = jax.random.split(rng)
+    alpha_x = jax.random.normal(rng_x, shape=(3,)) * 2
+    alpha_y = jax.random.normal(rng_y, shape=(3,)) * 2
+    return alpha_x, alpha_y
 
 
 def _prepare__sigma_k_integrand_expected(radius):
-  if radius == 1e-2:
-    return jnp.array([
-        3.127437136035467e-06, 1.820751207237299e-05, 5.2541055609899326e-05,
-        0.00010306278738724287, 0.00010775093864571042
-    ])
-  elif radius == 1e-1:
-    return jnp.array([
-        3.127323720477859e-06, 1.8185657328033008e-05, 5.04866815266023e-05,
-        2.3545351279405076e-05, 1.5387621097147294e-08
-    ])
-  elif radius == 1e0:
-    return jnp.array([
-        3.116000043964669e-06, 1.611047820558932e-05, 1.534716000730611e-10,
-        1.2524276267847e-09, 1.6400363103856087e-12
-    ])
-  else:
-    raise ValueError(f'Unsupported radius = {radius}')
+    if radius == 1e-2:
+        return jnp.array([
+                3.127437136035467e-06, 1.820751207237299e-05, 5.2541055609899326e-05,
+                0.00010306278738724287, 0.00010775093864571042
+        ])
+    elif radius == 1e-1:
+        return jnp.array([
+                3.127323720477859e-06, 1.8185657328033008e-05, 5.04866815266023e-05,
+                2.3545351279405076e-05, 1.5387621097147294e-08
+        ])
+    elif radius == 1e0:
+        return jnp.array([
+                3.116000043964669e-06, 1.611047820558932e-05, 1.534716000730611e-10,
+                1.2524276267847e-09, 1.6400363103856087e-12
+        ])
+    else:
+        raise ValueError(f'Unsupported radius = {radius}')
 
 
 def _prepare_reduced_to_physical_expected(z_lens, z_source):
-  if z_lens == 0.1 and z_source == 0.5:
-    return jnp.array([-0.6492834, 4.845596, -3.8608627])
-  elif z_lens == 0.5 and z_source == 1.0:
-    return jnp.array([-1.1832292, 8.830428, -7.035888])
-  elif z_lens == 0.5 and z_source == 2.0:
-    return jnp.array([-0.7973697, 5.9507623, -4.7414346])
-  else:
-    raise ValueError(f'Unsupported z_lens and z_source = ({z_lens},{z_source})')
+    if z_lens == 0.1 and z_source == 0.5:
+        return jnp.array([-0.6492834, 4.845596, -3.8608627])
+    elif z_lens == 0.5 and z_source == 1.0:
+        return jnp.array([-1.1832292, 8.830428, -7.035888])
+    elif z_lens == 0.5 and z_source == 2.0:
+        return jnp.array([-0.7973697, 5.9507623, -4.7414346])
+    else:
+        raise ValueError(f'Unsupported z_lens and z_source = ({z_lens},{z_source})')
 
 
 def _preapre_comoving_to_angle_expected(z_lens):
-  if z_lens == 0.1:
-    return jnp.array([[-3.1144106, -1.2556583, 0.8742881],
-                      [1.524639, -4.0299582, 3.8287995]])
-  elif z_lens == 0.5:
-    return jnp.array([[-0.6921371, -0.2790537, 0.19429912],
-                      [0.33883113, -0.8956057, 0.85090077]])
-  elif z_lens == 1.0:
-    return jnp.array([[-0.39674094, -0.15995677, 0.11137448],
-                      [0.19422188, -0.51337147, 0.48774606]])
+    if z_lens == 0.1:
+        return jnp.array([[-3.1144106, -1.2556583, 0.8742881],
+                                            [1.524639, -4.0299582, 3.8287995]])
+    elif z_lens == 0.5:
+        return jnp.array([[-0.6921371, -0.2790537, 0.19429912],
+                                            [0.33883113, -0.8956057, 0.85090077]])
+    elif z_lens == 1.0:
+        return jnp.array([[-0.39674094, -0.15995677, 0.11137448],
+                                            [0.19422188, -0.51337147, 0.48774606]])
 
 
 def _prepare_expected_comoving_distance(z_min, z_max):
-  if z_min == 0.0 and z_max == 0.0:
-    expected = 0.0
-  elif z_min == 0.0 and z_max == 1.283:
-    expected = 2.740512775683e3
-  elif z_min == 0.0 and z_max == 20.0:
-    expected = 7.432211484621e3
-  elif z_min == 0.5 and z_max == 1.283:
-    expected = 1.4225931620767858e3
-  else:
-    raise ValueError(f'Unsupported z_min and z_max = ({z_min},{z_max})')
+    if z_min == 0.0 and z_max == 0.0:
+        expected = 0.0
+    elif z_min == 0.0 and z_max == 1.283:
+        expected = 2.740512775683e3
+    elif z_min == 0.0 and z_max == 20.0:
+        expected = 7.432211484621e3
+    elif z_min == 0.5 and z_max == 1.283:
+        expected = 1.4225931620767858e3
+    else:
+        raise ValueError(f'Unsupported z_min and z_max = ({z_min},{z_max})')
 
-  # Expected is in units of /h
-  return expected / (COSMOLOGY_PARAMS_INIT['hubble_constant'] / 100.)
+    # Expected is in units of /h
+    return expected / (COSMOLOGY_PARAMS_INIT['hubble_constant'] / 100.)
 
 
 def _prepare_expected_angular_diameter_distance(z):
-  return _prepare_expected_comoving_distance(0.0, z) / (1 + z)
+    return _prepare_expected_comoving_distance(0.0, z) / (1 + z)
 
 
 def _prepare_expected_kpc_per_arcsecond(z):
-  if z == 0.0:
-    return 0.0
-  elif z == 1.283:
-    return 8.591241981530676
-  elif z == 20.0:
-    return 2.5329606272180145
-  else:
-    raise ValueError(f'Unsupported z = {z}')
+    if z == 0.0:
+        return 0.0
+    elif z == 1.283:
+        return 8.591241981530676
+    elif z == 20.0:
+        return 2.5329606272180145
+    else:
+        raise ValueError(f'Unsupported z = {z}')
 
 
 class CosmologyUtilsTest(chex.TestCase, parameterized.TestCase):
-  """Runs tests of various cosmology functions."""
-
-  @chex.all_variants
-  def test__e_z(self):
-    z_values = jnp.array([0.0, 1.283, 20.0])
-    expected = jnp.array([1.0, 2.090250748388e+00, 5.365766383626e+01])
-
-    e_z = self.variant(cosmology_utils._e_z)
-
-    np.testing.assert_allclose(
-        e_z(dict(COSMOLOGY_PARAMS_INIT), z_values), expected)
-
-  @chex.all_variants
-  def test__e_z_rad_to_dark(self):
-    z_values = jnp.array([0.0, 1.283, 20.0])
-    expected = jnp.array([1.0, 2.0896800302592498, 53.49218634529721])
-
-    e_z_rad_to_dark = self.variant(cosmology_utils._e_z_rad_to_dark)
-
-    np.testing.assert_allclose(
-        e_z_rad_to_dark(dict(COSMOLOGY_PARAMS_INIT), z_values), expected)
-
-  @chex.all_variants
-  @parameterized.named_parameters([
-      (f'_z_min_{z_min}_z_max_{z_max}', z_min, z_max)
-      for z_min, z_max in zip([0.0, 0.0, 0.0, 0.5], [0.0, 1.283, 20.0, 1.283])
-  ])
-  def test__comoving_distance_numerical(self, z_min, z_max):
-    expected = _prepare_expected_comoving_distance(z_min, z_max)
-    comoving_distance = self.variant(
-        cosmology_utils._comoving_distance_numerical)
-
-    np.testing.assert_allclose(
-        comoving_distance(dict(COSMOLOGY_PARAMS_INIT), z_min, z_max),
-        expected,
-        rtol=1e-4)
-
-  @chex.all_variants
-  @parameterized.named_parameters([(f'_r_{r}', r) for r in [1e-2, 1e-1, 1e0]])
-  def test__sigma_k_integrand(self, r):
-    log_k = jnp.linspace(-2, 5, 5)
-    expected = _prepare__sigma_k_integrand_expected(r)
-
-    sigma_k_integrand = self.variant(cosmology_utils._sigma_k_integrand)
-
-    np.testing.assert_allclose(
-        sigma_k_integrand(dict(COSMOLOGY_PARAMS_INIT), log_k, r),
-        expected,
-        rtol=1e-4)
-
-  @chex.all_variants
-  @parameterized.named_parameters([
-      (f'_r_{r}', r, expected)
-      for r, expected in zip([1e-2, 1e-1, 1e0], [0.004892, 0.002956, 0.001396])
-  ])
-  def test__sigma_numerical(self, r, expected):
-    sigma_numerical = self.variant(cosmology_utils._sigma_numerical)
-
-    np.testing.assert_allclose(
-        sigma_numerical(dict(COSMOLOGY_PARAMS_INIT), r), expected, rtol=1e-3)
-
-  @chex.all_variants
-  @parameterized.named_parameters([
-      (f'_se_{se}', se, expected) for se, expected in zip(
-          [0.82, 0.88], [8.60876269756018, 9.238672163235316])
-  ])
-  def test__sigma_norm(self, se, expected):
-    # The comparison is to the normalzied sigma_numerical calculation.
-    radius = 1e-2
-    cosmology_params = dict(COSMOLOGY_PARAMS_INIT)
-    cosmology_params['sigma_eight'] = se
-    sigma_norm = self.variant(cosmology_utils._sigma_norm)
-    sigma_numerical = self.variant(cosmology_utils._sigma_numerical)
-
-    np.testing.assert_allclose((sigma_norm(cosmology_params) *
-                                sigma_numerical(cosmology_params, radius)),
-                               expected,
-                               rtol=1e-3)
-
-  @chex.all_variants
-  @parameterized.named_parameters([
-      (f'_z_{z}', z, expected) for z, expected in zip(
-          [0.0, 0.1, 0.2, 0.3], [0.7843, 0.7442, 0.7060, 0.6698])
-  ])
-  def test___growth_factor_exact_unormalized(self, z, expected):
-    growth_factor = self.variant(
-        cosmology_utils._growth_factor_exact_unormalized)
-
-    self.assertAlmostEqual(
-        growth_factor(dict(COSMOLOGY_PARAMS_INIT), z), expected, places=3)
-
-  @parameterized.named_parameters([
-      (f'_z_l_max_{z_l_max}_dz_{dz}_r_max_{r_max}', z_l_max, dz, r_max)
-      for z_l_max, dz, r_max in zip([1.01, 2.02], [0.01, 0.02], [1e-2, 1e1])
-  ])
-  def test_add_lookup_tables_to_cosmology_params(self, z_l_max, dz, r_max):
-    cosmology_params = (
-        cosmology_utils.add_lookup_tables_to_cosmology_params(
-            dict(COSMOLOGY_PARAMS_INIT), z_l_max, dz, 1e-3, r_max, 2))
-
-    # For equivalent precision must also be jitted.
-    comoving_distance_numerical = jax.jit(
-        cosmology_utils._comoving_distance_numerical)
-    sigma_numerical = jax.jit(cosmology_utils._sigma_numerical)
-    sigma_norm = jax.jit(cosmology_utils._sigma_norm)
-    growth_factor = jax.jit(cosmology_utils._growth_factor_exact_unormalized)
-
-    # Tests for comoving distance lookup table.
-
-    # Diagonal should always be zero
-    np.testing.assert_allclose(
-        jnp.diag(cosmology_params['comoving_lookup_table']),
-        jnp.zeros(len(cosmology_params['comoving_lookup_table'])),
-        rtol=0,
-        atol=1e-7)
-
-    # Test a few values. Low precision from float 32.
-    i = 0
-    j = 30
-    self.assertAlmostEqual(
-        cosmology_params['comoving_lookup_table'][i, j],
-        comoving_distance_numerical(
-            dict(COSMOLOGY_PARAMS_INIT), dz * i, dz * j),
-        places=3)
-    i = 12
-    j = 37
-    self.assertAlmostEqual(
-        cosmology_params['comoving_lookup_table'][i, j],
-        comoving_distance_numerical(cosmology_params, dz * i, dz * j),
-        places=3)
-
-    # Check z_min > z_max
-    i = 42
-    j = 15
-    self.assertAlmostEqual(
-        cosmology_params['comoving_lookup_table'][i, j],
-        comoving_distance_numerical(cosmology_params, dz * i, dz * j),
-        places=3)
-
-    # Test sigma lookup table.
-    self.assertAlmostEqual(
-        cosmology_params['sigma_lookup_table'][1],
-        sigma_numerical(cosmology_params, r_max) * sigma_norm(cosmology_params),
-        places=3)
-
-    # Test growth factor.
-    i = 12
-    growth_norm = 1.0 / growth_factor(cosmology_params, 0.0)
-    self.assertAlmostEqual(
-        cosmology_params['growth_lookup_table'][i],
-        growth_factor(cosmology_params, dz * i) * growth_norm,
-        places=3)
-
-    i = 42
-    self.assertAlmostEqual(
-        cosmology_params['growth_lookup_table'][i],
-        growth_factor(cosmology_params, dz * i) * growth_norm,
-        places=3)
-
-  @chex.all_variants
-  @parameterized.named_parameters([
-      (f'_z_min_{z_min}_z_max_{z_max}', z_min, z_max)
-      for z_min, z_max in zip([0.1110, 0.2102], [0.2001, 0.3004])
-  ])
-  def test_comoving_distance(self, z_min, z_max):
-    # Test that inerpolation returns approximately the correct result even when
-    # the sampling doesn't fall perfectly on the grid.
-    cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, 0.31,
-                                                 0.01)
-
-    comoving_distance = self.variant(cosmology_utils.comoving_distance)
-    comoving_distance_numerical = jax.jit(
-        cosmology_utils._comoving_distance_numerical)
-
-    np.testing.assert_allclose(
-        comoving_distance(cosmology_params, z_min, z_max),
-        comoving_distance_numerical(cosmology_params, z_min, z_max),
-        rtol=1e-4)
-
-  @chex.all_variants
-  @parameterized.named_parameters([(f'_z_{z}', z) for z in [0.0, 1.283, 20.0]])
-  def test_angular_diameter_distance(self, z):
-    cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z, z)
-    expected = _prepare_expected_angular_diameter_distance(z)
-    angular_diameter_distance = self.variant(
-        cosmology_utils.angular_diameter_distance)
-
-    np.testing.assert_allclose(
-        angular_diameter_distance(cosmology_params, z), expected, rtol=1e-4)
-
-  @chex.all_variants
-  @parameterized.named_parameters([
-      (f'_z_min_{z_min}_z_max_{z_max}', z_min, z_max)
-      for z_min, z_max in zip([0.0, 0.0, 0.0], [0.0, 1.283, 20.0])
-  ])
-  def test_angular_diameter_distance_between(self, z_min, z_max):
-    cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z_max,
-                                                 z_max)
-    expected = _prepare_expected_comoving_distance(z_min, z_max) / (1 + z_max)
-    angular_diameter_distance_between = self.variant(
-        cosmology_utils.angular_diameter_distance_between)
-
-    np.testing.assert_allclose(
-        angular_diameter_distance_between(cosmology_params, z_min, z_max),
-        expected,
-        rtol=1e-4)
-
-  @chex.all_variants
-  @parameterized.named_parameters([(f'_z_{z}', z) for z in [0.0, 1.283, 20.0]])
-  def test_kpc_per_arcsecond(self, z):
-    cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z, z)
-    expected = _prepare_expected_kpc_per_arcsecond(z)
-    kpc_per_arcsecond = self.variant(cosmology_utils.kpc_per_arcsecond)
-
-    np.testing.assert_allclose(
-        kpc_per_arcsecond(cosmology_params, z), expected, rtol=1e-4)
-
-  @chex.all_variants
-  @parameterized.named_parameters([
-      (f'_z_lens_{z_lens}', z_lens) for z_lens in [0.1, 0.5, 1.0]
-  ])
-  def test_comoving_to_angle(self, z_lens):
-    x, y = _prepare_x_y()
-    cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_LENSTRONOMY,
-                                                 z_lens, z_lens)
-    expected = _preapre_comoving_to_angle_expected(z_lens)
-
-    comoving_to_angle = self.variant(cosmology_utils.comoving_to_angle)
-
-    np.testing.assert_allclose(
-        jnp.array(comoving_to_angle(x, y, cosmology_params, z_lens)),
-        expected,
-        rtol=1e-3)
-
-  @chex.all_variants
-  @parameterized.named_parameters([
-      (f'_z_l_{z_lens}_z_s_{z_source}', z_lens, z_source)
-      for z_lens, z_source in zip([0.1, 0.5, 0.5], [0.5, 1.0, 2.0])
-  ])
-  def test_reduced_to_physical(self, z_lens, z_source):
-    cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_LENSTRONOMY,
-                                                 z_source, z_lens)
-    reduced, _ = _prepare_alpha_x_alpha_y()
-    expected = _prepare_reduced_to_physical_expected(z_lens, z_source)
-
-    reduced_to_physical = self.variant(cosmology_utils.reduced_to_physical)
-
-    # Division and multiplication by two quantities with errors in 1e-6, so need
-    # to increase rtol for this test significantly.
-    np.testing.assert_allclose(
-        reduced_to_physical(reduced, cosmology_params, z_lens, z_source),
-        expected,
-        rtol=1e-3)
-
-  @chex.all_variants
-  @parameterized.named_parameters([
-      (f'_omega_m_{omega_m_zero}_z_{z}', omega_m_zero, z, expected)
-      for omega_m_zero, z, expected in zip([0.2, 0.3, 0.4], [0.2, 0.4, 1.2],
-                                           [95.917, 228.468, 1182.084])
-  ])
-  def test_rho_matter(self, omega_m_zero, z, expected):
-    rho_matter = self.variant(cosmology_utils.rho_matter)
-
-    self.assertAlmostEqual(
-        rho_matter({'omega_m_zero': omega_m_zero}, z) / 1e9, expected, places=3)
-
-  @chex.all_variants
-  @parameterized.named_parameters([
-      (f'_z_{z}', z, expected)
-      for z, expected in zip([0.2, 0.3, 0.4], [1.6788, 1.6800, 1.6809])
-  ])
-  def test_collapse_overdensity(self, z, expected):
-    cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z, z)
-    collapse_overdensity = self.variant(cosmology_utils.collapse_overdensity)
-
-    self.assertAlmostEqual(
-        collapse_overdensity(cosmology_params, z), expected, places=3)
-
-  @chex.all_variants
-  @parameterized.named_parameters([
-      (f'_m_log10_{m_log10}', m_log10, expected)
-      for m_log10, expected in zip([7, 8, 9], [0.0303, 0.06530, 0.14069])
-  ])
-  def test_lagrangian_radius(self, m_log10, expected):
-    cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, 0.1,
-                                                 0.1)
-    lagrangian_radius = self.variant(cosmology_utils.lagrangian_radius)
-
-    self.assertAlmostEqual(
-        lagrangian_radius(cosmology_params, 10**m_log10), expected, places=3)
-
-  @chex.all_variants
-  @parameterized.named_parameters([
-      (f'_r_{r}_z_{z}', r, z, expected) for r, z, expected in zip(
-          [1e-2, 1e-1, 1e0], [0.1, 0.2, 0.3],
-          [8.128724384821487, 4.660020562550168, 2.08730668354650])
-  ])
-  def test_sigma_tophat(self, r, z, expected):
-    # Minimal lookup tables to reproduce desired results.
-    cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z, z,
-                                                 1e-3, r, 2)
-
-    sigma_tophat = self.variant(cosmology_utils.sigma_tophat)
-
-    # Test that inerpolation behaves as expected.
-    self.assertAlmostEqual(
-        sigma_tophat(cosmology_params, r, z), expected, places=2)
-    self.assertAlmostEqual(
-        sigma_tophat(cosmology_params, r + 1e-5, z), expected, places=2)
-    self.assertAlmostEqual(
-        sigma_tophat(cosmology_params, r - 1e-5, z), expected, places=2)
-
-  @chex.all_variants
-  @parameterized.named_parameters([
-      (f'_r_{r}_z_{z}', r, z, expected) for r, z, expected in zip(
-          [1e-2, 1e-1, 1e0], [0.1, 0.2, 0.3],
-          [-0.1840887668734233, -0.26100267465205956, -0.4077157354139687])
-  ])
-  def test_derivative_log_sigma_log_r(self, r, z, expected):
-    # Minimal lookup tables to reproduce desired results.
-    dr = 1e-4
-    cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z, z,
-                                                 r - 2 * dr, r + 2 * dr, 5)
-
-    derivative_log_sigma_log_r = self.variant(
-        cosmology_utils.derivative_log_sigma_log_r)
-
-    self.assertAlmostEqual(
-        derivative_log_sigma_log_r(cosmology_params, r, z), expected, places=2)
+    """Runs tests of various cosmology functions."""
+
+    @chex.all_variants
+    def test__e_z(self):
+        z_values = jnp.array([0.0, 1.283, 20.0])
+        expected = jnp.array([1.0, 2.090250748388e+00, 5.365766383626e+01])
+
+        e_z = self.variant(cosmology_utils._e_z)
+
+        np.testing.assert_allclose(
+                e_z(dict(COSMOLOGY_PARAMS_INIT), z_values), expected)
+
+    @chex.all_variants
+    def test__e_z_rad_to_dark(self):
+        z_values = jnp.array([0.0, 1.283, 20.0])
+        expected = jnp.array([1.0, 2.0896800302592498, 53.49218634529721])
+
+        e_z_rad_to_dark = self.variant(cosmology_utils._e_z_rad_to_dark)
+
+        np.testing.assert_allclose(
+                e_z_rad_to_dark(dict(COSMOLOGY_PARAMS_INIT), z_values), expected)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+            (f'_z_min_{z_min}_z_max_{z_max}', z_min, z_max)
+            for z_min, z_max in zip([0.0, 0.0, 0.0, 0.5], [0.0, 1.283, 20.0, 1.283])
+    ])
+    def test__comoving_distance_numerical(self, z_min, z_max):
+        expected = _prepare_expected_comoving_distance(z_min, z_max)
+        comoving_distance = self.variant(
+                cosmology_utils._comoving_distance_numerical)
+
+        np.testing.assert_allclose(
+                comoving_distance(dict(COSMOLOGY_PARAMS_INIT), z_min, z_max),
+                expected,
+                rtol=1e-4)
+
+    @chex.all_variants
+    @parameterized.named_parameters([(f'_r_{r}', r) for r in [1e-2, 1e-1, 1e0]])
+    def test__sigma_k_integrand(self, r):
+        log_k = jnp.linspace(-2, 5, 5)
+        expected = _prepare__sigma_k_integrand_expected(r)
+
+        sigma_k_integrand = self.variant(cosmology_utils._sigma_k_integrand)
+
+        np.testing.assert_allclose(
+                sigma_k_integrand(dict(COSMOLOGY_PARAMS_INIT), log_k, r),
+                expected,
+                rtol=1e-4)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+            (f'_r_{r}', r, expected)
+            for r, expected in zip([1e-2, 1e-1, 1e0], [0.004892, 0.002956, 0.001396])
+    ])
+    def test__sigma_numerical(self, r, expected):
+        sigma_numerical = self.variant(cosmology_utils._sigma_numerical)
+
+        np.testing.assert_allclose(
+                sigma_numerical(dict(COSMOLOGY_PARAMS_INIT), r), expected, rtol=1e-3)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+            (f'_se_{se}', se, expected) for se, expected in zip(
+                    [0.82, 0.88], [8.60876269756018, 9.238672163235316])
+    ])
+    def test__sigma_norm(self, se, expected):
+        # The comparison is to the normalzied sigma_numerical calculation.
+        radius = 1e-2
+        cosmology_params = dict(COSMOLOGY_PARAMS_INIT)
+        cosmology_params['sigma_eight'] = se
+        sigma_norm = self.variant(cosmology_utils._sigma_norm)
+        sigma_numerical = self.variant(cosmology_utils._sigma_numerical)
+
+        np.testing.assert_allclose((sigma_norm(cosmology_params) *
+                                                                sigma_numerical(cosmology_params, radius)),
+                                                             expected,
+                                                             rtol=1e-3)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+            (f'_z_{z}', z, expected) for z, expected in zip(
+                    [0.0, 0.1, 0.2, 0.3], [0.7843, 0.7442, 0.7060, 0.6698])
+    ])
+    def test___growth_factor_exact_unormalized(self, z, expected):
+        growth_factor = self.variant(
+                cosmology_utils._growth_factor_exact_unormalized)
+
+        self.assertAlmostEqual(
+                growth_factor(dict(COSMOLOGY_PARAMS_INIT), z), expected, places=3)
+
+    @parameterized.named_parameters([
+            (f'_z_l_max_{z_l_max}_dz_{dz}_r_max_{r_max}', z_l_max, dz, r_max)
+            for z_l_max, dz, r_max in zip([1.01, 2.02], [0.01, 0.02], [1e-2, 1e1])
+    ])
+    def test_add_lookup_tables_to_cosmology_params(self, z_l_max, dz, r_max):
+        cosmology_params = (
+                cosmology_utils.add_lookup_tables_to_cosmology_params(
+                        dict(COSMOLOGY_PARAMS_INIT), z_l_max, dz, 1e-3, r_max, 2))
+
+        # For equivalent precision must also be jitted.
+        comoving_distance_numerical = jax.jit(
+                cosmology_utils._comoving_distance_numerical)
+        sigma_numerical = jax.jit(cosmology_utils._sigma_numerical)
+        sigma_norm = jax.jit(cosmology_utils._sigma_norm)
+        growth_factor = jax.jit(cosmology_utils._growth_factor_exact_unormalized)
+
+        # Tests for comoving distance lookup table.
+
+        # Diagonal should always be zero
+        np.testing.assert_allclose(
+                jnp.diag(cosmology_params['comoving_lookup_table']),
+                jnp.zeros(len(cosmology_params['comoving_lookup_table'])),
+                rtol=0,
+                atol=1e-7)
+
+        # Test a few values. Low precision from float 32.
+        i = 0
+        j = 30
+        self.assertAlmostEqual(
+                cosmology_params['comoving_lookup_table'][i, j],
+                comoving_distance_numerical(
+                        dict(COSMOLOGY_PARAMS_INIT), dz * i, dz * j),
+                places=3)
+        i = 12
+        j = 37
+        self.assertAlmostEqual(
+                cosmology_params['comoving_lookup_table'][i, j],
+                comoving_distance_numerical(cosmology_params, dz * i, dz * j),
+                places=3)
+
+        # Check z_min > z_max
+        i = 42
+        j = 15
+        self.assertAlmostEqual(
+                cosmology_params['comoving_lookup_table'][i, j],
+                comoving_distance_numerical(cosmology_params, dz * i, dz * j),
+                places=3)
+
+        # Test sigma lookup table.
+        self.assertAlmostEqual(
+                cosmology_params['sigma_lookup_table'][1],
+                sigma_numerical(cosmology_params, r_max) * sigma_norm(cosmology_params),
+                places=3)
+
+        # Test growth factor.
+        i = 12
+        growth_norm = 1.0 / growth_factor(cosmology_params, 0.0)
+        self.assertAlmostEqual(
+                cosmology_params['growth_lookup_table'][i],
+                growth_factor(cosmology_params, dz * i) * growth_norm,
+                places=3)
+
+        i = 42
+        self.assertAlmostEqual(
+                cosmology_params['growth_lookup_table'][i],
+                growth_factor(cosmology_params, dz * i) * growth_norm,
+                places=3)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+            (f'_z_min_{z_min}_z_max_{z_max}', z_min, z_max)
+            for z_min, z_max in zip([0.1110, 0.2102], [0.2001, 0.3004])
+    ])
+    def test_comoving_distance(self, z_min, z_max):
+        # Test that inerpolation returns approximately the correct result even when
+        # the sampling doesn't fall perfectly on the grid.
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, 0.31,
+                                                                                                 0.01)
+
+        comoving_distance = self.variant(cosmology_utils.comoving_distance)
+        comoving_distance_numerical = jax.jit(
+                cosmology_utils._comoving_distance_numerical)
+
+        np.testing.assert_allclose(
+                comoving_distance(cosmology_params, z_min, z_max),
+                comoving_distance_numerical(cosmology_params, z_min, z_max),
+                rtol=1e-4)
+
+    @chex.all_variants
+    @parameterized.named_parameters([(f'_z_{z}', z) for z in [0.0, 1.283, 20.0]])
+    def test_angular_diameter_distance(self, z):
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z, z)
+        expected = _prepare_expected_angular_diameter_distance(z)
+        angular_diameter_distance = self.variant(
+                cosmology_utils.angular_diameter_distance)
+
+        np.testing.assert_allclose(
+                angular_diameter_distance(cosmology_params, z), expected, rtol=1e-4)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+            (f'_z_min_{z_min}_z_max_{z_max}', z_min, z_max)
+            for z_min, z_max in zip([0.0, 0.0, 0.0], [0.0, 1.283, 20.0])
+    ])
+    def test_angular_diameter_distance_between(self, z_min, z_max):
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z_max,
+                                                                                                 z_max)
+        expected = _prepare_expected_comoving_distance(z_min, z_max) / (1 + z_max)
+        angular_diameter_distance_between = self.variant(
+                cosmology_utils.angular_diameter_distance_between)
+
+        np.testing.assert_allclose(
+                angular_diameter_distance_between(cosmology_params, z_min, z_max),
+                expected,
+                rtol=1e-4)
+
+    @chex.all_variants
+    @parameterized.named_parameters([(f'_z_{z}', z) for z in [0.0, 1.283, 20.0]])
+    def test_kpc_per_arcsecond(self, z):
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z, z)
+        expected = _prepare_expected_kpc_per_arcsecond(z)
+        kpc_per_arcsecond = self.variant(cosmology_utils.kpc_per_arcsecond)
+
+        np.testing.assert_allclose(
+                kpc_per_arcsecond(cosmology_params, z), expected, rtol=1e-4)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+            (f'_z_lens_{z_lens}', z_lens) for z_lens in [0.1, 0.5, 1.0]
+    ])
+    def test_comoving_to_angle(self, z_lens):
+        x, y = _prepare_x_y()
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_LENSTRONOMY,
+                                                                                                 z_lens, z_lens)
+        expected = _preapre_comoving_to_angle_expected(z_lens)
+
+        comoving_to_angle = self.variant(cosmology_utils.comoving_to_angle)
+
+        np.testing.assert_allclose(
+                jnp.array(comoving_to_angle(x, y, cosmology_params, z_lens)),
+                expected,
+                rtol=1e-3)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+            (f'_z_l_{z_lens}_z_s_{z_source}', z_lens, z_source)
+            for z_lens, z_source in zip([0.1, 0.5, 0.5], [0.5, 1.0, 2.0])
+    ])
+    def test_reduced_to_physical(self, z_lens, z_source):
+        cosmology_params = _prepare_cosmology_params(
+                COSMOLOGY_PARAMS_LENSTRONOMY, z_source, z_lens)
+        reduced, _ = _prepare_alpha_x_alpha_y()
+        expected = _prepare_reduced_to_physical_expected(z_lens, z_source)
+
+        reduced_to_physical = self.variant(cosmology_utils.reduced_to_physical)
+
+        # Division and multiplication by two quantities with errors in 1e-6, so need
+        # to increase rtol for this test significantly.
+        np.testing.assert_allclose(
+                reduced_to_physical(reduced, cosmology_params, z_lens, z_source),
+                expected,
+                rtol=1e-3)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+            (f'_omega_m_{omega_m_zero}_z_{z}', omega_m_zero, z, expected)
+            for omega_m_zero, z, expected in zip([0.2, 0.3, 0.4],
+                [0.2, 0.4, 1.2], [95.917, 228.468, 1182.084])
+    ])
+    def test_rho_matter(self, omega_m_zero, z, expected):
+        rho_matter = self.variant(cosmology_utils.rho_matter)
+
+        self.assertAlmostEqual(
+                rho_matter({'omega_m_zero': omega_m_zero}, z) / 1e9, expected, places=3)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+            (f'_z_{z}', z, expected)
+            for z, expected in zip([0.2, 0.3, 0.4], [1.6788, 1.6800, 1.6809])
+    ])
+    def test_collapse_overdensity(self, z, expected):
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z, z)
+        collapse_overdensity = self.variant(cosmology_utils.collapse_overdensity)
+
+        self.assertAlmostEqual(
+                collapse_overdensity(cosmology_params, z), expected, places=3)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+            (f'_m_log10_{m_log10}', m_log10, expected)
+            for m_log10, expected in zip([7, 8, 9], [0.0303, 0.06530, 0.14069])
+    ])
+    def test_lagrangian_radius(self, m_log10, expected):
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, 0.1,
+                0.1)
+        lagrangian_radius = self.variant(cosmology_utils.lagrangian_radius)
+
+        self.assertAlmostEqual(
+                lagrangian_radius(cosmology_params, 10**m_log10), expected, places=3)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+            (f'_r_{r}_z_{z}', r, z, expected) for r, z, expected in zip(
+                    [1e-2, 1e-1, 1e0], [0.1, 0.2, 0.3],
+                    [8.128724384821487, 4.660020562550168, 2.08730668354650])
+    ])
+    def test_sigma_tophat(self, r, z, expected):
+        # Minimal lookup tables to reproduce desired results.
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z,
+                z, 1e-3, r, 2)
+
+        sigma_tophat = self.variant(cosmology_utils.sigma_tophat)
+
+        # Test that inerpolation behaves as expected.
+        self.assertAlmostEqual(
+                sigma_tophat(cosmology_params, r, z), expected, places=2)
+        self.assertAlmostEqual(
+                sigma_tophat(cosmology_params, r + 1e-5, z), expected, places=2)
+        self.assertAlmostEqual(
+                sigma_tophat(cosmology_params, r - 1e-5, z), expected, places=2)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+            (f'_r_{r}_z_{z}', r, z, expected) for r, z, expected in zip(
+                    [1e-2, 1e-1, 1e0], [0.1, 0.2, 0.3],
+                    [-0.1840887668734233, -0.26100267465205956,
+                        -0.4077157354139687])
+    ])
+    def test_derivative_log_sigma_log_r(self, r, z, expected):
+        # Minimal lookup tables to reproduce desired results.
+        dr = 1e-4
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z,
+                z, r - 2 * dr, r + 2 * dr, 5)
+
+        derivative_log_sigma_log_r = self.variant(
+                cosmology_utils.derivative_log_sigma_log_r)
+
+        self.assertAlmostEqual(
+                derivative_log_sigma_log_r(cosmology_params, r, z), expected,
+                        places=2)
+
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+        (f'_m_{np.log10(m):.0f}_z_{z}', m, z, expected) for m, z, expected in
+        zip([1e6, 1e7, 1e8], [0.1, 0.2, 0.3],
+                [0.22125715469476553, 0.2727592115385218, 0.3426138949041374])
+    ])
+    def test_peak_height(self, m, z, expected):
+        # Figure out what lagrangian radius we need for the peak height calculation.
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z,
+                z)
+        r = cosmology_utils.lagrangian_radius(cosmology_params, m)
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z,
+                z, 1e-3, r, 2)
+        peak_height = self.variant(cosmology_utils.peak_height)
+
+        self.assertAlmostEqual(peak_height(cosmology_params, m, z), expected,
+            places=2)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+        (f'_z_{z}', z, expected) for z, expected in zip(
+            [0.1, 0.2, 0.3],
+            [305.92535270755224, 339.97600749243884, 380.20367635428045])
+    ])
+    def test_rho_crit(self, z, expected):
+        # Test a few reference values of the critical overdensity.
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z,
+                z)
+        rho_crit = self.variant(cosmology_utils.rho_crit)
+        self.assertAlmostEqual(rho_crit(cosmology_params, z), expected,
+                places=4)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+        (f'_z_{z}_zs_{z_source}', z, z_source, expected) for
+            z, z_source, expected in zip([0.1, 0.2, 0.3], [0.2, 0.4, 0.6],
+                [8684081395.9067, 5004482667.430436, 3832776963.481256])
+    ])
+    def test_calculate_sigma_crit(self, z, z_source, expected):
+        # Test a few reference values of the critical surface overdensity.
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT,
+            z_source, 0.1)
+        calculate_sigma_crit = self.variant(
+                cosmology_utils.calculate_sigma_crit)
+        self.assertAlmostEqual(
+            calculate_sigma_crit(cosmology_params, z, z_source) / 1e9,
+                expected / 1e9, places=3)
 
 
 if __name__ == '__main__':
-  absltest.main()
+    absltest.main()
diff --git a/jaxstronomy/image_simulation_test.py b/jaxstronomy/image_simulation_test.py
index 63bbe51..a22ff52 100644
--- a/jaxstronomy/image_simulation_test.py
+++ b/jaxstronomy/image_simulation_test.py
@@ -1,19 +1,4 @@
 # coding=utf-8
-# Copyright 2022 The Google Research Authors.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Copyright 2022 Google LLC
 
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -34,6 +19,7 @@
 from absl.testing import parameterized
 import chex
 from immutabledict import immutabledict
+import itertools
 import jax
 import jax.numpy as jnp
 import numpy as np
@@ -57,438 +43,604 @@
 
 
 def _prepare_cosmology_params(cosmology_params_init, z_lookup_max, dz):
-  # Only generate a lookup table for values we need.
-  z_lookup_max = max(z_lookup_max, 1e-7)
-  dz = max(dz, 1e-7)
-  # We won't be using the linear density field lookup tables, so provide
-  # reasonable radial bounds and only 2 bins.
-  return cosmology_utils.add_lookup_tables_to_cosmology_params(
-      dict(cosmology_params_init), z_lookup_max, dz, 1e-3, 2e-3, 2)
+    # Only generate a lookup table for values we need.
+    z_lookup_max = max(z_lookup_max, 1e-7)
+    dz = max(dz, 1e-7)
+    # We won't be using the linear density field lookup tables, so provide
+    # reasonable radial bounds and only 2 bins.
+    return cosmology_utils.add_lookup_tables_to_cosmology_params(
+        dict(cosmology_params_init), z_lookup_max, dz, 1e-3, 2e-3, 2)
 
 
 def _prepare_kwargs_detector():
-  return {
-      'n_x': 2,
-      'n_y': 2,
-      'pixel_width': 0.04,
-      'supersampling_factor': 2,
-      'exposure_time': 1024,
-      'num_exposures': 2.0,
-      'sky_brightness': 22,
-      'magnitude_zero_point': 25,
-      'read_noise': 3.0
-  }
+    return {
+        'n_x': 2,
+        'n_y': 2,
+        'pixel_width': 0.04,
+        'supersampling_factor': 2,
+        'exposure_time': 1024,
+        'num_exposures': 2.0,
+        'sky_brightness': 22,
+        'magnitude_zero_point': 25,
+        'read_noise': 3.0
+    }
 
 
 def _prepare_kwargs_psf():
-  kwargs_detector = _prepare_kwargs_detector()
-  kwargs_psf = dict([
-      (param, 0.05) for param in image_simulation.ALL_PSF_MODEL_PARAMETERS
-  ])
-  kwargs_psf['pixel_width'] = (
-      kwargs_detector['pixel_width'] / kwargs_detector['supersampling_factor'])
-  x = jnp.arange(-0.5, 0.5, 0.04) / 0.02
-  kernel = jnp.outer(jnp.exp(-x**2), jnp.exp(-x**2))
-  kwargs_psf['kernel_point_source'] = kernel
-  kwargs_psf['model_index'] = psf_models.__all__.index('Pixel')
-  return kwargs_psf
+    kwargs_detector = _prepare_kwargs_detector()
+    all_psf_models = _prepare_all_psf_models()
+    params = tuple(sorted(set(itertools.chain(
+        *[model.parameters for model in all_psf_models]))))
+    kwargs_psf = {param: 0.05 for param in params}
+    kwargs_psf['pixel_width'] = (
+        kwargs_detector['pixel_width'] / kwargs_detector['supersampling_factor'])
+    x = jnp.arange(-0.5, 0.5, 0.04) / 0.02
+    kernel = jnp.outer(jnp.exp(-x**2), jnp.exp(-x**2))
+    kwargs_psf['kernel_point_source'] = kernel
+    kwargs_psf['model_index'] = psf_models.__all__.index('Pixel')
+    return kwargs_psf
 
 
 def _prepare_x_y():
-  rng = jax.random.PRNGKey(3)
-  rng_x, rng_y = jax.random.split(rng)
-  x = jax.random.normal(rng_x, shape=(3,)) * 1e3
-  y = jax.random.normal(rng_y, shape=(3,)) * 1e3
-  return x, y
+    rng = jax.random.PRNGKey(3)
+    rng_x, rng_y = jax.random.split(rng)
+    x = jax.random.normal(rng_x, shape=(3,)) * 1e3
+    y = jax.random.normal(rng_y, shape=(3,)) * 1e3
+    return x, y
 
 
 def _prepare_x_y_angular():
-  x, y = _prepare_x_y()
-  return x / 1e3, y / 1e3
+    x, y = _prepare_x_y()
+    return x / 1e3, y / 1e3
 
 
 def _prepare_alpha_x_alpha_y():
-  rng = jax.random.PRNGKey(2)
-  rng_x, rng_y = jax.random.split(rng)
-  alpha_x = jax.random.normal(rng_x, shape=(3,)) * 2
-  alpha_y = jax.random.normal(rng_y, shape=(3,)) * 2
-  return alpha_x, alpha_y
+    rng = jax.random.PRNGKey(2)
+    rng_x, rng_y = jax.random.split(rng)
+    alpha_x = jax.random.normal(rng_x, shape=(3,)) * 2
+    alpha_y = jax.random.normal(rng_y, shape=(3,)) * 2
+    return alpha_x, alpha_y
 
 
 def _prepare_image():
-  return jax.random.uniform(jax.random.PRNGKey(3), shape=(4, 4))
+    return jax.random.uniform(jax.random.PRNGKey(3), shape=(4, 4))
 
 
-def _prepare_kwargs_lens():
-  return dict([
-      (param, 0.5) for param in image_simulation.ALL_LENS_MODEL_PARAMETERS
-  ])
+def _prepare_all_psf_models():
+    models = psf_models.__all__
+    return tuple([psf_models.__getattribute__(model) for model in models])
 
+def _prepare_all_lens_models(model_group):
+    if model_group == 'los':
+        models = ['NFW']
+    elif model_group == 'subhalos':
+        models = ['TNFW']
+    elif model_group == 'main_deflector':
+        models = ['Shear', 'EPL']
+    elif model_group == 'all':
+        models = ['NFW', 'TNFW', 'Shear', 'EPL']
+    else:
+        raise ValueError(f'Unsupported lens_models specification {model_group}')
 
-def _prepare_kwargs_lens_all():
-  kwargs_lens_all = _prepare_kwargs_lens()
-  # Set all values to jnp.ndarrays.
-  for kwarg in kwargs_lens_all:
-    kwargs_lens_all[kwarg] = jnp.ones((12,)) * 0.5
-  kwargs_lens_all['model_index'] = jnp.array([
-      lens_models.__all__.index('NFW'),
-      lens_models.__all__.index('Shear'),
-      lens_models.__all__.index('TNFW'),
-      -1  #  Make sure that a padded model does not change the output.
-  ] * 3)
-  z_lens_array = jnp.array([0.1] * 4 + [0.2] * 4 + [0.3] * 4)
-  return kwargs_lens_all, z_lens_array
+    return tuple([lens_models.__getattribute__(model) for model in models])
 
 
-def _prepare_kwargs_source():
-  kwargs_source = dict([
-      (param, 0.5) for param in image_simulation.ALL_SOURCE_MODEL_PARAMETERS
-  ])
-  kwargs_source['image'] = jax.random.normal(
-      jax.random.PRNGKey(3), shape=(16, 16))
-  kwargs_source['model_index'] = source_models.__all__.index('Interpol')
-  return kwargs_source
+def _prepare_all_source_models():
+    all_source_models = tuple(
+        [source_models.__getattribute__(model)
+            for model in source_models.__all__])
+    return all_source_models
 
 
-def _prepare_kwargs_source_slice():
-  kwargs_source_slice = _prepare_kwargs_source()
-  # Modify kwargs_source_slice to deal with interpolation and give sersic
-  # profile meaningful brightness.
-  for kwarg in kwargs_source_slice:
-    kwargs_source_slice[kwarg] = jnp.ones((2,)) * 0.5
-  kwargs_source_slice['image'] = jax.random.uniform(
-      jax.random.PRNGKey(3), shape=(2, 16, 16))
-  kwargs_source_slice['n_sersic'] = jnp.ones((2,)) * 3.0
-  kwargs_source_slice['model_index'] = jnp.array([
-      source_models.__all__.index('Interpol'),
-      source_models.__all__.index('SersicElliptic')
-  ])
-  return kwargs_source_slice
+def _prepare_all_models():
+    model_group = 'all'
+    all_lens_models = _prepare_all_lens_models(model_group)
+    all_models = {}
+    all_models['all_los_models'] = all_lens_models
+    all_models['all_subhalo_models'] = all_lens_models
+    all_models['all_main_deflector_models'] = all_lens_models
+    all_models['all_source_models'] = _prepare_all_source_models()
+    all_models['all_psf_models'] = _prepare_all_psf_models()
 
-
-def _prepare__ray_shooting_step_expected(z_lens, z_source):
-  if z_lens == 0.1 and z_source == 0.5:
-    return (
-        [-1455.0688, 261.96478, -263.31418],
-        [311.67578, -1881.2992, 1816.3245],
-        [-0.42915863, 3.7673984, -2.9809818],
-        [-1.632507, -0.58538866, 0.6765887],
-        0.1,
-    )
-  elif z_lens == 0.5 and z_source == 1.0:
-    return (
-        [-2219.5464, 5967.2534, -4809.161],
-        [-2153.1443, -2859.7515, 2950.956],
-        [-0.34272414, 3.638215, -2.8426921],
-        [-1.4692892, -0.5461657, 0.694726],
-        0.5,
-    )
-  else:
-    raise ValueError(f'Unsupported z_lens and z_source = ({z_lens},{z_source})')
+    return all_models
 
 
-def _prepare__add_deflection_expected(z_lens, z_source):
-  if z_lens == 0.1 and z_source == 0.5:
-    return jnp.array([[-0.42978406, 3.7907193, -3.0125906],
-                      [-1.6494985, -0.5895851, 0.66275066]])
-  elif z_lens == 0.5 and z_source == 1.0:
-    return jnp.array([[-0.11232474, 3.921643, -2.5539021],
-                      [-1.5750934, -0.37239552, 0.2341882]])
-  else:
-    raise ValueError(f'Unsupported z_lens and z_source = ({z_lens},{z_source})')
+def _prepare_kwargs_lens(model_group):
+    all_lens_models = _prepare_all_lens_models(model_group)
+    params = tuple(sorted(set(itertools.chain(
+        *[model.parameters for model in all_lens_models]))))
+    return {param: 0.5 for param in params}
 
 
-class ImageSimulationTest(chex.TestCase, parameterized.TestCase):
-  """Runs tests of image simulation functions."""
-
-  @chex.all_variants(without_device=False)
-  def test_generate_image(self):
-    kwargs_lens_all, z_lens_array = _prepare_kwargs_lens_all()
-    kwargs_source_slice = _prepare_kwargs_source_slice()
-    kwargs_lens_light_slice = _prepare_kwargs_source_slice()
-    kwargs_detector = _prepare_kwargs_detector()
-    kwargs_psf = _prepare_kwargs_psf()
-    z_source = 0.5
-    cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_LENSTRONOMY,
-                                                 z_source, 0.1)
-    # Need to evaluate on coordinate grid to match lenstronomy.
-    grid_x, grid_y = utils.coordinates_evaluate(
-        kwargs_detector['n_x'], kwargs_detector['n_y'],
-        kwargs_detector['pixel_width'], kwargs_detector['supersampling_factor'])
-    expected = jnp.array([[0.00221893, 0.0030974], [0.00328375, 0.00458501]])
-
-    g_image = self.variant(
-        functools.partial(
-            image_simulation.generate_image, kwargs_detector=kwargs_detector))
-
-    np.testing.assert_allclose(
-        utils.downsample(
-            g_image(grid_x, grid_y, kwargs_lens_all, kwargs_source_slice,
-                    kwargs_lens_light_slice, kwargs_psf, cosmology_params,
-                    z_lens_array, z_source),
-            kwargs_detector['supersampling_factor']),
-        expected,
-        rtol=1e-3)
-
-  @chex.all_variants
-  def test_psf_convolution(self):
-    kwargs_psf = _prepare_kwargs_psf()
-    image = _prepare_image()
-
-    convolve = self.variant(image_simulation.psf_convolution)
-
-    # Default model is Pixel
-    expected = psf_models.Pixel.convolve(image, kwargs_psf)
-    np.testing.assert_allclose(convolve(image, kwargs_psf), expected, rtol=1e-5)
-
-    # Try Gaussian model
-    kwargs_psf['model_index'] = psf_models.__all__.index('Gaussian')
-    expected = psf_models.Gaussian.convolve(image, kwargs_psf)
-    np.testing.assert_allclose(convolve(image, kwargs_psf), expected, rtol=1e-5)
-
-  @chex.all_variants
-  def test_noise_realization(self):
-    kwargs_detector = _prepare_kwargs_detector()
-    image = _prepare_image()
-
-    # We can't just pull the expectation from lenstronomy, since the noise is a
-    # random realization and lenstronomy uses numpy for it's random functions.
-    # Instead we pull the magnitude of the normal draw from lenstronomy.
-    background_noise = 0.0040833212572525535
-    flux_noise = jnp.array([[0.02159131, 0.02026811, 0.00825452, 0.01600841],
-                            [0.02006242, 0.01973204, 0.01489258, 0.01015701],
-                            [0.01921518, 0.01649803, 0.00997489, 0.01330792],
-                            [0.01392349, 0.00903829, 0.01201941, 0.01346962]])
-
-    rng_noise = jax.random.PRNGKey(0)
-    rng_normal, rng_poisson = jax.random.split(rng_noise)
-    expected = (
-        background_noise * jax.random.normal(rng_normal, shape=image.shape) +
-        flux_noise * jax.random.normal(rng_poisson, shape=image.shape))
-
-    noise = self.variant(image_simulation.noise_realization)
-
-    # Noise matches very closely, but numbers are small so use atol.
-    np.testing.assert_allclose(
-        noise(image, rng_noise, kwargs_detector), expected, rtol=0, atol=1e-7)
-
-  @chex.all_variants
-  def test_lens_light_surface_brightness(self):
-    kwargs_lens_light_slice = _prepare_kwargs_source_slice()
-    kwargs_detector = _prepare_kwargs_detector()
-    # Need to evaluate on coordinate grid to match lenstronomy.
-    theta_x, theta_y = utils.coordinates_evaluate(
-        kwargs_detector['n_x'], kwargs_detector['n_y'],
-        kwargs_detector['pixel_width'], kwargs_detector['supersampling_factor'])
-    expected = jnp.array([[0.00074752, 0.00076053], [0.00083884, 0.00086197]])
-
-    ll_surface_brightness = self.variant(
-        image_simulation.lens_light_surface_brightness)
-
-    image = utils.downsample(
-        jnp.reshape(
-            ll_surface_brightness(theta_x, theta_y, kwargs_lens_light_slice,
-                                  kwargs_detector),
-            (
-                kwargs_detector['n_x'] *
-                kwargs_detector['supersampling_factor'],
-                kwargs_detector['n_y'] *
-                kwargs_detector['supersampling_factor'],
-            ),
-        ),
-        kwargs_detector['supersampling_factor'],
-    )
-
-    np.testing.assert_allclose(image, expected, rtol=1e-5)
-
-  @chex.all_variants(without_device=False)
-  def test_source_surface_brightness(self):
-    kwargs_lens_all, z_lens_array = _prepare_kwargs_lens_all()
-    kwargs_source_slice = _prepare_kwargs_source_slice()
-    z_source = 0.5
-    cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_LENSTRONOMY,
-                                                 z_source, 0.1)
-    kwargs_detector = _prepare_kwargs_detector()
-    # Need to evaluate on coordinate grid to match lenstronomy.
-    alpha_x, alpha_y = utils.coordinates_evaluate(
-        kwargs_detector['n_x'], kwargs_detector['n_y'],
-        kwargs_detector['pixel_width'], kwargs_detector['supersampling_factor'])
-    expected = jnp.array([[0.00325632, 0.00348312], [0.00369543, 0.00387607]])
-
-    s_surface_brightness = self.variant(
-        image_simulation.source_surface_brightness)
-
-    # Lenstronomy comparison is after downsampling.
-    image = utils.downsample(
-        jnp.reshape(
-            s_surface_brightness(alpha_x, alpha_y, kwargs_lens_all,
-                                 kwargs_source_slice, kwargs_detector,
-                                 cosmology_params, z_lens_array, z_source),
-            (
-                kwargs_detector['n_x'] *
-                kwargs_detector['supersampling_factor'],
-                kwargs_detector['n_y'] *
-                kwargs_detector['supersampling_factor'],
-            ),
-        ),
-        kwargs_detector['supersampling_factor'],
-    )
-    np.testing.assert_allclose(image, expected, rtol=1e-3)
-
-  @chex.all_variants(without_device=False)
-  def test__image_flux(self):
-    alpha_x, alpha_y = _prepare_alpha_x_alpha_y()
-    kwargs_lens_all, z_lens_array = _prepare_kwargs_lens_all()
-    kwargs_source_slice = _prepare_kwargs_source_slice()
-    z_source = 0.5
-    cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_LENSTRONOMY,
-                                                 z_source, 0.1)
-    expected = jnp.array([2.08355751e-01, 4.74442133e-05, 2.49857234e-04])
-
-    image_flux = self.variant(image_simulation._image_flux)
-
-    np.testing.assert_allclose(
-        image_flux(alpha_x, alpha_y, kwargs_lens_all, kwargs_source_slice,
-                   cosmology_params, z_lens_array, z_source),
-        expected,
-        rtol=1e-2)
-
-  @chex.all_variants(without_device=False)
-  def test_ray_shooting(self):
-    alpha_x, alpha_y = _prepare_alpha_x_alpha_y()
-    kwargs_lens_all, z_lens_array = _prepare_kwargs_lens_all()
-    z_source = 0.5
-    cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_LENSTRONOMY,
-                                                 z_source, 0.1)
-    expected = jnp.array([[4268.8287053, 6375.01729267, -2325.55000691],
-                          [-2971.59041917, -10392.18963907, 9207.33107629]])
-
-    ray_shooting = self.variant(image_simulation._ray_shooting)
-
-    np.testing.assert_allclose(
-        jnp.array(
-            ray_shooting(alpha_x, alpha_y, kwargs_lens_all, cosmology_params,
-                         z_lens_array, z_source)),
-        expected,
-        rtol=1e-2)
-
-  @chex.all_variants
-  @parameterized.named_parameters([
-      (f'_z_lens_{z_lens}_z_source_{z_source}', z_lens, z_source)
-      for z_lens, z_source in zip([0.1, 0.5], [0.5, 1.0])
-  ])
-  def test__ray_shooting_step(self, z_lens, z_source):
-    x, y, = _prepare_x_y()
-    alpha_x, alpha_y = _prepare_alpha_x_alpha_y()
-    z_lens_last = 0.05
-    cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_LENSTRONOMY,
-                                                 z_source, z_lens_last)
-    state = (x, y, alpha_x, alpha_y, z_lens_last)
-    kwargs_lens = _prepare_kwargs_lens()
-    kwargs_lens['model_index'] = lens_models.__all__.index('TNFW')
-
-    expected = _prepare__ray_shooting_step_expected(z_lens, z_source)
-
-    ray_shooting_step = self.variant(image_simulation._ray_shooting_step)
-
-    new_state, new_state_copy = ray_shooting_step(
-        state,
-        {
-            'z_lens': z_lens,
-            'kwargs_lens': kwargs_lens
-        },
-        cosmology_params,
-        z_source,
-    )
-
-    for si in range(len(new_state)):
-      np.testing.assert_allclose(new_state[si], expected[si], rtol=1e-2)
-      np.testing.assert_allclose(new_state_copy[si], expected[si], rtol=1e-2)
-
-  @chex.all_variants
-  def test__ray_step_add(self):
-    x, y = _prepare_x_y()
-    alpha_x, alpha_y = _prepare_alpha_x_alpha_y()
-    delta_t = 1e3
-    expected = jnp.array([[-1852.1759, 3225.5708, -2624.6494],
-                          [-968.6724, -2389.5552, 2405.7075]])
-
-    ray_step_add = self.variant(image_simulation._ray_step_add)
-
-    np.testing.assert_allclose(
-        jnp.array(ray_step_add(x, y, alpha_x, alpha_y, delta_t)),
-        expected,
-        rtol=1e-6)
-
-  @chex.all_variants
-  @parameterized.named_parameters([
-      (f'_z_lens_{z_lens}_z_source_{z_source}', z_lens, z_source)
-      for z_lens, z_source in zip([0.1, 0.5], [0.5, 1.0])
-  ])
-  def test__add_deflection(self, z_lens, z_source):
-    x, y = _prepare_x_y()
-    alpha_x, alpha_y = _prepare_alpha_x_alpha_y()
-    kwargs_lens = _prepare_kwargs_lens()
-    kwargs_lens['model_index'] = lens_models.__all__.index('TNFW')
-    cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_LENSTRONOMY,
-                                                 z_source, z_lens)
+def _prepare_kwargs_lens_all():
+    model_group = 'all'
 
-    expected = _prepare__add_deflection_expected(z_lens, z_source)
+    kwargs_lens_all = {}
+    kwargs_lens_all['kwargs_los_before'] = _prepare_kwargs_lens(model_group)
+    kwargs_lens_all['kwargs_subhalos'] = _prepare_kwargs_lens(model_group)
+    kwargs_lens_all['kwargs_main_deflector'] = _prepare_kwargs_lens(model_group)
+    kwargs_lens_all['kwargs_los_after'] = _prepare_kwargs_lens(model_group)
+    # Set all values to jnp.ndarrays.
+    for _, kwargs_lens in kwargs_lens_all.items():
+        for kwarg in kwargs_lens:
+            kwargs_lens[kwarg] = jnp.ones((4,)) * 0.5
+        # Should be NFW, Shear, TNFW, and then padded indices.
+        # Make sure that a padded model does not change the output.
+        kwargs_lens['model_index'] = jnp.array([0, 2, 1, -1])
 
-    add_deflection = self.variant(image_simulation._add_deflection)
+    kwargs_lens_all['z_array_los_before'] = jnp.array([0.1] * 4)
+    kwargs_lens_all['z_array_subhalos'] = jnp.array([0.2] * 4)
+    kwargs_lens_all['z_array_main_deflector'] = jnp.array([0.2] * 4)
+    kwargs_lens_all['z_array_los_after'] = jnp.array([0.3] * 4)
 
-    np.testing.assert_allclose(
-        jnp.array(
-            add_deflection(x, y, alpha_x, alpha_y, kwargs_lens,
-                           cosmology_params, z_lens, z_source)),
-        expected,
-        rtol=1e-2)
+    # Turn off the subhalo models.
+    kwargs_lens_all['kwargs_subhalos']['model_index'] = jnp.array([-1] * 4)
 
-  @chex.all_variants
-  def test__calculate_derivatives(self):
-    x, y = _prepare_x_y_angular()
-    kwargs_lens = _prepare_kwargs_lens()
-    kwargs_lens['model_index'] = lens_models.__all__.index('NFW')
-    expected = jnp.array([[-0.42939782, -0.16496612, -0.05045471],
-                          [0.03707895, -0.3547466, 0.47889906]])
+    return kwargs_lens_all
 
-    calculate_derivatives = self.variant(
-        image_simulation._calculate_derivatives)
-    a_stack = jnp.array(calculate_derivatives(kwargs_lens, x, y))
 
-    np.testing.assert_allclose(a_stack, expected, rtol=1e-5)
+def _prepare_kwargs_source():
+    all_source_models = _prepare_all_source_models()
+    all_source_model_parameters = tuple(sorted(set(itertools.chain(
+        *[model.parameters for model in all_source_models]))))
+    kwargs_source = dict([
+        (param, 0.5) for param in all_source_model_parameters
+    ])
+    kwargs_source['image'] = jax.random.normal(
+        jax.random.PRNGKey(3), shape=(16, 16))
+    kwargs_source['model_index'] = source_models.__all__.index('Interpol')
+    return kwargs_source
 
-    # Also check that model_index of -1 returns the identity
-    kwargs_lens['model_index'] = -1
-    a_stack = jnp.array(calculate_derivatives(kwargs_lens, x, y))
-    np.testing.assert_allclose(a_stack, jnp.zeros_like(expected), rtol=1e-5)
 
-  @chex.all_variants
-  def test__surface_brightness(self):
-    x, y = _prepare_x_y_angular()
-    kwargs_source_slice = _prepare_kwargs_source_slice()
-    expected = jnp.array([0.28131152, 0.26588708, 0.32980748])
+def _prepare_kwargs_source_slice():
+    kwargs_source_slice = _prepare_kwargs_source()
+    # Modify kwargs_source_slice to deal with interpolation and give sersic
+    # profile meaningful brightness.
+    for kwarg in kwargs_source_slice:
+        kwargs_source_slice[kwarg] = jnp.ones((2,)) * 0.5
+    kwargs_source_slice['image'] = jax.random.uniform(
+        jax.random.PRNGKey(3), shape=(2, 16, 16))
+    kwargs_source_slice['n_sersic'] = jnp.ones((2,)) * 3.0
+    kwargs_source_slice['model_index'] = jnp.array([
+        source_models.__all__.index('Interpol'),
+        source_models.__all__.index('SersicElliptic')
+    ])
+    return kwargs_source_slice
 
-    surface_brightness = self.variant(image_simulation._surface_brightness)
 
-    np.testing.assert_allclose(
-        surface_brightness(x, y, kwargs_source_slice), expected, rtol=1e-5)
+def _prepare__ray_shooting_step_expected(z_lens, z_source):
+    if z_lens == 0.1 and z_source == 0.5:
+        return (
+            [-1455.0688, 261.96478, -263.31418],
+            [311.67578, -1881.2992, 1816.3245],
+            [-0.42915863, 3.7673984, -2.9809818],
+            [-1.632507, -0.58538866, 0.6765887],
+            0.1,
+        )
+    elif z_lens == 0.5 and z_source == 1.0:
+        return (
+            [-2219.5464, 5967.2534, -4809.161],
+            [-2153.1443, -2859.7515, 2950.956],
+            [-0.34272414, 3.638215, -2.8426921],
+            [-1.4692892, -0.5461657, 0.694726],
+            0.5,
+        )
+    else:
+        raise ValueError(
+            f'Unsupported z_lens and z_source = ({z_lens},{z_source})')
 
-  @chex.all_variants
-  def test__add_surface_brightness(self):
-    x, y = _prepare_x_y_angular()
-    kwargs_source = _prepare_kwargs_source()
-    expected = jnp.array([-0.0836301, 0.06599196, 0.23668436])
 
-    add_surface_brightness = self.variant(
-        image_simulation._add_surface_brightness)
+def _prepare__add_deflection_expected(z_lens, z_source):
+    if z_lens == 0.1 and z_source == 0.5:
+        return jnp.array([[-0.42978406, 3.7907193, -3.0125906],
+                        [-1.6494985, -0.5895851, 0.66275066]])
+    elif z_lens == 0.5 and z_source == 1.0:
+        return jnp.array([[-0.11232474, 3.921643, -2.5539021],
+                        [-1.5750934, -0.37239552, 0.2341882]])
+    else:
+        raise ValueError(
+            f'Unsupported z_lens and z_source = ({z_lens},{z_source})')
 
-    # Use x as previous surface brightness.
-    tb, b = add_surface_brightness(x, kwargs_source, x, y)
 
-    np.testing.assert_allclose(tb, expected + x, rtol=1e-5)
-    np.testing.assert_allclose(b, expected, rtol=1e-5)
+class ImageSimulationTest(chex.TestCase, parameterized.TestCase):
+    """Runs tests of image simulation functions."""
+
+    @chex.all_variants(without_device=False)
+    def test_generate_image(self):
+        all_models = _prepare_all_models()
+        kwargs_lens_all = _prepare_kwargs_lens_all()
+        kwargs_source_slice = _prepare_kwargs_source_slice()
+        kwargs_lens_light_slice = _prepare_kwargs_source_slice()
+        kwargs_detector = _prepare_kwargs_detector()
+        kwargs_psf = _prepare_kwargs_psf()
+        z_source = 0.5
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_LENSTRONOMY,
+                                                    z_source, 0.1)
+        # Need to evaluate on coordinate grid to match lenstronomy.
+        grid_x, grid_y = utils.coordinates_evaluate(
+            kwargs_detector['n_x'], kwargs_detector['n_y'],
+            kwargs_detector['pixel_width'], kwargs_detector['supersampling_factor'])
+        expected = jnp.array([[0.00221893, 0.0030974], [0.00328375, 0.00458501]])
+
+        g_image = self.variant(
+            functools.partial(
+                image_simulation.generate_image, kwargs_detector=kwargs_detector,
+                all_models=all_models))
+
+        result = utils.downsample(
+        g_image(grid_x, grid_y, kwargs_lens_all, kwargs_source_slice,
+                kwargs_lens_light_slice, kwargs_psf, cosmology_params,
+                z_source),
+        kwargs_detector['supersampling_factor'])
+        np.testing.assert_allclose(result, expected, rtol=1e-3)
+
+    @chex.all_variants
+    def test_psf_convolution(self):
+        all_psf_models = _prepare_all_psf_models()
+        kwargs_psf = _prepare_kwargs_psf()
+        image = _prepare_image()
+
+        convolve = self.variant(functools.partial(
+            image_simulation.psf_convolution, all_psf_models=all_psf_models))
+
+        # Default model is Pixel
+        expected = psf_models.Pixel.convolve(image, kwargs_psf)
+        np.testing.assert_allclose(convolve(image, kwargs_psf), expected, rtol=1e-5)
+
+        # Try Gaussian model
+        kwargs_psf['model_index'] = psf_models.__all__.index('Gaussian')
+        expected = psf_models.Gaussian.convolve(image, kwargs_psf)
+        np.testing.assert_allclose(convolve(image, kwargs_psf), expected, rtol=1e-5)
+
+    @chex.all_variants
+    def test_noise_realization(self):
+        kwargs_detector = _prepare_kwargs_detector()
+        image = _prepare_image()
+
+        # We can't just pull the expectation from lenstronomy, since the noise is a
+        # random realization and lenstronomy uses numpy for it's random functions.
+        # Instead we pull the magnitude of the normal draw from lenstronomy.
+        background_noise = 0.0040833212572525535
+        flux_noise = jnp.array([[0.02159131, 0.02026811, 0.00825452, 0.01600841],
+                                [0.02006242, 0.01973204, 0.01489258, 0.01015701],
+                                [0.01921518, 0.01649803, 0.00997489, 0.01330792],
+                                [0.01392349, 0.00903829, 0.01201941, 0.01346962]])
+
+        rng_noise = jax.random.PRNGKey(0)
+        rng_normal, rng_poisson = jax.random.split(rng_noise)
+        expected = (
+            background_noise * jax.random.normal(rng_normal, shape=image.shape) +
+            flux_noise * jax.random.normal(rng_poisson, shape=image.shape))
+
+        noise = self.variant(image_simulation.noise_realization)
+
+        # Noise matches very closely, but numbers are small so use atol.
+        np.testing.assert_allclose(
+            noise(image, rng_noise, kwargs_detector), expected, rtol=0, atol=1e-7)
+
+    @chex.all_variants
+    def test_lens_light_surface_brightness(self):
+        all_source_models = _prepare_all_source_models()
+        kwargs_lens_light_slice = _prepare_kwargs_source_slice()
+        kwargs_detector = _prepare_kwargs_detector()
+        # Need to evaluate on coordinate grid to match lenstronomy.
+        theta_x, theta_y = utils.coordinates_evaluate(
+            kwargs_detector['n_x'], kwargs_detector['n_y'],
+            kwargs_detector['pixel_width'], kwargs_detector['supersampling_factor'])
+        expected = jnp.array([[0.00074752, 0.00076053], [0.00083884, 0.00086197]])
+
+        ll_surface_brightness = self.variant(functools.partial(
+            image_simulation.lens_light_surface_brightness,
+            all_source_models=all_source_models))
+
+        image = utils.downsample(
+            jnp.reshape(
+                ll_surface_brightness(theta_x, theta_y, kwargs_lens_light_slice,
+                                    kwargs_detector),
+                (
+                    kwargs_detector['n_x'] *
+                    kwargs_detector['supersampling_factor'],
+                    kwargs_detector['n_y'] *
+                    kwargs_detector['supersampling_factor'],
+                ),
+            ),
+            kwargs_detector['supersampling_factor'],
+        )
+
+        np.testing.assert_allclose(image, expected, rtol=1e-5)
+
+    @chex.all_variants(without_device=False)
+    def test_source_surface_brightness(self):
+        all_models = _prepare_all_models()
+        kwargs_lens_all = _prepare_kwargs_lens_all()
+        kwargs_source_slice = _prepare_kwargs_source_slice()
+        z_source = 0.5
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_LENSTRONOMY,
+                                                    z_source, 0.1)
+        kwargs_detector = _prepare_kwargs_detector()
+        # Need to evaluate on coordinate grid to match lenstronomy.
+        alpha_x, alpha_y = utils.coordinates_evaluate(
+            kwargs_detector['n_x'], kwargs_detector['n_y'],
+            kwargs_detector['pixel_width'], kwargs_detector['supersampling_factor'])
+        expected = jnp.array([[0.00325632, 0.00348312], [0.00369543, 0.00387607]])
+
+        s_surface_brightness = self.variant(
+            functools.partial(image_simulation.source_surface_brightness,
+                all_models=all_models))
+
+        # Lenstronomy comparison is after downsampling.
+        image = utils.downsample(
+            jnp.reshape(
+                s_surface_brightness(alpha_x, alpha_y, kwargs_lens_all,
+                                    kwargs_source_slice, kwargs_detector,
+                                    cosmology_params, z_source),
+                (
+                    kwargs_detector['n_x'] *
+                    kwargs_detector['supersampling_factor'],
+                    kwargs_detector['n_y'] *
+                    kwargs_detector['supersampling_factor'],
+                ),
+            ),
+            kwargs_detector['supersampling_factor'],
+        )
+        np.testing.assert_allclose(image, expected, rtol=1e-3)
+
+    @chex.all_variants(without_device=False)
+    def test__image_flux(self):
+        alpha_x, alpha_y = _prepare_alpha_x_alpha_y()
+        all_models = _prepare_all_models()
+        kwargs_lens_all = _prepare_kwargs_lens_all()
+        kwargs_source_slice = _prepare_kwargs_source_slice()
+        z_source = 0.5
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_LENSTRONOMY,
+                                                    z_source, 0.1)
+        expected = jnp.array([2.08355751e-01, 4.74442133e-05, 2.49857234e-04])
+
+        image_flux = self.variant(functools.partial(
+            image_simulation._image_flux, all_models=all_models))
+
+        np.testing.assert_allclose(
+            image_flux(alpha_x, alpha_y, kwargs_lens_all, kwargs_source_slice,
+                    cosmology_params, z_source),
+            expected,
+            rtol=1e-2)
+
+    @chex.all_variants(without_device=False)
+    def test_ray_shooting(self):
+        alpha_x, alpha_y = _prepare_alpha_x_alpha_y()
+        all_models = _prepare_all_models()
+        kwargs_lens_all = _prepare_kwargs_lens_all()
+        z_source = 0.5
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_LENSTRONOMY,
+                                                    z_source, 0.1)
+        expected = jnp.array([[4268.8287053, 6375.01729267, -2325.55000691],
+                            [-2971.59041917, -10392.18963907, 9207.33107629]])
+
+        ray_shooting = self.variant(functools.partial(
+            image_simulation._ray_shooting, all_models=all_models))
+
+        np.testing.assert_allclose(
+            jnp.array(
+                ray_shooting(alpha_x, alpha_y, kwargs_lens_all,
+                cosmology_params, z_source)),
+            expected,
+            rtol=1e-2)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+        (f'_z_lens_{z_lens}_z_source_{z_source}', z_lens, z_source)
+        for z_lens, z_source in zip([0.1, 0.5], [0.5, 1.0])
+    ])
+    def test__ray_shooting_group(self, z_lens, z_source):
+        x, y, = _prepare_x_y()
+        alpha_x, alpha_y = _prepare_alpha_x_alpha_y()
+        z_lens_last = 0.05
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_LENSTRONOMY,
+                                                    z_source, z_lens_last)
+        state = (x, y, alpha_x, alpha_y, z_lens_last)
+        expected_state = state
+        model_group = 'subhalos'
+        all_lens_models = _prepare_all_lens_models(model_group)
+
+        # Make it four deflectors and make sure they add up correctly.
+        kwargs_lens = _prepare_kwargs_lens(model_group)
+        kwargs_lens_slice = {}
+        for kwarg in kwargs_lens:
+            kwargs_lens_slice[kwarg] = jnp.ones(4) * kwargs_lens[kwarg]
+        kwargs_lens['model_index'] = 0
+        kwargs_lens_slice['model_index'] = jnp.array([0] * 4)
+
+        for _ in range(4):
+            expected_state, _ = image_simulation._ray_shooting_step(
+                expected_state, {'z_lens': z_lens, 'kwargs_lens': kwargs_lens},
+                cosmology_params, z_source, all_lens_models)
+
+        ray_shooting_group = self.variant(functools.partial(
+            image_simulation._ray_shooting_group,
+            all_lens_models=all_lens_models))
+
+        new_state, new_state_copy = ray_shooting_group(
+            state, kwargs_lens_slice, cosmology_params, z_source, z_lens)
+
+        for si in range(len(new_state)):
+            np.testing.assert_allclose(new_state[si], expected_state[si],
+                                       rtol=1e-5)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+        (f'_z_lens_{z_lens}_z_source_{z_source}', z_lens, z_source)
+        for z_lens, z_source in zip([0.1, 0.5], [0.5, 1.0])
+    ])
+    def test__ray_shooting_step(self, z_lens, z_source):
+        x, y, = _prepare_x_y()
+        alpha_x, alpha_y = _prepare_alpha_x_alpha_y()
+        z_lens_last = 0.05
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_LENSTRONOMY,
+                                                    z_source, z_lens_last)
+        state = (x, y, alpha_x, alpha_y, z_lens_last)
+        model_group = 'subhalos'
+        all_lens_models = _prepare_all_lens_models(model_group)
+        kwargs_lens = _prepare_kwargs_lens(model_group)
+        kwargs_lens['model_index'] = 0
+
+        expected = _prepare__ray_shooting_step_expected(z_lens, z_source)
+
+        ray_shooting_step = self.variant(functools.partial(
+            image_simulation._ray_shooting_step,
+            all_lens_models=all_lens_models))
+
+        new_state, new_state_copy = ray_shooting_step(
+            state,
+            {
+                'z_lens': z_lens,
+                'kwargs_lens': kwargs_lens
+            },
+            cosmology_params,
+            z_source,
+        )
+
+        for si in range(len(new_state)):
+            np.testing.assert_allclose(new_state[si], expected[si],
+                                       rtol=1e-2)
+            np.testing.assert_allclose(new_state_copy[si], expected[si],
+                                       rtol=1e-2)
+
+    @chex.all_variants
+    def test__ray_step_add(self):
+        x, y = _prepare_x_y()
+        alpha_x, alpha_y = _prepare_alpha_x_alpha_y()
+        delta_t = 1e3
+        expected = jnp.array([[-1852.1759, 3225.5708, -2624.6494],
+                            [-968.6724, -2389.5552, 2405.7075]])
+
+        ray_step_add = self.variant(image_simulation._ray_step_add)
+
+        np.testing.assert_allclose(
+            jnp.array(ray_step_add(x, y, alpha_x, alpha_y, delta_t)),
+            expected,
+            rtol=1e-6)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+        (f'_z_lens_{z_lens}_z_source_{z_source}', z_lens, z_source)
+        for z_lens, z_source in zip([0.1, 0.5], [0.5, 1.0])
+    ])
+    def test__add_deflection_group(self, z_lens, z_source):
+        x, y = _prepare_x_y()
+        alpha_x, alpha_y = _prepare_alpha_x_alpha_y()
+        model_group = 'subhalos'
+        all_lens_models = _prepare_all_lens_models(model_group)
+
+        # Make it four deflectors and make sure they add up correctly.
+        kwargs_lens = _prepare_kwargs_lens(model_group)
+        kwargs_lens_slice = {}
+        for kwarg in kwargs_lens:
+            kwargs_lens_slice[kwarg] = jnp.ones(4) * kwargs_lens[kwarg]
+        kwargs_lens['model_index'] = 0
+        kwargs_lens_slice['model_index'] = jnp.array([0] * 4)
+
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_LENSTRONOMY,
+                                                    z_source, z_lens)
+
+        alpha_x_expected = jnp.copy(alpha_x)
+        alpha_y_expected = jnp.copy(alpha_y)
+        for _ in range(4):
+            alpha_x_expected, alpha_y_expected = (
+                image_simulation._add_deflection(
+                    x, y, alpha_x_expected, alpha_y_expected, kwargs_lens,
+                    cosmology_params, z_lens, z_source, all_lens_models))
+
+        add_deflection_group = self.variant(
+            functools.partial(image_simulation._add_deflection_group,
+            all_lens_models=all_lens_models))
+
+        np.testing.assert_allclose(
+            jnp.array(
+                add_deflection_group(x, y, alpha_x, alpha_y, kwargs_lens_slice,
+                            cosmology_params, z_lens, z_source)),
+            jnp.array([alpha_x_expected, alpha_y_expected]), rtol=1e-5)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+        (f'_z_lens_{z_lens}_z_source_{z_source}', z_lens, z_source)
+        for z_lens, z_source in zip([0.1, 0.5], [0.5, 1.0])
+    ])
+    def test__add_deflection(self, z_lens, z_source):
+        x, y = _prepare_x_y()
+        alpha_x, alpha_y = _prepare_alpha_x_alpha_y()
+        model_group = 'subhalos'
+        all_lens_models = _prepare_all_lens_models(model_group)
+        kwargs_lens = _prepare_kwargs_lens(model_group)
+        kwargs_lens['model_index'] = 0
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_LENSTRONOMY,
+                                                    z_source, z_lens)
+
+        expected = _prepare__add_deflection_expected(z_lens, z_source)
+
+        add_deflection = self.variant(
+            functools.partial(image_simulation._add_deflection,
+            all_lens_models=all_lens_models))
+
+        np.testing.assert_allclose(
+            jnp.array(
+                add_deflection(x, y, alpha_x, alpha_y, kwargs_lens,
+                            cosmology_params, z_lens, z_source)),
+            expected,
+            rtol=1e-2)
+
+    @chex.all_variants
+    def test__calculate_derivatives(self):
+        x, y = _prepare_x_y_angular()
+        model_group = 'los'
+        all_lens_models = _prepare_all_lens_models(model_group)
+        kwargs_lens = _prepare_kwargs_lens(model_group)
+        kwargs_lens['model_index'] = 0
+        expected = jnp.array([[-0.42939782, -0.16496612, -0.05045471],
+                            [0.03707895, -0.3547466, 0.47889906]])
+
+        calculate_derivatives = self.variant(functools.partial(
+            image_simulation._calculate_derivatives,
+            all_lens_models=all_lens_models))
+        a_stack = jnp.array(calculate_derivatives(kwargs_lens, x, y))
+
+        np.testing.assert_allclose(a_stack, expected, rtol=1e-5)
+
+        # Also check that model_index of -1 returns the identity
+        kwargs_lens['model_index'] = -1
+        a_stack = jnp.array(calculate_derivatives(kwargs_lens, x, y))
+        np.testing.assert_allclose(a_stack, jnp.zeros_like(expected), rtol=1e-5)
+
+    @chex.all_variants
+    def test__surface_brightness(self):
+        x, y = _prepare_x_y_angular()
+        all_source_models = _prepare_all_source_models()
+        kwargs_source_slice = _prepare_kwargs_source_slice()
+        expected = jnp.array([0.28131152, 0.26588708, 0.32980748])
+
+        surface_brightness = self.variant(functools.partial(
+            image_simulation._surface_brightness,
+            all_source_models=all_source_models))
+
+        np.testing.assert_allclose(
+            surface_brightness(x, y, kwargs_source_slice), expected, rtol=1e-5)
+
+    @chex.all_variants
+    def test__add_surface_brightness(self):
+        x, y = _prepare_x_y_angular()
+        all_source_models = _prepare_all_source_models()
+        kwargs_source = _prepare_kwargs_source()
+        expected = jnp.array([-0.0836301, 0.06599196, 0.23668436])
+
+        add_surface_brightness = self.variant(
+            functools.partial(image_simulation._add_surface_brightness,
+            all_source_models=all_source_models))
+
+        # Use x as previous surface brightness.
+        tb, b = add_surface_brightness(x, kwargs_source, x, y)
+
+        np.testing.assert_allclose(tb, expected + x, rtol=1e-5)
+        np.testing.assert_allclose(b, expected, rtol=1e-5)
 
 
 if __name__ == '__main__':
-  absltest.main()
+    absltest.main()
diff --git a/jaxstronomy/input_pipeline_test.py b/jaxstronomy/input_pipeline_test.py
new file mode 100644
index 0000000..2c66a08
--- /dev/null
+++ b/jaxstronomy/input_pipeline_test.py
@@ -0,0 +1,472 @@
+# coding=utf-8
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     https://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for input_pipeline.py."""
+
+import functools
+
+from absl.testing import parameterized
+import chex
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from jaxstronomy import cosmology_utils
+from jaxstronomy import input_pipeline
+from jaxstronomy import lens_models
+from jaxstronomy import psf_models
+from jaxstronomy import source_models
+
+class InputPipelineTests(chex.TestCase, parameterized.TestCase):
+    """Runs tests of image simulation functions."""
+
+    def setUp(self):
+        # Set up the baseline configuration.
+        self.config = {}
+        encode_constant = input_pipeline.encode_constant
+        self.config['cosmology_params'] = {
+            'omega_m_zero': encode_constant(0.3089),
+            'omega_b_zero': encode_constant(0.0486),
+            'omega_de_zero': encode_constant(0.6910088292453472),
+            'omega_rad_zero': encode_constant(9.117075466e-5),
+            'temp_cmb_zero': encode_constant(2.7255),
+            'hubble_constant': encode_constant(67.74),
+            'n_s': encode_constant(0.9667),
+            'sigma_eight': encode_constant(0.8159),
+        }
+        return super().setUp()
+
+    @chex.all_variants
+    def test_encode_normal(self):
+        # Test that the encoding behaves as expected
+        encode_normal = self.variant(input_pipeline.encode_normal)
+        mean = 1.2
+        std = 0.4
+        encoded = encode_normal(mean=mean, std=std)
+        expected = np.array([0.0, 0.0, 0.0, 1.0, mean, std, 0.0])
+        np.testing.assert_array_almost_equal(encoded, expected)
+
+    @chex.all_variants
+    def test_encode_uniform(self):
+        # Test that the encoding behaves as expected
+        encode_uniform = self.variant(input_pipeline.encode_uniform)
+        minimum = 0.2
+        maximum = 2.01
+        encoded = encode_uniform(minimum=minimum, maximum=maximum)
+        expected = np.array([1.0, minimum, maximum, 0.0, 0.0, 0.0, 0.0])
+        np.testing.assert_array_almost_equal(encoded, expected)
+
+    @chex.all_variants
+    def test_encode_constant(self):
+        # That that the encoding behaves as expected.
+        encode_constant = self.variant(input_pipeline.encode_constant)
+        constant = 1.43
+        encoded = encode_constant(constant=constant)
+        expected = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, constant])
+        np.testing.assert_array_almost_equal(encoded, expected)
+
+    @chex.all_variants
+    def test_decode_maximum(self):
+        # Test that for the constant, uniform, and gaussian case it gives
+        # the constant, the maximum, and the mean plus five sigma.
+        mean = 1.8
+        std = 1.0
+        minimum = 0.8
+        maximum = 12.0
+        constant = 13.4
+        constant_encoding = input_pipeline.encode_constant(constant)
+        uniform_encoding = input_pipeline.encode_uniform(minimum=minimum,
+                                                         maximum=maximum)
+        normal_encoding = input_pipeline.encode_normal(mean=mean, std=std)
+        decode_maximum = self.variant(input_pipeline.decode_maximum)
+
+        self.assertAlmostEqual(decode_maximum(constant_encoding), constant,
+                               places=6)
+        self.assertAlmostEqual(decode_maximum(uniform_encoding), maximum,
+                               places=6)
+        self.assertAlmostEqual(decode_maximum(normal_encoding), mean + 5 * std,
+                               places=6)
+
+    @chex.all_variants
+    def test_decode_minimum(self):
+        # Test that for the constant, uniform, and gaussian case it gives
+        # the constant, the maximum, and the mean plus five sigma.
+        mean = 1.8
+        std = 1.0
+        minimum = 0.8
+        maximum = 12.0
+        constant = 13.4
+        constant_encoding = input_pipeline.encode_constant(constant)
+        uniform_encoding = input_pipeline.encode_uniform(minimum=minimum,
+                                                         maximum=maximum)
+        normal_encoding = input_pipeline.encode_normal(mean=mean, std=std)
+        decode_minimum = self.variant(input_pipeline.decode_minimum)
+
+        self.assertAlmostEqual(decode_minimum(constant_encoding), constant,
+                               places=6)
+        self.assertAlmostEqual(decode_minimum(uniform_encoding), minimum,
+                               places=6)
+        self.assertAlmostEqual(decode_minimum(normal_encoding), mean - 5 * std,
+                               places=6)
+
+    @chex.all_variants
+    def test_draw_from_encoding(self):
+        # Test that the encoding gives reasonable draws from the distribution
+        # for the normal, uniform, and constant encoding.
+        draw_from_encoding = self.variant(input_pipeline.draw_from_encoding)
+        draw_from_encoding_vmap = jax.vmap(draw_from_encoding,
+                                           in_axes=[None, 0])
+        n_draws = 100000
+        rng_list = jax.random.split(jax.random.PRNGKey(0), n_draws)
+
+        constant = 12.0
+        constant_encoding = input_pipeline.encode_constant(constant)
+        constant_draws = draw_from_encoding_vmap(constant_encoding, rng_list)
+        np.testing.assert_array_almost_equal(constant_draws,
+                                             np.ones(n_draws) * constant)
+
+        mean = 1.8
+        std = 1.0
+        normal_encoding = input_pipeline.encode_normal(mean=mean, std=std)
+        normal_draws = draw_from_encoding_vmap(normal_encoding, rng_list)
+        self.assertAlmostEqual(jnp.mean(normal_draws), mean, places=2)
+        self.assertAlmostEqual(jnp.std(normal_draws), std, places=2)
+
+        minimum = 0.0
+        maximum = 4.0
+        uniform_encoding = input_pipeline.encode_uniform(minimum=minimum,
+                                                         maximum=maximum)
+        uniform_draws = draw_from_encoding_vmap(uniform_encoding, rng_list)
+        self.assertAlmostEqual(jnp.mean(uniform_draws < 1.0), 0.25, places=2)
+        self.assertAlmostEqual(jnp.mean(uniform_draws < 2.0), 0.5, places=2)
+        self.assertAlmostEqual(jnp.mean(uniform_draws < 3.0), 0.75, places=2)
+        self.assertAlmostEqual(jnp.mean(uniform_draws < 4.0), 1.0, places=2)
+
+    @chex.all_variants
+    def test_normalize_params(self):
+        # Test that the draws are normalized as expected.
+        normalize_param = self.variant(input_pipeline.normalize_param)
+        normalize_param_vmap = jax.vmap(normalize_param, in_axes=[0, None])
+        draw_from_encoding_vmap = jax.jit(jax.vmap(
+            input_pipeline.draw_from_encoding, in_axes=[None, 0]))
+
+        n_draws = 100000
+        rng_list = jax.random.split(jax.random.PRNGKey(0), n_draws)
+
+        constant = 12.0
+        constant_encoding = input_pipeline.encode_constant(constant)
+        constant_draws = draw_from_encoding_vmap(constant_encoding, rng_list)
+        normalized_draws = normalize_param_vmap(constant_draws,
+                                                constant_encoding)
+        np.testing.assert_array_almost_equal(normalized_draws, np.zeros(n_draws))
+
+        mean = 1.8
+        std = 1.0
+        normal_encoding = input_pipeline.encode_normal(mean=mean, std=std)
+        normal_draws = draw_from_encoding_vmap(normal_encoding, rng_list)
+        normalized_draws = normalize_param_vmap(normal_draws, normal_encoding)
+        self.assertAlmostEqual(jnp.mean(normalized_draws), 0.0, places=2)
+        self.assertAlmostEqual(jnp.std(normalized_draws), 1.0, places=2)
+
+        minimum = 0.0
+        maximum = 4.0
+        uniform_encoding = input_pipeline.encode_uniform(minimum=minimum,
+                                                         maximum=maximum)
+        uniform_draws = draw_from_encoding_vmap(uniform_encoding, rng_list)
+        normalized_draws = normalize_param_vmap(uniform_draws, uniform_encoding)
+        self.assertAlmostEqual(jnp.mean(normalized_draws < 0.25), 0.25,
+                               places=2)
+        self.assertAlmostEqual(jnp.mean(normalized_draws < 0.50), 0.5, places=2)
+        self.assertAlmostEqual(jnp.mean(normalized_draws < 0.75), 0.75,
+                               places=2)
+        self.assertAlmostEqual(jnp.mean(normalized_draws < 1.00), 1.0, places=2)
+
+    def test_generate_grids(self):
+        # Test that the grid has the expected dimensions.
+        kd = {
+            'n_x': 128, 'n_y': 128, 'pixel_width': 0.04,
+            'supersampling_factor': 2.0
+        }
+        self.config['kwargs_detector'] = kd
+
+        grid_x, grid_y = input_pipeline.generate_grids(self.config)
+
+        size = kd['n_x'] * kd['n_y'] * kd['supersampling_factor'] ** 2
+        self.assertTupleEqual(grid_x.shape, (size,))
+        self.assertTupleEqual(grid_y.shape, (size,))
+        angular_size = ((kd['n_x'] * kd['supersampling_factor'] - 1) *
+                        kd['pixel_width'] / kd['supersampling_factor'])
+        self.assertAlmostEqual(jnp.max(grid_x) - jnp.min(grid_x), angular_size)
+
+    def test_initialize_cosmology_params(self):
+        # Test that the cosmology params generated can handle the range of
+        # values defined by the config.
+        encode_constant = input_pipeline.encode_constant
+        dz = 0.01
+        self.config['lensing_config'] = {
+            'source_params': {'z_source': encode_constant(2.0)},
+            'subhalo_params': {
+                'm_max': encode_constant(1e10),
+                'm_min': encode_constant(1e8)
+                },
+            'los_params': {
+                'dz': encode_constant(dz),
+                'm_max': encode_constant(1e9),
+                'm_min': encode_constant(1e7)
+            }
+        }
+        rng = jax.random.PRNGKey(0)
+
+        cosmology_params = input_pipeline.intialize_cosmology_params(
+            self.config, rng)
+
+        # Test that the cosmology params do not cause issues if the masses
+        # are within the specified range.
+        r_min = cosmology_utils.lagrangian_radius(cosmology_params, 1e6)
+        r_max = cosmology_utils.lagrangian_radius(cosmology_params, 1e11)
+        # Cosmological values are calculated in incriments of dz / 2.
+        self.assertAlmostEqual(cosmology_params['dz'], dz / 2)
+        self.assertAlmostEqual(cosmology_params['r_min'], r_min, places=6)
+        self.assertAlmostEqual(cosmology_params['r_max'], r_max, places=6)
+
+    @chex.all_variants
+    def test_draw_sample(self):
+        # Test that the samples drawn match the expected distribution.
+        mean = 1.0
+        std = 0.5
+        normal_encoding = input_pipeline.encode_normal(mean=mean, std=std)
+        constant = 12.0
+        constant_encoding = input_pipeline.encode_constant(constant=constant)
+        encoded_configuration = {
+            'node_a': {'node_aa': normal_encoding, 'node_ab': normal_encoding},
+            'node_b': constant_encoding
+        }
+        n_draws = 100000
+        rng_list = jax.random.split(jax.random.PRNGKey(0), n_draws)
+
+        draw_sample = self.variant(input_pipeline.draw_sample)
+        draw_sample_vmap = jax.vmap(draw_sample, in_axes=[None, 0])
+
+        sampled_configuration = draw_sample_vmap(encoded_configuration,
+                                                 rng_list)
+        np.testing.assert_array_almost_equal(
+            sampled_configuration['node_b'], np.ones(n_draws) * constant
+        )
+        self.assertAlmostEqual(
+            jnp.mean(sampled_configuration['node_a']['node_aa']), mean,
+            places=2)
+        self.assertAlmostEqual(
+            jnp.mean(sampled_configuration['node_a']['node_ab']), mean,
+            places=2)
+        self.assertAlmostEqual(
+            jnp.std(sampled_configuration['node_a']['node_aa']), std,
+            places=2)
+        self.assertAlmostEqual(
+            jnp.std(sampled_configuration['node_a']['node_ab']), std,
+            places=2)
+
+    @chex.all_variants
+    def test_extract_multiple_models(self):
+        # Test that the code scales well to multiple models and a single model.
+        constant = 12.0
+        mean = 0.9
+        std = 0.3
+        encoded_configuration = {
+            'param_1': input_pipeline.encode_constant(constant),
+            'param_2': input_pipeline.encode_normal(mean=mean, std=std),
+        }
+        n_models = 100000
+        rng = jax.random.PRNGKey(0)
+
+        extract_multiple_models = self.variant(functools.partial(
+            input_pipeline.extract_multiple_models, n_models=n_models
+        ))
+        sampled_configuration = extract_multiple_models(encoded_configuration,
+                                                        rng)
+
+        self.assertTupleEqual(sampled_configuration['param_1'].shape, (100000,))
+        np.testing.assert_array_almost_equal(sampled_configuration['param_1'],
+                                             np.ones(n_models) * constant)
+        self.assertAlmostEqual(jnp.mean(sampled_configuration['param_2']),
+                               mean, places=2)
+        self.assertAlmostEqual(jnp.std(sampled_configuration['param_2']),
+                               std, places=2)
+        np.testing.assert_array_almost_equal(
+            sampled_configuration['model_index'], np.arange(n_models))
+
+        n_models = 1
+        extract_multiple_models = self.variant(functools.partial(
+            input_pipeline.extract_multiple_models, n_models=n_models
+        ))
+        sampled_configuration = extract_multiple_models(encoded_configuration,
+                                                        rng)
+        self.assertTupleEqual(sampled_configuration['param_1'].shape, (1,))
+        np.testing.assert_array_almost_equal(
+            sampled_configuration['model_index'], np.arange(n_models))
+
+
+    @chex.all_variants(without_device=False)
+    def test_extract_truth_values(self):
+        # Test that the extracted parameters match the values fed into the
+        # dictionary.
+        all_params = {
+            'a': {'a': 0.0, 'b': 0.5,
+                  'c': 12.2},
+            'b': {'a': 1.0, 'b': 0.7,
+                  'c': 10.2}}
+        constant = 1.0
+        mean = 2.0
+        std = 1.0
+        minimum = -1.0
+        maximum = 12.0
+        lensing_config = {
+            'a': {'a': input_pipeline.encode_constant(constant),
+                  'b': input_pipeline.encode_normal(mean, std),
+                  'c': input_pipeline.encode_uniform(minimum, maximum)},
+            'b': {'a': input_pipeline.encode_constant(constant),
+                  'b': input_pipeline.encode_normal(mean, std),
+                  'c': input_pipeline.encode_uniform(minimum, maximum)}}
+        extract_objects = ['a','a','b']
+        extract_keys = ['a','b','c']
+        truth_parameters = (extract_objects, extract_keys)
+        extract_truth_values = self.variant(functools.partial(
+            input_pipeline.extract_truth_values,
+            truth_parameters=truth_parameters))
+
+        parameter_array = extract_truth_values(all_params, lensing_config)
+        self.assertTupleEqual(parameter_array.shape, (3,))
+        np.testing.assert_array_almost_equal(
+            parameter_array, np.array([0.0, (0.5 - mean) / std,
+                                       (10.2 - minimum) / (maximum - minimum)])
+        )
+
+
+    @chex.all_variants(without_device=False)
+    def test_draw_image_and_truth(self):
+        # Test that the images have reasonable shape and that the truth values
+        # are drawn correctly.
+        encode_normal = input_pipeline.encode_normal
+        encode_uniform = input_pipeline.encode_uniform
+        encode_constant = input_pipeline.encode_constant
+        self.config['lensing_config'] = {
+            'los_params':{
+                'delta_los': encode_constant(10.0),
+                'r_min': encode_constant(0.5),
+                'r_max': encode_constant(10.0),
+                'm_min': encode_constant(1e8),
+                'm_max': encode_constant(1e10),
+                'dz': encode_constant(0.1),
+                'cone_angle': encode_constant(8.0),
+                'angle_buffer': encode_constant(0.8),
+                'c_zero': encode_constant(18),
+                'conc_zeta': encode_constant(-0.2),
+                'conc_beta': encode_constant(0.8),
+                'conc_m_ref': encode_constant(1e8),
+                'conc_dex_scatter': encode_constant(0.0)
+            },
+            'main_deflector_params': {
+                'mass': encode_constant(1e13),
+                'z_lens': encode_constant(0.5),
+                'theta_e': encode_constant(1.1),
+                'slope': encode_normal(mean=2.0, std=0.1),
+                'center_x': encode_normal(mean=0.0, std=0.16),
+                'center_y': encode_normal(mean=0.0, std=0.16),
+                'axis_ratio': encode_normal(mean=1.0, std=0.05),
+                'angle': encode_uniform(minimum=0.0, maximum=2 * jnp.pi),
+                'gamma_ext': encode_normal(mean=0.0, std=0.05)
+            },
+            'subhalo_params':{
+                'sigma_sub': encode_normal(mean=2.0e-3, std=1.1e-3),
+                'shmf_plaw_index': encode_uniform(minimum=-2.02, maximum=-1.92),
+                'm_pivot': encode_constant(1e10),
+                'm_min': encode_constant(1e6),
+                'm_max': encode_constant(1e9),
+                'k_one': encode_constant(0.0),
+                'k_two': encode_constant(0.0),
+                'c_zero': encode_constant(18),
+                'conc_zeta': encode_constant(-0.2),
+                'conc_beta': encode_constant(0.8),
+                'conc_m_ref': encode_constant(1e8),
+                'conc_dex_scatter': encode_constant(0.0)
+            },
+            'source_params':{
+                'z_source': encode_constant(1.5),
+                'amp': encode_uniform(minimum=1.0, maximum=10.0),
+                'sersic_radius': encode_uniform(minimum=1.0, maximum=4.0),
+                'n_sersic': encode_uniform(minimum=1.0, maximum=4.0),
+                'axis_ratio': encode_normal(mean=1.0, std=0.05),
+                'angle': encode_uniform(minimum=0.0, maximum=2 * jnp.pi),
+                'center_x': encode_normal(mean=0.0, std=0.16),
+                'center_y': encode_normal(mean=0.0, std=0.16)
+            },
+            'lens_light_params':{
+                'amp': encode_constant(0.0),
+                'sersic_radius': encode_uniform(minimum=1.0, maximum=3.0),
+                'n_sersic': encode_uniform(minimum=1.0, maximum=4.0),
+                'axis_ratio': encode_normal(mean=1.0, std=0.05),
+                'angle': encode_uniform(minimum=0.0, maximum=2 * jnp.pi),
+                'center_x': encode_normal(mean=0.0, std=0.16),
+                'center_y': encode_normal(mean=0.0, std=0.16)
+            }
+        }
+        rng = jax.random.PRNGKey(0)
+
+        cosmology_params = input_pipeline.intialize_cosmology_params(
+            self.config, rng)
+        n_x = 16
+        n_y = 16
+        self.config['kwargs_detector'] = {
+            'n_x': n_x, 'n_y': n_y, 'pixel_width': 0.4,
+            'supersampling_factor': 2, 'exposure_time': 1024,
+            'num_exposures': 1.0, 'sky_brightness': 22,
+            'magnitude_zero_point': 25, 'read_noise': 3.0
+        }
+        grid_x, grid_y = input_pipeline.generate_grids(self.config)
+        all_models = {
+            'all_los_models': (lens_models.NFW,),
+            'all_subhalo_models': (lens_models.TNFW,),
+            'all_main_deflector_models': (lens_models.Shear, lens_models.EPL),
+            'all_source_models': (source_models.SersicElliptic,),
+            'all_psf_models': (psf_models.Gaussian,)
+        }
+        principal_md_index = 0
+        principal_source_index = 0
+        kwargs_simulation = {
+            'num_z_bins': 1000,
+            'los_pad_length': 10,
+            'subhalos_pad_length': 10,
+            'sampling_pad_length': 1000,
+        }
+        kwargs_psf = {'model_index': 0, 'fwhm': 0.04, 'pixel_width': 0.02}
+        truth_parameters = (
+            ['main_deflector_params', 'subhalo_params'],
+            ['theta_e', 'sigma_sub']
+        )
+
+        draw_image_and_truth = self.variant(functools.partial(
+            input_pipeline.draw_image_and_truth, all_models=all_models,
+            principal_md_index=principal_md_index,
+            principal_source_index=principal_source_index,
+            kwargs_simulation=kwargs_simulation,
+            kwargs_detector=self.config['kwargs_detector'],
+            kwargs_psf=kwargs_psf, truth_parameters=truth_parameters))
+        image, truth = draw_image_and_truth(self.config['lensing_config'],
+                                            cosmology_params, grid_x, grid_y,
+                                            rng)
+
+        # Just test that the data and truths vary and that they are normalized
+        # as expected.
+        self.assertTupleEqual(image.shape, (n_x, n_y))
+        self.assertAlmostEqual(1.0, jnp.std(image), places=6)
+        self.assertTupleEqual(truth.shape, (2,))
+        self.assertEqual(truth[0], 0.0)
diff --git a/jaxstronomy/lens_models_test.py b/jaxstronomy/lens_models_test.py
index 24059f9..3b99e3f 100644
--- a/jaxstronomy/lens_models_test.py
+++ b/jaxstronomy/lens_models_test.py
@@ -1,5 +1,4 @@
 # coding=utf-8
-# Copyright 2022 The Google Research Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -7,20 +6,6 @@
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Copyright 2022 Google LLC
-
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-
-#     https://www.apache.org/licenses/LICENSE-2.0
-
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -34,7 +19,6 @@
 
 import inspect
 
-from absl.testing import absltest
 from absl.testing import parameterized
 import chex
 import jax
@@ -44,302 +28,299 @@
 
 
 def _prepare_epl_parameters():
-  return {
-      'theta_e': 1.1,
-      'slope': 2.0,
-      'axis_ratio': 0.5,
-      'angle': np.pi / 4,
-      'center_x': 0.2,
-      'center_y': -0.2,
-  }
+	return {
+			'theta_e': 1.1,
+			'slope': 2.0,
+			'axis_ratio': 0.5,
+			'angle': np.pi / 4,
+			'center_x': 0.2,
+			'center_y': -0.2,
+	}
 
 
 def _prepare_nfw_parameters():
-  return {
-      'scale_radius': 1.2,
-      'alpha_rs': 4.0,
-      'center_x': 0.2,
-      'center_y': -0.2,
-  }
+	return {
+			'scale_radius': 1.2,
+			'alpha_rs': 4.0,
+			'center_x': 0.2,
+			'center_y': -0.2,
+	}
 
 
 def _prepare_shear_parameters():
-  return {
-      'gamma_ext': 1.2,
-      'angle': np.pi / 7,
-      'center_x': 0.2,
-      'center_y': -0.2,
-  }
+	return {
+			'gamma_ext': 1.2,
+			'angle': np.pi / 7,
+			'center_x': 0.2,
+			'center_y': -0.2,
+	}
 
 
 def _prepare_tnfw_parameters():
-  tnfw_parameters = _prepare_nfw_parameters()
-  tnfw_parameters['trunc_radius'] = 2.0
-  return tnfw_parameters
+	tnfw_parameters = _prepare_nfw_parameters()
+	tnfw_parameters['trunc_radius'] = 2.0
+	return tnfw_parameters
 
 
 def _prepare_x_y():
-  rng = jax.random.PRNGKey(0)
-  rng_x, rng_y = jax.random.split(rng)
-  x = jax.random.normal(rng_x, shape=(3,))
-  y = jax.random.normal(rng_y, shape=(3,))
-  return x, y
+	rng = jax.random.PRNGKey(0)
+	rng_x, rng_y = jax.random.split(rng)
+	x = jax.random.normal(rng_x, shape=(3,))
+	y = jax.random.normal(rng_y, shape=(3,))
+	return x, y
 
 
 def _prepare_reduced_radii():
-  rng = jax.random.PRNGKey(0)
-  rng_rr, rng_trr = jax.random.split(rng)
-  reduced_radius = jax.random.uniform(rng_rr, shape=(3,))
-  truncated_reduced_radius = jax.random.uniform(rng_trr, shape=(3,))
-  return reduced_radius, truncated_reduced_radius
+	rng = jax.random.PRNGKey(0)
+	rng_rr, rng_trr = jax.random.split(rng)
+	reduced_radius = jax.random.uniform(rng_rr, shape=(3,))
+	truncated_reduced_radius = jax.random.uniform(rng_trr, shape=(3,))
+	return reduced_radius, truncated_reduced_radius
 
 
 def _prepare__polar_to_cartesian_expected(gamma_ext, angle):
-  if gamma_ext == 0.1 and angle == np.pi / 6:
-    return jnp.array([0.05, 0.0866025404])
-  elif gamma_ext == 2.0 and angle == -np.pi / 3:
-    return jnp.array([-1.0, -1.732050808])
-  elif gamma_ext == 20.3 and angle == np.pi / 22:
-    return jnp.array([19.4777074, 5.719170904])
-  else:
-    raise ValueError(
-        f'gamma_ext={gamma_ext} angle={angle} are not a supported parameter'
-        ' combination.')
+	if gamma_ext == 0.1 and angle == np.pi / 6:
+		return jnp.array([0.05, 0.0866025404])
+	elif gamma_ext == 2.0 and angle == -np.pi / 3:
+		return jnp.array([-1.0, -1.732050808])
+	elif gamma_ext == 20.3 and angle == np.pi / 22:
+		return jnp.array([19.4777074, 5.719170904])
+	else:
+		raise ValueError(
+				f'gamma_ext={gamma_ext} angle={angle} are not a supported parameter'
+				' combination.')
 
 
-class AllTest(absltest.TestCase):
-  """Runs tests of __all__ property of lens_models module."""
+class AllTest(chex.TestCase):
+	"""Runs tests of __all__ property of lens_models module."""
 
-  def test_all(self):
-    all_present = sorted(lens_models.__all__)
-    all_required = []
-    for name, value in inspect.getmembers(lens_models):
-      if inspect.isclass(value):
-        all_required.append(name)
+	def test_all(self):
+		all_present = sorted(lens_models.__all__)
+		all_required = []
+		for name, value in inspect.getmembers(lens_models):
+			if inspect.isclass(value):
+				all_required.append(name)
 
-    self.assertListEqual(all_present, sorted(all_required))
+		self.assertListEqual(all_present, sorted(all_required))
 
 
 class EPLTest(chex.TestCase, parameterized.TestCase):
-  """Runs tests of EPL derivative functions."""
-
-  def test_parameters(self):
-    annotated_parameters = sorted(lens_models.EPL.parameters)
-    correct_parameters = sorted(_prepare_epl_parameters().keys())
-    self.assertListEqual(annotated_parameters, correct_parameters)
-
-  @chex.all_variants
-  def test_derivatives(self):
-    x, y = _prepare_x_y()
-    parameters = _prepare_epl_parameters()
-    expected = jnp.array([[-0.28093671, 0.50107465, -0.92706789],
-                          [1.07278208, -1.03556666, -0.31412716]])
-
-    derivatives = self.variant(lens_models.EPL.derivatives)
-
-    np.testing.assert_allclose(
-        jnp.asarray(derivatives(x, y, **parameters)), expected, rtol=1e-6)
-
-  @chex.all_variants
-  def test__complex_derivative(self):
-    x, y = _prepare_x_y()
-    scale_length = 2.0
-    parameters = _prepare_epl_parameters()
-    expected = jnp.array([
-        0.24284143 + 3.02660008j,
-        0.78504067 - 2.88261008j,
-        -1.42493992 - 2.47468632j,
-    ])
-
-    complex_derivative = self.variant(lens_models.EPL._complex_derivative)
-
-    np.testing.assert_allclose(
-        complex_derivative(x, y, scale_length, parameters['axis_ratio'],
-                           parameters['slope']),
-        expected,
-        rtol=1e-6)
-
-  @chex.all_variants
-  def test__hypergeometric_series(self):
-    ellip_angle = jax.random.normal(jax.random.PRNGKey(0), shape=(3,)) * np.pi
-    parameters = _prepare_epl_parameters()
-    expected = jnp.array([
-        0.83773771 - 0.4450688j,
-        0.0816033 - 1.1360697j,
-        0.60248699 + 0.86020766j,
-    ])
-
-    hypergeometric_series = self.variant(lens_models.EPL._hypergeometric_series)
-
-    np.testing.assert_allclose(
-        hypergeometric_series(ellip_angle, parameters['slope'],
-                              parameters['axis_ratio']), expected)
+	"""Runs tests of EPL derivative functions."""
+
+	def test_parameters(self):
+		annotated_parameters = sorted(lens_models.EPL.parameters)
+		correct_parameters = sorted(_prepare_epl_parameters().keys())
+		self.assertListEqual(annotated_parameters, correct_parameters)
+
+	@chex.all_variants
+	def test_derivatives(self):
+		x, y = _prepare_x_y()
+		parameters = _prepare_epl_parameters()
+		expected = jnp.array([[-0.28093671, 0.50107465, -0.92706789],
+													[1.07278208, -1.03556666, -0.31412716]])
+
+		derivatives = self.variant(lens_models.EPL.derivatives)
+
+		np.testing.assert_allclose(
+				jnp.asarray(derivatives(x, y, **parameters)), expected, rtol=1e-6)
+
+	@chex.all_variants
+	def test__complex_derivative(self):
+		x, y = _prepare_x_y()
+		scale_length = 2.0
+		parameters = _prepare_epl_parameters()
+		expected = jnp.array([
+				0.24284143 + 3.02660008j,
+				0.78504067 - 2.88261008j,
+				-1.42493992 - 2.47468632j,
+		])
+
+		complex_derivative = self.variant(lens_models.EPL._complex_derivative)
+
+		np.testing.assert_allclose(
+				complex_derivative(x, y, scale_length, parameters['axis_ratio'],
+													 parameters['slope']),
+				expected,
+				rtol=1e-6)
+
+	@chex.all_variants
+	def test__hypergeometric_series(self):
+		ellip_angle = jax.random.normal(jax.random.PRNGKey(0), shape=(3,)) * np.pi
+		parameters = _prepare_epl_parameters()
+		expected = jnp.array([
+				0.83773771 - 0.4450688j,
+				0.0816033 - 1.1360697j,
+				0.60248699 + 0.86020766j,
+		])
+
+		hypergeometric_series = self.variant(lens_models.EPL._hypergeometric_series)
+
+		np.testing.assert_allclose(
+				hypergeometric_series(ellip_angle, parameters['slope'],
+															parameters['axis_ratio']), expected)
 
 
 class NFWTest(chex.TestCase, parameterized.TestCase):
-  """Runs tests of NFW derivative functions."""
-
-  def test_parameters(self):
-    annotated_parameters = sorted(lens_models.NFW.parameters)
-    correct_parameters = sorted(_prepare_nfw_parameters().keys())
-    self.assertListEqual(annotated_parameters, correct_parameters)
-
-  @chex.all_variants
-  def test_derivatives(self):
-    x, y = _prepare_x_y()
-    nfw_parameters = _prepare_nfw_parameters()
-    expected = jnp.array([[-0.18378359, 1.1456927, -3.3447392],
-                          [4.0263114, -3.7813387, -1.7911061]])
-
-    derivatives = self.variant(lens_models.NFW.derivatives)
-
-    np.testing.assert_allclose(
-        jnp.array(derivatives(x, y, **nfw_parameters)), expected, rtol=1e-4)
-
-  @chex.all_variants
-  @parameterized.named_parameters([
-      (f'_ars_{ars}_sr_{sr}', ars, sr, expected) for ars, sr, expected in zip(
-          [0.5, 1.2], [0.1, 1.5], [40.736141915886606, 0.43451884710279054])
-  ])
-  def test__alpha_to_rho(self, ars, sr, expected):
-    self.assertAlmostEqual(
-        self.variant(lens_models.NFW._alpha_to_rho)(ars, sr), expected)
-
-  @chex.all_variants
-  def test__nfw_derivatives(self):
-    x, y = _prepare_x_y()
-    radius = jnp.sqrt(x**2 + y**2)
-    rho_input = 2.0
-    nfw_parameters = _prepare_nfw_parameters()
-    expected = jnp.array([[0.42659888, 1.3704396, -2.2209157],
-                          [3.4939308, -3.2851558, -2.4732482]])
-
-    nfw_derivatives = self.variant(lens_models.NFW._nfw_derivatives)
-
-    np.testing.assert_allclose(
-        jnp.array(
-            nfw_derivatives(radius, nfw_parameters['scale_radius'], rho_input,
-                            x, y)),
-        expected,
-        rtol=1e-4)
-
-  @chex.all_variants
-  def test__nfw_integral(self):
-    reduced_radius = jax.random.uniform(jax.random.PRNGKey(0), shape=(3,))
-    expected = jnp.array([0.2952541, 0.07187331, 0.18085146])
-
-    integral = self.variant(lens_models.NFW._nfw_integral)
-
-    np.testing.assert_allclose(
-        jnp.asarray(integral(reduced_radius)), expected, rtol=1e-5)
+	"""Runs tests of NFW derivative functions."""
+
+	def test_parameters(self):
+		annotated_parameters = sorted(lens_models.NFW.parameters)
+		correct_parameters = sorted(_prepare_nfw_parameters().keys())
+		self.assertListEqual(annotated_parameters, correct_parameters)
+
+	@chex.all_variants
+	def test_derivatives(self):
+		x, y = _prepare_x_y()
+		nfw_parameters = _prepare_nfw_parameters()
+		expected = jnp.array([[-0.18378359, 1.1456927, -3.3447392],
+													[4.0263114, -3.7813387, -1.7911061]])
+
+		derivatives = self.variant(lens_models.NFW.derivatives)
+
+		np.testing.assert_allclose(
+				jnp.array(derivatives(x, y, **nfw_parameters)), expected, rtol=1e-4)
+
+	@chex.all_variants
+	@parameterized.named_parameters([
+			(f'_ars_{ars}_sr_{sr}', ars, sr, expected) for ars, sr, expected in zip(
+					[0.5, 1.2], [0.1, 1.5], [40.736141915886606, 0.43451884710279054])
+	])
+	def test__alpha_to_rho(self, ars, sr, expected):
+		self.assertAlmostEqual(
+				self.variant(lens_models.NFW._alpha_to_rho)(ars, sr), expected)
+
+	@chex.all_variants
+	def test__nfw_derivatives(self):
+		x, y = _prepare_x_y()
+		radius = jnp.sqrt(x**2 + y**2)
+		rho_input = 2.0
+		nfw_parameters = _prepare_nfw_parameters()
+		expected = jnp.array([[0.42659888, 1.3704396, -2.2209157],
+													[3.4939308, -3.2851558, -2.4732482]])
+
+		nfw_derivatives = self.variant(lens_models.NFW._nfw_derivatives)
+
+		np.testing.assert_allclose(
+				jnp.array(
+						nfw_derivatives(radius, nfw_parameters['scale_radius'], rho_input,
+														x, y)),
+				expected,
+				rtol=1e-4)
+
+	@chex.all_variants
+	def test__nfw_integral(self):
+		reduced_radius = jax.random.uniform(jax.random.PRNGKey(0), shape=(3,))
+		expected = jnp.array([0.2952541, 0.07187331, 0.18085146])
+
+		integral = self.variant(lens_models.NFW._nfw_integral)
+
+		np.testing.assert_allclose(
+				jnp.asarray(integral(reduced_radius)), expected, rtol=1e-5)
 
 
 class ShearTest(chex.TestCase, parameterized.TestCase):
-  """Runs tests of Shear derivative functions."""
+	"""Runs tests of Shear derivative functions."""
 
-  def test_parameters(self):
-    annotated_parameters = sorted(lens_models.Shear.parameters)
-    correct_parameters = sorted(_prepare_shear_parameters().keys())
-    self.assertListEqual(annotated_parameters, correct_parameters)
+	def test_parameters(self):
+		annotated_parameters = sorted(lens_models.Shear.parameters)
+		correct_parameters = sorted(_prepare_shear_parameters().keys())
+		self.assertListEqual(annotated_parameters, correct_parameters)
 
-  @chex.all_variants
-  def test_derivatives(self):
-    x, y = _prepare_x_y()
-    parameters = _prepare_shear_parameters()
-    expected = jnp.array([[1.2095041, -0.7264168, -0.9143843],
-                          [-1.0582784, 1.0540832, -0.393031]])
+	@chex.all_variants
+	def test_derivatives(self):
+		x, y = _prepare_x_y()
+		parameters = _prepare_shear_parameters()
+		expected = jnp.array([[1.2095041, -0.7264168, -0.9143843],
+													[-1.0582784, 1.0540832, -0.393031]])
 
-    derivatives = self.variant(lens_models.Shear.derivatives)
+		derivatives = self.variant(lens_models.Shear.derivatives)
 
-    np.testing.assert_allclose(
-        jnp.asarray(derivatives(x, y, **parameters)), expected, rtol=1e-5)
+		np.testing.assert_allclose(
+				jnp.asarray(derivatives(x, y, **parameters)), expected, rtol=1e-5)
 
-  @chex.all_variants
-  @parameterized.named_parameters([
-      (f'_ge_{ge}_ang_{ang}', ge, ang)
-      for ge, ang in zip([0.1, 2.0, 20.3], [np.pi / 6, -np.pi / 3, np.pi / 22])
-  ])
-  def test__polar_to_cartesian(self, gamma_ext, angle):
-    expected = _prepare__polar_to_cartesian_expected(gamma_ext, angle)
+	@chex.all_variants
+	@parameterized.named_parameters([
+			(f'_ge_{ge}_ang_{ang}', ge, ang)
+			for ge, ang in zip([0.1, 2.0, 20.3], [np.pi / 6, -np.pi / 3, np.pi / 22])
+	])
+	def test__polar_to_cartesian(self, gamma_ext, angle):
+		expected = _prepare__polar_to_cartesian_expected(gamma_ext, angle)
 
-    np.testing.assert_allclose(
-        jnp.array(
-            self.variant(lens_models.Shear._polar_to_cartesian)(gamma_ext,
-                                                                angle)),
-        expected,
-        rtol=1e-5)
+		np.testing.assert_allclose(
+				jnp.array(
+						self.variant(lens_models.Shear._polar_to_cartesian)(gamma_ext,
+																																angle)),
+				expected,
+				rtol=1e-5)
 
 
 class TNFWTest(chex.TestCase, parameterized.TestCase):
-  """Runs tests of TNFW derivative functions."""
-
-  def test_parameters(self):
-    annotated_parameters = sorted(lens_models.TNFW.parameters)
-    correct_parameters = sorted(_prepare_tnfw_parameters().keys())
-    self.assertListEqual(annotated_parameters, correct_parameters)
+	"""Runs tests of TNFW derivative functions."""
 
-  @chex.all_variants
-  def test_derivatives(self):
-    x, y = _prepare_x_y()
-    tnfw_parameters = _prepare_tnfw_parameters()
-    expected = jnp.array([[-0.13036814, 0.8646701, -2.6658154],
-                          [2.8560917, -2.8538284, -1.4275427]])
+	def test_parameters(self):
+		annotated_parameters = sorted(lens_models.TNFW.parameters)
+		correct_parameters = sorted(_prepare_tnfw_parameters().keys())
+		self.assertListEqual(annotated_parameters, correct_parameters)
 
-    derivatives = self.variant(lens_models.TNFW.derivatives)
+	@chex.all_variants
+	def test_derivatives(self):
+		x, y = _prepare_x_y()
+		tnfw_parameters = _prepare_tnfw_parameters()
+		expected = jnp.array([[-0.13036814, 0.8646701, -2.6658154],
+													[2.8560917, -2.8538284, -1.4275427]])
 
-    np.testing.assert_allclose(
-        jnp.asarray(derivatives(x, y, **tnfw_parameters)), expected, rtol=1e-5)
+		derivatives = self.variant(lens_models.TNFW.derivatives)
 
-  @chex.all_variants
-  def test__tnfw_derivatives(self):
-    x, y = _prepare_x_y()
-    r = jnp.sqrt(x**2 + y**2)
-    tnfw_parameters = _prepare_tnfw_parameters()
-    rho_input = 2.0
-    expected = jnp.array([[0.3161475, 0.97571015, -1.7840269],
-                          [2.5893118, -2.3389282, -1.9868112]])
+		np.testing.assert_allclose(
+				jnp.asarray(derivatives(x, y, **tnfw_parameters)), expected, rtol=1e-5)
 
-    tnfw_derivatives = self.variant(lens_models.TNFW._tnfw_derivatives)
+	@chex.all_variants
+	def test__tnfw_derivatives(self):
+		x, y = _prepare_x_y()
+		r = jnp.sqrt(x**2 + y**2)
+		tnfw_parameters = _prepare_tnfw_parameters()
+		rho_input = 2.0
+		expected = jnp.array([[0.3161475, 0.97571015, -1.7840269],
+													[2.5893118, -2.3389282, -1.9868112]])
 
-    np.testing.assert_allclose(
-        jnp.asarray(
-            tnfw_derivatives(r, tnfw_parameters['scale_radius'], rho_input,
-                             tnfw_parameters['trunc_radius'], x, y)),
-        expected,
-        rtol=1e-5)
+		tnfw_derivatives = self.variant(lens_models.TNFW._tnfw_derivatives)
 
-  @chex.all_variants
-  def test__tnfw_integral(self):
-    reduced_radius, truncated_reduced_radius = _prepare_reduced_radii()
-    expected = jnp.array([0.10208127, 0.01536825, 0.02873538])
+		np.testing.assert_allclose(
+				jnp.asarray(
+						tnfw_derivatives(r, tnfw_parameters['scale_radius'], rho_input,
+														 tnfw_parameters['trunc_radius'], x, y)),
+				expected,
+				rtol=1e-5)
 
-    tnfw_integral = self.variant(lens_models.TNFW._tnfw_integral)
+	@chex.all_variants
+	def test__tnfw_integral(self):
+		reduced_radius, truncated_reduced_radius = _prepare_reduced_radii()
+		expected = jnp.array([0.10208127, 0.01536825, 0.02873538])
 
-    np.testing.assert_allclose(
-        tnfw_integral(reduced_radius, truncated_reduced_radius),
-        expected,
-        rtol=1e-5)
+		tnfw_integral = self.variant(lens_models.TNFW._tnfw_integral)
 
-  @chex.all_variants
-  def test__tnfw_log(self):
-    reduced_radius, truncated_reduced_radius = _prepare_reduced_radii()
-    expected = jnp.array([-1.2336285, -0.15917292, -0.8316463])
+		np.testing.assert_allclose(
+				tnfw_integral(reduced_radius, truncated_reduced_radius),
+				expected,
+				rtol=1e-5)
 
-    tnfw_log = self.variant(lens_models.TNFW._tnfw_log)
+	@chex.all_variants
+	def test__tnfw_log(self):
+		reduced_radius, truncated_reduced_radius = _prepare_reduced_radii()
+		expected = jnp.array([-1.2336285, -0.15917292, -0.8316463])
 
-    np.testing.assert_allclose(
-        tnfw_log(reduced_radius, truncated_reduced_radius), expected, rtol=1e-5)
+		tnfw_log = self.variant(lens_models.TNFW._tnfw_log)
 
-  @chex.all_variants
-  def test__nfw_function(self):
-    reduced_radius = jnp.asarray([0.2, 1.0, 1.5])
-    expected = jnp.array([2.33970328, 1.0, 0.75227469])
+		np.testing.assert_allclose(
+				tnfw_log(reduced_radius, truncated_reduced_radius), expected, rtol=1e-5)
 
-    nfw_function = self.variant(lens_models.TNFW._nfw_function)
+	@chex.all_variants
+	def test__nfw_function(self):
+		reduced_radius = jnp.asarray([0.2, 1.0, 1.5])
+		expected = jnp.array([2.33970328, 1.0, 0.75227469])
 
-    np.testing.assert_allclose(
-        nfw_function(reduced_radius), expected, rtol=1e-5)
+		nfw_function = self.variant(lens_models.TNFW._nfw_function)
 
+		np.testing.assert_allclose(
+				nfw_function(reduced_radius), expected, rtol=1e-5)
 
-if __name__ == '__main__':
-  absltest.main()
diff --git a/jaxstronomy/los_test.py b/jaxstronomy/los_test.py
new file mode 100644
index 0000000..7b043d6
--- /dev/null
+++ b/jaxstronomy/los_test.py
@@ -0,0 +1,469 @@
+# coding=utf-8
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     https://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for los.py."""
+
+import functools
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import chex
+from immutabledict import immutabledict
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from jaxstronomy import cosmology_utils
+from jaxstronomy import los
+from jaxstronomy import power_law
+
+
+COSMOLOGY_PARAMS_INIT = immutabledict({
+    'omega_m_zero': 0.3089,
+    'omega_b_zero': 0.0486,
+    'omega_de_zero': 0.6910088292453472,
+    'omega_rad_zero': 9.117075466e-5,
+    'temp_cmb_zero': 2.7255,
+    'hubble_constant': 67.74,
+    'n_s': 0.9667,
+    'sigma_eight': 0.8159,
+})
+
+
+def _prepare_cosmology_params(cosmology_params_init, z_lookup_max, dz,
+    r_min=1e-4, r_max=1e3, n_r_bins=2):
+    # Only generate a lookup table for values we need.
+    # When 0,0 is specified for the two z values, need to select a small non-zero
+    # values to generate a non-empty table.
+    z_lookup_max = max(z_lookup_max, 1e-7)
+    dz = max(dz, 1e-7)
+    return cosmology_utils.add_lookup_tables_to_cosmology_params(
+        dict(cosmology_params_init), z_lookup_max, dz, r_min, r_max, n_r_bins)
+
+
+def _prepare_cosmology_params_los(los_params, cosmology_params_init,
+    z_lookup_max, dz, r_min=1e-4, r_max=1e3, n_r_bins=2):
+    cosmology_params = _prepare_cosmology_params(cosmology_params_init,
+        z_lookup_max, dz, r_min, r_max, n_r_bins)
+    return los.add_los_lookup_tables_to_cosmology_params(los_params,
+        cosmology_params, z_lookup_max)
+
+
+def _prepare_main_deflector_params():
+    main_deflector_params = {'mass': 1e13, 'z_lens': 0.5, 'theta_e': 2.38,
+        'center_x': 0.0, 'center_y': 0.0}
+    return main_deflector_params
+
+
+def _prepare_source_params():
+    main_deflector_params = {'z_source': 1.5}
+    return main_deflector_params
+
+
+def _prepare_los_params():
+    los_params = {'delta_los': 1.1, 'r_min':0.5, 'r_max':10.0, 'm_min': 1e6,
+        'm_max': 1e10, 'dz':0.1, 'cone_angle': 8.0, 'angle_buffer': 0.8,
+        'c_zero': 18, 'conc_zeta': -0.2, 'conc_beta': 0.8, 'conc_m_ref': 1e8,
+        'conc_dex_scatter': 0.0}
+    return los_params
+
+
+def _prepare_nu_function_expected():
+    return jnp.array([1.53490537e-02, 2.89663922e-02, 5.63195644e-02,
+        1.16683673e-01, 2.56467729e-01, 1.53451398e-01, 1.22804756e-15,
+        0.00000000e+00, 0.00000000e+00, 0.00000000e+00])
+
+
+def _prepare_mass_function_exact_expected(z):
+    if z == 0.1:
+        return jnp.array([5.76805863e-13, 1.38454033e-20])
+    elif z == 0.2:
+        return jnp.array([7.70313721e-13, 1.84294010e-20])
+    elif z == 0.3:
+        return jnp.array([1.00724901e-12, 2.40028314e-20])
+    else:
+        raise ValueError(f'{z} is not a valid input redshift.')
+
+
+class LosTests(chex.TestCase, parameterized.TestCase):
+    """Runs tests of image simulation functions."""
+
+    @chex.all_variants
+    def test_nu_function(self):
+        # Test the nu function for precomputed values.
+        nu = jnp.logspace(-3, 3, 10)
+        expected = _prepare_nu_function_expected()
+
+        nu_function =  self.variant(los.nu_function)
+
+        np.testing.assert_array_almost_equal(nu_function(nu), expected)
+
+    @chex.all_variants(without_device=False)
+    @parameterized.named_parameters([(f'_z_{z}', z) for z in [0.1, 0.2, 0.3]])
+    def test_mass_function_exact(self, z):
+        # Test that the mass function gives the correct exact evaluation.
+        masses = jnp.array([1e6, 1e6])
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z,
+            z)
+        dr = 1e-6
+        lag_r = cosmology_utils.lagrangian_radius(cosmology_params, masses[0])
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z,
+                z, lag_r - dr * 2, lag_r + dr * 2, 5)
+
+        # First expected value is for 1e-6. Second expected value is for 1e-10.
+        expected = _prepare_mass_function_exact_expected(z)
+
+        mass_function = self.variant(los.mass_function_exact)
+
+        self.assertAlmostEqual(mass_function(cosmology_params, masses, z)[0],
+            expected[0])
+
+        masses = jnp.array([1e10, 1e10])
+        lag_r = cosmology_utils.lagrangian_radius(cosmology_params, masses[0])
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z,
+                z, lag_r - dr * 2, lag_r + dr * 2, 5)
+
+        self.assertAlmostEqual(mass_function(cosmology_params, masses, z)[0],
+            expected[1])
+
+    @chex.all_variants(without_device=False)
+    @parameterized.named_parameters([(f'_z_{z}', z) for z in [0.1, 0.2, 0.3]])
+    def test__mass_function_power_law_numerical(self, z):
+        # Make sure that the power law fit is in decent agreement with the los
+        # values.
+        m_min = 1e6
+        m_max = 1e10
+        masses = jnp.logspace(jnp.log10(m_min) , jnp.log10(m_max), 100)
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z,
+            z)
+        r_min = cosmology_utils.lagrangian_radius(cosmology_params, m_min / 10)
+        r_max = cosmology_utils.lagrangian_radius(cosmology_params, m_max * 10)
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z,
+                z, r_min, r_max, 1000)
+        mass_function_exact = los.mass_function_exact(cosmology_params, masses,
+            z)
+
+        slope, norm = self.variant(los._mass_function_power_law_numerical)(
+            cosmology_params, z, m_min, m_max)
+        estimate = norm * masses ** slope
+
+        np.testing.assert_array_almost_equal(estimate / estimate,
+            mass_function_exact / estimate, decimal = 1)
+
+    def test_add_los_lookup_tables_to_cosmology_params(self):
+        # Make sure that the power law fit is in decent agreement with the los
+        # values.
+        los_params = _prepare_los_params()
+        m_min = los_params['m_min']
+        m_max = los_params['m_max']
+        z_lookup_max = 0.5
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT,
+            z_lookup_max, los_params['dz'])
+        r_min = cosmology_utils.lagrangian_radius(cosmology_params, m_min / 10)
+        r_max = cosmology_utils.lagrangian_radius(cosmology_params, m_max * 10)
+        cosmology_params = _prepare_cosmology_params_los(los_params,
+            COSMOLOGY_PARAMS_INIT, z_lookup_max, los_params['dz'], r_min, r_max,
+            1000)
+
+        # Manually test two values
+        i = 2
+        slope_expected, norm_expected = los._mass_function_power_law_numerical(
+            cosmology_params, los_params['dz'] / 2 * i, los_params['m_min'],
+            los_params['m_max']
+        )
+        self.assertAlmostEqual(
+            cosmology_params['mass_function_slope_lookup_table'][i],
+            slope_expected, places=3
+        )
+        self.assertAlmostEqual(
+            cosmology_params['mass_function_norm_lookup_table'][i],
+            norm_expected, places=3
+        )
+
+        i = 4
+        slope_expected, norm_expected = los._mass_function_power_law_numerical(
+            cosmology_params, los_params['dz'] / 2 * i, los_params['m_min'],
+            los_params['m_max']
+        )
+        self.assertAlmostEqual(
+            cosmology_params['mass_function_slope_lookup_table'][i],
+            slope_expected, places=3
+        )
+        self.assertAlmostEqual(
+            cosmology_params['mass_function_norm_lookup_table'][i],
+            norm_expected, places=3
+        )
+
+
+    @chex.all_variants(without_device=False)
+    def test_mass_function_power_law(self):
+        # Check that it agrees with the numerical function.
+        los_params = _prepare_los_params()
+        m_min = los_params['m_min']
+        m_max = los_params['m_max']
+        z_lookup_max = 0.5
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT,
+            z_lookup_max, los_params['dz'])
+        r_min = cosmology_utils.lagrangian_radius(cosmology_params, m_min / 10)
+        r_max = cosmology_utils.lagrangian_radius(cosmology_params, m_max * 10)
+        cosmology_params = _prepare_cosmology_params_los(los_params,
+            COSMOLOGY_PARAMS_INIT, z_lookup_max, los_params['dz'], r_min, r_max,
+            1000)
+
+        z = 0.1
+        mass_function_power_law = self.variant(los.mass_function_power_law)
+        test_slope, test_norm = mass_function_power_law(cosmology_params, z)
+        expected_slope, expected_norm = los._mass_function_power_law_numerical(
+            cosmology_params, z, m_min, m_max)
+        self.assertAlmostEqual(test_slope, expected_slope, places=3)
+        self.assertAlmostEqual(test_norm, expected_norm, places=3)
+
+        z = 0.2-1e-6
+        test_slope, test_norm = mass_function_power_law(cosmology_params, z)
+        expected_slope, expected_norm = los._mass_function_power_law_numerical(
+            cosmology_params, z, m_min, m_max)
+        self.assertAlmostEqual(test_slope, expected_slope, places=3)
+        self.assertAlmostEqual(test_norm, expected_norm, places=3)
+
+
+    @chex.all_variants(without_device=False)
+    def test_cone_angle_to_radius(self):
+        # Test that the conditional statements in jax behave as one would
+        # expect given a manual calculation in python.
+        main_deflector_params = _prepare_main_deflector_params()
+        source_params = _prepare_source_params()
+        los_params = _prepare_los_params()
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, 1.5,
+            0.25)
+
+        cone_angle_to_radius = self.variant(los.cone_angle_to_radius)
+
+        for z in [0.25, 0.5]:
+            hand_calc = cosmology_utils.kpc_per_arcsecond(cosmology_params, z)
+            hand_calc *= los_params['cone_angle'] * 0.5
+            self.assertAlmostEqual(hand_calc, cone_angle_to_radius(
+                    main_deflector_params, source_params, los_params,
+                    cosmology_params, z))
+
+        for z in [0.75, 1.0, 1.25, 1.5]:
+            # These should be shrinking from the main deflector redshift
+            hand_calc = los_params['angle_buffer']
+            hand_calc *= (
+                cosmology_utils.comoving_distance(cosmology_params, 0.0, 1.5) /
+                cosmology_utils.comoving_distance(cosmology_params, 0.5, 1.5))
+            hand_calc *= (
+                cosmology_utils.comoving_distance(cosmology_params, 0.5, z) /
+                cosmology_utils.comoving_distance(cosmology_params, 0.0, z))
+            hand_calc = 1-hand_calc
+            hand_calc *= cosmology_utils.kpc_per_arcsecond(cosmology_params,
+                z)
+            hand_calc *= los_params['cone_angle'] * 0.5
+            self.assertAlmostEqual(hand_calc, cone_angle_to_radius(
+                    main_deflector_params, source_params, los_params,
+                    cosmology_params, z), places=5)
+
+    @chex.all_variants(without_device=False)
+    @parameterized.named_parameters([(f'_z_{z}', z) for z in [0.1, 0.4, 0.8]])
+    def test_volume_element(self, z):
+        # Test that the volume element
+        main_deflector_params = _prepare_main_deflector_params()
+        source_params = _prepare_source_params()
+        los_params = _prepare_los_params()
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT,
+            source_params['z_source'], los_params['dz'] / 2)
+
+        volume_element = self.variant(los.volume_element)
+        los_radius = los.cone_angle_to_radius(main_deflector_params,
+            source_params, los_params, cosmology_params,
+            z + los_params['dz'] / 2)
+        dz_in_kpc = cosmology_utils.comoving_distance(cosmology_params, z,
+            z + los_params['dz'])
+        dz_in_kpc /= (1 + z + los_params['dz'] / 2) / 1e3
+        hand_calc = los_radius ** 2 * np.pi * dz_in_kpc
+
+        self.assertAlmostEqual(volume_element(main_deflector_params,
+            source_params, los_params, cosmology_params, z) / hand_calc,
+            hand_calc/hand_calc, places = 5)
+
+    @chex.all_variants(without_device=False)
+    @parameterized.named_parameters([(f'_z_{z}', z) for z in [0.1, 0.4, 0.8]])
+    def test_expected_num_halos(self, z):
+        # Check that the expected halo counts matches a calculation by hand.
+        main_deflector_params = _prepare_main_deflector_params()
+        source_params = _prepare_source_params()
+        los_params = _prepare_los_params()
+        m_min = los_params['m_min']
+        m_max = los_params['m_max']
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT,
+            source_params['z_source'], los_params['dz'] / 2)
+        r_min = cosmology_utils.lagrangian_radius(cosmology_params, m_min / 10)
+        r_max = cosmology_utils.lagrangian_radius(cosmology_params, m_max * 10)
+        cosmology_params = _prepare_cosmology_params_los(los_params,
+            COSMOLOGY_PARAMS_INIT, source_params['z_source'],
+            los_params['dz'] / 2, r_min, r_max, 1000)
+
+        slope, norm = los.mass_function_power_law(cosmology_params, z)
+        norm *= los.volume_element(main_deflector_params, source_params,
+            los_params, cosmology_params, z)
+        norm *= power_law.power_law_integrate(m_min, m_max, slope)
+        hand_calc = norm * los_params['delta_los']
+
+        expected_num_halos = self.variant(los.expected_num_halos)
+
+        # Low precision for jitted call from float32.
+        self.assertAlmostEqual(expected_num_halos(main_deflector_params,
+            source_params, los_params, cosmology_params, z) / hand_calc,
+            hand_calc / hand_calc, places=3)
+
+    @chex.all_variants(without_device=False)
+    def test_draw_redshifts(self):
+        # Check that the ratio of drawn redshifts approximates the expected
+        # number of los halos in each slice.
+        z_min = 0.3
+        z_max = 0.5
+        main_deflector_params = _prepare_main_deflector_params()
+        source_params = _prepare_source_params()
+        los_params = _prepare_los_params()
+        m_min = los_params['m_min']
+        m_max = los_params['m_max']
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT,
+            source_params['z_source'], los_params['dz'] / 2)
+        r_min = cosmology_utils.lagrangian_radius(cosmology_params, m_min / 10)
+        r_max = cosmology_utils.lagrangian_radius(cosmology_params, m_max * 10)
+        cosmology_params = _prepare_cosmology_params_los(los_params,
+            COSMOLOGY_PARAMS_INIT, source_params['z_source'],
+            los_params['dz'] / 2, r_min, r_max, 1000)
+        num_expected_min = los.expected_num_halos(main_deflector_params,
+            source_params, los_params, cosmology_params, z_min)
+        num_expected_max = los.expected_num_halos(main_deflector_params,
+            source_params, los_params, cosmology_params,
+            z_max - los_params['dz'])
+
+        num_z_bins = 1000
+        pad_length = 3000
+        draw_redshifts = self.variant(functools.partial(los.draw_redshifts,
+            num_z_bins=num_z_bins, pad_length=pad_length))
+
+        rng = jax.random.PRNGKey(0)
+        z_draws = draw_redshifts(main_deflector_params, source_params,
+            los_params, cosmology_params, z_min, z_max, rng)
+
+        # Doing many draws is expensive, so we won't be picky.
+        self.assertAlmostEqual(
+            num_expected_max / (num_expected_max + num_expected_min),
+            jnp.sum(z_draws > 0.4) / jnp.sum(z_draws > 0.3), places = 1)
+        self.assertAlmostEqual(1.0,
+            jnp.sum(z_draws > 0.3) / (num_expected_max + num_expected_min),
+            places = 1)
+
+    @chex.all_variants(without_device=False)
+    def test_draw_masses(self):
+        # Testing that the masses are roughly drawn from the desired power law.
+        source_params = _prepare_source_params()
+        los_params = _prepare_los_params()
+        m_min = los_params['m_min']
+        m_max = los_params['m_max']
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT,
+            source_params['z_source'], los_params['dz'] / 2)
+        r_min = cosmology_utils.lagrangian_radius(cosmology_params, m_min / 10)
+        r_max = cosmology_utils.lagrangian_radius(cosmology_params, m_max * 10)
+        cosmology_params = _prepare_cosmology_params_los(los_params,
+            COSMOLOGY_PARAMS_INIT, source_params['z_source'],
+            los_params['dz'] / 2, r_min, r_max, 1000)
+
+        rng = jax.random.PRNGKey(0)
+        z = 0.3
+        z_values = jnp.full((1000,), z)
+
+        draw_masses = self.variant(los.draw_masses)
+
+        slope, _ = los.mass_function_power_law(cosmology_params, z)
+        integrate_below = power_law.power_law_integrate(m_min, 5e6, slope)
+        integrate_above = power_law.power_law_integrate(5e6, m_max, slope)
+
+        masses = draw_masses(los_params, cosmology_params, z_values, rng)
+        # Draws are expensive in the non-jit case, so don't be picky.
+        self.assertAlmostEqual(jnp.sum(masses > 5e6) / jnp.sum(masses < 5e6),
+            integrate_above / integrate_below, places=1)
+
+    @chex.all_variants(without_device=False)
+    @parameterized.named_parameters([(f'_z_{z}', z) for z in [0.1, 0.4, 0.8]])
+    def test_draw_positions(self, z):
+        # Test that the positions are uniformly distributed within a disk.
+        main_deflector_params = _prepare_main_deflector_params()
+        source_params = _prepare_source_params()
+        los_params = _prepare_los_params()
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z,
+            z)
+
+        rng = jax.random.PRNGKey(0)
+        z_values = jnp.full((10000,), z)
+        draw_positions = self.variant(los.draw_positions)
+
+        los_pos = draw_positions(main_deflector_params, source_params,
+            los_params, cosmology_params, z_values, rng)
+
+        # Check the angle distribution
+        angles = jnp.arctan(los_pos[:,1] / los_pos[:,0])
+        self.assertAlmostEqual(np.mean(angles<0),np.mean(angles>0),places=1)
+
+        # Make sure that within a square they are uniformly distributed
+        bound = jnp.sqrt(2) * los.cone_angle_to_radius(main_deflector_params,
+            source_params, los_params, cosmology_params, z)
+        mask = jnp.logical_and(jnp.abs(los_pos[:,0]) < bound / 2,
+            jnp.abs(los_pos[:,1]) < bound / 2)
+        los_pos = los_pos[mask]
+
+        self.assertAlmostEqual(2 * jnp.mean(los_pos[:,0] < -bound / 4),
+            jnp.mean(los_pos[:,0] > 0), places=1)
+        self.assertAlmostEqual(2 * jnp.mean(los_pos[:,1] < -bound / 4),
+            jnp.mean(los_pos[:,1] > 0), places=1)
+
+    @chex.all_variants(without_device=False)
+    def test_draw_los(self):
+        # Some basic tests that the final output of draw_los is sensible.
+        main_deflector_params = _prepare_main_deflector_params()
+        source_params = _prepare_source_params()
+        los_params = _prepare_los_params()
+        # Don't want to draw too many halos
+        los_params['m_min'] = 1e8
+        m_min = los_params['m_min']
+        m_max = los_params['m_max']
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT,
+            source_params['z_source'], los_params['dz'] / 2)
+        r_min = cosmology_utils.lagrangian_radius(cosmology_params, m_min / 10)
+        r_max = cosmology_utils.lagrangian_radius(cosmology_params, m_max * 10)
+        cosmology_params = _prepare_cosmology_params_los(los_params,
+            COSMOLOGY_PARAMS_INIT, source_params['z_source'],
+            los_params['dz'] / 2, r_min, r_max, 1000)
+
+        rng = jax.random.PRNGKey(0)
+        num_z_bins = 1000
+        los_pad_length = 100
+        draw_los = self.variant(functools.partial(los.draw_los,
+            num_z_bins=num_z_bins, los_pad_length=los_pad_length))
+
+        los_tuple = draw_los(main_deflector_params, source_params, los_params,
+            cosmology_params, rng)
+
+        # Simple tests on the redshifts. TODO add more tests on the parameters.
+        los_before_z = los_tuple[0][0]
+        los_after_z = los_tuple[1][0]
+
+        # Check bounds on redshifts that are drawn.
+        self.assertEqual(jnp.sum(los_before_z >= 0.5), 0.0)
+        self.assertEqual(jnp.sum(los_before_z < 0.0),
+            jnp.sum(los_before_z == -1.0))
+        self.assertEqual(jnp.sum(los_after_z >= 1.5), 0.0)
+        self.assertEqual(jnp.sum(los_after_z < 0.5),
+            jnp.sum(los_after_z == -1.0))
diff --git a/jaxstronomy/nfw_functions_test.py b/jaxstronomy/nfw_functions_test.py
new file mode 100644
index 0000000..44ce092
--- /dev/null
+++ b/jaxstronomy/nfw_functions_test.py
@@ -0,0 +1,239 @@
+# coding=utf-8
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     https://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for nfw_functions.py."""
+
+import functools
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import chex
+from immutabledict import immutabledict
+import jax
+import jax.numpy as jnp
+import numpy as np
+from scipy.integrate import quad
+
+from jaxstronomy import cosmology_utils
+from jaxstronomy import nfw_functions
+
+
+COSMOLOGY_PARAMS_INIT = immutabledict({
+    'omega_m_zero': 0.3089,
+    'omega_b_zero': 0.0486,
+    'omega_de_zero': 0.6910088292453472,
+    'omega_rad_zero': 9.117075466e-5,
+    'temp_cmb_zero': 2.7255,
+    'hubble_constant': 67.74,
+    'n_s': 0.9667,
+    'sigma_eight': 0.8159,
+})
+
+
+def _prepare_cosmology_params(cosmology_params_init, z_lookup_max, dz,
+    r_min=1e-4, r_max=1e3, n_r_bins=2):
+    # Only generate a lookup table for values we need.
+    # When 0,0 is specified for the two z values, need to select a small non-zero
+    # values to generate a non-empty table.
+    z_lookup_max = max(z_lookup_max, 1e-7)
+    dz = max(dz, 1e-7)
+    return cosmology_utils.add_lookup_tables_to_cosmology_params(
+        dict(cosmology_params_init), z_lookup_max, dz, r_min, r_max, n_r_bins)
+
+
+def _prepare_main_deflector_params():
+    main_deflector_params = {'mass': 1e13, 'z_lens': 0.5, 'theta_e': 2.38,
+        'center_x': 0.0, 'center_y': 0.0}
+    return main_deflector_params
+
+
+def _prepare_substructure_params():
+    substructure_params = {'c_zero': 18, 'conc_zeta': -0.2, 'conc_beta': 0.8, 
+        'conc_m_ref': 1e8, 'conc_dex_scatter': 0.0}
+    return substructure_params
+
+
+class NfwFuntionsTests(chex.TestCase, parameterized.TestCase):
+    """Runs tests of image simulation functions."""
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+        (f'_m_{np.log10(m):.0f}_z_{z}', m, z, expected) for m, z, expected in
+        zip([1e7, 1e8, 1e9], [0.1, 0.2, 0.3], [4.39736, 9.14635, 18.98417])
+    ])
+    def test_r_two_hund_from_m(self, m, z, expected):
+        # Test that for scalars and arrays the correct overdensity radius is
+        # returned.
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z,
+            z)
+
+        r_two_hund_from_m =  self.variant(nfw_functions.r_two_hund_from_m)
+
+        self.assertAlmostEqual(r_two_hund_from_m(cosmology_params, m, z),
+            expected, places=4)
+
+        np.testing.assert_array_almost_equal(
+            r_two_hund_from_m(cosmology_params,jnp.full(10, m), z),
+            jnp.full(10, expected), decimal=4)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+        (f'_c_{c}_z_{z}', c, z, expected) for c, z, expected in
+        zip([8, 12, 14], [0.1, 0.2, 0.3],
+        [3662396.563509418, 10945906.187043915, 17983344.644410152])
+    ])
+    def test_rho_nfw_from_c(self, c, z, expected):
+        # Test that for scalars and arrays the correct normalization is
+        # returned.
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, z,
+            z)
+
+        rho_nfw_from_c =  self.variant(nfw_functions.rho_nfw_from_c)
+
+        # Divide by the order of magntiude since testing is in terms of decimal.
+        self.assertAlmostEqual(rho_nfw_from_c(cosmology_params, c, z) / 100000,
+            expected / 100000, places=4)
+
+        np.testing.assert_array_almost_equal(
+            rho_nfw_from_c(cosmology_params,jnp.full(10, c), z) / 100000,
+            jnp.full(10, expected) / 100000, decimal=4)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+        (f'_rt_{rt}_rho_{rho}_rs_{rs}', rt, rho, rs) for rt, rho, rs in
+        zip([0.1, 1.5, 12.2], [1, 10, 100], [2, 4, 8])
+    ])
+    def test__cored_nfw_integral(self, rt, rho, rs):
+        # Compare analytic answer to numerical integral
+        r_upper = jnp.linspace(0.1, 4, 100)
+
+        _cored_nfw_integral =  self.variant(nfw_functions._cored_nfw_integral)
+
+        def cored_nfw_func(r):
+            if r<rt:
+                x_tidal = rt / rs
+                return 4 * np.pi * rho / (x_tidal * (1 + x_tidal) ** 2) * r ** 2
+            else:
+                x = r / rs
+                return 4 * np.pi * r ** 2 * rho / (x * (1 + x) ** 2)
+
+        analytic_values = _cored_nfw_integral(rt, rho, rs, r_upper)
+
+        for i, r in enumerate(r_upper):
+            self.assertAlmostEqual(1.0,
+                quad(cored_nfw_func, 0, r, epsabs=1.49e-11,
+				    epsrel=1.49e-11)[0] / analytic_values[i], places=4)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+        (f'_rt_{rt}_rho_{rho}_rs_{rs}', rt, rho, rs) for rt, rho, rs in
+        zip([0.1, 1.5, 12.2], [1, 10, 100], [2, 4, 8])
+    ])
+    def test_cored_nfw_draws(self, rt, rho, rs):
+        # Compare to expected distributions.
+        r_max = 4
+        n_draws = int(2e5)
+        rng = jax.random.PRNGKey(0)
+
+        cored_nfw_draws =  self.variant(functools.partial(
+            nfw_functions.cored_nfw_draws, n_draws=n_draws))
+        r_draws = cored_nfw_draws(rt, rho, rs, r_max, rng)
+
+        n_test_points = 100
+        r_test = jnp.linspace(0, r_max, n_test_points)
+        analytic_values = nfw_functions._cored_nfw_integral(rt, rho, rs, r_test)
+
+        for i, r in enumerate(r_test):
+            self.assertAlmostEqual(jnp.mean(r_draws < r),
+                analytic_values[i] / jnp.max(analytic_values), places=2)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+        (f'_rs_{rs}_rho_{rho}', rs, rho, eo, et) for rs, rho, eo, et in
+        zip([0.1, 1.5, 12.2], [1e13, 2e13, 3e13],
+            [0.015902785701266995, 0.23854178551900487, 1.940139855554573],
+            [6.503977420097886, 2926.7898390440478, 290415.5997622107])
+    ])
+    def test_convert_to_lensing_nfw(self, rs, rho, eo, et):
+        # Compare to expected distributions.
+        z = 0.5
+        z_source = 1.0
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, 1.0,
+            0.5)
+
+        convert_to_lensing_nfw =  self.variant(
+            nfw_functions.convert_to_lensing_nfw)
+
+        r_s_ang, alpha_rs = convert_to_lensing_nfw(cosmology_params, rs, z, rho,
+            z_source)
+        self.assertAlmostEqual(r_s_ang, eo, places=4)
+        self.assertAlmostEqual(1.0, et / alpha_rs, places=4)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+        (f'_rt_{rt}', rs, rho, rt, expected) for rs, rho, rt, expected in
+        zip([0.1, 1.5, 12.2], [1e13, 2e13, 3e13], [0.1, 1.0, 1.5],
+            [0.015902785701266995, 0.15902785701266992, 0.23854178551900487])
+    ])
+    def test_convert_to_lensing_tnfw(self, rs, rho, rt, expected):
+        # Compare to expected distributions.
+        z = 0.5
+        z_source = 1.0
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT, 1.0,
+            0.5)
+
+        convert_to_lensing_tnfw =  self.variant(
+            nfw_functions.convert_to_lensing_tnfw)
+
+        _, _, rt_ang = convert_to_lensing_tnfw(cosmology_params, rs, z, rho, rt,
+            z_source)
+        self.assertAlmostEqual(rt_ang, expected, places=4)
+
+    @chex.all_variants(without_device=False)
+    @parameterized.named_parameters([
+        (f'_m_{m}_z_{z}', m, z, expected) for m, z, expected in zip(
+            [1e9, 1e10], [0.2, 0.3], [13.7, 10.8])
+    ])
+    def test_mass_concentration(self, m, z, expected):
+        # Test that the mass draws follow the desired distribution.
+        main_deflector_params = _prepare_main_deflector_params()
+        subhalo_params = _prepare_substructure_params()
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT,
+            main_deflector_params['z_lens'], main_deflector_params['z_lens'])
+        h = cosmology_params['hubble_constant'] / 100
+        r_min = cosmology_utils.lagrangian_radius(cosmology_params,
+            subhalo_params['conc_m_ref'] * h)
+        r_max = cosmology_utils.lagrangian_radius(cosmology_params, m * h)
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT,
+            main_deflector_params['z_lens'], main_deflector_params['z_lens'],
+            r_min, r_max)
+
+        mass_concentration = self.variant(nfw_functions.mass_concentration)
+
+        rng = jax.random.PRNGKey(0)
+        subhalo_params['conc_dex_scatter'] = 0.0
+        conc = mass_concentration(subhalo_params, cosmology_params,
+            jnp.full((10,), m), z, rng)
+        # Peak height is the limit on the precision of this comparison.
+        np.testing.assert_array_almost_equal(conc, jnp.full((10,), expected),
+            decimal=0)
+
+        # Check the scatter does something
+        subhalo_params['conc_dex_scatter'] = 0.1
+        conc = mass_concentration(subhalo_params, cosmology_params,
+            jnp.full((10000,), m), z, rng)
+        scatter = jnp.log10(conc) - expected
+        self.assertAlmostEqual(jnp.std(scatter), 0.1, places=2)
+
+if __name__ == '__main__':
+    absltest.main()
diff --git a/jaxstronomy/power_law_test.py b/jaxstronomy/power_law_test.py
new file mode 100644
index 0000000..22662cb
--- /dev/null
+++ b/jaxstronomy/power_law_test.py
@@ -0,0 +1,90 @@
+# coding=utf-8
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     https://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for power_law.py."""
+
+import functools
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import chex
+import jax
+import jax.numpy as jnp
+from scipy.integrate import quad
+
+from jaxstronomy import power_law
+
+
+class PowerLawTests(chex.TestCase, parameterized.TestCase):
+    """Runs tests of image simulation functions."""
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+        (f'_p_min_{p_min}_p_max_{p_max}_slope_{slope}', p_min, p_max, slope)
+        for p_min, p_max, slope in zip([3, 6, 4], [7, 9, 4], [-1.9, -1.7, 5])
+    ])
+    def test_power_law_integrate(self, p_min, p_max, slope):
+        # Check that analytic integral agrees with numerical results.
+        def p_func(x, slope):
+            return x**slope
+
+        pl_integrate =  self.variant(power_law.power_law_integrate)
+
+        self.assertAlmostEqual(pl_integrate(p_min, p_max, slope),
+            quad(p_func, p_min, p_max, args=(slope))[0])
+
+    @chex.all_variants
+    def test_power_law_draw(self):
+        # Check that draws roughly follow the power law and the norm
+        # we expect
+        p_min = 1e6
+        p_max = 1e9
+        slope = -1.9
+        desired_count = 100
+        norm = desired_count / power_law.power_law_integrate(p_min, p_max,
+            slope)
+
+        pl_draw = self.variant(functools.partial(power_law.power_law_draw,
+            pad_length = 200))
+
+        total_subs = 0
+        rng = jax.random.PRNGKey(0)
+        n_loops = 500
+        for _ in range(n_loops):
+            rng_draw, rng = jax.random.split(rng)
+            draws = pl_draw(p_min, p_max, slope, norm, rng_draw)
+            total_subs += jnp.sum(draws > 0)
+        self.assertEqual(jnp.round(total_subs/n_loops) , desired_count)
+
+        # Check for a specific draw that it follows the cdf
+        desired_count = 1e6
+        pad_length = int(2e6)
+        norm = desired_count / power_law.power_law_integrate(p_min, p_max,
+            slope)
+        pl_draw = self.variant(functools.partial(power_law.power_law_draw,
+            pad_length = pad_length))
+        draws = pl_draw(p_min, p_max, slope, norm, rng)
+
+        # Check the cdf at a handful of test points
+        test_points = jnp.logspace(6, 9, 20)
+        for test_point in test_points:
+            self.assertAlmostEqual(
+                jnp.sum(jnp.logical_and(draws < test_point , draws > 0)) /
+                jnp.sum(draws > 0),
+                power_law.power_law_integrate(p_min, test_point, slope) /
+                power_law.power_law_integrate(p_min, p_max, slope), places=2
+            )
+
+
+if __name__ == '__main__':
+    absltest.main()
diff --git a/jaxstronomy/source_models_test.py b/jaxstronomy/source_models_test.py
index 65d30f1..3d3dc45 100644
--- a/jaxstronomy/source_models_test.py
+++ b/jaxstronomy/source_models_test.py
@@ -1,5 +1,4 @@
 # coding=utf-8
-# Copyright 2022 The Google Research Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -7,20 +6,6 @@
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Copyright 2022 Google LLC
-
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-
-#     https://www.apache.org/licenses/LICENSE-2.0
-
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -40,7 +25,7 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from third_party.google_research.google_research.jaxstronomy import source_models
+from jaxstronomy import source_models
 
 
 def _prepare_x_y():
diff --git a/jaxstronomy/subhalos_test.py b/jaxstronomy/subhalos_test.py
new file mode 100644
index 0000000..585ee12
--- /dev/null
+++ b/jaxstronomy/subhalos_test.py
@@ -0,0 +1,290 @@
+# coding=utf-8
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     https://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for subhalos.py."""
+
+import functools
+
+from absl.testing import parameterized
+import chex
+from immutabledict import immutabledict
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from jaxstronomy import cosmology_utils
+from jaxstronomy import nfw_functions
+from jaxstronomy import power_law
+from jaxstronomy import subhalos
+
+
+COSMOLOGY_PARAMS_INIT = immutabledict({
+    'omega_m_zero': 0.3089,
+    'omega_b_zero': 0.0486,
+    'omega_de_zero': 0.6910088292453472,
+    'omega_rad_zero': 9.117075466e-5,
+    'temp_cmb_zero': 2.7255,
+    'hubble_constant': 67.74,
+    'n_s': 0.9667,
+    'sigma_eight': 0.8159,
+})
+
+
+def _prepare_cosmology_params(cosmology_params_init, z_lookup_max, dz,
+    r_min=1e-4, r_max=1e3, n_r_bins=2):
+    # Only generate a lookup table for values we need.
+    # When 0,0 is specified for the two z values, need to select a small non-zero
+    # values to generate a non-empty table.
+    z_lookup_max = max(z_lookup_max, 1e-7)
+    dz = max(dz, 1e-7)
+    return cosmology_utils.add_lookup_tables_to_cosmology_params(
+        dict(cosmology_params_init), z_lookup_max, dz, r_min, r_max, n_r_bins)
+
+
+def _prepare_main_deflector_params():
+    main_deflector_params = {'mass': 1e13, 'z_lens': 0.5, 'theta_e': 2.38,
+        'center_x': jnp.array([0.01]), 'center_y': jnp.array([-0.03])}
+    return main_deflector_params
+
+
+def _prepare_subhalo_params():
+    subhalo_params = {'sigma_sub': 5e-4, 'shmf_plaw_index': -2.0,
+        'm_pivot': 1e8, 'm_min': 1e6, 'm_max': 1e9, 'k_one': 0.0, 'k_two': 0.0,
+        'c_zero': 18, 'conc_zeta': -0.2, 'conc_beta': 0.8, 'conc_m_ref': 1e8,
+        'conc_dex_scatter': 0.0}
+    return subhalo_params
+
+
+def _prepare_test_truncation_radius_expected():
+    return jnp.array([0.22223615, 0.74981341, 1.4133784, 2.32006694,
+        3.57303672,5.30341642,7.68458502,10.94766194,15.40115448, 21.45666411])
+
+
+class SubhalosTests(chex.TestCase, parameterized.TestCase):
+    """Runs tests of image simulation functions."""
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+        (f'_m_{m}_z_{z}_ko_{ko}_kt_{kt}', m, z, ko, kt, expected) for
+            m, z, ko, kt, expected in zip([1e13, 2e13, 1e13], [0.1, 0.2, 0.3],
+                [0.0, 0.5, 1.0], [0.0, 0.5, 1.0], [1.0, 1.183215956619923, 0.8])
+    ])
+    def test_host_scaling_function(self, m, z, ko, kt, expected):
+        # Test the scaling function for different parameter values.
+
+        host_scaling_function =  self.variant(subhalos.host_scaling_function)
+
+        self.assertAlmostEqual(host_scaling_function(m, z, ko, kt),
+            expected, places=4)
+
+    @chex.all_variants
+    def test_draw_nfw_masses(self):
+        # Test that the mass draws follow the desired distribution.
+        main_deflector_params = _prepare_main_deflector_params()
+        subhalo_params = _prepare_subhalo_params()
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT,
+            main_deflector_params['z_lens'], main_deflector_params['z_lens'])
+
+        pad_length = 1000
+        draw_nfw_masses = self.variant(functools.partial(
+            subhalos.draw_nfw_masses, pad_length=pad_length))
+
+        # Calculate the norm by hand and make sure the statistics agree
+        kpa = cosmology_utils.kpc_per_arcsecond(cosmology_params,
+            main_deflector_params['z_lens'])
+        area_elem = jnp.pi * (3 * kpa * main_deflector_params['theta_e']) ** 2
+        f_host = subhalos.host_scaling_function(main_deflector_params['mass'],
+            main_deflector_params['z_lens'], subhalo_params['k_one'],
+            subhalo_params['k_two'])
+
+        e_counts =  power_law.power_law_integrate(subhalo_params['m_min'],
+            subhalo_params['m_max'],subhalo_params['shmf_plaw_index'])
+        norm = area_elem * f_host*(subhalo_params['m_pivot']**(-
+                subhalo_params['shmf_plaw_index']-1))
+        e_total = e_counts * norm * subhalo_params['sigma_sub']
+
+        rng = jax.random.PRNGKey(0)
+        n_loops = 100
+        total = 0
+        for _ in range(n_loops):
+            rng_draw, rng = jax.random.split(rng)
+            total += jnp.sum(draw_nfw_masses(main_deflector_params,
+                subhalo_params, cosmology_params, rng_draw) > 0)
+        self.assertAlmostEqual(total / n_loops, e_total, places = 0)
+
+    @chex.all_variants
+    def test_rejection_sampling(self):
+        # Test that the distributions are respected and that the keep
+        # calculation is correct
+        rng = jax.random.PRNGKey(0)
+        radial_coord = jax.random.uniform(rng, (int(1e6),))
+        r_two_hund = 0.5
+        r_bound = 0.7
+
+        rejection_sampling = self.variant(subhalos.rejection_sampling)
+
+        is_inside, cart_pos = rejection_sampling(radial_coord, r_two_hund,
+            r_bound, rng)
+
+        np.testing.assert_array_almost_equal(jnp.sqrt(jnp.sum(cart_pos ** 2,
+            axis = -1)), radial_coord)
+        phi = jnp.arccos(cart_pos[:, 2] / radial_coord)
+        theta = jnp.arctan(cart_pos[:, 1] / cart_pos[:, 0])
+        self.assertAlmostEqual(jnp.mean(phi), jnp.pi/2, places=2)
+        self.assertAlmostEqual(jnp.mean(theta), 0.0, places=2)
+
+        # Manually confirm bounds
+        self.assertEqual(
+            jnp.sum(is_inside[radial_coord * jnp.sin(phi) > r_bound]), 0.0)
+        self.assertEqual(
+            jnp.sum(is_inside[radial_coord * jnp.cos(phi) > r_two_hund]), 0.0)
+
+    @chex.all_variants(without_device=False)
+    def test_sample_cored_nfw(self):
+        # Test that for large enough padding the bounds are met and the
+        # distribution of points is correct.
+        main_deflector_params = _prepare_main_deflector_params()
+        z_lens = main_deflector_params['z_lens']
+        subhalo_params = _prepare_subhalo_params()
+
+        # Initialize our cosmology dictionary
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT,
+            z_lens, z_lens)
+        dr = 1e-6
+        lag_r = cosmology_utils.lagrangian_radius(cosmology_params,
+            main_deflector_params['mass'])
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT,
+            z_lens, z_lens, lag_r - dr * 2, lag_r + dr * 2, 5)
+
+        n_subs = 1000
+        sampling_pad_length = 200000
+        sample_cored_nfw = self.variant(functools.partial(
+            subhalos.sample_cored_nfw, n_subs=n_subs,
+            sampling_pad_length=sampling_pad_length))
+
+        rng = jax.random.PRNGKey(0)
+        cart_pos = sample_cored_nfw(main_deflector_params, subhalo_params,
+            cosmology_params, rng)
+
+        # Check that all the positions are within the desired volume.
+        radial_coord = jnp.sqrt(jnp.sum(jnp.square(cart_pos), axis=-1))
+        phi = jnp.arccos(cart_pos[:, 2] / radial_coord)
+        theta = jnp.arctan(cart_pos[:, 1] / cart_pos[:, 0])
+        kpa = cosmology_utils.kpc_per_arcsecond(cosmology_params, z_lens)
+        r_bound = main_deflector_params['theta_e'] * 3 * kpa
+        r_two_hund = nfw_functions.r_two_hund_from_m(cosmology_params,
+            main_deflector_params['mass'], z_lens)
+        self.assertAlmostEqual(jnp.mean(phi), jnp.pi/2, places=1)
+        self.assertAlmostEqual(jnp.mean(theta), 0.0, places=1)
+        self.assertEqual(jnp.mean(radial_coord * jnp.sin(phi) > r_bound), 0.0)
+        self.assertEqual(jnp.mean(radial_coord * jnp.cos(phi) > r_two_hund),
+            0.0)
+
+    @chex.all_variants
+    def test_truncation_radius(self):
+        # Run some precomputed values to check agreement.
+        m_two_hund = jnp.logspace(6, 9, 10)
+        radii = jnp.linspace(10, 300, 10)
+        expected = _prepare_test_truncation_radius_expected()
+
+        get_truncation_radius = self.variant(subhalos.get_truncation_radius)
+
+        np.testing.assert_array_almost_equal(get_truncation_radius(m_two_hund,
+            radii), expected, decimal=5)
+
+        # Make sure a truncation radius of 0 cannot be returned.
+        self.assertGreater(get_truncation_radius(0.0, 1e-4), 0.0)
+        self.assertGreater(get_truncation_radius(1e-4, 0.0), 0.0)
+
+    @chex.all_variants(without_device=False)
+    def test_convert_to_lensing(self):
+        main_deflector_params = _prepare_main_deflector_params()
+        z_lens = main_deflector_params['z_lens']
+        z_source = 1.5
+        source_params = {'z_source': z_source}
+        subhalo_params = _prepare_subhalo_params()
+        rng = jax.random.PRNGKey(0)
+
+        # Initialize our cosmology dictionary
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT,
+            z_source, 0.01)
+        r_min = cosmology_utils.lagrangian_radius(cosmology_params, 1e7)
+        r_max = cosmology_utils.lagrangian_radius(cosmology_params,
+            main_deflector_params['mass'] * 10)
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT,
+            z_source, 0.01, r_min, r_max, 1000)
+
+        subhalo_masses = jnp.logspace(8, 10, 100)
+        subhalo_masses *= jnp.arange(0, len(subhalo_masses)) > 2
+        subhalo_cart_pos = jax.random.normal(rng, shape=(100, 3))
+
+        convert_to_lensing = self.variant(subhalos.convert_to_lensing)
+
+        subhalos_z, subhalos_kwargs = convert_to_lensing(
+            main_deflector_params, source_params, subhalo_params,
+            cosmology_params, subhalo_masses, subhalo_cart_pos, rng)
+
+        np.testing.assert_array_almost_equal(subhalos_z,
+            jnp.ones(subhalos_z.shape) * z_lens)
+
+        # Test that the lensing quantities follow manual calculations.
+        kpa = cosmology_utils.kpc_per_arcsecond(cosmology_params, z_lens)
+        np.testing.assert_array_almost_equal(subhalos_kwargs['center_x'],
+            subhalo_cart_pos[:,0] / kpa + main_deflector_params['center_x'])
+        np.testing.assert_array_almost_equal(subhalos_kwargs['center_y'],
+            subhalo_cart_pos[:,1] / kpa + main_deflector_params['center_y'])
+        # Check that the model indexing worked correctly.
+        self.assertAlmostEqual(jnp.mean(
+            subhalos_kwargs['model_index'][subhalo_masses > 0]), 0.0)
+        self.assertAlmostEqual(jnp.mean(
+            subhalos_kwargs['model_index'][subhalo_masses == 0]), -1.0)
+
+        self.assertTrue('alpha_rs' in subhalos_kwargs)
+        self.assertTrue('scale_radius' in subhalos_kwargs)
+        self.assertTrue('trunc_radius' in subhalos_kwargs)
+
+    @chex.all_variants(without_device=False)
+    def test_draw_subhalos(self):
+        main_deflector_params = _prepare_main_deflector_params()
+        z_lens = main_deflector_params['z_lens']
+        z_source = 1.5
+        source_params = {'z_source': z_source}
+        subhalo_params = _prepare_subhalo_params()
+        rng = jax.random.PRNGKey(0)
+
+        # Initialize our cosmology dictionary
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT,
+            z_source, 0.01)
+        r_min = cosmology_utils.lagrangian_radius(cosmology_params, 1e7)
+        r_max = cosmology_utils.lagrangian_radius(cosmology_params,
+            main_deflector_params['mass'] * 10)
+        cosmology_params = _prepare_cosmology_params(COSMOLOGY_PARAMS_INIT,
+            z_source, 0.01, r_min, r_max, 1000)
+
+        subhalos_pad_length = 100
+        sampling_pad_length = 10000
+        draw_subhalos = self.variant(functools.partial(subhalos.draw_subhalos,
+            subhalos_pad_length=subhalos_pad_length,
+            sampling_pad_length=sampling_pad_length))
+
+        subhalos_z, subhalos_kwargs = draw_subhalos(main_deflector_params,
+            source_params, subhalo_params, cosmology_params, rng)
+
+        np.testing.assert_array_almost_equal(subhalos_z,
+            jnp.ones(subhalos_z.shape) * z_lens)
+        self.assertTrue('alpha_rs' in subhalos_kwargs)
+        self.assertTrue('scale_radius' in subhalos_kwargs)
+        self.assertTrue('trunc_radius' in subhalos_kwargs)
+        self.assertTrue('model_index' in subhalos_kwargs)
+        self.assertTrue('center_y' in subhalos_kwargs)
+        self.assertTrue('center_x' in subhalos_kwargs)
diff --git a/jaxstronomy/train_test.py b/jaxstronomy/train_test.py
new file mode 100644
index 0000000..6f91ca6
--- /dev/null
+++ b/jaxstronomy/train_test.py
@@ -0,0 +1,95 @@
+# coding=utf-8
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     https://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for train.py."""
+
+import functools
+
+from absl.testing import parameterized
+import chex
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from jaxstronomy import models
+from jaxstronomy import train
+
+class TrainTests(chex.TestCase, parameterized.TestCase):
+    """Runs tests of image simulation functions."""
+
+    def test_initialized(self):
+        # Test the the model includes initialized weights and biases and
+        # can be applied.
+        rng = jax.random.PRNGKey(0)
+        model = models.ResNet50(num_outputs=2, dtype=jnp.float32)
+        image_size = 128
+        params, batch_stats = train.initialized(rng, image_size, model)
+
+        self.assertTupleEqual(params['conv_init']['kernel'].shape,
+            (7, 7, 1, 64))
+        self.assertTupleEqual(params['bn_init']['scale'].shape,
+            (64,))
+        self.assertTupleEqual(batch_stats['bn_init']['mean'].shape,
+            (64,))
+
+        output, _ = model.apply({'params':params, 'batch_stats':batch_stats},
+            jax.random.normal(rng, shape=(3, image_size, image_size, 1)),
+            mutable=['batch_stats'])
+        self.assertTupleEqual(output.shape, (3, 2))
+        np.testing.assert_array_less(jnp.zeros(2), jnp.std(output, axis=0))
+
+    @chex.all_variants
+    def test_gaussian_loss(self):
+        # Test that the loss function is equivalent to applying a Gaussian.
+        batch_size = int(1e7)
+        rng = jax.random.PRNGKey(0)
+        outputs = jnp.stack([jax.random.normal(rng, (batch_size,)),
+            jnp.zeros((batch_size,))], axis=-1)
+        truth = jnp.zeros((batch_size,1))
+        gaussian_loss = self.variant(train.gaussian_loss)
+
+        # With variance 1 and dropping the factor of 2 pi, the negative entropy
+        # should simply be 0.5
+        self.assertAlmostEqual(gaussian_loss(outputs, truth), 0.5, places=2)
+        outputs = jnp.stack([
+            jax.random.normal(rng, (batch_size,)) * jnp.sqrt(jnp.e),
+            jnp.ones((batch_size,))], axis=-1)
+        self.assertAlmostEqual(gaussian_loss(outputs, truth), 1.0, places=2)
+
+        # We can repeat this experiment with a two dimensional output.
+        outputs = jnp.concatenate([jax.random.normal(rng, (batch_size, 2)),
+            jnp.zeros((batch_size, 2))], axis=-1)
+        truth = jnp.zeros((batch_size, 2))
+        self.assertAlmostEqual(gaussian_loss(outputs, truth), 1.0, places=2)
+
+
+    @chex.all_variants
+    def test_compute_metrics(self):
+        # Test the the computed metrics match expectations.
+        batch_size = int(1e3)
+        rng = jax.random.PRNGKey(0)
+        outputs = jnp.stack([jax.random.normal(rng, (batch_size,)),
+            jnp.zeros((batch_size,))], axis=-1)
+        truth = jnp.zeros((batch_size,1))
+        gaussian_loss = train.gaussian_loss(outputs, truth)
+        rmse = jnp.sqrt(jnp.mean(jnp.square(outputs[:,0] - truth)))
+
+        outputs = jnp.expand_dims(outputs, axis=0)
+        truth = jnp.expand_dims(truth, axis=0)
+
+        compute_metrics = jax.pmap(self.variant(train.compute_metrics),
+                                   axis_name='batch')
+        metrics = compute_metrics(outputs, truth)
+
+        self.assertAlmostEqual(metrics['rmse'], rmse, places=4)
+        self.assertAlmostEqual(metrics['loss'], gaussian_loss, places=4)
diff --git a/jaxstronomy/utils_test.py b/jaxstronomy/utils_test.py
index 93fc24f..b28cfe1 100644
--- a/jaxstronomy/utils_test.py
+++ b/jaxstronomy/utils_test.py
@@ -1,26 +1,11 @@
 # coding=utf-8
-# Copyright 2022 The Google Research Authors.
-#
+
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Copyright 2022 Google LLC
-
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-
-#     https://www.apache.org/licenses/LICENSE-2.0
-
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -38,85 +23,88 @@
 from jaxstronomy import utils
 
 COSMOLOGY_DICT = {
-    'omega_m_zero': 0.30966,
-    'omega_de_zero': 0.688846306,
-    'omega_rad_zero': 0.0014937,
-    'hubble_constant': 67.66,
+        'omega_m_zero': 0.30966,
+        'omega_de_zero': 0.688846306,
+        'omega_rad_zero': 0.0014937,
+        'hubble_constant': 67.66,
 }
 
 
 def _prepare_x_y_angular():
-  rng = jax.random.PRNGKey(3)
-  rng_x, rng_y = jax.random.split(rng)
-  x = jax.random.normal(rng_x, shape=(3,))
-  y = jax.random.normal(rng_y, shape=(3,))
-  return x, y
+    rng = jax.random.PRNGKey(3)
+    rng_x, rng_y = jax.random.split(rng)
+    x = jax.random.normal(rng_x, shape=(3,))
+    y = jax.random.normal(rng_y, shape=(3,))
+    return x, y
 
 
 def _prepare_kwargs_detector():
-  return {'n_x': 2, 'n_y': 2, 'pixel_width': 0.04, 'supersampling_factor': 2}
+    return {'n_x': 2, 'n_y': 2, 'pixel_width': 0.04, 'supersampling_factor': 2}
 
 
 class UtilsTest(chex.TestCase, parameterized.TestCase):
-  """Runs tests of utility functions."""
-
-  def test_coordinates_evaluate(self):
-    parameters = _prepare_kwargs_detector()
-    expected_x = jnp.array([
-        -0.03, -0.01, 0.01, 0.03, -0.03, -0.01, 0.01, 0.03, -0.03, -0.01, 0.01,
-        0.03, -0.03, -0.01, 0.01, 0.03
-    ])
-    expected_y = jnp.array([
-        -0.03, -0.03, -0.03, -0.03, -0.01, -0.01, -0.01, -0.01, 0.01, 0.01,
-        0.01, 0.01, 0.03, 0.03, 0.03, 0.03
+    """Runs tests of utility functions."""
+
+    def test_coordinates_evaluate(self):
+        parameters = _prepare_kwargs_detector()
+        expected_x = jnp.array([
+                -0.03, -0.01, 0.01, 0.03, -0.03, -0.01, 0.01, 0.03, -0.03,
+                -0.01, 0.01, 0.03, -0.03, -0.01, 0.01, 0.03
+        ])
+        expected_y = jnp.array([
+                -0.03, -0.03, -0.03, -0.03, -0.01, -0.01, -0.01, -0.01, 0.01,
+                0.01, 0.01, 0.01, 0.03, 0.03, 0.03, 0.03
+        ])
+
+        np.testing.assert_allclose(
+                jnp.array(utils.coordinates_evaluate(**parameters)),
+                jnp.array([expected_x, expected_y]))
+
+    @chex.all_variants
+    def test_unpack_parameters_xy(self):
+        x, y = _prepare_x_y_angular()
+        kwargs_lens = {
+                'alpha_rs': 1.0,
+                'scale_radius': 1.0,
+                'center_x': 0.0,
+                'center_y': 0.0,
+                'fake_param': 19.2
+        }
+        expected = jnp.array([[-0.90657, -0.29612964, 0.22304466],
+                              [0.44380534, -0.9504099, 0.97678715]])
+
+        unpack_parameters_derivatives = self.variant(
+                utils.unpack_parameters_xy(lens_models.NFW.derivatives,
+                                           lens_models.NFW.parameters))
+
+        np.testing.assert_allclose(
+                unpack_parameters_derivatives(x, y, kwargs_lens), expected,
+                rtol=1e-5)
+
+    def test_downsampling(self):
+        downsample = utils.downsample
+
+        image = jnp.ones((12, 12))
+        np.testing.assert_allclose(downsample(image, 3), jnp.ones((4, 4)))
+        np.testing.assert_allclose(downsample(image, 4), jnp.ones((3, 3)))
+
+        image = jax.random.normal(jax.random.PRNGKey(0), shape=(4, 4))
+        expected = jnp.array([[0.37571156, 0.7770451],
+                              [-0.67193794, 0.014301866]])
+        np.testing.assert_allclose(downsample(image, 2), expected)
+
+    @chex.all_variants
+    @parameterized.named_parameters([
+            (f'_mag_{mag}_mzp_{mzp}', mag, mzp, expected)
+            for mag, mzp, expected in zip(
+            [20.35, -2, 15], [18.0, 5.6, 8.8],
+            [0.11481536214968811, 1096.4781961431852, 0.0033113112148259144])
     ])
+    def test_magnitude_to_cps(self, mag, mzp, expected):
+        magnitude_to_cps = self.variant(utils.magnitude_to_cps)
 
-    np.testing.assert_allclose(
-        jnp.array(utils.coordinates_evaluate(**parameters)),
-        jnp.array([expected_x, expected_y]))
-
-  @chex.all_variants
-  def test_unpack_parameters_xy(self):
-    x, y = _prepare_x_y_angular()
-    kwargs_lens = {
-        'alpha_rs': 1.0,
-        'scale_radius': 1.0,
-        'center_x': 0.0,
-        'center_y': 0.0,
-        'fake_param': 19.2
-    }
-    expected = jnp.array([[-0.90657, -0.29612964, 0.22304466],
-                          [0.44380534, -0.9504099, 0.97678715]])
-
-    unpack_parameters_derivatives = self.variant(
-        utils.unpack_parameters_xy(lens_models.NFW.derivatives,
-                                   lens_models.NFW.parameters))
-
-    np.testing.assert_allclose(
-        unpack_parameters_derivatives(x, y, kwargs_lens), expected, rtol=1e-5)
-
-  def test_downsampling(self):
-    downsample = utils.downsample
-
-    image = jnp.ones((12, 12))
-    np.testing.assert_allclose(downsample(image, 3), jnp.ones((4, 4)))
-    np.testing.assert_allclose(downsample(image, 4), jnp.ones((3, 3)))
-
-    image = jax.random.normal(jax.random.PRNGKey(0), shape=(4, 4))
-    expected = jnp.array([[0.37571156, 0.7770451], [-0.67193794, 0.014301866]])
-    np.testing.assert_allclose(downsample(image, 2), expected)
-
-  @chex.all_variants
-  @parameterized.named_parameters([
-      (f'_mag_{mag}_mzp_{mzp}', mag, mzp, expected) for mag, mzp, expected in
-      zip([20.35, -2, 15], [18.0, 5.6, 8.8],
-          [0.11481536214968811, 1096.4781961431852, 0.0033113112148259144])
-  ])
-  def test_magnitude_to_cps(self, mag, mzp, expected):
-    magnitude_to_cps = self.variant(utils.magnitude_to_cps)
-
-    self.assertAlmostEqual(magnitude_to_cps(mag, mzp), expected, places=5)
+        self.assertAlmostEqual(magnitude_to_cps(mag, mzp), expected, places=5)
 
 
 if __name__ == '__main__':
-  absltest.main()
+    absltest.main()