search

chromatix-team / chromatix

Differentiable wave optics using JAX! Documentation can be found at https://chromatix.readthedocs.io
MIT License
77 stars 13 forks source link

Force the amplitude of a field to be len 3 for vector fields #112

Closed gschlafly closed 5 months ago

gschlafly commented 6 months ago

A vectorfield plane wave can be defined as follows:

field = cf.plane_wave(
        (1, 1),
        1.0,
        0.532,
        1.0,
        amplitude=[1, 1, 1],
        scalar=False,
    )

If the amplitude=[1, 1] with scalar=False then the field is not compatible with the polarization functions.

diptodip commented 6 months ago

Should be possible in a JAX-friendly way with https://chex.readthedocs.io/en/latest/api.html#chex.assert_axis_dimension

GJBoth commented 5 months ago

Closed by #138