Fixed scipy deprecating interp2 function, create animations using CPU multithreading, save as image, encode with ffmpeg, Python 3.12 code convention...

diffractsim/colour_functions.py

@@ -3,7 +3,14 @@
 from scipy.interpolate import CubicSpline
 from .util.backend_functions import backend as bd
 
-illuminant_d65 = np.loadtxt(Path(__file__).parent / "./data/illuminant_d65.txt", usecols=(1))
+illuminant_d65 = np.loadtxt(Path(__file__).parent / "./data/illuminant_d65.txt", usecols = 1)
+high_pressure_sodium = np.loadtxt(Path(__file__).parent / "./data/high_pressure_sodium.txt", usecols = 1)
+incandescent_tugsten = np.loadtxt(Path(__file__).parent / "./data/incandescent_tugsten.txt", usecols = 1)
+compact_fluorescent_lamp = np.loadtxt(Path(__file__).parent / "./data/compact_fluorescent_lamp.txt", usecols = 1)
+mercury_vapor = np.loadtxt(Path(__file__).parent / "./data/mercury_vapor.txt", usecols = 1)
+LED_6770K = np.loadtxt(Path(__file__).parent / "./data/LED_6770K.txt", usecols = 1)
+ceramic_metal_halide = np.loadtxt(Path(__file__).parent / "./data/ceramic_metal_halide.txt", usecols = 1)
+cie_cmf = np.loadtxt(Path(__file__).parent / "./data/cie-cmf.txt", usecols = 1)
 
 """
 MPL 2.0 Clause License 
@@ -12,28 +19,29 @@
 All rights reserved.
 """
 
+
 class ColourSystem:
-    def __init__(self, spectrum_size = 400, spec_divisions = 40, clip_method = 1):
+    def __init__(self, spectrum_size=400, spec_divisions=40, clip_method=1):
         global bd
         from .util.backend_functions import backend as bd
 
         self.spectrum_size = spectrum_size
         # import CIE XYZ standard observer color matching functions
 
-        cmf = np.loadtxt(Path(__file__).parent / "./data/cie-cmf.txt", usecols=(1, 2, 3))
-
-        self.Δλ = (779-380)/spectrum_size
-        self.λ_list = np.linspace(380,779, spectrum_size)
+        cmf = np.loadtxt(Path(__file__).parent / "./data/cie-cmf.txt", usecols = (1, 2, 3))
+
+        self.Δλ = (779 - 380) / spectrum_size
+        self.λ_list = np.linspace(380, 779, spectrum_size)
+
+        if spectrum_size == 400:
 
-        if spectrum_size == 400: 
-
             # CIE XYZ standard observer color matching functions
-            self.cie_x =  cmf.T[0]
-            self.cie_y =  cmf.T[1]
-            self.cie_z =  cmf.T[2]
+            self.cie_x = cmf.T[0]
+            self.cie_y = cmf.T[1]
+            self.cie_z = cmf.T[2]
 
-        else: #by default spectrum has a size of 400. If new size, we interpolate
-            λ_list_old = np.linspace(380,779, 400)
+        else:  # by default spectrum has a size of 400. If new size, we interpolate
+            λ_list_old = np.linspace(380, 779, 400)
             self.cie_x = np.interp(self.λ_list, λ_list_old, cmf.T[0])
             self.cie_y = np.interp(self.λ_list, λ_list_old, cmf.T[1])
             self.cie_z = np.interp(self.λ_list, λ_list_old, cmf.T[2])
@@ -54,8 +62,8 @@ def __init__(self, spectrum_size = 400, spec_divisions = 40, clip_method = 1):
         self.T = bd.vstack(
             [[3.2406, -1.5372, -0.4986], [-0.9689, 1.8758, 0.0415], [0.0557, -0.2040, 1.0570]]
         )
-        
-        #clip methods for negative sRGB values:
+
+        # clip methods for negative sRGB values:
 
         self.CLIP_CLAMP_TO_ZERO = 0
         self.CLIP_ADD_WHITE = 1
