treams - npm Package Compare versions

Comparing version 0.3.0 to 0.3.1

.github/workflows/build.yml

+name: build
+
+on:
+  push:
+    tags:
+      - v*
+
+jobs:
+  build_wheels:
+    name: Build wheels on ${{ matrix.os }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest, macos-latest, windows-latest]
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+      - name: Setup msys2
+        if: matrix.os == 'windows-latest'
+        uses: msys2/setup-msys2@v2
+        with:
+          update: true
+          install: mingw-w64-x86_64-gcc
+      - name: Set path
+        if: matrix.os == 'windows-latest'
+        run: echo "$env:RUNNER_TEMP\msys64\mingw64\bin" | Out-File -FilePath $env:GITHUB_PATH -Encoding utf8 -Append
+      - name: Build wheels
+        uses: pypa/cibuildwheel@v2.13.1
+      - uses: actions/upload-artifact@v3
+        with:
+          path: ./wheelhouse/*.whl
+
+  build_sdist:
+    name: Build source distribution
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+      - name: Build sdist
+        run: pipx run build --sdist
+      - uses: actions/upload-artifact@v3
+        with:
+          path: dist/*.tar.gz

.github/workflows/docs.yml

+name: docs
+
+on:
+  push:
+    branches:
+      - main
+
+jobs:
+  docs:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v3
+      with:
+        fetch-depth: 0
+    - uses: actions/setup-python@v4
+      with:
+        python-version: '3.11'
+    - name: Build treams
+      run: python -m pip install -e .
+    - name: Build html
+      run: |
+        python -m pip install treams[docs]
+        sphinx-build -b html docs docs/_build/html
+    - name: Upload artifacts
+      uses: actions/upload-artifact@v3
+      with:
+        name: html-docs
+        path: docs/_build/html/
+    - name: Deploy docs
+      uses: peaceiris/actions-gh-pages@v3
+      if: github.ref == 'refs/heads/main'
+      with:
+        github_token: ${{ secrets.GITHUB_TOKEN }}
+        publish_dir: docs/_build/html

.github/workflows/doctests.yml

+name: doctests
+
+on:
+  push:
+    branches:
+      - main
+
+jobs:
+  doctests:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v3
+      with:
+        fetch-depth: 0
+    - uses: actions/setup-python@v4
+      with:
+        python-version: '3.11'
+    - name: Build treams
+      run: python -m pip install -e .
+    - name: Run doctests
+      run: |
+        python -m pip install treams[docs,io]
+        sphinx-build -b doctest docs docs/_build/doctest

.github/workflows/tests.yml

+name: tests
+
+on:
+  push:
+    branches:
+      - main
+
+jobs:
+  tests:
+    runs-on: ubuntu-latest
+    permissions:
+      pull-requests: write
+      contents: write
+    steps:
+    - uses: actions/checkout@v3
+      with:
+        fetch-depth: 0
+    - uses: actions/setup-python@v4
+      with:
+        python-version: '3.11'
+    - name: Build treams with tracing
+      run: CYTHON_COVERAGE=1 python -m pip install -e .
+    - name: Run tests
+      run: |
+        python -m pip install treams[test,coverage,io]
+        python -m pytest tests/unit --cov src/treams --cov-report html
+    - name: Get coverage report
+      run: |
+        python -m coverage json
+        export TOTAL=$(python -c "import json;print(json.load(open('coverage.json'))['totals']['percent_covered_display'])")
+        echo -n "{ \"schemaVersion\": 1, \"label\": \"coverage\", \"message\": \"$TOTAL%\", \"color\": \"" > htmlcov/endpoint.json
+        if [ "$TOTAL" -lt 60 ]; then echo -n "red" >> htmlcov/endpoint.json; \
+        elif [ "$TOTAL" -lt 70 ]; then echo -n "orange" >> htmlcov/endpoint.json; \
+        elif [ "$TOTAL" -lt 80 ]; then echo -n "yellow" >> htmlcov/endpoint.json; \
+        elif [ "$TOTAL" -lt 90 ]; then echo -n "yellowgreen" >> htmlcov/endpoint.json; \
+        elif [ "$TOTAL" -lt 95 ]; then echo -n "green" >> htmlcov/endpoint.json; \
+        else echo -n "brightgreen" >> htmlcov/endpoint.json; fi
+        echo "\" }" >> htmlcov/endpoint.json
+        rm htmlcov/.gitignore
+
+    - name: Deploy report
+      uses: s0/git-publish-subdir-action@develop
+      env:
+        REPO: self
+        BRANCH: htmlcov
+        FOLDER: htmlcov
+        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

docs/examples/array_spheres_tmatrixc.py

+import matplotlib.pyplot as plt
+import numpy as np
+
+import treams
+
+k0s = 2 * np.pi * np.linspace(1 / 600, 1 / 350, 100)
+material_slab = 3
+thickness = 10
+material_sphere = (4, 1, 0.05)
+pitch = 500
+lattice = treams.Lattice.square(pitch)
+radius = 100
+lmax = mmax = 3
+
+tr = np.zeros((len(k0s), 2))
+for i, k0 in enumerate(k0s):
+    kpar = [0, 0.3 * k0]
+
+    spheres = treams.TMatrix.sphere(
+        lmax, k0, radius, [material_sphere, 1]
+    ).latticeinteraction.solve(pitch, kpar[0])
+    cwb = treams.CylindricalWaveBasis.diffr_orders(kpar[0], mmax, pitch, 0.02)
+    spheres_cw = treams.TMatrixC.from_array(spheres, cwb)
+    chain_cw = spheres_cw.latticeinteraction.solve(pitch, kpar[1])
+
+    pwb = treams.PlaneWaveBasisByComp.diffr_orders(kpar, lattice, 0.02)
+    plw = treams.plane_wave(kpar, [1, 0, 0], k0=k0, basis=pwb, material=1)
+    slab = treams.SMatrices.slab(thickness, pwb, k0, [1, material_slab, 1])
+    dist = treams.SMatrices.propagation([0, 0, radius], pwb, k0, 1)
+    array = treams.SMatrices.from_array(chain_cw, pwb)
+    total = treams.SMatrices.stack([slab, dist, array])
+    tr[i, :] = total.tr(plw)
+
+fig, ax = plt.subplots()
+ax.set_xlabel("Frequency (THz)")
+ax.plot(299792.458 * k0s / (2 * np.pi), tr[:, 0])
+ax.plot(299792.458 * k0s / (2 * np.pi), tr[:, 1])
+ax.legend(["$T$", "$R$"])
+fig.show()

docs/examples/array_spheres.py

+import matplotlib.pyplot as plt
+import numpy as np
+
+import treams
+
+k0s = 2 * np.pi * np.linspace(1 / 600, 1 / 350, 100)
+material_slab = 3
+thickness = 10
+material_sphere = (4, 1, 0.05)
+lattice = treams.Lattice.square(500)
+radius = 100
+lmax = 3
+
+tr = np.zeros((len(k0s), 2))
+for i, k0 in enumerate(k0s):
+    kpar = [0, 0.3 * k0]
+    spheres = treams.TMatrix.sphere(
+        lmax, k0, radius, [material_sphere, 1]
+    ).latticeinteraction.solve(lattice, kpar)
+
+    pwb = treams.PlaneWaveBasisByComp.diffr_orders(kpar, lattice, 0.02)
+    plw = treams.plane_wave(kpar, [1, 0, 0], k0=k0, basis=pwb, material=1)
+    slab = treams.SMatrices.slab(thickness, pwb, k0, [1, material_slab, 1])
+    dist = treams.SMatrices.propagation([0, 0, radius], pwb, k0, 1)
+    array = treams.SMatrices.from_array(spheres, pwb)
+    total = treams.SMatrices.stack([slab, dist, array])
+    tr[i, :] = total.tr(plw)
+
+fig, ax = plt.subplots()
+ax.set_xlabel("Frequency (THz)")
+ax.plot(299792.458 * k0s / (2 * np.pi), tr[:, 0])
+ax.plot(299792.458 * k0s / (2 * np.pi), tr[:, 1])
+ax.legend(["$T$", "$R$"])
+fig.show()

docs/examples/band_structure.py

+import matplotlib.pyplot as plt
+import numpy as np
+
+import treams
+
+k0s = 2 * np.pi * np.linspace(1 / 10_000, 1 / 350, 200)
+material_slab = 3
+thickness = 10
+material_sphere = (4 + 0.1j, 1, 0.05)
+lattice = treams.Lattice.square(500)
+radius = 100
+lmax = 3
+az = 210
+
+res = []
+for i, k0 in enumerate(k0s):
+    kpar = [0, 0]
+    spheres = treams.TMatrix.sphere(
+        lmax, k0, radius, [material_sphere, 1]
+    ).latticeinteraction.solve(lattice, kpar)
+
+    pwb = treams.PlaneWaveBasisByComp.diffr_orders(kpar, lattice, 0.02)
+    plw = treams.plane_wave(kpar, [1, 0, 0], k0=k0, basis=pwb, material=1)
+    slab = treams.SMatrices.slab(thickness, pwb, k0, [1, material_slab, 1])
+    dist = treams.SMatrices.propagation([0, 0, radius], pwb, k0, 1)
+    array = treams.SMatrices.from_array(spheres, pwb)
+    total = treams.SMatrices.stack([slab, dist, array, dist])
+    x, _ = total.bands_kz(az)
+    x = x * az / np.pi
+    sel = x[np.abs(np.imag(x)) < 0.1]
+    res.append(sel)
+
+
+fig, ax = plt.subplots()
+for k0, sel in zip(k0s, res):
+    ax.scatter(sel.real, len(sel) * [299792.458 * k0 / (2 * np.pi)], 0.2, c="C0")
+    ax.scatter(sel.imag, len(sel) * [299792.458 * k0 / (2 * np.pi)], 0.2, c="C1")
+ax.set_xlabel("$k_z a_z / \\pi$")
+ax.set_ylabel("Frequency (THz)")
+ax.set_xlim([-1, 1])
+ax.set_ylim(ymin=0)
+ax.legend(["$Real$", "$Imag$"])
+fig.show()

docs/examples/chain_tmatrixc.py

+import matplotlib.pyplot as plt
+import numpy as np
+
+import treams
+
+k0 = 2 * np.pi / 700
+materials = [treams.Material(-16.5 + 1j), treams.Material()]
+lmax = mmax = 3
+radii = [75, 65]
+positions = [[-30, 0, -75], [30, 0, 75]]
+lattice = treams.Lattice(300)
+kz = 0.005
+
+spheres = [treams.TMatrix.sphere(lmax, k0, r, materials) for r in radii]
+chain = treams.TMatrix.cluster(spheres, positions).latticeinteraction.solve(lattice, kz)
+
+bmax = 3.1 * lattice.reciprocal
+cwb = treams.CylindricalWaveBasis.diffr_orders(kz, mmax, lattice, bmax, 2, positions)
+chain_tmc = treams.TMatrixC.from_array(chain, cwb)
+
+inc = treams.plane_wave(
+    [np.sqrt(k0 ** 2 - kz ** 2), 0, kz],
+    [np.sqrt(0.5), np.sqrt(0.5)],
+    k0=chain.k0,
+    material=chain.material,
+)
+sca = chain_tmc @ inc.expand(chain_tmc.basis)
+
+grid = (
+    np.mgrid[-300:300:61j, 0:1, -150:150:31j].squeeze().transpose((1, 2, 0))[:, :-1, :]
+)
+ez = np.zeros_like(grid[..., 0], complex)
+valid = chain_tmc.valid_points(grid, radii)
+ez[valid] = (inc.efield(grid[valid]) + sca.efield(grid[valid]))[..., 2]
+ez[~valid] = np.nan
+sca = chain @ inc.expand(chain.basis)
+valid = ~valid & chain.valid_points(grid, radii)
+vals = []
+for i, r in enumerate(grid[valid]):
+    swb = treams.SphericalWaveBasis.default(1, positions=[r])
+    field = sca.expandlattice(basis=swb).efield(r)
+    vals.append(inc.efield(r)[2] + field[2])
+ez[valid] = vals
+
+ez = np.concatenate(
+    (
+        np.real(ez.T * np.exp(-1j * kz * lattice[...])),
+        ez.T.real,
+        np.real(ez.T * np.exp(1j * kz * lattice[...])),
+    )
+)
+zaxis = np.concatenate((grid[0, :, 2] - 300, grid[0, :, 2], grid[0, :, 2] + 300,))
+
+fig, ax = plt.subplots(figsize=(10, 20))
+pcm = ax.pcolormesh(grid[:, 0, 0], zaxis, ez, shading="nearest", vmin=-1, vmax=1,)
+cb = plt.colorbar(pcm)
+cb.set_label("$E_z$")
+ax.set_xlabel("x (nm)")
+ax.set_ylabel("z (nm)")
+ax.set_aspect("equal")
+fig.show()

docs/examples/cluster_tmatrixc.py

+import matplotlib.pyplot as plt
+import numpy as np
+
+import treams
+
+k0 = 2 * np.pi / 700
+materials = [treams.Material(-16.5 + 1j), treams.Material()]
+lmax = mmax = 3
+radii = [65, 55]
+positions = [[-40, -50, 0], [40, 50, 0]]
+lattice = treams.Lattice(300)
+kz = 0.005
+
+sphere = treams.TMatrix.sphere(lmax, k0, radii[0], materials)
+chain = sphere.latticeinteraction.solve(lattice, kz)
+bmax = 3.1 * lattice.reciprocal
+cwb = treams.CylindricalWaveBasis.diffr_orders(kz, mmax, lattice, bmax)
+chain_tmc = treams.TMatrixC.from_array(chain, cwb)
+
+cylinder = treams.TMatrixC.cylinder(np.unique(cwb.kz), mmax, k0, radii[1], materials)
+
+cluster = treams.TMatrixC.cluster([chain_tmc, cylinder], positions).interaction.solve()
+inc = treams.plane_wave(
+    [np.sqrt(k0 ** 2 - kz ** 2), 0, kz],
+    [np.sqrt(0.5), np.sqrt(0.5)],
+    k0=chain.k0,
+    material=chain.material,
+)
+sca = cluster @ inc.expand(cluster.basis)
+
+grid = (
+    np.mgrid[-300:300:61j, 0:1, -150:150:31j].squeeze().transpose((1, 2, 0))[:, :-1, :]
+)
+ez = np.zeros_like(grid[..., 0], complex)
+valid = cluster.valid_points(grid, radii)
+ez[valid] = (inc.efield(grid[valid]) + sca.efield(grid[valid]))[..., 2]
+
+ez = np.concatenate(
+    (
+        np.real(ez.T * np.exp(-1j * kz * lattice[...])),
+        ez.T.real,
+        np.real(ez.T * np.exp(1j * kz * lattice[...])),
+    )
+)
+zaxis = np.concatenate((grid[0, :, 2] - 300, grid[0, :, 2], grid[0, :, 2] + 300,))
+
+fig, ax = plt.subplots(figsize=(10, 20))
+pcm = ax.pcolormesh(grid[:, 0, 0], zaxis, ez, shading="nearest", vmin=-1, vmax=1,)
+cb = plt.colorbar(pcm)
+cb.set_label("$E_z$")
+ax.set_xlabel("x (nm)")
+ax.set_ylabel("z (nm)")
+ax.set_aspect("equal")
+fig.show()

docs/examples/crystal_tmatrixc.py

+import matplotlib.pyplot as plt
+import numpy as np
+
+import treams
+
+k0s = 2 * np.pi * np.linspace(0.01, 1, 200)
+materials = [treams.Material(12.4), treams.Material()]
+mmax = 3
+radius = 0.2
+lattice = treams.Lattice.square(0.3)
+kpar = [0, 0, 0]
+
+res = []
+for k0 in k0s:
+    cyl = treams.TMatrixC.cylinder(kpar[2], mmax, k0, radius, materials)
+    svd = np.linalg.svd(cyl.latticeinteraction(lattice, kpar[:2]), compute_uv=False)
+    res.append(svd[-1])
+
+fig, ax = plt.subplots()
+ax.set_xlabel("Frequency (THz)")
+ax.set_ylabel("Smallest singular value")
+ax.semilogy(299.792458 * k0s / (2 * np.pi), res)
+fig.show()

docs/examples/cylinder_tmatrixc.py

+import matplotlib.pyplot as plt
+import numpy as np
+
+import treams
+
+k0s = 2 * np.pi * np.linspace(1 / 6000, 1 / 300, 200)
+materials = [treams.Material(16 + 0.5j), treams.Material()]
+mmax = 4
+radius = 75
+kzs = [0]
+cylinders = [treams.TMatrixC.cylinder(kzs, mmax, k0, radius, materials) for k0 in k0s]
+
+xw_sca = np.array([tm.xw_sca_avg for tm in cylinders]) / (2 * radius)
+xw_ext = np.array([tm.xw_ext_avg for tm in cylinders]) / (2 * radius)
+
+cwb_mmax0 = treams.CylindricalWaveBasis.default(kzs, 0)
+cylinders_mmax0 = [tm[cwb_mmax0] for tm in cylinders]
+xw_sca_mmax0 = np.array([tm.xw_sca_avg for tm in cylinders_mmax0]) / (2 * radius)
+xw_ext_mmax0 = np.array([tm.xw_ext_avg for tm in cylinders_mmax0]) / (2 * radius)
+
+tm = cylinders[-1]
+inc = treams.plane_wave([tm.k0, 0, 0], 1, k0=tm.k0, material=tm.material)
+sca = tm @ inc.expand(tm.basis)
+grid = np.mgrid[-100:100:101j, -100:100:101j, 0:1].squeeze().transpose((1, 2, 0))
+intensity = np.zeros_like(grid[..., 0])
+valid = tm.valid_points(grid, [radius])
+intensity[~valid] = np.nan
+intensity[valid] = 0.5 * np.sum(
+    np.abs(inc.efield(grid[valid]) + sca.efield(grid[valid])) ** 2, -1
+)
+
+fig, ax = plt.subplots()
+ax.plot(k0s * 299792.458 / (2 * np.pi), xw_ext)
+ax.plot(k0s * 299792.458 / (2 * np.pi), xw_sca)
+ax.plot(k0s * 299792.458 / (2 * np.pi), xw_ext_mmax0, color="C0", linestyle=":")
+ax.plot(k0s * 299792.458 / (2 * np.pi), xw_sca_mmax0, color="C1", linestyle=":")
+ax.set_xlabel("Frequency (THz)")
+ax.set_ylabel("Efficiency")
+ax.legend(["Extinction", "Scattering", "Extinction $m=1$", "Scattering $m=1$"])
+fig.show()
+
+fig, ax = plt.subplots()
+pcm = ax.pcolormesh(
+    grid[0, :, 1], grid[:, 0, 0], intensity.T, shading="nearest", vmin=0, vmax=1,
+)
+cb = plt.colorbar(pcm)
+cb.set_label("Intensity")
+ax.set_xlabel("x (nm)")
+ax.set_ylabel("y (nm)")
+fig.show()

docs/examples/grating_tmatrixc.py

+import matplotlib.pyplot as plt
+import numpy as np
+
+import treams
+
+k0 = 2 * np.pi / 700
+materials = [treams.Material(-16.5 + 1j), treams.Material()]
+lmax = mmax = 3
+radii = [65, 55]
+positions = [[-40, -50, 0], [40, 50, 0]]
+lattice_z = treams.Lattice(300)
+lattice_x = treams.Lattice(300, "x")
+kz = 0.005
+
+sphere = treams.TMatrix.sphere(lmax, k0, radii[0], materials)
+chain = sphere.latticeinteraction.solve(lattice_z, kz)
+bmax = 1.1 * lattice_z.reciprocal
+cwb = treams.CylindricalWaveBasis.diffr_orders(kz, mmax, lattice_z, bmax)
+kzs = np.unique(cwb.kz)
+chain_tmc = treams.TMatrixC.from_array(chain, cwb)
+cylinder = treams.TMatrixC.cylinder(kzs, mmax, k0, radii[1], materials)
+
+cluster = treams.TMatrixC.cluster(
+    [chain_tmc, cylinder], positions
+).latticeinteraction.solve(lattice_x, 0)
+inc = treams.plane_wave(
+    [0, np.sqrt(k0 ** 2 - kz ** 2), kz],
+    [np.sqrt(0.5), -np.sqrt(0.5)],
+    k0=chain.k0,
+    material=chain.material,
+)
+sca = cluster @ inc.expand(cluster.basis)
+
+grid = np.mgrid[0:1, -150:150:16j, -150:150:16j].squeeze().transpose((1, 2, 0))
+ex = np.zeros_like(grid[..., 0])
+valid = cluster.valid_points(grid, radii)
+vals = []
+for i, r in enumerate(grid[valid]):
+    cwb = treams.CylindricalWaveBasis.default(kzs, 1, positions=[r])
+    field = sca.expandlattice(basis=cwb).efield(r)
+    vals.append(np.real(inc.efield(r)[0] + field[0]))
+ex[valid] = vals
+ex[~valid] = np.nan
+
+
+fig, ax = plt.subplots(figsize=(10, 20))
+pcm = ax.pcolormesh(
+    grid[0, :, 2], grid[:, 0, 1], ex.T, shading="nearest", vmin=-1, vmax=1,
+)
+cb = plt.colorbar(pcm)
+cb.set_label("$E_x$")
+ax.set_xlabel("y (nm)")
+ax.set_ylabel("z (nm)")
+ax.set_aspect("equal")
+fig.show()

docs/examples/slab.py

+import matplotlib.pyplot as plt
+import numpy as np
+
+import treams
+
+k0s = 2 * np.pi * np.linspace(1 / 1000, 1 / 300, 50)
+materials = [(), (12.4 + 1j, 1 + 0.1j, 0.5 + 0.05j), (2, 2)]
+thickness = 50
+tr = np.zeros((2, len(k0s), 2))
+for i, k0 in enumerate(k0s):
+    pwb = treams.PlaneWaveBasisByComp.default([0, 0.5 * k0])
+    slab = treams.SMatrices.slab(thickness, pwb, k0, materials)
+    tr[0, i, :] = slab.tr([1, 0])
+    tr[1, i, :] = slab.tr([0, 1])
+
+fig, ax = plt.subplots()
+ax.set_xlabel("Frequency (THz)")
+ax.plot(299792.458 * k0s / (2 * np.pi), tr[0, :, 0])
+ax.plot(299792.458 * k0s / (2 * np.pi), tr[0, :, 1])
+ax.plot(299792.458 * k0s / (2 * np.pi), tr[1, :, 0], c="C0", linestyle="--")
+ax.plot(299792.458 * k0s / (2 * np.pi), tr[1, :, 1], c="C1", linestyle="--")
+ax.legend(["$T_+$", "$R_+$", "$T_-$", "$R_-$"])
+fig.show()

src/treams.egg-info/dependency_links.txt

+

src/treams.egg-info/PKG-INFO

+Metadata-Version: 2.1
+Name: treams
+Version: 0.3.1
+Summary: "T-matrix scattering code for nanophotonic computations"
+Home-page: https://git.scc.kit.edu/tfp-photonics/treams
+Author: Dominik Beutel
+Author-email: [e-mail]
+License: MIT
+Platform: Linux
+Platform: Windows
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Natural Language :: English
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Cython
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Astronomy
+Classifier: Topic :: Scientific/Engineering :: Atmospheric Science
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Provides-Extra: coverage
+Provides-Extra: docs
+Provides-Extra: io
+Provides-Extra: test
+License-File: LICENSE
+
+![Version](https://img.shields.io/github/v/tag/tfp-photonics/treams)
+[![PyPI](https://img.shields.io/pypi/v/treams)](https://pypi.org/project/treams)
+![License](https://img.shields.io/github/license/tfp-photonics/treams)
+![build](https://github.com/tfp-photonics/treams/actions/workflows/build.yml/badge.svg)
+[![docs](https://github.com/tfp-photonics/treams/actions/workflows/docs.yml/badge.svg)](https://tfp-photonics.github.io/treams)
+![doctests](https://github.com/tfp-photonics/treams/actions/workflows/doctests.yml/badge.svg)
+![tests](https://github.com/tfp-photonics/treams/actions/workflows/tests.yml/badge.svg)
+[![coverage](https://img.shields.io/endpoint?url=https%3A%2F%2Fraw.githubusercontent.com%2Ftfp-photonics%2Ftreams%2Fhtmlcov%2Fendpoint.json)](https://htmlpreview.github.io/?https://github.com/tfp-photonics/treams/blob/htmlcov/index.html)
+
+# treams
+
+The package `treams` provides a framework to simplify computations of the
+electromagnetic scattering of waves at finite and at periodic arrangements of particles
+based on the T-matrix method.
+
+## Installation
+
+### Installation using pip
+
+To install the package with pip, use
+
+```sh
+pip install treams
+```
+
+If you're using the system wide installed version of python, you might consider the
+``--user`` option.
+
+## Documentation
+
+The documentation can be found at https://tfp-photonics.github.io/treams.
+
+## Publications
+
+When using this code please cite:
+
+[D. Beutel, A. Groner, C. Rockstuhl, C. Rockstuhl, and I. Fernandez-Corbaton, Efficient Simulation of Biperiodic, Layered Structures Based on the T-Matrix Method, J. Opt. Soc. Am. B, JOSAB 38, 1782 (2021).](https://doi.org/10.1364/JOSAB.419645)
+
+Other relevant publications are
+* [D. Beutel, I. Fernandez-Corbaton, and C. Rockstuhl, Unified Lattice Sums Accommodating Multiple Sublattices for Solutions of the Helmholtz Equation in Two and Three Dimensions, Phys. Rev. A 107, 013508 (2023).](https://doi.org/10.1103/PhysRevA.107.013508)
+* [D. Beutel, P. Scott, M. Wegener, C. Rockstuhl, and I. Fernandez-Corbaton, Enhancing the Optical Rotation of Chiral Molecules Using Helicity Preserving All-Dielectric Metasurfaces, Appl. Phys. Lett. 118, 221108 (2021).](https://doi.org/10.1063/5.0050411)
+
+
+## Features
+
+* [x] T-matrix calculations using a spherical or cylindrical wave basis set
+* [x] Calculations in helicity and parity (TE/TM) basis
+* [x] Scattering from clusters of particles
+* [x] Scattering from particles and clusters arranged in 3d-, 2d-, and 1d-lattices
+* [x] Calculation of light propagation in stratified media
+* [x] Band calculation in crystal structures

src/treams.egg-info/requires.txt

+numpy
+scipy>=1.6
+
+[coverage]
+Cython
+pytest-cov
+
+[docs]
+matplotlib
+sphinx
+
+[io]
+h5py
+
+[test]
+pytest

src/treams.egg-info/SOURCES.txt

+.coveragerc
+LICENSE
+MANIFEST.in
+README.md
+conftest.py
+pyproject.toml
+setup.cfg
+setup.py
+.github/workflows/build.yml
+.github/workflows/docs.yml
+.github/workflows/doctests.yml
+.github/workflows/tests.yml
+docs/Makefile
+docs/about.rst
+docs/conf.py
+docs/dev.rst
+docs/gettingstarted.rst
+docs/index.rst
+docs/intro.rst
+docs/make.bat
+docs/maxwell.rst
+docs/operators.rst
+docs/params.rst
+docs/physicsarray.rst
+docs/smatrix.rst
+docs/theory.rst
+docs/tmatrix.rst
+docs/tmatrixc.rst
+docs/treams.rst
+docs/_static/custom.css
+docs/_templates/autosummary/attribute.rst
+docs/_templates/autosummary/bare-module.rst
+docs/_templates/autosummary/class.rst
+docs/_templates/autosummary/cython-module.rst
+docs/_templates/autosummary/module.rst
+docs/examples/array_spheres.py
+docs/examples/array_spheres_tmatrixc.py
+docs/examples/band_structure.py
+docs/examples/chain.py
+docs/examples/chain_tmatrixc.py
+docs/examples/cluster.py
+docs/examples/cluster_tmatrixc.py
+docs/examples/crystal.py
+docs/examples/crystal_tmatrixc.py
+docs/examples/cylinder_tmatrixc.py
+docs/examples/grating_tmatrixc.py
+docs/examples/grid.py
+docs/examples/slab.py
+docs/examples/sphere.py
+src/treams/__init__.py
+src/treams/_core.py
+src/treams/_lattice.py
+src/treams/_material.py
+src/treams/_operators.py
+src/treams/_smatrix.py
+src/treams/_tmatrix.py
+src/treams/coeffs.c
+src/treams/config.c
+src/treams/cw.c
+src/treams/ebcm.py
+src/treams/io.py
+src/treams/misc.py
+src/treams/pw.c
+src/treams/sw.c
+src/treams/util.py
+src/treams.egg-info/PKG-INFO
+src/treams.egg-info/SOURCES.txt
+src/treams.egg-info/dependency_links.txt
+src/treams.egg-info/requires.txt
+src/treams.egg-info/top_level.txt
+src/treams/lattice/__init__.py
+src/treams/lattice/_dsum.c
+src/treams/lattice/_esum.c
+src/treams/lattice/_gufuncs.c
+src/treams/lattice/_misc.c
+src/treams/lattice/cython_lattice.c
+src/treams/lattice/cython_lattice.pxd
+src/treams/special/__init__.py
+src/treams/special/_bessel.c
+src/treams/special/_coord.c
+src/treams/special/_gufuncs.c
+src/treams/special/_integrals.c
+src/treams/special/_misc.c
+src/treams/special/_ufuncs.c
+src/treams/special/_waves.c
+src/treams/special/_wigner3j.c
+src/treams/special/_wignerd.c
+src/treams/special/cython_special.c
+src/treams/special/cython_special.pxd
+tests/integration/test_compare_direct_sum.py
+tests/unit/test_coeffs.py
+tests/unit/test_core.py
+tests/unit/test_cw.py
+tests/unit/test_io.py
+tests/unit/test_lattice.py
+tests/unit/test_lattice_module.py
+tests/unit/test_misc.py
+tests/unit/test_operators.py
+tests/unit/test_pw.py
+tests/unit/test_smatrix.py
+tests/unit/test_special.py
+tests/unit/test_sw.py
+tests/unit/test_tmatrix.py
+tests/unit/test_tmatrixc.py
+tests/unit/test_util.py

src/treams.egg-info/top_level.txt

+treams

src/treams/__init__.py

+"""TREAMS: T-Matrix scattering code for nanophotonic computations.
+
+.. currentmodule:: treams
+
+Classes
+=======
+
+The top-level classes and functions allow a high-level access to the functionality.
+
+Basis sets
+----------
+
+.. autosummary::
+   :toctree: generated/
+
+   CylindricalWaveBasis
+   PlaneWaveBasisByUnitVector
+   PlaneWaveBasisByComp
+   SphericalWaveBasis
+
+Matrices and Arrays
+-------------------
+
+.. autosummary::
+   :toctree: generated/
+
+   PhysicsArray
+   SMatrix
+   TMatrix
+   TMatrixC
+
+Other
+-----
+
+.. autosummary::
+   :toctree: generated/
+
+   Lattice
+   Material
+
+Functions
+=========
+
+.. autosummary::
+   :toctree: generated/
+
+   bfield
+   changepoltype
+   dfield
+   efield
+   expand
+   expandlattice
+   hfield
+   permute
+   rotate
+   translate
+
+"""
+
+from treams._core import (  # noqa: F401
+    CylindricalWaveBasis,
+    PhysicsArray,
+    PlaneWaveBasisByComp,
+    PlaneWaveBasisByUnitVector,
+    SphericalWaveBasis,
+)
+from treams._lattice import Lattice, WaveVector  # noqa: F401
+from treams._material import Material  # noqa: F401
+from treams._operators import (  # noqa: F401
+    BField,
+    ChangePoltype,
+    DField,
+    EField,
+    Expand,
+    ExpandLattice,
+    FField,
+    GField,
+    HField,
+    Permute,
+    Rotate,
+    Translate,
+    bfield,
+    changepoltype,
+    dfield,
+    efield,
+    expand,
+    expandlattice,
+    ffield,
+    gfield,
+    hfield,
+    permute,
+    rotate,
+    translate,
+)
+from treams._smatrix import (  # noqa: F401
+    SMatrices,
+    SMatrix,
+    chirality_density,
+    poynting_avg_z,
+)
+from treams._tmatrix import (  # noqa: F401
+    TMatrix,
+    TMatrixC,
+    cylindrical_wave,
+    plane_wave,
+    plane_wave_angle,
+    spherical_wave,
+)

src/treams/_core.py

+"""Basis sets and core array functionalities."""
+
+import abc
+from collections import namedtuple
+
+import numpy as np
+
+import treams._operators as op
+import treams.lattice as la
+from treams import util
+from treams._lattice import Lattice, WaveVector
+from treams._material import Material
+
+
+class BasisSet(util.OrderedSet, metaclass=abc.ABCMeta):
+    """Basis set base class.
+
+    It is the base class for all basis sets used. They are expected to be an ordered
+    sequence of the modes, that are included in a expansion. Basis sets are expected to
+    be immutable.
+    """
+
+    _names = ()
+    """Names of the relevant parameters"""
+
+    def __repr__(self):
+        """String representation.
+
+        Automatically generated when the attribute ``_names`` is defined.
+
+        Returns:
+            str
+        """
+        string = ",\n    ".join(f"{name}={i}" for name, i in zip(self._names, self[()]))
+        return f"{self.__class__.__name__}(\n    {string},\n)"
+
+    def __len__(self):
+        """Number of modes."""
+        return len(self.pol)
+
+    @classmethod
+    @abc.abstractmethod
+    def default(cls, *args, **kwargs):
+        """Construct a default basis from parameters.
+
+        Construct a basis set in a default order by giving few parameters.
+        """
+        raise NotImplementedError
+
+
+class SphericalWaveBasis(BasisSet):
+    r"""Basis of spherical waves.
+
+    Functions of the spherical wave basis are defined by their angular momentum ``l``,
+    its projection onto the z-axis ``m``, and the polarization ``pol``. If the basis
+    is defined with respect to a single origin it is referred to as "global", if it
+    contains multiple origins it is referred to as "local". In a local basis an
+    additional position index ``pidx`` is used to link the modes to one of the
+    specified ``positions``.
+
+    For spherical waves there exist multiple :ref:`params:Polarizations` and they
+    can be separated into incident and scattered fields. Depending on these combinations
+    the basis modes refer to one of the functions :func:`~treams.special.vsw_A`,
+    :func:`~treams.special.vsw_rA`, :func:`~treams.special.vsw_M`,
+    :func:`~treams.special.vsw_rM`, :func:`~treams.special.vsw_N`, or
+    :func:`~treams.special.vsw_rN`.
+
+    Args:
+        modes (array-like): A tuple containing a list for each of ``l``, ``m``, and
+            ``pol`` or ``pidx``, ``l``, ``m``, and``pol``.
+        positions (array-like, optional): The positions of the origins for the specified
+            modes. Defaults to ``[[0, 0, 0]]``.
+
+    Attributes:
+        pidx (array-like): Integer referring to a row in :attr:`positions`.
+        l (array-like): Angular momentum as an integer :math:`l > 0`
+        m (array-like): Angular momentum projection onto the z-axis, it is an integer
+            with :math:`m \leq |l|`
+        pol (array-like): Polarization, see also :ref:`params:Polarizations`.
+    """
+
+    _names = ("pidx", "l", "m", "pol")
+
+    def __init__(self, modes, positions=None):
+        """Initalization."""
+        tmp = []
+        if isinstance(modes, np.ndarray):
+            modes = modes.tolist()
+        for m in modes:
+            if m not in tmp:
+                tmp.append(m)
+        modes = tmp
+        if len(modes) == 0:
+            pidx = []
+            l = []  # noqa: E741
+            m = []
+            pol = []
+        elif len(modes[0]) == 3:
+            l, m, pol = (*zip(*modes),)
+            pidx = np.zeros_like(l)
+        elif len(modes[0]) == 4:
+            pidx, l, m, pol = (*zip(*modes),)
+        else:
+            raise ValueError("invalid shape of modes")
+
+        if positions is None:
+            positions = np.zeros((1, 3))
+        positions = np.array(positions, float)
+        if positions.ndim == 1:
+            positions = positions[None, :]
+        if positions.ndim != 2 or positions.shape[1] != 3:
+            raise ValueError(f"invalid shape of positions {positions.shape}")
+
+        self.pidx, self.l, self.m, self.pol = [
+            np.array(i, int) for i in (pidx, l, m, pol)
+        ]
+        for i, j in ((self.pidx, pidx), (self.l, l), (self.m, m), (self.pol, pol)):
+            i.flags.writeable = False
+            if np.any(i != j):
+                raise ValueError("parameters must be integer")
+        if np.any(self.l < 1):
+            raise ValueError("'l' must be a strictly positive integer")
+        if np.any(self.l < np.abs(self.m)):
+            raise ValueError("'|m|' cannot be larger than 'l'")
+        if np.any(self.pol > 1) or np.any(self.pol < 0):
+            raise ValueError("polarization must be '0' or '1'")
+        if np.any(self.pidx >= len(positions)):
+            raise ValueError("undefined position is indexed")
+
+        self._positions = positions
+        self._positions.flags.writeable = False
+        self.lattice = self.kpar = None
+
+    @property
+    def positions(self):
+        """Positions of the modes' origins.
+
+        The positions are an immutable (N, 3)-array. Each row corresponds to a point in
+        the three-dimensional Cartesian space.
+        """
+        return self._positions
+
+    def __repr__(self):
+        """String representation."""
+        positions = "positions=" + str(self.positions).replace("\n", ",")
+        return f"{super().__repr__()[:-1]}    {positions},\n)"
+
+    @property
+    def isglobal(self):
+        """Basis is defined with respect to a single (global) origin.
+
+        Returns:
+            bool
+        """
+        return len(self) == 0 or np.all(self.pidx == self.pidx[0])
+
+    def __getattr__(self, key):
+        dct = {"l": "l", "m": "m", "p": "pidx", "s": "pol"}
+        try:
+            return tuple(getattr(self, dct[k]) for k in key)
+        except KeyError:
+            raise AttributeError(
+                f"'{type(self).__name__}' object has no attribute '{key}'"
+            ) from None
+
+    def __getitem__(self, idx):
+        """Get a subset of the basis.
+
+        This function allows index into the basis set by an integer, a slice, a sequence
+        of integers or bools, an ellipsis, or an empty tuple. All of them except the
+        integer and the empty tuple results in another basis set being returned. In case
+        of the two exceptions a tuple is returned.
+
+        Alternatively, the string "plmp", "lmp", or "lm" can be used to access only a
+        subset of :attr:`pidx`, :attr:`l`, :attr:`m`, and :attr:`pol`.
+        """
+        res = self.pidx[idx], self.l[idx], self.m[idx], self.pol[idx]
+        if isinstance(idx, (int, np.integer)) or (
+            isinstance(idx, tuple) and len(idx) == 0
+        ):
+            return res
+        return type(self)(zip(*res), self.positions)
+
+    def __eq__(self, other):
+        """Compare basis sets.
+
+        Basis sets are considered equal, when they have the same modes in the same order
+        and the specified origin :attr:`positions` are equal.
+        """
+        try:
+            return self is other or (
+                np.array_equal(self.pidx, other.pidx)
+                and np.array_equal(self.l, other.l)
+                and np.array_equal(self.m, other.m)
+                and np.array_equal(self.pol, other.pol)
+                and np.array_equal(self.positions, other.positions)
+            )
+        except AttributeError:
+            return False
+
+    @classmethod
+    def default(cls, lmax, nmax=1, positions=None):
+        """Default basis for the given maximal multipolar order.
+
+        The default order contains separate blocks for each position index which are in
+        ascending order. Within each block the modes are sorted by angular momentum
+        :math:`l`, with the lowest angular momentum coming first. For each angular
+        momentum its z-projection is in ascending order from :math:`m = -l` to
+        :math:`m = l`. Finally, the polarization index is the fastest changing index
+        which iterates between 1 and 0.
+
+        Example:
+            >>> SphericalWaveBasis.default(2)
+            SphericalWaveBasis(
+                pidx=[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0],
+                l=[1 1 1 1 1 1 2 2 2 2 2 2 2 2 2 2],
+                m=[-1 -1  0  0  1  1 -2 -2 -1 -1  0  0  1  1  2  2],
+                pol=[1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0],
+                positions=[[0. 0. 0.]],
+            )
+            >>> SphericalWaveBasis.default(1, 2, [[0, 0, 1.], [0, 0, -1.]])
+            SphericalWaveBasis(
+                pidx=[0 0 0 0 0 0 1 1 1 1 1 1],
+                l=[1 1 1 1 1 1 1 1 1 1 1 1],
+                m=[-1 -1  0  0  1  1 -1 -1  0  0  1  1],
+                pol=[1 0 1 0 1 0 1 0 1 0 1 0],
+                positions=[[ 0.  0.  1.], [ 0.  0. -1.]],
+            )
+
+        Args:
+            lmax (int): Maximal multipolar order.
+            nmax (int, optional): Number of positions, defaults to 1.
+            positions (array-like, optional): Positions of the origins.
+        """
+        modes = [
+            [n, l, m, p]
+            for n in range(0, nmax)
+            for l in range(1, lmax + 1)  # noqa: E741
+            for m in range(-l, l + 1)
+            for p in range(1, -1, -1)
+        ]
+        return cls(modes, positions=positions)
+
+    @classmethod
+    def ebcm(cls, lmax, nmax=1, mmax=-1, positions=None):
+        """Order of modes suited for ECBM.
+
+        In comparison to :meth:`default` this order prioritises blocks of the
+        z-projection of the angular momentum :math:`m` over the angular momentum
+        :math:`l`. This is useful to get block-diagonal matrices for the extended
+        boundary condition method (EBCM).
+
+        Example:
+            >>> SphericalWaveBasis.ebcm(2)
+            SphericalWaveBasis(
+                pidx=[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0],
+                l=[2 2 1 1 2 2 1 1 2 2 1 1 2 2 2 2],
+                m=[-2 -2 -1 -1 -1 -1  0  0  0  0  1  1  1  1  2  2],
+                pol=[1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0],
+                positions=[[0. 0. 0.]],
+            )
+
+        Args:
+            lmax (int): Maximal multipolar order.
+            nmax (int, optional): Number of particles, defaults to 1.
+            mmax (int, optional): Maximal value of `|m|`. If ommitted or set to -1 it
+                is taken equal to `lmax`.
+            positions (array-like, optional): Positions of the origins.
+
+        Returns:
+            tuple
+        """
+        mmax = lmax if mmax == -1 else mmax
+        modes = [
+            [n, l, m, p]
+            for n in range(0, nmax)
+            for m in range(-mmax, mmax + 1)
+            for l in range(max(abs(m), 1), lmax + 1)  # noqa: E741
+            for p in range(1, -1, -1)
+        ]
+        return cls(modes, positions=positions)
+
+    @staticmethod
+    def defaultlmax(dim, nmax=1):
+        """Calculate the default mode order for a given length.
+
+        Given the dimension of the T-matrix return the estimated maximal value of `l`.
+        This is the inverse of :meth:`defaultdim`. A value of zero is allowed for empty
+        T-matrices.
+
+        Example:
+            >>> SphericalWaveBasis.defaultlmax(len(SphericalWaveBasis.default(3)))
+            3
+
+        Args:
+            dim (int): Dimension of the T-matrix, respectively number of modes.
+            nmax (int, optional): Number of particles, defaults to 1.
+
+        Returns:
+            int
+        """
+        res = np.sqrt(1 + dim * 0.5 / nmax) - 1
+        res_int = int(np.rint(res))
+        if np.abs(res - res_int) > 1e-8 * np.maximum(np.abs(res), np.abs(res_int)):
+            raise ValueError("cannot estimate the default lmax")
+        return res_int
+
+    @staticmethod
+    def defaultdim(lmax, nmax=1):
+        """Default number of modes for a given mulipolar order.
+
+        Given the maximal value of `l` return the size of the corresponding T-matrix.
+        This is the inverse of :meth:`defaultlmax`. A value of zero is allowed.
+
+        Args:
+            lmax (int): Maximal multipolar order
+            nmax (int, optional): Number of particles, defaults to `1`
+
+        Returns:
+            int
+        """
+        # zero is allowed and won't give an error
+        if lmax < 0 or nmax < 0:
+            raise ValueError("maximal order must be positive")
+        return 2 * lmax * (lmax + 2) * nmax
+
+    @classmethod
+    def _from_iterable(cls, it, positions=None):
+        if isinstance(cls, SphericalWaveBasis):
+            positions = cls.positions if positions is None else positions
+            cls = type(cls)
+        obj = cls(it, positions=positions)
+        return obj
+
+
+class CylindricalWaveBasis(BasisSet):
+    r"""Basis of cylindrical waves.
+
+    Functions of the cylindrical wave basis are defined by the z-components of the wave
+    vector ``kz`` and the angular momentum  ``m`` as well as the polarization ``pol``.
+    If the basis is defined with respect to a single origin it is referred to as
+    "global", if it contains multiple origins it is referred to as "local". In a local
+    basis an additional position index ``pidx`` is used to link the modes to one of the
+    specified ``positions``.
+
+    For cylindrical waves there exist multiple :ref:`params:Polarizations` and
+    they can be separated into incident and scattered fields. Depending on these
+    combinations the basis modes refer to one of the functions
+    :func:`~treams.special.vcw_A`, :func:`~treams.special.vcw_rA`,
+    :func:`~treams.special.vcw_M`, :func:`~treams.special.vcw_rM`,
+    :func:`~treams.special.vcw_N`, or :func:`~treams.special.vcw_rN`.
+
+    Args:
+        modes (array-like): A tuple containing a list for each of ``kz``, ``m``, and
+            ``pol`` or ``pidx``, ``kz``, ``m``, and``pol``.
+        positions (array-like, optional): The positions of the origins for the specified
+            modes. Defaults to ``[[0, 0, 0]]``.
+
+    Attributes:
+        pidx (array-like): Integer referring to a row in :attr:`positions`.
+        kz (array-like): Real valued z-component of the wave vector.
+        m (array-like): Integer angular momentum projection onto the z-axis.
+        pol (array-like): Polarization, see also :ref:`params:Polarizations`.
+    """
+
+    _names = ("pidx", "kz", "m", "pol")
+
+    def __init__(self, modes, positions=None):
+        """Initalization."""
+        tmp = []
+        if isinstance(modes, np.ndarray):
+            modes = modes.tolist()
+        for m in modes:
+            if m not in tmp:
+                tmp.append(m)
+        modes = tmp
+        if len(modes) == 0:
+            pidx = []
+            kz = []
+            m = []
+            pol = []
+        elif len(modes[0]) == 3:
+            kz, m, pol = (*zip(*modes),)
+            pidx = np.zeros_like(m)
+        elif len(modes[0]) == 4:
+            pidx, kz, m, pol = (*zip(*modes),)
+        else:
+            raise ValueError("invalid shape of modes")
+
+        if positions is None:
+            positions = np.zeros((1, 3))
+        positions = np.array(positions, float)
+        if positions.ndim == 1:
+            positions = positions[None, :]
+        if positions.ndim != 2 or positions.shape[1] != 3:
+            raise ValueError(f"invalid shape of positions {positions.shape}")
+
+        self.pidx, self.m, self.pol = [np.array(i, int) for i in (pidx, m, pol)]
+        self.kz = np.array(kz, float)
+        self.kz.flags.writeable = False
+        for i, j in ((self.pidx, pidx), (self.m, m), (self.pol, pol)):
+            i.flags.writeable = False
+            if np.any(i != j):
+                raise ValueError("parameters must be integer")
+        if np.any(self.pol > 1) or np.any(self.pol < 0):
+            raise ValueError("polarization must be '0' or '1'")
+        if np.any(self.pidx >= len(positions)):
+            raise ValueError("undefined position is indexed")
+
+        self._positions = positions
+        self._positions.flags.writeable = False
+
+        self.lattice = self.kpar = None
+        if len(self.kz) > 0 and np.all(self.kz == self.kz[0]):
+            self.kpar = WaveVector(self.kz[0])
+
+    @property
+    def positions(self):
+        """Positions of the modes' origins.
+
+        The positions are an immutable (N, 3)-array. Each row corresponds to a point in
+        the three-dimensional Cartesian space.
+        """
+        return self._positions
+
+    def __repr__(self):
+        """String representation."""
+        positions = "positions=" + str(self.positions).replace("\n", ",")
+        return f"{super().__repr__()[:-1]}    {positions},\n)"
+
+    @property
+    def isglobal(self):
+        """Basis is defined with respect to a single (global) origin.
+
+        Returns:
+            bool
+        """
+        return len(self) == 0 or np.all(self.pidx == self.pidx[0])
+
+    def __getattr__(self, key):
+        dct = {"z": "kz", "m": "m", "p": "pidx", "s": "pol"}
+        try:
+            return tuple(getattr(self, dct[k]) for k in key)
+        except KeyError:
+            raise AttributeError(
+                f"'{type(self).__name__}' object has no attribute '{key}'"
+            ) from None
+
+    def __getitem__(self, idx):
+        """Get a subset of the basis.
+
+        This function allows index into the basis set by an integer, a slice, a sequence
+        of integers or bools, an ellipsis, or an empty tuple. All of them except the
+        integer and the empty tuple results in another basis set being returned. In case
+        of the two exceptions a tuple is returned.
+
+        Alternatively, the string "pkzmp", "kzmp", or "kzm" can be used to access only a
+        subset of :attr:`pidx`, :attr:`kz`, :attr:`m`, and :attr:`pol`.
+        """
+        res = self.pidx[idx], self.kz[idx], self.m[idx], self.pol[idx]
+        if isinstance(idx, (int, np.integer)) or (isinstance(idx, tuple) and idx == ()):
+            return res
+        return type(self)(zip(*res), self.positions)
+
+    def __eq__(self, other):
+        """Compare basis sets.
+
+        Basis sets are considered equal, when they have the same modes in the same order
+        and the specified origin :attr:`positions` are equal.
+        """
+        try:
+            return self is other or (
+                np.array_equal(self.pidx, other.pidx)
+                and np.array_equal(self.kz, other.kz)
+                and np.array_equal(self.m, other.m)
+                and np.array_equal(self.pol, other.pol)
+                and np.array_equal(self.positions, other.positions)
+            )
+        except AttributeError:
+            return False
+
+    @classmethod
+    def default(cls, kzs, mmax, nmax=1, positions=None):
+        """Default basis for the given z-components of wave vector and angular momentum.
+
+        The default order contains separate blocks for each position index which are in
+        ascending order. Within each block the modes are sorted by the z-component of
+        the wave vector :math:`k_z`. For each of those values the z-projection of the
+        angular momentum is placed in ascending order. Finally, the polarization index
+        is the fastest changing index which iterates between 1 and 0.
+
+        Example:
+            >>> CylindricalWaveBasis.default([-0.5, 0.5], 1)
+            CylindricalWaveBasis(
+                pidx=[0 0 0 0 0 0 0 0 0 0 0 0],
+                kz=[-0.5 -0.5 -0.5 -0.5 -0.5 -0.5  0.5  0.5  0.5  0.5  0.5  0.5],
+                m=[-1 -1  0  0  1  1 -1 -1  0  0  1  1],
+                pol=[1 0 1 0 1 0 1 0 1 0 1 0],
+                positions=[[0. 0. 0.]],
+            )
+            >>> CylindricalWaveBasis.default([0], 1, 2, [[1., 0, 0], [-1., 0, 0]])
+            CylindricalWaveBasis(
+                pidx=[0 0 0 0 0 0 1 1 1 1 1 1],
+                kz=[0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.],
+                m=[-1 -1  0  0  1  1 -1 -1  0  0  1  1],
+                pol=[1 0 1 0 1 0 1 0 1 0 1 0],
+                positions=[[ 1.  0.  0.], [-1.  0.  0.]],
+            )
+
+        Args:
+            kzs (array-like, float): Maximal multipolar order.
+            mmax (int): Maximal value of the angular momentum z-component.
+            nmax (int, optional): Number of positions, defaults to 1.
+            positions (array-like, optional): Positions of the origins.
+        """
+        kzs = np.atleast_1d(kzs)
+        if kzs.ndim > 1:
+            raise ValueError(f"kzs has dimension larger than one: '{kzs.ndim}'")
+        modes = [
+            [n, kz, m, p]
+            for n in range(nmax)
+            for kz in kzs
+            for m in range(-mmax, mmax + 1)
+            for p in range(1, -1, -1)
+        ]
+        return cls(modes, positions=positions)
+
+    @classmethod
+    def diffr_orders(cls, kz, mmax, lattice, bmax, nmax=1, positions=None):
+        """Create a basis set for a system periodic in the z-direction.
+
+        Example:
+           >>> CylindricalWaveBasis.diffr_orders(0.1, 1, 2 * np.pi, 1)
+           CylindricalWaveBasis(
+               pidx=[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0],
+               kz=[-0.9 -0.9 -0.9 -0.9 -0.9 -0.9  0.1  0.1  0.1  0.1  0.1  0.1  1.1  1.1
+             1.1  1.1  1.1  1.1],
+               m=[-1 -1  0  0  1  1 -1 -1  0  0  1  1 -1 -1  0  0  1  1],
+               pol=[1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0],
+               positions=[[0. 0. 0.]],
+           )
+
+        Args:
+            kz (float): Wave vector z-component. Ideally it is in the first Brillouin
+                zone (use :func:`misc.firstbrillouin1d`).
+            mmax (int): Maximal value for the z-component of the angular momentum.
+            lattice (:class:`treams.Lattice` or float): Lattice definition or pitch.
+            bmax (float): Maximal change of the z-component of the wave vector. So,
+                this defines a maximal momentum transfer from the given value `kz`.
+            nmax (int, optional): Number of positions.
+            positions (array-like, optional): Positions of the origins.
+        """
+        lattice = Lattice(lattice)
+        lattice_z = Lattice(lattice, "z")
+        nkz = np.floor(np.abs(bmax / lattice_z.reciprocal))
+        kzs = kz + np.arange(-nkz, nkz + 1) * lattice_z.reciprocal
+        res = cls.default(kzs, mmax, nmax, positions=positions)
+        res.lattice = lattice
+        res.kpar = WaveVector(kz)
+        return res
+
+    @classmethod
+    def _from_iterable(cls, it, positions=None):
+        if isinstance(cls, CylindricalWaveBasis):
+            positions = cls.positions if positions is None else positions
+            lattice = cls.lattice
+            kpar = cls.kpar
+            cls = type(cls)
+        else:
+            lattice = kpar = None
+        obj = cls(it, positions=positions)
+        obj.lattice = lattice
+        obj.kpar = kpar
+        return obj
+
+    @staticmethod
+    def defaultmmax(dim, nkz=1, nmax=1):
+        """Calculate the default mode order for a given length.
+
+        Given the dimension of the T-matrix return the estimated maximal value of `m`.
+        This is the inverse of :meth:`defaultdim`. A value of zero is allowed for empty
+        T-matrices.
+
+        Example:
+            >>> CylindricalWaveBasis.defaultmmax(len(CylindricalWaveBasis.default([0], 2)), 1)
+            2
+
+        Args:
+            dim (int): Dimension of the T-matrix, respectively number of modes
+            nkz (int, optional): Number of z-components of the wave vector.
+            nmax (int, optional): Number of particles, defaults to 1.
+
+        Returns:
+            int
+        """
+        if dim % (2 * nkz * nmax) != 0:
+            raise ValueError("cannot estimate the default mmax")
+        dim = dim // (2 * nkz * nmax)
+        if (dim - 1) % 2 != 0:
+            raise ValueError("cannot estimate the default mmax")
+        return (dim - 1) // 2
+
+    @staticmethod
+    def defaultdim(nkz, mmax, nmax=1):
+        """Default number of modes for a given mulipolar order.
+
+        Given the maximal value of `l` return the size of the corresponding T-matrix.
+        This is the inverse of :meth:`defaultlmax`. A value of zero is allowed.
+
+        Args:
+            nkz (int): Number of z-components of the wave vector.
+            mmax (int): Maximal value of the angular momentum's z-component.
+            nmax (int, optional): Number of particles, defaults to 1.
+
+        Returns:
+            int
+        """
+        if nkz < 0 or mmax < 0:
+            raise ValueError("maximal order must be positive")
+        return (2 * mmax + 1) * nkz * 2 * nmax
+
+
+class PlaneWaveBasis(BasisSet):
+    """Plane wave basis parent class."""
+
+    isglobal = True
+
+
+class PlaneWaveBasisByUnitVector(PlaneWaveBasis):
+    """Plane wave basis.
+
+    A plane wave basis is defined by a collection of wave vectors specified by the
+    Cartesian wave vector components ``qx``, ``qy``, and ``qz`` normalized to
+    :math:`q_x^2 + q_y^2 + q_z^2 = 1` and the polarizations ``pol``.
+
+    For plane waves there exist multiple :ref:`params:Polarizations`, such that
+    these modes can refer to either :func:`~treams.special.vpw_A` or
+    :func:`~treams.special.vpw_M` and :func:`~treams.special.vpw_N`.
+
+    Args:
+        modes (array-like): A tuple containing a list for each of ``qx``, ``qy``,
+            ``qz``, and ``pol``.
+
+    Attributes:
+        qx (array-like): X-component of the normalized wave vector.
+        qy (array-like): Y-component of the normalized wave vector.
+        qz (array-like): Z-component of the normalized wave vector.
+        pol (array-like): Polarization, see also :ref:`params:Polarizations`.
+    """
+
+    _names = ("qx", "qy", "qz", "pol")
+    """A plane wave basis is always global."""
+
+    def __init__(self, modes):
+        """Initialization."""
+        tmp = []
+        if isinstance(modes, np.ndarray):
+            modes = modes.tolist()
+        for m in modes:
+            if m not in tmp:
+                tmp.append(m)
+        modes = tmp
+        if len(modes) == 0:
+            qx = []
+            qy = []
+            qz = []
+            pol = []
+        elif len(modes[0]) == 4:
+            qx, qy, qz, pol = (*zip(*modes),)
+        else:
+            raise ValueError("invalid shape of modes")
+
+        qx, qy, qz, pol = (np.array(i) for i in (qx, qy, qz, pol))
+        norm = np.emath.sqrt(qx * qx + qy * qy + qz * qz)
+        norm[np.abs(norm - 1) < 1e-14] = 1
+        qx, qy, qz = (np.true_divide(i, norm) for i in (qx, qy, qz))
+        for i in (qx, qy, qz, pol):
+            i.flags.writeable = False
+            if i.ndim > 1:
+                raise ValueError("invalid shape of parameters")
+        self.qx, self.qy, self.qz = qx, qy, qz
+        self.pol = np.array(pol.real, int)
+        if np.any(self.pol != pol):
+            raise ValueError("polarizations must be integer")
+        if np.any(self.pol > 1) or np.any(self.pol < 0):
+            raise ValueError("polarization must be '0' or '1'")
+        self.lattice = self.kpar = None
+
+    def __getattr__(self, key):
+        dct = {"x": "qx", "y": "qy", "z": "qz", "s": "pol"}
+        try:
+            return tuple(getattr(self, dct[k]) for k in key)
+        except KeyError:
+            raise AttributeError(
+                f"'{type(self).__name__}' object has no attribute '{key}'"
+            ) from None
+
+    def __getitem__(self, idx):
+        """Get a subset of the basis.
+
+        This function allows index into the basis set by an integer, a slice, a sequence
+        of integers or bools, an ellipsis, or an empty tuple. All of them except the
+        integer and the empty tuple results in another basis set being returned. In case
+        of the two exceptions a tuple is returned.
+
+        Alternatively, the string "xyzp", "xyp", or "zp" can be used to access only a
+        subset of :attr:`qx`, :attr:`qy`, :attr:`qz`, and :attr:`pol`.
+        """
+        res = self.qx[idx], self.qy[idx], self.qz[idx], self.pol[idx]
+        if isinstance(idx, (int, np.integer)) or (isinstance(idx, tuple) and idx == ()):
+            return res
+        return type(self)(zip(*res))
+
+    @classmethod
+    def default(cls, kvecs):
+        """Default basis from the given wave vectors.
+
+        For each wave vector the two polarizations 1 and 0 are taken.
+
+        Example:
+            >>> PlaneWaveBasisByUnitVector.default([[0, 0, 5], [0, 3, 4]])
+            PlaneWaveBasisByUnitVector(
+                qx=[0. 0. 0. 0.],
+                qy=[0.  0.  0.6 0.6],
+                qz=[1.  1.  0.8 0.8],
+                pol=[1 0 1 0],
+            )
+
+        Args:
+            kvecs (array-like): Wave vectors in Cartesian coordinates.
+        """
+        kvecs = np.atleast_2d(kvecs)
+        modes = np.empty((2 * kvecs.shape[0], 4), kvecs.dtype)
+        modes[::2, :3] = kvecs
+        modes[1::2, :3] = kvecs
+        modes[::2, 3] = 1
+        modes[1::2, 3] = 0
+        return cls(modes)
+
+    @classmethod
+    def _from_iterable(cls, it):
+        if isinstance(cls, PlaneWaveBasisByUnitVector):
+            lattice = cls.lattice
+            kpar = cls.kpar
+            cls = type(cls)
+        else:
+            lattice = kpar = None
+        obj = cls(it)
+        obj.lattice = lattice
+        obj.kpar = kpar
+        return obj
+
+    def __eq__(self, other):
+        """Compare basis sets.
+
+        Basis sets are considered equal, when they have the same modes in the same
+        order.
+        """
+        try:
+            return self is other or (
+                np.array_equal(self.qx, other.qx)
+                and np.array_equal(self.qy, other.qy)
+                and np.array_equal(self.qz, other.qz)
+                and np.array_equal(self.pol, other.pol)
+            )
+        except AttributeError:
+            return False
+
+    def bycomp(self, k0, alignment="xy", material=Material()):
+        """Create a :class:`PlaneWaveBasisByComp`.
+
+        The plane wave basis is changed to a partial basis, where only two (real-valued)
+        wave vector components are defined the third component is then inferred from the
+        dispersion relation, which depends on the wave number and material, and a
+        ``modetype`` that specifies the sign of the third component.
+
+        Args:
+            alignment (str, optional): Wave vector components that are part of the
+                partial basis. Defaults to "xy", other permitted values are "yz" and
+                "zx".
+            k0 (float, optional): Wave number. If given, it is checked that the current
+                basis fulfils the dispersion relation.
+            material (:class:`~treams.Material` or tuple): Material definition. Defaults
+                to vacuum/air.
+        """
+        ks = material.ks(k0)[self.pol]
+        if alignment in ("xy", "yz", "zx"):
+            kpars = [ks * getattr(self, "q" + s) for s in alignment]
+        else:
+            raise ValueError(f"invalid alignment '{alignment}'")
+        obj = PlaneWaveBasisByComp(zip(*kpars, self.pol), alignment)
+        obj.lattice = self.lattice
+        obj.kpar = self.kpar
+        return obj
+
+    def kvecs(self, k0, material=Material(), modetype=None):
+        """Wave vectors.
+
+        Args:
+            k0 (float): Wave number.
+            material (:class:`~treams.Material` or tuple, optional): Material
+                definition. Defaults to vacuum/air.
+            modetype (optional): Currently unused for this class.
+        """
+        # TODO: check kz depending on modetype (alignment?)
+        ks = Material(material).ks(k0)[self.pol]
+        return ks * self.qx, ks * self.qy, ks * self.qz
+
+    def permute(self, n=1):
+        n = n % 3
+        lattice = None if self.lattice is None else self.lattice.permute(n)
+        kpar = None if self.kpar is None else self.kpar.permute(n)
+        qx, qy, qz = self.qx, self.qy, self.qz
+        for _ in range(n):
+            qx, qy, qz = qz, qx, qy
+        obj = type(self)(zip(*(qx, qy, qz, self.pol)))
+        obj.lattice = lattice
+        obj.kpar = kpar
+        return obj
+
+
+class PlaneWaveBasisByComp(PlaneWaveBasis):
+    """Partial plane wave basis.
+
+    A partial plane wave basis is defined by two wave vector components out of all three
+    Cartesian wave vector components ``kx``, ``ky``, and ``kz`` and the polarizations
+    ``pol``. Which two components are given is specified in the :attr:`alignment`. This
+    basis is mostly used for stratified media that are periodic or uniform in the two
+    alignment directions, such that the given wave vector components correspond to the
+    diffraction orders.
+
+    For plane waves there exist multiple :ref:`params:Polarizations`, such that
+    these modes can refer to either :func:`~treams.special.vpw_A` or
+    :func:`~treams.special.vpw_M` and :func:`~treams.special.vpw_N`.
+
+    Args:
+        modes (array-like): A tuple containing a list for each of ``k1``, ``k2``, and
+            ``pol``.
+        alignment (str, optional): Definition which wave vector components are given.
+            Defaults to "xy", other possible values are "yz" and "zx".
+
+    Attributes:
+        alignment (str): Alignment of the partial basis.
+        pol (array-like): Polarization, see also :ref:`params:Polarizations`.
+    """
+
+    def __init__(self, modes, alignment="xy"):
+        """Initialization."""
+        if isinstance(modes, np.ndarray):
+            modes = modes.tolist()
+        tmp = []
+        for m in modes:
+            if m not in tmp:
+                tmp.append(m)
+        modes = tmp
+        if len(modes) == 0:
+            kx = []
+            ky = []
+            pol = []
+        elif len(modes[0]) == 3:
+            kx, ky, pol = (*zip(*modes),)
+        else:
+            raise ValueError("invalid shape of modes")
+
+        self._kx, self._ky = [np.real(i) for i in (kx, ky)]
+        self.pol = np.array(np.real(pol), int)
+        for i, j in [(self._kx, kx), (self._ky, ky), (self.pol, pol)]:
+            i.flags.writeable = False
+            if i.ndim > 1:
+                raise ValueError("invalid shape of parameters")
+            if np.any(i != j):
+                raise ValueError("invalid value for parameter, must be real")
+        if np.any(self.pol > 1) or np.any(self.pol < 0):
+            raise ValueError("polarization must be '0' or '1'")
+
+        if alignment in ("xy", "yz", "zx"):
+            self._names = (*("k" + i for i in alignment),) + ("pol",)
+        else:
+            raise ValueError(f"invalid alignment '{alignment}'")
+
+        self.alignment = alignment
+        self.lattice = self.kpar = None
+
+    def permute(self, n=1):
+        n = n % 3
+        lattice = None if self.lattice is None else self.lattice.permute(n)
+        kpar = None if self.kpar is None else self.kpar.permute(n)
+        alignments = {"xy": "yz", "yz": "zx", "zx": "xy"}
+        alignment = self.alignment
+        for _ in range(n):
+            alignment = alignments[alignment]
+        obj = self._from_iterable(self, alignment=alignment)
+        obj.lattice = lattice
+        obj.kpar = kpar
+        return obj
+
+    @property
+    def kx(self):
+        """X-components of the wave vector.
+
+        If the components are not specified `None` is returned.
+        """
+        if self.alignment == "xy":
+            return self._kx
+        if self.alignment == "zx":
+            return self._ky
+        return None
+
+    @property
+    def ky(self):
+        """Y-components of the wave vector.
+
+        If the components are not specified `None` is returned.
+        """
+        if self.alignment == "xy":
+            return self._ky
+        if self.alignment == "yz":
+            return self._kx
+        return None
+
+    @property
+    def kz(self):
+        """Z-components of the wave vector.
+
+        If the components are not specified `None` is returned.
+        """
+        if self.alignment == "yz":
+            return self._ky
+        if self.alignment == "zx":
+            return self._kx
+        return None
+
+    def __getattr__(self, key):
+        dct = {"x": "kx", "y": "ky", "z": "kz", "s": "pol"}
+        try:
+            return tuple(getattr(self, dct[k]) for k in key)
+        except KeyError:
+            raise AttributeError(
+                f"'{type(self).__name__}' object has no attribute '{key}'"
+            ) from None
+
+    def __getitem__(self, idx):
+        """Get a subset of the basis.
+
+        This function allows index into the basis set by an integer, a slice, a sequence
+        of integers or bools, an ellipsis, or an empty tuple. All of them except the
+        integer and the empty tuple results in another basis set being returned. In case
+        of the two exceptions a tuple is returned.
+        """
+        res = self._kx[idx], self._ky[idx], self.pol[idx]
+        if isinstance(idx, (int, np.integer)) or (isinstance(idx, tuple) and idx == ()):
+            return res
+        return type(self)(zip(*res))
+
+    @classmethod
+    def default(cls, kpars, alignment="xy"):
+        """Default basis from the given wave vectors.
+
+        For each wave vector the two polarizations 1 and 0 are taken.
+
+        Example:
+            >>> PlaneWaveBasisByComp.default([[0, 0], [0, 3]])
+            PlaneWaveBasisByComp(
+                kx=[0. 0. 0. 0.],
+                ky=[0. 0. 3. 3.],
+                pol=[1 0 1 0],
+            )
+
+        Args:
+            kpars (array-like): Wave vector components in Cartesian coordinates.
+            alignment (str, optional): Definition which wave vector components are
+                given. Defaults to "xy", other possible values are "yz" and "zx".
+
+        """
+        kpars = np.atleast_2d(kpars)
+        modes = np.empty((2 * kpars.shape[0], 3), float)
+        modes[::2, :2] = kpars
+        modes[1::2, :2] = kpars
+        modes[::2, 2] = 1
+        modes[1::2, 2] = 0
+        return cls(modes, alignment=alignment)
+
+    @classmethod
+    def diffr_orders(cls, kpar, lattice, bmax):
+        """Create a basis set for a two-dimensional periodic system.
+
+        The reciprocal lattice to the given lattice is taken to consider all diffraction
+        orders that lie within the defined maximal radius (in reciprocal space).
+
+        Example:
+            >>> PlaneWaveBasisByComp.diffr_orders([0, 0], Lattice.square(2 * np.pi), 1)
+            PlaneWaveBasisByComp(
+                kx=[ 0.  0.  0.  0.  0.  0.  1.  1. -1. -1.],
+                ky=[ 0.  0.  1.  1. -1. -1.  0.  0.  0.  0.],
+                pol=[1 0 1 0 1 0 1 0 1 0],
+            )
+
+        Args:
+            kpar (float): Tangential wave vector components. Ideally they are in the
+                first Brillouin zone (use :func:`misc.firstbrillouin2d`).
+            lattice (:class:`treams.Lattice` or float): Lattice definition or pitch.
+            bmax (float): Maximal change of tangential wave vector components. So,
+                this defines a maximal momentum transfer.
+        """
+        lattice = Lattice(lattice)
+        if lattice.dim != 2:
+            raise ValueError("invalid lattice dimensions")
+        latrec = lattice.reciprocal
+        kpars = kpar + la.diffr_orders_circle(latrec, bmax) @ latrec
+        obj = cls.default(kpars, alignment=lattice.alignment)
+        obj.lattice = lattice
+        obj.kpar = WaveVector(kpar, alignment=lattice.alignment)
+        return obj
+
+    @classmethod
+    def _from_iterable(cls, it, alignment="xy"):
+        if isinstance(cls, PlaneWaveBasisByComp):
+            alignment = cls.alignment if alignment is None else alignment
+            lattice = cls.lattice
+            kpar = cls.kpar
+            cls = type(cls)
+        else:
+            lattice = kpar = None
+        obj = cls(it, alignment)
+        obj.lattice = lattice
+        obj.kpar = kpar
+        return obj
+
+    def __eq__(self, other):
+        """Compare basis sets.
+
+        Basis sets are considered equal, when they have the same modes in the same
+        order.
+        """
+        try:
+            skx, sky, skz = self.kx, self.ky, self.kz
+            okx, oky, okz = other.kx, other.ky, other.kz
+            return self is other or (
+                (np.array_equal(skx, okx) or (skx is None and okx is None))
+                and (np.array_equal(sky, oky) or (sky is None and oky is None))
+                and (np.array_equal(skz, okz) or (skz is None and okz is None))
+                and np.array_equal(self.pol, other.pol)
+            )
+        except AttributeError:
+            return False
+
+    def byunitvector(self, k0, material=Material(), modetype="up"):
+        """Create a complete basis :class:`PlaneWaveBasis`.
+
+        A plane wave basis is considered complete, when all three Cartesian components
+        and the polarization is defined for each mode. So, the specified wave number,
+        material, and modetype is taken to calculate the third Cartesian wave vector.
+        The modetype "up" ("down") is for waves propagating in the positive (negative)
+        direction with respect to the Cartesian axis that is orthogonal to the
+        alignment.
+
+        Args:
+            k0 (float): Wave number.
+            material (:class:`~treams.Material` or tuple, optional): Material
+                definition. Defaults to vacuum/air.
+            modetype (str, optional): Propagation direction. Defaults to "up".
+        """
+        if modetype not in ("up", "down"):
+            raise ValueError("modetype not recognized")
+        material = Material(material)
+        kx = self._kx
+        ky = self._ky
+        kz = material.kzs(k0, kx, ky, self.pol) * (2 * (modetype == "up") - 1)
+        if self.alignment == "yz":
+            kx, ky, kz = kz, kx, ky
+        elif self.alignment == "zy":
+            kx, ky, kz = ky, kz, kx
+        obj = PlaneWaveBasisByUnitVector(zip(kx, ky, kz, self.pol))
+        obj.lattice = self.lattice
+        obj.kpar = self.kpar
+        return obj
+
+    def kvecs(self, k0, material=Material(), modetype="up"):
+        """Wave vectors.
+
+        Args:
+            k0 (float): Wave number.
+            material (:class:`~treams.Material` or tuple, optional): Material
+                definition. Defaults to vacuum/air.
+            modetype (str, optional): Propagation direction. Defaults to "up".
+        """
+        if modetype not in ("up", "down"):
+            raise ValueError("modetype not recognized")
+        material = Material(material)
+        kx = self._kx
+        ky = self._ky
+        kz = material.kzs(k0, kx, ky, self.pol) * (2 * (modetype == "up") - 1)
+        if self.alignment == "yz":
+            return kz, kx, ky
+        if self.alignment == "zx":
+            return ky, kz, kx
+        return kx, ky, kz
+
+
+def _raise_basis_error(*args):
+    raise TypeError("'basis' must be BasisSet")
+
+
+class PhysicsDict(util.AnnotationDict):
+    """Physics dictionary.
+
+    Derives from :class:`treams.util.AnnotationDict`. This dictionary has additionally
+    several properties defined.
+
+    Attributes:
+        basis (:class:`BasisSet`): Basis set.
+        k0 (float): Wave number.
+        kpar (list): Parallel wave vector components. Usually, this is a list of length
+            3 with its items corresponding to the Cartesian axes. Unspecified items are
+            set to `nan`.
+        lattice (:class:`~treams.Lattice`): Lattice definition.
+        material (:class:`~treams.Material`): Material definition.
+        modetype (str): Mode type, for spherical and cylindrical waves this can be
+            "incident" and "scattered", for partial plane waves it can be "up" or
+            "down".
+        poltype (str): Polarization, see also :ref:`params:Polarizations`.
+    """
+
+    properties = {
+        "basis": (
+            ":class:`~treams.BasisSet`.",
+            lambda x: isinstance(x, BasisSet),
+            _raise_basis_error,
+        ),
+        "k0": ("Wave number.", lambda x: isinstance(x, float), float),
+        "kpar": (
+            "Wave vector components tangential to the lattice.",
+            lambda x: isinstance(x, WaveVector),
+            WaveVector,
+        ),
+        "lattice": (
+            ":class:`~treams.Lattice`.",
+            lambda x: isinstance(x, Lattice),
+            Lattice,
+        ),
+        "material": (
+            ":class:`~treams.Material`.",
+            lambda x: isinstance(x, Material),
+            Material,
+        ),
+        "modetype": ("Mode type.", lambda x: isinstance(x, str), str),
+        "poltype": (":ref:`params:Polarizations`.", lambda x: isinstance(x, str), str,),
+    }
+    """Special properties tracked by the PhysicsDict."""
+
+    def __setitem__(self, key, val):
+        """Set item specified by key to the defined value.
+
+        When overwriting an existing key an :class:`AnnotationWarning` is emitted.
+        Avoid the warning by explicitly deleting the key first. The special attributes
+        are cast to their corresponding types.
+
+        Args:
+            key (hashable): Key
+            val : Value
+
+        Warns:
+            AnnotationWarning
+        """
+        if key not in self.properties:
+            raise AttributeError(f"invalid key '{key}'")
+        _, testfunc, castfunc = self.properties[key]
+        if not testfunc(val):
+            val = castfunc(val)
+        super().__setitem__(key, val)
+
+
+class PhysicsArray(util.AnnotatedArray):
+    """Physics-aware array.
+
+    A physics aware array is a special type of :class`~treams.util.AnnotatedArray`.
+    Additionally to keeping track of the annotations, it is enhanced by the ability to
+    create suiting linear operators to perform tasks like rotations, translations, or
+    expansions into different basis sets and by applying special rules for the
+    physical properties of :class:`PhysicsDict` upon matrix multiplications, see also
+    :meth:`__matmul__`.
+    """
+
+    _scales = {"basis"}
+
+    changepoltype = op.OperatorAttribute(op.ChangePoltype)
+    """Polarization change matrix, see also :class:`treams.operators.ChangePoltype`."""
+    efield = op.OperatorAttribute(op.EField)
+    hfield = op.OperatorAttribute(op.HField)
+    dfield = op.OperatorAttribute(op.DField)
+    bfield = op.OperatorAttribute(op.BField)
+    gfield = op.OperatorAttribute(op.GField)
+    ffield = op.OperatorAttribute(op.FField)
+    """Field evaluation matrix, see also :class:`treams.operators.EField`."""
+    expand = op.OperatorAttribute(op.Expand)
+    """Expansion matrix, see also :class:`treams.operators.Expand`."""
+    expandlattice = op.OperatorAttribute(op.ExpandLattice)
+    """Lattice expansion matrix, see also :class:`treams.operators.ExpandLattice`."""
+    permute = op.OperatorAttribute(op.Permute)
+    """Permutation matrix, see also :class:`treams.operators.Permute`."""
+    rotate = op.OperatorAttribute(op.Rotate)
+    """Rotation matrix, see also :class:`treams.operators.Rotate`."""
+    translate = op.OperatorAttribute(op.Translate)
+    """Translation matrix, see also :class:`treams.operators.Translate`."""
+
+    def __init__(self, arr, ann=(), /, **kwargs):
+        """Initialization."""
+        super().__init__(arr, ann, **kwargs)
+        self._check()
+
+    @property
+    def ann(self):
+        """Array annotations."""
+        # necessary to define the setter below
+        return super().ann
+
+    def __getattr__(self, key):
+        try:
+            return super().__getattr__(key)
+        except AttributeError as err:
+            if key in PhysicsDict.properties:
+                return None
+            raise err from None
+
+    def __setattr__(self, key, val):
+        if key in PhysicsDict.properties:
+            val = (val,) * self.ndim if not isinstance(val, tuple) else val
+            self.ann.as_dict[key] = val
+        else:
+            super().__setattr__(key, val)
+
+    @ann.setter
+    def ann(self, ann):
+        """Set array annotations.
+
+        See also :meth:`treams.util.AnnotatedArray.__setitem__`.
+        """
+        self._ann = util.AnnotationSequence(*(({},) * self.ndim), mapping=PhysicsDict)
+        self._ann.update(ann)
+
+    def __repr__(self):
+        """String representation.
+
+        For a more managable output only the special physics properties are shown
+        alongside the array itself.
+        """
+        repr_arr = "    " + repr(self._array)[6:-1].replace("\n  ", "\n") + ","
+        for key in PhysicsDict.properties:
+            if getattr(self, key) is not None:
+                repr_arr += f"\n    {key}={repr(getattr(self, key))},"
+        return f"{self.__class__.__name__}(\n{repr_arr}\n)"
+
+    def _check(self):
+        """Run checks to validate the physical properties.
+
+        The checks are run on the last two dimensions. They include:
+
+            * Dispersion relation checks for :class:`PlaneWaveBasis` if `basis`, `k0`
+              and `material` is defined`
+            * `parity` polarizations are not physically permitted in chiral media.
+            * All lattices explicitly given and in the basis hints must be compatible.
+            * All tangential wave vector compontents must be compatible.
+        """
+        for a in self.ann[-2:]:
+            k0 = a.get("k0")
+            material = a.get("material")
+            modetype = a.get("modetype")
+            poltype = a.get("poltype")
+            basis = a.get("basis", namedtuple("_basis", "lattice kpar")(None, None))
+            if poltype == "parity" and getattr(material, "ischiral", False):
+                raise ValueError("poltype 'parity' not permitted for chiral material")
+
+    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
+        """Implement ufunc API.
+
+        Additionally to keeping track of the annotations the special properties of a
+        PhysicsArray are also "transparent" in matrix multiplications.
+        """
+        res = super().__array_ufunc__(ufunc, method, *inputs, **kwargs)
+        if (
+            ufunc is np.matmul
+            and method == "__call__"
+            and not isinstance(res, np.generic)
+        ):
+            axes = kwargs.get(
+                "axes", [tuple(range(-min(np.ndim(i), 2), 0)) for i in inputs + (res,)]
+            )
+            anns = [
+                i.ann[ax] if hasattr(i, "ann") else [{}, {}]
+                for i, ax in zip(inputs, axes)
+            ]
+            for name in PhysicsDict.properties:
+                if name in anns[0][-1] and all(name not in a for a in anns[1]):
+                    res.ann[axes[-1][-1]].setdefault(name, anns[0][-1][name])
+                if name in anns[1][0] and all(name not in a for a in anns[0]):
+                    res.ann[axes[-1][0]].setdefault(name, anns[1][0][name])
+        return res

src/treams/_lattice.py

+import collections
+
+import numpy as np
+
+import treams.lattice as la
+
+
+class Lattice:
+    """Lattice definition.
+
+    The lattice can be one-, two-, or three-dimensional. If it is not three-dimensional
+    it is required to be embedded into a lower dimensional subspace that is aligned with
+    the Cartesian axes. The default alignment for one-dimensional lattices is along the
+    z-axis and for two-dimensional lattices it is in the x-y-plane.
+
+    For a one-dimensional lattice the definition consists of simply the value of the
+    lattice pitch. Higher-dimensional lattices are defined by (2, 2)- or (3, 3)-arrays.
+    If the arrays are diagonal it is sufficient to specify the (2,)- or (3,)-array
+    diagonals. Alternatively, if another instance of a Lattice is given, the defined
+    alignment is extracted, which can be used to separate a lower-dimensional
+    sublattice.
+
+    Lattices are immutable objects.
+
+    Example:
+        >>> Lattice([1, 2])
+        Lattice([[1. 0.]
+                 [0. 2.]], alignment='xy')
+        >>> Lattice(_, 'x')
+        Lattice(1.0, alignment='x')
+
+    Args:
+        arr (float, array, Lattice): Lattice definition. Each row corresponds to one
+            lattice vector, each column to the axis defined in `alignment`.
+        alignment (str, optional): Alignment of the lattice. Possible values are 'x',
+            'y', 'z', 'xy', 'yz', 'zx', and 'xyz'.
+    """
+
+    _allowed_alignments = {
+        "x": "y",
+        "y": "z",
+        "z": "x",
+        "xy": "yz",
+        "yz": "zx",
+        "zx": "xy",
+        "xyz": "xyz",
+    }
+
+    def __init__(self, arr, alignment=None):
+        """Initialization."""
+        if isinstance(arr, Lattice):
+            self._alignment = arr.alignment if alignment is None else alignment
+            self._lattice = arr._sublattice(self._alignment)[...]
+            self._reciprocal = (
+                2 * np.pi / self[...] if self.dim == 1 else la.reciprocal(self[...])
+            )
+            return
+        if arr is None:
+            raise ValueError("Lattice cannot be 'None'")
+
+        arr = np.array(arr, float)
+        arr.flags.writeable = False
+        alignments = {1: ("z", "x", "y"), 4: ("xy", "yz", "zx"), 9: ("xyz",)}
+        if arr.ndim < 3 and arr.size == 1:
+            self._lattice = np.squeeze(arr)
+        elif arr.ndim == 1 and arr.size in (2, 3):
+            self._lattice = np.diag(arr)
+        elif arr.ndim == 2 and arr.shape[0] == arr.shape[1] and arr.shape[0] in (2, 3):
+            self._lattice = arr
+        else:
+            raise ValueError(f"invalid shape '{arr.shape}'")
+
+        alignment = (
+            alignments[self._lattice.size][0] if alignment is None else alignment
+        )
+        if alignment not in alignments[self._lattice.size]:
+            raise ValueError(f"invalid alignment '{alignment}'")
+        self._alignment = alignment
+        if self.volume == 0:
+            raise ValueError("linearly dependent lattice vectors")
+        self._reciprocal = (
+            2 * np.pi / self[...] if self.dim == 1 else la.reciprocal(self[...])
+        )
+
+    @property
+    def alignment(self):
+        """The alignment of the lattice in three-dimensional space.
+
+        For three-dimensional lattices it is always 'xyz' but lower-dimensional lattices
+        have to be aligned with a subset of the axes.
+
+        Returns:
+            str
+        """
+        return self._alignment
+
+    @classmethod
+    def square(cls, pitch, alignment=None):
+        """Create a two-dimensional square lattice.
+
+        Args:
+            pitch (float): Lattice constant.
+            alignment (str, optional): Alignment of the two-dimensional lattice in the
+                three-dimensional space. Defaults to 'xy'.
+
+        Returns:
+            Lattice
+        """
+        return cls(2 * [pitch], alignment)
+
+    @classmethod
+    def cubic(cls, pitch, alignment=None):
+        """Create a three-dimensional cubic lattice.
+
+        Args:
+            pitch (float): Lattice constant.
+            alignment (str, optional): Alignment of the lattice. Defaults to 'xyz'.
+
+        Returns:
+            Lattice
+        """
+        return cls(3 * [pitch], alignment)
+
+    @classmethod
+    def rectangular(cls, x, y, alignment=None):
+        """Create a two-dimensional rectangular lattice.
+
+        Args:
+            x (float): Lattice constant along the first dimension. For the default
+                alignment this corresponds to the x-axis.
+            y (float): Lattice constant along the second dimension. For the default
+                alignment this corresponds to the y-axis.
+            alignment (str, optional): Alignment of the two-dimensional lattice in the
+                three-dimensional space. Defaults to 'xy'.
+
+        Returns:
+            Lattice
+        """
+        return cls([x, y], alignment)
+
+    @classmethod
+    def orthorhombic(cls, x, y, z, alignment=None):
+        """Create a three-dimensional orthorhombic lattice.
+
+        Args:
+            x (float): Lattice constant along the first dimension. For the default
+                alignment this corresponds to the x-axis.
+            y (float): Lattice constant along the second dimension. For the default
+                alignment this corresponds to the y-axis.
+            z (float): Lattice constant along the third dimension. For the default
+                alignment this corresponds to the z-axis.
+            alignment (str, optional): Alignment of the lattice. Defaults to 'xyz'.
+
+        Returns:
+            Lattice
+        """
+        return cls([x, y, z], alignment)
+
+    @classmethod
+    def hexagonal(cls, pitch, height=None, alignment=None):
+        """Create a hexagonal lattice.
+
+        The lattice is two-dimensional if no height is specified else it is
+        three-dimensional
+
+        Args:
+            pitch (float): Lattice constant.
+            height (float, optional): Separation along the third axis for a
+                three-dimensional lattice.
+            alignment (str, optional): Alignment of the two-dimensional lattice in the
+                three-dimensional space. Defaults to either 'xy' or 'xyz' in the two-
+                or three-dimensional case, respectively.
+
+        Returns:
+            Lattice
+        """
+        if height is None:
+            return cls(
+                np.array([[pitch, 0], [0.5 * pitch, np.sqrt(0.75) * pitch]]), alignment
+            )
+        return cls(
+            np.array(
+                [[pitch, 0, 0], [0.5 * pitch, np.sqrt(0.75) * pitch, 0], [0, 0, height]]
+            ),
+            alignment,
+        )
+
+    def __eq__(self, other):
+        """Equality.
+
+        Two lattices are considered equal when they have the same dimension, alignment,
+        and lattice vectors.
+
+        Args:
+            other (Lattice): Lattice to compare with.
+
+        Returns:
+            bool
+        """
+        return other is not None and (
+            self is other
+            or (
+                self.alignment == other.alignment
+                and self.dim == other.dim
+                and np.all(self[...] == other[...])
+            )
+        )
+
+    @property
+    def dim(self):
+        """Dimension of the lattice.
+
+        Returns:
+            int
+        """
+        if self._lattice.ndim == 0:
+            return 1
+        return self._lattice.shape[0]
+
+    @property
+    def volume(self):
+        """(Generalized) volume of the lattice.
+
+        The value gives the lattice pitch, area, or volume depending on its dimension.
+        The volume is signed.
+
+        Returns:
+            float
+        """
+        if self.dim == 1:
+            return self._lattice
+        return la.volume(self._lattice)
+
+    @property
+    def reciprocal(self):
+        r"""Reciprocal lattice.
+
+        The reciprocal lattice to a given lattice with dimension :math:`d` and lattice
+        vectors :math:`\boldsymbol a_i` for :math:`i \in \{1, \dots, d\}` is defined by
+        lattice vectors :math:`\boldsymbol b_j` with :math:`j \in \{1, \dots, d\}` such
+        that :math:`\boldsymbol a_i \boldsymbol b_j = 2 \pi \delta_{ij}` is fulfilled.
+
+        Returns:
+            Lattice
+        """
+        return self._reciprocal
+
+    def __getitem__(self, idx):
+        """Index into the lattice.
+
+        Indexing can be used to obtain entries from the lattice vector definitions or
+        by using the Ellipsis or empty tuple to obtain the full array.
+
+        Returns:
+            float
+        """
+        return self._lattice[idx]
+
+    def __str__(self):
+        """String representation.
+
+        Simply returns the lattice pitch or lattice vectors.
+
+        Returns:
+            str
+        """
+        return str(self._lattice)
+
+    def __repr__(self):
+        """Representation.
+
+        The result can be used to recreate an equal instance.
+
+        Returns:
+            str
+        """
+        string = str(self._lattice).replace("\n", "\n        ")
+        return f"Lattice({string}, alignment='{self.alignment}')"
+
+    def _sublattice(self, key):
+        """Get the sublattice defined by key.
+
+        This function is called when one gives another instance of Lattice to the
+        constructor.
+
+        Args:
+            key (str): Alignment of the sublattice to extract.
+
+        Returns:
+            Lattice
+        """
+        key = key.lower()
+        if self.dim == 1:
+            if key == self.alignment:
+                return Lattice(self[...], key)
+            raise ValueError(f"sublattice with key '{key}' not availale")
+        idx = []
+        for c in key:
+            idx.append(self.alignment.find(c))
+        if -1 in idx or key not in self._allowed_alignments:
+            raise ValueError(f"sublattice with key '{key}' not availale")
+        idx_opp = [i for i in range(self.dim) if i not in idx]
+        mask = np.any(self[:, idx] != 0, axis=-1)
+        if (
+            (self[mask, :][:, idx_opp] != 0).any()
+            or (self[np.logical_not(mask), :][:, idx] != 0).any()
+            or len(idx) != sum(mask)
+        ):
+            raise ValueError("cannot determine sublattice")
+        return Lattice(self[mask, :][:, idx], key)
+
+    def permute(self, n=1):
+        """Permute the lattice orientation.
+
+        Get a new lattice with the alignment permuted.
+
+        Examples:
+            >>> Lattice.hexagonal(1, 2).permute()
+            Lattice([[0.    1.    0.   ]
+                     [0.    0.5   0.866]
+                     [2.    0.    0.   ]], alignment='xyz')
+            >>> Lattice.hexagonal(1).permute()
+            Lattice([[1.    0.   ]
+                     [0.5   0.866]], alignment='yz')
+
+        Args:
+            n (int, optional): Number of repeated permutations. Defaults to `1`.
+
+        Returns:
+            Lattice
+        """
+        if n != int(n):
+            raise ValueError("'n' must be integer")
+        n = n % 3
+        lattice = self[...]
+        alignment = self.alignment
+        if self.dim == 3:
+            while n > 0:
+                lattice = lattice[:, [2, 0, 1]]
+                n -= 1
+        else:
+            while n > 0:
+                alignment = self._allowed_alignments[alignment]
+                n -= 1
+        return Lattice(lattice, alignment)
+
+    def __bool__(self):
+        """Lattice instances always equate to True.
+
+        Returns:
+            bool
+        """
+        return True
+
+    def __or__(self, other):
+        """Merge two lattices.
+
+        Two lattices are combined into one if possible.
+
+        Example:
+            >>> Lattice(1, 'x') | Lattice(2)
+            Lattice([[2. 0.]
+                     [0. 1.]], alignment='zx')
+
+        Args:
+            other (Lattice): Lattice to merge.
+
+        Returns:
+            Lattice
+        """
+        if other is None or self == other:
+            return Lattice(self)
+
+        alignment = list({c for lat in (self, other) for c in lat.alignment})
+        alignment = "".join(sorted(alignment))
+        if alignment == "xz":
+            alignment = "zx"
+
+        ndim = len(alignment)
+
+        if ndim == 2:
+            if self.dim == 1 == other.dim:
+                if alignment.find(self.alignment) == 0:
+                    return Lattice([self[...], other[...]], alignment)
+                return Lattice([other[...], self[...]], alignment)
+            if self.dim == 2 and Lattice(self, other.alignment) == other:
+                return Lattice(self)
+            if other.dim == 2 and Lattice(other, self.alignment) == self:
+                return Lattice(other)
+        elif ndim == 3:
+            if self.dim == 3 and Lattice(self, other.alignment) == other:
+                return Lattice(self)
+            if other.dim == 3 and Lattice(other, self.alignment) == self:
+                return Lattice(other)
+            if self.dim == 2 == other.dim:
+                arr = np.zeros(3)
+                for i, c in enumerate("xyz"):
+                    if c in self.alignment:
+                        la0 = Lattice(self, c)
+                    else:
+                        la0 = None
+                    if c in other.alignment:
+                        la1 = Lattice(other, c)
+                    else:
+                        la1 = None
+                    if None not in (la0, la1) and la0 != la1:
+                        raise ValueError("cannot combine lattices")
+                    arr[i] = la0[...] if la0 is not None else la1[...]
+                return Lattice(arr)
+            arr = np.zeros((3, 3))
+            la0, la1 = (self, other) if self.dim == 1 else (other, self)
+            if la0.alignment == "x":
+                arr[0, 0] = la0[...]
+                arr[1:, 1:] = la1[...]
+                return Lattice(arr)
+            if la0.alignment == "y":
+                arr[1, 1] = la0[...]
+                arr[[[2], [0]], [2, 0]] = la1[...]
+                return Lattice(arr)
+            if la0.alignment == "z":
+                arr[2, 2] = la0[...]
+                arr[:2, :2] = la1[...]
+                return Lattice(arr)
+        raise ValueError("cannot combine lattices")
+
+    def __and__(self, other):
+        """Intersect two lattices.
+
+        The intersection of two lattices is taken if possible.
+
+        Example:
+            >>> Lattice([1, 2]) & Lattice([2, 3], 'yz')
+            Lattice(2.0, alignment='y')
+
+        Args:
+            other (Lattice): Lattice to intersect.
+
+        Returns:
+            Lattice
+        """
+        if other is None:
+            return None
+        if self == other:
+            return Lattice(self)
+
+        alignment = list({c for c in self.alignment if c in other.alignment})
+        if len(alignment) == 0:
+            raise ValueError("cannot combine lattices")
+        alignment = "".join(sorted(alignment))
+        if alignment == "xz":
+            alignment = "zx"
+
+        a, b = Lattice(self, alignment), Lattice(other, alignment)
+        if a == b:
+            return a
+        raise ValueError("cannot combine lattices")
+
+    def __le__(self, other):
+        """Test if one lattice includes another.
+
+        Example:
+            >>> Lattice(3) <= Lattice([1, 2, 3])
+            True
+
+        Args:
+            other (Lattice): Lattice to compare with.
+
+        Returns:
+            bool
+        """
+        try:
+            lat = Lattice(other, self.alignment)
+        except ValueError:
+            return False
+        return lat == self
+
+    def isdisjoint(self, other):
+        """Test if lattices are disjoint.
+
+        Lattices are considered disjoint if their alignments are disjoint.
+
+        Example:
+            >>> Lattice([1, 2, 3]).isdisjoint(Lattice(1))
+            False
+
+        Args:
+            other (Lattice): Lattice to compare with.
+
+        Returns:
+            bool
+        """
+        for c in other.alignment:
+            if c in self.alignment:
+                return False
+        return True
+
+
+class WaveVector(collections.abc.Sequence):
+    def __init__(self, seq=(), alignment=None):
+        try:
+            length = len(seq)
+        except TypeError:
+            length = 1
+            seq = [seq]
+        if length == 3:
+            self._vec = tuple(seq)
+        elif length == 2:
+            if alignment in ("xy", None):
+                self._vec = (seq[0], seq[1], np.nan)
+            elif alignment == "yz":
+                self._vec = (np.nan, seq[0], seq[1])
+            elif alignment == "zx":
+                self._vec = (seq[1], np.nan, seq[0])
+            else:
+                raise ValueError(f"invalid alignment: {alignment}")
+        elif length == 1:
+            if alignment in ("z", None):
+                self._vec = (np.nan, np.nan, seq[0])
+            elif alignment == "y":
+                self._vec = (np.nan, seq[0], np.nan)
+            elif alignment == "x":
+                self._vec = (seq[0], np.nan, np.nan)
+            else:
+                raise ValueError(f"invalid alignment: {alignment}")
+        elif length == 0:
+            self._vec = (np.nan,) * 3
+        else:
+            raise ValueError("invalid sequence")
+
+    def __str__(self):
+        return str(self._vec)
+
+    def __repr__(self):
+        return self.__class__.__name__ + str(self._vec)
+
+    def __eq__(self, other):
+        other = WaveVector(other)
+        for a, b in zip(self, other):
+            if a != b and not (np.isnan(a) and np.isnan(b)):
+                return False
+        return True
+
+    def __len__(self):
+        return 3
+
+    def __getitem__(self, key):
+        return self._vec[key]
+
+    def __or__(self, other):
+        other = WaveVector(other)
+        seq = ()
+        for a, b in zip(self, other):
+            isnan = np.isnan(a) or np.isnan(b)
+            if a != b and not isnan:
+                raise ValueError("non-matching WaveVector")
+            seq += (np.nan,) if isnan else (a,)
+        return WaveVector(seq)
+
+    def __and__(self, other):
+        other = WaveVector(other)
+        seq = ()
+        for a, b in zip(self, other):
+            if a != b and not (np.isnan(a) or np.isnan(b)):
+                raise ValueError("non-matching WaveVector")
+            seq += (b,) if np.isnan(a) else (a,)
+        return WaveVector(seq)
+
+    def permute(self, n=1):
+        x, y, z = self
+        n = n % 3
+        for _ in range(n):
+            x, y, z = z, x, y
+        return WaveVector((x, y, z))
+
+    def __le__(self, other):
+        other = WaveVector(other)
+        for a, b in zip(self, other):
+            if a != b and not np.isnan(b):
+                return False
+        return True
+
+    def isdisjoint(self, other):
+        other = WaveVector(other)
+        for a, b in zip(self, other):
+            if not (np.isnan(a) or np.isnan(b)):
+                return False
+        return True

src/treams/_material.py

+import numpy as np
+
+from treams import misc
+
+
+class Material:
+    r"""Material definition.
+
+    The material properties are defined in the frequency domain through scalar values
+    for permittivity :math:`\epsilon`, permeability :math:`\mu`, and the chirality
+    parameter  :math:`\kappa`. Materials are, thus, assumed to be linear,
+    time-invariant, homogeneous, isotropic, and local. Also, it is assumed that they
+    have no gain. The relation of the electric and magnetic fields is defined by
+
+    .. math::
+
+        \begin{pmatrix}
+            \frac{1}{\epsilon_0} \boldsymbol D \\
+            c \boldsymbol B
+        \end{pmatrix}
+        =
+        \begin{pmatrix}
+            \epsilon & \mathrm i \kappa \\
+            -\mathrm i \kappa & \mu
+        \end{pmatrix}
+        \begin{pmatrix}
+            \boldsymbol E \\
+            Z_0 \boldsymbol H
+        \end{pmatrix}
+
+    for a fixed vacuum wave number :math:`k_0` and spatial position
+    :math:`\boldsymbol r` with :math:`\epsilon_0` the vacuum permittivity, :math:`c` the
+    speed of light in vacuum, and :math:`Z_0` the vacuum impedance.
+
+    Args:
+        epsilon (optional, complex): Relative permittivity. Defaults to 1.
+        mu (optional, complex): Relative permeability. Defaults to 1.
+        kappa (optional, complex): Chirality parameter. Defaults to 0.
+    """
+
+    def __init__(self, epsilon=1, mu=1, kappa=0):
+        """Initialization."""
+        if isinstance(epsilon, Material):
+            epsilon, mu, kappa = epsilon()
+        elif isinstance(epsilon, (tuple, list, np.ndarray)):
+            if len(epsilon) == 0:
+                epsilon = 1
+            elif len(epsilon) == 1:
+                epsilon = epsilon[0]
+            elif len(epsilon) == 2:
+                epsilon, mu = epsilon
+            elif len(epsilon) == 3:
+                epsilon, mu, kappa = epsilon
+            else:
+                raise ValueError("invalid material definition")
+        self._epsilon = epsilon
+        self._mu = mu
+        self._kappa = kappa
+
+    @property
+    def epsilon(self):
+        """Relative permittivity.
+
+        Returns:
+            float or complex
+        """
+        return self._epsilon
+
+    @property
+    def mu(self):
+        """Relative permeability.
+
+        Returns:
+            float or complex
+        """
+        return self._mu
+
+    @property
+    def kappa(self):
+        """Chirality parameter.
+
+        Returns:
+            float or complex
+        """
+        return self._kappa
+
+    def __iter__(self):
+        """Iterator for a tuple containing the material parameters.
+
+        Useful for unpacking the material parameters into a function that takes these
+        parameters separately, e.g. ``foo(*material)``.
+        """
+        return iter((self.epsilon, self.mu, self.kappa))
+
+    @classmethod
+    def from_n(cls, n=1, impedance=None, kappa=0):
+        r"""Create material from refractive index and relative impedance.
+
+        This function calculates the relative permeability and permittivity with
+        :math:`\epsilon = \frac{n}{Z}` and :math:`\mu = nZ`. The chirality parameter is
+        considered separately.
+
+        Note:
+            The refractive index is defined independently from the chirality parameter
+            here. For an alternative definition of the refractive index, see also
+            :func:`Material.from_nmp`.
+
+        Args:
+            n (complex, optional): Refractive index. Defaults to 1.
+            impedance(complex, optional): Relative impedance. Defaults to the inverse
+                of the refractive index.
+            kappa (complex, optional): Chirality parameter. Defaults to 0.
+
+        Returns:
+            Material
+        """
+        impedance = 1 / n if impedance is None else impedance
+        epsilon = n / impedance
+        mu = n * impedance
+        return cls(epsilon, mu, kappa)
+
+    @classmethod
+    def from_nmp(cls, ns=(1, 1), impedance=None):
+        r"""Create material from refractive indices of both helicities.
+
+        This function calculates the relative permeability and permittivity and the
+        chirality parameter with :math:`\epsilon = \frac{n_+ + n_-}{2Z}`,
+        :math:`\mu = \frac{(n_+ + n_-)Z}{2}` and :math:`\mu = \frac{(n_+ - n_-)}{2}`.
+
+        Note:
+            Two refractive indices are defined here that depend on the chirality
+            parameter. For an alternative definition of the refractive index, see also
+            :func:`Material.from_n`.
+
+        Args:
+            ns ((2,)-array-like, optional): Negative and positive helicity refractive
+                index. Defaults to (1, 1).
+            impedance(complex, optional): Relative impedance. Defaults to the inverse of
+                the refractive index.
+
+        Returns:
+            Material
+        """
+        n = sum(ns) * 0.5
+        kappa = (ns[1] - ns[0]) * 0.5
+        return cls.from_n(n, impedance, kappa)
+
+    @property
+    def n(self):
+        r"""Refractive index.
+
+        The refractive index is defined by :math:`n = \sqrt{\epsilon \mu}`, with an
+        enforced non-negative imaginary part.
+
+        Note:
+            The refractive index returned is independent from the chirality parameter.
+            For an alternative definition of the refractive index, see also
+            :func:`Material.nmp`.
+
+        Returns:
+            complex
+        """
+        n = np.sqrt(self.epsilon * self.mu)
+        if n.imag < 0:
+            n = -n
+        return n
+
+    @property
+    def nmp(self):
+        r"""Refractive indices of both helicities.
+
+        The refractive indices are defined by
+        :math:`n_\pm = \sqrt{\epsilon \mu} \pm \kappa`, with an enforced non-negative
+        imaginary part.
+
+        Note:
+            The refractive indices returned depend on the chirality parameter. For an
+            alternative definition of the refractive index, see also :func:`Material.n`.
+
+        Returns:
+            tuple
+        """
+        return misc.refractive_index(self.epsilon, self.mu, self.kappa)
+
+    @property
+    def impedance(self):
+        r"""Relative impedance.
+
+        The relative impedance is defined by :math:`Z = \sqrt{\frac{\epsilon}{\mu}}`.
+
+        Returns:
+            complex
+        """
+        return np.sqrt(self.mu / self.epsilon)
+
+    def __call__(self):
+        """Return a tuple containing all material parameters.
+
+        Returns:
+            tuple
+        """
+        return self.epsilon, self.mu, self.kappa
+
+    def __eq__(self, other):
+        """Compare material parameters.
+
+        Materials are considered equal, when all material parameters are equal. Also,
+        compares with objects that contain at most three values.
+
+        Returns:
+            bool
+        """
+        if other is None:
+            return False
+        if not isinstance(other, Material):
+            other = Material(*other)
+        return (
+            self.epsilon == other.epsilon
+            and self.mu == other.mu
+            and self.kappa == other.kappa
+        )
+
+    @property
+    def ischiral(self):
+        """Test if the material is chiral.
+
+        Returns:
+            bool
+        """
+        return self.kappa != 0
+
+    @property
+    def isreal(self):
+        """Test if the material has purely real parameters.
+
+        Returns:
+            bool
+        """
+        return all(i.imag == 0 for i in self)
+
+    def __str__(self):
+        """All three material parameters.
+
+        Returns:
+            str
+        """
+        return "(" + ", ".join([str(i) for i in self()]) + ")"
+
+    def __repr__(self):
+        """Representation that allows recreating the object.
+
+        Returns:
+            str
+        """
+        return self.__class__.__name__ + str(self)
+
+    def ks(self, k0):
+        """Return the wave numbers in the medium for both polarizations.
+
+        The first value corresponds to negative helicity and the second to positive
+        helicity. For achiral materials where parity polarizations can be used both
+        values are equal.
+
+        Args:
+            k0 (float): Wave number in vacuum.
+
+        Returns:
+            tuple
+        """
+        return k0 * self.nmp
+
+    def krhos(self, k0, kz, pol):
+        r"""The (cylindrically) radial part of the wave vector.
+
+        The cylindrically radial part is defined by :math:`k_\rho = \sqrt(k^2 - k_z^2)`.
+        In case of chiral materials :math:`k` and so :math`k_\rho` depends on the
+        polarization. The returned values have non-negative imaginary parts.
+
+        Args:
+            k0 (float): Wave number in vacuum.
+            kz (float, array-like): z-component of the wave vector
+            pol (int, array-like): Polarization indices
+                (:ref:`params:Polarizations`).
+
+        Returns:
+            complex, array-like
+        """
+        ks = self.ks(k0)[pol]
+        return misc.wave_vec_z(kz, 0, ks)
+
+    def kzs(self, k0, kx, ky, pol):
+        r"""The z-component of the wave vector.
+
+        The z-component of the wave vector is defined by
+        :math:`k_z = \sqrt(k^2 - k_x^2 - k_y^2)`. In case of chiral materials :math:`k`
+        and so :math`k_z` depends on the polarization. The returned values have
+        non-negative imaginary parts.
+
+        Args:
+            k0 (float): Wave number in vacuum.
+            kx (float, array-like): x-component of the wave vector
+            ky (float, array-like): y-component of the wave vector
+            pol (int, array-like): Polarization indices
+                (:ref:`params:Polarizations`).
+
+        Returns:
+            complex, array-like
+        """
+        ks = self.ks(k0)[pol]
+        return misc.wave_vec_z(kx, ky, ks)

src/treams/_smatrix.py

+import copy
+
+import numpy as np
+
+from treams import config, util
+from treams._core import PhysicsArray
+from treams._core import PlaneWaveBasisByComp as PWBC
+from treams._material import Material
+from treams._operators import translate
+from treams._tmatrix import TMatrixC
+from treams.coeffs import fresnel
+
+
+class SMatrix(PhysicsArray):
+    """S-matrix for a plane wave."""
+
+    def _check(self):
+        super()._check()
+        shape = np.shape(self)
+        if len(shape) != 2 or shape[0] != shape[1]:
+            raise util.AnnotationError(f"invalid shape: '{shape}'")
+        if not isinstance(self.k0, (int, float, np.floating, np.integer)):
+            raise util.AnnotationError("invalid k0")
+        if self.poltype is None:
+            self.poltype = config.POLTYPE
+        if self.poltype not in ("parity", "helicity"):
+            raise util.AnnotationError("invalid poltype")
+        material = (None, None) if self.material is None else self.material
+        if isinstance(material, tuple):
+            self.material = tuple(Material() if m is None else m for m in material)
+        if self.basis is None:
+            raise util.AnnotationError("basis not set")
+        if not isinstance(self.basis, PWBC):
+            self.basis = self.basis.bycomp(self.k0, self.material)
+
+
+class OperatorAttributeSMatrices:
+    def __init__(self, name):
+        self._name = name
+        self._obj = self._objtype = None
+
+    def __get__(self, obj, objtype=None):
+        self._obj = obj
+        self._objtype = type(obj) if objtype is None else objtype
+        return self
+
+    def __call__(self, *args, **kwargs):
+        res = [
+            [getattr(ii, self._name)(*args, **kwargs) for ii in i] for i in self._obj
+        ]
+        return self._objtype(res)
+
+    def apply_left(self, *args, **kwargs):
+        res = [
+            [getattr(ii, self._name).apply_left(*args, **kwargs) for ii in i]
+            for i in self._obj
+        ]
+        return self._objtype(res)
+
+    def apply_right(self, *args, **kwargs):
+        res = [
+            [getattr(ii, self._name).apply_right(*args, **kwargs) for ii in i]
+            for i in self._obj
+        ]
+        return self._objtype(res)
+
+
+class SMatrices:
+    r"""Collection of four S-matrices with a plane wave basis.
+
+    The S-matrix describes the scattering of incoming into outgoing modes using a plane
+    wave basis, with functions :func:`treams.special.vsw_A`,
+    :func:`treams.special.vsw_M`, and :func:`treams.special.vsw_N`. The primary
+    direction of propagation is parallel or anti-parallel to the z-axis. The scattering
+    object itself is infinitely extended in the x- and y-directions. The S-matrix is
+    divided into four submatrices :math:`S_{\uparrow \uparrow}`,
+    :math:`S_{\uparrow \downarrow}`, :math:`S_{\downarrow \uparrow}`, and
+    :math:`S_{\downarrow \downarrow}`:
+
+    .. math::
+
+        S = \begin{pmatrix}
+            S_{\uparrow \uparrow} & S_{\uparrow \downarrow} \\
+            S_{\downarrow \uparrow} & S_{\downarrow \downarrow}
+        \end{pmatrix}\,.
+
+    These matrices describe the transmission of plane waves propagating into positive
+    z-direction, reflection of plane waves into the positive z-direction, reflection of
+    plane waves into negative z-direction, and the transmission of plane waves
+    propagating in negative z-direction, respectively. Each of those for matrices
+    contain different diffraction orders and different polarizations. The polarizations
+    can be in either helicity of parity basis.
+
+    The wave number :attr:`k0` and, if not vacuum, the material :attr:`material` are
+    also required.
+
+    Args:
+        smats (SMatrix): S-matrices.
+        k0 (float): Wave number in vacuum.
+        basis (PlaneWaveBasisByComp): The basis for the S-matrices.
+        material (tuple, Material, optional): Material definition, if a tuple of length
+            two is specified, this refers to the materials above and below the S-matrix.
+        poltype (str, optional): Polarization type (:ref:`params:Polarizations`).
+    """
+    permute = OperatorAttributeSMatrices("permute")
+    translate = OperatorAttributeSMatrices("translate")
+    rotate = OperatorAttributeSMatrices("rotate")
+    changepoltype = OperatorAttributeSMatrices("changepoltype")
+
+    def __init__(self, smats, **kwargs):
+        """Initialization."""
+        if len(smats) != 2 or len(smats[0]) != 2 or len(smats[1]) != 2:
+            raise ValueError("invalid shape of S-matrices")
+        if "material" in kwargs:
+            material = kwargs["material"]
+            del kwargs["material"]
+        elif hasattr(smats[0][0], "material"):
+            material = smats[0][0].material
+        elif hasattr(smats[1][1], "material"):
+            material = smats[1][1].material
+            if isinstance(material, tuple):
+                material = material[::-1]
+        elif hasattr(smats[0][1], "material"):
+            material = smats[0][1].material
+        elif hasattr(smats[1][0], "material"):
+            material = smats[1][0].material
+        else:
+            material = Material()
+        if isinstance(material, tuple):
+            ma, mb = material
+        else:
+            ma = mb = Material(material)
+        material = [[(ma, mb), (ma, ma)], [(mb, mb), (mb, ma)]]
+        modetype = [[("up", "up"), ("up", "down")], [("down", "up"), ("down", "down")]]
+        self._sms = [
+            [
+                SMatrix(s, material=m, modetype=t, **kwargs)
+                for s, m, t in zip(ar, mr, tr)
+            ]
+            for ar, mr, tr in zip(smats, material, modetype)
+        ]
+        self.material = self[0, 0].material
+        self.basis = self[0, 0].basis
+        self.k0 = self[0, 0].k0
+        self.poltype = self[0, 0].poltype
+        for i in (self[1, 0], self[0, 1], self[1, 1]):
+            i.k0 = self.k0
+            i.basis = self.basis
+            i.poltype = self.poltype
+
+    def __eq__(self, other):
+        return all(
+            (np.abs(self[i, j] - other[i][j]) < 1e-14).all()
+            for i in range(2)
+            for j in range(2)
+        )
+
+    def __getitem__(self, key):
+        keys = {0: 0, "up": 0, 1: 1, "down": 1}
+        if key in keys:
+            return self._sms[keys[key]]
+        if not isinstance(key, tuple) or len(key) not in (1, 2):
+            raise KeyError(f"invalid key: '{key}'")
+        if len(key) == 1:
+            return self[key[0]]
+        key = tuple(keys[k] for k in key)
+        return self._sms[key[0]][key[1]]
+
+    def __len__(self):
+        return 2
+
+    def __iter__(self):
+        return iter(self._sms)
+
+    @classmethod
+    def interface(cls, basis, k0, materials):
+        """Planar interface between two media.
+
+        Args:
+            basis (PlaneWaveBasisByComp): Basis definitions.
+            k0 (float): Wave number in vacuum
+            materials (Sequence[Material]): Material definitions.
+
+        Returns:
+            SMatrix
+        """
+        materials = tuple(map(Material, materials))
+        ks = np.array([m.ks(k0) for m in materials])
+        choice = basis.pol == 0
+        kxs = basis.kx[choice], basis.kx[~choice]
+        kys = basis.ky[choice], basis.ky[~choice]
+        qs = np.zeros((2, 2, len(basis), len(basis)), complex)
+        if all(kxs[0] == kxs[1]) and all(kys[0] == kys[1]):
+            kzs = np.stack(
+                [
+                    m.kzs(k0, kxs[0][:, None], kys[0][:, None], np.array([[0, 1]]))
+                    for m in materials
+                ],
+                -2,
+            )
+            vals = fresnel(ks, kzs, [m.impedance for m in materials])
+            qs[:, :, choice, choice] = np.moveaxis(vals[:, :, :, 0, 0], 0, -1)
+            qs[:, :, choice, ~choice] = np.moveaxis(vals[:, :, :, 0, 1], 0, -1)
+            qs[:, :, ~choice, choice] = np.moveaxis(vals[:, :, :, 1, 0], 0, -1)
+            qs[:, :, ~choice, ~choice] = np.moveaxis(vals[:, :, :, 1, 1], 0, -1)
+        else:
+            for i, (kx, ky, pol) in enumerate(basis):
+                for ii, (kx2, ky2, pol2) in enumerate(basis):
+                    if kx != kx2 or ky != ky2:
+                        continue
+                    kzs = np.array([m.kzs(k0, kx, ky) for m in materials])
+                    vals = fresnel(ks, kzs, [m.impedance for m in materials])
+                    qs[:, :, ii, i] = vals[:, :, pol2, pol]
+        return cls(qs, k0=k0, basis=basis, material=materials[::-1], poltype="helicity")
+
+    @classmethod
+    def slab(cls, thickness, basis, k0, materials):
+        """Slab of material.
+
+        A slab of material, defined by a thickness and three materials. Consecutive
+        slabs of material can be defined by `n` thicknesses and and `n + 2` material
+        parameters.
+
+        Args:
+            k0 (float): Wave number in vacuum.
+            basis (PlaneWaveBasisByComp): Basis definition.
+            tickness (Sequence[Float]): Thickness of the slab or the thicknesses of all
+                slabs in order from negative to positive z.
+            materials (Sequenze[Material]): Material definitions from negative to
+                positive z.
+
+        Returns:
+            SMatrix
+        """
+        try:
+            iter(thickness)
+        except TypeError:
+            thickness = [thickness]
+        res = cls.interface(basis, k0, materials[:2])
+        for d, ma, mb in zip(thickness, materials[1:-1], materials[2:]):
+            if np.ndim(d) == 0:
+                d = [0, 0, d]
+            x = cls.propagation(d, basis, k0, ma)
+            res = res.add(x)
+            res = res.add(cls.interface(basis, k0, (ma, mb)))
+        return res
+
+    @classmethod
+    def stack(cls, items):
+        """Stack of S-matrices.
+
+        Electromagnetically couple multiple S-matrices in the order given. Before
+        coupling it can be checked for matching materials and modes.
+
+        Args:
+            items (Sequence[SMatrix]): An array of S-matrices in their intended order
+                from negative to positive z.
+
+        Returns:
+            SMatrix
+        """
+        acc = items[0]
+        for item in items[1:]:
+            acc = acc.add(item)
+        return acc
+
+    @classmethod
+    def propagation(cls, r, basis, k0, material=Material()):
+        """S-matrix for the propagation along a distance.
+
+        This S-matrix translates the reference origin along `r`.
+
+        Args:
+            r (float, (3,)-array): Translation vector.
+            k0 (float): Wave number in vacuum.
+            basis (PlaneWaveBasis): Basis definition.
+            material (Material, optional): Material definition.
+
+        Returns:
+            SMatrix
+        """
+        sup = translate(r, basis=basis, k0=k0, material=material, modetype="up")
+        sdown = translate(
+            np.negative(r), basis=basis, k0=k0, material=material, modetype="down"
+        )
+        zero = np.zeros_like(sup)
+        material = Material(material)
+        return cls([[sup, zero], [zero, sdown]], basis=basis, k0=k0, material=material)
+
+    @classmethod
+    def from_array(cls, tm, basis, *, eta=0):
+        """S-matrix from an array of (cylindrical) T-matrices.
+
+        Create a S-matrix for a two-dimensional array of objects described by the
+        T-Matrix or an one-dimensional array of objects described by a cylindrical
+        T-matrix.
+
+        Args:
+            tm (TMatrix or TMatrixC): (Cylindrical) T-matrix to place in the array.
+            basis (PlaneWaveBasisByComp): Basis definition.
+            eta (float or complex, optional): Splitting parameter in the lattice sum.
+
+        Returns:
+            SMatrix
+        """
+        if isinstance(tm, TMatrixC):
+            basis = basis.permute(-1)
+        pu, pd = (
+            tm.expandlattice(basis=basis, modetype=i, eta=eta) for i in ("up", "down")
+        )
+        au, ad = (tm.expand.eval_inv(basis, i) for i in ("up", "down"))
+        eye = np.eye(len(basis))
+        res = cls([[eye + pu @ au, pu @ ad], [pd @ au, eye + pd @ ad]])
+        if isinstance(tm, TMatrixC):
+            res = res.permute()
+        return res
+
+    def add(self, sm):
+        """Couple another S-matrix on top of the current one.
+
+        See also :func:`treams.SMatrix.stack` for a function that does not change the
+        current S-matrix but creates a new one.
+
+        Args:
+            sm (SMatrix): S-matrix to add.
+
+        Returns:
+            SMatrix
+        """
+        dim = len(self.basis)
+        snew = [[None, None], [None, None]]
+        s_tmp = np.linalg.solve(np.eye(dim) - self[0, 1] @ sm[1][0], self[0, 0])
+        snew[0][0] = sm[0][0] @ s_tmp
+        snew[1][0] = self[1, 0] + self[1, 1] @ sm[1][0] @ s_tmp
+        s_tmp = np.linalg.solve(np.eye(dim) - sm[1][0] @ self[0, 1], sm[1][1])
+        snew[1][1] = self[1, 1] @ s_tmp
+        snew[0][1] = sm[0][1] + sm[0][0] @ self[0, 1] @ s_tmp
+        return SMatrices(snew)
+
+    def double(self, n=1):
+        """Double the S-matrix.
+
+        By default this function doubles the S-matrix but it can also create a
+        :math:`2^n`-fold repetition of itself:
+
+        Args:
+            n (int, optional): Number of times to double itself. Defaults to 1.
+
+        Returns:
+            SMatrix
+        """
+        res = self
+        for _ in range(n):
+            res = res.add(self)
+        return res
+
+    def illuminate(self, illu, illu2=None, /, *, smat=None):
+        """Field coefficients above and below the S-matrix.
+
+        Given an illumination defined by the coefficients of each incoming mode
+        calculate the coefficients for the outgoing field above and below the S-matrix.
+        If a second SMatrix is given, the field expansions in between are also
+        calculated.
+
+        Args:
+            illu (array-like): Illumination, if `modetype` is specified, the direction
+                will be chosen accordingly.
+            illu2 (array-like, optional): Second illumination. If used, the first
+                argument is taken to be coming from below and this one to be coming from
+                above.
+            smat (SMatrix, optional): Second S-matrix for the calculation of the
+                field expansion between two S-matrices.
+
+        Returns:
+            tuple
+        """
+        modetype = getattr(illu, "modetype", "up")
+        if isinstance(modetype, tuple):
+            modetype = modetype[max(-2, -len(tuple))]
+        illu2 = np.zeros(np.shape(illu)[-2:]) if illu2 is None else illu2
+        if modetype == "down":
+            illu, illu2 = illu2, illu
+        if smat is None:
+            field_up = self[0, 0] @ illu + self[0, 1] @ illu2
+            field_down = self[1, 0] @ illu + self[1, 1] @ illu2
+            return field_up, field_down
+        stmp = np.eye(len(self.basis)) - self[0, 1] @ smat[1, 0]
+        field_in_up = np.linalg.solve(
+            stmp, self[0, 0] @ illu + self[0, 1] @ smat[1, 1] @ illu2
+        )
+        field_in_down = smat[1, 0] @ field_in_up + smat[1, 1] @ illu2
+        field_up = smat[0, 1] @ illu2 + smat[0, 0] @ field_in_up
+        field_down = self[1, 0] @ illu + self[1, 1] @ field_in_down
+        return field_up, field_down, field_in_up, field_in_down
+
+    def tr(self, illu):
+        """Transmittance and reflectance.
+
+        Calculate the transmittance and reflectance for one S-matrix with the given
+        illumination and direction.
+
+        Args:
+            illu (complex, array_like): Expansion coefficients for the incoming light
+
+        Returns:
+            tuple
+        """
+        modetype = getattr(illu, "modetype", "up")
+        if isinstance(modetype, tuple):
+            modetype = modetype[max(-2, -len(tuple))]
+        trans, refl = self.illuminate(illu)
+        material = self.material
+        if not isinstance(material, tuple):
+            material = material, material
+        paz = [poynting_avg_z(self.basis, self.k0, m, self.poltype) for m in material]
+        if modetype == "down":
+            trans, refl = refl, trans
+            paz.reverse()
+        illu = np.asarray(illu)
+        s_t = np.real(trans.conjugate().T @ paz[0][0] @ trans)
+        s_r = np.real(refl.conjugate().T @ paz[1][0] @ refl)
+        s_i = np.real(np.conjugate(illu).T @ paz[1][0] @ illu)
+        s_ir = np.real(
+            refl.conjugate().T @ paz[1][1] @ illu
+            - np.conjugate(illu).T @ paz[1][1] @ refl
+        )
+        return s_t / (s_i + s_ir), s_r / (s_i + s_ir)
+
+    def cd(self, illu):
+        """Transmission and absorption circular dichroism.
+
+        Calculate the transmission and absorption CD for one S-matrix with the given
+        illumination and direction.
+
+        Args:
+            illu (complex, PhysicsArray): Expansion coefficients for the incoming light
+
+        Returns:
+            tuple
+        """
+        minus, plus = self.basis.pol == 0, self.basis.pol == 1
+        if self.poltype == "helicity":
+            illuopposite = np.zeros_like(illu)
+            illuopposite[minus] = illu[plus]
+            illuopposite[plus] = illu[minus]
+        else:
+            illuopposite = copy.deepcopy(illu)
+            illuopposite[minus] *= -1
+        tm, rm = self.tr(illu)
+        tp, rp = self.tr(illuopposite)
+        return (tp - tm) / (tp + tm), (tp + rp - tm - rm) / (tp + rp + tm + rm)
+
+    def periodic(self):
+        r"""Periodic repetition of the S-matrix.
+
+        Transform the S-matrix to an infinite periodic arrangement of itself defined
+        by
+
+        .. math::
+
+            \begin{pmatrix}
+                S_{\uparrow \uparrow} & S_{\uparrow \downarrow} \\
+                -S_{\downarrow \downarrow}^{-1}
+                S_{\downarrow \uparrow} S_{\uparrow \uparrow} &
+                S_{\downarrow \downarrow}^{-1} (\mathbb{1} - S_{\downarrow \uparrow}
+                S_{\uparrow \downarrow})
+            \end{pmatrix}\,.
+
+        Returns:
+            complex, array_like
+
+        """
+        dim = len(self.basis)
+        res = np.empty((2 * dim, 2 * dim), dtype=complex)
+        res[0:dim, 0:dim] = self[0, 0]
+        res[0:dim, dim:] = self[0, 1]
+        res[dim:, 0:dim] = -np.linalg.solve(self[1, 1], self[1, 0] @ self[0, 0])
+        res[dim:, dim:] = np.linalg.solve(
+            self[1, 1], np.eye(dim) - self[1, 0] @ self[0, 1]
+        )
+        return res
+
+    def bands_kz(self, az):
+        r"""Band structure calculation.
+
+        Calculate the band structure for the given S-matrix, assuming it is periodically
+        repeated along the z-axis. The function returns the z-components of the wave
+        vector :math:`k_z` and the corresponding eigenvectors :math:`v` of
+
+        .. math::
+
+            \begin{pmatrix}
+                S_{\uparrow \uparrow} & S_{\uparrow \downarrow} \\
+                -S_{\downarrow \downarrow}^{-1} S_{\downarrow \uparrow}
+                    S_{\uparrow \uparrow} &
+                S_{\downarrow \downarrow}^{-1} (\mathbb{1} - S_{\downarrow \uparrow}
+                    S_{\uparrow \downarrow})
+            \end{pmatrix}
+            \boldsymbol v
+            =
+            \mathrm{e}^{\mathrm{i}k_z a_z}
+            \boldsymbol v\,.
+
+        Args:
+            az (float): Lattice pitch along the z direction
+
+        Returns:
+            tuple
+
+        """
+        w, v = np.linalg.eig(self.periodic())
+        return -1j * np.log(w) / az, v
+
+    def __repr__(self):
+        return f"""{type(self).__name__}(...,
+    basis={self.basis},
+    k0={self.k0},
+    material={self.material},
+    poltype='{self.poltype}',
+)"""
+
+
+def chirality_density(basis, k0, material=Material(), poltype=None, z=(0, 0)):
+    poltype = config.POLTYPE if poltype is None else poltype
+    material = Material(material)
+    kx, ky, kz = basis.kvecs(k0, material)
+    k = material.ks(k0)[basis.pol]
+    re, im = np.real(kz), np.imag(kz)
+    prefuu = 1 + np.real((k * k - kz * 2j * im) / (k * k.conjugate()))
+    prefdu = 1 + np.real((k * k - kz * 2 * re) / (k * k.conjugate()))
+    if np.abs(z[1] - z[0]) > 1e-16:
+        prefuu *= np.sinh(re * (z[1] - z[0])) / (re * (z[1] - z[0]))
+        prefdd = np.exp(im * (z[1] + z[0])) * prefuu
+        prefuu *= np.exp(-im * (z[1] + z[0]))
+        prefdu *= (
+            np.sin(re * (z[1] - z[0]))
+            * np.cos(re * (z[1] + z[0]))
+            / (re * (z[1] - z[0]))
+        )
+    else:
+        prefdd = prefuu
+    if poltype == "helicity":
+        return (
+            np.diag(prefuu * (2 * basis.pol - 1)),
+            np.diag(prefdd * (2 * basis.pol - 1)),
+            np.diag(2 * prefdu * (2 * basis.pol - 1)),
+        )
+    if poltype == "parity":
+        where = (
+            (kx[:, None] == kx)
+            & (ky[:, None] == ky)
+            & (basis.pol[:, None] != basis.pol)
+        )
+        return where * prefuu, where * prefdd, where * 2 * prefdu
+    raise ValueError(f"invalid poltype: '{poltype}'")
+
+
+def poynting_avg_z(basis, k0, material=Material(), poltype=None):
+    r"""Time-averaged z-component of the Poynting vector.
+
+    Calculate the time-averaged Poynting vector's z-component
+
+    .. math::
+
+        \langle S_z \rangle
+        = \frac{1}{2}
+        \Re (\boldsymbol E \times \boldsymbol H^\ast) \boldsymbol{\hat{z}}
+
+    on one side of the S-matrix with the given coefficients.
+
+    Returns:
+        tuple
+    """
+    poltype = config.POLTYPE if poltype is None else poltype
+    material = Material(material)
+    kx, ky, kzs = basis.kvecs(k0, material)
+    gamma = kzs / (material.ks(k0)[basis.pol] * material.impedance)
+    selection = (kx[:, None] == kx) & (ky[:, None] == ky)
+    if poltype == "parity":
+        pol = basis.pol
+        selection = selection & (pol[:, None] == pol)
+        a = selection * ((1 - pol) * gamma.conjugate() + pol * gamma) * 0.25
+        b = selection * ((1 - pol) * gamma.conjugate() - pol * gamma) * 0.25
+        return a, b
+    if poltype == "helicity":
+        pol = 2 * basis.pol - 1
+        a = selection * (pol[:, None] * pol * gamma[:, None].conjugate() + gamma) * 0.25
+        b = selection * (pol[:, None] * pol * gamma[:, None].conjugate() - gamma) * 0.25
+        return a, b
+    raise ValueError(f"invalid poltype: '{poltype}'")

src/treams/_tmatrix.py

+import warnings
+
+import numpy as np
+
+import treams._operators as op
+import treams.special as sc
+from treams import config
+from treams._core import CylindricalWaveBasis as CWB
+from treams._core import PhysicsArray
+from treams._core import PlaneWaveBasisByComp as PWBC
+from treams._core import PlaneWaveBasisByUnitVector as PWBUV
+from treams._core import SphericalWaveBasis as SWB
+from treams._material import Material
+from treams.coeffs import mie, mie_cyl
+from treams.util import AnnotationError
+
+
+class _Interaction:
+    def __init__(self):
+        self._obj = self._objtype = None
+
+    def __get__(self, obj, objtype=None):
+        self._obj = obj
+        self._objtype = objtype
+        return self
+
+    def __call__(self):
+        basis = self._obj.basis
+        return (
+            np.eye(self._obj.shape[-1]) - self._obj @ op.Expand(basis, "singular").inv
+        )
+
+    def solve(self):
+        return np.linalg.solve(self(), self._obj)
+
+
+class _LatticeInteraction:
+    def __init__(self):
+        self._obj = self._objtype = None
+
+    def __get__(self, obj, objtype=None):
+        self._obj = obj
+        self._objtype = objtype
+        return self
+
+    def __call__(self, lattice, kpar):
+        return np.eye(self._obj.shape[-1]) - self._obj @ op.ExpandLattice(
+            lattice=lattice, kpar=kpar
+        )
+
+    def solve(self, lattice, kpar):
+        return np.linalg.solve(self(lattice, kpar), self._obj)
+
+
+class TMatrix(PhysicsArray):
+    """T-matrix with a spherical basis.
+
+    The T-matrix is square relating incident (regular) fields
+    :func:`treams.special.vsw_rA` (helical polarizations) or
+    :func:`treams.special.vsw_rN` and :func:`treams.special.vsw_rM` (parity
+    polarizations) to the corresponding scattered fields :func:`treams.special.vsw_A` or
+    :func:`treams.special.vsw_N` and :func:`treams.special.vsw_M`. The modes themselves
+    are defined in :attr:`basis`, the polarization type in :attr:`poltype`. Also, the
+    wave number :attr:`k0` and, if not vacuum, the material :attr:`material` are
+    specified.
+
+    Args:
+        arr (float or complex, array-like): T-matrix itself.
+        k0 (float): Wave number in vacuum.
+        basis (SphericalWaveBasis, optional): Basis definition.
+        material (Material, optional): Embedding material, defaults to vacuum.
+        poltype (str, optional): Polarization type (:ref:`params:Polarizations`).
+        lattice (Lattice, optional): Lattice definition. If specified the T-Matrix is
+            assumed to be periodically repeated in the defined lattice.
+        kpar (list, optional): Phase factor for the periodic T-Matrix.
+    """
+
+    interaction = _Interaction()
+    latticeinteraction = _LatticeInteraction()
+
+    def _check(self):
+        """Fill in default values or raise errors for missing attributes."""
+        super()._check()
+        shape = np.shape(self)
+        if len(shape) != 2 or shape[0] != shape[1]:
+            raise AnnotationError(f"invalid shape: '{shape}'")
+        if not isinstance(self.k0, (int, float, np.floating, np.integer)):
+            raise AnnotationError("invalid k0")
+        if self.poltype is None:
+            self.poltype = config.POLTYPE
+        if self.poltype not in ("parity", "helicity"):
+            raise AnnotationError("invalid poltype")
+        modetype = self.modetype
+        if modetype is None or (
+            modetype[0] in (None, "singular") and modetype[1] in (None, "regular")
+        ):
+            self.modetype = ("singular", "regular")
+        else:
+            raise AnnotationError("invalid modetype")
+        if self.basis is None:
+            self.basis = SWB.default(SWB.defaultlmax(shape[0]))
+        if self.material is None:
+            self.material = Material()
+
+    @property
+    def ks(self):
+        """Wave numbers (in medium).
+
+        The wave numbers for both polarizations.
+        """
+        return self.material.ks(self.k0)
+
+    @classmethod
+    def sphere(cls, lmax, k0, radii, materials):
+        """T-Matrix of a (multi-layered) sphere.
+
+        Construct the T-matrix of the given order and material for a sphere. The object
+        can also consist of multiple concentric spherical shells with an arbitrary
+        number of layers. The calculation is always done in helicity basis.
+
+        Args:
+            lmax (int): Positive integer for the maximum degree of the T-matrix.
+            k0 (float): Wave number in vacuum.
+            radii (float or array): Radii from inside to outside of the sphere. For a
+                simple sphere the radius can be given as a single number, for a multi-
+                layered sphere it is a list of increasing radii for all shells.
+            material (list[Material]): The material parameters from the inside to the
+                outside. The last material in the list specifies the embedding medium.
+
+        Returns:
+            TMatrix
+        """
+        materials = [Material(m) for m in materials]
+        radii = np.atleast_1d(radii)
+        if radii.size != len(materials) - 1:
+            raise ValueError("incompatible lengths of radii and materials")
+        dim = SWB.defaultdim(lmax)
+        tmat = np.zeros((dim, dim), np.complex128)
+        for l in range(1, lmax + 1):  # noqa: E741
+            miecoeffs = mie(l, k0 * radii, *zip(*materials))
+            pos = SWB.defaultdim(l - 1)
+            for i in range(2 * l + 1):
+                tmat[
+                    pos + 2 * i : pos + 2 * i + 2, pos + 2 * i : pos + 2 * i + 2
+                ] = miecoeffs[::-1, ::-1]
+        return cls(tmat, k0=k0, basis=SWB.default(lmax), material=materials[-1])
+
+    @classmethod
+    def cluster(cls, tmats, positions):
+        r"""Block-diagonal T-matrix of multiple objects.
+
+        Construct the initial block-diagonal T-matrix for a cluster of objects. The
+        T-matrices in the list are placed together into a block-diagonal matrix and the
+        complete (local) basis is defined based on the individual T-matrices and their
+        bases together with the defined positions. In mathematical terms the matrix
+
+        .. math::
+
+            \begin{pmatrix}
+                T_0 & 0 & \dots & 0 \\
+                0 & T_1 & \ddots & \vdots \\
+                \vdots & \ddots & \ddots & 0 \\
+                0 & \dots & 0 & T_{N-1} \\
+            \end{pmatrix}
+
+        is created from the list of T-matrices :math:`(T_0, \dots, T_{N-1})`. Only
+        T-matrices of the same wave number, embedding material, and polarization type
+        can be combined.
+
+        Args:
+            tmats (Sequence): List of T-matrices.
+            positions (array): The positions of all individual objects in the cluster.
+
+        Returns:
+            TMatrix
+        """
+        for tm in tmats:
+            if not tm.basis.isglobal:
+                raise ValueError("global basis required")
+        positions = np.array(positions)
+        if len(tmats) < positions.shape[0]:
+            warnings.warn("specified more positions than T-matrices")
+        elif len(tmats) > positions.shape[0]:
+            raise ValueError(
+                f"'{len(tmats)}' T-matrices "
+                f"but only '{positions.shape[0]}' positions given"
+            )
+        mat = tmats[0].material
+        k0 = tmats[0].k0
+        poltype = tmats[0].poltype
+        modes = [], [], []
+        pidx = []
+        dim = sum(tmat.shape[0] for tmat in tmats)
+        tres = np.zeros((dim, dim), complex)
+        i = 0
+        for j, tm in enumerate(tmats):
+            if tm.material != mat:
+                raise ValueError(f"incompatible materials: '{mat}' and '{tm.material}'")
+            if tm.k0 != k0:
+                raise ValueError(f"incompatible k0: '{k0}' and '{tm.k0}'")
+            if tm.poltype != poltype:
+                raise ValueError(f"incompatible modetypes: '{poltype}', '{tm.poltype}'")
+            dim = tm.shape[0]
+            for m, n in zip(modes, tm.basis.lms):
+                m.extend(list(n))
+            pidx += [j] * dim
+            tres[i : i + dim, i : i + dim] = tm
+            i += dim
+        basis = SWB(zip(pidx, *modes), positions)
+        return cls(tres, k0=k0, material=mat, basis=basis, poltype=poltype)
+
+    @property
+    def isglobal(self):
+        """Test if a T-matrix is global.
+
+        A T-matrix is considered global, when its basis refers to only a single point
+        and it is not placed periodically in a lattice.
+        """
+        return self.basis.isglobal and self.lattice is None and self.kpar is None
+
+    @property
+    def xs_ext_avg(self):
+        r"""Rotation and polarization averaged extinction cross section.
+
+        The average is calculated as
+
+        .. math::
+
+            \langle \sigma_\mathrm{ext} \rangle
+            = -2 \pi \sum_{slm} \frac{\Re(T_{slm,slm})}{k_s^2}
+
+        where :math:`k_s` is the wave number in the embedding medium for the
+        polarization :math:`s`. It is only implemented for global T-matrices.
+
+        Returns:
+            float or complex
+        """
+        if not self.isglobal or not self.material.isreal:
+            raise NotImplementedError
+        if not self.material.ischiral:
+            k = self.ks[0]
+            res = -2 * np.pi * self.trace().real / (k * k)
+        else:
+            res = 0
+            modetype = self.modetype
+            del self.modetype
+            diag = np.diag(self)
+            self.modetype = modetype
+            for pol in [0, 1]:
+                choice = self.basis.pol == pol
+                k = self.ks[pol]
+                res += -2 * np.pi * diag[choice].sum().real / (k * k)
+        if res.imag == 0:
+            return res.real
+        return res
+
+    @property
+    def xs_sca_avg(self):
+        r"""Rotation and polarization averaged scattering cross section.
+
+        The average is calculated as
+
+        .. math::
+
+            \langle \sigma_\mathrm{sca} \rangle
+            = 2 \pi \sum_{slm} \sum_{s'l'm'}
+            \frac{|T_{slm,s'l'm'}|^2}{k_s^2}
+
+        where :math:`k_s` is the wave number in the embedding medium for the
+        polarization :math:`s`. It is only implemented for global T-matrices.
+
+        Returns:
+            float or complex
+        """
+        if not self.isglobal or not self.material.isreal:
+            raise NotImplementedError
+        if not self.material.ischiral:
+            ks = self.ks[0]
+        else:
+            ks = self.ks[self.basis.pol, None]
+        re, im = self.real, self.imag
+        res = 2 * np.pi * np.sum((re * re + im * im) / (ks * ks))
+        return res.real
+
+    @property
+    def cd(self):
+        r"""Circular dichroism (CD).
+
+        The CD is calculated as
+
+        .. math::
+
+            CD
+            = \frac{\langle \sigma_\mathrm{abs} \rangle_+
+            - \langle \sigma_\mathrm{abs} \rangle_-}
+            {\langle \sigma_\mathrm{abs} \rangle_+
+            + \langle \sigma_\mathrm{abs} \rangle_-}
+
+        where :math:`\langle \sigma \rangle_s` is the rotationally averaged absorption
+        cross section under the illumination with the polarization :math:`s`.
+        It is only implemented for global T-matrices.
+
+        Returns:
+            float or complex
+        """
+        if not (self.isglobal and self.poltype == "helicity" and self.material.isreal):
+            raise NotImplementedError
+        sel = np.array(self.basis.pol, bool)
+        re, im = self.real, self.imag
+        plus = -np.sum(re[sel[:, None] & sel]) / (self.ks[1] * self.ks[1])
+        re_part = re[:, sel] / self.ks[self.basis.pol, None]
+        im_part = im[:, sel] / self.ks[self.basis.pol, None]
+        plus -= np.sum(re_part * re_part + im_part * im_part)
+        sel = ~sel
+        minus = -np.sum(re[sel[:, None] & sel]) / (self.ks[0] * self.ks[0])
+        re_part = re[:, sel] / self.ks[self.basis.pol, None]
+        im_part = im[:, sel] / self.ks[self.basis.pol, None]
+        minus -= np.sum(re_part * re_part + im_part * im_part)
+        return np.real((plus - minus) / (plus + minus))
+
+    @property
+    def chi(self):
+        """Electromagnetic chirality.
+
+        Only implemented for global T-matrices.
+
+        Returns:
+            float
+        """
+        if not (self.isglobal and self.poltype == "helicity"):
+            raise NotImplementedError
+        sel = self.basis.pol == 0, self.basis.pol == 1
+        spp = np.linalg.svd(self[np.ix_(sel[1], sel[1])], compute_uv=False)
+        spm = np.linalg.svd(self[np.ix_(sel[1], sel[0])], compute_uv=False)
+        smp = np.linalg.svd(self[np.ix_(sel[0], sel[1])], compute_uv=False)
+        smm = np.linalg.svd(self[np.ix_(sel[0], sel[0])], compute_uv=False)
+        plus = np.concatenate((np.asarray(spp), np.asarray(spm)))
+        minus = np.concatenate((np.asarray(smm), np.asarray(smp)))
+        return np.linalg.norm(plus - minus) / np.sqrt(np.sum(np.power(np.abs(self), 2)))
+
+    @property
+    def db(self):
+        """Duality breaking.
+
+        Only implemented for global T-matrices.
+
+        Returns:
+            float
+        """
+        if not (self.isglobal and self.poltype == "helicity"):
+            raise NotImplementedError
+        sel = self.basis.pol == 0, self.basis.pol == 1
+        tpm = np.asarray(self[np.ix_(sel[1], sel[0])])
+        tmp = np.asarray(self[np.ix_(sel[0], sel[1])])
+        return np.sum(
+            tpm.real * tpm.real
+            + tpm.imag * tpm.imag
+            + tmp.real * tmp.real
+            + tmp.imag * tmp.imag
+        ) / (np.sum(np.power(np.abs(self), 2)))
+
+    def xs(self, illu, flux=0.5):
+        r"""Scattering and extinction cross section.
+
+        Possible for all T-matrices (global and local) in non-absorbing embedding. The
+        values are calculated by
+
+        .. math::
+
+            \sigma_\mathrm{sca}
+            = \frac{1}{2 I}
+            a_{slm}^\ast T_{s'l'm',slm}^\ast k_{s'}^{-2} C_{s'l'm',s''l''m''}^{(1)}
+            T_{s''l''m'',s'''l'''m'''} a_{s'''l'''m'''} \\
+            \sigma_\mathrm{ext}
+            = \frac{1}{2 I}
+            a_{slm}^\ast k_s^{-2} T_{slm,s'l'm'} a_{s'l'm'}
+
+        where :math:`a_{slm}` are the expansion coefficients of the illumination,
+        :math:`T` is the T-matrix, :math:`C^{(1)}` is the (regular) translation
+        matrix and :math:`k_s` are the wave numbers in the medium. All repeated indices
+        are summed over. The incoming flux is :math:`I`.
+
+        Args:
+            illu (complex, array): Illumination coefficients
+            flux (optional): Ingoing flux corresponding to the illumination. Used for
+                the result's normalization. The flux is given in units of
+                :math:`\frac{\text{V}^2}{{l^2}} \frac{1}{Z_0 Z}` where :math:`l` is the
+                unit of length used in the wave number (and positions). A plane wave
+                has the flux `0.5` in this normalization, which is used as default.
+
+        Returns:
+            tuple[float]
+        """
+        if not self.material.isreal:
+            raise NotImplementedError
+        illu = PhysicsArray(illu)
+        illu_basis = illu.basis
+        illu_basis = illu_basis[-2] if isinstance(illu_basis, tuple) else illu_basis
+        if not isinstance(illu_basis, SWB):
+            illu = illu.expand(self.basis)
+        p = self @ illu
+        p_invksq = p * np.power(self.ks[self.basis.pol], -2)
+        del illu.modetype
+        return (
+            0.5 * np.real(p.conjugate().T @ p_invksq.expand(p.basis)) / flux,
+            -0.5 * np.real(illu.conjugate().T @ p_invksq) / flux,
+        )
+
+    def valid_points(self, grid, radii):
+        """Points on the grid where the expansion is valid.
+
+        The expansion of the electromagnetic field is valid outside of the
+        circumscribing spheres of each object. From a given set of coordinates mark
+        those that are outside of the given radii.
+
+        Args:
+            grid (array-like): Points to assess. The last dimension needs length three
+                and corresponds to the Cartesian coordinates.
+            radii (Sequence[float]): Radii of the circumscribing spheres. Each radius
+                corresponds to a position of the basis.
+
+        Returns:
+            array
+        """
+        grid = np.asarray(grid)
+        if grid.shape[-1] != 3:
+            raise ValueError("invalid grid")
+        if len(radii) != len(self.basis.positions):
+            raise ValueError("invalid length of 'radii'")
+        res = np.ones(grid.shape[:-1], bool)
+        for r, p in zip(radii, self.basis.positions):
+            res &= np.sum(np.power(grid - p, 2), axis=-1) > r * r
+        return res
+
+    def __getitem__(self, key):
+        if isinstance(key, SWB):
+            key = np.array([self.basis.index(i) for i in key])
+            key = (key[:, None], key)
+        return super().__getitem__(key)
+
+
+class TMatrixC(PhysicsArray):
+    """T-matrix with a cylindrical basis.
+
+    The T-matrix is square relating incident (regular) fields
+    :func:`treams.special.vcw_rA` (helical polarizations) or
+    :func:`treams.special.vcw_rN` and :func:`treams.special.vcw_rM` (parity
+    polarizations) to the corresponding scattered fields :func:`treams.special.vcw_A` or
+    :func:`treams.special.vcw_N` and :func:`treams.special.vcw_M`. The modes themselves
+    are defined in :attr:`basis`, the polarization type in :attr:`poltype`. Also, the
+    wave number :attr:`k0` and, if not vacuum, the material :attr:`material` are
+    specified.
+
+    Args:
+        arr (float or complex, array-like): T-matrix itself.
+        k0 (float): Wave number in vacuum.
+        basis (SphericalWaveBasis, optional): Basis definition.
+        material (Material, optional): Embedding material, defaults to vacuum.
+        poltype (str, optional): Polarization type (:ref:`params:Polarizations`).
+        lattice (Lattice, optional): Lattice definition. If specified the T-Matrix is
+            assumed to be periodically repeated in the defined lattice.
+        kpar (list, optional): Phase factor for the periodic T-Matrix.
+    """
+
+    interaction = _Interaction()
+    latticeinteraction = _LatticeInteraction()
+
+    def _check(self):
+        """Fill in default values or raise errors for missing attributes."""
+        super()._check()
+        shape = np.shape(self)
+        if len(shape) != 2 or shape[0] != shape[1]:
+            raise AnnotationError(f"invalid shape: '{shape}'")
+        if not isinstance(self.k0, (int, float, np.floating, np.integer)):
+            raise AnnotationError("invalid k0")
+        if self.poltype is None:
+            self.poltype = config.POLTYPE
+        if self.poltype not in ("parity", "helicity"):
+            raise AnnotationError("invalid poltype")
+        modetype = self.modetype
+        if modetype is None or (
+            modetype[0] in (None, "singular") and modetype[1] in (None, "regular")
+        ):
+            self.modetype = ("singular", "regular")
+        else:
+            raise AnnotationError("invalid modetype")
+        if self.basis is None:
+            self.basis = CWB.default([0], CWB.defaultmmax(shape[0]))
+        if self.material is None:
+            self.material = Material()
+
+    @property
+    def ks(self):
+        """Wave numbers (in medium).
+
+        The wave numbers for both polarizations.
+        """
+        return self.material.ks(self.k0)
+
+    @property
+    def krhos(self):
+        r"""Radial part of the wave.
+
+        Calculate :math:`\sqrt{k^2 - k_z^2}`, where :math:`k` is the wave number in the
+        medium for each illumination
+
+        Returns:
+            Sequence[complex]
+        """
+        return self.material.krhos(self.k0, self.basis.kz, self.basis.pol)
+
+    @classmethod
+    def cylinder(cls, kzs, mmax, k0, radii, materials):
+        """T-Matrix of a (multi-layered) cylinder.
+
+        Construct the T-matrix of the given order and material for an infinitely
+        extended cylinder. The object can also consist of multiple concentric
+        cylindrical shells with an arbitrary number of layers. The calculation is always
+        done in helicity basis.
+
+        Args:
+            kzs (float, array_like): Z component of the cylindrical wave.
+            mmax (int): Positive integer for the maximum order of the T-matrix.
+            k0 (float): Wave number in vacuum.
+            radii (float or array): Radii from inside to outside of the cylinder. For a
+                simple cylinder the radius can be given as a single number, for a multi-
+                layered cylinder it is a list of increasing radii for all shells.
+            material (list[Material]): The material parameters from the inside to the
+                outside. The last material in the list specifies the embedding medium.
+
+        Returns:
+            T-Matrix object
+        """
+        materials = [Material(m) for m in materials]
+        kzs = np.atleast_1d(kzs)
+        radii = np.atleast_1d(radii)
+        if radii.size != len(materials) - 1:
+            raise ValueError("incompatible lengths of radii and materials")
+        dim = CWB.defaultdim(len(kzs), mmax)
+        tmat = np.zeros((dim, dim), np.complex128)
+        idx = 0
+        for kz in kzs:
+            for m in range(-mmax, mmax + 1):
+                miecoeffs = mie_cyl(kz, m, k0, radii, *zip(*materials))
+                tmat[idx : idx + 2, idx : idx + 2] = miecoeffs[::-1, ::-1]
+                idx += 2
+        return cls(tmat, k0=k0, basis=CWB.default(kzs, mmax), material=materials[-1])
+
+    @classmethod
+    def cluster(cls, tmats, positions):
+        r"""Block-diagonal T-matrix of multiple objects.
+
+        Construct the initial block-diagonal T-matrix for a cluster of objects. The
+        T-matrices in the list are placed together into a block-diagonal matrix and the
+        complete (local) basis is defined based on the individual T-matrices and their
+        bases together with the defined positions. In mathematical terms the matrix
+
+        .. math::
+
+            \begin{pmatrix}
+                T_0 & 0 & \dots & 0 \\
+                0 & T_1 & \ddots & \vdots \\
+                \vdots & \ddots & \ddots & 0 \\
+                0 & \dots & 0 & T_{N-1} \\
+            \end{pmatrix}
+
+        is created from the list of T-matrices :math:`(T_0, \dots, T_{N-1})`. Only
+        T-matrices of the same wave number, embedding material, and polarization type
+        can be combined.
+
+        Args:
+            tmats (Sequence): List of T-matrices.
+            positions (array): The positions of all individual objects in the cluster.
+
+        Returns:
+            TMatrix
+        """
+        for tm in tmats:
+            if not tm.basis.isglobal:
+                raise ValueError("global basis required")
+        positions = np.array(positions)
+        if len(tmats) < positions.shape[0]:
+            warnings.warn("specified more positions than T-matrices")
+        elif len(tmats) > positions.shape[0]:
+            raise ValueError(
+                f"'{len(tmats)}' T-matrices "
+                f"but only '{positions.shape[0]}' positions given"
+            )
+        mat = tmats[0].material
+        k0 = tmats[0].k0
+        poltype = tmats[0].poltype
+        modes = [], [], []
+        pidx = []
+        dim = sum(tmat.shape[0] for tmat in tmats)
+        tres = np.zeros((dim, dim), complex)
+        i = 0
+        for j, tm in enumerate(tmats):
+            if tm.material != mat:
+                raise ValueError(f"incompatible materials: '{mat}' and '{tm.material}'")
+            if tm.k0 != k0:
+                raise ValueError(f"incompatible k0: '{k0}' and '{tm.k0}'")
+            if tm.poltype != poltype:
+                raise ValueError(f"incompatible modetypes: '{poltype}', '{tm.poltype}'")
+            dim = tm.shape[0]
+            for m, n in zip(modes, tm.basis.zms):
+                m.extend(list(n))
+            pidx += [j] * dim
+            tres[i : i + dim, i : i + dim] = tm
+            i += dim
+        basis = CWB(zip(pidx, *modes), positions)
+        return cls(tres, k0=k0, material=mat, basis=basis, poltype=poltype)
+
+    @classmethod
+    def from_array(cls, tm, basis, *, eta=0):
+        """1d array of spherical T-matrices."""
+        return cls(
+            (tm @ op.Expand(basis).inv).expandlattice(basis=basis, eta=eta),
+            lattice=tm.lattice,
+            kpar=tm.kpar,
+        )
+
+    @property
+    def xw_ext_avg(self):
+        r"""Rotation and polarization averaged extinction cross width.
+
+        The average is calculated as
+
+        .. math::
+
+            \langle \lambda_\mathrm{ext} \rangle
+            = -\frac{2 \pi}{n_{k_z}} \sum_{sk_zm} \frac{\Re(T_{sk_zm,sk_zm})}{k_s}
+
+        where :math:`k_s` is the wave number in the embedding medium for the
+        polarization :math:`s` and :math:`n_{k_z}` is the number of wave components
+        :math:`k_z` included in the T-matrix. The average is taken over all given
+        z-components of the wave vector and rotations around the z-axis. It is only
+        implemented for global T-matrices.
+
+        Returns:
+            float or complex
+        """
+        if not self.isglobal or not self.material.isreal:
+            raise NotImplementedError
+        nk = np.unique(self.basis.kz).size
+        if not self.material.ischiral:
+            res = -2 * np.real(np.trace(self)) / (self.ks[0] * nk)
+        else:
+            res = 0
+            modetype = self.modetype
+            del self.modetype
+            diag = np.diag(self)
+            self.modetype = modetype
+            for pol in [0, 1]:
+                choice = self.basis.pol == pol
+                res += -2 * np.real(diag[choice].sum()) / (self.ks[pol] * nk)
+        if res.imag == 0:
+            return res.real
+        return res
+
+    @property
+    def xw_sca_avg(self):
+        r"""Rotation and polarization averaged scattering cross width.
+
+        The average is calculated as
+
+        .. math::
+
+            \langle \lambda_\mathrm{sca} \rangle
+            = \frac{2 \pi}{n_{k_z}} \sum_{sk_zm} \sum_{s'{k_z}'m'}
+            \frac{|T_{sk_zm,s'{k_z}'m'}|^2}{k_s}
+
+        where :math:`k_s` is the wave number in the embedding medium for the
+        polarization :math:`s`. and :math:`n_{k_z}` is the number of wave components
+        :math:`k_z` included in the T-matrix. The average is taken over all given
+        z-components of the wave vector and rotations around the z-axis. It is only
+        implemented for global T-matrices.
+
+        Returns:
+            float or complex
+        """
+        if not self.isglobal or not self.material.isreal:
+            raise NotImplementedError
+        if not self.material.ischiral:
+            ks = self.ks[0]
+        else:
+            ks = self.ks[self.basis.pol, None]
+        re, im = self.real, self.imag
+        nk = np.unique(self.basis.kz).size
+        res = 2 * np.sum((re * re + im * im) / (ks * nk))
+        return res.real
+
+    @property
+    def isglobal(self):
+        """Test if a T-matrix is global.
+
+        A T-matrix is considered global, when its basis refers to only a single point
+        and it is not placed periodically in a lattice.
+        """
+        return self.basis.isglobal and self.lattice is None and self.kpar is None
+
+    def xw(self, illu, flux=0.5):
+        r"""Scattering and extinction cross width.
+
+        Possible for all T-matrices (global and local) in non-absorbing embedding. The
+        values are calculated by
+
+        .. math::
+
+            \lambda_\mathrm{sca}
+            = \frac{1}{2 I}
+            a_{sk_zm}^\ast T_{s'{k_z}'m',sk_zm}^\ast k_{s'}^{-2}
+            C_{s'l'm',s''l''m''}^{(1)}
+            T_{s''l''m'',s'''l'''m'''} a_{s'''l'''m'''} \\
+            \sigma_\mathrm{ext}
+            = \frac{1}{2 I}
+            a_{slm}^\ast k_s^{-2} T_{slm,s'l'm'} a_{s'l'm'}
+
+        where :math:`a_{slm}` are the expansion coefficients of the illumination,
+        :math:`T` is the T-matrix, :math:`C^{(1)}` is the (regular) translation
+        matrix and :math:`k_s` are the wave numbers in the medium. All repeated indices
+        are summed over. The incoming flux is :math:`I`.
+
+        Args:
+            illu (complex, array): Illumination coefficients
+            flux (optional): Ingoing flux corresponding to the illumination. Used for
+                the result's normalization. The flux is given in units of
+                :math:`\frac{\text{V}^2}{{l^2}} \frac{1}{Z_0 Z}` where :math:`l` is the
+                unit of length used in the wave number (and positions). A plane wave
+                has the flux `0.5` in this normalization, which is used as default.
+
+        Returns:
+            tuple[float]
+        """
+        if not self.material.isreal:
+            raise NotImplementedError
+        illu = PhysicsArray(illu)
+        illu_basis = illu.basis
+        illu_basis = illu_basis[-2] if isinstance(illu_basis, tuple) else illu_basis
+        if not isinstance(illu_basis, CWB):
+            illu = illu.expand(self.basis) @ illu
+        p = self @ illu
+        m = self.expand() / self.ks[self.basis.pol]
+        del illu.modetype
+        return (
+            2 * np.real(p.conjugate().T @ (m @ p)) / flux,
+            -2 * np.real(illu.conjugate().T @ (p / self.ks[self.basis.pol])) / flux,
+        )
+
+    def valid_points(self, grid, radii):
+        grid = np.asarray(grid)
+        if grid.shape[-1] not in (2, 3):
+            raise ValueError("invalid grid")
+        if len(radii) != len(self.basis.positions):
+            raise ValueError("invalid length of 'radii'")
+        res = np.ones(grid.shape[:-1], bool)
+        for r, p in zip(radii, self.basis.positions):
+            res &= np.sum(np.power(grid[..., :2] - p[:2], 2), axis=-1) > r * r
+        return res
+
+    def __getitem__(self, key):
+        if isinstance(key, CWB):
+            key = np.array([self.basis.index(i) for i in key])
+            key = (key[:, None], key)
+        return super().__getitem__(key)
+
+
+def _plane_wave_partial(
+    kpar, pol, *, k0=None, basis=None, material=None, modetype=None, poltype=None,
+):
+    if basis is None:
+        basis = PWBC.default([kpar])
+    if pol in (0, -1):
+        pol = [1, 0]
+    elif pol == 1:
+        pol = [0, 1]
+    elif len(pol) == 3:
+        modetype = "up" if modetype is None else modetype
+        if modetype not in ("up", "down"):
+            raise ValueError(f"invalid 'modetype': {modetype}")
+        kzs = Material(material).kzs(k0, *kpar, [0, 1])
+        poltype = config.POLTYPE if poltype is None else poltype
+        if poltype == "parity":
+            pol = [
+                -sc.vpw_M(*kpar, kzs[0], 0, 0, 0) @ pol,
+                sc.vpw_N(*kpar, kzs[1], 0, 0, 0) @ pol,
+            ]
+        elif poltype == "helicity":
+            pol = sc.vpw_A(*kpar, kzs[::-1], 0, 0, 0, [1, 0]) @ pol
+        else:
+            raise ValueError(f"invalid 'poltype': {poltype}")
+    res = [pol[x[2]] * (np.abs(np.array(kpar) - x[:2]) < 1e-14).all() for x in basis]
+    return PhysicsArray(
+        res, basis=basis, k0=k0, material=material, modetype=modetype, poltype=poltype,
+    )
+
+
+def _plane_wave(
+    kvec, pol, *, k0=None, basis=None, material=None, modetype=None, poltype=None,
+):
+    if basis is None:
+        basis = PWBUV.default([kvec])
+    norm = np.sqrt(np.sum(np.power(kvec, 2)))
+    qvec = kvec / norm
+    if pol in (0, -1):
+        pol = [1, 0]
+    elif pol == 1:
+        pol = [0, 1]
+    elif len(pol) == 3:
+        if None not in (k0, material):
+            kvec = Material(material).ks(k0) * qvec[:, None]
+        else:
+            kvec = qvec
+        poltype = config.POLTYPE if poltype is None else poltype
+        if poltype == "parity":
+            pol = [
+                -sc.vpw_M(*kvec[:, 0], 0, 0, 0) @ pol,
+                sc.vpw_N(*kvec[:, 1], 0, 0, 0) @ pol,
+            ]
+        elif poltype == "helicity":
+            pol = sc.vpw_A(*kvec, 0, 0, 0, [1, 0]) @ pol
+        else:
+            raise ValueError(f"invalid 'poltype': {poltype}")
+    res = [pol[x[3]] * (np.abs(qvec - x[:3]) < 1e-14).all() for x in basis]
+    return PhysicsArray(
+        res, basis=basis, k0=k0, material=material, modetype=modetype, poltype=poltype,
+    )
+
+
+def plane_wave(
+    kvec, pol, *, k0=None, basis=None, material=None, modetype=None, poltype=None,
+):
+    """Array describing a plane wave.
+
+    Args:
+        kvec (Sequence): Wave vector.
+        pol (int or Sequence): Polarization index (see
+            :ref:`params:Polarizations`) to have a unit amplitude wave of the
+            corresponding wave, two values to specify the amplitude for each
+            polarization, or three values in a sequence to specify the Cartesian
+            electric field components. In the latter case, if the electric field has
+            longitudinal components they are neglected.
+        basis (PlaneWaveBasis, optional): Basis definition.
+        k0 (float, optional): Wave number in vacuum.
+        material (Material, optional): Material definition.
+        modetype (str, optional): Mode type (see :ref:`params:Mode types`).
+        poltype (str, optional): Polarization type (see
+            :ref:`polarization:Polarizations`).
+    """
+    if len(kvec) == 2:
+        return _plane_wave_partial(
+            kvec,
+            pol,
+            k0=k0,
+            basis=basis,
+            material=material,
+            modetype=modetype,
+            poltype=poltype,
+        )
+    if len(kvec) == 3:
+        return _plane_wave(
+            kvec,
+            pol,
+            k0=k0,
+            basis=basis,
+            material=material,
+            modetype=modetype,
+            poltype=poltype,
+        )
+    raise ValueError(f"invalid length of 'kvec': {len(kvec)}")
+
+
+def plane_wave_angle(theta, phi, pol, **kwargs):
+    qvec = [np.sin(theta) * np.cos(phi), np.sin(theta) * np.sin(phi), np.cos(theta)]
+    return plane_wave(qvec, pol, **kwargs)
+
+
+def spherical_wave(
+    l,  # noqa: E741
+    m,
+    pol,
+    *,
+    k0=None,
+    basis=None,
+    material=None,
+    modetype=None,
+    poltype=None,
+):
+    if basis is None:
+        basis = SWB.default(l)
+    if not basis.isglobal:
+        raise ValueError("basis must be global")
+    res = [0] * len(basis)
+    res[basis.index((0, l, m, pol))] = 1
+    return PhysicsArray(
+        res, basis=basis, k0=k0, material=material, modetype=modetype, poltype=poltype
+    )
+
+
+def cylindrical_wave(
+    kz, m, pol, *, k0=None, basis=None, material=None, modetype=None, poltype=None,
+):
+    if basis is None:
+        basis = CWB.default([kz], abs(m))
+    if not basis.isglobal:
+        raise ValueError("basis must be global")
+    res = [0] * len(basis)
+    res[basis.index((0, kz, m, pol))] = 1
+    return PhysicsArray(
+        res, basis=basis, k0=k0, material=material, modetype=modetype, poltype=poltype
+    )

src/treams/ebcm.py

+import itertools
+
+import numpy as np
+import scipy.integrate
+
+import treams.special as sc
+
+
+def _j_real(l0, l1, m, r, dr, pol0, pol1, ks, k, zs, z):
+    def fun(theta):
+        return np.real(
+            np.sin(theta)
+            * ((2 * pol0 - 1) * z + (2 * pol1 - 1) * zs)
+            * np.dot(
+                [r(theta), -dr(theta), 0],
+                np.cross(
+                    sc.vsw_rA(l0, -m, ks[pol0] * r(theta), theta, 0, pol0),
+                    sc.vsw_A(l1, m, k[pol1] * r(theta), theta, 0, pol1),
+                ),
+            )
+        )
+
+    return fun
+
+
+def _j_imag(l0, l1, m, r, dr, pol0, pol1, ks, k, zs, z):
+    def fun(theta):
+        return np.imag(
+            np.sin(theta)
+            * ((2 * pol0 - 1) * z + (2 * pol1 - 1) * zs)
+            * np.dot(
+                [r(theta), -dr(theta), 0],
+                np.cross(
+                    sc.vsw_rA(l0, -m, ks[pol0] * r(theta), theta, 0, pol0),
+                    sc.vsw_A(l1, m, k[pol1] * r(theta), theta, 0, pol1),
+                ),
+            )
+        )
+
+    return fun
+
+
+def _rj_real(l0, l1, m, r, dr, pol0, pol1, ks, k, zs, z):
+    def fun(theta):
+        return np.real(
+            np.sin(theta)
+            * ((2 * pol0 - 1) * z + (2 * pol1 - 1) * zs)
+            * np.dot(
+                [r(theta), -dr(theta), 0],
+                np.cross(
+                    sc.vsw_rA(l0, -m, ks[pol0] * r(theta), theta, 0, pol0),
+                    sc.vsw_rA(l1, m, k[pol1] * r(theta), theta, 0, pol1),
+                ),
+            )
+        )
+
+    return fun
+
+
+def _rj_imag(l0, l1, m, r, dr, pol0, pol1, ks, k, zs, z):
+    def fun(theta):
+        return np.imag(
+            np.sin(theta)
+            * ((2 * pol0 - 1) * z + (2 * pol1 - 1) * zs)
+            * np.dot(
+                [r(theta), -dr(theta), 0],
+                np.cross(
+                    sc.vsw_rA(l0, -m, ks[pol0] * r(theta), theta, 0, pol0),
+                    sc.vsw_rA(l1, m, k[pol1] * r(theta), theta, 0, pol1),
+                ),
+            )
+        )
+
+    return fun
+
+
+def qmat(r, dr, ks, zs, out, in_=None, singular=True):
+    l_out, m_out, pol_out = out
+    in_ = out if in_ is None else in_
+    l_in, m_in, pol_in = in_
+    if singular:
+        fr = _j_real
+        fi = _j_imag
+    else:
+        fr = _rj_real
+        fi = _rj_imag
+    size_out = len(l_out)
+    size_in = len(l_in)
+    res = np.zeros((size_out, size_in), complex)
+    for i, j in itertools.product(range(size_out), range(size_in)):
+        if m_out[i] != m_in[j]:
+            continue
+        res[i, j] = (
+            scipy.integrate.quad(
+                fr(l_out[i], l_in[j], m_out[i], r, dr, pol_out[i], pol_in[j], *ks, *zs),
+                0,
+                np.pi,
+            )[0]
+            + 1j
+            * scipy.integrate.quad(
+                fi(l_out[i], l_in[j], m_out[i], r, dr, pol_out[i], pol_in[j], *ks, *zs),
+                0,
+                np.pi,
+            )[0]
+        )
+    return res

src/treams/io.py

+"""Loading and storing data.
+
+Most functions rely on at least one of the external packages `h5py` or `gmsh`.
+"""
+
+import sys
+import uuid as _uuid
+from importlib.metadata import version
+
+import numpy as np
+
+import treams
+
+try:
+    import h5py
+except ImportError:
+    h5py = None
+
+LENGTHS = {
+    "ym": 1e-24,
+    "zm": 1e-21,
+    "am": 1e-18,
+    "fm": 1e-15,
+    "pm": 1e-12,
+    "nm": 1e-9,
+    "um": 1e-6,
+    "µm": 1e-6,
+    "mm": 1e-3,
+    "cm": 1e-2,
+    "dm": 1e-1,
+    "m": 1,
+    "dam": 1e1,
+    "hm": 1e2,
+    "km": 1e3,
+    "Mm": 1e6,
+    "Gm": 1e9,
+    "Tm": 1e12,
+    "Pm": 1e15,
+    "Em": 1e18,
+    "Zm": 1e21,
+    "Ym": 1e24,
+}
+
+INVLENGTHS = {
+    r"ym^{-1}": 1e24,
+    r"zm^{-1}": 1e21,
+    r"am^{-1}": 1e18,
+    r"fm^{-1}": 1e15,
+    r"pm^{-1}": 1e12,
+    r"nm^{-1}": 1e9,
+    r"um^{-1}": 1e6,
+    r"µm^{-1}": 1e6,
+    r"mm^{-1}": 1e3,
+    r"cm^{-1}": 1e2,
+    r"dm^{-1}": 1e1,
+    r"m^{-1}": 1,
+    r"dam^{-1}": 1e-1,
+    r"hm^{-1}": 1e-2,
+    r"km^{-1}": 1e-3,
+    r"Mm^{-1}": 1e-6,
+    r"Gm^{-1}": 1e-9,
+    r"Tm^{-1}": 1e-12,
+    r"Pm^{-1}": 1e-15,
+    r"Em^{-1}": 1e-18,
+    r"Zm^{-1}": 1e-21,
+    r"Ym^{-1}": 1e-24,
+}
+
+FREQUENCIES = {
+    "yHz": 1e-24,
+    "zHz": 1e-21,
+    "aHz": 1e-18,
+    "fHz": 1e-15,
+    "pHz": 1e-12,
+    "nHz": 1e-9,
+    "uHz": 1e-6,
+    "µHz": 1e-6,
+    "mHz": 1e-3,
+    "cHz": 1e-2,
+    "dHz": 1e-1,
+    "s": 1,
+    "daHz": 1e1,
+    "hHz": 1e2,
+    "kHz": 1e3,
+    "MHz": 1e6,
+    "GHz": 1e9,
+    "THz": 1e12,
+    "PHz": 1e15,
+    "EHz": 1e18,
+    "ZHz": 1e21,
+    "YHz": 1e24,
+    r"ys^{-1}": 1e24,
+    r"zs^{-1}": 1e21,
+    r"as^{-1}": 1e18,
+    r"fs^{-1}": 1e15,
+    r"ps^{-1}": 1e12,
+    r"ns^{-1}": 1e9,
+    r"us^{-1}": 1e6,
+    r"µs^{-1}": 1e6,
+    r"ms^{-1}": 1e3,
+    r"cs^{-1}": 1e2,
+    r"ds^{-1}": 1e1,
+    r"s^{-1}": 1,
+    r"das^{-1}": 1e-1,
+    r"hs^{-1}": 1e-2,
+    r"ks^{-1}": 1e-3,
+    r"Ms^{-1}": 1e-6,
+    r"Gs^{-1}": 1e-9,
+    r"Ts^{-1}": 1e-12,
+    r"Ps^{-1}": 1e-15,
+    r"Es^{-1}": 1e-18,
+    r"Zs^{-1}": 1e-21,
+    r"Ys^{-1}": 1e-24,
+}
+
+
+def mesh_spheres(radii, positions, model, meshsize=None, meshsize_boundary=None):
+    """Generate a mesh of multiple spheres.
+
+    This function facilitates generating a mesh for a cluster spheres using gmsh. It
+    requires the package `gmsh` to be installed.
+
+    Examples:
+        >>> import gmsh
+        >>> import treams.io
+        >>> gmsh.initialize()
+        >>> gmsh.model.add("spheres")
+        >>> treams.io.mesh_spheres([1, 2], [[0, 0, 2], [0, 0, -2]], gmsh.model)
+        <class 'gmsh.model'>
+        >>> gmsh.write("spheres.msh")  # doctest: +SKIP
+        >>> gmsh.finalize()  # doctest: +SKIP
+
+    Args:
+        radii (float, array_like): Radii of the spheres.
+        positions (float, (N, 3)-array): Positions of the spheres.
+        model (gmsh.model): Gmsh model to modify.
+        meshsize (float, optional): Mesh size, if None a fifth of the largest radius is
+            used.
+        meshsize (float, optional): Mesh size of the surfaces, if left empty it is set
+            equal to the general mesh size.
+
+    Returns:
+        gmsh.model
+    """
+    if meshsize is None:
+        meshsize = np.max(radii) * 0.2
+    if meshsize_boundary is None:
+        meshsize_boundary = meshsize
+
+    spheres = []
+    for i, (radius, position) in enumerate(zip(radii, positions)):
+        tag = i + 1
+        model.occ.addSphere(*position, radius, tag)
+        spheres.append((3, tag))
+
+    model.occ.synchronize()
+    for _, tag in spheres:
+        model.addPhysicalGroup(3, [tag], tag)
+        # Add surfaces for other mesh formats like stl, ...
+        model.addPhysicalGroup(2, [tag], tag)
+
+    model.mesh.setSize(model.getEntities(0), meshsize)
+    model.mesh.setSize(
+        model.getBoundary(spheres, False, False, True), meshsize_boundary
+    )
+    return model
+
+
+def _translate_polarizations(pols, poltype=None):
+    """Translate the polarization index into words.
+
+    The indices 0 and 1 are translated to "negative" and "positive", respectively, when
+    helicity modes are chosen. For parity modes they are translated to "magnetic" and
+    "electric".
+
+    Args:
+        pols (int, array_like): Array of indices 0 and 1.
+        poltype (str, optional): Polarization type (:ref:`polarizations.Polarizations`).
+
+    Returns:
+        list[str]
+    """
+    poltype = treams.config.POLTYPE if poltype is None else poltype
+    if poltype == "helicity":
+        names = ["negative", "positive"]
+    elif poltype == "parity":
+        names = ["magnetic", "electric"]
+    else:
+        raise ValueError("unrecognized poltype")
+    return [names[i] for i in pols]
+
+
+def _translate_polarizations_inv(pols):
+    """Translate the polarization into indices.
+
+    This function is the inverse of :func:`treams.io._translate_polarizations`. The
+    words "negative" and "minus" are translated to 0 and the words "positive" and "plus"
+    are translated to 1, if helicity modes are chosen. For parity modes, modes
+    "magnetic" or "te" are translated to 0 and the modes "electric" or "tm" to 1.
+
+    Args:
+        pols (string, array_like): Array of strings
+
+    Returns:
+        tuple[list[int], str]
+    """
+    helicity = {"plus": 1, "positive": 1, "minus": 0, "negative": 0}
+    parity = {"te": 0, "magnetic": 0, "tm": 1, "electric": 1, "M": 0, "N": 1}
+    if pols[0].decode() in helicity:
+        dct = helicity
+        poltype = "helicity"
+    elif pols[0].decode() in parity:
+        dct = parity
+        poltype = "parity"
+    else:
+        raise ValueError(f"unrecognized polarization '{pols[0].decode()}'")
+    return [dct[i.decode()] for i in pols], poltype
+
+
+def _remove_leading_ones(shape):
+    while len(shape) > 0 and shape[0] == 1:
+        shape = shape[1:]
+    return shape
+
+
+def _collapse_dims(arr):
+    for i in range(arr.ndim):
+        test = arr[(slice(None),) * i + (slice(1),)]
+        if np.all(arr == test):
+            arr = test
+    return arr.reshape(_remove_leading_ones(arr.shape))
+
+
+def save_hdf5(
+    h5file,
+    tms,
+    name="",
+    description="",
+    keywords="",
+    embedding_group=None,
+    embedding_name="",
+    embedding_description="",
+    embedding_keywords="",
+    uuid=None,
+    uuid_version=4,
+    lunit="nm",
+):
+    """Save a set of T-matrices in a HDF5 file.
+
+    With an open and writeable datafile, this function stores the main parts of as
+    T-matrix in the file. It is left open for the user to add additional metadata.
+
+    Args:
+        h5file (h5py.Group): A HDF5 file opened with h5py.
+        tms (TMatrix, array_like): Array of T-matrix instances.
+        name (str): Name to add to the file as attribute.
+        description (str): Description to add to file as attribute.
+        keywords (str): Keywords to add to file as attribute.
+        embedding_group (h5py.Group, optional): Group object for the embedding material,
+            defaults to "/materials/embedding/".
+        embedding_name (string, optional): Name of the embedding material.
+        embedding_description (string, optional): Description of the embedding material.
+        embedding_keywords (string, optional): Keywords for the embedding material.
+        uuid (bytes, optional): UUID of the file, a new one is created if omitted.
+        uuid_version (int, optional): UUID version.
+        lunit (string, optional): Length unit used for the positions and (as
+            inverse) for the wave number.
+    """
+    tms_arr = np.array(tms)
+    if tms_arr.dtype == object:
+        raise ValueError("can only save T-matrices of the same size")
+    tms_obj = np.empty(tms_arr.shape[:-2], object)
+    tms_obj[:] = tms
+    tm = tms_obj.flat[0]
+    basis = tm.basis
+    poltype = tm.poltype
+
+    k0s = np.zeros((tms_arr.shape)[:-2])
+    epsilon = np.zeros((tms_arr.shape)[:-2], complex)
+    mu = np.zeros((tms_arr.shape)[:-2], complex)
+    if poltype == "helicity":
+        kappa = np.zeros((tms_arr.shape)[:-2], complex)
+
+    for i, tm in enumerate(tms_obj.flat):
+        if poltype != tm.poltype or basis != tm.basis:
+            raise ValueError(
+                "incompatible T-matrices: mixed poltypes or different bases"
+            )
+        k0s.flat[i] = tm.k0
+        epsilon.flat[i] = tm.material.epsilon
+        mu.flat[i] = tm.material.mu
+        if poltype == "helicity":
+            kappa.flat[i] = tm.material.kappa
+
+    h5file["tmatrix"] = tms_arr
+    if uuid is None:
+        h5file["uuid"] = np.void(_uuid.uuid4().bytes)
+        h5file["uuid"].attrs["version"] = 4
+    else:
+        h5file["uuid"] = uuid
+        h5file["uuid"].attrs["version"] = uuid_version
+
+    _name_descr_kw(h5file, name, description, keywords)
+    h5file["angular_vacuum_wavenumber"] = _collapse_dims(k0s)
+    h5file["angular_vacuum_wavenumber"].attrs["unit"] = lunit + r"^{-1}"
+    h5file["modes/l"] = basis.l
+    h5file["modes/m"] = basis.m
+    h5file["modes/polarization"] = _translate_polarizations(basis.pol, poltype)
+    if any(basis.pidx != 0):
+        h5file["modes/pidx"] = basis.pidx
+    if not np.array_equiv(basis.positions, [[0, 0, 0]]):
+        h5file["modes/positions"] = basis.positions
+        h5file["modes/positions"].attrs["unit"] = lunit
+    if embedding_group is None:
+        embedding_group = h5file.create_group("materials/embedding")
+    embedding_group["relative_permittivity"] = _collapse_dims(epsilon)
+    embedding_group["relative_permeability"] = _collapse_dims(mu)
+    if poltype == "helicity":
+        embedding_group["chirality"] = _collapse_dims(kappa)
+    _name_descr_kw(
+        embedding_group, embedding_name, embedding_description, embedding_keywords
+    )
+    h5file["embedding"] = h5py.SoftLink(embedding_group.name)
+
+    h5file.attrs["created_with"] = (
+        f"python={sys.version.split()[0]},"
+        f"h5py={version('h5py')},"
+        f"treams={version('treams')}"
+    )
+    h5file.attrs["storage_format_version"] = "0.0.1-4-g1266244"
+
+
+def _name_descr_kw(fobj, name, description="", keywords=""):
+    for key, val in [
+        ("name", name),
+        ("description", description),
+        ("keywords", keywords),
+    ]:
+        val = str(val)
+        if val != "":
+            fobj.attrs[key] = val
+
+
+def _convert_to_k0(x, xtype, xunit, k0unit=r"nm^{-1}"):
+    c = 299792458.0
+    k0unit = INVLENGTHS[k0unit]
+    if xtype == "frequency":
+        xunit = FREQUENCIES[xunit]
+        return 2 * np.pi * x / c * (xunit / k0unit)
+    if xtype == "angular_frequency":
+        xunit = FREQUENCIES[xunit]
+        return x / c * (xunit / k0unit)
+    if xtype == "angular_vacuum_wavelength":
+        xunit = INVLENGTHS[xunit]
+        return x * (xunit / k0unit)
+    if xtype == "angular_vacuum_wavenumber":
+        xunit = LENGTHS[xunit]
+        return 2 * np.pi / (x * xunit * k0unit)
+    raise ValueError(f"unrecognized frequency/wavenumber/wavelength type: {xtype}")
+
+
+def load_hdf5(filename, lunit="nm"):
+    """Load a T-matrix stored in a HDF4 file.
+
+    Args:
+        filename (str or h5py.Group): Name of the h5py file or a handle to a h5py group.
+        lunit (str, optional): Unit of length to be used in the T-matrices.
+
+    Returns:
+        np.ndarray[TMatrix]
+    """
+    if isinstance(filename, h5py.Group):
+        return _load_hdf5(filename, lunit)
+    with h5py.File(filename, "r") as f:
+        return _load_hdf5(f, lunit)
+
+
+def _load_hdf5(h5file, lunit=None):
+    for freq_type in (
+        "frequency",
+        "angular_frequency",
+        "vacuum_wavelength",
+        "vacuum_wavenumber",
+        "angular_vacuum_wavenumber",
+    ):
+        if freq_type in h5file:
+            ld_freq = h5file[freq_type][()]
+            break
+    else:
+        raise ValueError("no definition of frequency found")
+    if "modes/positions" in h5file:
+        k0unit = h5file["modes/positions"].attrs.get("unit", lunit) + r"^{-1}"
+    else:
+        k0unit = lunit + r"^{-1}"
+    tms = h5file["tmatrix"][...]
+    k0s = _convert_to_k0(ld_freq, freq_type, h5file[freq_type].attrs["unit"], k0unit)
+
+    epsilon = h5file.get("embedding/relative_permittivity", np.array(None))[()]
+    mu = h5file.get("embedding/relative_permeability", np.array(None))[()]
+    if epsilon is None is mu:
+        n = h5file.get("embedding/refractive_index", np.array(1))[...]
+        z = h5file.get("embedding/relative_impedance", 1 / n)[...]
+        epsilon = n / z
+        mu = n * z
+    epsilon = 1 if epsilon is None else epsilon
+    mu = 1 if mu is None else mu
+
+    kappa = z = h5file.get("embedding/chirality_parameter", np.array(0))[...]
+
+    positions = h5file.get("modes/positions", np.zeros((1, 3)))[...]
+
+    l_inc = h5file.get("modes/l", np.array(None))[...]
+    l_inc = h5file.get("modes/l_incident", l_inc)[()]
+    m_inc = h5file.get("modes/m", np.array(None))[...]
+    m_inc = h5file.get("modes/m_incident", m_inc)[()]
+    pol_inc = h5file.get("modes/polarization", np.array(None))[...]
+    pol_inc = h5file.get("modes/polarization_incident", pol_inc)[()]
+    l_sca = h5file.get("modes/l", np.array(None))[...]
+    l_sca = h5file.get("modes/l_scattered", l_sca)[()]
+    m_sca = h5file.get("modes/m", np.array(None))[...]
+    m_sca = h5file.get("modes/m_scattered", m_sca)[()]
+    pol_sca = h5file.get("modes/polarization", np.array(None))[...]
+    pol_sca = h5file.get("modes/polarization_scattered", pol_sca)[()]
+
+    if any(x is None for x in (l_inc, l_sca, m_inc, m_sca, pol_inc, pol_sca)):
+        raise ValueError("mode definition missing")
+
+    pol_inc, poltype = _translate_polarizations_inv(pol_inc)
+    pol_sca, poltype_sca = _translate_polarizations_inv(pol_sca)
+
+    if poltype_sca != poltype:
+        raise ValueError("different modetypes")
+
+    pidx_inc = h5file.get("modes/position_index", np.zeros_like(l_inc))[()]
+    pidx_inc = h5file.get("modes/positions_index_scattered", pidx_inc)[()]
+    pidx_sca = h5file.get("modes/position_index", np.zeros_like(l_sca))[()]
+    pidx_sca = h5file.get("modes/positions_index_scattered", pidx_sca)[()]
+
+    shape = tms.shape[:-2]
+    k0s = np.broadcast_to(k0s, shape)
+    epsilon = np.broadcast_to(epsilon, shape)
+    mu = np.broadcast_to(mu, shape)
+    kappa = np.broadcast_to(kappa, shape)
+
+    basis_inc = treams.SphericalWaveBasis(
+        zip(pidx_inc, l_inc, m_inc, pol_inc), positions
+    )
+    basis_sca = treams.SphericalWaveBasis(
+        zip(pidx_sca, l_sca, m_sca, pol_sca), positions
+    )
+    basis = basis_inc | basis_sca
+
+    ix_inc = [basis.index(b) for b in basis_inc]
+    ix_sca = [[basis.index(b)] for b in basis_sca]
+
+    res = np.empty(shape, object)
+    for i in np.ndindex(*shape):
+        res[i] = treams.TMatrix(
+            np.zeros((len(basis),) * 2, complex),
+            k0=k0s[i],
+            basis=basis,
+            poltype=poltype,
+            material=treams.Material(epsilon[i], mu[i], kappa[i]),
+        )
+        res[i][ix_sca, ix_inc] = tms[i]
+    if not res.shape:
+        res = res.item()
+    return res

src/treams/lattice/__init__.py

+r"""Lattice sums.
+
+.. currentmodule:: treams.lattice
+
+Calculates the lattice sums of the forms
+
+.. math::
+   D_{lm}(k, \boldsymbol k_\parallel, \boldsymbol r, \Lambda_d)
+   = \sideset{}{'}{\sum_{\boldsymbol R \in \Lambda_d}}
+   h_l^{(1)}(k|\boldsymbol r + \boldsymbol R|)
+   Y_{lm}
+   (\theta_{-\boldsymbol r - \boldsymbol R}, \varphi_{-\boldsymbol r - \boldsymbol R})
+   \mathrm e^{\mathrm i \boldsymbol k_\parallel \boldsymbol R}
+
+and
+
+.. math::
+   D_{m}(k, \boldsymbol k_\parallel, \boldsymbol r, \Lambda_d)
+   = \sideset{}{'}{\sum_{\boldsymbol R \in \Lambda_d}}
+   H_m^{(1)}(k|\boldsymbol r + \boldsymbol R|)
+   \mathrm e^{\mathrm i m \varphi_{-\boldsymbol r - \boldsymbol R}}
+   \mathrm e^{\mathrm i \boldsymbol k_\parallel \boldsymbol R}
+
+that arise when translating and summing the spherical and cylindrical solutions
+of the Helmholtz equation in a periodic arrangement. These sums have a notoriously slow
+convergence and at least for lattices with a dimension :math:`d > 1` it is not advisable
+to use the direct approach. Fortunately, it is possible to convert them to exponentially
+convergent series, which are implemented here. For details on the math, see the
+references below.
+
+The lattice of dimension :math:`d \leq 3` (:math:`d \leq 2` in the second case) is
+denoted with :math:`\Lambda_d` and consists of all vectors
+:math:`\boldsymbol R = \sum_{i=1}^{d} n_i \boldsymbol a_i` with :math:`\boldsymbol a_i`
+being basis vectors of the lattice and :math:`n_i \in \mathbb Z`. For :math:`d = 2` the
+lattice is in the `z = 0` plane and for :math:`d = 1` it is along the z-axis. The vector
+:math:`\boldsymbol r` is arbitrary but for best convergence it should be reduced to the
+Wigner-Seitz cell of the lattice. The summation excludes the point
+:math:`\boldsymbol R + \boldsymbol r = 0` if it exists, which is indicated by the prime
+next to the summation sign.
+
+The wave in the lattice is defined by its -- possibly complex-valued -- wave number
+:math:`k` and the (real) components of the wave vector
+:math:`\boldsymbol k_\parallel \in \mathbb R^d` that are parallel to the lattice.
+
+The expressions include the (spherical) Hankel functions of first kind :math:`H_m^{(1)}`
+(:math:`h_l^{(1)}`) and the spherical harmonics :math:`Y_{lm}`. The angles
+:math:`\theta` and :math:`\varphi` are the polar and azimuthal angle, when expressing
+the points in spherical coordinates. In the first case, the degree is
+:math:`l \in \mathbb N_0` and the order is :math:`\mathbb Z \ni m \leq l`. In the second
+case :math:`m \in \mathbb Z`.
+
+
+Available functions
+===================
+
+Accelerated lattice summations
+------------------------------
+
+.. autosummary::
+   :toctree:
+
+   lsumcw1d
+   lsumcw1d_shift
+   lsumcw2d
+   lsumsw1d
+   lsumsw1d_shift
+   lsumsw2d
+   lsumsw2d_shift
+   lsumsw3d
+
+
+Direct summations
+-----------------
+
+The functions are almost only here for benchmarking and comparison.
+
+.. autosummary::
+   :toctree:
+
+   dsumcw1d
+   dsumcw1d_shift
+   dsumcw2d
+   dsumsw1d
+   dsumsw1d_shift
+   dsumsw2d
+   dsumsw2d_shift
+   dsumsw3d
+
+
+Miscellaneous functions
+-----------------------
+
+.. autosummary::
+   :toctree:
+
+   area
+   cube
+   cubeedge
+   diffr_orders_circle
+   reciprocal
+   volume
+
+
+Cython module
+-------------
+
+.. autosummary::
+   :toctree:
+   :template: cython-module
+
+   cython_lattice
+
+
+References
+==========
+
+* `[1] K. Kambe, Zeitschrift Für Naturforschung A 22, 3 (1967). <https://doi.org/10.1515/zna-1967-0305>`_
+* `[2] K. Kambe, Zeitschrift Für Naturforschung A 23, 9 (1968). <https://doi.org/10.1515/zna-1968-0908>`_
+* `[3] C. M. Linton, SIAM Rev. 52, 630 (2010). <https://doi.org/10.1137/09075130X>`_
+* `[4] D. Beutel el al., J. Opt. Soc. Am. B (2021). <https://doi.org/10.1364/JOSAB.419645>`_
+"""
+
+import numpy as np
+
+from treams.lattice import _misc
+from treams.lattice._gufuncs import *  # noqa: F403
+from treams.lattice._misc import cube, cubeedge  # noqa: F401
+
+
+def diffr_orders_circle(b, rmax):
+    """Diffraction orders in a circle.
+
+    Given the reciprocal lattice defined by the vectors that make up the rows of `b`,
+    return all diffraction orders within a circle of radius `rmax`.
+
+    Args:
+        b (float, (2, 2)-array): Reciprocal lattice vectors
+        rmax (float): Maximal radius
+
+    Returns:
+        float array
+    """
+    return _misc.diffr_orders_circle(np.array(b), rmax)
+
+
+def lsumsw(dim, l, m, k, kpar, a, r, eta, out=None, **kwargs):  # noqa: E741
+    if dim == 1:
+        return lsumsw1d_shift(l, m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
+    if dim == 2:
+        return lsumsw2d_shift(l, m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
+    if dim == 3:
+        return lsumsw3d(l, m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
+    raise ValueError(f"invalid dimension '{dim}'")
+
+
+def realsumsw(dim, l, m, k, kpar, a, r, eta, out=None, **kwargs):  # noqa: E741
+    if dim == 1:
+        return realsumsw1d_shift(  # noqa: F405
+            l, m, k, kpar, a, r, eta, out=out, **kwargs
+        )
+    if dim == 2:
+        return realsumsw2d_shift(  # noqa: F405
+            l, m, k, kpar, a, r, eta, out=out, **kwargs
+        )
+    if dim == 3:
+        return realsumsw3d(l, m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
+    raise ValueError(f"invalid dimension '{dim}'")
+
+
+def recsumsw(dim, l, m, k, kpar, a, r, eta, out=None, **kwargs):  # noqa: E741
+    if dim == 1:
+        return recsumsw1d_shift(  # noqa: F405
+            l, m, k, kpar, a, r, eta, out=out, **kwargs
+        )
+    if dim == 2:
+        return recsumsw2d_shift(  # noqa: F405
+            l, m, k, kpar, a, r, eta, out=out, **kwargs
+        )
+    if dim == 3:
+        return recsumsw3d(l, m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
+    raise ValueError(f"invalid dimension '{dim}'")
+
+
+def lsumcw(dim, m, k, kpar, a, r, eta, out=None, **kwargs):  # noqa: E741
+    if dim == 1:
+        return lsumcw1d_shift(m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
+    if dim == 2:
+        return lsumcw2d(m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
+    raise ValueError(f"invalid dimension '{dim}'")
+
+
+def realsumcw(dim, m, k, kpar, a, r, eta, out=None, **kwargs):
+    if dim == 1:
+        return realsumcw1d_shift(m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
+    if dim == 2:
+        return realsumcw2d(m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
+    raise ValueError(f"invalid dimension '{dim}'")
+
+
+def recsumcw(dim, m, k, kpar, a, r, eta, out=None, **kwargs):
+    if dim == 1:
+        return recsumcw1d_shift(m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
+    if dim == 2:
+        return recsumcw2d(m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
+    raise ValueError(f"invalid dimension '{dim}'")
+
+
+def dsumsw(dim, l, m, k, kpar, a, r, i, out=None, **kwargs):  # noqa: E741
+    if dim == 1:
+        return dsumsw1d_shift(l, m, k, kpar, a, r, i, out=out, **kwargs)  # noqa: F405
+    if dim == 2:
+        return dsumsw2d_shift(l, m, k, kpar, a, r, i, out=out, **kwargs)  # noqa: F405
+    if dim == 3:
+        return dsumsw3d(l, m, k, kpar, a, r, i, out=out, **kwargs)  # noqa: F405
+    raise ValueError(f"invalid dimension '{dim}'")
+
+
+def dsumcw(dim, m, k, kpar, a, r, i, out=None, **kwargs):
+    if dim == 1:
+        return dsumcw1d_shift(m, k, kpar, a, r, i, out=out, **kwargs)  # noqa: F405
+    if dim == 2:
+        return dsumcw2d(m, k, kpar, a, r, i, out=out, **kwargs)  # noqa: F405
+    raise ValueError(f"invalid dimension '{dim}'")

src/treams/lattice/cython_lattice.pxd

+ctypedef fused real_t:
+    double
+    long
+
+ctypedef fused number_t:
+    double complex
+    double
+
+cdef real_t area(real_t *a, real_t *b) nogil
+cdef long cube_next(long *r, long d, long n) nogil
+cdef long cubeedge_next(long *r, long d, long n) nogil
+cpdef diffr_orders_circle(real_t[:, :] b, double rmax)
+cpdef long ncube(long d, long n) nogil
+cpdef long nedge(long d, long n) nogil
+cdef void recvec2(double *a0, double *a1, double *b0, double *b1) nogil
+cdef void recvec3(double *a0, double *a1, double *a2, double *b0, double *b1, double *b2) nogil
+cdef real_t volume(real_t *a, real_t *b, real_t *c) nogil
+
+cpdef double complex lsumcw1d(long l, number_t k, double kpar, double a, double r, number_t eta) nogil
+cdef double complex lsumcw1d_shift(long l, number_t k, double kpar, double a, double *r, number_t eta) nogil
+cdef double complex lsumcw2d(long l, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
+cpdef double complex lsumsw1d(long l, number_t k, double kpar, double a, double r, number_t eta) nogil
+cdef double complex lsumsw1d_shift(long l, long m, number_t k, double kpar, double a, double *r, number_t eta) nogil
+cdef double complex lsumsw2d(long l, long m, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
+cdef double complex lsumsw2d_shift(long l, long m, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
+cdef double complex lsumsw3d(long l, long m, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
+
+cpdef double complex realsumcw1d(long l, number_t k, double kpar, double a, double r, number_t eta) nogil
+cdef double complex realsumcw1d_shift(long l, number_t k, double kpar, double a, double *r, number_t eta) nogil
+cdef double complex realsumcw2d(long l, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
+cpdef double complex realsumsw1d(long l, number_t k, double kpar, double a, double r, number_t eta) nogil
+cdef double complex realsumsw1d_shift(long l, long m, number_t k, double kpar, double a, double *r, number_t eta) nogil
+cdef double complex realsumsw2d(long l, long m, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
+cdef double complex realsumsw2d_shift(long l, long m, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
+cdef double complex realsumsw3d(long l, long m, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
+
+cpdef double complex recsumcw1d(long l, number_t k, double kpar, double a, double r, number_t eta) nogil
+cdef double complex recsumcw1d_shift(long l, number_t k, double kpar, double a, double *r, number_t eta) nogil
+cdef double complex recsumcw2d(long l, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
+cpdef double complex recsumsw1d(long l, number_t k, double kpar, double a, double r, number_t eta) nogil
+cdef double complex recsumsw1d_shift(long l, long m, number_t k, double kpar, double a, double *r, number_t eta) nogil
+cdef double complex recsumsw2d(long l, long m, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
+cdef double complex recsumsw2d_shift(long l, long m, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
+cdef double complex recsumsw3d(long l, long m, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
+
+cpdef double complex zero3d(double complex eta) nogil
+cpdef double complex zero2d(double complex eta) nogil
+
+cpdef double complex dsumcw1d(long l, number_t k, double kpar, double a, double r, long i) nogil
+cdef double complex dsumcw1d_shift(long l, number_t k, double kpar, double a, double *r, long i) nogil
+cdef double complex dsumcw2d(long l, number_t k, double *kpar, double *a, double *r, long i) nogil
+cpdef double complex dsumsw1d(long l, number_t k, double kpar, double a, double r, long i) nogil
+cdef double complex dsumsw1d_shift(long l, long m, number_t k, double kpar, double a, double *r, long i) nogil
+cdef double complex dsumsw2d(long l, long m, number_t k, double *kpar, double *a, double *r, long i) nogil
+cdef double complex dsumsw2d_shift(long l, long m, number_t k, double *kpar, double *a, double *r, long i) nogil
+cdef double complex dsumsw3d(long l, long m, number_t k, double *kpar, double *a, double *r, long i) nogil

src/treams/misc.py

+"""Miscellaneous functions.
+
+.. autosummary::
+   :toctree:
+
+   basischange
+   firstbrillouin1d
+   firstbrillouin2d
+   firstbrillouin3d
+   pickmodes
+   refractive_index
+   wave_vec_z
+"""
+
+import numpy as np
+
+import treams.lattice as la
+
+
+def refractive_index(epsilon=1, mu=1, kappa=0):
+    r"""Refractive index of a (chiral) medium.
+
+    The refractive indeces in a chiral medium :math:`\sqrt{\epsilon\mu} \mp \kappa` are
+    returned with the negative helicity result first.
+
+    Args:
+        epsilon (float or complex, array_like, optional): Relative permittivity,
+            defaults to 1.
+        mu (float or complex, array_like, optional): Relative permeability,
+            defaults to 1.
+        kappa (float or complex, array_like, optional): Chirality parameter,
+            defaults to 0.
+
+    Returns:
+        float or complex, (2,)-array
+
+    """
+    epsilon = np.array(epsilon)
+    n = np.sqrt(epsilon * mu)
+    res = np.stack((n - kappa, n + kappa), axis=-1)
+    res[np.imag(res) < 0] *= -1
+    return res
+
+
+def basischange(out, in_=None):
+    """Coefficients for the basis change between helicity and parity modes.
+
+    Args:
+        out (3- or 4-tuple of (M,)-arrays): Output modes, the last array is taken as
+            polarization.
+        in_ (3- or 4-tuple of (N,)-arrays, optional): Input modes, if none are given,
+            equal to the output modes
+
+    Returns:
+        float, ((M, N)-array
+    """
+    if in_ is None:
+        in_ = out
+    out = np.array([*zip(*out)])
+    in_ = np.array([*zip(*in_)])
+    res = np.zeros((out.shape[0], in_.shape[0]))
+    out = out[:, None, :]
+    sqhalf = np.sqrt(0.5)
+    equal = (out[:, :, :-1] == in_[:, :-1]).all(axis=-1)
+    minus = np.logical_and(out[:, :, -1] == in_[:, -1], in_[:, -1] == 0)
+    res[equal] = sqhalf
+    res[np.logical_and(equal, minus)] = -sqhalf
+    return res
+
+
+def pickmodes(out, in_):
+    """Coefficients to pick modes.
+
+    Args:
+        out (3- or 4-tuple of (M,)-arrays): Output modes, the last array is taken as
+            polarization.
+        in_ (3- or 4-tuple of (N,)-arrays, optional): Input modes, if none are given,
+            equal to the output modes
+
+    Returns:
+        float, ((M, N)-array
+    """
+    out = np.array([*zip(*out)])
+    in_ = np.array([*zip(*in_)])
+    return np.all(out[:, None, :] == in_, axis=-1)
+
+
+# def wave_vec_zs(kpars, ks):
+#     kpars = np.array(kpars)
+#     ks = np.array(ks)
+#     if ks.ndim == 0:
+#         nmat = npol = 1
+#     elif ks.ndim == 1:
+#         npol = ks.shape[0]
+#         nmat = 1
+#     elif ks.ndim == 2:
+#         nmat, npol = ks.shape
+#     else:
+#         raise ValueError("ks has invalid shape")
+#     res = np.zeros((nmat, kpars.shape[0], npol), np.complex128)
+#     for i, kpar in enumerate(kpars):
+#         res[:, i, :] = wave_vec_z(kpar[0], kpar[1], ks)
+#     return res
+
+
+def wave_vec_z(kx, ky, k):
+    r"""Z component of the wave vector with positive imaginary part.
+
+    The result is :math:`k_z = \sqrt{k^2 - k_x^2 - k_y^2}` with
+    :math:`\arg k_z \in \[ 0, \pi )`.
+
+    Args:
+        kx (float, array_like): X component of the wave vector
+        ky (float, array_like): Y component of the wave vector
+        k (float or complex, array_like): Wave number
+
+    Returns:
+        complex
+    """
+    kx = np.asarray(kx)
+    ky = np.asarray(ky)
+    k = np.asarray(k, complex)
+    res = np.sqrt(k * k - kx * kx - ky * ky)
+    if res.ndim == 0 and res.imag < 0:
+        res = -res
+    elif res.ndim > 0:
+        res[np.imag(res) < 0] *= -1
+    return res
+
+
+def firstbrillouin1d(kpar, b):
+    """Map wave vector to first Brillouin zone in 1D.
+
+    Reduce the 1d wave vector (actually just a number) to the first Brillouin zone, i.e.
+    the range `(-b/2, b/2]`
+
+    Args:
+        kpar (float64): (parallel) wave vector
+        b (float64): reciprocal lattice vector
+
+    Returns:
+        float64
+    """
+    kpar -= b * np.round(kpar / b)
+    if kpar > 0.5 * b:
+        kpar -= b
+    if kpar <= -0.5 * b:
+        kpar += b
+    return kpar
+
+
+def firstbrillouin2d(kpar, b, n=2):
+    """Map wave vector to first Brillouin zone in 2D.
+
+    The reduction to the first Brillouin zone is first approximated roughly. From this
+    approximated vector and its 8 neighbours, the shortest one is picked. As a
+    sufficient approximation is not guaranteed (especially for extreme geometries),
+    this process is iterated `n` times.
+
+    Args:
+        kpar (1d-array): parallel wave vector
+        b (2d-array): reciprocal lattice vectors
+        n (int): number of iterations
+
+    Returns:
+        (1d-array)
+    """
+    kparstart = kpar
+    b1 = b[0, :]
+    b2 = b[1, :]
+    normsq1 = b1 @ b1
+    normsq2 = b2 @ b2
+    normsqp = (b1 + b2) @ (b1 + b2)
+    normsqm = (b1 - b2) @ (b1 - b2)
+    if (
+        normsqp < normsq1 - 1e-14
+        or normsqp < normsq2 - 1e-14
+        or normsqm < normsq1 - 1e-14
+        or normsqm < normsq2 - 1e-14
+    ):
+        raise ValueError("Lattice vectors are not of minimal length")
+    # Rough estimate
+    kpar -= b1 * np.round((kpar @ b1) / normsq1)
+    kpar -= b2 * np.round((kpar @ b2) / normsq2)
+    # todo: precise
+    options = kpar + la.cube(2, 1) @ b
+    for option in options:
+        if option @ option < kpar @ kpar:
+            kpar = option
+    if n == 0 or np.array_equal(kpar, kparstart):
+        return kpar
+    return firstbrillouin2d(kpar, b, n - 1)
+
+
+def firstbrillouin3d(kpar, b, n=2):
+    """Map wave vector to first Brillouin zone in 3D.
+
+    The reduction to the first Brillouin zone is first approximated roughly. From this
+    approximated vector and its 26 neighbours, the shortest one is picked. As a
+    sufficient approximation is not guaranteed (especially for extreme geometries),
+    this process is iterated `n` times.
+
+    Args:
+        kpar (1d-array): parallel wave vector
+        b (2d-array): reciprocal lattice vectors
+        n (int): number of iterations
+
+    Returns:
+        (1d-array)
+    """
+    kparstart = kpar
+    b1 = b[0, :]
+    b2 = b[1, :]
+    b3 = b[2, :]
+    normsq1 = b1 @ b1
+    normsq2 = b2 @ b2
+    normsq3 = b3 @ b3
+    # todo: Minimal length
+    # Rough estimate
+    kpar -= b1 * np.round((kpar @ b1) / normsq1)
+    kpar -= b2 * np.round((kpar @ b2) / normsq2)
+    kpar -= b3 * np.round((kpar @ b3) / normsq3)
+    # todo: precise
+    options = kpar + la.cube(3, 1) @ b
+    for option in options:
+        if option @ option < kpar @ kpar:
+            kpar = option
+    if n == 0 or np.array_equal(kpar, kparstart):
+        return kpar
+    return firstbrillouin3d(kpar, b, n - 1)

src/treams/special/__init__.py

+r"""Special (mathematical) functions.
+
+.. currentmodule:: treams.special
+
+Special mathematical functions used in :mod:`treams`. Some functions are reexported from
+:py:mod:`scipy.special`. Most functions are available as Numpy universal functions
+(:py:class:`numpy.ufunc`) or as generalized universal functions
+(:ref:`c-api.generalized-ufuncs`).
+
+Available functions
+===================
+
+Bessel and Hankel functions, with their spherical counterparts, derivatives
+---------------------------------------------------------------------------
+
+.. autosummary::
+   :toctree:
+
+   hankel1_d
+   hankel2_d
+   jv_d
+   yv_d
+   spherical_hankel1
+   spherical_hankel2
+   spherical_hankel1_d
+   spherical_hankel2_d
+
+
+Those functions are just reexported from Scipy. So, one only needs to import this
+subpackage within treams.
+
++------------------------------------------------------------+-------------------------+
+| :py:data:`~scipy.special.hankel1`\(v, z[, out])            | Hankel function of the  |
+|                                                            | first kind.             |
++------------------------------------------------------------+-------------------------+
+| :py:data:`~scipy.special.hankel2`\(v, z[, out])            | Hankel function of the  |
+|                                                            | second kind.            |
++------------------------------------------------------------+-------------------------+
+| :py:data:`~scipy.special.jv`\(v, z[, out])                 | Bessel function of the  |
+|                                                            | first kind of real      |
+|                                                            | order and complex       |
+|                                                            | argument.               |
++------------------------------------------------------------+-------------------------+
+| :py:data:`~scipy.special.yv`\(v, z[, out])                 | Bessel function of the  |
+|                                                            | second kind of real     |
+|                                                            | order and complex       |
+|                                                            | argument.               |
++------------------------------------------------------------+-------------------------+
+| | :py:func:`spherical_jn <scipy.special.spherical_jn>`\(n, | Spherical Bessel        |
+|   z[, derivative])                                         | function of the first   |
+|                                                            | kind or its derivative. |
++------------------------------------------------------------+-------------------------+
+| | :py:func:`spherical_yn <scipy.special.spherical_yn>`\(n, | Spherical Bessel        |
+|   z[, derivative])                                         | function of the second  |
+|                                                            | kind or its derivative. |
++------------------------------------------------------------+-------------------------+
+
+Those functions just wrap Scipy functions with special optional arguments to be able to
+analogously access them like their non-spherical counterparts:
+
+.. autosummary::
+   :toctree:
+
+   spherical_jn_d
+   spherical_yn_d
+
+
+Scipy functions with enhanced domain
+------------------------------------
+
+.. autosummary::
+   :toctree:
+
+   sph_harm
+   lpmv
+
+
+Integrals for the Ewald summation
+---------------------------------
+
+.. autosummary::
+   :toctree:
+
+   incgamma
+   intkambe
+
+
+Wigner d- and Wigner D-matrix elements
+--------------------------------------
+
+.. autosummary::
+   :toctree:
+
+   wignersmalld
+   wignerd
+
+
+Wigner 3j-symbols
+-----------------
+
+.. autosummary::
+   :toctree:
+
+   wigner3j
+
+
+Vector wave functions
+--------------------------
+
+.. autosummary::
+   :toctree:
+
+   pi_fun
+   tau_fun
+
+Spherical waves and translation coefficients
+
+.. autosummary::
+   :toctree:
+
+   vsh_X
+   vsh_Y
+   vsh_Z
+   vsw_M
+   vsw_N
+   vsw_A
+   vsw_rM
+   vsw_rN
+   vsw_rA
+   tl_vsw_A
+   tl_vsw_B
+   tl_vsw_rA
+   tl_vsw_rB
+
+
+Cylindrical waves
+
+.. autosummary::
+   :toctree:
+
+   vcw_M
+   vcw_N
+   vcw_A
+   vcw_rM
+   vcw_rN
+   vcw_rA
+   tl_vcw
+   tl_vcw_r
+
+
+Plane waves
+
+.. autosummary::
+   :toctree:
+
+   vpw_M
+   vpw_N
+   vpw_A
+
+
+Coordinate system transformations
+---------------------------------
+
+.. autosummary::
+   :toctree:
+
+   car2cyl
+   car2sph
+   cyl2car
+   cyl2sph
+   sph2car
+   sph2cyl
+   vcar2cyl
+   vcar2sph
+   vcyl2car
+   vcyl2sph
+   vsph2car
+   vsph2cyl
+   car2pol
+   pol2car
+   vcar2pol
+   vpol2car
+
+
+Cython module
+-------------
+
+.. autosummary::
+   :toctree:
+   :template: cython-module
+
+   cython_special
+"""
+
+from scipy.special import (  # noqa: F401
+    hankel1,
+    hankel2,
+    jv,
+    spherical_jn,
+    spherical_yn,
+    yv,
+)
+
+from treams.special import _gufuncs, _ufuncs  # noqa: F401
+from treams.special._gufuncs import *  # noqa: F401, F403
+from treams.special._ufuncs import *  # noqa: F401, F403
+
+
+def spherical_jn_d(n, z):
+    """Derivative of the spherical Bessel function of the first kind.
+
+    This is simply a wrapper for `scipy.special.spherical_jn(n, z, True)`, see
+    :py:func:`scipy.special.spherical_jn`. It's here to have a consistent way of
+    calling the derivative of a (spherical) Bessel or Hankel function.
+
+    Args:
+        n (int, array_like): Order
+        z (float or complex, array_like): Argument
+
+    Returns:
+        float or complex
+
+    References:
+        - `DLMF 10.47 <https://dlmf.nist.gov/10.47>`_
+        - `DLMF 10.51 <https://dlmf.nist.gov/10.51>`_
+    """
+    return spherical_jn(n, z, True)
+
+
+def spherical_yn_d(n, z):
+    """Derivative of the spherical Bessel function of the second kind.
+
+    This is simply a wrapper for `scipy.special.spherical_yn(n, z, True)`, see
+    :py:func:`scipy.special.spherical_jn`. It's here to have a consistent way of
+    calling the derivative of a (spherical) Bessel or Hankel function.
+
+    Args:
+        n (int, array_like): Order
+        z (float or complex, array_like): Argument
+
+    Returns:
+        float or complex
+
+    References:
+        - `DLMF 10.47 <https://dlmf.nist.gov/10.47>`_
+        - `DLMF 10.51 <https://dlmf.nist.gov/10.51>`_
+    """
+    return spherical_yn(n, z, True)

src/treams/special/cython_special.pxd

+ctypedef fused number_t:
+    double complex
+    double
+
+
+cpdef double complex hankel1_d(double l, double complex z) nogil
+cpdef double complex hankel2_d(double l, double complex z) nogil
+cpdef number_t jv_d(double l, number_t z) nogil
+cpdef double complex spherical_hankel1(double n, number_t z) nogil
+cpdef double complex spherical_hankel1_d(double l, double complex z) nogil
+cpdef double complex spherical_hankel2(double n, number_t z) nogil
+cpdef double complex spherical_hankel2_d(double l, double complex z) nogil
+cpdef number_t yv_d(double l, number_t z) nogil
+
+
+cdef void car2cyl(double *input, double *output) nogil
+cdef void car2pol(double *input, double *output) nogil
+cdef void car2sph(double *input, double *output) nogil
+cdef void cyl2car(double *input, double *output) nogil
+cdef void cyl2sph(double *input, double *output) nogil
+cdef void pol2car(double *input, double *output) nogil
+cdef void sph2car(double *input, double *output) nogil
+cdef void sph2cyl(double *input, double *output) nogil
+
+cdef void vcar2cyl(number_t *iv, double *ip, number_t *ov) nogil
+cdef void vcar2pol(number_t *iv, double *ip, number_t *ov) nogil
+cdef void vcar2sph(number_t *iv, double *ip, number_t *ov) nogil
+cdef void vcyl2car(number_t *iv, double *ip, number_t *ov) nogil
+cdef void vcyl2sph(number_t *iv, double *ip, number_t *ov) nogil
+cdef void vpol2car(number_t *iv, double *ip, number_t *ov) nogil
+cdef void vsph2car(number_t *iv, double *ip, number_t *ov) nogil
+cdef void vsph2cyl(number_t *iv, double *ip, number_t *ov) nogil
+
+
+cpdef number_t incgamma(double l, number_t z) nogil
+cpdef number_t intkambe(long n, number_t z, number_t eta) nogil
+
+
+cpdef number_t lpmv(double m, double l, number_t z) nogil
+cpdef number_t pi_fun(double l, double m, number_t x) nogil
+cpdef double complex sph_harm(double m, double l, double phi, number_t theta) nogil
+cpdef number_t tau_fun(double l, double m, number_t x) nogil
+
+cpdef double complex tl_vcw(double kz1, long mu, double kz2, long m, double complex krr, double phi, double z) nogil
+cpdef double complex tl_vcw_r(double kz1, long mu, double kz2, long m, number_t krr, double phi, double z) nogil
+
+cpdef double complex _tl_vsw_helper(long l, long m, long lambda_, long mu, long p, long q) nogil
+
+cpdef double complex tl_vsw_A(long lambda_, long mu, long l, long m, double complex kr, number_t theta, double phi) nogil
+cpdef double complex tl_vsw_B(long lambda_, long mu, long l, long m, double complex kr, number_t theta, double phi) nogil
+cpdef double complex tl_vsw_rA(long lambda_, long mu, long l, long m, number_t kr, number_t theta, double phi) nogil
+cpdef double complex tl_vsw_rB(long lambda_, long mu, long l, long m, number_t kr, number_t theta, double phi) nogil
+
+cdef void vcw_A(double kz, long m, double complex krr, double phi, double z, double complex k, long pol, double complex *out) nogil
+cdef void vcw_M(double kz, long m, double complex krr, double phi, double z, double complex *out) nogil
+cdef void vcw_N(double kz, long m, double complex krr, double phi, double z, double complex k, double complex *out) nogil
+cdef void vcw_rA(double kz, long m, number_t krr, double phi, double z, double complex k, long pol, double complex *out) nogil
+cdef void vcw_rM(double kz, long m, number_t krr, double phi, double z, double complex *out, long i) nogil
+cdef void vcw_rN(double kz, long m, number_t krr, double phi, double z, double complex k, double complex *out) nogil
+
+cdef void vpw_A(number_t kx, number_t ky, number_t kz, double x, double y, double z, long pol, double complex *out) nogil
+cdef void vpw_M(number_t kx, number_t ky, number_t kz, double x, double y, double z, double complex *out) nogil
+cdef void vpw_N(number_t kx, number_t ky, number_t kz, double x, double y, double z, double complex *out) nogil
+
+cdef void vsh_X(long l, long m, number_t theta, double phi, double complex *out) nogil
+cdef void vsh_Y(long l, long m, number_t theta, double phi, double complex *out) nogil
+cdef void vsh_Z(long l, long m, number_t theta, double phi, double complex *out) nogil
+
+cdef void vsw_A(long l, long m, double complex kr, number_t theta, double phi, long pol, double complex *out) nogil
+cdef void vsw_M(long l, long m, double complex kr, number_t theta, double phi, double complex *out) nogil
+cdef void vsw_N(long l, long m, double complex kr, number_t theta, double phi, double complex *out) nogil
+cdef void vsw_rA(long l, long m, number_t kr, number_t theta, double phi, long pol, double complex *out) nogil
+cdef void vsw_rM(long l, long m, number_t kr, number_t theta, double phi, double complex *out) nogil
+cdef void vsw_rN(long l, long m, number_t kr, number_t theta, double phi, double complex *out) nogil
+
+
+cpdef number_t wignersmalld(long l, long m, long k, number_t theta) nogil
+cpdef double complex wignerd(long l, long m, long k, double phi, number_t theta, double psi) nogil
+
+
+cpdef double wigner3j(long j1, long j2, long j3, long m1, long m2, long m3) nogil

src/treams/util.py

+"""Utilities.
+
+Collection of functions to make :class:`AnnotatedArray` work. The implementation is
+aimed to be quite general, while still providing a solid basis for the rest of the code.
+"""
+
+import abc
+import collections
+import copy
+import itertools
+import warnings
+
+import numpy as np
+
+import treams._operators as _op
+
+__all__ = [
+    "AnnotatedArray",
+    "AnnotationDict",
+    "AnnotationError",
+    "AnnotationSequence",
+    "AnnotationWarning",
+    "implements",
+    "OrderedSet",
+]
+
+
+class OrderedSet(collections.abc.Sequence, collections.abc.Set):
+    """Ordered set.
+
+    A abstract base class that combines a sequence and set. In contrast to regular sets
+    it is expected that the equality comparison only returns `True` if all entries are
+    in the same order.
+    """
+
+    @abc.abstractmethod
+    def __eq__(self, other):
+        """Equality test."""
+        raise NotImplementedError
+
+
+class AnnotationWarning(UserWarning):
+    """Custom warning for Annotations.
+
+    By default the warning filter is set to 'always'.
+    """
+
+
+warnings.simplefilter("always", AnnotationWarning)
+
+
+class AnnotationError(Exception):
+    """Custom exception for Annotations."""
+
+
+class AnnotationDict(collections.abc.MutableMapping):
+    """Dictionary that notifies the user when overwriting keys.
+
+    Behaves mostly similar to regular dictionaries, except that when overwriting
+    existing keys a :class:`AnnotationWarning` is emmitted.
+
+    Examples:
+        An :class:`AnnotationDict` can be created from other mappings, from a list of
+        key-value pairs (note how a warning is emmitted for the duplicate key), and from
+        keyword arguments.
+
+        .. code-block:: python
+
+            >>> AnnotationDict({"a": 1, "b": 2})
+            AnnotationDict({'a': 1, 'b': 2})
+            >>> AnnotationDict([("a", 1), ("b", 2), ("a", 3)])
+            treams/util.py:74: AnnotationWarning: overwriting key 'a'
+            warnings.warn(f"overwriting key '{key}'", AnnotationWarning)
+            AnnotationDict({'a': 3, 'b': 2})
+            >>> AnnotationDict({"a": 1, "b": 2}, c=3)
+            AnnotationDict({'a': 1, 'b': 2, 'c': 3})
+
+    Warns:
+        AnnotationWarning
+    """
+
+    def __init__(self, items=(), /, **kwargs):
+        """Initialization."""
+        self._dct = {}
+        for i in (items.items() if hasattr(items, "items") else items, kwargs.items()):
+            for key, val in i:
+                self[key] = val
+
+    def __getitem__(self, key):
+        """Get a value by its key.
+
+        Args:
+            key (hashable): Key
+        """
+        return self._dct[key]
+
+    def __setitem__(self, key, val):
+        """Set item specified by key to the defined value.
+
+        When overwriting an existing key an :class:`AnnotationWarning` is emitted.
+        Avoid the warning by explicitly deleting the key first.
+
+        Args:
+            key (hashable): Key
+            val : Value
+
+        Warns:
+            AnnotationWarning
+        """
+        if key in self and self[key] != val:
+            warnings.warn(f"overwriting key '{key}'", AnnotationWarning)
+        self._dct[key] = val
+
+    def __delitem__(self, key):
+        """Delete the key."""
+        del self._dct[key]
+
+    def __iter__(self):
+        """Iterate over the keys.
+
+        Returns:
+            Iterator
+        """
+        return iter(self._dct)
+
+    def __len__(self):
+        """Number of keys contained.
+
+        Returns:
+            int
+        """
+        return len(self._dct)
+
+    def __repr__(self):
+        """String representation.
+
+        Returns:
+            str
+        """
+        return f"{self.__class__.__name__}({repr(self._dct)})"
+
+    def match(self, other):
+        """Compare the own keys to another dictionary.
+
+        This emits an :class:`AnnotationWarning` for each key that would be overwritten
+        by the given dictionary.
+
+        Args:
+            other (Mapping)
+
+        Returns:
+            None
+
+        Warns:
+            AnnotationWarning
+        """
+        for key, val in self.items():
+            if key in other and other[key] != val:
+                warnings.warn(f"incompatible key '{key}'", AnnotationWarning)
+
+
+class SequenceAsDict(collections.abc.MutableMapping):
+    def __get__(self, obj, objtype=None):
+        self._obj = obj
+        return self
+
+    def __set__(self, obj, dct):
+        self._obj = obj
+        for key, val in dct.items():
+            self[key] = val
+
+    def __getitem__(self, key):
+        res = tuple(i.get(key) for i in self._obj)
+        if all(i is None for i in res):
+            raise KeyError(key)
+        return res
+
+    def __setitem__(self, key, val):
+        val = (None,) * (len(self._obj) - len(val)) + val
+        for dct, value in zip(reversed(self._obj), reversed(val)):
+            if value is None:
+                dct.pop(key, None)
+            else:
+                dct[key] = value
+
+    def __delitem__(self, key):
+        found = False
+        for dct in self._obj:
+            if key in dct:
+                del dct[key]
+                found = True
+        if not found:
+            raise KeyError(key)
+
+    def __iter__(self):
+        return iter({key for dct in self._obj for key in dct})
+
+    def __len__(self):
+        return len({key for dct in self._obj for key in dct})
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({dict(i for i in self.items())})"
+
+
+class AnnotationSequence(collections.abc.Sequence):
+    """A Sequence of dictionaries.
+
+    This class is intended to work together with :class:`AnnotationDict`. It provides
+    convenience functions to interact with multiple of those dictionaries, which are
+    mainly used to keep track of the annotations made to each dimension of an
+    :class:`AnnotatedArray`. While the sequence itself is immutable the entries of each
+    dictionary is mutable.
+
+    Args:
+        *args: Items of the sequence
+        mapping (AnnotationDict): Type of mapping to use in the sequence.
+
+    Warns:
+        AnnotationWarning
+    """
+
+    as_dict = SequenceAsDict()
+
+    def __init__(self, *args, mapping=AnnotationDict):
+        """Initialization."""
+        self._ann = tuple(mapping(i) for i in args)
+
+    def __len__(self):
+        """Number of dictionaries in the sequence.
+
+        Returns:
+            int
+        """
+        return len(self._ann)
+
+    def __getitem__(self, key):
+        """Get an item or subsequence.
+
+        Indexing works with integers and slices like regular tuples. Additionally, it is
+        possible to get a copy of the object with `()`, or a new sequence of mappings in
+        a list (or other iterable) of integers.
+
+        Args:
+            key (iterable, slice, int)
+
+        Returns:
+            mapping
+        """
+        if isinstance(key, tuple) and key == ():
+            return copy.copy(self._ann)
+        if isinstance(key, slice):
+            return type(self)(*self._ann[key])
+        if isinstance(key, (int, np.integer)):
+            return self._ann[key]
+        res = []
+        for k in key:
+            if not isinstance(k, int):
+                raise TypeError(
+                    "sequence index must be integer, slice, list of integers, or '()'"
+                )
+            res.append(self[k])
+        return type(self)(*res)
+
+    def update(self, other):
+        """Update all mappings in the sequence at once.
+
+        The given sequence is aliged at the last entry and then updated pairwise.
+        Warnings for overwritten keys are extended by the information at which index
+        they occurred.
+
+        Args:
+            other (Sequence[Mapping]): Mappings to update with.
+
+        Warns:
+            AnnotationWarning
+        """
+        if len(other) > len(self):
+            warnings.warn(
+                f"argument of length {len(self)} given: "
+                f"ignore leading {len(other) - len(self)} entries",
+                AnnotationWarning,
+            )
+        for i, (dest, src) in enumerate(zip(reversed(self), reversed(other))):
+            with warnings.catch_warnings():
+                warnings.simplefilter("error", category=AnnotationWarning)
+                try:
+                    dest.update(src)
+                except AnnotationWarning as err:
+                    warnings.simplefilter("always", category=AnnotationWarning)
+                    warnings.warn(
+                        f"at index {len(self) - i - 1}: " + err.args[0],
+                        AnnotationWarning,
+                    )
+
+    def match(self, other):
+        """Match all mappings at once.
+
+        The given sequence is aliged at the last entry and then updated pairwise.
+        Warnings for overwritten keys are extended by the information at which index
+        they occurred, see also :func:`AnnotationDict.match`.
+
+        Args:
+            other (Sequence[Mappings]): Mappings to match.
+
+        Warns:
+            AnnotationWarning
+        """
+        for i, (dest, src) in enumerate(zip(reversed(self), reversed(other))):
+            with warnings.catch_warnings():
+                warnings.simplefilter("error", category=AnnotationWarning)
+                try:
+                    dest.match(src)
+                except AnnotationWarning as err:
+                    warnings.simplefilter("always", category=AnnotationWarning)
+                    warnings.warn(
+                        f"at dimension {len(self) + i}: " + err.args[0],
+                        AnnotationWarning,
+                    )
+
+    def __eq__(self, other):
+        """Equality test.
+
+        Two sequences of mappings are considered equal when they have equal length and
+        when they all mappings are equal.
+
+        Args:
+            other (Sequence[Mapping]): Mappings to compare with.
+
+        Returns:
+            bool
+        """
+        try:
+            lenother = len(other)
+        except TypeError:
+            return False
+        if len(self) != lenother:
+            return False
+        for a, b in zip(self, other):
+            if a != b:
+                return False
+        return True
+
+    def __repr__(self):
+        """String representation.
+
+        Returns:
+            str
+        """
+        if len(self) == 1:
+            return f"{type(self).__name__}({repr(self._ann[0])})"
+        return f"{type(self).__name__}{repr(self._ann)}"
+
+    def __add__(self, other):
+        return AnnotationSequence(*self, *other)
+
+    def __radd__(self, other):
+        return AnnotationSequence(*other, *self)
+
+
+def _cast_nparray(arr):
+    """Cast AnnotatedArray to numpy array."""
+    return np.asarray(arr) if isinstance(arr, AnnotatedArray) else arr
+
+
+def _cast_annarray(arr):
+    """Cast array to AnnotatedArray."""
+    return arr if isinstance(arr, np.generic) else AnnotatedArray(arr)
+
+
+def _parse_signature(signature, inputdims):
+    """Parse a ufunc signature based on the actual inputs.
+
+    The signature is matched with the input dimensions, to get the actual signature. It
+    is returned as two lists (inputs and outputs). Each list contains for each
+    individual input and output another list of the signature items.
+
+    Example:
+        >>> treams.util._parse_signature('(n?,k),(k,m?)->(n?,m?)', (2, 1))
+        ([['n?', 'k'], ['k']], [['n?']])
+
+    Args:
+        signature (str): Function signature
+        inputdims (Iterable[int]): Input dimensions
+
+    Returns:
+        tuple[list[list[str]]
+    """
+    signature = "".join(signature.split())  # remove whitespace
+    sigin, sigout = signature.split("->")  # split input and output
+    sigin = sigin[1:-1].split("),(")  # split input
+    sigin = [i.split(",") for i in sigin]
+    sigout = sigout[1:-1].split("),(")  # split output
+    sigout = [i.split(",") for i in sigout]
+    for i, idim in enumerate(inputdims):
+        j = 0
+        while j < len(sigin[i]):
+            d = sigin[i][j]
+            if d.endswith("?") and len(sigin[i]) > idim:
+                sigin = [[i for i in s if i != d] for s in sigin]
+                sigout = [[i for i in s if i != d] for s in sigout]
+            else:
+                j += 1
+    return sigin, sigout
+
+
+def _parse_key(key, ndim):
+    """Parse a key to index an array.
+
+    This function attempts to replicate the numpy indexing of arrays. It can handle
+    integers, slices, an Ellipsis, arrays of integers, arrays of bools, (bare) bools,
+    and None. It returns the key with an Ellipsis appended if the number of index
+    dimensions does not match the number of dimensions given. Additionally, it informs
+    about the number of dimension indexed by fancy indexing, if the fancy indexed
+    dimensions will be prepended, and the number of dimensions that the Ellipsis
+    contains.
+
+    Args:
+        key (tuple): The indexing key
+        ndim (int): Number of array dimensions
+
+    Returns:
+        tuple
+    """
+    consumed = 0
+    ellipsis = False
+    fancy_ndim = 0
+    # nfancy = 0
+    consecutive_intfancy = 0
+    # consecutive_intfancy = 0: no fancy/integer indexing
+    # consecutive_intfancy = 1: ongoing consecutive fancy/integer indexing
+    # consecutive_intfancy = 2: terminated consecutive fancy/integer indexing
+    # consecutive_intfancy >= 2: non-consecutive fancy/integer indexing
+
+    # The first pass gets the consumed dimensions, the presence of an ellipsis, and
+    # fancy index properties.
+    for k in key:
+        if k is not True and k is not False and isinstance(k, (int, np.integer)):
+            consumed += 1
+            consecutive_intfancy += (consecutive_intfancy + 1) % 2
+        elif k is None:
+            consecutive_intfancy += consecutive_intfancy % 2
+        elif isinstance(k, slice):
+            consumed += 1
+            consecutive_intfancy += consecutive_intfancy % 2
+        elif k is Ellipsis:
+            ellipsis = True
+            # consumed is determined at the end
+            consecutive_intfancy += consecutive_intfancy % 2
+        else:
+            arr = np.asanyarray(k)
+            if arr.dtype == bool:
+                consumed += arr.ndim
+                fancy_ndim = max(1, fancy_ndim)
+            else:
+                consumed += 1
+                fancy_ndim = max(arr.ndim, fancy_ndim)
+            # nfancy += arr.ndim
+            consecutive_intfancy += (consecutive_intfancy + 1) % 2
+
+    lenellipsis = ndim - consumed
+    if lenellipsis != 0 and not ellipsis:
+        key = key + (Ellipsis,)
+
+    return key, fancy_ndim, consecutive_intfancy >= 2, lenellipsis
+
+
+HANDLED_FUNCTIONS = {}
+"""Dictionary of numpy functions implemented for AnnotatedArrays."""
+
+
+def implements(np_func):
+    """Decorator to register an __array_function__ implementation to AnnotatedArrays."""
+
+    def decorator(func):
+        HANDLED_FUNCTIONS[np_func] = func
+        return func
+
+    return decorator
+
+
+class AnnotatedArray(np.lib.mixins.NDArrayOperatorsMixin):
+    """Array that keeps track of annotations for each dimension.
+
+    This class acts mostly like numpy arrays, but it is enhanced by the following
+    functionalities:
+
+        * Annotations are added to each dimension
+        * Annotations are compared and preserved for numpy (generalized)
+          :py:class:`numpy.ufunc` (like :py:data:`numpy.add`, :py:data:`numpy.exp`,
+          :py:data:`numpy.matmul` and many more "standard" mathematical functions)
+        * Special ufunc methods, like :py:meth:`numpy.ufunc.reduce`, are supported
+        * A growing subset of other numpy functions are supported, like
+          :py:func:`numpy.linalg.solve`
+        * Keywords can be specified as scale are also index into, when index the
+          AnnotatedArray
+        * Annotations can also be exposed as properties
+
+    .. testsetup::
+
+        from treams.util import AnnotatedArray
+
+    Example:
+        >>> a = AnnotatedArray([[0, 1], [2, 3]], ({"a": 1}, {"b": 2}))
+        >>> b = AnnotatedArray([1, 2], ({"b": 2},))
+        >>> a @ b
+        AnnotatedArray(
+            [2, 8],
+            AnnotationSequence(AnnotationDict({'a': 1})),
+        )
+
+    The interoperability with numpy is implemented using :ref:`basics.dispatch`
+    by defining :meth:`__array__`, :meth:`__array_ufunc__`, and
+    :meth:`__array_function__`.
+    """
+
+    _scales = set()
+
+    def __init__(self, array, ann=(), /, **kwargs):
+        """Initalization."""
+        self._array = np.asarray(array)
+        self.ann = getattr(array, "ann", ())
+        self.ann.update(ann)
+        for key, val in kwargs.items():
+            val = (val,) * self.ndim if not isinstance(val, tuple) else val
+            self.ann.as_dict[key] = val
+
+    @classmethod
+    def relax(cls, *args, mro=None, **kwargs):
+        """Try creating AnnotatedArray subclasses if possible.
+
+        Subclasses can impose stricter conditions on the Annotations. To allow a simple
+        "decaying" of those subclasses it is possible to create them with this
+        classmethod. It attempts array creations along the method resolution order until
+        it succeeds.
+
+        Args:
+            mro (array-like, optional): Method resolution order along which to create
+                the subclass. By default it takes the order of the calling class.
+
+        Note:
+            All other arguments are the same as for the default initialization.
+        """
+        mro = cls.__mro__[1:] if mro is None else mro
+        try:
+            return cls(*args, **kwargs)
+        except AnnotationError as err:
+            if cls == AnnotatedArray:
+                raise err from None
+        cls, *mro = mro
+        return cls.relax(*args, mro=mro, **kwargs)
+
+    def __str__(self):
+        """String of the array itself."""
+        return str(self._array)
+
+    def __repr__(self):
+        """String representation."""
+        repr_arr = "    " + repr(self._array)[6:-1].replace("\n  ", "\n")
+        return f"{self.__class__.__name__}(\n{repr_arr},\n    {self.ann},\n)"
+
+    def __int__(self):
+        return int(self._array)
+
+    def __float__(self):
+        return float(self._array)
+
+    def __array__(self, dtype=None):
+        """Convert to an numpy array.
+
+        This function returns the bare array without annotations. This function does not
+        necessarily make a copy of the array.
+
+        Args:
+            dtype (optional): Type of the returned array.
+        """
+        return np.asarray(self._array, dtype=dtype)
+
+    @property
+    def ann(self):
+        """Array annotations."""
+        return self._ann
+
+    @ann.setter
+    def ann(self, ann):
+        """Set array annotations.
+
+        This function copies the given sequence of dictionaries.
+        """
+        self._ann = AnnotationSequence(*(({},) * self.ndim))
+        self._ann.update(ann)
+
+    def __getattr__(self, key):
+        if key in self.ann.as_dict:
+            res = self.ann.as_dict[key]
+            if all(res[0] == i for i in res[1:]):
+                return res[0]
+            return res
+        raise AttributeError(
+            f"'{self.__class__.__name__}' object has no attribute '{key}'"
+        )
+
+    def __setattr__(self, key, val):
+        if key not in ("_array", "ann", "_ann") and key in self.ann.as_dict:
+            val = (val,) * self.ndim if not isinstance(val, tuple) else val
+            self.ann.as_dict[key] = val
+        else:
+            super().__setattr__(key, val)
+
+    def __delattr__(self, key):
+        try:
+            super().__delattr__(key)
+        except AttributeError:
+            del self.ann.as_dict[key]
+
+    def __len__(self):
+        return len(self._array)
+
+    def __copy__(self):
+        return type(self)(copy.copy(self._array), copy.copy(self._ann))
+
+    def __deepcopy__(self, memo):
+        cls = type(self)
+        return cls(copy.deepcopy(self._array, memo), copy.deepcopy(self._ann, memo))
+
+    def __bool__(self):
+        """Boolean value of the array."""
+        return bool(self._array)
+
+    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
+        """Implement ufunc API."""
+        # Compute result first to use numpy's comprehensive checks on the arguments
+        inputs_noaa = tuple(map(_cast_nparray, inputs))
+
+        ann_out = ()
+        out = kwargs.get("out")
+        out = out if isinstance(out, tuple) else (out,)
+        ann_out = tuple(getattr(i, "ann", None) for i in out)
+        if len(out) != 1 or out[0] is not None:
+            kwargs["out"] = tuple(map(_cast_nparray, out))
+
+        res = getattr(ufunc, method)(*inputs_noaa, **kwargs)
+        istuple, res = (True, res) if isinstance(res, tuple) else (False, (res,))
+        res = tuple(_cast_annarray(r) if o is None else o for r, o in zip(res, out))
+        for a, r in zip(ann_out, res):
+            if a is not None:
+                r.ann.update(a)
+
+        inputs_and_where = inputs + ((kwargs["where"],) if "where" in kwargs else ())
+
+        if ufunc.signature is None or all(
+            map(lambda x: x in " (),->", ufunc.signature)
+        ):
+            if method in ("__call__", "reduceat", "accumulate") or (
+                method == "reduce" and kwargs.get("keepdims", False)
+            ):
+                self._ufunc_call(inputs_and_where, res)
+            elif method == "reduce":
+                self._ufunc_reduce(inputs_and_where, res, kwargs.get("axis", 0))
+            elif method == "at":
+                self._ufunc_at(inputs)
+                return None
+            elif method == "outer":
+                self._ufunc_outer(inputs, res, kwargs.get("where", True))
+            else:
+                warnings.warn("unrecognized ufunc method", AnnotationWarning)
+        else:
+            res = self._gufunc_call(ufunc, inputs, kwargs, res)
+        res = tuple(r if isinstance(r, np.generic) else self.relax(r) for r in res)
+        return res if istuple else res[0]
+
+    @staticmethod
+    def _ufunc_call(inputs_and_where, res):
+        for out, in_ in itertools.product(res, inputs_and_where):
+            if not isinstance(out, AnnotatedArray):
+                continue
+            out.ann.update(getattr(in_, "ann", ()))
+
+    @staticmethod
+    def _ufunc_reduce(inputs_and_where, res, axis=0, keepdims=False):
+        out = res[0]  # reduce only allowed for single output functions
+        if not isinstance(out, AnnotatedArray):
+            return
+        axis = axis if isinstance(axis, tuple) else (axis,)
+        in_ = inputs_and_where[0]
+        axis = sorted(map(lambda x: x % np.ndim(in_) - np.ndim(in_), axis))
+        for in_ in inputs_and_where:
+            ann = list(getattr(in_, "ann", []))
+            if not keepdims:
+                for a in axis:
+                    try:
+                        del ann[a]
+                    except IndexError:
+                        pass
+            out.ann.update(ann)
+
+    @staticmethod
+    def _ufunc_at(inputs):
+        out = inputs[0]
+        if not isinstance(out, AnnotatedArray):
+            return
+        if any(d != {} for d in getattr(inputs[1], "ann", ())):
+            warnings.warn("annotations in indices are ignored", AnnotationWarning)
+        for in_ in inputs[2:]:
+            ann = getattr(in_, "ann", ())
+            out.ann.update(ann)
+
+    @staticmethod
+    def _ufunc_outer(inputs, res, where=True):
+        for out in res:
+            if not isinstance(out, AnnotatedArray):
+                continue
+            in_ann = tuple(
+                i for a in inputs for i in getattr(a, "ann", np.ndim(a) * ({},))
+            )
+            out.ann.update(in_ann)
+            where_ann = getattr(where, "ann", ())
+            out.ann.update(where_ann)
+
+    @staticmethod
+    def _gufunc_call(ufunc, inputs, kwargs, res):
+        sigin, sigout = _parse_signature(ufunc.signature, map(np.ndim, inputs))
+        if kwargs.get("keepdims", False):
+            sigout = [sigin[0] for _ in range(ufunc.nout)]
+        ndims = [np.ndim(i) for x in (inputs, res) for i in x]
+        axes = getattr(kwargs, "axes", None)
+        if axes is None:
+            axis = getattr(kwargs, "axis", None)
+            if axis is None:
+                axes = [tuple(range(-len(i), 0)) for i in sigin + sigout]
+            else:
+                axes = [(axis,) for _ in range(ufunc.nin)]
+        else:
+            axes = [
+                tuple(a) if isinstance(a, collections.abc.Iterable) else (a,)
+                for a in axes
+            ]
+        append = axes[0] if kwargs.get("keepdims", False) else ()
+        axes += [append] * (ufunc.nin + ufunc.nout - len(axes))
+        axes = [(*(i % ndim - ndim for i in a),) for a, ndim in zip(axes, ndims)]
+        iterdims = [
+            [i for i in range(-1, -1 - ndim, -1) if i not in a]
+            for a, ndim in zip(axes, ndims)
+        ]
+
+        # compare core dimensions
+        coredims = {}
+        for i, (ax, sig) in enumerate(zip(axes, sigin + sigout)):
+            for a, key in zip(ax, sig):
+                if key.isnumeric():
+                    continue
+                coredims.setdefault(key, []).append((i, a))
+        inout = tuple(inputs) + res
+        for val in coredims.values():
+            for (isrc, dimsrc), (idest, dimdest) in itertools.combinations(val, 2):
+                source = getattr(inout[isrc], "ann", {dimsrc: {}})[dimsrc]
+                dest = getattr(inout[idest], "ann", {dimdest: AnnotationDict()})[
+                    dimdest
+                ]
+                if isrc < ufunc.nin <= idest:
+                    dest.update(source)
+                else:
+                    dest.match(source)
+
+        # compare iteration dimensions
+        for iout, out in enumerate(res):
+            if not isinstance(out, AnnotatedArray):
+                continue
+            for idim, dim in enumerate(iterdims[ufunc.nin + iout]):
+                dest = out.ann[dim]
+                for in_, iterdim in zip(inputs, iterdims):
+                    if idim >= len(iterdim) or getattr(in_, "ann", None) is None:
+                        continue
+                    source = getattr(in_, "ann", {iterdim[idim]: {}})[iterdim[idim]]
+                    dest.update(source)
+        return res
+
+    def __array_function__(self, func, types, args, kwargs):
+        """Function calls on the array.
+
+        Calls defined function in :data:`HANDLED_FUNCTIONS` otherwise raises exception.
+        Add functions to it by using the decorator :func:`implements` for custom
+        implementations.
+        """
+        if func not in HANDLED_FUNCTIONS:
+            return NotImplemented
+        # if not all(issubclass(t, self.__class__) for t in types):
+        #     return NotImplemented
+        return HANDLED_FUNCTIONS[func](*args, **kwargs)
+
+    def __getitem__(self, key):
+        """Get an item from the AnnotatedArray.
+
+        The indexing supports most of numpys regular and fancy indexing.
+        """
+        res = AnnotatedArray(self._array[key])
+        if isinstance(res, np.generic) or res.ndim == 0:
+            return res
+
+        key = key if isinstance(key, tuple) else (key,)
+        key, fancy_ndim, prepend_fancy, lenellipsis = _parse_key(key, self.ndim)
+        source = 0
+        dest = fancy_ndim if prepend_fancy else 0
+        for k in key:
+            if k is not True and k is not False and isinstance(k, (int, np.integer)):
+                source += 1
+            elif k is None:
+                dest += 1
+            elif isinstance(k, slice):
+                for kk, val in self.ann[source].items():
+                    if kk in self._scales:
+                        res.ann[dest][kk] = val[k]
+                    else:
+                        res.ann[dest][kk] = val
+                dest += 1
+                source += 1
+            elif k is Ellipsis:
+                for _ in range(lenellipsis):
+                    res.ann[dest].update(self.ann[source])
+                    dest += 1
+                    source += 1
+            else:
+                k = np.asanyarray(k)
+                ksq = k.squeeze()
+                if ksq.ndim == 1:
+                    pos = (np.array(k.shape) == k.size).argmax()
+                    ann = (
+                        self.ann[source + pos] if k.dtype == bool else self.ann[source]
+                    )
+                    pos += int(not prepend_fancy) * dest + fancy_ndim - k.ndim
+                    for kk, val in ann.items():
+                        if kk in self._scales:
+                            res.ann[pos][kk] = val[ksq]
+                        else:
+                            res.ann[pos][kk] = val
+                source += k.ndim if k.dtype == bool else 1
+                if not prepend_fancy:
+                    dest += fancy_ndim
+                    fancy_ndim = 0
+        return self.relax(res)
+
+    def __setitem__(self, key, value):
+        """Set values.
+
+        If the provided value is an AnnotatedArray the annotations of corresponding
+        dimensions will be matched.
+        """
+        self._array[key] = value
+        if not hasattr(value, "ann"):
+            return
+        key = key if isinstance(key, tuple) else (key,)
+        key, fancy_ndim, prepend_fancy, lenellipsis = _parse_key(key, self.ndim)
+        source = 0
+        dest = fancy_ndim if prepend_fancy else 0
+        for k in key:
+            if k is not True and k is not False and isinstance(k, (int, np.integer)):
+                source += 1
+            elif k is None:
+                dest += 1
+            elif isinstance(k, slice):
+                for kk, val in self.ann[source].items():
+                    if kk not in value.ann[dest]:
+                        continue
+                    if kk in self._scales:
+                        val = val[k]
+                    if val != value.ann[dest][kk]:
+                        warnings.warn(
+                            f"incompatible annotations with key '{kk}' "
+                            f"comparing dimensions '{source}' and '{dest}'",
+                            AnnotationWarning,
+                        )
+                dest += 1
+                source += 1
+            elif k is Ellipsis:
+                for _ in range(lenellipsis):
+                    self.ann[source].match(value.ann[dest])
+                    dest += 1
+                    source += 1
+            else:
+                k = np.asanyarray(k)
+                ksq = k.squeeze()
+                if ksq.ndim == 1:
+                    pos = (np.array(k.shape) == k.size).argmax()
+                    ann = (
+                        self.ann[source + pos] if k.dtype == bool else self.ann[source]
+                    )
+                    pos += int(not prepend_fancy) * dest + fancy_ndim - k.ndim
+                    for kk, val in ann.items():
+                        if kk not in value.ann[pos]:
+                            continue
+                        if kk in self._scales:
+                            val = val[k]
+                        if val != value.ann[pos][kk]:
+                            warnings.warn(
+                                f"incompatible annotations with key '{kk}' "
+                                f"comparing dimensions '{source}' and '{pos}'",
+                                AnnotationWarning,
+                            )
+                source += k.ndim if k.dtype == bool else 1
+                if not prepend_fancy:
+                    dest += fancy_ndim
+                    fancy_ndim = 0
+
+    @property
+    def T(self):
+        """Transpose.
+
+        See also :py:attr:`numpy.ndarray.T`.
+        """
+        return self.transpose()
+
+    @implements(np.all)
+    def all(self, axis=None, dtype=None, out=None, keepdims=False, *, where=True):
+        """Test if all elements (along an axis) are True.
+
+        See also :py:meth:`numpy.ndarray.all`.
+        """
+        return np.logical_and.reduce(self, axis, dtype, out, keepdims, where=where)
+
+    @implements(np.any)
+    def any(self, axis=None, dtype=None, out=None, keepdims=False, *, where=True):
+        """Test if any element (along an axis) is True.
+
+        See also :py:meth:`numpy.ndarray.any`.
+        """
+        return np.logical_or.reduce(self, axis, dtype, out, keepdims, where=where)
+
+    @implements(np.max)
+    def max(self, axis=None, out=None, keepdims=False, initial=None, where=True):
+        """Maximum (along an axis).
+
+        See also :py:meth:`numpy.ndarray.max`.
+        """
+        return np.maximum.reduce(self, axis, None, out, keepdims, initial, where)
+
+    @implements(np.min)
+    def min(self, axis=None, out=None, keepdims=False, initial=None, where=True):
+        """Minimum (along an axis).
+
+        See also :py:meth:`numpy.ndarray.min`.
+        """
+        return np.minimum.reduce(self, axis, None, out, keepdims, initial, where)
+
+    @implements(np.sum)
+    def sum(
+        self, axis=None, dtype=None, out=None, keepdims=False, initial=0, where=True
+    ):
+        """Sum of elements (along an axis).
+
+        See also :py:meth:`numpy.ndarray.sum`.
+        """
+        return np.add.reduce(self, axis, dtype, out, keepdims, initial, where)
+
+    @implements(np.prod)
+    def prod(
+        self, axis=None, dtype=None, out=None, keepdims=False, initial=1, where=True
+    ):
+        """Product of elements (along an axis).
+
+        See also :py:meth:`numpy.ndarray.prod`.
+        """
+        return np.multiply.reduce(self, axis, dtype, out, keepdims, initial, where)
+
+    @implements(np.cumsum)
+    def cumsum(self, axis=None, dtype=None, out=None):
+        """Cumulative sum of elements (along an axis).
+
+        See also :py:meth:`numpy.ndarray.cumsum`.
+        """
+        if axis is None:
+            return np.add.accumulate(self.flatten(), 0, dtype, out)
+        return np.add.accumulate(self, axis, dtype, out)
+
+    @implements(np.cumprod)
+    def cumprod(self, axis=None, dtype=None, out=None):
+        """Cumulative product of elements (along an axis).
+
+        See also :py:meth:`numpy.ndarray.cumprod`.
+        """
+        if axis is None:
+            return np.multiply.accumulate(self.flatten(), 0, dtype, out)
+        return np.multiply.accumulate(self, axis, dtype, out)
+
+    def flatten(self, order="C"):
+        """Flatten array to one dimension.
+
+        See also :py:meth:`numpy.ndarray.flatten`.
+        """
+        res = self.relax(self._array.flatten(order))
+        if res.shape == self.shape:
+            res.ann.update(self.ann)
+        elif len(tuple(filter(lambda x: x != 1, self.shape))) == 1:
+            res.ann[0].update(self.ann[self.shape.index(res.size)])
+        return res
+
+    @implements(np.trace)
+    def trace(self, offset=0, axis1=0, axis2=1, dtype=None, out=None):
+        """Trace of an array.
+
+        See also :py:meth:`numpy.ndarray.trac`.
+        """
+        ann = tuple(a for i, a in enumerate(self.ann) if i not in (axis1, axis2))
+        return self.relax(self._array.trace(offset, axis1, axis2, dtype, out), ann)
+
+    def astype(self, *args, **kwargs):
+        """Return array as given type.
+
+        See also :py:meth:`numpy.ndarray.astype`.
+        """
+        return self.relax(self._array.astype(*args, **kwargs), self.ann)
+
+    @property
+    @implements(np.ndim)
+    def ndim(self):
+        """Number of array dimensions.
+
+        See also :py:attr:`numpy.ndarray.ndim`.
+        """
+        return self._array.ndim
+
+    @property
+    @implements(np.shape)
+    def shape(self):
+        """Array shape.
+
+        See also :py:attr:`numpy.ndarray.shape`.
+        """
+        return self._array.shape
+
+    @property
+    @implements(np.size)
+    def size(self):
+        """Array size.
+
+        See also :py:attr:`numpy.ndarray.size`.
+        """
+        return self._array.size
+
+    @property
+    @implements(np.imag)
+    def imag(self):
+        """Imaginary part of the array.
+
+        See also :py:attr:`numpy.ndarray.imag`.
+        """
+        return self.relax(self._array.imag, self.ann)
+
+    @property
+    @implements(np.real)
+    def real(self):
+        """Real part of the array.
+
+        See also :py:attr:`numpy.ndarray.real`.
+        """
+        return self.relax(self._array.real, self.ann)
+
+    def conjugate(self, *args, **kwargs):
+        """Complex conjugate elementwise.
+
+        See also :py:meth:`numpy.ndarray.conjugate`.
+        """
+        return np.conjugate(self, *args, **kwargs)
+
+    @implements(np.diagonal)
+    def diagonal(self, offset=0, axis1=0, axis2=1):
+        """Get the diagonal of the array.
+
+        See also :py:meth:`numpy.ndarray.diagonal`.
+        """
+        ann = tuple(a for i, a in enumerate(self.ann) if i not in (axis1, axis2)) + (
+            {},
+        )
+        return self.relax(self._array.diagonal(offset, axis1, axis2), ann)
+
+    @implements(np.transpose)
+    def transpose(self, axes=None):
+        """Transpose array dimensions.
+
+        See also :py:meth:`numpy.ndarray.transpose`.
+        """
+        axes = range(self.ndim - 1, -1, -1) if axes is None else axes
+        return self.relax(self._array.transpose(axes), self.ann[axes])
+
+    conj = conjugate
+
+    @implements(np.swapaxes)
+    def swapaxes(self, axis1, axis2):
+        axis1 = axis1 % self.ndim
+        axis2 = axis2 % self.ndim
+        opp = {axis1: axis2, axis2: axis1}
+        axes = [opp.get(i, i) for i in range(self.ndim)]
+        return self.relax(self._array.swapaxes(axis1, axis2), self.ann[axes])
+
+    def __matmul__(self, other):
+        if isinstance(other, _op.Operator):
+            return NotImplemented
+        return super().__matmul__(other)
+
+
+@implements(np.linalg.solve)
+def solve(a, b):
+    """Solve linear system.
+
+    See also :py:func:`numpy.linalg.solve`.
+    """
+    if issubclass(type(a), type(b)) or (
+        not issubclass(type(b), type(a)) and isinstance(a, AnnotatedArray)
+    ):
+        restype = type(a)
+    else:
+        restype = type(b)
+    res = AnnotatedArray(np.linalg.solve(np.asanyarray(a), np.asanyarray(b)))
+    a_ann = list(getattr(a, "ann", [{}, {}]))
+    b_ann = list(getattr(b, "ann", [{}, {}]))
+    if np.ndim(b) == np.ndim(a) - 1:
+        map(lambda x: x[0].match(x[1]), zip(a_ann[-2::-1], b_ann[-1::-1]))
+        del a_ann[-2]
+        del b_ann[-1]
+    else:
+        map(lambda x: x[0].match(x[1]), zip(a_ann[-2::-1], b_ann[-2::-1]))
+        del a_ann[-2]
+        del b_ann[-2]
+        a_ann += [{}]
+    res.ann.update(a_ann)
+    res.ann.update(b_ann)
+    return restype.relax(res)
+
+
+@implements(np.linalg.lstsq)
+def lstsq(a, b, rcond="warn"):
+    """Solve linear system using least squares.
+
+    See also :py:func:`numpy.linalg.lstsq`.
+    """
+    if issubclass(type(a), type(b)) or (
+        not issubclass(type(b), type(a)) and isinstance(a, AnnotatedArray)
+    ):
+        restype = type(a)
+    else:
+        restype = type(b)
+    res = list(np.linalg.lstsq(np.asanyarray(a), np.asanyarray(b), rcond))
+    res[0] = AnnotatedArray(res[0])
+    a_ann = getattr(a, "ann", ({},))
+    b_ann = getattr(b, "ann", ({},))
+    a_ann[0].match(b_ann[0])
+    res[0].ann[0].update(a_ann[-1])
+    if np.ndim(b) == 2:
+        res[0].ann[1].update(b_ann[-1])
+    res[0] = restype.relax(res[0])
+    return tuple(res)
+
+
+@implements(np.linalg.svd)
+def svd(a, full_matrices=True, compute_uv=True, hermitian=False):
+    """Compute the singular value decomposition.
+
+    See also :py:func:`numpy.linalg.svd`.
+    """
+    res = list(np.linalg.svd(np.asanyarray(a), full_matrices, compute_uv, hermitian))
+    ann = getattr(a, "ann", ({},))
+    if compute_uv:
+        res[0] = a.relax(res[0], ann[:-1] + ({},))
+        res[1] = a.relax(res[1], ann[:-2] + ({},))
+        res[2] = a.relax(res[2], ann[:-2] + ({}, ann[-1]))
+        return res
+    return a.relax(res, tuple(ann[:-2]) + ({},))
+
+
+@implements(np.diag)
+def diag(a, k=0):
+    """Extract diagonal from an array or create an diagonal array.
+
+    See also :py:func:`numpy.diag`.
+    """
+    res = np.diag(np.asanyarray(a), k)
+    ann = a.ann
+    if a.ndim == 1:
+        ann = (ann[0], copy.copy(ann[0]))
+    elif k == 0:
+        ann[0].match(ann[1])
+        ann = ({**ann[0], **ann[1]},)
+    else:
+        ann = ()
+    return a.relax(res, ann)
+
+
+@implements(np.tril)
+def tril(a, k=0):
+    """Get lower trinagular matrix.
+
+    See also :py:func:`numpy.tril`.
+    """
+    res = np.tril(np.asanyarray(a), k)
+    return a.relax(res, getattr(a, "ann", ({},)))
+
+
+@implements(np.triu)
+def triu(a, k=0):
+    """Get upper trinagular matrix.
+
+    See also :py:func:`numpy.triu`.
+    """
+    res = np.triu(np.asanyarray(a), k)
+    return a.relax(res, getattr(a, "ann", ({},)))
+
+
+@implements(np.zeros_like)
+def zeros_like(a, dtype=None, order="K", shape=None):
+    """Create a numpy array like the given array containing zeros."""
+    return np.zeros_like(np.asarray(a), dtype=dtype, order=order, shape=shape)
+
+
+@implements(np.ones_like)
+def ones_like(a, dtype=None, order="K", shape=None):
+    """Create a numpy array like the given array containing ones."""
+    return np.ones_like(np.asarray(a), dtype=dtype, order=order, shape=shape)

conftest.py

@@ -10,10 +10,14 @@ """Configuration for pytest.
 def pytest_addoption(parser):
-    """Add option '--runslow'."""
+    """Add option '--runslow' and '--rungmsh'."""
     parser.addoption(
         "--runslow", action="store_true", default=False, help="run slow tests"
     )
+    parser.addoption(
+        "--rungmsh", action="store_true", default=False, help="run tests needing gmsh"
+    )
 
 
 def pytest_configure(config):
-    """Add marker 'slow'."""
+    """Add marker 'slow' and 'gmsh'."""
+    config.addinivalue_line("markers", "gmsh: test needs gmsh")
     config.addinivalue_line("markers", "slow: mark test as slow to run")
@@ -23,9 +27,14 @@
 def pytest_collection_modifyitems(config, items):
-    """Skip slow tests without option '--runslow'."""
-    if config.getoption("--runslow"):
-        # --runslow given in cli: do not skip slow tests
-        return
-    skip_slow = pytest.mark.skip(reason="need --runslow option to run")
-    for item in items:
-        if "slow" in item.keywords:
-            item.add_marker(skip_slow)
+    """Skip tests without necessary options."""
+    if not config.getoption("--runslow"):
+        # --runslow not given in cli: skip slow tests
+        skip_slow = pytest.mark.skip(reason="need --runslow option to run")
+        for item in items:
+            if "slow" in item.keywords:
+                item.add_marker(skip_slow)
+    if not config.getoption("--rungmsh"):
+        # --rungmsh not given in cli: skip gmsh tests
+        skip_slow = pytest.mark.skip(reason="need --rungmsh option to run")
+        for item in items:
+            if "gmsh" in item.keywords:
+                item.add_marker(skip_slow)

docs/conf.py

@@ -57,2 +57,4 @@ """Sphinx config."""
 from treams import *
+
+np.set_printoptions(precision=3, suppress=True)
 """
@@ -59,0 +61,0 @@ intersphinx_mapping = {

docs/dev.rst

@@ -13,3 +13,3 @@ .. highlight:: console
 
-   git clone [e-mail]/tfp-photonics/treams.git
+   git clone [e-mail]:tfp-photonics/treams.git
 
@@ -28,3 +28,3 @@ or ::
 
-   conda env create --name treams-dev
+   conda env create --name treams-dev python
 
@@ -43,6 +43,9 @@ Activate the environment with::
 
-
 Running tests
 =============
 
+To install the required packages for testing use ::
+
+   pip install treams[test,coverage,io]
+
 Tests can be run using pytest with ::
@@ -60,7 +63,7 @@
 
-If coverage reports should be included one can use the option ``--cov treams`` (which
-requires pytest-cov to be installed). However, this will only report on the pure python
-files. To also get coverage reports for the cython part, it is necessary to compile it
-with linetracing support. This can be achieved by setting the environment variable
-``CYTHON_COVERAGE``, for example with ::
+If coverage reports should be included one can use the option ``--cov src/treams``
+(which requires pytest-cov to be installed). However, this will only report on the pure
+python files. To also get coverage reports for the cython part, it is necessary to
+compile it with linetracing support. This can be achieved by setting the environment
+variable ``CYTHON_COVERAGE``, for example with ::
 
@@ -73,5 +76,13 @@ CYTHON_COVERAGE=1 pip install -e .
 
+Some test for importing and exporting need gmsh. These tests can be activated with
+``--rungmsh`` analogous to the option ``--runslow`` above.
+
 Building the documentation
 ==========================
 
+
+To install the required packages for testing use ::
+
+   pip install treams[docs]
+
 The documentation is built with sphinx by using ::
@@ -84,2 +95,6 @@
 
+The doctests can be run with ::
+
+   sphinx-build -b doctest docs docs/_build/doctest
+
 Building the code on Windows
@@ -90,8 +105,5 @@ ============================
 is usually compiled with MSVC for Visual Studio. However, especially for calculations
-with complex numbers, Cython creates code that conforms to the (C99-) standard.Thus, it
+with complex numbers, Cython creates code that conforms to the (C99-) standard. Thus, it
 is not compatible with the non-standard implementation of complex numbers by Microsoft.
 
-As I understand, even large projects like numpy and scipy do not compile their code with
-MSVC (at least not completely) for the Windows distribution.
-
 Below you find three tested ways, how one can use treams with Windows. The first two
@@ -124,4 +136,4 @@ ways actually use a non-Windows version of Python, but have a more straightforward
 
-Compilation with mingw-w64 for MSVC Python
-------------------------------------------
+Compilation with mingw-w64
+--------------------------
 
@@ -131,23 +143,6 @@ This is approach is different from the others, since it finally combines binaries from
 guaranteed that it will work for all systems. The following part describes, how treams
-can be built for Windows. It was initially inspired by
-`<https://docs.scipy.org/doc/scipy/reference/building/windows.html>`_. However, it is
-not completely tested, which steps could possible be omitted.
+can be built for Windows.
 
-The first step is the installation of MSYS2 and components of Microsoft Visual Studio
-The installation of MSYS2 is pretty straightforward. Regarding the Microsoft Visual
-Studio components, it is unclear to me, which one are actually used, so for this part we
-just rely on the description for scipy, which is a sufficient set of components. Feel
-free to test this and adjust here accordingly. I suspect, that it might not be necessary
-to install most components at all. The only requirement so far seems to be the presence
-of `vcruntime140.dll`, which should come shipped with recent versions of Python
-(see also
-`Steve Dower's blog post <https://stevedower.id.au/blog/building-for-python-3-5-part-two>`_).
-If not present, they can additionally be installed with the pip package `msvc-runtime`.
-Obviously, an installation of Python on Windows is necessary. This can either be pure
-Python or can come with a distribution like Anaconda. In some cases, it might be
-necessary to patch distutils' `cygwinccompiler.py` to return `vcruntime140` instead of
-`msvcr140`.
+After installing MSYS2 use it to install ``mingw-w64-x86_64-gcc``.
 
-Within MSYS2 install `mingw-w64-x86_64-gcc`.
-
 The compilation is steered from the command line. First go into the directory of treams.
@@ -159,1 +154,10 @@ Then, set up your path by prepending the direction for MSYS2's mingw64 binaries with
 Windows Python. With this setup, building binaries should work with `python -m build`.
+
+Continuous integration
+======================
+
+Using Github actions the steps above (building the documentation, running the doctests,
+running the tests with a coverage report, and building for all platforms) are
+implemented to run automatically on specific triggers. The relevant files in the
+``.gihub/workflows`` directory can be also used as examples in addition to the
+description above.

docs/examples/chain.py

@@ -33,8 +33,3 @@ import matplotlib.pyplot as plt
 pcm = ax.pcolormesh(
-    grid[0, :, 2],
-    grid[:, 0, 0],
-    ez.T,
-    shading="nearest",
-    vmin=-0.5,
-    vmax=0.5,
+    grid[0, :, 2], grid[:, 0, 0], ez.T, shading="nearest", vmin=-0.5, vmax=0.5,
 )
@@ -41,0 +36,0 @@ cb = plt.colorbar(pcm)

docs/examples/cluster.py

@@ -36,8 +36,3 @@ import matplotlib.pyplot as plt
 pcm = ax.pcolormesh(
-    grid[0, :, 2],
-    grid[:, 0, 0],
-    intensity.T,
-    shading="nearest",
-    vmin=0,
-    vmax=2,
+    grid[0, :, 2], grid[:, 0, 0], intensity.T, shading="nearest", vmin=0, vmax=2,
 )
@@ -73,8 +68,3 @@ cb = plt.colorbar(pcm)
 pcm = ax.pcolormesh(
-    grid[0, :, 2],
-    grid[:, 0, 0],
-    intensity_global.T,
-    shading="nearest",
-    vmin=0,
-    vmax=2,
+    grid[0, :, 2], grid[:, 0, 0], intensity_global.T, shading="nearest", vmin=0, vmax=2,
 )
@@ -111,8 +101,3 @@ cb = plt.colorbar(pcm)
 pcm = ax.pcolormesh(
-    grid[0, :, 2],
-    grid[:, 0, 0],
-    intensity_global.T,
-    shading="nearest",
-    vmin=0,
-    vmax=2,
+    grid[0, :, 2], grid[:, 0, 0], intensity_global.T, shading="nearest", vmin=0, vmax=2,
 )
@@ -119,0 +104,0 @@ cb = plt.colorbar(pcm)

docs/examples/grid.py

@@ -35,8 +35,3 @@ import matplotlib.pyplot as plt
 pcm = ax.pcolormesh(
-    grid[0, :, 2],
-    grid[:, 0, 0],
-    ez.T,
-    shading="nearest",
-    vmin=-1,
-    vmax=1,
+    grid[0, :, 2], grid[:, 0, 0], ez.T, shading="nearest", vmin=-1, vmax=1,
 )
@@ -43,0 +38,0 @@ cb = plt.colorbar(pcm)

docs/examples/sphere.py

@@ -12,9 +12,9 @@ import matplotlib.pyplot as plt
 
-xs_sca = np.array([tm.xs_sca_avg for tm in spheres]) / (np.pi * radius**2)
-xs_ext = np.array([tm.xs_ext_avg for tm in spheres]) / (np.pi * radius**2)
+xs_sca = np.array([tm.xs_sca_avg for tm in spheres]) / (np.pi * radius ** 2)
+xs_ext = np.array([tm.xs_ext_avg for tm in spheres]) / (np.pi * radius ** 2)
 
 swb_lmax1 = treams.SphericalWaveBasis.default(1)
 spheres_lmax1 = [tm[swb_lmax1] for tm in spheres]
-xs_sca_lmax1 = np.array([tm.xs_sca_avg for tm in spheres_lmax1]) / (np.pi * radius**2)
-xs_ext_lmax1 = np.array([tm.xs_ext_avg for tm in spheres_lmax1]) / (np.pi * radius**2)
+xs_sca_lmax1 = np.array([tm.xs_sca_avg for tm in spheres_lmax1]) / (np.pi * radius ** 2)
+xs_ext_lmax1 = np.array([tm.xs_ext_avg for tm in spheres_lmax1]) / (np.pi * radius ** 2)
 
@@ -44,8 +44,3 @@ tm = spheres[-1]
 pcm = ax.pcolormesh(
-    grid[0, :, 2],
-    grid[:, 0, 0],
-    intensity.T,
-    shading="nearest",
-    vmin=0,
-    vmax=1,
+    grid[0, :, 2], grid[:, 0, 0], intensity.T, shading="nearest", vmin=0, vmax=1,
 )
@@ -52,0 +47,0 @@ cb = plt.colorbar(pcm)

docs/index.rst

@@ -36,3 +36,3 @@ =================================
 to the underlying functions operating directly on the spherical and cylindrical
-T-matrices or the Q-matrices based on the plane wave solutions.
+T-matrices or the S-matrices based on the plane wave solutions.
 
@@ -39,0 +39,0 @@ .. todo:: clean up intro

docs/operators.rst

@@ -65,14 +65,3 @@ =========
     PhysicsArray(
-        [[-1.+1.2246468e-16j,  0.+0.0000000e+00j,  0.+0.0000000e+00j,
-           0.+0.0000000e+00j,  0.+0.0000000e+00j,  0.+0.0000000e+00j],
-         [ 0.+0.0000000e+00j, -1.+1.2246468e-16j,  0.+0.0000000e+00j,
-           0.+0.0000000e+00j,  0.+0.0000000e+00j,  0.+0.0000000e+00j],
-         [ 0.+0.0000000e+00j,  0.+0.0000000e+00j,  1.+0.0000000e+00j,
-           0.+0.0000000e+00j,  0.+0.0000000e+00j,  0.+0.0000000e+00j],
-         [ 0.+0.0000000e+00j,  0.+0.0000000e+00j,  0.+0.0000000e+00j,
-           1.+0.0000000e+00j,  0.+0.0000000e+00j,  0.+0.0000000e+00j],
-         [ 0.-0.0000000e+00j,  0.+0.0000000e+00j,  0.-0.0000000e+00j,
-           0.+0.0000000e+00j, -1.-1.2246468e-16j,  0.+0.0000000e+00j],
-         [ 0.+0.0000000e+00j,  0.-0.0000000e+00j,  0.+0.0000000e+00j,
-           0.-0.0000000e+00j,  0.+0.0000000e+00j, -1.-1.2246468e-16j]],
+        [[...]],
         basis=SphericalWaveBasis(
@@ -95,8 +84,3 @@ pidx=[0 0 0 0 0 0],
     PhysicsArray(
-        [[1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
-         [0.+0.j, 1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
-         [0.+0.j, 0.+0.j, 1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
-         [0.+0.j, 0.+0.j, 0.+0.j, 1.+0.j, 0.+0.j, 0.+0.j],
-         [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 1.+0.j, 0.+0.j],
-         [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 1.+0.j]],
+        [[...]],
         basis=SphericalWaveBasis(
@@ -135,20 +119,3 @@ pidx=[0 0 0 0 0 0],
     PhysicsArray(
-        [[ 0.13702594-1.38777878e-17j,  0.        +0.00000000e+00j,
-          -0.02142403-4.28480668e-02j,  0.        +0.00000000e+00j,
-          -0.01514908+2.01987724e-02j,  0.        +0.00000000e+00j],
-         [ 0.        +0.00000000e+00j,  0.13702594+0.00000000e+00j,
-           0.        +0.00000000e+00j, -0.02142403-4.28480668e-02j,
-           0.        +0.00000000e+00j, -0.01514908+2.01987724e-02j],
-         [-0.02142403+4.28480668e-02j,  0.        +0.00000000e+00j,
-           0.07137993-6.84670061e-18j,  0.        +0.00000000e+00j,
-           0.02142403+4.28480668e-02j,  0.        +0.00000000e+00j],
-         [ 0.        +0.00000000e+00j, -0.02142403+4.28480668e-02j,
-           0.        +0.00000000e+00j,  0.07137993-1.52119906e-17j,
-           0.        +0.00000000e+00j,  0.02142403+4.28480668e-02j],
-         [-0.01514908-2.01987724e-02j,  0.        +0.00000000e+00j,
-           0.02142403-4.28480668e-02j,  0.        +0.00000000e+00j,
-           0.13702594+6.93889390e-18j,  0.        +0.00000000e+00j],
-         [ 0.        +0.00000000e+00j, -0.01514908-2.01987724e-02j,
-           0.        +0.00000000e+00j,  0.02142403-4.28480668e-02j,
-           0.        +0.00000000e+00j,  0.13702594-6.93889390e-18j]],
+        [[...]],
         basis=SphericalWaveBasis(
@@ -237,5 +204,4 @@ pidx=[0 0 0 0 0 0],
     PhysicsArray(
-        [ 3.06998012e-01-3.75964133e-17j, -2.76298211e+00+3.38367720e-16j,
-         -7.97540364e-17-1.30248226e+00j, -7.97540364e-17-1.30248226e+00j,
-         -2.76298211e+00+0.00000000e+00j,  3.06998012e-01+0.00000000e+00j],
+        [ 0.307-0.j   , -2.763+0.j   , -0.   -1.302j, -0.   -1.302j,
+         -2.763+0.j   ,  0.307+0.j   ],
         basis=SphericalWaveBasis(
@@ -271,5 +237,4 @@ pidx=[0 0 0 0 0 0],
     PhysicsArray(
-        [ 0.00000000e+00+0.00000000e+00j,  5.86797393e-17+3.19437623e-01j,
-          0.00000000e+00+0.00000000e+00j,  8.10453459e-01+3.79855139e-18j,
-          0.00000000e+00+0.00000000e+00j, -1.95599131e-17+3.19437623e-01j],
+        [ 0.  +0.j   ,  0.  +0.319j,  0.  +0.j   ,  0.81+0.j   ,
+          0.  +0.j   , -0.  +0.319j],
         basis=SphericalWaveBasis(
@@ -299,5 +264,4 @@ pidx=[0 0 0 0 0 0],
     PhysicsArray(
-        [0.        +0.j        , 1.4655919 +0.31943762j,
-         0.        +0.j        , 0.81045346+1.26220648j,
-         0.        +0.j        , 1.4655919 +0.31943762j],
+        [0.   +0.j   , 1.466+0.319j, 0.   +0.j   , 0.81 +1.262j,
+         0.   +0.j   , 1.466+0.319j],
         basis=SphericalWaveBasis(
@@ -327,8 +291,4 @@ pidx=[0 0 0 0 0 0],
     PhysicsArray(
-        [0.        +0.00000000e+00j, 0.        +0.00000000e+00j,
-         0.        +0.00000000e+00j, 0.90350604-7.59710279e-18j,
-         0.        +0.00000000e+00j, 0.        +0.00000000e+00j,
-         0.        +0.00000000e+00j, 0.        +0.00000000e+00j,
-         0.        +0.00000000e+00j, 0.90350604-7.59710279e-18j,
-         0.        +0.00000000e+00j, 0.        +0.00000000e+00j],
+        [0.   +0.j, 0.   +0.j, 0.   +0.j, 0.904-0.j, 0.   +0.j, 0.   +0.j,
+         0.   +0.j, 0.   +0.j, 0.   +0.j, 0.904-0.j, 0.   +0.j, 0.   +0.j],
         basis=SphericalWaveBasis(
@@ -411,5 +371,4 @@ pidx=[0 0 0 0 0 0 1 1 1 1 1 1],
     PhysicsArray(
-        [ 0.        +0.00000000e+00j,  0.11490348-1.40716185e-17j,
-          0.        +0.00000000e+00j, -0.44005059+5.38906541e-17j,
-          0.        +0.00000000e+00j,  0.76519769+0.00000000e+00j],
+        [ 0.   +0.j,  0.115-0.j,  0.   +0.j, -0.44 +0.j,  0.   +0.j,
+          0.765+0.j],
         basis=CylindricalWaveBasis(
@@ -436,4 +395,3 @@ pidx=[0 0 0 0 0 0],
     PhysicsArray(
-        [0.        +0.j, 0.        +0.j, 0.        +0.j, 0.        +0.j,
-         0.        +0.j, 3.06998012+0.j],
+        [0.  +0.j, 0.  +0.j, 0.  +0.j, 0.  +0.j, 0.  +0.j, 3.07+0.j],
         basis=SphericalWaveBasis(
@@ -462,3 +420,3 @@ pidx=[0 0 0 0 0 0],
 when using spherical or cylindrical waves. These expansions are needed to compute the
-electromagnetic interaction between particles within a lattice. It is assumed the the
+electromagnetic interaction between particles within a lattice. It is assumed that the
 given basis with singular modes are repeated periodically in the given lattice
@@ -469,7 +427,6 @@ structure. Then, these fields are expanded as regular fields in a single unit cell.
 
-    >>> cyw = cylindrical_wave(0, 1, 0, k0=1, material=1, modetype="regular")
+    >>> cyw = cylindrical_wave(0, 1, 0, k0=1, material=1, modetype="singular")
     >>> cyw.expandlattice(1, 0)
     PhysicsArray(
-        [0.+0.j        , 2.-3.8655259j , 0.+0.j        , 0.+0.j        ,
-         0.+0.j        , 1.+1.23397896j],
+        [0.+0.j   , 2.-3.866j, 0.+0.j   , 0.+0.j   , 0.+0.j   , 1.+1.234j],
         basis=CylindricalWaveBasis(
@@ -491,4 +448,4 @@ pidx=[0 0 0 0 0 0],
     PhysicsArray(
-        [ 0.+0.j        ,  0.+0.j        ,  0.+0.j        , -1.+7.72202545j,
-          0.+0.j        ,  0.+0.j        ],
+        [ 0.+0.j   ,  0.+0.j   ,  0.+0.j   , -1.+7.722j,  0.+0.j   ,
+          0.+0.j   ],
         basis=SphericalWaveBasis(
@@ -520,8 +477,7 @@ pidx=[0 0 0 0 0 0],
     PhysicsArray(
-        [ 0.00000000e+00+0.j        , -5.72699981e-18+0.093529j  ,
-          0.00000000e+00+0.j        , -9.44693623e-18+0.15428018j,
-          0.00000000e+00+0.j        , -6.57692473e-19+0.01074093j],
+        [ 0.+0.j   , -0.+0.094j,  0.+0.j   , -0.+0.154j,  0.+0.j   ,
+         -0.+0.011j],
         basis=CylindricalWaveBasis(
         pidx=[0 0 0 0 0 0],
-        kz=[-0.7975979 -0.7975979  0.1        0.1        0.9975979  0.9975979],
+        kz=[-0.798 -0.798  0.1    0.1    0.998  0.998],
         m=[0 0 0 0 0 0],
@@ -548,12 +504,7 @@ pol=[1 0 1 0 1 0],
     PhysicsArray(
-        [ 0.00000000e+00+0.j        , -2.72638852e-19+0.00445253j,
-          0.00000000e+00+0.j        , -5.70665864e-18+0.09319681j,
-          0.00000000e+00+0.j        , -5.70665864e-18+0.09319681j,
-          0.00000000e+00+0.j        , -3.90671188e-17+0.63801447j,
-          0.00000000e+00+0.j        , -3.58703457e-18+0.05858072j],
+        [ 0.+0.j   , -0.+0.004j,  0.+0.j   , -0.+0.093j,  0.+0.j   ,
+         -0.+0.093j,  0.+0.j   , -0.+0.638j,  0.+0.j   , -0.+0.059j],
         basis=PlaneWaveBasisByComp(
-        kx=[ 0.1        0.1        0.1        0.1        0.1        0.1
-      0.9975979  0.9975979 -0.7975979 -0.7975979],
-        ky=[ 0.         0.         0.8975979  0.8975979 -0.8975979 -0.8975979
-      0.         0.         0.         0.       ],
+        kx=[ 0.1    0.1    0.1    0.1    0.1    0.1    0.998  0.998 -0.798 -0.798],
+        ky=[ 0.     0.     0.898  0.898 -0.898 -0.898  0.     0.     0.     0.   ],
         pol=[1 0 1 0 1 0 1 0 1 0],
@@ -579,12 +530,8 @@ ),
     PhysicsArray(
-        [0.        +0.j        , 0.04081633-0.0041022j ,
-         0.        +0.j        , 0.04081633-0.58781404j,
-         0.        +0.j        , 0.04081633+0.05397146j,
-         0.        +0.j        , 0.        +0.j        ,
-         0.        +0.j        , 0.        +0.j        ],
+        [0.   +0.j   , 0.286-0.029j, 0.   +0.j   , 0.286-4.115j,
+         0.   +0.j   , 0.286+0.378j, 0.   +0.j   , 0.   +0.j   ,
+         0.   +0.j   , 0.   +0.j   ],
         basis=PlaneWaveBasisByComp(
-        kz=[ 0.         0.         0.         0.         0.         0.
-      0.8975979  0.8975979 -0.8975979 -0.8975979],
-        kx=[ 0.1        0.1        0.9975979  0.9975979 -0.7975979 -0.7975979
-      0.1        0.1        0.1        0.1      ],
+        kz=[ 0.     0.     0.     0.     0.     0.     0.898  0.898 -0.898 -0.898],
+        kx=[ 0.1    0.1    0.998  0.998 -0.798 -0.798  0.1    0.1    0.1    0.1  ],
         pol=[1 0 1 0 1 0 1 0 1 0],
@@ -594,4 +541,3 @@ ),
         kpar=WaveVector(nan, nan, 0.0),
-        lattice=Lattice([[7. 0.]
-             [0. 7.]], alignment='zx'),
+        lattice=Lattice(7.0, alignment='x'),
         material=Material(1, 1, 0),
@@ -614,4 +560,3 @@ modetype='up',
     PhysicsArray(
-        [ 0.        ,  0.        ,  0.70710678, -0.70710678,  0.        ,
-          0.        ],
+        [ 0.   ,  0.   ,  0.707, -0.707,  0.   ,  0.   ],
         basis=SphericalWaveBasis(
@@ -635,3 +580,3 @@ pidx=[0 0 0 0 0 0],
 :math:`(x', y', z') = (z, x, y)`. This operation is implemented separately as
-permutation, meaning the the axes labels get permuted.
+permutation, meaning the axes labels get permuted.
 
@@ -645,5 +590,5 @@ .. doctest::
         basis=PlaneWaveBasisByUnitVector(
-        qx=[0.28571429 0.28571429],
-        qy=[0.42857143 0.42857143],
-        qz=[0.85714286 0.85714286],
+        qx=[0.286 0.286],
+        qy=[0.429 0.429],
+        qz=[0.857 0.857],
         pol=[1 0],
@@ -654,7 +599,7 @@ ),
     PhysicsArray(
-        [ 0.        +0.j        , -0.49613894-0.86824314j],
+        [ 0.   +0.j   , -0.789+0.614j],
         basis=PlaneWaveBasisByUnitVector(
-        qx=[0.85714286 0.85714286],
-        qy=[0.28571429 0.28571429],
-        qz=[0.42857143 0.42857143],
+        qx=[0.857 0.857],
+        qy=[0.286 0.286],
+        qz=[0.429 0.429],
         pol=[1 0],
@@ -683,4 +628,4 @@ ),
     PhysicsArray(
-        [[0.+0.00000000e+00j, 0.+0.00000000e+00j, 0.+1.62867504e-01j],
-         [0.-8.08245652e-18j, 0.-7.35758865e-02j, 0.+1.31996532e-01j]],
+        [[0.+0.j   , 0.+0.j   , 0.+0.163j],
+         [0.-0.j   , 0.-0.074j, 0.+0.132j]],
     )

docs/params.rst

@@ -25,3 +25,3 @@ .. testsetup::
 precision when truncated to a finite number of modes. The chosen finite number of
-modes is given in the the the classes :class:`~treams.SphericalWaveBasis`,
+modes is given in the classes :class:`~treams.SphericalWaveBasis`,
 :class:`~treams.CylindricalWaveBasis`, and :class:`~treams._core.PlaneWaveBasis`, which
@@ -200,6 +200,4 @@ are all children of the base call :class:`~treams._core.BasisSet`.
     PlaneWaveBasisByComp(
-        kx=[ 0.          0.          0.          0.          0.          0.
-      6.28318531  6.28318531 -6.28318531 -6.28318531],
-        ky=[ 0.          0.          6.28318531  6.28318531 -6.28318531 -6.28318531
-      0.          0.          0.          0.        ],
+        kx=[ 0.     0.     0.     0.     0.     0.     6.283  6.283 -6.283 -6.283],
+        ky=[ 0.     0.     6.283  6.283 -6.283 -6.283  0.     0.     0.     0.   ],
         pol=[1 0 1 0 1 0 1 0 1 0],
@@ -323,4 +321,4 @@ )
     >>> treams.Lattice.hexagonal(2)
-    Lattice([[2.         0.        ]
-             [1.         1.73205081]], alignment='xy')
+    Lattice([[2.    0.   ]
+             [1.    1.732]], alignment='xy')
 
@@ -364,4 +362,4 @@ creates a hexagonal lattice with sidelength 2. It's also possible to extract a
     >>> treams.Lattice([1, 1]).reciprocal
-    array([[ 6.28318531, -0.        ],
-           [-0.        ,  6.28318531]])
+    array([[ 6.283, -0.   ],
+           [-0.   ,  6.283]])
 
@@ -368,0 +366,0 @@ Phase vector

docs/smatrix.rst

@@ -0,1 +1,5 @@
+.. highlight:: python
+
+.. only:: builder_html
+
 ==========================
@@ -39,8 +43,83 @@ S-Matrices for plane waves
 
+.. literalinclude:: examples/array_spheres.py
+   :language: python
+   :lines: 6-12
 
+Beforehand, we define all the necessary parameters. First the wave numbers, then the
+parameters of the slab, and finally those for the lattice and the spheres. Then, we can
+use a simple loop to solve the system for all wave numbers.
+
+.. literalinclude:: examples/array_spheres.py
+   :language: python
+   :lines: 14-27
+
+We set some oblique incidence and the array of spheres. Then, we define a linearly
+polarized plane wave and the needed S-matrices: a slab, the distance between the
+top interface of the slab to the center of the sphere array, and the array in the
+S-matrix representation itself.
+
+.. plot:: examples/array_spheres.py
+
 From cylindrical T-matrix gratings
 ==================================
 
+Now, we want to perform a small test of the methods. Instead of creating the
+two-dimensional sphere array right away, we intermediately create a one-dimensional
+array, then calculate cylindrical T-matrices, and place them in a second
+one-dimensional lattice, thereby, obtaining the S-matrix from the previous section.
 
-Band structures
-===============
+.. literalinclude:: examples/array_spheres_tmatrixc.py
+   :language: python
+   :lines: 6-13
+
+The definition of the parameters is quite similar. We store the lattice pitch for later
+use separately and define the maximal order :code:`mmax` separately.
+
+.. literalinclude:: examples/array_spheres_tmatrixc.py
+   :language: python
+   :lines: 15-32
+
+The first half of the loop is now a little bit different. After creating the spheres
+we solve the interaction along the z-direction, then create the cylindrical T-matrix
+and finally calculate the interaction along the x-direction. The second half is the
+same as in the previous calculation.
+
+The most important aspect to note here, is that the method
+:meth:`treams.SMatrix.from_array` implicitly converts the lattice in the z-x-plane to
+a lattice in the x-y-plane.
+
+.. plot:: examples/array_spheres_tmatrixc.py
+
+
+Band structure
+==============
+
+Finally, we want to compute the band structure of a system consisting of the periodic
+repetition of an S-matrix along the z-direction. In principle, one can obtain this
+band structure also from the lattice interaction in the T-matrix, but calculating it
+from the S-matrix has two benefits. First, more complex systems can be analyzed, because
+slabs and objects described by cylindrical T-matrices can be included. Second, one only
+defines :math:`k_0`, :math:`k_x`, and :math:`k_y`. Then, the result of the calculation
+are all components :math:`k_z` and the plane wave decomposition of the polarization from
+an eigenvalue decomposition. So, instead of a 4-dimensional parameter sweep only a
+3-dimensional sweep is necessary decreasing the computation time. The downside is, that
+one is restricted to unit cells, where one vector points along the z-axis and is
+perpendicular to the other two.
+
+We take the array of spheres on top of a slab and continue this one infinitely along the
+z-axis. Thus, the setup is
+
+.. literalinclude:: examples/band_structure.py
+   :language: python
+   :lines: 6-13
+
+where :code:`az` is the length of the lattice vector pointing in the z-direction. With a
+simple loop we can get the band structure for :math:`k_x = 0 = k_y`.
+
+.. literalinclude:: examples/band_structure.py
+   :language: python
+   :lines: 15-31
+
+which looks as follows, after a cut on the imaginary part of the :math:`k_z` component.
+
+.. plot:: examples/band_structure.py

docs/tmatrixc.rst

@@ -79,3 +79,3 @@ .. highlight:: python
 
-.. literalinclude:: examples/cylinder_tmatrixc.py
+.. literalinclude:: examples/chain_tmatrixc.py
     :language: python
@@ -87,3 +87,3 @@ :lines: 6-15
 
-.. literalinclude:: examples/cylinder_tmatrixc.py
+.. literalinclude:: examples/chain_tmatrixc.py
     :language: python
@@ -98,3 +98,3 @@ :lines: 17-19
 
-.. literalinclude:: examples/cylinder_tmatrixc.py
+.. literalinclude:: examples/chain_tmatrixc.py
     :language: python
@@ -153,3 +153,3 @@ :lines: 21-27
 
-.. literalinclude:: examples/cluster_tmatrixc.py
+.. literalinclude:: examples/grating_tmatrixc.py
     :language: python
@@ -161,3 +161,3 @@ :lines: 6-21
 
-.. literalinclude:: examples/cluster_tmatrixc.py
+.. literalinclude:: examples/grating_tmatrixc.py
     :language: python
@@ -169,3 +169,3 @@ :lines: 23-32
 
-.. literalinclude:: examples/cluster_tmatrixc.py
+.. literalinclude:: examples/grating_tmatrixc.py
     :language: python
@@ -172,0 +172,0 @@ :lines: 34-43

MANIFEST.in

@@ -1,3 +0,3 @@
-global-exclude *.pyx *.pxd .coveragerc .gitignore .gitlab-ci.yml environment.yml
-include treams/lattice/cython_lattice.pxd
-include treams/special/cython_special.pxd
+global-exclude *.pyx *.pxd .gitignore
+include src/treams/lattice/cython_lattice.pxd
+include src/treams/special/cython_special.pxd

PKG-INFO

 Metadata-Version: 2.1
 Name: treams
-Version: 0.3.0
-Summary: "Periodic T-matrix scattering algorithms"
-Home-page: https://git.scc.kit.edu/photonics/treams
+Version: 0.3.1
+Summary: "T-matrix scattering code for nanophotonic computations"
+Home-page: https://git.scc.kit.edu/tfp-photonics/treams
 Author: Dominik Beutel
@@ -19,5 +19,6 @@ Author-email: [e-mail]
 Classifier: Programming Language :: Python :: 3 :: Only
-Classifier: Programming Language :: Python :: 3.7
 Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Cython
@@ -29,9 +30,19 @@ Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Topic :: Scientific/Engineering :: Physics
-Requires-Python: >=3.7
+Requires-Python: >=3.8
 Description-Content-Type: text/markdown
-Provides-Extra: test
+Provides-Extra: coverage
 Provides-Extra: docs
 Provides-Extra: io
+Provides-Extra: test
 License-File: LICENSE
 
+![Version](https://img.shields.io/github/v/tag/tfp-photonics/treams)
+[![PyPI](https://img.shields.io/pypi/v/treams)](https://pypi.org/project/treams)
+![License](https://img.shields.io/github/license/tfp-photonics/treams)
+![build](https://github.com/tfp-photonics/treams/actions/workflows/build.yml/badge.svg)
+[![docs](https://github.com/tfp-photonics/treams/actions/workflows/docs.yml/badge.svg)](https://tfp-photonics.github.io/treams)
+![doctests](https://github.com/tfp-photonics/treams/actions/workflows/doctests.yml/badge.svg)
+![tests](https://github.com/tfp-photonics/treams/actions/workflows/tests.yml/badge.svg)
+[![coverage](https://img.shields.io/endpoint?url=https%3A%2F%2Fraw.githubusercontent.com%2Ftfp-photonics%2Ftreams%2Fhtmlcov%2Fendpoint.json)](https://htmlpreview.github.io/?https://github.com/tfp-photonics/treams/blob/htmlcov/index.html)
+
 # treams
@@ -50,3 +61,3 @@
 ```sh
-pip install git+https://github.com/tfp-photonics/treams.git
+pip install treams
 ```
@@ -57,11 +68,2 @@
 
-### Running on Windows
-
-For Windows, there are currently two tested ways how to install treams. The first option
-is using the
-[Windows Subsystem for Linux (WSL)](https://docs.microsoft.com/en-us/windows/wsl/install).
-Within WSL treams can be installed just as described above. The second option, that was
-tested is using [MSYS2](https://www.msys2.org/) with the ``mingw64`` environment.
-Likely, other python versions based on ``mingw-w64`` might also work.
-
 ## Documentation
@@ -68,0 +70,0 @@

pyproject.toml

@@ -19,1 +19,7 @@ [build-system]
 extension-pkg-whitelist = "treams"
+
+[tool.cibuildwheel]
+archs = ["auto64"]
+skip = ["pp*", "*musllinux*"]
+test-command = "python -m pytest {project}/tests/unit"
+test-extras = ["test", "io"]

README.md

@@ -0,1 +1,10 @@
+![Version](https://img.shields.io/github/v/tag/tfp-photonics/treams)
+[![PyPI](https://img.shields.io/pypi/v/treams)](https://pypi.org/project/treams)
+![License](https://img.shields.io/github/license/tfp-photonics/treams)
+![build](https://github.com/tfp-photonics/treams/actions/workflows/build.yml/badge.svg)
+[![docs](https://github.com/tfp-photonics/treams/actions/workflows/docs.yml/badge.svg)](https://tfp-photonics.github.io/treams)
+![doctests](https://github.com/tfp-photonics/treams/actions/workflows/doctests.yml/badge.svg)
+![tests](https://github.com/tfp-photonics/treams/actions/workflows/tests.yml/badge.svg)
+[![coverage](https://img.shields.io/endpoint?url=https%3A%2F%2Fraw.githubusercontent.com%2Ftfp-photonics%2Ftreams%2Fhtmlcov%2Fendpoint.json)](https://htmlpreview.github.io/?https://github.com/tfp-photonics/treams/blob/htmlcov/index.html)
+
 # treams
@@ -14,3 +23,3 @@
 ```sh
-pip install git+https://github.com/tfp-photonics/treams.git
+pip install treams
 ```
@@ -21,11 +30,2 @@
 
-### Running on Windows
-
-For Windows, there are currently two tested ways how to install treams. The first option
-is using the
-[Windows Subsystem for Linux (WSL)](https://docs.microsoft.com/en-us/windows/wsl/install).
-Within WSL treams can be installed just as described above. The second option, that was
-tested is using [MSYS2](https://www.msys2.org/) with the ``mingw64`` environment.
-Likely, other python versions based on ``mingw-w64`` might also work.
-
 ## Documentation
@@ -32,0 +32,0 @@

setup.cfg

@@ -5,4 +5,4 @@ [metadata]
 author_email = [e-mail]
-url = https://git.scc.kit.edu/photonics/treams
-description = "Periodic T-matrix scattering algorithms"
+url = https://git.scc.kit.edu/tfp-photonics/treams
+description = "T-matrix scattering code for nanophotonic computations"
 license = MIT
@@ -21,5 +21,6 @@ long_description = file: README.md
 	Programming Language :: Python :: 3 :: Only
-	Programming Language :: Python :: 3.7
 	Programming Language :: Python :: 3.8
 	Programming Language :: Python :: 3.9
+	Programming Language :: Python :: 3.10
+	Programming Language :: Python :: 3.11
 	Programming Language :: Cython
@@ -33,3 +34,5 @@ Programming Language :: Python :: Implementation :: CPython
 [options]
-python_requires = >= 3.7
+python_requires = >= 3.8
+package_dir = 
+	=src
 packages = 
@@ -44,10 +47,12 @@ treams
 [options.extras_require]
-test = 
-	pytest
+coverage = 
+	Cython
 	pytest-cov
 docs = 
+	matplotlib
 	sphinx
-	matplotlib
 io = 
 	h5py
+test = 
+	pytest
 
@@ -54,0 +59,0 @@ [sdist]

setup.py

@@ -15,3 +15,2 @@ """Packaging of treams."""
 if os.name == "nt":
-
     link_args = [
@@ -43,3 +42,2 @@ "-static-libgcc",
 
-
 else:
@@ -96,3 +94,3 @@ build_ext = _build_ext
 extensions = [
-    Extension(name, [f"{name.replace('.', '/')}.pyx"], **keys)
+    Extension(name, [f"src/{name.replace('.', '/')}.pyx"], **keys)
     for name in extension_names
@@ -99,0 +97,0 @@ ]

tests/unit/test_coeffs.py

@@ -1,2 +0,1 @@
-import pytest
 import numpy as np
@@ -64,8 +63,8 @@
             [
-                [[1.042449234640745, 0], [0, 1.005929062176551],],
-                [[-0.042449234640745, 0], [0, -0.005929062176551],],
+                [[1.042449234640745, 0], [0, 1.005929062176551]],
+                [[-0.042449234640745, 0], [0, -0.005929062176551]],
             ],
             [
-                [[0.042449234640745, 0], [0, 0.005929062176551],],
-                [[0.957550765359255, 0], [0, 0.994070937823449],],
+                [[0.042449234640745, 0], [0, 0.005929062176551]],
+                [[0.957550765359255, 0], [0, 0.994070937823449]],
             ],
@@ -92,3 +91,3 @@ ]
                 ],
-                [[0, 0.6], [0.6, 0.235 + 0.595294044989533j],],
+                [[0, 0.6], [0.6, 0.235 + 0.595294044989533j]],
             ],
@@ -95,0 +94,0 @@ [

tests/unit/test_core.py

@@ -138,5 +138,3 @@ import numpy as np
             zip(
-                3 * [1, 1, 2, 2],
-                [-1, -1, -1, -1, 0, 0, 0, 0, 1, 1, 1, 1],
-                6 * [1, 0],
+                3 * [1, 1, 2, 2], [-1, -1, -1, -1, 0, 0, 0, 0, 1, 1, 1, 1], 6 * [1, 0],
             ),
@@ -515,7 +513,2 @@ )
 
-    def test_lattice_fail(self):
-        b = treams.PlaneWaveBasisByComp.diffr_orders([0, 0], np.eye(2), 4)
-        with pytest.raises(ValueError):
-            treams.PhysicsArray([1, 2], lattice=treams.Lattice(2, "x"), basis=b)
-
     def test_matmul(self):
@@ -537,9 +530,5 @@ b = treams.SphericalWaveBasis([[1, 0, 0], [1, 0, 1]])
         x = p.changepoltype.apply_left()
-        print(repr(x))
         assert (
             x == np.sqrt(0.5) * np.array([[-1, 1], [1, 1]])
-        ).all() and x.poltype == (
-            "parity",
-            "helicity",
-        )
+        ).all() and x.poltype == ("parity", "helicity",)
 
@@ -552,6 +541,3 @@ def test_changepoltype_inv(self):
             x == np.sqrt(0.5) * np.array([[-1, 1], [1, 1]])
-        ).all() and x.poltype == (
-            "helicity",
-            "parity",
-        )
+        ).all() and x.poltype == ("helicity", "parity",)
 
@@ -558,0 +544,0 @@ def test_efield(self):

tests/unit/test_io.py

 import os
 import tempfile
 
-import gmsh
 import h5py
 import numpy as np
+import pytest
 
@@ -12,3 +12,6 @@ import treams
 
[e-mail]
 def test_meshspheres():
+    import gmsh
+
     gmsh.initialize()
@@ -15,0 +18,0 @@ gmsh.model.add("spheres")

tests/unit/test_lattice_module.py

@@ -450,11 +450,3 @@ import numpy as np
             la.diffr_orders_circle([[1, 0], [0.5, np.sqrt(0.75)]], 1.1),
-            [
-                [0, 0],
-                [0, 1],
-                [0, -1],
-                [1, 0],
-                [-1, 0],
-                [1, -1],
-                [-1, 1],
-            ],
+            [[0, 0], [0, 1], [0, -1], [1, 0], [-1, 0], [1, -1], [-1, 1],],
         )
@@ -461,0 +453,0 @@

tests/unit/test_lattice.py

@@ -0,5 +1,4 @@
+import numpy as np
 import pytest
 
-import numpy as np
-
 from treams import Lattice, WaveVector
@@ -44,3 +43,3 @@
         assert Lattice.hexagonal(1, 2) == Lattice(
-            [[1, 0, 0], [0.5, 0.75**0.5, 0], [0, 0, 2]]
+            [[1, 0, 0], [0.5, 0.75 ** 0.5, 0], [0, 0, 2]]
         )
@@ -47,0 +46,0 @@

tests/unit/test_operators.py

@@ -255,6 +255,3 @@ import numpy as np
         y = treams.PhysicsArray(
-            [[1, 0], [0, 0]],
-            basis=(a, b),
-            k0=2.5,
-            material=treams.Material(2, 2, 0),
+            [[1, 0], [0, 0]], basis=(a, b), k0=2.5, material=treams.Material(2, 2, 0),
         )
@@ -492,5 +489,3 @@ assert np.all(np.abs(x - y) < 1e-14) and x.ann == y.ann
         y = treams.PhysicsArray(
-            [[-np.sqrt(0.5), 0], [0, 0]],
-            basis=(b, b),
-            poltype=("helicity", "parity"),
+            [[-np.sqrt(0.5), 0], [0, 0]], basis=(b, b), poltype=("helicity", "parity"),
         )
@@ -505,5 +500,3 @@ assert np.all(np.abs(x - y) < 1e-14) and x.ann == y.ann
         y = treams.PhysicsArray(
-            [[-np.sqrt(0.5), 0], [0, 0]],
-            basis=(a, b),
-            poltype=("helicity", "parity"),
+            [[-np.sqrt(0.5), 0], [0, 0]], basis=(a, b), poltype=("helicity", "parity"),
         )
@@ -517,5 +510,3 @@ assert np.all(np.abs(x - y) < 1e-14) and x.ann == y.ann
         y = treams.PhysicsArray(
-            [[-np.sqrt(0.5), 0], [0, 0]],
-            basis=(b, b),
-            poltype=("helicity", "parity"),
+            [[-np.sqrt(0.5), 0], [0, 0]], basis=(b, b), poltype=("helicity", "parity"),
         )
@@ -529,5 +520,3 @@ assert np.all(np.abs(x - y) < 1e-14) and x.ann == y.ann
         y = treams.PhysicsArray(
-            [[-np.sqrt(0.5), 0], [0, 0]],
-            basis=(b, b),
-            poltype=("helicity", "parity"),
+            [[-np.sqrt(0.5), 0], [0, 0]], basis=(b, b), poltype=("helicity", "parity"),
         )
@@ -605,7 +594,3 @@ assert np.all(np.abs(x - y) < 1e-14) and x.ann == y.ann
         y = sc.vsw_rM(
-            [3, 1],
-            [-2, 1],
-            2 * k0 * rsph[..., 0],
-            rsph[..., 1],
-            rsph[..., 2],
+            [3, 1], [-2, 1], 2 * k0 * rsph[..., 0], rsph[..., 1], rsph[..., 2],
         )
@@ -622,16 +607,7 @@ assert np.all(sc.vsph2car(y, rsph).swapaxes(-1, -2) == x)
         x = treams.efield(
-            r,
-            basis=b,
-            k0=k0,
-            poltype="parity",
-            material=material,
-            modetype="singular",
+            r, basis=b, k0=k0, poltype="parity", material=material, modetype="singular",
         )
         rsph = sc.car2sph(r[None] - positions)
         y = sc.vsw_N(
-            [3, 1],
-            [-2, 1],
-            2 * k0 * rsph[..., 0],
-            rsph[..., 1],
-            rsph[..., 2],
+            [3, 1], [-2, 1], 2 * k0 * rsph[..., 0], rsph[..., 1], rsph[..., 2],
         )
@@ -652,3 +628,3 @@ assert np.all(sc.vsph2car(y, rsph).swapaxes(-1, -2) == x)
             [-2, 1],
-            rcyl[..., 0] * [np.sqrt(16 - 0.3**2), np.sqrt(144 - 0.1**2)],
+            rcyl[..., 0] * [np.sqrt(16 - 0.3 ** 2), np.sqrt(144 - 0.1 ** 2)],
             rcyl[..., 1],
@@ -680,3 +656,3 @@ rcyl[..., 2],
             [-2, 1],
-            rcyl[..., 0] * [np.sqrt(16 - 0.3**2), np.sqrt(144 - 0.1**2)],
+            rcyl[..., 0] * [np.sqrt(16 - 0.3 ** 2), np.sqrt(144 - 0.1 ** 2)],
             rcyl[..., 1],
@@ -701,3 +677,3 @@ rcyl[..., 2],
             [-2, 1],
-            rcyl[..., 0] * [np.sqrt(64 - 0.3**2), np.sqrt(64 - 0.1**2)],
+            rcyl[..., 0] * [np.sqrt(64 - 0.3 ** 2), np.sqrt(64 - 0.1 ** 2)],
             rcyl[..., 1],
@@ -716,8 +692,3 @@ rcyl[..., 2],
         x = treams.efield(
-            r,
-            basis=b,
-            k0=k0,
-            poltype="parity",
-            material=material,
-            modetype="singular",
+            r, basis=b, k0=k0, poltype="parity", material=material, modetype="singular",
         )
@@ -728,3 +699,3 @@ rcyl = sc.car2cyl(r[:, None] - positions)
             [-2, 1],
-            rcyl[..., 0] * [np.sqrt(64 - 0.3**2), np.sqrt(64 - 0.1**2)],
+            rcyl[..., 0] * [np.sqrt(64 - 0.3 ** 2), np.sqrt(64 - 0.1 ** 2)],
             rcyl[..., 1],
@@ -731,0 +702,0 @@ rcyl[..., 2],

tests/unit/test_pw.py

@@ -53,6 +53,10 @@ import numpy as np
     def test_h(self):
-        assert pw.permute_xyz(1, 2, 3, 0, 0) == (-3 - 2j * np.sqrt(14)) / np.sqrt(65)
+        assert isclose(
+            pw.permute_xyz(1, 2, 3, 0, 0), (-6 + 1j * np.sqrt(14)) / np.sqrt(50)
+        )
 
     def test_h_inv(self):
-        assert pw.permute_xyz(1, 2, 1j, 0, 0, inverse=True) == 3j / np.sqrt(15)
+        assert isclose(
+            pw.permute_xyz(1, 2, 1j, 0, 0, inverse=True), -1j * np.sqrt(5 / 3)
+        )
 
@@ -63,6 +67,8 @@ def test_h_opposite(self):
     def test_h_kxy_zero(self):
-        assert pw.permute_xyz(0, 0, -1, 0, 0) == 1
+        assert pw.permute_xyz(0, 0, -1, 0, 0) == 1j
 
     def test_p(self):
-        assert pw.permute_xyz(1, 2, 3, 0, 0, poltype="parity") == -3 / np.sqrt(65)
+        assert isclose(
+            pw.permute_xyz(1, 2, 3, 0, 0, poltype="parity"), -6 / np.sqrt(50)
+        )
 
@@ -75,5 +81,5 @@ def test_p_inv(self):
     def test_p_opposite(self):
-        assert pw.permute_xyz(1, 2, 3, 0, 1, poltype="parity") == 2j * np.sqrt(
+        assert pw.permute_xyz(1, 2, 3, 0, 1, poltype="parity") == -1j * np.sqrt(
             14
-        ) / np.sqrt(65)
+        ) / np.sqrt(50)
 
@@ -83,8 +89,8 @@ def test_p_inv_opposite(self):
             1, 2, 1j, 0, 1, poltype="parity", inverse=True
-        ) == -4j / np.sqrt(15)
+        ) == 4j / np.sqrt(15)
 
     def test_p_kxy_zero(self):
-        assert pw.permute_xyz(0, 0, -1, 0, 0, poltype="parity") == 1
+        assert pw.permute_xyz(0, 0, -1, 0, 0, poltype="parity") == 0
 
     def test_p_kxy_zero_opposite(self):
-        assert pw.permute_xyz(0, 0, -1, 0, 1, poltype="parity") == 0
+        assert pw.permute_xyz(0, 0, -1, 0, 1, poltype="parity") == -1j

tests/unit/test_special.py

@@ -0,4 +1,3 @@
+import numpy as np
 import pytest
-
-import numpy as np
 import scipy.special as ssc
@@ -180,3 +179,3 @@
             -9.748868790182211e181 - 3.9232159891109815e181j,
-            rel_tol=EPSSQ,
+            rel_tol=EPS,
         )
@@ -198,3 +197,3 @@
     def test_gamma(self):
-        assert sc.incgamma(4, 0) == 6
+        assert isclose(sc.incgamma(4, 0), 6)
 
@@ -355,3 +354,3 @@
         assert isclose(
-            sc.lpmv(3, 3, 0.1 + 0j), ssc.clpmn(3, 3, -0.1 + 0j, type=2)[0][3, 3,]
+            sc.lpmv(3, 3, 0.1 + 0j), ssc.clpmn(3, 3, -0.1 + 0j, type=2)[0][3, 3]
         )
@@ -1196,4 +1195,8 @@
     def test_1(self):
-        assert np.array_equal(
-            sc.car2sph([3, -4, 1]), [np.sqrt(26), np.arctan2(5, 1), -0.9272952180016122]
+        assert np.all(
+            np.abs(
+                sc.car2sph([3, -4, 1])
+                - [np.sqrt(26), np.arctan2(5, 1), -0.9272952180016122]
+            )
+            < EPSSQ
         )
@@ -1200,0 +1203,0 @@

tests/unit/test_sw.py

@@ -10,5 +10,5 @@ from treams import sw
     def test_h(self):
-        assert (
-            sw.periodic_to_cw(1, 2, 0, 3, 2, 0, 3, 2)
-            == -1.1890388773999045e-17 + 0.19418478507072567j
+        assert isclose(
+            sw.periodic_to_cw(1, 2, 0, 3, 2, 0, 3, 2),
+            -1.1890388773999045e-17 + 0.19418478507072567j,
         )
@@ -15,0 +15,0 @@

tests/unit/test_tmatrix.py

@@ -110,6 +110,3 @@ import copy
         xs = tm.xs(illu, 0.125)
-        assert isclose(
-            xs[0],
-            3.194830855171616,
-        ) and isclose(xs[1], 5.63547158)
+        assert isclose(xs[0], 3.194830855171616,) and isclose(xs[1], 5.63547158)
 
@@ -116,0 +113,0 @@

tests/unit/test_util.py

 import copy
 
+import numpy as np
 import pytest
-import numpy as np
 
@@ -406,3 +406,3 @@ from treams import util
         a = np.linalg.solve(m, b)
-        assert (a == [-1, 4]).all() and a.ann == ({"a": 1, "b": 1},)
+        assert np.all(np.abs(a - [-1, 4]) < 1e-14) and a.ann == ({"a": 1, "b": 1},)
 
@@ -413,3 +413,6 @@ def test_solve_multiple(self):
         a = np.linalg.solve(m, b)
-        assert (a == [[-1, 1], [4, 1]]).all() and a.ann == ({"a": 1, "b": 1}, {"a": 2})
+        assert np.all(np.abs(a - [[-1, 1], [4, 1]]) < 1e-14) and a.ann == (
+            {"a": 1, "b": 1},
+            {"a": 2},
+        )
 
@@ -416,0 +419,0 @@ def test_lstsq(self):

.github/workflows/ci.yml

-name: CI
-
-on:
-  push:
-    branches:    
-      - main
-
-jobs:
-  build:
-    runs-on: ubuntu-latest
-    steps:
-    - uses: actions/checkout@v3
-      with:
-        fetch-depth: 0
-    - uses: actions/setup-python@v4
-      with:
-        python-version: '3.10' 
-    - name: Build treams
-      run: python -m pip install -e .
-    - name: Build html
-      run: |
-        python -m pip install treams[docs]
-        sphinx-build -b html docs docs/_build/html
-    - name: Upload artifacts
-      uses: actions/upload-artifact@v3
-      with:
-        name: html-docs
-        path: docs/_build/html/
-    - name: Deploy
-      uses: peaceiris/actions-gh-pages@v3
-      if: github.ref == 'refs/heads/main'
-      with:
-        github_token: ${{ secrets.GITHUB_TOKEN }}
-        publish_dir: docs/_build/html

treams.egg-info/dependency_links.txt

-

treams.egg-info/PKG-INFO

-Metadata-Version: 2.1
-Name: treams
-Version: 0.3.0
-Summary: "Periodic T-matrix scattering algorithms"
-Home-page: https://git.scc.kit.edu/photonics/treams
-Author: Dominik Beutel
-Author-email: [e-mail]
-License: MIT
-Platform: Linux
-Platform: Windows
-Classifier: Development Status :: 4 - Beta
-Classifier: Intended Audience :: Science/Research
-Classifier: Operating System :: Microsoft :: Windows
-Classifier: Operating System :: POSIX :: Linux
-Classifier: Natural Language :: English
-Classifier: Programming Language :: Python
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3 :: Only
-Classifier: Programming Language :: Python :: 3.7
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Cython
-Classifier: Programming Language :: Python :: Implementation :: CPython
-Classifier: Topic :: Scientific/Engineering
-Classifier: Topic :: Scientific/Engineering :: Astronomy
-Classifier: Topic :: Scientific/Engineering :: Atmospheric Science
-Classifier: Topic :: Scientific/Engineering :: Physics
-Requires-Python: >=3.7
-Description-Content-Type: text/markdown
-Provides-Extra: test
-Provides-Extra: docs
-Provides-Extra: io
-License-File: LICENSE
-
-# treams
-
-The package `treams` provides a framework to simplify computations of the
-electromagnetic scattering of waves at finite and at periodic arrangements of particles
-based on the T-matrix method.
-
-## Installation
-
-### Installation using pip
-
-To install the package with pip, use
-
-```sh
-pip install git+https://github.com/tfp-photonics/treams.git
-```
-
-If you're using the system wide installed version of python, you might consider the
-``--user`` option.
-
-### Running on Windows
-
-For Windows, there are currently two tested ways how to install treams. The first option
-is using the
-[Windows Subsystem for Linux (WSL)](https://docs.microsoft.com/en-us/windows/wsl/install).
-Within WSL treams can be installed just as described above. The second option, that was
-tested is using [MSYS2](https://www.msys2.org/) with the ``mingw64`` environment.
-Likely, other python versions based on ``mingw-w64`` might also work.
-
-## Documentation
-
-The documentation can be found at https://tfp-photonics.github.io/treams.
-
-## Publications
-
-When using this code please cite:
-
-[D. Beutel, A. Groner, C. Rockstuhl, C. Rockstuhl, and I. Fernandez-Corbaton, Efficient Simulation of Biperiodic, Layered Structures Based on the T-Matrix Method, J. Opt. Soc. Am. B, JOSAB 38, 1782 (2021).](https://doi.org/10.1364/JOSAB.419645)
-
-Other relevant publications are
-* [D. Beutel, I. Fernandez-Corbaton, and C. Rockstuhl, Unified Lattice Sums Accommodating Multiple Sublattices for Solutions of the Helmholtz Equation in Two and Three Dimensions, Phys. Rev. A 107, 013508 (2023).](https://doi.org/10.1103/PhysRevA.107.013508)
-* [D. Beutel, P. Scott, M. Wegener, C. Rockstuhl, and I. Fernandez-Corbaton, Enhancing the Optical Rotation of Chiral Molecules Using Helicity Preserving All-Dielectric Metasurfaces, Appl. Phys. Lett. 118, 221108 (2021).](https://doi.org/10.1063/5.0050411)
-
-
-## Features
-
-* [x] T-matrix calculations using a spherical or cylindrical wave basis set
-* [x] Calculations in helicity and parity (TE/TM) basis
-* [x] Scattering from clusters of particles
-* [x] Scattering from particles and clusters arranged in 3d-, 2d-, and 1d-lattices
-* [x] Calculation of light propagation in stratified media
-* [x] Band calculation in crystal structures

treams.egg-info/requires.txt

-numpy
-scipy>=1.6
-
-[docs]
-sphinx
-matplotlib
-
-[io]
-h5py
-
-[test]
-pytest
-pytest-cov

treams.egg-info/SOURCES.txt

-LICENSE
-MANIFEST.in
-README.md
-conftest.py
-pyproject.toml
-setup.cfg
-setup.py
-.github/workflows/ci.yml
-docs/Makefile
-docs/about.rst
-docs/conf.py
-docs/dev.rst
-docs/gettingstarted.rst
-docs/index.rst
-docs/intro.rst
-docs/make.bat
-docs/maxwell.rst
-docs/operators.rst
-docs/params.rst
-docs/physicsarray.rst
-docs/smatrix.rst
-docs/theory.rst
-docs/tmatrix.rst
-docs/tmatrixc.rst
-docs/treams.rst
-docs/_static/custom.css
-docs/_templates/autosummary/attribute.rst
-docs/_templates/autosummary/bare-module.rst
-docs/_templates/autosummary/class.rst
-docs/_templates/autosummary/cython-module.rst
-docs/_templates/autosummary/module.rst
-docs/examples/chain.py
-docs/examples/cluster.py
-docs/examples/crystal.py
-docs/examples/grid.py
-docs/examples/sphere.py
-tests/integration/test_compare_direct_sum.py
-tests/unit/test_coeffs.py
-tests/unit/test_core.py
-tests/unit/test_cw.py
-tests/unit/test_io.py
-tests/unit/test_lattice.py
-tests/unit/test_lattice_module.py
-tests/unit/test_misc.py
-tests/unit/test_operators.py
-tests/unit/test_pw.py
-tests/unit/test_smatrix.py
-tests/unit/test_special.py
-tests/unit/test_sw.py
-tests/unit/test_tmatrix.py
-tests/unit/test_tmatrixc.py
-tests/unit/test_util.py
-treams/__init__.py
-treams/_core.py
-treams/_lattice.py
-treams/_material.py
-treams/_operators.py
-treams/_smatrix.py
-treams/_tmatrix.py
-treams/coeffs.c
-treams/config.c
-treams/cw.c
-treams/ebcm.py
-treams/io.py
-treams/misc.py
-treams/pw.c
-treams/sw.c
-treams/util.py
-treams.egg-info/PKG-INFO
-treams.egg-info/SOURCES.txt
-treams.egg-info/dependency_links.txt
-treams.egg-info/requires.txt
-treams.egg-info/top_level.txt
-treams/lattice/__init__.py
-treams/lattice/_dsum.c
-treams/lattice/_esum.c
-treams/lattice/_gufuncs.c
-treams/lattice/_misc.c
-treams/lattice/cython_lattice.c
-treams/lattice/cython_lattice.pxd
-treams/special/__init__.py
-treams/special/_bessel.c
-treams/special/_coord.c
-treams/special/_gufuncs.c
-treams/special/_integrals.c
-treams/special/_misc.c
-treams/special/_ufuncs.c
-treams/special/_waves.c
-treams/special/_wigner3j.c
-treams/special/_wignerd.c
-treams/special/cython_special.c
-treams/special/cython_special.pxd

treams.egg-info/top_level.txt

-treams

treams/__init__.py

-"""TREAMS: T-Matrix scattering code for nanophotonic computations.
-
-.. currentmodule:: treams
-
-Classes
-=======
-
-The top-level classes and functions allow a high-level access to the functionality.
-
-Basis sets
-----------
-
-.. autosummary::
-   :toctree: generated/
-
-   CylindricalWaveBasis
-   PlaneWaveBasisByUnitVector
-   PlaneWaveBasisByComp
-   SphericalWaveBasis
-
-Matrices and Arrays
--------------------
-
-.. autosummary::
-   :toctree: generated/
-
-   PhysicsArray
-   SMatrix
-   TMatrix
-   TMatrixC
-
-Other
------
-
-.. autosummary::
-   :toctree: generated/
-
-   Lattice
-   Material
-
-Functions
-=========
-
-.. autosummary::
-   :toctree: generated/
-
-   bfield
-   changepoltype
-   dfield
-   efield
-   expand
-   expandlattice
-   hfield
-   permute
-   rotate
-   translate
-
-"""
-
-from treams._core import (  # noqa: F401
-    CylindricalWaveBasis,
-    PhysicsArray,
-    PlaneWaveBasisByUnitVector,
-    PlaneWaveBasisByComp,
-    SphericalWaveBasis,
-)
-from treams._lattice import Lattice, WaveVector  # noqa: F401
-from treams._material import Material  # noqa: F401
-from treams._operators import (  # noqa: F401
-    BField,
-    bfield,
-    ChangePoltype,
-    changepoltype,
-    DField,
-    dfield,
-    EField,
-    efield,
-    Expand,
-    expand,
-    ExpandLattice,
-    expandlattice,
-    FField,
-    ffield,
-    GField,
-    gfield,
-    HField,
-    hfield,
-    Permute,
-    permute,
-    Rotate,
-    rotate,
-    Translate,
-    translate,
-)
-from treams._smatrix import (  # noqa: F401
-    SMatrices,
-    SMatrix,
-    chirality_density,
-    poynting_avg_z,
-)
-from treams._tmatrix import (  # noqa: F401
-    TMatrix,
-    TMatrixC,
-    plane_wave,
-    plane_wave_angle,
-    cylindrical_wave,
-    spherical_wave,
-)

treams/_core.py

-"""Basis sets and core array functionalities."""
-
-import abc
-from collections import namedtuple
-
-import numpy as np
-
-import treams._operators as op
-import treams.lattice as la
-from treams import util
-from treams._lattice import Lattice, WaveVector
-from treams._material import Material
-
-
-class BasisSet(util.OrderedSet, metaclass=abc.ABCMeta):
-    """Basis set base class.
-
-    It is the base class for all basis sets used. They are expected to be an ordered
-    sequence of the modes, that are included in a expansion. Basis sets are expected to
-    be immutable.
-    """
-
-    _names = ()
-    """Names of the relevant parameters"""
-
-    def __repr__(self):
-        """String representation.
-
-        Automatically generated when the attribute ``_names`` is defined.
-
-        Returns:
-            str
-        """
-        string = ",\n    ".join(f"{name}={i}" for name, i in zip(self._names, self[()]))
-        return f"{self.__class__.__name__}(\n    {string},\n)"
-
-    def __len__(self):
-        """Number of modes."""
-        return len(self.pol)
-
-    @classmethod
-    @abc.abstractmethod
-    def default(cls, *args, **kwargs):
-        """Construct a default basis from parameters.
-
-        Construct a basis set in a default order by giving few parameters.
-        """
-        raise NotImplementedError
-
-
-class SphericalWaveBasis(BasisSet):
-    r"""Basis of spherical waves.
-
-    Functions of the spherical wave basis are defined by their angular momentum ``l``,
-    its projection onto the z-axis ``m``, and the polarization ``pol``. If the basis
-    is defined with respect to a single origin it is referred to as "global", if it
-    contains multiple origins it is referred to as "local". In a local basis an
-    additional position index ``pidx`` is used to link the modes to one of the
-    specified ``positions``.
-
-    For spherical waves there exist multiple :ref:`params:Polarizations` and they
-    can be separated into incident and scattered fields. Depending on these combinations
-    the basis modes refer to one of the functions :func:`~treams.special.vsw_A`,
-    :func:`~treams.special.vsw_rA`, :func:`~treams.special.vsw_M`,
-    :func:`~treams.special.vsw_rM`, :func:`~treams.special.vsw_N`, or
-    :func:`~treams.special.vsw_rN`.
-
-    Args:
-        modes (array-like): A tuple containing a list for each of ``l``, ``m``, and
-            ``pol`` or ``pidx``, ``l``, ``m``, and``pol``.
-        positions (array-like, optional): The positions of the origins for the specified
-            modes. Defaults to ``[[0, 0, 0]]``.
-
-    Attributes:
-        pidx (array-like): Integer referring to a row in :attr:`positions`.
-        l (array-like): Angular momentum as an integer :math:`l > 0`
-        m (array-like): Angular momentum projection onto the z-axis, it is an integer
-            with :math:`m \leq |l|`
-        pol (array-like): Polarization, see also :ref:`params:Polarizations`.
-    """
-
-    _names = ("pidx", "l", "m", "pol")
-
-    def __init__(self, modes, positions=None):
-        """Initalization."""
-        tmp = []
-        if isinstance(modes, np.ndarray):
-            modes = modes.tolist()
-        for m in modes:
-            if m not in tmp:
-                tmp.append(m)
-        modes = tmp
-        if len(modes) == 0:
-            pidx = []
-            l = []  # noqa: E741
-            m = []
-            pol = []
-        elif len(modes[0]) == 3:
-            l, m, pol = (*zip(*modes),)
-            pidx = np.zeros_like(l)
-        elif len(modes[0]) == 4:
-            pidx, l, m, pol = (*zip(*modes),)
-        else:
-            raise ValueError("invalid shape of modes")
-
-        if positions is None:
-            positions = np.zeros((1, 3))
-        positions = np.array(positions, float)
-        if positions.ndim == 1:
-            positions = positions[None, :]
-        if positions.ndim != 2 or positions.shape[1] != 3:
-            raise ValueError(f"invalid shape of positions {positions.shape}")
-
-        self.pidx, self.l, self.m, self.pol = [
-            np.array(i, int) for i in (pidx, l, m, pol)
-        ]
-        for i, j in ((self.pidx, pidx), (self.l, l), (self.m, m), (self.pol, pol)):
-            i.flags.writeable = False
-            if np.any(i != j):
-                raise ValueError("parameters must be integer")
-        if np.any(self.l < 1):
-            raise ValueError("'l' must be a strictly positive integer")
-        if np.any(self.l < np.abs(self.m)):
-            raise ValueError("'|m|' cannot be larger than 'l'")
-        if np.any(self.pol > 1) or np.any(self.pol < 0):
-            raise ValueError("polarization must be '0' or '1'")
-        if np.any(self.pidx >= len(positions)):
-            raise ValueError("undefined position is indexed")
-
-        self._positions = positions
-        self._positions.flags.writeable = False
-        self.lattice = self.kpar = None
-
-    @property
-    def positions(self):
-        """Positions of the modes' origins.
-
-        The positions are an immutable (N, 3)-array. Each row corresponds to a point in
-        the three-dimensional Cartesian space.
-        """
-        return self._positions
-
-    def __repr__(self):
-        """String representation."""
-        positions = "positions=" + str(self.positions).replace("\n", ",")
-        return f"{super().__repr__()[:-1]}    {positions},\n)"
-
-    @property
-    def isglobal(self):
-        """Basis is defined with respect to a single (global) origin.
-
-        Returns:
-            bool
-        """
-        return len(self) == 0 or np.all(self.pidx == self.pidx[0])
-
-    def __getattr__(self, key):
-        dct = {"l": "l", "m": "m", "p": "pidx", "s": "pol"}
-        try:
-            return tuple(getattr(self, dct[k]) for k in key)
-        except KeyError:
-            raise AttributeError(
-                f"'{type(self).__name__}' object has no attribute '{key}'"
-            ) from None
-
-    def __getitem__(self, idx):
-        """Get a subset of the basis.
-
-        This function allows index into the basis set by an integer, a slice, a sequence
-        of integers or bools, an ellipsis, or an empty tuple. All of them except the
-        integer and the empty tuple results in another basis set being returned. In case
-        of the two exceptions a tuple is returned.
-
-        Alternatively, the string "plmp", "lmp", or "lm" can be used to access only a
-        subset of :attr:`pidx`, :attr:`l`, :attr:`m`, and :attr:`pol`.
-        """
-        res = self.pidx[idx], self.l[idx], self.m[idx], self.pol[idx]
-        if isinstance(idx, (int, np.integer)) or (
-            isinstance(idx, tuple) and len(idx) == 0
-        ):
-            return res
-        return type(self)(zip(*res), self.positions)
-
-    def __eq__(self, other):
-        """Compare basis sets.
-
-        Basis sets are considered equal, when they have the same modes in the same order
-        and the specified origin :attr:`positions` are equal.
-        """
-        try:
-            return self is other or (
-                np.array_equal(self.pidx, other.pidx)
-                and np.array_equal(self.l, other.l)
-                and np.array_equal(self.m, other.m)
-                and np.array_equal(self.pol, other.pol)
-                and np.array_equal(self.positions, other.positions)
-            )
-        except AttributeError:
-            return False
-
-    @classmethod
-    def default(cls, lmax, nmax=1, positions=None):
-        """Default basis for the given maximal multipolar order.
-
-        The default order contains separate blocks for each position index which are in
-        ascending order. Within each block the modes are sorted by angular momentum
-        :math:`l`, with the lowest angular momentum coming first. For each angular
-        momentum its z-projection is in ascending order from :math:`m = -l` to
-        :math:`m = l`. Finally, the polarization index is the fastest changing index
-        which iterates between 1 and 0.
-
-        Example:
-            >>> SphericalWaveBasis.default(2)
-            SphericalWaveBasis(
-                pidx=[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0],
-                l=[1 1 1 1 1 1 2 2 2 2 2 2 2 2 2 2],
-                m=[-1 -1  0  0  1  1 -2 -2 -1 -1  0  0  1  1  2  2],
-                pol=[1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0],
-                positions=[[0. 0. 0.]],
-            )
-            >>> SphericalWaveBasis.default(1, 2, [[0, 0, 1.], [0, 0, -1.]])
-            SphericalWaveBasis(
-                pidx=[0 0 0 0 0 0 1 1 1 1 1 1],
-                l=[1 1 1 1 1 1 1 1 1 1 1 1],
-                m=[-1 -1  0  0  1  1 -1 -1  0  0  1  1],
-                pol=[1 0 1 0 1 0 1 0 1 0 1 0],
-                positions=[[ 0.  0.  1.], [ 0.  0. -1.]],
-            )
-
-        Args:
-            lmax (int): Maximal multipolar order.
-            nmax (int, optional): Number of positions, defaults to 1.
-            positions (array-like, optional): Positions of the origins.
-        """
-        modes = [
-            [n, l, m, p]
-            for n in range(0, nmax)
-            for l in range(1, lmax + 1)  # noqa: E741
-            for m in range(-l, l + 1)
-            for p in range(1, -1, -1)
-        ]
-        return cls(modes, positions=positions)
-
-    @classmethod
-    def ebcm(cls, lmax, nmax=1, mmax=-1, positions=None):
-        """Order of modes suited for ECBM.
-
-        In comparison to :meth:`default` this order prioritises blocks of the
-        z-projection of the angular momentum :math:`m` over the angular momentum
-        :math:`l`. This is useful to get block-diagonal matrices for the extended
-        boundary condition method (EBCM).
-
-        Example:
-            >>> SphericalWaveBasis.ebcm(2)
-            SphericalWaveBasis(
-                pidx=[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0],
-                l=[2 2 1 1 2 2 1 1 2 2 1 1 2 2 2 2],
-                m=[-2 -2 -1 -1 -1 -1  0  0  0  0  1  1  1  1  2  2],
-                pol=[1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0],
-                positions=[[0. 0. 0.]],
-            )
-
-        Args:
-            lmax (int): Maximal multipolar order.
-            nmax (int, optional): Number of particles, defaults to 1.
-            mmax (int, optional): Maximal value of `|m|`. If ommitted or set to -1 it
-                is taken equal to `lmax`.
-            positions (array-like, optional): Positions of the origins.
-
-        Returns:
-            tuple
-        """
-        mmax = lmax if mmax == -1 else mmax
-        modes = [
-            [n, l, m, p]
-            for n in range(0, nmax)
-            for m in range(-mmax, mmax + 1)
-            for l in range(max(abs(m), 1), lmax + 1)  # noqa: E741
-            for p in range(1, -1, -1)
-        ]
-        return cls(modes, positions=positions)
-
-    @staticmethod
-    def defaultlmax(dim, nmax=1):
-        """Calculate the default mode order for a given length.
-
-        Given the dimension of the T-matrix return the estimated maximal value of `l`.
-        This is the inverse of :meth:`defaultdim`. A value of zero is allowed for empty
-        T-matrices.
-
-        Example:
-            >>> SphericalWaveBasis.defaultlmax(len(SphericalWaveBasis.default(3)))
-            3
-
-        Args:
-            dim (int): Dimension of the T-matrix, respectively number of modes.
-            nmax (int, optional): Number of particles, defaults to 1.
-
-        Returns:
-            int
-        """
-        res = np.sqrt(1 + dim * 0.5 / nmax) - 1
-        res_int = int(np.rint(res))
-        if np.abs(res - res_int) > 1e-8 * np.maximum(np.abs(res), np.abs(res_int)):
-            raise ValueError("cannot estimate the default lmax")
-        return res_int
-
-    @staticmethod
-    def defaultdim(lmax, nmax=1):
-        """Default number of modes for a given mulipolar order.
-
-        Given the maximal value of `l` return the size of the corresponding T-matrix.
-        This is the inverse of :meth:`defaultlmax`. A value of zero is allowed.
-
-        Args:
-            lmax (int): Maximal multipolar order
-            nmax (int, optional): Number of particles, defaults to `1`
-
-        Returns:
-            int
-        """
-        # zero is allowed and won't give an error
-        if lmax < 0 or nmax < 0:
-            raise ValueError("maximal order must be positive")
-        return 2 * lmax * (lmax + 2) * nmax
-
-    @classmethod
-    def _from_iterable(cls, it, positions=None):
-        if isinstance(cls, SphericalWaveBasis):
-            positions = cls.positions if positions is None else positions
-            cls = type(cls)
-        obj = cls(it, positions=positions)
-        return obj
-
-
-class CylindricalWaveBasis(BasisSet):
-    r"""Basis of cylindrical waves.
-
-    Functions of the cylindrical wave basis are defined by the z-components of the wave
-    vector ``kz`` and the angular momentum  ``m`` as well as the polarization ``pol``.
-    If the basis is defined with respect to a single origin it is referred to as
-    "global", if it contains multiple origins it is referred to as "local". In a local
-    basis an additional position index ``pidx`` is used to link the modes to one of the
-    specified ``positions``.
-
-    For cylindrical waves there exist multiple :ref:`params:Polarizations` and
-    they can be separated into incident and scattered fields. Depending on these
-    combinations the basis modes refer to one of the functions
-    :func:`~treams.special.vcw_A`, :func:`~treams.special.vcw_rA`,
-    :func:`~treams.special.vcw_M`, :func:`~treams.special.vcw_rM`,
-    :func:`~treams.special.vcw_N`, or :func:`~treams.special.vcw_rN`.
-
-    Args:
-        modes (array-like): A tuple containing a list for each of ``kz``, ``m``, and
-            ``pol`` or ``pidx``, ``kz``, ``m``, and``pol``.
-        positions (array-like, optional): The positions of the origins for the specified
-            modes. Defaults to ``[[0, 0, 0]]``.
-
-    Attributes:
-        pidx (array-like): Integer referring to a row in :attr:`positions`.
-        kz (array-like): Real valued z-component of the wave vector.
-        m (array-like): Integer angular momentum projection onto the z-axis.
-        pol (array-like): Polarization, see also :ref:`params:Polarizations`.
-    """
-
-    _names = ("pidx", "kz", "m", "pol")
-
-    def __init__(self, modes, positions=None):
-        """Initalization."""
-        tmp = []
-        if isinstance(modes, np.ndarray):
-            modes = modes.tolist()
-        for m in modes:
-            if m not in tmp:
-                tmp.append(m)
-        modes = tmp
-        if len(modes) == 0:
-            pidx = []
-            kz = []
-            m = []
-            pol = []
-        elif len(modes[0]) == 3:
-            kz, m, pol = (*zip(*modes),)
-            pidx = np.zeros_like(m)
-        elif len(modes[0]) == 4:
-            pidx, kz, m, pol = (*zip(*modes),)
-        else:
-            raise ValueError("invalid shape of modes")
-
-        if positions is None:
-            positions = np.zeros((1, 3))
-        positions = np.array(positions, float)
-        if positions.ndim == 1:
-            positions = positions[None, :]
-        if positions.ndim != 2 or positions.shape[1] != 3:
-            raise ValueError(f"invalid shape of positions {positions.shape}")
-
-        self.pidx, self.m, self.pol = [np.array(i, int) for i in (pidx, m, pol)]
-        self.kz = np.array(kz, float)
-        self.kz.flags.writeable = False
-        for i, j in ((self.pidx, pidx), (self.m, m), (self.pol, pol)):
-            i.flags.writeable = False
-            if np.any(i != j):
-                raise ValueError("parameters must be integer")
-        if np.any(self.pol > 1) or np.any(self.pol < 0):
-            raise ValueError("polarization must be '0' or '1'")
-        if np.any(self.pidx >= len(positions)):
-            raise ValueError("undefined position is indexed")
-
-        self._positions = positions
-        self._positions.flags.writeable = False
-
-        self.lattice = self.kpar = None
-        if len(self.kz) > 0 and np.all(self.kz == self.kz[0]):
-            self.kpar = WaveVector(self.kz[0])
-
-    @property
-    def positions(self):
-        """Positions of the modes' origins.
-
-        The positions are an immutable (N, 3)-array. Each row corresponds to a point in
-        the three-dimensional Cartesian space.
-        """
-        return self._positions
-
-    def __repr__(self):
-        """String representation."""
-        positions = "positions=" + str(self.positions).replace("\n", ",")
-        return f"{super().__repr__()[:-1]}    {positions},\n)"
-
-    @property
-    def isglobal(self):
-        """Basis is defined with respect to a single (global) origin.
-
-        Returns:
-            bool
-        """
-        return len(self) == 0 or np.all(self.pidx == self.pidx[0])
-
-    def __getattr__(self, key):
-        dct = {"z": "kz", "m": "m", "p": "pidx", "s": "pol"}
-        try:
-            return tuple(getattr(self, dct[k]) for k in key)
-        except KeyError:
-            raise AttributeError(
-                f"'{type(self).__name__}' object has no attribute '{key}'"
-            ) from None
-
-    def __getitem__(self, idx):
-        """Get a subset of the basis.
-
-        This function allows index into the basis set by an integer, a slice, a sequence
-        of integers or bools, an ellipsis, or an empty tuple. All of them except the
-        integer and the empty tuple results in another basis set being returned. In case
-        of the two exceptions a tuple is returned.
-
-        Alternatively, the string "pkzmp", "kzmp", or "kzm" can be used to access only a
-        subset of :attr:`pidx`, :attr:`kz`, :attr:`m`, and :attr:`pol`.
-        """
-        res = self.pidx[idx], self.kz[idx], self.m[idx], self.pol[idx]
-        if isinstance(idx, (int, np.integer)) or (isinstance(idx, tuple) and idx == ()):
-            return res
-        return type(self)(zip(*res), self.positions)
-
-    def __eq__(self, other):
-        """Compare basis sets.
-
-        Basis sets are considered equal, when they have the same modes in the same order
-        and the specified origin :attr:`positions` are equal.
-        """
-        try:
-            return self is other or (
-                np.array_equal(self.pidx, other.pidx)
-                and np.array_equal(self.kz, other.kz)
-                and np.array_equal(self.m, other.m)
-                and np.array_equal(self.pol, other.pol)
-                and np.array_equal(self.positions, other.positions)
-            )
-        except AttributeError:
-            return False
-
-    @classmethod
-    def default(cls, kzs, mmax, nmax=1, positions=None):
-        """Default basis for the given z-components of wave vector and angular momentum.
-
-        The default order contains separate blocks for each position index which are in
-        ascending order. Within each block the modes are sorted by the z-component of
-        the wave vector :math:`k_z`. For each of those values the z-projection of the
-        angular momentum is placed in ascending order. Finally, the polarization index
-        is the fastest changing index which iterates between 1 and 0.
-
-        Example:
-            >>> CylindricalWaveBasis.default([-0.5, 0.5], 1)
-            CylindricalWaveBasis(
-                pidx=[0 0 0 0 0 0 0 0 0 0 0 0],
-                kz=[-0.5 -0.5 -0.5 -0.5 -0.5 -0.5  0.5  0.5  0.5  0.5  0.5  0.5],
-                m=[-1 -1  0  0  1  1 -1 -1  0  0  1  1],
-                pol=[1 0 1 0 1 0 1 0 1 0 1 0],
-                positions=[[0. 0. 0.]],
-            )
-            >>> CylindricalWaveBasis.default([0], 1, 2, [[1., 0, 0], [-1., 0, 0]])
-            CylindricalWaveBasis(
-                pidx=[0 0 0 0 0 0 1 1 1 1 1 1],
-                kz=[0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.],
-                m=[-1 -1  0  0  1  1 -1 -1  0  0  1  1],
-                pol=[1 0 1 0 1 0 1 0 1 0 1 0],
-                positions=[[ 1.  0.  0.], [-1.  0.  0.]],
-            )
-
-        Args:
-            kzs (array-like, float): Maximal multipolar order.
-            mmax (int): Maximal value of the angular momentum z-component.
-            nmax (int, optional): Number of positions, defaults to 1.
-            positions (array-like, optional): Positions of the origins.
-        """
-        kzs = np.atleast_1d(kzs)
-        if kzs.ndim > 1:
-            raise ValueError(f"kzs has dimension larger than one: '{kzs.ndim}'")
-        modes = [
-            [n, kz, m, p]
-            for n in range(nmax)
-            for kz in kzs
-            for m in range(-mmax, mmax + 1)
-            for p in range(1, -1, -1)
-        ]
-        return cls(modes, positions=positions)
-
-    @classmethod
-    def diffr_orders(cls, kz, mmax, lattice, bmax, nmax=1, positions=None):
-        """Create a basis set for a system periodic in the z-direction.
-
-        Example:
-           >>> CylindricalWaveBasis.diffr_orders(0.1, 1, 2 * np.pi, 1)
-           CylindricalWaveBasis(
-               pidx=[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0],
-               kz=[-0.9 -0.9 -0.9 -0.9 -0.9 -0.9  0.1  0.1  0.1  0.1  0.1  0.1  1.1  1.1
-             1.1  1.1  1.1  1.1],
-               m=[-1 -1  0  0  1  1 -1 -1  0  0  1  1 -1 -1  0  0  1  1],
-               pol=[1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0],
-               positions=[[0. 0. 0.]],
-           )
-
-        Args:
-            kz (float): Wave vector z-component. Ideally it is in the first Brillouin
-                zone (use :func:`misc.firstbrillouin1d`).
-            mmax (int): Maximal value for the z-component of the angular momentum.
-            lattice (:class:`treams.Lattice` or float): Lattice definition or pitch.
-            bmax (float): Maximal change of the z-component of the wave vector. So,
-                this defines a maximal momentum transfer from the given value `kz`.
-            nmax (int, optional): Number of positions.
-            positions (array-like, optional): Positions of the origins.
-        """
-        lattice = Lattice(lattice)
-        lattice_z = Lattice(lattice, "z")
-        nkz = np.floor(np.abs(bmax / lattice_z.reciprocal))
-        kzs = kz + np.arange(-nkz, nkz + 1) * lattice_z.reciprocal
-        res = cls.default(kzs, mmax, nmax, positions=positions)
-        res.lattice = lattice
-        res.kpar = WaveVector(kz)
-        return res
-
-    @classmethod
-    def _from_iterable(cls, it, positions=None):
-        if isinstance(cls, CylindricalWaveBasis):
-            positions = cls.positions if positions is None else positions
-            lattice = cls.lattice
-            kpar = cls.kpar
-            cls = type(cls)
-        else:
-            lattice = kpar = None
-        obj = cls(it, positions=positions)
-        obj.lattice = lattice
-        obj.kpar = kpar
-        return obj
-
-    @staticmethod
-    def defaultmmax(dim, nkz=1, nmax=1):
-        """Calculate the default mode order for a given length.
-
-        Given the dimension of the T-matrix return the estimated maximal value of `m`.
-        This is the inverse of :meth:`defaultdim`. A value of zero is allowed for empty
-        T-matrices.
-
-        Example:
-            >>> CylindricalWaveBasis.defaultmmax(len(CylindricalWaveBasis.default([0], 2)), 1)
-            2
-
-        Args:
-            dim (int): Dimension of the T-matrix, respectively number of modes
-            nkz (int, optional): Number of z-components of the wave vector.
-            nmax (int, optional): Number of particles, defaults to 1.
-
-        Returns:
-            int
-        """
-        if dim % (2 * nkz * nmax) != 0:
-            raise ValueError("cannot estimate the default mmax")
-        dim = dim // (2 * nkz * nmax)
-        if (dim - 1) % 2 != 0:
-            raise ValueError("cannot estimate the default mmax")
-        return (dim - 1) // 2
-
-    @staticmethod
-    def defaultdim(nkz, mmax, nmax=1):
-        """Default number of modes for a given mulipolar order.
-
-        Given the maximal value of `l` return the size of the corresponding T-matrix.
-        This is the inverse of :meth:`defaultlmax`. A value of zero is allowed.
-
-        Args:
-            nkz (int): Number of z-components of the wave vector.
-            mmax (int): Maximal value of the angular momentum's z-component.
-            nmax (int, optional): Number of particles, defaults to 1.
-
-        Returns:
-            int
-        """
-        if nkz < 0 or mmax < 0:
-            raise ValueError("maximal order must be positive")
-        return (2 * mmax + 1) * nkz * 2 * nmax
-
-
-class PlaneWaveBasis(BasisSet):
-    """Plane wave basis parent class."""
-
-    isglobal = True
-
-
-class PlaneWaveBasisByUnitVector(PlaneWaveBasis):
-    """Plane wave basis.
-
-    A plane wave basis is defined by a collection of wave vectors specified by the
-    Cartesian wave vector components ``qx``, ``qy``, and ``qz`` normalized to
-    :math:`q_x^2 + q_y^2 + q_z^2 = 1` and the polarizations ``pol``.
-
-    For plane waves there exist multiple :ref:`params:Polarizations`, such that
-    these modes can refer to either :func:`~treams.special.vpw_A` or
-    :func:`~treams.special.vpw_M` and :func:`~treams.special.vpw_N`.
-
-    Args:
-        modes (array-like): A tuple containing a list for each of ``qx``, ``qy``,
-            ``qz``, and ``pol``.
-
-    Attributes:
-        qx (array-like): X-component of the normalized wave vector.
-        qy (array-like): Y-component of the normalized wave vector.
-        qz (array-like): Z-component of the normalized wave vector.
-        pol (array-like): Polarization, see also :ref:`params:Polarizations`.
-    """
-
-    _names = ("qx", "qy", "qz", "pol")
-    """A plane wave basis is always global."""
-
-    def __init__(self, modes):
-        """Initialization."""
-        tmp = []
-        if isinstance(modes, np.ndarray):
-            modes = modes.tolist()
-        for m in modes:
-            if m not in tmp:
-                tmp.append(m)
-        modes = tmp
-        if len(modes) == 0:
-            qx = []
-            qy = []
-            qz = []
-            pol = []
-        elif len(modes[0]) == 4:
-            qx, qy, qz, pol = (*zip(*modes),)
-        else:
-            raise ValueError("invalid shape of modes")
-
-        qx, qy, qz, pol = (np.array(i) for i in (qx, qy, qz, pol))
-        norm = np.emath.sqrt(qx * qx + qy * qy + qz * qz)
-        norm[np.abs(norm - 1) < 1e-14] = 1
-        qx, qy, qz = (np.true_divide(i, norm) for i in (qx, qy, qz))
-        for i in (qx, qy, qz, pol):
-            i.flags.writeable = False
-            if i.ndim > 1:
-                raise ValueError("invalid shape of parameters")
-        self.qx, self.qy, self.qz = qx, qy, qz
-        self.pol = np.array(pol.real, int)
-        if np.any(self.pol != pol):
-            raise ValueError("polarizations must be integer")
-        if np.any(self.pol > 1) or np.any(self.pol < 0):
-            raise ValueError("polarization must be '0' or '1'")
-        self.lattice = self.kpar = None
-
-    def __getattr__(self, key):
-        dct = {"x": "qx", "y": "qy", "z": "qz", "s": "pol"}
-        try:
-            return tuple(getattr(self, dct[k]) for k in key)
-        except KeyError:
-            raise AttributeError(
-                f"'{type(self).__name__}' object has no attribute '{key}'"
-            ) from None
-
-    def __getitem__(self, idx):
-        """Get a subset of the basis.
-
-        This function allows index into the basis set by an integer, a slice, a sequence
-        of integers or bools, an ellipsis, or an empty tuple. All of them except the
-        integer and the empty tuple results in another basis set being returned. In case
-        of the two exceptions a tuple is returned.
-
-        Alternatively, the string "xyzp", "xyp", or "zp" can be used to access only a
-        subset of :attr:`qx`, :attr:`qy`, :attr:`qz`, and :attr:`pol`.
-        """
-        res = self.qx[idx], self.qy[idx], self.qz[idx], self.pol[idx]
-        if isinstance(idx, (int, np.integer)) or (isinstance(idx, tuple) and idx == ()):
-            return res
-        return type(self)(zip(*res))
-
-    @classmethod
-    def default(cls, kvecs):
-        """Default basis from the given wave vectors.
-
-        For each wave vector the two polarizations 1 and 0 are taken.
-
-        Example:
-            >>> PlaneWaveBasisByUnitVector.default([[0, 0, 5], [0, 3, 4]])
-            PlaneWaveBasisByUnitVector(
-                qx=[0. 0. 0. 0.],
-                qy=[0.  0.  0.6 0.6],
-                qz=[1.  1.  0.8 0.8],
-                pol=[1 0 1 0],
-            )
-
-        Args:
-            kvecs (array-like): Wave vectors in Cartesian coordinates.
-        """
-        kvecs = np.atleast_2d(kvecs)
-        modes = np.empty((2 * kvecs.shape[0], 4), kvecs.dtype)
-        modes[::2, :3] = kvecs
-        modes[1::2, :3] = kvecs
-        modes[::2, 3] = 1
-        modes[1::2, 3] = 0
-        return cls(modes)
-
-    @classmethod
-    def _from_iterable(cls, it):
-        if isinstance(cls, PlaneWaveBasisByUnitVector):
-            lattice = cls.lattice
-            kpar = cls.kpar
-            cls = type(cls)
-        else:
-            lattice = kpar = None
-        obj = cls(it)
-        obj.lattice = lattice
-        obj.kpar = kpar
-        return obj
-
-    def __eq__(self, other):
-        """Compare basis sets.
-
-        Basis sets are considered equal, when they have the same modes in the same
-        order.
-        """
-        try:
-            return self is other or (
-                np.array_equal(self.qx, other.qx)
-                and np.array_equal(self.qy, other.qy)
-                and np.array_equal(self.qz, other.qz)
-                and np.array_equal(self.pol, other.pol)
-            )
-        except AttributeError:
-            return False
-
-    def bycomp(self, k0, alignment="xy", material=Material()):
-        """Create a :class:`PlaneWaveBasisByComp`.
-
-        The plane wave basis is changed to a partial basis, where only two (real-valued)
-        wave vector components are defined the third component is then inferred from the
-        dispersion relation, which depends on the wave number and material, and a
-        ``modetype`` that specifies the sign of the third component.
-
-        Args:
-            alignment (str, optional): Wave vector components that are part of the
-                partial basis. Defaults to "xy", other permitted values are "yz" and
-                "zx".
-            k0 (float, optional): Wave number. If given, it is checked that the current
-                basis fulfils the dispersion relation.
-            material (:class:`~treams.Material` or tuple): Material definition. Defaults
-                to vacuum/air.
-        """
-        ks = material.ks(k0)[self.pol]
-        if alignment in ("xy", "yz", "zx"):
-            kpars = [ks * getattr(self, "q" + s) for s in alignment]
-        else:
-            raise ValueError(f"invalid alignment '{alignment}'")
-        obj = PlaneWaveBasisByComp(zip(*kpars, self.pol), alignment)
-        obj.lattice = self.lattice
-        obj.kpar = self.kpar
-        return obj
-
-    def kvecs(self, k0, material=Material(), modetype=None):
-        """Wave vectors
-
-
-        Args:
-            k0 (float): Wave number.
-            material (:class:`~treams.Material` or tuple, optional): Material
-                definition. Defaults to vacuum/air.
-            modetype (optional): Currently unused for this class.
-        """
-        # TODO: check kz depending on modetype (alignment?)
-        ks = Material(material).ks(k0)[self.pol]
-        return ks * self.qx, ks * self.qy, ks * self.qz
-
-
-class PlaneWaveBasisByComp(PlaneWaveBasis):
-    """Partial plane wave basis.
-
-    A partial plane wave basis is defined by two wave vector components out of all three
-    Cartesian wave vector components ``kx``, ``ky``, and ``kz`` and the polarizations
-    ``pol``. Which two components are given is specified in the :attr:`alignment`. This
-    basis is mostly used for stratified media that are periodic or uniform in the two
-    alignment directions, such that the given wave vector components correspond to the
-    diffraction orders.
-
-    For plane waves there exist multiple :ref:`params:Polarizations`, such that
-    these modes can refer to either :func:`~treams.special.vpw_A` or
-    :func:`~treams.special.vpw_M` and :func:`~treams.special.vpw_N`.
-
-    Args:
-        modes (array-like): A tuple containing a list for each of ``k1``, ``k2``, and
-            ``pol``.
-        alignment (str, optional): Definition which wave vector components are given.
-            Defaults to "xy", other possible values are "yz" and "zx".
-
-    Attributes:
-        alignment (str): Alignment of the partial basis.
-        pol (array-like): Polarization, see also :ref:`params:Polarizations`.
-    """
-
-    def __init__(self, modes, alignment="xy"):
-        """Initialization."""
-        if isinstance(modes, np.ndarray):
-            modes = modes.tolist()
-        tmp = []
-        for m in modes:
-            if m not in tmp:
-                tmp.append(m)
-        modes = tmp
-        if len(modes) == 0:
-            kx = []
-            ky = []
-            pol = []
-        elif len(modes[0]) == 3:
-            kx, ky, pol = (*zip(*modes),)
-        else:
-            raise ValueError("invalid shape of modes")
-
-        self._kx, self._ky = [np.real(i) for i in (kx, ky)]
-        self.pol = np.array(np.real(pol), int)
-        for i, j in [(self._kx, kx), (self._ky, ky), (self.pol, pol)]:
-            i.flags.writeable = False
-            if i.ndim > 1:
-                raise ValueError("invalid shape of parameters")
-            if np.any(i != j):
-                raise ValueError("invalid value for parameter, must be real")
-        if np.any(self.pol > 1) or np.any(self.pol < 0):
-            raise ValueError("polarization must be '0' or '1'")
-
-        if alignment in ("xy", "yz", "zx"):
-            self._names = (*("k" + i for i in alignment),) + ("pol",)
-        else:
-            raise ValueError(f"invalid alignment '{alignment}'")
-
-        self.alignment = alignment
-        self.lattice = self.kpar = None
-
-    def permute(self, n=1):
-        n = n % 3
-        lattice = None if self.lattice is None else self.lattice.permute(n)
-        kpar = None if self.kpar is None else self.kpar.permute(n)
-        alignments = {"xy": "yz", "yz": "zx", "zx": "xy"}
-        alignment = self.alignment
-        for _ in range(n):
-            alignment = alignments[alignment]
-        obj = self._from_iterable(self, alignment=alignment)
-        obj.lattice = lattice
-        obj.kpar = kpar
-        return obj
-
-    @property
-    def kx(self):
-        """X-components of the wave vector.
-
-        If the components are not specified `None` is returned.
-        """
-        if self.alignment == "xy":
-            return self._kx
-        elif self.alignment == "zx":
-            return self._ky
-        return None
-
-    @property
-    def ky(self):
-        """Y-components of the wave vector.
-
-        If the components are not specified `None` is returned.
-        """
-        if self.alignment == "xy":
-            return self._ky
-        elif self.alignment == "yz":
-            return self._kx
-        return None
-
-    @property
-    def kz(self):
-        """Z-components of the wave vector.
-
-        If the components are not specified `None` is returned.
-        """
-        if self.alignment == "yz":
-            return self._ky
-        elif self.alignment == "zx":
-            return self._kx
-        return None
-
-    def __getattr__(self, key):
-        dct = {"x": "kx", "y": "ky", "z": "kz", "s": "pol"}
-        try:
-            return tuple(getattr(self, dct[k]) for k in key)
-        except KeyError:
-            raise AttributeError(
-                f"'{type(self).__name__}' object has no attribute '{key}'"
-            ) from None
-
-    def __getitem__(self, idx):
-        """Get a subset of the basis.
-
-        This function allows index into the basis set by an integer, a slice, a sequence
-        of integers or bools, an ellipsis, or an empty tuple. All of them except the
-        integer and the empty tuple results in another basis set being returned. In case
-        of the two exceptions a tuple is returned.
-        """
-        res = self._kx[idx], self._ky[idx], self.pol[idx]
-        if isinstance(idx, (int, np.integer)) or (isinstance(idx, tuple) and idx == ()):
-            return res
-        return type(self)(zip(*res))
-
-    @classmethod
-    def default(cls, kpars, alignment="xy"):
-        """Default basis from the given wave vectors.
-
-        For each wave vector the two polarizations 1 and 0 are taken.
-
-        Example:
-            >>> PlaneWaveBasisByComp.default([[0, 0], [0, 3]])
-            PlaneWaveBasisByComp(
-                kx=[0. 0. 0. 0.],
-                ky=[0. 0. 3. 3.],
-                pol=[1 0 1 0],
-            )
-
-        Args:
-            kpars (array-like): Wave vector components in Cartesian coordinates.
-            alignment (str, optional): Definition which wave vector components are
-                given. Defaults to "xy", other possible values are "yz" and "zx".
-
-        """
-        kpars = np.atleast_2d(kpars)
-        modes = np.empty((2 * kpars.shape[0], 3), float)
-        modes[::2, :2] = kpars
-        modes[1::2, :2] = kpars
-        modes[::2, 2] = 1
-        modes[1::2, 2] = 0
-        return cls(modes, alignment=alignment)
-
-    @classmethod
-    def diffr_orders(cls, kpar, lattice, bmax):
-        """Create a basis set for a two-dimensional periodic system.
-
-        The reciprocal lattice to the given lattice is taken to consider all diffraction
-        orders that lie within the defined maximal radius (in reciprocal space).
-
-        Example:
-            >>> PlaneWaveBasisByComp.diffr_orders([0, 0], Lattice.square(2 * np.pi), 1)
-            PlaneWaveBasisByComp(
-                kx=[ 0.  0.  0.  0.  0.  0.  1.  1. -1. -1.],
-                ky=[ 0.  0.  1.  1. -1. -1.  0.  0.  0.  0.],
-                pol=[1 0 1 0 1 0 1 0 1 0],
-            )
-
-        Args:
-            kpar (float): Tangential wave vector components. Ideally they are in the
-                first Brillouin zone (use :func:`misc.firstbrillouin2d`).
-            lattice (:class:`treams.Lattice` or float): Lattice definition or pitch.
-            bmax (float): Maximal change of tangential wave vector components. So,
-                this defines a maximal momentum transfer.
-        """
-        lattice = Lattice(lattice)
-        if lattice.dim != 2:
-            raise ValueError("invalid lattice dimensions")
-        latrec = lattice.reciprocal
-        kpars = kpar + la.diffr_orders_circle(latrec, bmax) @ latrec
-        obj = cls.default(kpars, alignment=lattice.alignment)
-        obj.lattice = lattice
-        obj.kpar = WaveVector(kpar, alignment=lattice.alignment)
-        return obj
-
-    @classmethod
-    def _from_iterable(cls, it, alignment="xy"):
-        if isinstance(cls, PlaneWaveBasisByComp):
-            alignment = cls.alignment if alignment is None else alignment
-            lattice = cls.lattice
-            kpar = cls.kpar
-            cls = type(cls)
-        else:
-            lattice = kpar = None
-        obj = cls(it, alignment)
-        obj.lattice = lattice
-        obj.kpar = kpar
-        return obj
-
-    def __eq__(self, other):
-        """Compare basis sets.
-
-        Basis sets are considered equal, when they have the same modes in the same
-        order.
-        """
-        try:
-            skx, sky, skz = self.kx, self.ky, self.kz
-            okx, oky, okz = other.kx, other.ky, other.kz
-            return self is other or (
-                (np.array_equal(skx, okx) or (skx is None and okx is None))
-                and (np.array_equal(sky, oky) or (sky is None and oky is None))
-                and (np.array_equal(skz, okz) or (skz is None and okz is None))
-                and np.array_equal(self.pol, other.pol)
-            )
-        except AttributeError:
-            return False
-
-    def byunitvector(self, k0, material=Material(), modetype="up"):
-        """Create a complete basis :class:`PlaneWaveBasis`.
-
-        A plane wave basis is considered complete, when all three Cartesian components
-        and the polarization is defined for each mode. So, the specified wave number,
-        material, and modetype is taken to calculate the third Cartesian wave vector.
-        The modetype "up" ("down") is for waves propagating in the positive (negative)
-        direction with respect to the Cartesian axis that is orthogonal to the
-        alignment.
-
-        Args:
-            k0 (float): Wave number.
-            material (:class:`~treams.Material` or tuple, optional): Material
-                definition. Defaults to vacuum/air.
-            modetype (str, optional): Propagation direction. Defaults to "up".
-        """
-        if modetype not in ("up", "down"):
-            raise ValueError("modetype not recognized")
-        material = Material(material)
-        kx = self._kx
-        ky = self._ky
-        kz = material.kzs(k0, kx, ky, self.pol) * (2 * (modetype == "up") - 1)
-        if self.alignment == "yz":
-            kx, ky, kz = kz, kx, ky
-        elif self.alignment == "zy":
-            kx, ky, kz = ky, kz, kx
-        obj = PlaneWaveBasisByUnitVector(zip(kx, ky, kz, self.pol))
-        obj.lattice = self.lattice
-        obj.kpar = self.kpar
-        return obj
-
-    def kvecs(self, k0, material=Material(), modetype="up"):
-        """Wave vectors.
-
-        Args:
-            k0 (float): Wave number.
-            material (:class:`~treams.Material` or tuple, optional): Material
-                definition. Defaults to vacuum/air.
-            modetype (str, optional): Propagation direction. Defaults to "up".
-        """
-        if modetype not in ("up", "down"):
-            raise ValueError("modetype not recognized")
-        material = Material(material)
-        kx = self._kx
-        ky = self._ky
-        kz = material.kzs(k0, kx, ky, self.pol) * (2 * (modetype == "up") - 1)
-        if self.alignment == "yz":
-            return kz, kx, ky
-        elif self.alignment == "zx":
-            return ky, kz, kx
-        return kx, ky, kz
-
-
-def _raise_basis_error(*args):
-    raise TypeError("'basis' must be BasisSet")
-
-
-class PhysicsDict(util.AnnotationDict):
-    """Physics dictionary.
-
-    Derives from :class:`treams.util.AnnotationDict`. This dictionary has additionally
-    several properties defined.
-
-    Attributes:
-        basis (:class:`BasisSet`): Basis set.
-        k0 (float): Wave number.
-        kpar (list): Parallel wave vector components. Usually, this is a list of length
-            3 with its items corresponding to the Cartesian axes. Unspecified items are
-            set to `nan`.
-        lattice (:class:`~treams.Lattice`): Lattice definition.
-        material (:class:`~treams.Material`): Material definition.
-        modetype (str): Mode type, for spherical and cylindrical waves this can be
-            "incident" and "scattered", for partial plane waves it can be "up" or
-            "down".
-        poltype (str): Polarization, see also :ref:`params:Polarizations`.
-    """
-
-    properties = {
-        "basis": (
-            ":class:`~treams.BasisSet`.",
-            lambda x: isinstance(x, BasisSet),
-            _raise_basis_error,
-        ),
-        "k0": ("Wave number.", lambda x: isinstance(x, float), float),
-        "kpar": (
-            "Wave vector components tangential to the lattice.",
-            lambda x: isinstance(x, WaveVector),
-            WaveVector,
-        ),
-        "lattice": (
-            ":class:`~treams.Lattice`.",
-            lambda x: isinstance(x, Lattice),
-            Lattice,
-        ),
-        "material": (
-            ":class:`~treams.Material`.",
-            lambda x: isinstance(x, Material),
-            Material,
-        ),
-        "modetype": ("Mode type.", lambda x: isinstance(x, str), str),
-        "poltype": (
-            ":ref:`params:Polarizations`.",
-            lambda x: isinstance(x, str),
-            str,
-        ),
-    }
-    """Special properties tracked by the PhysicsDict."""
-
-    def __setitem__(self, key, val):
-        """Set item specified by key to the defined value.
-
-        When overwriting an existing key an :class:`AnnotationWarning` is emitted.
-        Avoid the warning by explicitly deleting the key first. The special attributes
-        are cast to their corresponding types.
-
-        Args:
-            key (hashable): Key
-            val : Value
-
-        Warns:
-            AnnotationWarning
-        """
-        if key not in self.properties:
-            raise AttributeError(f"invalid key '{key}'")
-        _, testfunc, castfunc = self.properties[key]
-        if not testfunc(val):
-            val = castfunc(val)
-        super().__setitem__(key, val)
-
-
-class PhysicsArray(util.AnnotatedArray):
-    """Physics-aware array.
-
-    A physics aware array is a special type of :class`~treams.util.AnnotatedArray`.
-    Additionally to keeping track of the annotations, it is enhanced by the ability to
-    create suiting linear operators to perform tasks like rotations, translations, or
-    expansions into different basis sets and by applying special rules for the
-    physical properties of :class:`PhysicsDict` upon matrix multiplications, see also
-    :meth:`__matmul__`.
-    """
-
-    _scales = {"basis"}
-
-    changepoltype = op.OperatorAttribute(op.ChangePoltype)
-    """Polarization change matrix, see also :class:`treams.operators.ChangePoltype`."""
-    efield = op.OperatorAttribute(op.EField)
-    hfield = op.OperatorAttribute(op.HField)
-    dfield = op.OperatorAttribute(op.DField)
-    bfield = op.OperatorAttribute(op.BField)
-    gfield = op.OperatorAttribute(op.GField)
-    ffield = op.OperatorAttribute(op.FField)
-    """Field evaluation matrix, see also :class:`treams.operators.EField`."""
-    expand = op.OperatorAttribute(op.Expand)
-    """Expansion matrix, see also :class:`treams.operators.Expand`."""
-    expandlattice = op.OperatorAttribute(op.ExpandLattice)
-    """Lattice expansion matrix, see also :class:`treams.operators.ExpandLattice`."""
-    permute = op.OperatorAttribute(op.Permute)
-    """Permutation matrix, see also :class:`treams.operators.Permute`."""
-    rotate = op.OperatorAttribute(op.Rotate)
-    """Rotation matrix, see also :class:`treams.operators.Rotate`."""
-    translate = op.OperatorAttribute(op.Translate)
-    """Translation matrix, see also :class:`treams.operators.Translate`."""
-
-    def __init__(self, arr, ann=(), /, **kwargs):
-        """Initialization."""
-        super().__init__(arr, ann, **kwargs)
-        self._check()
-
-    @property
-    def ann(self):
-        """Array annotations."""
-        # necessary to define the setter below
-        return super().ann
-
-    def __getattr__(self, key):
-        try:
-            return super().__getattr__(key)
-        except AttributeError as err:
-            if key in PhysicsDict.properties:
-                return
-            raise err from None
-
-    def __setattr__(self, key, val):
-        if key in PhysicsDict.properties:
-            val = (val,) * self.ndim if not isinstance(val, tuple) else val
-            self.ann.as_dict[key] = val
-        else:
-            super().__setattr__(key, val)
-
-    @ann.setter
-    def ann(self, ann):
-        """Set array annotations.
-
-        See also :meth:`treams.util.AnnotatedArray.__setitem__`.
-        """
-        self._ann = util.AnnotationSequence(*(({},) * self.ndim), mapping=PhysicsDict)
-        self._ann.update(ann)
-
-    def __repr__(self):
-        """String representation.
-
-        For a more managable output only the special physics properties are shown
-        alongside the array itself.
-        """
-        repr_arr = "    " + repr(self._array)[6:-1].replace("\n  ", "\n") + ","
-        for key in PhysicsDict.properties:
-            if getattr(self, key) is not None:
-                repr_arr += f"\n    {key}={repr(getattr(self, key))},"
-        return f"{self.__class__.__name__}(\n{repr_arr}\n)"
-
-    def _check(self):
-        """Run checks to validate the physical properties.
-
-        The checks are run on the last two dimensions. They include:
-
-            * Dispersion relation checks for :class:`PlaneWaveBasis` if `basis`, `k0`
-              and `material` is defined`
-            * `parity` polarizations are not physically permitted in chiral media.
-            * All lattices explicitly given and in the basis hints must be compatible.
-            * All tangential wave vector compontents must be compatible.
-        """
-        total_lat = None
-        total_kpar = WaveVector()
-        for a in self.ann[-2:]:
-            k0 = a.get("k0")
-            material = a.get("material")
-            modetype = a.get("modetype")
-            poltype = a.get("poltype")
-            basis = a.get("basis", namedtuple("_basis", "lattice kpar")(None, None))
-            lattice = a.get("lattice")
-            for lat in (lattice, basis.lattice):
-                if lat is not None:
-                    total_lat = lat | total_lat
-            for kpar in (a.get("kpar"), basis.kpar):
-                if kpar is not None:
-                    total_kpar = total_kpar & kpar
-            if type(basis) == PlaneWaveBasis and None not in (k0, material):
-                basis.complete(k0, material, modetype)
-            if poltype == "parity" and getattr(material, "ischiral", False):
-                raise ValueError("poltype 'parity' not permitted for chiral material")
-
-    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
-        """Implement ufunc API.
-
-        Additionally to keeping track of the annotations the special properties of a
-        PhysicsArray are also "transparent" in matrix multiplications.
-        """
-        res = super().__array_ufunc__(ufunc, method, *inputs, **kwargs)
-        if (
-            ufunc is np.matmul
-            and method == "__call__"
-            and not isinstance(res, np.generic)
-        ):
-            axes = kwargs.get(
-                "axes", [tuple(range(-min(np.ndim(i), 2), 0)) for i in inputs + (res,)]
-            )
-            anns = [
-                i.ann[ax] if hasattr(i, "ann") else [{}, {}]
-                for i, ax in zip(inputs, axes)
-            ]
-            for name in PhysicsDict.properties:
-                if name in anns[0][-1] and all(name not in a for a in anns[1]):
-                    res.ann[axes[-1][-1]].setdefault(name, anns[0][-1][name])
-                if name in anns[1][0] and all(name not in a for a in anns[0]):
-                    res.ann[axes[-1][0]].setdefault(name, anns[1][0][name])
-        return res

treams/_lattice.py

-import collections
-
-import numpy as np
-
-import treams.lattice as la
-
-
-class Lattice:
-    """Lattice definition.
-
-    The lattice can be one-, two-, or three-dimensional. If it is not three-dimensional
-    it is required to be embedded into a lower dimensional subspace that is aligned with
-    the Cartesian axes. The default alignment for one-dimensional lattices is along the
-    z-axis and for two-dimensional lattices it is in the x-y-plane.
-
-    For a one-dimensional lattice the definition consists of simply the value of the
-    lattice pitch. Higher-dimensional lattices are defined by (2, 2)- or (3, 3)-arrays.
-    If the arrays are diagonal it is sufficient to specify the (2,)- or (3,)-array
-    diagonals. Alternatively, if another instance of a Lattice is given, the defined
-    alignment is extracted, which can be used to separate a lower-dimensional
-    sublattice.
-
-    Lattices are immutable objects.
-
-    Example:
-        >>> Lattice([1, 2])
-        Lattice([[1. 0.]
-                 [0. 2.]], alignment='xy')
-        >>> Lattice(_, 'x')
-        Lattice(1.0, alignment='x')
-
-    Args:
-        arr (float, array, Lattice): Lattice definition. Each row corresponds to one
-            lattice vector, each column to the axis defined in `alignment`.
-        alignment (str, optional): Alignment of the lattice. Possible values are 'x',
-            'y', 'z', 'xy', 'yz', 'zx', and 'xyz'.
-    """
-
-    _allowed_alignments = {
-        "x": "y",
-        "y": "z",
-        "z": "x",
-        "xy": "yz",
-        "yz": "zx",
-        "zx": "xy",
-        "xyz": "xyz",
-    }
-
-    def __init__(self, arr, alignment=None):
-        """Initialization."""
-        if isinstance(arr, Lattice):
-            self._alignment = arr.alignment if alignment is None else alignment
-            self._lattice = arr._sublattice(self._alignment)[...]
-            self._reciprocal = (
-                2 * np.pi / self[...] if self.dim == 1 else la.reciprocal(self[...])
-            )
-            return
-        if arr is None:
-            raise ValueError("Lattice cannot be 'None'")
-
-        arr = np.array(arr, float)
-        arr.flags.writeable = False
-        alignments = {1: ("z", "x", "y"), 4: ("xy", "yz", "zx"), 9: ("xyz",)}
-        if arr.ndim < 3 and arr.size == 1:
-            self._lattice = np.squeeze(arr)
-        elif arr.ndim == 1 and arr.size in (2, 3):
-            self._lattice = np.diag(arr)
-        elif arr.ndim == 2 and arr.shape[0] == arr.shape[1] and arr.shape[0] in (2, 3):
-            self._lattice = arr
-        else:
-            raise ValueError(f"invalid shape '{arr.shape}'")
-
-        alignment = (
-            alignments[self._lattice.size][0] if alignment is None else alignment
-        )
-        if alignment not in alignments[self._lattice.size]:
-            raise ValueError(f"invalid alignment '{alignment}'")
-        self._alignment = alignment
-        if self.volume == 0:
-            raise ValueError("linearly dependent lattice vectors")
-        self._reciprocal = (
-            2 * np.pi / self[...] if self.dim == 1 else la.reciprocal(self[...])
-        )
-
-    @property
-    def alignment(self):
-        """The alignment of the lattice in three-dimensional space.
-
-        For three-dimensional lattices it is always 'xyz' but lower-dimensional lattices
-        have to be aligned with a subset of the axes.
-
-        Returns:
-            str
-        """
-        return self._alignment
-
-    @classmethod
-    def square(cls, pitch, alignment=None):
-        """Create a two-dimensional square lattice.
-
-        Args:
-            pitch (float): Lattice constant.
-            alignment (str, optional): Alignment of the two-dimensional lattice in the
-                three-dimensional space. Defaults to 'xy'.
-
-        Returns:
-            Lattice
-        """
-        return cls(2 * [pitch], alignment)
-
-    @classmethod
-    def cubic(cls, pitch, alignment=None):
-        """Create a three-dimensional cubic lattice.
-
-        Args:
-            pitch (float): Lattice constant.
-            alignment (str, optional): Alignment of the lattice. Defaults to 'xyz'.
-
-        Returns:
-            Lattice
-        """
-        return cls(3 * [pitch], alignment)
-
-    @classmethod
-    def rectangular(cls, x, y, alignment=None):
-        """Create a two-dimensional rectangular lattice.
-
-        Args:
-            x (float): Lattice constant along the first dimension. For the default
-                alignment this corresponds to the x-axis.
-            y (float): Lattice constant along the second dimension. For the default
-                alignment this corresponds to the y-axis.
-            alignment (str, optional): Alignment of the two-dimensional lattice in the
-                three-dimensional space. Defaults to 'xy'.
-
-        Returns:
-            Lattice
-        """
-        return cls([x, y], alignment)
-
-    @classmethod
-    def orthorhombic(cls, x, y, z, alignment=None):
-        """Create a three-dimensional orthorhombic lattice.
-
-        Args:
-            x (float): Lattice constant along the first dimension. For the default
-                alignment this corresponds to the x-axis.
-            y (float): Lattice constant along the second dimension. For the default
-                alignment this corresponds to the y-axis.
-            z (float): Lattice constant along the third dimension. For the default
-                alignment this corresponds to the z-axis.
-            alignment (str, optional): Alignment of the lattice. Defaults to 'xyz'.
-
-        Returns:
-            Lattice
-        """
-        return cls([x, y, z], alignment)
-
-    @classmethod
-    def hexagonal(cls, pitch, height=None, alignment=None):
-        """Create a hexagonal lattice.
-
-        The lattice is two-dimensional if no height is specified else it is
-        three-dimensional
-
-        Args:
-            pitch (float): Lattice constant.
-            height (float, optional): Separation along the third axis for a
-                three-dimensional lattice.
-            alignment (str, optional): Alignment of the two-dimensional lattice in the
-                three-dimensional space. Defaults to either 'xy' or 'xyz' in the two-
-                or three-dimensional case, respectively.
-
-        Returns:
-            Lattice
-        """
-        if height is None:
-            return cls(
-                np.array([[pitch, 0], [0.5 * pitch, np.sqrt(0.75) * pitch]]), alignment
-            )
-        return cls(
-            np.array(
-                [[pitch, 0, 0], [0.5 * pitch, np.sqrt(0.75) * pitch, 0], [0, 0, height]]
-            ),
-            alignment,
-        )
-
-    def __eq__(self, other):
-        """Equality.
-
-        Two lattices are considered equal when they have the same dimension, alignment,
-        and lattice vectors.
-
-        Args:
-            other (Lattice): Lattice to compare with.
-
-        Returns:
-            bool
-        """
-        return other is not None and (
-            self is other
-            or (
-                self.alignment == other.alignment
-                and self.dim == other.dim
-                and np.all(self[...] == other[...])
-            )
-        )
-
-    @property
-    def dim(self):
-        """Dimension of the lattice.
-
-        Returns:
-            int
-        """
-        if self._lattice.ndim == 0:
-            return 1
-        return self._lattice.shape[0]
-
-    @property
-    def volume(self):
-        """(Generalized) volume of the lattice.
-
-        The value gives the lattice pitch, area, or volume depending on its dimension.
-        The volume is signed.
-
-        Returns:
-            float
-        """
-        if self.dim == 1:
-            return self._lattice
-        return la.volume(self._lattice)
-
-    @property
-    def reciprocal(self):
-        r"""Reciprocal lattice.
-
-        The reciprocal lattice to a given lattice with dimension :math:`d` and lattice
-        vectors :math:`\boldsymbol a_i` for :math:`i \in \{1, \dots, d\}` is defined by
-        lattice vectors :math:`\boldsymbol b_j` with :math:`j \in \{1, \dots, d\}` such
-        that :math:`\boldsymbol a_i \boldsymbol b_j = 2 \pi \delta_{ij}` is fulfilled.
-
-        Returns:
-            Lattice
-        """
-        return self._reciprocal
-
-    def __getitem__(self, idx):
-        """Index into the lattice.
-
-        Indexing can be used to obtain entries from the lattice vector definitions or
-        by using the Ellipsis or empty tuple to obtain the full array.
-
-        Returns:
-            float
-        """
-        return self._lattice[idx]
-
-    def __str__(self):
-        """String representation.
-
-        This just prints the lattice pitch or lattice vectors.
-
-        Returns:
-            str
-        """
-        return str(self._lattice)
-
-    def __repr__(self):
-        """Representation.
-
-        The result can be used to recreate an equal instance.
-
-        Returns:
-            str
-        """
-        string = str(self._lattice).replace("\n", "\n        ")
-        return f"Lattice({string}, alignment='{self.alignment}')"
-
-    def _sublattice(self, key):
-        """Get the sublattice defined by key.
-
-        This function is called when one gives another instance of Lattice to the
-        constructor.
-
-        Args:
-            key (str): Alignment of the sublattice to extract.
-
-        Returns:
-            Lattice
-        """
-        key = key.lower()
-        if self.dim == 1:
-            if key == self.alignment:
-                return Lattice(self[...], key)
-            raise ValueError(f"sublattice with key '{key}' not availale")
-        idx = []
-        for c in key:
-            idx.append(self.alignment.find(c))
-        if -1 in idx or key not in self._allowed_alignments:
-            raise ValueError(f"sublattice with key '{key}' not availale")
-        idx_opp = [i for i in range(self.dim) if i not in idx]
-        mask = np.any(self[:, idx] != 0, axis=-1)
-        if (
-            (self[mask, :][:, idx_opp] != 0).any()
-            or (self[np.logical_not(mask), :][:, idx] != 0).any()
-            or len(idx) != sum(mask)
-        ):
-            raise ValueError("cannot determine sublattice")
-        return Lattice(self[mask, :][:, idx], key)
-
-    def permute(self, n=1):
-        """Permute the lattice orientation.
-
-        Get a new lattice with the alignment permuted.
-
-        Examples:
-            >>> Lattice.hexagonal(1, 2).permute()
-            Lattice([[0.        1.        0.       ]
-                     [0.        0.5       0.8660254]
-                     [2.        0.        0.       ]], alignment='xyz')
-            >>> Lattice.hexagonal(1).permute()
-            Lattice([[1.        0.       ]
-                     [0.5       0.8660254]], alignment='yz')
-
-        Args:
-            n (int, optional): Number of repeated permutations. Defaults to `1`.
-
-        Returns:
-            Lattice
-        """
-        if n != int(n):
-            raise ValueError("'n' must be integer")
-        n = n % 3
-        lattice = self[...]
-        alignment = self.alignment
-        if self.dim == 3:
-            while n > 0:
-                lattice = lattice[:, [2, 0, 1]]
-                n -= 1
-        else:
-            while n > 0:
-                alignment = self._allowed_alignments[alignment]
-                n -= 1
-        return Lattice(lattice, alignment)
-
-    def __bool__(self):
-        """Lattice instances always equate to True.
-
-        Returns:
-            bool
-        """
-        return True
-
-    def __or__(self, other):
-        """Merge two lattices.
-
-        Two lattices are combined into one if possible.
-
-        Example:
-            >>> Lattice(1, 'x') | Lattice(2)
-            Lattice([[2. 0.]
-                     [0. 1.]], alignment='zx')
-
-        Args:
-            other (Lattice): Lattice to merge.
-
-        Returns:
-            Lattice
-        """
-        if other is None or self == other:
-            return Lattice(self)
-
-        alignment = list({c for lat in (self, other) for c in lat.alignment})
-        alignment = "".join(sorted(alignment))
-        if alignment == "xz":
-            alignment = "zx"
-
-        ndim = len(alignment)
-
-        if ndim == 2:
-            if self.dim == 1 == other.dim:
-                if alignment.find(self.alignment) == 0:
-                    return Lattice([self[...], other[...]], alignment)
-                return Lattice([other[...], self[...]], alignment)
-            if self.dim == 2 and Lattice(self, other.alignment) == other:
-                return Lattice(self)
-            if other.dim == 2 and Lattice(other, self.alignment) == self:
-                return Lattice(other)
-        elif ndim == 3:
-            if self.dim == 3 and Lattice(self, other.alignment) == other:
-                return Lattice(self)
-            if other.dim == 3 and Lattice(other, self.alignment) == self:
-                return Lattice(other)
-            if self.dim == 2 == other.dim:
-                arr = np.zeros(3)
-                for i, c in enumerate("xyz"):
-                    if c in self.alignment:
-                        la0 = Lattice(self, c)
-                    else:
-                        la0 = None
-                    if c in other.alignment:
-                        la1 = Lattice(other, c)
-                    else:
-                        la1 = None
-                    if None not in (la0, la1) and la0 != la1:
-                        raise ValueError("cannot combine lattices")
-                    arr[i] = la0[...] if la0 is not None else la1[...]
-                return Lattice(arr)
-            arr = np.zeros((3, 3))
-            la0, la1 = (self, other) if self.dim == 1 else (other, self)
-            if la0.alignment == "x":
-                arr[0, 0] = la0[...]
-                arr[1:, 1:] = la1[...]
-                return Lattice(arr)
-            if la0.alignment == "y":
-                arr[1, 1] = la0[...]
-                arr[[[2], [0]], [2, 0]] = la1[...]
-                return Lattice(arr)
-            if la0.alignment == "z":
-                arr[2, 2] = la0[...]
-                arr[:2, :2] = la1[...]
-                return Lattice(arr)
-        raise ValueError("cannot combine lattices")
-
-    def __and__(self, other):
-        """Intersect two lattices.
-
-        The intersection of two lattices is taken if possible.
-
-        Example:
-            >>> Lattice([1, 2]) & Lattice([2, 3], 'yz')
-            Lattice(2.0, alignment='y')
-
-        Args:
-            other (Lattice): Lattice to intersect.
-
-        Returns:
-            Lattice
-        """
-        if other is None:
-            return None
-        if self == other:
-            return Lattice(self)
-
-        alignment = list({c for c in self.alignment if c in other.alignment})
-        if len(alignment) == 0:
-            raise ValueError("cannot combine lattices")
-        alignment = "".join(sorted(alignment))
-        if alignment == "xz":
-            alignment = "zx"
-
-        a, b = Lattice(self, alignment), Lattice(other, alignment)
-        if a == b:
-            return a
-        raise ValueError("cannot combine lattices")
-
-    def __le__(self, other):
-        """Test if one lattice includes another.
-
-        Example:
-            >>> Lattice(3) <= Lattice([1, 2, 3])
-            True
-
-        Args:
-            other (Lattice): Lattice to compare with.
-
-        Returns:
-            bool
-        """
-        try:
-            lat = Lattice(other, self.alignment)
-        except ValueError:
-            return False
-        return lat == self
-
-    def isdisjoint(self, other):
-        """Test if lattices are disjoint.
-
-        Lattices are considered disjoint if their alignments are disjoint.
-
-        Example:
-            >>> Lattice([1, 2, 3]).isdisjoint(Lattice(1))
-            False
-
-        Args:
-            other (Lattice): Lattice to compare with.
-
-        Returns:
-            bool
-        """
-        for c in other.alignment:
-            if c in self.alignment:
-                return False
-        return True
-
-
-class WaveVector(collections.abc.Sequence):
-    def __init__(self, seq=(), alignment=None):
-        try:
-            length = len(seq)
-        except TypeError:
-            length = 1
-            seq = [seq]
-        if length == 3:
-            self._vec = tuple(seq)
-        elif length == 2:
-            if alignment in ("xy", None):
-                self._vec = (seq[0], seq[1], np.nan)
-            elif alignment == "yz":
-                self._vec = (np.nan, seq[0], seq[1])
-            elif alignment == "zx":
-                self._vec = (seq[1], np.nan, seq[0])
-            else:
-                raise ValueError(f"invalid alignment: {alignment}")
-        elif length == 1:
-            if alignment in ("z", None):
-                self._vec = (np.nan, np.nan, seq[0])
-            elif alignment == "y":
-                self._vec = (np.nan, seq[0], np.nan)
-            elif alignment == "x":
-                self._vec = (seq[0], np.nan, np.nan)
-            else:
-                raise ValueError(f"invalid alignment: {alignment}")
-        elif length == 0:
-            self._vec = (np.nan,) * 3
-        else:
-            raise ValueError("invalid sequence")
-
-    def __str__(self):
-        return str(self._vec)
-
-    def __repr__(self):
-        return self.__class__.__name__ + str(self._vec)
-
-    def __eq__(self, other):
-        other = WaveVector(other)
-        for a, b in zip(self, other):
-            if a != b and not (np.isnan(a) and np.isnan(b)):
-                return False
-        return True
-
-    def __len__(self):
-        return 3
-
-    def __getitem__(self, key):
-        return self._vec[key]
-
-    def __or__(self, other):
-        other = WaveVector(other)
-        seq = ()
-        for a, b in zip(self, other):
-            isnan = np.isnan(a) or np.isnan(b)
-            if a != b and not isnan:
-                raise ValueError("non-matching WaveVector")
-            seq += (np.nan,) if isnan else (a,)
-        return WaveVector(seq)
-
-    def __and__(self, other):
-        other = WaveVector(other)
-        seq = ()
-        for a, b in zip(self, other):
-            if a != b and not (np.isnan(a) or np.isnan(b)):
-                raise ValueError("non-matching WaveVector")
-            seq += (b,) if np.isnan(a) else (a,)
-        return WaveVector(seq)
-
-    def permute(self, n=1):
-        x, y, z = self
-        n = n % 3
-        for i in range(n):
-            x, y, z = z, x, y
-        return WaveVector((x, y, z))
-
-    def __le__(self, other):
-        other = WaveVector(other)
-        for a, b in zip(self, other):
-            if a != b and not np.isnan(b):
-                return False
-        return True
-
-    def isdisjoint(self, other):
-        other = WaveVector(other)
-        for a, b in zip(self, other):
-            if not (np.isnan(a) or np.isnan(b)):
-                return False
-        return True

treams/_material.py

-import numpy as np
-
-from treams import misc
-
-
-class Material:
-    r"""Material definition.
-
-    The material properties are defined in the frequency domain through scalar values
-    for permittivity :math:`\epsilon`, permeability :math:`\mu`, and the chirality
-    parameter  :math:`\kappa`. Materials are, thus, assumed to be linear,
-    time-invariant, homogeneous, isotropic, and local. Also, it is assumed that they
-    have no gain. The relation of the electric and magnetic fields is defined by
-
-    .. math::
-
-        \begin{pmatrix}
-            \frac{1}{\epsilon_0} \boldsymbol D \\
-            c \boldsymbol B
-        \end{pmatrix}
-        =
-        \begin{pmatrix}
-            \epsilon & \mathrm i \kappa \\
-            -\mathrm i \kappa & \mu
-        \end{pmatrix}
-        \begin{pmatrix}
-            \boldsymbol E \\
-            Z_0 \boldsymbol H
-        \end{pmatrix}
-
-    for a fixed vacuum wave number :math:`k_0` and spatial position
-    :math:`\boldsymbol r` with :math:`\epsilon_0` the vacuum permittivity, :math:`c` the
-    speed of light in vacuum, and :math:`Z_0` the vacuum impedance.
-
-    Args:
-        epsilon (optional, complex): Relative permittivity. Defaults to 1.
-        mu (optional, complex): Relative permeability. Defaults to 1.
-        kappa (optional, complex): Chirality parameter. Defaults to 0.
-    """
-
-    def __init__(self, epsilon=1, mu=1, kappa=0):
-        """Initialization."""
-        if isinstance(epsilon, Material):
-            epsilon, mu, kappa = epsilon()
-        elif isinstance(epsilon, (tuple, list, np.ndarray)):
-            if len(epsilon) == 0:
-                epsilon = 1
-            elif len(epsilon) == 1:
-                epsilon = epsilon[0]
-            elif len(epsilon) == 2:
-                epsilon, mu = epsilon
-            elif len(epsilon) == 3:
-                epsilon, mu, kappa = epsilon
-            else:
-                raise ValueError("invalid material definition")
-        self._epsilon = epsilon
-        self._mu = mu
-        self._kappa = kappa
-
-    @property
-    def epsilon(self):
-        """Relative permittivity.
-
-        Returns:
-            float or complex
-        """
-        return self._epsilon
-
-    @property
-    def mu(self):
-        """Relative permeability.
-
-        Returns:
-            float or complex
-        """
-        return self._mu
-
-    @property
-    def kappa(self):
-        """Chirality parameter.
-
-        Returns:
-            float or complex
-        """
-        return self._kappa
-
-    def __iter__(self):
-        """Iterator for a tuple containing the material parameters.
-
-        Useful for unpacking the material parameters into a function that takes these
-        parameters separately, e.g. ``foo(*material)``.
-        """
-        return iter((self.epsilon, self.mu, self.kappa))
-
-    @classmethod
-    def from_n(cls, n=1, impedance=None, kappa=0):
-        r"""Create material from refractive index and relative impedance.
-
-        This function calculates the relative permeability and permittivity with
-        :math:`\epsilon = \frac{n}{Z}` and :math:`\mu = nZ`. The chirality parameter is
-        considered separately.
-
-        Note:
-            The refractive index is defined independently from the chirality parameter
-            here. For an alternative definition of the refractive index, see also
-            :func:`Material.from_nmp`.
-
-        Args:
-            n (complex, optional): Refractive index. Defaults to 1.
-            impedance(complex, optional): Relative impedance. Defaults to the inverse
-                of the refractive index.
-            kappa (complex, optional): Chirality parameter. Defaults to 0.
-
-        Returns:
-            Material
-        """
-        impedance = 1 / n if impedance is None else impedance
-        epsilon = n / impedance
-        mu = n * impedance
-        return cls(epsilon, mu, kappa)
-
-    @classmethod
-    def from_nmp(cls, ns=(1, 1), impedance=None):
-        r"""Create material from refractive indices of both helicities.
-
-        This function calculates the relative permeability and permittivity and the
-        chirality parameter with :math:`\epsilon = \frac{n_+ + n_-}{2Z}`,
-        :math:`\mu = \frac{(n_+ + n_-)Z}{2}` and :math:`\mu = \frac{(n_+ - n_-)}{2}`.
-
-        Note:
-            Two refractive indices are defined here that depend on the chirality
-            parameter. For an alternative definition of the refractive index, see also
-            :func:`Material.from_n`.
-
-        Args:
-            ns ((2,)-array-like, optional): Negative and positive helicity refractive
-                index. Defaults to (1, 1).
-            impedance(complex, optional): Relative impedance. Defaults to the inverse of
-                the refractive index.
-
-        Returns:
-            Material
-        """
-        n = sum(ns) * 0.5
-        kappa = (ns[1] - ns[0]) * 0.5
-        return cls.from_n(n, impedance, kappa)
-
-    @property
-    def n(self):
-        r"""Refractive index.
-
-        The refractive index is defined by :math:`n = \sqrt{\epsilon \mu}`, with an
-        enforced non-negative imaginary part.
-
-        Note:
-            The refractive index returned is independent from the chirality parameter.
-            For an alternative definition of the refractive index, see also
-            :func:`Material.nmp`.
-
-        Returns:
-            complex
-        """
-        n = np.sqrt(self.epsilon * self.mu)
-        if n.imag < 0:
-            n = -n
-        return n
-
-    @property
-    def nmp(self):
-        r"""Refractive indices of both helicities.
-
-        The refractive indices are defined by
-        :math:`n_\pm = \sqrt{\epsilon \mu} \pm \kappa`, with an enforced non-negative
-        imaginary part.
-
-        Note:
-            The refractive indices returned depend on the chirality parameter. For an
-            alternative definition of the refractive index, see also :func:`Material.n`.
-
-        Returns:
-            tuple
-        """
-        return misc.refractive_index(self.epsilon, self.mu, self.kappa)
-
-    @property
-    def impedance(self):
-        r"""Relative impedance.
-
-        The relative impedance is defined by :math:`Z = \sqrt{\frac{\epsilon}{\mu}}`.
-
-        Returns:
-            complex
-        """
-        return np.sqrt(self.mu / self.epsilon)
-
-    def __call__(self):
-        """Return a tuple containing all material parameters.
-
-        Returns:
-            tuple
-        """
-        return self.epsilon, self.mu, self.kappa
-
-    def __eq__(self, other):
-        """Compare material parameters.
-
-        Materials are considered equal, when all material parameters are equal. Also,
-        compares with objects that contain at most three values.
-
-        Returns:
-            bool
-        """
-        if other is None:
-            return False
-        if not isinstance(other, Material):
-            other = Material(*other)
-        return (
-            self.epsilon == other.epsilon
-            and self.mu == other.mu
-            and self.kappa == other.kappa
-        )
-
-    @property
-    def ischiral(self):
-        """Test if the material is chiral.
-
-        Returns:
-            bool
-        """
-        return self.kappa != 0
-
-    @property
-    def isreal(self):
-        """Test if the material has purely real parameters.
-
-        Returns:
-            bool
-        """
-        return all(i.imag == 0 for i in self)
-
-    def __str__(self):
-        """All three material parameters.
-
-        Returns:
-            str
-        """
-        return "(" + ", ".join([str(i) for i in self()]) + ")"
-
-    def __repr__(self):
-        """Representation that allows recreating the object.
-
-        Returns:
-            str
-        """
-        return self.__class__.__name__ + str(self)
-
-    def ks(self, k0):
-        """Return the wave numbers in the medium for both polarizations.
-
-        The first value corresponds to negative helicity and the second to positive
-        helicity. For achiral materials where parity polarizations can be used both
-        values are equal.
-
-        Args:
-            k0 (float): Wave number in vacuum.
-
-        Returns:
-            tuple
-        """
-        return k0 * self.nmp
-
-    def krhos(self, k0, kz, pol):
-        r"""The (cylindrically) radial part of the wave vector.
-
-        The cylindrically radial part is defined by :math:`k_\rho = \sqrt(k^2 - k_z^2)`.
-        In case of chiral materials :math:`k` and so :math`k_\rho` depends on the
-        polarization. The returned values have non-negative imaginary parts.
-
-        Args:
-            k0 (float): Wave number in vacuum.
-            kz (float, array-like): z-component of the wave vector
-            pol (int, array-like): Polarization indices
-                (:ref:`params:Polarizations`).
-
-        Returns:
-            complex, array-like
-        """
-        ks = self.ks(k0)[pol]
-        return misc.wave_vec_z(kz, 0, ks)
-
-    def kzs(self, k0, kx, ky, pol):
-        r"""The z-component of the wave vector.
-
-        The z-component of the wave vector is defined by
-        :math:`k_z = \sqrt(k^2 - k_x^2 - k_y^2)`. In case of chiral materials :math:`k`
-        and so :math`k_z` depends on the polarization. The returned values have
-        non-negative imaginary parts.
-
-        Args:
-            k0 (float): Wave number in vacuum.
-            kx (float, array-like): x-component of the wave vector
-            ky (float, array-like): y-component of the wave vector
-            pol (int, array-like): Polarization indices
-                (:ref:`params:Polarizations`).
-
-        Returns:
-            complex, array-like
-        """
-        ks = self.ks(k0)[pol]
-        return misc.wave_vec_z(kx, ky, ks)

treams/_smatrix.py

-import copy
-import itertools
-
-import numpy as np
-
-from treams import config, util
-from treams._core import PhysicsArray
-from treams._core import PlaneWaveBasisByComp as PWBC
-from treams._core import PlaneWaveBasisByUnitVector as PWBUV
-from treams._material import Material
-from treams._operators import translate
-from treams._tmatrix import TMatrix, TMatrixC
-from treams.coeffs import fresnel
-
-
-class SMatrix(PhysicsArray):
-    """S-matrix for a plane wave."""
-
-    def _check(self):
-        super()._check()
-        shape = np.shape(self)
-        if len(shape) != 2 or shape[0] != shape[1]:
-            raise util.AnnotationError(f"invalid shape: '{shape}'")
-        if not isinstance(self.k0, (int, float, np.floating, np.integer)):
-            raise util.AnnotationError("invalid k0")
-        if self.poltype is None:
-            self.poltype = config.POLTYPE
-        if self.poltype not in ("parity", "helicity"):
-            raise util.AnnotationError("invalid poltype")
-        material = (None, None) if self.material is None else self.material
-        if isinstance(material, tuple):
-            self.material = tuple(Material() if m is None else m for m in material)
-        if self.basis is None:
-            raise util.AnnotationError("basis not set")
-        elif not isinstance(self.basis, PWBC):
-            self.basis = self.basis.bycomp(self.k0, self.material)
-
-
-class SMatrices:
-    r"""Collection of four S-matrices with a plane wave basis.
-
-    The S-matrix describes the scattering of incoming into outgoing modes using a plane
-    wave basis, with functions :func:`treams.special.vsw_A`,
-    :func:`treams.special.vsw_M`, and :func:`treams.special.vsw_N`. The primary
-    direction of propagation is parallel or anti-parallel to the z-axis. The scattering
-    object itself is infinitely extended in the x- and y-directions. The S-matrix is
-    divided into four submatrices :math:`S_{\uparrow \uparrow}`,
-    :math:`S_{\uparrow \downarrow}`, :math:`S_{\downarrow \uparrow}`, and
-    :math:`S_{\downarrow \downarrow}`:
-
-    .. math::
-
-        S = \begin{pmatrix}
-            S_{\uparrow \uparrow} & S_{\uparrow \downarrow} \\
-            S_{\downarrow \uparrow} & S_{\downarrow \downarrow}
-        \end{pmatrix}\,.
-
-    These matrices describe the transmission of plane waves propagating into positive
-    z-direction, reflection of plane waves into the positive z-direction, reflection of
-    plane waves into negative z-direction, and the transmission of plane waves
-    propagating in negative z-direction, respectively. Each of those for matrices
-    contain different diffraction orders and different polarizations. The polarizations
-    can be in either helicity of parity basis.
-
-    The wave number :attr:`k0` and, if not vacuum, the material :attr:`material` are
-    also required.
-
-    Args:
-        smats (SMatrix): S-matrices.
-        k0 (float): Wave number in vacuum.
-        basis (PlaneWaveBasisByComp): The basis for the S-matrices.
-        material (tuple, Material, optional): Material definition, if a tuple of length
-            two is specified, this refers to the materials above and below the S-matrix.
-        poltype (str, optional): Polarization type (:ref:`params:Polarizations`).
-    """
-
-    def __init__(self, smats, **kwargs):
-        """Initialization."""
-        if len(smats) != 2 or len(smats[0]) != 2 or len(smats[1]) != 2:
-            raise ValueError("invalid shape of S-matrices")
-        if "material" in kwargs:
-            material = kwargs["material"]
-            del kwargs["material"]
-        elif hasattr(smats[0][0], "material"):
-            material = smats[0][0].material
-        elif hasattr(smats[1][1], "material"):
-            material = smats[1][1].material
-            if isinstance(material, tuple):
-                material = material[::-1]
-        elif hasattr(smats[0][1], "material"):
-            material = smats[0][1].material
-        elif hasattr(smats[1][0], "material"):
-            material = smats[1][0].material
-        else:
-            material = Material()
-        if isinstance(material, tuple):
-            ma, mb = material
-        else:
-            ma = mb = Material(material)
-        material = [[(ma, mb), (ma, ma)], [(mb, mb), (mb, ma)]]
-        modetype = [[("up", "up"), ("up", "down")], [("down", "up"), ("down", "down")]]
-        self._sms = [
-            [
-                SMatrix(s, material=m, modetype=t, **kwargs)
-                for s, m, t in zip(ar, mr, tr)
-            ]
-            for ar, mr, tr in zip(smats, material, modetype)
-        ]
-        self.material = self[0, 0].material
-        self.basis = self[0, 0].basis
-        self.k0 = self[0, 0].k0
-        self.poltype = self[0, 0].poltype
-        for i in (self[1, 0], self[0, 1], self[1, 1]):
-            i.k0 = self.k0
-            i.basis = self.basis
-            i.poltype = self.poltype
-
-    def __eq__(self, other):
-        return all(
-            (np.abs(self[i, j] - other[i][j]) < 1e-14).all()
-            for i in range(2)
-            for j in range(2)
-        )
-
-    def __getitem__(self, key):
-        keys = {0: 0, "up": 0, 1: 1, "down": 1}
-        if key in keys:
-            return self._sms[keys[key]]
-        if not isinstance(key, tuple) or len(key) not in (1, 2):
-            raise KeyError("invalid key")
-        if len(key) == 1:
-            return self[key[0]]
-        key = tuple(keys[k] for k in key)
-        return self._sms[key[0]][key[1]]
-
-    @classmethod
-    def interface(cls, basis, k0, materials):
-        """Planar interface between two media.
-
-        Args:
-            basis (PlaneWaveBasisByComp): Basis definitions.
-            k0 (float): Wave number in vacuum
-            materials (Sequence[Material]): Material definitions.
-
-        Returns:
-            SMatrix
-        """
-        materials = tuple(map(Material, materials))
-        ks = np.array([m.ks(k0) for m in materials])
-        choice = basis.pol == 0
-        kxs = basis.kx[choice], basis.kx[~choice]
-        kys = basis.ky[choice], basis.ky[~choice]
-        qs = np.zeros((2, 2, len(basis), len(basis)), complex)
-        if all(kxs[0] == kxs[1]) and all(kys[0] == kys[1]):
-            kzs = np.stack(
-                [
-                    m.kzs(k0, kxs[0][:, None], kys[0][:, None], [[0, 1]])
-                    for m in materials
-                ],
-                -2,
-            )
-            vals = fresnel(ks, kzs, [m.impedance for m in materials])
-            qs[:, :, choice, choice] = np.moveaxis(vals[:, :, :, 0, 0], 0, -1)
-            qs[:, :, choice, ~choice] = np.moveaxis(vals[:, :, :, 0, 1], 0, -1)
-            qs[:, :, ~choice, choice] = np.moveaxis(vals[:, :, :, 1, 0], 0, -1)
-            qs[:, :, ~choice, ~choice] = np.moveaxis(vals[:, :, :, 1, 1], 0, -1)
-        else:
-            for i, (kx, ky, pol) in enumerate(basis):
-                for ii, (kx2, ky2, pol2) in enumerate(basis):
-                    if kx != kx2 or ky != ky2:
-                        continue
-                    kzs = np.array([m.kzs(k0, kx, ky) for m in materials])
-                    vals = fresnel(ks, kzs, [m.impedance for m in materials])
-                    qs[:, :, ii, i] = vals[:, :, pol2, pol]
-        return cls(qs, k0=k0, basis=basis, material=materials[::-1], poltype="helicity")
-
-    @classmethod
-    def slab(cls, thickness, basis, k0, materials):
-        """Slab of material.
-
-        A slab of material, defined by a thickness and three materials. Consecutive
-        slabs of material can be defined by `n` thicknesses and and `n + 2` material
-        parameters.
-
-        Args:
-            k0 (float): Wave number in vacuum.
-            basis (PlaneWaveBasisByComp): Basis definition.
-            tickness (Sequence[Float]): Thickness of the slab or the thicknesses of all
-                slabs in order from negative to positive z.
-            materials (Sequenze[Material]): Material definitions from negative to
-                positive z.
-
-        Returns:
-            SMatrix
-        """
-        try:
-            iter(thickness)
-        except TypeError:
-            thickness = [thickness]
-        res = cls.interface(basis, k0, materials[:2])
-        for d, ma, mb in zip(thickness, materials[1:-1], materials[2:]):
-            if np.ndim(d) == 0:
-                d = [0, 0, d]
-            x = cls.propagation(d, basis, k0, ma)
-            res = res.add(x)
-            res = res.add(cls.interface(basis, k0, (ma, mb)))
-        return res
-
-    @classmethod
-    def stack(cls, items):
-        """Stack of S-matrices.
-
-        Electromagnetically couple multiple S-matrices in the order given. Before
-        coupling it can be checked for matching materials and modes.
-
-        Args:
-            items (Sequence[SMatrix]): An array of S-matrices in their intended order
-                from negative to positive z.
-
-        Returns:
-            SMatrix
-        """
-        acc = items[0]
-        for item in items[1:]:
-            acc = acc.add(item)
-        return acc
-
-    @classmethod
-    def propagation(cls, r, basis, k0, material=Material()):
-        """S-matrix for the propagation along a distance.
-
-        This S-matrix translates the reference origin along `r`.
-
-        Args:
-            r (float, (3,)-array): Translation vector.
-            k0 (float): Wave number in vacuum.
-            basis (PlaneWaveBasis): Basis definition.
-            material (Material, optional): Material definition.
-
-        Returns:
-            SMatrix
-        """
-        sup = translate(r, basis=basis, k0=k0, material=material, modetype="up")
-        sdown = translate(
-            np.negative(r), basis=basis, k0=k0, material=material, modetype="down"
-        )
-        zero = np.zeros_like(sup)
-        material = Material(material)
-        return cls([[sup, zero], [zero, sdown]], basis=basis, k0=k0, material=material)
-
-    @classmethod
-    def from_array(cls, tm, basis, *, eta=0):
-        """S-matrix from an array of (cylindrical) T-matrices.
-
-        Create a S-matrix for a two-dimensional array of objects described by the
-        T-Matrix or an one-dimensional array of objects described by a cylindrical
-        T-matrix.
-
-        Args:
-            tm (TMatrix or TMatrixC): (Cylindrical) T-matrix to place in the array.
-            basis (PlaneWaveBasisByComp): Basis definition.
-            eta (float or complex, optional): Splitting parameter in the lattice sum.
-
-        Returns:
-            SMatrix
-        """
-        if isinstance(tm, TMatrixC):
-            basis = basis.permute()
-        pu, pd = (tm.expandlattice(basis=basis, modetype=i) for i in ("up", "down"))
-        au, ad = (tm.expand.eval_inv(basis, i) for i in ("up", "down"))
-        eye = np.eye(len(basis))
-        return cls([[eye + pu @ au, pu @ ad], [pd @ au, eye + pd @ ad]])
-
-    def add(self, sm):
-        """Couple another S-matrix on top of the current one.
-
-        See also :func:`treams.SMatrix.stack` for a function that does not change the
-        current S-matrix but creates a new one.
-
-        Args:
-            sm (SMatrix): S-matrix to add.
-
-        Returns:
-            SMatrix
-        """
-        dim = len(self.basis)
-        snew = [[None, None], [None, None]]
-        s_tmp = np.linalg.solve(np.eye(dim) - self[0, 1] @ sm[1][0], self[0, 0])
-        snew[0][0] = sm[0][0] @ s_tmp
-        snew[1][0] = self[1, 0] + self[1, 1] @ sm[1][0] @ s_tmp
-        s_tmp = np.linalg.solve(np.eye(dim) - sm[1][0] @ self[0, 1], sm[1][1])
-        snew[1][1] = self[1, 1] @ s_tmp
-        snew[0][1] = sm[0][1] + sm[0][0] @ self[0, 1] @ s_tmp
-        return SMatrices(snew)
-
-    def double(self, n=1):
-        """Double the S-matrix.
-
-        By default this function doubles the S-matrix but it can also create a
-        :math:`2^n`-fold repetition of itself:
-
-        Args:
-            n (int, optional): Number of times to double itself. Defaults to 1.
-
-        Returns:
-            SMatrix
-        """
-        res = self
-        for _ in range(n):
-            res = res.add(self)
-        return res
-
-    def illuminate(self, illu, illu2=None, /, *, smat=None):
-        """Field coefficients above and below the S-matrix.
-
-        Given an illumination defined by the coefficients of each incoming mode
-        calculate the coefficients for the outgoing field above and below the S-matrix.
-        If a second SMatrix is given, the field expansions in between are also
-        calculated.
-
-        Args:
-            illu (array-like): Illumination, if `modetype` is specified, the direction
-                will be chosen accordingly.
-            illu2 (array-like, optional): Second illumination. If used, the first
-                argument is taken to be coming from below and this one to be coming from
-                above.
-            smat (SMatrix, optional): Second S-matrix for the calculation of the
-                field expansion between two S-matrices.
-
-        Returns:
-            tuple
-        """
-        modetype = getattr(illu, "modetype", "up")
-        if isinstance(modetype, tuple):
-            modetype = modetype[max(-2, -len(tuple))]
-        illu2 = np.zeros(np.shape(illu)[-2:]) if illu2 is None else illu2
-        if modetype == "down":
-            illu, illu2 = illu2, illu
-        if smat is None:
-            field_up = self[0, 0] @ illu + self[0, 1] @ illu2
-            field_down = self[1, 0] @ illu + self[1, 1] @ illu2
-            return field_up, field_down
-        stmp = np.eye(len(self.basis)) - self[0, 1] @ smat[1, 0]
-        field_in_up = np.linalg.solve(
-            stmp, self[0, 0] @ illu + self[0, 1] @ smat[1, 1] @ illu2
-        )
-        field_in_down = smat[1, 0] @ field_in_up + smat[1, 1] @ illu2
-        field_up = smat[0, 1] @ illu2 + smat[0, 0] @ field_in_up
-        field_down = self[1, 0] @ illu + self[1, 1] @ field_in_down
-        return field_up, field_down, field_in_up, field_in_down
-
-    def tr(self, illu, direction=1):
-        """Transmittance and reflectance.
-
-        Calculate the transmittance and reflectance for one S-matrix with the given
-        illumination and direction.
-
-        Args:
-            illu (complex, array_like): Expansion coefficients for the incoming light
-            direction (int, optional): The direction of the field, options are `-1` and
-                `1`
-
-        Returns:
-            tuple
-        """
-        modetype = getattr(illu, "modetype", "up")
-        if isinstance(modetype, tuple):
-            modetype = modetype[max(-2, -len(tuple))]
-        trans, refl = self.illuminate(illu)
-        material = self.material
-        if not isinstance(material, tuple):
-            material = material, material
-        paz = [poynting_avg_z(self.basis, self.k0, m, self.poltype) for m in material]
-        if modetype == "down":
-            trans, refl = refl, trans
-            paz.reverse()
-        illu = np.asarray(illu)
-        s_t = np.real(trans.conjugate().T @ paz[0][0] @ trans)
-        s_r = np.real(refl.conjugate().T @ paz[1][0] @ refl)
-        s_i = np.real(np.conjugate(illu).T @ paz[1][0] @ illu)
-        s_ir = np.real(
-            refl.conjugate().T @ paz[1][1] @ illu
-            - np.conjugate(illu).T @ paz[1][1] @ refl
-        )
-        return s_t / (s_i + s_ir), s_r / (s_i + s_ir)
-
-    def cd(self, illu):
-        """Transmission and absorption circular dichroism.
-
-        Calculate the transmission and absorption CD for one S-matrix with the given
-        illumination and direction.
-
-        Args:
-            illu (complex, PhysicsArray): Expansion coefficients for the incoming light
-
-        Returns:
-            tuple
-        """
-        minus, plus = self.basis.pol == 0, self.basis.pol == 1
-        if self.poltype == "helicity":
-            illuopposite = np.zeros_like(illu)
-            illuopposite[minus] = illu[plus]
-            illuopposite[plus] = illu[minus]
-        else:
-            illuopposite = copy.deepcopy(illu)
-            illuopposite[minus] *= -1
-        tm, rm = self.tr(illu)
-        tp, rp = self.tr(illuopposite)
-        return (tp - tm) / (tp + tm), (tp + rp - tm - rm) / (tp + rp + tm + rm)
-
-    def periodic(self):
-        r"""Periodic repetition of the S-matrix.
-
-        Transform the S-matrix to an infinite periodic arrangement of itself defined
-        by
-
-        .. math::
-
-            \begin{pmatrix}
-                S_{\uparrow \uparrow} & S_{\uparrow \downarrow} \\
-                -S_{\downarrow \downarrow}^{-1}
-                S_{\downarrow \uparrow} S_{\uparrow \uparrow} &
-                S_{\downarrow \downarrow}^{-1} (\mathbb{1} - S_{\downarrow \uparrow}
-                S_{\uparrow \downarrow})
-            \end{pmatrix}\,.
-
-        Returns:
-            complex, array_like
-
-        """
-        dim = len(self.basis)
-        res = np.empty((2 * dim, 2 * dim), dtype=complex)
-        res[0:dim, 0:dim] = self[0, 0]
-        res[0:dim, dim:] = self[0, 1]
-        res[dim:, 0:dim] = -np.linalg.solve(self[1, 1], self[1, 0] @ self[0, 0])
-        res[dim:, dim:] = np.linalg.solve(
-            self[1, 1], np.eye(dim) - self[1, 0] @ self[0, 1]
-        )
-        return res
-
-    def bands_kz(self, az):
-        r"""Band structure calculation.
-
-        Calculate the band structure for the given S-matrix, assuming it is periodically
-        repeated along the z-axis. The function returns the z-components of the wave
-        vector :math:`k_z` and the corresponding eigenvectors :math:`v` of
-
-        .. math::
-
-            \begin{pmatrix}
-                S_{\uparrow \uparrow} & S_{\uparrow \downarrow} \\
-                -S_{\downarrow \downarrow}^{-1} S_{\downarrow \uparrow}
-                    S_{\uparrow \uparrow} &
-                S_{\downarrow \downarrow}^{-1} (\mathbb{1} - S_{\downarrow \uparrow}
-                    S_{\uparrow \downarrow})
-            \end{pmatrix}
-            \boldsymbol v
-            =
-            \mathrm{e}^{\mathrm{i}k_z a_z}
-            \boldsymbol v\,.
-
-        Args:
-            az (float): Lattice pitch along the z direction
-
-        Returns:
-            tuple
-
-        """
-        w, v = np.linalg.eig(self.periodic())
-        return -1j * np.log(w) / az, v
-
-    def __repr__(self):
-        return f"""{type(self).__name__}(...,
-    basis={self.basis},
-    k0={self.k0},
-    material={self.material},
-    poltype="{self.poltype}",
-)"""
-
-
-def chirality_density(basis, k0, material=Material(), poltype=None, z=(0, 0)):
-    poltype = config.POLTYPE if poltype is None else poltype
-    material = Material(material)
-    kx, ky, kz = basis.kvecs(k0, material)
-    k = material.ks(k0)[basis.pol]
-    re, im = np.real(kz), np.imag(kz)
-    prefuu = 1 + np.real((k * k - kz * 2j * im) / (k * k.conjugate()))
-    prefdu = 1 + np.real((k * k - kz * 2 * re) / (k * k.conjugate()))
-    if np.abs(z[1] - z[0]) > 1e-16:
-        prefuu *= np.sinh(re * (z[1] - z[0])) / (re * (z[1] - z[0]))
-        prefdd = np.exp(im * (z[1] + z[0])) * prefuu
-        prefuu *= np.exp(-im * (z[1] + z[0]))
-        prefdu *= (
-            np.sin(re * (z[1] - z[0]))
-            * np.cos(re * (z[1] + z[0]))
-            / (re * (z[1] - z[0]))
-        )
-    else:
-        prefdd = prefuu
-    if poltype == "helicity":
-        return (
-            np.diag(prefuu * (2 * basis.pol - 1)),
-            np.diag(prefdd * (2 * basis.pol - 1)),
-            np.diag(2 * prefdu * (2 * basis.pol - 1)),
-        )
-    elif poltype == "parity":
-        where = (
-            (kx[:, None] == kx)
-            & (ky[:, None] == ky)
-            & (basis.pol[:, None] != basis.pol)
-        )
-        return where * prefuu, where * prefdd, where * 2 * prefdu
-    raise ValueError(f"invalid poltype: '{poltype}'")
-
-
-def poynting_avg_z(basis, k0, material=Material(), poltype=None):
-    r"""Time-averaged z-component of the Poynting vector.
-
-    Calculate the time-averaged Poynting vector's z-component
-
-    .. math::
-
-        \langle S_z \rangle
-        = \frac{1}{2}
-        \Re (\boldsymbol E \times \boldsymbol H^\ast) \boldsymbol{\hat{z}}
-
-    on one side of the S-matrix with the given coefficients.
-
-    Returns:
-        tuple
-    """
-    poltype = config.POLTYPE if poltype is None else poltype
-    material = Material(material)
-    kx, ky, kzs = basis.kvecs(k0, material)
-    gamma = kzs / (material.ks(k0)[basis.pol] * material.impedance)
-    selection = (kx[:, None] == kx) & (ky[:, None] == ky)
-    if poltype == "parity":
-        pol = basis.pol
-        selection = selection & (pol[:, None] == pol)
-        a = selection * ((1 - pol) * gamma.conjugate() + pol * gamma) * 0.25
-        b = selection * ((1 - pol) * gamma.conjugate() - pol * gamma) * 0.25
-        return a, b
-    if poltype == "helicity":
-        pol = 2 * basis.pol - 1
-        a = selection * (pol[:, None] * pol * gamma[:, None].conjugate() + gamma) * 0.25
-        b = selection * (pol[:, None] * pol * gamma[:, None].conjugate() - gamma) * 0.25
-        return a, b
-    raise ValueError(f"invalid poltype: '{poltype}'")

treams/_tmatrix.py

-import warnings
-
-import numpy as np
-
-import treams.special as sc
-from treams import config
-from treams._core import CylindricalWaveBasis as CWB
-from treams._core import PhysicsArray
-from treams._core import PlaneWaveBasisByUnitVector as PWBUV
-from treams._core import PlaneWaveBasisByComp as PWBC
-from treams._core import SphericalWaveBasis as SWB
-import treams._operators as op
-from treams._material import Material
-from treams.coeffs import mie, mie_cyl
-from treams.util import AnnotationError
-
-
-class _Interaction:
-    def __init__(self):
-        self._obj = self._objtype = None
-
-    def __get__(self, obj, objtype=None):
-        self._obj = obj
-        self._objtype = objtype
-        return self
-
-    def __call__(self):
-        basis = self._obj.basis
-        return (
-            np.eye(self._obj.shape[-1]) - self._obj @ op.Expand(basis, "singular").inv
-        )
-
-    def solve(self):
-        return np.linalg.solve(self(), self._obj)
-
-
-class _LatticeInteraction:
-    def __init__(self):
-        self._obj = self._objtype = None
-
-    def __get__(self, obj, objtype=None):
-        self._obj = obj
-        self._objtype = objtype
-        return self
-
-    def __call__(self, lattice, kpar):
-        return np.eye(self._obj.shape[-1]) - self._obj @ op.ExpandLattice(
-            lattice=lattice, kpar=kpar
-        )
-
-    def solve(self, lattice, kpar):
-        return np.linalg.solve(self(lattice, kpar), self._obj)
-
-
-class TMatrix(PhysicsArray):
-    """T-matrix with a spherical basis.
-
-    The T-matrix is square relating incident (regular) fields
-    :func:`treams.special.vsw_rA` (helical polarizations) or
-    :func:`treams.special.vsw_rN` and :func:`treams.special.vsw_rM` (parity
-    polarizations) to the corresponding scattered fields :func:`treams.special.vsw_A` or
-    :func:`treams.special.vsw_N` and :func:`treams.special.vsw_M`. The modes themselves
-    are defined in :attr:`basis`, the polarization type in :attr:`poltype`. Also, the
-    wave number :attr:`k0` and, if not vacuum, the material :attr:`material` are
-    specified.
-
-    Args:
-        arr (float or complex, array-like): T-matrix itself.
-        k0 (float): Wave number in vacuum.
-        basis (SphericalWaveBasis, optional): Basis definition.
-        material (Material, optional): Embedding material, defaults to vacuum.
-        poltype (str, optional): Polarization type (:ref:`params:Polarizations`).
-        lattice (Lattice, optional): Lattice definition. If specified the T-Matrix is
-            assumed to be periodically repeated in the defined lattice.
-        kpar (list, optional): Phase factor for the periodic T-Matrix.
-    """
-
-    interaction = _Interaction()
-    latticeinteraction = _LatticeInteraction()
-
-    def _check(self):
-        """Fill in default values or raise errors for missing attributes."""
-        super()._check()
-        shape = np.shape(self)
-        if len(shape) != 2 or shape[0] != shape[1]:
-            raise AnnotationError(f"invalid shape: '{shape}'")
-        if not isinstance(self.k0, (int, float, np.floating, np.integer)):
-            raise AnnotationError("invalid k0")
-        if self.poltype is None:
-            self.poltype = config.POLTYPE
-        if self.poltype not in ("parity", "helicity"):
-            raise AnnotationError("invalid poltype")
-        modetype = self.modetype
-        if modetype is None or (
-            modetype[0] in (None, "singular") and modetype[1] in (None, "regular")
-        ):
-            self.modetype = ("singular", "regular")
-        else:
-            raise AnnotationError("invalid modetype")
-        if self.basis is None:
-            self.basis = SWB.default(SWB.defaultlmax(shape[0]))
-        if self.material is None:
-            self.material = Material()
-
-    @property
-    def ks(self):
-        """Wave numbers (in medium).
-
-        The wave numbers for both polarizations.
-        """
-        return self.material.ks(self.k0)
-
-    @classmethod
-    def sphere(cls, lmax, k0, radii, materials):
-        """T-Matrix of a (multi-layered) sphere.
-
-        Construct the T-matrix of the given order and material for a sphere. The object
-        can also consist of multiple concentric spherical shells with an arbitrary
-        number of layers. The calculation is always done in helicity basis.
-
-        Args:
-            lmax (int): Positive integer for the maximum degree of the T-matrix.
-            k0 (float): Wave number in vacuum.
-            radii (float or array): Radii from inside to outside of the sphere. For a
-                simple sphere the radius can be given as a single number, for a multi-
-                layered sphere it is a list of increasing radii for all shells.
-            material (list[Material]): The material parameters from the inside to the
-                outside. The last material in the list specifies the embedding medium.
-
-        Returns:
-            TMatrix
-        """
-        materials = [Material(m) for m in materials]
-        radii = np.atleast_1d(radii)
-        if radii.size != len(materials) - 1:
-            raise ValueError("incompatible lengths of radii and materials")
-        dim = SWB.defaultdim(lmax)
-        tmat = np.zeros((dim, dim), np.complex128)
-        for l in range(1, lmax + 1):  # noqa: E741
-            miecoeffs = mie(l, k0 * radii, *zip(*materials))
-            pos = SWB.defaultdim(l - 1)
-            for i in range(2 * l + 1):
-                tmat[
-                    pos + 2 * i : pos + 2 * i + 2, pos + 2 * i : pos + 2 * i + 2
-                ] = miecoeffs[::-1, ::-1]
-        return cls(tmat, k0=k0, basis=SWB.default(lmax), material=materials[-1])
-
-    @classmethod
-    def cluster(cls, tmats, positions):
-        r"""Block-diagonal T-matrix of multiple objects.
-
-        Construct the initial block-diagonal T-matrix for a cluster of objects. The
-        T-matrices in the list are placed together into a block-diagonal matrix and the
-        complete (local) basis is defined based on the individual T-matrices and their
-        bases together with the defined positions. In mathematical terms the matrix
-
-        .. math::
-
-            \begin{pmatrix}
-                T_0 & 0 & \dots & 0 \\
-                0 & T_1 & \ddots & \vdots \\
-                \vdots & \ddots & \ddots & 0 \\
-                0 & \dots & 0 & T_{N-1} \\
-            \end{pmatrix}
-
-        is created from the list of T-matrices :math:`(T_0, \dots, T_{N-1})`. Only
-        T-matrices of the same wave number, embedding material, and polarization type
-        can be combined.
-
-        Args:
-            tmats (Sequence): List of T-matrices.
-            positions (array): The positions of all individual objects in the cluster.
-
-        Returns:
-            TMatrix
-        """
-        for tm in tmats:
-            if not tm.basis.isglobal:
-                raise ValueError("global basis required")
-        positions = np.array(positions)
-        if len(tmats) < positions.shape[0]:
-            warnings.warn("specified more positions than T-matrices")
-        elif len(tmats) > positions.shape[0]:
-            raise ValueError(
-                f"'{len(tmats)}' T-matrices "
-                f"but only '{positions.shape[0]}' positions given"
-            )
-        mat = tmats[0].material
-        k0 = tmats[0].k0
-        poltype = tmats[0].poltype
-        modes = [], [], []
-        pidx = []
-        dim = sum([tmat.shape[0] for tmat in tmats])
-        tres = np.zeros((dim, dim), complex)
-        i = 0
-        for j, tm in enumerate(tmats):
-            if tm.material != mat:
-                raise ValueError(f"incompatible materials: '{mat}' and '{tm.material}'")
-            if tm.k0 != k0:
-                raise ValueError(f"incompatible k0: '{k0}' and '{tm.k0}'")
-            if tm.poltype != poltype:
-                raise ValueError(f"incompatible modetypes: '{poltype}', '{tm.poltype}'")
-            dim = tm.shape[0]
-            for m, n in zip(modes, tm.basis.lms):
-                m.extend(list(n))
-            pidx += [j] * dim
-            tres[i : i + dim, i : i + dim] = tm
-            i += dim
-        basis = SWB(zip(pidx, *modes), positions)
-        return cls(tres, k0=k0, material=mat, basis=basis, poltype=poltype)
-
-    @property
-    def isglobal(self):
-        """Test if a T-matrix is global.
-
-        A T-matrix is considered global, when its basis refers to only a single point
-        and it is not placed periodically in a lattice.
-        """
-        return self.basis.isglobal and self.lattice is None and self.kpar is None
-
-    @property
-    def xs_ext_avg(self):
-        r"""Rotation and polarization averaged extinction cross section.
-
-        The average is calculated as
-
-        .. math::
-
-            \langle \sigma_\mathrm{ext} \rangle
-            = -2 \pi \sum_{slm} \frac{\Re(T_{slm,slm})}{k_s^2}
-
-        where :math:`k_s` is the wave number in the embedding medium for the
-        polarization :math:`s`. It is only implemented for global T-matrices.
-
-        Returns:
-            float or complex
-        """
-        if not self.isglobal or not self.material.isreal:
-            raise NotImplementedError
-        if not self.material.ischiral:
-            k = self.ks[0]
-            res = -2 * np.pi * self.trace().real / (k * k)
-        else:
-            res = 0
-            modetype = self.modetype
-            del self.modetype
-            diag = np.diag(self)
-            self.modetype = modetype
-            for pol in [0, 1]:
-                choice = self.basis.pol == pol
-                k = self.ks[pol]
-                res += -2 * np.pi * diag[choice].sum().real / (k * k)
-        if res.imag == 0:
-            return res.real
-        return res
-
-    @property
-    def xs_sca_avg(self):
-        r"""Rotation and polarization averaged scattering cross section.
-
-        The average is calculated as
-
-        .. math::
-
-            \langle \sigma_\mathrm{sca} \rangle
-            = 2 \pi \sum_{slm} \sum_{s'l'm'}
-            \frac{|T_{slm,s'l'm'}|^2}{k_s^2}
-
-        where :math:`k_s` is the wave number in the embedding medium for the
-        polarization :math:`s`. It is only implemented for global T-matrices.
-
-        Returns:
-            float or complex
-        """
-        if not self.isglobal or not self.material.isreal:
-            raise NotImplementedError
-        if not self.material.ischiral:
-            ks = self.ks[0]
-        else:
-            ks = self.ks[self.basis.pol, None]
-        re, im = self.real, self.imag
-        res = 2 * np.pi * np.sum((re * re + im * im) / (ks * ks))
-        return res.real
-
-    @property
-    def cd(self):
-        r"""Circular dichroism (CD).
-
-        The CD is calculated as
-
-        .. math::
-
-            CD
-            = \frac{\langle \sigma_\mathrm{abs} \rangle_+
-            - \langle \sigma_\mathrm{abs} \rangle_-}
-            {\langle \sigma_\mathrm{abs} \rangle_+
-            + \langle \sigma_\mathrm{abs} \rangle_-}
-
-        where :math:`\langle \sigma \rangle_s` is the rotationally averaged absorption
-        cross section under the illumination with the polarization :math:`s`.
-        It is only implemented for global T-matrices.
-
-        Returns:
-            float or complex
-        """
-        if not (self.isglobal and self.poltype == "helicity" and self.material.isreal):
-            raise NotImplementedError
-        sel = np.array(self.basis.pol, bool)
-        re, im = self.real, self.imag
-        plus = -np.sum(re[sel[:, None] & sel]) / (self.ks[1] * self.ks[1])
-        re_part = re[:, sel] / self.ks[self.basis.pol, None]
-        im_part = im[:, sel] / self.ks[self.basis.pol, None]
-        plus -= np.sum(re_part * re_part + im_part * im_part)
-        sel = ~sel
-        minus = -np.sum(re[sel[:, None] & sel]) / (self.ks[0] * self.ks[0])
-        re_part = re[:, sel] / self.ks[self.basis.pol, None]
-        im_part = im[:, sel] / self.ks[self.basis.pol, None]
-        minus -= np.sum(re_part * re_part + im_part * im_part)
-        return np.real((plus - minus) / (plus + minus))
-
-    @property
-    def chi(self):
-        """Electromagnetic chirality.
-
-        Only implemented for global T-matrices.
-
-        Returns:
-            float
-        """
-        if not (self.isglobal and self.poltype == "helicity"):
-            raise NotImplementedError
-        sel = self.basis.pol == 0, self.basis.pol == 1
-        spp = np.linalg.svd(self[np.ix_(sel[1], sel[1])], compute_uv=False)
-        spm = np.linalg.svd(self[np.ix_(sel[1], sel[0])], compute_uv=False)
-        smp = np.linalg.svd(self[np.ix_(sel[0], sel[1])], compute_uv=False)
-        smm = np.linalg.svd(self[np.ix_(sel[0], sel[0])], compute_uv=False)
-        plus = np.concatenate((np.asarray(spp), np.asarray(spm)))
-        minus = np.concatenate((np.asarray(smm), np.asarray(smp)))
-        return np.linalg.norm(plus - minus) / np.sqrt(np.sum(np.power(np.abs(self), 2)))
-
-    @property
-    def db(self):
-        """Duality breaking.
-
-        Only implemented for global T-matrices.
-
-        Returns:
-            float
-        """
-        if not (self.isglobal and self.poltype == "helicity"):
-            raise NotImplementedError
-        sel = self.basis.pol == 0, self.basis.pol == 1
-        tpm = np.asarray(self[np.ix_(sel[1], sel[0])])
-        tmp = np.asarray(self[np.ix_(sel[0], sel[1])])
-        return np.sum(
-            tpm.real * tpm.real
-            + tpm.imag * tpm.imag
-            + tmp.real * tmp.real
-            + tmp.imag * tmp.imag
-        ) / (np.sum(np.power(np.abs(self), 2)))
-
-    def xs(self, illu, flux=0.5):
-        r"""Scattering and extinction cross section.
-
-        Possible for all T-matrices (global and local) in non-absorbing embedding. The
-        values are calculated by
-
-        .. math::
-
-            \sigma_\mathrm{sca}
-            = \frac{1}{2 I}
-            a_{slm}^\ast T_{s'l'm',slm}^\ast k_{s'}^{-2} C_{s'l'm',s''l''m''}^{(1)}
-            T_{s''l''m'',s'''l'''m'''} a_{s'''l'''m'''} \\
-            \sigma_\mathrm{ext}
-            = \frac{1}{2 I}
-            a_{slm}^\ast k_s^{-2} T_{slm,s'l'm'} a_{s'l'm'}
-
-        where :math:`a_{slm}` are the expansion coefficients of the illumination,
-        :math:`T` is the T-matrix, :math:`C^{(1)}` is the (regular) translation
-        matrix and :math:`k_s` are the wave numbers in the medium. All repeated indices
-        are summed over. The incoming flux is :math:`I`.
-
-        Args:
-            illu (complex, array): Illumination coefficients
-            flux (optional): Ingoing flux corresponding to the illumination. Used for
-                the result's normalization. The flux is given in units of
-                :math:`\frac{\text{V}^2}{{l^2}} \frac{1}{Z_0 Z}` where :math:`l` is the
-                unit of length used in the wave number (and positions). A plane wave
-                has the flux `0.5` in this normalization, which is used as default.
-
-        Returns:
-            tuple[float]
-        """
-        if not self.material.isreal:
-            raise NotImplementedError
-        illu = PhysicsArray(illu)
-        illu_basis = illu.basis
-        illu_basis = illu_basis[-2] if isinstance(illu_basis, tuple) else illu_basis
-        if not isinstance(illu_basis, SWB):
-            illu = illu.expand(self.basis)
-        p = self @ illu
-        p_invksq = p * np.power(self.ks[self.basis.pol], -2)
-        del illu.modetype
-        return (
-            0.5 * np.real(p.conjugate().T @ p_invksq.expand(p.basis)) / flux,
-            -0.5 * np.real(illu.conjugate().T @ p_invksq) / flux,
-        )
-
-    def valid_points(self, grid, radii):
-        """Points on the grid where the expansion is valid.
-
-        The expansion of the electromagnetic field is valid outside of the
-        circumscribing spheres of each object. From a given set of coordinates mark
-        those that are outside of the given radii.
-
-        Args:
-            grid (array-like): Points to assess. The last dimension needs length three
-                and corresponds to the Cartesian coordinates.
-            radii (Sequence[float]): Radii of the circumscribing spheres. Each radius
-                corresponds to a position of the basis.
-
-        Returns:
-            array
-        """
-        grid = np.asarray(grid)
-        if grid.shape[-1] != 3:
-            raise ValueError("invalid grid")
-        if len(radii) != len(self.basis.positions):
-            raise ValueError("invalid length of 'radii'")
-        res = np.ones(grid.shape[:-1], bool)
-        for r, p in zip(radii, self.basis.positions):
-            res &= np.sum(np.power(grid - p, 2), axis=-1) > r * r
-        return res
-
-    def __getitem__(self, key):
-        if isinstance(key, SWB):
-            key = np.array([self.basis.index(i) for i in key])
-            key = (key[:, None], key)
-        return super().__getitem__(key)
-
-
-class TMatrixC(PhysicsArray):
-    """T-matrix with a cylindrical basis.
-
-    The T-matrix is square relating incident (regular) fields
-    :func:`treams.special.vcw_rA` (helical polarizations) or
-    :func:`treams.special.vcw_rN` and :func:`treams.special.vcw_rM` (parity
-    polarizations) to the corresponding scattered fields :func:`treams.special.vcw_A` or
-    :func:`treams.special.vcw_N` and :func:`treams.special.vcw_M`. The modes themselves
-    are defined in :attr:`basis`, the polarization type in :attr:`poltype`. Also, the
-    wave number :attr:`k0` and, if not vacuum, the material :attr:`material` are
-    specified.
-
-    Args:
-        arr (float or complex, array-like): T-matrix itself.
-        k0 (float): Wave number in vacuum.
-        basis (SphericalWaveBasis, optional): Basis definition.
-        material (Material, optional): Embedding material, defaults to vacuum.
-        poltype (str, optional): Polarization type (:ref:`params:Polarizations`).
-        lattice (Lattice, optional): Lattice definition. If specified the T-Matrix is
-            assumed to be periodically repeated in the defined lattice.
-        kpar (list, optional): Phase factor for the periodic T-Matrix.
-    """
-
-    interaction = _Interaction()
-    latticeinteraction = _LatticeInteraction()
-
-    def _check(self):
-        """Fill in default values or raise errors for missing attributes."""
-        super()._check()
-        shape = np.shape(self)
-        if len(shape) != 2 or shape[0] != shape[1]:
-            raise AnnotationError(f"invalid shape: '{shape}'")
-        if not isinstance(self.k0, (int, float, np.floating, np.integer)):
-            raise AnnotationError("invalid k0")
-        if self.poltype is None:
-            self.poltype = config.POLTYPE
-        if self.poltype not in ("parity", "helicity"):
-            raise AnnotationError("invalid poltype")
-        modetype = self.modetype
-        if modetype is None or (
-            modetype[0] in (None, "singular") and modetype[1] in (None, "regular")
-        ):
-            self.modetype = ("singular", "regular")
-        else:
-            raise AnnotationError("invalid modetype")
-        if self.basis is None:
-            self.basis = CWB.default([0], CWB.defaultmmax(shape[0]))
-        if self.material is None:
-            self.material = Material()
-
-    @property
-    def ks(self):
-        """Wave numbers (in medium).
-
-        The wave numbers for both polarizations.
-        """
-        return self.material.ks(self.k0)
-
-    @property
-    def krhos(self):
-        r"""Radial part of the wave.
-
-        Calculate :math:`\sqrt{k^2 - k_z^2}`, where :math:`k` is the wave number in the
-        medium for each illumination
-
-        Returns:
-            Sequence[complex]
-        """
-        return self.material.krhos(self.k0, self.basis.kz, self.basis.pol)
-
-    @classmethod
-    def cylinder(cls, kzs, mmax, k0, radii, materials):
-        """T-Matrix of a (multi-layered) cylinder.
-
-        Construct the T-matrix of the given order and material for an infinitely
-        extended cylinder. The object can also consist of multiple concentric
-        cylindrical shells with an arbitrary number of layers. The calculation is always
-        done in helicity basis.
-
-        Args:
-            kzs (float, array_like): Z component of the cylindrical wave.
-            mmax (int): Positive integer for the maximum order of the T-matrix.
-            k0 (float): Wave number in vacuum.
-            radii (float or array): Radii from inside to outside of the cylinder. For a
-                simple cylinder the radius can be given as a single number, for a multi-
-                layered cylinder it is a list of increasing radii for all shells.
-            material (list[Material]): The material parameters from the inside to the
-                outside. The last material in the list specifies the embedding medium.
-
-        Returns:
-            T-Matrix object
-        """
-        materials = [Material(m) for m in materials]
-        kzs = np.atleast_1d(kzs)
-        radii = np.atleast_1d(radii)
-        if radii.size != len(materials) - 1:
-            raise ValueError("incompatible lengths of radii and materials")
-        dim = CWB.defaultdim(len(kzs), mmax)
-        tmat = np.zeros((dim, dim), np.complex128)
-        idx = 0
-        for kz in kzs:
-            for m in range(-mmax, mmax + 1):
-                miecoeffs = mie_cyl(kz, m, k0, radii, *zip(*materials))
-                tmat[idx : idx + 2, idx : idx + 2] = miecoeffs[::-1, ::-1]
-                idx += 2
-        return cls(tmat, k0=k0, basis=CWB.default(kzs, mmax), material=materials[-1])
-
-    @classmethod
-    def cluster(cls, tmats, positions):
-        r"""Block-diagonal T-matrix of multiple objects.
-
-        Construct the initial block-diagonal T-matrix for a cluster of objects. The
-        T-matrices in the list are placed together into a block-diagonal matrix and the
-        complete (local) basis is defined based on the individual T-matrices and their
-        bases together with the defined positions. In mathematical terms the matrix
-
-        .. math::
-
-            \begin{pmatrix}
-                T_0 & 0 & \dots & 0 \\
-                0 & T_1 & \ddots & \vdots \\
-                \vdots & \ddots & \ddots & 0 \\
-                0 & \dots & 0 & T_{N-1} \\
-            \end{pmatrix}
-
-        is created from the list of T-matrices :math:`(T_0, \dots, T_{N-1})`. Only
-        T-matrices of the same wave number, embedding material, and polarization type
-        can be combined.
-
-        Args:
-            tmats (Sequence): List of T-matrices.
-            positions (array): The positions of all individual objects in the cluster.
-
-        Returns:
-            TMatrix
-        """
-        for tm in tmats:
-            if not tm.basis.isglobal:
-                raise ValueError("global basis required")
-        positions = np.array(positions)
-        if len(tmats) < positions.shape[0]:
-            warnings.warn("specified more positions than T-matrices")
-        elif len(tmats) > positions.shape[0]:
-            raise ValueError(
-                f"'{len(tmats)}' T-matrices "
-                f"but only '{positions.shape[0]}' positions given"
-            )
-        mat = tmats[0].material
-        k0 = tmats[0].k0
-        poltype = tmats[0].poltype
-        modes = [], [], []
-        pidx = []
-        dim = sum([tmat.shape[0] for tmat in tmats])
-        tres = np.zeros((dim, dim), complex)
-        i = 0
-        for j, tm in enumerate(tmats):
-            if tm.material != mat:
-                raise ValueError(f"incompatible materials: '{mat}' and '{tm.material}'")
-            if tm.k0 != k0:
-                raise ValueError(f"incompatible k0: '{k0}' and '{tm.k0}'")
-            if tm.poltype != poltype:
-                raise ValueError(f"incompatible modetypes: '{poltype}', '{tm.poltype}'")
-            dim = tm.shape[0]
-            for m, n in zip(modes, tm.basis.zms):
-                m.extend(list(n))
-            pidx += [j] * dim
-            tres[i : i + dim, i : i + dim] = tm
-            i += dim
-        basis = CWB(zip(pidx, *modes), positions)
-        return cls(tres, k0=k0, material=mat, basis=basis, poltype=poltype)
-
-    @classmethod
-    def from_array(cls, tm, basis, *, eta=0):
-        """1d array of spherical T-matrices."""
-        return cls(
-            (tm @ op.Expand(basis).inv).expandlattice(basis=basis, eta=eta),
-            lattice=tm.lattice,
-            kpar=tm.kpar,
-        )
-
-    @property
-    def xw_ext_avg(self):
-        r"""Rotation and polarization averaged extinction cross width.
-
-        The average is calculated as
-
-        .. math::
-
-            \langle \lambda_\mathrm{ext} \rangle
-            = -\frac{2 \pi}{n_{k_z}} \sum_{sk_zm} \frac{\Re(T_{sk_zm,sk_zm})}{k_s}
-
-        where :math:`k_s` is the wave number in the embedding medium for the
-        polarization :math:`s` and :math:`n_{k_z}` is the number of wave components
-        :math:`k_z` included in the T-matrix. The average is taken over all given
-        z-components of the wave vector and rotations around the z-axis. It is only
-        implemented for global T-matrices.
-
-        Returns:
-            float or complex
-        """
-        if not self.isglobal or not self.material.isreal:
-            raise NotImplementedError
-        nk = np.unique(self.basis.kz).size
-        if not self.material.ischiral:
-            res = -2 * np.real(np.trace(self)) / (self.ks[0] * nk)
-        else:
-            res = 0
-            modetype = self.modetype
-            del self.modetype
-            diag = np.diag(self)
-            self.modetype = modetype
-            for pol in [0, 1]:
-                choice = self.basis.pol == pol
-                res += -2 * np.real(diag[choice].sum()) / (self.ks[pol] * nk)
-        if res.imag == 0:
-            return res.real
-        return res
-
-    @property
-    def xw_sca_avg(self):
-        r"""Rotation and polarization averaged scattering cross width.
-
-        The average is calculated as
-
-        .. math::
-
-            \langle \lambda_\mathrm{sca} \rangle
-            = \frac{2 \pi}{n_{k_z}} \sum_{sk_zm} \sum_{s'{k_z}'m'}
-            \frac{|T_{sk_zm,s'{k_z}'m'}|^2}{k_s}
-
-        where :math:`k_s` is the wave number in the embedding medium for the
-        polarization :math:`s`. and :math:`n_{k_z}` is the number of wave components
-        :math:`k_z` included in the T-matrix. The average is taken over all given
-        z-components of the wave vector and rotations around the z-axis. It is only
-        implemented for global T-matrices.
-
-        Returns:
-            float or complex
-        """
-        if not self.isglobal or not self.material.isreal:
-            raise NotImplementedError
-        if not self.material.ischiral:
-            ks = self.ks[0]
-        else:
-            ks = self.ks[self.basis.pol, None]
-        re, im = self.real, self.imag
-        nk = np.unique(self.basis.kz).size
-        res = 2 * np.sum((re * re + im * im) / (ks * nk))
-        return res.real
-
-    @property
-    def isglobal(self):
-        """Test if a T-matrix is global.
-
-        A T-matrix is considered global, when its basis refers to only a single point
-        and it is not placed periodically in a lattice.
-        """
-        return self.basis.isglobal and self.lattice is None and self.kpar is None
-
-    def xw(self, illu, flux=0.5):
-        r"""Scattering and extinction cross width.
-
-        Possible for all T-matrices (global and local) in non-absorbing embedding. The
-        values are calculated by
-
-        .. math::
-
-            \lambda_\mathrm{sca}
-            = \frac{1}{2 I}
-            a_{sk_zm}^\ast T_{s'{k_z}'m',sk_zm}^\ast k_{s'}^{-2}
-            C_{s'l'm',s''l''m''}^{(1)}
-            T_{s''l''m'',s'''l'''m'''} a_{s'''l'''m'''} \\
-            \sigma_\mathrm{ext}
-            = \frac{1}{2 I}
-            a_{slm}^\ast k_s^{-2} T_{slm,s'l'm'} a_{s'l'm'}
-
-        where :math:`a_{slm}` are the expansion coefficients of the illumination,
-        :math:`T` is the T-matrix, :math:`C^{(1)}` is the (regular) translation
-        matrix and :math:`k_s` are the wave numbers in the medium. All repeated indices
-        are summed over. The incoming flux is :math:`I`.
-
-        Args:
-            illu (complex, array): Illumination coefficients
-            flux (optional): Ingoing flux corresponding to the illumination. Used for
-                the result's normalization. The flux is given in units of
-                :math:`\frac{\text{V}^2}{{l^2}} \frac{1}{Z_0 Z}` where :math:`l` is the
-                unit of length used in the wave number (and positions). A plane wave
-                has the flux `0.5` in this normalization, which is used as default.
-
-        Returns:
-            tuple[float]
-        """
-        if not self.material.isreal:
-            raise NotImplementedError
-        illu = PhysicsArray(illu)
-        illu_basis = illu.basis
-        illu_basis = illu_basis[-2] if isinstance(illu_basis, tuple) else illu_basis
-        if not isinstance(illu_basis, CWB):
-            illu = illu.expand(self.basis) @ illu
-        p = self @ illu
-        m = self.expand() / self.ks[self.basis.pol]
-        del illu.modetype
-        return (
-            2 * np.real(p.conjugate().T @ (m @ p)) / flux,
-            -2 * np.real(illu.conjugate().T @ (p / self.ks[self.basis.pol])) / flux,
-        )
-
-    def globalmat(self, basis=None):
-        """Global T-matrix.
-
-        Calculate the global T-matrix starting from a local one. This changes the
-        T-matrix.
-
-        Args:
-            origin (array, optional): The origin of the new T-matrix
-            modes (array, optional): The modes that are considered for the global
-                T-matrix
-
-        """
-        if not basis.isglobal:
-            raise ValueError("global basis required")
-        if basis is None:
-            basis = CWB.default(np.unique(self.kz), max(self.basis.l))
-        return TMatrix(self.expand(basis) @ self @ self.expand.inv(basis))
-
-    def valid_points(self, grid, radii):
-        grid = np.asarray(grid)
-        if grid.shape[-1] not in (2, 3):
-            raise ValueError("invalid grid")
-        if len(radii) != len(self.basis.positions):
-            raise ValueError("invalid length of 'radii'")
-        res = np.ones(grid.shape[:-1], bool)
-        for r, p in zip(radii, self.basis.positions):
-            res &= np.sum(np.power(grid[..., :2] - p[:2], 2), axis=-1) > r * r
-        return res
-
-    def __getitem__(self, key):
-        if isinstance(key, CWB):
-            key = np.array([self.basis.index(i) for i in key])
-            key = (key[:, None], key)
-        return super().__getitem__(key)
-
-
-def _plane_wave_partial(
-    kpar,
-    pol,
-    *,
-    k0=None,
-    basis=None,
-    material=None,
-    modetype=None,
-    poltype=None,
-):
-    if basis is None:
-        basis = PWBC.default([kpar])
-    if pol == 0 or pol == -1:
-        pol = [1, 0]
-    elif pol == 1:
-        pol = [0, 1]
-    elif len(pol) == 3:
-        modetype = "up" if modetype is None else modetype
-        if modetype not in ("up", "down"):
-            raise ValueError(f"invalid 'modetype': {modetype}")
-        kzs = Material(material).kzs(k0, *kpar, [0, 1])
-        poltype = config.POLTYPE if poltype is None else poltype
-        if poltype == "parity":
-            pol = [
-                -sc.vpw_M(*kpar, kzs[0], 0, 0, 0) @ pol,
-                sc.vpw_N(*kpar, kzs[1], 0, 0, 0) @ pol,
-            ]
-        elif poltype == "helicity":
-            pol = sc.vpw_A(*kpar, kzs[::-1], 0, 0, 0, [1, 0]) @ pol
-        else:
-            raise ValueError(f"invalid 'poltype': {poltype}")
-    res = [pol[x[2]] * (np.abs(np.array(kpar) - x[:2]) < 1e-14).all() for x in basis]
-    return PhysicsArray(
-        res,
-        basis=basis,
-        k0=k0,
-        material=material,
-        modetype=modetype,
-        poltype=poltype,
-    )
-
-
-def _plane_wave(
-    kvec,
-    pol,
-    *,
-    k0=None,
-    basis=None,
-    material=None,
-    modetype=None,
-    poltype=None,
-):
-    if basis is None:
-        basis = PWBUV.default([kvec])
-    norm = np.sqrt(np.sum(np.power(kvec, 2)))
-    qvec = kvec / norm
-    if pol == 0 or pol == -1:
-        pol = [1, 0]
-    elif pol == 1:
-        pol = [0, 1]
-    elif len(pol) == 3:
-        if None not in (k0, material):
-            kvec = Material(material).ks(k0) * qvec[:, None]
-        else:
-            kvec = qvec
-        poltype = config.POLTYPE if poltype is None else poltype
-        if poltype == "parity":
-            pol = [
-                -sc.vpw_M(*kvec[:, 0], 0, 0, 0) @ pol,
-                sc.vpw_N(*kvec[:, 1], 0, 0, 0) @ pol,
-            ]
-        elif poltype == "helicity":
-            pol = sc.vpw_A(*kvec, 0, 0, 0, [1, 0]) @ pol
-        else:
-            raise ValueError(f"invalid 'poltype': {poltype}")
-    res = [pol[x[3]] * (np.abs(qvec - x[:3]) < 1e-14).all() for x in basis]
-    return PhysicsArray(
-        res,
-        basis=basis,
-        k0=k0,
-        material=material,
-        modetype=modetype,
-        poltype=poltype,
-    )
-
-
-def plane_wave(
-    kvec,
-    pol,
-    *,
-    k0=None,
-    basis=None,
-    material=None,
-    modetype=None,
-    poltype=None,
-):
-    """Array describing a plane wave.
-
-    Args:
-        kvec (Sequence): Wave vector.
-        pol (int or Sequence): Polarization index (see
-            :ref:`params:Polarizations`) to have a unit amplitude wave of the
-            corresponding wave, two values to specify the amplitude for each
-            polarization, or three values in a sequence to specify the Cartesian
-            electric field components. In the latter case, if the electric field has
-            longitudinal components they are neglected.
-        basis (PlaneWaveBasis, optional): Basis definition.
-        k0 (float, optional): Wave number in vacuum.
-        material (Material, optional): Material definition.
-        modetype (str, optional): Mode type (see :ref:`params:Mode types`).
-        poltype (str, optional): Polarization type (see
-            :ref:`polarization:Polarizations`).
-    """
-    if len(kvec) == 2:
-        return _plane_wave_partial(
-            kvec,
-            pol,
-            k0=k0,
-            basis=basis,
-            material=material,
-            modetype=modetype,
-            poltype=poltype,
-        )
-    elif len(kvec) == 3:
-        return _plane_wave(
-            kvec,
-            pol,
-            k0=k0,
-            basis=basis,
-            material=material,
-            modetype=modetype,
-            poltype=poltype,
-        )
-    raise ValueError(f"invalid length of 'kvec': {len(kvec)}")
-
-
-def plane_wave_angle(theta, phi, pol, **kwargs):
-    qvec = [np.sin(theta) * np.cos(phi), np.sin(theta) * np.sin(phi), np.cos(theta)]
-    return plane_wave(qvec, pol, **kwargs)
-
-
-def spherical_wave(
-    l,  # noqa: E741
-    m,
-    pol,
-    *,
-    k0=None,
-    basis=None,
-    material=None,
-    modetype=None,
-    poltype=None,
-):
-    if basis is None:
-        basis = SWB.default(l)
-    if not basis.isglobal:
-        raise ValueError("basis must be global")
-    res = [0] * len(basis)
-    res[basis.index((0, l, m, pol))] = 1
-    return PhysicsArray(
-        res, basis=basis, k0=k0, material=material, modetype=modetype, poltype=poltype
-    )
-
-
-def cylindrical_wave(
-    kz,
-    m,
-    pol,
-    *,
-    k0=None,
-    basis=None,
-    material=None,
-    modetype=None,
-    poltype=None,
-):
-    if basis is None:
-        basis = CWB.default([kz], abs(m))
-    if not basis.isglobal:
-        raise ValueError("basis must be global")
-    res = [0] * len(basis)
-    res[basis.index((0, kz, m, pol))] = 1
-    return PhysicsArray(
-        res, basis=basis, k0=k0, material=material, modetype=modetype, poltype=poltype
-    )

treams/ebcm.py

-import itertools
-
-import numpy as np
-import scipy.integrate
-
-import treams.special as sc
-
-
-def _j_real(l0, l1, m, r, dr, pol0, pol1, ks, k, zs, z):
-    def fun(theta):
-        return np.real(
-            np.sin(theta)
-            * ((2 * pol0 - 1) * z + (2 * pol1 - 1) * zs)
-            * np.dot(
-                [r(theta), -dr(theta), 0],
-                np.cross(
-                    sc.vsw_rA(l0, -m, ks[pol0] * r(theta), theta, 0, pol0),
-                    sc.vsw_A(l1, m, k[pol1] * r(theta), theta, 0, pol1),
-                ),
-            )
-        )
-
-    return fun
-
-
-def _j_imag(l0, l1, m, r, dr, pol0, pol1, ks, k, zs, z):
-    def fun(theta):
-        return np.imag(
-            np.sin(theta)
-            * ((2 * pol0 - 1) * z + (2 * pol1 - 1) * zs)
-            * np.dot(
-                [r(theta), -dr(theta), 0],
-                np.cross(
-                    sc.vsw_rA(l0, -m, ks[pol0] * r(theta), theta, 0, pol0),
-                    sc.vsw_A(l1, m, k[pol1] * r(theta), theta, 0, pol1),
-                ),
-            )
-        )
-
-    return fun
-
-
-def _rj_real(l0, l1, m, r, dr, pol0, pol1, ks, k, zs, z):
-    def fun(theta):
-        return np.real(
-            np.sin(theta)
-            * ((2 * pol0 - 1) * z + (2 * pol1 - 1) * zs)
-            * np.dot(
-                [r(theta), -dr(theta), 0],
-                np.cross(
-                    sc.vsw_rA(l0, -m, ks[pol0] * r(theta), theta, 0, pol0),
-                    sc.vsw_rA(l1, m, k[pol1] * r(theta), theta, 0, pol1),
-                ),
-            )
-        )
-
-    return fun
-
-
-def _rj_imag(l0, l1, m, r, dr, pol0, pol1, ks, k, zs, z):
-    def fun(theta):
-        return np.imag(
-            np.sin(theta)
-            * ((2 * pol0 - 1) * z + (2 * pol1 - 1) * zs)
-            * np.dot(
-                [r(theta), -dr(theta), 0],
-                np.cross(
-                    sc.vsw_rA(l0, -m, ks[pol0] * r(theta), theta, 0, pol0),
-                    sc.vsw_rA(l1, m, k[pol1] * r(theta), theta, 0, pol1),
-                ),
-            )
-        )
-
-    return fun
-
-
-def qmat(r, dr, ks, zs, out, in_=None, singular=True):
-    l_out, m_out, pol_out = out
-    in_ = out if in_ is None else in_
-    l_in, m_in, pol_in = in_
-    if singular:
-        fr = _j_real
-        fi = _j_imag
-    else:
-        fr = _rj_real
-        fi = _rj_imag
-    size_out = len(l_out)
-    size_in = len(l_in)
-    res = np.zeros((size_out, size_in), complex)
-    for i, j in itertools.product(range(size_out), range(size_in)):
-        if m_out[i] != m_in[j]:
-            continue
-        res[i, j] = (
-            scipy.integrate.quad(
-                fr(l_out[i], l_in[j], m_out[i], r, dr, pol_out[i], pol_in[j], *ks, *zs),
-                0,
-                np.pi,
-            )[0]
-            + 1j
-            * scipy.integrate.quad(
-                fi(l_out[i], l_in[j], m_out[i], r, dr, pol_out[i], pol_in[j], *ks, *zs),
-                0,
-                np.pi,
-            )[0]
-        )
-    return res

treams/io.py

-"""Loading and storing data.
-
-Most functions rely on at least one of the external packages `h5py` or `gmsh`.
-"""
-
-import sys
-import uuid as _uuid
-from importlib.metadata import version
-
-import numpy as np
-
-import treams
-
-try:
-    import h5py
-except ImportError:
-    h5py = None
-
-LENGTHS = {
-    "ym": 1e-24,
-    "zm": 1e-21,
-    "am": 1e-18,
-    "fm": 1e-15,
-    "pm": 1e-12,
-    "nm": 1e-9,
-    "um": 1e-6,
-    "µm": 1e-6,
-    "mm": 1e-3,
-    "cm": 1e-2,
-    "dm": 1e-1,
-    "m": 1,
-    "dam": 1e1,
-    "hm": 1e2,
-    "km": 1e3,
-    "Mm": 1e6,
-    "Gm": 1e9,
-    "Tm": 1e12,
-    "Tm": 1e12,
-    "Pm": 1e15,
-    "Em": 1e18,
-    "Zm": 1e21,
-    "Ym": 1e24,
-}
-
-INVLENGTHS = {
-    r"ym^{-1}": 1e24,
-    r"zm^{-1}": 1e21,
-    r"am^{-1}": 1e18,
-    r"fm^{-1}": 1e15,
-    r"pm^{-1}": 1e12,
-    r"nm^{-1}": 1e9,
-    r"um^{-1}": 1e6,
-    r"µm^{-1}": 1e6,
-    r"mm^{-1}": 1e3,
-    r"cm^{-1}": 1e2,
-    r"dm^{-1}": 1e1,
-    r"m^{-1}": 1,
-    r"dam^{-1}": 1e-1,
-    r"hm^{-1}": 1e-2,
-    r"km^{-1}": 1e-3,
-    r"Mm^{-1}": 1e-6,
-    r"Gm^{-1}": 1e-9,
-    r"Tm^{-1}": 1e-12,
-    r"Pm^{-1}": 1e-15,
-    r"Em^{-1}": 1e-18,
-    r"Zm^{-1}": 1e-21,
-    r"Ym^{-1}": 1e-24,
-}
-
-FREQUENCIES = {
-    "yHz": 1e-24,
-    "zHz": 1e-21,
-    "aHz": 1e-18,
-    "fHz": 1e-15,
-    "pHz": 1e-12,
-    "nHz": 1e-9,
-    "uHz": 1e-6,
-    "µHz": 1e-6,
-    "mHz": 1e-3,
-    "cHz": 1e-2,
-    "dHz": 1e-1,
-    "s": 1,
-    "daHz": 1e1,
-    "hHz": 1e2,
-    "kHz": 1e3,
-    "MHz": 1e6,
-    "GHz": 1e9,
-    "THz": 1e12,
-    "PHz": 1e15,
-    "EHz": 1e18,
-    "ZHz": 1e21,
-    "YHz": 1e24,
-    r"ys^{-1}": 1e24,
-    r"zs^{-1}": 1e21,
-    r"as^{-1}": 1e18,
-    r"fs^{-1}": 1e15,
-    r"ps^{-1}": 1e12,
-    r"ns^{-1}": 1e9,
-    r"us^{-1}": 1e6,
-    r"µs^{-1}": 1e6,
-    r"ms^{-1}": 1e3,
-    r"cs^{-1}": 1e2,
-    r"ds^{-1}": 1e1,
-    r"s^{-1}": 1,
-    r"das^{-1}": 1e-1,
-    r"hs^{-1}": 1e-2,
-    r"ks^{-1}": 1e-3,
-    r"Ms^{-1}": 1e-6,
-    r"Gs^{-1}": 1e-9,
-    r"Ts^{-1}": 1e-12,
-    r"Ps^{-1}": 1e-15,
-    r"Es^{-1}": 1e-18,
-    r"Zs^{-1}": 1e-21,
-    r"Ys^{-1}": 1e-24,
-}
-
-
-def mesh_spheres(radii, positions, model, meshsize=None, meshsize_boundary=None):
-    """Generate a mesh of multiple spheres.
-
-    This function facilitates generating a mesh for a cluster spheres using gmsh. It
-    requires the package `gmsh` to be installed.
-
-    Examples:
-        >>> import gmsh
-        >>> import treams.io
-        >>> gmsh.initialize()
-        >>> gmsh.model.add("spheres")
-        >>> treams.io.mesh_spheres([1, 2], [[0, 0, 2], [0, 0, -2]], gmsh.model)
-        <class 'gmsh.model'>
-        >>> gmsh.write("spheres.msh")
-        >>> gmsh.finalize()
-
-    Args:
-        radii (float, array_like): Radii of the spheres.
-        positions (float, (N, 3)-array): Positions of the spheres.
-        model (gmsh.model): Gmsh model to modify.
-        meshsize (float, optional): Mesh size, if None a fifth of the largest radius is
-            used.
-        meshsize (float, optional): Mesh size of the surfaces, if left empty it is set
-            equal to the general mesh size.
-
-    Returns:
-        gmsh.model
-    """
-    if meshsize is None:
-        meshsize = np.max(radii) * 0.2
-    if meshsize_boundary is None:
-        meshsize_boundary = meshsize
-
-    spheres = []
-    for i, (radius, position) in enumerate(zip(radii, positions)):
-        tag = i + 1
-        model.occ.addSphere(*position, radius, tag)
-        spheres.append((3, tag))
-
-    model.occ.synchronize()
-    for _, tag in spheres:
-        model.addPhysicalGroup(3, [tag], tag)
-        # Add surfaces for other mesh formats like stl, ...
-        model.addPhysicalGroup(2, [tag], tag)
-
-    model.mesh.setSize(model.getEntities(0), meshsize)
-    model.mesh.setSize(
-        model.getBoundary(spheres, False, False, True), meshsize_boundary
-    )
-    return model
-
-
-def _translate_polarizations(pols, poltype=None):
-    """Translate the polarization index into words.
-
-    The indices 0 and 1 are translated to "negative" and "positive", respectively, when
-    helicity modes are chosen. For parity modes they are translated to "magnetic" and
-    "electric".
-
-    Args:
-        pols (int, array_like): Array of indices 0 and 1.
-        poltype (str, optional): Polarization type (:ref:`polarizations.Polarizations`).
-
-    Returns:
-        list[str]
-    """
-    poltype = treams.config.POLTYPE if poltype is None else poltype
-    if poltype == "helicity":
-        names = ["negative", "positive"]
-    elif poltype == "parity":
-        names = ["magnetic", "electric"]
-    else:
-        raise ValueError("unrecognized poltype")
-    return [names[i] for i in pols]
-
-
-def _translate_polarizations_inv(pols):
-    """Translate the polarization into indices.
-
-    This function is the inverse of :func:`treams.io._translate_polarizations`. The
-    words "negative" and "minus" are translated to 0 and the words "positive" and "plus"
-    are translated to 1, if helicity modes are chosen. For parity modes, modes
-    "magnetic" or "te" are translated to 0 and the modes "electric" or "tm" to 1.
-
-    Args:
-        pols (string, array_like): Array of strings
-
-    Returns:
-        tuple[list[int], str]
-    """
-    helicity = {"plus": 1, "positive": 1, "minus": 0, "negative": 0}
-    parity = {"te": 0, "magnetic": 0, "tm": 1, "electric": 1, "M": 0, "N": 1}
-    if pols[0].decode() in helicity:
-        dct = helicity
-        poltype = "helicity"
-    elif pols[0].decode() in parity:
-        dct = parity
-        poltype = "parity"
-    else:
-        raise ValueError(f"unrecognized polarization '{pols[0].decode()}'")
-    return [dct[i.decode()] for i in pols], poltype
-
-
-def _remove_leading_ones(shape):
-    while len(shape) > 0 and shape[0] == 1:
-        shape = shape[1:]
-    return shape
-
-
-def _collapse_dims(arr):
-    for i in range(arr.ndim):
-        test = arr[(slice(None),) * i + (slice(1),)]
-        if np.all(arr == test):
-            arr = test
-    return arr.reshape(_remove_leading_ones(arr.shape))
-
-
-def save_hdf5(
-    h5file,
-    tms,
-    name="",
-    description="",
-    keywords="",
-    embedding_group=None,
-    embedding_name="",
-    embedding_description="",
-    embedding_keywords="",
-    uuid=None,
-    uuid_version=4,
-    lunit="nm",
-):
-    """Save a set of T-matrices in a HDF5 file.
-
-    With an open and writeable datafile, this function stores the main parts of as
-    T-matrix in the file. It is left open for the user to add additional metadata.
-
-    Args:
-        h5file (h5py.Group): A HDF5 file opened with h5py.
-        tms (TMatrix, array_like): Array of T-matrix instances.
-        name (str): Name to add to the file as attribute.
-        description (str): Description to add to file as attribute.
-        keywords (str): Keywords to add to file as attribute.
-        embedding_group (h5py.Group, optional): Group object for the embedding material,
-            defaults to "/materials/embedding/".
-        embedding_name (string, optional): Name of the embedding material.
-        embedding_description (string, optional): Description of the embedding material.
-        embedding_keywords (string, optional): Keywords for the embedding material.
-        uuid (bytes, optional): UUID of the file, a new one is created if omitted.
-        uuid_version (int, optional): UUID version.
-        lunit (string, optional): Length unit used for the positions and (as
-            inverse) for the wave number.
-    """
-    tms_arr = np.array(tms)
-    if tms_arr.dtype == object:
-        raise ValueError("can only save T-matrices of the same size")
-    tms_obj = np.empty(tms_arr.shape[:-2], object)
-    tms_obj[:] = tms
-    tm = tms_obj.flat[0]
-    basis = tm.basis
-    poltype = tm.poltype
-
-    k0s = np.zeros((tms_arr.shape)[:-2])
-    epsilon = np.zeros((tms_arr.shape)[:-2], complex)
-    mu = np.zeros((tms_arr.shape)[:-2], complex)
-    if poltype == "helicity":
-        kappa = np.zeros((tms_arr.shape)[:-2], complex)
-
-    for i, tm in enumerate(tms_obj.flat):
-        if poltype != tm.poltype or basis != tm.basis:
-            raise ValueError(
-                "incompatible T-matrices: mixed poltypes or different bases"
-            )
-        k0s.flat[i] = tm.k0
-        epsilon.flat[i] = tm.material.epsilon
-        mu.flat[i] = tm.material.mu
-        if poltype == "helicity":
-            kappa.flat[i] = tm.material.kappa
-
-    h5file["tmatrix"] = tms_arr
-    if uuid is None:
-        h5file["uuid"] = np.void(_uuid.uuid4().bytes)
-        h5file["uuid"].attrs["version"] = 4
-    else:
-        h5file["uuid"] = uuid
-        h5file["uuid"].attrs["version"] = uuid_version
-
-    _name_descr_kw(h5file, name, description, keywords)
-    h5file["angular_vacuum_wavenumber"] = _collapse_dims(k0s)
-    h5file["angular_vacuum_wavenumber"].attrs["unit"] = lunit + r"^{-1}"
-    h5file["modes/l"] = basis.l
-    h5file["modes/m"] = basis.m
-    h5file["modes/polarization"] = _translate_polarizations(basis.pol, poltype)
-    if any(basis.pidx != 0):
-        h5file["modes/pidx"] = basis.pidx
-    if not np.array_equiv(basis.positions, [[0, 0, 0]]):
-        h5file["modes/positions"] = basis.positions
-        h5file["modes/positions"].attrs["unit"] = lunit
-    if embedding_group is None:
-        embedding_group = h5file.create_group("materials/embedding")
-    embedding_group["relative_permittivity"] = _collapse_dims(epsilon)
-    embedding_group["relative_permeability"] = _collapse_dims(mu)
-    if poltype == "helicity":
-        embedding_group["chirality"] = _collapse_dims(kappa)
-    _name_descr_kw(
-        embedding_group, embedding_name, embedding_description, embedding_keywords
-    )
-    h5file["embedding"] = h5py.SoftLink(embedding_group.name)
-
-    h5file.attrs["created_with"] = (
-        f"python={sys.version.split()[0]},"
-        f"h5py={version('h5py')},"
-        f"treams={version('treams')}"
-    )
-    h5file.attrs["storage_format_version"] = "0.0.1-4-g1266244"
-
-
-def _name_descr_kw(fobj, name, description="", keywords=""):
-    for key, val in [
-        ("name", name),
-        ("description", description),
-        ("keywords", keywords),
-    ]:
-        val = str(val)
-        if val != "":
-            fobj.attrs[key] = val
-
-
-def _convert_to_k0(x, xtype, xunit, k0unit=r"nm^{-1}"):
-    c = 299792458.0
-    k0unit = INVLENGTHS[k0unit]
-    if xtype == "frequency":
-        xunit = FREQUENCIES[xunit]
-        return 2 * np.pi * x / c * (xunit / k0unit)
-    elif xtype == "angular_frequency":
-        xunit = FREQUENCIES[xunit]
-        return x / c * (xunit / k0unit)
-    elif xtype == "angular_vacuum_wavelength":
-        xunit = INVLENGTHS[xunit]
-        return x * (xunit / k0unit)
-    elif xtype == "angular_vacuum_wavenumber":
-        xunit = LENGTHS[xunit]
-        return 2 * np.pi / (x * xunit * k0unit)
-    raise ValueError(f"unrecognized frequency/wavenumber/wavelength type: {xtype}")
-
-
-def load_hdf5(filename, lunit="nm"):
-    """Load a T-matrix stored in a HDF4 file.
-
-    Args:
-        filename (str or h5py.Group): Name of the h5py file or a handle to a h5py group.
-        lunit (str, optional): Unit of length to be used in the T-matrices.
-
-    Returns:
-        np.ndarray[TMatrix]
-    """
-    if isinstance(filename, h5py.Group):
-        return _load_hdf5(filename, lunit)
-    with h5py.File(filename, "r") as f:
-        return _load_hdf5(f, lunit)
-
-
-def _load_hdf5(h5file, lunit=None):
-    for freq_type in (
-        "frequency",
-        "angular_frequency",
-        "vacuum_wavelength",
-        "vacuum_wavenumber",
-        "angular_vacuum_wavenumber",
-    ):
-        if freq_type in h5file:
-            ld_freq = h5file[freq_type][()]
-            break
-    else:
-        raise ValueError("no definition of frequency found")
-    if "modes/positions" in h5file:
-        k0unit = h5file["modes/positions"].attrs.get("unit", lunit) + r"^{-1}"
-    else:
-        k0unit = lunit + r"^{-1}"
-    tms = h5file["tmatrix"][...]
-    k0s = _convert_to_k0(ld_freq, freq_type, h5file[freq_type].attrs["unit"], k0unit)
-
-    epsilon = h5file.get("embedding/relative_permittivity", np.array(None))[()]
-    mu = h5file.get("embedding/relative_permeability", np.array(None))[()]
-    if epsilon is None is mu:
-        n = h5file.get("embedding/refractive_index", np.array(1))[...]
-        z = h5file.get("embedding/relative_impedance", 1 / n)[...]
-        epsilon = n / z
-        mu = n * z
-    epsilon = 1 if epsilon is None else epsilon
-    mu = 1 if mu is None else mu
-
-    kappa = z = h5file.get("embedding/chirality_parameter", np.array(0))[...]
-
-    positions = h5file.get("modes/positions", np.zeros((1, 3)))[...]
-
-    l_inc = h5file.get("modes/l", np.array(None))[...]
-    l_inc = h5file.get("modes/l_incident", l_inc)[()]
-    m_inc = h5file.get("modes/m", np.array(None))[...]
-    m_inc = h5file.get("modes/m_incident", m_inc)[()]
-    pol_inc = h5file.get("modes/polarization", np.array(None))[...]
-    pol_inc = h5file.get("modes/polarization_incident", pol_inc)[()]
-    l_sca = h5file.get("modes/l", np.array(None))[...]
-    l_sca = h5file.get("modes/l_scattered", l_sca)[()]
-    m_sca = h5file.get("modes/m", np.array(None))[...]
-    m_sca = h5file.get("modes/m_scattered", m_sca)[()]
-    pol_sca = h5file.get("modes/polarization", np.array(None))[...]
-    pol_sca = h5file.get("modes/polarization_scattered", pol_sca)[()]
-
-    if any(x is None for x in (l_inc, l_sca, m_inc, m_sca, pol_inc, pol_sca)):
-        raise ValueError("mode definition missing")
-
-    pol_inc, poltype = _translate_polarizations_inv(pol_inc)
-    pol_sca, poltype_sca = _translate_polarizations_inv(pol_sca)
-
-    if poltype_sca != poltype:
-        raise ValueError("different modetypes")
-
-    pidx_inc = h5file.get("modes/position_index", np.zeros_like(l_inc))[()]
-    pidx_inc = h5file.get("modes/positions_index_scattered", pidx_inc)[()]
-    pidx_sca = h5file.get("modes/position_index", np.zeros_like(l_sca))[()]
-    pidx_sca = h5file.get("modes/positions_index_scattered", pidx_sca)[()]
-
-    shape = tms.shape[:-2]
-    k0s = np.broadcast_to(k0s, shape)
-    epsilon = np.broadcast_to(epsilon, shape)
-    mu = np.broadcast_to(mu, shape)
-    kappa = np.broadcast_to(kappa, shape)
-
-    basis_inc = treams.SphericalWaveBasis(
-        zip(pidx_inc, l_inc, m_inc, pol_inc), positions
-    )
-    basis_sca = treams.SphericalWaveBasis(
-        zip(pidx_sca, l_sca, m_sca, pol_sca), positions
-    )
-    basis = basis_inc | basis_sca
-
-    ix_inc = [basis.index(b) for b in basis_inc]
-    ix_sca = [[basis.index(b)] for b in basis_sca]
-
-    res = np.empty(shape, object)
-    for i in np.ndindex(*shape):
-        res[i] = treams.TMatrix(
-            np.zeros((len(basis),) * 2, complex),
-            k0=k0s[i],
-            basis=basis,
-            poltype=poltype,
-            material=treams.Material(epsilon[i], mu[i], kappa[i]),
-        )
-        res[i][ix_sca, ix_inc] = tms[i]
-    if res.shape == ():
-        res = res.item()
-    return res

treams/lattice/__init__.py

-r"""Lattice sums.
-
-.. currentmodule:: treams.lattice
-
-Calculates the lattice sums of the forms
-
-.. math::
-   D_{lm}(k, \boldsymbol k_\parallel, \boldsymbol r, \Lambda_d)
-   = \sideset{}{'}{\sum_{\boldsymbol R \in \Lambda_d}}
-   h_l^{(1)}(k|\boldsymbol r + \boldsymbol R|)
-   Y_{lm}
-   (\theta_{-\boldsymbol r - \boldsymbol R}, \varphi_{-\boldsymbol r - \boldsymbol R})
-   \mathrm e^{\mathrm i \boldsymbol k_\parallel \boldsymbol R}
-
-and
-
-.. math::
-   D_{m}(k, \boldsymbol k_\parallel, \boldsymbol r, \Lambda_d)
-   = \sideset{}{'}{\sum_{\boldsymbol R \in \Lambda_d}}
-   H_m^{(1)}(k|\boldsymbol r + \boldsymbol R|)
-   \mathrm e^{\mathrm i m \varphi_{-\boldsymbol r - \boldsymbol R}}
-   \mathrm e^{\mathrm i \boldsymbol k_\parallel \boldsymbol R}
-
-that arise when translating and summing the spherical and cylindrical solutions
-of the Helmholtz equation in a periodic arrangement. These sums have a notoriously slow
-convergence and at least for lattices with a dimension :math:`d > 1` it is not advisable
-to use the direct approach. Fortunately, it is possible to convert them to exponentially
-convergent series, which are implemented here. For details on the math, see the
-references below.
-
-The lattice of dimension :math:`d \leq 3` (:math:`d \leq 2` in the second case) is
-denoted with :math:`\Lambda_d` and consists of all vectors
-:math:`\boldsymbol R = \sum_{i=1}^{d} n_i \boldsymbol a_i` with :math:`\boldsymbol a_i`
-being basis vectors of the lattice and :math:`n_i \in \mathbb Z`. For :math:`d = 2` the
-lattice is in the `z = 0` plane and for :math:`d = 1` it is along the z-axis. The vector
-:math:`\boldsymbol r` is arbitrary but for best convergence it should be reduced to the
-Wigner-Seitz cell of the lattice. The summation excludes the point
-:math:`\boldsymbol R + \boldsymbol r = 0` if it exists, which is indicated by the prime
-next to the summation sign.
-
-The wave in the lattice is defined by its -- possibly complex-valued -- wave number
-:math:`k` and the (real) components of the wave vector
-:math:`\boldsymbol k_\parallel \in \mathbb R^d` that are parallel to the lattice.
-
-The expressions include the (spherical) Hankel functions of first kind :math:`H_m^{(1)}`
-(:math:`h_l^{(1)}`) and the spherical harmonics :math:`Y_{lm}`. The angles
-:math:`\theta` and :math:`\varphi` are the polar and azimuthal angle, when expressing
-the points in spherical coordinates. In the first case, the degree is
-:math:`l \in \mathbb N_0` and the order is :math:`\mathbb Z \ni m \leq l`. In the second
-case :math:`m \in \mathbb Z`.
-
-
-Available functions
-===================
-
-Accelerated lattice summations
-------------------------------
-
-.. autosummary::
-   :toctree:
-
-   lsumcw1d
-   lsumcw1d_shift
-   lsumcw2d
-   lsumsw1d
-   lsumsw1d_shift
-   lsumsw2d
-   lsumsw2d_shift
-   lsumsw3d
-
-
-Direct summations
------------------
-
-The functions are almost only here for benchmarking and comparison.
-
-.. autosummary::
-   :toctree:
-
-   dsumcw1d
-   dsumcw1d_shift
-   dsumcw2d
-   dsumsw1d
-   dsumsw1d_shift
-   dsumsw2d
-   dsumsw2d_shift
-   dsumsw3d
-
-
-Miscellaneous functions
------------------------
-
-.. autosummary::
-   :toctree:
-
-   area
-   cube
-   cubeedge
-   diffr_orders_circle
-   reciprocal
-   volume
-
-
-Cython module
--------------
-
-.. autosummary::
-   :toctree:
-   :template: cython-module
-
-   cython_lattice
-
-
-References
-==========
-
-* `[1] K. Kambe, Zeitschrift Für Naturforschung A 22, 3 (1967). <https://doi.org/10.1515/zna-1967-0305>`_
-* `[2] K. Kambe, Zeitschrift Für Naturforschung A 23, 9 (1968). <https://doi.org/10.1515/zna-1968-0908>`_
-* `[3] C. M. Linton, SIAM Rev. 52, 630 (2010). <https://doi.org/10.1137/09075130X>`_
-* `[4] D. Beutel el al., J. Opt. Soc. Am. B (2021). <https://doi.org/10.1364/JOSAB.419645>`_
-"""
-
-import numpy as np
-
-from treams.lattice import _misc
-from treams.lattice._gufuncs import *  # noqa: F403
-from treams.lattice._misc import cube, cubeedge  # noqa: F401
-
-
-def diffr_orders_circle(b, rmax):
-    """Diffraction orders in a circle.
-
-    Given the reciprocal lattice defined by the vectors that make up the rows of `b`,
-    return all diffraction orders within a circle of radius `rmax`.
-
-    Args:
-        b (float, (2, 2)-array): Reciprocal lattice vectors
-        rmax (float): Maximal radius
-
-    Returns:
-        float array
-    """
-    return _misc.diffr_orders_circle(np.array(b), rmax)
-
-
-def lsumsw(dim, l, m, k, kpar, a, r, eta, out=None, **kwargs):  # noqa: E741
-    if dim == 1:
-        return lsumsw1d_shift(l, m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
-    elif dim == 2:
-        return lsumsw2d_shift(l, m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
-    elif dim == 3:
-        return lsumsw3d(l, m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
-    raise ValueError(f"invalid dimension '{dim}'")
-
-
-def realsumsw(dim, l, m, k, kpar, a, r, eta, out=None, **kwargs):  # noqa: E741
-    if dim == 1:
-        return realsumsw1d_shift(  # noqa: F405
-            l, m, k, kpar, a, r, eta, out=out, **kwargs
-        )
-    elif dim == 2:
-        return realsumsw2d_shift(  # noqa: F405
-            l, m, k, kpar, a, r, eta, out=out, **kwargs
-        )
-    elif dim == 3:
-        return realsumsw3d(l, m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
-    raise ValueError(f"invalid dimension '{dim}'")
-
-
-def recsumsw(dim, l, m, k, kpar, a, r, eta, out=None, **kwargs):  # noqa: E741
-    if dim == 1:
-        return recsumsw1d_shift(  # noqa: F405
-            l, m, k, kpar, a, r, eta, out=out, **kwargs
-        )
-    elif dim == 2:
-        return recsumsw2d_shift(  # noqa: F405
-            l, m, k, kpar, a, r, eta, out=out, **kwargs
-        )
-    elif dim == 3:
-        return recsumsw3d(l, m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
-    raise ValueError(f"invalid dimension '{dim}'")
-
-
-def lsumcw(dim, m, k, kpar, a, r, eta, out=None, **kwargs):  # noqa: E741
-    if dim == 1:
-        return lsumcw1d_shift(m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
-    elif dim == 2:
-        return lsumcw2d(m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
-    raise ValueError(f"invalid dimension '{dim}'")
-
-
-def realsumcw(dim, m, k, kpar, a, r, eta, out=None, **kwargs):
-    if dim == 1:
-        return realsumcw1d_shift(m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
-    elif dim == 2:
-        return realsumcw2d(m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
-    raise ValueError(f"invalid dimension '{dim}'")
-
-
-def recsumcw(dim, m, k, kpar, a, r, eta, out=None, **kwargs):
-    if dim == 1:
-        return recsumcw1d_shift(m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
-    elif dim == 2:
-        return recsumcw2d(m, k, kpar, a, r, eta, out=out, **kwargs)  # noqa: F405
-    raise ValueError(f"invalid dimension '{dim}'")
-
-
-def dsumsw(dim, l, m, k, kpar, a, r, i, out=None, **kwargs):  # noqa: E741
-    if dim == 1:
-        return dsumsw1d_shift(l, m, k, kpar, a, r, i, out=out, **kwargs)  # noqa: F405
-    elif dim == 2:
-        return dsumsw2d_shift(l, m, k, kpar, a, r, i, out=out, **kwargs)  # noqa: F405
-    elif dim == 3:
-        return dsumsw3d(l, m, k, kpar, a, r, i, out=out, **kwargs)  # noqa: F405
-    raise ValueError(f"invalid dimension '{dim}'")
-
-
-def dsumcw(dim, m, k, kpar, a, r, i, out=None, **kwargs):
-    if dim == 1:
-        return dsumcw1d_shift(m, k, kpar, a, r, i, out=out, **kwargs)  # noqa: F405
-    elif dim == 2:
-        return dsumcw2d(m, k, kpar, a, r, i, out=out, **kwargs)  # noqa: F405
-    raise ValueError(f"invalid dimension '{dim}'")

treams/lattice/cython_lattice.pxd

-ctypedef fused real_t:
-    double
-    long
-
-ctypedef fused number_t:
-    double complex
-    double
-
-cdef real_t area(real_t *a, real_t *b) nogil
-cdef long cube_next(long *r, long d, long n) nogil
-cdef long cubeedge_next(long *r, long d, long n) nogil
-cpdef diffr_orders_circle(real_t[:, :] b, double rmax)
-cpdef long ncube(long d, long n) nogil
-cpdef long nedge(long d, long n) nogil
-cdef void recvec2(double *a0, double *a1, double *b0, double *b1) nogil
-cdef void recvec3(double *a0, double *a1, double *a2, double *b0, double *b1, double *b2) nogil
-cdef real_t volume(real_t *a, real_t *b, real_t *c) nogil
-
-cpdef double complex lsumcw1d(long l, number_t k, double kpar, double a, double r, number_t eta) nogil
-cdef double complex lsumcw1d_shift(long l, number_t k, double kpar, double a, double *r, number_t eta) nogil
-cdef double complex lsumcw2d(long l, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
-cpdef double complex lsumsw1d(long l, number_t k, double kpar, double a, double r, number_t eta) nogil
-cdef double complex lsumsw1d_shift(long l, long m, number_t k, double kpar, double a, double *r, number_t eta) nogil
-cdef double complex lsumsw2d(long l, long m, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
-cdef double complex lsumsw2d_shift(long l, long m, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
-cdef double complex lsumsw3d(long l, long m, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
-
-cpdef double complex realsumcw1d(long l, number_t k, double kpar, double a, double r, number_t eta) nogil
-cdef double complex realsumcw1d_shift(long l, number_t k, double kpar, double a, double *r, number_t eta) nogil
-cdef double complex realsumcw2d(long l, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
-cpdef double complex realsumsw1d(long l, number_t k, double kpar, double a, double r, number_t eta) nogil
-cdef double complex realsumsw1d_shift(long l, long m, number_t k, double kpar, double a, double *r, number_t eta) nogil
-cdef double complex realsumsw2d(long l, long m, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
-cdef double complex realsumsw2d_shift(long l, long m, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
-cdef double complex realsumsw3d(long l, long m, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
-
-cpdef double complex recsumcw1d(long l, number_t k, double kpar, double a, double r, number_t eta) nogil
-cdef double complex recsumcw1d_shift(long l, number_t k, double kpar, double a, double *r, number_t eta) nogil
-cdef double complex recsumcw2d(long l, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
-cpdef double complex recsumsw1d(long l, number_t k, double kpar, double a, double r, number_t eta) nogil
-cdef double complex recsumsw1d_shift(long l, long m, number_t k, double kpar, double a, double *r, number_t eta) nogil
-cdef double complex recsumsw2d(long l, long m, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
-cdef double complex recsumsw2d_shift(long l, long m, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
-cdef double complex recsumsw3d(long l, long m, number_t k, double *kpar, double *a, double *r, number_t eta) nogil
-
-cpdef double complex zero3d(double complex eta) nogil
-cpdef double complex zero2d(double complex eta) nogil
-
-cpdef double complex dsumcw1d(long l, number_t k, double kpar, double a, double r, long i) nogil
-cdef double complex dsumcw1d_shift(long l, number_t k, double kpar, double a, double *r, long i) nogil
-cdef double complex dsumcw2d(long l, number_t k, double *kpar, double *a, double *r, long i) nogil
-cpdef double complex dsumsw1d(long l, number_t k, double kpar, double a, double r, long i) nogil
-cdef double complex dsumsw1d_shift(long l, long m, number_t k, double kpar, double a, double *r, long i) nogil
-cdef double complex dsumsw2d(long l, long m, number_t k, double *kpar, double *a, double *r, long i) nogil
-cdef double complex dsumsw2d_shift(long l, long m, number_t k, double *kpar, double *a, double *r, long i) nogil
-cdef double complex dsumsw3d(long l, long m, number_t k, double *kpar, double *a, double *r, long i) nogil

treams/misc.py

-"""Miscellaneous functions.
-
-.. autosummary::
-   :toctree:
-
-   basischange
-   firstbrillouin1d
-   firstbrillouin2d
-   firstbrillouin3d
-   pickmodes
-   refractive_index
-   wave_vec_z
-"""
-
-import numpy as np
-
-import treams.lattice as la
-
-
-def refractive_index(epsilon=1, mu=1, kappa=0):
-    r"""Refractive index of a (chiral) medium.
-
-    The refractive indeces in a chiral medium :math:`\sqrt{\epsilon\mu} \mp \kappa` are
-    returned with the negative helicity result first.
-
-    Args:
-        epsilon (float or complex, array_like, optional): Relative permittivity,
-            defaults to 1.
-        mu (float or complex, array_like, optional): Relative permeability,
-            defaults to 1.
-        kappa (float or complex, array_like, optional): Chirality parameter,
-            defaults to 0.
-
-    Returns:
-        float or complex, (2,)-array
-
-    """
-    epsilon = np.array(epsilon)
-    n = np.sqrt(epsilon * mu)
-    res = np.stack((n - kappa, n + kappa), axis=-1)
-    res[np.imag(res) < 0] *= -1
-    return res
-
-
-def basischange(out, in_=None):
-    """Coefficients for the basis change between helicity and parity modes.
-
-    Args:
-        out (3- or 4-tuple of (M,)-arrays): Output modes, the last array is taken as
-            polarization.
-        in_ (3- or 4-tuple of (N,)-arrays, optional): Input modes, if none are given,
-            equal to the output modes
-
-    Returns:
-        float, ((M, N)-array
-    """
-    if in_ is None:
-        in_ = out
-    out = np.array([*zip(*out)])
-    in_ = np.array([*zip(*in_)])
-    res = np.zeros((out.shape[0], in_.shape[0]))
-    out = out[:, None, :]
-    sqhalf = np.sqrt(0.5)
-    equal = (out[:, :, :-1] == in_[:, :-1]).all(axis=-1)
-    minus = np.logical_and(out[:, :, -1] == in_[:, -1], in_[:, -1] == 0)
-    res[equal] = sqhalf
-    res[np.logical_and(equal, minus)] = -sqhalf
-    return res
-
-
-def pickmodes(out, in_):
-    """Coefficients to pick modes.
-
-    Args:
-        out (3- or 4-tuple of (M,)-arrays): Output modes, the last array is taken as
-            polarization.
-        in_ (3- or 4-tuple of (N,)-arrays, optional): Input modes, if none are given,
-            equal to the output modes
-
-    Returns:
-        float, ((M, N)-array
-    """
-    out = np.array([*zip(*out)])
-    in_ = np.array([*zip(*in_)])
-    return np.all(out[:, None, :] == in_, axis=-1)
-
-
-# def wave_vec_zs(kpars, ks):
-#     kpars = np.array(kpars)
-#     ks = np.array(ks)
-#     if ks.ndim == 0:
-#         nmat = npol = 1
-#     elif ks.ndim == 1:
-#         npol = ks.shape[0]
-#         nmat = 1
-#     elif ks.ndim == 2:
-#         nmat, npol = ks.shape
-#     else:
-#         raise ValueError("ks has invalid shape")
-#     res = np.zeros((nmat, kpars.shape[0], npol), np.complex128)
-#     for i, kpar in enumerate(kpars):
-#         res[:, i, :] = wave_vec_z(kpar[0], kpar[1], ks)
-#     return res
-
-
-def wave_vec_z(kx, ky, k):
-    r"""Z component of the wave vector with positive imaginary part.
-
-    The result is :math:`k_z = \sqrt{k^2 - k_x^2 - k_y^2}` with
-    :math:`\arg k_z \in \[ 0, \pi )`.
-
-    Args:
-        kx (float, array_like): X component of the wave vector
-        ky (float, array_like): Y component of the wave vector
-        k (float or complex, array_like): Wave number
-
-    Returns:
-        complex
-    """
-    kx = np.asarray(kx)
-    ky = np.asarray(ky)
-    k = np.asarray(k, complex)
-    res = np.sqrt(k * k - kx * kx - ky * ky)
-    if res.ndim == 0 and res.imag < 0:
-        res = -res
-    elif res.ndim > 0:
-        res[np.imag(res) < 0] *= -1
-    return res
-
-
-def firstbrillouin1d(kpar, b):
-    """Map wave vector to first Brillouin zone in 1D.
-
-    Reduce the 1d wave vector (actually just a number) to the first Brillouin zone, i.e.
-    the range `(-b/2, b/2]`
-
-    Args:
-        kpar (float64): (parallel) wave vector
-        b (float64): reciprocal lattice vector
-
-    Returns:
-        float64
-    """
-    kpar -= b * np.round(kpar / b)
-    if kpar > 0.5 * b:
-        kpar -= b
-    if kpar <= -0.5 * b:
-        kpar += b
-    return kpar
-
-
-def firstbrillouin2d(kpar, b, n=2):
-    """Map wave vector to first Brillouin zone in 2D.
-
-    The reduction to the first Brillouin zone is first approximated roughly. From this
-    approximated vector and its 8 neighbours, the shortest one is picked. As a
-    sufficient approximation is not guaranteed (especially for extreme geometries),
-    this process is iterated `n` times.
-
-    Args:
-        kpar (1d-array): parallel wave vector
-        b (2d-array): reciprocal lattice vectors
-        n (int): number of iterations
-
-    Returns:
-        (1d-array)
-    """
-    kparstart = kpar
-    b1 = b[0, :]
-    b2 = b[1, :]
-    normsq1 = b1 @ b1
-    normsq2 = b2 @ b2
-    normsqp = (b1 + b2) @ (b1 + b2)
-    normsqm = (b1 - b2) @ (b1 - b2)
-    if (
-        normsqp < normsq1 - 1e-14
-        or normsqp < normsq2 - 1e-14
-        or normsqm < normsq1 - 1e-14
-        or normsqm < normsq2 - 1e-14
-    ):
-        raise ValueError("Lattice vectors are not of minimal length")
-    # Rough estimate
-    kpar -= b1 * np.round((kpar @ b1) / normsq1)
-    kpar -= b2 * np.round((kpar @ b2) / normsq2)
-    # todo: precise
-    options = kpar + la.cube(2, 1) @ b
-    for i, option in enumerate(options):
-        if option @ option < kpar @ kpar:
-            kpar = option
-    if n == 0 or np.array_equal(kpar, kparstart):
-        return kpar
-    return firstbrillouin2d(kpar, b, n - 1)
-
-
-def firstbrillouin3d(kpar, b, n=2):
-    """Map wave vector to first Brillouin zone in 3D.
-
-    The reduction to the first Brillouin zone is first approximated roughly. From this
-    approximated vector and its 26 neighbours, the shortest one is picked. As a
-    sufficient approximation is not guaranteed (especially for extreme geometries),
-    this process is iterated `n` times.
-
-    Args:
-        kpar (1d-array): parallel wave vector
-        b (2d-array): reciprocal lattice vectors
-        n (int): number of iterations
-
-    Returns:
-        (1d-array)
-    """
-    kparstart = kpar
-    b1 = b[0, :]
-    b2 = b[1, :]
-    b3 = b[2, :]
-    normsq1 = b1 @ b1
-    normsq2 = b2 @ b2
-    normsq3 = b3 @ b3
-    # todo: Minimal length
-    # Rough estimate
-    kpar -= b1 * np.round((kpar @ b1) / normsq1)
-    kpar -= b2 * np.round((kpar @ b2) / normsq2)
-    kpar -= b3 * np.round((kpar @ b3) / normsq3)
-    # todo: precise
-    options = kpar + la.cube(3, 1) @ b
-    for i, option in enumerate(options):
-        if option @ option < kpar @ kpar:
-            kpar = option
-    if n == 0 or np.array_equal(kpar, kparstart):
-        return kpar
-    return firstbrillouin3d(kpar, b, n - 1)

treams/special/__init__.py

-r"""Special (mathematical) functions.
-
-.. currentmodule:: treams.special
-
-Special mathematical functions used in :mod:`treams`. Some functions are reexported from
-:py:mod:`scipy.special`. Most functions are available as Numpy universal functions
-(:py:class:`numpy.ufunc`) or as generalized universal functions
-(:ref:`c-api.generalized-ufuncs`).
-
-Available functions
-===================
-
-Bessel and Hankel functions, with their spherical counterparts, derivatives
----------------------------------------------------------------------------
-
-.. autosummary::
-   :toctree:
-
-   hankel1_d
-   hankel2_d
-   jv_d
-   yv_d
-   spherical_hankel1
-   spherical_hankel2
-   spherical_hankel1_d
-   spherical_hankel2_d
-
-
-Those functions are just reexported from Scipy. So, one only needs to import this
-subpackage within treams.
-
-+------------------------------------------------------------+-------------------------+
-| :py:data:`~scipy.special.hankel1`\(v, z[, out])            | Hankel function of the  |
-|                                                            | first kind.             |
-+------------------------------------------------------------+-------------------------+
-| :py:data:`~scipy.special.hankel2`\(v, z[, out])            | Hankel function of the  |
-|                                                            | second kind.            |
-+------------------------------------------------------------+-------------------------+
-| :py:data:`~scipy.special.jv`\(v, z[, out])                 | Bessel function of the  |
-|                                                            | first kind of real      |
-|                                                            | order and complex       |
-|                                                            | argument.               |
-+------------------------------------------------------------+-------------------------+
-| :py:data:`~scipy.special.yv`\(v, z[, out])                 | Bessel function of the  |
-|                                                            | second kind of real     |
-|                                                            | order and complex       |
-|                                                            | argument.               |
-+------------------------------------------------------------+-------------------------+
-| | :py:func:`spherical_jn <scipy.special.spherical_jn>`\(n, | Spherical Bessel        |
-|   z[, derivative])                                         | function of the first   |
-|                                                            | kind or its derivative. |
-+------------------------------------------------------------+-------------------------+
-| | :py:func:`spherical_yn <scipy.special.spherical_yn>`\(n, | Spherical Bessel        |
-|   z[, derivative])                                         | function of the second  |
-|                                                            | kind or its derivative. |
-+------------------------------------------------------------+-------------------------+
-
-Those functions just wrap Scipy functions with special optional arguments to be able to
-analogously access them like their non-spherical counterparts:
-
-.. autosummary::
-   :toctree:
-
-   spherical_jn_d
-   spherical_yn_d
-
-
-Scipy functions with enhanced domain
-------------------------------------
-
-.. autosummary::
-   :toctree:
-
-   sph_harm
-   lpmv
-
-
-Integrals for the Ewald summation
----------------------------------
-
-.. autosummary::
-   :toctree:
-
-   incgamma
-   intkambe
-
-
-Wigner d- and Wigner D-matrix elements
---------------------------------------
-
-.. autosummary::
-   :toctree:
-
-   wignersmalld
-   wignerd
-
-
-Wigner 3j-symbols
------------------
-
-.. autosummary::
-   :toctree:
-
-   wigner3j
-
-
-Vector wave functions
---------------------------
-
-.. autosummary::
-   :toctree:
-
-   pi_fun
-   tau_fun
-
-Spherical waves and translation coefficients
-
-.. autosummary::
-   :toctree:
-
-   vsh_X
-   vsh_Y
-   vsh_Z
-   vsw_M
-   vsw_N
-   vsw_A
-   vsw_rM
-   vsw_rN
-   vsw_rA
-   tl_vsw_A
-   tl_vsw_B
-   tl_vsw_rA
-   tl_vsw_rB
-
-
-Cylindrical waves
-
-.. autosummary::
-   :toctree:
-
-   vcw_M
-   vcw_N
-   vcw_A
-   vcw_rM
-   vcw_rN
-   vcw_rA
-   tl_vcw
-   tl_vcw_r
-
-
-Plane waves
-
-.. autosummary::
-   :toctree:
-
-   vpw_M
-   vpw_N
-   vpw_A
-
-
-Coordinate system transformations
----------------------------------
-
-.. autosummary::
-   :toctree:
-
-   car2cyl
-   car2sph
-   cyl2car
-   cyl2sph
-   sph2car
-   sph2cyl
-   vcar2cyl
-   vcar2sph
-   vcyl2car
-   vcyl2sph
-   vsph2car
-   vsph2cyl
-   car2pol
-   pol2car
-   vcar2pol
-   vpol2car
-
-
-Cython module
--------------
-
-.. autosummary::
-   :toctree:
-   :template: cython-module
-
-   cython_special
-"""
-
-from scipy.special import (  # noqa: F401
-    hankel1,
-    hankel2,
-    jv,
-    spherical_jn,
-    spherical_yn,
-    yv,
-)
-
-from treams.special import _gufuncs, _ufuncs  # noqa: F401
-from treams.special._gufuncs import *  # noqa: F401, F403
-from treams.special._ufuncs import *  # noqa: F401, F403
-
-
-def spherical_jn_d(n, z):
-    """Derivative of the spherical Bessel function of the first kind.
-
-    This is simply a wrapper for `scipy.special.spherical_jn(n, z, True)`, see
-    :py:func:`scipy.special.spherical_jn`. It's here to have a consistent way of
-    calling the derivative of a (spherical) Bessel or Hankel function.
-
-    Args:
-        n (int, array_like): Order
-        z (float or complex, array_like): Argument
-
-    Returns:
-        float or complex
-
-    References:
-        - `DLMF 10.47 <https://dlmf.nist.gov/10.47>`_
-        - `DLMF 10.51 <https://dlmf.nist.gov/10.51>`_
-    """
-    return spherical_jn(n, z, True)
-
-
-def spherical_yn_d(n, z):
-    """Derivative of the spherical Bessel function of the second kind.
-
-    This is simply a wrapper for `scipy.special.spherical_yn(n, z, True)`, see
-    :py:func:`scipy.special.spherical_jn`. It's here to have a consistent way of
-    calling the derivative of a (spherical) Bessel or Hankel function.
-
-    Args:
-        n (int, array_like): Order
-        z (float or complex, array_like): Argument
-
-    Returns:
-        float or complex
-
-    References:
-        - `DLMF 10.47 <https://dlmf.nist.gov/10.47>`_
-        - `DLMF 10.51 <https://dlmf.nist.gov/10.51>`_
-    """
-    return spherical_yn(n, z, True)

treams/special/cython_special.pxd

-ctypedef fused number_t:
-    double complex
-    double
-
-
-cpdef double complex hankel1_d(double l, double complex z) nogil
-cpdef double complex hankel2_d(double l, double complex z) nogil
-cpdef number_t jv_d(double l, number_t z) nogil
-cpdef double complex spherical_hankel1(double n, number_t z) nogil
-cpdef double complex spherical_hankel1_d(double l, double complex z) nogil
-cpdef double complex spherical_hankel2(double n, number_t z) nogil
-cpdef double complex spherical_hankel2_d(double l, double complex z) nogil
-cpdef number_t yv_d(double l, number_t z) nogil
-
-
-cdef void car2cyl(double *input, double *output) nogil
-cdef void car2pol(double *input, double *output) nogil
-cdef void car2sph(double *input, double *output) nogil
-cdef void cyl2car(double *input, double *output) nogil
-cdef void cyl2sph(double *input, double *output) nogil
-cdef void pol2car(double *input, double *output) nogil
-cdef void sph2car(double *input, double *output) nogil
-cdef void sph2cyl(double *input, double *output) nogil
-
-cdef void vcar2cyl(number_t *iv, double *ip, number_t *ov) nogil
-cdef void vcar2pol(number_t *iv, double *ip, number_t *ov) nogil
-cdef void vcar2sph(number_t *iv, double *ip, number_t *ov) nogil
-cdef void vcyl2car(number_t *iv, double *ip, number_t *ov) nogil
-cdef void vcyl2sph(number_t *iv, double *ip, number_t *ov) nogil
-cdef void vpol2car(number_t *iv, double *ip, number_t *ov) nogil
-cdef void vsph2car(number_t *iv, double *ip, number_t *ov) nogil
-cdef void vsph2cyl(number_t *iv, double *ip, number_t *ov) nogil
-
-
-cpdef number_t incgamma(double l, number_t z) nogil
-cpdef number_t intkambe(long n, number_t z, number_t eta) nogil
-
-
-cpdef number_t lpmv(double m, double l, number_t z) nogil
-cpdef number_t pi_fun(double l, double m, number_t x) nogil
-cpdef double complex sph_harm(double m, double l, double phi, number_t theta) nogil
-cpdef number_t tau_fun(double l, double m, number_t x) nogil
-
-cpdef double complex tl_vcw(double kz1, long mu, double kz2, long m, double complex krr, double phi, double z) nogil
-cpdef double complex tl_vcw_r(double kz1, long mu, double kz2, long m, number_t krr, double phi, double z) nogil
-
-cpdef double complex _tl_vsw_helper(long l, long m, long lambda_, long mu, long p, long q) nogil
-
-cpdef double complex tl_vsw_A(long lambda_, long mu, long l, long m, double complex kr, number_t theta, double phi) nogil
-cpdef double complex tl_vsw_B(long lambda_, long mu, long l, long m, double complex kr, number_t theta, double phi) nogil
-cpdef double complex tl_vsw_rA(long lambda_, long mu, long l, long m, number_t kr, number_t theta, double phi) nogil
-cpdef double complex tl_vsw_rB(long lambda_, long mu, long l, long m, number_t kr, number_t theta, double phi) nogil
-
-cdef void vcw_A(double kz, long m, double complex krr, double phi, double z, double complex k, long pol, double complex *out) nogil
-cdef void vcw_M(double kz, long m, double complex krr, double phi, double z, double complex *out) nogil
-cdef void vcw_N(double kz, long m, double complex krr, double phi, double z, double complex k, double complex *out) nogil
-cdef void vcw_rA(double kz, long m, number_t krr, double phi, double z, double complex k, long pol, double complex *out) nogil
-cdef void vcw_rM(double kz, long m, number_t krr, double phi, double z, double complex *out, long i) nogil
-cdef void vcw_rN(double kz, long m, number_t krr, double phi, double z, double complex k, double complex *out) nogil
-
-cdef void vpw_A(number_t kx, number_t ky, number_t kz, double x, double y, double z, long pol, double complex *out) nogil
-cdef void vpw_M(number_t kx, number_t ky, number_t kz, double x, double y, double z, double complex *out) nogil
-cdef void vpw_N(number_t kx, number_t ky, number_t kz, double x, double y, double z, double complex *out) nogil
-
-cdef void vsh_X(long l, long m, number_t theta, double phi, double complex *out) nogil
-cdef void vsh_Y(long l, long m, number_t theta, double phi, double complex *out) nogil
-cdef void vsh_Z(long l, long m, number_t theta, double phi, double complex *out) nogil
-
-cdef void vsw_A(long l, long m, double complex kr, number_t theta, double phi, long pol, double complex *out) nogil
-cdef void vsw_M(long l, long m, double complex kr, number_t theta, double phi, double complex *out) nogil
-cdef void vsw_N(long l, long m, double complex kr, number_t theta, double phi, double complex *out) nogil
-cdef void vsw_rA(long l, long m, number_t kr, number_t theta, double phi, long pol, double complex *out) nogil
-cdef void vsw_rM(long l, long m, number_t kr, number_t theta, double phi, double complex *out) nogil
-cdef void vsw_rN(long l, long m, number_t kr, number_t theta, double phi, double complex *out) nogil
-
-
-cpdef number_t wignersmalld(long l, long m, long k, number_t theta) nogil
-cpdef double complex wignerd(long l, long m, long k, double phi, number_t theta, double psi) nogil
-
-
-cpdef double wigner3j(long j1, long j2, long j3, long m1, long m2, long m3) nogil

treams/util.py

-"""Utilities.
-
-Collection of functions to make :class:`AnnotatedArray` work. The implementation is
-aimed to be quite general, while still providing a solid basis for the rest of the code.
-"""
-
-import abc
-import collections
-import copy
-import functools
-import itertools
-import warnings
-
-import numpy as np
-
-import treams._operators as _op
-
-__all__ = [
-    "AnnotatedArray",
-    "AnnotationDict",
-    "AnnotationError",
-    "AnnotationSequence",
-    "AnnotationWarning",
-    "implements",
-    "OrderedSet",
-]
-
-
-class OrderedSet(collections.abc.Sequence, collections.abc.Set):
-    """Ordered set.
-
-    A abstract base class that combines a sequence and set. In contrast to regular sets
-    it is expected that the equality comparison only returns `True` if all entries are
-    in the same order.
-    """
-
-    @abc.abstractmethod
-    def __eq__(self, other):
-        """Equality test."""
-        raise NotImplementedError
-
-
-class AnnotationWarning(UserWarning):
-    """Custom warning for Annotations.
-
-    By default the warning filter is set to 'always'.
-    """
-
-    pass
-
-
-warnings.simplefilter("always", AnnotationWarning)
-
-
-class AnnotationError(Exception):
-    """Custom exception for Annotations."""
-
-    pass
-
-
-class AnnotationDict(collections.abc.MutableMapping):
-    """Dictionary that notifies the user when overwriting keys.
-
-    Behaves mostly similar to regular dictionaries, except that when overwriting
-    existing keys a :class:`AnnotationWarning` is emmitted.
-
-    Examples:
-        An :class:`AnnotationDict` can be created from other mappings, from a list of
-        key-value pairs (note how a warning is emmitted for the duplicate key), and from
-        keyword arguments.
-
-        .. code-block:: python
-
-            >>> AnnotationDict({"a": 1, "b": 2})
-            AnnotationDict({'a': 1, 'b': 2})
-            >>> AnnotationDict([("a", 1), ("b", 2), ("a", 3)])
-            treams/util.py:74: AnnotationWarning: overwriting key 'a'
-            warnings.warn(f"overwriting key '{key}'", AnnotationWarning)
-            AnnotationDict({'a': 3, 'b': 2})
-            >>> AnnotationDict({"a": 1, "b": 2}, c=3)
-            AnnotationDict({'a': 1, 'b': 2, 'c': 3})
-
-    Warns:
-        AnnotationWarning
-    """
-
-    def __init__(self, items=(), /, **kwargs):
-        """Initialization."""
-        self._dct = {}
-        for i in (items.items() if hasattr(items, "items") else items, kwargs.items()):
-            for key, val in i:
-                self[key] = val
-
-    def __getitem__(self, key):
-        """Get a value by its key.
-
-        Args:
-            key (hashable): Key
-        """
-        return self._dct[key]
-
-    def __setitem__(self, key, val):
-        """Set item specified by key to the defined value.
-
-        When overwriting an existing key an :class:`AnnotationWarning` is emitted.
-        Avoid the warning by explicitly deleting the key first.
-
-        Args:
-            key (hashable): Key
-            val : Value
-
-        Warns:
-            AnnotationWarning
-        """
-        if key in self and self[key] != val:
-            warnings.warn(f"overwriting key '{key}'", AnnotationWarning)
-        self._dct[key] = val
-
-    def __delitem__(self, key):
-        """Delete the key."""
-        del self._dct[key]
-
-    def __iter__(self):
-        """Iterate over the keys.
-
-        Returns:
-            Iterator
-        """
-        return iter(self._dct)
-
-    def __len__(self):
-        """Number of keys contained.
-
-        Returns:
-            int
-        """
-        return len(self._dct)
-
-    def __repr__(self):
-        """String representation.
-
-        Returns:
-            str
-        """
-        return f"{self.__class__.__name__}({repr(self._dct)})"
-
-    def match(self, other):
-        """Compare the own keys to another dictionary.
-
-        This emits an :class:`AnnotationWarning` for each key that would be overwritten
-        by the given dictionary.
-
-        Args:
-            other (Mapping)
-
-        Returns:
-            None
-
-        Warns:
-            AnnotationWarning
-        """
-        for key, val in self.items():
-            if key in other and other[key] != val:
-                warnings.warn(f"incompatible key '{key}'", AnnotationWarning)
-
-
-class SequenceAsDict(collections.abc.MutableMapping):
-    def __get__(self, obj, objtype=None):
-        self._obj = obj
-        return self
-
-    def __set__(self, obj, dct):
-        self._obj = obj
-        for key, val in dct.items():
-            self[key] = val
-
-    def __getitem__(self, key):
-        res = tuple(i.get(key) for i in self._obj)
-        if all(i is None for i in res):
-            raise KeyError(key)
-        return res
-
-    def __setitem__(self, key, val):
-        val = (None,) * (len(self._obj) - len(val)) + val
-        for dct, val in zip(reversed(self._obj), reversed(val)):
-            if val is None:
-                dct.pop(key, None)
-            else:
-                dct[key] = val
-
-    def __delitem__(self, key):
-        found = False
-        for dct in self._obj:
-            if key in dct:
-                del dct[key]
-                found = True
-        if not found:
-            raise KeyError(key)
-
-    def __iter__(self):
-        return iter({key for dct in self._obj for key in dct})
-
-    def __len__(self):
-        return len({key for dct in self._obj for key in dct})
-
-    def __repr__(self):
-        return f"{self.__class__.__name__}({dict(i for i in self.items())})"
-
-
-class AnnotationSequence(collections.abc.Sequence):
-    """A Sequence of dictionaries.
-
-    This class is intended to work together with :class:`AnnotationDict`. It provides
-    convenience functions to interact with multiple of those dictionaries, which are
-    mainly used to keep track of the annotations made to each dimension of an
-    :class:`AnnotatedArray`. While the sequence itself is immutable the entries of each
-    dictionary is mutable.
-
-    Args:
-        *args: Items of the sequence
-        mapping (AnnotationDict): Type of mapping to use in the sequence.
-
-    Warns:
-        AnnotationWarning
-    """
-
-    as_dict = SequenceAsDict()
-
-    def __init__(self, *args, mapping=AnnotationDict):
-        """Initialization."""
-        self._ann = tuple(mapping(i) for i in args)
-
-    def __len__(self):
-        """Number of dictionaries in the sequence.
-
-        Returns:
-            int
-        """
-        return len(self._ann)
-
-    def __getitem__(self, key):
-        """Get an item or subsequence.
-
-        Indexing works with integers and slices like regular tuples. Additionally, it is
-        possible to get a copy of the object with `()`, or a new sequence of mappings in
-        a list (or other iterable) of integers.
-
-        Args:
-            key (iterable, slice, int)
-
-        Returns:
-            mapping
-        """
-        if isinstance(key, tuple) and key == ():
-            return copy.copy(self._ann)
-        if isinstance(key, slice):
-            return type(self)(*self._ann[key])
-        if isinstance(key, (int, np.integer)):
-            return self._ann[key]
-        res = []
-        for k in key:
-            if not isinstance(k, int):
-                raise TypeError(
-                    "sequence index must be integer, slice, list of integers, or '()'"
-                )
-            res.append(self[k])
-        return type(self)(*res)
-
-    def update(self, other):
-        """Update all mappings in the sequence at once.
-
-        The given sequence is aliged at the last entry and then updated pairwise.
-        Warnings for overwritten keys are extended by the information at which index
-        they occurred.
-
-        Args:
-            other (Sequence[Mapping]): Mappings to update with.
-
-        Warns:
-            AnnotationWarning
-        """
-        if len(other) > len(self):
-            warnings.warn(
-                f"argument of length {len(self)} given: "
-                f"ignore leading {len(other) - len(self)} entries",
-                AnnotationWarning,
-            )
-        for i, (dest, src) in enumerate(zip(reversed(self), reversed(other))):
-            with warnings.catch_warnings():
-                warnings.simplefilter("error", category=AnnotationWarning)
-                try:
-                    dest.update(src)
-                except AnnotationWarning as err:
-                    warnings.simplefilter("always", category=AnnotationWarning)
-                    warnings.warn(
-                        f"at index {len(self) - i - 1}: " + err.args[0],
-                        AnnotationWarning,
-                    )
-
-    def match(self, other):
-        """Match all mappings at once.
-
-        The given sequence is aliged at the last entry and then updated pairwise.
-        Warnings for overwritten keys are extended by the information at which index
-        they occurred, see also :func:`AnnotationDict.match`.
-
-        Args:
-            other (Sequence[Mappings]): Mappings to match.
-
-        Warns:
-            AnnotationWarning
-        """
-        for i, (dest, src) in enumerate(zip(reversed(self), reversed(other))):
-            with warnings.catch_warnings():
-                warnings.simplefilter("error", category=AnnotationWarning)
-                try:
-                    dest.match(src)
-                except AnnotationWarning as err:
-                    warnings.simplefilter("always", category=AnnotationWarning)
-                    warnings.warn(
-                        f"at dimension {len(self) + i}: " + err.args[0],
-                        AnnotationWarning,
-                    )
-
-    def __eq__(self, other):
-        """Equality test.
-
-        Two sequences of mappings are considered equal when they have equal length and
-        when they all mappings are equal.
-
-        Args:
-            other (Sequence[Mapping]): Mappings to compare with.
-
-        Returns:
-            bool
-        """
-        try:
-            lenother = len(other)
-        except TypeError:
-            return False
-        if len(self) != lenother:
-            return False
-        for a, b in zip(self, other):
-            if a != b:
-                return False
-        return True
-
-    def __repr__(self):
-        """String representation.
-
-        Returns:
-            str
-        """
-        if len(self) == 1:
-            return f"{type(self).__name__}({repr(self._ann[0])})"
-        return f"{type(self).__name__}{repr(self._ann)}"
-
-    def __add__(self, other):
-        return AnnotationSequence(*self, *other)
-
-    def __radd__(self, other):
-        return AnnotationSequence(*other, *self)
-
-
-def _cast_nparray(arr):
-    """Cast AnnotatedArray to numpy array."""
-    return np.asarray(arr) if isinstance(arr, AnnotatedArray) else arr
-
-
-def _cast_annarray(arr):
-    """Cast array to AnnotatedArray."""
-    return arr if isinstance(arr, np.generic) else AnnotatedArray(arr)
-
-
-def _parse_signature(signature, inputdims):
-    """Parse a ufunc signature based on the actual inputs.
-
-    The signature is matched with the input dimensions, to get the actual signature. It
-    is returned as two lists (inputs and outputs). Each list contains for each
-    individual input and output another list of the signature items.
-
-    Example:
-        >>> treams.util._parse_signature('(n?,k),(k,m?)->(n?,m?)', (2, 1))
-        ([['n?', 'k'], ['k']], [['n?']])
-
-    Args:
-        signature (str): Function signature
-        inputdims (Iterable[int]): Input dimensions
-
-    Returns:
-        tuple[list[list[str]]
-    """
-    signature = "".join(signature.split())  # remove whitespace
-    sigin, sigout = signature.split("->")  # split input and output
-    sigin = sigin[1:-1].split("),(")  # split input
-    sigin = [i.split(",") for i in sigin]
-    sigout = sigout[1:-1].split("),(")  # split output
-    sigout = [i.split(",") for i in sigout]
-    for i, idim in enumerate(inputdims):
-        j = 0
-        while j < len(sigin[i]):
-            d = sigin[i][j]
-            if d.endswith("?") and len(sigin[i]) > idim:
-                sigin = [[i for i in s if i != d] for s in sigin]
-                sigout = [[i for i in s if i != d] for s in sigout]
-            else:
-                j += 1
-    return sigin, sigout
-
-
-def _parse_key(key, ndim):
-    """Parse a key to index an array.
-
-    This function attempts to replicate the numpy indexing of arrays. It can handle
-    integers, slices, an Ellipsis, arrays of integers, arrays of bools, (bare) bools,
-    and None. It returns the key with an Ellipsis appended if the number of index
-    dimensions does not match the number of dimensions given. Additionally, it informs
-    about the number of dimension indexed by fancy indexing, if the fancy indexed
-    dimensions will be prepended, and the number of dimensions that the Ellipsis
-    contains.
-
-    Args:
-        key (tuple): The indexing key
-        ndim (int): Number of array dimensions
-
-    Returns:
-        tuple
-    """
-    consumed = 0
-    ellipsis = False
-    fancy_ndim = 0
-    # nfancy = 0
-    consecutive_intfancy = 0
-    # consecutive_intfancy = 0: no fancy/integer indexing
-    # consecutive_intfancy = 1: ongoing consecutive fancy/integer indexing
-    # consecutive_intfancy = 2: terminated consecutive fancy/integer indexing
-    # consecutive_intfancy >= 2: non-consecutive fancy/integer indexing
-
-    # The first pass gets the consumed dimensions, the presence of an ellipsis, and
-    # fancy index properties.
-    for k in key:
-        if k is not True and k is not False and isinstance(k, (int, np.integer)):
-            consumed += 1
-            consecutive_intfancy += (consecutive_intfancy + 1) % 2
-        elif k is None:
-            consecutive_intfancy += consecutive_intfancy % 2
-        elif isinstance(k, slice):
-            consumed += 1
-            consecutive_intfancy += consecutive_intfancy % 2
-        elif k is Ellipsis:
-            ellipsis = True
-            # consumed is determined at the end
-            consecutive_intfancy += consecutive_intfancy % 2
-        else:
-            arr = np.asanyarray(k)
-            if arr.dtype == bool:
-                consumed += arr.ndim
-                fancy_ndim = max(1, fancy_ndim)
-            else:
-                consumed += 1
-                fancy_ndim = max(arr.ndim, fancy_ndim)
-            # nfancy += arr.ndim
-            consecutive_intfancy += (consecutive_intfancy + 1) % 2
-
-    lenellipsis = ndim - consumed
-    if lenellipsis != 0 and not ellipsis:
-        key = key + (Ellipsis,)
-
-    return key, fancy_ndim, consecutive_intfancy >= 2, lenellipsis
-
-
-HANDLED_FUNCTIONS = {}
-"""Dictionary of numpy functions implemented for AnnotatedArrays."""
-
-
-def implements(np_func):
-    """Decorator to register an __array_function__ implementation to AnnotatedArrays."""
-
-    def decorator(func):
-        HANDLED_FUNCTIONS[np_func] = func
-        return func
-
-    return decorator
-
-
-class AnnotatedArray(np.lib.mixins.NDArrayOperatorsMixin):
-    """Array that keeps track of annotations for each dimension.
-
-    This class acts mostly like numpy arrays, but it is enhanced by the following
-    functionalities:
-
-        * Annotations are added to each dimension
-        * Annotations are compared and preserved for numpy (generalized)
-          :py:class:`numpy.ufunc` (like :py:data:`numpy.add`, :py:data:`numpy.exp`,
-          :py:data:`numpy.matmul` and many more "standard" mathematical functions)
-        * Special ufunc methods, like :py:meth:`numpy.ufunc.reduce`, are supported
-        * A growing subset of other numpy functions are supported, like
-          :py:func:`numpy.linalg.solve`
-        * Keywords can be specified as scale are also index into, when index the
-          AnnotatedArray
-        * Annotations can also be exposed as properties
-
-    .. testsetup::
-
-        from treams.util import AnnotatedArray
-
-    Example:
-        >>> a = AnnotatedArray([[0, 1], [2, 3]], ({"a": 1}, {"b": 2}))
-        >>> b = AnnotatedArray([1, 2], ({"b": 2},))
-        >>> a @ b
-        AnnotatedArray(
-            [2, 8],
-            AnnotationSequence(AnnotationDict({'a': 1})),
-        )
-
-    The interoperability with numpy is implemented using :ref:`basics.dispatch`
-    by defining :meth:`__array__`, :meth:`__array_ufunc__`, and
-    :meth:`__array_function__`.
-    """
-
-    _scales = set()
-
-    def __init__(self, array, ann=(), /, **kwargs):
-        """Initalization."""
-        self._array = np.asarray(array)
-        self.ann = getattr(array, "ann", ())
-        self.ann.update(ann)
-        for key, val in kwargs.items():
-            val = (val,) * self.ndim if not isinstance(val, tuple) else val
-            self.ann.as_dict[key] = val
-
-    @classmethod
-    def relax(cls, *args, mro=None, **kwargs):
-        """Try creating AnnotatedArray subclasses if possible.
-
-        Subclasses can impose stricter conditions on the Annotations. To allow a simple
-        "decaying" of those subclasses it is possible to create them with this
-        classmethod. It attempts array creations along the method resolution order until
-        it succeeds.
-
-        Args:
-            mro (array-like, optional): Method resolution order along which to create
-                the subclass. By default it takes the order of the calling class.
-
-        Note:
-            All other arguments are the same as for the default initialization.
-        """
-        mro = cls.__mro__[1:] if mro is None else mro
-        try:
-            return cls(*args, **kwargs)
-        except AnnotationError as err:
-            if cls == AnnotatedArray:
-                raise err from None
-        cls, *mro = mro
-        return cls.relax(*args, mro=mro, **kwargs)
-
-    def __str__(self):
-        """String of the array itself."""
-        return str(self._array)
-
-    def __repr__(self):
-        """String representation."""
-        repr_arr = "    " + repr(self._array)[6:-1].replace("\n  ", "\n")
-        return f"{self.__class__.__name__}(\n{repr_arr},\n    {self.ann},\n)"
-
-    def __int__(self):
-        return int(self._array)
-
-    def __float__(self):
-        return float(self._array)
-
-    def __array__(self, dtype=None):
-        """Convert to an numpy array.
-
-        This function returns the bare array without annotations. This function does not
-        necessarily make a copy of the array.
-
-        Args:
-            dtype (optional): Type of the returned array.
-        """
-        return np.asarray(self._array, dtype=dtype)
-
-    @property
-    def ann(self):
-        """Array annotations."""
-        return self._ann
-
-    @ann.setter
-    def ann(self, ann):
-        """Set array annotations.
-
-        This function copies the given sequence of dictionaries.
-        """
-        self._ann = AnnotationSequence(*(({},) * self.ndim))
-        self._ann.update(ann)
-
-    def __getattr__(self, key):
-        if key in self.ann.as_dict:
-            res = self.ann.as_dict[key]
-            if all(res[0] == i for i in res[1:]):
-                return res[0]
-            return res
-        raise AttributeError(
-            f"'{self.__class__.__name__}' object has no attribute '{key}'"
-        )
-
-    def __setattr__(self, key, val):
-        if key not in ("_array", "ann", "_ann") and key in self.ann.as_dict:
-            val = (val,) * self.ndim if not isinstance(val, tuple) else val
-            self.ann.as_dict[key] = val
-        else:
-            super().__setattr__(key, val)
-
-    def __delattr__(self, key):
-        try:
-            super().__delattr__(key)
-        except AttributeError:
-            del self.ann.as_dict[key]
-
-    def __len__(self):
-        return len(self._array)
-
-    def __copy__(self):
-        return type(self)(copy.copy(self._array), copy.copy(self._ann))
-
-    def __deepcopy__(self, memo):
-        cls = type(self)
-        return cls(copy.deepcopy(self._array, memo), copy.deepcopy(self._ann, memo))
-
-    def __bool__(self):
-        """Boolean value of the array."""
-        return bool(self._array)
-
-    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
-        """Implement ufunc API."""
-        # Compute result first to use numpy's comprehensive checks on the arguments
-        inputs_noaa = tuple(map(_cast_nparray, inputs))
-
-        ann_out = ()
-        out = kwargs.get("out")
-        out = out if isinstance(out, tuple) else (out,)
-        ann_out = tuple(getattr(i, "ann", None) for i in out)
-        if len(out) != 1 or out[0] is not None:
-            kwargs["out"] = tuple(map(_cast_nparray, out))
-
-        res = getattr(ufunc, method)(*inputs_noaa, **kwargs)
-        istuple, res = (True, res) if isinstance(res, tuple) else (False, (res,))
-        res = tuple(_cast_annarray(r) if o is None else o for r, o in zip(res, out))
-        for a, r in zip(ann_out, res):
-            if a is not None:
-                r.ann.update(a)
-
-        inputs_and_where = inputs + ((kwargs["where"],) if "where" in kwargs else ())
-
-        if ufunc.signature is None or all(
-            map(lambda x: x in " (),->", ufunc.signature)
-        ):
-            if method in ("__call__", "reduceat", "accumulate") or (
-                method == "reduce" and kwargs.get("keepdims", False)
-            ):
-                self._ufunc_call(inputs_and_where, res)
-            elif method == "reduce":
-                self._ufunc_reduce(inputs_and_where, res, kwargs.get("axis", 0))
-            elif method == "at":
-                self._ufunc_at(inputs)
-                return
-            elif method == "outer":
-                self._ufunc_outer(inputs, res, kwargs.get("where", True))
-            else:
-                warnings.warn("unrecognized ufunc method", AnnotationWarning)
-        else:
-            res = self._gufunc_call(ufunc, inputs, kwargs, res)
-        res = tuple(r if isinstance(r, np.generic) else self.relax(r) for r in res)
-        return res if istuple else res[0]
-
-    @staticmethod
-    def _ufunc_call(inputs_and_where, res):
-        for out, in_ in itertools.product(res, inputs_and_where):
-            if not isinstance(out, AnnotatedArray):
-                continue
-            out.ann.update(getattr(in_, "ann", ()))
-
-    @staticmethod
-    def _ufunc_reduce(inputs_and_where, res, axis=0, keepdims=False):
-        out = res[0]  # reduce only allowed for single output functions
-        if not isinstance(out, AnnotatedArray):
-            return
-        axis = axis if isinstance(axis, tuple) else (axis,)
-        in_ = inputs_and_where[0]
-        axis = sorted(map(lambda x: x % np.ndim(in_) - np.ndim(in_), axis))
-        for in_ in inputs_and_where:
-            ann = list(getattr(in_, "ann", []))
-            if not keepdims:
-                for a in axis:
-                    try:
-                        del ann[a]
-                    except IndexError:
-                        pass
-            out.ann.update(ann)
-
-    @staticmethod
-    def _ufunc_at(inputs):
-        out = inputs[0]
-        if not isinstance(out, AnnotatedArray):
-            return
-        if any(d != {} for d in getattr(inputs[1], "ann", ())):
-            warnings.warn("annotations in indices are ignored", AnnotationWarning)
-        for in_ in inputs[2:]:
-            ann = getattr(in_, "ann", ())
-            out.ann.update(ann)
-
-    @staticmethod
-    def _ufunc_outer(inputs, res, where=True):
-        for out in res:
-            if not isinstance(out, AnnotatedArray):
-                continue
-            in_ann = tuple(
-                i for a in inputs for i in getattr(a, "ann", np.ndim(a) * ({},))
-            )
-            out.ann.update(in_ann)
-            where_ann = getattr(where, "ann", ())
-            out.ann.update(where_ann)
-
-    @staticmethod
-    def _gufunc_call(ufunc, inputs, kwargs, res):
-        sigin, sigout = _parse_signature(ufunc.signature, map(np.ndim, inputs))
-        if kwargs.get("keepdims", False):
-            sigout = [sigin[0] for _ in range(ufunc.nout)]
-        ndims = [np.ndim(i) for x in (inputs, res) for i in x]
-        axes = getattr(kwargs, "axes", None)
-        if axes is None:
-            axis = getattr(kwargs, "axis", None)
-            if axis is None:
-                axes = [tuple(range(-len(i), 0)) for i in sigin + sigout]
-            else:
-                axes = [(axis,) for _ in range(ufunc.nin)]
-        else:
-            axes = [
-                tuple(a) if isinstance(a, collections.abc.Iterable) else (a,)
-                for a in axes
-            ]
-        append = axes[0] if kwargs.get("keepdims", False) else ()
-        axes += [append] * (ufunc.nin + ufunc.nout - len(axes))
-        axes = [(*(i % ndim - ndim for i in a),) for a, ndim in zip(axes, ndims)]
-        iterdims = [
-            [i for i in range(-1, -1 - ndim, -1) if i not in a]
-            for a, ndim in zip(axes, ndims)
-        ]
-
-        # compare core dimensions
-        coredims = {}
-        for i, (ax, sig) in enumerate(zip(axes, sigin + sigout)):
-            for a, key in zip(ax, sig):
-                if key.isnumeric():
-                    continue
-                coredims.setdefault(key, []).append((i, a))
-        inout = tuple(inputs) + res
-        for val in coredims.values():
-            for (isrc, dimsrc), (idest, dimdest) in itertools.combinations(val, 2):
-                source = getattr(inout[isrc], "ann", {dimsrc: {}})[dimsrc]
-                dest = getattr(inout[idest], "ann", {dimdest: AnnotationDict()})[
-                    dimdest
-                ]
-                if isrc < ufunc.nin <= idest:
-                    dest.update(source)
-                else:
-                    dest.match(source)
-
-        # compare iteration dimensions
-        for iout, out in enumerate(res):
-            if not isinstance(out, AnnotatedArray):
-                continue
-            for idim, dim in enumerate(iterdims[ufunc.nin + iout]):
-                dest = out.ann[dim]
-                for in_, iterdim in zip(inputs, iterdims):
-                    if idim >= len(iterdim) or getattr(in_, "ann", None) is None:
-                        continue
-                    source = getattr(in_, "ann", {iterdim[idim]: {}})[iterdim[idim]]
-                    dest.update(source)
-        return res
-
-    def __array_function__(self, func, types, args, kwargs):
-        """Function calls on the array.
-
-        Calls defined function in :data:`HANDLED_FUNCTIONS` otherwise raises exception.
-        Add functions to it by using the decorator :func:`implements` for custom
-        implementations.
-        """
-        if func not in HANDLED_FUNCTIONS:
-            return NotImplemented
-        # if not all(issubclass(t, self.__class__) for t in types):
-        #     return NotImplemented
-        return HANDLED_FUNCTIONS[func](*args, **kwargs)
-
-    def __getitem__(self, key):
-        """Get an item from the AnnotatedArray.
-
-        The indexing supports most of numpys regular and fancy indexing.
-        """
-        res = AnnotatedArray(self._array[key])
-        if isinstance(res, np.generic) or res.ndim == 0:
-            return res
-
-        key = key if isinstance(key, tuple) else (key,)
-        key, fancy_ndim, prepend_fancy, lenellipsis = _parse_key(key, self.ndim)
-        source = 0
-        dest = fancy_ndim if prepend_fancy else 0
-        for k in key:
-            if k is not True and k is not False and isinstance(k, (int, np.integer)):
-                source += 1
-            elif k is None:
-                dest += 1
-            elif isinstance(k, slice):
-                for kk, val in self.ann[source].items():
-                    if kk in self._scales:
-                        res.ann[dest][kk] = val[k]
-                    else:
-                        res.ann[dest][kk] = val
-                dest += 1
-                source += 1
-            elif k is Ellipsis:
-                for _ in range(lenellipsis):
-                    res.ann[dest].update(self.ann[source])
-                    dest += 1
-                    source += 1
-            else:
-                k = np.asanyarray(k)
-                ksq = k.squeeze()
-                if ksq.ndim == 1:
-                    pos = (np.array(k.shape) == k.size).argmax()
-                    ann = (
-                        self.ann[source + pos] if k.dtype == bool else self.ann[source]
-                    )
-                    pos += int(not prepend_fancy) * dest + fancy_ndim - k.ndim
-                    for kk, val in ann.items():
-                        if kk in self._scales:
-                            res.ann[pos][kk] = val[ksq]
-                        else:
-                            res.ann[pos][kk] = val
-                source += k.ndim if k.dtype == bool else 1
-                if not prepend_fancy:
-                    dest += fancy_ndim
-                    fancy_ndim = 0
-        return self.relax(res)
-
-    def __setitem__(self, key, value):
-        """Set values.
-
-        If the provided value is an AnnotatedArray the annotations of corresponding
-        dimensions will be matched.
-        """
-        self._array[key] = value
-        if not hasattr(value, "ann"):
-            return
-        key = key if isinstance(key, tuple) else (key,)
-        key, fancy_ndim, prepend_fancy, lenellipsis = _parse_key(key, self.ndim)
-        source = 0
-        dest = fancy_ndim if prepend_fancy else 0
-        for k in key:
-            if k is not True and k is not False and isinstance(k, (int, np.integer)):
-                source += 1
-            elif k is None:
-                dest += 1
-            elif isinstance(k, slice):
-                for kk, val in self.ann[source].items():
-                    if kk not in value.ann[dest]:
-                        continue
-                    if kk in self._scales:
-                        val = val[k]
-                    if val != value.ann[dest][kk]:
-                        warnings.warn(
-                            f"incompatible annotations with key '{kk}' "
-                            f"comparing dimensions '{source}' and '{dest}'",
-                            AnnotationWarning,
-                        )
-                dest += 1
-                source += 1
-            elif k is Ellipsis:
-                for _ in range(lenellipsis):
-                    self.ann[source].match(value.ann[dest])
-                    dest += 1
-                    source += 1
-            else:
-                k = np.asanyarray(k)
-                ksq = k.squeeze()
-                if ksq.ndim == 1:
-                    pos = (np.array(k.shape) == k.size).argmax()
-                    ann = (
-                        self.ann[source + pos] if k.dtype == bool else self.ann[source]
-                    )
-                    pos += int(not prepend_fancy) * dest + fancy_ndim - k.ndim
-                    for kk, val in ann.items():
-                        if kk not in value.ann[pos]:
-                            continue
-                        if kk in self._scales:
-                            val = val[k]
-                        if val != value.ann[pos][kk]:
-                            warnings.warn(
-                                f"incompatible annotations with key '{kk}' "
-                                f"comparing dimensions '{source}' and '{pos}'",
-                                AnnotationWarning,
-                            )
-                source += k.ndim if k.dtype == bool else 1
-                if not prepend_fancy:
-                    dest += fancy_ndim
-                    fancy_ndim = 0
-
-    @property
-    def T(self):
-        """Transpose.
-
-        See also :py:attr:`numpy.ndarray.T`.
-        """
-        return self.transpose()
-
-    @implements(np.all)
-    def all(self, axis=None, dtype=None, out=None, keepdims=False, *, where=True):
-        """Test if all elements (along an axis) are True.
-
-        See also :py:meth:`numpy.ndarray.all`.
-        """
-        return np.logical_and.reduce(self, axis, dtype, out, keepdims, where=where)
-
-    @implements(np.any)
-    def any(self, axis=None, dtype=None, out=None, keepdims=False, *, where=True):
-        """Test if any element (along an axis) is True.
-
-        See also :py:meth:`numpy.ndarray.any`.
-        """
-        return np.logical_or.reduce(self, axis, dtype, out, keepdims, where=where)
-
-    @implements(np.max)
-    def max(self, axis=None, out=None, keepdims=False, initial=None, where=True):
-        """Maximum (along an axis).
-
-        See also :py:meth:`numpy.ndarray.max`.
-        """
-        return np.maximum.reduce(self, axis, None, out, keepdims, initial, where)
-
-    @implements(np.min)
-    def min(self, axis=None, out=None, keepdims=False, initial=None, where=True):
-        """Minimum (along an axis).
-
-        See also :py:meth:`numpy.ndarray.min`.
-        """
-        return np.minimum.reduce(self, axis, None, out, keepdims, initial, where)
-
-    @implements(np.sum)
-    def sum(
-        self, axis=None, dtype=None, out=None, keepdims=False, initial=0, where=True
-    ):
-        """Sum of elements (along an axis).
-
-        See also :py:meth:`numpy.ndarray.sum`.
-        """
-        return np.add.reduce(self, axis, dtype, out, keepdims, initial, where)
-
-    @implements(np.prod)
-    def prod(
-        self, axis=None, dtype=None, out=None, keepdims=False, initial=1, where=True
-    ):
-        """Product of elements (along an axis).
-
-        See also :py:meth:`numpy.ndarray.prod`.
-        """
-        return np.multiply.reduce(self, axis, dtype, out, keepdims, initial, where)
-
-    @implements(np.cumsum)
-    def cumsum(self, axis=None, dtype=None, out=None):
-        """Cumulative sum of elements (along an axis).
-
-        See also :py:meth:`numpy.ndarray.cumsum`.
-        """
-        if axis is None:
-            return np.add.accumulate(self.flatten(), 0, dtype, out)
-        return np.add.accumulate(self, axis, dtype, out)
-
-    @implements(np.cumprod)
-    def cumprod(self, axis=None, dtype=None, out=None):
-        """Cumulative product of elements (along an axis).
-
-        See also :py:meth:`numpy.ndarray.cumprod`.
-        """
-        if axis is None:
-            return np.multiply.accumulate(self.flatten(), 0, dtype, out)
-        return np.multiply.accumulate(self, axis, dtype, out)
-
-    def flatten(self, order="C"):
-        """Flatten array to one dimension.
-
-        See also :py:meth:`numpy.ndarray.flatten`.
-        """
-        res = self.relax(self._array.flatten(order))
-        if res.shape == self.shape:
-            res.ann.update(self.ann)
-        elif len(tuple(filter(lambda x: x != 1, self.shape))) == 1:
-            res.ann[0].update(self.ann[self.shape.index(res.size)])
-        return res
-
-    @implements(np.trace)
-    def trace(self, offset=0, axis1=0, axis2=1, dtype=None, out=None):
-        """Trace of an array.
-
-        See also :py:meth:`numpy.ndarray.trac`.
-        """
-        ann = tuple(a for i, a in enumerate(self.ann) if i not in (axis1, axis2))
-        return self.relax(self._array.trace(offset, axis1, axis2, dtype, out), ann)
-
-    def astype(self, *args, **kwargs):
-        """Return array as given type.
-
-        See also :py:meth:`numpy.ndarray.astype`.
-        """
-        return self.relax(self._array.astype(*args, **kwargs), self.ann)
-
-    @property
-    @implements(np.ndim)
-    def ndim(self):
-        """Number of array dimensions.
-
-        See also :py:attr:`numpy.ndarray.ndim`.
-        """
-        return self._array.ndim
-
-    @property
-    @implements(np.shape)
-    def shape(self):
-        """Array shape.
-
-        See also :py:attr:`numpy.ndarray.shape`.
-        """
-        return self._array.shape
-
-    @property
-    @implements(np.size)
-    def size(self):
-        """Array size.
-
-        See also :py:attr:`numpy.ndarray.size`.
-        """
-        return self._array.size
-
-    @property
-    @implements(np.imag)
-    def imag(self):
-        """Imaginary part of the array.
-
-        See also :py:attr:`numpy.ndarray.imag`.
-        """
-        return self.relax(self._array.imag, self.ann)
-
-    @property
-    @implements(np.real)
-    def real(self):
-        """Real part of the array.
-
-        See also :py:attr:`numpy.ndarray.real`.
-        """
-        return self.relax(self._array.real, self.ann)
-
-    def conjugate(self, *args, **kwargs):
-        """Complex conjugate elementwise.
-
-        See also :py:meth:`numpy.ndarray.conjugate`.
-        """
-        return np.conjugate(self, *args, **kwargs)
-
-    @implements(np.diagonal)
-    def diagonal(self, offset=0, axis1=0, axis2=1):
-        """Get the diagonal of the array.
-
-        See also :py:meth:`numpy.ndarray.diagonal`.
-        """
-        ann = tuple(a for i, a in enumerate(self.ann) if i not in (axis1, axis2)) + (
-            {},
-        )
-        return self.relax(self._array.diagonal(offset, axis1, axis2), ann)
-
-    @implements(np.transpose)
-    def transpose(self, axes=None):
-        """Transpose array dimensions.
-
-        See also :py:meth:`numpy.ndarray.transpose`.
-        """
-        axes = range(self.ndim - 1, -1, -1) if axes is None else axes
-        return self.relax(self._array.transpose(axes), self.ann[axes])
-
-    conj = conjugate
-
-    @implements(np.swapaxes)
-    def swapaxes(self, axis1, axis2):
-        axis1 = axis1 % self.ndim
-        axis2 = axis2 % self.ndim
-        opp = {axis1: axis2, axis2: axis1}
-        axes = [opp.get(i, i) for i in range(self.ndim)]
-        return self.relax(self._array.swapaxes(axis1, axis2), self.ann[axes])
-
-    def __matmul__(self, other):
-        if isinstance(other, _op.Operator):
-            return NotImplemented
-        return super().__matmul__(other)
-
-
-@implements(np.linalg.solve)
-def solve(a, b):
-    """Solve linear system.
-
-    See also :py:func:`numpy.linalg.solve`.
-    """
-    if issubclass(type(a), type(b)) or (
-        not issubclass(type(b), type(a)) and isinstance(a, AnnotatedArray)
-    ):
-        restype = type(a)
-    else:
-        restype = type(b)
-    res = AnnotatedArray(np.linalg.solve(np.asanyarray(a), np.asanyarray(b)))
-    a_ann = list(getattr(a, "ann", [{}, {}]))
-    b_ann = list(getattr(b, "ann", [{}, {}]))
-    if np.ndim(b) == np.ndim(a) - 1:
-        map(lambda x: x[0].match(x[1]), zip(a_ann[-2::-1], b_ann[-1::-1]))
-        del a_ann[-2]
-        del b_ann[-1]
-    else:
-        map(lambda x: x[0].match(x[1]), zip(a_ann[-2::-1], b_ann[-2::-1]))
-        del a_ann[-2]
-        del b_ann[-2]
-        a_ann += [{}]
-    res.ann.update(a_ann)
-    res.ann.update(b_ann)
-    return restype.relax(res)
-
-
-@implements(np.linalg.lstsq)
-def lstsq(a, b, rcond="warn"):
-    """Solve linear system using least squares.
-
-    See also :py:func:`numpy.linalg.lstsq`.
-    """
-    if issubclass(type(a), type(b)) or (
-        not issubclass(type(b), type(a)) and isinstance(a, AnnotatedArray)
-    ):
-        restype = type(a)
-    else:
-        restype = type(b)
-    res = list(np.linalg.lstsq(np.asanyarray(a), np.asanyarray(b), rcond))
-    res[0] = AnnotatedArray(res[0])
-    a_ann = getattr(a, "ann", ({},))
-    b_ann = getattr(b, "ann", ({},))
-    a_ann[0].match(b_ann[0])
-    res[0].ann[0].update(a_ann[-1])
-    if np.ndim(b) == 2:
-        res[0].ann[1].update(b_ann[-1])
-    res[0] = restype.relax(res[0])
-    return tuple(res)
-
-
-@implements(np.linalg.svd)
-def svd(a, full_matrices=True, compute_uv=True, hermitian=False):
-    """Compute the singular value decomposition.
-
-    See also :py:func:`numpy.linalg.svd`.
-    """
-    res = list(np.linalg.svd(np.asanyarray(a), full_matrices, compute_uv, hermitian))
-    ann = getattr(a, "ann", ({},))
-    if compute_uv:
-        res[0] = a.relax(res[0], ann[:-1] + ({},))
-        res[1] = a.relax(res[1], ann[:-2] + ({},))
-        res[2] = a.relax(res[2], ann[:-2] + ({}, ann[-1]))
-        return res
-    return a.relax(res, tuple(ann[:-2]) + ({},))
-
-
-@implements(np.diag)
-def diag(a, k=0):
-    """Extract diagonal from an array or create an diagonal array.
-
-    See also :py:func:`numpy.diag`.
-    """
-    res = np.diag(np.asanyarray(a), k)
-    ann = a.ann
-    if a.ndim == 1:
-        ann = (ann[0], copy.copy(ann[0]))
-    elif k == 0:
-        ann[0].match(ann[1])
-        ann = ({**ann[0], **ann[1]},)
-    else:
-        ann = ()
-    return a.relax(res, ann)
-
-
-@implements(np.tril)
-def tril(a, k=0):
-    """Get lower trinagular matrix.
-
-    See also :py:func:`numpy.tril`.
-    """
-    res = np.tril(np.asanyarray(a), k)
-    return a.relax(res, getattr(a, "ann", ({},)))
-
-
-@implements(np.triu)
-def triu(a, k=0):
-    """Get upper trinagular matrix.
-
-    See also :py:func:`numpy.triu`.
-    """
-    res = np.triu(np.asanyarray(a), k)
-    return a.relax(res, getattr(a, "ann", ({},)))
-
-
-@implements(np.zeros_like)
-def zeros_like(a, dtype=None, order="K", shape=None):
-    """Create a numpy array like the given array containing zeros."""
-    return np.zeros_like(np.asarray(a), dtype=dtype, order=order, shape=shape)
-
-
-@implements(np.ones_like)
-def ones_like(a, dtype=None, order="K", shape=None):
-    """Create a numpy array like the given array containing ones."""
-    return np.ones_like(np.asarray(a), dtype=dtype, order=order, shape=shape)

Alert delta unavailable

Currently unable to show alert delta for PyPI packages.

Improved metrics

Number of package files

106

13.98%

Lines of code

12927

1.59%

Worsened metrics

Total package byte prevSize

13414205

-0.68%