fix tests, examples

examples/Si_optimize_diffraction_angle.py

@@ -1,18 +1,12 @@
 from rayflare.rigorous_coupled_wave_analysis import rcwa_structure
-from rayflare.transfer_matrix_method import tmm_structure
 from solcore import material, si
 from solcore.structure import Layer
 from rayflare.options import default_options
 import numpy as np
 import matplotlib.pyplot as plt
-from solcore.light_source import LightSource
-from solcore.constants import q
-import seaborn as sns
-from matplotlib import rc
-from joblib import Parallel, delayed
-from time import time
 import pygmo as pg
 from rayflare.analytic.diffraction import get_order_directions
+import seaborn as sns
 
 # Paper: https://doi.org/10.1002/adom.201700585
 MgF2 = material("MgF2")()  # MgF2 (SOPRA database)
@@ -30,6 +24,8 @@ def __init__(self,
         self.wavelengths = wavelengths
         # self.angle = target_diffraction_angle
         # self.tolerance = angle_tolerance
+        self.target = target_diffraction_angle
+        self.tolerance = angle_tolerance
         self.min_angle = np.pi - target_diffraction_angle - angle_tolerance
         self.max_angle = np.pi - target_diffraction_angle + angle_tolerance
         self.pol = pol
@@ -38,9 +34,13 @@ def calculate(self, x):
         # wavelengths = np.array([1e-6])
 
         # lattice vectors for the grating. Units are in nm!
+        # Note: annoyingly have to set options here internally every time, since the options object
+        # doesn't like the parallel implementation for some reason. Hopefully will be fixed at some
+        # point
+
         options = default_options()
         options.wavelength = self.wavelengths
-        options.orders = 50
+        options.orders = 43
         options.detailed_rcwa = True
         options.parallel = True
         # options.pol = (1/np.sqrt(2), 1j/np.sqrt(2))
@@ -95,14 +95,13 @@ def fitness(self, x):
 
     def get_bounds(self):
         # [lower bounds for all parameters], [upper bounds for all parameters]
-        return [(100, 0, 200), (15000, 1, 3000)]
+        return [(100, 0, 200), (10000, 0.5, 3000)]
 
 
 
 crit_angle = np.arcsin(Air.n(5e-6)/Si.n(5e-6))
 test_wl = np.linspace(4.5,  5.5,10)*1e-6
 
-
 obj_class = optimize_surface(wavelengths=test_wl, target_diffraction_angle=45*np.pi/180,
                                    angle_tolerance=10*np.pi/180)
 prob = pg.problem(obj_class)
@@ -114,7 +113,10 @@ def get_bounds(self):
 # pop = pg.population(prob, pop_size)
 # pop = algo.evolve(pop)
 
-algo = pg.algorithm(pg.de(gen=1))
+algo = pg.algorithm(pg.de(gen=1)) # note: I am running this one generation at a time so I can get the
+# best population at each step and plot how it evolves. If you just care about the final result,
+# you can just set gen=n_generations here and not have the for loop below.
+
 n_params = len(prob.get_bounds()[0])
 
 best_f = np.zeros(n_generations)
@@ -123,101 +125,112 @@ def get_bounds(self):
 
 pop = pg.population(prob, pop_size)
 
-results_table = np.zeros((5, 4))
+for i1 in range(n_generations):
 
-for j1 in range(5):
+    pop = algo.evolve(pop)
+    print(i1, pop.champion_f[0])
+    best_f[i1] = pop.champion_f[0]
+    mean_f[i1] = np.mean(pop.get_f())
+    best_x[i1] = pop.champion_x
 
-    for i1 in range(n_generations):
+print(pop.champion_x, pop.champion_f)
 
-        pop = algo.evolve(pop)
-        print(i1, pop.champion_f[0])
-        best_f[i1] = pop.champion_f[0]
-        mean_f[i1] = np.mean(pop.get_f())
-        best_x[i1] = pop.champion_x
-    results_table[j1] = [pop.champion_x[0], pop.champion_x[1], pop.champion_x[2], pop.champion_f[0]]
 
-print(pop.champion_x, pop.champion_f)
+fig, ax1 = plt.subplots(1, 1, figsize=(5, 4))
+ax1.plot(np.arange(n_generations) + 1, best_f, 'ro-', label="Champion fitness")
+ax1.plot(np.arange(n_generations) + 1, mean_f, 'kx-',  label="Mean fitness")
 
-fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))
-ax1.plot(np.arange(n_generations) + 1, best_f, 'ro-', label="Best fitness")
-ax2.plot(np.arange(n_generations) + 1, mean_f, 'kx-',  label="Mean fitness")
+ax1.axhline(-1, color='k', linestyle='--', label="Target")
 ax1.set_xlabel("Generation")
 ax1.set_ylabel("Fitness")
-ax2.set_xlabel("Generation")
-
-ax1.set_title("Champion fitness")
-ax2.set_title("Mean population fitness")
+ax1.legend()
+plt.tight_layout()
 plt.show()
 
-fig, axes = plt.subplots(1, n_params, figsize=(12, 6))
+fig, axes = plt.subplots(1, n_params, figsize=(10, 4))
 axes = axes.flatten()
 for i1 in range(n_params):
     axes[i1].plot(np.arange(n_generations) + 1, best_x[:, i1], 'ro-', label="Best x")
     axes[i1].set_xlabel("Generation")
     axes[i1].set_ylabel("x")
     axes[i1].set_title("Parameter {}".format(i1))
     axes[i1].set_ylim(prob.get_bounds()[0][i1], prob.get_bounds()[1][i1])
-plt.show()
 