@@ -73,7 +81,7 @@ def XYZ_to_sRGB_linear(self, XYZ):
                                                                          [Z1,Z2,Z3...]])
         """
 
-        rgb = bd.tensordot(self.T, XYZ, axes=([1, 0]))
+        rgb = bd.tensordot(self.T, XYZ, axes = ([1, 0]))
 
         if self.clip_method == self.CLIP_CLAMP_TO_ZERO:
             # set negative rgb values to zero
@@ -83,9 +91,9 @@ def XYZ_to_sRGB_linear(self, XYZ):
         if self.clip_method == self.CLIP_ADD_WHITE:
             # add enough white to make all rgb values nonnegative
             # find max negative rgb (or 0.0 if all non-negative), we need that much white
-            rgb_min = bd.amin(rgb, axis=0)
+            rgb_min = bd.amin(rgb, axis = 0)
             # get max positive component
-            rgb_max = bd.amax(rgb, axis=0)
+            rgb_max = bd.amax(rgb, axis = 0)
 
             # get scaling factor to maintain max rgb after adding white
             scaling = bd.where(rgb_max > 0.0, rgb_max / (rgb_max - rgb_min + 0.00001), bd.ones(rgb.shape))
@@ -94,8 +102,7 @@ def XYZ_to_sRGB_linear(self, XYZ):
             rgb = bd.where(rgb_min < 0.0, scaling * (rgb - rgb_min), rgb)
             return rgb
 
-
-    def sRGB_linear_to_sRGB(self,rgb_linear):
+    def sRGB_linear_to_sRGB(self, rgb_linear):
 
         """
         Convert a linear RGB color to a non linear RGB color (gamma correction).
@@ -114,14 +121,13 @@ def sRGB_linear_to_sRGB(self,rgb_linear):
         )
 
         # clip intensity if needed (rgb values > 1.0) by scaling
-        rgb_max = bd.amax(rgb, axis=0) + 0.00001  # avoid division by zero
+        rgb_max = bd.amax(rgb, axis = 0) + 0.00001  # avoid division by zero
         intensity_cutoff = 1.0
         rgb = bd.where(rgb_max > intensity_cutoff, rgb * intensity_cutoff / (rgb_max), rgb)
 
         return rgb
 
-
-    def sRGB_to_sRGB_linear(self,rgb):
+    def sRGB_to_sRGB_linear(self, rgb):
         """
         Convert a RGB color to a linear RGB color.
 
@@ -132,7 +138,6 @@ def sRGB_to_sRGB_linear(self,rgb):
         """
         return bd.where(rgb <= 0.03928, rgb / 12.92, bd.power((rgb + 0.055) / 1.055, 2.4))
 
-
     def XYZ_to_sRGB(self, XYZ):
         """
         Convert a XYZ to an RGB color.
@@ -148,7 +153,6 @@ def XYZ_to_sRGB(self, XYZ):
 
         return rgb
 
-
     def spec_to_XYZ(self, spec):
         """
         Convert a spectrum to an XYZ color.
@@ -158,8 +162,6 @@ def spec_to_XYZ(self, spec):
         Number of samples of each spectral intensity list doesn't matter, but they must be equally spaced.
         """
 
-
-
         if spec.ndim == 1:
 
             X = bd.dot(spec, self.cie_x) * self.Δλ * 0.003975 * 683.002
@@ -168,10 +170,9 @@ def spec_to_XYZ(self, spec):
             return bd.array([X, Y, Z])
 
         else:
-            return bd.tensordot(spec, self.cie_xyz, axes=([1, 1])).T * self.Δλ * 0.003975 * 683.002
-
+            return bd.tensordot(spec, self.cie_xyz, axes = ([1, 1])).T * self.Δλ * 0.003975 * 683.002
 
-    def spec_partition_to_XYZ(self, spec_partition, index = 0):
+    def spec_partition_to_XYZ(self, spec_partition, index=0):
         """
         Convert a spectrum to an XYZ color.
 
@@ -187,9 +188,7 @@ def spec_partition_to_XYZ(self, spec_partition, index = 0):
             return bd.array([X, Y, Z])
 
         else:
