Simulate transmission spectrum of extreme ultraviolet laser pulse through laser-dressed finite sample (Python Version)

Question

I am currently participating in research on transient absorption spectroscopy and four wave mixing. In the experimental design, an extreme ultraviolet (XUV) laser pulse and infrared (IR) laser pulse are sent through a finite gas sample with a certain time delay between them. My current task was to develop a function that, given data on the system and input pulses, provides the transmission spectrum of the XUV pulse electric field after exiting the sample.

I had originally completed said task in MATLAB, and my post on that version of the program can be found here: Simulate transmission spectrum of extreme ultraviolet laser pulse through laser-dressed finite sample. However, I figured it would be a useful exercise to translate the code into python, so as to get a better grasp of python and start the ball rolling on learning more. This python version is thus almost identical to the first MATLAB version in its structure.

The Procedure

The gist of the procedure is to take two matrices describing the sample susceptibility, combine them into a single matrix, diagonalize said matrix, normalize the left and right eigenvector matrices according to each other, apply an exponential function to the eigenvalue matrix, and finally compute the exiting spectrum given an input spectrum.

The Code

import numpy as np
import matplotlib.pyplot as plt
import scipy.linalg as linalg



def FiniteSampleTransmissionXUV(w, chi_0, chi_nd, E_in, a, N, tau):

# Function takes as input:
# > w - A vector that describes the frequency spectrum range
# > X_0 - A vector that describes the susceptibility 
#   without the IR pulse
# > X_nd - A matrix that describes the susceptibility with the
#   IR pulse
# > E_in - A vector that describes the frequency spectrum of the 
#   incoming XUV pulse electric field
# > a - A constant that describes the intensity of the IR pulse
# > N - A constant that describes the optical depth of the
#   sample
# > tau - A constant that describes the time delay between the 
#   IR pulse and the XUV pulse
#   Function provides output:
# > A vector that describes the frequency spectrum of the
#   XUV pulse electric field at the end of the sample


    #determines number of frequency steps from input frequency range
    n = len(w)    # constant

    # determines frequency step size
    delta_w = 1 / (n - 1)   # constant
    
    # create matrix sqrt(w_i*w_j)
    sqrt_w = np.outer( np.sqrt(w) , np.sqrt(w).T )     # nxn matrix

    # combine X_0 and X_nd into total suscptibility matrix
    chi_ij = (a * delta_w * sqrt_w * chi_nd) + \
            np.diag( np.diag(sqrt_w) * (chi_0) )      # nxn matrix
    
    
    # diagonalize X_ij where sum(U_R_i^2) = 1
    (Lambda, u_l, u_r) = linalg.eig(chi_ij, left = True, right = True)     # vector and nxn matrices
    
    # attain the function that scales U_L'*U_R
    sqrt_F = np.sqrt( np.sum( u_l.T @ u_r , axis = 0 ) )     # row vector
    
    # scale U_R and U_L so that U_L'*U_R = 1
    u_r_bar = u_r / sqrt_F    # nxn matrix
    u_l_bar = u_l / sqrt_F    # nxn matrix

    
    # apply exponential transform to eigenvalues      # diagonal nxn matrix
    exp_lambda = np.diag( np.exp(1j * (2*np.pi*N / (3e8)) * Lambda) )
    
    # create phase shift due to the time delay
    tau_phase = np.exp(1j * w * tau)  # vector
    
    # recombine transformed susceptibility matrix
    ulambdau = u_r_bar @ exp_lambda @ u_l_bar.T     # nxn matrix
    
    # apply effects of finite propagation and pulse
    # time delay to input electric field spectrum
    E_out = ulambdau @ (tau_phase * E_in).T   # vector

    return E_out



def GaussianSpectrum(w, E0, w_0, sigma):
# creates a quasi-gaussian spectrum defined by the input
# frequency range w, amplitude E_0, center frequency w_0,
# and spectral width sigma
    E_in = E0 * np.sin(np.pi*w) * \
           (np.exp( - ((w - w_0) / (np.sqrt(2) * sigma))**2))

    return E_in




# Testing the FiniteSampleTransmissionXUV function

