Skip to content

Core

Spectrum

MonoSpectrum

Bases: Spectrum

A spectrum with a scalar wavelength (density is a delta function).

Attributes:

Name Type Description
wavelength Float[Array, '1']

A scalar representing the wavelength in the same units of distance as the sampling of the corresponding field.
density Float[Array, '1']

A scalar set to 1.0 because there is only one wavelength in the spectrum.
Source code in src/chromatix/core/spectrum.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
class MonoSpectrum(Spectrum):
    """
    A spectrum with a scalar wavelength (density is a delta function).

    Attributes:
        wavelength: A scalar representing the wavelength in the same units of
            distance as the sampling of the corresponding field.
        density: A scalar set to 1.0 because there is only one wavelength in
            the spectrum.
    """

    wavelength: Float[Array, "1"]
    density: Float[Array, "1"]

    def __init__(self, wavelength: float | Array):
        self.wavelength = jnp.atleast_1d(jnp.asarray(wavelength))
        assert self.wavelength.shape == (1,), (
            f"MonoSpectrum requires a scalar wavelength, got array of shape {self.wavelength.shape}"
        )
        self.density = jnp.ones((1,))

density = jnp.ones((1,)) instance-attribute

wavelength = jnp.atleast_1d(jnp.asarray(wavelength)) instance-attribute

__init__(wavelength)

Source code in src/chromatix/core/spectrum.py 
123
124
125
126
127
128
def __init__(self, wavelength: float | Array):
    self.wavelength = jnp.atleast_1d(jnp.asarray(wavelength))
    assert self.wavelength.shape == (1,), (
        f"MonoSpectrum requires a scalar wavelength, got array of shape {self.wavelength.shape}"
    )
    self.density = jnp.ones((1,))

Spectrum

Bases: Module

A spectrum defined by wavelengths and their corresponding densities.

Attributes:

Name Type Description
wavelength Float[Array, 'wv']

A 1D array representing the wavelengths in the discretized spectrum in the same units of distance as the sampling of the corresponding field.
density Float[Array, 'wv']

A 1D array representing the weight of each wavelength in the discretized spectrum.
Source code in src/chromatix/core/spectrum.py 
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
class Spectrum(eqx.Module):
    """
    A spectrum defined by wavelengths and their corresponding densities.

    Attributes:
        wavelength: A 1D array representing the wavelengths in the discretized
            spectrum in the same units of distance as the sampling of the
            corresponding field.
        density: A 1D array representing the weight of each wavelength in the
            discretized spectrum.
    """

    wavelength: Float[Array, "wv"]
    density: Float[Array, "wv"]

    def __init__(self, wavelength: float | Array, density: float | Array | None = None):
        wavelength = jnp.atleast_1d(jnp.asarray(wavelength))
        self.wavelength = eqx.error_if(
            wavelength,
            wavelength.ndim != 1,
            "Wavelength must be 1D for polychromatic spectrum.",
        )

        if density is None:
            self.density = jnp.full_like(self.wavelength, 1 / self.wavelength.size)
        else:
            density = jnp.atleast_1d(jnp.asarray(density))
            density = density / density.sum()
            self.density = eqx.error_if(
                density,
                density.shape != self.wavelength.shape,
                f"Density shape {density.shape} must match wavelength shape {self.wavelength.shape}.",
            )

    @classmethod
    def build(
        cls,
        spectrum: Spectrum
        | ScalarLike
        | Float[Array, "wv"]
        | tuple[Float[Array, "wv"], Float[Array, "wv"]],
    ) -> Spectrum | MonoSpectrum:
        """
        Function that creates a spectrum type from either a single float
        (MonoSpectrum), an array of multiple wavelengths, or a tuple of two
        arrays: one array of wavelengths and another array of the relative power
        of each wavelength in the spectrum.

        Args:
            spectrum: If this is a scalar float value, a single wavelength
                ``MonoSpectrum`` will be created. If this is a single array
                of floats, a ``Spectrum`` of multiple wavelengths that each
                have equivalent intensity will be created. If this is a tuple,
                the first element of the tuple will be used as the array of
                wavelengths and the second element of the tuple will be used
                as the array of weights representing the relative power of each
                wavelength in the spectrum.
        Returns:
            An object representing a spectrum (or a single wavelength in the
            monochrome case).
        """
        if isinstance(spectrum, Spectrum):
            return spectrum
        elif isinstance(spectrum, float) or isinstance(spectrum, int):
            return MonoSpectrum(wavelength=spectrum)
        elif isinstance(spectrum, tuple) or isinstance(spectrum, list):
            return Spectrum(wavelength=spectrum[0], density=spectrum[1])
        elif isinstance(spectrum, ArrayLike) and spectrum.size == 1:
            return MonoSpectrum(wavelength=spectrum.squeeze())
        else:
            return Spectrum(wavelength=spectrum)

    @property
    def size(self) -> int:
        """Returns the number of wavelengths in the discretized spectrum."""
        return self.wavelength.size

    @property
    def central_wavelength(self) -> float:
        """
        The central wavelength of the spectrum (defined as the first wavelength
        provided to the spectrum because you could construct the complex field
        with multiple wavelengths in any order).
        """
        return self.wavelength[0]

    @property
    def spectral_modulation(self) -> Array:
        """
        Returns the ratio of the central wavelength to all the wavelengths
        in the spectrum. This is useful when calculating phase responses of a
        material that scales approximately linearly with wavelength.
        """
        return self.central_wavelength / self.wavelength

    def __repr__(self) -> str:
        return eqx.tree_pformat(self, short_arrays=False)

