Repository: esa/NIDN (language: python)

Issue:
Initial structure

Find a way to store the best physical (i.e. unnormalised) model
Merge #1 first.

- Only saving model if not renormalized
- Added some logging


Validated it in debug

![image](https://user-images.githubusercontent.com/2656685/128314020-21ce73d5-e19e-426d-996b-6f0092c39b63.png)



Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index b6e4761..02affe1 100644
--- a/.gitignore
+++ b/.gitignore
@@ -127,3 +127,7 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+my_notebooks
+.vscode
+results
diff --git a/.readthedocs.yml b/.readthedocs.yml
new file mode 100644
index 0000000..5069375
--- /dev/null
+++ b/.readthedocs.yml
@@ -0,0 +1,23 @@
+# .readthedocs.yml
+# Read the Docs configuration file
+# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
+
+# Required
+version: 2
+
+# Build documentation in the docs/ directory with Sphinx
+sphinx:
+  configuration: docs/source/conf.py
+
+# Optionally build your docs in additional formats such as PDF
+formats: all
+
+# Optionally set the version of Python and requirements required to build your docs
+python:
+  version: 3.8
+  install:
+    - method: setuptools
+      path: .
+
+conda:
+  environment: environment.yml
\ No newline at end of file
diff --git a/environment.yml b/environment.yml
new file mode 100644
index 0000000..f7954be
--- /dev/null
+++ b/environment.yml
@@ -0,0 +1,17 @@
+name: nidn
+channels:
+- conda-forge
+- pytorch
+dependencies:
+- cudatoolkit>=11.0.221
+- dotmap>=1.3.24
+- loguru>=0.5.3
+- matplotlib>=3.3.3
+- pandas>=1.3.1
+- pytest>=6.2.1
+- python>=3.8
+- pytorch>=1.9
+- scipy>=1.6.0
+- sphinx>=3.4.3
+- sphinx_rtd_theme>=0.5.1
+- tqdm>=4.56.0
\ No newline at end of file
diff --git a/nidn/__init__.py b/nidn/__init__.py
new file mode 100644
index 0000000..1befd89
--- /dev/null
+++ b/nidn/__init__.py
@@ -0,0 +1,32 @@
+import os
+
+# Set main device by default to cpu if no other choice was made before
+if "TORCH_DEVICE" not in os.environ:
+    os.environ["TORCH_DEVICE"] = "cpu"
+
+# Add exposed features here
+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
+from .training.run_training import run_training
+from .trcwa.compute_target_frequencies import compute_target_frequencies
+from .trcwa.get_frequency_points import get_frequency_points
+from .utils.convert_units import freq_to_wl, wl_to_phys_wl, phys_freq_to_phys_wl
+from .utils.fix_random_seeds import fix_random_seeds
+from .utils.load_default_cfg import load_default_cfg
+from .utils.print_cfg import print_cfg
+
+__all__ = [
+    "compute_target_frequencies",
+    "get_frequency_points",
+    "fix_random_seeds",
+    "freq_to_wl",
+    "load_default_cfg",
+    "run_training",
+    "phys_freq_to_phys_wl",
+    "plot_model_grid",
+    "plot_model_grid_per_freq",
+    "plot_spectra",
+    "print_cfg",
+    "wl_to_phys_wl",
+]
diff --git a/nidn/plots/__init__.py b/nidn/plots/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/nidn/plots/plot_model_grid.py b/nidn/plots/plot_model_grid.py
new file mode 100644
index 0000000..5e97b5e
--- /dev/null
+++ b/nidn/plots/plot_model_grid.py
@@ -0,0 +1,61 @@
+import torch
+import numpy as np
+from matplotlib import pyplot as plt
+
+from ..training.model.model_to_eps_grid import model_to_eps_grid
+
+
+def plot_model_grid(model, run_cfg):
+    """Plots the absolute value of the epsilon over all frequencies for each 3-D 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
+    eps = model_to_eps_grid(model, run_cfg)
+    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))
+
+    # 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()
+
+    X = X.cpu().numpy()
+    Y = Y.cpu().numpy()
+    Z = Z.cpu().numpy()
+
+    # Here we plot it
+    fig = plt.figure(figsize=(6, 6), dpi=150)
+    fig.patch.set_facecolor("white")
+
+    ax = fig.add_subplot(111, 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,
+        alpha=1.0,
+        c=material_id,
+    )
+    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 ( #TODO confirm this)
+    cbar.set_label("Absolute value of the complex permittivity", labelpad=10)
diff --git a/nidn/plots/plot_model_grid_per_freq.py b/nidn/plots/plot_model_grid_per_freq.py
new file mode 100644
index 0000000..35954b8
--- /dev/null
+++ b/nidn/plots/plot_model_grid_per_freq.py
@@ -0,0 +1,85 @@
+import torch
+import numpy as np
+from matplotlib import pyplot as plt
+from ..utils.convert_units import freq_to_wl
+from ..training.model.model_to_eps_grid import model_to_eps_grid
+
+
+def plot_model_grid_per_freq(model, run_cfg, freq_idx=[0, 1, 2, 3]):
+    """Plots the real and imaginary part of the permittivity in two separate plots for frequencies in target_frequencies defined by freq_idx.
+
+    Args:
+        model (torch.model): The model to be plotted.
+        run_cfg (dict): The run configuration.
+        freq_idx (list of int): Which of the frequency indices in target_frequencies we want to plot. Defaults to [0, 1, 2, 3].
+    """
+    Nx, Ny, N_layers = run_cfg.Nx, run_cfg.Ny, run_cfg.N_layers
+    eps = model_to_eps_grid(model, run_cfg)
+    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))
+
+    X = X.cpu().numpy()
+    Y = Y.cpu().numpy()
+    Z = Z.cpu().numpy()
+
+    # Here we plot it
+    fig = plt.figure(figsize=(8, 12), dpi=150)
+    fig.patch.set_facecolor("white")
+
+    min_eps_real, max_eps_real = eps.real.min(), eps.real.max()
+    min_eps_imag, max_eps_imag = eps.imag.min(), eps.imag.max()
+
+    for idx in range(len(freq_idx)):
+        eps_real = eps[:, :, :, freq_idx[idx]].real
+        eps_imag = eps[:, :, :, freq_idx[idx]].imag
+        material_id_real = eps_real.detach().cpu().numpy()
+        material_id_imag = eps_imag.detach().cpu().numpy()
+
+        re_or_im_idx = 0
+        for real_or_imag in [material_id_real, material_id_imag]:
+            ax = fig.add_subplot(421 + re_or_im_idx + idx * 2, projection="3d")
+            ax.view_init(elev=25, azim=100)
+            vmin, vmax = min_eps_real, max_eps_real
+            if re_or_im_idx > 0:
+                vmin, vmax = min_eps_imag, max_eps_imag
+            p = ax.scatter(
+                X.reshape(-1, 1),
+                Y.reshape(-1, 1),
+                Z.reshape(-1, 1),
+                marker="s",
+                linewidths=0,
+                s=140,
+                alpha=1.0,
+                c=real_or_imag,
+                vmin=vmin,
+                vmax=vmax,
+            )
+            ax.set_title(
+                "WL="
+                + str(round(freq_to_wl(run_cfg.target_frequencies[freq_idx[idx]]), 2))
+                + "um",
+                fontsize=8,
+            )
+            # 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.grid(False)  # Hide grid lines
+            ax.set_xticks([])
+            ax.set_yticks([])
+            ax.set_zticks(np.linspace(-1, 1, N_layers))
+            z_labels = [""] * N_layers
+            for layer in range(N_layers):
+                z_labels[layer] = "L" + str(layer + 1)
+            ax.set_zticklabels(
+                z_labels
+            )  # Where L1 is (seemingly) the bottom one ( #TODO confirm this)
+            cb = plt.colorbar(p, ax=ax)
+            cb.set_label("Real permittivity, $\epsilon'$", rotation=270, labelpad=15)
+            if re_or_im_idx == 1:
+                cb.set_label(
+                    "Imaginary permittivity, $\epsilon''$", rotation=270, labelpad=15
+                )
+            cb.ax.tick_params(labelsize=7)
+            re_or_im_idx += 1
diff --git a/nidn/plots/plot_spectra.py b/nidn/plots/plot_spectra.py
new file mode 100644
index 0000000..f773859
--- /dev/null
+++ b/nidn/plots/plot_spectra.py
@@ -0,0 +1,100 @@
+import torch
+import numpy as np
+from matplotlib import pyplot as plt
+from matplotlib.ticker import FormatStrFormatter
+from ..utils.convert_units import freq_to_wl
+from ..trcwa.compute_spectrum import compute_spectrum
+from ..training.model.model_to_eps_grid import model_to_eps_grid
+
+
+def plot_spectra(model, run_cfg,target_R_spectrum,target_T_spectrum):
+    """Plots the produced RTA spectra together with the target spectra.
+
+    Args:
+        model (torch.model): The model to be plotted.
+        run_cfg (dict): The run configuration.
+        target_R_spectrum (torch.tensor): The target reflection spectrum.
+        target_T_spectrum (torch.tensor): The target transmission spectrum.
+    """
+    
+    eps = model_to_eps_grid(model, run_cfg)
+    prod_R_spectrum,prod_T_spectrum = compute_spectrum(eps,run_cfg)
+    target_frequencies = run_cfg.target_frequencies
+
+    prod_R_spectrum = torch.tensor(prod_R_spectrum).detach().cpu().numpy()
+    prod_T_spectrum = torch.tensor(prod_T_spectrum).detach().cpu().numpy()
+    target_R_spectrum = torch.tensor(target_R_spectrum).detach().cpu().numpy()
+    target_T_spectrum = torch.tensor(target_T_spectrum).detach().cpu().numpy()
+
+    prod_A_spectrum = np.ones_like(np.asarray(prod_R_spectrum)) - (
+        np.asarray(prod_T_spectrum) + np.asarray(prod_R_spectrum)
+    )
+    target_A_spectrum = np.ones_like(np.asarray(target_R_spectrum)) - (
+        np.asarray(target_T_spectrum) + np.asarray(target_R_spectrum)
+    )
+
+    # To align all plots
+    ylimits = [-0.225, 1]
+    if (
+        (max(prod_A_spectrum) > 1 or min(prod_A_spectrum) < 0)
+        or (max(prod_T_spectrum) > 1 or min(prod_T_spectrum) < 0)
+        or (max(prod_R_spectrum) > 1 or min(prod_R_spectrum) < 0)
+    ):
+        ylimits = [
+            min(min(prod_A_spectrum), min(prod_T_spectrum), min(prod_R_spectrum))
+            - 0.225,
+            max(max(prod_A_spectrum), max(prod_T_spectrum), max(prod_R_spectrum)) + 0.1,
+        ]
+
+    fig = plt.figure(figsize=(15, 5), dpi=150)
+    fig.patch.set_facecolor("white")
+
+    ax = fig.add_subplot(131)
+    ax.plot(
+        freq_to_wl(target_frequencies), prod_R_spectrum, marker=6, c="cornflowerblue"
+    )
+    ax.plot(freq_to_wl(target_frequencies), target_R_spectrum, marker=7, c="limegreen")
+    ax.legend(["Produced Reflectance", "Target Reflectance"], loc="lower center")
+    ax.set_xlabel("Wavelength [µm]")
+    ax.set_ylabel("Reflectance")
+    ax.set_xscale("log")
+    ax.xaxis.set_major_formatter(FormatStrFormatter("%.1f"))
+    # ax.xaxis.set_minor_formatter(FormatStrFormatter("%.1f"))
+    ax.axhspan(-6, 0, facecolor="gray", alpha=0.3)
+    ax.axhspan(1, 5, facecolor="gray", alpha=0.3)
+    ax.set_ylim(ylimits)
+
+    ax = fig.add_subplot(132)
+    ax.plot(
+        freq_to_wl(target_frequencies), prod_T_spectrum, marker=6, c="cornflowerblue"
+    )
+    ax.plot(freq_to_wl(target_frequencies), target_T_spectrum, marker=7, c="limegreen")
+    ax.legend(["Produced Transmittance", "Target Transmittance"], loc="lower center")
+    ax.set_xlabel("Wavelength [µm]")
+    ax.set_ylabel("Transmittance")
+    ax.set_xscale("log")
+    ax.xaxis.set_major_formatter(FormatStrFormatter("%.1f"))
+
+    # ax.xaxis.set_minor_formatter(FormatStrFormatter("%.1f"))
+    ax.axhspan(-6, 0, facecolor="gray", alpha=0.3)
+    ax.axhspan(1, 5, facecolor="gray", alpha=0.3)
+    ax.set_ylim(ylimits)
+
+    ax = fig.add_subplot(133)
+    ax.plot(
+        freq_to_wl(target_frequencies), prod_A_spectrum, marker=6, c="cornflowerblue"
+    )
+    ax.plot(freq_to_wl(target_frequencies), target_A_spectrum, marker=7, c="limegreen")
+
+    ax.legend(["Produced Absorptance", "Target Absorptance"], loc="lower center")
+    ax.set_xlabel("Wavelength [µm]")
+    ax.set_ylabel("Absorptance")
+    ax.set_xscale("log")
+    ax.xaxis.set_major_formatter(FormatStrFormatter("%.1f"))
+
+    # ax.xaxis.set_minor_formatter(FormatStrFormatter("%.1f"))
+    ax.axhspan(-6, 0, facecolor="gray", alpha=0.3)
+    ax.axhspan(1, 5, facecolor="gray", alpha=0.3)
+    ax.set_ylim(ylimits)
+
+    plt.tight_layout()
diff --git a/nidn/training/__init__.py b/nidn/training/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/nidn/training/losses/spectrum_loss.py b/nidn/training/losses/spectrum_loss.py
new file mode 100644
index 0000000..41db3ef
--- /dev/null
+++ b/nidn/training/losses/spectrum_loss.py
@@ -0,0 +1,66 @@
+from loguru import logger
+import torch
+
+
+def _spectrum_loss_fn(
+    produced_R_spectrum,
+    produced_T_spectrum,
+    target_R_spectrum,
+    target_T_spectrum,
+    target_frequencies,
+    L=0.5,
+    include_absorption=False,
+):
+    """Computes the loss for the difference in reflectance,
+    transmittance and optionally absorption spectra.
+
+    Args:
+        produced_R_spectrum (torch.tensor): Produced reflectance spectrum.
+        produced_T_spectrum (torch.tensor): Produced transmittance spectrum.
+        target_R_spectrum (torch.tensor): Target reflectance spectrum.
+        target_T_spectrum (torch.tensor): Target transmittance spectrum.
+        target_frequencies (list of float): Target frequencies.
+        L (float, optional): Exponent of the loss, e.g. L=2 for squared loss values. Defaults to 0.5.
+        include_absorption (bool, optional): If absorption should be included in loss term. Defaults to False.
+
+    Returns:
+        [torch.tensor],renormalized: Loss values for reflectance, transmittance and optionally absorption and if renormalization happend.
+    """
+    logger.trace("Computing spectrum loss..")
+
+    R_loss, T_loss, A_loss = 0, 0, 0
+    # Tracks if renormization happened
+    renormalized = False
+
+    # Iterate over all frequencies
+    for prod_R, prod_T, target_R, target_T, freq in zip(
+        produced_R_spectrum,
+        produced_T_spectrum,
+        target_R_spectrum,
+        target_T_spectrum,
+        target_frequencies,
+    ):
+
+        # If the unphysical scenario of negative absorption, i.e.
+        # reflection + transmission > 1, then renormalize them to 1 and warn
+        if prod_R + prod_T > 1.0:
+            logger.warning(
+                f"R+T>1 for freq={freq:.4f}. Renormalizing this point in the loss..."
+            )
+            prod_R = prod_R / (prod_R + prod_T)
+            prod_T = prod_T / (prod_R + prod_T)
+            renormalized = True
+
+        # Accumulate losses
+        R_loss += torch.abs(prod_R - target_R) ** L
+        T_loss += torch.abs(prod_T - target_T) ** L
+
+        if include_absorption:
+            prod_A = 1 - prod_T - prod_R
+            target_A = 1 - target_T - target_R
+            A_loss += torch.abs(prod_A - target_A) ** L
+
+    if include_absorption:
+        return (A_loss + R_loss + T_loss) / (3 * len(produced_R_spectrum)), renormalized
+    else:
+        return (R_loss + T_loss) / (2 * len(produced_R_spectrum)), renormalized
diff --git a/nidn/training/model/__init__.py b/nidn/training/model/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/nidn/training/model/init_network.py b/nidn/training/model/init_network.py
new file mode 100644
index 0000000..ca83c00
--- /dev/null
+++ b/nidn/training/model/init_network.py
@@ -0,0 +1,35 @@
+from dotmap import DotMap
+from loguru import logger
+import torch.nn as nn
+
+from ..utils.nerf import NERF
+from ..utils.siren import Siren
+
+
+def init_network(
+    run_cfg: DotMap,
+):
+    """Network architecture. Note that the dimensionality of the first linear layer must match the output of the encoding chosen in the config file.
+    Returns:
+        torch model: Initialized model
+    """
+    logger.debug("Initializing model..."+run_cfg.model_type)
+    if run_cfg.model_type == "nerf":
+        return NERF(
+            in_features=run_cfg.encoding_dim,
+            n_neurons=run_cfg.n_neurons,
+            activation=nn.Sigmoid(),
+            skip=[4],
+            hidden_layers=run_cfg.hidden_layers,
+        )
+    elif run_cfg.model_type == "siren":
+        return Siren(
+            in_features=run_cfg.encoding_dim,
+            out_features=run_cfg.out_features,
+            hidden_features=run_cfg.n_neurons,
+            hidden_layers=run_cfg.hidden_layers,
+            outermost_linear=True,
+            outermost_activation=nn.Sigmoid(),
+            first_omega_0=run_cfg.siren_omega,
+            hidden_omega_0=run_cfg.siren_omega,
+        )
diff --git a/nidn/training/model/model_to_eps_grid.py b/nidn/training/model/model_to_eps_grid.py
new file mode 100644
index 0000000..701e8c2
--- /dev/null
+++ b/nidn/training/model/model_to_eps_grid.py
@@ -0,0 +1,68 @@
+import torch
+from dotmap import DotMap
+
+
+def model_to_eps_grid(model, run_cfg: DotMap):
+    """Computes a 4D grid of epsilons for the given model.
+
+    Args:
+        model (torch.model): Trained neural network model. Should map one 4D input to a [real,imag] epsilon.
+        run_cfg (DotMap): Configuration for the run.
+
+    Returns:
+        [torch.tensor]: Resulting 4D [real,imag] epsilon grid
+    """
+
+    # Get the grid ticks
+    x = torch.linspace(-1, 1, run_cfg.Nx)
+    y = torch.linspace(-1, 1, run_cfg.Ny)
+    z = torch.linspace(-1, 1, run_cfg.N_layers)
+
+    # Linearly spaced frequency points
+    #     freq = torch.linspace(-1,1,len(run_cfg.target_frequencies))
+
+    # Logspaced frequency points
+    # Normalize to max val = 1
+    freq = torch.tensor(run_cfg.target_frequencies / max(run_cfg.target_frequencies))
+    # Transform to -1 to 1
+    freq = (freq * 2) - 1
+
+    # Create a meshgrid from the grid ticks
+    X, Y, Z, FREQ = torch.meshgrid((x, y, z, freq))
+
+    # Reshape the grid to have a Nx4 tensor
+    nn_inputs = torch.cat(
+        (X.reshape(-1, 1), Y.reshape(-1, 1), Z.reshape(-1, 1), FREQ.reshape(-1, 1)),
+        dim=1,
+    )
+
+    # Compute model output at each grid point
+    out = model(nn_inputs)
+
+    # Reshape the output to have a 4D tensor again
+    # Note we ouput real and imaginary parts separatly, hence 2*N_freq
+    out = out.reshape(run_cfg.Nx, run_cfg.Ny, run_cfg.N_layers, 2 * run_cfg.N_freq)
+
+    # Oversample the created eps grid
+    out = torch.repeat_interleave(out, run_cfg.eps_oversampling, dim=0)
+    out = torch.repeat_interleave(out, run_cfg.eps_oversampling, dim=1)
+
+    # Initialize the epsilon grid
+    eps = torch.zeros(
+        [run_cfg.Nx, run_cfg.Ny, run_cfg.N_layers, run_cfg.N_freq],
+        dtype=torch.cfloat,
+    )
+
+    # Net out is [0,1] thus we transform to desired real and imaginary ranges
+    # first half contains real entries
+    eps.real = out[:, :, :, 0 : run_cfg.N_freq]
+    eps.real = (eps.real * (run_cfg.real_max_eps - run_cfg.real_min_eps)).clip(
+        run_cfg.real_min_eps, run_cfg.real_max_eps
+    )
+    # second half imaginary
+    eps.imag = out[:, :, :, run_cfg.N_freq :]
+    eps.imag = (eps.imag * (run_cfg.imag_max_eps - run_cfg.imag_min_eps)).clip(
+        run_cfg.imag_min_eps, run_cfg.imag_max_eps
+    )
+
+    return eps
diff --git a/nidn/training/run_training.py b/nidn/training/run_training.py
new file mode 100644
index 0000000..da78717
--- /dev/null
+++ b/nidn/training/run_training.py
@@ -0,0 +1,141 @@
+from dotmap import DotMap
+
+import torch
+import numpy as np
+import numpy.typing as npt
+from collections import deque
+from copy import deepcopy
+from loguru import logger
+
+from ..utils.fix_random_seeds import fix_random_seeds
+from .model.init_network import init_network
+from ..trcwa.compute_spectrum import compute_spectrum
+from ..trcwa.compute_target_frequencies import compute_target_frequencies
+from .model.model_to_eps_grid import model_to_eps_grid
+from .losses.spectrum_loss import _spectrum_loss_fn
+from .utils.validate_config import _validate_config
+
+
+def run_training(
+    run_cfg: DotMap,
+    target_reflectance_spectrum: npt.NDArray,
+    target_transmittance_spectrum: npt.NDArray,
+    model=None,
+):
+    """Runs a training run with the passed config, target reflectance and transmittance spectra. Optionally a model can be passed to continue training
+
+    Args:
+        run_cfg (DotMap): Run configuration.
+        target_reflectance_spectrum (np.array): Target reflectance spectrum.
+        target_transmittance_spectrum (np.array): Target transmittance spectrum.
+        model (torch.model, optional): Model to continue training. If None, a new model will be created according to the run configuration. Defaults to None.
+
+    Returns:
+        torch.model,DotMap: The best model achieved in the training run, and the loss results of the training run.
+    """
+    logger.trace("Initializing training...")
+
+    # Validate config
+    _validate_config(run_cfg)
+
+    # Fix random seed for reproducibility
+    fix_random_seeds(run_cfg.seed)
+
+    # Determine target frequencies
+    run_cfg.target_frequencies = compute_target_frequencies(
+        run_cfg.physical_wavelength_range[0],
+        run_cfg.physical_wavelength_range[1],
+        run_cfg.N_freq,
+    )
+
+    logger.debug("Computed target frequencies:")
+    logger.debug(run_cfg.target_frequencies)
+
+    # Init model
+    if model is None:
+        model = init_network(run_cfg)
+
+    # Initialize some utility
+    optimizer = torch.optim.Adam(model.parameters(), lr=run_cfg.learning_rate)
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+        optimizer, factor=0.75, patience=200, min_lr=1e-6, verbose=True
+    )
+
+    results = DotMap()
+
+    # When a new network is created we init empty training logs
+    results.loss_log = []
+    results.weighted_average_log = []
+    weighted_average = deque([], maxlen=20)
+
+    # And store the best results
+    best_loss = np.inf
+    best_model_state_dict = model.state_dict()
+
+    logger.trace("Starting training...")
+    # Training Loop
+    for it in range(run_cfg.iterations):
+        torch.cuda.empty_cache()
+
+        # Compute the epsilon values from the model
+        eps_grid = model_to_eps_grid(model, run_cfg)
+
+        # Compute the spectrum using TRCWA for this grid
+        produced_R_spectrum, produced_T_spectrum = compute_spectrum(eps_grid, run_cfg)
+
+        # Compute loss between target spectrum and
+        # the one from the current network structure
+        spectrum_loss, renormalized = _spectrum_loss_fn(
+            produced_R_spectrum,
+            produced_T_spectrum,
+            target_reflectance_spectrum,
+            target_transmittance_spectrum,
+            run_cfg.target_frequencies,
+            run_cfg.L,
+            run_cfg.absorption_loss,
+        )
+
+        loss = spectrum_loss
+
+        # We store the model if it has the lowest loss yet
+        # (this is to avoid losing good results during a run that goes wild)
+        if loss < best_loss:
+            best_loss = loss
+            logger.info(
+                f"New Best={loss.item():.4f} SpectrumLoss={spectrum_loss.detach().item():4f}"
+            )
+            if not renormalized:
+                logger.debug("Saving model state...")
+                best_model_state_dict = deepcopy(model.state_dict())
+
+        # Update the loss trend indicators
+        weighted_average.append(loss.item())
+
+        # Update the logs
+        results.weighted_average_log.append(np.mean(weighted_average))
+        results.loss_log.append(loss.item())
+
+        # Print every i iterations
+        if it % 5 == 0:
+            wa_out = np.mean(weighted_average)
+            logger.info(
+                f"It={it}\t loss={loss.item():.3e}\t  weighted_average={wa_out:.3e}\t SpectrumLoss={spectrum_loss.detach().item():4f}"
+            )
+
+        # Zeroes the gradient (otherwise would accumulate)
+        optimizer.zero_grad()
+
+        # Compute gradient of the loss with respect to model parameters
+        loss.backward()
+
+        # Update model parameters
+        optimizer.step()
+
+        # Perform a step in LR scheduler to update LR if necessary
+        scheduler.step(loss.item())
+
+    logger.trace("Reloading best model state...")
+    # Return best model in the end
+    model.load_state_dict(best_model_state_dict)
+
+    return model, results
diff --git a/nidn/training/utils/__init__.py b/nidn/training/utils/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/nidn/training/utils/abs_layer.py b/nidn/training/utils/abs_layer.py
new file mode 100644
index 0000000..195a651
--- /dev/null
+++ b/nidn/training/utils/abs_layer.py
@@ -0,0 +1,13 @@
+from torch import nn
+import torch
+
+
+class AbsLayer(nn.Module):
+    """Very simple activation layer to allow different abs layer activations of the siren.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, input):
+        return torch.abs(input)
diff --git a/nidn/training/utils/nerf.py b/nidn/training/utils/nerf.py
new file mode 100644
index 0000000..4f58a66
--- /dev/null
+++ b/nidn/training/utils/nerf.py
@@ -0,0 +1,61 @@
+from torch import nn
+import torch
+
+
+class NERFLayer(nn.Module):
+    """Layer as used in NERF."""
+
+    def __init__(self, in_features, out_features, bias=True, activation=nn.ReLU()):
+        super().__init__()
+        self.linear = nn.Linear(in_features, out_features, bias=bias)
+        self.init_weights()
+        self.activation = activation
+
+    def init_weights(self):
+        nn.init.xavier_uniform_(self.linear.weight)
+        nn.init.uniform_(self.linear.bias.data, -0.0, 0.0)
+
+    def forward(self, input):
+        return self.activation(self.linear(input))
+
+
+class NERF(nn.Module):
+    """NERF architecture as described by Mildenhall et al. (2020)"""
+
+    def __init__(
+        self,
+        in_features,
+        skip=[],
+        n_neurons=100,
+        activation=nn.Sigmoid(),
+        hidden_layers=8,
+    ):
+        super().__init__()
+        self.in_features = in_features
+        self.skip = skip
+
+        self.net = nn.ModuleList()
+        self.net.append(NERFLayer(in_features, n_neurons))
+
+        for i in range(hidden_layers):
+            if i in self.skip:
+                self.net.append(NERFLayer(n_neurons + in_features, n_neurons))
+            else:
+                self.net.append(NERFLayer(n_neurons, n_neurons))
+
+        self.net.append(NERFLayer(n_neurons, 1, activation=activation))
+
+    def forward(self, x):
+        # save for skip connection
+        identity = x
+
+        # compute first layer
+        out = self.net[0].forward(x)
+
+        # compute all other layers and apply skip where requested
+        for layer_idx in range(1, len(self.net)):
+            out = self.net[layer_idx].forward(out)
+            if layer_idx in self.skip:
+                out = torch.cat([out, identity], dim=1)
+
+        return out
diff --git a/nidn/training/utils/siren.py b/nidn/training/utils/siren.py
new file mode 100644
index 0000000..9dcdde1
--- /dev/null
+++ b/nidn/training/utils/siren.py
@@ -0,0 +1,146 @@
+# From https://github.com/vsitzmann/siren/blob/master/explore_siren.ipynb
+from torch import nn
+import torch
+import numpy as np
+
+from .abs_layer import AbsLayer
+
+
+class SineLayer(nn.Module):
+    """Sinuosidal layer for the SIREN model."""
+
+    # See paper sec. 3.2, final paragraph, and supplement Sec. 1.5 for discussion of omega_0.
+
+    # If is_first=True, omega_0 is a frequency factor which simply multiplies the activations before the
+    # nonlinearity. Different signals may require different omega_0 in the first layer - this is a
+    # hyperparameter.
+
+    # If is_first=False, then the weights will be divided by omega_0 so as to keep the magnitude of
+    # activations constant, but boost gradients to the weight matrix (see supplement Sec. 1.5)
+
+    def __init__(
+        self, in_features, out_features, bias=True, is_first=False, omega_0=30
+    ):
+        super().__init__()
+        self.omega_0 = omega_0
+        self.is_first = is_first
+
+        self.in_features = in_features
+        self.linear = nn.Linear(in_features, out_features, bias=bias)
+
+        self.init_weights()
+
+    def init_weights(self):
+        with torch.no_grad():
+            if self.is_first:
+                self.linear.weight.uniform_(-1 / self.in_features, 1 / self.in_features)
+            else:
+                self.linear.weight.uniform_(
+                    -np.sqrt(6 / self.in_features) / self.omega_0,
+                    np.sqrt(6 / self.in_features) / self.omega_0,
+                )
+
+    def forward(self, input):
+        return torch.sin(self.omega_0 * self.linear(input))
+
+    def forward_with_intermediate(self, input):
+        # For visualization of activation distributions
+        intermediate = self.omega_0 * self.linear(input)
+        return torch.sin(intermediate), intermediate
+
+
+class Siren(nn.Module):
+    """Siren model as described in the paper and provided notebooks."""
+
+    def __init__(
+        self,
+        in_features,
+        hidden_features,
+        hidden_layers,
+        out_features,
+        outermost_linear=False,
+        outermost_activation=AbsLayer(),
+        first_omega_0=30,
+        hidden_omega_0=30.0,
+    ):
+
+        super().__init__()
+
+        self.in_features = in_features
+        self.net = []
+        self.net.append(
+            SineLayer(
+                in_features, hidden_features, is_first=True, omega_0=first_omega_0
+            )
+        )
+
+        for i in range(hidden_layers):
+            self.net.append(
+                SineLayer(
+                    hidden_features,
+                    hidden_features,
+                    is_first=False,
+                    omega_0=hidden_omega_0,
+                )
+            )
+
+        if outermost_linear:
+            final_linear = nn.Linear(hidden_features, out_features)
+
+            with torch.no_grad():
+                final_linear.weight.uniform_(
+                    -np.sqrt(6 / hidden_features) / hidden_omega_0,
+                    np.sqrt(6 / hidden_features) / hidden_omega_0,
+                )
+
+            self.net.append(final_linear)
+            self.net.append(outermost_activation)
+        else:
+            self.net.append(
+                SineLayer(
+                    hidden_features,
+                    out_features,
+                    is_first=False,
+                    omega_0=hidden_omega_0,
+                )
+            )
+
+        self.net = nn.Sequential(*self.net)
+
+    def forward(self, coords):
+        coords = (
+            coords.clone().detach().requires_grad_(True)
+        )  # allows to take derivative w.r.t. input
+        output = self.net(coords)
+        return output
+
+    def forward_with_activations(self, coords, retain_grad=False):
+        """Returns not only model output, but also intermediate activations.
+        Only used for visualizing activations later!"""
+        activations = OrderedDict()
+
+        activation_count = 0
+        x = coords.clone().detach().requires_grad_(True)
+        activations["input"] = x
+        for i, layer in enumerate(self.net):
+            if isinstance(layer, SineLayer):
+                x, intermed = layer.forward_with_intermediate(x)
+
+                if retain_grad:
+                    x.retain_grad()
+                    intermed.retain_grad()
+
+                activations[
+                    "_".join((str(layer.__class__), "%d" % activation_count))
+                ] = intermed
+                activation_count += 1
+            else:
+                x = layer(x)
+
+                if retain_grad:
+                    x.retain_grad()
+
+            activations["_".join((str(layer.__class__), "%d" % activation_count))] = x
+            activation_count += 1
+
+        return activations
diff --git a/nidn/training/utils/validate_config.py b/nidn/training/utils/validate_config.py
new file mode 100644
index 0000000..1ddd9e7
--- /dev/null
+++ b/nidn/training/utils/validate_config.py
@@ -0,0 +1,36 @@
+from dotmap import DotMap
+
+
+def _validate_config(cfg: DotMap):
+    """This function validates that all required entries are in the config.
+
+    Args:
+        cfg (DotMap): Run config you intend to use
+    """
+    required_keys = [
+        "seed",
+        "eps_oversampling",
+        "real_min_eps",
+        "real_max_eps",
+        "imag_min_eps",
+        "imag_max_eps",
+        "Nx",
+        "Ny",
+        "N_layers",
+        "physical_wavelength_range",
+        "N_freq",
+        "model_type",
+        "encoding_dim",
+        "out_features",
+        "siren_omega",
+        "n_neurons",
+        "hidden_layers",
+        "learning_rate",
+        "iterations",
+        "L",
+        "absorption_loss",
+    ]
+
+    for key in required_keys:
+        if key not in cfg:
+            raise ValueError(f"Missing required key: {key}")
diff --git a/nidn/trcwa/__init__.py b/nidn/trcwa/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/nidn/trcwa/compute_spectrum.py b/nidn/trcwa/compute_spectrum.py
new file mode 100644
index 0000000..0efc727
--- /dev/null
+++ b/nidn/trcwa/compute_spectrum.py
@@ -0,0 +1,46 @@
+from dotmap import DotMap
+import torch
+from loguru import logger
+
+from .init_trcwa import _init_trcwa
+
+
+def compute_spectrum(eps_grid, run_cfg: DotMap):
+    """Evaluates TRCWA for the given epsilon values and run configuration.
+
+    Args:
+        eps_grid (torch.tensor): Epsilon values to evaluate TRCWA.
+        run_cfg (DotMap): Run configuration.
+
+    Raises:
+        ValueError: Raises an error epsilon contains NaN values.
+
+    Returns:
+        tuple of lists: Reflection coefficients, transmission coefficients for the given epsilon values and frequencies.
+    """
+    logger.trace("Computing spectrum")
+    produced_R_spectrum = []
+    produced_T_spectrum = []
+
+    logger.debug("Testing input for NaNs")
+    # Check for NaNs in eps_grid
+    if torch.any(torch.isnan(eps_grid)):
+        raise ValueError("Model output became NaN...", eps_grid)
+
+    logger.debug("Checking frequencies were passed")
+    assert "target_frequencies" in run_cfg, "No frequencies were passed."
+
+    logger.debug("Iterating over passed frequencies")
+    # For each frequency, evaluate TRCWA and with the grid of epsilon values
+    for idx, freq in enumerate(run_cfg.target_frequencies):
+
+        # Create TRCWA Object for this frequency
+        trcwa = _init_trcwa(eps_grid[:, :, :, idx], freq)
+
+        # Compute the spectrum
+        reflectance, transmittance = trcwa.RT_Solve(normalize=1)
+
+        produced_R_spectrum.append(reflectance)
+        produced_T_spectrum.append(transmittance)
+
+    return produced_R_spectrum, produced_T_spectrum
diff --git a/nidn/trcwa/compute_target_frequencies.py b/nidn/trcwa/compute_target_frequencies.py
new file mode 100644
index 0000000..2aa6540
--- /dev/null
+++ b/nidn/trcwa/compute_target_frequencies.py
@@ -0,0 +1,28 @@
+import numpy as np
+
+from ..utils.convert_units import phys_wl_to_wl, freq_to_wl
+
+
+def compute_target_frequencies(
+    min_physical_wl: float, max_physical_wl: float, N_freq: int
+):
+    """Computes the target frequencies for a given set of physical frequencies.
+
+    Args:
+        min_physical_wl (float): Minimum physical frequency in the target spectrum.
+        max_physical_wl (float): Maximum physical frequency in the target spectrum.
+        N_freq (int): The number of target frequencies.
+
+    Returns:
+        np.array: Corresponding target frequencies
+    """
+    max_physical_freq = freq_to_wl(phys_wl_to_wl(min_physical_wl))
+    min_physical_freq = freq_to_wl(phys_wl_to_wl(max_physical_wl))
+
+    logval_min = np.log10(min_physical_freq)
+    logval_max = np.log10(max_physical_freq)
+
+    target_frequencies = np.flip(
+        np.logspace(logval_min, logval_max, num=N_freq, base=10)
+    )
+    return target_frequencies
diff --git a/nidn/trcwa/constants.py b/nidn/trcwa/constants.py
new file mode 100644
index 0000000..fe4fe2d
--- /dev/null
+++ b/nidn/trcwa/constants.py
@@ -0,0 +1,37 @@
+import torch
+
+# NB: Here, we assume that the standard is micrometre (um).
+TRCWA_UNIT_MAGNITUDE = 10 ** (-6)
+
+# m /s
+TRCWA_SPEED_OF_LIGHT = 299792458.0
+
+TRCWA_Q_ABS = 1e5
+
+# Truncation order (actual number might be smaller)
+TRCWA_NG = 11
+
+TRCWA_L1 = [1, 0]  # everything is set relative to this
+TRCWA_L2 = [0, 1]  # it is easiest to say 1 = 1 micron
+
+# Dielectric for top and bottom layer
+TRCWA_VACUUM_EPS = torch.tensor(1.0)
+
+TRCWA_PER_LAYER_THICKNESS = 1.0
+
+# Planewave excitation
+TRCWA_PLANEWAVE = {"p_amp": 0, "s_amp": 1, "p_phase": 0, "s_phase": 0}
+
+# Angles
+TRCWA_THETA = 0.0
+TRCWA_PHI = 0.0
+
+# elementary charge, C
+TRCWA_E_CHARGE = 1.602176634 * 10 ** (-19)
+# weight of electron, kg
+TRCWA_M_E = 9.1093837015 * 10 ** (-31)
+# vacuum permittivity, F/m
+TRCWA_EPS_0 = 8.8541878128 * 10 ** (-12)
+
+# There is no torch.pi so we define it here
+TRCWA_PI = torch.acos(torch.zeros(1)).item() * 2  # which is 3.1415927410125732
diff --git a/nidn/trcwa/get_frequency_points.py b/nidn/trcwa/get_frequency_points.py
new file mode 100644
index 0000000..5735376
--- /dev/null
+++ b/nidn/trcwa/get_frequency_points.py
@@ -0,0 +1,28 @@
+from dotmap import DotMap
+from .compute_target_frequencies import compute_target_frequencies
+from ..utils.convert_units import freq_to_wl, wl_to_phys_wl
+
+
+def get_frequency_points(run_cfg: DotMap):
+    """Computes both the physical and normalized frequency points which will be used.
+    This can be helpful to determine the points at which you need to define the target
+    spectrum.
+
+    Args:
+        run_cfg (DotMap): Config for the run.
+
+    Returns:
+        physical_target_frequencies (list): List of physical frequencies.
+        target_frequencies (list): List of target frequencies.
+    """
+
+    # Determine target frequencies
+    run_cfg.target_frequencies = compute_target_frequencies(
+        run_cfg.physical_wavelength_range[0],
+        run_cfg.physical_wavelength_range[1],
+        run_cfg.N_freq,
+    )
+
+    physical_target_frequencies = wl_to_phys_wl(freq_to_wl(run_cfg.target_frequencies))
+
+    return physical_target_frequencies, run_cfg.target_frequencies
diff --git a/nidn/trcwa/init_trcwa.py b/nidn/trcwa/init_trcwa.py
new file mode 100644
index 0000000..cd9f980
--- /dev/null
+++ b/nidn/trcwa/init_trcwa.py
@@ -0,0 +1,65 @@
+from .constants import *
+from .trcwa import TRCWA
+
+
+def _init_trcwa(eps_grid, target_frequency):
+    """Creates a TRCWA object matching the given eps_grid and target_frequency.
+
+    Args:
+        eps_grid (torch.tensor): Grid of epsilon values. Should be [Nx, Ny, N_layers].
+        target_frequency (float): Target frequency for this simulation run.
+
+    Returns:
+        TRCWA obj: The created object which is ready to compute the spectrum.
+    """
+
+    Nx, Ny, N_layers = eps_grid.shape[0:3]
+
+    # Squeeze out Nx=1, Ny=1 dimension (for uniform layer)
+    eps_grid = eps_grid.squeeze(0)
+    eps_grid = eps_grid.squeeze(1)
+
+    # Adding a small imaginary part to the frequency to avoid singularities in RCWA
+    # TODO Add a nice reference for this?
+    freqcmp = target_frequency * (1 + (1j / (2.0 * TRCWA_Q_ABS)))
+
+    # Initialize TRCWA object
+    trcwa = TRCWA(
+        TRCWA_NG,
+        TRCWA_L1,
+        TRCWA_L2,
+        freqcmp,
+        TRCWA_THETA,
+        TRCWA_PHI,
+        verbose=0,
+    )
+
+    # Add vacuum layer at the top
+    trcwa.Add_LayerUniform(TRCWA_PER_LAYER_THICKNESS, TRCWA_VACUUM_EPS)
+
+    # Add material layers (homogeneous if Nx and Ny are 1)
+    for layer in range(N_layers):
+        if Nx > 1 or Ny > 1:
+            trcwa.Add_LayerGrid(TRCWA_PER_LAYER_THICKNESS, Nx, Ny)
+        else:
+            trcwa.Add_LayerUniform(TRCWA_PER_LAYER_THICKNESS, eps_grid[layer])
+
+    # Add vacuum layer at the bottom
+    trcwa.Add_LayerUniform(TRCWA_PER_LAYER_THICKNESS, TRCWA_VACUUM_EPS)
+
+    # Initialize the object properly
+    trcwa.Init_Setup()
+
+    trcwa.MakeExcitationPlanewave(
+        TRCWA_PLANEWAVE["p_amp"],
+        TRCWA_PLANEWAVE["p_phase"],
+        TRCWA_PLANEWAVE["s_amp"],
+        TRCWA_PLANEWAVE["s_phase"],
+        order=0,
+    )
+
+    # Apply eps if heterogeneous layers
+    if Nx > 1 or Ny > 1:
+        trcwa.GridLayer_geteps(eps_grid)
+
+    return trcwa
diff --git a/nidn/trcwa/load_material_data.py b/nidn/trcwa/load_material_data.py
new file mode 100644
index 0000000..a3b85ee
--- /dev/null
+++ b/nidn/trcwa/load_material_data.py
@@ -0,0 +1,27 @@
+import pandas
+import torch
+
+
+def _load_material_data(name, target_frequencies=[1.0]):
+    """Loads the passed wavelength, n, k data from the passed csv file for the closest frequencies and returns epsilon (permittivity).
+
+    Args:
+        name (str): Path to csv.
+        target_frequencies (list, optional): Target frequencies to compute epsilon for. Defaults to [1.0].
+
+    Returns:
+        torch.tensor: Epsilon (complex permittivity) for the material
+    """
+    csv_data = pandas.read_csv(name, delimiter="\t")
+
+    eps = []
+    for freq in target_frequencies:
+        wl = 1.0 / freq
+        result_index = csv_data["Wavelength"].sub(wl).abs().idxmin()
+        entry = csv_data.iloc[result_index]
+        real = entry.n * entry.n - entry.k * entry.k
+        imag = 2 * entry.n * entry.k
+        eps.append([real + imag * 1.0j])
+
+    eps = torch.tensor(eps)
+    return eps
diff --git a/nidn/trcwa/trcwa.py b/nidn/trcwa/trcwa.py
new file mode 100644
index 0000000..1ba015b
--- /dev/null
+++ b/nidn/trcwa/trcwa.py
@@ -0,0 +1,812 @@
+import torch
+
+from .constants import *
+
+from .utils.fft_funs import Epsilon_fft, get_ifft
+from .utils.kbloch import Lattice_Reciprocate, Lattice_getG, Lattice_SetKs
+from .utils.torch_functions import torch_transpose, torch_dot, torch_eye, torch_zeros
+
+
+class TRCWA:
+    def __init__(self, nG, L1, L2, freq, theta, phi, verbose=1):
+        """The time harmonic convention is exp(-i omega t), speed of light = 1
+
+        Two kinds of layers are currently supported: uniform layer,
+        patterned layer from grids. Interface for patterned layer by
+        direct analytic expression of Fourier series is included, but
+        no examples inclded so far.
+
+        nG: truncation order, but the actual truncation order might not be nG
+        L1,L2: lattice vectors, in the list format, (x,y)
+
+        """
+        self.freq = freq
+        self.omega = 2 * TRCWA_PI * freq + 0.0j
+        self.L1 = L1
+        self.L2 = L2
+        self.phi = torch.tensor(phi)
+        self.theta = torch.tensor(theta)
+        self.nG = nG
+        self.verbose = verbose
+        self.Layer_N = 0  # total number of layers
+
+        # the length of the following variables = number of total layers
+        self.thickness_list = []
+        self.id_list = (
+            []
+        )  # [type, No., No. in patterned/uniform, No. in its family] starting from 0
+        # type:0 for uniform, 1 for Grids, 2 for Fourier
+
+        self.kp_list = []
+        self.q_list = []  # eigenvalues
+        self.phi_list = []  # eigenvectors
+
+        # Uniform layer
+        self.Uniform_ep_list = []
+        self.Uniform_N = 0
+
+        # Patterned layer
+        self.Patterned_N = 0  # total number of patterned layers
+        self.Patterned_epinv_list = []
+        self.Patterned_ep2_list = []
+
+        # patterned layer from Grids
+        self.GridLayer_N = 0
+        self.GridLayer_Nxy_list = []
+
+        # layers of analytic Fourier series (e.g. circles)
+        self.FourierLayer_N = 0
+        self.FourierLayer_params = []
+
+    def Add_LayerUniform(self, thickness, epsilon):
+        # assert type(thickness) == float, 'thickness should be a float'
+
+        self.id_list.append([0, self.Layer_N, self.Uniform_N])
+        self.Uniform_ep_list.append(epsilon)
+        self.thickness_list.append(torch.tensor(thickness))
+
+        self.Layer_N += 1
+        self.Uniform_N += 1
+
+    def Add_LayerGrid(self, thickness, Nx, Ny):
+        self.thickness_list.append(thickness)
+        self.GridLayer_Nxy_list.append([Nx, Ny])
+        self.id_list.append([1, self.Layer_N, self.Patterned_N, self.GridLayer_N])
+
+        self.Layer_N += 1
+        self.GridLayer_N += 1
+        self.Patterned_N += 1
+
+    def Add_LayerFourier(self, thickness, params):
+        self.thickness_list.append(thickness)
+        self.FourierLayer_params.append(params)
+        self.id_list.append([2, self.Layer_N, self.Patterned_N, self.FourierLayer_N])
+
+        self.Layer_N += 1
+        self.Patterned_N += 1
+        self.FourierLayer_N += 1
+
+    def Init_Setup(self, Pscale=1.0, Gmethod=0):
+        """
+        Set up reciprocal lattice (Gmethod:truncation scheme, 0 for circular, 1 for rectangular)
+        Pscale: scale the period
+        Compute eigenvalues for uniform layers
+        Initialize vectors for patterned layers
+        """
+        kx0 = (
+            self.omega
+            * torch.sin(self.theta)
+            * torch.cos(self.phi)
+            * torch.sqrt(self.Uniform_ep_list[0])
+        )
+        ky0 = (
+            self.omega
+            * torch.sin(self.theta)
+            * torch.sin(self.phi)
+            * torch.sqrt(self.Uniform_ep_list[0])
+        )
+
+        # set up reciprocal lattice
+        self.Lk1, self.Lk2 = Lattice_Reciprocate(self.L1, self.L2)
+        self.G, self.nG = Lattice_getG(self.nG, self.Lk1, self.Lk2, method=Gmethod)
+
+        self.Lk1 = self.Lk1 / Pscale
+        self.Lk2 = self.Lk2 / Pscale
+        # self.kx = kx0 + 2*torch.pi*(self.Lk1[0]*self.G[:,0]+self.Lk2[0]*self.G[:,1])
+        # self.ky = ky0 + 2*torch.pi*(self.Lk1[1]*self.G[:,0]+self.Lk2[1]*self.G[:,1])
+        self.kx, self.ky = Lattice_SetKs(self.G, kx0, ky0, self.Lk1, self.Lk2)
+
+        # normalization factor for energies off normal incidence
+        self.normalization = torch.sqrt(self.Uniform_ep_list[0]) / torch.cos(self.theta)
+
+        # if comm.rank == 0 and verbose>0:
+        if self.verbose > 0:
+            print("Total nG = ", self.nG)
+
+        self.Patterned_ep2_list = [None] * self.Patterned_N
+        self.Patterned_epinv_list = [None] * self.Patterned_N
+        for i in range(self.Layer_N):
+            if self.id_list[i][0] == 0:
+                ep = self.Uniform_ep_list[self.id_list[i][2]]
+                kp = MakeKPMatrix(self.omega, 0, 1.0 / ep, self.kx, self.ky)
+                self.kp_list.append(kp)
+
+                q, phi = SolveLayerEigensystem_uniform(self.omega, self.kx, self.ky, ep)
+                self.q_list.append(q)
+                self.phi_list.append(phi)
+            else:
+                self.kp_list.append(None)
+                self.q_list.append(None)
+                self.phi_list.append(None)
+
+    def MakeExcitationPlanewave(
+        self, p_amp, p_phase, s_amp, s_phase, order=0, direction="forward"
+    ):
+        """
+        Front incidence
+        """
+        self.direction = direction
+        theta = self.theta
+        phi = self.phi
+        a0 = torch_zeros(2 * self.nG, dtype=complex)
+        bN = torch_zeros(2 * self.nG, dtype=complex)
+        if direction == "forward":
+            tmp1 = torch_zeros(2 * self.nG, dtype=complex)
+            tmp1[order] = 1.0
+            a0 = a0 + tmp1 * (
+                -s_amp
+                * torch.cos(theta)
+                * torch.cos(phi)
+                * torch.exp(torch.tensor(1j * s_phase))
+                - p_amp * torch.sin(phi) * torch.exp(torch.tensor(1j * p_phase))
+            )
+
+            tmp2 = torch_zeros(2 * self.nG, dtype=complex)
+            tmp2[order + self.nG] = 1.0
+            a0 = a0 + tmp2 * (
+                -s_amp
+                * torch.cos(theta)
+                * torch.sin(phi)
+                * torch.exp(torch.tensor(1j * s_phase))
+                + p_amp * torch.cos(phi) * torch.exp(torch.tensor(1j * p_phase))
+            )
+        elif direction == "backward":
+            tmp1 = torch_zeros(2 * self.nG, dtype=complex)
+            tmp1[order] = 1.0
+            bN = bN + tmp1 * (
+                -s_amp
+                * torch.cos(theta)
+                * torch.cos(phi)
+                * torch.exp(torch.tensor(1j * s_phase))
+                - p_amp * torch.sin(phi) * torch.exp(torch.tensor(1j * p_phase))
+            )
+
+            tmp2 = torch_zeros(2 * self.nG, dtype=complex)
+            tmp2[order + self.nG] = 1.0
+            bN = bN + tmp2 * (
+                -s_amp
+                * torch.cos(theta)
+                * torch.sin(phi)
+                * torch.exp(torch.tensor(1j * s_phase))
+                + p_amp * torch.cos(phi) * torch.exp(torch.tensor(1j * p_phase))
+            )
+
+        self.a0 = a0
+        self.bN = bN
+
+    def GridLayer_geteps(self, ep_all):
+        """
+        Fourier transform + eigenvalue for grid layer
+        """
+        ptri = 0
+        ptr = 0
+        for i in range(self.Layer_N):
+            if self.id_list[i][0] != 1:
+                continue
+
+            Nx = self.GridLayer_Nxy_list[ptri][0]
+            Ny = self.GridLayer_Nxy_list[ptri][1]
+            dN = 1.0 / Nx / Ny
+
+            if len(ep_all) == 3 and ep_all[0].ndim > 0:
+                ep_grid = [
+                    torch.reshape(ep_all[0][ptr : ptr + Nx * Ny], [Nx, Ny]),
+                    torch.reshape(ep_all[1][ptr : ptr + Nx * Ny], [Nx, Ny]),
+                    torch.reshape(ep_all[2][ptr : ptr + Nx * Ny], [Nx, Ny]),
+                ]
+            elif len(ep_all.shape) == 3:
+                ep_grid = ep_all[:, :, i - 1]
+            else:
+                ep_grid = torch.reshape(
+                    torch.tensor(ep_all[ptr : ptr + Nx * Ny]), [Nx, Ny]
+                )
+
+            epinv, ep2 = Epsilon_fft(dN, ep_grid, self.G)
+
+            self.Patterned_epinv_list[self.id_list[i][2]] = epinv
+            self.Patterned_ep2_list[self.id_list[i][2]] = ep2
+
+            kp = MakeKPMatrix(self.omega, 1, epinv, self.kx, self.ky)
+            self.kp_list[self.id_list[i][1]] = kp
+
+            q, phi = SolveLayerEigensystem(self.omega, self.kx, self.ky, kp, ep2)
+            self.q_list[self.id_list[i][1]] = q
+            self.phi_list[self.id_list[i][1]] = phi
+
+            ptr += Nx * Ny
+            ptri += 1
+
+    def Return_eps(self, which_layer, Nx, Ny, component="xx"):
+        """
+        For patterned layer component = 'xx','xy','yx','yy','zz'
+        For uniform layer, currently it's assumed to be isotropic
+        """
+        i = which_layer
+        # uniform layer
+        if self.id_list[i][0] == 0:
+            ep = self.Uniform_ep_list[self.id_list[i][2]]
+            return torch.ones((Nx, Ny)) * ep
+
+        # patterned layer
+        elif self.id_list[i][0] == 1:
+            if component == "zz":
+                epk = torch.inv(self.Patterned_epinv_list[self.id_list[i][2]])
+            elif component == "xx":
+                epk = self.Patterned_ep2_list[self.id_list[i][2]][: self.nG, : self.nG]
+            elif component == "xy":
+                epk = self.Patterned_ep2_list[self.id_list[i][2]][: self.nG, self.nG :]
+            elif component == "yx":
+                epk = self.Patterned_ep2_list[self.id_list[i][2]][self.nG :, : self.nG]
+            elif component == "yy":
+                epk = self.Patterned_ep2_list[self.id_list[i][2]][self.nG :, self.nG :]
+
+            return get_ifft(Nx, Ny, epk[0, :], self.G)
+
+    def RT_Solve(self, normalize=0, byorder=0):
+        """
+        Reflection and transmission power computation
+        Returns 2R and 2T, following Victor's notation
+        Maybe because 2* makes S_z = 1 for H=1 in vacuum
+
+        if normalize = 1, it will be divided by n[0]*cos(theta)
+        """
+        aN, b0 = SolveExterior(
+            self.a0,
+            self.bN,
+            self.q_list,
+            self.phi_list,
+            self.kp_list,
+            self.thickness_list,
+        )
+
+        fi, bi = GetZPoyntingFlux(
+            self.a0,
+            b0,
+            self.omega,
+            self.kp_list[0],
+            self.phi_list[0],
+            self.q_list[0],
+            byorder=byorder,
+        )
+
+        fe, be = GetZPoyntingFlux(
+            aN,
+            self.bN,
+            self.omega,
+            self.kp_list[-1],
+            self.phi_list[-1],
+            self.q_list[-1],
+            byorder=byorder,
+        )
+
+        if self.direction == "forward":
+            R = torch.real(-bi)
+            T = torch.real(fe)
+        elif self.direction == "backward":
+            R = torch.real(fe)
+            T = torch.real(-bi)
+
+        if normalize == 1:
+            R = R * self.normalization
+            T = T * self.normalization
+        return R, T
+
+    def GetAmplitudes_noTranslate(self, which_layer):
+        """
+        returns fourier amplitude
+        """
+        if which_layer == 0:
+            aN, b0 = SolveExterior(
+                self.a0,
+                self.bN,
+                self.q_list,
+                self.phi_list,
+                self.kp_list,
+                self.thickness_list,
+            )
+            ai = self.a0
+            bi = b0
+
+        elif which_layer == self.Layer_N - 1:
+            aN, b0 = SolveExterior(
+                self.a0,
+                self.bN,
+                self.q_list,
+                self.phi_list,
+                self.kp_list,
+                self.thickness_list,
+            )
+            ai = aN
+            bi = self.bN
+
+        else:
+            ai, bi = SolveInterior(
+                which_layer,
+                self.a0,
+                self.bN,
+                self.q_list,
+                self.phi_list,
+                self.kp_list,
+                self.thickness_list,
+            )
+        return ai, bi
+
+    def GetAmplitudes(self, which_layer, z_offset):
+        """
+        returns fourier amplitude
+        """
+        if which_layer == 0:
+            aN, b0 = SolveExterior(
+                self.a0,
+                self.bN,
+                self.q_list,
+                self.phi_list,
+                self.kp_list,
+                self.thickness_list,
+            )
+            ai = self.a0
+            bi = b0
+
+        elif which_layer == self.Layer_N - 1:
+            aN, b0 = SolveExterior(
+                self.a0,
+                self.bN,
+                self.q_list,
+                self.phi_list,
+                self.kp_list,
+                self.thickness_list,
+            )
+            ai = aN
+            bi = self.bN
+
+        else:
+            ai, bi = SolveInterior(
+                which_layer,
+                self.a0,
+                self.bN,
+                self.q_list,
+                self.phi_list,
+                self.kp_list,
+                self.thickness_list,
+            )
+
+        ai, bi = TranslateAmplitudes(
+            self.q_list[which_layer], self.thickness_list[which_layer], z_offset, ai, bi
+        )
+
+        return ai, bi
+
+    def Solve_FieldFourier(self, which_layer, z_offset):
+        """
+        returns field amplitude in fourier space: [ex,ey,ez], [hx,hy,hz]
+        """
+        ai0, bi0 = self.GetAmplitudes_noTranslate(which_layer)
+        # ai, bi = self.GetAmplitudes(which_layer,z_offset)
+
+        if torch.isinstance(z_offset, float) or torch.isinstance(z_offset, int):
+            zl = [z_offset]
+        else:
+            zl = z_offset
+
+        eh = []
+        for zoff in zl:
+            ai, bi = TranslateAmplitudes(
+                self.q_list[which_layer],
+                self.thickness_list[which_layer],
+                zoff,
+                ai0,
+                bi0,
+            )
+            # hx, hy in Fourier space
+            fhxy = torch_dot(self.phi_list[which_layer], ai + bi)
+            fhx = fhxy[: self.nG]
+            fhy = fhxy[self.nG :]
+
+            # ex,ey in Fourier space
+            tmp1 = (ai - bi) / self.omega / self.q_list[which_layer]
+            tmp2 = torch_dot(self.phi_list[which_layer], tmp1)
+            fexy = torch_dot(self.kp_list[which_layer], tmp2)
+            fey = -fexy[: self.nG]
+            fex = fexy[self.nG :]
+
+            # hz in Fourier space
+            fhz = (self.kx * fey - self.ky * fex) / self.omega
+
+            # ez in Fourier space
+            fez = (self.ky * fhx - self.kx * fhy) / self.omega
+            if self.id_list[which_layer][0] == 0:
+                fez = fez / self.Uniform_ep_list[self.id_list[which_layer][2]]
+            else:
+                fez = torch_dot(
+                    self.Patterned_epinv_list[self.id_list[which_layer][2]], fez
+                )
+            eh.append([[fex, fey, fez], [fhx, fhy, fhz]])
+        return eh
+
+    def Solve_FieldOnGrid(self, which_layer, z_offset, Nxy=None):
+        # Nxy = [Nx,Ny], if not supplied, will use the number in patterned layer
+        # if single z_offset:  output [[ex,ey,ez],[hx,hy,hz]]
+        # if z_offset is list: output [[[ex1,ey1,ez1],[hx1,hy1,hz1]],  [[ex2,ey2,ez2],[hx2,hy2,hz2]] ...]
+
+        if torch.isinstance(Nxy, type(None)):
+            Nxy = self.GridLayer_Nxy_list[self.id_list[which_layer][3]]
+        Nx = Nxy[0]
+        Ny = Nxy[1]
+
+        # e,h in Fourier space
+        fehl = self.Solve_FieldFourier(which_layer, z_offset)
+
+        eh = []
+        for feh in fehl:
+            fe = feh[0]
+            fh = feh[1]
+            ex = get_ifft(Nx, Ny, fe[0], self.G)
+            ey = get_ifft(Nx, Ny, fe[1], self.G)
+            ez = get_ifft(Nx, Ny, fe[2], self.G)
+
+            hx = get_ifft(Nx, Ny, fh[0], self.G)
+            hy = get_ifft(Nx, Ny, fh[1], self.G)
+            hz = get_ifft(Nx, Ny, fh[2], self.G)
+            eh.append([[ex, ey, ez], [hx, hy, hz]])
+        if torch.isinstance(z_offset, float) or torch.isinstance(z_offset, int):
+            eh = eh[0]
+        return eh
+
+    def Volume_integral(self, which_layer, Mx, My, Mz, normalize=0):
+        """Mxyz is convolution matrix.
+        This function computes 1/A\int_V Mx|Ex|^2+My|Ey|^2+Mz|Ez|^2
+        To be consistent with Poynting vector defintion here, the absorbed power will be just omega*output
+        """
+        kp = self.kp_list[which_layer]
+        q = self.q_list[which_layer]
+        phi = self.phi_list[which_layer]
+
+        if self.id_list[which_layer][0] == 0:
+            epinv = 1.0 / self.Uniform_ep_list[self.id_list[which_layer][2]]
+        else:
+            epinv = self.Patterned_epinv_list[self.id_list[which_layer][2]]
+
+        # un-translated amplitdue
+        ai, bi = SolveInterior(
+            which_layer,
+            self.a0,
+            self.bN,
+            self.q_list,
+            self.phi_list,
+            self.kp_list,
+            self.thickness_list,
+        )
+        ab = torch.hstack((ai, bi))
+        abMatrix = torch.outer(ab, torch.conj(ab))
+
+        Mt = Matrix_zintegral(q, self.thickness_list[which_layer])
+        # overall
+        abM = abMatrix * Mt
+
+        # F matrix
+        Faxy = torch_dot(torch_dot(kp, phi), torch.diag(1.0 / self.omega / q))
+        Faz1 = 1.0 / self.omega * torch_dot(epinv, torch.diag(self.ky))
+        Faz2 = -1.0 / self.omega * torch_dot(epinv, torch.diag(self.kx))
+        Faz = torch_dot(torch.hstack((Faz1, Faz2)), phi)
+
+        tmp1 = torch.vstack((Faxy, Faz))
+        tmp2 = torch.vstack((-Faxy, Faz))
+        F = torch.hstack((tmp1, tmp2))
+
+        # consider Mtotal
+        Mzeros = torch.zeros_like(Mx)
+        Mtotal = torch.vstack(
+            (
+                torch.hstack((Mx, Mzeros, Mzeros)),
+                torch.hstack((Mzeros, My, Mzeros)),
+                torch.hstack((Mzeros, Mzeros, Mz)),
+            )
+        )
+
+        # integral = Tr[ abMatrix * F^\dagger *  Matconv *F ]
+        tmp = torch_dot(torch_dot(torch.conj(torch_transpose(F)), Mtotal), F)
+        val = torch.trace(torch_dot(abM, tmp))
+
+        if normalize == 1:
+            val = val * self.normalization
+        return val
+
+    def Solve_ZStressTensorIntegral(self, which_layer):
+        """
+        returns 2F_x,2F_y,2F_z, integrated over z-plane
+        """
+        z_offset = 0.0
+        eh = self.Solve_FieldFourier(which_layer, z_offset)
+        e = eh[0][0]
+        h = eh[0][1]
+        ex = e[0]
+        ey = e[1]
+        ez = e[2]
+
+        hx = h[0]
+        hy = h[1]
+        hz = h[2]
+
+        # compute D = epsilon E
+        ## Dz = epsilon_z E_z = (ky*hx - kx*hy)/self.omega
+        dz = (self.ky * hx - self.kx * hy) / self.omega
+
+        ## Dxy = epsilon2 * Exy
+        if self.id_list[which_layer][0] == 0:
+            dx = ex * self.Uniform_ep_list[self.id_list[which_layer][2]]
+            dy = ey * self.Uniform_ep_list[self.id_list[which_layer][2]]
+        else:
+            exy = torch.hstack((-ey, ex))
+            dxy = torch_dot(self.Patterned_ep2_list[self.id_list[which_layer][2]], exy)
+            dx = dxy[self.nG :]
+            dy = -dxy[: self.nG]
+
+        Tx = torch.sum(ex * torch.conj(dz) + hx * torch.conj(hz))
+        Ty = torch.sum(ey * torch.conj(dz) + hy * torch.conj(hz))
+        Tz = 0.5 * torch.sum(
+            ez * torch.conj(dz)
+            + hz * torch.conj(hz)
+            - ey * torch.conj(dy)
+            - ex * torch.conj(dx)
+            - torch.abs(hx) ** 2
+            - torch.abs(hy) ** 2
+        )
+
+        Tx = torch.real(Tx)
+        Ty = torch.real(Ty)
+        Tz = torch.real(Tz)
+
+        return Tx, Ty, Tz
+
+
+def MakeKPMatrix(omega, layer_type, epinv, kx, ky):
+    nG = len(kx)
+
+    # uniform layer, epinv has length 1
+    if layer_type == 0:
+        # JkkJT = np.block([[np.diag(ky*ky), np.diag(-ky*kx)],
+        #                 [np.diag(-kx*ky),np.diag(kx*kx)]])
+
+        Jk = torch.vstack((torch.diag(-ky), torch.diag(kx)))
+        JkkJT = torch_dot(Jk, torch_transpose(Jk))
+
+        kp = omega ** 2 * torch_eye(2 * nG) - epinv * JkkJT
+    # patterned layer
+    else:
+        Jk = torch.vstack((torch.diag(-ky), torch.diag(kx)))
+        tmp = torch_dot(Jk, epinv)
+        kp = omega ** 2 * torch_eye(2 * nG) - torch_dot(tmp, torch_transpose(Jk))
+
+    return kp
+
+
+def SolveLayerEigensystem_uniform(omega, kx, ky, epsilon):
+    nG = len(kx)
+    q = torch.sqrt(epsilon * omega ** 2 - kx ** 2 - ky ** 2)
+    # branch cut choice
+    q = torch.where(torch.imag(q) < 0.0, -q, q)
+
+    q = torch.cat((q, q))
+    phi = torch_eye(2 * nG)
+    return q, phi
+
+
+def SolveLayerEigensystem(omega, kx, ky, kp, ep2):
+    nG = len(kx)
+
+    k = torch.vstack((torch.diag(kx), torch.diag(ky)))
+    kkT = torch_dot(k, torch_transpose(k))
+    M = torch_dot(ep2, kp) - kkT
+
+    q, phi = torch.linalg.eig(M)
+
+    q = torch.sqrt(q)
+    # branch cut choice
+    q = torch.where(torch.imag(q) < 0.0, -q, q)
+    return q, phi
+
+
+def GetSMatrix(indi, indj, q_list, phi_list, kp_list, thickness_list):
+    """S_ij: size 4n*4n"""
+    # assert type(indi) == int, 'layer index i must be integar'
+    # assert type(indj) == int, 'layer index j must be integar'
+
+    nG2 = len(q_list[0])
+    S11 = torch_eye(nG2, dtype=complex)
+    S12 = torch.zeros_like(S11)
+    S21 = torch.zeros_like(S11)
+    S22 = torch_eye(nG2, dtype=complex)
+    if indi == indj:
+        return S11, S12, S21, S22
+    elif indi > indj:
+        raise Exception("indi must be < indj")
+
+    for l in range(indi, indj):
+        ## next layer
+        lp1 = l + 1
+        ## Q = inv(phi_l) * phi_lp1
+        Q = torch_dot(torch.inverse(phi_list[l]), phi_list[lp1])
+        ## P = ql*inv(kp_l*phi_l) * kp_lp1*phi_lp1*q_lp1^-1
+        P1 = torch_dot(
+            torch.diag(q_list[l]), torch.inverse(torch_dot(kp_list[l], phi_list[l]))
+        )
+        P2 = torch_dot(
+            torch_dot(kp_list[lp1], phi_list[lp1]), torch.diag(1.0 / q_list[lp1])
+        )
+        P = torch_dot(P1, P2)
+        # P1 = torch_dot(kp_list[l],phi_list[l])
+        # P2 = torch_dot(torch_dot(kp_list[lp1],phi_list[lp1]),   torch.diag(1./q_list[lp1]))
+        # P = np.linalg.solve(P1,P2)
+        # P = np.dot(np.diag(q_list[l]),P)
+
+        # T11=T22, T12=T21
+        T11 = 0.5 * (Q + P)
+        T12 = 0.5 * (Q - P)
+
+        # phase
+        d1 = torch.diag(torch.exp(1j * q_list[l] * thickness_list[l]))
+        d2 = torch.diag(torch.exp(1j * q_list[lp1] * thickness_list[lp1]))
+
+        # S11 = inv(T11-d1*S12o*T12)*d1*S11o
+        P1 = T11 - torch_dot(torch_dot(d1, S12), T12)
+        P1 = torch.inverse(P1)  # hold for further use
+        S11 = torch_dot(torch_dot(P1, d1), S11)
+
+        # S12 = P1*(d1*S12o*T11-T12)*d2
+        P2 = torch_dot(d1, torch_dot(S12, T11)) - T12
+        S12 = torch_dot(torch_dot(P1, P2), d2)
+
+        # S21 = S22o*T12*S11+S21o
+        S21 = S21 + torch_dot(S22, torch_dot(T12, S11))
+
+        # S22 = S22o*T12*S12+S22o*T11*d2
+        P2 = torch_dot(S22, torch_dot(T12, S12))
+        P1 = torch_dot(S22, torch_dot(T11, d2))
+        S22 = P1 + P2
+
+    return S11, S12, S21, S22
+
+
+def SolveExterior(a0, bN, q_list, phi_list, kp_list, thickness_list):
+    """
+    Given a0, bN, solve for b0, aN
+    """
+
+    Nlayer = len(thickness_list)  # total number of layers
+    S11, S12, S21, S22 = GetSMatrix(
+        0, Nlayer - 1, q_list, phi_list, kp_list, thickness_list
+    )
+
+    aN = torch_dot(S11, a0) + torch_dot(S12, bN)
+    b0 = torch_dot(S21, a0) + torch_dot(S22, bN)
+
+    return aN, b0
+
+
+def SolveInterior(which_layer, a0, bN, q_list, phi_list, kp_list, thickness_list):
+    """
+    Given a0, bN, solve for ai, bi
+    Layer numbering starts from 0
+    """
+    Nlayer = len(thickness_list)  # total number of layers
+    nG2 = len(q_list[0])
+
+    S11, S12, S21, S22 = GetSMatrix(
+        0, which_layer, q_list, phi_list, kp_list, thickness_list
+    )
+    pS11, pS12, pS21, pS22 = GetSMatrix(
+        which_layer, Nlayer - 1, q_list, phi_list, kp_list, thickness_list
+    )
+
+    # tmp = inv(1-S12*pS21)
+    tmp = torch.inv(torch_eye(nG2) - torch_dot(S12, pS21))
+    # ai = tmp * (S11 a0 + S12 pS22 bN)
+    ai = torch_dot(tmp, torch_dot(S11, a0) + torch_dot(S12, torch_dot(pS22, bN)))
+    # bi = pS21 ai + pS22 bN
+    bi = torch_dot(pS21, ai) + torch_dot(pS22, bN)
+
+    return ai, bi
+
+
+def TranslateAmplitudes(q, thickness, dz, ai, bi):
+    aim = ai * torch.exp(1j * q * dz)
+    bim = bi * torch.exp(1j * q * (thickness - dz))
+    return aim, bim
+
+
+def GetZPoyntingFlux(ai, bi, omega, kp, phi, q, byorder=0):
+    """
+    Returns 2S_z/A, following Victor's notation
+    Maybe because 2* makes S_z = 1 for H=1 in vacuum
+    """
+    n2 = len(ai)
+    n = int(n2 / 2)
+    # A = kp phi inv(omega*q)
+    A = torch_dot(torch_dot(kp, phi), torch.diag(1.0 / omega / q))
+
+    pa = torch_dot(phi, ai)
+    pb = torch_dot(phi, bi)
+    Aa = torch_dot(A, ai)
+    Ab = torch_dot(A, bi)
+
+    # diff = 0.5*(pb* Aa - Ab* pa)
+    diff = 0.5 * (torch.conj(pb) * Aa - torch.conj(Ab) * pa)
+    # forward = real(Aa* pa) + diff
+    forward_xy = torch.real(torch.conj(Aa) * pa) + diff
+    backward_xy = -torch.real(torch.conj(Ab) * pb) + torch.conj(diff)
+
+    forward = forward_xy[:n] + forward_xy[n:]
+    backward = backward_xy[:n] + backward_xy[n:]
+    if byorder == 0:
+        forward = torch.sum(forward)
+        backward = torch.sum(backward)
+
+    return forward, backward
+
+
+def Matrix_zintegral(q, thickness, shift=1e-12):
+    """Generate matrix for z-integral"""
+    nG2 = len(q)
+    qi, qj = Gmeshgrid(q)
+
+    # # Maa = \int exp(i q_i z)^* exp(i q_j z)
+    # #     = [exp(i(q_j-q_i^*)t)-1]/i(q_j-q_i^*)
+    # # Mbb = \int exp(i q_i (t-z))^* exp(i q_j (t-z))
+    # #     = exp(i(q_j-q_i^*)t)  [exp(i(q_i^*-q_j)t)-1]/i(q_i^*-q_j)
+    # #     = Maa
+    # # Mab = \int exp(i q_i z)^* exp(i q_j (t-z))
+    # #     = exp(iq_j t) [1-exp(-i(q_i^*+q_j)t)]/i(q_i^*+q_j)
+    # #     = [exp(iq_j t)-exp(-i q_i^* t)]/i(q_i^*+q_j)
+    # # Mba = \int exp(i q_i (t-z))^* exp(i q_j z)
+    # #     = exp(-iq_i^* t) [exp(i(q_j+q_i^*)t)-1]/i(q_j+q_i^*)
+    # #     = Mab
+    # Maa = (np.exp(1j*(qj-np.conj(qi))*thickness)-1)/1j/(qj-np.conj(qi))
+    # Mab = (np.exp(1j*qj*thickness)-np.exp(-1j*np.conj(qi)*thickness))/1j/(qj+np.conj(qi))
+
+    # M = t exp(0.5it (qj-qi^*)) sinc(0.5d (sjqj-siqi^*), note in python sinc = sin(pi x)/pi x
+    qij = qj - torch.conj(qi) + torch_eye(nG2) * shift
+    Maa = (torch.exp(1j * qij * thickness) - 1) / 1j / qij  # this is more robust
+    # Maa = thickness * np.exp(0.5j*thickness*qij) * np.sinc(0.5*thickness*qij/np.pi)
+
+    qij2 = qj + torch.conj(qi)
+    # Mab = thickness * np.exp(0.5j*thickness*qij) * np.sinc(0.5*thickness*qij2/np.pi)
+    Mab = (
+        (torch.exp(1j * qj * thickness) - torch.exp(-1j * torch.conj(qi) * thickness))
+        / 1j
+        / qij2
+    )
+
+    tmp1 = torch.vstack((Maa, Mab))
+    tmp2 = torch.vstack((Mab, Maa))
+    Mt = torch.hstack((tmp1, tmp2))
+    return Mt
+
+
+def Gmeshgrid(x):
+    N = len(x)
+    qj = []
+    for i in range(N):
+        qj.append(x)
+    qj = torch.tensor(qj)
+    qi = torch_transpose(qj)
+    return qi, qj
diff --git a/nidn/trcwa/utils/__init__.py b/nidn/trcwa/utils/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/nidn/trcwa/utils/fft_funs.py b/nidn/trcwa/utils/fft_funs.py
new file mode 100644
index 0000000..df91d0b
--- /dev/null
+++ b/nidn/trcwa/utils/fft_funs.py
@@ -0,0 +1,87 @@
+import torch
+
+
+def Epsilon_fft(dN, eps_grid, G):
+    """dN = 1/Nx/Ny
+    For now, assume epsilon is diagonal; if epsilon has xz,yz component, just simply add them to off-diagonal eps2
+
+    eps_grid is  (1) for isotropic, a numpy 2d array in the format of (Nx,Ny),
+                 (2) for anisotropic, a list of numpy 2d array [(Nx,Ny),(Nx,Ny),(Nx,Ny)]
+    """
+
+    if len(eps_grid) == 3 and eps_grid[0].ndim == 2:
+        epsx_fft = get_conv(dN, eps_grid[0], G)
+        epsy_fft = get_conv(dN, eps_grid[1], G)
+        epsz_fft = get_conv(dN, eps_grid[2], G)
+        epsinv = torch.inverse(epsz_fft)
+
+        tmp1 = torch.vstack((epsx_fft, torch.zeros_like(epsx_fft)))
+        tmp2 = torch.vstack((torch.zeros_like(epsx_fft), epsy_fft))
+        eps2 = torch.hstack((tmp1, tmp2))
+
+    elif eps_grid[0].ndim == 1:
+        eps_fft = get_conv(dN, eps_grid, G)
+        epsinv = torch.inverse(eps_fft)
+
+        tmp1 = torch.vstack((eps_fft, torch.zeros_like(eps_fft)))
+        tmp2 = torch.vstack((torch.zeros_like(eps_fft), eps_fft))
+        eps2 = torch.hstack((tmp1, tmp2))
+    else:
+        raise ValueError("Wrong eps_grid type")
+
+    # Set very small components to zero for torch
+    for i in range(eps2.shape[0]):
+        for j in range(eps2.shape[1]):
+            if eps2[i][j].real < 1e-14 and eps2[i][j].imag < 1e-14:
+                eps2[i][j] = 0 + 0j
+
+    return epsinv, eps2
+
+
+def get_conv(dN, s_in, G):
+    """Attain convolution matrix
+    dN = 1/Nx/Ny
+    s_in: np.array of length Nx*Ny
+    G: shape (nG,2), 2 for Lk1,Lk2
+    s_out: 1/N sum a_m exp(-2pi i mk/n), shape (nGx*nGy)
+    """
+    nG, _ = G.shape
+    sfft = torch.fft.fft2(s_in) * dN
+
+    ix = range(nG)
+    ii, jj = torch.meshgrid(torch.tensor(ix), torch.tensor(ix))
+    s_out = sfft[G[ii, 0] - G[jj, 0], G[ii, 1] - G[jj, 1]]
+    return s_out
+
+
+def get_fft(dN, s_in, G):
+    """
+    FFT to get Fourier components
+
+    s_in: np.2d array of size (Nx,Ny)
+    G: shape (nG,2), 2 for Gx,Gy
+    s_out: 1/N sum a_m exp(-2pi i mk/n), shape (nGx*nGy)
+    """
+
+    sfft = torch.fft.fft2(s_in) * dN
+    return sfft[G[:, 0], G[:, 1]]
+
+
+def get_ifft(Nx, Ny, s_in, G):
+    """
+    Reconstruct real-space fields
+    """
+    dN = 1.0 / Nx / Ny
+    nG, _ = G.shape
+
+    s0 = torch.zeros((Nx, Ny), dtype=complex)
+    for i in range(nG):
+        x = G[i, 0]
+        y = G[i, 1]
+
+        stmp = torch.zeros((Nx, Ny), dtype=complex)
+        stmp[x, y] = 1.0
+        s0 = s0 + s_in[i] * stmp
+
+    s_out = torch.ifft2(s0) / dN
+    return s_out
diff --git a/nidn/trcwa/utils/kbloch.py b/nidn/trcwa/utils/kbloch.py
new file mode 100644
index 0000000..b044361
--- /dev/null
+++ b/nidn/trcwa/utils/kbloch.py
@@ -0,0 +1,144 @@
+import numpy as np
+import torch
+
+from ..constants import *
+
+
+def Lattice_Reciprocate(L1, L2):
+    """Given two lattice vectors L1,L2 in the form of (Lx,Ly), returns the
+    reciprocate vectors Lk/(2*pi)
+    """
+
+    assert (
+        type(L1) == list and type(L2) == list
+    ), "Lattice vectors should be in list format."
+    assert len(L1) == 2, "Both x,y components of Lattice vector L1 are required."
+    assert len(L2) == 2, "Both x,y components of Lattice vector L2 are required."
+
+    d = L1[0] * L2[1] - L1[1] * L2[0]
+
+    Lk1 = torch.tensor([L2[1] / d, -L2[0] / d])
+    Lk2 = torch.tensor([-L1[1] / d, L1[0] / d])
+
+    return Lk1, Lk2
+
+
+def Lattice_getG(nG, Lk1, Lk2, method=0):
+    """
+    The G is defined to produce the following reciprocal vector:
+    k = G[:,0] Lk1 + G[:,1] Lk2 (both k and Lk don't include the 2pi factor)
+
+    method:0 for circular truncation, 1 for parallelogramic truncation
+    """
+    assert type(nG) == int, "nG must be integar"
+
+    if method == 0:
+        G, nG = Gsel_circular(nG, Lk1, Lk2)
+    elif method == 1:
+        G, nG = Gsel_parallelogramic(nG, Lk1, Lk2)
+    else:
+        raise Exception("Truncation scheme is not included")
+
+    return G, nG
+
+
+def Lattice_SetKs(G, kx0, ky0, Lk1, Lk2):
+    """
+    Construct kx,ky including all relevant orders, given initial scalar kx,ky
+    2pi factor is now included in the returned kx,ky
+    """
+
+    kx = kx0 + 2 * TRCWA_PI * (Lk1[0] * G[:, 0] + Lk2[0] * G[:, 1])
+    ky = ky0 + 2 * TRCWA_PI * (Lk1[1] * G[:, 0] + Lk2[1] * G[:, 1])
+
+    return kx, ky
+
+
+def Gsel_parallelogramic(nG, Lk1, Lk2):
+    """From Liu's gsel.c"""
+    u = torch.norm(Lk1)
+    v = torch.norm(Lk2)
+    uv = torch.dot(Lk1, Lk2)
+
+    NGroot = int(torch.sqrt(nG))
+    if torch.mod(NGroot, 2) == 0:
+        NGroot -= 1
+
+    M = NGroot // 2
+
+    xG = range(-M, NGroot - M)
+    G1, G2 = torch.meshgrid(xG, xG, indexing="ij")
+    G1 = G1.flatten()
+    G2 = G2.flatten()
+
+    # sorting
+    Gl2 = G1 ** 2 * u ** 2 + G2 ** 2 * v ** 2 + 2 * G2 * G1 * uv
+    sort = torch.argsort(Gl2)
+    G1 = G1[sort]
+    G2 = G2[sort]
+
+    # final G
+    nG = NGroot * NGroot
+    G = torch.zeros((nG, 2), dtype=int)
+    G[:, 0] = G1[:nG]
+    G[:, 1] = G2[:nG]
+
+    return G, nG
+
+
+def Gsel_circular(nG, Lk1, Lk2):
+    """From Liu's gsel.c.
+    NG * |u x v| is approximately the area in k-space we will need
+    cover with a circular disc. (u and v are the 2 shortest lattice
+    vectors) From the area, we can find the radius (and round it
+    up). Then, we can find the minimum extends in each of the two
+    lattice directions.
+    """
+    u = torch.norm(Lk1)
+    v = torch.norm(Lk2)
+    uv = torch.dot(Lk1, Lk2)
+    uxv = Lk1[0] * Lk2[1] - Lk1[1] * Lk2[0]
+    circ_area = nG * torch.abs(uxv)
+    circ_radius = torch.sqrt(circ_area / TRCWA_PI) + u + v
+
+    u_extent = 1 + int(circ_radius / (u * torch.sqrt(1.0 - uv ** 2 / (u * v) ** 2)))
+    v_extent = 1 + int(circ_radius / (v * torch.sqrt(1.0 - uv ** 2 / (u * v) ** 2)))
+
+    uext21 = 2 * u_extent + 1
+    vext21 = 2 * v_extent + 1
+
+    xG = range(-u_extent, uext21 - u_extent)
+    yG = range(-v_extent, vext21 - v_extent)
+    G1, G2 = torch.meshgrid(torch.tensor(xG), torch.tensor(yG))
+    G1 = G1.flatten()
+    G2 = G2.flatten()
+
+    # sorting
+    Gl2 = G1 ** 2 * u ** 2 + G2 ** 2 * v ** 2 + 2 * G2 * G1 * uv
+
+    # Torch's sorting algorithm is not stable, so we sort with numpy for reproducibility
+    sort = np.argsort(Gl2.cpu().numpy())
+    # sort = torch.argsort(Gl2)
+
+    G1 = G1[sort]
+    G2 = G2[sort]
+    Gl2 = Gl2[sort]
+
+    nGtmp = uext21 * vext21
+
+    if nG < nGtmp:
+        nGtmp = nG
+
+    # removing the part outside the cycle
+    tol = 1e-10 * max(u ** 2, v ** 2)
+    for i in range(nGtmp - 1, -1, -1):
+        if torch.abs(Gl2[i] - Gl2[i - 1]) > tol:
+            break
+    nG = i
+
+    # final G
+    G = torch.zeros((nG, 2), dtype=int)
+    G[:, 0] = G1[:nG]
+    G[:, 1] = G2[:nG]
+
+    return G, nG
diff --git a/nidn/trcwa/utils/torch_functions.py b/nidn/trcwa/utils/torch_functions.py
new file mode 100644
index 0000000..cbd388d
--- /dev/null
+++ b/nidn/trcwa/utils/torch_functions.py
@@ -0,0 +1,79 @@
+"""Utility functions to make torch operations compatible with numpy syntax."""
+import torch
+
+
+def torch_zeros(shape, dtype):
+    """Implementation of torch.zeros that returns a complex tensor. Else falls back on torch.zeros.
+
+    Args:
+        shape (list): Shape vector for zeros tensor.
+        dtype (dtype): Requested datatype.
+
+    Returns:
+        torch.tensor: Initialized tensor with real and imaginary part set to zero
+    """
+    if dtype == complex:
+        real = torch.zeros(shape)
+        imag = torch.zeros(shape)
+        return torch.complex(real, imag)
+    else:
+        return torch.zeros(shape, dtype=dtype)
+
+
+def torch_eye(size, dtype=None):
+    """Implementation of torch.eye that returns a complex tensor. Else falls back on torch.eye.
+
+    Args:
+        shape (list): Shape vector for identity matrix.
+        dtype (dtype): Requested datatype.
+
+    Returns:
+        torch.tensor: Initialized tensor with real part of the diagonal set to one and imaginary part of the diagonal set to zero, and zeros elsewhere
+    """
+    if dtype == complex:
+        real = torch.eye(size)
+        imag = torch.zeros([size, size])
+        return torch.complex(real, imag)
+    else:
+        return torch.eye(size, dtype=dtype)
+
+
+def torch_transpose(n_dimensional_tensor):
+    """Custom transpose for torch as torch.transpose diverges from numpy, see https://github.com/pytorch/pytorch/issues/51280
+
+    Args:
+        n_dimensional_tensor (torch.tensor): Tensor to transpose.
+
+    Returns:
+        torch.tensor: Transposed tensor
+    """
+    # Construct reverse index (e.g. [3,2,1,0])
+    indices = [
+        len(n_dimensional_tensor.shape) - i - 1
+        for i in range(len(n_dimensional_tensor.shape))
+    ]
+    return n_dimensional_tensor.permute(indices)
+
+
+def torch_dot(a, b):
+    """Computes a dot product in torch that is compatible with numpy style syntax.
+
+    Args:
+        a (torch.tensor): Tensor a in the dot product of a and b.
+        b (torch.tensor): Tensor b in the dot product of a and b.
+
+    Returns:
+        torch.tensor: Dot product of a and b
+    """
+    if a.dtype != b.dtype:  # in case data types are different
+        if a.dtype == torch.cfloat or a.dtype == torch.cdouble:
+            b = b.type(a.dtype)
+        elif b.dtype == torch.cfloat or b.dtype == torch.cdouble:
+            a = a.type(b.dtype)
+    # to be compatible with numpy broadcasting
+    if len(a.shape) == 2 and len(b.shape) == 1:
+        return torch.mv(a, b)
+    if len(a.shape) > 1:
+        return torch.mm(a, b)
+    else:
+        return torch.dot(a, b)
diff --git a/nidn/utils/README.md b/nidn/utils/README.md
new file mode 100644
index 0000000..ad6187b
--- /dev/null
+++ b/nidn/utils/README.md
@@ -0,0 +1,1 @@
+This folder contains utility functionality which is exposed to the user. Other utility should be in the subfolder where it is needed.
\ No newline at end of file
diff --git a/nidn/utils/convert_units.py b/nidn/utils/convert_units.py
new file mode 100644
index 0000000..24da15b
--- /dev/null
+++ b/nidn/utils/convert_units.py
@@ -0,0 +1,51 @@
+from nidn.trcwa.constants import TRCWA_UNIT_MAGNITUDE, TRCWA_SPEED_OF_LIGHT
+
+
+def freq_to_wl(freq):
+    """Gives the wavelength in units of the lattice constants
+    from the unitless frequency and vice versa. In TRCWA the magnitude is set to micrometres (um).
+
+    Args:
+        freq (float): The unitless frequency, or the wavelength in units of TRCWA_UNIT_MAGNITUDE.
+
+    Returns:
+        float: The wavelength with units of TRCWA_UNIT_MAGNITUDE, or the unitless frequency.
+    """
+    return 1.0 / freq
+
+
+def phys_wl_to_wl(phys_wl):
+    """Gives the wavelength from the physical wavelength and vice versa.
+
+    Args:
+        phys_wl (float): The physical wavelength in metres or the physical frequency in Hz.
+
+    Returns:
+        float: The physical frequency or the physical wavelength.
+    """
+    return phys_wl / TRCWA_UNIT_MAGNITUDE
+
+
+def wl_to_phys_wl(wl):
+    """Gives the physical wavelength from the wavelength in units of TRCWA_UNIT_MAGNITUDE.
+
+    Args:
+        wl (float): The wavelength in units of TRCWA_UNIT_MAGNITUDE.
+
+    Returns:
+        float: The physical wavelength.
+    """
+    return wl * TRCWA_UNIT_MAGNITUDE
+
+
+def phys_freq_to_phys_wl(freq):
+    """Gives the physical frequency from the physical wavelength and vice versa.
+
+    Args:
+        input (float): The physical wavelength in metres or the physical frequency in Hz.
+
+    Returns:
+        float: The physical frequency or the physical wavelength.
+    """
+    c = TRCWA_SPEED_OF_LIGHT
+    return c / freq
diff --git a/nidn/utils/fix_random_seeds.py b/nidn/utils/fix_random_seeds.py
new file mode 100644
index 0000000..3517aea
--- /dev/null
+++ b/nidn/utils/fix_random_seeds.py
@@ -0,0 +1,12 @@
+import torch
+import numpy as np
+
+
+def fix_random_seeds(seed: int = 42):
+    """This function sets the random seeds in torch and numpy to enable reproducible behavior.
+
+    Args:
+        seed (int): Seed to use.
+    """
+    torch.manual_seed(seed)
+    np.random.seed(seed)
diff --git a/nidn/utils/load_default_cfg.py b/nidn/utils/load_default_cfg.py
new file mode 100644
index 0000000..cfed381
--- /dev/null
+++ b/nidn/utils/load_default_cfg.py
@@ -0,0 +1,12 @@
+import os
+import toml
+from dotmap import DotMap
+
+
+def load_default_cfg():
+    """Loads the default toml config file from the cfg folder."""
+    path = os.path.join(
+        os.path.dirname(__file__) + "/resources/", "default_config.toml"
+    )
+    with open(path) as cfg:
+        return DotMap(toml.load(cfg))
diff --git a/nidn/utils/print_cfg.py b/nidn/utils/print_cfg.py
new file mode 100644
index 0000000..dd05c9a
--- /dev/null
+++ b/nidn/utils/print_cfg.py
@@ -0,0 +1,24 @@
+from dotmap import DotMap
+import numpy as np
+
+
+def print_cfg(cfg: DotMap):
+    """Prints the config in a more readable way.
+
+    Args:
+        cfg (DotMap): Config to print.
+    """
+    # Print the config three values per line
+    idx = 0
+    for key, value in cfg.items():
+        if isinstance(value, list) or isinstance(value, np.ndarray):
+            print()
+            print(f"{key}: {value}")
+            idx = 0
+        else:
+            if idx % 3 == 2:
+                print(f"{key:<20}: {value:<20}")
+            else:
+                print(f"{key:<20}: {value:<20}", end="")
+        idx += 1
+    print()
diff --git a/nidn/utils/resources/default_config.toml b/nidn/utils/resources/default_config.toml
new file mode 100644
index 0000000..43b6dd3
--- /dev/null
+++ b/nidn/utils/resources/default_config.toml
@@ -0,0 +1,32 @@
+# Training Properties
+seed = 42
+model_type = "siren"
+iterations = 2000
+learning_rate = 6.5e-5
+
+# Loss
+L = 0.5
+absorption_loss = false
+
+# Model Parameters
+n_neurons = 100
+hidden_layers = 9
+encoding_dim = 4
+out_features = 2
+siren_omega = 1
+
+# Epsilon Properties
+eps_oversampling = 1
+real_min_eps = 1e-3
+real_max_eps = 20
+imag_min_eps = 0
+imag_max_eps = 5
+
+# Grid dimensions
+Nx = 5
+Ny = 5
+N_layers = 1
+
+# Target frequencies
+physical_wavelength_range = [1000e-9,3e-6]
+N_freq = 20
diff --git a/notebooks/Training.ipynb b/notebooks/Training.ipynb
new file mode 100644
index 0000000..45306d6
--- /dev/null
+++ b/notebooks/Training.ipynb
@@ -0,0 +1,139 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "dd0c102f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "### Imports\n",
+    "%load_ext autoreload\n",
+    "%autoreload 2\n",
+    "\n",
+    "# Append main folder\n",
+    "import sys\n",
+    "sys.path.append(\"../\")\n",
+    "\n",
+    "import nidn\n",
+    "\n",
+    "# Set precision (and potentially GPU)\n",
+    "import torch\n",
+    "torch.set_default_tensor_type(torch.DoubleTensor)\n",
+    "\n",
+    "# Set up some logging\n",
+    "from loguru import logger\n",
+    "logger.remove()\n",
+    "logger.add(sys.stderr, format=\"{level} {message}\", level=\"DEBUG\");"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c4e7ae8b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cfg = nidn.load_default_cfg()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "dd459f38",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cfg.Nx = 9\n",
+    "cfg.Ny = 9\n",
+    "cfg.N_layers = 2\n",
+    "cfg.iterations = 50\n",
+    "nidn.print_cfg(cfg)\n",
+    "physical_freqs, normalized_freqs = nidn.get_frequency_points(cfg)\n",
+    "print(\"Physical frequencies are:\",physical_freqs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "dc935c7a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "target_R_spectrum =  [1.0]*6+[0.0]*14\n",
+    "target_T_spectrum =  [0.0]*6+[0.0]*14"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8e92c124",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "model, results = nidn.run_training(cfg,target_R_spectrum,target_T_spectrum)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e36e6250",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "nidn.plot_spectra(model,cfg,target_R_spectrum,target_T_spectrum)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0e7363a7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "nidn.plot_model_grid_per_freq(model,cfg)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1a628937",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "nidn.plot_model_grid(model,cfg)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bae8066f",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..bd437d5
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,5 @@
+matplotlib>=3.3.3
+pandas>=1.3.1
+torch==1.9.0
+tqdm>=4.56.0
+loguru>=0.5.3
\ No newline at end of file