n = 100
w = np.linspace(0, 1, n)


chi_0 = (1 + 1j) * ( np.sin(5*np.pi*w) )**2


chi_sin = np.sin(1*np.pi*w)
chi_nd = (1 + 1j) * np.outer( (chi_sin**2), (chi_sin**2).T )


E_in = GaussianSpectrum(w, 1, 0.5, 0.1)


a = 1
N = 1e8
tau = 0


E_out = FiniteSampleTransmissionXUV(w, chi_0, chi_nd, E_in, a, N, tau)


plt.plot(w, E_in, w, abs(E_out))
plt.show()

In plotting the output electric field spectrum, we see the frequency spectrum of the XUV pulse before entering the sample in blue with the exiting spectrum in orange. One sees that frequencies of the XUV pulse have been absorbed, while new frequencies have been generated.

As I imagine this is a very MATLAB looking python program, please let me know how to make it more python-ey. Any comments on violations of best practices and ways to refactor for improved functionality and cleanliness would be greatly appreciated. Thank you!

Answer 1

Overall it's not bad.

• Where possible, move code into functions, among other reasons to increase testability and maintainability and keep the global namespace clean
• Prefer from x import y if you don't need to alias a submodule
• Prefer lower_snake_case for method and variable names. Exceptions apply, particularly here where E_x has scientific meaning.
• Avoid line continuation \; favour parentheses for multi-line expressions
• Can drop some redundant parens due to order of operations
• Consider adding a parametric transfer loop plot, since you have an input and an output. This is a Lissajous-like figure that depicts a different way of interpreting the relationship of your input and output spectra
• Consider plotting your introduced phase shift
• Add axis and figure titles
• Example code for Seaborn and Pandas added for the first subplot
• Particularly for function scalar value parameters, prefer named argument notation
• Add some PEP484 type hinting
• For u_r and u_l, in-place division is simpler and the bar variables are not needed
• Use standard triple-quoted docstring syntax
• Rather than 3e8 use the scipy-provided c

Example code

I'm pretty sure (?) that this is numerically equivalent, and looks identical to your results.

from typing import Tuple

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from matplotlib.collections import LineCollection
from scipy import constants, linalg


def finite_sample_transmission_xuv(
    w: np.ndarray,
    chi_0: np.ndarray,
    chi_nd: np.ndarray,
    E_in: np.ndarray,
    a: float,
    N: float,
    tau: float,
) -> np.ndarray:
    """
    Function takes as input:
    > w - A vector that describes the frequency spectrum range
    > X_0 - A vector that describes the susceptibility
      without the IR pulse
    > X_nd - A matrix that describes the susceptibility with the
      IR pulse
    > E_in - A vector that describes the frequency spectrum of the
      incoming XUV pulse electric field
    > a - A constant that describes the intensity of the IR pulse
    > N - A constant that describes the optical depth of the
      sample
    > tau - A constant that describes the time delay between the
      IR pulse and the XUV pulse
      Function provides output:
    > A vector that describes the frequency spectrum of the
      XUV pulse electric field at the end of the sample
    """

    # determines number of frequency steps from input frequency range
    n = len(w)  # constant

    # determines frequency step size
    delta_w = 1 / (n - 1)  # constant

    # create matrix sqrt(w_i*w_j)
    sqrt_w = np.outer(np.sqrt(w), np.sqrt(w).T)  # nxn matrix

    # combine X_0 and X_nd into total susceptibility matrix
    chi_ij = (  # nxn matrix
         a * delta_w * sqrt_w * chi_nd
         + np.diag(np.diag(sqrt_w) * chi_0)
    )

    # diagonalize X_ij where sum(U_R_i^2) = 1
    Lambda, u_l, u_r = linalg.eig(chi_ij, left=True, right=True)  # vector and nxn matrices

    # attain the function that scales U_L'*U_R
    sqrt_F = np.sqrt(np.sum(u_l.T @ u_r, axis=0))  # row vector

    # scale U_R and U_L so that U_L'*U_R = 1
    u_r /= sqrt_F  # nxn matrix
    u_l /= sqrt_F  # nxn matrix

    # apply exponential transform to eigenvalues
    exp_lambda = np.diag(  # diagonal nxn matrix
        np.exp(2j * np.pi * N / constants.c * Lambda)
    )

    # create phase shift due to the time delay
    tau_phase = np.exp(1j * w * tau)  # vector

    # recombine transformed susceptibility matrix
    u_lambda_u = u_r @ exp_lambda @ u_l.T  # nxn matrix

    # apply effects of finite propagation and pulse
    # time delay to input electric field spectrum
    E_out = u_lambda_u @ (tau_phase * E_in).T  # vector

    return E_out


