From 3db2d0cdf6e24d90aa679d077f4131201dfd001a Mon Sep 17 00:00:00 2001
From: Anastasiia Kutakh <anastkut@mit.edu>
Date: Thu, 5 Aug 2021 22:55:43 -0400
Subject: [PATCH] .

---
 CDTools/tools/polarization/polarization.py |  5 ++--
 tests/tools/test_polarization.py           | 27 ++++++++++++++++++++++
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/CDTools/tools/polarization/polarization.py b/CDTools/tools/polarization/polarization.py
index ad4a27f..f6352b9 100644
--- a/CDTools/tools/polarization/polarization.py
+++ b/CDTools/tools/polarization/polarization.py
@@ -12,7 +12,8 @@ __all__ = ['apply_linear_polarizer',
            'apply_half_wave_plate',
            'apply_quarter_wave_plate',
            'apply_circular_polarizer',
-           'apply_jones_matrix']
+           'apply_jones_matrix',
+           'generate_linear_polarizer']
 
 
 # Abe - split these into two functions
@@ -29,7 +30,7 @@ def generate_linear_polarizer(pol_angle):
     pol_angle_rad = t.deg2rad(pol_angle)
     jones_matrices = t.stack([t.tensor([[(t.cos(p)) ** 2, t.sin(p) * t.cos(p)],
                                         [t.sin(p) * t.cos(p), (t.sin(p)) ** 2]])
-                              for p in polarizer])
+                              for p in pol_angle_rad])
     if single_angle:
         return jones_matrices[0].to(dtype=t.cfloat) 
     else:
diff --git a/tests/tools/test_polarization.py b/tests/tools/test_polarization.py
index 0f37b1a..74b8f94 100644
--- a/tests/tools/test_polarization.py
+++ b/tests/tools/test_polarization.py
@@ -1,7 +1,9 @@
 import numpy as np
 import torch as t
+from CDTools.tools.polarization import apply_linear_polarizer, generate_linear_polarizer
 from CDTools.tools.polarization import apply_jones_matrix as jones
 
+
 # Abe - I removed all the imports that didn't need to be here.
 
 # Abe - A few issues. First, you could just write "from math import cos, sin"
@@ -77,6 +79,31 @@ def test_apply_jones_matrix_no_modes_no_mult_patterns_one_jones_matr():
     assert out.shape == t.Size((2, 3, 4))
 
 
+def test_generate_linear_polarizer():
+    pol_angles = [0, 45, 90]
+    pol_angle1 = 45
+    pol_angle2 = t.tensor(90)
+    pol_angle3 = t.tensor([0])
+    pols = generate_linear_polarizer(pol_angles)
+    pol1 = generate_linear_polarizer(pol_angle1)
+    pol2 = generate_linear_polarizer(pol_angle2)
+    pol3 = generate_linear_polarizer(pol_angle3)
+    print('polarizers 0, 45, 90 (1D tensor) shape:', pols.shape)
+    print('shape of the polarizer generated from int:', pol1.shape)
+    print('shape of the polarizer generated from 0D tensor:', pol2.shape)
+    print('shape of the linear polarizer generated from t.Size(0) tensor:', pol3.shape)
+    print('90', pol2)
+    print(jones90)
+    probe = t.ones(4, 4)
+    jones_m = [jones0, jones45, jones90]
+    jones_m = t.stack([matr for matr in jones_m])
+    assert pols.shape == t.Size((3, 2, 2))
+    assert pol1.shape == t.Size((2, 2))
+    assert pol2.shape == t.Size((2, 2))
+    assert pol3.shape == t.Size((1, 2, 2))
+    assert t.allclose(pol1, jones45)
+
+
 def test_apply_jones_matrix_no_modes_no_mult_patterns_diff_jones_matr():    
     '''probe: no multiple modes, 1 diffr pattern
         2xMxL = 2x4x4