Test patch:
diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml
new file mode 100644
index 0000000..a58439f
--- /dev/null
+++ b/.github/workflows/run_tests.yml
@@ -0,0 +1,36 @@
+# This workflow will install Python dependencies, run tests and lint with a single version of Python
+# For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions
+
+name: Running tests
+
+on:
+  push:
+    branches:
+    - main
+  pull_request:
+    branches:
+    - main
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python 3.8
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.8
+    - name: Install dependencies
+      run: |
+        $CONDA/bin/conda env update --file environment.yml --name base
+    # - name: Lint with flake8
+    #   run: |
+    #     # stop the build if there are Python syntax errors or undefined names
+    #     flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
+    #     # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
+    #     flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+    - name: Test with pytest
+      run: |
+        conda install pytest
+        $CONDA/bin/pytest
diff --git a/nidn/tests/__init__.py b/nidn/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/nidn/tests/target_frequency_test.py b/nidn/tests/target_frequency_test.py
new file mode 100644
index 0000000..eeffcd7
--- /dev/null
+++ b/nidn/tests/target_frequency_test.py
@@ -0,0 +1,29 @@
+from loguru import logger
+
+from ..trcwa.compute_target_frequencies import compute_target_frequencies
+from ..trcwa.constants import TRCWA_UNIT_MAGNITUDE
+
+
+def test_conversion():
+    """Test whether conversion from real wavelengths into target frequencies
+    is correctly donewl_to_phys_wl
+    """
+    logger.trace("Running frequency conversion test...")
+    min_wl = 300e-9  # 300 nm
+    max_wl = 3e-6  # 3 µm
+
+    target_frequencies = compute_target_frequencies(min_wl, max_wl, 5)
+    logger.debug("Proposed frequencies are")
+    logger.debug(target_frequencies)
+
+    assert len(target_frequencies) == 5
+
+    prod_min_wl = 1.0 / target_frequencies[0] * TRCWA_UNIT_MAGNITUDE
+    prod_max_wl = 1.0 / target_frequencies[-1] * TRCWA_UNIT_MAGNITUDE
+
+    assert abs(prod_min_wl - min_wl) < 1e-8
+    assert abs(prod_max_wl - max_wl) < 1e-8
+
+
+if __name__ == "__main__":
+    test_conversion()
diff --git a/nidn/tests/trcwa_test.py b/nidn/tests/trcwa_test.py
new file mode 100644
index 0000000..c47d33b
--- /dev/null
+++ b/nidn/tests/trcwa_test.py
@@ -0,0 +1,148 @@
+import torch
+from loguru import logger
+from dotmap import DotMap
+import sys
+
+from ..trcwa.compute_spectrum import compute_spectrum
+
+_TRCWA_TEST_TOLERANCE = 1e-7
+torch.set_printoptions(precision=12)
+torch.set_default_tensor_type(torch.DoubleTensor)
+
+
+def test_single_layer():
+    """Tests a single patterned layer with a single frequency. TRCWA vs. GRCWA"""
+    logger.info("Running single layer test...")
+    run_cfg = DotMap()
+    run_cfg.Nx = 5
+    run_cfg.Ny = 5
+    run_cfg.N_layers = 1
+    run_cfg.target_frequencies = [0.5]
+    run_cfg.N_freq = 1
+
+    # Get eps shape
+    shape = [
+        run_cfg.Nx,
+        run_cfg.Ny,
+        run_cfg.N_layers,
+        run_cfg.N_freq,
+    ]
+
+    # Get a grid of 1 + 1j
+    eps_grid = (
+        torch.ones(
+            shape,
+            dtype=torch.cfloat,
+        )
+        * (1.0 + 1j)
+    )
+
+    logger.debug("Computing spectrum...")
+
+    R, T = compute_spectrum(eps_grid, run_cfg)
+
+    # Computed with GRCWA
+    R_error = torch.abs(R[0] - 0.04256584753869945)
+    T_error = torch.abs(T[0] - 0.06064642157274551)
+
+    logger.info(f"R_error = {R_error}")
+    logger.info(f"T_error = {T_error}")
+
+    assert R_error < _TRCWA_TEST_TOLERANCE
+    assert T_error < _TRCWA_TEST_TOLERANCE
+
+
+def test_uniform_layer():
+    """Tests a single uniform layer at two frequencies. TRCWA vs. GRCWA"""
+    logger.info("Running uniform layer test...")
+    run_cfg = DotMap()
+    run_cfg.Nx = 1
+    run_cfg.Ny = 1
+    run_cfg.N_layers = 1
+    run_cfg.target_frequencies = [0.05, 0.1]
+    run_cfg.N_freq = 2
+
+    # Get eps shape
+    shape = [
+        run_cfg.Nx,
+        run_cfg.Ny,
+        run_cfg.N_layers,
+        run_cfg.N_freq,
+    ]
+
+    # Get a grid of 3 + 1.5j
+    eps_grid = torch.ones(
+        shape,
+        dtype=torch.cfloat,
+    )
+
+    eps_grid = eps_grid * (3.0 + 1.5j)
+
+    logger.debug("Computing spectrum...")
+
+    R, T = compute_spectrum(eps_grid, run_cfg)
+
+    # Computed with GRCWA
+    R_error, T_error = [], []
+    R_error.append(torch.abs(R[0] - 0.08805902003703245))
+    R_error.append(torch.abs(R[1] - 0.18524218068012266))
+    T_error.append(torch.abs(T[0] - 0.6307193983893462))
+    T_error.append(torch.abs(T[1] - 0.4510135948077737))
+
+    logger.info(f"R_error = {R_error}")
+    logger.info(f"T_error = {T_error}")
+
+    assert all(torch.tensor(R_error) < _TRCWA_TEST_TOLERANCE)
+    assert all(torch.tensor(T_error) < _TRCWA_TEST_TOLERANCE)
+
+
+def test_three_layer():
+    """Tests a three stacked patterned layer with a single frequency. TRCWA vs. GRCWA"""
+    logger.info("Running three layer test...")
+    run_cfg = DotMap()
+    run_cfg.Nx = 9
+    run_cfg.Ny = 9
+    run_cfg.N_layers = 3
+    run_cfg.target_frequencies = [3.0]
+    run_cfg.N_freq = 2
+
+    # Get eps shape
+    shape = [
+        run_cfg.Nx,
+        run_cfg.Ny,
+        run_cfg.N_layers,
+        run_cfg.N_freq,
+    ]
+
+    # Get a somewhat interesting grid
+    eps_grid = (
+        torch.ones(
+            shape,
+            dtype=torch.cfloat,
+        )
+        * (-4.2 + 0.42j)
+    )
+
+    eps_grid[0:3, 0:3, 0, :] = eps_grid[0:3, 0:3, 0, :] * 1.0
+    eps_grid[0:3, 0:3, 1, :] = eps_grid[0:3, 0:3, 1, :] * 2.0
+    eps_grid[0:3, 0:3, 2, :] = eps_grid[0:3, 0:3, 2, :] * 3.0
+
+    logger.debug("Computing spectrum...")
+
+    R, T = compute_spectrum(eps_grid, run_cfg)
+
+    # Computed with GRCWA
+    R_error = torch.abs(R[0] - 0.9243750060320585)
+    T_error = torch.abs(T[0] - 0.0)
+
+    logger.info(f"R_error = {R_error}")
+    logger.info(f"T_error = {T_error}")
+
+    assert R_error < _TRCWA_TEST_TOLERANCE
+    assert T_error < _TRCWA_TEST_TOLERANCE
+
+
+if __name__ == "__main__":
+    test_single_layer()
+    test_uniform_layer()
+    test_three_layer()