central_wavelength property

The central wavelength of the spectrum (defined as the first wavelength provided to the spectrum because you could construct the complex field with multiple wavelengths in any order).

density instance-attribute

size property

Returns the number of wavelengths in the discretized spectrum.

spectral_modulation property

Returns the ratio of the central wavelength to all the wavelengths in the spectrum. This is useful when calculating phase responses of a material that scales approximately linearly with wavelength.

wavelength = eqx.error_if(wavelength, wavelength.ndim != 1, 'Wavelength must be 1D for polychromatic spectrum.') instance-attribute

__init__(wavelength, density=None)

Source code in src/chromatix/core/spectrum.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
def __init__(self, wavelength: float | Array, density: float | Array | None = None):
    wavelength = jnp.atleast_1d(jnp.asarray(wavelength))
    self.wavelength = eqx.error_if(
        wavelength,
        wavelength.ndim != 1,
        "Wavelength must be 1D for polychromatic spectrum.",
    )

    if density is None:
        self.density = jnp.full_like(self.wavelength, 1 / self.wavelength.size)
    else:
        density = jnp.atleast_1d(jnp.asarray(density))
        density = density / density.sum()
        self.density = eqx.error_if(
            density,
            density.shape != self.wavelength.shape,
            f"Density shape {density.shape} must match wavelength shape {self.wavelength.shape}.",
        )

__repr__()

Source code in src/chromatix/core/spectrum.py 
105
106
def __repr__(self) -> str:
    return eqx.tree_pformat(self, short_arrays=False)

build(spectrum) classmethod

Function that creates a spectrum type from either a single float (MonoSpectrum), an array of multiple wavelengths, or a tuple of two arrays: one array of wavelengths and another array of the relative power of each wavelength in the spectrum.

Parameters:

Name Type Description Default
spectrum Spectrum | ScalarLike | Float[Array, 'wv'] | tuple[Float[Array, 'wv'], Float[Array, 'wv']]

If this is a scalar float value, a single wavelength MonoSpectrum will be created. If this is a single array of floats, a Spectrum of multiple wavelengths that each have equivalent intensity will be created. If this is a tuple, the first element of the tuple will be used as the array of wavelengths and the second element of the tuple will be used as the array of weights representing the relative power of each wavelength in the spectrum.

 required

Returns: An object representing a spectrum (or a single wavelength in the monochrome case).

