Fix material matching plot (try plotting ground truth and see if right material comes out)

nidn/__init__.py

@@ -7,6 +7,7 @@
 from .materials.material_collection import MaterialCollection
 from .plots.plot_epsilon_grid import plot_epsilon_grid
 from .plots.plot_eps_per_point import plot_eps_per_point
+from .plots.plot_material_grid import plot_material_grid
 from .plots.plot_model_grid import plot_model_grid
 from .plots.plot_model_grid_per_freq import plot_model_grid_per_freq
 from .plots.plot_spectra import plot_spectra
@@ -45,6 +46,7 @@
     "phys_freq_to_phys_wl",
     "plot_epsilon_grid",
     "plot_eps_per_point",
+    "plot_material_grid",
     "plot_model_grid",
     "plot_model_grid_per_freq",
     "plot_spectra",

nidn/materials/find_closest_material.py

@@ -0,0 +1,36 @@
+from .material_collection import MaterialCollection
+
+import torch
+
+
+def _find_closest_material(eps, run_cfg):
+    """Finds the closest matching material and distance to that (in epsilon)
+    from a given epsilon grid.
+
+    Args:
+        eps (torch.tensor): A tensor of epsilon values (Nx,Ny,N_layers,N_freq).
+        run_cfg (DotMap): Run configuration.
+
+    Returns:
+        tuple: The closest materials and distances to that.
+    """
+    Nx, Ny, N_layers, N_freq = run_cfg.Nx, run_cfg.Ny, run_cfg.N_layers, run_cfg.N_freq
+
+    # Initiate material collection
+    material_collection = MaterialCollection(run_cfg.target_frequencies)
+
+    comparisons = torch.zeros(
+        [Nx, Ny, N_layers, N_freq, material_collection.N_materials,]
+    )
+
+    # Compute differences for all materials
+    for idx, material_name in enumerate(material_collection.material_names):
+        material_eps = material_collection[material_name]
+
+        comparisons[..., idx] = torch.abs(eps - material_eps)
+
+    # Find minimal entries
+    comparisons = comparisons.mean(dim=3)  # Average over frequencies
+    minimal_comparisons, indices = torch.min(comparisons, dim=-1)
+
+    return minimal_comparisons, indices

nidn/plots/plot_eps_per_point.py

@@ -1,12 +1,12 @@
-import os
-import inspect
-
+import torch
 from matplotlib import pyplot as plt
 from matplotlib.ticker import FormatStrFormatter
+
 from ..utils.convert_units import freq_to_wl
 from ..training.model.model_to_eps_grid import model_to_eps_grid
 from ..trcwa.load_material_data import _load_material_data
 from ..materials.material_collection import MaterialCollection
+from ..materials.find_closest_material import _find_closest_material
 
 
 def plot_eps_per_point(model, run_cfg, compare_to_material=None):
