Introduction

This method assumes the goal is to calculate $L_z$, the $z$ component of the orbital angular momentum (henceforth just called angular momentum) of a system and that you have the wavefunction, $\psi$, of the system in a cartesian position basis on some 2D grid with equal spacing between lattice points and equal numbers of lattice points in each direction (though it should be fairly straightforward to generalize this to rectangular lattices).

The operator for the angular momentum in the position basis is $-i\hbar(x \partial_y - y \partial_x)$. However, by Fourier transforming $\psi$ along the $x$ axis one obtains a wavefunction in the basis of $k_x$ and $y$ where each of these operators is simply a multiplication by a number, and mutatis mutandis for $k_y$ and $x$. I.e. $$\int \psi^*(x, y)(x \partial_y - y\partial_x)\psi(x, y) dx dy \propto\int \psi^*(x, k_y) xk_y\psi(x, k_y) dk_y dx - \int\psi^*(k_x, y) yk_x \psi(k_x, y) dk_x dy$$

In Code

In Python (feeling a bit too lazy to do this in Rust), using numpy, this can be numerically calulated with the following:

import numpy as np
from numba import vectorize, float64, complex128 # not strictly necessary, could
                                                 # also use numpy's inbuilt vectorizer
from numpy.fft import fft, fftshift


# should be slightly faster than x.conj()*x and produces a purely real number
@vectorize([float64(complex64)])
def norm_sqr(x):
    return x.real*x.real + x.imag*x.imag


