diff --git a/CDTools/tools/polarization/polarization.py b/CDTools/tools/polarization/polarization.py index c6764d6..1efc5e9 100644 --- a/CDTools/tools/polarization/polarization.py +++ b/CDTools/tools/polarization/polarization.py @@ -29,7 +29,8 @@ def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True linearly polarized probe: t.Tensor (N)(P)x2x1xMxL """ - + if len(polarizer.shape) == 0: + polarizer = t.tensor([polarizer]) pol_cos = lambda idx: cos(math.radians(polarizer[idx])) pol_sin = lambda idx: sin(math.radians(polarizer[idx])) jones_matrices = t.stack(([t.tensor([[(pol_cos(idx)) ** 2, pol_sin(idx) * pol_cos(idx)], [pol_sin(idx) * pol_cos(idx), (pol_sin(idx)) ** 2]]).to(dtype=t.cfloat) for idx in range(len(polarizer))]))