@@ -19,20 +19,16 @@ def plot_eps_per_point(model, run_cfg, compare_to_material=None):
     """
     # Create epsilon grid from the model
     eps, _ = model_to_eps_grid(model, run_cfg)
-    eps = eps.detach().cpu().numpy()
+    eps_np = eps.detach().cpu().numpy()
+
+    material_collection = MaterialCollection(run_cfg.target_frequencies)
 
     # Load material data for comparison
     if compare_to_material is not None:
-        material_data = _load_material_data(
-            os.path.dirname(inspect.getfile(MaterialCollection))
-            + "/data/"
-            + compare_to_material
-            + ".csv",
-            run_cfg.target_frequencies,
-        )
+        material_data = material_collection[compare_to_material]
 
     # Create figure
-    fig = plt.figure(figsize=(10, 5), dpi=150)
+    fig = plt.figure(figsize=(10, 6), dpi=150)
     fig.patch.set_facecolor("white")
 
     # Plot epsilon
@@ -41,8 +37,8 @@ def plot_eps_per_point(model, run_cfg, compare_to_material=None):
     wl = freq_to_wl(run_cfg.target_frequencies)
 
     # Add some horizontal space
-    ax.set_xlim(wl[0], 2.5 * wl[-1])
-    ax2.set_xlim(wl[0], 2.5 * wl[-1])
+    ax.set_xlim(wl.min(), 2.5 * wl.max())
+    ax2.set_xlim(wl.min(), 2.5 * wl.max())
 
     ax.set_xlabel("Wavelength [µm]")
     ax.set_ylabel("Epsilon real part")
@@ -59,15 +55,41 @@ def plot_eps_per_point(model, run_cfg, compare_to_material=None):
     for x in range(eps.shape[0]):
         for y in range(eps.shape[1]):
             for N_layer in range(eps.shape[2]):
-                eps_point_real = eps[x, y, N_layer].real
-                eps_point_imag = eps[x, y, N_layer].imag
+                eps_point_real = eps_np[x, y, N_layer].real
+                eps_point_imag = eps_np[x, y, N_layer].imag
                 ax.plot(wl, eps_point_real, linewidth=1)
                 ax2.plot(wl, eps_point_imag, linewidth=1)
 
-                ax.text(wl[-1], eps_point_real[-1], f" {N_layer},{x},{y}", va="center")
-                ax2.text(wl[-1], eps_point_imag[-1], f" {N_layer},{x},{y}", va="center")
+                ax.text(
+                    wl.max(),
+                    eps_point_real[0],
+                    f" {N_layer},{x},{y}",
+                    va="center",
+                    fontsize=7,
+                )
+                ax2.text(
+                    wl.max(),
+                    eps_point_imag[0],
+                    f" {N_layer},{x},{y}",
+                    va="center",
+                    fontsize=7,
+                )
 
     # Plot material data
     if compare_to_material is not None:
         ax.plot(wl, material_data.real, "--", color="black", linewidth=1.5)
         ax2.plot(wl, material_data.imag, "--", color="black", linewidth=1.5)
+    else:
+        _, indices = _find_closest_material(eps, run_cfg)
+        unique_indices = torch.unique(indices)
+        names = [material_collection.material_names[i] for i in unique_indices]
+        for name in names:
+            material_data = material_collection[name]
+            ax.plot(wl, material_data.real, "--", color="black", linewidth=1.5)
+            ax2.plot(wl, material_data.imag, "--", color="black", linewidth=1.5)
+            ax.text(
+                wl.max(), material_data.real[0], " " + name, va="center", fontsize=7
+            )
+            ax2.text(
+                wl.max(), material_data.imag[0], " " + name, va="center", fontsize=7
+            )

nidn/plots/plot_material_grid.py

@@ -0,0 +1,98 @@
+import torch
+import numpy as np
+from matplotlib import pyplot as plt
+
+from ..materials.material_collection import MaterialCollection
+from ..materials.find_closest_material import _find_closest_material
+from ..training.model.model_to_eps_grid import model_to_eps_grid
+
+
+def plot_material_grid(model, run_cfg):
+    """Plots the materials closest to the used ones for each grid point. 
+
+    Args:
+        model (torch.model): The model to be plotted.
+        run_cfg (dict): The run configuration.
+    """
+    Nx, Ny, N_layers = run_cfg.Nx, run_cfg.Ny, run_cfg.N_layers
+    # Create epsilon grid from the model
+    eps, _ = model_to_eps_grid(model, run_cfg)
+
+    # Setup grid
+    x = torch.linspace(-1, 1, Nx)
+    y = torch.linspace(-1, 1, Ny)
+    z = torch.linspace(-1, 1, N_layers)
+    X, Y, Z = torch.meshgrid((x, y, z))
+
+    # Load material data
+    material_collection = MaterialCollection(run_cfg.target_frequencies)
+
+    # Get closest materials
+    errors, material_id = _find_closest_material(eps, run_cfg)
+
+    cmap = plt.get_cmap("rainbow", material_collection.N_materials)
+
+    # Here we plot it
+    fig = plt.figure(figsize=(10, 5), dpi=150)
+    fig.patch.set_facecolor("white")
+    ax = fig.add_subplot(121, projection="3d")
+    ax.view_init(elev=25, azim=100)
+    p = ax.scatter(
+        X.reshape(-1, 1),
+        Y.reshape(-1, 1),
+        Z.reshape(-1, 1),
+        marker="s",
+        s=120,
+        linewidths=0,
+        cmap=cmap,
+        vmin=1 - 0.5,  # This is where we get the discrete colormap
+        vmax=material_collection.N_materials
+        + 0.5,  # This is where we get the discrete colormap
+        alpha=1.0,
+        c=material_id.detach().cpu().numpy() + 1,
+    )
+    cbar = fig.colorbar(
+        p, ticks=np.arange(1, material_collection.N_materials + 1)
+    )  # This is where we get the discrete colormap
+    cbar.set_ticklabels(material_collection.material_names)
+    cbar.ax.tick_params(labelsize=7)
+    # cbar.set_label("Materials", labelpad=-1)
+
+    ax.grid(False)  # Hide grid lines
+    ax.set_xlabel("$N_x =$" + str(Nx))
+    ax.set_ylabel("$N_y =$" + str(Ny))
+    # ax.set_zlabel("# of layers", rotation=60)  # TODO Fix rotation
+    ax.set_xticks([])
+    ax.set_yticks([])
+    ax.set_zticks(np.linspace(-1, 1, N_layers))
+
+    z_labels = [""] * N_layers
+    for idx in range(N_layers):
+        z_labels[idx] = "L" + str(idx + 1)
+    ax.set_zticklabels(z_labels)  # Where L1 is (seemingly) the bottom one
+
+    ax = fig.add_subplot(122, projection="3d")
+    ax.view_init(elev=25, azim=100)
+    p = ax.scatter(
+        X.reshape(-1, 1),
+        Y.reshape(-1, 1),
+        Z.reshape(-1, 1),
+        marker="s",
+        s=120,
+        linewidths=0,
+        c=errors.detach().cpu().numpy(),
+    )
+    cbar = plt.colorbar(p, ax=ax)
+    cbar.ax.tick_params(labelsize=7)
+    ax.grid(False)  # Hide grid lines
+    ax.set_xlabel("$N_x =$" + str(Nx))
+    ax.set_ylabel("$N_y =$" + str(Ny))
+    # ax.set_zlabel("# of layers", rotation=60)  # TODO Fix rotation
+    ax.set_xticks([])
+    ax.set_yticks([])
+    ax.set_zticks(np.linspace(-1, 1, N_layers))
+
+    z_labels = [""] * N_layers
+    for idx in range(N_layers):
+        z_labels[idx] = "L" + str(idx + 1)
+    ax.set_zticklabels(z_labels)  # Where L1 is (seemingly) the bottom one

nidn/plots/plot_model_grid.py

@@ -21,7 +21,7 @@ def plot_model_grid(model, run_cfg):
     # Here we calculate the absolute value of the permittivity over all frequencies for each grid point
     eps = torch.norm(eps, dim=3)
 
-    material_id = eps.detach().cpu().numpy()
+    abs_values = eps.detach().cpu().numpy()
 
     X = X.cpu().numpy()
     Y = Y.cpu().numpy()
@@ -41,7 +41,7 @@ def plot_model_grid(model, run_cfg):
         s=120,
         linewidths=0,
         alpha=1.0,
-        c=material_id,
+        c=abs_values,
     )
     cbar = plt.colorbar(p, ax=ax)
     cbar.ax.tick_params(labelsize=7)

notebooks/Training.ipynb

@@ -37,9 +37,11 @@
    "source": [
     "cfg.Nx = 9\n",
     "cfg.Ny = 9\n",
-    "cfg.N_layers = 3\n",
+    "cfg.N_layers = 2\n",
+    "cfg.L = 0.5\n",
     "cfg.eps_oversampling = 3\n",
     "cfg.iterations = 100\n",
+    "cfg.type = \"classification\"\n",
     "nidn.print_cfg(cfg)\n",
     "physical_freqs, normalized_freqs = nidn.get_frequency_points(cfg)\n",
     "print(\"Physical frequencies are:\")\n",
@@ -107,8 +109,26 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "nidn.plot_eps_per_point(model,cfg,\"aluminium_nitride\")"
+    "nidn.plot_eps_per_point(model,cfg)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d29d4b07",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "nidn.plot_material_grid(model,cfg)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "222cc2c5",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {