update functions

CITATION.cff

@@ -0,0 +1,15 @@
+cff-version: 1.2.0
+title: Tmmax
+message: >-
+  If you use this software, please cite it using the
+  metadata from this file.
+type: software
+authors:
+  - given-names: Bahrem Serhat
+    family-names: Danis
+    orcid: 'https://orcid.org/0009-0002-9880-0446'
+repository-code: 'https://github.com/bahremsd/tmmax'
+url: 'https://tmmax.readthedocs.io/'
+license: MIT
+version: 1.0.0
+date-released: '2025-02-01'

benchmarks/layer_size_benchmark.py

@@ -0,0 +1,102 @@
+import numpy as np
+import jax
+import jax.numpy as jnp
+import timeit
+import os
+os.environ["JAX_PLATFORMS"] = "cpu"  # Force JAX to use CPU
+
+from utils import generate_material_distribution_indices
+from utils import generate_material_list_with_air
+
+from utils import generate_tmm_args, tmm_coh_tmm_wrapper
+from utils import generate_vtmm_args
+from utils import generate_tmm_fast_args
+from utils import generate_tmmax_args
+
+from tmmax.tmm import vectorized_coh_tmm
+from vtmm import tmm_rt
+from tmm_fast import coh_tmm
+
+
+# Number of layers to test from 2 to 400 (exclusive), using integer values
+number_of_layers = np.arange(2, 400, 20, dtype=int)
+
+# Number of repetitions for the timeit function to get an average execution time
+timeit_repetition = 50
+
+# List of materials used in the simulations
+material_set = ["SiO2", "TiO2", "MgF2", "MgO", "SiO", "Al2O3", "CdS"]
+
+# Polarization type, 's' polarization in this case
+polarization = 's'
+
+# Array of angles of incidence, linearly spaced from 0 to pi/2, with 20 points
+angle_of_incidences = np.linspace(0, np.pi/2, 20)
+angle_of_incidences_tmmax = jnp.array(angle_of_incidences)
+
+# Wavelength array from 500 nm to 1000 nm, linearly spaced with 20 points
+wavelength_arr = np.linspace(500e-9, 1000e-9, 20)
+wavelength_arr_tmmax = jnp.array(wavelength_arr)
+
+# Lists to store execution times for each method
+time_tmm = []
+time_vtmm = []
+time_tmm_fast = []
+time_tmmax = []
+
+# Loop through different numbers of layers
+for N in number_of_layers:
+
+    # Generate random material distribution indices and the corresponding material list with air layers
+    indices = generate_material_distribution_indices(N, low=0, high=len(material_set))
+    material_list = generate_material_list_with_air(indices, material_set) # np.load(f"material_distribution_with_layer_num_{N}.npy", allow_pickle=True)
+
+    # Save the material distribution list as a .npy file for reference
+    np.save(f"material_distribution_with_layer_num_{N}.npy", material_list)
+
+    # Randomly generate thicknesses for each layer between 100 nm and 500 nm
+    thickness_list = np.random.uniform(100, 500, N) * 1e-9 # np.load(f"thickness_list_with_layer_num_{N}.npy")
+    np.save(f"thickness_list_with_layer_num_{N}.npy", thickness_list)  # Save the thickness list
+
+    thickness_list_tmm, nk_list_tmm = generate_tmm_args(material_list = material_list, 
+                                                    thickness_list = thickness_list,
+                                                    wavelength_arr = wavelength_arr)
+
+    omega_vtmm, kx_vtmm, nk_list_vtmm, thickness_list_vtmm = generate_vtmm_args(wavelength_arr = wavelength_arr, 
+                                                                                angle_of_incidences = angle_of_incidences,
+                                                                                thickness_list = thickness_list,
+                                                                                material_list = material_list)
+
+    t_tmm = timeit.timeit( lambda: tmm_coh_tmm_wrapper(polarization, nk_list_tmm, thickness_list_tmm, angle_of_incidences, wavelength_arr),  number=timeit_repetition )
+    t_vtmm = timeit.timeit( lambda: tmm_rt(polarization, omega_vtmm, kx_vtmm, nk_list_vtmm, thickness_list_vtmm),  number = timeit_repetition )
+
+    M_tmm_fast, T_tmm_fast = generate_tmm_fast_args(material_list = material_list, 
+                                                    thickness_list = thickness_list,
+                                                    wavelength_arr = wavelength_arr)
+
+    t_tmm_fast = timeit.timeit( lambda: coh_tmm(polarization, M_tmm_fast, T_tmm_fast, angle_of_incidences, wavelength_arr, device='cpu'),  number = timeit_repetition )
+
+
+    thickness_list_tmmax = jnp.array(thickness_list)
+
+    data_tmmax, material_distribution_tmmax, polarization_tmmax = generate_tmmax_args(material_list = material_list, polarization = polarization)
+
+    t_tmmax = timeit.timeit( lambda: jax.block_until_ready(vectorized_coh_tmm(data_tmmax, material_distribution_tmmax, thickness_list_tmmax, wavelength_arr_tmmax, angle_of_incidences_tmmax, polarization_tmmax)),  number = timeit_repetition)
+
+    time_tmm.append(t_tmm)
+    time_vtmm.append(t_vtmm)
+    time_tmm_fast.append(t_tmm_fast)
+    time_tmmax.append(t_tmmax)
+
+
+
+    print(f"{N} tmm took ", t_tmm)
+    print(f"{N} vtmm took ", t_vtmm)
+    print(f"{N} tmm-fast took ", t_tmm_fast)
+    print(f"{N} tmmax took ", t_tmmax)
+
+# Save the time measurements for each method into .npy files
+np.save("time_of_tmm.npy", time_tmm)
+np.save("time_of_vtmmbur.npy", time_vtmm)
+np.save("time_of_tmm_fastbur.npy", time_tmm_fast)
+np.save("time_of_tmmaxbur.npy", time_tmmax)