Source code in src/chromatix/core/spectrum.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
@classmethod
def build(
    cls,
    spectrum: Spectrum
    | ScalarLike
    | Float[Array, "wv"]
    | tuple[Float[Array, "wv"], Float[Array, "wv"]],
) -> Spectrum | MonoSpectrum:
    """
    Function that creates a spectrum type from either a single float
    (MonoSpectrum), an array of multiple wavelengths, or a tuple of two
    arrays: one array of wavelengths and another array of the relative power
    of each wavelength in the spectrum.

    Args:
        spectrum: If this is a scalar float value, a single wavelength
            ``MonoSpectrum`` will be created. If this is a single array
            of floats, a ``Spectrum`` of multiple wavelengths that each
            have equivalent intensity will be created. If this is a tuple,
            the first element of the tuple will be used as the array of
            wavelengths and the second element of the tuple will be used
            as the array of weights representing the relative power of each
            wavelength in the spectrum.
    Returns:
        An object representing a spectrum (or a single wavelength in the
        monochrome case).
    """
    if isinstance(spectrum, Spectrum):
        return spectrum
    elif isinstance(spectrum, float) or isinstance(spectrum, int):
        return MonoSpectrum(wavelength=spectrum)
    elif isinstance(spectrum, tuple) or isinstance(spectrum, list):
        return Spectrum(wavelength=spectrum[0], density=spectrum[1])
    elif isinstance(spectrum, ArrayLike) and spectrum.size == 1:
        return MonoSpectrum(wavelength=spectrum.squeeze())
    else:
        return Spectrum(wavelength=spectrum)

Base

Field = TypeVar('Field') module-attribute

__all__ = ['Monochromatic', 'Chromatic', 'Scalar', 'Vector', 'Sample', 'Absorbing', 'Scattering', 'Fluorescent', 'Volume', 'Sensor', 'Resampler'] module-attribute

Absorbing

Bases: Module

Source code in src/chromatix/core/base.py 
143
144
class Absorbing(eqx.Module, strict=_strict_config):
    absorption: eqx.AbstractVar[Array]

absorption instance-attribute

Chromatic

Bases: Module

Source code in src/chromatix/core/base.py 
50
51
class Chromatic(eqx.Module, strict=_strict_config):
    spectrum: eqx.AbstractVar[Spectrum]

spectrum instance-attribute

Fluorescent

Bases: Module

Source code in src/chromatix/core/base.py 
151
152
class Fluorescent(eqx.Module, strict=_strict_config):
    fluorescence: eqx.AbstractVar[Array]

fluorescence instance-attribute

Monochromatic

Bases: Module

Source code in src/chromatix/core/base.py 
42
43
44
45
46
47
class Monochromatic(eqx.Module, strict=_strict_config):
    spectrum: eqx.AbstractVar[MonoSpectrum]

    @property
    def wavelength(self) -> MonoSpectrum:
        return self.spectrum

spectrum instance-attribute

wavelength property

Resampler

Bases: Module

Source code in src/chromatix/core/base.py 
175
176
177
178
179
180
181
class Resampler(eqx.Module, strict=_strict_config):
    out_shape: eqx.AbstractVar[tuple[int, ...]]
    out_spacing: eqx.AbstractVar[Array]

    @abc.abstractmethod
    def __call__(self, resample_input: Array, in_spacing: Array) -> Array:
        pass

out_shape instance-attribute

out_spacing instance-attribute

__call__(resample_input, in_spacing) abstractmethod

Source code in src/chromatix/core/base.py 
179
180
181
@abc.abstractmethod
def __call__(self, resample_input: Array, in_spacing: Array) -> Array:
    pass

Sample

Bases: Module

Source code in src/chromatix/core/base.py 
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
class Sample(eqx.Module, strict=_strict_config):
    dx: eqx.AbstractVar[ScalarLike | Float[Array, "2"]]
    thickness: eqx.AbstractVar[ScalarLike | Array]

    def _verify_matching_spacing(self, field: Field):
        # NOTE(dd/2025-07-17): A field may have different spacings for each
        # wavelength of its spectrum along with different numbers of dimensions
        # depending on the type of the field, so this subtraction allows us
        # to broadcast all the potential spacings from the field to the single
        # spacing of the sample (just in case). This is why we don't just do a
        # simple equality check.
        assert jnp.all((field.dx - self.dx) == 0), (
            "Incoming field must have same sampling as sample"
        )

    def _verify_scalar_thickness(self):
        assert np.asarray(self.thickness).size == 1, (
            "Thickness must be a scalar (uniform thickness of each plane of the sample)"
        )

    @property
    @abc.abstractmethod
    def shape(self) -> tuple[int, ...]:
        pass

    @property
    @abc.abstractmethod
    def ndim(self) -> int:
        pass

    @property
    @abc.abstractmethod
    def grid(self) -> Array:
        pass

    @abc.abstractmethod
    def __call__(self, field: Field) -> Field:
        pass

dx instance-attribute