-ax1.set_title("Champion fitness")
-ax2.set_title("Mean population fitness")
+plt.tight_layout()
 plt.show()
 
 
 RAT, options, size = obj_class.calculate(pop.champion_x)
 
+diffracted_directions = get_order_directions(options.wavelength * 1e9, size, 4, Air, Si,
+                                             options.theta_in,
+                                             0,
+                                             np.pi / 3)
+orders_analytical = diffracted_directions['order_index']
+
 orders = np.array(RAT['basis_set'])  # dimensions: (orders, 2)
 
-# avg power per order over wl:
+# avg power per order over wl, determine which orders need to be included on x-axis
 ind_power = np.mean(RAT['T_per_order'], axis=0) > 1e-7
+all_labels = np.array([f'({x[0]}, {x[1]})' for x in RAT['basis_set']])[ind_power]
+
+relevant_orders = np.array(RAT['basis_set'])[ind_power]
+
+orders_analytical_ind = [np.where(
+    np.all((orders_analytical[0].flatten() == x[0], orders_analytical[1].flatten() == x[1]),
+           axis=0))[0][0] for x in relevant_orders]
+
+theta_analytical = diffracted_directions['theta_t'][:, orders_analytical_ind]
 
+# find angle for these orders:
+
+colors = sns.color_palette("viridis", len(options.wavelength))
 # plot power per order:
-all_labels = np.array([f'({x[0]}, {x[1]})' for x in RAT['basis_set']])[ind_power]
-plt.figure()
+
+plt.figure(figsize=(5, 3.5))
 for i1 in range(len(options.wavelength)):
-    plt.scatter(all_labels, RAT['T_per_order'][i1, ind_power])
+    theta_increasing_order = np.argsort(theta_analytical[i1])
+    theta_i1 = theta_analytical[i1][theta_increasing_order]
+    T_i1 = RAT['T_per_order'][i1, ind_power][theta_increasing_order]
+    # group together and add data points where theta is identical:
+    values, indices, counts = np.unique(theta_i1, return_counts=True, return_index=True)
+    subarrays = np.split(T_i1, indices)
+    T_per_angle = np.array([np.sum(x) for x in subarrays])[1:]
+
+    values = np.pi - values
+
+    plt.scatter(180*values/np.pi, T_per_angle, color=colors[i1], facecolor='none',
+                label=f'{options.wavelength[i1]*1e6:.2f} um')
 # plt.scatter([f'({x[0]}, {x[1]})' for x in RAT['basis_set']], RAT['T_per_order'])
 # rotate x axis labels 90 degrees:
+# degree symbol:
+plt.axvline(obj_class.target*180/np.pi, color='k', linestyle='-', alpha=0.5)
+plt.axvline((obj_class.target + obj_class.tolerance)*180/np.pi, color='k', linestyle='--', alpha=0.5)
+plt.axvline((obj_class.target - obj_class.tolerance)*180/np.pi, color='k', linestyle='--', alpha=0.5)
 plt.xticks(rotation=90)
+plt.xlabel("Direction in Si (°)")
+plt.ylabel("Power")
+plt.legend()
+plt.tight_layout()
 plt.show()
 
-diffracted_directions = get_order_directions(options.wavelength * 1e9, size, 4, Air, Si,
-                                             options.theta_in,
-                                             0,
-                                             np.pi / 3)
-orders_analytical = diffracted_directions['order_index']
-
-orders_analytical_ind = [np.where(
-    np.all((orders_analytical[0].flatten() == x[0], orders_analytical[1].flatten() == x[1]),
-           axis=0))[0][0] for x in orders]
 
-theta_analytical = diffracted_directions['theta_t'][:, orders_analytical_ind]
+# orders_analytical_ind = [np.where(
+#     np.all((orders_analytical[0].flatten() == x[0], orders_analytical[1].flatten() == x[1]),
+#            axis=0))[0][0] for x in orders]
+#
+# theta_analytical = diffracted_directions['theta_t'][:, orders_analytical_ind]
 
 mask = np.all((theta_analytical > obj_class.min_angle, theta_analytical < obj_class.max_angle), axis=0)
 
-T_diff_per_wl = [np.sum(RAT['T_per_order'][i1][mask[i1]]) for i1 in range(len(options.wavelength))]
+T_diff_per_wl = [np.sum(RAT['T_per_order'][i1, ind_power][mask[i1]]) for i1 in range(len(options.wavelength))]
 
 plt.figure()
 plt.plot(options.wavelength*1e6, RAT['T'], '--k', label='Total transmission')
-plt.plot(options.wavelength*1e6, T_diff_per_wl, label='Diffracted into target angle')
+plt.plot(options.wavelength*1e6, T_diff_per_wl, '-k', label='Diffracted into target angle')
 plt.xlabel("Wavelength (um)")
+plt.ylabel("Transmission")
+plt.legend()
+plt.tight_layout()
 plt.show()
 
-# what happens if we try inverse: start in diffracted direction and see how much gets coupled out?
-
-x = pop.champion_x
-
-Air = material("Air")()
-Si = material("Si")()
-
-size = ((x[0], 0), (x[0] / 2, np.sin(np.pi / 3) * x[0]))
-
-grating_circles = [{"type": "circle", "mat": Si, "center": (0, 0), "radius": x[1] * x[0]}]
-
-layers = [Layer(material=Air, width=si(x[2], "nm"), geometry=grating_circles)]
-
-S4_setup = rcwa_structure(layers, size=size, options=options, incidence=Si,
-                          transmission=Air)
-options.theta_in = np.pi - 70*np.pi/180
-
-RAT = S4_setup.calculate(options)