-            return bd.tensordot(spec_partition, self.cie_xyz_partitions[index], axes=([1, 1])).T * self.Δλ * 0.003975 * 683.002
-
-
+            return bd.tensordot(spec_partition, self.cie_xyz_partitions[index], axes = ([1, 1])).T * self.Δλ * 0.003975 * 683.002
 
     def spec_to_sRGB(self, spec):
         """
@@ -204,12 +203,10 @@ def spec_to_sRGB(self, spec):
         XYZ = self.spec_to_XYZ(spec)
         return self.XYZ_to_sRGB(XYZ)
 
-
-    def wavelength_to_XYZ(self,wavelength, intensity):
-
+    def wavelength_to_XYZ(self, wavelength, intensity):
 
         if (wavelength > 380) and (wavelength < 780):
-            index = int(wavelength-380)
+            index = int(wavelength - 380)
             X = intensity * self.cie_x[index] * self.Δλ * 0.003975 * 683.002
             Y = intensity * self.cie_y[index] * self.Δλ * 0.003975 * 683.002
             Z = intensity * self.cie_z[index] * self.Δλ * 0.003975 * 683.002
@@ -220,7 +217,6 @@ def wavelength_to_XYZ(self,wavelength, intensity):
 
         return bd.array([X, Y, Z])
 
-
     def wavelength_to_sRGB(self, wavelength, intensity):
 
         XYZ = self.wavelength_to_XYZ(wavelength, intensity)
@@ -229,4 +225,4 @@ def wavelength_to_sRGB(self, wavelength, intensity):
     def wavelength_to_sRGB_linear(self, wavelength, intensity):
 
         XYZ = self.wavelength_to_XYZ(wavelength, intensity)
-        return self.XYZ_to_sRGB_linear(XYZ)
+        return self.XYZ_to_sRGB_linear(XYZ)

diffractsim/monochromatic_simulator.py

@@ -32,9 +32,7 @@ def __init__(self,  wavelength, extent_x, extent_y, Nx, Ny, intensity = 0.1 * W
         intensity: intensity of the field
         """
         global bd
-        global backend_name
         from .util.backend_functions import backend as bd
-        from .util.backend_functions import backend_name
 
         self.extent_x = extent_x
         self.extent_y = extent_y
@@ -109,8 +107,6 @@ def zoom_propagate(self, z, x_interval, y_interval):
         self.E = bluestein_method(self, self.E, z, self.λ, x_interval, y_interval)
 
 