grid abstractmethod property

ndim abstractmethod property

shape abstractmethod property

thickness instance-attribute

__call__(field) abstractmethod

Source code in src/chromatix/core/base.py 
138
139
140
@abc.abstractmethod
def __call__(self, field: Field) -> Field:
    pass

Scalar

Bases: Module

Source code in src/chromatix/core/base.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
class Scalar(eqx.Module, strict=_strict_config):
    u: eqx.AbstractVar[Array]
    dx: eqx.AbstractVar[Array]
    spectrum: eqx.AbstractVar[Spectrum]
    dims: eqx.AbstractClassVar[IntEnum]

    @property
    def wavelength(self) -> Array:
        return self.spectrum.wavelength

    @property
    def power(self):
        area = jnp.prod(self.dx, axis=-1)
        total_intensity = self.spectrum.density * jnp.sum(
            jnp.abs(self.u) ** 2, axis=(self.dims.p, self.dims.y, self.dims.x)
        )
        return area * total_intensity

dims instance-attribute

dx instance-attribute

power property

spectrum instance-attribute

u instance-attribute

wavelength property

Scattering

Bases: Module

Source code in src/chromatix/core/base.py 
147
148
class Scattering(eqx.Module, strict=_strict_config):
    dn: eqx.AbstractVar[Array]

dn instance-attribute

Sensor

Bases: Module

Source code in src/chromatix/core/base.py 
162
163
164
165
166
167
168
169
170
171
172
class Sensor(eqx.Module, strict=_strict_config):
    shape: eqx.AbstractVar[tuple[int, ...]]
    spacing: eqx.AbstractVar[Array]

    @abc.abstractmethod
    def __call__(self, sensor_input: Field | Array) -> Array:
        pass

    @abc.abstractmethod
    def resample(self, resample_input: Array, input_spacing: ScalarLike) -> Array:
        pass

shape instance-attribute

spacing instance-attribute

__call__(sensor_input) abstractmethod

Source code in src/chromatix/core/base.py 
166
167
168
@abc.abstractmethod
def __call__(self, sensor_input: Field | Array) -> Array:
    pass

resample(resample_input, input_spacing) abstractmethod

Source code in src/chromatix/core/base.py 
170
171
172
@abc.abstractmethod
def resample(self, resample_input: Array, input_spacing: ScalarLike) -> Array:
    pass

Vector

Bases: Module

Source code in src/chromatix/core/base.py 
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
class Vector(eqx.Module, strict=_strict_config):
    u: eqx.AbstractVar[Array]
    dx: eqx.AbstractVar[Array]
    spectrum: eqx.AbstractVar[Spectrum]
    dims: eqx.AbstractClassVar[IntEnum]

    @property
    def jones_vector(self) -> Array:
        norm = jnp.linalg.norm(self.u, axis=self.dims.p, keepdims=True)
        norm = jnp.where(norm == 0, 1, norm)  # Set to 1 to avoid division by zero
        return self.u / norm

    @property
    def power(self):
        area = jnp.prod(self.dx, axis=-1)
        total_intensity = self.spectrum.density * jnp.sum(
            jnp.abs(self.u) ** 2, axis=(self.dims.p, self.dims.y, self.dims.x)
        )
        return area * total_intensity

    @property
    def wavelength(self) -> Array:
        return rearrange(self.spectrum.wavelength, "wv -> wv 1")

    @property
    def intensity(self):
        spectral_density = rearrange(self.spectrum.density, "... wv -> ... 1 1 wv")
        return spectral_density * jnp.sum(jnp.abs(self.u) ** 2, axis=self.dims.p)

dims instance-attribute

dx instance-attribute

intensity property

jones_vector property

power property

spectrum instance-attribute

u instance-attribute

wavelength property

Volume

Bases: Module

Source code in src/chromatix/core/base.py 
155
156
157
158
159
class Volume(eqx.Module, strict=_strict_config):
    @property
    @abc.abstractmethod
    def num_planes(self) -> int:
        pass

num_planes abstractmethod property

replace(tree, **kwargs)

Source code in src/chromatix/core/base.py 
36
37
38
39
def replace(tree: eqx.Module, **kwargs) -> eqx.Module:
    for key, value in kwargs.items():
        tree = eqx.tree_at(lambda t: getattr(t, key), tree, value)
    return tree