def gaussian_spectrum(
    w: np.ndarray,
    E0: float,
    w_0: float,
    sigma: float,
) -> np.ndarray:
    """
    creates a quasi-gaussian spectrum defined by the input
    frequency range w, amplitude E_0, center frequency w_0,
    and spectral width sigma
    """
    return (
        E0
        * np.sin(np.pi * w)
        * np.exp(
            -(
                (w - w_0) / np.sqrt(2) / sigma
            )**2
        )
    )


def generate_test_spectrum() -> Tuple[
    np.ndarray,  # w
    np.ndarray,  # E_in
    np.ndarray,  # E_out
]:
    """
    Testing the finite_sample_transmission_xuv function
    """

    w = np.linspace(start=0, stop=1, num=300)

    chi_0 = (1 + 1j) * np.sin(5 * np.pi * w)**2
    chi_sin = np.sin(1 * np.pi * w)
    chi_sin2 = chi_sin**2
    chi_nd = (1 + 1j) * np.outer(chi_sin2, chi_sin2.T)

    E_in = gaussian_spectrum(w=w, E0=1, w_0=0.5, sigma=0.1)
    E_out = finite_sample_transmission_xuv(
        w=w, chi_0=chi_0, chi_nd=chi_nd, E_in=E_in,
        a=1, N=1e8, tau=0,
    )
    return w, E_in, E_out


def test_plot() -> None:
    w, E_in, E_out = generate_test_spectrum()

    fig = plt.figure()
    fig.suptitle('Finite Sample Transmission XUV (Gaussian test)')
    grid = fig.add_gridspec(nrows=2, ncols=2)

    top = fig.add_subplot(grid[0, :])
    top.set_title('Input and output spectra')
    top.set_xlabel('ω normalized')
    top.set_ylabel('|E| normalized')
    top.grid()
    E_df = pd.DataFrame(index=w, data={'Entering': E_in, 'Exiting': np.abs(E_out)})
    sns.lineplot(ax=top, data=E_df, dashes=False)

    phi = np.fmod(np.angle(E_out) + 2*np.pi, 2*np.pi)
    left = fig.add_subplot(grid[1, 0])
    left.set_title('Output phase for real input')
    left.set_xlabel('ω normalized')
    left.set_ylabel('φ rad')
    left.set_yticks(np.linspace(0, 2*np.pi, 9))
    left.grid()
    left.plot(w, phi)

    # Adapted from https://stackoverflow.com/a/17241345/313768
    xy = E_df.to_numpy().reshape(-1, 1, 2)
    segments = np.hstack([xy[:-1], xy[1:]])
    # There's an argument to be made that a cyclic cmap like 'hsv' from
    # https://matplotlib.org/stable/tutorials/colors/colormaps.html#cyclic
    # should be used since this is a loop, but that's up to you
    lines = LineCollection(segments, cmap='viridis')
    lines.set_array(w)

    right = fig.add_subplot(grid[1, 1])
    right.set_title('Parametric transfer loop')
    right.set_xlabel('|E| in')
    right.set_ylabel('|E| out')
    right.add_collection(lines)
    right.grid()
    right.autoscale_view()
    plt.colorbar(lines, ax=right, label='ω normalized')

    plt.show()


if __name__ == '__main__':
    test_plot()