def fftangmom(psi, xv, yv, kxv, kyv):
    psiykx = fftshift(fft(fftshift(psi, axes=0), axis=0), axes=0)
    psixky = fftshift(fft(fftshift(psi, axes=1), axis=1), axes=1)
    return np.sum(norm_sqr(psixky * xv * kyv - norm_sqr(psiykx) * yv * kxv)/ np.shape(psi)[0]

This could of course be GPU accelerated with PyTorch, though in that case using z.conj()*z would probably perform better than the custom function.

To test the algorithm I created some wavefunctions which are analytically eigenstates of $L_z$, and compared the numerical results to the expected results. The following tests a variety of dipole combinations and a tripole.

import numpy as np
from numba import vectorize, float64, complex128
from numpy.fft import fft, fftshift
from tabulate import tabulate


# For table making purposes
class MetaData:
    __slots__ = ('pos', 'l', 'sigma', 'sigmainvsq') # Doesn't matter too much here but
                                                    # defining slots can significantly
                                                    # speed up class accesses

    def __init__(self, pos, l, sigma):
        self.pos = pos
        self.l = l
        self.sigma = sigma
        self.sigmainvsq = 1 / (sigma*sigma)


# should be slightly faster than x.conj()*x and produces a purely real number
@vectorize([float64(complex128)])
def norm_sqr(x):
    return x.real*x.real + x.imag*x.imag


def fftangmom(psi, xv, yv, kxv, kyv):
    psiykx = fftshift(fft(fftshift(psi, axes=0), axis=0), axes=0)
    psixky = fftshift(fft(fftshift(psi, axes=1), axis=1), axes=1)
    return np.sum(norm_sqr(psixky) * xv * kyv - norm_sqr(psiykx) * yv * kxv)/ np.shape(psi)[0]


# A simple ang. mom. eigenfunction has the form \psi(r, \theta) = f(r)*e^{im\theta}
def angulargaussian(x, y, data):
    return np.exp((-(x * x - data.pos[0]) - (y * y - data.pos[1]))*data.sigmainvsq)\
            * np.exp(1j * data.l * np.arctan2(x, y))


# A mistaken function which will not produce an ang. mom. eigenfunction
def faultyangulargaussian(x, y, data):
    return np.exp((-(x * x - data.pos[0]) - (y * y - data.pos[1]))*data.sigmainvsq)\
            * np.exp(1j * data.l * np.arctan2(x-data.pos[0], y-data.pos[1]))


n = 1024
norm_sqr = np.vectorize(norm_sqr)
lengths = [10, 20, 30, 50, 75, 100]
for length in lengths:
    dx = 2 * length / n
    x = np.arange(-length, length, dx)
    xv, yv = np.meshgrid(x, x)

    kmax = np.pi / dx
    dk = 2 * kmax / n
    kx = np.arange(-kmax, kmax, dk)
    kxv, kyv = np.meshgrid(kx, kx)


    psi0meta = MetaData((0, 0), 0, 0.1)
    psi0 = angulargaussian(xv, yv, psi0meta)
    psisq = np.sum(norm_sqr(psi0))
    psi0 /= np.sqrt(psisq)

    psi1meta = MetaData((6, 6), 1, 0.2)
    psi1 = angulargaussian(xv, yv, psi1meta)
    psisq = np.sum(norm_sqr(psi1))
    psi1 /= np.sqrt(psisq)

    faultymeta = MetaData((6, 6), 1, 0.3)
    faultypsi1 = faultyangulargaussian(xv, yv, faultymeta)
    psisq = np.sum(norm_sqr(psi1))
    faultypsi1 /= np.sqrt(psisq)

    psi2meta = MetaData((-5, -5), -1, 0.4)
    psi2 = angulargaussian(xv, yv, psi2meta)
    psisq = np.sum(norm_sqr(psi2))
    psi2 /= np.sqrt(psisq)

    psi3meta = MetaData((1, 1), 2, 4)
    psi3 = angulargaussian(xv, yv, psi3meta)
    psisq = np.sum(norm_sqr(psi3))
    psi3 /= np.sqrt(psisq)

    psi4meta = MetaData((40, 40), 1, 1)
    psi4 = angulargaussian(xv, yv, psi4meta)
    psisq = np.sum(norm_sqr(psi4))
    psi4 /= np.sqrt(psisq)

    dipole12 = psi1 + psi2
    dipole23 = psi2 + psi3
    dipole31 = psi3 + psi1
    dipole42 = psi4 + psi2

    tripole = psi1 + psi2 + psi3

    l0 = fftangmom(psi0, xv, yv, kxv, kyv)
    l1 = fftangmom(psi1, xv, yv, kxv, kyv)
    faultyl1 = fftangmom(faultypsi1, xv, yv, kxv, kyv)
    l2 = fftangmom(psi2, xv, yv, kxv, kyv)
    l3 = fftangmom(psi3, xv, yv, kxv, kyv)
    l4 = fftangmom(psi4, xv, yv, kxv, kyv)
    d12 = fftangmom(dipole12, xv, yv, kxv, kyv)
    d23 = fftangmom(dipole23, xv, yv, kxv, kyv)
    d31 = fftangmom(dipole31, xv, yv, kxv, kyv)
    d42 = fftangmom(dipole42, xv, yv, kxv, kyv)
    t = fftangmom(tripole, xv, yv, kxv, kyv)
    table = [
            ["Wavefunction: (center), stdev, unnormalized", "numerical L", "analytical L"],
            [f'psi0: {psi0meta.pos}, {psi0meta.sigma}', l0.real, psi0meta.l],
            [f'psi1: {psi1meta.pos}, {psi1meta.sigma}', l1.real, psi1meta.l],
            [f"incorrectpsi1: {faultymeta.pos}, {psi1meta.sigma}", faultyl1.real, '???'],
            [f"psi2: {psi2meta.pos}, {psi2meta.sigma}", l2.real, psi2meta.l],
            [f'psi3: {psi3meta.pos}, {psi3meta.sigma}', l3.real, psi3meta.l],
            [f'psi4: {psi4meta.pos}, {psi4meta.sigma}', l4.real, psi4meta.l],
            ['psi1 + psi2', d12.real, psi1meta.l + psi2meta.l],
            ["psi2 + psi3", d23.real, psi2meta.l + psi3meta.l],
            ["psi3 + psi1", d31.real, psi3meta.l + psi1meta.l],
            ["psi4 + psi2", d42.real, psi4meta.l + psi2meta.l],
            ["psi1 + psi2 + psi3", t.real, psi3meta.l + psi1meta.l + psi2meta.l]
            ]
    
    print(f"Numerical grid is {n} x {n}, with axes x = [-{length}, {length}[; y = [-{length}, {length}[")
    print(tabulate(table, headers='firstrow', tablefmt='simple'))

One point of concern was that an "eigenfunction" might have most of its bulk located way outside the defined grid and that this might mess with the calculation. Thankfully, as long as you avoid NaNs this doesn't affect the numerics too much, as evidenced by table 1. One thing that does affect the numerics quite heavily is the available resolution. One can see that the accuracy for $\psi_1$, with a FWHM of only ~0.47, goes down drastically as the spatial resolution decreases and is already practically unusable at a resolution of ~0.06 ppx. From the tables it seems that this algorithm performs decently at resolutions up to about 7% of a gaussian's FWHM. Reassuringly, though not surprisingly, the algorithm has the expected linearity, i.e. $L(\psi_1 + \psi_2) = L(\psi_1) + L(\psi_2)$ even when its accuracy has gone down the drain. Each table here is produced with a 1024x1024 grid.

Axes x = [-10, 10[; y = [-10, 10[

Axes x = [-20, 20[; y = [-20, 20[

Axes x = [-30, 30[; y = [-30, 30[

Axes x = [-50, 50[; y = [-50, 50[

Axes x = [-75, 75[; y = [-75, 75[

Axes x = [-100, 100[; y = [-100, 100[