-
-
     def propagate_to_image_plane(self, pupil, M, zi, z0, scale_factor = 1):
         """
         Assuming an optical system with linear response and assuming the system is only diffraction-limited by
@@ -237,21 +233,23 @@ def compute_colors_at(self, z):
 
     def interpolate(self, Nx, Ny):
         """Interpolate the field to the new shape (Nx,Ny)"""
-        from scipy.interpolate import RectBivariateSpline
+        from scipy.interpolate import interp2d
 
 
-        if backend_name == 'cupy':
+        if bd != np:
             self.E = self.E.get()
 
-        fun_real = RectBivariateSpline(
+        fun_real = interp2d(
                     self.dx*(np.arange(self.Nx)-self.Nx//2),
                     self.dy*(np.arange(self.Ny)-self.Ny//2),
-                    np.real(self.E))
+                    np.real(self.E),
+                    kind="cubic",)
 
-        fun_imag = RectBivariateSpline(
+        fun_imag = interp2d(
                     self.dx*(np.arange(self.Nx)-self.Nx//2),
                     self.dy*(np.arange(self.Ny)-self.Ny//2),
-                    np.imag(self.E))
+                    np.imag(self.E),
+                    kind="cubic",)
 
 
         self.Nx = Nx
@@ -306,12 +304,8 @@ def get_longitudinal_profile(self, start_distance, end_distance, steps, scale_fa
                 self.yy/=scale_factor
 
             rgb = self.get_colors()
-            if backend_name == 'jax':
-                longitudinal_profile_rgb = longitudinal_profile_rgb.at[i,:,:].set( rgb[self.Ny//2,:,:])
-                longitudinal_profile_E = longitudinal_profile_E.at[i,:].set(self.E[self.Ny//2,:])
-            else:
-                longitudinal_profile_rgb[i,:,:]  = rgb[self.Ny//2,:,:]
-                longitudinal_profile_E[i,:] = self.E[self.Ny//2,:]
+            longitudinal_profile_rgb[i,:,:]  = rgb[self.Ny//2,:,:]
+            longitudinal_profile_E[i,:] = self.E[self.Ny//2,:]
             self.E = np.copy(self.E0)
 
 
@@ -340,4 +334,4 @@ def __add__(self, Field):
             "The wavelength, dimensions and sampling of the interfering fields must be identical")
 
 
-    from .visualization import plot_colors, plot_phase, plot_intensity, plot_longitudinal_profile_colors, plot_longitudinal_profile_intensity
+    from .visualization import save_plot, plot_colors, plot_phase, plot_intensity, plot_longitudinal_profile_colors, plot_longitudinal_profile_intensity

diffractsim/polychromatic_simulator.py

@@ -23,9 +23,7 @@
 class PolychromaticField:
     def __init__(self, spectrum, extent_x, extent_y, Nx, Ny, spectrum_size = 180, spectrum_divisions = 30):
         global bd
-        global backend_name
         from .util.backend_functions import backend as bd
-        from .util.backend_functions import backend_name
 
         self.extent_x = extent_x
         self.extent_y = extent_y
@@ -97,7 +95,7 @@ def get_colors(self):
 
         bar = progressbar.ProgressBar()
 
-        # We compute the pattern of each wavelength separately, and associate it to small spectrum interval dλ = (780- 380)/spectrum_divisions . We approximately the final colour
+        # We compute the pattern of each wavelength separately, and associate it to small spectrum interval dλ = (780- 380)/spectrum_divisions . We approximate the final colour
         # by summing the contribution of each small spectrum interval converting its intensity distribution to a RGB space.
 
 
@@ -126,15 +124,20 @@ def get_colors(self):
                         sRGB_linear += self.cs.XYZ_to_sRGB_linear(XYZ)
 
 
-
-        if backend_name == 'cupy':
+        if bd != np:
             bd.cuda.Stream.null.synchronize()
         rgb = self.cs.sRGB_linear_to_sRGB(sRGB_linear)
         rgb = (rgb.T).reshape((self.Ny, self.Nx, 3))
         print ("Computation Took", time.time() - t0)
         return rgb
 
 
+    def get_field(self):
+        """get field of the cross-section profile at the current distance"""
+
+        return self.E
+
+
     def scale_propagate(self, z, scale_factor):
         """
         #raise NotImplementedError(self.__class__.__name__ + '.scale_propagate')
@@ -176,7 +179,6 @@ def get_colors_at_image_plane(self, pupil, M, zi, z0, scale_factor = 1):
         """
 
 
-
         for j in range(len(self.optical_elements)):
             self.E = self.E * self.optical_elements[j].get_transmittance(self.xx, self.yy, 0)
 
@@ -215,7 +217,7 @@ def get_colors_at_image_plane(self, pupil, M, zi, z0, scale_factor = 1):
             XYZ = self.cs.spec_partition_to_XYZ(bd.outer(Iλ, self.spec_partitions[i]),i)
             sRGB_linear += self.cs.XYZ_to_sRGB_linear(XYZ)
 
-        if backend_name == 'cupy':
+        if bd != np:
             bd.cuda.Stream.null.synchronize()
 
         self.xx = M_abs * self.xx
@@ -233,4 +235,4 @@ def get_colors_at_image_plane(self, pupil, M, zi, z0, scale_factor = 1):
 
 
 
-    from .visualization import plot_colors
+    from .visualization import save_plot, plot_colors

diffractsim/visualization/__init__.py

@@ -1,4 +1,4 @@
-from .plot_colors import plot_colors
+from .plot_colors import save_plot, plot_colors
 from .plot_intensity import plot_intensity
 from .plot_phase import plot_phase
 from .plot_longitudinal_profile import plot_longitudinal_profile_colors, plot_longitudinal_profile_intensity