Yet another way of visualising your output spectrum is on the polar-complex plane:

    xy = np.vstack((np.angle(E_out), np.abs(E_out))).T.reshape(-1, 1, 2)
    segments = np.hstack([xy[:-1], xy[1:]])
    lines = LineCollection(segments, cmap='viridis')
    lines.set_array(w)

    fig, ax = plt.subplots(subplot_kw={'projection': 'polar'})
    ax.set_title('Polar output spectrum')
    ax.add_collection(lines)
    ax.set_rmax(0.6)
    ax.set_rticks(np.linspace(0.1, 0.6, 6))
    ax.set_rlabel_position(-90)
    ax.grid(True)
    plt.colorbar(lines, ax=ax, label='ω normalized')
    plt.show()

Jupyter

This category of script is the ideal use-case for a Jupyter notebook. Think of it as a hybrid of a Mathematica-like interpreter and a paper-writing environment. You'll want either jupyterlab or notebook, and ipympl, all from pip. This will allow for a mix of "cells" (basically paragraphs), either

• Markdown, supporting $$ $$ formula notation, or
• Python, supporting matplotlib figures even with interactive cursor pan-and-zoom

I have not written a full notebook for you, but have covered examples of all of the above. If you import this JSON into xuv.ipynb:

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bef2a1f5",
   "metadata": {},
   "source": [
    "Simulate transmission spectrum of extreme ultraviolet laser pulse through laser-dressed finite sample\n",
    "==="
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0d3cd8f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %matplotlib widget for interactive mode, or\n",
    "# %matplotlib inline for print mode\n",
    "%matplotlib widget\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "from matplotlib.collections import LineCollection\n",
    "from scipy import constants, linalg\n",
    "from typing import Tuple\n",
    "\n",
    "ASPECT = 1.4\n",
    "figsize_inches = (8, 8/ASPECT)\n",
    "matplotlib.rc('figure', figsize=figsize_inches, dpi=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d9262f3",
   "metadata": {},
   "source": [
    "Introduction\n",
    "---------\n",
    "\n",
    "Based on [this CodeReview question](https://codereview.stackexchange.com/questions/261494/simulate-transmission-spectrum-of-extreme-ultraviolet-laser-pulse-through-laser/261517).\n",
    "\n",
    "I am currently participating in research on transient absorption spectroscopy and four wave mixing. In the experimental design, an extreme ultraviolet (XUV) laser pulse and infrared (IR) laser pulse are sent through a finite gas sample with a certain time delay between them. My current task was to develop a function that, given data on the system and input pulses, provides the transmission spectrum of the XUV pulse electric field after exiting the sample.\n",
    "\n",
    "The gist of the procedure is to take two matrices describing the sample susceptibility, combine them into a single matrix, diagonalize said matrix, normalize the left and right eigenvector matrices according to each other, apply an exponential function to the eigenvalue matrix, and finally compute the exiting spectrum given an input spectrum."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89c73f62",
   "metadata": {},
   "source": [
    "Given an input frequency range\n",
    "\n",
    "$$\n",
    "0 \\le \\omega_i \\le 1, 0 \\le i \\lt n\n",
    "$$\n",
    "\n",
    "blah blah... prepared by"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3e887232",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gaussian_spectrum(\n",
    "    w: np.ndarray,\n",
    "    E0: float,\n",
    "    w_0: float,\n",
    "    sigma: float,\n",
    ") -> np.ndarray:\n",
    "    \"\"\"\n",
    "    creates a quasi-gaussian spectrum defined by the input\n",
    "    frequency range w, amplitude E_0, center frequency w_0,\n",
    "    and spectral width sigma\n",
    "    \"\"\"\n",
    "    return (\n",
    "        E0\n",
    "        * np.sin(np.pi * w)\n",
    "        * np.exp(\n",
    "            -(\n",
    "                (w - w_0) / np.sqrt(2) / sigma\n",
    "            )**2\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e96b5104",
   "metadata": {},
   "source": [
    "The transmission transform is defined by\n",
    "\n",
    "... lots of formulas ...\n",
    "\n",
    "In Python:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9caf8046",
   "metadata": {},
   "outputs": [],
   "source": [
    "def finite_sample_transmission_xuv(\n",
    "    w: np.ndarray,\n",
    "    chi_0: np.ndarray,\n",
    "    chi_nd: np.ndarray,\n",
    "    E_in: np.ndarray,\n",
    "    a: float,\n",
    "    N: float,\n",
    "    tau: float,\n",
    ") -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Function takes as input:\n",
    "    > w - A vector that describes the frequency spectrum range\n",
    "    > X_0 - A vector that describes the susceptibility\n",
    "      without the IR pulse\n",
    "    > X_nd - A matrix that describes the susceptibility with the\n",
    "      IR pulse\n",
    "    > E_in - A vector that describes the frequency spectrum of the\n",
    "      incoming XUV pulse electric field\n",
    "    > a - A constant that describes the intensity of the IR pulse\n",
    "    > N - A constant that describes the optical depth of the\n",
    "      sample\n",
    "    > tau - A constant that describes the time delay between the\n",
    "      IR pulse and the XUV pulse\n",
    "      Function provides output:\n",
    "    > A vector that describes the frequency spectrum of the\n",
    "      XUV pulse electric field at the end of the sample\n",
    "    \"\"\"\n",
    "\n",
    "    # determines number of frequency steps from input frequency range\n",
    "    n = len(w)  # constant\n",
    "\n",
    "    # determines frequency step size\n",
    "    delta_w = 1 / (n - 1)  # constant\n",
    "\n",
    "    # create matrix sqrt(w_i*w_j)\n",
    "    sqrt_w = np.outer(np.sqrt(w), np.sqrt(w).T)  # nxn matrix\n",
    "\n",
    "    # combine X_0 and X_nd into total susceptibility matrix\n",
    "    chi_ij = (  # nxn matrix\n",
    "         a * delta_w * sqrt_w * chi_nd\n",
    "         + np.diag(np.diag(sqrt_w) * chi_0)\n",
    "    )\n",
    "\n",
    "    # diagonalize X_ij where sum(U_R_i^2) = 1\n",
    "    Lambda, u_l, u_r = linalg.eig(chi_ij, left=True, right=True)  # vector and nxn matrices\n",
    "\n",
    "    # attain the function that scales U_L'*U_R\n",
    "    sqrt_F = np.sqrt(np.sum(u_l.T @ u_r, axis=0))  # row vector\n",
    "\n",
    "    # scale U_R and U_L so that U_L'*U_R = 1\n",
    "    u_r /= sqrt_F  # nxn matrix\n",
    "    u_l /= sqrt_F  # nxn matrix\n",
    "\n",
    "    # apply exponential transform to eigenvalues\n",
    "    exp_lambda = np.diag(  # diagonal nxn matrix\n",
    "        np.exp(2j * np.pi * N / constants.c * Lambda)\n",
    "    )\n",
    "\n",
    "    # create phase shift due to the time delay\n",
    "    tau_phase = np.exp(1j * w * tau)  # vector\n",
    "\n",
    "    # recombine transformed susceptibility matrix\n",
    "    u_lambda_u = u_r @ exp_lambda @ u_l.T  # nxn matrix\n",
    "\n",
    "    # apply effects of finite propagation and pulse\n",
    "    # time delay to input electric field spectrum\n",
    "    E_out = u_lambda_u @ (tau_phase * E_in).T  # vector\n",
    "\n",
    "    return E_out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8250f881",
   "metadata": {},
   "source": [
    "To test this, run the simulation with\n",
    "\n",
    "$$ n = 500 $$\n",
    "$$ E_0 = 1 $$\n",
    "$$ \\omega_0 = 0.5 $$\n",
    "$$ \\sigma = 0.1 $$\n",
    "$$ \\tau = 0 $$\n",
    "$$ N = 1 \\times 10^8 $$\n",
    "$$ a = 1 $$\n",
    "\n",
    "via"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a9802320",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_test_spectrum() -> Tuple[\n",
    "    np.ndarray,  # w\n",
    "    np.ndarray,  # E_in\n",
    "    np.ndarray,  # E_out\n",
    "]:\n",
    "    \"\"\"\n",
    "    Testing the finite_sample_transmission_xuv function\n",
    "    \"\"\"\n",
    "\n",
    "    w = np.linspace(start=0, stop=1, num=500)\n",
    "\n",
    "    chi_0 = (1 + 1j) * np.sin(5 * np.pi * w)**2\n",
    "    chi_sin = np.sin(1 * np.pi * w)\n",
    "    chi_sin2 = chi_sin**2\n",
    "    chi_nd = (1 + 1j) * np.outer(chi_sin2, chi_sin2.T)\n",
    "\n",
    "    E_in = gaussian_spectrum(w=w, E0=1, w_0=0.5, sigma=0.1)\n",
    "    E_out = finite_sample_transmission_xuv(\n",
    "        w=w, chi_0=chi_0, chi_nd=chi_nd, E_in=E_in,\n",
    "        a=1, N=1e8, tau=0,\n",
    "    )\n",
    "    return w, E_in, E_out\n",
    "\n",
    "\n",
    "w, E_in, E_out = generate_test_spectrum()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a9d2291",
   "metadata": {},
   "source": [
    "Plots\n",
    "------\n",
    "\n",
    "First we look at the input and output spectral magnitudes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8b18f532",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "53acc1e145394bc4b0650a1492851d22",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def plot_spectra(w: np.ndarray, E_in: np.ndarray, E_out: np.ndarray) -> None:\n",
    "    fig, ax = plt.subplots(num=1)\n",
    "    ax.set_title('Input and output spectra')\n",
    "    ax.set_xlabel('ω normalized')\n",
    "    ax.set_ylabel('|E| normalized')\n",
    "    ax.grid()\n",
    "    E_df = pd.DataFrame(index=w, data={'Entering': E_in, 'Exiting': np.abs(E_out)})\n",
    "    sns.lineplot(ax=ax, data=E_df, dashes=False)\n",
    "\n",
    "    \n",
    "plot_spectra(w, E_in, E_out)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f6eb7f1",
   "metadata": {},
   "source": [
    "The input electric field spectrum is real. The output spectrum has an introduced phase shift:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "264e5b17",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e96a25d0be4b44d1ad6d287bb89334b2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def plot_phase(w: np.ndarray, E_in: np.ndarray, E_out: np.ndarray) -> None:\n",
    "    fig, ax = plt.subplots(num=2)\n",
    "    phi = np.fmod(np.angle(E_out) + 2*np.pi, 2*np.pi)\n",
    "    ax.set_title('Output phase for real input')\n",
    "    ax.set_xlabel('ω normalized')\n",
    "    ax.set_ylabel('φ rad')\n",
    "    ax.set_yticks(np.linspace(0, 2*np.pi, 9))\n",
    "    ax.grid()\n",
    "    ax.plot(w, phi)\n",
    "\n",
    "    \n",
    "plot_phase(w, E_in, E_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "245ca66b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}

and then open it in Jupyter, you'll see the results; some screenshots:

[...]

This supports a paper-like narrative with interactive figures, pretty formulae, and full source interspersal.

Answer 2

PEP-8 recommends using spaces around binary operators, such as * and **, and no spaces after a ( or before a ). This makes statements like:

chi_0 = (1 + 1j) * ( np.sin(5*np.pi*w) )**2

into something a little more readable:

chi_0 = (1 + 1j) * np.sin(5 * np.pi * w) ** 2

---

You ran into Python's keyword lambda, and had to revert to using Lambda as the variable name, to keep the variable name "clean", but end up deviating from PEP-8's preference of using snake_case for variable names.

Here is a controversial solution. Just go ahead and use lambda. It is a greek letter ... which is a Unicode "letter" ... which means it is valid as a Python identifier.

(λ, u_l, u_r) = ...

This can continue to many of your other Python identifiers, like Δω, τ, ϕ, χ. Certain subscripts are even possible (ₐₑₒₓₔₕₖₗₘₙₚₛₜ), allowing you to use τᵩ instead of tau_phase.

>>> τᵩ = 1
>>> f"{τᵩ=}"
'τᵩ=1'
>>> 

While it may make editing you program harder (having a collection of various unicode characters in a comment for easy copy/paste is helpful), using the proper mathematical symbols might actually make the program easier to understand.

from numpy import pi as π
from numpy import sin

...

    χ₀ = (1 + 1j) * sin(5 * π * ω) ** 2

Doesn't that look nice?